diff --git a/packages/cli/src/modules/agents/__tests__/agent-workflow-execution.service.test.ts b/packages/cli/src/modules/agents/__tests__/agent-workflow-execution.service.test.ts index d0cb79d3943..e469ea114bb 100644 --- a/packages/cli/src/modules/agents/__tests__/agent-workflow-execution.service.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agent-workflow-execution.service.test.ts @@ -110,11 +110,12 @@ function makeRuntime(chunks: StreamChunk[] = [{ type: 'finish', finishReason: 's .fn() .mockResolvedValue({ runId: 'runtime-run-1', stream: makeReadableStream(chunks) }), structuredOutput: vi.fn(), - close: vi.fn(), + close: vi.fn().mockResolvedValue(undefined), } as unknown as RuntimeAgent & { stream: Mock; resume: Mock; structuredOutput: Mock; + close: Mock; }, toolRegistry: mock(), mcpServerAttributions: new Map(), @@ -159,8 +160,9 @@ function makeService() { credentialsService.findAllCredentialIdsForProject.mockResolvedValue([]); credentialsService.findAllGlobalCredentialIds.mockResolvedValue([]); + const logger = mockLogger(); const service = new AgentWorkflowExecutionService( - mockLogger(), + logger, agentRepository, new AgentTurnExecutionService( mockLogger(), @@ -194,6 +196,7 @@ function makeService() { executionLevelTracer, nodeToolAiGatewayService, integrationMessageContextService, + logger, }; } @@ -1616,4 +1619,102 @@ describe('AgentWorkflowExecutionService', () => { expect(passedInstrumentation).toBeUndefined(); }); }); + + describe('isolated agent close', () => { + const inlinePayload = { + config: { + name: 'Inline Agent', + model: 'anthropic/claude-sonnet-4-5', + credential: 'cred-1', + instructions: 'Help users', + }, + }; + + function arrange(source: 'stored' | 'inline') { + const context = makeService(); + const runtime = makeRuntime([{ type: 'finish', finishReason: 'stop' }]); + if (source === 'stored') { + context.agentRepository.findByIdAndProjectId.mockResolvedValue(makeAgent()); + context.reconstructionService.reconstructFromAgentEntity.mockResolvedValue(runtime); + } else { + context.reconstructionService.reconstructFromResolvedSource.mockResolvedValue(runtime); + } + return { ...context, runtime }; + } + + async function execute(source: 'stored' | 'inline', service: AgentWorkflowExecutionService) { + return source === 'stored' + ? await service.executeForWorkflow(agentId, 'hello', 'execution-1', 'thread-1', projectId) + : await service.executeInlineForWorkflow( + inlinePayload, + 'hello', + 'execution-1', + 'thread-1', + projectId, + ); + } + + it.each(['stored', 'inline'] as const)( + 'closes the %s agent once after a successful run', + async (source) => { + const { service, runtime } = arrange(source); + + const result = await execute(source, service); + + expect(result.finishReason).toBe('stop'); + expect(runtime.agent.close).toHaveBeenCalledTimes(1); + }, + ); + + it.each(['stored', 'inline'] as const)( + 'closes the %s agent when the stream fails and keeps the original error', + async (source) => { + const { service, runtime } = arrange(source); + runtime.agent.stream.mockRejectedValue(new Error('stream failed')); + + await expect(execute(source, service)).rejects.toThrow('stream failed'); + expect(runtime.agent.close).toHaveBeenCalledTimes(1); + }, + ); + + it.each(['stored', 'inline'] as const)( + 'keeps the %s run result when close fails', + async (source) => { + const { service, runtime, logger } = arrange(source); + runtime.agent.close.mockRejectedValue(new Error('close failed')); + + const result = await execute(source, service); + + expect(result.finishReason).toBe('stop'); + expect(logger.warn).toHaveBeenCalledWith('Failed to close workflow agent after execution', { + error: 'close failed', + }); + }, + ); + + it.each(['stored', 'inline'] as const)( + 'returns the %s run result when close does not settle', + async (source) => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + try { + const { service, runtime, logger } = arrange(source); + runtime.agent.close.mockReturnValue(new Promise(() => {})); + + const pending = execute(source, service); + await vi.advanceTimersByTimeAsync(10_000); + const result = await pending; + + expect(result.finishReason).toBe('stop'); + expect(logger.warn).toHaveBeenCalledWith( + 'Failed to close workflow agent after execution', + { + error: 'Agent close timed out after 10000ms', + }, + ); + } finally { + vi.useRealTimers(); + } + }, + ); + }); }); diff --git a/packages/cli/src/modules/agents/agent-workflow-execution.service.ts b/packages/cli/src/modules/agents/agent-workflow-execution.service.ts index d85a3018808..d6103983985 100644 --- a/packages/cli/src/modules/agents/agent-workflow-execution.service.ts +++ b/packages/cli/src/modules/agents/agent-workflow-execution.service.ts @@ -165,6 +165,8 @@ function createWorkflowAgentExecutionError(message: string, cause?: Error): Oper }); } +const WORKFLOW_AGENT_CLOSE_TIMEOUT_MS = 10_000; + /** * Executes agents invoked from inside a workflow execution (the AI Agent node * or a "Message an Agent" tool call): non-streaming runs against isolated @@ -657,28 +659,60 @@ export class AgentWorkflowExecutionService { }); } + /** + * Release the isolated agent after the run. `close()` disconnects its MCP + * clients. Stop waiting after 10 seconds so a hung close cannot hold the + * run result or the original error. A close failure is logged and does not + * replace that result. + */ + private async closeAgentQuietly(agent: BuiltAgent): Promise { + const close = agent.close(); + void close.catch(() => {}); + + let timeoutId: ReturnType | undefined; + const timeout = new Promise((_, reject) => { + timeoutId = setTimeout(() => { + reject(new Error(`Agent close timed out after ${WORKFLOW_AGENT_CLOSE_TIMEOUT_MS}ms`)); + }, WORKFLOW_AGENT_CLOSE_TIMEOUT_MS); + }); + + try { + await Promise.race([close, timeout]); + } catch (error) { + this.logger.warn('Failed to close workflow agent after execution', { + error: error instanceof Error ? error.message : String(error), + }); + } finally { + clearTimeout(timeoutId); + } + } + private async executeForWorkflowInternal( params: StoredWorkflowExecutionContext, ): Promise { const { agentId, projectId, threadId, outputSchema, sandboxScope } = params; const { agentInstance, recordingParams, runType, budget } = await this.prepareStoredWorkflowRun(params); - const run = await this.streamCompiledWorkflowAgent(agentInstance, params, { - telemetryAgentId: agentId, - savedAgentId: agentId, - runType, - budget, - recordingParams: { ...recordingParams, agentName: agentInstance.name }, - ...(sandboxScope - ? { sandboxScope: { projectId, principalHash: sandboxScope.principalHash } } - : {}), - }); - return this.buildWorkflowResult({ - run, - // The caller remaps the scoped thread key to its session id. - session: { agentId, projectId, sessionId: threadId, threadId }, - outputSchema, - }); + try { + const run = await this.streamCompiledWorkflowAgent(agentInstance, params, { + telemetryAgentId: agentId, + savedAgentId: agentId, + runType, + budget, + recordingParams: { ...recordingParams, agentName: agentInstance.name }, + ...(sandboxScope + ? { sandboxScope: { projectId, principalHash: sandboxScope.principalHash } } + : {}), + }); + return this.buildWorkflowResult({ + run, + // The caller remaps the scoped thread key to its session id. + session: { agentId, projectId, sessionId: threadId, threadId }, + outputSchema, + }); + } finally { + await this.closeAgentQuietly(agentInstance); + } } private async prepareStoredWorkflowRun(params: StoredWorkflowExecutionContext) { @@ -880,13 +914,18 @@ export class AgentWorkflowExecutionService { workflowContext, streamObserver, }; - const run = await this.streamCompiledWorkflowAgent(compiled.agent, params, { - telemetryAgentId: syntheticAgentId, - runType, - budget: compiled.budget, - }); - this.trackInlineRun(params, syntheticAgentId, runType, runtimeConfig, run); - return this.buildWorkflowResult({ run, session: null, outputSchema }); + const agentInstance = compiled.agent; + try { + const run = await this.streamCompiledWorkflowAgent(agentInstance, params, { + telemetryAgentId: syntheticAgentId, + runType, + budget: compiled.budget, + }); + this.trackInlineRun(params, syntheticAgentId, runType, runtimeConfig, run); + return this.buildWorkflowResult({ run, session: null, outputSchema }); + } finally { + await this.closeAgentQuietly(agentInstance); + } } private async prepareInlineRuntime(