@arcote.tech/arc-chat 0.5.1 → 0.5.5

package/src/listeners/ai-generation-listener.ts

@@ -1,293 +1,569 @@
 /// <reference path="../arc.d.ts" />
-import { listener, type ArcContextElement } from "@arcote.tech/arc";
-import type { ArcToolAny, LLMProvider, Message, ToolContext } from "@arcote.tech/arc-ai";
-import type { PrepareContext, PrepareParams, PrepareResult } from "../chat-builder";
-import {
-  createStreamSession,
-  getStreamSession,
-  deleteStreamSession,
-} from "../streaming/stream-registry";
+import { listener, type ArcContextElement, type ArcFunction } from "@arcote.tech/arc";
+import type {
+  ArcToolAny,
+  AssistantContentBlock,
+  Conversation,
+  ConversationTurn,
+  LLMProvider,
+  ToolCall,
+} from "@arcote.tech/arc-ai";
+import { broadcast, endStream } from "../streaming/stream-registry";
 
 // ─── Config ─────────────────────────────────────────────────────
 
 export interface AiGenerationListenerConfig {
   name: string;
   messageElement: any;
-  resolveProvider: (model: string) => LLMProvider | undefined;
-  prepare?: (ctx: PrepareContext, params: PrepareParams) => Promise<PrepareResult>;
-  tools: ArcToolAny[];
-  clientTools: ArcToolAny[];
-  toolMutationElements: ArcContextElement<any>[];
+  resolveProvider: (model: string, scopeId?: string) => LLMProvider | undefined;
+  instruction?: ArcFunction<any>;
+  serverTools: ArcToolAny[];
+  interactiveTools: ArcToolAny[];
+  allQueryElements: ArcContextElement<any>[];
+  allMutationElements: ArcContextElement<any>[];
   maxExecutionCount: number;
+  toolChoice?: "auto" | "required" | { type: "function"; name: string };
 }
 
-// ─── Factory ────────────────────────────────────────────────────
+// ─── History reconstruction ─────────────────────────────────────
+
+/**
+ * Convert DB message rows into a `ConversationTurn[]` that the provider
+ * adapter can consume directly. Each `assistant` row already carries its full
+ * blocks (text + tool_call) as JSON. Tool results are separate rows.
+ */
+function buildHistory(
+  messages: any[],
+  skipMessageId?: string,
+): ConversationTurn[] {
+  const turns: ConversationTurn[] = [];
+
+  for (const msg of messages) {
+    if (msg._id === skipMessageId) continue;
+
+    // System messages are developer-injected priming prompts (stage welcome,
+    // startStage, etc.). The UI hides them, but the LLM needs to see them as
+    // conversational turns. Most providers reserve `role: "system"` for the
+    // top-level instruction prompt which we inject separately via
+    // `buildInstructions` — so for history purposes we map system rows onto
+    // `user` turns.
+    if (msg.role === "user" || msg.role === "system") {
+      if (typeof msg.content === "string" && msg.content.length > 0) {
+        turns.push({ role: "user", content: msg.content });
+      }
+      continue;
+    }
+
+    if (msg.role === "assistant") {
+      if (msg.isGenerating && !msg.blocks) continue;
+      let blocks: AssistantContentBlock[] = [];
+      if (typeof msg.blocks === "string" && msg.blocks.length > 0) {
+        try {
+          blocks = JSON.parse(msg.blocks);
+        } catch {
+          blocks = [];
+        }
+      }
+      turns.push({
+        role: "assistant",
+        blocks,
+        responseId: msg.previousResponseId,
+      });
+      continue;
+    }
+
+    if (msg.role === "tool_result") {
+      turns.push({
+        role: "tool_result",
+        toolCallId: msg.toolCallId,
+        name: msg.toolName,
+        content: msg.content ?? "",
+      });
+    }
+  }
+
+  return turns;
+}
 
-export function createAiGenerationListener(config: AiGenerationListenerConfig) {
+// ─── Instructions ───────────────────────────────────────────────
+
+/**
+ * Render the system prompt by invoking the consumer's `instruction()` handler
+ * with a thin wrapper around the listener's ctx. Always called fresh — never
+ * cached — so dynamic state (e.g. identity just updated by a tool call) shows
+ * up in the next provider call.
+ */
+async function buildInstructions(
+  instruction: ArcFunction<any> | undefined,
+  ctx: any,
+  scopeId: string,
+): Promise<string> {
+  if (!instruction?.handler) return "";
+  const instructionCtx = {
+    query: (element: ArcContextElement<any>) => ctx.query(element),
+    mutate: (element: ArcContextElement<any>) => ctx.mutate(element),
+    // The chat's `identifyBy` value (= scopeId of the conversation thread).
+    // Lets the consumer prompt scope its queries to the current entity.
+    identifyBy: scopeId,
+    scopeId,
+  };
+  const result = await (instruction.handler as Function)(instructionCtx);
+  return typeof result === "string" ? result : "";
+}
+
+// ─── Conversation mode selection ────────────────────────────────
+
+/**
+ * Decide whether to ask the provider for a continuation (delta) or send the
+ * full conversation. Continuation is only used when the provider supports it
+ * AND we have a known `responseId` to anchor the request.
+ *
+ * @param history    Full conversation history including any new turns appended
+ *                   for this call.
+ * @param newTurnsStartIdx  Index in `history` where "new" turns begin
+ *                          (everything before is "already known" by the model).
+ */
+function makeConversation(
+  provider: LLMProvider,
+  history: ConversationTurn[],
+  lastResponseId: string | undefined,
+  newTurnsStartIdx: number,
+): Conversation {
+  if (provider.supportsContinuation && lastResponseId) {
+    return {
+      mode: "continuation",
+      previousResponseId: lastResponseId,
+      newTurns: history.slice(newTurnsStartIdx),
+    };
+  }
+  return { mode: "full", turns: history };
+}
+
+/** Find the most recent `assistant` turn that carries a responseId. */
+function findLastResponseId(history: ConversationTurn[]): string | undefined {
+  for (let i = history.length - 1; i >= 0; i--) {
+    const t = history[i];
+    if (t.role === "assistant" && t.responseId) return t.responseId;
+  }
+  return undefined;
+}
+
+// ─── Generation loop ────────────────────────────────────────────
+
+interface RunLoopConfig {
+  ctx: any;
+  messageElement: any;
+  provider: LLMProvider;
+  model: string;
+  /** Full conversation history at the start of the loop (already includes the
+   *  user/tool_result turn that triggered this generation). */
+  history: ConversationTurn[];
+  /** Index in `history` where the FIRST iteration's "new turns" start. After
+   *  the first iteration this is recomputed from the latest responseId. */
+  initialNewTurnsStartIdx: number;
+  toolDefs: any[] | undefined;
+  serverToolsMap: Map<string, ArcToolAny>;
+  interactiveToolNames: Set<string>;
+  generationMessageId: string;
+  scopeId: string;
+  sessionId: string;
+  maxExecutionCount: number;
+  toolChoice?: "auto" | "required" | { type: "function"; name: string };
+  instruction?: ArcFunction<any>;
+}
+
+async function runGenerationLoop(config: RunLoopConfig) {
   const {
-    name,
+    ctx,
     messageElement,
-    resolveProvider,
-    prepare,
-    tools: defaultTools,
-    clientTools: defaultClientTools,
-    toolMutationElements,
-    maxExecutionCount: defaultMaxExecution,
+    provider,
+    model,
+    toolDefs,
+    serverToolsMap,
+    interactiveToolNames,
+    generationMessageId,
+    scopeId,
+    sessionId,
+    maxExecutionCount,
+    toolChoice,
+    instruction,
   } = config;
 
-  const messageSentEvent = messageElement.getEvent("messageSent");
-
-  return listener(`${name}AiGeneration`)
-    .listenTo([messageSentEvent])
-    .async()
-    .query([messageElement])
-    .mutate([messageElement, ...toolMutationElements])
-    .handle(async (ctx, event) => {
-      const payload = event.payload;
-      const {
-        sessionId,
-        scopeId,
-        content: userContent,
-        model: modelName,
-      } = payload;
+  let history = config.history;
+  let newTurnsStartIdx = config.initialNewTurnsStartIdx;
+  let executionCount = 0;
+
+  try {
+    while (executionCount <= maxExecutionCount) {
+      // Always re-render instructions — picks up state mutated by tool calls
+      // in the previous iteration (e.g. updateIdentity).
+      const instructions = await buildInstructions(instruction, ctx, scopeId);
+
+      const lastResponseId = findLastResponseId(history);
+      const conversation = makeConversation(
+        provider,
+        history,
+        lastResponseId,
+        newTurnsStartIdx,
+      );
+
+      const result = await provider.streamComplete(
+        {
+          model,
+          instructions,
+          conversation,
+          tools: toolDefs,
+          toolChoice,
+        },
+        (chunk) => {
+          if (chunk.type === "content_delta" && chunk.content) {
+            broadcast(sessionId, {
+              type: "content_delta",
+              sessionId,
+              content: chunk.content,
+            });
+          } else if (chunk.type === "usage_update") {
+            broadcast(sessionId, {
+              type: "usage_update",
+              sessionId,
+              usage: chunk.usage,
+            });
+          }
+        },
+      );
 
-      // 1. Get or create stream session
-      let session = getStreamSession(sessionId);
-      if (!session) {
-        session = createStreamSession(sessionId);
+      // Persist this turn's assistant blocks as a single message row.
+      if (result.blocks.length > 0) {
+        await ctx.mutate(messageElement).saveAssistantMessage({
+          scopeId,
+          sessionId,
+          blocks: JSON.stringify(result.blocks),
+          model,
+          previousResponseId: result.responseId,
+        });
       }
 
-      // 2. Resolve provider
-      const model = modelName ?? "gpt-4o";
-      const provider = resolveProvider(model);
-      if (!provider) {
-        session.push({
-          type: "error",
+      // Append to local history so the next iteration sees this turn.
+      const assistantTurn: ConversationTurn = {
+        role: "assistant",
+        blocks: result.blocks,
+        responseId: result.responseId,
+      };
+      history.push(assistantTurn);
+
+      // Pull out tool calls from the blocks (preserves order, but for
+      // execution we only care about the set).
+      const toolCalls: ToolCall[] = result.blocks
+        .filter((b): b is Extract<AssistantContentBlock, { type: "tool_call" }> =>
+          b.type === "tool_call",
+        )
+        .map((b) => ({ id: b.id, name: b.name, arguments: b.arguments }));
+
+      const hasToolCalls =
+        result.finishReason === "tool_call" && toolCalls.length > 0;
+
+      if (!hasToolCalls) {
+        await ctx.mutate(messageElement).completeGeneration({
+          generationMessageId,
           sessionId,
-          error: `Provider not found for model: ${model}`,
+          usage: JSON.stringify(result.usage),
         });
-        session.close();
-        deleteStreamSession(sessionId);
+        broadcast(sessionId, {
+          type: "done",
+          sessionId,
+          usage: result.usage,
+          finishReason: result.finishReason,
+          executionCount,
+        });
+        endStream(sessionId);
         return;
       }
 
-      // 3. Call prepare callback to get instructions, tools, clientTools
-      let instructions = "";
-      let serverTools = defaultTools;
-      let clientTools = defaultClientTools;
-
-      if (prepare) {
-        const prepareCtx: PrepareContext = {
-          query: (element) => ctx.query(element),
-          mutate: (element) => ctx.mutate(element),
-        };
-        const prepareResult = await prepare(prepareCtx, {
-          content: userContent,
-          identifyBy: scopeId,
-          model,
+      const serverCalls = toolCalls.filter((tc) => serverToolsMap.has(tc.name));
+      const interactiveCalls = toolCalls.filter((tc) =>
+        interactiveToolNames.has(tc.name),
+      );
+
+      // Execute server tools — append each result to history as a separate turn
+      const newToolResults: ConversationTurn[] = [];
+      for (const tc of serverCalls) {
+        broadcast(sessionId, {
+          type: "server_tool_start",
+          sessionId,
+          toolCall: tc,
+          executionCount,
         });
-        instructions = prepareResult.instructions;
-        if (prepareResult.tools) serverTools = prepareResult.tools;
-        if (prepareResult.clientTools) clientTools = prepareResult.clientTools;
-      }
 
-      // Build server tools map
-      const serverToolsMap = new Map(serverTools.map((t) => [t.name, t]));
-      const serverToolNames = [...serverToolsMap.keys()];
+        const tool = serverToolsMap.get(tc.name);
+        let resultContent: string;
+        let isError = false;
 
-      // Build tool defs for LLM (server + client)
-      const allToolsForLLM = [...serverTools, ...clientTools];
-      const toolDefs = allToolsForLLM.length > 0
-        ? allToolsForLLM.map((t) => t.toJsonSchema())
-        : undefined;
+        if (tool) {
+          try {
+            resultContent = await tool.executeWithContext(
+              tc.arguments,
+              ctx,
+              scopeId,
+            );
+          } catch (err) {
+            resultContent = `Tool error: ${err instanceof Error ? err.message : String(err)}`;
+            isError = true;
+          }
+        } else {
+          resultContent = `Tool "${tc.name}" not found`;
+          isError = true;
+        }
 
-      // 4. Load conversation history
-      const history = await ctx.query(messageElement).getByScope({ scopeId });
+        await ctx.mutate(messageElement).saveToolResult({
+          scopeId,
+          sessionId,
+          toolName: tc.name,
+          toolCallId: tc.id,
+          content: resultContent,
+          isError,
+        });
 
-      // 5. Build messages array
-      const messages: Message[] = [];
+        broadcast(sessionId, {
+          type: "server_tool_result",
+          sessionId,
+          toolCall: tc,
+          toolResult: {
+            toolCallId: tc.id,
+            name: tc.name,
+            content: resultContent,
+            isError,
+          },
+          executionCount,
+        });
 
-      if (instructions) {
-        messages.push({ role: "system", content: instructions });
+        newToolResults.push({
+          role: "tool_result",
+          toolCallId: tc.id,
+          name: tc.name,
+          content: resultContent,
+          isError,
+        });
       }
 
-      for (const msg of history) {
-        if (msg.role === "user" && msg.content === userContent && msg._id === payload.messageId) {
-          continue;
-        }
-        messages.push({
-          role: msg.role as Message["role"],
-          content: msg.content,
+      // Interactive tools — stop the loop, wait for userResponded.
+      // The assistant turn (with the interactive tool_call) is already
+      // persisted above. Listener B will resume.
+      if (interactiveCalls.length > 0) {
+        broadcast(sessionId, {
+          type: "interactive_tool_request",
+          sessionId,
+          toolCalls: interactiveCalls,
+          executionCount,
         });
+        return;
       }
 
-      messages.push({ role: "user", content: userContent });
-
-      // 6. Build tool context for server tool execution
-      const toolCtx: ToolContext = {
-        mutate: (element) => ctx.mutate(element),
-        query: (element) => ctx.query(element),
-        identifyBy: scopeId,
-      };
+      // Append tool results to history; mark them as the "new turns" for the
+      // next iteration's continuation request.
+      const assistantTurnIdx = history.length - 1;
+      history.push(...newToolResults);
+      newTurnsStartIdx = assistantTurnIdx + 1;
+
+      executionCount++;
+    }
+  } catch (err) {
+    broadcast(sessionId, {
+      type: "error",
+      sessionId,
+      error: `AI error: ${err instanceof Error ? err.message : String(err)}`,
+      executionCount,
+    });
+    try {
+      await ctx.mutate(messageElement).completeGeneration({
+        generationMessageId,
+        sessionId,
+      });
+    } catch {}
+    endStream(sessionId);
+  }
+}
 
-      // 7. AI generation loop
-      let executionCount = 0;
-      let fullContent = "";
-      let previousResponseId: string | undefined;
-
-      try {
-        while (executionCount <= defaultMaxExecution) {
-          const result = await provider.streamComplete(
-            { model, messages, tools: toolDefs, previousResponseId },
-            (chunk) => {
-              switch (chunk.type) {
-                case "content_delta":
-                  if (chunk.content) {
-                    fullContent += chunk.content;
-                    session!.push({
-                      type: "content_delta",
-                      sessionId,
-                      content: chunk.content,
-                    });
-                  }
-                  break;
-                case "usage_update":
-                  session!.push({
-                    type: "usage_update",
-                    sessionId,
-                    usage: chunk.usage,
-                  });
-                  break;
-              }
-            },
-          );
-
-          if (result.content) {
-            fullContent = result.content;
-          }
-          previousResponseId = result.responseId;
-
-          // No tool calls — generation complete
-          if (
-            result.finishReason !== "tool_call" ||
-            result.toolCalls.length === 0
-          ) {
-            await ctx.mutate(messageElement).completeAssistantMessage({
-              scopeId,
-              sessionId,
-              content: fullContent,
-              model,
-              usage: JSON.stringify(result.usage),
-            });
+// ─── Listener A: messageSent → AI generation ────────────────────
 
-            session.push({
-              type: "done",
-              sessionId,
-              usage: result.usage,
-              finishReason: result.finishReason,
-              executionCount,
-            });
-            break;
-          }
+export function createAiGenerationListener(config: AiGenerationListenerConfig) {
+  const {
+    name,
+    messageElement,
+    resolveProvider,
+    instruction,
+    serverTools,
+    interactiveTools,
+    allQueryElements,
+    allMutationElements,
+    maxExecutionCount,
+  } = config;
 
-          // Separate server vs client tool calls
-          const serverCalls = result.toolCalls.filter((tc) =>
-            serverToolNames.includes(tc.name),
-          );
-          const clientCalls = result.toolCalls.filter(
-            (tc) => !serverToolNames.includes(tc.name),
-          );
-
-          // Execute server tools with aggregate context
-          for (const tc of serverCalls) {
-            session.push({
-              type: "server_tool_start",
-              sessionId,
-              toolCall: tc,
-              executionCount,
-            });
+  const messageSentEvent = messageElement.getEvent("messageSent");
+  const serverToolsMap = new Map(serverTools.map((t) => [t.name, t]));
+  const interactiveToolNames = new Set(interactiveTools.map((t) => t.name));
+  const allToolsForLLM = [...serverTools, ...interactiveTools];
+  const toolDefs =
+    allToolsForLLM.length > 0
+      ? allToolsForLLM.map((t) => t.toJsonSchema())
+      : undefined;
 
-            const tool = serverToolsMap.get(tc.name);
-            let resultContent: string;
-            let isError = false;
-
-            if (tool) {
-              try {
-                resultContent = await tool.executeWithContext(tc.arguments, toolCtx);
-              } catch (err) {
-                resultContent = `Tool execution error: ${err instanceof Error ? err.message : String(err)}`;
-                isError = true;
-              }
-            } else {
-              resultContent = `Tool "${tc.name}" not found on server`;
-              isError = true;
-            }
-
-            session.push({
-              type: "server_tool_result",
-              sessionId,
-              toolResult: {
-                toolCallId: tc.id,
-                name: tc.name,
-                content: resultContent,
-                isError,
-              },
-              executionCount,
-            });
+  return listener(`${name}AiGeneration`)
+    .listenTo([messageSentEvent])
+    .async()
+    .query([messageElement, ...allQueryElements])
+    .mutate([messageElement, ...allMutationElements])
+    .handle(async (ctx, event) => {
+      const {
+        sessionId,
+        scopeId,
+        content: userContent,
+        model: modelName,
+      } = event.payload;
+
+      const model = modelName ?? "gpt-5.4-mini";
+      const provider = resolveProvider(model, scopeId);
+      if (!provider) return;
+
+      const dbMessages = await ctx
+        .query(messageElement)
+        .getByScope({ scopeId });
+
+      // Build the conversation history from DB. Skip the just-emitted user
+      // message — we'll append it explicitly so we know exactly where the
+      // "new turn" boundary is.
+      const history = buildHistory(dbMessages, event.payload.messageId);
+      const newTurnsStartIdx = history.length;
+      history.push({ role: "user", content: userContent });
+
+      // Placeholder assistant message so the UI can render "AI is typing".
+      // Empty blocks; the real one is saved by the loop after streaming.
+      const generationResult = await ctx
+        .mutate(messageElement)
+        .saveAssistantMessage({
+          scopeId,
+          sessionId,
+          blocks: "[]",
+          model,
+          isGenerating: true,
+        });
 
-            messages.push({
-              role: "tool",
-              content: resultContent,
-              toolCallId: tc.id,
-              name: tc.name,
-            });
-          }
+      await runGenerationLoop({
+        ctx,
+        messageElement,
+        provider,
+        model,
+        history,
+        initialNewTurnsStartIdx: newTurnsStartIdx,
+        toolDefs,
+        serverToolsMap,
+        interactiveToolNames,
+        generationMessageId: generationResult.messageId,
+        scopeId,
+        sessionId,
+        maxExecutionCount,
+        toolChoice: config.toolChoice,
+        instruction,
+      });
+    });
+}
 
-          // Request client tool execution
-          if (clientCalls.length > 0) {
-            session.push({
-              type: "client_tool_request",
-              sessionId,
-              toolCalls: clientCalls,
-              executionCount,
-            });
+// ─── Listener B: userResponded → AI resume ──────────────────────
 
-            try {
-              const clientResults =
-                await session.waitForClientToolResults();
-
-              for (const tr of clientResults) {
-                messages.push({
-                  role: "tool",
-                  content: tr.content,
-                  toolCallId: tr.toolCallId,
-                  name: tr.name,
-                });
-              }
-            } catch (err) {
-              session.push({
-                type: "error",
-                sessionId,
-                error: `Client tool execution failed: ${err instanceof Error ? err.message : String(err)}`,
-                executionCount,
-              });
-              break;
-            }
-          }
+export function createAiResumeListener(config: AiGenerationListenerConfig) {
+  const {
+    name,
+    messageElement,
+    resolveProvider,
+    instruction,
+    serverTools,
+    interactiveTools,
+    allQueryElements,
+    allMutationElements,
+    maxExecutionCount,
+  } = config;
 
-          fullContent = "";
-          executionCount++;
+  const userRespondedEvent = messageElement.getEvent("userResponded");
+  const serverToolsMap = new Map(serverTools.map((t) => [t.name, t]));
+  const interactiveToolNames = new Set(interactiveTools.map((t) => t.name));
+  const allToolsForLLM = [...serverTools, ...interactiveTools];
+  const toolDefs =
+    allToolsForLLM.length > 0
+      ? allToolsForLLM.map((t) => t.toJsonSchema())
+      : undefined;
+
+  return listener(`${name}AiResume`)
+    .listenTo([userRespondedEvent])
+    .async()
+    .query([messageElement, ...allQueryElements])
+    .mutate([messageElement, ...allMutationElements])
+    .handle(async (ctx, event) => {
+      const {
+        sessionId,
+        scopeId,
+        toolCallId,
+        toolName,
+        content: toolResult,
+      } = event.payload;
+
+      const dbMessages = await ctx
+        .query(messageElement)
+        .getByScope({ scopeId });
+
+      // Build full history. The userResponded event already created a
+      // tool_result row in DB, so it's part of the history naturally.
+      const history = buildHistory(dbMessages);
+
+      // Compute "new turns start" — index of the just-arrived tool_result.
+      // Anything before it is "already known" (assistant emitted the matching
+      // tool_call earlier and OpenAI has it server-side).
+      let newTurnsStartIdx = history.length;
+      for (let i = history.length - 1; i >= 0; i--) {
+        const t = history[i];
+        if (t.role === "tool_result" && t.toolCallId === toolCallId) {
+          newTurnsStartIdx = i;
+          break;
         }
-      } catch (err) {
-        session.push({
-          type: "error",
+      }
+
+      // Determine the model from the most recent assistant row in DB
+      const lastAssistantRow = [...dbMessages]
+        .reverse()
+        .find((m: any) => m.role === "assistant" && m.model);
+      const model = lastAssistantRow?.model ?? "gpt-5.4-mini";
+
+      const provider = resolveProvider(model, scopeId);
+      if (!provider) return;
+
+      // Placeholder assistant message for "AI is typing"
+      const generationResult = await ctx
+        .mutate(messageElement)
+        .saveAssistantMessage({
+          scopeId,
           sessionId,
-          error: `AI generation error: ${err instanceof Error ? err.message : String(err)}`,
-          executionCount,
+          blocks: "[]",
+          model,
+          isGenerating: true,
         });
-      } finally {
-        session.close();
-        deleteStreamSession(sessionId);
-      }
+
+      void toolName;
+      void toolResult;
+
+      await runGenerationLoop({
+        ctx,
+        messageElement,
+        provider,
+        model,
+        history,
+        initialNewTurnsStartIdx: newTurnsStartIdx,
+        toolDefs,
+        serverToolsMap,
+        interactiveToolNames,
+        generationMessageId: generationResult.messageId,
+        scopeId,
+        sessionId,
+        maxExecutionCount,
+        toolChoice: config.toolChoice,
+        instruction,
+      });
     });
 }