kaiban-distributed - change:Fixed DistributedStateMiddleware Bug to publish to global Redis pub/sub channel

Author: andreibesleaga

package-lock.json

@@ -1,7 +1,5 @@
-import {
-  IMessagingDriver,
-  MessagePayload,
-} from "../../infrastructure/messaging/interfaces";
+import { Redis } from 'ioredis';
+import type { MessagePayload } from "../../infrastructure/messaging/interfaces";
 
 const PII_DENYLIST: ReadonlySet<string> = new Set([
   'email', 'name', 'phone', 'ip', 'password', 'token', 'secret', 'ssn', 'dob',
@@ -23,11 +21,11 @@ interface ZustandStore {
 }
 
 export class DistributedStateMiddleware {
-  private driver: IMessagingDriver;
+  private redis: Redis;
   private channelName: string;
 
-  constructor(driver: IMessagingDriver, channelName = 'kaiban-state-events') {
-    this.driver = driver;
+  constructor(redisUrl: string, channelName = 'kaiban-state-events') {
+    this.redis = new Redis(redisUrl, { lazyConnect: false });
     this.channelName = channelName;
   }
 
@@ -46,7 +44,7 @@ export class DistributedStateMiddleware {
       };
 
       try {
-        await this.driver.publish(this.channelName, payload);
+        await this.redis.publish(this.channelName, JSON.stringify(payload));
       } catch (err) {
         console.error("[DistributedStateMiddleware] Failed to publish state delta:", err);
       }
@@ -56,11 +54,21 @@ export class DistributedStateMiddleware {
   }
 
   public async listen(onStateChange: (delta: Record<string, unknown>) => void): Promise<void> {
-    await this.driver.subscribe(
-      this.channelName,
-      async (payload: MessagePayload) => {
-        onStateChange(payload.data['stateUpdate'] as Record<string, unknown>);
-      },
-    );
+    const sub = new Redis(this.redis.options);
+    await sub.subscribe(this.channelName);
+    sub.on('message', (channel, message) => {
+      if (channel === this.channelName) {
+        try {
+          const payload = JSON.parse(message) as MessagePayload;
+          onStateChange(payload.data['stateUpdate'] as Record<string, unknown>);
+        } catch (e) {
+          console.error("[DistributedStateMiddleware] Failed to parse message:", e);
+        }
+      }
+    });
+  }
+
+  public async disconnect(): Promise<void> {
+    await this.redis.quit();
   }
 }

src/adapters/state/distributedMiddleware.ts

@@ -1,7 +1,6 @@
 import { Team } from 'kaibanjs';
 import type { Agent, Task } from 'kaibanjs';
 import { DistributedStateMiddleware } from '../../adapters/state/distributedMiddleware';
-import type { IMessagingDriver } from '../messaging/interfaces';
 
 export interface KaibanTeamConfig {
   name: string;
@@ -17,7 +16,7 @@ export interface KaibanTeamConfig {
  * SocketGateway → kaiban-board in real-time.
  *
  * Usage:
- *   const bridge = new KaibanTeamBridge({ name, agents, tasks }, driver);
+ *   const bridge = new KaibanTeamBridge({ name, agents, tasks }, 'redis://localhost:6379');
  *   const result = await bridge.start({ topic: 'AI news' });
  */
 export class KaibanTeamBridge {
@@ -26,11 +25,11 @@ export class KaibanTeamBridge {
 
   constructor(
     config: KaibanTeamConfig,
-    driver: IMessagingDriver,
+    redisUrl: string,
     stateChannel = 'kaiban-state-events',
   ) {
     this.team = new Team({ name: config.name, agents: config.agents, tasks: config.tasks, env: config.env ?? {} });
-    this.middleware = new DistributedStateMiddleware(driver, stateChannel);
+    this.middleware = new DistributedStateMiddleware(redisUrl, stateChannel);
 
     const store = this.team.getStore() as unknown as { setState: (p: Record<string, unknown>) => void };
     this.middleware.attach(store);
@@ -50,4 +49,8 @@ export class KaibanTeamBridge {
   ): () => void {
     return this.team.subscribeToChanges(listener, properties);
   }
+
+  async disconnect(): Promise<void> {
+    await this.middleware.disconnect();
+  }
 }

src/infrastructure/kaibanjs/kaiban-team-bridge.ts

@@ -1,6 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { KaibanTeamBridge } from '../../../src/infrastructure/kaibanjs/kaiban-team-bridge';
-import type { IMessagingDriver } from '../../../src/infrastructure/messaging/interfaces';
 
 const mockGetStore = vi.fn().mockReturnValue({
   setState: vi.fn(),
@@ -18,49 +17,51 @@ vi.mock('kaibanjs', () => ({
   Task: vi.fn().mockImplementation(function (params: Record<string, unknown>) { return params; }),
 }));
 
-function makeMockDriver(): IMessagingDriver {
-  return {
-    publish: vi.fn().mockResolvedValue(undefined),
-    subscribe: vi.fn().mockResolvedValue(undefined),
-    unsubscribe: vi.fn().mockResolvedValue(undefined),
-    disconnect: vi.fn().mockResolvedValue(undefined),
-  };
-}
+const mockRedis = {
+  publish: vi.fn().mockResolvedValue(undefined),
+  subscribe: vi.fn().mockResolvedValue(undefined),
+  quit: vi.fn().mockResolvedValue(undefined),
+  on: vi.fn(),
+};
+
+vi.mock('ioredis', () => ({
+  Redis: vi.fn().mockImplementation(function() { return mockRedis; }),
+}));
 
 describe('KaibanTeamBridge', () => {
   beforeEach(() => { vi.clearAllMocks(); });
 
   it('calls team.getStore() to attach DistributedStateMiddleware', () => {
     const bridge = new KaibanTeamBridge(
       { name: 'Blog Team', agents: [], tasks: [] },
-      makeMockDriver(),
+      'redis://localhost:6379',
     );
     expect(bridge).toBeDefined();
     expect(mockGetStore).toHaveBeenCalledOnce();
   });
 
   it('getTeam() returns the underlying KaibanJS Team', () => {
-    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, makeMockDriver());
+    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, 'redis://localhost:6379');
     const team = bridge.getTeam();
     expect(team).toBeDefined();
     expect(typeof team.start).toBe('function');
   });
 
   it('start() delegates to team.start() with inputs', async () => {
-    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, makeMockDriver());
+    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, 'redis://localhost:6379');
     const result = await bridge.start({ topic: 'AI trends' });
     expect(mockStart).toHaveBeenCalledWith({ topic: 'AI trends' });
     expect(result.status).toBe('FINISHED');
   });
 
   it('start() with no inputs passes empty object', async () => {
-    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, makeMockDriver());
+    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, 'redis://localhost:6379');
     await bridge.start();
     expect(mockStart).toHaveBeenCalledWith({});
   });
 
   it('subscribeToChanges() sets up a store listener', () => {
-    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, makeMockDriver());
+    const bridge = new KaibanTeamBridge({ name: 'T', agents: [], tasks: [] }, 'redis://localhost:6379');
     const listener = vi.fn();
     bridge.subscribeToChanges(listener, ['teamWorkflowStatus']);
     expect(mockGetStore).toHaveBeenCalled();

tests/unit/kaibanjs/kaiban-team-bridge.test.ts

@@ -1,15 +1,18 @@
 import { describe, it, expect, vi } from "vitest";
 import { DistributedStateMiddleware } from "../../../src/adapters/state/distributedMiddleware";
-import { IMessagingDriver, MessagePayload } from "../../../src/infrastructure/messaging/interfaces";
+import { MessagePayload } from "../../../src/infrastructure/messaging/interfaces";
 
-function makeMockDriver(): IMessagingDriver {
-  return {
-    publish: vi.fn().mockResolvedValue(undefined),
-    subscribe: vi.fn().mockResolvedValue(undefined),
-    unsubscribe: vi.fn().mockResolvedValue(undefined),
-    disconnect: vi.fn().mockResolvedValue(undefined),
-  };
-}
+const mockRedis = {
+  publish: vi.fn().mockResolvedValue(undefined),
+  subscribe: vi.fn(),
+  quit: vi.fn().mockResolvedValue(undefined),
+  on: vi.fn(),
+  options: {},
+};
+
+vi.mock('ioredis', () => ({
+  Redis: vi.fn().mockImplementation(function() { return mockRedis; }),
+}));
 
 interface MockStore {
   state: Record<string, unknown>;
@@ -24,65 +27,68 @@ function makeStore(initial: Record<string, unknown> = {}): MockStore {
 }
 
 describe("DistributedStateMiddleware", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
   it("intercepts setState and publishes to the driver", async () => {
-    const mockDriver = makeMockDriver();
-    const mw = new DistributedStateMiddleware(mockDriver);
+    const mw = new DistributedStateMiddleware('redis://localhost');
     const store = makeStore({ count: 0 });
     mw.attach(store);
     await store.setState({ count: 1 });
     expect(store.state['count']).toBe(1);
-    expect(mockDriver.publish).toHaveBeenCalledWith("kaiban-state-events", expect.objectContaining({ data: { stateUpdate: { count: 1 } } }));
+    expect(mockRedis.publish).toHaveBeenCalledWith(
+      "kaiban-state-events",
+      expect.stringContaining('"stateUpdate":{"count":1}')
+    );
   });
 
   it("sanitizeDelta strips PII keys", async () => {
-    const mockDriver = makeMockDriver();
-    const mw = new DistributedStateMiddleware(mockDriver);
+    const mw = new DistributedStateMiddleware('redis://localhost');
     const store = makeStore();
     mw.attach(store);
     await store.setState({ count: 2, email: "x@y.com", token: "abc", password: "pw" });
-    const call = (mockDriver.publish as ReturnType<typeof vi.fn>).mock.calls[0];
-    const data = call[1].data.stateUpdate as Record<string, unknown>;
+    const call = mockRedis.publish.mock.calls[0];
+    const parsed = JSON.parse(call[1] as string) as MessagePayload;
+    const data = parsed.data['stateUpdate'] as Record<string, unknown>;
     expect(data['count']).toBe(2);
     expect(data['email']).toBeUndefined();
     expect(data['token']).toBeUndefined();
     expect(data['password']).toBeUndefined();
   });
 
   it("sanitizeDelta handles null partial (covers null branch)", async () => {
-    const mockDriver = makeMockDriver();
-    const mw = new DistributedStateMiddleware(mockDriver);
+    const mw = new DistributedStateMiddleware('redis://localhost');
     const store = makeStore();
     mw.attach(store);
     // Trigger with null via coercion to cover `if (partial === null) return {}`
     await (store.setState as (p: unknown) => Promise<void>)(null);
-    const call = (mockDriver.publish as ReturnType<typeof vi.fn>).mock.calls[0];
-    expect(call[1].data.stateUpdate).toEqual({});
+    const call = mockRedis.publish.mock.calls[0];
+    const parsed = JSON.parse(call[1] as string) as MessagePayload;
+    expect(parsed.data['stateUpdate']).toEqual({});
   });
 
   it("listen() subscribes and delivers state deltas via callback", async () => {
-    let capturedHandler!: (payload: MessagePayload) => Promise<void>;
-    const mockDriver: IMessagingDriver = {
-      publish: vi.fn().mockResolvedValue(undefined),
-      subscribe: vi.fn((_q, handler) => { capturedHandler = handler; return Promise.resolve(); }),
-      unsubscribe: vi.fn().mockResolvedValue(undefined),
-      disconnect: vi.fn().mockResolvedValue(undefined),
-    };
-    const mw = new DistributedStateMiddleware(mockDriver);
+    let capturedHandler!: (channel: string, message: string) => void;
+    mockRedis.on.mockImplementation((event, handler) => {
+      if (event === 'message') capturedHandler = handler;
+    });
+    
+    const mw = new DistributedStateMiddleware('redis://localhost');
     const onStateChange = vi.fn();
     await mw.listen(onStateChange);
-    await capturedHandler({ taskId: "g", agentId: "system", timestamp: 0, data: { stateUpdate: { x: 1 } } });
+    
+    // Simulate incoming Redis pub/sub message
+    const msg = JSON.stringify({ taskId: "g", agentId: "system", timestamp: 0, data: { stateUpdate: { x: 1 } } });
+    capturedHandler("kaiban-state-events", msg);
+    
     expect(onStateChange).toHaveBeenCalledWith({ x: 1 });
   });
 
   it("publish error is caught and logged without throwing", async () => {
-    const mockDriver: IMessagingDriver = {
-      publish: vi.fn().mockRejectedValue(new Error("redis down")),
-      subscribe: vi.fn().mockResolvedValue(undefined),
-      unsubscribe: vi.fn().mockResolvedValue(undefined),
-      disconnect: vi.fn().mockResolvedValue(undefined),
-    };
+    mockRedis.publish.mockRejectedValueOnce(new Error("redis down"));
     const errSpy = vi.spyOn(console, "error").mockImplementation(() => {});
-    const mw = new DistributedStateMiddleware(mockDriver);
+    const mw = new DistributedStateMiddleware('redis://localhost');
     const store = makeStore();
     mw.attach(store);
     await expect(store.setState({ x: 1 })).resolves.not.toThrow();
@@ -91,11 +97,10 @@ describe("DistributedStateMiddleware", () => {
   });
 
   it("sanitizeDelta handles empty state update", async () => {
-    const mockDriver = makeMockDriver();
-    const mw = new DistributedStateMiddleware(mockDriver);
+    const mw = new DistributedStateMiddleware('redis://localhost');
     const store = makeStore();
     mw.attach(store);
     await store.setState({});
-    expect(mockDriver.publish).toHaveBeenCalledOnce();
+    expect(mockRedis.publish).toHaveBeenCalledOnce();
   });
 });