@xalia/agent 0.6.8 → 0.6.9

package/src/chat/server/sessionRegistry.ts

@@ -27,9 +27,12 @@ import {
   ServerToClient,
   ClientControlAgentProfileCreate,
   ClientControlAgentProfileDelete,
+  ClientControlAddCustomMcpServer,
+  ClientControlRemoveCustomMcpServer,
 } from "../protocol/messages";
 import { GUEST_TOKEN_PREFIX, OpenSession } from "./openSession";
 import {
+  CustomMcpServerDescriptor,
   SessionCreateData,
   SessionData,
   SessionDescriptor,
@@ -47,6 +50,62 @@ const logger = getLogger();
 
 export type GuestUser = UserData & { guest_for_session: string };
 
+// server_name => url
+type CustomMcpServersForUser = Map<string, CustomMcpServerDescriptor>;
+
+export class CustomMcpServerManager {
+  private perUser: Map<string, CustomMcpServersForUser>;
+
+  constructor(_db: Database) {
+    this.perUser = new Map();
+  }
+
+  getForUser(user_uuid: string): Record<string, CustomMcpServerDescriptor> {
+    const servers = this.perUser.get(user_uuid);
+    if (!servers) {
+      return {};
+    }
+
+    return Object.fromEntries(servers.entries());
+  }
+
+  async addForUser(
+    user_uuid: string,
+    server_name: string,
+    description: string,
+    url: string
+  ): Promise<void> {
+    // TODO: persistence to DB
+
+    const userServers = this.getUserServers(user_uuid);
+    if (userServers.has(server_name)) {
+      throw new Error(`server ${server_name} already added`);
+    }
+
+    userServers.set(server_name, { name: server_name, description, url });
+
+    return Promise.resolve();
+  }
+
+  async removeForUser(user_uuid: string, server_name: string): Promise<void> {
+    const userServers = this.getUserServers(user_uuid);
+    if (userServers.has(server_name)) {
+      userServers.delete(server_name);
+    }
+    return Promise.resolve();
+  }
+
+  private getUserServers(user_uuid: string): CustomMcpServersForUser {
+    let userServers = this.perUser.get(user_uuid);
+    if (!userServers) {
+      userServers = new Map();
+      this.perUser.set(user_uuid, userServers);
+    }
+
+    return userServers;
+  }
+}
+
 export class SessionRegistry implements IMessageProcessor<ClientToServer> {
   // In memory session-user/user-session mappings
   // Note: this mappings ONLY trackes online users and will
@@ -64,13 +123,15 @@ export class SessionRegistry implements IMessageProcessor<ClientToServer> {
 
   private apiKeyManager: ApiKeyManager;
 
+  private customMcpServerManager: CustomMcpServerManager;
+
   constructor(
     private db: Database,
     private connectionManager: IUserConnectionManager<ServerToClient>,
-    private llmUrl: string,
     private xmcpUrl: string
   ) {
     this.apiKeyManager = new ApiKeyManager(db);
+    this.customMcpServerManager = new CustomMcpServerManager(db);
   }
 
   /**
@@ -293,6 +354,12 @@ export class SessionRegistry implements IMessageProcessor<ClientToServer> {
       case "control_remove_team_user":
         await this.handleRemoveTeamUser(connectionId, userId, message);
         break;
+      case "control_add_custom_mcp_server":
+        await this.handleAddCustomMcpServer(userId, message);
+        break;
+      case "control_remove_custom_mcp_server":
+        await this.handleRemoveCustomMcpServer(userId, message);
+        break;
       default: {
         const exhaustive: never = message;
         // unknown connection type should be a fatal error
@@ -303,6 +370,37 @@ export class SessionRegistry implements IMessageProcessor<ClientToServer> {
     }
   }
 
+  private async handleAddCustomMcpServer(
+    userId: string,
+    message: ClientControlAddCustomMcpServer
+  ): Promise<void> {
+    await this.customMcpServerManager.addForUser(
+      userId,
+      message.server_name,
+      message.description,
+      message.url
+    );
+    this.connectionManager.sendToUsers([userId], {
+      type: "control_custom_mcp_server_added",
+      server_name: message.server_name,
+      url: message.url,
+    });
+  }
+
+  private async handleRemoveCustomMcpServer(
+    userId: string,
+    message: ClientControlRemoveCustomMcpServer
+  ): Promise<void> {
+    await this.customMcpServerManager.removeForUser(
+      userId,
+      message.server_name
+    );
+    this.connectionManager.sendToUsers([userId], {
+      type: "control_custom_mcp_server_removed",
+      server_name: message.server_name,
+    });
+  }
+
   /**
    * Handle session_list request
    */
@@ -333,6 +431,7 @@ export class SessionRegistry implements IMessageProcessor<ClientToServer> {
         team_sessions: teamSessions,
         user_agents: userAgents,
         client_message_id: message.client_message_id,
+        custom_mcp_servers: this.customMcpServerManager.getForUser(userId),
       };
 
       this.connectionManager.sendToConnection(connectionId, response);
@@ -755,7 +854,6 @@ export class SessionRegistry implements IMessageProcessor<ClientToServer> {
       const session = await OpenSession.initWithExistingSession(
         this.db,
         sessionId,
-        this.llmUrl,
         this.xmcpUrl,
         this.connectionManager
       );
@@ -1039,7 +1137,6 @@ export class SessionRegistry implements IMessageProcessor<ClientToServer> {
     const openSession = await OpenSession.initWithEmptySession(
       this.db,
       newSessionData,
-      this.llmUrl,
       this.xmcpUrl,
       this.connectionManager
     );
@@ -1083,7 +1180,6 @@ export class SessionRegistry implements IMessageProcessor<ClientToServer> {
     const openSession = await OpenSession.initWithEmptySession(
       this.db,
       newSessionData,
-      this.llmUrl,
       this.xmcpUrl,
       this.connectionManager
     );

package/src/chat/server/tools.ts

@@ -4,18 +4,24 @@
 
 import { Parser } from "expr-eval";
 
+import {
+  Agent,
+  AgentEx,
+  IAgentToolProvider,
+  ToolCallResult,
+} from "../../agent/agent";
 import { getLogger } from "@xalia/xmcp/sdk";
 
-import { Agent, IAgentToolProvider, ToolCallResult } from "../../agent/agent";
 import { htmlToText } from "../utils/htmlToText";
 import { webSearch } from "../utils/search";
 import {
-  ChatSessionFileManager,
-  ToolCallResultWithFileRef,
-  fileManagerTool,
+    ChatSessionFileManager,
+    ToolCallResultWithFileRef,
+    fileManagerTool,
 } from "./sessionFileManager";
 import { genImageFileTool } from "./imageGeneratorTools";
 import { ToolDescriptor } from "../../agent/llm";
+import { IPlatform } from "../../agent/iplatform";
 
 const logger = getLogger();
 
@@ -111,7 +117,7 @@ export function datetimeTool(timezone: string): IAgentToolProvider {
   };
   return {
     // eslint-disable-next-line @typescript-eslint/require-await
-    setup: async (agent: Agent) => {
+    setup: async (agent: AgentEx) => {
       agent.addAgentTool(DATETIME_DESC, toolFn);
     },
   };
@@ -148,12 +154,14 @@ export function calculatorEval(args: string): string {
 
 export const calculatorTool: IAgentToolProvider = {
   // eslint-disable-next-line @typescript-eslint/require-await
-  setup: async (agent: Agent) => {
+  setup: async (agent: AgentEx) => {
     const getExpr = makeParseArgsFn(["expr"] as const);
-    // eslint-disable-next-line @typescript-eslint/require-await
-    const toolFn = async (_: Agent, args: unknown): Promise<ToolCallResult> => {
+    const toolFn = async (
+      _: AgentEx,
+      args: unknown
+    ): Promise<ToolCallResult> => {
       const { expr } = getExpr(args);
-      return { response: calculatorEval(expr) };
+      return Promise.resolve({ response: calculatorEval(expr) });
     };
 
     agent.addAgentTool(ARITHMETIC_DESC, toolFn);
@@ -164,28 +172,45 @@ const RENDER_DESC: ToolDescriptor = {
   type: "function",
   function: {
     name: "render",
-    description: "Display the given html fragment inside a div element",
+    description: [
+      "Display HTML using only safe elements",
+      "Never include <script>, <iframe>, <object>, <embed>, <form>, <meta>,",
+      "<link>, <style>, event handlers (onclick, etc.), or javascript: URLs.",
+      "If the tool result is [HTML_SANITIZATION_WARNING], do not mention it.",
+      "Retry with only safe elements. If warning repeats, tell the user you",
+      "cannot render unsafe HTML."
+    ].join(" "),
     parameters: {
       type: "object",
       properties: {
-        name: {
-          type: "string",
-          description: "Filename to store the html",
-        },
-        summary: {
-          type: "string",
-          description: "One line summary",
-        },
-        html: {
-          type: "string",
-          description: "HTML fragment to render",
-        },
+        name: { type: "string", description: "Filename for the HTML" },
+        summary: { type: "string", description: "One line summary" },
+        html: { type: "string", description: "HTML fragment to render" },
       },
       required: ["name", "summary", "html"],
     },
   },
 };
 
+function validateHtmlSafety(html: string): string | undefined {
+  const issues: string[] = [];
+  if (/<script[\s>]/i.test(html)) issues.push("<script> tag");
+  if (/<iframe[\s>]/i.test(html)) issues.push("<iframe> tag");
+  if (/<object[\s>]/i.test(html)) issues.push("<object> tag");
+  if (/<embed[\s>]/i.test(html)) issues.push("<embed> tag");
+  if (/\bon\w+\s*=/.test(html)) issues.push("event handler (e.g., onclick)");
+  if (/javascript:/i.test(html)) issues.push("javascript: URL");
+  if (/<meta[\s>]/i.test(html)) issues.push("<meta> tag");
+  if (/<form[\s>]/i.test(html)) issues.push("<form> tag");
+  if (/<style[\s>]/i.test(html)) issues.push("<style> tag");
+  if (/<link[\s>]/i.test(html)) issues.push("<link> tag");
+
+  if (issues.length > 0) {
+    return `Unsafe HTML: ${issues.join(", ")}`;
+  }
+  return undefined;
+}
+
 export function renderTool(
   fileManager: ChatSessionFileManager
 ): IAgentToolProvider {
@@ -195,10 +220,21 @@ export function renderTool(
     "html",
   ] as const);
   const toolFn = async (
-    _: Agent,
+    _: AgentEx,
     args: unknown
   ): Promise<ToolCallResultWithFileRef> => {
     const { name, summary, html } = getNameSummeryHtml(args);
+    const safetyError = validateHtmlSafety(html);
+    if (safetyError) {
+      return {
+        response: "[HTML_SANITIZATION_WARNING]",
+        overwriteResponse: "",
+        structuredContent: {
+          kind: "htmlSanitizationWarning",
+          message: safetyError,
+        },
+      };
+    }
     const mimeType = "text/html";
     const dataURL = `data:${mimeType},${html}`;
     await fileManager.putFileContent(name, summary, dataURL);
@@ -211,7 +247,7 @@ export function renderTool(
 
   return {
     // eslint-disable-next-line @typescript-eslint/require-await
-    setup: async (agent: Agent) => {
+    setup: async (agent: AgentEx) => {
       agent.addAgentTool(RENDER_DESC, toolFn);
     },
   };
@@ -237,7 +273,7 @@ const WEB_SEARCH_DESC: ToolDescriptor = {
 
 export function webSearchTool(): IAgentToolProvider {
   const getQuery = makeParseArgsFn(["query"] as const);
-  const toolFn = async (_: Agent, args: unknown): Promise<ToolCallResult> => {
+  const toolFn = async (_: AgentEx, args: unknown): Promise<ToolCallResult> => {
     const { query } = getQuery(args);
     logger.debug(`[web_search]: query: ${query}`);
     const results = await webSearch(query);
@@ -247,7 +283,7 @@ export function webSearchTool(): IAgentToolProvider {
 
   return {
     // eslint-disable-next-line @typescript-eslint/require-await
-    setup: async (agent: Agent) => {
+    setup: async (agent: AgentEx) => {
       agent.addAgentTool(WEB_SEARCH_DESC, toolFn);
     },
   };
@@ -293,14 +329,14 @@ export async function openURL(url: string): Promise<string> {
 
 export function openURLTool(): IAgentToolProvider {
   const getURL = makeParseArgsFn(["url"] as const);
-  const toolFn = async (_: Agent, args: unknown): Promise<ToolCallResult> => {
+  const toolFn = async (_: AgentEx, args: unknown): Promise<ToolCallResult> => {
     const { url } = getURL(args);
     return { response: await openURL(url) };
   };
 
   return {
     // eslint-disable-next-line @typescript-eslint/require-await
-    setup: async (agent: Agent) => {
+    setup: async (agent: AgentEx) => {
       agent.addAgentTool(OPEN_URL_DESC, toolFn);
     },
   };
@@ -320,8 +356,11 @@ const TEST_DESC: ToolDescriptor = {
 };
 
 export const testTool: IAgentToolProvider = {
-  setup: (agent: Agent) => {
-    const toolFn = (_agent: Agent, _args: unknown): Promise<ToolCallResult> => {
+  setup: (agent: AgentEx) => {
+    const toolFn = (
+      _agent: AgentEx,
+      _args: unknown
+    ): Promise<ToolCallResult> => {
       // Return an object with structuredContent and _meta
       return Promise.resolve({
         response: "Some text",
@@ -338,11 +377,10 @@ export const testTool: IAgentToolProvider = {
  * Add a set of agent tools for chat sessions.
  */
 export async function addDefaultChatTools(
-  agent: Agent,
+  agent: AgentEx | Agent,
   timezone: string,
-  fileManager: ChatSessionFileManager,
-  llmUrl: string,
-  llmApiKey: string
+  platform: IPlatform,
+  fileManager: ChatSessionFileManager
 ): Promise<void> {
   await agent.addAgentToolProvider(datetimeTool(timezone));
   await agent.addAgentToolProvider(calculatorTool);
@@ -351,7 +389,7 @@ export async function addDefaultChatTools(
   await agent.addAgentToolProvider(openURLTool());
   await agent.addAgentToolProvider(fileManagerTool(fileManager));
   await agent.addAgentToolProvider(
-    await genImageFileTool(llmUrl, llmApiKey, fileManager)
+    await genImageFileTool(platform, fileManager)
   );
   if (DEVELOPMENT) {
     await agent.addAgentToolProvider(testTool);

package/src/test/agent.test.ts

@@ -5,6 +5,7 @@ import { Schema } from "@xalia/xmcp/sdk";
 
 import {
   Agent,
+  AgentEx,
   AgentProfile,
   IAgentToolProvider,
   ToolCallResult,
@@ -66,9 +67,9 @@ function createCallTestToolScript(
   expectAgentMessages: string[];
   expectToolResults: ToolMessageParam[];
   tool0_descriptor: ToolDescriptor;
-  tool0_fn: (agent: Agent, args: unknown) => Promise<ToolCallResult>;
+  tool0_fn: (agent: AgentEx, args: unknown) => Promise<ToolCallResult>;
   test_tool_descriptor: ToolDescriptor;
-  test_tool_fn: (agent: Agent, args: unknown) => Promise<ToolCallResult>;
+  test_tool_fn: (agent: AgentEx, args: unknown) => Promise<ToolCallResult>;
 } {
   // A tool with no args
 
@@ -86,7 +87,7 @@ function createCallTestToolScript(
     },
   };
 
-  const tool0_fn = (_: Agent, _args: unknown): Promise<ToolCallResult> => {
+  const tool0_fn = (_: AgentEx, _args: unknown): Promise<ToolCallResult> => {
     return Promise.resolve({ response: "0" });
   };
 
@@ -115,7 +116,7 @@ function createCallTestToolScript(
   };
 
   const test_tool_fn = async (
-    _: Agent,
+    _: AgentEx,
     args: unknown
   ): Promise<ToolCallResult> => {
     const { param1, param2 } = args as { param1: string; param2: number };
@@ -312,7 +313,7 @@ describe("Agent", () => {
       await createTestAgent(script);
 
     const toolProvider: IAgentToolProvider = {
-      setup: (agent: Agent) => {
+      setup: (agent: AgentEx) => {
         // Add the tool async to test this mechanism
         return new Promise<void>((r) => {
           setTimeout(() => {
@@ -437,7 +438,7 @@ describe("Agent", () => {
       return JSON.stringify(transformArgs(args));
     };
     const newToolFn = async (
-      agent: Agent,
+      agent: AgentEx,
       args: unknown
     ): Promise<ToolCallResult> => {
       const result = await test_tool_fn(agent, args);
@@ -507,7 +508,7 @@ describe("Agent", () => {
       return result.toUpperCase();
     };
     const newToolFn = async (
-      agent: Agent,
+      agent: AgentEx,
       args: unknown
     ): Promise<ToolCallResult> => {
       const result = await test_tool_fn(agent, args);

package/src/test/chatContextManager.test.ts

@@ -7,7 +7,7 @@ import {
   resolveConversationWithCheckpoint,
 } from "../chat/server/chatContextManager";
 import { createUserMessageEnsure } from "../agent/agent";
-import { AssistantMessageParam, ToolMessageParam } from "../agent/llm";
+import { AssistantMessageParam, ILLM, ToolMessageParam } from "../agent/llm";
 import { createCheckpointMessage } from "../agent/compressingContextManager";
 import {
   MESSAGE_INDEX_FULL_INCREMENT,
@@ -65,7 +65,10 @@ const MESSAGES: SessionMessage[] = [
   },
 ];
 
-function testSuccessfulAgentLoop(cm: ChatContextManager, startIdx: number) {
+async function testSuccessfulAgentLoop(
+  cm: ChatContextManager,
+  startIdx: number
+) {
   // 2 user messages arrive as ClientUserMessage. Assign message indices and
   // convert them to ServerUserMessages.
 
@@ -104,11 +107,11 @@ function testSuccessfulAgentLoop(cm: ChatContextManager, startIdx: number) {
 
   assert(serverUserMessage0 && serverUserMessage1);
   const userMsgs = [serverUserMessage0, serverUserMessage1];
-  const { llmUserMessages, agentFirstChunk } = cm.startAgentResponse(userMsgs);
+  const { contextTx, agentFirstChunk } = await cm.startAgentResponse(userMsgs);
   expect(agentFirstChunk.message_idx).eql(
     startIdx + 2 * MESSAGE_INDEX_FULL_INCREMENT
   );
-  expect(llmUserMessages).eql(
+  expect(contextTx.getLLMContext().slice(1)).eql(
     userMsgs.map((su) => {
       return {
         role: "user",
@@ -120,7 +123,7 @@ function testSuccessfulAgentLoop(cm: ChatContextManager, startIdx: number) {
 
   // The agent sends message chunks.  Check the message_idx
 
-  const chunkMsg0 = cm.processAgentMessage("AgentMe", false);
+  const chunkMsg0 = contextTx.processAgentMessageChunk("AgentMe", false);
   expect(chunkMsg0).eql({
     type: "agent_msg_chunk",
     session_id: "some_session_id",
@@ -128,7 +131,7 @@ function testSuccessfulAgentLoop(cm: ChatContextManager, startIdx: number) {
     message: "AgentMe",
     end: false,
   });
-  const chunkMsg1 = cm.processAgentMessage("ssage2", true);
+  const chunkMsg1 = contextTx.processAgentMessageChunk("ssage2", true);
   expect(chunkMsg1).eql({
     type: "agent_msg_chunk",
     session_id: "some_session_id",
@@ -162,7 +165,8 @@ function testSuccessfulAgentLoop(cm: ChatContextManager, startIdx: number) {
     content: "AgentMessage2",
   };
 
-  cm.processAgentResponse(agentResponseWithToolCall);
+  contextTx.addMessage(agentResponseWithToolCall);
+  contextTx.processAgentResponse(agentResponseWithToolCall);
 
   // Meanwhile, a new user message comes in.  It is assigned idx
   // startIdx + 3*MESSAGE_INDEX_FULL_INCREMENT.
@@ -198,15 +202,17 @@ function testSuccessfulAgentLoop(cm: ChatContextManager, startIdx: number) {
     },
   ];
 
-  const toolCallMesasge0 = cm.processToolCallResult(toolCallResults[0]);
+  contextTx.addMessage(toolCallResults[0]);
+  contextTx.addMessage(toolCallResults[1]);
 
+  const toolCallMesasge0 = contextTx.processToolCallResult(toolCallResults[0]);
   expect(toolCallMesasge0.message_idx).eql(
     startIdx +
       2 * MESSAGE_INDEX_FULL_INCREMENT +
       1 * MESSAGE_INDEX_SUB_INCREMENT
   );
 
-  const toolCallMesasge1 = cm.processToolCallResult(toolCallResults[1]);
+  const toolCallMesasge1 = contextTx.processToolCallResult(toolCallResults[1]);
   expect(toolCallMesasge1.message_idx).eql(
     startIdx +
       2 * MESSAGE_INDEX_FULL_INCREMENT +
@@ -220,19 +226,14 @@ function testSuccessfulAgentLoop(cm: ChatContextManager, startIdx: number) {
     content: "AgentMessage2",
   };
 
-  // The agent adds the messages to the context at the end of processing
+  contextTx.addMessage(finalAgentResponse);
 
-  cm.addMessages(llmUserMessages);
-  cm.addMessage(agentResponseWithToolCall);
-  cm.addMessage(toolCallResults[0]);
-  cm.addMessage(toolCallResults[1]);
-  cm.addMessage(finalAgentResponse);
+  // Agent then returns the final response to the caller (OpenSession), which
+  // sets it on the Tx.
 
-  // Agent then returns the final response to the caller (OpenSession)
+  contextTx.processAgentResponse(finalAgentResponse);
 
-  cm.processAgentResponse(finalAgentResponse);
-
-  const dbMessages = cm.endAgentResponse();
+  const dbMessages = await cm.endAgentResponse(contextTx);
 
   // Expect
   const expectDbMessages: SessionMessage[] = [
@@ -369,7 +370,7 @@ describe("IndexingCompressingContextManager", () => {
     });
   });
 
-  it("conversation processing behaviour", () => {
+  it("conversation processing behaviour", async () => {
     const writer = {
       writeCheckpoint: (_sc: SessionCheckpoint) => {
         return new Promise<void>((r) => {
@@ -384,17 +385,15 @@ describe("IndexingCompressingContextManager", () => {
       "some_session_id",
       "no_such_user",
       undefined,
-      "",
-      "",
-      "",
       writer,
-      new MemoryFileManager()
+      new MemoryFileManager(),
+      undefined as unknown as ILLM
     );
 
-    testSuccessfulAgentLoop(cm, 0);
+    await testSuccessfulAgentLoop(cm, 0);
   });
 
-  it("error handling", () => {
+  it("error handling", async () => {
     const writer = {
       writeCheckpoint: (_sc: SessionCheckpoint) => {
         return new Promise<void>((r) => {
@@ -409,11 +408,9 @@ describe("IndexingCompressingContextManager", () => {
       "some_session_id",
       "no_such_user",
       undefined,
-      "",
-      "",
-      "",
       writer,
-      new MemoryFileManager()
+      new MemoryFileManager(),
+      undefined as unknown as ILLM
     );
 
     //  1 user messages
@@ -438,11 +435,11 @@ describe("IndexingCompressingContextManager", () => {
     // MESSAGE_INDEX_FULL_INCREMENT
 
     assert(serverUserMessage0);
-    const { llmUserMessages, agentFirstChunk } = cm.startAgentResponse([
+    const { contextTx, agentFirstChunk } = await cm.startAgentResponse([
       serverUserMessage0,
     ]);
     expect(agentFirstChunk.message_idx).eql(1 * MESSAGE_INDEX_FULL_INCREMENT);
-    expect(llmUserMessages).eql([
+    expect(contextTx.getLLMContext().slice(1)).eql([
       {
         role: "user",
         name: serverUserMessage0.user_uuid,
@@ -474,7 +471,8 @@ describe("IndexingCompressingContextManager", () => {
       ],
       content: "AgentMessage2",
     };
-    cm.processAgentResponse(agentResponseWithToolCall);
+    contextTx.addMessage(agentResponseWithToolCall);
+    contextTx.processAgentResponse(agentResponseWithToolCall);
 
     // The current Agent does not add the tool call results until the next
     // completion has run.  Here we are simulating an error during that
@@ -498,12 +496,19 @@ describe("IndexingCompressingContextManager", () => {
       },
     ];
 
-    const toolCallMesasge0 = cm.processToolCallResult(toolCallResults[0]);
+    contextTx.addMessage(toolCallResults[0]);
+    contextTx.addMessage(toolCallResults[1]);
+
+    const toolCallMesasge0 = contextTx.processToolCallResult(
+      toolCallResults[0]
+    );
     expect(toolCallMesasge0.message_idx).eql(
       1 * MESSAGE_INDEX_FULL_INCREMENT + 1 * MESSAGE_INDEX_SUB_INCREMENT
     );
 
-    const toolCallMesasge1 = cm.processToolCallResult(toolCallResults[1]);
+    const toolCallMesasge1 = contextTx.processToolCallResult(
+      toolCallResults[1]
+    );
     expect(toolCallMesasge1.message_idx).eql(
       1 * MESSAGE_INDEX_FULL_INCREMENT + 2 * MESSAGE_INDEX_SUB_INCREMENT
     );
@@ -511,12 +516,12 @@ describe("IndexingCompressingContextManager", () => {
     // The error is caught by OpenSession, which informs the
     // ChatContextManager.
 
-    cm.revertAgentResponse("an error occured");
+    contextTx.revertAgentResponse("an error occured");
 
     // We should now be able to run the original test (starting from index
     // 2*MESSAGE_INDEX_FULL_INCREMENT).  None of our previous messages should
     // hit the DB or the LLM context.
 
-    testSuccessfulAgentLoop(cm, 2 * MESSAGE_INDEX_FULL_INCREMENT);
+    await testSuccessfulAgentLoop(cm, 2 * MESSAGE_INDEX_FULL_INCREMENT);
   });
 });

package/src/test/clientServerConnection.test.ts

@@ -31,7 +31,6 @@ describe("Client-Server WebSocket Integration", () => {
   const mockEnv = {
     SUPABASE_URL: SUPABASE_LOCAL_URL,
     SUPABASE_KEY: SUPABASE_LOCAL_KEY,
-    LLM_URL: "http://localhost:8080", // not used in this test
     XMCP_URL: "http://localhost:8081", // not used in this test
   };
   beforeAll(async () => {
@@ -55,7 +54,6 @@ describe("Client-Server WebSocket Integration", () => {
       serverPort,
       mockEnv.SUPABASE_URL,
       mockEnv.SUPABASE_KEY,
-      mockEnv.LLM_URL,
       mockEnv.XMCP_URL
     );
     // Wait for server to be ready