@axiom-lattice/gateway 2.1.37 → 2.1.39

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/gateway",
-  "version": "2.1.37",
+  "version": "2.1.39",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
   "types": "dist/index.d.ts",
@@ -36,9 +36,9 @@
     "pino-roll": "^3.1.0",
     "redis": "^5.0.1",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "2.1.31",
-    "@axiom-lattice/protocols": "2.1.18",
-    "@axiom-lattice/queue-redis": "1.0.17"
+    "@axiom-lattice/core": "2.1.33",
+    "@axiom-lattice/protocols": "2.1.20",
+    "@axiom-lattice/queue-redis": "1.0.19"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/agent_service.test.ts

@@ -0,0 +1,238 @@
+/**
+ * Agent Service Tests
+ * 
+ * Tests for agent_service.ts to ensure proper async/await usage
+ * and error handling when working with agent clients.
+ */
+
+import { describe, it, expect, beforeEach, jest } from '@jest/globals';
+import { agent_state, agent_messages, draw_graph, agent_invoke, agent_stream } from '../services/agent_service';
+import { getAgentClient, agentLatticeManager } from '@axiom-lattice/core';
+
+// Mock the dependencies
+jest.mock('@axiom-lattice/core', () => ({
+  getAgentClient: jest.fn(),
+  agentLatticeManager: {
+    getAgentLatticeWithTenant: jest.fn(),
+  },
+  InMemoryChunkBuffer: jest.fn().mockImplementation(() => ({
+    addChunk: jest.fn(),
+    completeThread: jest.fn(),
+    abortThread: jest.fn(),
+  })),
+  registerChunkBuffer: jest.fn(),
+  getChunkBuffer: jest.fn(),
+  hasChunkBuffer: jest.fn().mockReturnValue(false),
+}));
+
+describe('Agent Service - getAgentClient async handling', () => {
+  const mockTenantId = 'test-tenant';
+  const mockAssistantId = 'test-assistant';
+  const mockThreadId = 'test-thread';
+  
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('agent_state', () => {
+    it('should properly await getAgentClient before calling getState', async () => {
+      const mockAgent = {
+        getState: jest.fn().mockResolvedValue({
+          values: { messages: [] },
+          next: [],
+        }),
+      };
+      
+      // Mock getAgentClient to return a promise that resolves to mockAgent
+      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
+
+      const result = await agent_state({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+      });
+
+      // Verify getAgentClient was called with correct arguments
+      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
+      
+      // Verify getState was called on the resolved agent (not on a promise)
+      expect(mockAgent.getState).toHaveBeenCalledWith({
+        configurable: { thread_id: mockThreadId, subgraphs: false },
+      });
+      
+      expect(result).toEqual({
+        values: { messages: [] },
+        next: [],
+      });
+    });
+
+    it('should throw error when agent is not found', async () => {
+      (getAgentClient as jest.Mock).mockResolvedValue(null);
+
+      await expect(agent_state({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+      })).rejects.toThrow(`Agent ${mockAssistantId} not found`);
+    });
+
+    it('should throw error when getAgentClient returns undefined', async () => {
+      (getAgentClient as jest.Mock).mockResolvedValue(undefined);
+
+      await expect(agent_state({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+      })).rejects.toThrow(`Agent ${mockAssistantId} not found`);
+    });
+  });
+
+  describe('agent_messages', () => {
+    it('should properly await getAgentClient before calling getState', async () => {
+      const mockMessages = [
+        { id: '1', getType: () => 'human', content: 'Hello', lc_kwargs: {} },
+        { id: '2', getType: () => 'ai', content: 'Hi there', lc_kwargs: {} },
+      ];
+      
+      const mockAgent = {
+        getState: jest.fn().mockResolvedValue({
+          values: { messages: mockMessages },
+        }),
+      };
+      
+      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
+
+      const result = await agent_messages({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+      });
+
+      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
+      expect(mockAgent.getState).toHaveBeenCalledWith({
+        configurable: { thread_id: mockThreadId, subgraphs: false },
+      });
+      
+      expect(result).toHaveLength(2);
+      expect(result[0].role).toBe('human');
+      expect(result[1].role).toBe('ai');
+    });
+  });
+
+  describe('draw_graph', () => {
+    it('should properly await getAgentClient before calling getGraphAsync', async () => {
+      const mockGraph = {
+        drawMermaid: jest.fn().mockResolvedValue('graph TD; A-->B;'),
+      };
+      
+      const mockAgent = {
+        getGraphAsync: jest.fn().mockResolvedValue(mockGraph),
+      };
+      
+      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
+
+      const result = await draw_graph(mockAssistantId, mockTenantId);
+
+      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
+      expect(mockAgent.getGraphAsync).toHaveBeenCalled();
+      expect(mockGraph.drawMermaid).toHaveBeenCalled();
+      expect(result).toBe('graph TD; A-->B;');
+    });
+
+    it('should throw error when agent is not found', async () => {
+      (getAgentClient as jest.Mock).mockResolvedValue(null);
+
+      await expect(draw_graph(mockAssistantId, mockTenantId))
+        .rejects.toThrow(`Agent ${mockAssistantId} not found`);
+    });
+  });
+
+  describe('agent_stream', () => {
+    it('should properly await getAgentClient before calling stream', async () => {
+      const mockStream = async function* () {
+        yield ['updates', { messages: [{ toDict: () => ({ type: 'ai', content: 'Hello' }) }] }];
+      }();
+      
+      const mockAgent = {
+        stream: jest.fn().mockResolvedValue(mockStream),
+      };
+      
+      (getAgentClient as jest.Mock).mockResolvedValue(mockAgent);
+      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
+        config: { runConfig: {} },
+      });
+
+      const result = await agent_stream({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+        input: { message: 'Test' },
+      });
+
+      expect(getAgentClient).toHaveBeenCalledWith(mockTenantId, mockAssistantId);
+      expect(mockAgent.stream).toHaveBeenCalled();
+      expect(typeof result[Symbol.asyncIterator]).toBe('function');
+    });
+
+    it('should throw error when agent is not found', async () => {
+      (getAgentClient as jest.Mock).mockResolvedValue(null);
+      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
+        config: { runConfig: {} },
+      });
+
+      // The function will throw but also try to call chunkBuffer.abortThread
+      // which might fail in test environment, so we just check it throws
+      await expect(agent_stream({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+        input: { message: 'Test' },
+      })).rejects.toThrow();
+    });
+  });
+
+  describe('agent_invoke', () => {
+    it('should use agentLattice.client directly without calling getAgentClient', async () => {
+      const mockResult = {
+        messages: [
+          { toDict: () => ({ type: 'ai', content: 'Response' }) },
+        ],
+      };
+      
+      const mockAgent = {
+        invoke: jest.fn().mockResolvedValue(mockResult),
+      };
+      
+      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
+        client: mockAgent,
+        config: { runConfig: {} },
+      });
+
+      const result = await agent_invoke({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+        input: { message: 'Hello' },
+      });
+
+      // agent_invoke uses agentLattice.client directly, not getAgentClient
+      expect(getAgentClient).not.toHaveBeenCalled();
+      expect(mockAgent.invoke).toHaveBeenCalled();
+      expect(result.messages).toHaveLength(1);
+    });
+
+    it('should throw error when agent lattice client is not found', async () => {
+      (agentLatticeManager.getAgentLatticeWithTenant as jest.Mock).mockReturnValue({
+        client: null,
+        config: { runConfig: {} },
+      });
+
+      await expect(agent_invoke({
+        tenant_id: mockTenantId,
+        assistant_id: mockAssistantId,
+        thread_id: mockThreadId,
+        input: { message: 'Hello' },
+      })).rejects.toThrow(`Agent ${mockAssistantId} not found`);
+    });
+  });
+});

package/src/__tests__/workspace.test.ts

@@ -0,0 +1,253 @@
+/**
+ * WorkspaceController Tenant Isolation Tests
+ *
+ * Tests to verify that file operations in WorkspaceController
+ * properly use tenant-specific paths.
+ */
+
+import { describe, it, expect, beforeEach, jest } from "@jest/globals";
+import type { FastifyRequest, FastifyReply } from "fastify";
+
+// Mock store lattice
+const mockGetWorkspaceById = jest.fn();
+const mockGetAllWorkspaces = jest.fn();
+const mockCreateWorkspace = jest.fn();
+const mockUpdateWorkspace = jest.fn();
+const mockDeleteWorkspace = jest.fn();
+const mockGetProjectsByWorkspace = jest.fn();
+const mockCreateProject = jest.fn();
+const mockGetProjectById = jest.fn();
+const mockUpdateProject = jest.fn();
+const mockDeleteProject = jest.fn();
+
+const mockStoreLattice = {
+  store: {
+    getAllWorkspaces: mockGetAllWorkspaces,
+    createWorkspace: mockCreateWorkspace,
+    getWorkspaceById: mockGetWorkspaceById,
+    updateWorkspace: mockUpdateWorkspace,
+    deleteWorkspace: mockDeleteWorkspace,
+    getProjectsByWorkspace: mockGetProjectsByWorkspace,
+    createProject: mockCreateProject,
+    getProjectById: mockGetProjectById,
+    updateProject: mockUpdateProject,
+    deleteProject: mockDeleteProject,
+  },
+};
+
+// Mock the dependencies
+jest.mock("@axiom-lattice/core", () => ({
+  getStoreLattice: jest.fn().mockReturnValue(mockStoreLattice),
+  SandboxFilesystem: jest.fn().mockImplementation(() => ({
+    lsInfo: jest.fn(),
+    read: jest.fn(),
+  })),
+  FilesystemBackend: jest.fn().mockImplementation(() => ({
+    lsInfo: jest.fn(),
+    read: jest.fn(),
+  })),
+  getSandBoxManager: jest.fn().mockReturnValue({
+    createSandbox: jest.fn().mockResolvedValue({
+      file: {
+        downloadFile: jest.fn(),
+        uploadFile: jest.fn(),
+      },
+    }),
+  }),
+}));
+
+// Import after mocks are set up
+import { WorkspaceController } from "../controllers/workspace";
+
+describe("WorkspaceController - Tenant Isolation", () => {
+  let controller: WorkspaceController;
+  const mockTenantId = "test-tenant-123";
+  const mockWorkspaceId = "workspace-456";
+  const mockProjectId = "project-789";
+
+  // Helper to create mock request with tenant
+  const createMockRequest = (overrides: any = {}) => ({
+    params: { workspaceId: mockWorkspaceId, projectId: mockProjectId },
+    query: {},
+    headers: { "x-tenant-id": mockTenantId },
+    ...overrides,
+  }) as unknown as FastifyRequest;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    controller = new WorkspaceController();
+  });
+
+  describe("getBackend path construction", () => {
+    it("should use tenant-specific path for filesystem storage", async () => {
+      const { FilesystemBackend } = await import("@axiom-lattice/core");
+      
+      const mockWorkspace = {
+        id: mockWorkspaceId,
+        tenantId: mockTenantId,
+        storageType: "filesystem",
+      };
+
+      mockGetWorkspaceById.mockResolvedValue(mockWorkspace);
+
+      // Access private method for testing
+      const getBackend = (controller as any).getBackend.bind(controller);
+      await getBackend(mockTenantId, mockWorkspaceId, mockProjectId);
+
+      // Verify FilesystemBackend was called with tenant-specific path
+      expect(FilesystemBackend).toHaveBeenCalledWith({
+        rootDir: `/lattice_store/tenants/${mockTenantId}/workspaces/${mockWorkspaceId}/${mockProjectId}`,
+        virtualMode: true,
+      });
+    });
+
+    it("should use tenant-specific path for sandbox storage", async () => {
+      const { getSandBoxManager } = await import("@axiom-lattice/core");
+
+      const mockWorkspace = {
+        id: mockWorkspaceId,
+        tenantId: mockTenantId,
+        storageType: "sandbox",
+      };
+
+      mockGetWorkspaceById.mockResolvedValue(mockWorkspace);
+
+      const getBackend = (controller as any).getBackend.bind(controller);
+      await getBackend(mockTenantId, mockWorkspaceId, mockProjectId);
+
+      // Verify getSandBoxManager was called (uses default key, sandbox manager is not tenant-specific)
+      expect(getSandBoxManager).toHaveBeenCalled();
+    });
+
+    it("should pass tenantId to workspace store lookup", async () => {
+      const mockWorkspace = {
+        id: mockWorkspaceId,
+        tenantId: mockTenantId,
+        storageType: "filesystem",
+      };
+
+      mockGetWorkspaceById.mockResolvedValue(mockWorkspace);
+
+      const getBackend = (controller as any).getBackend.bind(controller);
+      await getBackend(mockTenantId, mockWorkspaceId, mockProjectId);
+
+      // Verify workspace lookup used correct tenantId
+      expect(mockGetWorkspaceById).toHaveBeenCalledWith(mockTenantId, mockWorkspaceId);
+    });
+  });
+
+  describe("listPath tenant isolation", () => {
+    it("should extract tenantId from request headers", async () => {
+      const { FilesystemBackend } = await import("@axiom-lattice/core");
+      
+      const mockWorkspace = {
+        id: mockWorkspaceId,
+        tenantId: mockTenantId,
+        storageType: "filesystem",
+      };
+
+      mockGetWorkspaceById.mockResolvedValue(mockWorkspace);
+
+      const request = createMockRequest();
+      await controller.listPath(request);
+
+      // Verify tenant-specific path was used
+      expect(FilesystemBackend).toHaveBeenCalledWith({
+        rootDir: `/lattice_store/tenants/${mockTenantId}/workspaces/${mockWorkspaceId}/${mockProjectId}`,
+        virtualMode: true,
+      });
+    });
+
+    it("should use different paths for different tenants", async () => {
+      const { FilesystemBackend } = await import("@axiom-lattice/core");
+      
+      const tenantA = "tenant-a";
+      const tenantB = "tenant-b";
+
+      const mockWorkspaceA = {
+        id: mockWorkspaceId,
+        tenantId: tenantA,
+        storageType: "filesystem",
+      };
+
+      const mockWorkspaceB = {
+        id: mockWorkspaceId,
+        tenantId: tenantB,
+        storageType: "filesystem",
+      };
+
+      // First call with tenant A
+      mockGetWorkspaceById.mockResolvedValue(mockWorkspaceA);
+      await controller.listPath(createMockRequest({ headers: { "x-tenant-id": tenantA } }));
+
+      expect(FilesystemBackend).toHaveBeenCalledWith({
+        rootDir: `/lattice_store/tenants/${tenantA}/workspaces/${mockWorkspaceId}/${mockProjectId}`,
+        virtualMode: true,
+      });
+
+      // Second call with tenant B
+      mockGetWorkspaceById.mockResolvedValue(mockWorkspaceB);
+      await controller.listPath(createMockRequest({ headers: { "x-tenant-id": tenantB } }));
+
+      expect(FilesystemBackend).toHaveBeenCalledWith({
+        rootDir: `/lattice_store/tenants/${tenantB}/workspaces/${mockWorkspaceId}/${mockProjectId}`,
+        virtualMode: true,
+      });
+    });
+  });
+
+  describe("readFile tenant isolation", () => {
+    it("should use tenant-specific backend for reading files", async () => {
+      const { FilesystemBackend } = await import("@axiom-lattice/core");
+      
+      const mockWorkspace = {
+        id: mockWorkspaceId,
+        tenantId: mockTenantId,
+        storageType: "filesystem",
+      };
+
+      mockGetWorkspaceById.mockResolvedValue(mockWorkspace);
+
+      const request = createMockRequest({
+        query: { path: "/test.txt", offset: 0, limit: 100 },
+      });
+      await controller.readFile(request);
+
+      // Verify tenant-specific path was used
+      expect(FilesystemBackend).toHaveBeenCalledWith({
+        rootDir: `/lattice_store/tenants/${mockTenantId}/workspaces/${mockWorkspaceId}/${mockProjectId}`,
+        virtualMode: true,
+      });
+    });
+  });
+
+  describe("getTenantId helper", () => {
+    it("should extract tenantId from user context when available", () => {
+      const request = {
+        user: { tenantId: "user-context-tenant" },
+        headers: {},
+      } as unknown as FastifyRequest;
+
+      const tenantId = (controller as any).getTenantId(request);
+      expect(tenantId).toBe("user-context-tenant");
+    });
+
+    it("should fallback to x-tenant-id header when no user context", () => {
+      const request = {
+        headers: { "x-tenant-id": "header-tenant" },
+      } as unknown as FastifyRequest;
+
+      const tenantId = (controller as any).getTenantId(request);
+      expect(tenantId).toBe("header-tenant");
+    });
+
+    it("should use default tenant when no context or header", () => {
+      const request = {
+        headers: {},
+      } as unknown as FastifyRequest;
+
+      const tenantId = (controller as any).getTenantId(request);
+      expect(tenantId).toBe("default");
+    });
+  });
+});