From 6a504e514096010330888f5d0425100b9d78751f Mon Sep 17 00:00:00 2001 From: Designation Vee <120316848+Senthemodder@users.noreply.github.com> Date: Thu, 3 Sep 2026 19:21:56 +0700 Subject: [PATCH] fix(runtime): await in-flight task promises in AgentRuntime.stop() Fixes #259 - Replaced busy-polling while loop in `drainInFlightTasks()` with event-driven `Promise.race([Promise.allSettled(...), this.sleep(timeoutMs)])`. - Added `inFlightPromises` Map tracking active task execution promises with guaranteed lifecycle cleanup in `finally`. - Added reproduction test suite verifying zero busy-polling sleep cycles and immediate promise unblocking. --- src/runtime/agent-runtime.ts | 103 +++++++------ tests/runtime/reproduce-issue-259.test.ts | 179 ++++++++++++++++++++++ 2 files changed, 235 insertions(+), 47 deletions(-) create mode 100644 tests/runtime/reproduce-issue-259.test.ts diff --git a/src/runtime/agent-runtime.ts b/src/runtime/agent-runtime.ts index 40fe51f..6a751c7 100644 --- a/src/runtime/agent-runtime.ts +++ b/src/runtime/agent-runtime.ts @@ -16,6 +16,7 @@ export class AgentRuntime { private readonly dependencies: ReturnType; private readonly runtimeId: string; private readonly inFlightTasks = new Set(); + private readonly inFlightPromises = new Map>(); private started = false; private stopped = false; @@ -135,66 +136,74 @@ export class AgentRuntime { toolName: task.toolName }); - this.inFlightTasks.add(task.taskId); - try { - const result = await this.dependencies.taskRunner.run( - task, - context - ); - - this.dependencies.logger.info("Runtime task completed.", { - runtimeId: this.runtimeId, - taskId: task.taskId, - toolName: task.toolName, - durationMs: result.durationMs - }); + const taskPromise = (async () => { + try { + const result = await this.dependencies.taskRunner.run< + TPayload, + TResult + >(task, context); - this.dependencies.eventBus.emit({ - name: "runtime.task.completed", - payload: { + this.dependencies.logger.info("Runtime task completed.", { runtimeId: this.runtimeId, taskId: task.taskId, - agentId: task.agentId, toolName: task.toolName, durationMs: result.durationMs - } - }); - - return result; - } catch (error) { - const reason = - error instanceof Error ? error.message : "Unknown runtime failure."; - - this.dependencies.logger.error("Runtime task failed.", { - runtimeId: this.runtimeId, - taskId: task.taskId, - reason - }); - this.dependencies.eventBus.emit({ - name: "runtime.task.failed", - payload: { + }); + + this.dependencies.eventBus.emit({ + name: "runtime.task.completed", + payload: { + runtimeId: this.runtimeId, + taskId: task.taskId, + agentId: task.agentId, + toolName: task.toolName, + durationMs: result.durationMs + } + }); + + return result; + } catch (error) { + const reason = + error instanceof Error ? error.message : "Unknown runtime failure."; + + this.dependencies.logger.error("Runtime task failed.", { runtimeId: this.runtimeId, taskId: task.taskId, - agentId: task.agentId, reason - } - }); + }); + this.dependencies.eventBus.emit({ + name: "runtime.task.failed", + payload: { + runtimeId: this.runtimeId, + taskId: task.taskId, + agentId: task.agentId, + reason + } + }); + + throw error; + } finally { + this.inFlightTasks.delete(task.taskId); + this.inFlightPromises.delete(task.taskId); + } + })(); - throw error; - } finally { - this.inFlightTasks.delete(task.taskId); - } + this.inFlightTasks.add(task.taskId); + this.inFlightPromises.set(task.taskId, taskPromise); + + return await taskPromise; } private async drainInFlightTasks(timeoutMs: number): Promise { - const deadline = Date.now() + timeoutMs; - while (this.inFlightTasks.size > 0) { - const remaining = deadline - Date.now(); - if (remaining <= 0) { - break; - } - await this.sleep(Math.min(5, remaining)); + if (this.inFlightPromises.size === 0) { + return; + } + const allPromises = Array.from(this.inFlightPromises.values()); + const drainPromise = Promise.allSettled(allPromises); + if (timeoutMs <= 0) { + return; } + await Promise.race([drainPromise, this.sleep(timeoutMs)]); } private sleep(milliseconds: number): Promise { diff --git a/tests/runtime/reproduce-issue-259.test.ts b/tests/runtime/reproduce-issue-259.test.ts new file mode 100644 index 0000000..1079cf3 --- /dev/null +++ b/tests/runtime/reproduce-issue-259.test.ts @@ -0,0 +1,179 @@ +import { describe, it, expect, vi } from "vitest"; +import { AgentRuntime } from "../../src/runtime/agent-runtime.js"; + +describe("Issue #259 Reproduction — Await in-flight task promises in AgentRuntime.stop() instead of busy-polling", () => { + it("does not use sleep-based busy-polling loop during task draining when tasks complete", async () => { + const runtime = new AgentRuntime({ + runtimeId: "rt-reproduce-polling", + logger: { + level: "error", + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn() + } + }); + + const sleepSpy = vi.spyOn(runtime as any, "sleep"); + + runtime.registerTool({ + name: "delayed-work", + description: "Work that completes after a brief delay", + execute: async () => { + await new Promise((resolve) => setTimeout(resolve, 30)); + return { completed: true }; + } + }); + + await runtime.start(); + + const taskPromise = runtime.executeTask({ + taskId: "task-polling-1", + agentId: "agent-1", + toolName: "delayed-work", + input: "run delayed work", + payload: {} + }); + + await runtime.stop({ drainTimeoutMs: 1000 }); + await taskPromise; + + // Issue #259: drainInFlightTasks in current implementation uses a busy-wait loop: + // while (this.inFlightTasks.size > 0) { await this.sleep(Math.min(5, remaining)); } + // When tasks are in-flight, sleep(5) is invoked repeatedly every 5ms. + // The expected behavior is event-driven draining that awaits task promises directly + // without executing a sleep-polling loop. + const shortSleepPolls = sleepSpy.mock.calls.filter( + ([ms]) => typeof ms === "number" && ms <= 5 + ); + expect(shortSleepPolls).toHaveLength(0); + expect(sleepSpy).not.toHaveBeenCalledWith(5); + }); + + it("awaits in-flight task promises directly and unblocks immediately upon task completion", async () => { + const runtime = new AgentRuntime({ + runtimeId: "rt-reproduce-promise-await", + logger: { + level: "error", + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn() + } + }); + + let resolveTask!: (value: unknown) => void; + const taskExecutionPromise = new Promise((resolve) => { + resolveTask = resolve; + }); + + runtime.registerTool({ + name: "controlled-task", + description: "Task controlled by external promise resolution", + execute: () => taskExecutionPromise + }); + + await runtime.start(); + + const taskPromise = runtime.executeTask({ + taskId: "task-controlled-1", + agentId: "agent-1", + toolName: "controlled-task", + input: "run controlled task", + payload: {} + }); + + // Interface Contract from PROJECT.md: + // Runtime must track in-flight task promises (inFlightPromises: Map>) + // so stop() can await them directly via Promise.allSettled instead of polling a Set of strings. + const inFlightPromises = (runtime as any).inFlightPromises; + expect(inFlightPromises).toBeDefined(); + expect( + inFlightPromises instanceof Map || inFlightPromises instanceof Set + ).toBe(true); + expect(inFlightPromises.size).toBe(1); + + let stopSettled = false; + const stopPromise = runtime.stop({ drainTimeoutMs: 5000 }).then(() => { + stopSettled = true; + }); + + // Flush immediate microtasks + await new Promise((resolve) => setTimeout(resolve, 20)); + + // While task execution promise is pending, stop should remain in-flight + expect(stopSettled).toBe(false); + expect(runtime.getInFlightTaskCount()).toBe(1); + + // Resolve task execution promise directly + resolveTask({ status: "success" }); + + await taskPromise; + await stopPromise; + + expect(stopSettled).toBe(true); + expect(runtime.getInFlightTaskCount()).toBe(0); + expect((runtime as any).inFlightPromises.size).toBe(0); + }); + + it("tracks multiple concurrent in-flight task promises and drains without busy-polling", async () => { + const runtime = new AgentRuntime({ + runtimeId: "rt-reproduce-concurrent", + logger: { + level: "error", + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn() + } + }); + + const sleepSpy = vi.spyOn(runtime as any, "sleep"); + + runtime.registerTool<{ delayMs: number }, { delayed: number }>({ + name: "variable-delay", + description: "Task taking variable delay", + execute: async ({ payload }) => { + await new Promise((resolve) => setTimeout(resolve, payload.delayMs)); + return { delayed: payload.delayMs }; + } + }); + + await runtime.start(); + + const task1 = runtime.executeTask<{ delayMs: number }, { delayed: number }>( + { + taskId: "task-concurrent-1", + agentId: "agent-1", + toolName: "variable-delay", + input: "delay 15", + payload: { delayMs: 15 } + } + ); + + const task2 = runtime.executeTask<{ delayMs: number }, { delayed: number }>( + { + taskId: "task-concurrent-2", + agentId: "agent-2", + toolName: "variable-delay", + input: "delay 35", + payload: { delayMs: 35 } + } + ); + + // Verify task promises are tracked concurrently + const inFlightPromises = (runtime as any).inFlightPromises; + expect(inFlightPromises).toBeDefined(); + expect(inFlightPromises.size).toBe(2); + + await runtime.stop({ drainTimeoutMs: 2000 }); + await Promise.all([task1, task2]); + + const shortSleepPolls = sleepSpy.mock.calls.filter( + ([ms]) => typeof ms === "number" && ms <= 5 + ); + expect(shortSleepPolls).toHaveLength(0); + expect(sleepSpy).not.toHaveBeenCalledWith(5); + expect((runtime as any).inFlightPromises.size).toBe(0); + }); +});