mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-12 12:54:07 -07:00
feat(scheduler): support multi-scheduler tool aggregation and nested call IDs (#17429)
This commit is contained in:
@@ -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,
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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'];
|
||||||
|
|||||||
Reference in New Issue
Block a user