@xalia/agent 0.5.7 → 0.6.0

package/src/chat/client/sessionClient.ts

@@ -0,0 +1,443 @@
+import { Tool } from "@modelcontextprotocol/sdk/types.js";
+import { strict as assert } from "assert";
+import { v4 as uuidv4 } from "uuid";
+
+import {
+  AgentProfile,
+  McpServerBrief,
+  SavedAgentProfile,
+  getLogger,
+} from "@xalia/xmcp/sdk";
+
+import { ISkillManager } from "../../agent/sudoMcpServerManager";
+import { McpServerInfo, McpServerInfoRW } from "../../agent/mcpServerManager";
+import { ChatCompletionMessageParam, IConversation } from "../../agent/agent";
+
+import {
+  ClientSessionMessage,
+  ClientSessionMessageData,
+  ServerSessionScopedMessage,
+  ServerSessionInfo,
+  ClientSetWorkspace,
+  ClientToServer,
+} from "../protocol/messages";
+import { SessionParticipantMap } from "../data/dataModels";
+import { createSessionParticipantMap } from "../data/database";
+import { ISessionMessageSender } from "./interfaces";
+import { IMessageSender } from "./connection";
+
+const logger = getLogger();
+
+class RemoteSudoMcpServerManager implements ISkillManager {
+  private sender: ISessionMessageSender;
+  private briefs: McpServerBrief[];
+  private briefsMap = new Map<string, McpServerBrief>();
+  private mcpServers = new Map<string, McpServerInfoRW>();
+
+  constructor(sender: ISessionMessageSender, briefs: McpServerBrief[]) {
+    this.sender = sender;
+    this.briefs = briefs;
+    briefs.forEach((b) => {
+      this.briefsMap.set(b.name, b);
+    });
+  }
+
+  hasMcpServer(mcpServerName: string): boolean {
+    return this.mcpServers.has(mcpServerName);
+  }
+
+  getMcpServerNames(): string[] {
+    return Array.from(this.mcpServers.keys());
+  }
+
+  getMcpServer(mcpServerName: string): McpServerInfo {
+    const server = this.mcpServers.get(mcpServerName);
+    if (server) {
+      return server;
+    }
+    throw Error(`[getMcpServer] unknown server ${mcpServerName}`);
+  }
+
+  // eslint-disable-next-line @typescript-eslint/require-await
+  async removeMcpServer(mcpServerName: string): Promise<void> {
+    if (!this.mcpServers.has(mcpServerName)) {
+      logger.info(
+        `[removeMcpServer] mcpServers: ${JSON.stringify(this.mcpServers)}`
+      );
+      throw Error(`no server ${mcpServerName} (removeMcpServer)`);
+    }
+
+    this.sender.sendSessionMessage({
+      type: "remove_mcp_server",
+      server_name: mcpServerName,
+    });
+  }
+
+  enableAllTools(mcpServerName: string): void {
+    if (!this.mcpServers.has(mcpServerName)) {
+      throw Error(`no server ${mcpServerName} (enableAllTools)`);
+    }
+
+    this.sender.sendSessionMessage({
+      type: "enable_all_mcp_server_tools",
+      server_name: mcpServerName,
+    });
+  }
+
+  disableAllTools(mcpServerName: string): void {
+    if (!this.mcpServers.has(mcpServerName)) {
+      throw Error(`no server ${mcpServerName} (disableAllTools)`);
+    }
+
+    this.sender.sendSessionMessage({
+      type: "disable_all_mcp_server_tools",
+      server_name: mcpServerName,
+    });
+  }
+
+  enableTool(mcpServerName: string, toolName: string): void {
+    const server = this.mcpServers.get(mcpServerName);
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (enableTool)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    this.sender.sendSessionMessage({
+      type: "enable_mcp_server_tool",
+      server_name: mcpServerName,
+      tool: toolName,
+    });
+  }
+
+  disableTool(mcpServerName: string, toolName: string): void {
+    const server = this.mcpServers.get(mcpServerName);
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (disableTool)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    this.sender.sendSessionMessage({
+      type: "disable_mcp_server_tool",
+      server_name: mcpServerName,
+      tool: toolName,
+    });
+  }
+
+  onMcpServerAdded(
+    mcpServerName: string,
+    tools: Tool[],
+    enabled_tools: string[]
+  ) {
+    logger.debug(
+      `[onMcpServerAdded]: ${mcpServerName}, tools: ${JSON.stringify(tools)}` +
+        `, enabled: ${JSON.stringify(enabled_tools)}`
+    );
+
+    const mcpServerInfo = new McpServerInfoRW(tools);
+    for (const tool of enabled_tools) {
+      mcpServerInfo.enableTool(tool);
+    }
+    this.mcpServers.set(mcpServerName, mcpServerInfo);
+  }
+
+  onMcpServerRemoved(mcpServerName: string) {
+    this.mcpServers.delete(mcpServerName);
+  }
+
+  onMcpServerToolEnabled(mcpServerName: string, toolName: string) {
+    const server = this.mcpServers.get(mcpServerName);
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (onMcpServerToolEnabled)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    server.enableTool(toolName);
+  }
+
+  onMcpServerToolDisabled(mcpServerName: string, toolName: string) {
+    const server = this.mcpServers.get(mcpServerName);
+    if (!server) {
+      throw Error(`no server ${mcpServerName} (onMcpServerToolDisabled)`);
+    }
+    const tools = server.getTool(toolName);
+    if (!tools) {
+      throw Error(`no tool ${toolName} on server ${mcpServerName}`);
+    }
+
+    server.disableTool(toolName);
+  }
+
+  getServerBriefs(): McpServerBrief[] {
+    return this.briefs;
+  }
+
+  // eslint-disable-next-line @typescript-eslint/require-await
+  async addMcpServer(server_name: string, enable_all: boolean): Promise<void> {
+    if (!this.briefsMap.has(server_name)) {
+      throw Error(`no such server ${server_name} (addMcpServer)`);
+    }
+
+    this.sender.sendSessionMessage({
+      type: "add_mcp_server",
+      server_name,
+      enable_all,
+    });
+  }
+
+  hasServer(server_name: string): boolean {
+    return this.briefsMap.has(server_name);
+  }
+
+  async shutdown(): Promise<void> {}
+}
+
+export class SessionClient implements ISessionMessageSender, IConversation {
+  private smsm: RemoteSudoMcpServerManager;
+  private systemPrompt: string;
+  private model: string;
+
+  public constructor(
+    private sessionUUID: string,
+    private savedAgentProfile: SavedAgentProfile,
+    private sender: IMessageSender<ClientToServer>,
+    serverBriefs: McpServerBrief[],
+    private participants: SessionParticipantMap
+  ) {
+    this.sender = sender;
+    this.smsm = new RemoteSudoMcpServerManager(this, serverBriefs);
+    this.systemPrompt = "";
+    this.model = "";
+    this.participants = participants;
+  }
+
+  public getSessionUUID(): string {
+    return this.sessionUUID;
+  }
+
+  public getSudoMcpServerManager(): ISkillManager {
+    return this.smsm;
+  }
+
+  public getConversation(): ChatCompletionMessageParam[] {
+    throw new Error("unimpl: getConversation");
+  }
+
+  public getAgentProfile(): AgentProfile {
+    return this.savedAgentProfile.profile;
+  }
+
+  public getAgentUuid(): string {
+    return this.savedAgentProfile.uuid;
+  }
+
+  getSavedAgentProfile(): SavedAgentProfile {
+    return this.savedAgentProfile;
+  }
+
+  getSystemPrompt(): string {
+    return this.systemPrompt;
+  }
+
+  setSystemPrompt(system_prompt: string): void {
+    // Don't set system prompt here.  Wait until we get confirmation from the
+    // server.
+    this.sendSessionMessage({ type: "set_system_prompt", system_prompt });
+  }
+
+  getModel(): string {
+    return this.model;
+  }
+
+  setModel(model: string): void {
+    // Don't set model here.  Wait until we get confirmation from the server.
+    this.sendSessionMessage({ type: "set_model", model });
+  }
+
+  userMessage(msg?: string, imageB64?: string): void {
+    assert(msg || imageB64, "Either message or image must be provided");
+
+    this.sendSessionMessage({ type: "msg", message: msg, imageB64 });
+  }
+
+  resetConversation(): void {
+    throw new Error("resetConversation not implemented for ChatClient");
+  }
+
+  shutdown(): Promise<void> {
+    throw new Error("shutdown not implemented for ChatClient");
+  }
+
+  getParticipants(): SessionParticipantMap {
+    return this.participants;
+  }
+
+  /**
+   * Sets auto-approval preference for a specific tool.
+   */
+  setAutoApproval(
+    serverName: string,
+    toolName: string,
+    autoApprove: boolean
+  ): void {
+    const smsm = this.getSudoMcpServerManager();
+    if (!smsm.hasMcpServer(serverName)) {
+      throw Error(`server ${serverName} not added (setAutoApproval)`);
+    }
+    const server = smsm.getMcpServer(serverName);
+    const tool = server.getTool(toolName);
+    if (!tool) {
+      throw Error(`no tool ${toolName} on server ${serverName}`);
+    }
+
+    this.sendSessionMessage({
+      type: "set_auto_approval",
+      server_name: serverName,
+      tool: toolName,
+      auto_approve: autoApprove,
+    });
+  }
+
+  /**
+   * Send a generic message (for compatibility)
+   */
+  public sendMessage(message: { type: string; message: string }): void {
+    throw new Error(`Unsupported message type: ${message.type}`);
+  }
+
+  /**
+   * The caller sees a message `ServerApproveToolCall`, at which point it
+   * should prompt the user for permission to run the tool. The result of
+   * that prompt should be passed here.
+   */
+  toolCallApprovalResult(
+    id: string,
+    result: boolean,
+    auto_approve: boolean
+  ): void {
+    this.sendSessionMessage({
+      type: "tool_call_approval_result",
+      id,
+      result,
+      auto_approve,
+    });
+  }
+
+  // Implementation of IMessageSender interface
+  public sendSessionMessage(message: ClientSessionMessageData): void {
+    const enrichedMessage: ClientSessionMessage = {
+      ...message,
+      client_message_id: uuidv4(),
+      session_id: this.sessionUUID,
+    };
+
+    this.sender.send(enrichedMessage);
+  }
+
+  public updateSessionInfo(sessionInfo: ServerSessionInfo): void {
+    // TODO: Implement this
+    // throw new Error("[SessionClient.updateSessionInfo] not implemented");
+    const infoStr = JSON.stringify(sessionInfo.workspace);
+    logger.debug(`[SessionClient] ignoring session info: ${infoStr}`);
+  }
+
+  public setWorkspace(
+    message: string | undefined,
+    imageB64: string | undefined
+  ): void {
+    const msg: ClientSetWorkspace = {
+      type: "set_workspace",
+      message,
+      imageB64,
+    };
+    this.sendSessionMessage(msg);
+  }
+
+  /**
+   * Set MCP server configuration.
+   */
+  public setMcpServerConfig(
+    server_name: string,
+    config: Record<string, string>
+  ): void {
+    this.sendSessionMessage({
+      type: "set_mcp_server_config",
+      server_name,
+      config,
+    });
+  }
+
+  public deleteMcpServerConfig(server_name: string): void {
+    this.sendSessionMessage({
+      type: "set_mcp_server_config",
+      server_name,
+    });
+  }
+
+  handleMessage(message: ServerSessionScopedMessage): void {
+    switch (message.type) {
+      //
+      // State updates
+      //
+
+      case "mcp_server_added":
+        this.smsm.onMcpServerAdded(
+          message.server_name,
+          message.tools,
+          message.enabled_tools
+        );
+        break;
+      case "mcp_server_removed":
+        this.smsm.onMcpServerRemoved(message.server_name);
+        break;
+      case "mcp_server_tool_enabled":
+        this.smsm.onMcpServerToolEnabled(message.server_name, message.tool);
+        break;
+      case "mcp_server_tool_disabled":
+        this.smsm.onMcpServerToolDisabled(message.server_name, message.tool);
+        break;
+      case "system_prompt_updated":
+        this.systemPrompt = message.system_prompt;
+        break;
+      case "model_updated":
+        this.model = message.model;
+        break;
+      case "session_info":
+        // Update participants if we receive session_info again
+        if ("participants" in message) {
+          this.participants = createSessionParticipantMap(message.participants);
+        }
+        break;
+      case "user_added":
+        this.participants.set(message.user_uuid, {
+          user_uuid: message.user_uuid,
+          nickname: message.nickname,
+          email: message.email,
+          role: message.role,
+        });
+        break;
+      case "user_removed":
+        this.participants.delete(message.user_uuid);
+        break;
+
+      //
+      // Ignore other messages - the owner (the UI layer) can handle them at
+      // its discretion.
+      //
+
+      default:
+        logger.debug(
+          `[handleMessageInternal]: ignoring message: ${message.type}`
+        );
+        break;
+    }
+  }
+}

package/src/chat/client/teamManager.ts

@@ -0,0 +1,29 @@
+import { ClientTeamInfo } from "./chatClient";
+import { AgentSessionData } from "../data/dataModels";
+
+export interface ITeamManager {
+  /**
+   * Get the map of teams, keys are team UUIDs
+   */
+  getTeams(): Map<string, ClientTeamInfo>;
+
+  /**
+   * Get the current team ID
+   */
+  getCurrentTeamId(): string | undefined;
+
+  /**
+   * Get the current team info
+   */
+  getCurrentTeamInfo(): ClientTeamInfo | undefined;
+
+  /**
+   * Set the current team ID
+   */
+  setCurrentTeamId(teamUuid: string): void;
+
+  /**
+   * Get the agent session map of a user
+   */
+  getUserAgentSessionMap(): Map<string, AgentSessionData>;
+}

package/src/chat/{apiKeyManager.ts → data/apiKeyManager.ts}

@@ -1,15 +1,19 @@
 // TODO:
 // - lru-cache
 
-import { Database, UserData } from "./db";
+import { Database, UserData } from "./database";
+import { getLogger } from "@xalia/xmcp/sdk";
+
+const logger = getLogger();
 
 export class ApiKeyManager {
   constructor(private db: Database) {}
 
   public async verifyApiKey(apiKey: string): Promise<UserData | undefined> {
     // TODO: Cache this
-
+    logger.info(`[ApiKeyManager] Verifying API key: ${apiKey}`);
     const userInfo = await this.db.getUserDataFromApiKey(apiKey);
+    logger.info(`[ApiKeyManager] User info: ${JSON.stringify(userInfo)}`);
     return userInfo;
 
     // if (apiKey.startsWith("dummy_key")) {

package/src/chat/data/dataModels.ts

@@ -0,0 +1,85 @@
+import { SavedAgentProfile } from "@xalia/xmcp/sdk";
+import { ChatCompletionMessageParam } from "../../agent/agent";
+
+export type UserMessageData = {
+  message?: string | undefined;
+  imageB64?: string | undefined;
+};
+
+export type TeamRole = "owner" | "participant";
+
+export type TeamParticipant = {
+  user_uuid: string;
+  nickname: string;
+  email: string;
+  role: TeamRole;
+};
+
+export type SessionMessage = {
+  message_idx: number;
+  sender_uuid?: string;
+  is_for_llm: boolean;
+  content: ChatCompletionMessageParam;
+};
+
+export type LLMSessionMessage = {
+  message_idx: number;
+  sender_uuid: string;
+  content: ChatCompletionMessageParam;
+};
+
+/**
+ * Checkpoint data for a specific conversation.
+ */
+export type SessionCheckpoint = {
+  /**
+   * Index of last message included in this checkpoint.
+   */
+  message_idx: number;
+  /**
+   * Summary of the conversation up to and including `message_idx`.
+   */
+  summary: string;
+};
+
+/**
+ * Session data required to create a new session.
+ */
+export type SessionCreateData = {
+  session_uuid: string;
+  title: string;
+  agent_profile_uuid: string;
+  user_uuid: string;
+  team_uuid?: string | undefined;
+  workspace?: UserMessageData | undefined;
+};
+
+/**
+ * Data for existing session.
+ */
+export type SessionData = SessionCreateData & {
+  updated_at: string;
+};
+
+export type AgentSessionData = {
+  agent_profile: SavedAgentProfile;
+  sessions: SessionData[];
+  updated_at: number;
+};
+
+/**
+ * Team sessions for ServerSessionList response
+ */
+export type TeamInfo = {
+  team_uuid: string;
+  team_name: string;
+  owner_uuid: string;
+  participants: Array<TeamParticipant>;
+  sessions: Array<SessionData>;
+  agents: Array<SavedAgentProfile>;
+};
+
+/**
+ * Map from user_uuid to role
+ */
+export type SessionParticipantMap = Map<string, TeamParticipant>;