basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Qwen4 * SPDX-License-Identifier: Apache-2.05 */6 7import { describe, it, expect, vi, beforeEach } from 'vitest';8import { SendMessageTool } from './send-message.js';9import { BackgroundTaskRegistry } from '../agents/background-tasks.js';10import { ToolErrorType } from './tool-error.js';11import type { ApprovalMode, Config } from '../config/config.js';12import { runWithTeammateIdentity } from '../agents/team/identity.js';13 14const DEFAULT_MODE = 'default' as ApprovalMode;15const PLAN_MODE = 'plan' as ApprovalMode;16 17function makeTeamConfig(opts?: {18  teamManager?: {19    sendMessage: (...args: unknown[]) => Promise<void>;20    broadcast: (...args: unknown[]) => Promise<void>;21    requestShutdown?: (...args: unknown[]) => Promise<void>;22  } | null;23  approvalMode?: ApprovalMode;24}) {25  return {26    getTeamManager: () => opts?.teamManager ?? null,27    getBackgroundTaskRegistry: () => new BackgroundTaskRegistry(),28    getApprovalMode: () => opts?.approvalMode ?? DEFAULT_MODE,29  } as unknown as Config;30}31 32describe('SendMessageTool — team mode', () => {33  it('has the correct name', () => {34    const tool = new SendMessageTool(makeTeamConfig());35    expect(tool.name).toBe('send_message');36  });37 38  it('sends a message via TeamManager', async () => {39    const sendMessage = vi.fn().mockResolvedValue(undefined);40    const tool = new SendMessageTool(41      makeTeamConfig({42        teamManager: {43          sendMessage,44          broadcast: vi.fn(),45        },46      }),47    );48 49    const invocation = tool.build({50      to: 'alice',51      message: 'hello',52    });53    const result = await invocation.execute(new AbortController().signal);54    expect(result.error).toBeUndefined();55    expect(result.llmContent).toContain('alice');56    expect(sendMessage).toHaveBeenCalledWith(57      'alice',58      'hello',59      'leader',60      undefined,61    );62  });63 64  it('broadcasts with "*"', async () => {65    const broadcast = vi.fn().mockResolvedValue(undefined);66    const tool = new SendMessageTool(67      makeTeamConfig({68        teamManager: {69          sendMessage: vi.fn(),70          broadcast,71        },72      }),73    );74 75    const invocation = tool.build({76      to: '*',77      message: 'hey all',78    });79    const result = await invocation.execute(new AbortController().signal);80    expect(result.error).toBeUndefined();81    expect(result.llmContent).toContain('broadcast');82    expect(broadcast).toHaveBeenCalledWith('hey all', 'leader');83  });84 85  it('returns error when no team is active and no task_id given', async () => {86    const tool = new SendMessageTool(makeTeamConfig());87    const invocation = tool.build({88      to: 'alice',89      message: 'hello',90    });91    const result = await invocation.execute(new AbortController().signal);92    expect(result.error).toBeDefined();93    expect(result.llmContent).toContain('No active team');94  });95 96  it('routes shutdown_request via requestShutdown', async () => {97    const requestShutdown = vi.fn().mockResolvedValue(undefined);98    const tool = new SendMessageTool(99      makeTeamConfig({100        teamManager: {101          sendMessage: vi.fn(),102          broadcast: vi.fn(),103          requestShutdown,104        },105      }),106    );107 108    const invocation = tool.build({109      to: 'bob',110      message: 'Please shut down.',111      type: 'shutdown_request',112    });113    const result = await invocation.execute(new AbortController().signal);114    expect(result.error).toBeUndefined();115    expect(result.llmContent).toContain('Shutdown');116    expect(result.llmContent).toContain('bob');117    expect(requestShutdown).toHaveBeenCalledWith('bob');118  });119 120  it('rejects shutdown_request from a teammate (leader-only)', async () => {121    // A teammate calling shutdown_request would impersonate the122    // leader, since requestShutdown writes the mailbox entry with123    // `from: LEADER_NAME` and arms shutdown_approved tracking.124    const requestShutdown = vi.fn().mockResolvedValue(undefined);125    const tool = new SendMessageTool(126      makeTeamConfig({127        teamManager: {128          sendMessage: vi.fn(),129          broadcast: vi.fn(),130          requestShutdown,131        },132      }),133    );134 135    const invocation = tool.build({136      to: 'bob',137      message: 'Please shut down.',138      type: 'shutdown_request',139    });140    const result = await runWithTeammateIdentity(141      {142        agentName: 'attacker',143        teamName: 'team',144        agentId: 'attacker@team',145        isTeamLead: false,146      },147      () => invocation.execute(new AbortController().signal),148    );149    expect(result.error).toBeDefined();150    expect(result.llmContent).toContain('Only the team leader');151    expect(requestShutdown).not.toHaveBeenCalled();152  });153 154  it('blocks plan-required teammates before leader approval', async () => {155    const sendMessage = vi.fn().mockResolvedValue(undefined);156    const tool = new SendMessageTool(157      makeTeamConfig({158        approvalMode: PLAN_MODE,159        teamManager: {160          sendMessage,161          broadcast: vi.fn(),162        },163      }),164    );165 166    const invocation = tool.build({167      to: 'alice',168      message: 'execute this before approval',169    });170    const result = await runWithTeammateIdentity(171      {172        agentName: 'planner',173        teamName: 'team',174        agentId: 'planner@team',175        isTeamLead: false,176        planModeRequired: true,177      },178      () => invocation.execute(new AbortController().signal),179    );180 181    expect(result.error).toBeDefined();182    expect(result.llmContent).toContain('waiting for leader approval');183    expect(sendMessage).not.toHaveBeenCalled();184  });185 186  it('validates required params', () => {187    const tool = new SendMessageTool(makeTeamConfig());188    // `message` is required.189    expect(() => tool.build({} as never)).toThrow();190    expect(() => tool.build({ to: 'alice' } as never)).toThrow();191  });192});193 194describe('SendMessageTool — background-task mode', () => {195  let registry: BackgroundTaskRegistry;196  let config: Config;197  let tool: SendMessageTool;198  let resumeBackgroundAgent: ReturnType<typeof vi.fn>;199  let reviveCompletedBackgroundAgent: ReturnType<typeof vi.fn>;200 201  beforeEach(() => {202    registry = new BackgroundTaskRegistry();203    resumeBackgroundAgent = vi.fn();204    reviveCompletedBackgroundAgent = vi.fn();205    config = {206      getBackgroundTaskRegistry: () => registry,207      getTeamManager: () => null,208      resumeBackgroundAgent,209      reviveCompletedBackgroundAgent,210    } as unknown as Config;211    tool = new SendMessageTool(config);212  });213 214  it('queues a message for a running task', async () => {215    registry.register({216      agentId: 'agent-1',217      description: 'test agent',218      status: 'running',219      startTime: Date.now(),220      abortController: new AbortController(),221      isBackgrounded: true,222      outputFile: '/tmp/test.jsonl',223    });224 225    const result = await tool.validateBuildAndExecute(226      { task_id: 'agent-1', message: 'do more work' },227      new AbortController().signal,228    );229 230    expect(result.error).toBeUndefined();231    expect(result.llmContent).toContain('Message queued');232    expect(registry.get('agent-1')!.pendingMessages).toEqual(['do more work']);233  });234 235  it('queues multiple messages in order', async () => {236    registry.register({237      agentId: 'agent-1',238      description: 'test agent',239      status: 'running',240      startTime: Date.now(),241      abortController: new AbortController(),242      isBackgrounded: true,243      outputFile: '/tmp/test.jsonl',244    });245 246    await tool.validateBuildAndExecute(247      { task_id: 'agent-1', message: 'first' },248      new AbortController().signal,249    );250    await tool.validateBuildAndExecute(251      { task_id: 'agent-1', message: 'second' },252      new AbortController().signal,253    );254 255    expect(registry.get('agent-1')!.pendingMessages).toEqual([256      'first',257      'second',258    ]);259  });260 261  it('returns error for non-existent task', async () => {262    const result = await tool.validateBuildAndExecute(263      { task_id: 'nope', message: 'hello' },264      new AbortController().signal,265    );266 267    expect(result.error?.type).toBe(ToolErrorType.SEND_MESSAGE_NOT_FOUND);268    expect(result.llmContent).toContain('No background task found');269  });270 271  it('returns error for a failed (non-running, non-revivable) task', async () => {272    registry.register({273      agentId: 'agent-1',274      description: 'test agent',275      status: 'running',276      startTime: Date.now(),277      abortController: new AbortController(),278      isBackgrounded: true,279      outputFile: '/tmp/test.jsonl',280    });281    registry.fail('agent-1', 'boom');282 283    const result = await tool.validateBuildAndExecute(284      { task_id: 'agent-1', message: 'hello' },285      new AbortController().signal,286    );287 288    expect(result.error?.type).toBe(ToolErrorType.SEND_MESSAGE_NOT_RUNNING);289    expect(result.llmContent).toContain('not running');290    expect(reviveCompletedBackgroundAgent).not.toHaveBeenCalled();291  });292 293  it('rejects messages for a cancelled task', async () => {294    // Once task_stop fires, the reasoning loop is winding down — there is295    // no next tool-round boundary to drain into, so the message would be296    // silently dropped. Reject instead of accepting a message that will297    // never be delivered.298    registry.register({299      agentId: 'agent-1',300      description: 'test agent',301      status: 'running',302      startTime: Date.now(),303      abortController: new AbortController(),304      isBackgrounded: true,305      outputFile: '/tmp/test.jsonl',306    });307    registry.cancel('agent-1');308 309    const result = await tool.validateBuildAndExecute(310      { task_id: 'agent-1', message: 'too late' },311      new AbortController().signal,312    );313 314    expect(result.error?.type).toBe(ToolErrorType.SEND_MESSAGE_NOT_RUNNING);315    expect(registry.get('agent-1')!.pendingMessages).toEqual([]);316  });317 318  it('resumes a paused task and injects the message as continuation input', async () => {319    registry.register({320      agentId: 'agent-1',321      description: 'test agent',322      status: 'paused',323      startTime: Date.now(),324      abortController: new AbortController(),325      isBackgrounded: true,326      outputFile: '/tmp/test.jsonl',327    });328    resumeBackgroundAgent.mockResolvedValue(registry.get('agent-1'));329 330    const result = await tool.validateBuildAndExecute(331      { task_id: 'agent-1', message: 'pick up from the TODO list' },332      new AbortController().signal,333    );334 335    expect(resumeBackgroundAgent).toHaveBeenCalledWith(336      'agent-1',337      'pick up from the TODO list',338    );339    expect(result.error).toBeUndefined();340    expect(result.llmContent).toContain('resumed');341  });342 343  it('revives a completed task with the message as the next instruction', async () => {344    registry.register({345      agentId: 'agent-1',346      description: 'test agent',347      status: 'completed',348      startTime: Date.now(),349      abortController: new AbortController(),350      isBackgrounded: true,351      outputFile: '/tmp/test.jsonl',352      metaPath: '/tmp/test.meta.json',353    });354    reviveCompletedBackgroundAgent.mockResolvedValue(registry.get('agent-1'));355 356    const result = await tool.validateBuildAndExecute(357      { task_id: 'agent-1', message: 'now refactor the helper' },358      new AbortController().signal,359    );360 361    expect(reviveCompletedBackgroundAgent).toHaveBeenCalledWith(362      'agent-1',363      'now refactor the helper',364    );365    expect(resumeBackgroundAgent).not.toHaveBeenCalled();366    expect(result.error).toBeUndefined();367    expect(result.llmContent).toContain('revived');368    expect(result.returnDisplay).toContain('Revived');369  });370 371  it('returns error when a completed task cannot be revived', async () => {372    registry.register({373      agentId: 'agent-1',374      description: 'test agent',375      status: 'completed',376      startTime: Date.now(),377      abortController: new AbortController(),378      isBackgrounded: true,379      outputFile: '/tmp/test.jsonl',380      metaPath: '/tmp/test.meta.json',381    });382    reviveCompletedBackgroundAgent.mockResolvedValue(undefined);383 384    const result = await tool.validateBuildAndExecute(385      { task_id: 'agent-1', message: 'try again' },386      new AbortController().signal,387    );388 389    expect(result.error?.type).toBe(ToolErrorType.SEND_MESSAGE_NOT_RUNNING);390    expect(result.llmContent).toContain('could not be revived');391  });392 393  it('includes task description in success display', async () => {394    registry.register({395      agentId: 'agent-1',396      description: 'Search for auth code',397      status: 'running',398      startTime: Date.now(),399      abortController: new AbortController(),400      isBackgrounded: true,401      outputFile: '/tmp/test.jsonl',402    });403 404    const result = await tool.validateBuildAndExecute(405      { task_id: 'agent-1', message: 'focus on login' },406      new AbortController().signal,407    );408 409    expect(result.returnDisplay).toContain('Search for auth code');410  });411});412