feat: preempt active runs in interrupt queue mode
This commit is contained in:
@@ -367,4 +367,64 @@ describe('createAgentHandlers queue policy resolution', () => {
|
||||
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((sent[0] as GatewayEvent).event).toBe('done');
|
||||
});
|
||||
});
|
||||
|
||||
@@ -58,6 +58,20 @@ export function createAgentHandlers(deps: AgentHandlerDeps) {
|
||||
const laneId = sessionId ?? connectionId;
|
||||
const channel = sessionId?.split(':', 1)[0] ?? 'ws';
|
||||
const resolvedPolicy = deps.resolveQueuePolicy?.({ laneId, sessionId, connectionId, channel });
|
||||
const laneQueueWithProcessing = deps.laneQueue as LaneQueue & { isProcessing?: (lane: string) => boolean };
|
||||
const laneIsProcessing = typeof laneQueueWithProcessing.isProcessing === 'function'
|
||||
? laneQueueWithProcessing.isProcessing(laneId)
|
||||
: false;
|
||||
|
||||
// Interrupt mode should preempt active work when a newer request arrives.
|
||||
// LaneQueue itself only rejects queued entries, so we also request agent cancellation.
|
||||
if (resolvedPolicy?.mode === 'interrupt' && laneIsProcessing) {
|
||||
if (sessionId) {
|
||||
deps.sessionBridge.cancelSession(sessionId);
|
||||
} else {
|
||||
deps.sessionBridge.cancel(connectionId);
|
||||
}
|
||||
}
|
||||
|
||||
// Enqueue the work — if the lane is idle it runs immediately,
|
||||
// otherwise it waits for earlier requests on the same session to finish.
|
||||
|
||||
@@ -185,6 +185,30 @@ describe('SessionBridge', () => {
|
||||
expect(cancelSpy).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('cancelSession requests cancellation when session has a busy connection', () => {
|
||||
const bridge = createBridge();
|
||||
bridge.connect('conn-1');
|
||||
bridge.connect('conn-2');
|
||||
bridge.switchSession('conn-2', 'ws:conn-1');
|
||||
|
||||
const agent = bridge.getAgent('conn-1');
|
||||
if (!agent) {
|
||||
throw new Error('Expected agent for conn-1');
|
||||
}
|
||||
const cancelSpy = vi.spyOn(agent, 'cancel');
|
||||
|
||||
bridge.setBusy('conn-2', true);
|
||||
expect(bridge.cancelSession('ws:conn-1')).toBe(true);
|
||||
expect(cancelSpy).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('cancelSession returns false when no active operation exists', () => {
|
||||
const bridge = createBridge();
|
||||
bridge.connect('conn-1');
|
||||
|
||||
expect(bridge.cancelSession('ws:conn-1')).toBe(false);
|
||||
});
|
||||
|
||||
it('switchSession changes session for a connection', () => {
|
||||
const bridge = createBridge();
|
||||
bridge.connect('conn-1');
|
||||
|
||||
@@ -132,6 +132,30 @@ export class SessionBridge {
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Request cancellation for the active operation in a session.
|
||||
* Returns true if at least one connection in the session is currently busy.
|
||||
*/
|
||||
cancelSession(sessionId: string): boolean {
|
||||
const clients = Array.from(this.clients.values()).filter((client) => client.sessionId === sessionId);
|
||||
if (clients.length === 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const hasBusyClient = clients.some((client) => client.busy);
|
||||
if (!hasBusyClient) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const agent = this.agents.get(sessionId);
|
||||
if (!agent) {
|
||||
return false;
|
||||
}
|
||||
|
||||
agent.cancel();
|
||||
return true;
|
||||
}
|
||||
|
||||
/** Set onToolUse callback for a connection's agent. */
|
||||
setOnToolUse(connectionId: string, callback: ((event: ToolUseEvent) => void) | undefined): void {
|
||||
const client = this.clients.get(connectionId);
|
||||
|
||||
Reference in New Issue
Block a user