@xalia/agent 0.5.7 → 0.6.0

package/src/chat/server/server.ts

@@ -0,0 +1,177 @@
+#!/usr/bin/env node
+// -*- typescript -*-
+
+import * as dotenv from "dotenv";
+import { getLogger } from "@xalia/xmcp/sdk";
+import { WebSocketServer } from "ws";
+import * as ws from "ws";
+import { IncomingMessage } from "http";
+import { ConnectionManager, IUserConnectionManager } from "./connectionManager";
+import { Database, resolveCompoundName, UserData } from "../data/database";
+import { SessionData } from "../data/dataModels";
+import { ApiKeyManager } from "../data/apiKeyManager";
+import { ChatFatalError } from "../protocol/errors";
+import { ServerToClient } from "../protocol/messages";
+import { SessionRegistry } from "./sessionRegistry";
+
+dotenv.config();
+
+const logger = getLogger();
+
+/**
+ * Extract error message from unknown error type
+ */
+function extractErrorMessage(e: unknown): string {
+  if (typeof e === "string") return e;
+  if (e instanceof Error) return e.message;
+  return "Unknown connection error";
+}
+
+/**
+ * Send error response and close WebSocket connection
+ */
+function sendErrorAndClose(ws: ws.WebSocket, message: string): void {
+  try {
+    sendErrorResponse(ws, message);
+  } catch (closeError) {
+    logger.error(`[server] Error closing connection:`, closeError);
+  } finally {
+    ws.close(4000, message);
+  }
+}
+
+/**
+ * Send error response but keep connection open
+ */
+function sendErrorResponse(ws: ws.WebSocket, message: string): void {
+  try {
+    ws.send(
+      JSON.stringify({
+        t: "error",
+        e: message,
+      })
+    );
+  } catch (sendError) {
+    logger.error(`[server] Error sending error response:`, sendError);
+  }
+}
+
+export async function runServer(
+  port: number,
+  supabaseUrl: string,
+  supabaseKey: string,
+  llmUrl: string,
+  xmcpUrl: string
+): Promise<ws.Server> {
+  return new Promise((r, _e) => {
+    const wss = new WebSocketServer({ port });
+    const db = new Database(supabaseUrl, supabaseKey);
+    const apiKeyManager = new ApiKeyManager(db);
+    const createSessionRegistry = (
+      connManager: IUserConnectionManager<ServerToClient>
+    ) => {
+      return new SessionRegistry(db, connManager, llmUrl, xmcpUrl);
+    };
+    const connectionManager = ConnectionManager.init(
+      apiKeyManager,
+      createSessionRegistry
+    );
+
+    // eslint-disable-next-line @typescript-eslint/no-misused-promises
+    wss.on("connection", async (ws: ws.WebSocket, req: IncomingMessage) => {
+      try {
+        logger.info(`[server] connection: ${req.url || "unknown"}`);
+        logger.info(`[server] headers: ${JSON.stringify(req.headers)}`);
+
+        // Extract API key and clientMessageId from WebSocket subprotocol header
+        // Format: "apiKey, clientMessageId" (comma-separated)
+        const subprotocols = req.headers["sec-websocket-protocol"];
+        if (!subprotocols) {
+          throw new ChatFatalError("empty api key");
+        }
+
+        // Parse the subprotocols - they come as comma-separated string
+        if (typeof subprotocols !== "string") {
+          throw new ChatFatalError("subprotocols was not a string");
+        }
+        const protocols = subprotocols.split(",").map((p) => p.trim());
+        const apiKey = protocols[0];
+        if (!apiKey) {
+          throw new ChatFatalError("empty api key");
+        }
+
+        // Handle via ConnectionManager (multi-session protocol only)
+        await connectionManager.handleConnection(ws, apiKey, req);
+      } catch (e) {
+        logger.error(`[server] Connection error:`, e);
+
+        // Extract error message consistently
+        const errorMessage = extractErrorMessage(e);
+
+        if (e instanceof ChatFatalError) {
+          // Fatal errors: send error response and close connection
+          sendErrorAndClose(ws, errorMessage);
+        } else {
+          // Non-fatal errors: send error response but keep connection open
+          sendErrorResponse(ws, errorMessage);
+          logger.error(`[server] Client error:`, errorMessage);
+        }
+      }
+    });
+
+    logger.info(`[server] started: ws://localhost:${String(port)}/`);
+
+    r(wss);
+  });
+}
+
+export async function resolveSessionIdFromIdentifier(
+  db: Database,
+  sessionIdentifier: string
+): Promise<string | undefined> {
+  logger.info(`[resolveSessionIdFromIdentifier] ${sessionIdentifier}`);
+  let session: SessionData | undefined = undefined;
+  const compound = resolveCompoundName(sessionIdentifier);
+  logger.info(
+    `[resolveSessionIdFromIdentifier] compound: ${JSON.stringify(compound)}`
+  );
+  if (typeof compound === "string") {
+    // Interpret as an id
+    session = await db.sessionGetById(compound);
+  } else {
+    session = await db.sessionGetByName(compound[0], compound[1]);
+  }
+  logger.info(`[resolveSessionIdFromIdentifier] ${JSON.stringify(session)}`);
+  return session?.session_uuid;
+}
+
+/**
+ * Try as id, then by user/name and return the AgentProfile uuid.
+ */
+export async function resolveAgentProfileId(
+  db: Database,
+  userData: UserData,
+  agentProfileIdentifier: string
+): Promise<string | undefined> {
+  let ap = await db.getSavedAgentProfileById(agentProfileIdentifier);
+  logger.debug(`[resolveAgentProfileId]: by id: {JSON.stringify(ap)}`);
+  if (ap) {
+    return agentProfileIdentifier;
+  }
+
+  ap = await db.getSavedAgentProfileByName(
+    userData.uuid,
+    agentProfileIdentifier
+  );
+  logger.debug(`[resolveAgentProfileId]: by name: {JSON.stringify(ap)}`);
+  if (ap) {
+    return ap.uuid;
+  }
+
+  logger.debug("[resolveAgentProfileId]: agent profile not found");
+  return undefined;
+}
+
+// Export main classes for external use
+export { ConnectionManager } from "./connectionManager";
+export { OpenSession } from "./openSession";

package/src/chat/server/sessionFileManager.ts

@@ -0,0 +1,239 @@
+import { OpenAI } from "openai";
+import { v4 as uuidv4 } from "uuid";
+
+import {
+  Agent,
+  IAgentToolProvider,
+  ToolCallResult,
+  ToolHandler,
+} from "../../agent/agent";
+import { makeParseArgsFn } from "./tools";
+import { Database } from "../data/database";
+
+export const SESSION_FILE_TYPES = ["text", "image", "pdf"] as const;
+
+// SessionFileType = "text" |"image" | "pdf";
+export type SessionFileType = (typeof SESSION_FILE_TYPES)[number];
+
+export type SessionFileDescriptor = {
+  name: string;
+  file_type: SessionFileType;
+  summary?: string;
+};
+
+export type SessionFileEntry = SessionFileDescriptor & { content: string };
+
+/**
+ * Calling code implements this to be informed of file changes.
+ */
+export interface ISessionFileManagerEventHandler {
+  onFileDescriptorChange(desc: SessionFileDescriptor): void;
+  onFileChange(entry: SessionFileEntry): void;
+}
+
+/**
+ * Interface to a set of files in a session.
+ */
+export interface ISessionFileManager {
+  /**
+   * List file names, types and summaries,
+   */
+  listFiles(): SessionFileDescriptor[];
+  /**
+   * Retrieve file contents.
+   */
+  getFileContent(name: string): Promise<string>;
+  /**
+   * Create or update a file with the given name.  Returns the name (which is
+   * created if one is not passed in).
+   */
+  putFileContent(
+    name: string | undefined,
+    file_type: SessionFileType,
+    summary: string | undefined,
+    content: string
+  ): Promise<string>;
+
+  setEventHandler(eventHandler: ISessionFileManagerEventHandler): void;
+}
+
+/**
+ * In-memory implementation of ISessionFileManager
+ */
+export class MemoryFileManager implements ISessionFileManager {
+  private readonly files: Map<string, SessionFileEntry>;
+  private eventHandler: ISessionFileManagerEventHandler | undefined;
+
+  constructor() {
+    this.files = new Map();
+    this.eventHandler = undefined;
+  }
+
+  // ISessionFileManager.listFiles
+  listFiles(): SessionFileDescriptor[] {
+    return Array.from(this.files.values());
+  }
+
+  // ISessionFileManager.getFileContent
+  getFileContent(name: string): Promise<string> {
+    return new Promise((r, e) => {
+      const entry = this.files.get(name);
+      if (entry) {
+        r(entry.content);
+      } else {
+        e(new Error(`no such file ${name}`));
+      }
+    });
+  }
+
+  // ISessionFileManager.addFile
+  putFileContent(
+    name: string | undefined,
+    file_type: SessionFileType,
+    summary: string | undefined,
+    content: string
+  ): Promise<string> {
+    return new Promise((r, e) => {
+      if (!name) {
+        name = uuidv4();
+      }
+      if (!summary) {
+        summary = "";
+      }
+      if (!isSingleLine(summary)) {
+        e(new Error("summary must no contain new-lines"));
+      } else {
+        const entry = { name, file_type, summary, content };
+        this.eventHandler?.onFileChange(entry);
+        this.files.set(name, entry);
+        r(name);
+      }
+    });
+  }
+
+  // ISessionFileManager.setEventHandler
+  setEventHandler(eventHandler: ISessionFileManagerEventHandler) {
+    this.eventHandler = eventHandler;
+  }
+}
+
+/**
+ * Implementation of ISessionFileManager which can store files in the DB.
+ */
+export class ChatSessionFileManager extends MemoryFileManager {
+  // private readonly db: Database;
+  // private readonly sessionUUID: string;
+
+  constructor(_sessionUUID: string, _db: Database) {
+    super();
+    // this.db = db;
+    // this.sessionUUID = sessionUUID;
+  }
+}
+
+/**
+ * Return the file list in a form easily parsable by the LLM.
+ */
+export function listFilesForLLM(fm: ISessionFileManager): string {
+  const files = fm.listFiles();
+  if (files.length === 0) {
+    return "";
+  }
+  let summary = "Available files:\nname,type,summary\n";
+  for (const f of files) {
+    summary += `${f.name},${f.file_type},${f.summary || ""}\n`;
+  }
+  return summary;
+}
+
+const GET_FILE_CONTENT_TOOL: OpenAI.ChatCompletionTool = {
+  type: "function",
+  function: {
+    name: "get_file_content",
+    description: "Obtain the contents of a file listed in the system prompt",
+    parameters: {
+      type: "object",
+      properties: {
+        name: {
+          type: "string",
+          description: "File name",
+        },
+      },
+      required: ["name"],
+    },
+  },
+};
+
+const PUT_FILE_CONTENT_TOOL: OpenAI.ChatCompletionTool = {
+  type: "function",
+  function: {
+    name: "put_file_content",
+    description: "Create or update file content",
+    parameters: {
+      type: "object",
+      properties: {
+        name: {
+          type: "string",
+          description: "File name",
+        },
+        file_type: {
+          type: "string",
+          enum: SESSION_FILE_TYPES,
+        },
+        summary: {
+          type: "string",
+          description: "Content summary",
+        },
+        content: {
+          type: "string",
+          description: "Content (text/dataurl)",
+        },
+      },
+      required: ["name", "type", "summary", "content"],
+    },
+  },
+};
+
+export function fileManagerTool(
+  fileManager: ChatSessionFileManager
+): IAgentToolProvider {
+  // get_file_content
+  const getName = makeParseArgsFn(["name"] as const);
+  const getFileContentFn: ToolHandler = async (
+    _agent: Agent,
+    args: unknown
+  ): Promise<ToolCallResult> => {
+    const { name } = getName(args);
+    const response = await fileManager.getFileContent(name);
+    return { response };
+  };
+
+  // set_file_content
+  const putArgs = ["name", "file_type", "summary", "content"] as const;
+  const getNameSummaryContent = makeParseArgsFn(putArgs);
+  const putFileContentFn: ToolHandler = async (_: Agent, args: unknown) => {
+    const { name, file_type, summary, content } = getNameSummaryContent(args);
+    const response = await fileManager.putFileContent(
+      name,
+      file_type as SessionFileType,
+      summary,
+      content
+    );
+    return { response };
+  };
+
+  const provider = {
+    setup: (agent: Agent) => {
+      agent.addAgentTool(GET_FILE_CONTENT_TOOL, getFileContentFn);
+      agent.addAgentTool(PUT_FILE_CONTENT_TOOL, putFileContentFn);
+      return new Promise<void>((r) => {
+        r();
+      });
+    },
+  };
+  return provider;
+}
+
+function isSingleLine(str: string): boolean {
+  return !/[\r\n]/.test(str);
+}

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

@@ -0,0 +1,138 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { SessionRegistry } from "./sessionRegistry";
+import {
+  createMockDatabase,
+  createMockUserConnectionManager,
+  createMockOpenSession,
+  setupStandardMockBehaviors,
+  MOCK_USERS,
+} from "./test-utils/mockFactories";
+
+// Mock the complex agent creation to avoid network calls
+const mockCreateAgentWithoutSkills = vi.hoisted(() => vi.fn());
+vi.mock("../../agent/agentUtils", () => ({
+  createAgentWithoutSkills: mockCreateAgentWithoutSkills,
+}));
+
+describe("SessionRegistry", () => {
+  let sessionRegistry: SessionRegistry;
+  let mockDatabase: ReturnType<typeof createMockDatabase>;
+  let mockUserConnectionManager: ReturnType<
+    typeof createMockUserConnectionManager
+  >;
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+
+    // Setup the mock to return array with agent and skill manager
+    mockCreateAgentWithoutSkills.mockResolvedValue([
+      {
+        /* mock agent */
+      },
+      {
+        getServerBriefs: vi.fn().mockReturnValue([]),
+        /* mock skill manager */
+      },
+    ]);
+
+    // Create mocks using factories
+    mockDatabase = createMockDatabase();
+    mockUserConnectionManager = createMockUserConnectionManager();
+
+    // Setup standard behaviors
+    setupStandardMockBehaviors({
+      database: mockDatabase,
+      userConnectionManager: mockUserConnectionManager,
+    });
+
+    // Create SessionRegistry instance
+    sessionRegistry = new SessionRegistry(
+      mockDatabase.mock,
+      mockUserConnectionManager.mock,
+      "http://llm-api.test",
+      "http://xmcp-api.test"
+    );
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe("Session Membership Tracking", () => {
+    it("should correctly track user-session mapping", () => {
+      // Arrange - Setup session state using public API
+      const sessionId1 = "session-1";
+      const sessionId2 = "session-2";
+      const userId1 = MOCK_USERS.owner.uuid;
+      const userId2 = MOCK_USERS.participant.uuid;
+
+      // Mock database for session validation
+      mockDatabase.spies.sessionGetById.mockImplementation(
+        (sessionId: string) => {
+          if (sessionId === sessionId1 || sessionId === sessionId2) {
+            return Promise.resolve({
+              session_uuid: sessionId,
+              title: `Test Session`,
+              agent_profile_uuid: "agent-profile-1",
+              owner_uuid: userId1,
+              updated_at: "2025-01-01T00:00:00Z",
+            });
+          }
+          return Promise.resolve(null);
+        }
+      );
+
+      // Create mock sessions in openSessions map
+      const mockSession1 = createMockOpenSession(sessionId1);
+      const mockSession2 = createMockOpenSession(sessionId2);
+      mockSession1.spies.getSessionState.mockReturnValue([]);
+      mockSession2.spies.getSessionState.mockReturnValue([]);
+      // Use bracket notation to access private property for testing
+      const registryWithSessions = sessionRegistry as unknown as {
+        openSessions: Map<string, unknown>;
+      };
+      registryWithSessions.openSessions.set(sessionId1, mockSession1.mock);
+      registryWithSessions.openSessions.set(sessionId2, mockSession2.mock);
+
+      // Act - Use private method via bracket notation for testing
+      type SessionRegistryWithPrivate = {
+        addUserToSessionMemory: (userId: string, sessionId: string) => void;
+      };
+      const registryWithPrivate =
+        sessionRegistry as unknown as SessionRegistryWithPrivate;
+      registryWithPrivate.addUserToSessionMemory(userId1, sessionId1);
+      registryWithPrivate.addUserToSessionMemory(userId2, sessionId1);
+      registryWithPrivate.addUserToSessionMemory(userId1, sessionId2);
+
+      // Assert - Test getUserSessions for both users
+      const user1Sessions = sessionRegistry.getUserSessions(userId1);
+      expect(user1Sessions).toEqual(new Set([sessionId1, sessionId2]));
+
+      const user2Sessions = sessionRegistry.getUserSessions(userId2);
+      expect(user2Sessions).toEqual(new Set([sessionId1]));
+
+      // Assert - Test getInMemorySessionUsers for both sessions
+      const session1Users = sessionRegistry.getInMemorySessionUsers(sessionId1);
+      expect(session1Users).toEqual(new Set([userId1, userId2]));
+      expect(session1Users.size).toBe(2);
+
+      const session2Users = sessionRegistry.getInMemorySessionUsers(sessionId2);
+      expect(session2Users).toEqual(new Set([userId1]));
+      expect(session2Users.size).toBe(1);
+    });
+
+    it("should return empty sets for non-existent sessions or users", () => {
+      // Act & Assert - Test methods with non-existent data
+      const nonExistentSessionUsers = sessionRegistry.getInMemorySessionUsers(
+        "non-existent-session"
+      );
+      expect(nonExistentSessionUsers).toEqual(new Set());
+      expect(nonExistentSessionUsers.size).toBe(0);
+
+      const nonExistentUserSessions =
+        sessionRegistry.getUserSessions("non-existent-user");
+      expect(nonExistentUserSessions).toEqual(new Set());
+      expect(nonExistentUserSessions.size).toBe(0);
+    });
+  });
+});