generate.service.ts

generate.service.tsapps/api/src/modules/ai/generator/generate.service.tsraw ← generator22,562 B← Back to the app
1/**
2 * TASK-0126 — AI page generation (prompt → structured tool calls executed
3 * through the domain services with the user's permissions).
4 *
5 *   POST /workspaces/:workspaceId/ai/generate   (body `{ prompt, parentId? }`,
6 *   header `Idempotency-Key?`)   → SSE progress stream.
7 *
8 * The wire contract (docs/api/ai.md §10):
9 *
10 *   tool* (per step: {type:"tool", name, status:"started"|"done", ref?})
11 *     → done { pageId, created:{pages,blocks,databases,relations}, tokensIn, tokensOut }
12 *   or an in-band error (a step failed / the plan was invalid / timeout /
13 *   budget / rate-limit) + a terminal done { reason }.
14 *
15 * Security (spec §Security, AGENTS #7/#18/#19/#20/#27):
16 *   - pre-stream: non-member → 404 (fail-closed, like the pages API + the
17 *     RAG fix), no-EDIT → 403 FORBIDDEN, budget → 402 AI_BUDGET_EXCEEDED
18 *     (the guard also fires the 24 h-throttled admin notify + audit),
19 *     rate-limit → 429, validation → 400;
20 *   - the LLM emits a PLAN (structured tool calls) — its UNTRUSTED output,
21 *     validated by `parseGeneratePlan` (zod) before ANY write;
22 *   - the executor runs the plan through the DOMAIN services with the
23 *     REQUESTING user's context (the permission-context hook rejects an
24 *     elevated context); a failing step rolls back the partial tree;
25 *   - total capped at 120 s (injectable `timeoutMs`; the test overrides to
26 *     2 s) → in-band `ai_timeout`; the client's SSE abort → the executor stops
27 *     after the current step, rolls back, and emits `cancelled`;
28 *   - cost: one hidden per-call conversation (`metadata.feature="generate"`)
29 *     with full token counts; audit `ai.page_generated` / `ai.generation_failed`;
30 *   - idempotency: the `Idempotency-Key` header memoizes the created page id
31 *     (Redis, 24 h) so a retried cancelled generation does not double-create.
32 */
33import { Inject, Injectable } from "@nestjs/common";
34import type { Request, Response } from "express";
35 
36import type { PrismaClient } from "@prisma/client";
37import { parseGeneratePlan } from "@factory/ai";
38import { PrismaService } from "../../../common/prisma.service.js";
39import { AuditService } from "../../../common/audit.service.js";
40import { ApiError } from "../../../common/errors.js";
41import { logger } from "../../../common/logger.js";
42import { actorMembership, assertCan } from "../../../lib/authorize.js";
43import { RedisService } from "../../../common/redis.service.js";
44import {
45  AI_GATEWAY,
46  mapAiError,
47  primeWorkspaceCache,
48} from "../gateway-provider.js";
49import { GENERATE_PLAN_SYSTEM_PROMPT, GENERATE_PLAN_TOOL, normalizeGeneratePlanArgs } from "../prompts.js";
50import { AiBudgetGuardService } from "../budget-guard.service.js";
51import { AiRateLimitService } from "../ai-rate-limit.js";
52import { PageService } from "../../pages/page.service.js";
53import { BlockService } from "../../blocks/block.service.js";
54import { DatabaseService } from "../../databases/database.service.js";
55import { DatabaseQueryService } from "../../databases/query.service.js";
56import type { AiGateway } from "@factory/ai";
57import {
58  GenerationAborted,
59  GenerationExecutor,
60  GenerationStepFailed,
61  type StepName,
62} from "./executor.js";
63import type { AiGenerateInput } from "./generate.dto.js";
64 
65/** The hidden per-call conversation's title (never lists in the panel). */
66export const GENERATE_CONV_PREFIX = "generate:";
67/** The 24 h idempotency memo TTL (a cancelled run's key lives a day). */
68const IDEM_TTL_SECONDS = 24 * 60 * 60;
69/** The server-side total cap (the spec's 120 s). Read PER CALL (not at boot)
70 *  so the integration/E2E timeout suite can override `AI_GENERATE_TIMEOUT_MS`
71 *  to 2 s without a service-instance seam. */
72const DEFAULT_TIMEOUT_MS = 120_000;
73/** SSE comment-keepalive interval — well under any proxy/browser idle cap. */
74const KEEPALIVE_MS = 15_000;
75 
76/** The SSE frame contract (mirrors the web client's `GenerateFrame` — the
77 *  frontend treats exactly one of `done` / `cancelled` / `error` as the
78 *  terminal frame per stream; `tool` frames precede it). */
79export type AiGenerateSseEvent =
80  | { type: "tool"; name: StepName; status: "started" | "done"; ref?: string }
81  | {
82      type: "done";
83      pageId: string;
84      created: { pages: string[]; blocks: number; databases: string[]; relations: number };
85      tokensIn?: number;
86      tokensOut?: number;
87    }
88  | { type: "cancelled" }
89  | { type: "error"; code: string; message: string; reason?: string; retryable?: boolean };
90 
91function sseFrame(event: AiGenerateSseEvent): string {
92  return `data: ${JSON.stringify(event)}\n\n`;
93}
94 
95/** True when the error is the gateway's caller-abort (a mid-stream abort
96 *  maps to a shared `AI_ERROR` with `details.reason === "ai_aborted"`). Used
97 *  to distinguish a TIMEOUT / client-cancel that hit the plan fetch (before
98 *  the executor created anything) from a genuine provider failure. */
99function isAiAborted(err: unknown): boolean {
100  const d = (err as { details?: unknown } | null)?.details;
101  if (d && typeof d === "object") {
102    if (Array.isArray(d)) return d.some((x) => (x as { reason?: unknown })?.reason === "ai_aborted");
103    return (d as { reason?: unknown }).reason === "ai_aborted";
104  }
105  return false;
106}
107 
108@Injectable()
109export class AiGenerateService {
110  /** The total generation cap (per call; the timeout suite overrides the env). */
111  get timeoutMs(): number {
112    const envMs = Number(process.env.AI_GENERATE_TIMEOUT_MS);
113    return Number.isFinite(envMs) && envMs > 0 ? envMs : DEFAULT_TIMEOUT_MS;
114  }
115 
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(RedisService) private readonly redis: RedisService,
122    @Inject(PageService) private readonly pages: PageService,
123    @Inject(BlockService) private readonly blocks: BlockService,
124    @Inject(DatabaseService) private readonly databases: DatabaseService,
125    @Inject(DatabaseQueryService) private readonly records: DatabaseQueryService,
126    @Inject(AuditService) private readonly audit: AuditService,
127  ) {}
128 
129  /**
130   * The generate endpoint. Pre-stream failures throw (→ JSON envelope); the
131   * SSE stream carries the progress `tool` frames + the terminal `done` /
132   * in-band `error`.
133   */
134  async generate(
135    workspaceId: string,
136    userId: string,
137    input: AiGenerateInput,
138    idempotencyKey: string | undefined,
139    req: Request,
140    res: Response,
141  ): Promise<void> {
142    // ---- pre-stream gates (before any SSE byte) ----
143    // Non-member → 404 (fail-closed: hide the workspace's existence, like the
144    // pages API + the RAG gate-order fix). A member without EDIT → 403.
145    // TASK-0127: actor-aware — an API key of the workspace is a member-
146    // equivalent (its scope, not the creator's membership, decides).
147    const membership = await actorMembership(this.prisma, userId, workspaceId, "EDIT");
148    if (!membership) throw ApiError.notFound("Workspace not found.");
149    await assertCan(this.prisma, { workspaceId, userId }, "EDIT");
150 
151    await this.rateLimitCheck(userId);
152 
153    // Budget pre-check (before ANY model call) + the 24 h admin notify/audit.
154    await primeWorkspaceCache(this.prisma, workspaceId);
155    try {
156      await this.guardService.guard.check(workspaceId, undefined, req);
157    } catch (err) {
158      throw mapAiError(err);
159    }
160 
161    // Idempotency: a memoized (cancelled/failed-then-retried) key replays the
162    // created page without re-running the executor (no double-create).
163    if (idempotencyKey) {
164      const memo = await this.readIdem(workspaceId, userId, idempotencyKey);
165      if (memo) {
166        this.sseHeaders(res);
167        const finish = (e: AiGenerateSseEvent): void => {
168          if (!res.writableEnded) {
169            res.write(sseFrame(e));
170            res.end();
171          }
172        };
173        finish({
174          type: "done",
175          pageId: memo.pageId,
176          created: memo.created,
177          tokensIn: memo.tokensIn,
178          tokensOut: memo.tokensOut,
179        });
180        return;
181      }
182    }
183 
184    // ---- SSE ----
185    this.sseHeaders(res);
186    const ac = new AbortController();
187    let clientClosed = false;
188    // The total-timeout watchdog fired (distinct from a client cancel): the
189    // gateway surfaces a mid-stream abort as a shared `AI_ERROR`/`ai_aborted`
190    // (NOT via `signal.reason` — that's getter-only in strict mode), so the
191    // service classifies timeouts by THIS flag, not the error type.
192    let timedOut = false;
193    const onClose = (): void => {
194      if (res.writableEnded) return;
195      clientClosed = true;
196      ac.abort(); // the executor stops after the current step + rolls back
197    };
198    res.on("close", onClose);
199    // The total cap (the spec's 120 s; the test overrides `this.timeoutMs`).
200    const timer = setTimeout(() => {
201      timedOut = true;
202      ac.abort();
203    }, this.timeoutMs);
204    // SSE keepalive: the provider's thinking phase can emit ZERO bytes for
205    // minutes (a successful run has been observed silent for ~150 s), and a
206    // silent stream looks dead to the client (the dialog sat on "Drafting the
207    // page…" with no way to distinguish "working" from "stuck"). A comment
208    // frame every 15 s is invisible to the client's SSE parser (it only
209    // reads `data:` lines) but resets every proxy/browser idle timer and the
210    // client's own idle watchdog can now detect a GENUINELY dead stream.
211    const keepAlive = setInterval(() => {
212      if (!clientClosed && !res.writableEnded) res.write(": keepalive\n\n");
213    }, KEEPALIVE_MS);
214    const cleanup = (): void => {
215      clearTimeout(timer);
216      clearInterval(keepAlive);
217      res.off("close", onClose);
218    };
219    const finish = (event: AiGenerateSseEvent): void => {
220      cleanup();
221      if (!res.writableEnded) {
222        res.write(sseFrame(event));
223        res.end();
224      }
225    };
226 
227    // The hidden per-call conversation (cost ledger; never lists in the panel).
228    const conversation = await this.prisma.aiConversation.create({
229      data: {
230        workspaceId,
231        userId,
232        title: GENERATE_CONV_PREFIX,
233        context: { feature: "generate", parentId: input.parentId ?? null } as object,
234      },
235    });
236    await this.prisma.aiMessage.create({
237      data: { conversationId: conversation.id, role: "user", content: `generate: "${input.prompt.slice(0, 200)}"` },
238    });
239 
240    let usage = { in: 0, out: 0 };
241    // The resolved chat model, captured ONCE now — before the (minutes-long)
242    // model call. Recorded decision (2025, root cause of an API crash): the
243    // earlier code re-resolved via `this.gateway.resolve(workspaceId)` in the
244    // post-stream persist calls, which re-reads the 30 s TTL workspace-row
245    // cache; a model call longer than 30 s expired it, `resolve` threw 404
246    // ("Workspace not found"), and inside the catch block that throw escaped
247    // into the exception filter AFTER the SSE headers were already sent →
248    // uncaught `ERR_HTTP_HEADERS_SENT` → the API process died. The model name
249    // is a static property of the resolved config for the lifetime of this
250    // call, so capturing it early is exact.
251    const chatModel = this.gateway.resolve(workspaceId).chatModel;
252    try {
253      // 1) The model emits the PLAN (one tool call; the mock is deterministic).
254      const plan = await this.chatPlan(workspaceId, userId, input.prompt, ac.signal);
255      usage = plan.usage;
256 
257      // 2) Validate the UNTRUSTED plan (zod) before ANY write.
258      const parsed = parseGeneratePlan(plan.args);
259      if (!parsed.ok) {
260        // Observability (AGENTS #26): the plan is untrusted model output —
261        // log WHY it failed so a model-shape regression is diagnosable.
262        logger.warn(
263          `ai generate: invalid plan from model ${chatModel}: ` +
264            parsed.errors.issues.map((i) => `${i.path.join(".") || "(root)"} ${i.message}`).slice(0, 8).join("; "),
265        );
266        await this.persistAssistant(conversation.id, chatModel, usage, {
267          feature: "generate",
268          error: "ai_invalid_plan",
269        });
270        // Terminal in-band error (no resource was created — nothing to roll
271        // back; the plan never ran).
272        finish({
273          type: "error",
274          code: "UNPROCESSABLE_ENTITY",
275          message: "The AI returned an invalid page plan.",
276          reason: "ai_invalid_plan",
277        });
278        return;
279      }
280 
281      // 3) Execute the plan through the domain services (the user's context).
282      const executor = new GenerationExecutor(
283        { pages: this.pages, blocks: this.blocks, databases: this.databases, records: this.records, prisma: this.prisma },
284        {
285          workspaceId,
286          userId,
287          expectedUserId: userId,
288          req,
289          signal: ac.signal,
290          onStep: (name, status, ref) => {
291            if (!clientClosed && !res.writableEnded) res.write(sseFrame({ type: "tool", name, status, ...(ref ? { ref } : {}) }));
292          },
293        },
294      );
295      const result = await executor.execute(parsed.plan, input.parentId);
296 
297      // 4) Success: audit + memo the idempotency key + persist usage.
298      await this.audit.audit(req, req?.auth, {
299        action: "ai.page_generated",
300        resourceType: "page",
301        resourceId: result.pageId,
302        workspaceId,
303        metadata: { steps: parsed.plan.tools.length, created: result.created },
304      });
305      await this.persistAssistant(conversation.id, chatModel, usage, {
306        feature: "generate",
307        pageId: result.pageId,
308        created: result.created,
309      });
310      await this.writeIdem(workspaceId, userId, idempotencyKey, {
311        pageId: result.pageId,
312        created: result.created,
313        tokensIn: usage.in,
314        tokensOut: usage.out,
315      });
316      finish({
317        type: "done",
318        pageId: result.pageId,
319        created: result.created,
320        tokensIn: usage.in,
321        tokensOut: usage.out,
322      });
323    } catch (err) {
324      // ---- terminal classification (exactly ONE frame per path) ----
325      // Total-timeout watchdog fired (or the gateway's abort surfaced before
326      // the executor ran, while the client is still here): the partial tree
327      // is rolled back (by the executor when it was running, by the abort
328      // otherwise) and the run did NOT complete — a retryable in-band error.
329      if (timedOut || (isAiAborted(err) && !clientClosed)) {
330        await this.audit.audit(req, req?.auth, {
331          action: "ai.generation_failed",
332          resourceType: "page",
333          resourceId: conversation.id,
334          workspaceId,
335          metadata: { failedStep: "timeout", reason: "ai_timeout" },
336        });
337        if (!clientClosed) {
338          finish({
339            type: "error",
340            code: "SERVICE_UNAVAILABLE",
341            message: "Generation took too long and was cancelled (partial work rolled back).",
342            reason: "ai_timeout",
343            retryable: true,
344          });
345        } else {
346          cleanup();
347        }
348        return;
349      }
350      if (err instanceof GenerationAborted) {
351        // The client went away (or cancelled): if the socket is still open
352        // emit the terminal `cancelled`; otherwise just release resources.
353        if (!clientClosed) finish({ type: "cancelled" });
354        else cleanup();
355        return;
356      }
357      if (err instanceof GenerationStepFailed) {
358        // The executor already rolled back the partial tree.
359        // Observability (AGENTS #26): the step error (usually the model's
360        // non-conforming arg shape) is the thing to diagnose next.
361        const stepMsg = (err.error as Error | null)?.message ?? String(err.error);
362        logger.warn(`ai generate: step ${err.failedStep} failed: ${stepMsg.slice(0, 500)}`);
363        await this.audit.audit(req, req?.auth, {
364          action: "ai.generation_failed",
365          resourceType: "page",
366          resourceId: conversation.id,
367          workspaceId,
368          // Truncated: the audit row is also user-visible; keep it short.
369          metadata: { failedStep: err.failedStep, reason: stepMsg.slice(0, 300) },
370        });
371        await this.persistAssistant(conversation.id, chatModel, usage, {
372          feature: "generate",
373          error: "generation_failed",
374          failedStep: err.failedStep,
375        });
376        if (!clientClosed) {
377          finish({
378            type: "error",
379            code: "UNPROCESSABLE_ENTITY",
380            message: `Page generation failed at step ${err.failedStep} (partial work rolled back).`,
381            reason: "generation_failed",
382            retryable: true,
383          });
384        } else {
385          cleanup();
386        }
387        return;
388      }
389      if (err instanceof ApiError && err.code === "RATE_LIMITED") {
390        // Provider 429 (the gateway's normalizeProviderError). Pre-fix this
391        // fell through to the generic `INTERNAL` frame below, which the
392        // client rendered as an opaque failure with no retry guidance; now the
393        // dialog shows the rate-limit banner + a Retry affordance.
394        logger.warn(`ai generate: provider rate-limited: ${(err as Error).message}`);
395        if (!clientClosed) {
396          finish({
397            type: "error",
398            code: "RATE_LIMITED",
399            message: err.message || "AI provider rate limit reached. Try again shortly.",
400            reason: "rate_limited",
401            retryable: true,
402          });
403        } else {
404          cleanup();
405        }
406        return;
407      }
408      // A gateway/network error that slipped through (defensive).
409      logger.warn(`ai generate failed: ${(err as Error).message}`);
410      if (!clientClosed) {
411        finish({ type: "error", code: "INTERNAL", message: "Page generation failed.", reason: "ai_error", retryable: true });
412      } else {
413        cleanup();
414      }
415    }
416  }
417 
418  // ------------------------------------------------------------- helpers ---
419 
420  private async rateLimitCheck(userId: string): Promise<void> {
421    const rl = await this.rateLimit.check(userId);
422    if (!rl.allowed) {
423      throw new ApiError("RATE_LIMITED", 429, "AI call limit reached.", [
424        { reason: "ai_call_rate_limited", retryAfterSeconds: rl.retryAfterSeconds },
425      ]);
426    }
427  }
428 
429  private sseHeaders(res: Response): void {
430    res.status(200);
431    res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
432    res.setHeader("Cache-Control", "no-cache, no-transform");
433    res.setHeader("Connection", "keep-alive");
434    res.setHeader("X-Accel-Buffering", "no");
435    res.flushHeaders?.();
436    if (typeof res.socket?.setNoDelay === "function") res.socket.setNoDelay(true);
437  }
438 
439  /**
440   * The LLM's plan (one tool call; the mock is the deterministic plan).
441   *
442   * The `generatePlan` tool is FORCED (`toolChoice`) — recorded decision
443   * (2025, root cause of the "Generation took too long" failure): with only a
444   * system-prompt instruction, the OSS reasoning model streams minutes of
445   * prose and never emits the plan, so the 120 s watchdog always fired. The
446   * adapter sends `tool_choice: {type:"function", function:{name}}` for this.
447   * The model's args are UNTRUSTED and get a small, total normalization pass
448   * (`normalizeGeneratePlanArgs`) before the strict zod validation in the
449   * caller (repairs the observed tools-as-string + `$page`-placeholder
450   * shapes; see prompts.ts).
451   */
452  private async chatPlan(
453    workspaceId: string,
454    userId: string,
455    prompt: string,
456    signal: AbortSignal,
457  ): Promise<{ args: unknown; usage: { in: number; out: number } }> {
458    const messages: Array<{ role: "system" | "user"; content: string }> = [
459      { role: "system", content: GENERATE_PLAN_SYSTEM_PROMPT },
460      { role: "user", content: prompt },
461    ];
462    let args: unknown;
463    let usage = { in: 0, out: 0 };
464    for await (const chunk of this.gateway.chat({
465      workspaceId,
466      userId,
467      capability: "chat",
468      messages,
469      tools: [GENERATE_PLAN_TOOL],
470      toolChoice: { name: "generatePlan" },
471      metadata: { mode: "generate" },
472      signal,
473    })) {
474      if (chunk.toolCallDelta) args = chunk.toolCallDelta.arguments;
475      if (chunk.done && chunk.usage) usage = { in: chunk.usage.in, out: chunk.usage.out };
476    }
477    if (args === undefined) {
478      throw ApiError.unprocessable("The model returned no plan.", [{ reason: "ai_invalid_plan", issue: "no tool call" }]);
479    }
480    return { args: normalizeGeneratePlanArgs(args), usage };
481  }
482 
483  private async persistAssistant(
484    conversationId: string,
485    model: string,
486    usage: { in: number; out: number },
487    metadata: Record<string, unknown>,
488  ): Promise<string | undefined> {
489    try {
490      const row = await this.prisma.aiMessage.create({
491        data: {
492          conversationId,
493          role: "assistant",
494          content: metadata.error ? `error: ${String(metadata.error)}` : "ok",
495          model,
496          ...(usage.in > 0 ? { tokensIn: usage.in } : {}),
497          ...(usage.out > 0 ? { tokensOut: usage.out } : {}),
498          metadata: metadata as object,
499        },
500      });
501      return row.id;
502    } catch (err) {
503      logger.error(`ai generate assistant persist failed: ${(err as Error).message}`);
504      return undefined;
505    }
506  }
507 
508  private idemKey(workspaceId: string, userId: string, key: string): string {
509    return `ai:generate:idem:${workspaceId}:${userId}:${key}`;
510  }
511 
512  private async readIdem(
513    workspaceId: string,
514    userId: string,
515    key: string,
516  ): Promise<{ pageId: string; created: { pages: string[]; blocks: number; databases: string[]; relations: number }; tokensIn: number; tokensOut: number } | null> {
517    try {
518      const raw = await this.redis.getRedis().get(this.idemKey(workspaceId, userId, key));
519      if (!raw) return null;
520      return JSON.parse(raw);
521    } catch {
522      return null;
523    }
524  }
525 
526  private async writeIdem(
527    workspaceId: string,
528    userId: string,
529    key: string | undefined,
530    value: { pageId: string; created: { pages: string[]; blocks: number; databases: string[]; relations: number }; tokensIn: number; tokensOut: number },
531  ): Promise<void> {
532    if (!key) return;
533    try {
534      await this.redis.getRedis().set(this.idemKey(workspaceId, userId, key), JSON.stringify(value), "EX", IDEM_TTL_SECONDS);
535    } catch (err) {
536      logger.warn(`ai generate idem store failed: ${(err as Error).message}`);
537    }
538  }
539}
540 
541