@xalia/agent 0.6.8 → 0.6.9

package/src/test/compressingContextManager.test.ts

@@ -4,50 +4,45 @@ import {
   CompressingContextManager,
   createCheckpointMessage,
 } from "../agent/compressingContextManager";
-import { MessageParam } from "../agent/llm";
-
-const MESSAGES: MessageParam[] = [
-  { role: "user", content: "msg200" },
-  { role: "assistant", content: "msg300" },
-  { role: "user", content: "msg400" },
-  {
-    role: "assistant",
-    content: "msg500",
-    tool_calls: [
-      {
-        id: "tool_call_0",
-        type: "function",
-        function: { name: "tool1", arguments: "" },
-      },
-    ],
-  },
-  {
-    role: "tool",
-    content: "msg501",
-    tool_call_id: "tool_call_0",
-  },
-  { role: "assistant", content: "msg502" },
-];
+import { ILLM, MessageParam, UserMessageParam } from "../agent/llm";
+import { MESSAGES, basicContextTest } from "./contextTestTools";
+import { createLLM } from "../agent/agentUtils";
+import { IPlatform } from "../agent/iplatform";
 
 describe("Compression context", () => {
-  it("context compression", async function () {
+  it("context compression (basic)", async function () {
     const ccm = new CompressingContextManager(
       "sys_prompt",
       [],
-      "", // llmUrl
-      "repeat", // llmModel
-      "" // llmApiKey
+      () => undefined as unknown as Promise<ILLM> // getLLM
+    );
+    await basicContextTest(ccm);
+  });
+
+  it("context compression", async function () {
+    const llm = await createLLM(
+      "",
+      "",
+      "repeat",
+      false,
+      undefined as unknown as IPlatform
+    );
+    const ccm = new CompressingContextManager("sys_prompt", [], () =>
+      Promise.resolve(llm)
     );
 
     // Add first 2 messages
-    ccm.addMessages(MESSAGES.slice(0, 2));
-    ccm.commit();
+    let tx = await ccm.startTx([MESSAGES[0] as UserMessageParam]);
+    tx.addMessage(MESSAGES[1]);
+    await ccm.commit(tx);
 
     // Start the compression
     const compressionP = ccm.compress();
 
     // Add more messages and update system prompt
-    ccm.addMessages(MESSAGES.slice(2));
+    tx = await ccm.startTx([MESSAGES[2] as UserMessageParam]);
+    tx.addMessages(MESSAGES.slice(3));
+    await ccm.commit(tx);
     ccm.setAgentPrompt("New system prompt");
 
     // Wait for compression to complete, then get the new context
@@ -74,9 +69,9 @@ describe("Compression context", () => {
     expect(summary).eql(expectSummary);
     expect(context).eql(expectContext);
 
-    // Pending messages shoudl still be intact
+    // // Pending messages should still be intact
 
-    const pending = ccm.getPending();
-    expect(pending).eql(MESSAGES.slice(2));
+    // const pending = ccm.getPending();
+    // expect(pending).eql(MESSAGES.slice(2));
   });
 });

package/src/test/context.test.ts

@@ -1,10 +1,12 @@
 import { expect } from "vitest";
 import { strict as assert } from "assert";
 
-import { ContextManagerWithWorkspace } from "../agent/context";
+import { ContextManagerWithWorkspace } from "../agent/contextWithWorkspace";
 import { createUserMessage } from "../agent/agent";
-import { MessageParam } from "../agent/llm";
+import { MessageParam, UserMessageParam } from "../agent/llm";
 import { SystemPromptProvider } from "../agent/promptProvider";
+import { ContextManager } from "../agent/context";
+import { basicContextTest } from "./contextTestTools";
 
 describe("ContextManagerWithWorkspace", () => {
   const msgs: MessageParam[] = [
@@ -14,34 +16,61 @@ describe("ContextManagerWithWorkspace", () => {
     { role: "assistant", content: "message D" },
   ];
 
-  it("should correctly manage messages", function () {
+  it("context (basic)", async function () {
+    const cm = new ContextManager("sys_prompt", []);
+    await basicContextTest(cm);
+  });
+
+  it("context with workspace (basic)", async function () {
+    const cm = new ContextManagerWithWorkspace("sys_prompt", []);
+    await basicContextTest(cm);
+  });
+
+  it("should correctly manage messages", async function () {
     const cm = new ContextManagerWithWorkspace("sys_msg", msgs.slice(0, 2));
 
-    // set workspace and add new messages
+    // set workspace and add new user and assistant messages
     cm.setWorkspace({ role: "user", content: "workspace1" });
-    cm.addMessages(msgs.slice(2));
 
-    // context with workspace
-    const ctx1 = cm.getLLMContext();
-    expect(ctx1).eql([
+    const ctxTx1 = await cm.startTx([msgs[2] as UserMessageParam]);
+    ctxTx1.addMessage(msgs[3]);
+
+    // context writer.  workspace should be before first assistant message
+    expect(ctxTx1.getLLMContext()).eql([
       { role: "system", content: "sys_msg" },
-      ...msgs,
+      ...msgs.slice(0, 3),
+      { role: "user", content: "workspace1" },
+      ...msgs.slice(3),
+    ]);
+
+    // original cm unchanged
+    expect(cm.getLLMContext()).eql([
+      { role: "system", content: "sys_msg" },
+      ...msgs.slice(0, 2),
       { role: "user", content: "workspace1" },
     ]);
 
-    // context with new message and workspace
+    // Commit
+    await cm.commit(ctxTx1);
+
+    // write new messages
     const lastMsg: MessageParam = {
       role: "user",
       content: "hello",
     };
-    cm.addMessage(lastMsg);
-    const ctx2 = cm.getLLMContext();
-    expect(ctx2).eql([
+    const ctxTx2 = await cm.startTx([lastMsg]);
+    expect(ctxTx2.getLLMContext()).eql([
       { role: "system", content: "sys_msg" },
       ...msgs,
       lastMsg,
       { role: "user", content: "workspace1" },
     ]);
+    expect(cm.getLLMContext()).eql([
+      { role: "system", content: "sys_msg" },
+      ...msgs,
+      { role: "user", content: "workspace1" },
+    ]);
+    await cm.commit(ctxTx2);
 
     // context without workspace
     cm.setWorkspace(undefined);
@@ -72,7 +101,7 @@ describe("ContextManagerWithWorkspace", () => {
     expect(llmContext[0].content).eql("global\nagent_prompt\nFRAG1");
   });
 
-  it("handles rich workspace messages", function () {
+  it("handles rich workspace messages", async function () {
     const cm = new ContextManagerWithWorkspace("sys_msg", msgs.slice(0, 2));
 
     expect(cm.getLLMContext()).eql([
@@ -97,7 +126,22 @@ describe("ContextManagerWithWorkspace", () => {
       wsMsg,
     ]);
 
-    cm.addMessages(msgs.slice(2));
+    const ctxTx = await cm.startTx([msgs[2] as UserMessageParam]);
+    expect(ctxTx.getLLMContext()).eql([
+      { role: "system", content: "sys_msg" },
+      ...msgs.slice(0, 3),
+      wsMsg,
+    ]);
+    ctxTx.addMessage(msgs[3]);
+    expect(ctxTx.getLLMContext()).eql([
+      { role: "system", content: "sys_msg" },
+      ...msgs.slice(0, 3),
+      wsMsg,
+      msgs[3],
+    ]);
+    expect(ctxTx.newMessages()).eql(msgs.slice(2, 4));
+
+    await cm.commit(ctxTx);
 
     // context with workspace
     const ctx1 = cm.getLLMContext();

package/src/test/contextTestTools.ts

@@ -0,0 +1,95 @@
+import { expect } from "vitest";
+
+import { IContextManager } from "../agent/context";
+import { MessageParam, UserMessageParam } from "../agent/llm";
+
+export const MESSAGES: MessageParam[] = [
+  { role: "user", content: "msg200" },
+  { role: "assistant", content: "msg300" },
+  { role: "user", content: "msg400" },
+  {
+    role: "assistant",
+    content: "msg500",
+    tool_calls: [
+      {
+        id: "tool_call_0",
+        type: "function",
+        function: { name: "tool1", arguments: "" },
+      },
+    ],
+  },
+  {
+    role: "tool",
+    content: "msg501",
+    tool_call_id: "tool_call_0",
+  },
+  { role: "assistant", content: "msg502" },
+];
+
+export async function basicContextTest(cm: IContextManager) {
+  expect(cm.getAgentPrompt()).eql("sys_prompt");
+
+  // Add first 2 messages
+  let tx = await cm.startTx([MESSAGES[0] as UserMessageParam]);
+  tx.addMessage(MESSAGES[1]);
+  expect(tx.getLLMContext()).eql([
+    {
+      role: "system",
+      content: "sys_prompt",
+    },
+    MESSAGES[0],
+    MESSAGES[1],
+  ]);
+  await cm.commit(tx);
+
+  // Add more messages
+  tx = await cm.startTx([MESSAGES[2] as UserMessageParam]);
+  expect(tx.getLLMContext()).eql([
+    {
+      role: "system",
+      content: "sys_prompt",
+    },
+    MESSAGES[0],
+    MESSAGES[1],
+    MESSAGES[2],
+  ]);
+  tx.addMessages(MESSAGES.slice(3));
+  expect(tx.getLLMContext()).eql([
+    {
+      role: "system",
+      content: "sys_prompt",
+    },
+    MESSAGES[0],
+    MESSAGES[1],
+    MESSAGES[2],
+    MESSAGES[3],
+    MESSAGES[4],
+    MESSAGES[5],
+  ]);
+  await cm.commit(tx);
+
+  // Update system prompt
+  cm.setAgentPrompt("New system prompt");
+
+  // Wait for compression to complete, then get the new context
+  const context = cm.getLLMContext();
+
+  // Expect
+  //   - callback happened
+  //   - new context is as expected
+
+  const expectContext: MessageParam[] = [
+    {
+      role: "system",
+      content: "New system prompt",
+    },
+    MESSAGES[0], // 400
+    MESSAGES[1], // 500
+    MESSAGES[2], // 400
+    MESSAGES[3], // 500
+    MESSAGES[4], // 501
+    MESSAGES[5], // 502
+  ];
+
+  expect(context).eql(expectContext);
+}

package/src/tool/chatMain.ts

@@ -121,7 +121,6 @@ const server = command({
     }),
     supabaseUrl: options.supabaseUrl,
     supabaseKey: options.supabaseKey,
-    llmUrl: options.llmUrl,
     xmcpUrl: options.xmcpUrl,
     pidFile: options.pidFile,
   },
@@ -130,7 +129,6 @@ const server = command({
     port,
     supabaseUrl,
     supabaseKey,
-    llmUrl,
     xmcpUrl,
     pidFile,
   }): Promise<void> => {
@@ -138,7 +136,7 @@ const server = command({
       fs.writeFileSync(pidFile, process.pid.toString());
     }
 
-    void runServer(port, supabaseUrl, supabaseKey, llmUrl, xmcpUrl);
+    void runServer(port, supabaseUrl, supabaseKey, xmcpUrl);
   },
 });
 
@@ -224,7 +222,9 @@ const client = command({
   ${chalk.yellow("/get-file <name>")}    Get file contents
   ${chalk.yellow("/share-session <file>")} Share and write access token
   ${chalk.yellow("/pause-agent <0|1>")}  Pause (or unpause) the agent
-  ${chalk.yellow("/add-mcp-server-url <server-name> <url>")} Add MCP server
+  ${chalk.yellow("/add-custom-mcp-server <server-name> <url>")} Add MCP server
+  ${chalk.yellow("/remove-custom-mcp-server <server-name>")} Remove MCP server
+  ${chalk.yellow("/list-custom-mcp-servers")} list custom mcp servers
   ${chalk.yellow("/get-resource <server-name> <uri>")} get MCP server resource
 `;
 
@@ -302,8 +302,14 @@ const client = command({
         case "pause-agent":
           pauseAgent(sessionClient, cmds[1]);
           break;
-        case "add-mcp-server-url":
-          addMcpServerFromUrl(sessionClient, cmds[1], cmds[2]);
+        case "add-custom-mcp-server":
+          addCustomMcpServer(client, cmds[1], cmds[2]);
+          break;
+        case "remove-custom-mcp-server":
+          removeCustomMcpServer(client, cmds[1]);
+          break;
+        case "list-custom-mcp-servers":
+          listCustomMcpServers(client);
           break;
         case "get-resource":
           await getResource(sessionClient, cmds[1], cmds[2]);
@@ -392,12 +398,20 @@ function pauseAgent(sessionClient: SessionClient, pause: string) {
   sessionClient.setAgentPaused(!!parseInt(pause));
 }
 
-function addMcpServerFromUrl(
-  sessionClient: SessionClient,
-  name: string,
-  url: string
-) {
-  sessionClient.addMcpServerFromUrl(name, url);
+function addCustomMcpServer(chatClient: ChatClient, name: string, url: string) {
+  chatClient.addCustomMcpServer(name, "Temporary description", url);
+}
+
+function removeCustomMcpServer(chatClient: ChatClient, name: string) {
+  chatClient.removeCustomMcpServer(name);
+}
+
+function listCustomMcpServers(chatClient: ChatClient) {
+  const customMcpServers: Record<string, string> =
+    chatClient.getCustomMcpServers();
+  console.log(
+    `Custom mcp serevrs:\n${JSON.stringify(customMcpServers, undefined, 2)}`
+  );
 }
 
 async function getResource(