ai-chat.service.ts

ai-chat.service.tsapps/api/src/modules/ai/ai-chat.service.tsraw ← ai17,311 B← Back to the app
1/**
2 * TASK-0122 — Ask AI streaming service (SSE).
3 *
4 * `POST /workspaces/:workspaceId/ai/conversations/:conversationId/messages`
5 *
6 * Contract (documented in `docs/api/ai.md`):
7 *
8 *   - response: `text/event-stream` of JSON `data:` frames:
9 *       `{type:"delta",  text}`                — assistant content increments
10 *       `{type:"tool_call", id, name, arguments}`
11 *       `{type:"usage",  tokensIn, tokensOut}`
12 *       `{type:"error",  code, message, retryable?, retryAfterSeconds?}`
13 *       `{type:"done",   assistantMessageId?, userMessageId?, reason?}`
14 *     (plus `: keepalive` comment frames every 10 s).
15 *   - context: the client sends RESOURCE IDS only (`{pageId?, blockId?,
16 *     databaseId?}` — the conversation's stored context, optionally
17 *     overridden per message). `prompts.ts` is the single choke point that
18 *     re-checks permissions and assembles the system-prompt section (caps:
19 *     40 blocks / 20 rows; unviewable resources are dropped, never leaked).
20 *   - budget: checked BEFORE any SSE header is written, so an exhausted
21 *     budget is the canonical JSON `402 AI_BUDGET_EXCEEDED` envelope (the
22 *     client's error branch renders the banner from it); errors that arise
23 *     MID-stream (provider outage / timeout) are SSE `error` deltas with the
24 *     same `{code, message, retryable}` shape — raw provider text is never
25 *     surfaced (gateway normalization + `mapAiError`).
26 *   - rate limit: 30 messages / min per user (Redis sliding window) →
27 *     `429 RATE_LIMITED` JSON (pre-stream) carrying `retryAfterSeconds`.
28 *   - cancellation: the client's `AbortController` (fetch abort on panel
29 *     close / the "Stop" button) propagates to `req.signal` → the provider
30 *     stream is cancelled (TASK-0121 abort plumbing) and a PARTIAL assistant
31 *     message is persisted with `metadata.truncated = true`.
32 *   - retry: `body.retryOf = <userMessageId>` re-runs the turn WITHOUT
33 *     duplicating the persisted user message (the UI's Retry button sends it
34 *     after a mid-stream error).
35 *   - auto-title: the first user message (first 60 chars) titles the
36 *     conversation when it has no title yet.
37 *   - persistence: user row (except on retry) + assistant row carrying
38 *     `tokensIn/tokensOut/latencyMs/toolCalls` (the usage-ledger row).
39 *
40 * Pre-stream failures (ownership / budget / rate limit) throw `ApiError`
41 * BEFORE a single byte is written → the global filter's JSON envelope.
42 * Mid-stream failures are in-band `error` deltas.
43 */
44import { Inject, Injectable } from "@nestjs/common";
45import type { Request, Response } from "express";
46 
47import type { PrismaClient } from "@prisma/client";
48import { PrismaService } from "../../common/prisma.service.js";
49import { ApiError } from "../../common/errors.js";
50import { logger } from "../../common/logger.js";
51import { assertCan } from "../../lib/authorize.js";
52import {
53  AI_GATEWAY,
54  mapAiError,
55  primeWorkspaceCache,
56  type AiBudgetGuard,
57} from "./gateway-provider.js";
58import { AiBudgetGuardService } from "./budget-guard.service.js";
59import { AI_CALLS_PER_USER_LIMIT, AiRateLimitService } from "./ai-rate-limit.js";
60import { assembleContext, BASE_SYSTEM_PROMPT, type AiContextRef } from "./prompts.js";
61import type { AiGateway, ChatChunk } from "@factory/ai";
62 
63/** 30 AI messages / min per user (spec §Security; 429 RATE_LIMITED).
64 *  TASK-0123: the cap now lives in the shared `ai.calls` limiter (chat and
65 *  writing count against the SAME window) — re-exported for the integration
66 *  tests that pin the number. */
67export const AI_CHAT_PER_USER_LIMIT = AI_CALLS_PER_USER_LIMIT;
68/** Client-side watchdog: "took too long" if no frame arrives within 30 s. */
69const AI_CHAT_TIMEOUT_MS = 30_000;
70const KEEPALIVE_MS = 10_000;
71 
72export interface AiChatStreamInput {
73  content: string;
74  context?: AiContextRef;
75  /** Re-send this (already-persisted) user message without duplicating it. */
76  retryOf?: string;
77}
78 
79/** One SSE frame (the documented wire contract). */
80export type AiSseEvent =
81  | { type: "delta"; text: string }
82  | { type: "tool_call"; id: string; name: string; arguments: Record<string, unknown> }
83  | { type: "usage"; tokensIn: number; tokensOut: number }
84  | { type: "error"; code: string; message: string; retryable?: boolean; retryAfterSeconds?: number }
85  | { type: "done"; assistantMessageId?: string; userMessageId?: string; reason?: "ok" | "aborted" | "error" };
86 
87function sseFrame(event: AiSseEvent): string {
88  return `data: ${JSON.stringify(event)}\n\n`;
89}
90 
91@Injectable()
92export class AiChatService {
93  constructor(
94    @Inject(PrismaService) private readonly prisma: PrismaClient,
95    @Inject(AI_GATEWAY) private readonly gateway: AiGateway,
96    @Inject(AiBudgetGuardService) private readonly guardService: AiBudgetGuardService,
97    @Inject(AiRateLimitService) private readonly rateLimit: AiRateLimitService,
98  ) {}
99 
100  private get guard(): AiBudgetGuard {
101    return this.guardService.guard;
102  }
103 
104  /**
105   * Stream one Ask AI turn. Pre-stream errors throw (JSON envelope);
106   * mid-stream errors are `error` deltas followed by a terminal `done`.
107   */
108  async sendMessage(
109    workspaceId: string,
110    userId: string,
111    conversationId: string,
112    input: AiChatStreamInput,
113    req: Request,
114    res: Response,
115  ): Promise<void> {
116    // ---- pre-stream gates (canonical JSON envelope on failure) ----
117    await assertCan(this.prisma, { workspaceId, userId }, "VIEW");
118    const conversation = await this.prisma.aiConversation.findUnique({ where: { id: conversationId } });
119    // Own-conversation scoping server-side: another member's conversation is
120    // a 404 (existence hidden across the boundary — error-model convention).
121    if (!conversation || conversation.userId !== userId || conversation.workspaceId !== workspaceId || conversation.deletedAt) {
122      throw ApiError.notFound("Conversation not found.");
123    }
124    await primeWorkspaceCache(this.prisma, workspaceId);
125    try {
126      await this.guard.check(workspaceId, undefined, req);
127    } catch (err) {
128      throw mapAiError(err);
129    }
130 
131    // per-user message rate limit (sliding window; degrades OPEN when Redis
132    // is down — availability over strictness, same policy as the form limiter)
133    const rl = await this.checkRateLimit(userId);
134    if (!rl.allowed) {
135      throw new ApiError("RATE_LIMITED", 429, "AI message limit reached.", [
136        { reason: "ai_message_rate_limited", retryAfterSeconds: rl.retryAfterSeconds },
137      ]);
138    }
139 
140    // ---- context assembly (permission re-check + caps; never leaks) ----
141    const ref: AiContextRef = {
142      ...(input.context ?? (conversation.context as AiContextRef | null) ?? {}),
143    };
144    let ctx;
145    try {
146      ctx = await assembleContext(this.prisma, workspaceId, userId, ref);
147    } catch (err) {
148      logger.warn(`ai context assembly failed for ${userId}: ${(err as Error).message}`);
149      ctx = { systemPrompt: undefined, note: "Context is unavailable for this message.", dropped: ["context"] };
150    }
151 
152    // ---- persist the user message (or resolve the retry) ----
153    let userRow;
154    if (input.retryOf) {
155      const found = await this.prisma.aiMessage.findFirst({
156        where: { conversationId: conversation.id, id: input.retryOf, role: "user" },
157      });
158      if (!found) throw ApiError.notFound("Message not found.");
159      userRow = found;
160    } else {
161      userRow = await this.prisma.aiMessage.create({
162        data: { conversationId: conversation.id, role: "user", content: input.content },
163      });
164      const needsTitle = !conversation.title;
165      const updates: { updatedAt: Date; title?: string } = { updatedAt: new Date() };
166      if (needsTitle) updates.title = input.content.slice(0, 60);
167      await this.prisma.aiConversation.update({ where: { id: conversation.id }, data: updates });
168    }
169 
170    // ---- SSE header (the first byte — after this, errors are in-band) ----
171    res.status(200);
172    res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
173    res.setHeader("Cache-Control", "no-cache, no-transform");
174    res.setHeader("Connection", "keep-alive");
175    res.setHeader("X-Accel-Buffering", "no");
176    res.flushHeaders?.();
177    if (typeof res.socket?.setNoDelay === "function") res.socket.setNoDelay(true);
178 
179    // abort plumbing: client fetch abort → cancel the provider stream.
180    // `res`'s "close" fires on BOTH normal completion and client departure —
181    // the `writableEnded` check distinguishes them (on a clean end,
182    // "finish" already flipped it; on a hang-up it is still false). This is
183    // the reliable signal: `req`'s "close" can fire late/never for a
184    // mid-body fetch abort.
185    const ac = new AbortController();
186    let clientClosed = false;
187    let reason = "ok";
188    const onResClose = (): void => {
189      if (res.writableEnded) return; // normal completion
190      clientClosed = true;
191      ac.abort();
192    };
193    res.on("close", onResClose);
194    const keepalive = setInterval(() => {
195      if (!res.writableEnded) res.write(": keepalive\n\n");
196    }, KEEPALIVE_MS);
197    // client-side watchdog: no provider frame for 30 s → "took too long".
198    let timer: ReturnType<typeof setTimeout> | undefined;
199    const armTimeout = (): void => {
200      if (timer) clearTimeout(timer);
201      timer = setTimeout(() => {
202        ac.abort();
203        reason = "timeout";
204      }, AI_CHAT_TIMEOUT_MS);
205    };
206    armTimeout();
207 
208    let content = "";
209    const toolCalls: Array<{ id: string; name: string; arguments: Record<string, unknown> }> = [];
210    let usage = { in: 0, out: 0 };
211    let finished = false;
212    const chatModel = this.gateway.resolve(workspaceId).chatModel;
213    const startedAt = Date.now();
214 
215    const messages: Array<{ role: "system" | "user" | "assistant"; content: string }> = [
216      { role: "system", content: BASE_SYSTEM_PROMPT },
217    ];
218    if (ctx.systemPrompt) messages.push({ role: "system", content: ctx.systemPrompt });
219    const history = await this.prisma.aiMessage.findMany({
220      where: { conversationId: conversation.id, id: { not: userRow.id } },
221      orderBy: { createdAt: "desc" },
222      take: 20,
223      select: { role: true, content: true },
224    });
225    for (const m of [...history].reverse()) {
226      if (m.role === "user" || m.role === "assistant") messages.push({ role: m.role, content: m.content });
227    }
228    messages.push({ role: "user", content: input.content });
229 
230    // context-unavailable note (permission drop) — the FIRST delta, so the
231    // panel can render it even when the provider call fails before deltas.
232    if (ctx.note) res.write(sseFrame({ type: "delta", text: `> ${ctx.note}\n\n` }));
233 
234    const persistAssistant = async (opts: { content: string; truncated: boolean; error?: string }) => {
235      try {
236        // Only persist when there is visible content — an empty assistant
237        // row (abort before the first delta / provider error pre-output)
238        // would pollute the history; the user row + the error banner cover
239        // the UI state in those cases.
240        if (!opts.content) return undefined;
241        return await this.prisma.aiMessage.create({
242          data: {
243            conversationId: conversation.id,
244            role: "assistant",
245            content: opts.content,
246            model: chatModel,
247            ...(usage.in > 0 ? { tokensIn: usage.in } : {}),
248            ...(usage.out > 0 ? { tokensOut: usage.out } : {}),
249            latencyMs: Date.now() - startedAt,
250            ...(toolCalls.length > 0 ? { toolCalls: toolCalls as unknown as object } : {}),
251            metadata: {
252              feature: "chat",
253              ...(opts.truncated ? { truncated: true } : {}),
254              ...(opts.error ? { error: opts.error } : {}),
255            } as object,
256          },
257        });
258      } catch (err) {
259        // persistence must never crash the stream teardown
260        logger.error(`ai chat assistant persist failed: ${(err as Error).message}`);
261        return undefined;
262      }
263    };
264    const finish = (event: AiSseEvent): void => {
265      if (finished) return;
266      finished = true;
267      if (timer) clearTimeout(timer);
268      clearInterval(keepalive);
269      res.off("close", onResClose);
270      if (!res.writableEnded) {
271        res.write(sseFrame(event));
272        res.end();
273      }
274    };
275    const sendErrorDelta = (err: unknown): void => {
276      const mapped = mapAiError(err);
277      const detail = (mapped.details?.[0] ?? {}) as Record<string, unknown>;
278      // `retryable` rides on the mapped details (the gateway's normalization
279      // sets it for 429/5xx/timeout; provider 4xx and budget 402 → false).
280      // A rate limit is inherently transient → retryable (the UI shows a
281      // countdown from `retryAfterSeconds`).
282      const retryable = detail.retryable === true || detail.retryAfterSeconds != null || mapped.code === "RATE_LIMITED";
283      const event: AiSseEvent = {
284        type: "error",
285        code: mapped.code,
286        // raw provider text is never surfaced — mapped.message is the
287        // normalized (no-secret) message from the gateway error map.
288        message: mapped.message,
289        ...(retryable ? { retryable: true } : {}),
290        ...(typeof detail.retryAfterSeconds === "number" ? { retryAfterSeconds: detail.retryAfterSeconds } : {}),
291      };
292      res.write(sseFrame(event));
293      logger.info(
294        `ai chat stream error workspace=${workspaceId} conversation=${conversation.id} code=${mapped.code} retryable=${retryable}`,
295        { event: "ai.chat.error", workspaceId, conversationId, code: mapped.code, retryable },
296      );
297    };
298 
299    try {
300      for await (const chunk of this.gateway.chat({
301        workspaceId,
302        userId,
303        capability: "chat",
304        messages,
305        metadata: { conversationId: conversation.id },
306        signal: ac.signal,
307      })) {
308        if (timer) armTimeout();
309        if (chunk.delta) {
310          content += chunk.delta;
311          res.write(sseFrame({ type: "delta", text: chunk.delta }));
312        }
313        if (chunk.toolCallDelta) {
314          toolCalls.push({
315            id: chunk.toolCallDelta.id,
316            name: chunk.toolCallDelta.name,
317            arguments: chunk.toolCallDelta.arguments,
318          });
319          res.write(sseFrame({ type: "tool_call", ...chunk.toolCallDelta }));
320        }
321        if (chunk.done) {
322          if (chunk.usage) {
323            usage = { in: chunk.usage.in, out: chunk.usage.out };
324            res.write(sseFrame({ type: "usage", tokensIn: usage.in, tokensOut: usage.out }));
325          }
326          const row = await persistAssistant({ content, truncated: false });
327          finish({ type: "done", assistantMessageId: row?.id, userMessageId: userRow.id, reason: "ok" });
328        }
329      }
330      // provider stream ended without a done chunk (defensive) — persist
331      // once (the done-chunk branch above already finishes when it fires).
332      if (!finished) {
333        const row = await persistAssistant({ content, truncated: false });
334        finish({ type: "done", assistantMessageId: row?.id, userMessageId: userRow.id, reason: "ok" });
335      }
336    } catch (err) {
337      if (timer) clearTimeout(timer);
338      clearInterval(keepalive);
339      res.off("close", onResClose);
340 
341      if (reason === "timeout") {
342        // idle watchdog: the provider stopped sending frames (30 s) — the
343        // client's fetch is still connected; tell it in-band and persist the
344        // partial (metadata.truncated).
345        const row = await persistAssistant({ content, truncated: content.length > 0, error: "timeout" });
346        if (!clientClosed) {
347          res.write(sseFrame({ type: "error", code: "TIMEOUT", message: "The AI request took too long.", retryable: true }));
348        }
349        finish({ type: "done", assistantMessageId: row?.id, userMessageId: userRow.id, reason: "error" });
350        return;
351      }
352      if (clientClosed || ac.signal.aborted) {
353        // client hang-up (panel closed / tab switched) or the "Stop" button:
354        // the provider stream was cancelled; persist the partial with
355        // metadata.truncated (the "Stop" case still gets the terminal done).
356        const row = await persistAssistant({ content, truncated: content.length > 0, error: "aborted" });
357        if (!clientClosed && !res.writableEnded) {
358          res.write(sseFrame({ type: "done", assistantMessageId: row?.id, userMessageId: userRow.id, reason: "aborted" }));
359          res.end();
360        }
361        return;
362      }
363      // provider failure (normalized — raw provider text never surfaces):
364      // error delta + terminal done; persist any partial content.
365      sendErrorDelta(err);
366      const row = await persistAssistant({ content, truncated: content.length > 0, error: "provider_error" });
367      finish({ type: "done", assistantMessageId: row?.id, userMessageId: userRow.id, reason: "error" });
368    }
369  }
370 
371  /**
372   * 30 AI messages / min per user — the SHARED `ai.calls` window
373   * (`ai:calls:rl:<userId>`, TASK-0123 `AiRateLimitService`): chat and the
374   * writing tools count against the SAME limiter. Degrades OPEN when Redis is
375   * unavailable (same policy as the form limiter).
376   */
377  private checkRateLimit(userId: string): Promise<{ allowed: boolean; remaining: number; retryAfterSeconds: number }> {
378    return this.rateLimit.check(userId);
379  }
380}
381