@vellumai/assistant 0.7.0 → 0.7.1

package/src/__tests__/events-client-registration.test.ts

@@ -3,14 +3,16 @@
  * headers on the /events endpoint.
  *
  * Validates:
- *   - Client is registered in ClientRegistry on SSE connect
+ *   - Client is registered as a hub subscriber on SSE connect
  *   - Client is unregistered on SSE disconnect (abort)
  *   - Client is touched on heartbeat interval
  *   - Missing interfaceId with clientId throws BadRequestError
  *   - Invalid interfaceId throws BadRequestError
  *   - Missing both headers skips registration (backwards compat)
+ *   - Duplicate clientId subscribers are deduplicated on reconnect
+ *   - disposeClient() force-disconnects all subscribers for a clientId
  */
-import { beforeEach, describe, expect, mock, test } from "bun:test";
+import { describe, expect, mock, test } from "bun:test";
 
 mock.module("../util/logger.js", () => ({
   getLogger: () =>
@@ -32,10 +34,6 @@ mock.module("../config/loader.js", () => ({
 
 import { initializeDb } from "../memory/db-init.js";
 import { AssistantEventHub } from "../runtime/assistant-event-hub.js";
-import {
-  __resetClientRegistryForTests,
-  getClientRegistry,
-} from "../runtime/client-registry.js";
 import {
   BadRequestError,
   ServiceUnavailableError,
@@ -45,10 +43,6 @@ import { handleSubscribeAssistantEvents } from "../runtime/routes/events-routes.
 initializeDb();
 
 describe("events client registration", () => {
-  beforeEach(() => {
-    __resetClientRegistryForTests();
-  });
-
   // ── Registration on connect ───────────────────────────────────────────────
 
   test("registers client when both headers are provided", () => {
@@ -68,11 +62,12 @@ describe("events client registration", () => {
 
     expect(stream).toBeInstanceOf(ReadableStream);
 
-    const registry = getClientRegistry();
-    const entry = registry.get("test-mac-001");
+    const clients = hub.listClients();
+    const entry = clients.find((c) => c.clientId === "test-mac-001");
     expect(entry).toBeDefined();
-    expect(entry!.interfaceId).toBe("macos");
-    expect(entry!.capabilities).toContain("host_bash");
+    expect(entry?.interfaceId).toBe("macos");
+    expect(entry?.capabilities).toContain("host_bash");
+    expect(entry?.type).toBe("client");
 
     ac.abort();
   });
@@ -83,7 +78,7 @@ describe("events client registration", () => {
 
     handleSubscribeAssistantEvents({ abortSignal: ac.signal }, { hub });
 
-    expect(getClientRegistry().size).toBe(0);
+    expect(hub.listClients()).toHaveLength(0);
 
     ac.abort();
   });
@@ -100,7 +95,7 @@ describe("events client registration", () => {
       { hub },
     );
 
-    expect(getClientRegistry().size).toBe(0);
+    expect(hub.listClients()).toHaveLength(0);
 
     ac.abort();
   });
@@ -118,7 +113,7 @@ describe("events client registration", () => {
         { hub },
       ),
     ).toThrow(BadRequestError);
-    expect(getClientRegistry().size).toBe(0);
+    expect(hub.listClients()).toHaveLength(0);
   });
 
   test("throws BadRequestError when interfaceId is invalid", () => {
@@ -135,7 +130,7 @@ describe("events client registration", () => {
         { hub },
       ),
     ).toThrow(BadRequestError);
-    expect(getClientRegistry().size).toBe(0);
+    expect(hub.listClients()).toHaveLength(0);
   });
 
   // ── Unregistration on disconnect ──────────────────────────────────────────
@@ -155,21 +150,20 @@ describe("events client registration", () => {
       { hub },
     );
 
-    const registry = getClientRegistry();
-    expect(registry.get("test-mac-002")).toBeDefined();
+    expect(hub.listClients().some((c) => c.clientId === "test-mac-002")).toBe(
+      true,
+    );
 
-    // Start reading so start() runs and the abort listener is installed
     const reader = stream.getReader();
-    // Consume the initial heartbeat
     await reader.read();
 
-    // Abort the request — simulates client disconnect
     ac.abort();
 
-    // Allow microtasks to flush
     await new Promise((r) => setTimeout(r, 10));
 
-    expect(registry.get("test-mac-002")).toBeUndefined();
+    expect(hub.listClients().some((c) => c.clientId === "test-mac-002")).toBe(
+      false,
+    );
   });
 
   test("unregisters client when stream is cancelled", async () => {
@@ -187,22 +181,23 @@ describe("events client registration", () => {
       { hub },
     );
 
-    const registry = getClientRegistry();
-    expect(registry.get("test-mac-003")).toBeDefined();
+    expect(hub.listClients().some((c) => c.clientId === "test-mac-003")).toBe(
+      true,
+    );
 
-    // Cancel the stream directly
     await stream.cancel();
 
-    expect(registry.get("test-mac-003")).toBeUndefined();
+    expect(hub.listClients().some((c) => c.clientId === "test-mac-003")).toBe(
+      false,
+    );
   });
 
   // ── Heartbeat touch ───────────────────────────────────────────────────────
 
-  test("touches client registry on heartbeat", async () => {
+  test("touches client on heartbeat", async () => {
     const ac = new AbortController();
     const hub = new AssistantEventHub();
 
-    // Use a very short heartbeat interval for testing
     const stream = handleSubscribeAssistantEvents(
       {
         headers: {
@@ -214,20 +209,19 @@ describe("events client registration", () => {
       { hub, heartbeatIntervalMs: 50 },
     );
 
-    const registry = getClientRegistry();
-    const entry = registry.get("test-mac-004");
+    const clients = hub.listClients();
+    const entry = clients.find((c) => c.clientId === "test-mac-004");
     expect(entry).toBeDefined();
-    const initialActive = entry!.lastActiveAt;
+    const initialActive = entry?.lastActiveAt.getTime() ?? 0;
 
-    // Read the initial heartbeat to ensure start() has run
     const reader = stream.getReader();
     await reader.read();
 
-    // Wait for at least one heartbeat cycle
     await new Promise((r) => setTimeout(r, 100));
 
-    // lastActiveAt should have been touched
-    expect(entry!.lastActiveAt).toBeGreaterThanOrEqual(initialActive);
+    // Re-query — the entry object is the same reference, so lastActiveAt
+    // should have been bumped by touchClient().
+    expect(entry?.lastActiveAt.getTime()).toBeGreaterThanOrEqual(initialActive);
 
     ac.abort();
   });
@@ -235,8 +229,6 @@ describe("events client registration", () => {
   // ── Eviction cleanup ──────────────────────────────────────────────────────
 
   test("unregisters client when evicted by hub capacity limit", async () => {
-    // Create a hub with capacity 1 — the first subscriber is evicted when
-    // the second subscribes.
     const hub = new AssistantEventHub({ maxSubscribers: 1 });
 
     const ac1 = new AbortController();
@@ -251,10 +243,8 @@ describe("events client registration", () => {
       { hub },
     );
 
-    const registry = getClientRegistry();
-    expect(registry.get("evict-me")).toBeDefined();
+    expect(hub.listClients().some((c) => c.clientId === "evict-me")).toBe(true);
 
-    // Second subscriber evicts the first
     const ac2 = new AbortController();
     handleSubscribeAssistantEvents(
       {
@@ -267,8 +257,10 @@ describe("events client registration", () => {
       { hub },
     );
 
-    expect(registry.get("evict-me")).toBeUndefined();
-    expect(registry.get("i-stay")).toBeDefined();
+    expect(hub.listClients().some((c) => c.clientId === "evict-me")).toBe(
+      false,
+    );
+    expect(hub.listClients().some((c) => c.clientId === "i-stay")).toBe(true);
 
     ac1.abort();
     ac2.abort();
@@ -276,8 +268,7 @@ describe("events client registration", () => {
 
   // ── Capacity limit cleanup ────────────────────────────────────────────────
 
-  test("cleans up registration when hub subscribe throws RangeError", () => {
-    // A hub with 0 capacity cannot accept any subscribers.
+  test("throws ServiceUnavailableError when hub has zero capacity", () => {
     const hub = new AssistantEventHub({ maxSubscribers: 0 });
 
     expect(() =>
@@ -291,7 +282,160 @@ describe("events client registration", () => {
         { hub },
       ),
     ).toThrow(ServiceUnavailableError);
-    // Should have been cleaned up
-    expect(getClientRegistry().get("no-room")).toBeUndefined();
+    expect(hub.listClients()).toHaveLength(0);
+  });
+
+  // ── Client deduplication on reconnect ─────────────────────────────────────
+
+  test("deduplicates stale subscribers when same clientId reconnects", () => {
+    const hub = new AssistantEventHub();
+    const ac1 = new AbortController();
+    const ac2 = new AbortController();
+
+    handleSubscribeAssistantEvents(
+      {
+        headers: {
+          "x-vellum-client-id": "dedup-001",
+          "x-vellum-interface-id": "chrome-extension",
+        },
+        abortSignal: ac1.signal,
+      },
+      { hub },
+    );
+
+    expect(hub.listClients()).toHaveLength(1);
+    expect(hub.subscriberCount()).toBe(1);
+
+    handleSubscribeAssistantEvents(
+      {
+        headers: {
+          "x-vellum-client-id": "dedup-001",
+          "x-vellum-interface-id": "chrome-extension",
+        },
+        abortSignal: ac2.signal,
+      },
+      { hub },
+    );
+
+    expect(hub.listClients()).toHaveLength(1);
+    expect(hub.subscriberCount()).toBe(1);
+    expect(hub.listClients()[0]?.clientId).toBe("dedup-001");
+
+    ac1.abort();
+    ac2.abort();
+  });
+
+  test("deduplication evicts stale entry via onEvict callback", () => {
+    const hub = new AssistantEventHub();
+    const ac1 = new AbortController();
+    const ac2 = new AbortController();
+    let evicted = false;
+
+    hub.subscribe({
+      type: "client" as const,
+      clientId: "evict-cb-001",
+      interfaceId: "chrome-extension",
+      capabilities: ["host_browser"],
+      callback: () => {},
+      onEvict: () => {
+        evicted = true;
+      },
+    });
+
+    expect(evicted).toBe(false);
+
+    handleSubscribeAssistantEvents(
+      {
+        headers: {
+          "x-vellum-client-id": "evict-cb-001",
+          "x-vellum-interface-id": "chrome-extension",
+        },
+        abortSignal: ac2.signal,
+      },
+      { hub },
+    );
+
+    expect(evicted).toBe(true);
+    expect(hub.listClients()).toHaveLength(1);
+
+    ac1.abort();
+    ac2.abort();
+  });
+
+  test("different clientIds are not deduplicated", () => {
+    const hub = new AssistantEventHub();
+    const ac1 = new AbortController();
+    const ac2 = new AbortController();
+
+    handleSubscribeAssistantEvents(
+      {
+        headers: {
+          "x-vellum-client-id": "client-A",
+          "x-vellum-interface-id": "chrome-extension",
+        },
+        abortSignal: ac1.signal,
+      },
+      { hub },
+    );
+
+    handleSubscribeAssistantEvents(
+      {
+        headers: {
+          "x-vellum-client-id": "client-B",
+          "x-vellum-interface-id": "macos",
+        },
+        abortSignal: ac2.signal,
+      },
+      { hub },
+    );
+
+    expect(hub.listClients()).toHaveLength(2);
+
+    ac1.abort();
+    ac2.abort();
+  });
+
+  // ── disposeClient (force disconnect) ──────────────────────────────────────
+
+  test("disposeClient removes all subscribers for the clientId", () => {
+    const hub = new AssistantEventHub();
+
+    hub.subscribe({
+      type: "client" as const,
+      clientId: "force-dc-001",
+      interfaceId: "chrome-extension",
+      capabilities: ["host_browser"],
+      callback: () => {},
+    });
+
+    expect(hub.listClients()).toHaveLength(1);
+
+    const count = hub.disposeClient("force-dc-001");
+    expect(count).toBe(1);
+    expect(hub.listClients()).toHaveLength(0);
+  });
+
+  test("disposeClient returns 0 for unknown clientId", () => {
+    const hub = new AssistantEventHub();
+    expect(hub.disposeClient("nonexistent")).toBe(0);
+  });
+
+  test("disposeClient fires onEvict for each disposed entry", () => {
+    const hub = new AssistantEventHub();
+    let evictCount = 0;
+
+    hub.subscribe({
+      type: "client" as const,
+      clientId: "force-dc-evict",
+      interfaceId: "chrome-extension",
+      capabilities: ["host_browser"],
+      callback: () => {},
+      onEvict: () => {
+        evictCount++;
+      },
+    });
+
+    hub.disposeClient("force-dc-evict");
+    expect(evictCount).toBe(1);
   });
 });

package/src/__tests__/filing-service.test.ts

@@ -21,6 +21,9 @@ let mockConfig = {
     activeHoursStart: null as number | null,
     activeHoursEnd: null as number | null,
   },
+  memory: {
+    v2: { enabled: false },
+  },
 };
 
 mock.module("../config/loader.js", () => ({
@@ -31,6 +34,9 @@ mock.module("../config/loader.js", () => ({
   invalidateConfigCache: () => {},
 }));
 
+const { _setOverridesForTesting } =
+  await import("../config/assistant-feature-flags.js");
+
 // Mock conversation store
 const createdConversations: Array<{ title: string; conversationType: string }> =
   [];
@@ -83,16 +89,18 @@ mock.module("../memory/conversation-title-service.js", () => ({
 }));
 
 // Mock processMessage — FilingService now imports it directly.
-let _testProcessMessage: ((...args: unknown[]) => Promise<{ messageId: string }>) | undefined;
+let _testProcessMessage:
+  | ((...args: unknown[]) => Promise<{ messageId: string }>)
+  | undefined;
 
 mock.module("../daemon/process-message.js", () => ({
   processMessage: async (...args: unknown[]) => {
     if (_testProcessMessage) return _testProcessMessage(...args);
     return { messageId: `mock-msg-${Date.now()}` };
   },
+  processMessageInBackground: async () => ({ messageId: "mock-bg" }),
   resolveTurnChannel: () => "vellum",
   resolveTurnInterface: () => "vellum",
-  makePendingInteractionRegistrar: () => () => {},
   prepareConversationForMessage: async () => ({}),
 }));
 
@@ -120,6 +128,7 @@ describe("FilingService", () => {
     } catch {
       // best-effort
     }
+    _setOverridesForTesting({});
   });
 
   beforeEach(() => {
@@ -132,7 +141,9 @@ describe("FilingService", () => {
       processMessageCalls.push({
         conversationId: args[0] as string,
         content: args[1] as string,
-        options: (args[3] as { speed?: string; callSite?: string } | undefined) ?? undefined,
+        options:
+          (args[3] as { speed?: string; callSite?: string } | undefined) ??
+          undefined,
       });
       return { messageId: "msg-1" };
     });
@@ -147,6 +158,9 @@ describe("FilingService", () => {
         activeHoursStart: null,
         activeHoursEnd: null,
       },
+      memory: {
+        v2: { enabled: false },
+      },
     };
 
     // Seed buffer.md with content so runOnce doesn't skip
@@ -160,9 +174,7 @@ describe("FilingService", () => {
   });
 
   function createService(overrides?: {
-    processMessage?: (
-      ...args: unknown[]
-    ) => Promise<{ messageId: string }>;
+    processMessage?: (...args: unknown[]) => Promise<{ messageId: string }>;
   }) {
     if (overrides?.processMessage) {
       setTestProcessMessage(overrides.processMessage);
@@ -175,7 +187,9 @@ describe("FilingService", () => {
     await service.runOnce();
 
     expect(processMessageCalls).toHaveLength(1);
-    expect(processMessageCalls[0].options).toMatchObject({ callSite: "filingAgent" });
+    expect(processMessageCalls[0].options).toMatchObject({
+      callSite: "filingAgent",
+    });
     expect(processMessageCalls[0].options?.callSite).toBe("filingAgent");
   });
 
@@ -316,6 +330,43 @@ describe("FilingService", () => {
     });
   });
 
+  describe("memory v2 gate", () => {
+    test("start() does not schedule timers when v2 flag and config are both on", () => {
+      _setOverridesForTesting({ "memory-v2-enabled": true });
+      mockConfig.memory.v2.enabled = true;
+
+      const service = createService();
+      service.start();
+
+      expect(service.nextRunAt).toBeNull();
+      expect(service.nextCompactionAt).toBeNull();
+    });
+
+    test("start() schedules timers when only the flag is on", () => {
+      _setOverridesForTesting({ "memory-v2-enabled": true });
+      mockConfig.memory.v2.enabled = false;
+
+      const service = createService();
+      service.start();
+
+      expect(service.nextRunAt).not.toBeNull();
+      expect(service.nextCompactionAt).not.toBeNull();
+      service.stop();
+    });
+
+    test("start() schedules timers when only the config is on", () => {
+      _setOverridesForTesting({ "memory-v2-enabled": false });
+      mockConfig.memory.v2.enabled = true;
+
+      const service = createService();
+      service.start();
+
+      expect(service.nextRunAt).not.toBeNull();
+      expect(service.nextCompactionAt).not.toBeNull();
+      service.stop();
+    });
+  });
+
   describe("llm.callSites.filingAgent resolution", () => {
     // These tests verify that the call-site name used by FilingService
     // ('filingAgent') resolves through the unified `llm` config the way