From 1cdec9cdcb7195eefc2183ea89df9a7d1f163ee4 Mon Sep 17 00:00:00 2001
From: Alex Smolya <alexsmolya@gmail.com>
Date: Thu, 3 Sep 2026 10:31:52 +0200
Subject: [PATCH] fix(runtime): warn on stranded tasks after drain timeout

---
 src/runtime/agent-runtime.ts             |  18 ++++
 tests/runtime/agent-runtime-stop.test.ts | 111 +++++++++++++++++++++--
 2 files changed, 122 insertions(+), 7 deletions(-)

diff --git a/src/runtime/agent-runtime.ts b/src/runtime/agent-runtime.ts
index 40fe51f..89ae912 100644
--- a/src/runtime/agent-runtime.ts
+++ b/src/runtime/agent-runtime.ts
@@ -82,7 +82,25 @@ export class AgentRuntime {
     this.started = false;
 
     if (options.drainTimeoutMs !== undefined && options.drainTimeoutMs > 0) {
+      const drainStartedAt = Date.now();
       await this.drainInFlightTasks(options.drainTimeoutMs);
+
+      const strandedTaskIds = [...this.inFlightTasks];
+      if (strandedTaskIds.length > 0) {
+        const drainElapsedMs = Date.now() - drainStartedAt;
+        this.dependencies.logger.warn(
+          `Runtime stop timed out after ${drainElapsedMs}ms with tasks still in flight: ${strandedTaskIds.join(
+            ", "
+          )}.`,
+          {
+            runtimeId: this.runtimeId,
+            taskIds: strandedTaskIds,
+            inFlightTaskCount: strandedTaskIds.length,
+            drainTimeoutMs: options.drainTimeoutMs,
+            drainElapsedMs
+          }
+        );
+      }
     }
 
     if (options.clearListeners === true) {
diff --git a/tests/runtime/agent-runtime-stop.test.ts b/tests/runtime/agent-runtime-stop.test.ts
index 80cbe9b..85457d6 100644
--- a/tests/runtime/agent-runtime-stop.test.ts
+++ b/tests/runtime/agent-runtime-stop.test.ts
@@ -1,21 +1,53 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
+import { InMemoryRuntimeLogger } from "../../src/logger/runtime-logger.js";
 import { AgentRuntime } from "../../src/runtime/agent-runtime.js";
 import type { RuntimeOptions } from "../../src/runtime/types.js";
 
 describe("AgentRuntime.stop", () => {
   let runtime: AgentRuntime;
   let emitSpy: ReturnType<typeof vi.fn>;
+  let logger: InMemoryRuntimeLogger;
+
+  function deferred() {
+    let resolve!: () => void;
+    const promise = new Promise<void>((res) => {
+      resolve = res;
+    });
+    return { promise, resolve };
+  }
+
+  function taskInput(taskId: string) {
+    return {
+      taskId,
+      agentId: "agent-1",
+      toolName: "blocked",
+      input: "test",
+      payload: {}
+    };
+  }
+
+  function registerBlockedTask() {
+    const started = deferred();
+    const release = deferred();
+
+    runtime.registerTool({
+      name: "blocked",
+      description: "Waits until released",
+      execute: async () => {
+        started.resolve();
+        await release.promise;
+        return { done: true };
+      }
+    });
+
+    return { started, release };
+  }
 
   beforeEach(() => {
+    logger = new InMemoryRuntimeLogger();
     const options: RuntimeOptions = {
       runtimeId: "test-runtime-stop",
-      logger: {
-        level: "error",
-        info: vi.fn(),
-        warn: vi.fn(),
-        error: vi.fn(),
-        debug: vi.fn()
-      }
+      logger
     };
     runtime = new AgentRuntime(options);
     emitSpy = vi.fn();
@@ -54,4 +86,69 @@ describe("AgentRuntime.stop", () => {
     );
     expect(stoppedEvents.length).toBe(1);
   });
+
+  it("warns once when the drain timeout expires with a task in flight", async () => {
+    const { started, release } = registerBlockedTask();
+
+    await runtime.start();
+    const taskPromise = runtime.executeTask(taskInput("task-timeout"));
+    await started.promise;
+
+    await runtime.stop({ drainTimeoutMs: 15 });
+    await runtime.stop({ drainTimeoutMs: 15 });
+
+    const warnings = logger.entries.filter((entry) => entry.level === "warn");
+    expect(warnings).toHaveLength(1);
+    expect(warnings[0]?.message).toContain("task-timeout");
+    expect(warnings[0]?.metadata).toMatchObject({
+      runtimeId: "test-runtime-stop",
+      taskIds: ["task-timeout"],
+      inFlightTaskCount: 1,
+      drainTimeoutMs: 15,
+      drainElapsedMs: expect.any(Number)
+    });
+
+    release.resolve();
+    await taskPromise;
+  });
+
+  it("does not warn when all tasks drain before the timeout", async () => {
+    const { started, release } = registerBlockedTask();
+
+    await runtime.start();
+    const taskPromise = runtime.executeTask(taskInput("task-drained"));
+    await started.promise;
+    release.resolve();
+
+    await runtime.stop({ drainTimeoutMs: 50 });
+    await taskPromise;
+
+    expect(
+      logger.entries.filter((entry) => entry.level === "warn")
+    ).toHaveLength(0);
+    expect(runtime.getInFlightTaskCount()).toBe(0);
+  });
+
+  it.each([undefined, 0])(
+    "does not warn when drainTimeoutMs is %s",
+    async (drainTimeoutMs) => {
+      const { started, release } = registerBlockedTask();
+
+      await runtime.start();
+      const taskPromise = runtime.executeTask(taskInput("task-no-timeout"));
+      await started.promise;
+
+      const options = drainTimeoutMs === undefined ? {} : { drainTimeoutMs };
+      await runtime.stop(options);
+
+      expect(
+        logger.entries.filter((entry) => entry.level === "warn")
+      ).toHaveLength(0);
+      expect(runtime.getInFlightTaskCount()).toBe(1);
+
+      release.resolve();
+      await taskPromise;
+      expect(runtime.getInFlightTaskCount()).toBe(0);
+    }
+  );
 });