@xalia/agent 0.6.1 → 0.6.3

package/src/chat/server/test-utils/mockFactories.ts

@@ -128,21 +128,18 @@ export function createMockUserConnectionManager(): {
   mock: IUserConnectionManager<ServerToClient>;
   spies: {
     sendToUsers: ReturnType<typeof vi.fn>;
-    getLiveUserApiKey: ReturnType<typeof vi.fn>;
   };
 } {
   const spies = {
     sendToUsers: vi.fn(),
     sendToConnection: vi.fn(),
     sendServerError: vi.fn(),
-    getLiveUserApiKey: vi.fn(),
   };
 
   const mock = {
     sendToUsers: spies.sendToUsers,
     sendToConnection: spies.sendToConnection,
     sendServerError: spies.sendServerError,
-    getLiveUserApiKey: spies.getLiveUserApiKey,
   } as IUserConnectionManager<ServerToClient>;
 
   return { mock, spies };
@@ -154,16 +151,25 @@ export function createMockUserConnectionManager(): {
 export function createMockSessionRegistry(): {
   mock: IMessageProcessor<ClientToServer>;
   spies: {
+    authenticate: ReturnType<typeof vi.fn>;
     processMessage: ReturnType<typeof vi.fn>;
     handleUserDisconnect: ReturnType<typeof vi.fn>;
   };
 } {
   const spies = {
+    authenticate: vi.fn().mockImplementation((apiKey) => {
+      if (apiKey === "valid-api-key")
+        return Promise.resolve(MOCK_USERS.owner.uuid);
+      if (apiKey === "participant-api-key")
+        return Promise.resolve(MOCK_USERS.participant.uuid);
+      return Promise.resolve(null);
+    }),
     processMessage: vi.fn(),
     handleUserDisconnect: vi.fn(),
   };
 
   const mock = {
+    authenticate: spies.authenticate,
     processMessage: spies.processMessage,
     handleUserDisconnect: spies.handleUserDisconnect,
   } as IMessageProcessor<ClientToServer>;
@@ -267,6 +273,7 @@ export function createMockSessionList(): Array<SessionData> {
       workspace: undefined,
       updated_at: MOCK_SESSIONS.active.updated_at || "",
       user_uuid: MOCK_SESSIONS.active.user_uuid,
+      agent_paused: false,
     },
     {
       session_uuid: MOCK_SESSIONS.secondary.uuid,
@@ -276,6 +283,7 @@ export function createMockSessionList(): Array<SessionData> {
       workspace: undefined,
       updated_at: MOCK_SESSIONS.secondary.updated_at || "",
       user_uuid: MOCK_SESSIONS.secondary.user_uuid,
+      agent_paused: false,
     },
   ];
 }
@@ -409,14 +417,6 @@ export function setupStandardMockBehaviors(mocks: {
   }
 
   // Setup user connection manager mocks
-  if (mocks.userConnectionManager) {
-    mocks.userConnectionManager.spies.getLiveUserApiKey.mockImplementation(
-      (userId: string) => {
-        if (userId === MOCK_USERS.owner.uuid) return "valid-api-key";
-        if (userId === MOCK_USERS.participant.uuid)
-          return "participant-api-key";
-        return undefined;
-      }
-    );
-  }
+  // if (mocks.userConnectionManager) {
+  // }
 }

package/src/chat/server/tools.ts

@@ -10,6 +10,8 @@ import { IPlatform } from "../../agent/iplatform";
 import { htmlToText } from "../utils/htmlToText";
 import { getLogger } from "@xalia/xmcp/sdk";
 import { webSearch } from "../utils/search";
+import { ChatSessionFileManager, fileManagerTool } from "./sessionFileManager";
+import { genImageFileTool } from "./imageGeneratorTools";
 
 const logger = getLogger();
 
@@ -17,10 +19,19 @@ const logger = getLogger();
  * Returns a function which parses an `args` struct and attempts to extract
  * multiple string parameters with the given names.  e.g.
  *
- *   const parseFn = makeParseArgsFn(["arg0", "arg1"] as const)
+ *   const parseFn = makeParseArgsFn(
+ *     ["arg0", "arg1"] as const,
+ *     ["opt0"] as const)
  *
- * creates `parseFn: (args: unknown) => { arg0: string, arg1: string }` which
- * can be used to parse tool arguments.
+ * creates
+ *
+ *   parseFn: (args: unknown) => {
+ *     arg0: string,
+ *     arg1: string,
+ *     opt0: string|undefined
+ *   }
+ *
+ * which can be used to parse tool arguments.
  *
  * NOTE, the complex type parameters ensures that the name list is a
  * compile-time value, which in turn ensures that the return value of this
@@ -28,19 +39,36 @@ const logger = getLogger();
  */
 export function makeParseArgsFn<
   T extends readonly string[] & (string extends T[number] ? never : unknown),
->(names: T): (args: unknown) => { [K in T[number]]: string } {
+  U extends readonly string[] & (string extends U[number] ? never : unknown),
+>(
+  names: T,
+  optNames?: U
+): (
+  args: unknown
+) => { [K in T[number]]: string } & { [K in U[number]]: string | undefined } {
   return (args: unknown) => {
-    if (typeof args !== "object") {
+    if (!args || typeof args !== "object") {
       throw new Error(`invalid args: ${typeof args}`);
     }
-    const argsObj = args as Record<string, string>;
+    const argsObj = args as Record<string, string | undefined>;
     for (const name of names) {
       const val = argsObj[name];
       if (typeof val !== "string") {
         throw new Error(`invalid expr args.${name}: ${typeof val}`);
       }
     }
-    return argsObj as { [K in T[number]]: string };
+    if (optNames) {
+      for (const name of optNames) {
+        const val = argsObj[name];
+        if (typeof val !== "undefined" && typeof val !== "string") {
+          throw new Error(`invalid expr args.${name}: ${typeof val}`);
+        }
+      }
+    }
+
+    return argsObj as { [K in T[number]]: string } & {
+      [K in U[number]]: string | undefined;
+    };
   };
 }
 
@@ -255,11 +283,18 @@ export function openURLTool(): IAgentToolProvider {
 export async function addDefaultChatTools(
   agent: Agent,
   timezone: string,
-  platform: IPlatform
+  platform: IPlatform,
+  fileManager: ChatSessionFileManager,
+  llmUrl: string,
+  llmApiKey: string
 ): Promise<void> {
   await agent.addAgentToolProvider(datetimeTool(timezone));
   await agent.addAgentToolProvider(calculatorTool);
   await agent.addAgentToolProvider(renderTool(platform));
   await agent.addAgentToolProvider(webSearchTool());
   await agent.addAgentToolProvider(openURLTool());
+  await agent.addAgentToolProvider(fileManagerTool(fileManager));
+  await agent.addAgentToolProvider(
+    await genImageFileTool(llmUrl, llmApiKey, fileManager)
+  );
 }

package/src/chat/utils/agentSessionMap.ts

@@ -1,5 +1,5 @@
 import { SavedAgentProfile, getLogger } from "@xalia/xmcp/sdk";
-import { SessionData, AgentSessionData } from "../data/dataModels";
+import { AgentSessionData, SessionDescriptor } from "../data/dataModels";
 
 const logger = getLogger();
 
@@ -23,7 +23,7 @@ export function emptyAgentProfile(agentUuid: string): SavedAgentProfile {
 
 // build agentSessionMap from sessions and agents
 export function buildAgentSessionMap(
-  sessions: Map<string, SessionData>,
+  sessions: Map<string, SessionDescriptor>,
   agents: Map<string, SavedAgentProfile>
 ): Map<string, AgentSessionData> {
   const agentSessionMap: Map<string, AgentSessionData> = new Map();

package/src/chat/utils/approvalManager.ts

@@ -1,108 +1,180 @@
-import { getLogger } from "@xalia/xmcp/sdk";
+import {
+  AgentPreferences,
+  getLogger,
+  prefsGetAutoApprove,
+  prefsSetAutoApprove,
+} from "@xalia/xmcp/sdk";
+import { ResponseAwaiter } from "./responseAwaiter";
+import {
+  ClientToolCallApprovalResult,
+  ServerToClient,
+  ServerToolAutoApprovalSet,
+} from "../protocol/messages";
+import { Database } from "../data/database";
+import { ChatCompletionMessageToolCall } from "../../agent/llm";
+import { ISessionMessageSender } from "../server/openSessionMessageSender";
 
 const logger = getLogger();
 
-export type ApprovalResult = {
-  approved: boolean;
-  auto_approve: boolean;
-};
-
-/**
- * Thrown in the resultP promise when an approval times out.
- */
-export class ApprovalTimeout extends Error {
-  constructor(message: string) {
-    super(message);
-    this.name = "ApprovalTimeout";
-  }
+export interface IAgentPreferencesWriter {
+  updatePreferences(
+    agentProfileUUID: string,
+    settings: AgentPreferences
+  ): Promise<void>;
 }
 
-export class ApprovalCancelled extends Error {
-  constructor(message: string) {
-    super(message);
-    this.name = "ApprovalCancelled";
+export class DbAgentPreferencesWriter implements IAgentPreferencesWriter {
+  constructor(private db: Database) {}
+
+  updatePreferences(
+    agentProfileUUID: string,
+    preferences: AgentPreferences
+  ): Promise<void> {
+    return this.db.updateAgentProfilePreferences(agentProfileUUID, preferences);
   }
 }
 
-type ApprovalData = {
-  resolve: (result: ApprovalResult) => void;
-  error: (e: Error) => void;
-  timeoutId?: NodeJS.Timeout;
-};
-
 /**
- * The caller initiates an approval for a specific server, and a unique ID is
- * generated for it, along with a promise for the resolution of the approval.
- * The caller returns the ID to some client(s) and waits on the promise.
- *
- * When an approval or rejection (or timeout) is received, the promise is
- * resolved.
+ * Handles an in-memory caching / updating of the auto-approve settings for
+ * tool calls.  Also handles querying the client for approval and waiting for
+ * responses.
  */
-export class ApprovalManager {
-  private approvals = new Map<string, ApprovalData>();
-  private timeoutMs?: number;
+export class ToolApprovalManager {
+  private sessionUUID: string;
+  private agentProfileUUID: string;
+  private agentProfilePreferences: AgentPreferences;
+  private sender: ISessionMessageSender<ServerToClient>;
+  private writer: IAgentPreferencesWriter;
+  private responseAwaiter: ResponseAwaiter<ClientToolCallApprovalResult>;
 
-  constructor(timeoutMs?: number) {
-    this.timeoutMs = timeoutMs;
+  constructor(
+    sessionUUID: string,
+    agentProfileUUID: string,
+    agentProfilePreferences: AgentPreferences,
+    sender: ISessionMessageSender<ServerToClient>,
+    writer: IAgentPreferencesWriter,
+    timeoutMs?: number
+  ) {
+    this.sessionUUID = sessionUUID;
+    this.agentProfileUUID = agentProfileUUID;
+    this.agentProfilePreferences = agentProfilePreferences;
+    this.sender = sender;
+    this.writer = writer;
+    this.responseAwaiter = ResponseAwaiter.init(
+      undefined,
+      (msg) => msg.id,
+      timeoutMs
+    );
   }
 
-  public shutdown(): void {
-    for (const [id, approval] of this.approvals) {
-      if (approval.timeoutId) {
-        clearTimeout(approval.timeoutId);
-      }
-      approval.error(new ApprovalCancelled("shutdown"));
-      this.approvals.delete(id);
+  /**
+   * Check for auto-approval, or query the client.  Handle approval response
+   * (or timeout) and update auto-approval settings.
+   *
+   * The returned `requested` value indicates whether approval was requested.
+   */
+  public async getApproval(
+    serverName: string,
+    tool: string,
+    toolCall: ChatCompletionMessageToolCall
+  ): Promise<{ approved: boolean; requested: boolean }> {
+    const autoApproved = prefsGetAutoApprove(
+      this.agentProfilePreferences,
+      serverName,
+      tool
+    );
+    if (autoApproved) {
+      return { approved: true, requested: false };
     }
-  }
 
-  public startApproval(name: string): {
-    id: string;
-    resultP: Promise<ApprovalResult>;
-  } {
-    const id = this.generateUniqueId(name);
-    const resultP = new Promise<ApprovalResult>((resolve, error) => {
-      let timeoutId;
-      if (this.timeoutMs) {
-        timeoutId = setTimeout(() => {
-          const approval = this.approvals.get(id);
-          if (approval) {
-            this.approvals.delete(id);
-            error(new ApprovalTimeout(`approval ${id} (${name}) timed out`));
-          }
-        }, this.timeoutMs);
+    // Query the owner for approval
+
+    const id = this.generateUniqueId(toolCall.function.name);
+    try {
+      const approvalP = this.responseAwaiter.waitForResponse(id);
+
+      this.sender.broadcast({
+        type: "approve_tool_call",
+        id,
+        tool_call: toolCall,
+        session_id: this.sessionUUID,
+      });
+
+      logger.debug(`[ApprovalManager.getApproval] awaiting approval ${id}`);
+      const approval = await approvalP;
+      logger.debug(
+        `[ApprovalManager.getApproval] approval ${JSON.stringify(approval)}`
+      );
+
+      // Handle any auto-approve update, informing other clients.
+
+      if (approval.auto_approve) {
+        logger.debug("[ApprovalManager.getApproval] updated preferences");
+        const autoApprovalMsg = await this.setAutoApprove(
+          serverName,
+          tool,
+          true
+        );
+        if (autoApprovalMsg) {
+          this.sender.broadcast(autoApprovalMsg);
+        }
       }
 
-      this.approvals.set(id, { resolve, error, timeoutId });
-    });
+      // Broadcast the result of the approval
+
+      this.sender.broadcast({
+        type: "tool_call_approval_result",
+        id: approval.id,
+        result: approval.result,
+        session_id: this.sessionUUID,
+      });
 
-    logger.debug(`new approval ${id}`);
-    return { id, resultP };
+      return { approved: approval.result, requested: true };
+    } catch (e) {
+      logger.debug(
+        `[OpenSession.onToolCall] error waiting for approval ${id}: ` +
+          String(e)
+      );
+      throw e;
+    }
   }
 
   /**
-   * Returns true if this result was accepted.  False if the approval was
-   * already answered.
+   * Handle a request to set auto-approval for a given tool.  If there was a
+   * change, return the message to be broadcast.
    */
-  public approvalResult(
-    id: string,
-    approved: boolean,
-    auto_approve: boolean
-  ): boolean {
-    const approval = this.approvals.get(id);
-    if (approval) {
-      logger.debug(`approval ${id} present.  resolving ${String(approved)}`);
-
-      if (approval.timeoutId) {
-        clearTimeout(approval.timeoutId);
-      }
-      this.approvals.delete(id);
-      approval.resolve({ approved, auto_approve });
-      return true;
+  public async setAutoApprove(
+    serverName: string,
+    tool: string,
+    autoApprove: boolean
+  ): Promise<ServerToolAutoApprovalSet | undefined> {
+    if (
+      prefsSetAutoApprove(
+        this.agentProfilePreferences,
+        serverName,
+        tool,
+        autoApprove
+      )
+    ) {
+      await this.writer.updatePreferences(
+        this.agentProfileUUID,
+        this.agentProfilePreferences
+      );
+      return {
+        type: "tool_auto_approval_set",
+        server_name: serverName,
+        tool,
+        auto_approve: autoApprove,
+        session_id: this.sessionUUID,
+      };
     }
+  }
 
-    logger.debug(`approval ${id} not present`);
-    return false;
+  /**
+   * Forward all approval result messages here.
+   */
+  public onApprovalResult(msg: ClientToolCallApprovalResult): void {
+    this.responseAwaiter.onMessage(msg);
   }
 
   private generateUniqueId(tag: string): string {

package/src/chat/utils/multiAsyncQueue.ts

@@ -3,6 +3,7 @@ export class MultiAsyncQueue<T> {
   private process: (queueEntry: T[]) => Promise<void>;
   private maxBacklog: number;
   private running: boolean = false;
+  private paused: boolean = false;
 
   constructor(
     process: (queueEntry: T[]) => Promise<void>,
@@ -30,8 +31,17 @@ export class MultiAsyncQueue<T> {
     return true;
   }
 
+  public pause() {
+    this.paused = true;
+  }
+
+  public unpause() {
+    this.paused = false;
+    setTimeout(() => void this.tryNext(), 0);
+  }
+
   private async tryNext() {
-    if (this.running) {
+    if (this.running || this.paused) {
       return;
     }
 

package/src/chat/{client/responseHandler.ts → utils/responseAwaiter.ts}

@@ -2,12 +2,13 @@ import { getLogger } from "@xalia/xmcp/sdk";
 
 const DEFAULT_TIMEOUT_MS: number = 10000;
 
-type ClientMsg = { client_message_id: string };
-type ServerMsg = { type: string; client_message_id?: string };
+type ServerMsg = { type: string };
 // Messages of type S which also have a message field.  The message
 // representing an error should be of this type.
 type ServerErr<S> = Extract<S, { message: string }>;
 
+type IdExtractor<S extends ServerMsg> = (msg: S) => string | undefined;
+
 type WaitingEntry<ServerMessageT> = {
   resolve: (s: ServerMessageT) => void;
   error: (e: Error) => void;
@@ -16,13 +17,19 @@ type WaitingEntry<ServerMessageT> = {
 
 const logger = getLogger();
 
+function defaultImmediate<T>(x: T): T {
+  return x;
+}
+
 /**
  *
  * Handles response messages and timeouts for client request messages.
  *
- * Create a ResponseHandler for a specific class of queries
+ * Create a ResponseAwaiter for a specific class of queries
  *
- *   this.responseHandler = new ResponseHandler<SomeRequest, SomeResponse>()
+ *   this.responseHandler = new ResponseAwaiter<SomeRequest, SomeResponse>(
+ *     (msg: SomeResponse) => msg.client_message_id;
+ *   )
  *
  * Use as follows:
  *
@@ -36,14 +43,15 @@ const logger = getLogger();
  *
  *     // Get a Promise representing the response to this message, which
  *     // we can await.
- *     const response = await this.responseHandler.waitForResponse(msg);
+ *     const response =
+ *       await this.responseHandler.waitForResponse(client_message_id);
  *
  *     // Perform any processing on the response
  *     return response.response_data;
  *   }
  *   ```
  *
- * ResponseHandler must be informed of relevant messages in order to resolve
+ * ResponseAwaiter must be informed of relevant messages in order to resolve
  * responses:
  *
  *   ```
@@ -57,46 +65,82 @@ const logger = getLogger();
  *     }
  *   }
  *   ```
+ *
+ * If some actions need to happen immediately on receipt of the message,
+ * instead of when the `Promise.resolve` function is resolved, add an
+ * `immediateAction` callback, which can optionally transfrom `ServerMsgT`
+ * into `FinalResponseT` to be passed back by `waitForResponse`.
  */
-export class ResponseHandler<
-  ClientMsgT extends ClientMsg,
+export class ResponseAwaiter<
   ServerMsgT extends ServerMsg,
+  FinalResponseT = ServerMsgT,
 > {
-  readonly waiting: Map<string, WaitingEntry<ServerMsgT>>;
+  readonly waiting: Map<string, WaitingEntry<FinalResponseT>>;
+  readonly idExtractor: IdExtractor<ServerMsgT>;
+  readonly immediateAction: (x: ServerMsgT) => FinalResponseT;
   readonly timeoutMS: number;
   readonly errorType: ServerErr<ServerMsgT>["type"] | undefined;
 
   /**
    * errorType: the type field of the message which represents an error.
    */
-  constructor(
+  private constructor(
     // value of the `type` field of an error message type,
     // e.g. "session_error". Compiler should ensure that the `ServerMsgT` with
     // this `type` field at least has a `message` field.
     errorType: undefined | ServerErr<ServerMsgT>["type"],
+    idExtractor: IdExtractor<ServerMsgT>,
+    immediateAction: (x: ServerMsgT) => FinalResponseT,
     timeoutMS: number = DEFAULT_TIMEOUT_MS
   ) {
     this.waiting = new Map();
+    this.idExtractor = idExtractor;
+    this.immediateAction = immediateAction;
     this.timeoutMS = timeoutMS;
     this.errorType = errorType;
   }
 
+  static init<S extends ServerMsg>(
+    errorType: undefined | ServerErr<S>["type"],
+    idExtractor: IdExtractor<S>,
+    timeoutMS: number = DEFAULT_TIMEOUT_MS
+  ): ResponseAwaiter<S, S> {
+    return new ResponseAwaiter(
+      errorType,
+      idExtractor,
+      defaultImmediate,
+      timeoutMS
+    );
+  }
+
+  static initWithImmediate<S extends ServerMsg, F>(
+    errorType: undefined | ServerErr<S>["type"],
+    idExtractor: IdExtractor<S>,
+    immediateAction: (x: S) => F,
+    timeoutMS: number = DEFAULT_TIMEOUT_MS
+  ): ResponseAwaiter<S, F> {
+    return new ResponseAwaiter(
+      errorType,
+      idExtractor,
+      immediateAction,
+      timeoutMS
+    );
+  }
+
   /**
    * Given a request message, return a promise representing the corresponding
    * response message from the server.
    */
-  waitForResponse(msg: ClientMsgT): Promise<ServerMsgT> {
-    return new Promise<ServerMsgT>((resolve, error) => {
-      const msgId = msg.client_message_id;
-
+  waitForResponse(msgId: string): Promise<FinalResponseT> {
+    return new Promise<FinalResponseT>((resolve, error) => {
       const timeoutId = setTimeout(() => {
         const waiting = this.waiting.get(msgId);
         if (!waiting) {
-          logger.warn(`[ResponseHandler] timeout for ${msgId} with no entry`);
+          logger.warn(`[ResponseAwaiter] timeout for ${msgId} with no entry`);
           return;
         }
         this.waiting.delete(msgId);
-        logger.warn(`[ResponseHandler] timeout for client_msg_id ${msgId}`);
+        logger.warn(`[ResponseAwaiter] timeout for client_msg_id ${msgId}`);
         error(new Error(`timeout for client_msg_id ${msgId}`));
       }, this.timeoutMS);
 
@@ -104,20 +148,25 @@ export class ResponseHandler<
     });
   }
 
+  waitingForId(msgId: string): boolean {
+    return this.waiting.has(msgId);
+  }
+
   /**
-   * Pass response (and error) messages in here.
+   * Pass response (and error) messages in here.  Returns `true` if the
+   * message was consumed.  Otherwise `false`.
    */
-  onMessage(msg: ServerMsgT): void {
-    const msgId = msg.client_message_id;
+  onMessage(msg: ServerMsgT): boolean {
+    const msgId = this.idExtractor(msg);
     if (!msgId) {
-      return;
+      return false;
     }
     const waiting = this.waiting.get(msgId);
     if (!waiting) {
       logger.warn(
-        `[ResponseHandler] resolve for ${msgId} with no entry (timeout?)`
+        `[ResponseAwaiter] resolve for ${msgId} with no entry (timeout?)`
       );
-      return;
+      return false;
     }
 
     clearTimeout(waiting.timeoutId);
@@ -125,7 +174,8 @@ export class ResponseHandler<
     if (this.errorType && msg.type === this.errorType) {
       waiting.error(new Error((msg as ServerErr<ServerMsgT>).message));
     } else {
-      waiting.resolve(msg);
+      waiting.resolve(this.immediateAction(msg));
     }
+    return true;
   }
 }