feat(scheduler): support multi-scheduler tool aggregation and nested call IDs (#17429)

This commit is contained in:
Abhi
2026-01-26 13:38:11 -05:00
committed by GitHub
parent 3e1a377d78
commit d745d86af1
9 changed files with 241 additions and 23 deletions
@@ -19,6 +19,7 @@ import {
type ToolCallsUpdateMessage, type ToolCallsUpdateMessage,
type AnyDeclarativeTool, type AnyDeclarativeTool,
type AnyToolInvocation, type AnyToolInvocation,
ROOT_SCHEDULER_ID,
} from '@google/gemini-cli-core'; } from '@google/gemini-cli-core';
import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js'; import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js';
@@ -73,6 +74,10 @@ describe('useToolExecutionScheduler', () => {
} as unknown as Config; } as unknown as Config;
}); });
afterEach(() => {
vi.clearAllMocks();
});
it('initializes with empty tool calls', () => { it('initializes with empty tool calls', () => {
const { result } = renderHook(() => const { result } = renderHook(() =>
useToolExecutionScheduler( useToolExecutionScheduler(
@@ -112,6 +117,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({ void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE, type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall], toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage); } as ToolCallsUpdateMessage);
}); });
@@ -156,6 +162,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({ void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE, type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall], toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage); } as ToolCallsUpdateMessage);
}); });
@@ -212,6 +219,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({ void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE, type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall], toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage); } as ToolCallsUpdateMessage);
}); });
@@ -274,6 +282,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({ void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE, type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall], toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage); } as ToolCallsUpdateMessage);
}); });
@@ -290,6 +299,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({ void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE, type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall], toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage); } as ToolCallsUpdateMessage);
}); });
@@ -326,6 +336,7 @@ describe('useToolExecutionScheduler', () => {
invocation: createMockInvocation(), invocation: createMockInvocation(),
}, },
], ],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage); } as ToolCallsUpdateMessage);
}); });
@@ -412,4 +423,103 @@ describe('useToolExecutionScheduler', () => {
expect(completedResult).toEqual([completedToolCall]); expect(completedResult).toEqual([completedToolCall]);
expect(onComplete).toHaveBeenCalledWith([completedToolCall]); expect(onComplete).toHaveBeenCalledWith([completedToolCall]);
}); });
it('setToolCallsForDisplay re-groups tools by schedulerId (Multi-Scheduler support)', () => {
const { result } = renderHook(() =>
useToolExecutionScheduler(
vi.fn().mockResolvedValue(undefined),
mockConfig,
() => undefined,
),
);
const callRoot = {
status: 'success' as const,
request: {
callId: 'call-root',
name: 'test',
args: {},
isClientInitiated: false,
prompt_id: 'p1',
},
tool: createMockTool(),
invocation: createMockInvocation(),
response: {
callId: 'call-root',
responseParts: [],
resultDisplay: 'OK',
error: undefined,
errorType: undefined,
},
schedulerId: ROOT_SCHEDULER_ID,
};
const callSub = {
...callRoot,
request: { ...callRoot.request, callId: 'call-sub' },
schedulerId: 'subagent-1',
};
// 1. Populate state with multiple schedulers
act(() => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [callRoot],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [callSub],
schedulerId: 'subagent-1',
} as ToolCallsUpdateMessage);
});
let [toolCalls] = result.current;
expect(toolCalls).toHaveLength(2);
expect(
toolCalls.find((t) => t.request.callId === 'call-root')?.schedulerId,
).toBe(ROOT_SCHEDULER_ID);
expect(
toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
).toBe('subagent-1');
// 2. Call setToolCallsForDisplay (e.g., simulate a manual update or clear)
act(() => {
const [, , , setToolCalls] = result.current;
setToolCalls((prev) =>
prev.map((t) => ({ ...t, responseSubmittedToGemini: true })),
);
});
// 3. Verify that tools are still present and maintain their scheduler IDs
// The internal map should have been re-grouped.
[toolCalls] = result.current;
expect(toolCalls).toHaveLength(2);
expect(toolCalls.every((t) => t.responseSubmittedToGemini)).toBe(true);
const updatedRoot = toolCalls.find((t) => t.request.callId === 'call-root');
const updatedSub = toolCalls.find((t) => t.request.callId === 'call-sub');
expect(updatedRoot?.schedulerId).toBe(ROOT_SCHEDULER_ID);
expect(updatedSub?.schedulerId).toBe('subagent-1');
// 4. Verify that a subsequent update to ONE scheduler doesn't wipe the other
act(() => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [{ ...callRoot, status: 'executing' }],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
[toolCalls] = result.current;
expect(toolCalls).toHaveLength(2);
expect(
toolCalls.find((t) => t.request.callId === 'call-root')?.status,
).toBe('executing');
expect(
toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
).toBe('subagent-1');
});
}); });
@@ -16,6 +16,7 @@ import {
Scheduler, Scheduler,
type EditorType, type EditorType,
type ToolCallsUpdateMessage, type ToolCallsUpdateMessage,
ROOT_SCHEDULER_ID,
} from '@google/gemini-cli-core'; } from '@google/gemini-cli-core';
import { useCallback, useState, useMemo, useEffect, useRef } from 'react'; import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
@@ -54,8 +55,10 @@ export function useToolExecutionScheduler(
CancelAllFn, CancelAllFn,
number, number,
] { ] {
// State stores Core objects, not Display objects // State stores tool calls organized by their originating schedulerId
const [toolCalls, setToolCalls] = useState<TrackedToolCall[]>([]); const [toolCallsMap, setToolCallsMap] = useState<
Record<string, TrackedToolCall[]>
>({});
const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0); const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0);
const messageBus = useMemo(() => config.getMessageBus(), [config]); const messageBus = useMemo(() => config.getMessageBus(), [config]);
@@ -76,6 +79,7 @@ export function useToolExecutionScheduler(
config, config,
messageBus, messageBus,
getPreferredEditor: () => getPreferredEditorRef.current(), getPreferredEditor: () => getPreferredEditorRef.current(),
schedulerId: ROOT_SCHEDULER_ID,
}), }),
[config, messageBus], [config, messageBus],
); );
@@ -88,15 +92,21 @@ export function useToolExecutionScheduler(
useEffect(() => { useEffect(() => {
const handler = (event: ToolCallsUpdateMessage) => { const handler = (event: ToolCallsUpdateMessage) => {
setToolCalls((prev) => { // Update output timer for UI spinners (Side Effect)
const adapted = internalAdaptToolCalls(event.toolCalls, prev); if (event.toolCalls.some((tc) => tc.status === 'executing')) {
setLastToolOutputTime(Date.now());
}
// Update output timer for UI spinners setToolCallsMap((prev) => {
if (event.toolCalls.some((tc) => tc.status === 'executing')) { const adapted = internalAdaptToolCalls(
setLastToolOutputTime(Date.now()); event.toolCalls,
} prev[event.schedulerId] ?? [],
);
return adapted; return {
...prev,
[event.schedulerId]: adapted,
};
}); });
}; };
@@ -109,12 +119,14 @@ export function useToolExecutionScheduler(
const schedule: ScheduleFn = useCallback( const schedule: ScheduleFn = useCallback(
async (request, signal) => { async (request, signal) => {
// Clear state for new run // Clear state for new run
setToolCalls([]); setToolCallsMap({});
// 1. Await Core Scheduler directly // 1. Await Core Scheduler directly
const results = await scheduler.schedule(request, signal); const results = await scheduler.schedule(request, signal);
// 2. Trigger legacy reinjection logic (useGeminiStream loop) // 2. Trigger legacy reinjection logic (useGeminiStream loop)
// Since this hook instance owns the "root" scheduler, we always trigger
// onComplete when it finishes its batch.
await onCompleteRef.current(results); await onCompleteRef.current(results);
return results; return results;
@@ -131,13 +143,52 @@ export function useToolExecutionScheduler(
const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback( const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback(
(callIdsToMark: string[]) => { (callIdsToMark: string[]) => {
setToolCalls((prevCalls) => setToolCallsMap((prevMap) => {
prevCalls.map((tc) => const nextMap = { ...prevMap };
callIdsToMark.includes(tc.request.callId) for (const [sid, calls] of Object.entries(nextMap)) {
? { ...tc, responseSubmittedToGemini: true } nextMap[sid] = calls.map((tc) =>
: tc, callIdsToMark.includes(tc.request.callId)
), ? { ...tc, responseSubmittedToGemini: true }
); : tc,
);
}
return nextMap;
});
},
[],
);
// Flatten the map for the UI components that expect a single list of tools.
const toolCalls = useMemo(
() => Object.values(toolCallsMap).flat(),
[toolCallsMap],
);
// Provide a setter that maintains compatibility with legacy [].
const setToolCallsForDisplay = useCallback(
(action: React.SetStateAction<TrackedToolCall[]>) => {
setToolCallsMap((prev) => {
const currentFlattened = Object.values(prev).flat();
const nextFlattened =
typeof action === 'function' ? action(currentFlattened) : action;
if (nextFlattened.length === 0) {
return {};
}
// Re-group by schedulerId to preserve multi-scheduler state
const nextMap: Record<string, TrackedToolCall[]> = {};
for (const call of nextFlattened) {
// All tool calls should have a schedulerId from the core.
// Default to ROOT_SCHEDULER_ID as a safeguard.
const sid = call.schedulerId ?? ROOT_SCHEDULER_ID;
if (!nextMap[sid]) {
nextMap[sid] = [];
}
nextMap[sid].push(call);
}
return nextMap;
});
}, },
[], [],
); );
@@ -146,7 +197,7 @@ export function useToolExecutionScheduler(
toolCalls, toolCalls,
schedule, schedule,
markToolsAsSubmitted, markToolsAsSubmitted,
setToolCalls, setToolCallsForDisplay,
cancelAll, cancelAll,
lastToolOutputTime, lastToolOutputTime,
]; ];
@@ -26,6 +26,7 @@ export enum MessageBusType {
export interface ToolCallsUpdateMessage { export interface ToolCallsUpdateMessage {
type: MessageBusType.TOOL_CALLS_UPDATE; type: MessageBusType.TOOL_CALLS_UPDATE;
toolCalls: ToolCall[]; toolCalls: ToolCall[];
schedulerId: string;
} }
export interface ToolConfirmationRequest { export interface ToolConfirmationRequest {
@@ -29,6 +29,7 @@ import {
import type { SchedulerStateManager } from './state-manager.js'; import type { SchedulerStateManager } from './state-manager.js';
import type { ToolModificationHandler } from './tool-modifier.js'; import type { ToolModificationHandler } from './tool-modifier.js';
import type { ValidatingToolCall, WaitingToolCall } from './types.js'; import type { ValidatingToolCall, WaitingToolCall } from './types.js';
import { ROOT_SCHEDULER_ID } from './types.js';
import type { Config } from '../config/config.js'; import type { Config } from '../config/config.js';
import type { EditorType } from '../utils/editor.js'; import type { EditorType } from '../utils/editor.js';
import { randomUUID } from 'node:crypto'; import { randomUUID } from 'node:crypto';
@@ -52,7 +53,7 @@ describe('confirmation.ts', () => {
}); });
afterEach(() => { afterEach(() => {
vi.clearAllMocks(); vi.restoreAllMocks();
}); });
const emitResponse = (response: ToolConfirmationResponse) => { const emitResponse = (response: ToolConfirmationResponse) => {
@@ -188,6 +189,7 @@ describe('confirmation.ts', () => {
state: mockState, state: mockState,
modifier: mockModifier, modifier: mockModifier,
getPreferredEditor, getPreferredEditor,
schedulerId: ROOT_SCHEDULER_ID,
}); });
expect(result.outcome).toBe(ToolConfirmationOutcome.ProceedOnce); expect(result.outcome).toBe(ToolConfirmationOutcome.ProceedOnce);
@@ -217,6 +219,7 @@ describe('confirmation.ts', () => {
state: mockState, state: mockState,
modifier: mockModifier, modifier: mockModifier,
getPreferredEditor, getPreferredEditor,
schedulerId: ROOT_SCHEDULER_ID,
}); });
await listenerPromise; await listenerPromise;
@@ -252,6 +255,7 @@ describe('confirmation.ts', () => {
state: mockState, state: mockState,
modifier: mockModifier, modifier: mockModifier,
getPreferredEditor, getPreferredEditor,
schedulerId: ROOT_SCHEDULER_ID,
}); });
await waitForListener(MessageBusType.TOOL_CONFIRMATION_RESPONSE); await waitForListener(MessageBusType.TOOL_CONFIRMATION_RESPONSE);
@@ -293,6 +297,7 @@ describe('confirmation.ts', () => {
state: mockState, state: mockState,
modifier: mockModifier, modifier: mockModifier,
getPreferredEditor, getPreferredEditor,
schedulerId: ROOT_SCHEDULER_ID,
}); });
await listenerPromise1; await listenerPromise1;
@@ -351,6 +356,7 @@ describe('confirmation.ts', () => {
state: mockState, state: mockState,
modifier: mockModifier, modifier: mockModifier,
getPreferredEditor, getPreferredEditor,
schedulerId: ROOT_SCHEDULER_ID,
}); });
await listenerPromise; await listenerPromise;
@@ -397,6 +403,7 @@ describe('confirmation.ts', () => {
state: mockState, state: mockState,
modifier: mockModifier, modifier: mockModifier,
getPreferredEditor, getPreferredEditor,
schedulerId: ROOT_SCHEDULER_ID,
}); });
const result = await promise; const result = await promise;
@@ -420,6 +427,7 @@ describe('confirmation.ts', () => {
state: mockState, state: mockState,
modifier: mockModifier, modifier: mockModifier,
getPreferredEditor, getPreferredEditor,
schedulerId: ROOT_SCHEDULER_ID,
}), }),
).rejects.toThrow(/lost during confirmation loop/); ).rejects.toThrow(/lost during confirmation loop/);
}); });
@@ -103,6 +103,7 @@ export async function resolveConfirmation(
state: SchedulerStateManager; state: SchedulerStateManager;
modifier: ToolModificationHandler; modifier: ToolModificationHandler;
getPreferredEditor: () => EditorType | undefined; getPreferredEditor: () => EditorType | undefined;
schedulerId: string;
}, },
): Promise<ResolutionResult> { ): Promise<ResolutionResult> {
const { state } = deps; const { state } = deps;
@@ -66,6 +66,7 @@ import type {
CancelledToolCall, CancelledToolCall,
ToolCallResponseInfo, ToolCallResponseInfo,
} from './types.js'; } from './types.js';
import { ROOT_SCHEDULER_ID } from './types.js';
import { ToolErrorType } from '../tools/tool-error.js'; import { ToolErrorType } from '../tools/tool-error.js';
import * as ToolUtils from '../utils/tool-utils.js'; import * as ToolUtils from '../utils/tool-utils.js';
import type { EditorType } from '../utils/editor.js'; import type { EditorType } from '../utils/editor.js';
@@ -94,6 +95,8 @@ describe('Scheduler (Orchestrator)', () => {
args: { foo: 'bar' }, args: { foo: 'bar' },
isClientInitiated: false, isClientInitiated: false,
prompt_id: 'prompt-1', prompt_id: 'prompt-1',
schedulerId: ROOT_SCHEDULER_ID,
parentCallId: undefined,
}; };
const req2: ToolCallRequestInfo = { const req2: ToolCallRequestInfo = {
@@ -102,6 +105,8 @@ describe('Scheduler (Orchestrator)', () => {
args: { foo: 'baz' }, args: { foo: 'baz' },
isClientInitiated: false, isClientInitiated: false,
prompt_id: 'prompt-1', prompt_id: 'prompt-1',
schedulerId: ROOT_SCHEDULER_ID,
parentCallId: undefined,
}; };
const mockTool = { const mockTool = {
@@ -208,6 +213,7 @@ describe('Scheduler (Orchestrator)', () => {
config: mockConfig, config: mockConfig,
messageBus: mockMessageBus, messageBus: mockMessageBus,
getPreferredEditor, getPreferredEditor,
schedulerId: 'root',
}); });
// Reset Tool build behavior // Reset Tool build behavior
@@ -271,6 +277,8 @@ describe('Scheduler (Orchestrator)', () => {
request: req1, request: req1,
tool: mockTool, tool: mockTool,
invocation: mockInvocation, invocation: mockInvocation,
schedulerId: ROOT_SCHEDULER_ID,
startTime: expect.any(Number),
}), }),
]), ]),
); );
@@ -769,6 +777,7 @@ describe('Scheduler (Orchestrator)', () => {
config: mockConfig, config: mockConfig,
messageBus: mockMessageBus, messageBus: mockMessageBus,
state: mockStateManager, state: mockStateManager,
schedulerId: ROOT_SCHEDULER_ID,
}), }),
); );
+18 -3
View File
@@ -48,6 +48,8 @@ export interface SchedulerOptions {
config: Config; config: Config;
messageBus: MessageBus; messageBus: MessageBus;
getPreferredEditor: () => EditorType | undefined; getPreferredEditor: () => EditorType | undefined;
schedulerId: string;
parentCallId?: string;
} }
const createErrorResponse = ( const createErrorResponse = (
@@ -85,6 +87,8 @@ export class Scheduler {
private readonly config: Config; private readonly config: Config;
private readonly messageBus: MessageBus; private readonly messageBus: MessageBus;
private readonly getPreferredEditor: () => EditorType | undefined; private readonly getPreferredEditor: () => EditorType | undefined;
private readonly schedulerId: string;
private readonly parentCallId?: string;
private isProcessing = false; private isProcessing = false;
private isCancelling = false; private isCancelling = false;
@@ -94,7 +98,9 @@ export class Scheduler {
this.config = options.config; this.config = options.config;
this.messageBus = options.messageBus; this.messageBus = options.messageBus;
this.getPreferredEditor = options.getPreferredEditor; this.getPreferredEditor = options.getPreferredEditor;
this.state = new SchedulerStateManager(this.messageBus); this.schedulerId = options.schedulerId;
this.parentCallId = options.parentCallId;
this.state = new SchedulerStateManager(this.messageBus, this.schedulerId);
this.executor = new ToolExecutor(this.config); this.executor = new ToolExecutor(this.config);
this.modifier = new ToolModificationHandler(); this.modifier = new ToolModificationHandler();
@@ -228,16 +234,21 @@ export class Scheduler {
try { try {
const toolRegistry = this.config.getToolRegistry(); const toolRegistry = this.config.getToolRegistry();
const newCalls: ToolCall[] = requests.map((request) => { const newCalls: ToolCall[] = requests.map((request) => {
const enrichedRequest: ToolCallRequestInfo = {
...request,
schedulerId: this.schedulerId,
parentCallId: this.parentCallId,
};
const tool = toolRegistry.getTool(request.name); const tool = toolRegistry.getTool(request.name);
if (!tool) { if (!tool) {
return this._createToolNotFoundErroredToolCall( return this._createToolNotFoundErroredToolCall(
request, enrichedRequest,
toolRegistry.getAllToolNames(), toolRegistry.getAllToolNames(),
); );
} }
return this._validateAndCreateToolCall(request, tool); return this._validateAndCreateToolCall(enrichedRequest, tool);
}); });
this.state.enqueue(newCalls); this.state.enqueue(newCalls);
@@ -263,6 +274,7 @@ export class Scheduler {
ToolErrorType.TOOL_NOT_REGISTERED, ToolErrorType.TOOL_NOT_REGISTERED,
), ),
durationMs: 0, durationMs: 0,
schedulerId: this.schedulerId,
}; };
} }
@@ -278,6 +290,7 @@ export class Scheduler {
tool, tool,
invocation, invocation,
startTime: Date.now(), startTime: Date.now(),
schedulerId: this.schedulerId,
}; };
} catch (e) { } catch (e) {
return { return {
@@ -290,6 +303,7 @@ export class Scheduler {
ToolErrorType.INVALID_TOOL_PARAMS, ToolErrorType.INVALID_TOOL_PARAMS,
), ),
durationMs: 0, durationMs: 0,
schedulerId: this.schedulerId,
}; };
} }
} }
@@ -411,6 +425,7 @@ export class Scheduler {
state: this.state, state: this.state,
modifier: this.modifier, modifier: this.modifier,
getPreferredEditor: this.getPreferredEditor, getPreferredEditor: this.getPreferredEditor,
schedulerId: this.schedulerId,
}); });
outcome = result.outcome; outcome = result.outcome;
lastDetails = result.lastDetails; lastDetails = result.lastDetails;
+13 -1
View File
@@ -17,6 +17,7 @@ import type {
ExecutingToolCall, ExecutingToolCall,
ToolCallResponseInfo, ToolCallResponseInfo,
} from './types.js'; } from './types.js';
import { ROOT_SCHEDULER_ID } from './types.js';
import type { import type {
ToolConfirmationOutcome, ToolConfirmationOutcome,
ToolResultDisplay, ToolResultDisplay,
@@ -39,7 +40,10 @@ export class SchedulerStateManager {
private readonly queue: ToolCall[] = []; private readonly queue: ToolCall[] = [];
private _completedBatch: CompletedToolCall[] = []; private _completedBatch: CompletedToolCall[] = [];
constructor(private readonly messageBus: MessageBus) {} constructor(
private readonly messageBus: MessageBus,
private readonly schedulerId: string = ROOT_SCHEDULER_ID,
) {}
addToolCalls(calls: ToolCall[]): void { addToolCalls(calls: ToolCall[]): void {
this.enqueue(calls); this.enqueue(calls);
@@ -201,6 +205,7 @@ export class SchedulerStateManager {
void this.messageBus.publish({ void this.messageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE, type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: snapshot, toolCalls: snapshot,
schedulerId: this.schedulerId,
}); });
} }
@@ -321,6 +326,7 @@ export class SchedulerStateManager {
response, response,
durationMs: startTime ? Date.now() - startTime : undefined, durationMs: startTime ? Date.now() - startTime : undefined,
outcome: call.outcome, outcome: call.outcome,
schedulerId: call.schedulerId,
}; };
} }
@@ -336,6 +342,7 @@ export class SchedulerStateManager {
response, response,
durationMs: startTime ? Date.now() - startTime : undefined, durationMs: startTime ? Date.now() - startTime : undefined,
outcome: call.outcome, outcome: call.outcome,
schedulerId: call.schedulerId,
}; };
} }
@@ -364,6 +371,7 @@ export class SchedulerStateManager {
startTime: 'startTime' in call ? call.startTime : undefined, startTime: 'startTime' in call ? call.startTime : undefined,
outcome: call.outcome, outcome: call.outcome,
invocation: call.invocation, invocation: call.invocation,
schedulerId: call.schedulerId,
}; };
} }
@@ -388,6 +396,7 @@ export class SchedulerStateManager {
startTime: 'startTime' in call ? call.startTime : undefined, startTime: 'startTime' in call ? call.startTime : undefined,
outcome: call.outcome, outcome: call.outcome,
invocation: call.invocation, invocation: call.invocation,
schedulerId: call.schedulerId,
}; };
} }
@@ -442,6 +451,7 @@ export class SchedulerStateManager {
}, },
durationMs: startTime ? Date.now() - startTime : undefined, durationMs: startTime ? Date.now() - startTime : undefined,
outcome: call.outcome, outcome: call.outcome,
schedulerId: call.schedulerId,
}; };
} }
@@ -462,6 +472,7 @@ export class SchedulerStateManager {
startTime: 'startTime' in call ? call.startTime : undefined, startTime: 'startTime' in call ? call.startTime : undefined,
outcome: call.outcome, outcome: call.outcome,
invocation: call.invocation, invocation: call.invocation,
schedulerId: call.schedulerId,
}; };
} }
@@ -482,6 +493,7 @@ export class SchedulerStateManager {
invocation: call.invocation, invocation: call.invocation,
liveOutput, liveOutput,
pid, pid,
schedulerId: call.schedulerId,
}; };
} }
} }
+11
View File
@@ -16,6 +16,8 @@ import type { AnsiOutput } from '../utils/terminalSerializer.js';
import type { ToolErrorType } from '../tools/tool-error.js'; import type { ToolErrorType } from '../tools/tool-error.js';
import type { SerializableConfirmationDetails } from '../confirmation-bus/types.js'; import type { SerializableConfirmationDetails } from '../confirmation-bus/types.js';
export const ROOT_SCHEDULER_ID = 'root';
export interface ToolCallRequestInfo { export interface ToolCallRequestInfo {
callId: string; callId: string;
name: string; name: string;
@@ -24,6 +26,8 @@ export interface ToolCallRequestInfo {
prompt_id: string; prompt_id: string;
checkpoint?: string; checkpoint?: string;
traceId?: string; traceId?: string;
parentCallId?: string;
schedulerId?: string;
} }
export interface ToolCallResponseInfo { export interface ToolCallResponseInfo {
@@ -43,6 +47,7 @@ export type ValidatingToolCall = {
invocation: AnyToolInvocation; invocation: AnyToolInvocation;
startTime?: number; startTime?: number;
outcome?: ToolConfirmationOutcome; outcome?: ToolConfirmationOutcome;
schedulerId?: string;
}; };
export type ScheduledToolCall = { export type ScheduledToolCall = {
@@ -52,6 +57,7 @@ export type ScheduledToolCall = {
invocation: AnyToolInvocation; invocation: AnyToolInvocation;
startTime?: number; startTime?: number;
outcome?: ToolConfirmationOutcome; outcome?: ToolConfirmationOutcome;
schedulerId?: string;
}; };
export type ErroredToolCall = { export type ErroredToolCall = {
@@ -61,6 +67,7 @@ export type ErroredToolCall = {
tool?: AnyDeclarativeTool; tool?: AnyDeclarativeTool;
durationMs?: number; durationMs?: number;
outcome?: ToolConfirmationOutcome; outcome?: ToolConfirmationOutcome;
schedulerId?: string;
}; };
export type SuccessfulToolCall = { export type SuccessfulToolCall = {
@@ -71,6 +78,7 @@ export type SuccessfulToolCall = {
invocation: AnyToolInvocation; invocation: AnyToolInvocation;
durationMs?: number; durationMs?: number;
outcome?: ToolConfirmationOutcome; outcome?: ToolConfirmationOutcome;
schedulerId?: string;
}; };
export type ExecutingToolCall = { export type ExecutingToolCall = {
@@ -82,6 +90,7 @@ export type ExecutingToolCall = {
startTime?: number; startTime?: number;
outcome?: ToolConfirmationOutcome; outcome?: ToolConfirmationOutcome;
pid?: number; pid?: number;
schedulerId?: string;
}; };
export type CancelledToolCall = { export type CancelledToolCall = {
@@ -92,6 +101,7 @@ export type CancelledToolCall = {
invocation: AnyToolInvocation; invocation: AnyToolInvocation;
durationMs?: number; durationMs?: number;
outcome?: ToolConfirmationOutcome; outcome?: ToolConfirmationOutcome;
schedulerId?: string;
}; };
export type WaitingToolCall = { export type WaitingToolCall = {
@@ -113,6 +123,7 @@ export type WaitingToolCall = {
correlationId?: string; correlationId?: string;
startTime?: number; startTime?: number;
outcome?: ToolConfirmationOutcome; outcome?: ToolConfirmationOutcome;
schedulerId?: string;
}; };
export type Status = ToolCall['status']; export type Status = ToolCall['status'];