@arcote.tech/arc-chat 0.5.1 → 0.5.2

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

@@ -1,293 +1,366 @@
 /// <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, LLMProvider, Message } 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 ────────────────────────────────────────────────────
+// ─── Utilities ──────────────────────────────────────────────────
+
+function buildLlmMessages(
+  history: any[],
+  systemPrompt?: string,
+  skipMessageId?: string,
+): Message[] {
+  const messages: Message[] = [];
+  if (systemPrompt) {
+    messages.push({ role: "system", content: systemPrompt });
+  }
+  for (const msg of history) {
+    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,
+        toolCallId: msg.toolCallId,
+        name: msg.toolName,
+      });
+    }
+  }
+  return messages;
+}
 
-export function createAiGenerationListener(config: AiGenerationListenerConfig) {
-  const {
-    name,
-    messageElement,
-    resolveProvider,
-    prepare,
-    tools: defaultTools,
-    clientTools: defaultClientTools,
-    toolMutationElements,
-    maxExecutionCount: defaultMaxExecution,
-  } = config;
+async function buildSystemPrompt(
+  instruction: ArcFunction<any> | undefined,
+  ctx: any,
+): Promise<string> {
+  if (!instruction?.handler) return "";
+  const instructionCtx = {
+    query: (element: ArcContextElement<any>) => ctx.query(element),
+    mutate: (element: ArcContextElement<any>) => ctx.mutate(element),
+  };
+  return (instruction.handler as Function)(instructionCtx);
+}
 
-  const messageSentEvent = messageElement.getEvent("messageSent");
+// ─── AI generation loop ─────────────────────────────────────────
 
-  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;
-
-      // 1. Get or create stream session
-      let session = getStreamSession(sessionId);
-      if (!session) {
-        session = createStreamSession(sessionId);
-      }
+async function runGenerationLoop(config: {
+  ctx: any;
+  messageElement: any;
+  provider: LLMProvider;
+  model: string;
+  initialMessages: Message[];
+  toolDefs: any[] | undefined;
+  serverToolsMap: Map<string, ArcToolAny>;
+  interactiveToolNames: Set<string>;
+  generationMessageId: string;
+  scopeId: string;
+  sessionId: string;
+  maxExecutionCount: number;
+  initialPreviousResponseId?: string;
+  toolChoice?: "auto" | "required" | { type: "function"; name: string };
+}) {
+  const {
+    ctx, messageElement, provider, model, toolDefs,
+    serverToolsMap, interactiveToolNames,
+    generationMessageId, scopeId, sessionId, maxExecutionCount, toolChoice,
+  } = config;
 
-      // 2. Resolve provider
-      const model = modelName ?? "gpt-4o";
-      const provider = resolveProvider(model);
-      if (!provider) {
-        session.push({
-          type: "error",
+  let executionCount = 0;
+  let fullContent = "";
+  let previousResponseId = config.initialPreviousResponseId;
+  let currentMessages = config.initialMessages;
+
+  try {
+    while (executionCount <= maxExecutionCount) {
+      const result = await provider.streamComplete(
+        { model, messages: currentMessages, tools: toolDefs, previousResponseId, toolChoice },
+        (chunk) => {
+          if (chunk.type === "content_delta" && chunk.content) {
+            fullContent += 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,
+            });
+          }
+        },
+      );
+
+      if (result.content) fullContent = result.content;
+      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,
+          });
+        }
+        await ctx.mutate(messageElement).completeGeneration({
+          generationMessageId, sessionId,
+          usage: JSON.stringify(result.usage),
+        });
+        broadcast(sessionId, {
+          type: "done",
           sessionId,
-          error: `Provider not found for model: ${model}`,
+          usage: result.usage,
+          finishReason: result.finishReason,
+          executionCount,
         });
-        session.close();
-        deleteStreamSession(sessionId);
+        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,
+      // Save intermediate text
+      if (fullContent) {
+        await ctx.mutate(messageElement).saveAssistantMessage({
+          scopeId, sessionId, content: fullContent, model,
+          previousResponseId,
         });
-        instructions = prepareResult.instructions;
-        if (prepareResult.tools) serverTools = prepareResult.tools;
-        if (prepareResult.clientTools) clientTools = prepareResult.clientTools;
+        fullContent = "";
       }
 
-      // Build server tools map
-      const serverToolsMap = new Map(serverTools.map((t) => [t.name, t]));
-      const serverToolNames = [...serverToolsMap.keys()];
+      // Separate server vs interactive
+      const serverCalls = result.toolCalls.filter((tc) => serverToolsMap.has(tc.name));
+      const interactiveCalls = result.toolCalls.filter((tc) => interactiveToolNames.has(tc.name));
+
+      // Execute server tools — collect ONLY new results
+      const newToolResults: Message[] = [];
+      for (const tc of serverCalls) {
+        await ctx.mutate(messageElement).saveToolCall({
+          scopeId, sessionId,
+          toolName: tc.name, toolCallId: tc.id,
+          content: JSON.stringify(tc.arguments),
+          previousResponseId,
+        });
 
-      // Build tool defs for LLM (server + client)
-      const allToolsForLLM = [...serverTools, ...clientTools];
-      const toolDefs = allToolsForLLM.length > 0
-        ? allToolsForLLM.map((t) => t.toJsonSchema())
-        : undefined;
+        broadcast(sessionId, {
+          type: "server_tool_start",
+          sessionId,
+          toolCall: tc,
+          executionCount,
+        });
 
-      // 4. Load conversation history
-      const history = await ctx.query(messageElement).getByScope({ scopeId });
+        const tool = serverToolsMap.get(tc.name);
+        let resultContent: string;
+        let isError = false;
 
-      // 5. Build messages array
-      const messages: Message[] = [];
+        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;
+        }
+
+        await ctx.mutate(messageElement).saveToolResult({
+          scopeId, sessionId,
+          toolName: tc.name, toolCallId: tc.id,
+          content: resultContent, isError,
+        });
 
-      if (instructions) {
-        messages.push({ role: "system", content: instructions });
+        broadcast(sessionId, {
+          type: "server_tool_result",
+          sessionId,
+          toolCall: tc,
+          toolResult: { toolCallId: tc.id, name: tc.name, content: resultContent, isError },
+          executionCount,
+        });
+
+        newToolResults.push({
+          role: "tool", content: resultContent,
+          toolCallId: tc.id, name: tc.name,
+        });
       }
 
-      for (const msg of history) {
-        if (msg.role === "user" && msg.content === userContent && msg._id === payload.messageId) {
-          continue;
+      // Interactive tools — save and STOP
+      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,
+          });
         }
-        messages.push({
-          role: msg.role as Message["role"],
-          content: msg.content,
+        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;
       }
 
-      messages.push({ role: "user", content: userContent });
+      // Next iteration: ONLY new tool results (provider has rest via previousResponseId)
+      currentMessages = newToolResults;
+      fullContent = "";
+      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);
+  }
+}
 
-      // 6. Build tool context for server tool execution
-      const toolCtx: ToolContext = {
-        mutate: (element) => ctx.mutate(element),
-        query: (element) => ctx.query(element),
-        identifyBy: scopeId,
-      };
-
-      // 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;
 
-            messages.push({
-              role: "tool",
-              content: resultContent,
-              toolCallId: tc.id,
-              name: tc.name,
-            });
-          }
+      const model = modelName ?? "gpt-5.4-nano";
+      const provider = resolveProvider(model, scopeId);
+      if (!provider) return;
 
-          // Request client tool execution
-          if (clientCalls.length > 0) {
-            session.push({
-              type: "client_tool_request",
-              sessionId,
-              toolCalls: clientCalls,
-              executionCount,
-            });
+      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 });
 
-            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;
-            }
-          }
+      const generationResult = await ctx.mutate(messageElement).saveAssistantMessage({
+        scopeId, sessionId, content: "", model, isGenerating: true,
+      });
+
+      await runGenerationLoop({
+        ctx, messageElement, provider, model,
+        initialMessages: messages,
+        toolDefs, serverToolsMap, interactiveToolNames,
+        generationMessageId: generationResult.messageId,
+        scopeId, sessionId, maxExecutionCount,
+        toolChoice: config.toolChoice,
+      });
+    });
+}
+
+// ─── Listener B: userResponded → AI resume ──────────────────────
+
+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 history = await ctx.query(messageElement).getByScope({ scopeId });
+
+      // Find previousResponseId from the tool_call this responds to
+      const matchingToolCall = [...history]
+        .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";
+
+      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 });
         }
-      } catch (err) {
-        session.push({
-          type: "error",
-          sessionId,
-          error: `AI generation error: ${err instanceof Error ? err.message : String(err)}`,
-          executionCount,
-        });
-      } finally {
-        session.close();
-        deleteStreamSession(sessionId);
       }
+
+      const generationResult = await ctx.mutate(messageElement).saveAssistantMessage({
+        scopeId, sessionId, content: "", model, isGenerating: true,
+      });
+
+      await runGenerationLoop({
+        ctx, messageElement, provider, model,
+        initialMessages,
+        toolDefs, serverToolsMap, interactiveToolNames,
+        generationMessageId: generationResult.messageId,
+        scopeId, sessionId, maxExecutionCount,
+        initialPreviousResponseId: prevResponseId,
+        toolChoice: config.toolChoice,
+      });
     });
 }