@xalia/agent 0.5.4 → 0.5.6

package/src/agent/mcpServerManager.ts

@@ -54,17 +54,31 @@ export class McpServerInfo {
     return this.tools;
   }
 
-  public getTool(toolName: string): Tool {
+  public getTool(toolName: string): Tool | undefined {
     return this.toolsMap[toolName];
   }
 }
 
+/**
+ * Instance of McpServerInfo which supports setting tool state.  Intended for
+ * IMcpServerManager implementations, not for client code.
+ */
+export class McpServerInfoRW extends McpServerInfo {
+  public enableTool(toolName: string) {
+    this.enabledToolsMap[toolName] = true;
+  }
+
+  public disableTool(toolName: string) {
+    delete this.enabledToolsMap[toolName];
+  }
+}
+
 /**
  * The internal class holds server info and allows it to be updated. Managed
  * by McpServerManager.  Do not access these methods except via the
  * McpServerManager.
  */
-class McpServerInfoInternal extends McpServerInfo {
+class McpServerInfoInternal extends McpServerInfoRW {
   private readonly client: McpClient;
   private readonly callbacks: McpCallbacks;
 
@@ -114,25 +128,37 @@ class McpServerInfoInternal extends McpServerInfo {
     await this.client.close();
   }
 
-  public enableTool(toolName: string) {
-    this.enabledToolsMap[toolName] = true;
-  }
-
-  public disableTool(toolName: string) {
-    delete this.enabledToolsMap[toolName];
-  }
-
   public getCallback(toolName: string) {
     return this.callbacks[toolName];
   }
 }
 
+/**
+ * The client's interface to a manager which has mcp servers assigned to it,
+ * and can then query them for their tools, enable/disbale spercific tools and
+ * remove servers completely.
+ *
+ * This interface says nothing about communication with a specific agent.  It
+ * only defines the client-facing interactions.
+ */
+export interface IMcpServerManager {
+  hasMcpServer(mcpServerName: string): boolean;
+  getMcpServerNames(): string[];
+  getMcpServer(mcpServerName: string): McpServerInfo;
+
+  removeMcpServer(mcpServerName: string): Promise<void>;
+  enableAllTools(mcpServerName: string): void;
+  disableAllTools(mcpServerName: string): void;
+  enableTool(mcpServerName: string, toolName: string): void;
+  disableTool(mcpServerName: string, toolName: string): void;
+}
+
 /**
  * Manage a set of MCP servers, where the tools for each server have an
  * 'enabled' flag.  Tools are disabled by default.  The set of enabled tools
  * over all servers is exposed as a single list of OpenAI functions.
  */
-export class McpServerManager {
+export class McpServerManager implements IMcpServerManager {
   private mcpServers: { [serverName: string]: McpServerInfoInternal } = {};
   private enabledToolsDirty: boolean = true;
   private enabledOpenAITools: ChatCompletionTool[] = [];
@@ -148,6 +174,10 @@ export class McpServerManager {
     this.mcpServers = {};
   }
 
+  public hasMcpServer(mcpServerName: string): boolean {
+    return !!this.mcpServers[mcpServerName];
+  }
+
   public getMcpServerNames(): string[] {
     return Object.keys(this.mcpServers);
   }
@@ -208,7 +238,7 @@ export class McpServerManager {
     }
   }
 
-  public async removeMcpServer(mcpServerName: string) {
+  public async removeMcpServer(mcpServerName: string): Promise<void> {
     const server = this.getMcpServerInternal(mcpServerName);
     delete this.mcpServers[mcpServerName];
     await server.shutdown();
@@ -287,7 +317,7 @@ export class McpServerManager {
   public getMcpServerSettings(): McpServerSettings {
     const config: McpServerSettings = {};
 
-    // NOTE: on load, entires of the form:
+    // NOTE: on load, entries of the form:
     //
     //   <server>: []
     //

package/src/agent/openAILLM.ts

@@ -1,3 +1,4 @@
+import { DEFAULT_LLM_MODEL, XALIA_APP_HEADER } from "./agentUtils";
 import { ILLM } from "./llm";
 import { OpenAI } from "openai";
 
@@ -14,8 +15,9 @@ export class OpenAILLM implements ILLM {
       apiKey,
       baseURL: apiUrl,
       dangerouslyAllowBrowser: true,
+      defaultHeaders: XALIA_APP_HEADER,
     });
-    this.model = model || "gpt-4o-mini";
+    this.model = model || DEFAULT_LLM_MODEL;
   }
 
   public setModel(model: string) {

package/src/agent/openAILLMStreaming.ts

@@ -2,6 +2,7 @@ import { getLogger } from "@xalia/xmcp/sdk";
 import { ILLM } from "./llm";
 import { OpenAI } from "openai";
 import { strict as assert } from "assert";
+import { DEFAULT_LLM_MODEL, XALIA_APP_HEADER } from "./agentUtils";
 
 const logger = getLogger();
 
@@ -265,7 +266,8 @@ function initializeCompletionChoice(
   return {
     choice: {
       message,
-      finish_reason: chunkChoice.finish_reason || "stop",
+      // We use `null` to signal that `finish_reason` is unset
+      finish_reason: chunkChoice.finish_reason || (null as unknown as "stop"),
       index: chunkChoice.index,
       logprobs: chunkChoice.logprobs || null,
     },
@@ -300,6 +302,10 @@ function updateCompletionChoice(
   assert(completionChoice.index === chunkChoice.index);
   updateCompletionMessage(completionChoice.message, chunkChoice.delta);
   if (chunkChoice.finish_reason) {
+    assert(
+      completionChoice.finish_reason === null,
+      `finish_reason already set: (${completionChoice.finish_reason})`
+    );
     completionChoice.finish_reason = chunkChoice.finish_reason;
     return true;
   }
@@ -313,7 +319,7 @@ function initializeCompletionChoices(
   // content to stream.  We keep it simple for now and assume only single
   // choices, which allows us to mark everything as done if any choice we hit is
   // done.
-  assert(chunkChoices.length === 1);
+  assert(chunkChoices.length < 2);
 
   let msgDone = false;
   const choices: OpenAI.Chat.Completions.ChatCompletion.Choice[] = [];
@@ -354,7 +360,7 @@ function updateCompletionChoices(
   return msgDone;
 }
 
-function initializeCompletion(
+export function initializeCompletion(
   chunk: OpenAI.Chat.Completions.ChatCompletionChunk
 ): { initMessage: OpenAI.Chat.Completions.ChatCompletion; done: boolean } {
   //   export interface ChatCompletionChunk {
@@ -397,7 +403,7 @@ function initializeCompletion(
   };
 }
 
-function updateCompletion(
+export function updateCompletion(
   completion: OpenAI.Chat.Completions.ChatCompletion,
   chunk: OpenAI.Chat.Completions.ChatCompletionChunk
 ): boolean {
@@ -448,8 +454,9 @@ export class OpenAILLMStreaming implements ILLM {
       apiKey,
       baseURL: apiUrl,
       dangerouslyAllowBrowser: true,
+      defaultHeaders: XALIA_APP_HEADER,
     });
-    this.model = model || "gpt-4o-mini";
+    this.model = model || DEFAULT_LLM_MODEL;
   }
 
   public setModel(model: string) {
@@ -482,46 +489,40 @@ export class OpenAILLMStreaming implements ILLM {
     }
 
     let aggregatedMessage: OpenAI.Chat.Completions.ChatCompletion | undefined;
-    let done = false;
-
     for await (const chunk of chunks) {
       logger.debug(`[stream] chunk: ${JSON.stringify(chunk)}`);
-      assert(!done);
 
       if (chunk.object !== "chat.completion.chunk") {
         // logger.warn("[stream]: unexpected message");
         continue;
       }
 
-      done = (() => {
-        if (!aggregatedMessage) {
-          logger.debug(`[stream] first}`);
-          const { initMessage, done } = initializeCompletion(chunk);
-          aggregatedMessage = initMessage;
-          return done;
-        } else {
-          return updateCompletion(aggregatedMessage, chunk);
-        }
-      })();
+      if (!aggregatedMessage) {
+        logger.debug(`[stream] first}`);
+        const { initMessage } = initializeCompletion(chunk);
+        aggregatedMessage = initMessage;
+      } else {
+        updateCompletion(aggregatedMessage, chunk);
+      }
 
       if (onMessage) {
-        // Inform the call of a message fragment.  Note that even if there is
-        // no content, we must call `onMessage` once `done` is true.
-
-        const delta = chunk.choices[0].delta;
-        if (delta.content) {
-          await onMessage(delta.content, done);
-        } else if (done) {
-          await onMessage("", true);
+        // Inform the call of a message fragment if it contains any text.
+
+        const delta = chunk.choices[0]?.delta;
+        if (delta?.content) {
+          await onMessage(delta.content, false);
         }
       }
     }
 
+    if (onMessage) {
+      await onMessage("", true);
+    }
+
     logger.debug(
       `[stream] final message: ${JSON.stringify(aggregatedMessage)}`
     );
 
-    assert(done);
     assert(aggregatedMessage);
     return aggregatedMessage;
   }

package/src/agent/repeatLLM.ts

@@ -0,0 +1,51 @@
+import { ILLM } from "./llm";
+import { OpenAI } from "openai";
+import { strict as assert } from "assert";
+
+export class RepeatLLM implements ILLM {
+  private idx: number = 0;
+
+  public getModel(): string {
+    return "repeat";
+  }
+
+  public getUrl(): string {
+    throw "cannot get url for RepeatLLM";
+  }
+
+  public async getConversationResponse(
+    _messages: OpenAI.Chat.Completions.ChatCompletionMessageParam[],
+    _tools?: OpenAI.Chat.Completions.ChatCompletionTool[],
+    onMessage?: (msg: string, msgEnd: boolean) => Promise<void>
+  ): Promise<OpenAI.Chat.Completions.ChatCompletion> {
+    await new Promise((r) => setTimeout(r, 0));
+
+    const content = `Message number ${this.idx++}`;
+    const response: OpenAI.Chat.Completions.ChatCompletion.Choice = {
+      finish_reason: "stop",
+      index: 0,
+      logprobs: null,
+      message: {
+        content,
+        refusal: null,
+        role: "assistant",
+      },
+    };
+
+    if (onMessage) {
+      onMessage(content, true);
+    }
+
+    return {
+      id: "" + this.idx,
+      choices: [response],
+      created: Date.now(),
+      model: "dummyLlmModel",
+      object: "chat.completion",
+    };
+  }
+
+  public setModel(_model: string): void {
+    assert(false, "unexpected call to setModel");
+  }
+}

package/src/agent/sudoMcpServerManager.ts

@@ -1,5 +1,9 @@
 import { Tool } from "@modelcontextprotocol/sdk/types.js";
-import { McpServerManager, McpServerSettings } from "./mcpServerManager";
+import {
+  IMcpServerManager,
+  McpServerManager,
+  McpServerSettings,
+} from "./mcpServerManager";
 import {
   ApiClient,
   McpServerBrief,
@@ -35,11 +39,29 @@ class SanitizedServerBrief extends McpServerBrief {
   }
 }
 
+/**
+ * A manager which can expose a set of mcp servers and set them to an
+ * McpServerManager.
+ */
+export interface ISkillManager {
+  getMcpServerManager(): IMcpServerManager;
+  getServerBriefs(): McpServerBrief[];
+
+  /**
+   * In some cases (requests to a remote SkillManager), the promise does not
+   * correspond the server being added.  Hence the caller cannot be sure when
+   * it is safe to run the enable-all logic. `enableAll` solves this by having
+   * the SkillManager (local or remote) run the enableAll logic as the server
+   * is enabled.
+   */
+  addMcpServer(serverName: string, enableAll: boolean): Promise<void>;
+}
+
 /**
  * Manages access to the catalogue of servers hosted by sudomcp.  Supports
  * adding these servers to McpServerManager.
  */
-export class SkillManager {
+export class SkillManager implements ISkillManager {
   private constructor(
     private mcpServerManager: McpServerManager,
     private apiClient: ApiClient,
@@ -109,7 +131,7 @@ export class SkillManager {
     // Concurrently establish all server connections
     const addServer = async (serverName: string) => {
       logger.debug(`restoring ${serverName} ...`);
-      return this.addMcpServer(serverName);
+      return this.addMcpServer(serverName, false);
     };
     await Promise.all(Object.entries(mcpConfig).map((e) => addServer(e[0])));
 
@@ -144,6 +166,10 @@ export class SkillManager {
     this.toolCache = {};
   }
 
+  public hasServer(serverName: string): boolean {
+    return !!this.serverBriefsMap[serverName];
+  }
+
   public getServerBriefs(): McpServerBrief[] {
     return this.serverBriefs;
   }
@@ -171,11 +197,13 @@ export class SkillManager {
   }
 
   /**
-   * Add a server to the `McpServerManager`, using `ApiClient`
-   * to produce the transport. Validates the server's config
-   * schema, if applicable.
+   * Add a server to the `McpServerManager`, using `ApiClient` to produce the
+   * transport. Validates the server's config schema, if applicable.
    */
-  public async addMcpServer(serverName: string): Promise<void> {
+  public async addMcpServer(
+    serverName: string,
+    enableAll: boolean
+  ): Promise<void> {
     const tools = await this.getServerTools(serverName);
     const originalName = this.serverBriefsMap[serverName].originalName;
     const mcpserver = await this.apiClient.getDetails(originalName, "run");
@@ -190,11 +218,12 @@ export class SkillManager {
       this.openUrl,
       this.authorized_url
     );
-    await this.mcpServerManager.addMcpServerWithClient(
-      client,
-      serverName,
-      tools
-    );
+
+    const msm = this.mcpServerManager;
+    await msm.addMcpServerWithClient(client, serverName, tools);
+    if (enableAll) {
+      msm.enableAllTools(serverName);
+    }
   }
 
   public getOriginalName(serverName: string): string {