@xalia/agent 0.6.8 → 0.6.10

package/src/chat/server/chatContextManager.ts

@@ -2,14 +2,14 @@ import { strict as assert } from "assert";
 
 import { getLogger } from "@xalia/xmcp/sdk";
 
-import { createUserMessage } from "../../agent/agent";
 import {
   AssistantMessageParam,
+  ILLM,
   MessageParam,
   ToolMessageParam,
   UserMessageParam,
 } from "../../agent/llm";
-import { IContextManager } from "../../agent/context";
+import { IContextTransaction } from "../../agent/context";
 import {
   CompressingContextManager,
   createCheckpointMessage,
@@ -34,11 +34,14 @@ import {
   ISessionFileManager,
   ISessionFileManagerEventHandler,
   createSessionFilesManagerPrompt,
-} from "./sessionFileManager";
-import {
   SessionFileDescriptor,
   SessionFileEntry,
-} from "../data/dbSessionFileModels";
+} from "../../agent/tools/fileManager";
+// eslint-disable-next-line max-len
+import { ContextTransactionWithWorkspace } from "../../agent/contextWithWorkspace";
+import { getErrorString } from "./errorUtils";
+import { createUserMessage } from "../../agent/agent";
+import { TokenCounter } from "../../agent/tokenCounter";
 
 const logger = getLogger();
 
@@ -57,6 +60,144 @@ export interface ICheckpointWriter {
   writeCheckpoint(checkpoint: SessionCheckpoint): Promise<void>;
 }
 
+export class ChatContextTransaction implements IContextTransaction {
+  private readonly baseTx: ContextTransactionWithWorkspace;
+  private readonly sessionUUID: string;
+  /// Index of final message in the committed context.  If this has changed
+  /// before we try to commit this tx, the commit will fail.
+  private readonly baseMsgIdx: number | undefined;
+  private readonly startingLLMContextLength: number;
+  private readonly pendingMessages: ConversationMessage[];
+  private curAgentMsgIdx: number;
+
+  constructor(
+    baseTx: ContextTransactionWithWorkspace,
+    sessionUUID: string,
+    baseMsgIdx: number | undefined,
+    pendingUserMessages: ServerUserMessage[],
+    curAgentMsgIdx: number
+  ) {
+    assert(typeof curAgentMsgIdx !== "undefined");
+
+    this.sessionUUID = sessionUUID;
+    this.baseTx = baseTx;
+    this.baseMsgIdx = baseMsgIdx;
+    this.startingLLMContextLength = baseTx.getLLMContextLength();
+    this.pendingMessages = pendingUserMessages;
+    this.curAgentMsgIdx = curAgentMsgIdx;
+  }
+
+  // IContextTransaction.addMessages
+  addMessages(messages: MessageParam[]): number {
+    return this.baseTx.addMessages(messages);
+  }
+
+  // IContextTransaction.addMessage
+  addMessage(message: MessageParam): number {
+    return this.baseTx.addMessage(message);
+  }
+
+  // IContextTransaction.getMessage
+  getMessage(handle: number): MessageParam {
+    return this.baseTx.getMessage(handle);
+  }
+
+  // IContextTransaction.getLLMContext
+  getLLMContext(): MessageParam[] {
+    return this.baseTx.getLLMContext();
+  }
+
+  // IContextTransaction.getLLMContextLength
+  getLLMContextLength(): number {
+    return this.baseTx.getLLMContextLength();
+  }
+
+  getPending(): ConversationMessage[] {
+    return this.pendingMessages;
+  }
+
+  baseMessageIdx(): number | undefined {
+    return this.baseMsgIdx;
+  }
+
+  getBaseTx(): ContextTransactionWithWorkspace {
+    return this.baseTx;
+  }
+
+  /**
+   * Process a FULL Agent message (not chunks from stream).  No message is
+   * required for broadcast as the calling code is expected to broadcast this
+   * as chunks.
+   */
+  processAgentResponse(result: AssistantMessageParam) {
+    // Insert this (full) agent response into the list of agent messages
+    const msg: ServerAgentMessage = {
+      type: "agent_msg",
+      session_id: this.sessionUUID,
+      message_idx: this.getNextMessageSubIdx(),
+      message: result,
+    };
+    this.pendingMessages.push(msg);
+  }
+
+  processAgentMessageChunk(msg: string, end: boolean): ServerAgentMessageChunk {
+    const message: ServerAgentMessageChunk = {
+      type: "agent_msg_chunk",
+      session_id: this.sessionUUID,
+      message_idx: this.getCurrentAgentMessageIdx(),
+      message: msg,
+      end,
+    };
+    return message;
+  }
+
+  processToolCallResult(result: ToolMessageParam): ServerToolCallResult {
+    // Allocate the sub-index for this tool call result.  It should not
+    // have been used already.
+
+    const message_idx = this.getNextMessageSubIdx();
+    const numPending = this.pendingMessages.length;
+    assert(numPending > 0);
+    assert(this.pendingMessages[numPending - 1].message_idx < message_idx);
+
+    const msg: ServerToolCallResult = {
+      type: "tool_call_result",
+      session_id: this.sessionUUID,
+      message_idx,
+      result,
+    };
+    this.pendingMessages.push(msg);
+    return msg;
+  }
+
+  revertAgentResponse(errMsg: string): void {
+    logger.warn(`[ChatContextManager.revertAgentResponse] error: ${errMsg}`);
+
+    // Remove all messages since the user messages were placed on.
+
+    while (this.baseTx.getLLMContextLength() > this.startingLLMContextLength) {
+      const last = this.baseTx.popMessage();
+      assert(last);
+    }
+  }
+
+  newMessages(): MessageParam[] {
+    return this.baseTx.newMessages();
+  }
+
+  private getNextMessageSubIdx(): number {
+    const idx = this.curAgentMsgIdx;
+    this.curAgentMsgIdx += MESSAGE_INDEX_SUB_INCREMENT;
+    return idx;
+  }
+
+  /// Get the current index to use for streaming Agent chunks
+  private getCurrentAgentMessageIdx(): number {
+    assert(typeof this.curAgentMsgIdx !== "undefined");
+    return this.curAgentMsgIdx;
+  }
+}
+
 /**
  * A context manager for Agents interacting with the (potentially multi-user)
  * chat conversations.
@@ -68,40 +209,35 @@ export interface ICheckpointWriter {
  * - Maintain pending user messages, agent loop messages and messages to be
  *   committed to the DB
  */
-export class ChatContextManager
-  implements IContextManager, ISessionFileManagerEventHandler
-{
+export class ChatContextManager implements ISessionFileManagerEventHandler {
   // Including any pending user messages
-  private sessionUUID: string;
-  private conversationMessages: ConversationMessage[];
-  private pendingUserMessages: ServerUserMessage[];
-  private llmContext: CompressingContextManager;
-  private nextMessageIdx: number;
+  private readonly sessionUUID: string;
+  private readonly conversationMessages: ConversationMessage[];
+  private readonly llmContext: CompressingContextManager;
 
-  // State while processing agent messages
-  private startingLLMContextLength: number | undefined;
-  private curAgentMsgIdx: number | undefined; // active agent message
-  private pendingMessages: ConversationMessage[] | undefined;
+  private nextMessageIdx: number;
 
   // Compression state
-  private checkpointWriter: ICheckpointWriter;
+  private readonly checkpointWriter: ICheckpointWriter;
   private pendingCompression: boolean;
 
   // FileManager
-  private fileManager: ISessionFileManager;
+  private readonly fileManager: ISessionFileManager;
   private fileManagerDescriptionsDirty: boolean;
 
+  // LLM and token counting
+  private readonly llm: ILLM;
+  private tokenCounter: TokenCounter;
+
   constructor(
     systemPrompt: string,
     sessionMessages: SessionMessage[],
     sessionUUID: string,
     defaultUserName: string,
     checkpoint: SessionCheckpoint | undefined = undefined,
-    compressionAgentUrl: string,
-    compressionAgentModel: string,
-    compressionAgentApiKey: string,
     checkpointWriter: ICheckpointWriter,
-    fileManager: ISessionFileManager
+    fileManager: ISessionFileManager,
+    llm: ILLM
   ) {
     const nextMessageIdx = sessionMessagesToNextIndex(sessionMessages);
     const { messages: llmMessages } = resolveConversationWithCheckpoint(
@@ -113,7 +249,7 @@ export class ChatContextManager
       `[ChatContextManager]: llm messages: ${JSON.stringify(llmMessages)}`
     );
 
-    // Insert a system message placeholder into the context
+    const getLLM = () => Promise.resolve(llm);
 
     this.sessionUUID = sessionUUID;
     this.conversationMessages = sessionMessagesToConversationMessages(
@@ -121,40 +257,25 @@ export class ChatContextManager
       defaultUserName,
       sessionUUID
     );
-    this.pendingUserMessages = [];
     this.llmContext = new CompressingContextManager(
       systemPrompt,
       llmMessages,
-      compressionAgentUrl,
-      compressionAgentModel,
-      compressionAgentApiKey
+      getLLM
     );
     this.nextMessageIdx = nextMessageIdx;
-    this.startingLLMContextLength = undefined;
-    this.curAgentMsgIdx = undefined;
-    this.pendingMessages = undefined;
     this.pendingCompression = false;
     this.checkpointWriter = checkpointWriter;
     this.fileManager = fileManager;
     fileManager.addEventHandler(this);
     this.fileManagerDescriptionsDirty = true;
-  }
-
-  // IContextManager.addMessages
-  addMessages(messages: MessageParam[]): void {
-    this.llmContext.addMessages(messages);
-  }
-
-  // IContextManager.addMessage
-  addMessage(message: MessageParam): void {
-    this.llmContext.addMessage(message);
+    this.llm = llm;
+    this.tokenCounter = new TokenCounter(llm.getModel());
   }
 
   // IContextManager.getLLMContext
   getLLMContext(): MessageParam[] {
     if (this.fileManagerDescriptionsDirty) {
       const prompt = createSessionFilesManagerPrompt(this.fileManager);
-      logger.debug(`[ChatContextManager] filemanager prompt:\n${prompt}`);
       this.llmContext.setPromptFragment("file_manager", prompt);
       this.fileManagerDescriptionsDirty = false;
     }
@@ -166,6 +287,7 @@ export class ChatContextManager
   getAgentPrompt(): string {
     return this.llmContext.getAgentPrompt();
   }
+
   // IContextManager.setAgentPrompt
   setAgentPrompt(prompt: string): void {
     this.llmContext.setAgentPrompt(prompt);
@@ -205,7 +327,22 @@ export class ChatContextManager
 
   // Get the conversation (to send to clients)
   getConversationMessages(): ConversationMessage[] {
-    return this.conversationMessages.concat(this.pendingUserMessages);
+    return this.conversationMessages;
+  }
+
+  // Get current context usage (tokens)
+  getContextUsage(): { used: number; max: number } {
+    // Update tokenCounter if model changed
+    const currentModel = this.llm.getModel();
+    if (this.tokenCounter.getModel() !== currentModel) {
+      this.tokenCounter.free();
+      this.tokenCounter = new TokenCounter(currentModel);
+    }
+
+    const messages = this.getLLMContext();
+    const used = this.tokenCounter.countMessagesTokens(messages);
+    const max = this.tokenCounter.getContextWindow();
+    return { used, max };
   }
 
   processUserMessage(
@@ -217,7 +354,11 @@ export class ChatContextManager
     // pass in our generated messages back into `startAgentResponse`.
 
     // Filter out null messages immediately.
-    if (!msg.imageB64 && !msg.message) {
+    if (
+      !msg.imageB64 &&
+      !msg.message &&
+      (!msg.attachedFiles || msg.attachedFiles.length === 0)
+    ) {
       return undefined;
     }
 
@@ -233,8 +374,13 @@ export class ChatContextManager
     if (msg.imageB64) {
       userMessage.imageB64 = msg.imageB64;
     }
+    if (msg.attachedFiles) {
+      userMessage.attachedFiles = msg.attachedFiles;
+    }
+    if (msg.race_mode) {
+      userMessage.race_mode = msg.race_mode;
+    }
 
-    this.pendingUserMessages.push(userMessage);
     return userMessage;
   }
 
@@ -247,43 +393,25 @@ export class ChatContextManager
   // class manage the interaction with the agent and ensure this process only
   // happens one-at-a-time.
 
-  startAgentResponse(msgs: ServerUserMessage[]): {
-    llmUserMessages: UserMessageParam[];
+  async startAgentResponse(msgs: ServerUserMessage[]): Promise<{
+    contextTx: ChatContextTransaction;
     agentFirstChunk: ServerAgentMessageChunk;
-  } {
-    // Sanity check the state
-
-    assert(
-      typeof this.startingLLMContextLength === "undefined",
-      "already processing"
-    );
-    assert(typeof this.pendingMessages === "undefined", "already processing");
-    assert(typeof this.curAgentMsgIdx === "undefined", "already processing");
-
-    // Sanity check the messages
-
-    const numMessages = this.pendingUserMessages.length;
-    assert(numMessages > 0);
-    assert(msgs.length === this.pendingUserMessages.length);
-    assert(msgs[0].message_idx === this.pendingUserMessages[0].message_idx);
-    assert(
-      msgs[numMessages - 1].message_idx ===
-        this.pendingUserMessages[numMessages - 1].message_idx
-    );
+  }> {
+    // Sanity check the state - the incoming user messages should match the
+    // pending user messages.
 
+    const numMessages = msgs.length;
+    assert(numMessages > 0, "no messages");
     // Collect the pending user messages and allocate a starting index for
     // agent messages and tool calls.
 
-    const pendingUserMessages = this.pendingUserMessages;
-    this.pendingUserMessages = [];
-    this.startingLLMContextLength = this.llmContext.getCommittedLength();
-    this.curAgentMsgIdx = this.getNextMessageIdx();
-    this.pendingMessages = pendingUserMessages as ConversationMessage[];
+    const baseMsgIdx = this.lastCommittedMessageIdx();
+    const curAgentMsgIdx = this.getNextMessageIdx();
 
     // Compute the new llm messages
 
     const llmUserMessages: UserMessageParam[] = [];
-    for (const msg of pendingUserMessages) {
+    for (const msg of msgs) {
       const userMsg = createUserMessage(
         msg.message,
         msg.imageB64,
@@ -294,49 +422,60 @@ export class ChatContextManager
       }
     }
 
-    return {
-      llmUserMessages,
-      agentFirstChunk: {
-        type: "agent_msg_chunk",
-        session_id: this.sessionUUID,
-        message_idx: this.curAgentMsgIdx,
-        message: "",
-        end: false,
-      },
+    // Return the context tx and first ServerAgentMessageChunk
+
+    const agentFirstChunk: ServerAgentMessageChunk = {
+      type: "agent_msg_chunk",
+      session_id: this.sessionUUID,
+      message_idx: curAgentMsgIdx,
+      message: "",
+      end: false,
     };
-  }
 
-  endAgentResponse(): SessionMessage[] {
-    assert(
-      typeof this.startingLLMContextLength !== "undefined",
-      "agent response not started (startingLLMContextLength)"
-    );
-    assert(
-      typeof this.pendingMessages !== "undefined",
-      "agent response not started (pendingMessages)"
-    );
-    assert(
-      typeof this.curAgentMsgIdx !== "undefined",
-      "agent response not started (curAgentMsgIdx)"
+    // Update file manager fragment BEFORE starting the transaction
+    if (this.fileManagerDescriptionsDirty) {
+      const prompt = createSessionFilesManagerPrompt(this.fileManager);
+      this.llmContext.setPromptFragment("file_manager", prompt);
+      this.fileManagerDescriptionsDirty = false;
+    }
+
+    const baseTx = await this.llmContext.startTx(llmUserMessages);
+    const contextTx = new ChatContextTransaction(
+      baseTx,
+      this.sessionUUID,
+      baseMsgIdx,
+      msgs,
+      curAgentMsgIdx
     );
+    return { agentFirstChunk, contextTx };
+  }
 
-    const numPending = this.pendingMessages.length;
+  async endAgentResponse(tx: IContextTransaction): Promise<SessionMessage[]> {
+    assert(tx instanceof ChatContextTransaction);
+    if (tx.baseMessageIdx() !== this.lastCommittedMessageIdx()) {
+      throw new Error(
+        `Tx stale? tx.baseMessageIdx=${String(tx.baseMessageIdx())}, ` +
+          `this.conv: ${JSON.stringify(this.conversationMessages)}`
+      );
+    }
+
+    const pending = tx.getPending();
+    const numPending = pending.length;
     assert(numPending > 0, "no pending"); // at least 1 user message
 
     // Compute DB messages
 
-    const newSessionMessages = chatMessagesToSessionMessages(
-      this.pendingMessages
-    );
-    const newLLMMessages = this.llmContext.getPending();
+    const newSessionMessages = chatMessagesToSessionMessages(pending);
+    const newLLMMessages = tx.newMessages();
 
     const messageListError = (error: string) => {
-      throw new Error(
-        `${error}:` +
-          `\n  newSessionMessages: ${JSON.stringify(newSessionMessages)}` +
-          `\n  this.pendingMessages: ${JSON.stringify(this.pendingMessages)}` +
-          `\n  newLLMMessages: ${JSON.stringify(newLLMMessages)}`
-      );
+      const fullError =
+        `[endAgentResponse] Message list validation failed - ${error}:` +
+        `\n  newSessionMessages: ${JSON.stringify(newSessionMessages)}` +
+        `\n  pending: ${JSON.stringify(pending)}` +
+        `\n  newLLMMessages: ${JSON.stringify(newLLMMessages)}`;
+      logger.error(fullError);
+      throw new Error(fullError);
     };
 
     if (newSessionMessages.length !== numPending) {
@@ -353,7 +492,7 @@ export class ChatContextManager
 
     for (let i = 0; i < numPending; ++i) {
       const sMsg = newSessionMessages[i];
-      const pMsg = this.pendingMessages[i];
+      const pMsg = pending[i];
       const lMsg = newLLMMessages[i];
 
       if (sMsg.content.role !== lMsg.role) {
@@ -379,11 +518,8 @@ export class ChatContextManager
     // Update our internal state and return the SessionMessages to write to
     // the DB
 
-    this.llmContext.commit();
-    this.conversationMessages.push(...this.pendingMessages);
-    this.startingLLMContextLength = undefined;
-    this.pendingMessages = undefined;
-    this.curAgentMsgIdx = undefined;
+    await this.llmContext.commit(tx.getBaseTx());
+    this.conversationMessages.push(...pending);
 
     // Kick off a compression?
     this.checkCompression();
@@ -391,98 +527,12 @@ export class ChatContextManager
     return newSessionMessages;
   }
 
-  /**
-   * End the Agent message session with an error.  Caller should not call
-   * `endAgentResponse` after calling this function.
-   *
-   * This function checks that nothing has been entered into the LLM context,
-   * and drops any new user messages or responses before the error.
-   */
-  revertAgentResponse(errMsg: string): void {
-    logger.warn(`[ChatContextManager.revertAgentResponse] error: ${errMsg}`);
-
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    // Sanity check that no new messages were put into the context (The Agent
-    // is expected to only call `addMessage(s)` at the end of the Agent loop.
-    // (Note, we don't check for equality here, just in case the context was
-    // compressed while the Agent was executing).
-
-    const contextLength = this.llmContext.getCommittedLength();
-    if (contextLength > this.startingLLMContextLength) {
-      logger.error(
-        "[ChatContextManager.revertAgentResponse] llmContext has grown " +
-          `despite Agent error (${String(contextLength)}, ` +
-          `${String(this.startingLLMContextLength)})`
-      );
+  private lastCommittedMessageIdx(): number | undefined {
+    const numMsgs = this.conversationMessages.length;
+    if (numMsgs > 0) {
+      return this.conversationMessages[numMsgs - 1].message_idx;
     }
-
-    // We simply reset the state, dropping any pending messages.
-
-    this.startingLLMContextLength = undefined;
-    this.pendingMessages = undefined;
-    this.curAgentMsgIdx = undefined;
-  }
-
-  processAgentMessage(msg: string, end: boolean): ServerAgentMessageChunk {
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    const message: ServerAgentMessageChunk = {
-      type: "agent_msg_chunk",
-      session_id: this.sessionUUID,
-      message_idx: this.getCurrentAgentMessageIdx(),
-      message: msg,
-      end,
-    };
-    return message;
-  }
-
-  /**
-   * Process a FULL Agent message (not chunks from stream).  No message is
-   * required for broadcast as the calling code is expected to broadcast this
-   * as chunks.
-   */
-  processAgentResponse(result: AssistantMessageParam) {
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    // Insert this (full) agent response into the list of agent messages
-
-    const msg: ServerAgentMessage = {
-      type: "agent_msg",
-      session_id: this.sessionUUID,
-      message_idx: this.getNextMessageSubIdx(),
-      message: result,
-    };
-    this.pendingMessages.push(msg);
-  }
-
-  processToolCallResult(result: ToolMessageParam): ServerToolCallResult {
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    // Allocate the sub-index for this tool call result.  It should not
-    // have been used already.
-
-    const message_idx = this.getNextMessageSubIdx();
-    const numPending = this.pendingMessages.length;
-    assert(numPending > 0);
-    assert(this.pendingMessages[numPending - 1].message_idx < message_idx);
-
-    const msg: ServerToolCallResult = {
-      type: "tool_call_result",
-      session_id: this.sessionUUID,
-      message_idx,
-      result,
-    };
-    this.pendingMessages.push(msg);
-    return msg;
+    return undefined;
   }
 
   private getNextMessageIdx(): number {
@@ -503,22 +553,6 @@ export class ChatContextManager
     this.nextMessageIdx = messageIdx;
   }
 
-  /// Get the current index to use for streaming Agent chunks
-  private getCurrentAgentMessageIdx(): number {
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-    return this.curAgentMsgIdx;
-  }
-
-  private getNextMessageSubIdx(): number {
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    const idx = this.curAgentMsgIdx;
-    this.curAgentMsgIdx += MESSAGE_INDEX_SUB_INCREMENT;
-    return idx;
-  }
-
   private checkCompression(): void {
     if (this.pendingCompression) {
       return;
@@ -526,7 +560,7 @@ export class ChatContextManager
 
     // TODO: track tokens and use that to trigger compression
 
-    const numCommitted = this.llmContext.getCommittedLength();
+    const numCommitted = this.llmContext.numMessages();
     if (numCommitted < COMPRESSION_TRIGGER_NUM_MESSAGES) {
       return;
     }
@@ -551,7 +585,7 @@ export class ChatContextManager
       await this.checkpointWriter.writeCheckpoint(checkpoint);
     } catch (err: unknown) {
       logger.warn(
-        `[runCompression] error during compression: ${JSON.stringify(err)}`
+        `[runCompression] error during compression: ${getErrorString(err)}`
       );
     } finally {
       this.pendingCompression = false;

package/src/chat/server/connectionManager.ts

@@ -25,7 +25,7 @@ export interface IUserConnectionManager<ServerMsgT> {
    * Send message to all active connections of specific users.
    * Handles user-to-connection routing internally.
    */
-  sendToUsers(userIds: Set<string>, message: ServerMsgT): void;
+  sendToUsers(userIds: Set<string> | string[], message: ServerMsgT): void;
 
   /**
    * Send message to a specific connection.

package/src/chat/server/conversation.ts

@@ -127,6 +127,9 @@ export function userMessageToConversationMessage(
   if (userMsgData.imageB64) {
     msg.imageB64 = userMsgData.imageB64;
   }
+  if (userMsgData.attachedFiles) {
+    msg.attachedFiles = userMsgData.attachedFiles;
+  }
   return msg;
 }