@xalia/agent 0.6.9 → 0.6.10

package/src/chat/server/openSession.ts

@@ -1,6 +1,7 @@
 import OpenAI from "openai";
 import { strict as assert } from "assert";
 import { Tool } from "@modelcontextprotocol/sdk/types.js";
+import { v4 as uuidv4 } from "uuid";
 
 import {
   AgentPreferences,
@@ -23,6 +24,7 @@ import { McpServerInfo } from "../../agent/mcpServerManager";
 import { IAgentEventHandler } from "../../agent/iAgentEventHandler";
 import { createSpecializedLLM } from "../../agent/agentUtils";
 import { IPlatform } from "../../agent/iplatform";
+import { ITitleGenerator, createTitleGenerator } from "./titleGenerator";
 
 import type {
   ServerToClient,
@@ -47,6 +49,8 @@ import type {
   ClientSessionMessageBase,
   ServerAgentPaused,
   ClientGetMcpResource,
+  ClientRaceModeResult,
+  ServerRaceModeGetResult,
 } from "../protocol/messages";
 import { AsyncQueue } from "../utils/asyncQueue";
 import { MultiAsyncQueue } from "../utils/multiAsyncQueue";
@@ -80,20 +84,25 @@ import {
   llmUserMessageToUserMessageData,
   MESSAGE_INDEX_START_VALUE,
 } from "./conversation";
+import { ChatSessionFileManager } from "./sessionFileManager";
 import {
-  ChatSessionFileManager,
   ISessionFileManager,
   ISessionFileManagerEventHandler,
-} from "./sessionFileManager";
+  SessionFileEntry,
+  ParsedContent,
+  getMimeTypeFromDataUrl,
+} from "../../agent/tools/fileManager";
+import { pdfToText } from "../../agent/tools/contentExtractors/pdfToText";
+import { summarizeDocument } from "../../agent/documentSummarizer";
 import { IUserConnectionManager } from "./connectionManager";
 import { getErrorString } from "./errorUtils";
 import { NODE_PLATFORM } from "../../tool/nodePlatform";
 import { DbMcpServerConfigs } from "../data/dbMcpServerConfigs";
-import { SessionFileEntry } from "../data/dbSessionFileModels";
 import { ApiKeyManager } from "../data/apiKeyManager";
 import { DbSessionMessages } from "../data/dbSessionMessages";
 import { ISessionMessageSender } from "./openSessionMessageSender";
 import { OpenAIRouterLLM } from "./openAIRouterLLM";
+import { ResponseAwaiter } from "../utils/responseAwaiter";
 
 /**
  * The model to use when the AgentProfile does not specify one.
@@ -108,6 +117,8 @@ export const DEFAULT_NUM_MESSGAES = 500;
 
 export const GUEST_TOKEN_PREFIX = "guest";
 
+export const RACE_MODE_TIMEOUT_MS = 120_000;
+
 const logger = getLogger();
 
 type QueuedClientMessage<
@@ -214,22 +225,29 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
     private readonly sessionUUID: string,
     private readonly sender: ISessionMessageSender<ServerToClient>,
     private readonly approvalManager: ToolApprovalManager,
-    private readonly contextTx: ChatContextTransaction
+    private readonly contextTx: ChatContextTransaction,
+    private readonly alt?: string
   ) {}
 
   onCompletion(result: AssistantMessageParam): void {
-    logger.debug(`[OpenSession.onCompletion] : ${JSON.stringify(result)}`);
+    logger.debug(
+      `[OpenSession.onCompletion] ${this.alt || ""} : ${JSON.stringify(result)}`
+    );
     // Nothing to broadcast. Caller will receive this via onAgentMessage.
     this.contextTx.processAgentResponse(result);
   }
 
   onImage(image: OpenAI.Chat.Completions.ChatCompletionContentPartImage): void {
-    logger.debug(`[OpenSession.onImage] : ${image.image_url.url}`);
+    logger.debug(
+      `[OpenSession.onImage] ${this.alt || ""} : ${image.image_url.url}`
+    );
     throw new Error("[OpenSession.onImage] unimplemented");
   }
 
   onToolCallResult(result: ToolMessageParam): void {
-    logger.debug(`[onToolCallResult] : ${JSON.stringify(result)}`);
+    logger.debug(
+      `[onToolCallResult] ${this.alt || ""} : ${JSON.stringify(result)}`
+    );
     const toolCallMessage = this.contextTx.processToolCallResult(result);
     this.sender.broadcast(toolCallMessage);
   }
@@ -245,6 +263,7 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
         type: "tool_call",
         tool_call: toolCall,
         session_id: this.sessionUUID,
+        ...(this.alt ? { alt: this.alt } : {}),
       });
       return true;
     }
@@ -255,7 +274,8 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
     const { approved, requested } = await this.approvalManager.getApproval(
       serverName,
       tool,
-      toolCall
+      toolCall,
+      this.alt
     );
 
     // For now, the frontend uses the tool_call data in the
@@ -267,6 +287,7 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
         type: "tool_call",
         tool_call: toolCall,
         session_id: this.sessionUUID,
+        ...(this.alt ? { alt: this.alt } : {}),
       });
     }
 
@@ -275,23 +296,28 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
 
   onAgentMessage(msg: string, end: boolean): Promise<void> {
     logger.debug(
-      `[OpenSession.onAgentMessage] msg: ${msg}, end: ${String(end)}`
+      `[OpenSession.onAgentMessage] ${this.alt || ""} msg: ${msg}, ` +
+        `end: ${String(end)}`
     );
 
     // Inform the contextManager and broadcast the ServerAgentMessageChunk
     const agentMsgChunk = this.contextTx.processAgentMessageChunk(msg, end);
+    if (this.alt) {
+      agentMsgChunk.alt = this.alt;
+    }
     this.sender.broadcast(agentMsgChunk);
     return Promise.resolve();
   }
 
   onReasoning(reasoning: string): Promise<void> {
     return new Promise<void>((r) => {
-      logger.debug(`[OpenSession.onReasoning]${reasoning}`);
+      logger.debug(`[OpenSession.onReasoning] ${this.alt || ""} ${reasoning}`);
       if (reasoning.length > 0) {
         this.sender.broadcast({
           type: "agent_reasoning_chunk",
           reasoning,
           session_id: this.sessionUUID,
+          ...(this.alt ? { alt: this.alt } : {}),
         });
       }
       r();
@@ -326,11 +352,14 @@ export class OpenSession implements ISessionFileManagerEventHandler {
   private readonly approvalManager: ToolApprovalManager;
   private readonly savedAgentProfile: SavedAgentProfile;
   private readonly sessionFileManager: ISessionFileManager;
+  private readonly platform: IPlatform;
+  private readonly raceModeAwaiter: ResponseAwaiter<ClientRaceModeResult>;
   private isPersisted: boolean;
   private accessToken: string | undefined;
   private sessionTitle: string;
   private sessionUpdatedAt: string;
   private agentPaused: boolean;
+  private readonly titleGenerator: ITitleGenerator;
 
   private constructor(
     db: Database,
@@ -344,7 +373,8 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     contextManager: ChatContextManager,
     sender: ChatSessionMessageSender,
     approvalManager: ToolApprovalManager,
-    fileManager: ISessionFileManager
+    fileManager: ISessionFileManager,
+    platform: IPlatform
   ) {
     this.db = db;
     this.agent = agent;
@@ -368,10 +398,17 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     this.approvalManager = approvalManager;
     this.savedAgentProfile = savedAgentProfile;
     this.sessionFileManager = fileManager;
+    this.platform = platform;
+    this.raceModeAwaiter = ResponseAwaiter.init(
+      undefined,
+      (m: ClientRaceModeResult) => m.message_id,
+      RACE_MODE_TIMEOUT_MS
+    );
     this.isPersisted = isPersisted;
     this.sessionTitle = sessionData.title;
     this.sessionUpdatedAt = sessionData.updated_at;
     this.agentPaused = sessionData.agent_paused;
+    this.titleGenerator = createTitleGenerator();
 
     fileManager.addEventHandler(this);
     this.updateParticipantsPrompt();
@@ -419,7 +456,8 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       xmcpUrl,
       fileManager,
       platform,
-      checkpointWriter
+      checkpointWriter,
+      sender
     );
 
     const openSession = new OpenSession(
@@ -434,7 +472,8 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       contextManager,
       sender,
       toolApprovalManager,
-      fileManager
+      fileManager,
+      platform
     );
 
     // Note, MCP servers have not been enabled yet
@@ -582,7 +621,9 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     });
 
     // send conversation history
-    const conversationMessages = this.contextManager.getConversationMessages();
+    const conversationMessages = this.contextManager
+      .getConversationMessages()
+      .concat(this.userMessageQueue.getEntries());
     conversationMessages.forEach((message) => {
       connMgr.sendToConnection(connectionId, message);
     });
@@ -680,29 +721,28 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       }
     };
 
+    let errorMessage: string;
     if (err instanceof ChatFatalError) {
       // ChatFatalError in session context should send error to specific user
       // or broadcast if no specific user context
-      sendError({
-        type: "session_error",
-        message: err.message,
-        session_id: this.sessionUUID,
-      });
+      errorMessage = err.message;
     } else if (err instanceof ChatErrorMessage) {
-      sendError({
-        type: "session_error",
-        message: err.message,
-        session_id: this.sessionUUID,
-      });
+      errorMessage = err.message;
     } else {
-      const errString = getErrorString(err);
-      sendError({
-        type: "session_error",
-        message: errString,
-        session_id: this.sessionUUID,
-      });
+      errorMessage = getErrorString(err);
     }
 
+    logger.error(
+      `[OpenSession.handleError] sessionUUID=${this.sessionUUID} ` +
+        `from=${from ?? "broadcast"}: ${errorMessage}`
+    );
+
+    sendError({
+      type: "session_error",
+      message: errorMessage,
+      session_id: this.sessionUUID,
+    });
+
     return true;
   }
 
@@ -715,6 +755,16 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     });
   }
 
+  private broadcastContextUsage(): void {
+    const { used, max } = this.contextManager.getContextUsage();
+    this.sender.broadcast({
+      type: "context_usage",
+      session_id: this.sessionUUID,
+      used_tokens: used,
+      max_tokens: max,
+    });
+  }
+
   // ISessionFileManagerEventHandler.onFileChanged
   onFileChanged(entry: SessionFileEntry, new_file: boolean): void {
     this.sender.broadcast({
@@ -812,6 +862,9 @@ export class OpenSession implements ISessionFileManagerEventHandler {
         case "msg":
           broadcastMsg = await this.handleUserMessage(msg, queuedMessage.from);
           break;
+        case "stop":
+          this.agent.stop();
+          break;
         case "add_mcp_server":
           broadcastMsg = await this.handleAddMcpServer(
             msg.server_name,
@@ -883,6 +936,9 @@ export class OpenSession implements ISessionFileManagerEventHandler {
         case "get_mcp_resource":
           void this.handleGetMcpResource(msg, queuedMessage.from);
           break;
+        case "race_mode_result":
+          this.raceModeAwaiter.onMessage(msg);
+          break;
         default: {
           const exhaustive: never = msg; // Error => non-exhaustive switch-case.
           return exhaustive;
@@ -896,6 +952,10 @@ export class OpenSession implements ISessionFileManagerEventHandler {
           });
         } else {
           this.sender.broadcast(broadcastMsg);
+          // Broadcast context usage after user message
+          if (broadcastMsg.type === "user_msg") {
+            this.broadcastContextUsage();
+          }
         }
       }
     } catch (err: unknown) {
@@ -978,16 +1038,14 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     // they are available to the LLM once it is restarted.
 
     const { contextTx } = await this.contextManager.startAgentResponse(msgs);
-    return this.contextManager.endAgentResponse(contextTx);
+    const result = this.contextManager.endAgentResponse(contextTx);
+    this.broadcastContextUsage();
+    return result;
   }
 
   private async processUserMessagesActive(
     msgs: ServerUserMessage[]
   ): Promise<SessionMessage[]> {
-    // TODO: create the contextTx and store all new messages on this.  Event
-    // handler should accept the contextTx and forward messages to it, as well
-    // as sending the updates.
-
     // All accumulated messages (DB, Protocol and LLM) should be on the
     // specialized contextTx.
 
@@ -1017,14 +1075,105 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       // DB.  Should we keep them?
       throw new Error(errMsg);
     }
-    return this.contextManager.endAgentResponse(contextTx);
+    const result = this.contextManager.endAgentResponse(contextTx);
+    this.broadcastContextUsage();
+    return result;
+  }
+
+  /**
+   * `processUserMessage` logic for race-mode.
+   */
+  private async processUserMessagesRaceMode(
+    msgs: ServerUserMessage[]
+  ): Promise<SessionMessage[]> {
+    // Create a second agent, same skillManager
+
+    const modelB = msgs[0].race_mode;
+    assert(typeof modelB === "string");
+
+    const agentB = await (async () => {
+      try {
+        const llmB = await createLLM(modelB, this.platform);
+        return new AgentEx(this.skillManager, llmB);
+      } catch (e) {
+        logger.warn(
+          "[OpenSession.processUserMessages] error creating race agent: " +
+            String(e)
+        );
+        throw new Error(`error creating race: ${String(e)}`);
+      }
+    })();
+
+    const run = async (
+      alt: string,
+      msgs: ServerUserMessage[]
+    ): Promise<ChatContextTransaction> => {
+      const { contextTx: contextTx, agentFirstChunk: agentFirstChunk } =
+        await this.contextManager.startAgentResponse(msgs.slice());
+      this.sender.broadcast(agentFirstChunk);
+
+      const eventHandler = new ChatSessionAgentEventHandler(
+        this.sessionUUID,
+        this.sender,
+        this.approvalManager,
+        contextTx,
+        alt
+      );
+
+      try {
+        const agent = alt === "B" ? agentB : this.agent;
+        await agent.userMessagesRaw(contextTx, eventHandler);
+      } catch (e) {
+        logger.warn(
+          `[OpenSession.processUserMessages] agent ${alt} error: ${String(e)}`
+        );
+        const errMsg = `error from LLM: ${String(e)}`;
+        contextTx.revertAgentResponse(errMsg);
+        throw new Error(errMsg);
+      }
+
+      return contextTx;
+    };
+
+    const [txA, txB] = await Promise.all([
+      run("A", msgs.slice()),
+      run("B", msgs),
+    ]);
+
+    // Ask which fork to commit
+
+    const answer = await (async () => {
+      const message_id = uuidv4();
+      const answerP = this.raceModeAwaiter.waitForResponse(message_id);
+      const getResultMsg: ServerRaceModeGetResult = {
+        type: "race_mode_get_result",
+        message_id,
+        alts: ["A", "B"],
+        session_id: this.sessionUUID,
+      };
+      this.sender.broadcast(getResultMsg);
+
+      return answerP;
+    })();
+
+    if (answer.result === "A") {
+      return this.contextManager.endAgentResponse(txA);
+    } else {
+      return this.contextManager.endAgentResponse(txB);
+    }
+
+    // UI is responsible for switching model if the user wants to.
   }
 
   private async processUserMessages(msgs: ServerUserMessage[]): Promise<void> {
+    logger.debug(`[processUserMessages] msgs: ${JSON.stringify(msgs)}`);
+
     try {
       const newSessionMessages = this.agentPaused
         ? await this.processUserMessagePaused(msgs)
-        : await this.processUserMessagesActive(msgs);
+        : await (msgs[0].race_mode
+            ? this.processUserMessagesRaceMode(msgs)
+            : this.processUserMessagesActive(msgs));
 
       logger.debug(
         "[processUserMessages] newSessionMessages: " +
@@ -1035,6 +1184,10 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       const dbsm = this.db.createTypedClient(DbSessionMessages);
       await dbsm.append(this.sessionUUID, newSessionMessages);
     } catch (e) {
+      logger.error(
+        `[processUserMessages] ERROR session=${this.sessionUUID}: ` +
+          String(e)
+      );
       if (!this.handleError(e)) {
         throw e;
       }
@@ -1049,11 +1202,22 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     // on a queue to be dealt with in another loop.  This allows Agent
     // processing of user messages to depend on other messages.
 
-    assert(msg);
-    assert(from);
+    assert(msg, "undefined user message");
+    assert(from, "undefined user message sender");
 
-    // Assign the user message_idx and attempt to enqueue.
+    logger.info(
+      `[handleUserMessage] msg.attachedFiles: ${JSON.stringify(
+        msg.attachedFiles
+          ? msg.attachedFiles.map((f) => ({
+              name: f.name,
+              data_url_length: f.data_url.length,
+            }))
+          : "none"
+      )}`
+    );
 
+    // Assign the user message_idx first so we can check if this is
+    // the first message
     const user = this.sessionParticipants.get(from);
     if (!user) {
       throw new Error(`unrecognized user ${from}`);
@@ -1066,8 +1230,87 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     if (!userMessage) {
       return;
     }
-    // Special case for the first message of the session
 
+    // Handle attached files by adding them to session files
+    if (msg.attachedFiles && msg.attachedFiles.length > 0) {
+      const fileCount = String(msg.attachedFiles.length);
+      logger.info(`[handleUserMessage] Processing ${fileCount} attached files`);
+
+      // If this is the first message, set the session title before
+      // creating the session
+      if (
+        userMessage.message_idx === MESSAGE_INDEX_START_VALUE &&
+        !this.isPersisted
+      ) {
+        this.sessionTitle = userMessage.message?.slice(0, 128) || "New Chat";
+      }
+
+      // If the session hasn't been persisted, it must be written to
+      // the DB first
+      if (!this.isPersisted) {
+        await this.createSessionInDB();
+      }
+
+      // Add each attached file to session files
+      for (const file of msg.attachedFiles) {
+        logger.info(
+          `[handleUserMessage] Adding file ${file.name} to session files`
+        );
+
+        // Extract content and generate summary for PDFs
+        let parsed_content: ParsedContent | undefined;
+        let summary: string | undefined;
+
+        const mimeType = getMimeTypeFromDataUrl(file.data_url);
+        if (mimeType === "application/pdf") {
+          try {
+            // Extract base64 data and convert to ArrayBuffer
+            const base64Data = file.data_url.split(",")[1];
+            const binaryString = atob(base64Data);
+            const bytes = new Uint8Array(binaryString.length);
+            for (let i = 0; i < binaryString.length; i++) {
+              bytes[i] = binaryString.charCodeAt(i);
+            }
+            const arrayBuffer = bytes.buffer;
+
+            // Extract text from PDF
+            const extractedText = await pdfToText(arrayBuffer);
+            parsed_content = {
+              version: 1,
+              text: extractedText,
+            };
+
+            // Generate summary from extracted text
+            if (extractedText.trim().length > 0) {
+              summary = await summarizeDocument(extractedText);
+            }
+
+            logger.info(
+              `[handleUserMessage] Extracted ${String(extractedText.length)}` +
+                ` chars from PDF ${file.name}`
+            );
+          } catch (err) {
+            logger.warn(
+              `[handleUserMessage] Failed to extract PDF content: ` +
+                String(err)
+            );
+          }
+        }
+
+        await this.sessionFileManager.putFileContent(
+          file.name,
+          summary,
+          file.data_url,
+          parsed_content
+        );
+      }
+
+      // Remove attachedFiles from the message so they don't get
+      // included in the LLM context
+      msg = { ...msg, attachedFiles: undefined };
+    }
+
+    // Special case for the first message of the session
     if (userMessage.message_idx === MESSAGE_INDEX_START_VALUE) {
       // No need to wait for this to complete before broadcasting.
       await this.onFirstMessage(userMessage);
@@ -1118,9 +1361,11 @@ export class OpenSession implements ISessionFileManagerEventHandler {
   }
 
   private async onFirstMessage(userMsg: ServerUserMessage): Promise<void> {
-    // Update title on the class before writing to the DB
+    // Generate title using LLM with automatic fallback
 
-    this.sessionTitle = userMsg.message?.slice(0, 128) || "New Chat";
+    this.sessionTitle = await this.titleGenerator.generateTitle(
+      userMsg.message || ""
+    );
 
     // The session may already have been saved (e.g. if the workspace is
     // updated before any messages are sent).
@@ -1626,6 +1871,15 @@ async function loadSessionData(
   };
 }
 
+async function createLLM(model: string, platform: IPlatform): Promise<ILLM> {
+  let llm = await createSpecializedLLM(model, platform);
+  if (!llm) {
+    llm = new OpenAIRouterLLM(model);
+  }
+  assert(llm);
+  return llm;
+}
+
 async function createContextAndAgent(
   sessionUUID: string,
   systemPrompt: string,
@@ -1638,22 +1892,38 @@ async function createContextAndAgent(
   xmcpUrl: string,
   fileManager: ChatSessionFileManager,
   platform: IPlatform,
-  checkpointWriter: ICheckpointWriter
+  checkpointWriter: ICheckpointWriter,
+  messageSender: ISessionMessageSender<ServerToClient>
 ): Promise<{
   agent: AgentEx;
   skillManager: SkillManager;
   contextManager: ChatContextManager;
 }> {
+  // Create SkillManager
+
+  const xmcpConfig = Configuration.new(ownerApiKey, xmcpUrl, false);
+  const skillManager = await SkillManager.initialize(
+    (url: string, authResultP: Promise<boolean>, displayName: string) => {
+      platform.openUrl(url, authResultP, displayName);
+    },
+    xmcpConfig.backend_url,
+    xmcpConfig.api_key,
+    undefined /* authorizedUrl */
+  );
+
   // Fn to create the llm.  One invocation for the compression context, one
   // for the Agent.
-  const createLLM = async (): Promise<ILLM> => {
-    let llm = await createSpecializedLLM(model, platform);
-    if (!llm) {
-      llm = new OpenAIRouterLLM(model);
-    }
-    assert(llm);
-    return llm;
-  };
+
+  const llm = await createLLM(model, platform);
+  const agent = new AgentEx(skillManager, llm);
+
+  await addDefaultChatTools(
+    agent,
+    ownerData.timezone,
+    platform,
+    fileManager,
+    messageSender
+  );
 
   const contextManager = new ChatContextManager(
     systemPrompt,
@@ -1663,7 +1933,7 @@ async function createContextAndAgent(
     sessionCheckpoint,
     checkpointWriter,
     fileManager,
-    await createLLM()
+    await createLLM(model, platform)
   );
   if (workspace) {
     contextManager.setWorkspace(
@@ -1671,19 +1941,5 @@ async function createContextAndAgent(
     );
   }
 
-  const xmcpConfig = Configuration.new(ownerApiKey, xmcpUrl, false);
-  const skillManager = await SkillManager.initialize(
-    (url: string, authResultP: Promise<boolean>, displayName: string) => {
-      platform.openUrl(url, authResultP, displayName);
-    },
-    xmcpConfig.backend_url,
-    xmcpConfig.api_key,
-    undefined /* authorizedUrl */
-  );
-  const llm = await createLLM();
-  const agent = new AgentEx(skillManager, llm);
-
-  await addDefaultChatTools(agent, ownerData.timezone, platform, fileManager);
-
   return { agent, skillManager, contextManager };
 }