@xalia/agent 0.6.10 → 0.6.11

package/src/chat/server/chatContextManager.ts

@@ -1,660 +0,0 @@
-import { strict as assert } from "assert";
-
-import { getLogger } from "@xalia/xmcp/sdk";
-
-import {
-  AssistantMessageParam,
-  ILLM,
-  MessageParam,
-  ToolMessageParam,
-  UserMessageParam,
-} from "../../agent/llm";
-import { IContextTransaction } from "../../agent/context";
-import {
-  CompressingContextManager,
-  createCheckpointMessage,
-} from "../../agent/compressingContextManager";
-import { SessionCheckpoint, SessionMessage } from "../data/dataModels";
-import {
-  ClientUserMessage,
-  ServerAgentMessage,
-  ServerAgentMessageChunk,
-  ServerToolCallResult,
-  ServerUserMessage,
-} from "../protocol/messages";
-import {
-  ConversationMessage,
-  MESSAGE_INDEX_FULL_INCREMENT,
-  MESSAGE_INDEX_SUB_INCREMENT,
-  chatMessagesToSessionMessages,
-  sessionMessagesToConversationMessages,
-  sessionMessagesToNextIndex,
-} from "./conversation";
-import {
-  ISessionFileManager,
-  ISessionFileManagerEventHandler,
-  createSessionFilesManagerPrompt,
-  SessionFileDescriptor,
-  SessionFileEntry,
-} 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();
-
-/**
- * TODO: Until token-tracking is in place.
- */
-const COMPRESSION_TRIGGER_NUM_MESSAGES: number = parseInt(
-  process.env["COMPRESSION_TRIGGER_NUM_MESSAGES"] || "80",
-  10
-);
-
-/**
- * Interface to store checkpoints for a specific session.
- */
-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.
- *
- * - Recreate an LLM context from message history + checkpoints
- *
- * - Maintain a compressing LLM context, and generate new checkpoints
- *
- * - Maintain pending user messages, agent loop messages and messages to be
- *   committed to the DB
- */
-export class ChatContextManager implements ISessionFileManagerEventHandler {
-  // Including any pending user messages
-  private readonly sessionUUID: string;
-  private readonly conversationMessages: ConversationMessage[];
-  private readonly llmContext: CompressingContextManager;
-
-  private nextMessageIdx: number;
-
-  // Compression state
-  private readonly checkpointWriter: ICheckpointWriter;
-  private pendingCompression: boolean;
-
-  // FileManager
-  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,
-    checkpointWriter: ICheckpointWriter,
-    fileManager: ISessionFileManager,
-    llm: ILLM
-  ) {
-    const nextMessageIdx = sessionMessagesToNextIndex(sessionMessages);
-    const { messages: llmMessages } = resolveConversationWithCheckpoint(
-      sessionMessages,
-      checkpoint
-    );
-
-    logger.debug(
-      `[ChatContextManager]: llm messages: ${JSON.stringify(llmMessages)}`
-    );
-
-    const getLLM = () => Promise.resolve(llm);
-
-    this.sessionUUID = sessionUUID;
-    this.conversationMessages = sessionMessagesToConversationMessages(
-      sessionMessages,
-      defaultUserName,
-      sessionUUID
-    );
-    this.llmContext = new CompressingContextManager(
-      systemPrompt,
-      llmMessages,
-      getLLM
-    );
-    this.nextMessageIdx = nextMessageIdx;
-    this.pendingCompression = false;
-    this.checkpointWriter = checkpointWriter;
-    this.fileManager = fileManager;
-    fileManager.addEventHandler(this);
-    this.fileManagerDescriptionsDirty = true;
-    this.llm = llm;
-    this.tokenCounter = new TokenCounter(llm.getModel());
-  }
-
-  // IContextManager.getLLMContext
-  getLLMContext(): MessageParam[] {
-    if (this.fileManagerDescriptionsDirty) {
-      const prompt = createSessionFilesManagerPrompt(this.fileManager);
-      this.llmContext.setPromptFragment("file_manager", prompt);
-      this.fileManagerDescriptionsDirty = false;
-    }
-
-    return this.llmContext.getLLMContext();
-  }
-
-  // IContextManager.getAgentPrompt
-  getAgentPrompt(): string {
-    return this.llmContext.getAgentPrompt();
-  }
-
-  // IContextManager.setAgentPrompt
-  setAgentPrompt(prompt: string): void {
-    this.llmContext.setAgentPrompt(prompt);
-  }
-
-  // IContextManager.setPromptFragment
-  setPromptFragment(fragmentID: string, prompt: string): void {
-    this.llmContext.setPromptFragment(fragmentID, prompt);
-  }
-
-  // IContextManager.removePromptFragment
-  removePromptFragment(fragmentID: string): void {
-    this.llmContext.removePromptFragment(fragmentID);
-  }
-
-  // ISessionFileManagerEventHandler.onFileDescriptorChange
-  onFileDescriptorChange(_desc: SessionFileDescriptor): void {
-    this.fileManagerDescriptionsDirty = true;
-  }
-
-  // ISessionFileManagerEventHandler.onFileChange
-  onFileChanged(_entry: SessionFileEntry): void {
-    this.fileManagerDescriptionsDirty = true;
-  }
-
-  onFileDeleted(_name: string): void {
-    this.fileManagerDescriptionsDirty = true;
-  }
-
-  setWorkspace(userMessage: UserMessageParam | undefined): void {
-    this.llmContext.setWorkspace(userMessage);
-  }
-
-  getWorkspace(): UserMessageParam | undefined {
-    return this.llmContext.getWorkspace();
-  }
-
-  // Get the conversation (to send to clients)
-  getConversationMessages(): ConversationMessage[] {
-    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(
-    msg: ClientUserMessage,
-    from_uuid: string,
-    from_nickname: string
-  ): ServerUserMessage | undefined {
-    // TODO: maintain a queue internally instead of relying on the caller to
-    // pass in our generated messages back into `startAgentResponse`.
-
-    // Filter out null messages immediately.
-    if (
-      !msg.imageB64 &&
-      !msg.message &&
-      (!msg.attachedFiles || msg.attachedFiles.length === 0)
-    ) {
-      return undefined;
-    }
-
-    const message_idx = this.getNextMessageIdx();
-    const userMessage: ServerUserMessage = {
-      type: "user_msg",
-      session_id: this.sessionUUID,
-      message_idx,
-      message: msg.message,
-      user_uuid: from_uuid,
-      user_nickname: from_nickname,
-    };
-    if (msg.imageB64) {
-      userMessage.imageB64 = msg.imageB64;
-    }
-    if (msg.attachedFiles) {
-      userMessage.attachedFiles = msg.attachedFiles;
-    }
-    if (msg.race_mode) {
-      userMessage.race_mode = msg.race_mode;
-    }
-
-    return userMessage;
-  }
-
-  unprocessUserMessage(userMsg: ServerUserMessage): void {
-    // TODO: when we maintain a queue, remove the entry from the queue.
-    this.freeMessageIdx(userMsg.message_idx);
-  }
-
-  // TODO: Don't take the set of messages.  Instead, have OpenSession or this
-  // class manage the interaction with the agent and ensure this process only
-  // happens one-at-a-time.
-
-  async startAgentResponse(msgs: ServerUserMessage[]): Promise<{
-    contextTx: ChatContextTransaction;
-    agentFirstChunk: ServerAgentMessageChunk;
-  }> {
-    // 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 baseMsgIdx = this.lastCommittedMessageIdx();
-    const curAgentMsgIdx = this.getNextMessageIdx();
-
-    // Compute the new llm messages
-
-    const llmUserMessages: UserMessageParam[] = [];
-    for (const msg of msgs) {
-      const userMsg = createUserMessage(
-        msg.message,
-        msg.imageB64,
-        msg.user_uuid
-      );
-      if (userMsg) {
-        llmUserMessages.push(userMsg);
-      }
-    }
-
-    // Return the context tx and first ServerAgentMessageChunk
-
-    const agentFirstChunk: ServerAgentMessageChunk = {
-      type: "agent_msg_chunk",
-      session_id: this.sessionUUID,
-      message_idx: curAgentMsgIdx,
-      message: "",
-      end: false,
-    };
-
-    // 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 };
-  }
-
-  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(pending);
-    const newLLMMessages = tx.newMessages();
-
-    const messageListError = (error: string) => {
-      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) {
-      messageListError("newSessionMessages.length !== numPending");
-    }
-    if (newLLMMessages.length !== numPending) {
-      messageListError("newLLMMessages.length !== numPending");
-    }
-
-    // The SessionMessages should satisfy:
-    //  - sMsg.message_idx === pMsg.message_idx
-    //  - sMsg.content === llmMsg
-    // this ensures all representations are aligned.
-
-    for (let i = 0; i < numPending; ++i) {
-      const sMsg = newSessionMessages[i];
-      const pMsg = pending[i];
-      const lMsg = newLLMMessages[i];
-
-      if (sMsg.content.role !== lMsg.role) {
-        messageListError(
-          `newSessionMessages[${String(i)}].role !== ` +
-            `newLLMMessages[${String(i)}].role`
-        );
-      }
-      if (JSON.stringify(sMsg.content) !== JSON.stringify(lMsg)) {
-        messageListError(
-          `newSessionMessages[${String(i)}].content !== ` +
-            `newLLMMessages[${String(i)}]`
-        );
-      }
-      if (sMsg.message_idx !== pMsg.message_idx) {
-        messageListError(
-          `newSessionMessages[${String(i)}].message_idx !== ` +
-            `pendingMessages[${String(i)}].message_idx`
-        );
-      }
-    }
-
-    // Update our internal state and return the SessionMessages to write to
-    // the DB
-
-    await this.llmContext.commit(tx.getBaseTx());
-    this.conversationMessages.push(...pending);
-
-    // Kick off a compression?
-    this.checkCompression();
-
-    return newSessionMessages;
-  }
-
-  private lastCommittedMessageIdx(): number | undefined {
-    const numMsgs = this.conversationMessages.length;
-    if (numMsgs > 0) {
-      return this.conversationMessages[numMsgs - 1].message_idx;
-    }
-    return undefined;
-  }
-
-  private getNextMessageIdx(): number {
-    const idx = this.nextMessageIdx;
-    this.nextMessageIdx += MESSAGE_INDEX_FULL_INCREMENT;
-    return idx;
-  }
-
-  /**
-   * Only reuse it if no other indices have been allocated.  It's the callers
-   * responsibility to ensure this.
-   */
-  private freeMessageIdx(messageIdx: number) {
-    assert(
-      messageIdx === this.nextMessageIdx - MESSAGE_INDEX_FULL_INCREMENT,
-      "message idx cannot be free-ed"
-    );
-    this.nextMessageIdx = messageIdx;
-  }
-
-  private checkCompression(): void {
-    if (this.pendingCompression) {
-      return;
-    }
-
-    // TODO: track tokens and use that to trigger compression
-
-    const numCommitted = this.llmContext.numMessages();
-    if (numCommitted < COMPRESSION_TRIGGER_NUM_MESSAGES) {
-      return;
-    }
-
-    void this.runCompression();
-  }
-
-  private async runCompression() {
-    assert(!this.pendingCompression);
-
-    this.pendingCompression = true;
-    const checkpointIndex =
-      this.conversationMessages[this.conversationMessages.length - 1]
-        .message_idx;
-
-    try {
-      const summary = await this.llmContext.compress();
-      const checkpoint: SessionCheckpoint = {
-        message_idx: checkpointIndex,
-        summary,
-      };
-      await this.checkpointWriter.writeCheckpoint(checkpoint);
-    } catch (err: unknown) {
-      logger.warn(
-        `[runCompression] error during compression: ${getErrorString(err)}`
-      );
-    } finally {
-      this.pendingCompression = false;
-    }
-  }
-}
-
-export function resolveConversationWithCheckpoint(
-  messages: SessionMessage[],
-  checkpoint: SessionCheckpoint | undefined
-): {
-  messages: MessageParam[];
-  lastEntryIdx: number;
-} {
-  const numMessages = messages.length;
-  let lastEntryIdx = 0;
-  if (numMessages !== 0) {
-    lastEntryIdx = messages[numMessages - 1].message_idx;
-  }
-
-  // Only keep the messages we care about
-
-  messages = messages.filter((msg) => msg.is_for_llm);
-
-  // If no checkpoint, return all messages
-
-  if (!checkpoint) {
-    return {
-      messages: messages.map((msg) => msg.content),
-      lastEntryIdx,
-    };
-  }
-
-  // Find the first entry in messages s.t. entry.message_idx >
-  // checkpoint.message_idx
-  const checkpointIdx = checkpoint.message_idx;
-  const keepIdx = (() => {
-    let idx = numMessages - 1;
-    while (idx >= 0) {
-      if (messages[idx].message_idx <= checkpointIdx) {
-        break;
-      }
-      idx--;
-    }
-    // Idx points to the first entry (from the end) with message_idx <=
-    // checkpointIdx, so the first entry with message_idx > checkpointIdx is
-    // idx+1 and We can then extract the required messages with
-    //
-    //   message.slice(idx+1)
-    //
-    // Edge-cases:
-    //
-    // - loop exhausted, idx = -1, we want slice(0 = idx+1) (all messages)
-    //
-    // - loop exited on first entry (all entries are contained in the
-    //   checkpoint), idx = numMessages - 1, we want
-    //   slice(numMessages = idx + 1)
-
-    return idx + 1;
-  })();
-
-  const checkpointMessage = createCheckpointMessage(checkpoint.summary);
-  const llmMessages: MessageParam[] = [
-    checkpointMessage,
-    ...messages.slice(keepIdx).map((msg) => msg.content),
-  ];
-
-  return {
-    messages: llmMessages,
-    lastEntryIdx: Math.max(lastEntryIdx, checkpoint.message_idx),
-  };
-}