@k2b/cloud 0.6.0-rc.0

package/src/ai/stream.ts

@@ -0,0 +1,308 @@
+import { lazySync } from "../_internal/process-sync";
+import { logger } from "../services/logging";
+import { latestTopicCursor } from "../services/topic-cursor";
+import {
+  type AiStreamEvent,
+  type AiTurnSnapshot,
+  type AiWireEvent,
+  isNewerWireEvent,
+  reconcileResolvedTurnActions,
+  steerAppliedBlockId,
+  steerMessageBlockId,
+} from "./protocol";
+import { projectPublicAiStoredMessages, publicAiStoredMessages } from "./public-projection";
+import { aiConversations } from "./store";
+import type { AiConversation } from "./types";
+
+const log = logger("ai:stream");
+
+/**
+ * Live fanout for wire events. Events carry their full payload so the SSE hot
+ * path never touches Postgres; durable state lives in ai.messages plus the
+ * throttled ai.turns.live_blocks snapshot.
+ */
+export const aiStreamTopic = lazySync((sync) =>
+  sync.topic<AiWireEvent>({
+    id: "cloud-ai-stream",
+    owner: "cloud",
+    retention: { maxAgeMs: 15 * 60 * 1000, maxBytes: 256 * 1024 * 1024 },
+    maxPayloadBytes: 257 * 1024,
+  }),
+);
+
+export type AiTurnControlEvent = { type: "abort"; conversationId: string; turnId: string };
+
+export const aiTurnControlsTopic = lazySync((sync) =>
+  sync.topic<AiTurnControlEvent>({
+    id: "cloud-ai-turn-controls",
+    owner: "cloud",
+    // Abort signals are rare and tiny (5 KiB cap): 4 MiB keeps ~800 of them
+    // across the 15-minute window.
+    retention: { maxAgeMs: 15 * 60 * 1000, maxBytes: 4 * 1024 * 1024 },
+    maxPayloadBytes: 5 * 1024,
+  }),
+);
+
+export const publishAiWireEvent = async (event: AiWireEvent): Promise<void> => {
+  await aiStreamTopic().publish({
+    tenantId: event.conversationId,
+    orderingKey: event.turnId,
+    data: event,
+    idempotencyKey: `wire:${event.turnId}:${event.attempt}:${event.seq}`,
+  });
+};
+
+export const publishAiTurnAbort = async (input: { conversationId: string; turnId: string }): Promise<void> => {
+  await aiTurnControlsTopic().publish({
+    tenantId: input.conversationId,
+    orderingKey: input.turnId,
+    data: { type: "abort", conversationId: input.conversationId, turnId: input.turnId },
+    idempotencyKey: `turn-abort:${input.turnId}`,
+  });
+};
+
+const encoder = new TextEncoder();
+const DEFAULT_HEARTBEAT_MS = 5_000;
+
+export const sseHeaders = {
+  "Content-Type": "text/event-stream; charset=utf-8",
+  "Cache-Control": "no-cache, no-transform",
+  Connection: "keep-alive",
+  "X-Accel-Buffering": "no",
+} as const;
+
+export const encodeSseEvent = (event: AiStreamEvent): Uint8Array =>
+  encoder.encode(`event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`);
+
+export const encodeSseHeartbeat = (): Uint8Array => encoder.encode(": heartbeat\n\n");
+
+const turnSnapshotFromActive = (active: NonNullable<Awaited<ReturnType<typeof aiConversations.getActiveTurn>>>): AiTurnSnapshot => ({
+  turnId: active.turn.shortId,
+  attempt: active.turn.attempt,
+  status: active.turn.status,
+  seq: active.liveSeq,
+  blocks: [...active.liveBlocks],
+  modelProfileId: active.turn.modelProfileId,
+  createdAt: active.turn.createdAt,
+});
+
+/** Initial history window; older messages load on demand while scrolling up. */
+export const AI_STREAM_INITIAL_MESSAGE_LIMIT = 100;
+
+export const loadAiStreamState = async (conversation: AiConversation): Promise<Extract<AiStreamEvent, { type: "state" }>> => {
+  const [page, active] = await Promise.all([
+    aiConversations.listMessagesPage({ conversationId: conversation.id, limit: AI_STREAM_INITIAL_MESSAGE_LIMIT }),
+    aiConversations.getActiveTurn({ conversationId: conversation.id }),
+  ]);
+  const snapshot = active ? turnSnapshotFromActive(active) : null;
+  if (snapshot) {
+    const [steers, resolvedActions] = await Promise.all([
+      aiConversations.listTurnSteers({ conversationId: conversation.id, turnId: active!.turn.id }),
+      aiConversations.listResolvedPendingActions({ conversationId: conversation.id, turnId: active!.turn.id }),
+    ]);
+    snapshot.blocks = reconcileResolvedTurnActions(snapshot.blocks, resolvedActions);
+    const known = new Set(snapshot.blocks.map((block) => block.id));
+    for (const steer of steers) {
+      if (steer.status === "discarded" || known.has(steerMessageBlockId(steer.id))) continue;
+      snapshot.blocks.push({
+        id: steerMessageBlockId(steer.id),
+        kind: "steer_message",
+        steerId: steer.id,
+        text: steer.text,
+        status: steer.status === "pending" ? "pending" : "consumed",
+      });
+      if (steer.status === "consumed" && !known.has(steerAppliedBlockId(steer.id))) {
+        snapshot.blocks.push({ id: steerAppliedBlockId(steer.id), kind: "steer_applied", steerId: steer.id });
+      }
+    }
+  }
+  return {
+    type: "state",
+    conversation: { ...conversation, id: conversation.shortId },
+    messages: await publicAiStoredMessages(page.messages, conversation),
+    hasMoreMessages: page.hasMore,
+    activeTurn: snapshot,
+  };
+};
+
+type SnapshotTailItem<TSnapshot, TEvent> = { kind: "snapshot"; value: TSnapshot } | { kind: "event"; value: TEvent };
+
+async function* streamSnapshotThenTail<TCursor, TSnapshot, TEvent>(input: {
+  captureCursor: () => Promise<TCursor>;
+  loadSnapshot: () => Promise<TSnapshot>;
+  tail: (cursor: TCursor) => AsyncIterable<TEvent>;
+}): AsyncGenerator<SnapshotTailItem<TSnapshot, TEvent>> {
+  const cursor = await input.captureCursor();
+  yield { kind: "snapshot", value: await input.loadSnapshot() };
+  for await (const event of input.tail(cursor)) yield { kind: "event", value: event };
+}
+
+/**
+ * Transport-neutral conversation stream: one `state` snapshot, then the live
+ * tail. SSE and WebSocket adapters must both consume this feed.
+ *
+ * The topic cursor is grabbed before the snapshot is loaded, so every event
+ * that races the snapshot is replayed from the tail and deduplicated via
+ * (attempt, seq). Events of unknown turns are dropped until their
+ * `turn_started` arrives, which makes stale retention entries harmless.
+ * A memoized per-conversation hub preserves this snapshot race guarantee. While
+ * a conversation has subscribers it costs one full-topic follower per process
+ * (Sync filters tenants locally); the hub retires when the last subscriber
+ * leaves. live() has no subscription-ready barrier.
+ */
+export async function* streamAiConversationEvents(input: {
+  conversation: AiConversation;
+  signal: AbortSignal;
+}): AsyncGenerator<AiStreamEvent> {
+  let current: { turnId: string; publicTurnId: string; attempt: number; seq: number } | null = null;
+  for await (const item of streamSnapshotThenTail({
+    captureCursor: () => latestTopicCursor({ topic: aiStreamTopic(), resourceId: "cloud-ai-stream", tenantId: input.conversation.id }),
+    loadSnapshot: () => loadAiStreamState(input.conversation),
+    tail: (after) => aiStreamTopic().hub({ tenantId: input.conversation.id }).subscribe({ after, signal: input.signal }),
+  })) {
+    if (item.kind === "snapshot") {
+      const state = item.value;
+      yield state;
+      const activeTurn = state.activeTurn
+        ? await aiConversations.getTurnByShortId({
+            conversationId: input.conversation.id,
+            shortId: state.activeTurn.turnId,
+          })
+        : null;
+      current =
+        state.activeTurn && activeTurn
+          ? {
+              turnId: activeTurn.id,
+              publicTurnId: state.activeTurn.turnId,
+              attempt: state.activeTurn.attempt,
+              seq: state.activeTurn.seq,
+            }
+          : null;
+      continue;
+    }
+
+    const received = item.value;
+    const event = received.data;
+    if (current?.turnId === event.turnId) {
+      if (!isNewerWireEvent(event, current)) continue;
+    } else if (event.type !== "turn_started") {
+      continue;
+    }
+    const publicTurnId =
+      current?.turnId === event.turnId
+        ? current.publicTurnId
+        : (await aiConversations.getTurn({ conversationId: input.conversation.id, turnId: event.turnId }))?.shortId;
+    if (!publicTurnId) continue;
+    current = { turnId: event.turnId, publicTurnId, attempt: event.attempt, seq: event.seq };
+    if (event.type === "turn_finished") {
+      const messages = await aiConversations
+        .listTurnMessages({ conversationId: event.conversationId, loopId: event.turnId })
+        .catch(() => []);
+      yield {
+        ...event,
+        conversationId: input.conversation.shortId,
+        turnId: publicTurnId,
+        messages: projectPublicAiStoredMessages(messages, input.conversation.shortId, new Map([[event.turnId, publicTurnId]])),
+      };
+      continue;
+    }
+
+    yield { ...event, conversationId: input.conversation.shortId, turnId: publicTurnId };
+  }
+}
+
+export const __aiStreamTest = { streamSnapshotThenTail };
+
+/** Conversation-scoped SSE adapter for the shared event feed. */
+export const createAiConversationStreamResponse = (input: {
+  conversation: AiConversation;
+  signal?: AbortSignal;
+  heartbeatMs?: number;
+  /**
+   * Re-checked on every heartbeat. A stream can outlive the grant that opened
+   * it by hours, and authorizing once at connect time means a withdrawn grant
+   * keeps delivering model output until the client happens to disconnect.
+   * Returning false closes the stream; the client's reconnect then meets the
+   * ordinary 403.
+   */
+  revalidate?: () => Promise<boolean>;
+}): Response => {
+  const heartbeatMs = input.heartbeatMs ?? DEFAULT_HEARTBEAT_MS;
+  const liveAbort = new AbortController();
+  let heartbeat: ReturnType<typeof setInterval> | undefined;
+  let closed = false;
+
+  const abortLive = () => liveAbort.abort();
+  if (input.signal?.aborted) abortLive();
+  else input.signal?.addEventListener("abort", abortLive, { once: true });
+
+  const stream = new ReadableStream<Uint8Array>({
+    async start(controller) {
+      const enqueue = (chunk: Uint8Array): boolean => {
+        if (closed) return false;
+        try {
+          controller.enqueue(chunk);
+          return true;
+        } catch {
+          close();
+          return false;
+        }
+      };
+      const close = () => {
+        if (closed) return;
+        closed = true;
+        if (heartbeat) clearInterval(heartbeat);
+        heartbeat = undefined;
+        input.signal?.removeEventListener("abort", abortLive);
+        abortLive();
+        try {
+          controller.close();
+        } catch {
+          // The client may already have cancelled the stream.
+        }
+      };
+
+      if (heartbeatMs > 0) {
+        heartbeat = setInterval(() => {
+          if (!enqueue(encodeSseHeartbeat())) return;
+          if (!input.revalidate) return;
+          // Fail closed: a revalidation that throws is not a pass.
+          void input
+            .revalidate()
+            .catch((error) => {
+              log.warn("AI conversation stream revalidation failed", {
+                conversationId: input.conversation.id,
+                error: error instanceof Error ? error.message : "revalidation failed",
+              });
+              return false;
+            })
+            .then((allowed) => {
+              if (!allowed) close();
+            });
+        }, heartbeatMs);
+      }
+
+      try {
+        for await (const event of streamAiConversationEvents({ conversation: input.conversation, signal: liveAbort.signal })) {
+          if (!enqueue(encodeSseEvent(event))) return;
+        }
+      } catch (error) {
+        if (!liveAbort.signal.aborted) {
+          log.warn("AI conversation stream failed", {
+            conversationId: input.conversation.id,
+            error: error instanceof Error ? error.message : "AI conversation stream failed",
+          });
+        }
+      } finally {
+        close();
+      }
+    },
+    cancel() {
+      closed = true;
+      if (heartbeat) clearInterval(heartbeat);
+      abortLive();
+    },
+  });
+
+  return new Response(stream, { headers: sseHeaders });
+};

package/src/ai/structured-runs.ts

@@ -0,0 +1,66 @@
+import type { Usage } from "@k2b/nessi";
+import { sql } from "bun";
+import { logger } from "../services/logging";
+
+export type AiUsageAttribution = {
+  userId?: string;
+  conversationId?: string;
+  turnId?: string;
+  workflowRunId?: string;
+};
+
+export type AiStructuredRunRecord = {
+  attribution?: AiUsageAttribution;
+  traceId?: string;
+  task: string;
+  appId?: string;
+  modelProfileId?: string;
+  providerModel?: string;
+  status: "ok" | "failed";
+  durationMs: number;
+  usage?: Usage;
+  mode?: string;
+  repaired?: boolean;
+  attempts?: number;
+  errorCode?: string;
+  error?: string;
+};
+
+const usageNumber = (usage: Usage | undefined, key: "input" | "output" | "total" | "creditsUsed") => {
+  const value = (usage as Partial<Record<typeof key, unknown>> | undefined)?.[key];
+  return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : null;
+};
+
+/** Metadata-only durable accounting. Prompt, input, and output content never enter this ledger. */
+export const recordAiStructuredRun = async (record: AiStructuredRunRecord): Promise<void> => {
+  await sql`
+    INSERT INTO ai.structured_runs (
+      task, app_id, model_profile_id, provider_model, status, duration_ms,
+      input_tokens, output_tokens, total_tokens, credits_used,
+      mode, repaired, attempts, error_code, error, user_id, conversation_id, turn_id, workflow_run_id, trace_id
+    ) VALUES (
+      ${record.task}, ${record.appId ?? null}, ${record.modelProfileId ?? null}, ${record.providerModel ?? null},
+      ${record.status}, ${Math.max(0, Math.round(record.durationMs))},
+      ${usageNumber(record.usage, "input")}, ${usageNumber(record.usage, "output")},
+      ${usageNumber(record.usage, "total")}, ${usageNumber(record.usage, "creditsUsed")},
+      ${record.mode ?? null}, ${record.repaired ?? null}, ${record.attempts ?? null},
+      ${record.errorCode ?? null}, ${record.error?.slice(0, 2_000) ?? null},
+      COALESCE(${record.attribution?.userId ?? null}::uuid,
+        (SELECT created_by_user_id FROM ai.conversations WHERE id = ${record.attribution?.conversationId ?? null}::uuid)),
+      ${record.attribution?.conversationId ?? null}::uuid, ${record.attribution?.turnId ?? null}::uuid,
+      ${record.attribution?.workflowRunId ?? null}::uuid, ${record.traceId ?? null}
+    )
+  `;
+};
+
+/** Accounting failure must not turn a completed inference into a retry. */
+export const safelyRecordStructuredRun = async (record: AiStructuredRunRecord): Promise<void> => {
+  try {
+    await recordAiStructuredRun(record);
+  } catch (error) {
+    logger("ai.structured").error("Could not persist AI structured-run usage", {
+      task: record.task,
+      error: error instanceof Error ? error.message : String(error),
+    });
+  }
+};

package/src/ai/structured.ts

@@ -0,0 +1,186 @@
+import type { Input, LoopAggregate, StructuredMeta, Usage } from "@k2b/nessi";
+import { nessi, StructuredOutputError } from "@k2b/nessi";
+import type { z } from "zod";
+import { coreSettings } from "../services";
+import type { TraceContext } from "../services/logging";
+import { trace } from "../services/logging";
+import { resolveAiModel } from "./settings";
+import { type AiUsageAttribution, safelyRecordStructuredRun } from "./structured-runs";
+import type { AiResolvedModel } from "./types";
+
+export const AI_BACKGROUND_MODEL_SETTING_KEY = "ai.background_model_id";
+export const AI_WORKFLOW_MODEL_SETTING_KEY = "ai.workflow_model_id";
+
+export const selectAiWorkflowModelId = (input: {
+  requestedModelId?: string;
+  workflowModelId?: string;
+  backgroundModelId?: string;
+}): string | undefined => input.requestedModelId?.trim() || input.workflowModelId?.trim() || input.backgroundModelId?.trim() || undefined;
+
+/**
+ * Resolve the model for background inference: explicit request →
+ * `ai.background_model_id` setting → platform default. Throws when AI is
+ * disabled or the model is unavailable — callers skip their work then.
+ */
+export const resolveAiBackgroundModel = async (requestedModelId?: string): Promise<AiResolvedModel> => {
+  const backgroundModelId = String((await coreSettings.get<string>(AI_BACKGROUND_MODEL_SETTING_KEY)) ?? "").trim();
+  return resolveAiModel({ kind: "selectable" }, requestedModelId?.trim() || backgroundModelId || undefined);
+};
+
+/** Resolve a workflow model: action override → workflow setting → background setting → platform default. */
+export const resolveAiWorkflowModel = async (requestedModelId?: string): Promise<AiResolvedModel> => {
+  const [workflowModelId, backgroundModelId] = await Promise.all([
+    coreSettings.get<string>(AI_WORKFLOW_MODEL_SETTING_KEY),
+    coreSettings.get<string>(AI_BACKGROUND_MODEL_SETTING_KEY),
+  ]);
+  return resolveAiModel(
+    { kind: "selectable" },
+    selectAiWorkflowModelId({
+      requestedModelId,
+      workflowModelId: String(workflowModelId ?? ""),
+      backgroundModelId: String(backgroundModelId ?? ""),
+    }),
+  );
+};
+
+export type RunAiStructuredInput<TOutput extends z.ZodType> = {
+  /** Short machine name for tracing, e.g. "chat-enrich". */
+  task: string;
+  input: Input;
+  output: TOutput;
+  outputName?: string;
+  systemPrompt?: string;
+  requestedModelId?: string;
+  temperature?: number;
+  maxOutputTokens?: number;
+  signal?: AbortSignal;
+  /** Parent trace span when the caller already runs inside one (e.g. a sync job). */
+  traceParent?: TraceContext;
+  appId?: string;
+  /** Metadata-only attribution, supplied by the authorized caller. */
+  attribution?: AiUsageAttribution;
+  /** Model resolution seam — tests inject a fake so they never touch shared settings. */
+  resolveModel?: (requestedModelId?: string) => Promise<AiResolvedModel>;
+};
+
+export type RunAiStructuredResult<TOutput extends z.ZodType> = {
+  output: z.infer<TOutput>;
+  modelProfileId: string;
+  usage?: Usage;
+  structuredMeta: StructuredMeta;
+};
+
+/**
+ * One schema-valid background inference via nessi.structured, wrapped in a
+ * trace span (events: model.resolved, llm.completed — metadata only, never
+ * prompt or output content).
+ */
+export const runAiStructured = async <TOutput extends z.ZodType>(
+  input: RunAiStructuredInput<TOutput>,
+): Promise<RunAiStructuredResult<TOutput>> => {
+  return trace.withSpan(
+    {
+      name: `ai.structured.${input.task}`,
+      source: `ai:structured:${input.task}`,
+      appId: input.appId,
+      category: "ai",
+      parent: input.traceParent,
+    },
+    async (span) => {
+      const startedAt = Date.now();
+      let resolved: AiResolvedModel | undefined;
+      try {
+        resolved = await (input.resolveModel ?? resolveAiBackgroundModel)(input.requestedModelId);
+        await trace.record({
+          context: span,
+          event: "model.resolved",
+          attributes: { model: resolved.profile.id, providerModel: resolved.profile.model, provider: resolved.profile.provider },
+        });
+
+        const result = await nessi.structured({
+          agentId: "cloud-bg",
+          provider: resolved.provider,
+          systemPrompt: input.systemPrompt,
+          input: input.input,
+          output: input.output,
+          outputName: input.outputName,
+          temperature: input.temperature ?? 0,
+          maxOutputTokens: input.maxOutputTokens ?? resolved.profile.maxOutputTokens,
+          disableReasoning: true,
+          signal: input.signal,
+        });
+        const durationMs = Date.now() - startedAt;
+
+        await trace.record({
+          context: span,
+          event: "llm.completed",
+          attributes: {
+            model: resolved.profile.id,
+            durationMs,
+            mode: result.structuredMeta.mode,
+            repaired: result.structuredMeta.repaired,
+            attempts: result.structuredMeta.attempts,
+            inputTokens: result.usage?.input,
+            outputTokens: result.usage?.output,
+          },
+        });
+        await safelyRecordStructuredRun({
+          task: input.task,
+          attribution: input.attribution,
+          traceId: span.traceId,
+          appId: input.appId,
+          modelProfileId: resolved.profile.id,
+          providerModel: resolved.profile.model,
+          status: "ok",
+          durationMs,
+          usage: result.usage,
+          mode: result.structuredMeta.mode,
+          repaired: result.structuredMeta.repaired,
+          attempts: result.structuredMeta.attempts,
+        });
+
+        return {
+          output: result.output,
+          modelProfileId: resolved.profile.id,
+          usage: result.usage,
+          structuredMeta: result.structuredMeta,
+        };
+      } catch (error) {
+        const details =
+          error instanceof StructuredOutputError
+            ? (error.details as { attempts?: number; aggregate?: LoopAggregate } | undefined)
+            : undefined;
+        await safelyRecordStructuredRun({
+          task: input.task,
+          attribution: input.attribution,
+          traceId: span.traceId,
+          appId: input.appId,
+          modelProfileId: resolved?.profile.id,
+          providerModel: resolved?.profile.model,
+          status: "failed",
+          durationMs: Date.now() - startedAt,
+          usage: details?.aggregate?.usage,
+          attempts: details?.attempts,
+          errorCode: error instanceof StructuredOutputError ? error.code : error instanceof Error ? error.name : "unknown",
+          error: error instanceof Error ? error.message : String(error),
+        });
+        throw error;
+      }
+    },
+    {
+      summarize: (result) => ({ model: result.modelProfileId, mode: result.structuredMeta.mode, repaired: result.structuredMeta.repaired }),
+      onError: (error) => (error instanceof StructuredOutputError ? structuredFailureSummary(error) : undefined),
+    },
+  );
+};
+
+/** Metadata-only failure diagnostics: error code, attempts, and per-attempt stop reasons (catches max_tokens truncation). */
+const structuredFailureSummary = (error: StructuredOutputError): Record<string, unknown> => {
+  const details = error.details as { attempts?: number; aggregate?: LoopAggregate } | undefined;
+  return {
+    code: error.code,
+    attempts: details?.attempts,
+    stopReasons: details?.aggregate?.turns?.map((turn) => turn.stopReason ?? "unknown").join(","),
+    outputTokens: details?.aggregate?.usage?.output,
+  };
+};