Files
flynn/src/gateway/handlers/agent.test.ts
T
2026-02-25 10:22:44 -08:00

717 lines
24 KiB
TypeScript

import { describe, it, expect, vi, beforeEach } from 'vitest';
import type { GatewayEvent, GatewayRequest, OutboundMessage } from '../protocol.js';
import { LaneQueue, LaneQueueRejectedError } from '../lane-queue.js';
import { createAgentHandlers } from './agent.js';
import type { AgentHandlerDeps } from './agent.js';
import { CommandRegistry, registerBuiltinCommands } from '../../commands/index.js';
import { initAuditLogger } from '../../audit/index.js';
describe('createAgentHandlers command fast-path', () => {
const mockAgent = {
process: vi.fn(async () => 'agent response'),
consumeContextAlert: vi.fn(() => undefined as unknown),
getContextBudget: vi.fn(() => ({
estimatedTokens: 100,
contextWindow: 200000,
remainingTokens: 199900,
usagePct: 0.05,
thresholdPct: 80,
thresholdTokens: 160000,
shouldCompact: false,
})),
getUsage: vi.fn(() => ({
primary: { inputTokens: 10, outputTokens: 5, calls: 1 },
delegation: {},
total: { inputTokens: 10, outputTokens: 5, calls: 1, estimatedCost: 0 },
})),
getModelTier: vi.fn(() => 'default'),
setModelTier: vi.fn(),
compact: vi.fn(async () => null),
reset: vi.fn(),
};
const sessionBridge = {
getAgent: vi.fn(() => mockAgent),
getSessionId: vi.fn(() => 'ws:conn-1'),
setBusy: vi.fn(),
setOnToolUse: vi.fn(),
isBusy: vi.fn(() => false),
cancel: vi.fn(() => true),
};
const sessionManager = {
getSessionConfig: vi.fn(),
setSessionConfig: vi.fn(),
deleteSessionConfig: vi.fn(),
};
const commandRegistry = new CommandRegistry();
registerBuiltinCommands(commandRegistry);
const mockAuditLogger = {
userAction: vi.fn(),
queuePreempt: vi.fn(),
runState: vi.fn(),
runCancel: vi.fn(),
};
const handlers = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue: new LaneQueue(),
sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
commandRegistry,
});
beforeEach(() => {
vi.clearAllMocks();
initAuditLogger(mockAuditLogger as any);
mockAgent.process.mockResolvedValue('agent response');
mockAgent.compact.mockResolvedValue(null);
});
it('handles known commands without calling agent.process', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 1,
method: 'agent.send',
params: { message: '/usage', connectionId: 'conn-1' },
};
await handlers['agent.send'](req, send);
expect(mockAgent.process).not.toHaveBeenCalled();
expect(sent).toHaveLength(1);
const event = sent[0] as GatewayEvent;
expect(event.event).toBe('done');
expect((event.data as { content: string }).content).toContain('Token Usage');
});
it('handles /context command via fast-path', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 8,
method: 'agent.send',
params: { message: '/context', connectionId: 'conn-1' },
};
await handlers['agent.send'](req, send);
expect(mockAgent.process).not.toHaveBeenCalled();
expect(sent).toHaveLength(1);
const event = sent[0] as GatewayEvent;
expect(event.event).toBe('done');
expect((event.data as { content: string }).content).toContain('Context Usage');
expect((event.data as { content: string }).content).toContain('Compaction threshold');
});
it('handles metadata commands via fast-path', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 2,
method: 'agent.send',
params: {
message: '/reset',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'reset' },
},
};
await handlers['agent.send'](req, send);
expect(mockAgent.reset).toHaveBeenCalledOnce();
expect(sessionManager.deleteSessionConfig).toHaveBeenCalledWith('ws', 'ws:conn-1', 'modelTier');
expect(mockAgent.process).not.toHaveBeenCalled();
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Session reset.');
});
it('handles /model command via fast-path and persists session tier', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 4,
method: 'agent.send',
params: {
message: '/model fast',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'model', commandArgs: 'fast' },
},
};
await handlers['agent.send'](req, send);
expect(mockAgent.setModelTier).toHaveBeenCalledWith('fast');
expect(sessionManager.setSessionConfig).toHaveBeenCalledWith('ws', 'ws:conn-1', 'modelTier', 'fast');
expect(mockAgent.process).not.toHaveBeenCalled();
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Switched to model tier: fast');
});
it('handles /model <tier> <provider/model> in gateway sessions', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const modelRouter = {
setClient: vi.fn(),
setTierStrict: vi.fn(),
};
const handlersWithRouter = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue: new LaneQueue(),
sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
commandRegistry,
modelRouter: modelRouter as unknown as AgentHandlerDeps['modelRouter'],
runtimeConfig: {
models: {
default: { provider: 'anthropic', model: 'claude-sonnet-4', api_key: 'test-anthropic-key' },
fallback_chain: [],
},
} as unknown as AgentHandlerDeps['runtimeConfig'],
});
const req: GatewayRequest = {
id: 9,
method: 'agent.send',
params: {
message: '/model default github/gpt-5-mini',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'model', commandArgs: 'default github/gpt-5-mini' },
},
};
await handlersWithRouter['agent.send'](req, send);
expect(modelRouter.setClient).toHaveBeenCalledWith('default', expect.anything(), 'github/gpt-5-mini');
expect(modelRouter.setTierStrict).toHaveBeenCalledWith('default', true);
expect(mockAgent.setModelTier).toHaveBeenCalledWith('default');
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Set default to: github/gpt-5-mini');
});
it('handles /model <tier> reset in gateway sessions', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const modelRouter = {
setClient: vi.fn(),
setTierStrict: vi.fn(),
};
const handlersWithRouter = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue: new LaneQueue(),
sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
commandRegistry,
modelRouter: modelRouter as unknown as AgentHandlerDeps['modelRouter'],
runtimeConfig: {
models: {
default: { provider: 'anthropic', model: 'claude-sonnet-4', api_key: 'test-anthropic-key' },
fallback_chain: [],
},
} as unknown as AgentHandlerDeps['runtimeConfig'],
});
const req: GatewayRequest = {
id: 10,
method: 'agent.send',
params: {
message: '/model default reset',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'model', commandArgs: 'default reset' },
},
};
await handlersWithRouter['agent.send'](req, send);
expect(modelRouter.setClient).toHaveBeenCalledWith('default', expect.anything(), 'anthropic/claude-sonnet-4');
expect(modelRouter.setTierStrict).toHaveBeenCalledWith('default', false);
expect(mockAgent.setModelTier).toHaveBeenCalledWith('default');
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Reset default to: anthropic/claude-sonnet-4');
});
it('handles /stop command via fast-path and requests cancellation', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 12,
method: 'agent.send',
params: {
message: '/stop',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'stop' },
},
};
await handlers['agent.send'](req, send);
expect(sessionBridge.cancel).toHaveBeenCalledWith('conn-1');
expect(mockAuditLogger.runCancel).toHaveBeenCalledWith(
expect.objectContaining({
session_id: 'ws:conn-1',
source: 'gateway',
requested: true,
acknowledged: true,
}),
);
expect(mockAuditLogger.runState).toHaveBeenCalledWith(
expect.objectContaining({
session_id: 'ws:conn-1',
source: 'gateway',
state: 'cancel_requested',
}),
);
expect(mockAgent.process).not.toHaveBeenCalled();
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Cancellation requested');
});
it('handles /runtime status via command fast-path without calling agent.process', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 13,
method: 'agent.send',
params: {
message: '/runtime status',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'runtime', commandArgs: 'status' },
},
};
await handlers['agent.send'](req, send);
expect(mockAgent.process).not.toHaveBeenCalled();
expect((sent[0] as GatewayEvent).event).toBe('done');
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Backend mode:');
});
it('handles /runtime deactivate pi via shared backend mode service callbacks', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
let backendMode: 'config_default' | 'force_native' | 'force_pi_embedded' = 'force_pi_embedded';
const handlersWithBackendMode = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue: new LaneQueue(),
sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
commandRegistry,
runtimeConfig: {
backends: {
default: 'pi_embedded',
claude_code: { enabled: false },
opencode: { enabled: false },
codex: { enabled: false },
gemini: { enabled: false },
pi_embedded: { enabled: true },
},
} as unknown as AgentHandlerDeps['runtimeConfig'],
getBackendMode: () => backendMode,
setBackendMode: (mode) => {
backendMode = mode;
},
});
const req: GatewayRequest = {
id: 14,
method: 'agent.send',
params: {
message: '/runtime deactivate pi',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'runtime', commandArgs: 'deactivate pi' },
},
};
await handlersWithBackendMode['agent.send'](req, send);
expect(mockAgent.process).not.toHaveBeenCalled();
expect(backendMode).toBe('force_native');
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('deactivated globally');
});
it('falls through to agent.process for unknown commands', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 3,
method: 'agent.send',
params: { message: '/not-a-real-command', connectionId: 'conn-1' },
};
await handlers['agent.send'](req, send);
expect(mockAgent.process).toHaveBeenCalledWith('/not-a-real-command', undefined);
const doneEvent = sent.find((msg) => (msg as GatewayEvent).event === 'done') as GatewayEvent | undefined;
expect(doneEvent).toBeTruthy();
expect(((doneEvent as GatewayEvent).data as { content: string }).content).toBe('agent response');
});
it('handles /approvals command via fast-path when hook engine is available', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const hookEngine = {
getPendingConfirmations: vi.fn(() => []),
resolveConfirmation: vi.fn(() => false),
};
const handlersWithHooks = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue: new LaneQueue(),
sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
commandRegistry,
hookEngine: hookEngine as unknown as AgentHandlerDeps['hookEngine'],
});
const req: GatewayRequest = {
id: 11,
method: 'agent.send',
params: {
message: '/approvals',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'approvals' },
},
};
await handlersWithHooks['agent.send'](req, send);
expect(mockAgent.process).not.toHaveBeenCalled();
expect(hookEngine.getPendingConfirmations).toHaveBeenCalledWith({ sessionId: 'ws:conn-1' });
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('No pending approvals');
});
it('emits user.action audit events for gateway requests', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 7,
method: 'agent.send',
params: { message: 'hello there', connectionId: 'conn-1' },
};
await handlers['agent.send'](req, send);
expect(mockAuditLogger.userAction).toHaveBeenCalledWith(
expect.objectContaining({
source: 'gateway',
action_type: 'message',
sender: 'conn-1',
}),
);
expect(mockAuditLogger.runState).toHaveBeenCalledWith(
expect.objectContaining({
session_id: 'ws:conn-1',
source: 'gateway',
state: 'start',
}),
);
expect(mockAuditLogger.runState).toHaveBeenCalledWith(
expect.objectContaining({
session_id: 'ws:conn-1',
source: 'gateway',
state: 'complete',
}),
);
});
it('emits cancelled run state when agent returns cancellation text', async () => {
mockAgent.process.mockResolvedValueOnce('Operation cancelled by user.');
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 15,
method: 'agent.send',
params: { message: 'cancel me', connectionId: 'conn-1' },
};
await handlers['agent.send'](req, send);
expect(mockAuditLogger.runState).toHaveBeenCalledWith(
expect.objectContaining({
session_id: 'ws:conn-1',
source: 'gateway',
state: 'cancelled',
}),
);
expect(sent.some((msg) => (msg as GatewayEvent).event === 'done')).toBe(true);
});
it('emits run.cancel telemetry for agent.cancel requests', async () => {
const result = await handlers['agent.cancel']({
id: 16,
method: 'agent.cancel',
params: { connectionId: 'conn-1' },
}, vi.fn());
expect(sessionBridge.cancel).toHaveBeenCalledWith('conn-1');
expect(mockAuditLogger.runCancel).toHaveBeenCalledWith(
expect.objectContaining({
session_id: 'ws:conn-1',
source: 'gateway',
requested: true,
acknowledged: true,
}),
);
expect((result as { result: { cancelled: boolean } }).result.cancelled).toBe(true);
});
it('handles /queue command via fast-path and persists queue session config', async () => {
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
const req: GatewayRequest = {
id: 5,
method: 'agent.send',
params: {
message: '/queue set mode followup',
connectionId: 'conn-1',
metadata: { isCommand: true, command: 'queue', commandArgs: 'set mode followup' },
},
};
await handlers['agent.send'](req, send);
expect(sessionManager.setSessionConfig).toHaveBeenCalledWith('ws', 'ws:conn-1', 'queue.mode', 'followup');
expect(mockAgent.process).not.toHaveBeenCalled();
expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Set queue.mode=followup');
});
it('emits context_warning event before done when orchestrator reports an alert', async () => {
mockAgent.consumeContextAlert.mockReturnValueOnce({
level: 'checkpoint',
message: 'Context usage is 86.0% (172000/200000 estimated tokens).',
budget: {
estimatedTokens: 172000,
contextWindow: 200000,
remainingTokens: 28000,
usagePct: 86,
thresholdPct: 80,
thresholdTokens: 160000,
shouldCompact: true,
},
actions: {
checkpointSaved: true,
autoCompacted: false,
checkpointNamespace: 'session/checkpoints/ws/conn-1',
},
});
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
await handlers['agent.send']({
id: 6,
method: 'agent.send',
params: { message: 'hello', connectionId: 'conn-1' },
}, send);
const events = sent.map((msg) => (msg as GatewayEvent).event);
expect(events).toContain('context_warning');
expect(events).toContain('done');
expect(events.indexOf('context_warning')).toBeLessThan(events.indexOf('done'));
});
});
describe('createAgentHandlers queue policy resolution', () => {
const mockAuditLogger = {
userAction: vi.fn(),
queuePreempt: vi.fn(),
};
beforeEach(() => {
vi.clearAllMocks();
initAuditLogger(mockAuditLogger as any);
});
it('passes resolved per-request queue policy into lane enqueue', async () => {
const mockAgent = {
process: vi.fn(async () => 'ok'),
consumeContextAlert: vi.fn(() => undefined),
getContextBudget: vi.fn(() => ({
estimatedTokens: 0,
contextWindow: 128000,
remainingTokens: 128000,
usagePct: 0,
thresholdPct: 80,
thresholdTokens: 102400,
shouldCompact: false,
})),
getUsage: vi.fn(() => ({
primary: { inputTokens: 0, outputTokens: 0, calls: 0 },
delegation: {},
total: { inputTokens: 0, outputTokens: 0, calls: 0, estimatedCost: 0 },
})),
getModelTier: vi.fn(() => 'default'),
setModelTier: vi.fn(),
compact: vi.fn(async () => null),
reset: vi.fn(),
setOnToolUse: vi.fn(),
};
const sessionBridge = {
getAgent: vi.fn(() => mockAgent),
getSessionId: vi.fn(() => 'ws:s1'),
setBusy: vi.fn(),
setOnToolUse: vi.fn(),
isBusy: vi.fn(() => false),
};
const laneQueue = {
enqueue: vi.fn(async (_laneId: string, work: () => Promise<unknown>) => work()),
cancel: vi.fn(),
} as unknown as LaneQueue;
const resolveQueuePolicy = vi.fn(() => ({ mode: 'steer_backlog' as const, cap: 3, debounceMs: 25 }));
const handlers = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue,
resolveQueuePolicy,
});
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
await handlers['agent.send']({
id: 1,
method: 'agent.send',
params: { message: 'hello', connectionId: 'conn-1' },
}, send);
expect(resolveQueuePolicy).toHaveBeenCalledWith({
laneId: 'ws:s1',
sessionId: 'ws:s1',
connectionId: 'conn-1',
channel: 'ws',
});
expect((laneQueue.enqueue as unknown as ReturnType<typeof vi.fn>).mock.calls[0][2]).toEqual({
mode: 'steer_backlog',
cap: 3,
debounceMs: 25,
});
});
it('emits structured queue error events for lane rejections', async () => {
const sessionBridge = {
getAgent: vi.fn(() => ({
process: vi.fn(async () => 'ok'),
consumeContextAlert: vi.fn(() => undefined),
getContextBudget: vi.fn(() => ({
estimatedTokens: 0,
contextWindow: 128000,
remainingTokens: 128000,
usagePct: 0,
thresholdPct: 80,
thresholdTokens: 102400,
shouldCompact: false,
})),
getUsage: vi.fn(() => ({
primary: { inputTokens: 0, outputTokens: 0, calls: 0 },
delegation: {},
total: { inputTokens: 0, outputTokens: 0, calls: 0, estimatedCost: 0 },
})),
getModelTier: vi.fn(() => 'default'),
setModelTier: vi.fn(),
compact: vi.fn(async () => null),
reset: vi.fn(),
})),
getSessionId: vi.fn(() => 'ws:s1'),
setBusy: vi.fn(),
setOnToolUse: vi.fn(),
isBusy: vi.fn(() => false),
};
const laneQueue = {
enqueue: vi.fn(async () => {
throw new LaneQueueRejectedError({
code: 'overflow',
laneId: 'ws:s1',
mode: 'followup',
overflow: 'drop_new',
droppedCount: 1,
message: 'Lane queue full (drop_new)',
});
}),
cancel: vi.fn(),
} as unknown as LaneQueue;
const handlers = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue,
});
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
await handlers['agent.send']({
id: 6,
method: 'agent.send',
params: { message: 'hello', connectionId: 'conn-1' },
}, send);
expect(sent).toHaveLength(1);
const event = sent[0] as GatewayEvent;
expect(event.event).toBe('error');
expect((event.data as { code: number }).code).toBe(3);
expect((event.data as { queue?: { code: string } }).queue?.code).toBe('overflow');
});
it('requests active-session cancellation when interrupt mode receives a new message', async () => {
const mockAgent = {
process: vi.fn(async () => 'ok'),
consumeContextAlert: vi.fn(() => undefined),
getContextBudget: vi.fn(() => ({
estimatedTokens: 0,
contextWindow: 128000,
remainingTokens: 128000,
usagePct: 0,
thresholdPct: 80,
thresholdTokens: 102400,
shouldCompact: false,
})),
getUsage: vi.fn(() => ({
primary: { inputTokens: 0, outputTokens: 0, calls: 0 },
delegation: {},
total: { inputTokens: 0, outputTokens: 0, calls: 0, estimatedCost: 0 },
})),
getModelTier: vi.fn(() => 'default'),
setModelTier: vi.fn(),
compact: vi.fn(async () => null),
reset: vi.fn(),
};
const sessionBridge = {
getAgent: vi.fn(() => mockAgent),
getSessionId: vi.fn(() => 'ws:s1'),
setBusy: vi.fn(),
setOnToolUse: vi.fn(),
isBusy: vi.fn(() => false),
cancelSession: vi.fn(() => true),
cancel: vi.fn(() => true),
};
const laneQueue = {
enqueue: vi.fn(async (_laneId: string, work: () => Promise<unknown>) => work()),
cancel: vi.fn(),
isProcessing: vi.fn(() => true),
} as unknown as LaneQueue;
const handlers = createAgentHandlers({
sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
laneQueue,
resolveQueuePolicy: vi.fn(() => ({ mode: 'interrupt' as const })),
});
const sent: OutboundMessage[] = [];
const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
await handlers['agent.send']({
id: 7,
method: 'agent.send',
params: { message: 'newest', connectionId: 'conn-1' },
}, send);
expect(sessionBridge.cancelSession).toHaveBeenCalledWith('ws:s1');
expect(sessionBridge.cancel).not.toHaveBeenCalled();
expect(mockAuditLogger.queuePreempt).toHaveBeenCalledWith(expect.objectContaining({
session_id: 'ws:s1',
lane_id: 'ws:s1',
request_id: '7',
mode: 'interrupt',
cancelled_active_run: true,
}));
const contentEvent = sent.find((msg) => (msg as GatewayEvent).event === 'content') as GatewayEvent | undefined;
expect(contentEvent).toBeTruthy();
expect(((contentEvent as GatewayEvent).data as { text: string }).text).toContain('Interrupt mode');
expect(sent.some((msg) => (msg as GatewayEvent).event === 'done')).toBe(true);
});
});