@xalia/agent 0.6.8 → 0.6.9

package/src/chat/server/openAIRouterLLM.ts

@@ -0,0 +1,171 @@
+import { OpenAI } from "openai";
+
+import { getLogger } from "@xalia/xmcp/sdk";
+import { utils } from "@xalia/xmcp/tool";
+
+import { OpenAILLMStreaming } from "../../agent/openAILLMStreaming";
+import {
+  ILLM,
+  MessageParam,
+  ToolDescriptor,
+  Completion,
+  XALIA_APP_HEADER,
+} from "../../agent/llm";
+
+const logger = getLogger();
+
+// TODO: May be better to allow the calling code to determine this.
+
+const DEFAULT_PROVIDER_URLS: Record<string, string | undefined> = {
+  openrouter: "https://openrouter.ai/api/v1",
+  openai: "https://api.openai.com/v1",
+  anthropic: "https://api.anthropic.com/v1",
+  together: "https://api.together.xyz/v1",
+};
+
+const DEFAULT_MODEL_MAP: Record<string, string | undefined> = {
+  "gpt-4o-mini": "openai",
+  "gpt-4o": "openai",
+  "openai/gpt-4o-mini": "openrouter",
+  "openai/gpt-4o": "openrouter",
+  "google/gemini-2.5-flash": "openrouter",
+  "google/gemini-2.5-pro": "openrouter",
+  "google/gemini-2.5-flash-image-preview": "openrouter",
+  "anthropic/claude-3.7-sonnet": "openrouter",
+  "anthropic/claude-sonnet-4": "openrouter",
+  "anthropic/claude-sonnet-4.5": "openrouter",
+  "claude-3-7-sonnet-20250219": "anthropic",
+  "arcee-ai/AFM-4.5B-Preview": "together",
+  "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": "together",
+};
+
+type CachedMaps = {
+  modelToProvider: Record<string, string | undefined>;
+  providerUrls: Record<string, string | undefined>;
+  providerApiKeys: Record<string, string | undefined>;
+};
+
+let cached_maps: CachedMaps | undefined = undefined;
+
+// TODO: Move this to the CLI command args?
+
+function decodeOrUseDefault(
+  stringOrFile: string | undefined,
+  defaultValue: Record<string, string | undefined>
+): Record<string, string | undefined> {
+  const value =
+    utils.loadContentOrFileOrUndefined<Record<string, string | undefined>>(
+      stringOrFile
+    );
+  if (value) {
+    return value;
+  }
+  return defaultValue;
+}
+
+function getMaps(): CachedMaps {
+  // global cached_maps;
+  if (!cached_maps) {
+    const providerApiKeys = utils.loadContentOrFileOrUndefined<
+      Record<string, string | undefined>
+    >(process.env.LLM_API_KEY_MAP);
+    if (!providerApiKeys) {
+      throw new Error("LLM_API_KEY_MAP not not given");
+    }
+    const providerUrls = decodeOrUseDefault(
+      process.env.LLM_PROVIDER_URLS,
+      DEFAULT_PROVIDER_URLS
+    );
+    const modelToProvider = decodeOrUseDefault(
+      process.env.LLM_MODEL_MAP,
+      DEFAULT_MODEL_MAP
+    );
+    cached_maps = {
+      modelToProvider,
+      providerUrls,
+      providerApiKeys,
+    };
+  }
+
+  return cached_maps;
+}
+
+export function getOpenAIClientParams(model: string): {
+  apiKey: string | undefined;
+  baseURL: string;
+} {
+  const maps = getMaps();
+  const provider = maps.modelToProvider[model];
+  if (!provider) {
+    throw new Error(`no provider registered for model ${model}`);
+  }
+  const baseURL = maps.providerUrls[provider];
+  if (!baseURL) {
+    throw new Error(`no provider url for provider ${provider}`);
+  }
+  const apiKey = maps.providerApiKeys[provider];
+  if (!apiKey) {
+    logger.warn(`No apiKey registered for provider ${provider}`);
+  }
+
+  return { apiKey, baseURL };
+}
+
+export function getOpenAIClient(model: string): OpenAI {
+  const { apiKey, baseURL } = getOpenAIClientParams(model);
+  return new OpenAI({
+    apiKey,
+    baseURL,
+    dangerouslyAllowBrowser: true,
+    defaultHeaders: XALIA_APP_HEADER,
+  });
+}
+
+/**
+ * An ILLM implementation that accepts maps:
+ *
+ *   model-name => prover-name
+ *   provider-name => url
+ *   provider-name => api-key
+ *
+ * Requests are dynamically routed to the appropriate provider based on the
+ * currently active model.
+ */
+export class OpenAIRouterLLM implements ILLM {
+  private model: string;
+  private openai: OpenAI;
+
+  constructor(model: string) {
+    this.model = model;
+    this.openai = getOpenAIClient(model);
+  }
+
+  public setModel(model: string) {
+    this.model = model;
+    this.openai = getOpenAIClient(model);
+  }
+
+  public getModel(): string {
+    return this.model;
+  }
+
+  public getUrl(): string {
+    return this.openai.baseURL;
+  }
+
+  public async getConversationResponse(
+    messages: MessageParam[],
+    tools?: ToolDescriptor[],
+    onMessage?: (msg: string, end: boolean) => Promise<void>,
+    onReasoning?: (reasoning: string) => Promise<void>
+  ): Promise<Completion> {
+    return OpenAILLMStreaming.makeRequest(
+      this.openai,
+      this.model,
+      messages,
+      tools,
+      onMessage,
+      onReasoning
+    );
+  }
+}

package/src/chat/server/openSession.ts

@@ -4,22 +4,24 @@ import { Tool } from "@modelcontextprotocol/sdk/types.js";
 
 import {
   AgentPreferences,
+  AgentProfile,
   Configuration,
   getLogger,
   McpServerSettings,
   SavedAgentProfile,
 } from "@xalia/xmcp/sdk";
 
-import { Agent, createUserMessage } from "../../agent/agent";
+import { AgentEx, createUserMessage } from "../../agent/agent";
 import {
   AssistantMessageParam,
+  ILLM,
   MessageToolCall,
   ToolMessageParam,
 } from "../../agent/llm";
 import { SkillManager } from "../../agent/sudoMcpServerManager";
 import { McpServerInfo } from "../../agent/mcpServerManager";
 import { IAgentEventHandler } from "../../agent/iAgentEventHandler";
-import { createAgentWithoutSkills } from "../../agent/agentUtils";
+import { createSpecializedLLM } from "../../agent/agentUtils";
 import { IPlatform } from "../../agent/iplatform";
 
 import type {
@@ -44,7 +46,6 @@ import type {
   ClientShareSession,
   ClientSessionMessageBase,
   ServerAgentPaused,
-  ClientAddMcpServerFromUrl,
   ClientGetMcpResource,
 } from "../protocol/messages";
 import { AsyncQueue } from "../utils/asyncQueue";
@@ -70,7 +71,11 @@ import {
 } from "../utils/approvalManager";
 import { ChatErrorMessage, ChatFatalError } from "../protocol/errors";
 import { addDefaultChatTools } from "./tools";
-import { ChatContextManager, ICheckpointWriter } from "./chatContextManager";
+import {
+  ChatContextManager,
+  ChatContextTransaction,
+  ICheckpointWriter,
+} from "./chatContextManager";
 import {
   llmUserMessageToUserMessageData,
   MESSAGE_INDEX_START_VALUE,
@@ -88,6 +93,7 @@ import { SessionFileEntry } from "../data/dbSessionFileModels";
 import { ApiKeyManager } from "../data/apiKeyManager";
 import { DbSessionMessages } from "../data/dbSessionMessages";
 import { ISessionMessageSender } from "./openSessionMessageSender";
+import { OpenAIRouterLLM } from "./openAIRouterLLM";
 
 /**
  * The model to use when the AgentProfile does not specify one.
@@ -208,13 +214,13 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
     private readonly sessionUUID: string,
     private readonly sender: ISessionMessageSender<ServerToClient>,
     private readonly approvalManager: ToolApprovalManager,
-    private readonly contextManager: ChatContextManager
+    private readonly contextTx: ChatContextTransaction
   ) {}
 
   onCompletion(result: AssistantMessageParam): void {
     logger.debug(`[OpenSession.onCompletion] : ${JSON.stringify(result)}`);
     // Nothing to broadcast. Caller will receive this via onAgentMessage.
-    this.contextManager.processAgentResponse(result);
+    this.contextTx.processAgentResponse(result);
   }
 
   onImage(image: OpenAI.Chat.Completions.ChatCompletionContentPartImage): void {
@@ -224,7 +230,7 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
 
   onToolCallResult(result: ToolMessageParam): void {
     logger.debug(`[onToolCallResult] : ${JSON.stringify(result)}`);
-    const toolCallMessage = this.contextManager.processToolCallResult(result);
+    const toolCallMessage = this.contextTx.processToolCallResult(result);
     this.sender.broadcast(toolCallMessage);
   }
 
@@ -273,7 +279,7 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
     );
 
     // Inform the contextManager and broadcast the ServerAgentMessageChunk
-    const agentMsgChunk = this.contextManager.processAgentMessage(msg, end);
+    const agentMsgChunk = this.contextTx.processAgentMessageChunk(msg, end);
     this.sender.broadcast(agentMsgChunk);
     return Promise.resolve();
   }
@@ -305,7 +311,7 @@ export class ChatSessionAgentEventHandler implements IAgentEventHandler {
  */
 export class OpenSession implements ISessionFileManagerEventHandler {
   private readonly db: Database;
-  private /* readonly */ agent: Agent;
+  private /* readonly */ agent: AgentEx;
   private readonly sessionUUID: string;
   private readonly teamUUID: string | undefined;
   private readonly userUUID: string;
@@ -328,7 +334,7 @@ export class OpenSession implements ISessionFileManagerEventHandler {
 
   private constructor(
     db: Database,
-    agent: Agent,
+    agent: AgentEx,
     sessionData: SessionData,
     savedAgentProfile: SavedAgentProfile,
     isPersisted: boolean,
@@ -381,7 +387,6 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     ownerData: UserData,
     ownerApiKey: string,
     sessionCheckpoint: SessionCheckpoint | undefined,
-    llmUrl: string,
     xmcpUrl: string,
     connectionManager: IUserConnectionManager<ServerToClient>
   ): Promise<OpenSession> {
@@ -400,6 +405,8 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       sender,
       new DbAgentPreferencesWriter(db)
     );
+    const checkpointWriter = new DBCheckpointWriter(db, sessionId);
+
     const { agent, skillManager, contextManager } = await createContextAndAgent(
       sessionId,
       savedAgentProfile.profile.system_prompt,
@@ -409,12 +416,10 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       sessionCheckpoint,
       ownerData,
       ownerApiKey,
-      llmUrl,
       xmcpUrl,
       fileManager,
-      sender,
       platform,
-      toolApprovalManager
+      checkpointWriter
     );
 
     const openSession = new OpenSession(
@@ -440,7 +445,6 @@ export class OpenSession implements ISessionFileManagerEventHandler {
   static async initWithEmptySession(
     db: Database,
     sessionData: SessionData,
-    llmUrl: string,
     xmcpUrl: string,
     connectionManager: IUserConnectionManager<ServerToClient>
   ): Promise<OpenSession> {
@@ -481,7 +485,6 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       ownerData,
       ownerApiKey,
       sessionCheckpoint,
-      llmUrl,
       xmcpUrl,
       connectionManager
     );
@@ -490,7 +493,6 @@ export class OpenSession implements ISessionFileManagerEventHandler {
   static async initWithExistingSession(
     db: Database,
     sessionId: string,
-    llmUrl: string,
     xmcpUrl: string,
     connectionManager: IUserConnectionManager<ServerToClient>
   ): Promise<OpenSession> {
@@ -514,7 +516,6 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       ownerData,
       ownerApiKey,
       sessionCheckpoint,
-      llmUrl,
       xmcpUrl,
       connectionManager
     );
@@ -590,7 +591,7 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     this.sendMcpSettings(connectionId);
 
     // add system prompt and model
-    const agentProfile = this.agent.getAgentProfile();
+    const agentProfile = this.getAgentProfile();
     connMgr.sendToConnection(connectionId, {
       type: "system_prompt_updated",
       system_prompt: agentProfile.system_prompt,
@@ -879,9 +880,6 @@ export class OpenSession implements ISessionFileManagerEventHandler {
         case "share_session":
           await this.handleShareSession(msg, queuedMessage.from);
           break;
-        case "add_mcp_server_from_url":
-          void this.handleAddMcpServerFromUrl(msg, queuedMessage.from);
-          break;
         case "get_mcp_resource":
           void this.handleGetMcpResource(msg, queuedMessage.from);
           break;
@@ -928,34 +926,6 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     });
   }
 
-  private async handleAddMcpServerFromUrl(
-    msg: ClientAddMcpServerFromUrl,
-    from: string
-  ): Promise<void> {
-    const skillManager = this.skillManager;
-    try {
-      const serverName = msg.server_name;
-
-      await skillManager.addMcpServerWithStreamableHTTPUrl(serverName, msg.url);
-      this.skillManager.enableAllTools(serverName);
-
-      const server = skillManager.getMcpServer(serverName);
-      const tools = server.getTools();
-      const resources = server.getResources();
-      const enabled_tools = Array.from(server.getEnabledTools().keys());
-      this.sender.broadcast({
-        type: "mcp_server_added",
-        server_name: serverName,
-        tools,
-        resources,
-        enabled_tools,
-        session_id: this.sessionUUID,
-      });
-    } catch (err: unknown) {
-      this.handleError(err, from);
-    }
-  }
-
   private async handleGetMcpResource(
     msg: ClientGetMcpResource,
     from: string
@@ -1001,26 +971,39 @@ export class OpenSession implements ISessionFileManagerEventHandler {
    * `processUserMessage` logic when agent is paused.  Trigger the context,
    * add the user messages and then extract the new DB messages.
    */
-  private processUserMessagePaused(
+  private async processUserMessagePaused(
     msgs: ServerUserMessage[]
-  ): SessionMessage[] {
-    const { llmUserMessages } = this.contextManager.startAgentResponse(msgs);
-
-    // Just send the user LLM messages direct to the ContextManager, so they
-    // are available to the LLM once it is restarted.
+  ): Promise<SessionMessage[]> {
+    // Just send the user LLM user messages direct to the ContextManager, so
+    // they are available to the LLM once it is restarted.
 
-    this.contextManager.addMessages(llmUserMessages);
-    return this.contextManager.endAgentResponse();
+    const { contextTx } = await this.contextManager.startAgentResponse(msgs);
+    return this.contextManager.endAgentResponse(contextTx);
   }
 
   private async processUserMessagesActive(
     msgs: ServerUserMessage[]
   ): Promise<SessionMessage[]> {
-    const { llmUserMessages, agentFirstChunk } =
-      this.contextManager.startAgentResponse(msgs);
+    // TODO: create the contextTx and store all new messages on this.  Event
+    // handler should accept the contextTx and forward messages to it, as well
+    // as sending the updates.
+
+    // All accumulated messages (DB, Protocol and LLM) should be on the
+    // specialized contextTx.
+
+    const { contextTx, agentFirstChunk } =
+      await this.contextManager.startAgentResponse(msgs);
     this.sender.broadcast(agentFirstChunk);
+
+    const eventHandler = new ChatSessionAgentEventHandler(
+      this.sessionUUID,
+      this.sender,
+      this.approvalManager,
+      contextTx
+    );
+
     try {
-      await this.agent.userMessagesRaw(llmUserMessages);
+      await this.agent.userMessagesRaw(contextTx, eventHandler);
     } catch (e) {
       logger.warn(
         `[OpenSession.processUserMessages] agent error: ${String(e)}`
@@ -1029,16 +1012,18 @@ export class OpenSession implements ISessionFileManagerEventHandler {
       // Errors during agent replies must be turned into messages.
 
       const errMsg = `error from LLM: ${String(e)}`;
-      this.contextManager.revertAgentResponse(errMsg);
+      contextTx.revertAgentResponse(errMsg);
+      // TODO: This will prevent the user messages from being saved in the
+      // DB.  Should we keep them?
       throw new Error(errMsg);
     }
-    return this.contextManager.endAgentResponse();
+    return this.contextManager.endAgentResponse(contextTx);
   }
 
   private async processUserMessages(msgs: ServerUserMessage[]): Promise<void> {
     try {
       const newSessionMessages = this.agentPaused
-        ? this.processUserMessagePaused(msgs)
+        ? await this.processUserMessagePaused(msgs)
         : await this.processUserMessagesActive(msgs);
 
       logger.debug(
@@ -1328,7 +1313,7 @@ export class OpenSession implements ISessionFileManagerEventHandler {
   private async handleSetSystemPrompt(
     system_prompt: string
   ): Promise<ServerSystemPromptUpdated> {
-    this.agent.setSystemPrompt(system_prompt);
+    this.contextManager.setAgentPrompt(system_prompt);
     await this.updateAgentProfile();
     return {
       type: "system_prompt_updated",
@@ -1338,7 +1323,7 @@ export class OpenSession implements ISessionFileManagerEventHandler {
   }
 
   private async handleSetModel(model: string): Promise<ServerModelUpdated> {
-    this.agent.setModel(model);
+    this.agent.llm.setModel(model);
     await this.updateAgentProfile();
     return { type: "model_updated", model, session_id: this.sessionUUID };
   }
@@ -1444,8 +1429,16 @@ export class OpenSession implements ISessionFileManagerEventHandler {
     return tool;
   }
 
+  private getAgentProfile(): AgentProfile {
+    return new AgentProfile(
+      this.agent.llm.getModel(),
+      this.contextManager.getAgentPrompt(),
+      this.skillManager.getMcpServerSettings()
+    );
+  }
+
   private async updateAgentProfile(): Promise<void> {
-    const profile = this.agent.getAgentProfile();
+    const profile = this.getAgentProfile();
     logger.debug(
       `[updateAgentProfile]: uuid: ${this.agentProfileUUID} profile: ` +
         JSON.stringify(profile)
@@ -1642,28 +1635,35 @@ async function createContextAndAgent(
   sessionCheckpoint: SessionCheckpoint | undefined,
   ownerData: UserData,
   ownerApiKey: string,
-  llmUrl: string,
   xmcpUrl: string,
   fileManager: ChatSessionFileManager,
-  sender: ISessionMessageSender<ServerToClient>,
   platform: IPlatform,
-  approvalManager: ToolApprovalManager
+  checkpointWriter: ICheckpointWriter
 ): Promise<{
-  agent: Agent;
+  agent: AgentEx;
   skillManager: SkillManager;
   contextManager: ChatContextManager;
 }> {
+  // Fn to create the llm.  One invocation for the compression context, one
+  // for the Agent.
+  const createLLM = async (): Promise<ILLM> => {
+    let llm = await createSpecializedLLM(model, platform);
+    if (!llm) {
+      llm = new OpenAIRouterLLM(model);
+    }
+    assert(llm);
+    return llm;
+  };
+
   const contextManager = new ChatContextManager(
     systemPrompt,
     sessionMessages,
     sessionUUID,
     ownerData.uuid,
     sessionCheckpoint,
-    llmUrl,
-    model,
-    ownerApiKey,
-    undefined as unknown as DBCheckpointWriter, // TODO
-    fileManager
+    checkpointWriter,
+    fileManager,
+    await createLLM()
   );
   if (workspace) {
     contextManager.setWorkspace(
@@ -1671,32 +1671,19 @@ async function createContextAndAgent(
     );
   }
 
-  const eventHandler = new ChatSessionAgentEventHandler(
-    sessionUUID,
-    sender,
-    approvalManager,
-    contextManager
-  );
-
   const xmcpConfig = Configuration.new(ownerApiKey, xmcpUrl, false);
-  const [agent, skillManager] = await createAgentWithoutSkills(
-    llmUrl,
-    model,
-    eventHandler,
-    platform,
-    contextManager,
-    ownerApiKey,
-    xmcpConfig,
-    undefined,
-    true
-  );
-  await addDefaultChatTools(
-    agent,
-    ownerData.timezone,
-    fileManager,
-    llmUrl,
-    ownerApiKey
+  const skillManager = await SkillManager.initialize(
+    (url: string, authResultP: Promise<boolean>, displayName: string) => {
+      platform.openUrl(url, authResultP, displayName);
+    },
+    xmcpConfig.backend_url,
+    xmcpConfig.api_key,
+    undefined /* authorizedUrl */
   );
+  const llm = await createLLM();
+  const agent = new AgentEx(skillManager, llm);
+
+  await addDefaultChatTools(agent, ownerData.timezone, platform, fileManager);
 
   return { agent, skillManager, contextManager };
 }

package/src/chat/server/server.ts

@@ -13,7 +13,12 @@ import { ChatFatalError } from "../protocol/errors";
 import { ServerToClient } from "../protocol/messages";
 import { SessionRegistry } from "./sessionRegistry";
 
+const DEVELOPMENT: boolean = process.env.DEVELOPMENT === "1";
+
 dotenv.config();
+if (DEVELOPMENT) {
+  dotenv.config({ path: ".env.development" });
+}
 
 const logger = getLogger();
 
@@ -59,7 +64,6 @@ export async function runServer(
   port: number,
   supabaseUrl: string,
   supabaseKey: string,
-  llmUrl: string,
   xmcpUrl: string
 ): Promise<ws.Server> {
   return new Promise((r, _e) => {
@@ -68,7 +72,7 @@ export async function runServer(
     const createSessionRegistry = (
       connManager: IUserConnectionManager<ServerToClient>
     ) => {
-      return new SessionRegistry(db, connManager, llmUrl, xmcpUrl);
+      return new SessionRegistry(db, connManager, xmcpUrl);
     };
     const connectionManager = ConnectionManager.init(createSessionRegistry);
 

package/src/chat/server/sessionFileManager.ts

@@ -2,7 +2,7 @@ import { v4 as uuidv4 } from "uuid";
 import { strict as assert } from "assert";
 
 import {
-  Agent,
+  AgentEx,
   IAgentToolProvider,
   ToolCallResult,
   ToolHandler,
@@ -359,7 +359,7 @@ export function fileManagerTool(
 
   const parseName = makeParseArgsFn(["name"] as const);
   const getFileContentFn: ToolHandler = async (
-    _agent: Agent,
+    _agent: AgentEx,
     args: unknown
   ): Promise<ToolCallResult> => {
     const { name } = parseName(args);
@@ -376,7 +376,7 @@ export function fileManagerTool(
   const putArgs = ["name", "summary", "data_url"] as const;
   const parseNameSummaryDataUrl = makeParseArgsFn(putArgs);
   const putFileContentFn: ToolHandler = async (
-    _: Agent,
+    _: AgentEx,
     args: unknown
   ): Promise<ToolCallResultWithFileRef> => {
     const parsed = parseNameSummaryDataUrl(args);
@@ -399,14 +399,14 @@ export function fileManagerTool(
   // Allows LLM to write data (as a data-url) to the session file manager. The
   // data is replaced by a session file url after being saved.
 
-  const deleteFileFn: ToolHandler = async (_: Agent, args: unknown) => {
+  const deleteFileFn: ToolHandler = async (_: AgentEx, args: unknown) => {
     const { name } = parseName(args);
     await fileManager.deleteFile(name);
     return { response: "" };
   };
 
   return {
-    setup: (agent: Agent) => {
+    setup: (agent: AgentEx) => {
       agent.addAgentTool(GET_FILE_CONTENT_TOOL, getFileContentFn);
       agent.addAgentTool(PUT_FILE_CONTENT_TOOL, putFileContentFn);
       agent.addAgentTool(DELETE_FILE_CONTENT_TOOL, deleteFileFn);

package/src/chat/server/sessionRegistry.test.ts

@@ -49,7 +49,6 @@ describe("SessionRegistry", () => {
     sessionRegistry = new SessionRegistry(
       mockDatabase.mock,
       mockUserConnectionManager.mock,
-      "http://llm-api.test",
       "http://xmcp-api.test"
     );
   });