mirror of
https://github.com/Gitlawb/openclaude.git
synced 2026-08-24 10:14:19 -05:00
feat(query): add lifecycle identity and terminal reasons (#1682)
* feat(query): add lifecycle identity and terminal reasons * fix(query): isolate lifecycle tracking context * fix(query): guard lifecycle metadata updates * fix(query): track lifecycle during tool waits * test(query): cover bash lifecycle metadata * fix(query): scope lifecycle tracking to request attempts * fix(query): disambiguate lifecycle abort log reason * fix(query): preserve foreground subagent lifecycle tracking * fix(query): clean up timeout and fallback lifecycle events * fix(query): emit timeout end after cleanup
This commit is contained in:
Bogdan
@@ -64,6 +64,7 @@ import type { FileStateCache } from './utils/fileStateCache.js'
|
||||
import type { DenialTrackingState } from './utils/permissions/denialTracking.js'
|
||||
import type { SystemPrompt } from './utils/systemPromptType.js'
|
||||
import type { ContentReplacementState } from './utils/toolResultStorage.js'
|
||||
import type { QueryLifecycleOperationTracker } from './utils/queryLifecycle.js'
|
||||
|
||||
// Re-export progress types for backwards compatibility
|
||||
export type {
|
||||
@@ -291,6 +292,7 @@ export type ToolUseContext = {
|
||||
toolInputSummary?: string | null,
|
||||
) => (request: PromptRequest) => Promise<PromptResponse>
|
||||
toolUseId?: string
|
||||
queryLifecycle?: QueryLifecycleOperationTracker
|
||||
criticalSystemReminder_EXPERIMENTAL?: string
|
||||
/** When true, preserve toolUseResult on messages even for subagents.
|
||||
* Used by in-process teammates whose transcripts are viewable by the user. */
|
||||
|
||||
@@ -909,6 +909,7 @@ async function* queryLoop(
|
||||
c => c.type === 'pending',
|
||||
),
|
||||
queryTracking,
|
||||
queryLifecycle: toolUseContext.queryLifecycle,
|
||||
effortValue: appState.effortValue,
|
||||
advisorModel: appState.advisorModel,
|
||||
skipCacheWrite,
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
import { describe, expect, test } from 'bun:test'
|
||||
import { readFileSync } from 'node:fs'
|
||||
import { join } from 'node:path'
|
||||
|
||||
const source = readFileSync(join(import.meta.dirname, 'REPL.tsx'), 'utf8')
|
||||
|
||||
function getAbortTimedOutQueryBody(): string {
|
||||
const start = source.indexOf('const abortTimedOutQuery = useCallback')
|
||||
expect(start).toBeGreaterThan(-1)
|
||||
const end = source.indexOf('}, [mrOnTurnComplete, resetLoadingState])', start)
|
||||
expect(end).toBeGreaterThan(start)
|
||||
return source.slice(start, end)
|
||||
}
|
||||
|
||||
function getQueryFinallyBody(): string {
|
||||
const queryStart = source.indexOf('await onQueryImpl(')
|
||||
expect(queryStart).toBeGreaterThan(-1)
|
||||
const finallyStart = source.indexOf('} finally {', queryStart)
|
||||
expect(finallyStart).toBeGreaterThan(queryStart)
|
||||
const finallyEnd = source.indexOf('// Auto-restore:', finallyStart)
|
||||
expect(finallyEnd).toBeGreaterThan(finallyStart)
|
||||
return source.slice(finallyStart, finallyEnd)
|
||||
}
|
||||
|
||||
describe('REPL query lifecycle timeout logging', () => {
|
||||
test('does not emit terminal timeout end from timeout handler', () => {
|
||||
const body = getAbortTimedOutQueryBody()
|
||||
const queueMicrotaskIndex = body.indexOf('queueMicrotask(() => {')
|
||||
expect(queueMicrotaskIndex).toBeGreaterThan(-1)
|
||||
|
||||
const abortAcknowledgedIndex = body.indexOf(
|
||||
"logQueryLifecycle('abort_acknowledged'",
|
||||
queueMicrotaskIndex,
|
||||
)
|
||||
|
||||
expect(abortAcknowledgedIndex).toBeGreaterThan(queueMicrotaskIndex)
|
||||
expect(body).not.toContain("logQueryLifecycle('end'")
|
||||
})
|
||||
|
||||
test('emits timeout end from the query finally cleanup path', () => {
|
||||
const body = getQueryFinallyBody()
|
||||
|
||||
expect(body).toContain('const guardCompletedContext = queryGuard.lastContext')
|
||||
expect(body).toContain("guardCompletedContext?.terminalReason === 'query-timeout'")
|
||||
expect(body).toContain("guardCompletedContext?.terminalReason === 'hard-max-query-timeout'")
|
||||
expect(body).toContain('guardCompletedContext.queryGeneration === thisGeneration')
|
||||
expect(body).toContain('logCompletedLifecycle(guardCompletedContext)')
|
||||
})
|
||||
})
|
||||
+114
-10
@@ -36,6 +36,7 @@ import { updateLastInteractionTime, getLastInteractionTime, getOriginalCwd, getP
|
||||
import { asSessionId, asAgentId } from '../types/ids.js';
|
||||
import { logForDebugging } from '../utils/debug.js';
|
||||
import { QueryGuard } from '../utils/QueryGuard.js';
|
||||
import { QueryLifecycleOperationTracker, formatQueryLifecycleAbortSignalReason, formatQueryLifecycleLogMessage, type QueryActiveOperationSnapshot, type QueryGuardTimeoutInfo, type QueryLifecycleContext, type QueryTerminalReason } from '../utils/queryLifecycle.js';
|
||||
import { createCombinedAbortSignal } from '../utils/combinedAbortSignal.js';
|
||||
import { isEnvTruthy } from '../utils/envUtils.js';
|
||||
import { formatTokens, truncateToWidth } from '../utils/format.js';
|
||||
@@ -551,6 +552,34 @@ function _temp2(setFrame_0) {
|
||||
function _temp(f) {
|
||||
return (f + 1) % TITLE_ANIMATION_FRAMES.length;
|
||||
}
|
||||
function getAbortReasonLabel(reason: unknown): string | undefined {
|
||||
if (reason === undefined) return undefined;
|
||||
if (typeof reason === 'string') return reason;
|
||||
if (reason instanceof Error) return reason.name;
|
||||
return String(reason);
|
||||
}
|
||||
function getQueryTerminalReason(signal: AbortSignal, didThrow: boolean): QueryTerminalReason {
|
||||
if (!signal.aborted) return didThrow ? 'unknown' : 'ok';
|
||||
switch (getAbortReasonLabel(signal.reason)) {
|
||||
case 'query-timeout':
|
||||
return 'query-timeout';
|
||||
case 'user-cancel':
|
||||
case 'interrupt':
|
||||
return 'user-abort';
|
||||
case 'background':
|
||||
return 'parent-ended';
|
||||
default:
|
||||
return 'unknown';
|
||||
}
|
||||
}
|
||||
function summarizeActiveOperations(snapshot: QueryActiveOperationSnapshot): string {
|
||||
const apiIds = snapshot.apiCalls.map(call => call.requestId ?? call.clientRequestId ?? 'unknown').join(',');
|
||||
const toolIds = snapshot.toolUses.map(tool => `${tool.toolName}:${tool.toolUseId}`).join(',');
|
||||
return `activeApiCalls=${snapshot.apiCalls.length} activeToolUses=${snapshot.toolUses.length}` + (apiIds ? ` apiIds=${apiIds}` : '') + (toolIds ? ` toolIds=${toolIds}` : '');
|
||||
}
|
||||
function logQueryLifecycle(event: string, context: QueryLifecycleContext, extras = ''): void {
|
||||
logForDebugging(formatQueryLifecycleLogMessage(event, context, extras));
|
||||
}
|
||||
export type Props = {
|
||||
commands: Command[];
|
||||
debug: boolean;
|
||||
@@ -1449,6 +1478,7 @@ export function REPL({
|
||||
}, [setLocalCommands]);
|
||||
const [inProgressToolUseIDs, setInProgressToolUseIDs] = useState<Set<string>>(new Set());
|
||||
const hasInterruptibleToolInProgressRef = useRef(false);
|
||||
const queryLifecycleTrackerRef = useRef(new QueryLifecycleOperationTracker());
|
||||
|
||||
// Remote session hook - manages WebSocket connection and message handling for --remote mode
|
||||
const remoteSession = useRemoteSession({
|
||||
@@ -1726,11 +1756,17 @@ export function REPL({
|
||||
const mrOnBeforeQuery = useCallback(async (_input: string, _allMessages: MessageType[], _newMessageCount: number) => true, []);
|
||||
const mrOnTurnComplete = useCallback(async (_allMessages: MessageType[], _aborted: boolean) => { }, []);
|
||||
const mrRender = useCallback(() => null, []);
|
||||
const abortTimedOutQuery = useCallback(() => {
|
||||
const abortTimedOutQuery = useCallback((timeout: QueryGuardTimeoutInfo) => {
|
||||
const timeoutOperations = summarizeActiveOperations(timeout.activeOperations);
|
||||
logQueryLifecycle('timeout', timeout.context, timeoutOperations);
|
||||
const activeAbortController = abortControllerRef.current;
|
||||
if (activeAbortController && !activeAbortController.signal.aborted) {
|
||||
logQueryLifecycle('abort_requested', timeout.context, formatQueryLifecycleAbortSignalReason('query-timeout'));
|
||||
activeAbortController.abort('query-timeout');
|
||||
}
|
||||
if (timeout.activeOperations.apiCalls.length > 0) {
|
||||
logForDebugging(`api.call.active_on_abort queryId=${timeout.context.queryId} generation=${timeout.generation} ${timeoutOperations}`);
|
||||
}
|
||||
if (feature('TOKEN_BUDGET')) {
|
||||
snapshotOutputTokensForTurn(null);
|
||||
}
|
||||
@@ -1738,6 +1774,7 @@ export function REPL({
|
||||
// QueryGuard calls this before forceEnd(); defer UI cleanup until after
|
||||
// the guard has released so the normal stale-generation finally path skips.
|
||||
queueMicrotask(() => {
|
||||
logQueryLifecycle('abort_acknowledged', timeout.context, formatQueryLifecycleAbortSignalReason('query-timeout'));
|
||||
resetLoadingState();
|
||||
setAbortController(null);
|
||||
void mrOnTurnComplete(messagesRef.current, true);
|
||||
@@ -2229,7 +2266,17 @@ export function REPL({
|
||||
if (feature('PROACTIVE') || feature('KAIROS')) {
|
||||
proactiveModule?.pauseProactive();
|
||||
}
|
||||
queryGuard.forceEnd();
|
||||
const cancelContext = queryGuard.activeContext;
|
||||
const cancelOperations = queryLifecycleTrackerRef.current.snapshot();
|
||||
const completedCancelContext = cancelContext ? {
|
||||
...cancelContext,
|
||||
terminalReason: 'user-abort' as const,
|
||||
abortReason: 'user-cancel'
|
||||
} : null;
|
||||
if (cancelContext) {
|
||||
logQueryLifecycle('abort_requested', cancelContext, formatQueryLifecycleAbortSignalReason('user-cancel'));
|
||||
}
|
||||
queryGuard.forceEnd('user-abort', 'user-cancel');
|
||||
skipIdleCheckRef.current = false;
|
||||
|
||||
// Preserve partially-streamed text so the user can read what was
|
||||
@@ -2265,6 +2312,16 @@ export function REPL({
|
||||
} else {
|
||||
abortController?.abort('user-cancel');
|
||||
}
|
||||
if (cancelContext) {
|
||||
logQueryLifecycle('abort_acknowledged', cancelContext, formatQueryLifecycleAbortSignalReason('user-cancel'));
|
||||
}
|
||||
if (completedCancelContext) {
|
||||
const cancelOperationSummary = summarizeActiveOperations(cancelOperations);
|
||||
if (cancelOperations.apiCalls.length > 0) {
|
||||
logForDebugging(`api.call.active_on_abort queryId=${completedCancelContext.queryId} generation=${completedCancelContext.queryGeneration} ${cancelOperationSummary}`);
|
||||
}
|
||||
logQueryLifecycle('end', completedCancelContext, cancelOperationSummary);
|
||||
}
|
||||
|
||||
// Clear the controller so subsequent Escape presses don't see a stale
|
||||
// aborted signal. Without this, canCancelRunningTask is false (signal
|
||||
@@ -2503,7 +2560,7 @@ export function REPL({
|
||||
reject
|
||||
}]);
|
||||
}), []);
|
||||
const getToolUseContext = useCallback((messages: MessageType[], newMessages: MessageType[], abortController: AbortController, mainLoopModel: string, queryGeneration?: number): ProcessUserInputContext => {
|
||||
const getToolUseContext = useCallback((messages: MessageType[], newMessages: MessageType[], abortController: AbortController, mainLoopModel: string, queryGeneration?: number, queryLifecycle?: QueryLifecycleOperationTracker): ProcessUserInputContext => {
|
||||
// Read mutable values fresh from the store rather than closure-capturing
|
||||
// useAppState() snapshots. Same values today (closure is refreshed by the
|
||||
// render between turns); decouples freshness from React's render cycle for
|
||||
@@ -2530,6 +2587,7 @@ export function REPL({
|
||||
} satisfies ProcessUserInputContext['queryActivity'];
|
||||
return {
|
||||
abortController,
|
||||
...(queryLifecycle ? { queryLifecycle } : {}),
|
||||
options: {
|
||||
commands,
|
||||
tools: computeTools(),
|
||||
@@ -2793,7 +2851,7 @@ export function REPL({
|
||||
void removeTranscriptMessage(tombstonedMessage.uuid);
|
||||
}, setStreamingThinking, undefined, onStreamingText);
|
||||
}, [setMessages, setResponseLength, setStreamMode, setStreamingToolUses, setStreamingThinking, onStreamingText]);
|
||||
const onQueryImpl = useCallback(async (messagesIncludingNewMessages: MessageType[], newMessages: MessageType[], abortController: AbortController, shouldQuery: boolean, additionalAllowedTools: string[], mainLoopModelParam: string, queryGeneration: number, effort?: EffortValue) => {
|
||||
const onQueryImpl = useCallback(async (messagesIncludingNewMessages: MessageType[], newMessages: MessageType[], abortController: AbortController, shouldQuery: boolean, additionalAllowedTools: string[], mainLoopModelParam: string, queryGeneration: number, effort?: EffortValue, queryLifecycle?: QueryLifecycleOperationTracker) => {
|
||||
// Prepare IDE integration for new prompt. Read mcpClients fresh from
|
||||
// store — useManageMCPConnections may have populated it since the
|
||||
// render that captured this closure (same pattern as computeTools).
|
||||
@@ -2880,7 +2938,7 @@ export function REPL({
|
||||
setAbortController(null);
|
||||
return;
|
||||
}
|
||||
const toolUseContext = getToolUseContext(messagesIncludingNewMessages, newMessages, abortController, mainLoopModelParam, queryGeneration);
|
||||
const toolUseContext = getToolUseContext(messagesIncludingNewMessages, newMessages, abortController, mainLoopModelParam, queryGeneration, queryLifecycle);
|
||||
const querySessionId = getSessionId();
|
||||
const queryAutoCompactTracking = getAutoCompactTrackingForSession(querySessionId);
|
||||
// getToolUseContext reads tools/mcpClients fresh from store.getState()
|
||||
@@ -2977,10 +3035,17 @@ export function REPL({
|
||||
}
|
||||
|
||||
// Concurrent guard via state machine. tryStart() atomically checks
|
||||
// and transitions idle→running, returning the generation number.
|
||||
// and transitions idle→running, returning lifecycle context.
|
||||
// Returns null if already running — no separate check-then-set.
|
||||
const thisGeneration = queryGuard.tryStart();
|
||||
if (thisGeneration === null) {
|
||||
const lifecycleTracker = queryLifecycleTrackerRef.current;
|
||||
const querySource = getQuerySourceForREPL();
|
||||
const startResult = queryGuard.tryStart({
|
||||
queryId: randomUUID(),
|
||||
querySource,
|
||||
startedAt: Date.now(),
|
||||
getActiveOperations: () => lifecycleTracker.snapshot()
|
||||
});
|
||||
if (startResult === null) {
|
||||
logEvent('tengu_concurrent_onquery_detected', {});
|
||||
|
||||
// Extract and enqueue user message text, skipping meta messages
|
||||
@@ -2997,6 +3062,12 @@ export function REPL({
|
||||
});
|
||||
return;
|
||||
}
|
||||
lifecycleTracker.clear();
|
||||
const thisGeneration = startResult.generation;
|
||||
const queryContext = startResult.context;
|
||||
logQueryLifecycle('start', queryContext);
|
||||
logQueryLifecycle('guard_start', queryContext);
|
||||
let didThrow = false;
|
||||
try {
|
||||
// isLoading is derived from queryGuard — tryStart() above already
|
||||
// transitioned dispatching→running, so no setter call needed here.
|
||||
@@ -3034,12 +3105,36 @@ export function REPL({
|
||||
return;
|
||||
}
|
||||
}
|
||||
await onQueryImpl(latestMessages, newMessages, abortController, shouldQuery, additionalAllowedTools, mainLoopModelParam, thisGeneration, effort);
|
||||
await onQueryImpl(latestMessages, newMessages, abortController, shouldQuery, additionalAllowedTools, mainLoopModelParam, thisGeneration, effort, lifecycleTracker);
|
||||
} catch (error) {
|
||||
didThrow = true;
|
||||
throw error;
|
||||
} finally {
|
||||
const terminalReason = getQueryTerminalReason(abortController.signal, didThrow);
|
||||
const abortReason = getAbortReasonLabel(abortController.signal.reason);
|
||||
const activeOperations = lifecycleTracker.snapshot();
|
||||
const completedContext = {
|
||||
...queryContext,
|
||||
terminalReason,
|
||||
...(abortReason !== undefined && {
|
||||
abortReason
|
||||
})
|
||||
};
|
||||
const activeOperationSummary = summarizeActiveOperations(activeOperations);
|
||||
const logCompletedLifecycle = (context: QueryLifecycleContext) => {
|
||||
logQueryLifecycle('end', context, activeOperationSummary);
|
||||
if (activeOperations.apiCalls.length > 0) {
|
||||
logForDebugging(`api.call.orphaned_on_query_end queryId=${queryContext.queryId} generation=${thisGeneration} ${activeOperationSummary}`, {
|
||||
level: 'warn'
|
||||
});
|
||||
}
|
||||
};
|
||||
// queryGuard.end() atomically checks generation and transitions
|
||||
// running→idle. Returns false if a newer query owns the guard
|
||||
// (cancel+resubmit race where the stale finally fires as a microtask).
|
||||
if (queryGuard.end(thisGeneration)) {
|
||||
if (queryGuard.end(thisGeneration, terminalReason, abortReason)) {
|
||||
logCompletedLifecycle(completedContext);
|
||||
lifecycleTracker.clear();
|
||||
setLastQueryCompletionTime(Date.now());
|
||||
skipIdleCheckRef.current = false;
|
||||
// Always reset loading state in finally - this ensures cleanup even
|
||||
@@ -3125,6 +3220,15 @@ export function REPL({
|
||||
// controller makes ctrl+c fire onCancel() (aborting nothing) instead of
|
||||
// propagating to the double-press exit flow.
|
||||
setAbortController(null);
|
||||
} else {
|
||||
const guardCompletedContext = queryGuard.lastContext;
|
||||
if ((guardCompletedContext?.terminalReason === 'query-timeout' || guardCompletedContext?.terminalReason === 'hard-max-query-timeout') && guardCompletedContext.queryGeneration === thisGeneration) {
|
||||
logCompletedLifecycle(guardCompletedContext);
|
||||
setLastQueryCompletionTime(Date.now());
|
||||
lifecycleTracker.clear();
|
||||
} else if (!queryGuard.isActive) {
|
||||
lifecycleTracker.clear();
|
||||
}
|
||||
}
|
||||
|
||||
// Auto-restore: if the user interrupted before any meaningful response
|
||||
|
||||
@@ -0,0 +1,409 @@
|
||||
import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
|
||||
import type {
|
||||
BetaMessage,
|
||||
BetaMessageStreamParams,
|
||||
} from '@anthropic-ai/sdk/resources/beta/messages/messages.mjs'
|
||||
import { mkdtempSync, rmSync } from 'fs'
|
||||
import { tmpdir } from 'os'
|
||||
import { join } from 'path'
|
||||
import {
|
||||
acquireSharedMutationLock,
|
||||
releaseSharedMutationLock,
|
||||
} from '../../test/sharedMutationLock.js'
|
||||
import { getEmptyToolPermissionContext } from '../../Tool.js'
|
||||
import type { Message } from '../../types/message.js'
|
||||
import { QueryLifecycleOperationTracker } from '../../utils/queryLifecycle.js'
|
||||
import { asSystemPrompt } from '../../utils/systemPromptType.js'
|
||||
import {
|
||||
executeNonStreamingRequest,
|
||||
type Options,
|
||||
queryModelWithStreaming,
|
||||
} from './claude.js'
|
||||
import { EMPTY_USAGE } from './emptyUsage.js'
|
||||
|
||||
const envKeys = [
|
||||
'ANTHROPIC_AUTH_TOKEN',
|
||||
'ANTHROPIC_API_KEY',
|
||||
'ANTHROPIC_BASE_URL',
|
||||
'ANTHROPIC_MODEL',
|
||||
'CLAUDE_CODE_TEST_FIXTURES_ROOT',
|
||||
'CLAUDE_CODE_PROVIDER_PROFILE_ENV_APPLIED',
|
||||
'CLAUDE_CODE_PROVIDER_PROFILE_ENV_APPLIED_ID',
|
||||
'CLAUDE_CODE_USE_BEDROCK',
|
||||
'CLAUDE_CODE_USE_FOUNDRY',
|
||||
'CLAUDE_CODE_USE_GEMINI',
|
||||
'CLAUDE_CODE_USE_GITHUB',
|
||||
'CLAUDE_CODE_USE_MISTRAL',
|
||||
'CLAUDE_CODE_USE_OPENAI',
|
||||
'CLAUDE_CODE_USE_VERTEX',
|
||||
'GEMINI_API_KEY',
|
||||
'OPENAI_API_KEY',
|
||||
'OPENAI_BASE_URL',
|
||||
'OPENAI_MODEL',
|
||||
'OPENCLAUDE_MAX_RETRIES',
|
||||
'VCR_RECORD',
|
||||
] as const
|
||||
const originalEnv = { ...process.env }
|
||||
const originalFetch = globalThis.fetch
|
||||
const hadSavedMacro = Object.hasOwn(globalThis, 'MACRO')
|
||||
const savedMacro = (globalThis as Record<string, unknown>).MACRO
|
||||
let fixturesRoot: string | undefined
|
||||
|
||||
type FetchOverride = NonNullable<Options['fetchOverride']>
|
||||
type LifecycleSnapshot = ReturnType<QueryLifecycleOperationTracker['snapshot']>
|
||||
|
||||
function makeJsonResponse(body: unknown, status = 200): Response {
|
||||
return new Response(JSON.stringify(body), {
|
||||
status,
|
||||
headers: {
|
||||
'content-type': 'application/json',
|
||||
'request-id': `req-${status}`,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
function makeErrorResponse(status: number, message: string): Response {
|
||||
return makeJsonResponse(
|
||||
{
|
||||
type: 'error',
|
||||
error: {
|
||||
type: 'api_error',
|
||||
message,
|
||||
},
|
||||
},
|
||||
status,
|
||||
)
|
||||
}
|
||||
|
||||
function makeBetaMessage(): BetaMessage {
|
||||
return {
|
||||
id: 'msg-lifecycle-test',
|
||||
type: 'message',
|
||||
role: 'assistant',
|
||||
model: 'claude-lifecycle-test',
|
||||
content: [],
|
||||
container: null,
|
||||
context_management: null,
|
||||
stop_details: null,
|
||||
stop_reason: 'end_turn',
|
||||
stop_sequence: null,
|
||||
usage: {
|
||||
...EMPTY_USAGE,
|
||||
input_tokens: 1,
|
||||
output_tokens: 1,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
function makeOpenAIChatCompletionResponse(): Response {
|
||||
return makeJsonResponse({
|
||||
id: 'chatcmpl-lifecycle-fallback',
|
||||
object: 'chat.completion',
|
||||
created: 1_771_264_800,
|
||||
model: 'gpt-override',
|
||||
choices: [
|
||||
{
|
||||
index: 0,
|
||||
message: {
|
||||
role: 'assistant',
|
||||
content: 'fallback ok',
|
||||
},
|
||||
finish_reason: 'stop',
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 1,
|
||||
completion_tokens: 1,
|
||||
total_tokens: 2,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
function parseRequestBody(init: RequestInit | undefined): Record<string, unknown> {
|
||||
if (typeof init?.body !== 'string') return {}
|
||||
const parsed = JSON.parse(init.body) as unknown
|
||||
return parsed && typeof parsed === 'object'
|
||||
? (parsed as Record<string, unknown>)
|
||||
: {}
|
||||
}
|
||||
|
||||
async function drainGenerator<T>(
|
||||
generator: AsyncGenerator<unknown, T>,
|
||||
): Promise<T> {
|
||||
while (true) {
|
||||
const result = await generator.next()
|
||||
if (result.done) return result.value
|
||||
}
|
||||
}
|
||||
|
||||
function makeParams(context: { model: string }): BetaMessageStreamParams {
|
||||
return {
|
||||
model: context.model,
|
||||
max_tokens: 64,
|
||||
messages: [{ role: 'user', content: 'hello' }],
|
||||
} as BetaMessageStreamParams
|
||||
}
|
||||
|
||||
function makeOptions(
|
||||
queryLifecycle: QueryLifecycleOperationTracker,
|
||||
): Options {
|
||||
return {
|
||||
getToolPermissionContext: async () => getEmptyToolPermissionContext(),
|
||||
model: 'claude-lifecycle-test',
|
||||
isNonInteractiveSession: false,
|
||||
querySource: 'sdk',
|
||||
agents: [],
|
||||
hasAppendSystemPrompt: false,
|
||||
mcpTools: [],
|
||||
queryLifecycle,
|
||||
}
|
||||
}
|
||||
|
||||
function setTestMacro(): void {
|
||||
;(globalThis as Record<string, unknown>).MACRO = {
|
||||
VERSION: '0.0.0-test',
|
||||
DISPLAY_VERSION: '0.0.0-test',
|
||||
BUILD_TIME: 'test',
|
||||
ISSUES_EXPLAINER: 'test',
|
||||
PACKAGE_URL: 'test',
|
||||
NATIVE_PACKAGE_URL: undefined,
|
||||
}
|
||||
}
|
||||
|
||||
function setClientTestEnv(): void {
|
||||
setTestMacro()
|
||||
fixturesRoot = mkdtempSync(join(tmpdir(), 'claude-lifecycle-vcr-'))
|
||||
for (const key of envKeys) {
|
||||
delete process.env[key]
|
||||
}
|
||||
process.env.ANTHROPIC_API_KEY = 'sk-test-lifecycle'
|
||||
process.env.CLAUDE_CODE_TEST_FIXTURES_ROOT = fixturesRoot
|
||||
process.env.VCR_RECORD = '1'
|
||||
}
|
||||
|
||||
beforeEach(async () => {
|
||||
await acquireSharedMutationLock('claude.lifecycle.test.ts')
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
try {
|
||||
for (const key of envKeys) {
|
||||
if (originalEnv[key] === undefined) {
|
||||
delete process.env[key]
|
||||
} else {
|
||||
process.env[key] = originalEnv[key]
|
||||
}
|
||||
}
|
||||
if (hadSavedMacro) {
|
||||
;(globalThis as Record<string, unknown>).MACRO = savedMacro
|
||||
} else {
|
||||
delete (globalThis as Record<string, unknown>).MACRO
|
||||
}
|
||||
globalThis.fetch = originalFetch
|
||||
if (fixturesRoot) {
|
||||
rmSync(fixturesRoot, { force: true, recursive: true })
|
||||
fixturesRoot = undefined
|
||||
}
|
||||
} finally {
|
||||
releaseSharedMutationLock()
|
||||
}
|
||||
})
|
||||
|
||||
describe('Claude API lifecycle tracking', () => {
|
||||
test('ends a failed streaming dispatch before retry backoff is reported', async () => {
|
||||
setClientTestEnv()
|
||||
process.env.OPENCLAUDE_MAX_RETRIES = '1'
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const dispatchSnapshots: ReturnType<
|
||||
QueryLifecycleOperationTracker['snapshot']
|
||||
>[] = []
|
||||
const fetchOverride: FetchOverride = async () => {
|
||||
dispatchSnapshots.push(queryLifecycle.snapshot())
|
||||
return makeErrorResponse(500, 'stream dispatch failed')
|
||||
}
|
||||
|
||||
const generator = queryModelWithStreaming({
|
||||
messages: [
|
||||
{
|
||||
type: 'user',
|
||||
uuid: '00000000-0000-0000-0000-000000000001',
|
||||
timestamp: '2026-06-17T00:00:00.000Z',
|
||||
message: { role: 'user', content: 'hello' },
|
||||
} as Message,
|
||||
],
|
||||
systemPrompt: asSystemPrompt([]),
|
||||
thinkingConfig: { type: 'disabled' },
|
||||
tools: [],
|
||||
signal: new AbortController().signal,
|
||||
options: {
|
||||
...makeOptions(queryLifecycle),
|
||||
fetchOverride,
|
||||
},
|
||||
})
|
||||
|
||||
const first = await generator.next()
|
||||
expect(first.done).toBe(false)
|
||||
expect(first.value).toMatchObject({
|
||||
type: 'system',
|
||||
subtype: 'api_error',
|
||||
})
|
||||
expect(dispatchSnapshots.length).toBeGreaterThanOrEqual(1)
|
||||
expect(dispatchSnapshots.some(snapshot => snapshot.apiCalls.length === 1)).toBe(
|
||||
true,
|
||||
)
|
||||
expect(queryLifecycle.snapshot().apiCalls).toEqual([])
|
||||
|
||||
await generator.return(undefined)
|
||||
})
|
||||
|
||||
test('preserves provider override and query source during 404 non-streaming fallback', async () => {
|
||||
setClientTestEnv()
|
||||
process.env.OPENCLAUDE_MAX_RETRIES = '0'
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const providerBaseURL = 'https://provider.example/v1'
|
||||
const requests: {
|
||||
authorization: string | null
|
||||
snapshot: LifecycleSnapshot
|
||||
stream: unknown
|
||||
url: string
|
||||
}[] = []
|
||||
|
||||
globalThis.fetch = (async (input, init) => {
|
||||
const body = parseRequestBody(init)
|
||||
requests.push({
|
||||
authorization: new Headers(init?.headers).get('authorization'),
|
||||
snapshot: queryLifecycle.snapshot(),
|
||||
stream: body.stream,
|
||||
url: input instanceof Request ? input.url : String(input),
|
||||
})
|
||||
|
||||
if (body.stream === true) {
|
||||
return makeErrorResponse(404, 'streaming unavailable')
|
||||
}
|
||||
|
||||
return makeOpenAIChatCompletionResponse()
|
||||
}) as typeof fetch
|
||||
|
||||
const messages: unknown[] = []
|
||||
const generator = queryModelWithStreaming({
|
||||
messages: [
|
||||
{
|
||||
type: 'user',
|
||||
uuid: '00000000-0000-0000-0000-000000000002',
|
||||
timestamp: '2026-06-17T00:00:00.000Z',
|
||||
message: { role: 'user', content: 'hello' },
|
||||
} as Message,
|
||||
],
|
||||
systemPrompt: asSystemPrompt([]),
|
||||
thinkingConfig: { type: 'disabled' },
|
||||
tools: [],
|
||||
signal: new AbortController().signal,
|
||||
options: {
|
||||
...makeOptions(queryLifecycle),
|
||||
providerOverride: {
|
||||
model: 'gpt-override',
|
||||
baseURL: providerBaseURL,
|
||||
apiKey: 'provider-test-key',
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
for await (const message of generator) {
|
||||
messages.push(message)
|
||||
}
|
||||
|
||||
const streamingRequest = requests.find(request => request.stream === true)
|
||||
const fallbackRequest = requests.find(request => request.stream === false)
|
||||
|
||||
expect(
|
||||
messages.some(
|
||||
message =>
|
||||
typeof message === 'object' &&
|
||||
message !== null &&
|
||||
(message as { type?: unknown }).type === 'assistant',
|
||||
),
|
||||
).toBe(true)
|
||||
expect(streamingRequest?.url.startsWith(providerBaseURL)).toBe(true)
|
||||
expect(fallbackRequest?.url.startsWith(providerBaseURL)).toBe(true)
|
||||
expect(fallbackRequest?.authorization).toBe('Bearer provider-test-key')
|
||||
expect(fallbackRequest?.snapshot.apiCalls).toHaveLength(1)
|
||||
expect(fallbackRequest?.snapshot.apiCalls[0]).toMatchObject({
|
||||
querySource: 'sdk',
|
||||
})
|
||||
expect(queryLifecycle.snapshot().apiCalls).toEqual([])
|
||||
})
|
||||
|
||||
test('tracks each non-streaming fallback request and clears it on success', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const requestSnapshots: ReturnType<
|
||||
QueryLifecycleOperationTracker['snapshot']
|
||||
>[] = []
|
||||
setClientTestEnv()
|
||||
const fetchOverride: FetchOverride = async () => {
|
||||
requestSnapshots.push(queryLifecycle.snapshot())
|
||||
return makeJsonResponse(makeBetaMessage())
|
||||
}
|
||||
|
||||
const result = await drainGenerator(
|
||||
executeNonStreamingRequest(
|
||||
{ model: 'claude-lifecycle-test', source: 'sdk', fetchOverride },
|
||||
{
|
||||
model: 'claude-lifecycle-test',
|
||||
thinkingConfig: { type: 'disabled' },
|
||||
signal: new AbortController().signal,
|
||||
querySource: 'sdk',
|
||||
},
|
||||
makeParams,
|
||||
() => {},
|
||||
() => {},
|
||||
null,
|
||||
queryLifecycle,
|
||||
),
|
||||
)
|
||||
|
||||
expect(result.id).toBe('msg-lifecycle-test')
|
||||
expect(requestSnapshots).toHaveLength(1)
|
||||
expect(requestSnapshots[0]?.apiCalls).toHaveLength(1)
|
||||
expect(requestSnapshots[0]?.apiCalls[0]).toMatchObject({
|
||||
model: 'claude-lifecycle-test',
|
||||
querySource: 'sdk',
|
||||
})
|
||||
expect(queryLifecycle.snapshot().apiCalls).toEqual([])
|
||||
})
|
||||
|
||||
test('clears non-streaming fallback lifecycle entries after request errors', async () => {
|
||||
setClientTestEnv()
|
||||
process.env.OPENCLAUDE_MAX_RETRIES = '0'
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const requestSnapshots: ReturnType<
|
||||
QueryLifecycleOperationTracker['snapshot']
|
||||
>[] = []
|
||||
const fetchOverride: FetchOverride = async () => {
|
||||
requestSnapshots.push(queryLifecycle.snapshot())
|
||||
return makeErrorResponse(400, 'fallback failed')
|
||||
}
|
||||
|
||||
await expect(
|
||||
drainGenerator(
|
||||
executeNonStreamingRequest(
|
||||
{ model: 'claude-lifecycle-test', source: 'sdk', fetchOverride },
|
||||
{
|
||||
model: 'claude-lifecycle-test',
|
||||
thinkingConfig: { type: 'disabled' },
|
||||
signal: new AbortController().signal,
|
||||
querySource: 'sdk',
|
||||
},
|
||||
makeParams,
|
||||
() => {},
|
||||
() => {},
|
||||
null,
|
||||
queryLifecycle,
|
||||
),
|
||||
),
|
||||
).rejects.toThrow('fallback failed')
|
||||
|
||||
expect(requestSnapshots).toHaveLength(1)
|
||||
expect(requestSnapshots[0]?.apiCalls).toHaveLength(1)
|
||||
expect(queryLifecycle.snapshot().apiCalls).toEqual([])
|
||||
})
|
||||
})
|
||||
+71
-16
@@ -170,6 +170,7 @@ import { COMPACT_MAX_OUTPUT_TOKENS, getContextWindowForModel, getMaxThinkingToke
|
||||
import { logForDebugging } from 'src/utils/debug.js'
|
||||
import { logForDiagnosticsNoPII } from 'src/utils/diagLogs.js'
|
||||
import { type EffortValue, modelSupportsEffort } from 'src/utils/effort.js'
|
||||
import type { QueryLifecycleOperationTracker } from 'src/utils/queryLifecycle.js'
|
||||
import {
|
||||
isFastModeAvailable,
|
||||
isFastModeCooldown,
|
||||
@@ -716,6 +717,7 @@ export type Options = {
|
||||
// (query.ts decrements across the agentic loop).
|
||||
taskBudget?: { total: number; remaining?: number }
|
||||
providerOverride?: { model: string; baseURL: string; apiKey: string }
|
||||
queryLifecycle?: QueryLifecycleOperationTracker
|
||||
}
|
||||
|
||||
export async function queryModelWithoutStreaming({
|
||||
@@ -852,6 +854,7 @@ export async function* executeNonStreamingRequest(
|
||||
* from. Emitted in tengu_nonstreaming_fallback_error for funnel correlation.
|
||||
*/
|
||||
originatingRequestId?: string | null,
|
||||
queryLifecycle?: QueryLifecycleOperationTracker,
|
||||
): AsyncGenerator<SystemAPIErrorMessage, BetaMessage> {
|
||||
const fallbackTimeoutMs = getNonstreamingFallbackTimeoutMs()
|
||||
const generator = withRetry(
|
||||
@@ -874,6 +877,12 @@ export async function* executeNonStreamingRequest(
|
||||
retryParams,
|
||||
MAX_NON_STREAMING_TOKENS,
|
||||
)
|
||||
const activeApiCallKey =
|
||||
queryLifecycle?.startApiCall({
|
||||
model: context.model,
|
||||
querySource: retryOptions.querySource,
|
||||
startedAt: start,
|
||||
}) ?? null
|
||||
|
||||
try {
|
||||
// biome-ignore lint/plugin: non-streaming API call
|
||||
@@ -908,6 +917,10 @@ export async function* executeNonStreamingRequest(
|
||||
'unknown') as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
throw err
|
||||
} finally {
|
||||
if (activeApiCallKey) {
|
||||
queryLifecycle?.endApiCall(activeApiCallKey)
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -1537,9 +1550,24 @@ async function* queryModel(
|
||||
let stream: Stream<BetaRawMessageStreamEvent> | undefined = undefined
|
||||
let streamRequestId: string | null | undefined = undefined
|
||||
let clientRequestId: string | undefined = undefined
|
||||
let activeApiCallKey: string | null = null
|
||||
// eslint-disable-next-line eslint-plugin-n/no-unsupported-features/node-builtins -- Response is available in supported Node runtimes and is used by the SDK
|
||||
let streamResponse: Response | undefined = undefined
|
||||
|
||||
function endActiveApiCall(key = activeApiCallKey): void {
|
||||
if (!key) return
|
||||
options.queryLifecycle?.endApiCall(key)
|
||||
if (activeApiCallKey === key) {
|
||||
activeApiCallKey = null
|
||||
}
|
||||
}
|
||||
|
||||
function startActiveApiCall(call: Parameters<QueryLifecycleOperationTracker['startApiCall']>[0]): string | null {
|
||||
endActiveApiCall()
|
||||
activeApiCallKey = options.queryLifecycle?.startApiCall(call) ?? null
|
||||
return activeApiCallKey
|
||||
}
|
||||
|
||||
// Release all stream resources to prevent native memory leaks.
|
||||
// The Response object holds native TLS/socket buffers that live outside the
|
||||
// V8 heap (observed on the Node.js/npm path; see GH #32920), so we must
|
||||
@@ -1849,26 +1877,42 @@ async function* queryModel(
|
||||
getAPIProvider() === 'firstParty' && isFirstPartyAnthropicBaseUrl()
|
||||
? randomUUID()
|
||||
: undefined
|
||||
const attemptApiCallKey = startActiveApiCall({
|
||||
clientRequestId,
|
||||
model: options.model,
|
||||
querySource: options.querySource,
|
||||
startedAt: start,
|
||||
})
|
||||
|
||||
// Use raw stream instead of BetaMessageStream to avoid O(n²) partial JSON parsing
|
||||
// BetaMessageStream calls partialParse() on every input_json_delta, which we don't need
|
||||
// since we handle tool input accumulation ourselves
|
||||
// biome-ignore lint/plugin: main conversation loop handles attribution separately
|
||||
const result = await anthropic.beta.messages
|
||||
.create(
|
||||
{ ...params, stream: true },
|
||||
{
|
||||
signal,
|
||||
...(clientRequestId && {
|
||||
headers: { [CLIENT_REQUEST_ID_HEADER]: clientRequestId },
|
||||
}),
|
||||
},
|
||||
)
|
||||
.withResponse()
|
||||
queryCheckpoint('query_response_headers_received')
|
||||
streamRequestId = result.request_id
|
||||
streamResponse = result.response
|
||||
return result.data
|
||||
try {
|
||||
const result = await anthropic.beta.messages
|
||||
.create(
|
||||
{ ...params, stream: true },
|
||||
{
|
||||
signal,
|
||||
...(clientRequestId && {
|
||||
headers: { [CLIENT_REQUEST_ID_HEADER]: clientRequestId },
|
||||
}),
|
||||
},
|
||||
)
|
||||
.withResponse()
|
||||
queryCheckpoint('query_response_headers_received')
|
||||
streamRequestId = result.request_id
|
||||
if (attemptApiCallKey) {
|
||||
options.queryLifecycle?.updateApiCall(attemptApiCallKey, {
|
||||
requestId: streamRequestId,
|
||||
})
|
||||
}
|
||||
streamResponse = result.response
|
||||
return result.data
|
||||
} catch (err) {
|
||||
endActiveApiCall(attemptApiCallKey)
|
||||
throw err
|
||||
}
|
||||
},
|
||||
{
|
||||
model: options.model,
|
||||
@@ -2585,6 +2629,7 @@ async function* queryModel(
|
||||
? 'watchdog'
|
||||
: 'other') as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
})
|
||||
endActiveApiCall()
|
||||
const result = yield* executeNonStreamingRequest(
|
||||
{ model: options.model, source: options.querySource, providerOverride: options.providerOverride, effortValue: effort },
|
||||
{
|
||||
@@ -2603,6 +2648,7 @@ async function* queryModel(
|
||||
},
|
||||
params => captureAPIRequest(params, options.querySource),
|
||||
streamRequestId,
|
||||
options.queryLifecycle,
|
||||
)
|
||||
|
||||
const m: AssistantMessage = {
|
||||
@@ -2684,14 +2730,21 @@ async function* queryModel(
|
||||
|
||||
try {
|
||||
// Fall back to non-streaming mode
|
||||
endActiveApiCall()
|
||||
const result = yield* executeNonStreamingRequest(
|
||||
{ model: options.model, source: options.querySource, effortValue: effort },
|
||||
{
|
||||
model: options.model,
|
||||
source: options.querySource,
|
||||
providerOverride: options.providerOverride,
|
||||
effortValue: effort,
|
||||
},
|
||||
{
|
||||
model: options.model,
|
||||
fallbackModel: options.fallbackModel,
|
||||
thinkingConfig,
|
||||
...(isFastModeEnabled() && { fastMode: isFastMode }),
|
||||
signal,
|
||||
querySource: options.querySource,
|
||||
},
|
||||
paramsFromContext,
|
||||
(attempt, _startTime, tokens) => {
|
||||
@@ -2700,6 +2753,7 @@ async function* queryModel(
|
||||
},
|
||||
params => captureAPIRequest(params, options.querySource),
|
||||
failedRequestId,
|
||||
options.queryLifecycle,
|
||||
)
|
||||
|
||||
const m: AssistantMessage = {
|
||||
@@ -2841,6 +2895,7 @@ async function* queryModel(
|
||||
return
|
||||
}
|
||||
} finally {
|
||||
endActiveApiCall()
|
||||
stopSessionActivity('api_call')
|
||||
// Must be in the finally block: if the generator is terminated early
|
||||
// via .return() (e.g. consumer breaks out of for-await-of, or query.ts
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
import { describe, expect, test } from 'bun:test'
|
||||
import type { ToolUseBlock } from '@anthropic-ai/sdk/resources/index.mjs'
|
||||
import { z } from 'zod/v4'
|
||||
|
||||
import type { CanUseToolFn } from '../../hooks/useCanUseTool.js'
|
||||
import { createToolFixture } from '../../test/toolFixtures.js'
|
||||
import {
|
||||
getEmptyToolPermissionContext,
|
||||
type Tool,
|
||||
type ToolUseContext,
|
||||
} from '../../Tool.js'
|
||||
import type { AssistantMessage } from '../../types/message.js'
|
||||
import { QueryLifecycleOperationTracker } from '../../utils/queryLifecycle.js'
|
||||
import { StreamingToolExecutor } from './StreamingToolExecutor.js'
|
||||
|
||||
const assistantMessage = {
|
||||
uuid: 'assistant-message-1',
|
||||
requestId: 'request-1',
|
||||
message: {
|
||||
id: 'assistant-api-message-1',
|
||||
content: [],
|
||||
},
|
||||
} as unknown as AssistantMessage
|
||||
|
||||
function makeToolUseContext(
|
||||
tools: readonly Tool[],
|
||||
queryLifecycle: QueryLifecycleOperationTracker,
|
||||
inProgressToolUseIds: { current: Set<string> },
|
||||
hasInterruptibleToolInProgress: { current: boolean },
|
||||
): ToolUseContext {
|
||||
return {
|
||||
abortController: new AbortController(),
|
||||
messages: [],
|
||||
queryLifecycle,
|
||||
options: {
|
||||
tools,
|
||||
commands: [],
|
||||
debug: false,
|
||||
verbose: false,
|
||||
mainLoopModel: 'test-model',
|
||||
mcpClients: [],
|
||||
mcpResources: {},
|
||||
isNonInteractiveSession: false,
|
||||
},
|
||||
getAppState: () => ({
|
||||
toolPermissionContext: getEmptyToolPermissionContext(),
|
||||
sessionHooks: new Map(),
|
||||
}),
|
||||
setAppState: () => {},
|
||||
setInProgressToolUseIDs: updater => {
|
||||
inProgressToolUseIds.current = updater(inProgressToolUseIds.current)
|
||||
},
|
||||
setHasInterruptibleToolInProgress: value => {
|
||||
hasInterruptibleToolInProgress.current = value
|
||||
},
|
||||
setResponseLength: () => {},
|
||||
updateFileHistoryState: () => {},
|
||||
updateAttributionState: () => {},
|
||||
} as unknown as ToolUseContext
|
||||
}
|
||||
|
||||
describe('StreamingToolExecutor lifecycle tracking', () => {
|
||||
test('discard aborts in-flight tools and clears lifecycle tracking immediately', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const inProgressToolUseIds = { current: new Set<string>() }
|
||||
const hasInterruptibleToolInProgress = { current: false }
|
||||
let resolveStarted!: () => void
|
||||
const started = new Promise<void>(resolve => {
|
||||
resolveStarted = resolve
|
||||
})
|
||||
let observedAbortReason: unknown
|
||||
let resolveAborted!: () => void
|
||||
const aborted = new Promise<void>(resolve => {
|
||||
resolveAborted = resolve
|
||||
})
|
||||
const tool = createToolFixture(z.object({}), {
|
||||
name: 'SlowLifecycleTool',
|
||||
interruptBehavior: () => 'cancel',
|
||||
async call(_input, context) {
|
||||
resolveStarted()
|
||||
await new Promise<void>(resolve => {
|
||||
if (context.abortController.signal.aborted) {
|
||||
observedAbortReason = context.abortController.signal.reason
|
||||
resolveAborted()
|
||||
resolve()
|
||||
return
|
||||
}
|
||||
context.abortController.signal.addEventListener(
|
||||
'abort',
|
||||
() => {
|
||||
observedAbortReason = context.abortController.signal.reason
|
||||
resolveAborted()
|
||||
resolve()
|
||||
},
|
||||
{ once: true },
|
||||
)
|
||||
})
|
||||
return { data: 'aborted' }
|
||||
},
|
||||
})
|
||||
const toolUseContext = makeToolUseContext(
|
||||
[tool],
|
||||
queryLifecycle,
|
||||
inProgressToolUseIds,
|
||||
hasInterruptibleToolInProgress,
|
||||
)
|
||||
const executor = new StreamingToolExecutor(
|
||||
[tool],
|
||||
(async () => ({ behavior: 'allow' })) as CanUseToolFn,
|
||||
toolUseContext,
|
||||
)
|
||||
|
||||
executor.addTool(
|
||||
{
|
||||
type: 'tool_use',
|
||||
id: 'tool-use-1',
|
||||
name: tool.name,
|
||||
input: {},
|
||||
} as ToolUseBlock,
|
||||
assistantMessage,
|
||||
)
|
||||
await started
|
||||
|
||||
expect(queryLifecycle.snapshot().toolUses).toMatchObject([
|
||||
{
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: tool.name,
|
||||
},
|
||||
])
|
||||
expect(inProgressToolUseIds.current.has('tool-use-1')).toBe(true)
|
||||
expect(hasInterruptibleToolInProgress.current).toBe(true)
|
||||
|
||||
executor.discard()
|
||||
|
||||
expect(queryLifecycle.snapshot()).toEqual({ apiCalls: [], toolUses: [] })
|
||||
expect(inProgressToolUseIds.current.has('tool-use-1')).toBe(false)
|
||||
expect(hasInterruptibleToolInProgress.current).toBe(false)
|
||||
await aborted
|
||||
expect(observedAbortReason).toBe('streaming_fallback')
|
||||
expect([...executor.getCompletedResults()]).toEqual([])
|
||||
})
|
||||
})
|
||||
@@ -67,7 +67,21 @@ export class StreamingToolExecutor {
|
||||
* Queued tools won't start, and in-progress tools will receive synthetic errors.
|
||||
*/
|
||||
discard(): void {
|
||||
if (this.discarded) return
|
||||
this.discarded = true
|
||||
this.siblingAbortController.abort('streaming_fallback')
|
||||
for (const tool of this.tools) {
|
||||
if (tool.status === 'yielded') continue
|
||||
this.toolUseContext.queryLifecycle?.endToolUse(tool.id)
|
||||
markToolUseAsComplete(this.toolUseContext, tool.id)
|
||||
tool.pendingProgress.length = 0
|
||||
tool.status = 'yielded'
|
||||
}
|
||||
this.updateInterruptibleState()
|
||||
if (this.progressAvailableResolve) {
|
||||
this.progressAvailableResolve()
|
||||
this.progressAvailableResolve = undefined
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -1,13 +1,92 @@
|
||||
import { describe, expect, test } from 'bun:test'
|
||||
import { z } from 'zod/v4'
|
||||
|
||||
import { SkillTool } from '../../tools/SkillTool/SkillTool.js'
|
||||
import { AskUserQuestionTool } from '../../tools/AskUserQuestionTool/AskUserQuestionTool.js'
|
||||
import { BASH_TOOL_NAME } from '../../tools/BashTool/toolName.js'
|
||||
import type { CanUseToolFn } from '../../hooks/useCanUseTool.js'
|
||||
import { createToolFixture } from '../../test/toolFixtures.js'
|
||||
import {
|
||||
getEmptyToolPermissionContext,
|
||||
type Tool,
|
||||
type ToolUseContext,
|
||||
} from '../../Tool.js'
|
||||
import type { AssistantMessage } from '../../types/message.js'
|
||||
import { QueryLifecycleOperationTracker } from '../../utils/queryLifecycle.js'
|
||||
import {
|
||||
getSchemaValidationErrorOverride,
|
||||
getSchemaValidationToolUseResult,
|
||||
type MessageUpdateLazy,
|
||||
normalizeToolInputForValidation,
|
||||
runToolUse,
|
||||
} from './toolExecution.js'
|
||||
|
||||
const lifecycleToolInputSchema = z.object({
|
||||
command: z.string(),
|
||||
timeout: z.number().optional(),
|
||||
})
|
||||
|
||||
const assistantMessage = {
|
||||
uuid: 'assistant-message-1',
|
||||
requestId: 'request-1',
|
||||
message: {
|
||||
id: 'assistant-api-message-1',
|
||||
},
|
||||
} as unknown as AssistantMessage
|
||||
|
||||
function makeToolUseContext(
|
||||
tools: readonly Tool[],
|
||||
queryLifecycle: QueryLifecycleOperationTracker,
|
||||
): ToolUseContext {
|
||||
return {
|
||||
abortController: new AbortController(),
|
||||
messages: [],
|
||||
queryLifecycle,
|
||||
options: {
|
||||
tools,
|
||||
commands: [],
|
||||
debug: false,
|
||||
verbose: false,
|
||||
mainLoopModel: 'test-model',
|
||||
mcpClients: [],
|
||||
mcpResources: {},
|
||||
isNonInteractiveSession: false,
|
||||
},
|
||||
getAppState: () => ({
|
||||
toolPermissionContext: getEmptyToolPermissionContext(),
|
||||
sessionHooks: new Map(),
|
||||
}),
|
||||
setAppState: () => {},
|
||||
setInProgressToolUseIDs: () => {},
|
||||
setResponseLength: () => {},
|
||||
updateFileHistoryState: () => {},
|
||||
updateAttributionState: () => {},
|
||||
} as unknown as ToolUseContext
|
||||
}
|
||||
|
||||
async function collectToolUseUpdates(
|
||||
tool: Tool,
|
||||
input: Record<string, unknown>,
|
||||
canUseTool: CanUseToolFn,
|
||||
toolUseContext: ToolUseContext,
|
||||
) {
|
||||
const updates: MessageUpdateLazy[] = []
|
||||
for await (const update of runToolUse(
|
||||
{
|
||||
type: 'tool_use',
|
||||
id: 'tool-use-1',
|
||||
name: tool.name,
|
||||
input,
|
||||
} as Parameters<typeof runToolUse>[0],
|
||||
assistantMessage,
|
||||
canUseTool,
|
||||
toolUseContext,
|
||||
)) {
|
||||
updates.push(update)
|
||||
}
|
||||
return updates
|
||||
}
|
||||
|
||||
describe('getSchemaValidationErrorOverride', () => {
|
||||
test('returns actionable missing-skill error for SkillTool', () => {
|
||||
expect(getSchemaValidationErrorOverride(SkillTool, {})).toBe(
|
||||
@@ -34,6 +113,205 @@ describe('getSchemaValidationErrorOverride', () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe('runToolUse lifecycle tracking', () => {
|
||||
test('tracks the tool use while async input validation is pending and clears it on validation failure', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const snapshots: ReturnType<QueryLifecycleOperationTracker['snapshot']>[] =
|
||||
[]
|
||||
const tool = createToolFixture(lifecycleToolInputSchema, {
|
||||
name: 'LifecycleTestTool',
|
||||
async validateInput() {
|
||||
snapshots.push(queryLifecycle.snapshot())
|
||||
return {
|
||||
result: false,
|
||||
message: 'blocked by validation',
|
||||
errorCode: 123,
|
||||
}
|
||||
},
|
||||
async call() {
|
||||
throw new Error('call should not run after validation failure')
|
||||
},
|
||||
})
|
||||
const toolUseContext = makeToolUseContext([tool], queryLifecycle)
|
||||
const canUseTool = (async () => ({
|
||||
behavior: 'allow',
|
||||
})) as CanUseToolFn
|
||||
|
||||
const updates = await collectToolUseUpdates(
|
||||
tool,
|
||||
{ command: 'echo hi', timeout: 1234 },
|
||||
canUseTool,
|
||||
toolUseContext,
|
||||
)
|
||||
|
||||
expect(updates).toHaveLength(1)
|
||||
expect(snapshots).toHaveLength(1)
|
||||
expect(snapshots[0]?.apiCalls).toEqual([])
|
||||
expect(snapshots[0]?.toolUses).toHaveLength(1)
|
||||
expect(snapshots[0]?.toolUses[0]).toMatchObject({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: 'LifecycleTestTool',
|
||||
})
|
||||
expect(typeof snapshots[0]?.toolUses[0]?.startedAt).toBe('number')
|
||||
expect(queryLifecycle.snapshot()).toEqual({ apiCalls: [], toolUses: [] })
|
||||
})
|
||||
|
||||
test('tracks Bash timeout metadata while async input validation is pending', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const snapshots: ReturnType<QueryLifecycleOperationTracker['snapshot']>[] =
|
||||
[]
|
||||
const tool = createToolFixture(lifecycleToolInputSchema, {
|
||||
name: BASH_TOOL_NAME,
|
||||
async validateInput() {
|
||||
snapshots.push(queryLifecycle.snapshot())
|
||||
return {
|
||||
result: false,
|
||||
message: 'blocked by validation',
|
||||
errorCode: 123,
|
||||
}
|
||||
},
|
||||
async call() {
|
||||
throw new Error('call should not run after validation failure')
|
||||
},
|
||||
})
|
||||
const toolUseContext = makeToolUseContext([tool], queryLifecycle)
|
||||
const canUseTool = (async () => ({
|
||||
behavior: 'allow',
|
||||
})) as CanUseToolFn
|
||||
|
||||
await collectToolUseUpdates(
|
||||
tool,
|
||||
{ command: 'sleep 1', timeout: 4321 },
|
||||
canUseTool,
|
||||
toolUseContext,
|
||||
)
|
||||
|
||||
expect(snapshots).toHaveLength(1)
|
||||
expect(snapshots[0]?.toolUses).toHaveLength(1)
|
||||
expect(snapshots[0]?.toolUses[0]).toMatchObject({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: BASH_TOOL_NAME,
|
||||
isBash: true,
|
||||
timeoutMs: 4321,
|
||||
})
|
||||
expect(queryLifecycle.snapshot()).toEqual({ apiCalls: [], toolUses: [] })
|
||||
})
|
||||
|
||||
test('tracks the tool use while permission resolution is pending and clears it on denial', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const snapshots: ReturnType<QueryLifecycleOperationTracker['snapshot']>[] =
|
||||
[]
|
||||
const tool = createToolFixture(lifecycleToolInputSchema, {
|
||||
name: 'LifecyclePermissionTool',
|
||||
async validateInput() {
|
||||
return { result: true }
|
||||
},
|
||||
async call() {
|
||||
throw new Error('call should not run after permission denial')
|
||||
},
|
||||
})
|
||||
const toolUseContext = makeToolUseContext([tool], queryLifecycle)
|
||||
const canUseTool = (async () => {
|
||||
snapshots.push(queryLifecycle.snapshot())
|
||||
return {
|
||||
behavior: 'deny',
|
||||
message: 'denied by test',
|
||||
decisionReason: {
|
||||
type: 'other',
|
||||
reason: 'denied by test',
|
||||
},
|
||||
}
|
||||
}) as CanUseToolFn
|
||||
|
||||
const previousSimpleMode = process.env.CLAUDE_CODE_SIMPLE
|
||||
process.env.CLAUDE_CODE_SIMPLE = '1'
|
||||
let updates: Awaited<ReturnType<typeof collectToolUseUpdates>>
|
||||
try {
|
||||
updates = await collectToolUseUpdates(
|
||||
tool,
|
||||
{ command: 'echo hi' },
|
||||
canUseTool,
|
||||
toolUseContext,
|
||||
)
|
||||
} finally {
|
||||
if (previousSimpleMode === undefined) {
|
||||
delete process.env.CLAUDE_CODE_SIMPLE
|
||||
} else {
|
||||
process.env.CLAUDE_CODE_SIMPLE = previousSimpleMode
|
||||
}
|
||||
}
|
||||
|
||||
expect(updates).toHaveLength(1)
|
||||
expect(snapshots).toHaveLength(1)
|
||||
expect(snapshots[0]?.apiCalls).toEqual([])
|
||||
expect(snapshots[0]?.toolUses).toHaveLength(1)
|
||||
expect(snapshots[0]?.toolUses[0]).toMatchObject({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: 'LifecyclePermissionTool',
|
||||
})
|
||||
expect(typeof snapshots[0]?.toolUses[0]?.startedAt).toBe('number')
|
||||
expect(queryLifecycle.snapshot()).toEqual({ apiCalls: [], toolUses: [] })
|
||||
})
|
||||
|
||||
test('refreshes Bash timeout metadata after permission input rewrites', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const permissionSnapshots: ReturnType<
|
||||
QueryLifecycleOperationTracker['snapshot']
|
||||
>[] = []
|
||||
const callSnapshots: ReturnType<QueryLifecycleOperationTracker['snapshot']>[] =
|
||||
[]
|
||||
const tool = createToolFixture(lifecycleToolInputSchema, {
|
||||
name: BASH_TOOL_NAME,
|
||||
async validateInput() {
|
||||
return { result: true }
|
||||
},
|
||||
async call(input) {
|
||||
callSnapshots.push(queryLifecycle.snapshot())
|
||||
expect(input).toMatchObject({
|
||||
command: 'sleep 2',
|
||||
timeout: 2222,
|
||||
})
|
||||
return { data: 'ok' }
|
||||
},
|
||||
})
|
||||
const toolUseContext = makeToolUseContext([tool], queryLifecycle)
|
||||
const canUseTool = (async () => {
|
||||
permissionSnapshots.push(queryLifecycle.snapshot())
|
||||
return {
|
||||
behavior: 'allow',
|
||||
updatedInput: { command: 'sleep 2', timeout: 2222 },
|
||||
decisionReason: {
|
||||
type: 'other',
|
||||
reason: 'allowed by test',
|
||||
},
|
||||
}
|
||||
}) as CanUseToolFn
|
||||
|
||||
await collectToolUseUpdates(
|
||||
tool,
|
||||
{ command: 'sleep 1', timeout: 1111 },
|
||||
canUseTool,
|
||||
toolUseContext,
|
||||
)
|
||||
|
||||
expect(permissionSnapshots).toHaveLength(1)
|
||||
expect(permissionSnapshots[0]?.toolUses[0]).toMatchObject({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: BASH_TOOL_NAME,
|
||||
isBash: true,
|
||||
timeoutMs: 1111,
|
||||
})
|
||||
expect(callSnapshots).toHaveLength(1)
|
||||
expect(callSnapshots[0]?.toolUses[0]).toMatchObject({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: BASH_TOOL_NAME,
|
||||
isBash: true,
|
||||
timeoutMs: 2222,
|
||||
})
|
||||
expect(queryLifecycle.snapshot()).toEqual({ apiCalls: [], toolUses: [] })
|
||||
})
|
||||
})
|
||||
|
||||
describe('normalizeToolInputForValidation', () => {
|
||||
test('treats blank Read.pages as omitted', () => {
|
||||
expect(
|
||||
|
||||
+952
-918
File diff suppressed because it is too large
Load Diff
@@ -12,6 +12,7 @@ import { createUserMessage } from '../../utils/messages.js'
|
||||
import {
|
||||
resetSettingsCache,
|
||||
} from '../../utils/settings/settingsCache.js'
|
||||
import { QueryLifecycleOperationTracker } from '../../utils/queryLifecycle.js'
|
||||
import type { SettingsJson } from '../../utils/settings/types.js'
|
||||
import type { AgentDefinition } from './loadAgentsDir.js'
|
||||
import type { runAgent as runAgentFn } from './runAgent.js'
|
||||
@@ -167,6 +168,71 @@ describe('runAgent provider routing', () => {
|
||||
await expect(generator.next()).rejects.toBe(stop)
|
||||
expect(capturedContext?.options.providerOverride).toBeUndefined()
|
||||
})
|
||||
|
||||
test('preserves query lifecycle tracking for synchronous child contexts', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const parentContext = createToolUseContext('parent-model', queryLifecycle)
|
||||
const stop = new Error('stop after cache-safe params')
|
||||
let capturedContext: ToolUseContext | undefined
|
||||
const runAgent = await importRunAgent()
|
||||
|
||||
const generator = runAgent({
|
||||
agentDefinition: createAgentDefinition(),
|
||||
promptMessages: [createUserMessage({ content: 'inspect this' })],
|
||||
toolUseContext: parentContext,
|
||||
canUseTool: async () => ({ behavior: 'allow' }),
|
||||
isAsync: false,
|
||||
querySource: 'agent:builtin:general-purpose',
|
||||
availableTools: [],
|
||||
onCacheSafeParams: params => {
|
||||
capturedContext = params.toolUseContext
|
||||
throw stop
|
||||
},
|
||||
})
|
||||
|
||||
await expect(generator.next()).rejects.toBe(stop)
|
||||
|
||||
expect(capturedContext?.queryLifecycle).toBe(queryLifecycle)
|
||||
capturedContext?.queryLifecycle?.startToolUse({
|
||||
toolUseId: 'child-tool-use',
|
||||
toolName: 'Read',
|
||||
startedAt: 1,
|
||||
})
|
||||
expect(queryLifecycle.snapshot().toolUses).toEqual([
|
||||
{
|
||||
toolUseId: 'child-tool-use',
|
||||
toolName: 'Read',
|
||||
startedAt: 1,
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
test('keeps query lifecycle tracking out of asynchronous child contexts', async () => {
|
||||
const queryLifecycle = new QueryLifecycleOperationTracker()
|
||||
const parentContext = createToolUseContext('parent-model', queryLifecycle)
|
||||
const stop = new Error('stop after cache-safe params')
|
||||
let capturedContext: ToolUseContext | undefined
|
||||
const runAgent = await importRunAgent()
|
||||
|
||||
const generator = runAgent({
|
||||
agentDefinition: createAgentDefinition(),
|
||||
promptMessages: [createUserMessage({ content: 'inspect this' })],
|
||||
toolUseContext: parentContext,
|
||||
canUseTool: async () => ({ behavior: 'allow' }),
|
||||
isAsync: true,
|
||||
querySource: 'agent:builtin:general-purpose',
|
||||
availableTools: [],
|
||||
onCacheSafeParams: params => {
|
||||
capturedContext = params.toolUseContext
|
||||
throw stop
|
||||
},
|
||||
})
|
||||
|
||||
await expect(generator.next()).rejects.toBe(stop)
|
||||
|
||||
expect(capturedContext?.queryLifecycle).toBeUndefined()
|
||||
expect(queryLifecycle.snapshot()).toEqual({ apiCalls: [], toolUses: [] })
|
||||
})
|
||||
})
|
||||
|
||||
function createAgentDefinition(): AgentDefinition {
|
||||
@@ -185,7 +251,10 @@ async function importRunAgent(): Promise<typeof runAgentFn> {
|
||||
return module.runAgent
|
||||
}
|
||||
|
||||
function createToolUseContext(mainLoopModel: string): ToolUseContext {
|
||||
function createToolUseContext(
|
||||
mainLoopModel: string,
|
||||
queryLifecycle?: QueryLifecycleOperationTracker,
|
||||
): ToolUseContext {
|
||||
const appState = {
|
||||
mainLoopModel,
|
||||
mainLoopModelForSession: mainLoopModel,
|
||||
@@ -220,6 +289,7 @@ function createToolUseContext(mainLoopModel: string): ToolUseContext {
|
||||
},
|
||||
},
|
||||
abortController: new AbortController(),
|
||||
...(queryLifecycle ? { queryLifecycle } : {}),
|
||||
readFileState: createFileStateCacheWithSizeLimit(
|
||||
READ_FILE_STATE_CACHE_SIZE,
|
||||
),
|
||||
|
||||
@@ -743,6 +743,9 @@ export async function* runAgent({
|
||||
readFileState: agentReadFileState,
|
||||
abortController: agentAbortController,
|
||||
getAppState: agentGetAppState,
|
||||
...(!isAsync && toolUseContext.queryLifecycle
|
||||
? { queryLifecycle: toolUseContext.queryLifecycle }
|
||||
: {}),
|
||||
// Sync agents share these callbacks with parent
|
||||
shareSetAppState: !isAsync,
|
||||
shareSetResponseLength: true, // Both sync and async contribute to response metrics
|
||||
|
||||
+327
-42
@@ -1,5 +1,6 @@
|
||||
import { afterEach, describe, test, expect, vi } from 'vitest'
|
||||
import { QueryGuard } from './QueryGuard.js'
|
||||
import { QueryLifecycleOperationTracker } from './queryLifecycle.js'
|
||||
|
||||
describe('QueryGuard', () => {
|
||||
afterEach(() => {
|
||||
@@ -51,16 +52,14 @@ describe('QueryGuard', () => {
|
||||
guard.tryStart()
|
||||
expect(guard.isActive).toBe(true)
|
||||
|
||||
// Just before timeout
|
||||
vi.advanceTimersByTime(5 * 60 * 1000 - 1)
|
||||
expect(guard.isActive).toBe(true)
|
||||
|
||||
// At timeout
|
||||
vi.advanceTimersByTime(1)
|
||||
expect(guard.isActive).toBe(false)
|
||||
})
|
||||
|
||||
test('timeout notifies owner with the timed-out generation and reason', () => {
|
||||
test('timeout notifies owner with lifecycle context and timeout reason', () => {
|
||||
vi.useFakeTimers()
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {})
|
||||
const guard = new QueryGuard()
|
||||
@@ -71,10 +70,112 @@ describe('QueryGuard', () => {
|
||||
vi.advanceTimersByTime(5 * 60 * 1000)
|
||||
|
||||
expect(onTimeout).toHaveBeenCalledTimes(1)
|
||||
expect(onTimeout).toHaveBeenCalledWith(gen, 'idle')
|
||||
expect(onTimeout).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
generation: gen,
|
||||
reason: 'idle',
|
||||
timeoutMs: 5 * 60 * 1000,
|
||||
context: expect.objectContaining({
|
||||
queryGeneration: gen,
|
||||
terminalReason: 'query-timeout',
|
||||
abortReason: 'idle',
|
||||
}),
|
||||
}),
|
||||
)
|
||||
expect(guard.isActive).toBe(false)
|
||||
})
|
||||
|
||||
test('tryStart accepts explicit query identity and returns lifecycle context', () => {
|
||||
const guard = new QueryGuard()
|
||||
const start = guard.tryStart({
|
||||
queryId: 'query-1',
|
||||
querySource: 'repl_main_thread',
|
||||
parentQueryId: 'parent-query',
|
||||
subagentId: 'agent-1',
|
||||
startedAt: 1234,
|
||||
})
|
||||
|
||||
expect(start).toEqual({
|
||||
generation: 1,
|
||||
context: {
|
||||
queryId: 'query-1',
|
||||
queryGeneration: 1,
|
||||
querySource: 'repl_main_thread',
|
||||
parentQueryId: 'parent-query',
|
||||
subagentId: 'agent-1',
|
||||
startedAt: 1234,
|
||||
},
|
||||
})
|
||||
expect(guard.activeContext).toEqual(start!.context)
|
||||
})
|
||||
|
||||
test('timeout callback receives explicit query identity and active operation snapshot', () => {
|
||||
vi.useFakeTimers()
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {})
|
||||
const tracker = new QueryLifecycleOperationTracker()
|
||||
tracker.startApiCall({
|
||||
clientRequestId: 'client-request-1',
|
||||
requestId: 'server-request-1',
|
||||
model: 'model-name',
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 10,
|
||||
})
|
||||
tracker.startToolUse({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: 'Bash',
|
||||
startedAt: 20,
|
||||
isBash: true,
|
||||
timeoutMs: 120_000,
|
||||
})
|
||||
|
||||
const guard = new QueryGuard()
|
||||
const onTimeout = vi.fn()
|
||||
guard.setTimeoutHandler(onTimeout)
|
||||
const start = guard.tryStart({
|
||||
queryId: 'query-1',
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 1,
|
||||
getActiveOperations: () => tracker.snapshot(),
|
||||
})!
|
||||
|
||||
vi.advanceTimersByTime(5 * 60 * 1000)
|
||||
|
||||
expect(onTimeout).toHaveBeenCalledTimes(1)
|
||||
expect(onTimeout).toHaveBeenCalledWith({
|
||||
generation: start.generation,
|
||||
reason: 'idle',
|
||||
timeoutMs: 5 * 60 * 1000,
|
||||
elapsedMs: expect.any(Number),
|
||||
context: {
|
||||
...start.context,
|
||||
terminalReason: 'query-timeout',
|
||||
abortReason: 'idle',
|
||||
},
|
||||
activeOperations: {
|
||||
apiCalls: [
|
||||
{
|
||||
clientRequestId: 'client-request-1',
|
||||
requestId: 'server-request-1',
|
||||
model: 'model-name',
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 10,
|
||||
},
|
||||
],
|
||||
toolUses: [
|
||||
{
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: 'Bash',
|
||||
startedAt: 20,
|
||||
isBash: true,
|
||||
timeoutMs: 120_000,
|
||||
},
|
||||
],
|
||||
},
|
||||
})
|
||||
expect(guard.lastContext?.terminalReason).toBe('query-timeout')
|
||||
expect(guard.lastContext?.abortReason).toBe('idle')
|
||||
})
|
||||
|
||||
test('timeout handler cleanup prevents stale notification', () => {
|
||||
vi.useFakeTimers()
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {})
|
||||
@@ -103,7 +204,10 @@ describe('QueryGuard', () => {
|
||||
|
||||
expect(() => vi.advanceTimersByTime(5 * 60 * 1000)).not.toThrow()
|
||||
expect(guard.isActive).toBe(false)
|
||||
expect(consoleError).toHaveBeenCalledWith('[QueryGuard] Timeout handler failed', handlerError)
|
||||
expect(consoleError).toHaveBeenCalledWith(
|
||||
'[QueryGuard] Timeout handler failed',
|
||||
handlerError,
|
||||
)
|
||||
})
|
||||
|
||||
test('API stream activity extends the idle deadline only while progress continues', () => {
|
||||
@@ -147,11 +251,14 @@ describe('QueryGuard', () => {
|
||||
toolLeaseGraceMs: 10,
|
||||
})
|
||||
const gen = guard.tryStart()!
|
||||
const lease = guard.acquireLease({
|
||||
owner: 'bash',
|
||||
id: 'toolu_1',
|
||||
timeoutMs: 500,
|
||||
}, gen)
|
||||
const lease = guard.acquireLease(
|
||||
{
|
||||
owner: 'bash',
|
||||
id: 'toolu_1',
|
||||
timeoutMs: 500,
|
||||
},
|
||||
gen,
|
||||
)
|
||||
|
||||
vi.advanceTimersByTime(500)
|
||||
|
||||
@@ -171,18 +278,31 @@ describe('QueryGuard', () => {
|
||||
const onTimeout = vi.fn()
|
||||
guard.setTimeoutHandler(onTimeout)
|
||||
const gen = guard.tryStart()!
|
||||
guard.acquireLease({
|
||||
owner: 'bash',
|
||||
id: 'toolu_1',
|
||||
timeoutMs: 500,
|
||||
}, gen)
|
||||
guard.acquireLease(
|
||||
{
|
||||
owner: 'bash',
|
||||
id: 'toolu_1',
|
||||
timeoutMs: 500,
|
||||
},
|
||||
gen,
|
||||
)
|
||||
|
||||
vi.advanceTimersByTime(509)
|
||||
expect(guard.isActive).toBe(true)
|
||||
vi.advanceTimersByTime(1)
|
||||
|
||||
expect(guard.isActive).toBe(false)
|
||||
expect(onTimeout).toHaveBeenCalledWith(gen, 'lease_expired')
|
||||
expect(onTimeout).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
generation: gen,
|
||||
reason: 'lease_expired',
|
||||
timeoutMs: 510,
|
||||
context: expect.objectContaining({
|
||||
terminalReason: 'query-timeout',
|
||||
abortReason: 'lease_expired',
|
||||
}),
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
test('hard maximum aborts even with active leases and activity', () => {
|
||||
@@ -196,11 +316,14 @@ describe('QueryGuard', () => {
|
||||
const onTimeout = vi.fn()
|
||||
guard.setTimeoutHandler(onTimeout)
|
||||
const gen = guard.tryStart()!
|
||||
guard.acquireLease({
|
||||
owner: 'bash',
|
||||
id: 'toolu_1',
|
||||
timeoutMs: 5_000,
|
||||
}, gen)
|
||||
guard.acquireLease(
|
||||
{
|
||||
owner: 'bash',
|
||||
id: 'toolu_1',
|
||||
timeoutMs: 5_000,
|
||||
},
|
||||
gen,
|
||||
)
|
||||
|
||||
for (let elapsed = 0; elapsed < 900; elapsed += 90) {
|
||||
vi.advanceTimersByTime(90)
|
||||
@@ -211,7 +334,17 @@ describe('QueryGuard', () => {
|
||||
vi.advanceTimersByTime(100)
|
||||
|
||||
expect(guard.isActive).toBe(false)
|
||||
expect(onTimeout).toHaveBeenCalledWith(gen, 'hard_max')
|
||||
expect(onTimeout).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
generation: gen,
|
||||
reason: 'hard_max',
|
||||
timeoutMs: 1_000,
|
||||
context: expect.objectContaining({
|
||||
terminalReason: 'hard-max-query-timeout',
|
||||
abortReason: 'hard_max',
|
||||
}),
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
test('lease hard cap is relative to acquisition and capped by query remaining budget', () => {
|
||||
@@ -227,19 +360,28 @@ describe('QueryGuard', () => {
|
||||
const gen = guard.tryStart()!
|
||||
|
||||
vi.advanceTimersByTime(400)
|
||||
guard.acquireLease({
|
||||
owner: 'tool',
|
||||
id: 'toolu_late',
|
||||
timeoutMs: 500,
|
||||
hardCapMs: 300,
|
||||
}, gen)
|
||||
guard.acquireLease(
|
||||
{
|
||||
owner: 'tool',
|
||||
id: 'toolu_late',
|
||||
timeoutMs: 500,
|
||||
hardCapMs: 300,
|
||||
},
|
||||
gen,
|
||||
)
|
||||
|
||||
vi.advanceTimersByTime(299)
|
||||
expect(guard.isActive).toBe(true)
|
||||
vi.advanceTimersByTime(1)
|
||||
|
||||
expect(guard.isActive).toBe(false)
|
||||
expect(onTimeout).toHaveBeenCalledWith(gen, 'lease_expired')
|
||||
expect(onTimeout).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
generation: gen,
|
||||
reason: 'lease_expired',
|
||||
timeoutMs: 300,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
test('stale generations cannot extend or release a newer query', () => {
|
||||
@@ -254,19 +396,25 @@ describe('QueryGuard', () => {
|
||||
guard.setTimeoutHandler(onTimeout)
|
||||
|
||||
const gen1 = guard.tryStart()!
|
||||
const staleLease = guard.acquireLease({
|
||||
owner: 'bash',
|
||||
id: 'toolu_stale',
|
||||
timeoutMs: 500,
|
||||
}, gen1)
|
||||
const staleLease = guard.acquireLease(
|
||||
{
|
||||
owner: 'bash',
|
||||
id: 'toolu_stale',
|
||||
timeoutMs: 500,
|
||||
},
|
||||
gen1,
|
||||
)
|
||||
guard.forceEnd()
|
||||
|
||||
const gen2 = guard.tryStart()!
|
||||
const liveLease = guard.acquireLease({
|
||||
owner: 'bash',
|
||||
id: 'toolu_live',
|
||||
timeoutMs: 500,
|
||||
}, gen2)
|
||||
const liveLease = guard.acquireLease(
|
||||
{
|
||||
owner: 'bash',
|
||||
id: 'toolu_live',
|
||||
timeoutMs: 500,
|
||||
},
|
||||
gen2,
|
||||
)
|
||||
|
||||
staleLease.release()
|
||||
vi.advanceTimersByTime(100)
|
||||
@@ -276,7 +424,46 @@ describe('QueryGuard', () => {
|
||||
guard.registerActivity('stale_api_stream', gen1)
|
||||
vi.advanceTimersByTime(100)
|
||||
expect(guard.isActive).toBe(false)
|
||||
expect(onTimeout).toHaveBeenCalledWith(gen2, 'idle')
|
||||
expect(onTimeout).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
generation: gen2,
|
||||
reason: 'idle',
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
test('end stamps terminal reason on the completed lifecycle context', () => {
|
||||
const guard = new QueryGuard()
|
||||
const start = guard.tryStart({
|
||||
queryId: 'query-1',
|
||||
querySource: 'repl_main_thread',
|
||||
})!
|
||||
|
||||
expect(guard.end(start.generation, 'user-abort')).toBe(true)
|
||||
expect(guard.lastContext).toEqual({
|
||||
...start.context,
|
||||
terminalReason: 'user-abort',
|
||||
})
|
||||
})
|
||||
|
||||
test('query metadata from one start does not leak into the next start', () => {
|
||||
const guard = new QueryGuard()
|
||||
const child = guard.tryStart({
|
||||
queryId: 'child-query',
|
||||
querySource: 'agent:general-purpose',
|
||||
parentQueryId: 'parent-query',
|
||||
subagentId: 'agent-1',
|
||||
})!
|
||||
expect(guard.end(child.generation, 'ok')).toBe(true)
|
||||
|
||||
const parent = guard.tryStart({
|
||||
queryId: 'parent-query',
|
||||
querySource: 'repl_main_thread',
|
||||
})!
|
||||
|
||||
expect(parent.context.parentQueryId).toBeUndefined()
|
||||
expect(parent.context.subagentId).toBeUndefined()
|
||||
expect(parent.context.queryId).toBe('parent-query')
|
||||
})
|
||||
|
||||
test('timeout is cleared when end() is called normally', () => {
|
||||
@@ -285,11 +472,9 @@ describe('QueryGuard', () => {
|
||||
const gen = guard.tryStart()!
|
||||
guard.end(gen)
|
||||
|
||||
// Advance past timeout — should not affect anything
|
||||
vi.advanceTimersByTime(10 * 60 * 1000)
|
||||
expect(guard.isActive).toBe(false)
|
||||
|
||||
// Should be able to start a new query
|
||||
const gen2 = guard.tryStart()
|
||||
expect(gen2).not.toBeNull()
|
||||
expect(guard.isActive).toBe(true)
|
||||
@@ -297,3 +482,103 @@ describe('QueryGuard', () => {
|
||||
guard.forceEnd()
|
||||
})
|
||||
})
|
||||
|
||||
describe('QueryLifecycleOperationTracker', () => {
|
||||
test('tracks and cleans up active API and tool operations', () => {
|
||||
const tracker = new QueryLifecycleOperationTracker()
|
||||
const apiKey = tracker.startApiCall({
|
||||
clientRequestId: 'client-request-1',
|
||||
model: 'model-name',
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 100,
|
||||
})
|
||||
tracker.updateApiCall(apiKey, { requestId: 'server-request-1' })
|
||||
tracker.startToolUse({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: 'Bash',
|
||||
startedAt: 200,
|
||||
isBash: true,
|
||||
timeoutMs: 300_000,
|
||||
})
|
||||
|
||||
expect(tracker.snapshot()).toEqual({
|
||||
apiCalls: [
|
||||
{
|
||||
clientRequestId: 'client-request-1',
|
||||
requestId: 'server-request-1',
|
||||
model: 'model-name',
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 100,
|
||||
},
|
||||
],
|
||||
toolUses: [
|
||||
{
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: 'Bash',
|
||||
startedAt: 200,
|
||||
isBash: true,
|
||||
timeoutMs: 300_000,
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
tracker.endApiCall(apiKey)
|
||||
tracker.endToolUse('tool-use-1')
|
||||
|
||||
expect(tracker.snapshot()).toEqual({ apiCalls: [], toolUses: [] })
|
||||
})
|
||||
|
||||
test('snapshots contain only safe operation metadata', () => {
|
||||
const tracker = new QueryLifecycleOperationTracker()
|
||||
const apiKey = tracker.startApiCall({
|
||||
clientRequestId: 'client-request-1',
|
||||
model: 'model-name',
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 1,
|
||||
prompt: 'tool output',
|
||||
apiKey: 'ANTHROPIC_API_KEY=secret',
|
||||
} as Parameters<QueryLifecycleOperationTracker['startApiCall']>[0])
|
||||
tracker.updateApiCall(apiKey, {
|
||||
requestId: 'server-request-1',
|
||||
cwd: '/home/user/project',
|
||||
} as Partial<Parameters<QueryLifecycleOperationTracker['startApiCall']>[0]>)
|
||||
tracker.updateApiCall(apiKey, {
|
||||
requestId: undefined,
|
||||
model: undefined,
|
||||
querySource: undefined,
|
||||
startedAt: undefined,
|
||||
cwd: '/home/user/project',
|
||||
} as Partial<Parameters<QueryLifecycleOperationTracker['startApiCall']>[0]>)
|
||||
tracker.startToolUse({
|
||||
toolUseId: 'tool-use-1',
|
||||
toolName: 'Bash',
|
||||
startedAt: 1,
|
||||
isBash: true,
|
||||
timeoutMs: 300_000,
|
||||
command: 'cat /home/user/project/.env',
|
||||
output: 'tool output',
|
||||
} as Parameters<QueryLifecycleOperationTracker['startToolUse']>[0])
|
||||
|
||||
const snapshot = tracker.snapshot()
|
||||
|
||||
expect(snapshot.apiCalls).toEqual([
|
||||
{
|
||||
clientRequestId: 'client-request-1',
|
||||
requestId: 'server-request-1',
|
||||
model: 'model-name',
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 1,
|
||||
},
|
||||
])
|
||||
expect(Object.keys(snapshot.toolUses[0]!)).toEqual([
|
||||
'toolUseId',
|
||||
'toolName',
|
||||
'startedAt',
|
||||
'isBash',
|
||||
'timeoutMs',
|
||||
])
|
||||
expect(JSON.stringify(snapshot)).not.toContain('/home/')
|
||||
expect(JSON.stringify(snapshot)).not.toContain('ANTHROPIC_API_KEY')
|
||||
expect(JSON.stringify(snapshot)).not.toContain('tool output')
|
||||
})
|
||||
})
|
||||
|
||||
+157
-35
@@ -3,16 +3,16 @@
|
||||
* React's `useSyncExternalStore`.
|
||||
*
|
||||
* Three states:
|
||||
* idle → no query, safe to dequeue and process
|
||||
* dispatching → an item was dequeued, async chain hasn't reached onQuery yet
|
||||
* running → onQuery called tryStart(), query is executing
|
||||
* idle -> no query, safe to dequeue and process
|
||||
* dispatching -> an item was dequeued, async chain hasn't reached onQuery yet
|
||||
* running -> onQuery called tryStart(), query is executing
|
||||
*
|
||||
* Transitions:
|
||||
* idle → dispatching (reserve)
|
||||
* dispatching → running (tryStart)
|
||||
* idle → running (tryStart, for direct user submissions)
|
||||
* running → idle (end / forceEnd / timeout)
|
||||
* dispatching → idle (cancelReservation, when processQueueIfReady fails)
|
||||
* idle -> dispatching (reserve)
|
||||
* dispatching -> running (tryStart)
|
||||
* idle -> running (tryStart, for direct user submissions)
|
||||
* running -> idle (end / forceEnd / timeout)
|
||||
* dispatching -> idle (cancelReservation, when processQueueIfReady fails)
|
||||
*
|
||||
* `isActive` returns true for both dispatching and running, preventing
|
||||
* re-entry from the queue processor during the async gap.
|
||||
@@ -29,19 +29,22 @@
|
||||
* )
|
||||
*/
|
||||
import { createSignal } from './signal.js'
|
||||
import type {
|
||||
QueryActiveOperationSnapshot,
|
||||
QueryGuardMetadata,
|
||||
QueryGuardStart,
|
||||
QueryGuardTimeoutInfo,
|
||||
QueryGuardTimeoutReason,
|
||||
QueryLifecycleContext,
|
||||
QueryTerminalReason,
|
||||
} from './queryLifecycle.js'
|
||||
|
||||
export type { QueryGuardTimeoutReason } from './queryLifecycle.js'
|
||||
|
||||
export const DEFAULT_QUERY_IDLE_TIMEOUT_MS = 5 * 60 * 1000 // 5 minutes
|
||||
export const DEFAULT_QUERY_HARD_MAX_MS = 30 * 60 * 1000 // 30 minutes
|
||||
export const DEFAULT_TOOL_LEASE_GRACE_MS = 5_000
|
||||
|
||||
/**
|
||||
* Why QueryGuard force-ended the current query.
|
||||
* - `idle`: no progress and no valid lease existed for the idle timeout.
|
||||
* - `hard_max`: the query reached its absolute maximum lifetime.
|
||||
* - `lease_expired`: bounded active work exceeded its own lease deadline.
|
||||
*/
|
||||
export type QueryGuardTimeoutReason = 'idle' | 'hard_max' | 'lease_expired'
|
||||
|
||||
/**
|
||||
* Input for a bounded unit of active work.
|
||||
*
|
||||
@@ -73,10 +76,7 @@ export type QueryGuardLease = {
|
||||
release(): void
|
||||
}
|
||||
|
||||
type QueryTimeoutHandler = (
|
||||
generation: number,
|
||||
reason: QueryGuardTimeoutReason,
|
||||
) => void
|
||||
type QueryTimeoutHandler = (timeout: QueryGuardTimeoutInfo) => void
|
||||
|
||||
type QueryGuardOptions = {
|
||||
idleTimeoutMs?: number
|
||||
@@ -94,12 +94,25 @@ type LeaseRecord = {
|
||||
description?: string
|
||||
}
|
||||
|
||||
const EMPTY_ACTIVE_OPERATIONS: QueryActiveOperationSnapshot = {
|
||||
apiCalls: [],
|
||||
toolUses: [],
|
||||
}
|
||||
|
||||
function positiveOrDefault(value: number | undefined, fallback: number): number {
|
||||
return typeof value === 'number' && Number.isFinite(value) && value > 0
|
||||
? value
|
||||
: fallback
|
||||
}
|
||||
|
||||
function terminalReasonForTimeout(
|
||||
reason: QueryGuardTimeoutReason,
|
||||
): QueryTerminalReason {
|
||||
return reason === 'hard_max'
|
||||
? 'hard-max-query-timeout'
|
||||
: 'query-timeout'
|
||||
}
|
||||
|
||||
export class QueryGuard {
|
||||
private _status: 'idle' | 'dispatching' | 'running' = 'idle'
|
||||
private _generation = 0
|
||||
@@ -110,6 +123,10 @@ export class QueryGuard {
|
||||
private _lastActivityAt = 0
|
||||
private _leaseCounter = 0
|
||||
private _activeLeases = new Map<string, LeaseRecord>()
|
||||
private _context: QueryLifecycleContext | null = null
|
||||
private _lastContext: QueryLifecycleContext | null = null
|
||||
private _getActiveOperations: (() => QueryActiveOperationSnapshot) | null =
|
||||
null
|
||||
private readonly _idleTimeoutMs: number
|
||||
private readonly _hardMaxQueryMs: number
|
||||
private readonly _toolLeaseGraceMs: number
|
||||
@@ -130,7 +147,7 @@ export class QueryGuard {
|
||||
}
|
||||
|
||||
/**
|
||||
* Reserve the guard for queue processing. Transitions idle → dispatching.
|
||||
* Reserve the guard for queue processing. Transitions idle -> dispatching.
|
||||
* Returns false if not idle (another query or dispatch in progress).
|
||||
*/
|
||||
reserve(): boolean {
|
||||
@@ -142,7 +159,7 @@ export class QueryGuard {
|
||||
|
||||
/**
|
||||
* Cancel a reservation when processQueueIfReady had nothing to process.
|
||||
* Transitions dispatching → idle.
|
||||
* Transitions dispatching -> idle.
|
||||
*/
|
||||
cancelReservation(): void {
|
||||
if (this._status !== 'dispatching') return
|
||||
@@ -156,15 +173,26 @@ export class QueryGuard {
|
||||
* Accepts transitions from both idle (direct user submit)
|
||||
* and dispatching (queue processor path).
|
||||
*/
|
||||
tryStart(): number | null {
|
||||
tryStart(): number | null
|
||||
tryStart(metadata: QueryGuardMetadata): QueryGuardStart | null
|
||||
tryStart(metadata?: QueryGuardMetadata): number | QueryGuardStart | null {
|
||||
if (this._status === 'running') return null
|
||||
this._status = 'running'
|
||||
++this._generation
|
||||
this._activeLeases.clear()
|
||||
this._lastContext = null
|
||||
this._queryStartedAt = Date.now()
|
||||
this._lastActivityAt = this._queryStartedAt
|
||||
this._context = this._createContext(metadata)
|
||||
this._getActiveOperations = metadata?.getActiveOperations ?? null
|
||||
this._startTimeout()
|
||||
this._notify()
|
||||
if (metadata) {
|
||||
return {
|
||||
generation: this._generation,
|
||||
context: this._context,
|
||||
}
|
||||
}
|
||||
return this._generation
|
||||
}
|
||||
|
||||
@@ -173,12 +201,18 @@ export class QueryGuard {
|
||||
* (meaning the caller should perform cleanup). Returns false if a
|
||||
* newer query has started (stale finally block from a cancelled query).
|
||||
*/
|
||||
end(generation: number): boolean {
|
||||
end(
|
||||
generation: number,
|
||||
terminalReason: QueryTerminalReason = 'ok',
|
||||
abortReason?: string,
|
||||
): boolean {
|
||||
if (this._generation !== generation) return false
|
||||
if (this._status !== 'running') return false
|
||||
this._clearTimeout()
|
||||
this._activeLeases.clear()
|
||||
this._completeContext(terminalReason, abortReason)
|
||||
this._status = 'idle'
|
||||
this._getActiveOperations = null
|
||||
this._notify()
|
||||
return true
|
||||
}
|
||||
@@ -189,11 +223,16 @@ export class QueryGuard {
|
||||
* Increments generation so stale finally blocks from the cancelled
|
||||
* query's promise rejection will see a mismatch and skip cleanup.
|
||||
*/
|
||||
forceEnd(): void {
|
||||
forceEnd(
|
||||
terminalReason: QueryTerminalReason = 'unknown',
|
||||
abortReason?: string,
|
||||
): void {
|
||||
if (this._status === 'idle') return
|
||||
this._clearTimeout()
|
||||
this._activeLeases.clear()
|
||||
this._completeContext(terminalReason, abortReason)
|
||||
this._status = 'idle'
|
||||
this._getActiveOperations = null
|
||||
++this._generation
|
||||
this._notify()
|
||||
}
|
||||
@@ -283,7 +322,7 @@ export class QueryGuard {
|
||||
|
||||
/**
|
||||
* Is the guard active (dispatching or running)?
|
||||
* Always synchronous — not subject to React state batching delays.
|
||||
* Always synchronous - not subject to React state batching delays.
|
||||
*/
|
||||
get isActive(): boolean {
|
||||
return this._status !== 'idle'
|
||||
@@ -293,6 +332,14 @@ export class QueryGuard {
|
||||
return this._generation
|
||||
}
|
||||
|
||||
get activeContext(): QueryLifecycleContext | null {
|
||||
return this._context ? { ...this._context } : null
|
||||
}
|
||||
|
||||
get lastContext(): QueryLifecycleContext | null {
|
||||
return this._lastContext ? { ...this._lastContext } : null
|
||||
}
|
||||
|
||||
/**
|
||||
* Register a single owner callback for watchdog timeouts. The callback runs
|
||||
* before forceEnd(), so callers can abort in-flight work while the timed-out
|
||||
@@ -310,7 +357,7 @@ export class QueryGuard {
|
||||
// --
|
||||
// useSyncExternalStore interface
|
||||
|
||||
/** Subscribe to state changes. Stable reference — safe as useEffect dep. */
|
||||
/** Subscribe to state changes. Stable reference - safe as useEffect dep. */
|
||||
subscribe = this._changed.subscribe
|
||||
|
||||
/** Snapshot for useSyncExternalStore. Returns `isActive`. */
|
||||
@@ -322,6 +369,55 @@ export class QueryGuard {
|
||||
this._changed.emit()
|
||||
}
|
||||
|
||||
private _createContext(
|
||||
metadata: QueryGuardMetadata | undefined,
|
||||
): QueryLifecycleContext {
|
||||
return {
|
||||
queryId: metadata?.queryId ?? `generation-${this._generation}`,
|
||||
queryGeneration: this._generation,
|
||||
querySource: metadata?.querySource ?? 'unknown',
|
||||
...(metadata?.parentQueryId && { parentQueryId: metadata.parentQueryId }),
|
||||
...(metadata?.subagentId && { subagentId: metadata.subagentId }),
|
||||
startedAt: metadata?.startedAt ?? Date.now(),
|
||||
}
|
||||
}
|
||||
|
||||
private _completeContext(
|
||||
terminalReason: QueryTerminalReason,
|
||||
abortReason?: string,
|
||||
): void {
|
||||
if (!this._context) return
|
||||
const completed = {
|
||||
...this._context,
|
||||
terminalReason,
|
||||
...(abortReason !== undefined && { abortReason }),
|
||||
}
|
||||
this._context = null
|
||||
this._lastContext = completed
|
||||
}
|
||||
|
||||
private _activeContextWithTerminalReason(
|
||||
terminalReason: QueryTerminalReason,
|
||||
abortReason?: string,
|
||||
): QueryLifecycleContext {
|
||||
const context = this._context ?? this._createContext(undefined)
|
||||
return {
|
||||
...context,
|
||||
terminalReason,
|
||||
...(abortReason !== undefined && { abortReason }),
|
||||
}
|
||||
}
|
||||
|
||||
private _snapshotActiveOperations(): QueryActiveOperationSnapshot {
|
||||
if (!this._getActiveOperations) return EMPTY_ACTIVE_OPERATIONS
|
||||
try {
|
||||
return this._getActiveOperations()
|
||||
} catch (error) {
|
||||
console.error('[QueryGuard] Active operation snapshot failed', error)
|
||||
return EMPTY_ACTIVE_OPERATIONS
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Start a watchdog timer. Stuck work aborts after the idle timeout, active
|
||||
* bounded work can continue while its lease is valid, and the hard maximum
|
||||
@@ -354,21 +450,36 @@ export class QueryGuard {
|
||||
this._timeoutId = null
|
||||
if (this._status !== 'running') return
|
||||
|
||||
const reason = this._getTimeoutReason(Date.now())
|
||||
const now = Date.now()
|
||||
const reason = this._getTimeoutReason(now)
|
||||
if (!reason) {
|
||||
this._scheduleTimeout()
|
||||
return
|
||||
}
|
||||
|
||||
const terminalReason = terminalReasonForTimeout(reason)
|
||||
const context = this._activeContextWithTerminalReason(
|
||||
terminalReason,
|
||||
reason,
|
||||
)
|
||||
const timeout: QueryGuardTimeoutInfo = {
|
||||
generation: this._generation,
|
||||
reason,
|
||||
timeoutMs: this._getTimeoutMsForReason(reason, now),
|
||||
elapsedMs: now - context.startedAt,
|
||||
context,
|
||||
activeOperations: this._snapshotActiveOperations(),
|
||||
}
|
||||
|
||||
console.error(
|
||||
`[QueryGuard] Query ${reason} timeout — force-ending to prevent infinite spinner`,
|
||||
`[QueryGuard] Query ${reason} timeout - force-ending to prevent infinite spinner`,
|
||||
)
|
||||
try {
|
||||
this._timeoutHandler?.(this._generation, reason)
|
||||
this._timeoutHandler?.(timeout)
|
||||
} catch (error) {
|
||||
console.error('[QueryGuard] Timeout handler failed', error)
|
||||
} finally {
|
||||
this.forceEnd()
|
||||
this.forceEnd(terminalReason, reason)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -385,16 +496,27 @@ export class QueryGuard {
|
||||
hasValidLease = true
|
||||
}
|
||||
|
||||
if (
|
||||
!hasValidLease &&
|
||||
now >= this._lastActivityAt + this._idleTimeoutMs
|
||||
) {
|
||||
if (!hasValidLease && now >= this._lastActivityAt + this._idleTimeoutMs) {
|
||||
return 'idle'
|
||||
}
|
||||
|
||||
return null
|
||||
}
|
||||
|
||||
private _getTimeoutMsForReason(
|
||||
reason: QueryGuardTimeoutReason,
|
||||
now: number,
|
||||
): number {
|
||||
if (reason === 'hard_max') return this._hardMaxQueryMs
|
||||
if (reason === 'idle') return this._idleTimeoutMs
|
||||
|
||||
const expiredLease = [...this._activeLeases.values()].find(
|
||||
lease => lease.deadlineAt <= now,
|
||||
)
|
||||
if (!expiredLease) return this._idleTimeoutMs
|
||||
return Math.max(0, expiredLease.deadlineAt - expiredLease.startedAt)
|
||||
}
|
||||
|
||||
private _getNextDeadlineAt(now: number): number | null {
|
||||
if (this._status !== 'running') return null
|
||||
|
||||
|
||||
@@ -277,6 +277,8 @@ export type SubagentContextOverrides = {
|
||||
abortController?: AbortController
|
||||
/** Override the getAppState function */
|
||||
getAppState?: ToolUseContext['getAppState']
|
||||
/** Explicitly opt in to sharing a lifecycle tracker with this subagent. */
|
||||
queryLifecycle?: ToolUseContext['queryLifecycle']
|
||||
|
||||
/**
|
||||
* Explicit opt-in to share parent's setAppState callback.
|
||||
@@ -452,6 +454,9 @@ export function createSubagentContext(
|
||||
// Generate new agentId for subagents (each subagent should have its own ID)
|
||||
agentId: overrides?.agentId ?? createAgentId(),
|
||||
agentType: overrides?.agentType,
|
||||
...(overrides?.queryLifecycle
|
||||
? { queryLifecycle: overrides.queryLifecycle }
|
||||
: {}),
|
||||
|
||||
// Create new query tracking chain for subagent with incremented depth
|
||||
queryTracking: {
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
import { describe, expect, test } from 'vitest'
|
||||
import {
|
||||
formatQueryLifecycleAbortSignalReason,
|
||||
formatQueryLifecycleLogMessage,
|
||||
type QueryLifecycleContext,
|
||||
} from './queryLifecycle.js'
|
||||
|
||||
describe('query lifecycle log formatting', () => {
|
||||
test('keeps timeout context abort reason distinct from abort signal reason', () => {
|
||||
const context: QueryLifecycleContext = {
|
||||
queryId: 'query-1',
|
||||
queryGeneration: 1,
|
||||
querySource: 'repl_main_thread',
|
||||
startedAt: 1,
|
||||
terminalReason: 'query-timeout',
|
||||
abortReason: 'idle',
|
||||
}
|
||||
|
||||
const line = formatQueryLifecycleLogMessage(
|
||||
'abort_requested',
|
||||
context,
|
||||
formatQueryLifecycleAbortSignalReason('query-timeout'),
|
||||
)
|
||||
|
||||
expect(line).toContain('abortReason=idle')
|
||||
expect(line).toContain('abortSignalReason=query-timeout')
|
||||
expect(line).not.toContain('abortReason=query-timeout')
|
||||
expect(line.match(/\babortReason=/g)).toHaveLength(1)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,157 @@
|
||||
export type QueryTerminalReason =
|
||||
| 'ok'
|
||||
| 'query-timeout'
|
||||
| 'hard-max-query-timeout'
|
||||
| 'user-abort'
|
||||
| 'api-error'
|
||||
| 'tool-error'
|
||||
| 'budget-exhausted'
|
||||
| 'parent-ended'
|
||||
| 'unknown'
|
||||
|
||||
export type QueryGuardTimeoutReason = 'idle' | 'hard_max' | 'lease_expired'
|
||||
|
||||
export type QueryActiveApiCall = {
|
||||
clientRequestId?: string
|
||||
requestId?: string | null
|
||||
model?: string
|
||||
querySource?: string
|
||||
startedAt: number
|
||||
}
|
||||
|
||||
export type QueryActiveToolUse = {
|
||||
toolUseId: string
|
||||
toolName: string
|
||||
startedAt: number
|
||||
isBash?: boolean
|
||||
timeoutMs?: number
|
||||
}
|
||||
|
||||
export type QueryActiveOperationSnapshot = {
|
||||
apiCalls: QueryActiveApiCall[]
|
||||
toolUses: QueryActiveToolUse[]
|
||||
}
|
||||
|
||||
export type QueryLifecycleContext = {
|
||||
queryId: string
|
||||
queryGeneration: number
|
||||
querySource: string
|
||||
parentQueryId?: string
|
||||
subagentId?: string
|
||||
startedAt: number
|
||||
terminalReason?: QueryTerminalReason
|
||||
abortReason?: string
|
||||
}
|
||||
|
||||
export type QueryGuardMetadata = {
|
||||
queryId: string
|
||||
querySource: string
|
||||
parentQueryId?: string
|
||||
subagentId?: string
|
||||
startedAt?: number
|
||||
getActiveOperations?: () => QueryActiveOperationSnapshot
|
||||
}
|
||||
|
||||
export type QueryGuardStart = {
|
||||
generation: number
|
||||
context: QueryLifecycleContext
|
||||
}
|
||||
|
||||
export type QueryGuardTimeoutInfo = {
|
||||
generation: number
|
||||
reason: QueryGuardTimeoutReason
|
||||
timeoutMs: number
|
||||
elapsedMs: number
|
||||
context: QueryLifecycleContext
|
||||
activeOperations: QueryActiveOperationSnapshot
|
||||
}
|
||||
|
||||
export function formatQueryLifecycleAbortSignalReason(reason: string): string {
|
||||
return `abortSignalReason=${reason}`
|
||||
}
|
||||
|
||||
export function formatQueryLifecycleLogMessage(
|
||||
event: string,
|
||||
context: QueryLifecycleContext,
|
||||
extras = '',
|
||||
): string {
|
||||
const parent = context.parentQueryId ? ` parentQueryId=${context.parentQueryId}` : ''
|
||||
const subagent = context.subagentId ? ` subagentId=${context.subagentId}` : ''
|
||||
const terminal = context.terminalReason ? ` terminalReason=${context.terminalReason}` : ''
|
||||
const abort = context.abortReason ? ` abortReason=${context.abortReason}` : ''
|
||||
return `query.${event} queryId=${context.queryId} generation=${context.queryGeneration} source=${context.querySource}${parent}${subagent}${terminal}${abort}${extras ? ` ${extras}` : ''}`
|
||||
}
|
||||
|
||||
// Rebuild snapshots from an allowlist so debug logging cannot leak runtime extras.
|
||||
function toSafeApiCallSnapshot(call: QueryActiveApiCall): QueryActiveApiCall {
|
||||
return {
|
||||
...(call.clientRequestId !== undefined ? { clientRequestId: call.clientRequestId } : {}),
|
||||
...(call.requestId !== undefined ? { requestId: call.requestId } : {}),
|
||||
...(call.model !== undefined ? { model: call.model } : {}),
|
||||
...(call.querySource !== undefined ? { querySource: call.querySource } : {}),
|
||||
startedAt: call.startedAt,
|
||||
}
|
||||
}
|
||||
|
||||
function toDefinedApiCallUpdate(update: Partial<QueryActiveApiCall>): Partial<QueryActiveApiCall> {
|
||||
return {
|
||||
...(update.clientRequestId !== undefined ? { clientRequestId: update.clientRequestId } : {}),
|
||||
...(update.requestId !== undefined ? { requestId: update.requestId } : {}),
|
||||
...(update.model !== undefined ? { model: update.model } : {}),
|
||||
...(update.querySource !== undefined ? { querySource: update.querySource } : {}),
|
||||
...(update.startedAt !== undefined ? { startedAt: update.startedAt } : {}),
|
||||
}
|
||||
}
|
||||
|
||||
// Keep tool snapshots on the same explicit allowlist boundary as API calls.
|
||||
function toSafeToolUseSnapshot(toolUse: QueryActiveToolUse): QueryActiveToolUse {
|
||||
return {
|
||||
toolUseId: toolUse.toolUseId,
|
||||
toolName: toolUse.toolName,
|
||||
startedAt: toolUse.startedAt,
|
||||
...(toolUse.isBash !== undefined ? { isBash: toolUse.isBash } : {}),
|
||||
...(toolUse.timeoutMs !== undefined ? { timeoutMs: toolUse.timeoutMs } : {}),
|
||||
}
|
||||
}
|
||||
|
||||
export class QueryLifecycleOperationTracker {
|
||||
private apiCalls = new Map<string, QueryActiveApiCall>()
|
||||
private toolUses = new Map<string, QueryActiveToolUse>()
|
||||
private apiCallSeq = 0
|
||||
|
||||
startApiCall(call: QueryActiveApiCall): string {
|
||||
const key = call.clientRequestId ?? call.requestId ?? `api-call-${++this.apiCallSeq}`
|
||||
this.apiCalls.set(key, toSafeApiCallSnapshot(call))
|
||||
return key
|
||||
}
|
||||
|
||||
updateApiCall(key: string, update: Partial<QueryActiveApiCall>): void {
|
||||
const current = this.apiCalls.get(key)
|
||||
if (!current) return
|
||||
this.apiCalls.set(key, toSafeApiCallSnapshot({ ...current, ...toDefinedApiCallUpdate(update) }))
|
||||
}
|
||||
|
||||
endApiCall(key: string): void {
|
||||
this.apiCalls.delete(key)
|
||||
}
|
||||
|
||||
startToolUse(toolUse: QueryActiveToolUse): void {
|
||||
this.toolUses.set(toolUse.toolUseId, toSafeToolUseSnapshot(toolUse))
|
||||
}
|
||||
|
||||
endToolUse(toolUseId: string): void {
|
||||
this.toolUses.delete(toolUseId)
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.apiCalls.clear()
|
||||
this.toolUses.clear()
|
||||
}
|
||||
|
||||
snapshot(): QueryActiveOperationSnapshot {
|
||||
return {
|
||||
apiCalls: [...this.apiCalls.values()].map(toSafeApiCallSnapshot),
|
||||
toolUses: [...this.toolUses.values()].map(toSafeToolUseSnapshot),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user