fix(core): Close agent after run when invoked from workflow (#41128)

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<ToolRegistry>(),
 		mcpServerAttributions: new Map<string, string>(),
@@ -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();
+				}
+			},
+		);
+	});
 });

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<void> {
+		const close = agent.close();
+		void close.catch(() => {});
+
+		let timeoutId: ReturnType<typeof setTimeout> | undefined;
+		const timeout = new Promise<never>((_, 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<ExecuteAgentData> {
 		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(