moving

packages/worker/tests/handlers/chat-handler.test.ts

@@ -1,4 +1,7 @@
-import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
 import type { Job } from 'bullmq';
 import type { Db } from '../../../server/src/db/connection.js';
 import type { JobData } from '../../src/job-processor.js';
@@ -8,23 +11,34 @@ function makeJob(data: JobData): Job<JobData> {
   return { data } as unknown as Job<JobData>;
 }
 
-vi.mock('@waggle/agent', () => ({
-  runAgentLoop: vi.fn(async () => ({
-    content: 'Agent response here',
-    toolsUsed: ['search_memory'],
-    usage: { inputTokens: 100, outputTokens: 50 },
-  })),
-  createSystemTools: vi.fn(() => []),
-}));
+vi.mock('@waggle/agent', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@waggle/agent')>();
+  return {
+    ...actual,
+    runAgentLoop: vi.fn(async () => ({
+      content: 'Agent response here',
+      toolsUsed: ['read_file'],
+      usage: { inputTokens: 100, outputTokens: 50 },
+    })),
+  };
+});
 
 import { chatHandler } from '../../src/handlers/chat-handler.js';
-import { runAgentLoop, createSystemTools } from '@waggle/agent';
+import { runAgentLoop } from '@waggle/agent';
 
 describe('Chat Handler', () => {
   const mockDb = {} as Db;
+  let dataDir: string;
 
   beforeEach(() => {
     vi.clearAllMocks();
+    dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-worker-chat-'));
+    vi.stubEnv('WAGGLE_DATA_DIR', dataDir);
+  });
+
+  afterEach(() => {
+    vi.unstubAllEnvs();
+    fs.rmSync(dataDir, { recursive: true, force: true });
   });
 
   it('calls runAgentLoop and returns response with correct shape', async () => {
@@ -41,7 +55,7 @@ describe('Chat Handler', () => {
     expect(result.response).toBe('Agent response here');
     expect(result.userId).toBe('u1');
     expect(result.model).toBe('claude-sonnet');
-    expect(result.toolsUsed).toEqual(['search_memory']);
+    expect(result.toolsUsed).toEqual(['read_file']);
     expect(result.tokensUsed).toBe(150);
   });
 
@@ -61,21 +75,56 @@ describe('Chat Handler', () => {
     expect(callArgs.model).toBe('gpt-4o');
     expect(callArgs.messages).toEqual([{ role: 'user', content: 'test message' }]);
     expect(callArgs.systemPrompt).toContain('Waggle AI agent');
-    expect(callArgs.tools).toEqual([]);
+    expect(callArgs.systemPrompt).toContain('read-only');
+    expect(callArgs.systemPrompt).toContain('cannot run shell commands, execute code, or write, edit, or delete files');
+    expect(callArgs.tools.map(tool => tool.name)).toEqual([
+      'read_file',
+      'search_files',
+      'search_content',
+      'web_search',
+      'web_fetch',
+    ]);
   });
 
-  it('creates system tools with workspaceDir from input', async () => {
+  it('ignores caller workspaceDir and roots tools in the team data directory', async () => {
+    const callerWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-attacker-root-'));
+    const tenantRoot = path.join(dataDir, 'teams', 'team-1', 'files');
+    fs.mkdirSync(tenantRoot, { recursive: true });
+    fs.writeFileSync(path.join(tenantRoot, 'tenant.txt'), 'tenant-root-content');
+    fs.writeFileSync(path.join(callerWorkspace, 'tenant.txt'), 'caller-root-content');
+
     const mockJob = makeJob({
       jobId: 'j3',
-      teamId: 't1',
+      teamId: 'team-1',
       userId: 'u1',
       jobType: 'chat',
-      input: { message: 'test', workspaceDir: '/custom/workspace' },
+      input: { message: 'test', workspaceDir: callerWorkspace },
     });
 
-    await chatHandler(mockJob, mockDb);
+    try {
+      await chatHandler(mockJob, mockDb);
 
-    expect(createSystemTools).toHaveBeenCalledWith('/custom/workspace');
+      const callArgs = vi.mocked(runAgentLoop).mock.calls[0][0];
+      const readFile = callArgs.tools.find(tool => tool.name === 'read_file');
+      expect(readFile).toBeDefined();
+      await expect(readFile!.execute({ path: 'tenant.txt' })).resolves.toContain('tenant-root-content');
+    } finally {
+      fs.rmSync(callerWorkspace, { recursive: true, force: true });
+    }
+  });
+
+  it('fails closed when WAGGLE_DATA_DIR is unavailable', async () => {
+    vi.stubEnv('WAGGLE_DATA_DIR', '');
+    const mockJob = makeJob({
+      jobId: 'j-no-root',
+      teamId: 'team-1',
+      userId: 'u1',
+      jobType: 'chat',
+      input: { message: 'test' },
+    });
+
+    await expect(chatHandler(mockJob, mockDb)).rejects.toThrow('WAGGLE_DATA_DIR');
+    expect(runAgentLoop).not.toHaveBeenCalled();
   });
 
   it('uses default model when not specified in input', async () => {