Enforce timeout for subagents (#12232)

This commit is contained in:
Silvio Junior
2025-11-03 15:33:04 -05:00
committed by GitHub
parent c4377c1b1a
commit 1c18552426
2 changed files with 116 additions and 13 deletions

View File

@@ -908,7 +908,62 @@ describe('AgentExecutor', () => {
expect(mockSendMessageStream).toHaveBeenCalledTimes(MAX);
});
it('should terminate if timeout is reached', async () => {
it('should terminate with TIMEOUT if a model call takes too long', async () => {
const definition = createTestDefinition([LS_TOOL_NAME], {
max_time_minutes: 0.5, // 30 seconds
});
const executor = await AgentExecutor.create(
definition,
mockConfig,
onActivity,
);
// Mock a model call that is interruptible by an abort signal.
mockSendMessageStream.mockImplementationOnce(async (_model, params) => {
const signal = params?.config?.abortSignal;
// eslint-disable-next-line require-yield
return (async function* () {
await new Promise<void>((resolve) => {
// This promise resolves when aborted, ending the generator.
signal?.addEventListener('abort', () => {
resolve();
});
});
})();
});
const runPromise = executor.run({ goal: 'Timeout test' }, signal);
// Advance time past the timeout to trigger the abort.
await vi.advanceTimersByTimeAsync(31 * 1000);
const output = await runPromise;
expect(output.terminate_reason).toBe(AgentTerminateMode.TIMEOUT);
expect(output.result).toContain('Agent timed out after 0.5 minutes.');
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
// Verify activity stream reported the timeout
expect(activities).toContainEqual(
expect.objectContaining({
type: 'ERROR',
data: expect.objectContaining({
context: 'timeout',
error: 'Agent timed out after 0.5 minutes.',
}),
}),
);
// Verify telemetry
expect(mockedLogAgentFinish).toHaveBeenCalledWith(
mockConfig,
expect.objectContaining({
terminate_reason: AgentTerminateMode.TIMEOUT,
}),
);
});
it('should terminate with TIMEOUT if a tool call takes too long', async () => {
const definition = createTestDefinition([LS_TOOL_NAME], {
max_time_minutes: 1,
});

View File

@@ -159,6 +159,16 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
let terminateReason: AgentTerminateMode = AgentTerminateMode.ERROR;
let finalResult: string | null = null;
const { max_time_minutes } = this.definition.runConfig;
const timeoutController = new AbortController();
const timeoutId = setTimeout(
() => timeoutController.abort(new Error('Agent timed out.')),
max_time_minutes * 60 * 1000,
);
// Combine the external signal with the internal timeout signal.
const combinedSignal = AbortSignal.any([signal, timeoutController.signal]);
logAgentStart(
this.runtimeContext,
new AgentStartEvent(this.agentId, this.definition.name),
@@ -180,8 +190,11 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
terminateReason = reason;
break;
}
if (signal.aborted) {
terminateReason = AgentTerminateMode.ABORTED;
if (combinedSignal.aborted) {
// Determine which signal caused the abort.
terminateReason = timeoutController.signal.aborted
? AgentTerminateMode.TIMEOUT
: AgentTerminateMode.ABORTED;
break;
}
@@ -190,11 +203,19 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
const { functionCalls } = await promptIdContext.run(
promptId,
async () =>
this.callModel(chat, currentMessage, tools, signal, promptId),
this.callModel(
chat,
currentMessage,
tools,
combinedSignal,
promptId,
),
);
if (signal.aborted) {
terminateReason = AgentTerminateMode.ABORTED;
if (combinedSignal.aborted) {
terminateReason = timeoutController.signal.aborted
? AgentTerminateMode.TIMEOUT
: AgentTerminateMode.ABORTED;
break;
}
@@ -210,7 +231,11 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
}
const { nextMessage, submittedOutput, taskCompleted } =
await this.processFunctionCalls(functionCalls, signal, promptId);
await this.processFunctionCalls(
functionCalls,
combinedSignal,
promptId,
);
if (taskCompleted) {
finalResult = submittedOutput ?? 'Task completed successfully.';
@@ -221,6 +246,14 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
currentMessage = nextMessage;
}
if (terminateReason === AgentTerminateMode.TIMEOUT) {
finalResult = `Agent timed out after ${this.definition.runConfig.max_time_minutes} minutes.`;
this.emitActivity('ERROR', {
error: finalResult,
context: 'timeout',
});
}
if (terminateReason === AgentTerminateMode.GOAL) {
return {
result: finalResult || 'Task completed.',
@@ -234,9 +267,29 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
terminate_reason: terminateReason,
};
} catch (error) {
// Check if the error is an AbortError caused by our internal timeout.
if (
error instanceof Error &&
error.name === 'AbortError' &&
timeoutController.signal.aborted &&
!signal.aborted // Ensure the external signal was not the cause
) {
terminateReason = AgentTerminateMode.TIMEOUT;
finalResult = `Agent timed out after ${this.definition.runConfig.max_time_minutes} minutes.`;
this.emitActivity('ERROR', {
error: finalResult,
context: 'timeout',
});
return {
result: finalResult,
terminate_reason: terminateReason,
};
}
this.emitActivity('ERROR', { error: String(error) });
throw error; // Re-throw the error for the parent context to handle.
throw error; // Re-throw other errors or external aborts.
} finally {
clearTimeout(timeoutId);
logAgentFinish(
this.runtimeContext,
new AgentFinishEvent(
@@ -745,11 +798,6 @@ Important Rules:
return AgentTerminateMode.MAX_TURNS;
}
const elapsedMinutes = (Date.now() - startTime) / (1000 * 60);
if (elapsedMinutes >= runConfig.max_time_minutes) {
return AgentTerminateMode.TIMEOUT;
}
return null;
}