ai-writing.service.ts

ai-writing.service.tsapps/api/src/modules/ai/ai-writing.service.tsraw ← ai21,955 B← Back to the app
1/**
2 * TASK-0123 — AI writing tools (summarize / rewrite / improve / translate).
3 *
4 * Endpoints (documented in `docs/api/ai.md`):
5 *
6 *   POST /workspaces/:workspaceId/ai/writing         (SSE)
7 *       body: { blockIds: string[] (1..50), pageId, action, language? }
8 *       → `text/event-stream` frames (same contract as Ask AI chat):
9 *           delta / tool_call / usage / error / done
10 *         `tool_call` = the model's `generateBlocks` call (validated
11 *         server-side: unknown blockIds are DROPPED + noted, never applied);
12 *         `done.assistantMessageId` = the SUGGESTION id (the assistant row
13 *         whose `toolResults` hold the full suggestion).
14 *
15 *   POST /workspaces/:workspaceId/ai/writing/apply
16 *       body: { suggestionId, blockIds?: string[] }
17 *       → 200 { results: [{ blockId, status: "applied" | "failed", code?,
18 *         message?, actualVersion? }] } — per-block failure isolation (one
19 *         403/409 never aborts the rest). Each applied block goes through the
20 *         normal `BlockService.update` (EDIT re-check + `block.updated` audit
21 *         + `page.updated` realtime — the spec's "not a special AI write")
22 *         with the suggestion-time version as `baseVersion` (a stale block →
23 *         a per-result 409 entry carrying `actualVersion` for the rebase UI).
24 *
25 * Security (spec §Security, AGENTS.md #19):
26 *   - the client sends block IDs only; content is assembled server-side
27 *     (`assembleWritingContext`) with per-page VIEW re-checks (IDOR guard);
28 *   - the prompt frames block content as UNTRUSTED data;
29 *   - `generateBlocks` arguments are validated against the REQUESTED set —
30 *     unknown IDs are dropped (the mock `forceUnknownId` seam exercises it);
31 *   - apply re-validates every `{blockId, newText}` pair against the STORED
32 *     suggestion (no client-smuggled edits) and re-checks EDIT per block at
33 *     APPLY time (permissions may change between suggest and apply);
34 *   - hidden per-call conversations (`title: "writing:<action>"`) never
35 *     appear in the panel history (filtered in `listConversations`).
36 *
37 * Limits (spec §Limits): ≤50 blocks (422), ≤100k chars of selected text
38 * (413 — thrown by the assembler), language ≤64 chars, action enum; the
39 * per-user cap is the SHARED `ai.calls` limiter (30 chat+writing calls/min).
40 */
41import { Inject, Injectable } from "@nestjs/common";
42import type { Request, Response } from "express";
43 
44import type { PrismaClient } from "@prisma/client";
45import { PrismaService } from "../../common/prisma.service.js";
46import { AuditService } from "../../common/audit.service.js";
47import { ApiError, type ErrorCode, type ErrorDetail } from "../../common/errors.js";
48import { logger } from "../../common/logger.js";
49import { assertCan } from "../../lib/authorize.js";
50import { AI_GATEWAY, mapAiError, primeWorkspaceCache, type AiBudgetGuard } from "./gateway-provider.js";
51import { AiBudgetGuardService } from "./budget-guard.service.js";
52import { AiRateLimitService } from "./ai-rate-limit.js";
53import { type AiWritingInput, type ApplyWritingInput } from "./ai-writing.dto.js";
54import {
55  assembleWritingContext,
56  GENERATE_BLOCKS_TOOL,
57  parseGeneratedBlocks,
58  renderWritingSections,
59  WRITING_CONV_PREFIX,
60  WRITING_MAX_BLOCKS,
61  WRITING_SYSTEM_PROMPT,
62  writingTaskLine,
63  WritingContextError,
64  type WritingAction,
65} from "./prompts.js";
66import { BlockService } from "../blocks/block.service.js";
67import type { AiGateway } from "@factory/ai";
68 
69const AI_WRITING_TIMEOUT_MS = 60_000;
70const KEEPALIVE_MS = 10_000;
71 
72/** One SSE frame — the documented wire contract (same as Ask AI). */
73export type AiWritingSseEvent =
74  | { type: "delta"; text: string }
75  | { type: "tool_call"; id: string; name: string; arguments: Record<string, unknown> }
76  | { type: "usage"; tokensIn: number; tokensOut: number }
77  | { type: "error"; code: string; message: string; retryable?: boolean; retryAfterSeconds?: number }
78  | { type: "done"; assistantMessageId?: string; userMessageId?: string; reason?: "ok" | "aborted" | "error" };
79 
80function sseFrame(event: AiWritingSseEvent): string {
81  return `data: ${JSON.stringify(event)}\n\n`;
82}
83 
84/** The persisted suggestion (`assistant.toolResults`): the server-side
85 *  record that apply re-validates against. */
86export interface WritingSuggestion {
87  action: WritingAction;
88  language?: string;
89  pageId: string;
90  blockIds: string[];
91  blocks: Array<{
92    blockId: string;
93    newText: string;
94    /** The block's version at SUGGEST time (apply's `baseVersion`). */
95    version: number;
96    /** The block's full content snapshot (apply replaces `text` on it). */
97    content: Record<string, unknown>;
98    /** The original plain text (the UI's diff "old" side). */
99    text: string;
100  }>;
101  /** blockIds the model emitted that were NOT in the request (dropped). */
102  droppedIds: string[];
103}
104 
105/** One apply result row (the documented per-block isolation contract). */
106export interface ApplyResultRow {
107  blockId: string;
108  status: "applied" | "failed";
109  code?: string;
110  message?: string;
111  actualVersion?: number;
112}
113 
114@Injectable()
115export class AiWritingService {
116  constructor(
117    @Inject(PrismaService) private readonly prisma: PrismaClient,
118    @Inject(AI_GATEWAY) private readonly gateway: AiGateway,
119    @Inject(AiBudgetGuardService) private readonly guardService: AiBudgetGuardService,
120    @Inject(AiRateLimitService) private readonly rateLimit: AiRateLimitService,
121    @Inject(BlockService) private readonly blocks: BlockService,
122    @Inject(AuditService) private readonly audit: AuditService,
123  ) {}
124 
125  private get guard(): AiBudgetGuard {
126    return this.guardService.guard;
127  }
128 
129  // ---------------------------------------------------------------- write ---
130 
131  async generate(workspaceId: string, userId: string, input: AiWritingInput, req: Request, res: Response): Promise<void> {
132    // ---- pre-stream gates (canonical JSON envelope on failure) ----
133    // VIEW is the SUGGESTION's gate (reading the blocks' content); the
134    // spec's acceptance criteria run SUGGEST as a view-only user and expect
135    // the 403 only at APPLY (per-block EDIT). Apply keeps its EDIT gate.
136    await assertCan(this.prisma, { workspaceId, userId }, "VIEW");
137    if (input.blockIds.length > WRITING_MAX_BLOCKS) {
138      throw new ApiError("VALIDATION_ERROR", 422, `Too many blocks: maximum is ${WRITING_MAX_BLOCKS}.`, [
139        { reason: "too_many_blocks", count: input.blockIds.length, max: WRITING_MAX_BLOCKS },
140      ]);
141    }
142 
143    // shared ai.calls cap (chat + writing combined; degrades OPEN w/o Redis)
144    const rl = await this.rateLimit.check(userId);
145    if (!rl.allowed) {
146      throw new ApiError("RATE_LIMITED", 429, "AI call limit reached.", [
147        { reason: "ai_call_rate_limited", retryAfterSeconds: rl.retryAfterSeconds },
148      ]);
149    }
150 
151    // budget BEFORE any SSE header (402 JSON envelope — spec: "budget 402")
152    await primeWorkspaceCache(this.prisma, workspaceId);
153    try {
154      await this.guard.check(workspaceId, undefined, req);
155    } catch (err) {
156      throw mapAiError(err);
157    }
158 
159    // ---- context (permission re-check + ±5 + caps; 413 on >100k chars) ----
160    const uniqueIds = [...new Set(input.blockIds)];
161    let ctx;
162    try {
163      ctx = await assembleWritingContext(this.prisma, workspaceId, userId, uniqueIds);
164    } catch (err) {
165      // The assembler's 413 (PAYLOAD_TOO_LARGE) is a plain Error — map it to
166      // the canonical envelope BEFORE any SSE byte is written.
167      if (err instanceof WritingContextError) {
168        throw new ApiError(err.code as ErrorCode, err.status, err.message, (err.details ?? undefined) as ErrorDetail[] | undefined);
169      }
170      throw err;
171    }
172    const noteParts: string[] = [];
173    if (ctx.missing.length > 0) noteParts.push(`${ctx.missing.length} selected block(s) are no longer available.`);
174    if (ctx.droppedPages.length > 0) noteParts.push(`${ctx.droppedPages.length} selected page(s) are unavailable (no view access).`);
175    if (ctx.nonText.length > 0) noteParts.push(`${ctx.nonText.length} selected block(s) have no editable text.`);
176    // `sections` can be non-empty with ZERO selected text blocks (all the
177    // found blocks were non-text — the page is still listed, just with no
178    // selection). The tool call would come back empty → 422 up front.
179    if (ctx.sections.every((s) => s.selected.length === 0)) {
180      throw new ApiError("VALIDATION_ERROR", 422, "No editable blocks available for this request.", [
181        { reason: "no_editable_blocks" },
182      ]);
183    }
184 
185    // ---- hidden per-call conversation (never listed in the panel) ----
186    const conversation = await this.prisma.aiConversation.create({
187      data: {
188        workspaceId,
189        userId,
190        title: `${WRITING_CONV_PREFIX}${input.action}`,
191        context: { pageId: input.pageId, blockIds: uniqueIds, action: input.action, language: input.language ?? null } as object,
192      },
193    });
194 
195    // ---- SSE header (the first byte — after this, errors are in-band) ----
196    res.status(200);
197    res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
198    res.setHeader("Cache-Control", "no-cache, no-transform");
199    res.setHeader("Connection", "keep-alive");
200    res.setHeader("X-Accel-Buffering", "no");
201    res.flushHeaders?.();
202    if (typeof res.socket?.setNoDelay === "function") res.socket.setNoDelay(true);
203 
204    // abort plumbing (same as Ask AI: `res` close + writableEnded check)
205    const ac = new AbortController();
206    let clientClosed = false;
207    const onResClose = (): void => {
208      if (res.writableEnded) return;
209      clientClosed = true;
210      ac.abort();
211    };
212    res.on("close", onResClose);
213    const keepalive = setInterval(() => {
214      if (!res.writableEnded) res.write(": keepalive\n\n");
215    }, KEEPALIVE_MS);
216    let timer: ReturnType<typeof setTimeout> | undefined;
217    const armTimeout = (): void => {
218      if (timer) clearTimeout(timer);
219      timer = setTimeout(() => ac.abort(), AI_WRITING_TIMEOUT_MS);
220    };
221    armTimeout();
222 
223    // ---- the request (persisted as the user row — a deterministic
224    // description; the user never types free-form text here) ----
225    const userRow = await this.prisma.aiMessage.create({
226      data: {
227        conversationId: conversation.id,
228        role: "user",
229        content: `writing:${input.action}${input.language ? `:${input.language}` : ""} blocks=${uniqueIds.length} page=${input.pageId}`,
230      },
231    });
232 
233    const messages: Array<{ role: "system" | "user"; content: string }> = [
234      { role: "system", content: WRITING_SYSTEM_PROMPT },
235      {
236        role: "user",
237        content: `${writingTaskLine(input.action, input.language)}\n\n${renderWritingSections(ctx.sections)}`,
238      },
239    ];
240 
241    const chatModel = this.gateway.resolve(workspaceId).chatModel;
242    const startedAt = Date.now();
243    let usage = { in: 0, out: 0 };
244    let toolArgs: unknown;
245    let toolSeen = false;
246    let finished = false;
247 
248    const finish = (event: AiWritingSseEvent): void => {
249      if (finished) return;
250      finished = true;
251      if (timer) clearTimeout(timer);
252      clearInterval(keepalive);
253      res.off("close", onResClose);
254      if (!res.writableEnded) {
255        res.write(sseFrame(event));
256        res.end();
257      }
258    };
259    const sendErrorDelta = (err: unknown): void => {
260      const mapped = mapAiError(err);
261      const detail = (mapped.details?.[0] ?? {}) as Record<string, unknown>;
262      const retryable = detail.retryable === true || detail.retryAfterSeconds != null || mapped.code === "RATE_LIMITED";
263      const event: AiWritingSseEvent = {
264        type: "error",
265        code: mapped.code,
266        message: mapped.message,
267        ...(retryable ? { retryable: true } : {}),
268        ...(typeof detail.retryAfterSeconds === "number" ? { retryAfterSeconds: detail.retryAfterSeconds } : {}),
269      };
270      res.write(sseFrame(event));
271      logger.info(`ai writing stream error workspace=${workspaceId} code=${mapped.code} retryable=${retryable}`, {
272        event: "ai.writing.error",
273        workspaceId,
274        code: mapped.code,
275        retryable,
276      });
277    };
278 
279    /** Persist the assistant row = THE SUGGESTION. Only when the model
280     *  actually returned a tool call (an aborted/error stream before output
281     *  has no suggestion to persist — the error frame + user row cover it).
282     *  `done.assistantMessageId` = this row's id (the suggestion id). */
283    const persistSuggestion = async (errorTag?: string): Promise<string | undefined> => {
284      if (!toolSeen) return undefined;
285      try {
286        const requestedIds = ctx.sections.flatMap((s) => s.selected.map((b) => b.id));
287        const parsed = parseGeneratedBlocks(toolArgs, requestedIds);
288        const selById = new Map(ctx.sections.flatMap((s) => s.selected.map((b) => [b.id, b] as const)));
289        const suggestion: WritingSuggestion = {
290          action: input.action,
291          ...(input.language ? { language: input.language } : {}),
292          pageId: input.pageId,
293          blockIds: uniqueIds,
294          blocks:
295            parsed.ok
296              ? parsed.blocks
297                  .filter((b) => selById.has(b.blockId))
298                  .map((b) => {
299                    const sel = selById.get(b.blockId)!;
300                    return {
301                      blockId: b.blockId,
302                      newText: b.newText,
303                      version: sel.version,
304                      content: (typeof sel.content === "object" && sel.content !== null ? sel.content : { text: sel.text }) as Record<string, unknown>,
305                      text: sel.text,
306                    };
307                  })
308              : [],
309          droppedIds: parsed.ok ? parsed.dropped : [],
310        };
311        const row = await this.prisma.aiMessage.create({
312          data: {
313            conversationId: conversation.id,
314            role: "assistant",
315            content: suggestion.blocks.length > 0 ? `suggestion:${suggestion.blocks.length} block(s)` : "suggestion:0 blocks",
316            model: chatModel,
317            ...(usage.in > 0 ? { tokensIn: usage.in } : {}),
318            ...(usage.out > 0 ? { tokensOut: usage.out } : {}),
319            latencyMs: Date.now() - startedAt,
320            toolResults: suggestion as unknown as object,
321            metadata: {
322              feature: "writing",
323              action: input.action,
324              ...(input.language ? { language: input.language } : {}),
325              ...(parsed.ok && parsed.dropped.length > 0 ? { droppedIds: parsed.dropped } : {}),
326              ...(errorTag ? { error: errorTag } : {}),
327            } as object,
328          },
329        });
330        return row.id;
331      } catch (err) {
332        // persistence must never crash the stream teardown
333        logger.error(`ai writing suggestion persist failed: ${(err as Error).message}`);
334        return undefined;
335      }
336    };
337 
338    try {
339      for await (const chunk of this.gateway.chat({
340        workspaceId,
341        userId,
342        capability: "writing",
343        messages,
344        tools: [GENERATE_BLOCKS_TOOL],
345        toolChoice: { name: "generateBlocks" },
346        metadata: {
347          mode: "writing",
348          action: input.action,
349          ...(input.language ? { language: input.language } : {}),
350          conversationId: conversation.id,
351        },
352        signal: ac.signal,
353      })) {
354        if (timer) armTimeout();
355        if (chunk.delta) res.write(sseFrame({ type: "delta", text: chunk.delta }));
356        if (chunk.toolCallDelta) {
357          toolSeen = true;
358          toolArgs = chunk.toolCallDelta.arguments;
359          res.write(sseFrame({ type: "tool_call", ...chunk.toolCallDelta }));
360        }
361        if (chunk.done) {
362          if (chunk.usage) {
363            usage = { in: chunk.usage.in, out: chunk.usage.out };
364            res.write(sseFrame({ type: "usage", tokensIn: usage.in, tokensOut: usage.out }));
365          }
366          const note = noteParts.join(" ");
367          if (note && !clientClosed) res.write(sseFrame({ type: "delta", text: `> ${note}\n\n` }));
368          const suggestionId = await persistSuggestion();
369          finish({ type: "done", assistantMessageId: suggestionId, userMessageId: userRow.id, reason: "ok" });
370        }
371      }
372      // provider stream ended without a done chunk (defensive).
373      if (!finished) {
374        const suggestionId = await persistSuggestion();
375        finish({ type: "done", assistantMessageId: suggestionId, userMessageId: userRow.id, reason: "ok" });
376      }
377    } catch (err) {
378      if (timer) clearTimeout(timer);
379      clearInterval(keepalive);
380      res.off("close", onResClose);
381 
382      if (clientClosed || ac.signal.aborted) {
383        // client hang-up / Stop: cancelled — a partial suggestion (the tool
384        // call may already have arrived) is still persisted.
385        const suggestionId = await persistSuggestion("aborted");
386        if (!clientClosed && !res.writableEnded) {
387          res.write(sseFrame({ type: "done", assistantMessageId: suggestionId, userMessageId: userRow.id, reason: "aborted" }));
388          res.end();
389        }
390        return;
391      }
392      // provider failure (normalized — raw provider text never surfaces):
393      // error delta + terminal done; persist the suggestion if the tool
394      // call already arrived mid-stream.
395      sendErrorDelta(err);
396      const suggestionId = await persistSuggestion("provider_error");
397      finish({ type: "done", assistantMessageId: suggestionId, userMessageId: userRow.id, reason: "error" });
398    }
399  }
400 
401  // ---------------------------------------------------------------- apply ---
402 
403  async apply(workspaceId: string, userId: string, input: ApplyWritingInput, req: Request): Promise<{ results: ApplyResultRow[] }> {
404    // NO workspace-level EDIT gate here: the spec's contract is PER-BLOCK
405    // failure isolation (a view-only user's apply 200s with per-block
406    // `FORBIDDEN requires_EDIT` entries — "other blocks apply"). The real
407    // check is the per-page EDIT inside BlockService.update below. Ownership
408    // (the suggestion is the caller's, in this workspace) is enforced next.
409 
410    // 1) the suggestion (ownership-scoped 404 — existence hidden).
411    const suggestionRow = await this.prisma.aiMessage.findUnique({
412      where: { id: input.suggestionId },
413      include: { conversation: { select: { id: true, workspaceId: true, userId: true } } },
414    });
415    if (
416      !suggestionRow ||
417      !suggestionRow.conversation ||
418      suggestionRow.conversation.workspaceId !== workspaceId ||
419      suggestionRow.conversation.userId !== userId ||
420      suggestionRow.role !== "assistant"
421    ) {
422      throw ApiError.notFound("Suggestion not found.");
423    }
424    const suggestion = suggestionRow.toolResults as unknown as WritingSuggestion | null;
425    if (!suggestion || !Array.isArray(suggestion.blocks) || suggestion.blocks.length === 0) {
426      throw ApiError.notFound("Suggestion not found.");
427    }
428 
429    // 2) which blocks: the whole suggestion by default; a SUBSET is allowed
430    //    (the UI's "accept one") — but every requested id must be IN the
431    //    suggestion (no smuggling new blockIds / newText pairs).
432    const byId = new Map(suggestion.blocks.map((b) => [b.blockId, b] as const));
433    const targetIds = input.blockIds ? input.blockIds : suggestion.blocks.map((b) => b.blockId);
434    const smuggled = targetIds.filter((id) => !byId.has(id));
435    if (smuggled.length > 0) {
436      throw new ApiError("VALIDATION_ERROR", 422, "blockIds must be a subset of the suggestion.", [
437        { reason: "unknown_block_id", blockIds: smuggled },
438      ]);
439    }
440 
441    // 3) per-block apply with failure isolation: EDIT is re-checked per
442    //    block's page (permissions may have changed since the suggestion),
443    //    then the normal BlockService.update with the suggestion-time
444    //    version as baseVersion (stale → 409 entry with actualVersion).
445    const results: ApplyResultRow[] = [];
446    for (const blockId of targetIds) {
447      const entry = byId.get(blockId)!;
448      try {
449        // Preserve the block's content SHAPE (object canonical vs the web
450        // editor's HTML-string model): spreading a string primitive would
451        // corrupt it into an index-keyed object, and writing `{ text }` over
452        // a string block would leave the web editor showing the old markup.
453        const newContent: Record<string, unknown> | string =
454          typeof entry.content === "string"
455            ? `<p>${entry.newText}</p>`
456            : { ...(entry.content ?? {}), text: entry.newText };
457        await this.blocks.update(
458          blockId,
459          userId,
460          { content: newContent as unknown as Record<string, unknown>, baseVersion: entry.version },
461          req,
462        );
463        results.push({ blockId, status: "applied" });
464      } catch (err) {
465        const e = err as { code?: string; message?: string; details?: Array<Record<string, unknown>> };
466        const detail = e.details?.[0] ?? {};
467        const row: ApplyResultRow = {
468          blockId,
469          status: "failed",
470          code: e.code ?? "ERROR",
471          message: e.message ?? "Failed to apply.",
472        };
473        if (e.code === "CONFLICT" && typeof detail.actualVersion === "number") row.actualVersion = detail.actualVersion;
474        results.push(row);
475        logger.info(`ai writing apply failed block=${blockId} code=${row.code}`, {
476          event: "ai.writing.apply_failed",
477          workspaceId,
478          blockId,
479          code: row.code,
480        });
481      }
482    }
483 
484    const applied = results.filter((r) => r.status === "applied").length;
485    await this.audit.audit(req, req?.auth, {
486      action: "ai.writing_applied",
487      resourceType: "ai_suggestion",
488      resourceId: suggestionRow.id,
489      workspaceId,
490      pageId: suggestion.pageId,
491      metadata: { action: suggestion.action, applied, failed: results.length - applied },
492    });
493 
494    return { results };
495  }
496}
497