@arcote.tech/arc-chat 0.5.2 → 0.5.6

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

@@ -1,6 +1,13 @@
 /// <reference path="../arc.d.ts" />
 import { listener, type ArcContextElement, type ArcFunction } from "@arcote.tech/arc";
-import type { ArcToolAny, LLMProvider, Message } from "@arcote.tech/arc-ai";
+import type {
+  ArcToolAny,
+  AssistantContentBlock,
+  Conversation,
+  ConversationTurn,
+  LLMProvider,
+  ToolCall,
+} from "@arcote.tech/arc-ai";
 import { broadcast, endStream } from "../streaming/stream-registry";
 
 // ─── Config ─────────────────────────────────────────────────────
@@ -18,54 +25,160 @@ export interface AiGenerationListenerConfig {
   toolChoice?: "auto" | "required" | { type: "function"; name: string };
 }
 
-// ─── Utilities ──────────────────────────────────────────────────
+// ─── History reconstruction ─────────────────────────────────────
 
-function buildLlmMessages(
-  history: any[],
-  systemPrompt?: string,
+/**
+ * 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,
-): Message[] {
-  const messages: Message[] = [];
-  if (systemPrompt) {
-    messages.push({ role: "system", content: systemPrompt });
-  }
-  for (const msg of history) {
+): ConversationTurn[] {
+  const turns: ConversationTurn[] = [];
+
+  for (const msg of messages) {
     if (msg._id === skipMessageId) continue;
-    if (msg.role === "user" || msg.role === "assistant") {
-      if (msg.isGenerating && !msg.content) continue;
-      messages.push({ role: msg.role as Message["role"], content: msg.content });
-    } else if (msg.role === "tool_result") {
-      messages.push({
-        role: "tool",
-        content: msg.content,
+
+    // 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 messages;
+
+  return turns;
+}
+
+// ─── Instructions ───────────────────────────────────────────────
+
+/**
+ * Result from an instruction handler. Can be a plain string (prompt only)
+ * or an object with prompt + optional tool filtering.
+ */
+export interface InstructionResult {
+  /** System prompt text sent to the LLM. */
+  prompt: string;
+  /**
+   * If provided, only tools whose names appear in this array will be
+   * sent to the LLM for this generation call. Tools not listed are
+   * hidden — the LLM cannot call them. Omit to send all registered tools.
+   */
+  enabledTools?: string[];
 }
 
-async function buildSystemPrompt(
+/**
+ * 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.
+ *
+ * Returns `{ prompt, enabledTools? }`. When the handler returns a plain string,
+ * it's wrapped as `{ prompt: str }` with no tool filtering.
+ */
+async function buildInstructions(
   instruction: ArcFunction<any> | undefined,
   ctx: any,
-): Promise<string> {
-  if (!instruction?.handler) return "";
+  scopeId: string,
+): Promise<InstructionResult> {
+  if (!instruction?.handler) return { prompt: "" };
   const instructionCtx = {
     query: (element: ArcContextElement<any>) => ctx.query(element),
     mutate: (element: ArcContextElement<any>) => ctx.mutate(element),
+    identifyBy: scopeId,
+    scopeId,
   };
-  return (instruction.handler as Function)(instructionCtx);
+  const result = await (instruction.handler as Function)(instructionCtx);
+  if (typeof result === "string") return { prompt: result };
+  if (result && typeof result === "object" && "prompt" in result) return result as InstructionResult;
+  return { prompt: "" };
 }
 
-// ─── AI generation loop ─────────────────────────────────────────
+// ─── 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 };
+}
 
-async function runGenerationLoop(config: {
+/** 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;
-  initialMessages: Message[];
+  /** 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>;
@@ -73,27 +186,60 @@ async function runGenerationLoop(config: {
   scopeId: string;
   sessionId: string;
   maxExecutionCount: number;
-  initialPreviousResponseId?: string;
   toolChoice?: "auto" | "required" | { type: "function"; name: string };
-}) {
+  instruction?: ArcFunction<any>;
+}
+
+async function runGenerationLoop(config: RunLoopConfig) {
   const {
-    ctx, messageElement, provider, model, toolDefs,
-    serverToolsMap, interactiveToolNames,
-    generationMessageId, scopeId, sessionId, maxExecutionCount, toolChoice,
+    ctx,
+    messageElement,
+    provider,
+    model,
+    toolDefs,
+    serverToolsMap,
+    interactiveToolNames,
+    generationMessageId,
+    scopeId,
+    sessionId,
+    maxExecutionCount,
+    toolChoice,
+    instruction,
   } = config;
 
+  let history = config.history;
+  let newTurnsStartIdx = config.initialNewTurnsStartIdx;
   let executionCount = 0;
-  let fullContent = "";
-  let previousResponseId = config.initialPreviousResponseId;
-  let currentMessages = config.initialMessages;
 
   try {
     while (executionCount <= maxExecutionCount) {
+      // Always re-render instructions — picks up state mutated by tool calls
+      // in the previous iteration (e.g. updateIdentity).
+      const instructionResult = await buildInstructions(instruction, ctx, scopeId);
+
+      // Filter tools if instruction handler specified enabledTools
+      const effectiveToolDefs = instructionResult.enabledTools
+        ? toolDefs?.filter((td) => instructionResult.enabledTools!.includes(td.name))
+        : toolDefs;
+
+      const lastResponseId = findLastResponseId(history);
+      const conversation = makeConversation(
+        provider,
+        history,
+        lastResponseId,
+        newTurnsStartIdx,
+      );
+
       const result = await provider.streamComplete(
-        { model, messages: currentMessages, tools: toolDefs, previousResponseId, toolChoice },
+        {
+          model,
+          instructions: instructionResult.prompt,
+          conversation,
+          tools: effectiveToolDefs,
+          toolChoice,
+        },
         (chunk) => {
           if (chunk.type === "content_delta" && chunk.content) {
-            fullContent += chunk.content;
             broadcast(sessionId, {
               type: "content_delta",
               sessionId,
@@ -109,19 +255,40 @@ async function runGenerationLoop(config: {
         },
       );
 
-      if (result.content) fullContent = result.content;
-      previousResponseId = result.responseId;
+      // 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,
+        });
+      }
 
-      // No tool calls — done
-      if (result.finishReason !== "tool_call" || result.toolCalls.length === 0) {
-        if (fullContent) {
-          await ctx.mutate(messageElement).saveAssistantMessage({
-            scopeId, sessionId, content: fullContent, model,
-            previousResponseId,
-          });
-        }
+      // 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,
+          generationMessageId,
+          sessionId,
           usage: JSON.stringify(result.usage),
         });
         broadcast(sessionId, {
@@ -135,29 +302,14 @@ async function runGenerationLoop(config: {
         return;
       }
 
-      // Save intermediate text
-      if (fullContent) {
-        await ctx.mutate(messageElement).saveAssistantMessage({
-          scopeId, sessionId, content: fullContent, model,
-          previousResponseId,
-        });
-        fullContent = "";
-      }
-
-      // Separate server vs interactive
-      const serverCalls = result.toolCalls.filter((tc) => serverToolsMap.has(tc.name));
-      const interactiveCalls = result.toolCalls.filter((tc) => interactiveToolNames.has(tc.name));
+      const serverCalls = toolCalls.filter((tc) => serverToolsMap.has(tc.name));
+      const interactiveCalls = toolCalls.filter((tc) =>
+        interactiveToolNames.has(tc.name),
+      );
 
-      // Execute server tools — collect ONLY new results
-      const newToolResults: Message[] = [];
+      // Execute server tools — append each result to history as a separate turn
+      const newToolResults: ConversationTurn[] = [];
       for (const tc of serverCalls) {
-        await ctx.mutate(messageElement).saveToolCall({
-          scopeId, sessionId,
-          toolName: tc.name, toolCallId: tc.id,
-          content: JSON.stringify(tc.arguments),
-          previousResponseId,
-        });
-
         broadcast(sessionId, {
           type: "server_tool_start",
           sessionId,
@@ -171,7 +323,11 @@ async function runGenerationLoop(config: {
 
         if (tool) {
           try {
-            resultContent = await tool.executeWithContext(tc.arguments, ctx, scopeId);
+            resultContent = await tool.executeWithContext(
+              tc.arguments,
+              ctx,
+              scopeId,
+            );
           } catch (err) {
             resultContent = `Tool error: ${err instanceof Error ? err.message : String(err)}`;
             isError = true;
@@ -182,49 +338,55 @@ async function runGenerationLoop(config: {
         }
 
         await ctx.mutate(messageElement).saveToolResult({
-          scopeId, sessionId,
-          toolName: tc.name, toolCallId: tc.id,
-          content: resultContent, isError,
+          scopeId,
+          sessionId,
+          toolName: tc.name,
+          toolCallId: tc.id,
+          content: resultContent,
+          isError,
         });
 
         broadcast(sessionId, {
           type: "server_tool_result",
           sessionId,
           toolCall: tc,
-          toolResult: { toolCallId: tc.id, name: tc.name, content: resultContent, isError },
+          toolResult: {
+            toolCallId: tc.id,
+            name: tc.name,
+            content: resultContent,
+            isError,
+          },
           executionCount,
         });
 
         newToolResults.push({
-          role: "tool", content: resultContent,
-          toolCallId: tc.id, name: tc.name,
+          role: "tool_result",
+          toolCallId: tc.id,
+          name: tc.name,
+          content: resultContent,
+          isError,
         });
       }
 
-      // Interactive tools — save and STOP
+      // 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) {
-        for (const tc of interactiveCalls) {
-          await ctx.mutate(messageElement).saveToolCall({
-            scopeId, sessionId,
-            toolName: tc.name, toolCallId: tc.id,
-            content: JSON.stringify(tc.arguments),
-            previousResponseId,
-          });
-        }
         broadcast(sessionId, {
           type: "interactive_tool_request",
           sessionId,
           toolCalls: interactiveCalls,
           executionCount,
         });
-        // Don't endStream — client stays connected for possible updates
-        // Don't completeGeneration — Listener B will resume
         return;
       }
 
-      // Next iteration: ONLY new tool results (provider has rest via previousResponseId)
-      currentMessages = newToolResults;
-      fullContent = "";
+      // 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) {
@@ -236,7 +398,8 @@ async function runGenerationLoop(config: {
     });
     try {
       await ctx.mutate(messageElement).completeGeneration({
-        generationMessageId, sessionId,
+        generationMessageId,
+        sessionId,
       });
     } catch {}
     endStream(sessionId);
@@ -247,18 +410,25 @@ async function runGenerationLoop(config: {
 
 export function createAiGenerationListener(config: AiGenerationListenerConfig) {
   const {
-    name, messageElement, resolveProvider, instruction,
-    serverTools, interactiveTools,
-    allQueryElements, allMutationElements, maxExecutionCount,
+    name,
+    messageElement,
+    resolveProvider,
+    instruction,
+    serverTools,
+    interactiveTools,
+    allQueryElements,
+    allMutationElements,
+    maxExecutionCount,
   } = config;
 
   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 toolDefs =
+    allToolsForLLM.length > 0
+      ? allToolsForLLM.map((t) => t.toJsonSchema())
+      : undefined;
 
   return listener(`${name}AiGeneration`)
     .listenTo([messageSentEvent])
@@ -266,28 +436,56 @@ export function createAiGenerationListener(config: AiGenerationListenerConfig) {
     .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-nano";
+      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 systemPrompt = await buildSystemPrompt(instruction, ctx);
-      const history = await ctx.query(messageElement).getByScope({ scopeId });
-      const messages = buildLlmMessages(history, systemPrompt, event.payload.messageId);
-      messages.push({ role: "user", content: userContent });
-
-      const generationResult = await ctx.mutate(messageElement).saveAssistantMessage({
-        scopeId, sessionId, content: "", model, isGenerating: true,
-      });
+      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,
+        });
 
       await runGenerationLoop({
-        ctx, messageElement, provider, model,
-        initialMessages: messages,
-        toolDefs, serverToolsMap, interactiveToolNames,
+        ctx,
+        messageElement,
+        provider,
+        model,
+        history,
+        initialNewTurnsStartIdx: newTurnsStartIdx,
+        toolDefs,
+        serverToolsMap,
+        interactiveToolNames,
         generationMessageId: generationResult.messageId,
-        scopeId, sessionId, maxExecutionCount,
+        scopeId,
+        sessionId,
+        maxExecutionCount,
         toolChoice: config.toolChoice,
+        instruction,
       });
     });
 }
@@ -296,18 +494,25 @@ export function createAiGenerationListener(config: AiGenerationListenerConfig) {
 
 export function createAiResumeListener(config: AiGenerationListenerConfig) {
   const {
-    name, messageElement, resolveProvider, instruction,
-    serverTools, interactiveTools,
-    allQueryElements, allMutationElements, maxExecutionCount,
+    name,
+    messageElement,
+    resolveProvider,
+    instruction,
+    serverTools,
+    interactiveTools,
+    allQueryElements,
+    allMutationElements,
+    maxExecutionCount,
   } = config;
 
   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;
+  const toolDefs =
+    allToolsForLLM.length > 0
+      ? allToolsForLLM.map((t) => t.toJsonSchema())
+      : undefined;
 
   return listener(`${name}AiResume`)
     .listenTo([userRespondedEvent])
@@ -315,52 +520,73 @@ export function createAiResumeListener(config: AiGenerationListenerConfig) {
     .query([messageElement, ...allQueryElements])
     .mutate([messageElement, ...allMutationElements])
     .handle(async (ctx, event) => {
-      const { sessionId, scopeId, toolCallId, toolName, content: toolResult } = event.payload;
-
-      const history = await ctx.query(messageElement).getByScope({ scopeId });
+      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;
+        }
+      }
 
-      // Find previousResponseId from the tool_call this responds to
-      const matchingToolCall = [...history]
+      // Determine the model from the most recent assistant row in DB
+      const lastAssistantRow = [...dbMessages]
         .reverse()
-        .find((msg: any) => msg.role === "tool_call" && msg.toolCallId === toolCallId);
-      const prevResponseId = matchingToolCall?.previousResponseId;
-
-      const model = matchingToolCall?.model
-        ?? history.find((m: any) => m.model)?.model
-        ?? "gpt-5.4-nano";
+        .find((m: any) => m.role === "assistant" && m.model);
+      const model = lastAssistantRow?.model ?? "gpt-5.4-mini";
 
       const provider = resolveProvider(model, scopeId);
       if (!provider) return;
 
-      // Build initial messages for this iteration
-      let initialMessages: Message[];
-      if (prevResponseId) {
-        // Provider has context — send only the new tool result
-        initialMessages = [{
-          role: "tool", content: toolResult,
-          toolCallId, name: toolName,
-        }];
-      } else {
-        // Fallback: full history
-        const systemPrompt = await buildSystemPrompt(instruction, ctx);
-        initialMessages = buildLlmMessages(history, systemPrompt);
-        if (!history.some((m: any) => m.toolCallId === toolCallId && m.role === "tool_result")) {
-          initialMessages.push({ role: "tool", content: toolResult, toolCallId, name: toolName });
-        }
-      }
+      // Placeholder assistant message for "AI is typing"
+      const generationResult = await ctx
+        .mutate(messageElement)
+        .saveAssistantMessage({
+          scopeId,
+          sessionId,
+          blocks: "[]",
+          model,
+          isGenerating: true,
+        });
 
-      const generationResult = await ctx.mutate(messageElement).saveAssistantMessage({
-        scopeId, sessionId, content: "", model, isGenerating: true,
-      });
+      void toolName;
+      void toolResult;
 
       await runGenerationLoop({
-        ctx, messageElement, provider, model,
-        initialMessages,
-        toolDefs, serverToolsMap, interactiveToolNames,
+        ctx,
+        messageElement,
+        provider,
+        model,
+        history,
+        initialNewTurnsStartIdx: newTurnsStartIdx,
+        toolDefs,
+        serverToolsMap,
+        interactiveToolNames,
         generationMessageId: generationResult.messageId,
-        scopeId, sessionId, maxExecutionCount,
-        initialPreviousResponseId: prevResponseId,
+        scopeId,
+        sessionId,
+        maxExecutionCount,
         toolChoice: config.toolChoice,
+        instruction,
       });
     });
 }