fix(agents): avoid session event queue self-wait

luoyanglang · steipete · commit 552c794395e7 · 2026-05-27T20:29:36.000+01:00
diff --git a/src/agents/embedded-agent-runner/run/attempt.session-lock.test.ts b/src/agents/embedded-agent-runner/run/attempt.session-lock.test.ts
@@ -1189,6 +1189,147 @@ describe("embedded attempt session lock lifecycle", () => {
     ]);
   });
 
+  it("does not wait on the session event queue from a hook running during active event processing", async () => {
+    const events: string[] = [];
+    const session = {
+      _agentEventQueue: Promise.resolve(),
+      _extensionRunner: {
+        hasHandlers: vi.fn((eventType: string) => eventType === "tool_call"),
+      },
+      _processAgentEvent: vi.fn(async (event: { type?: string }) => {
+        events.push(`process:${event.type}`);
+        await session.agent.beforeToolCall();
+        events.push("process:end");
+      }),
+      _handleAgentEvent(event: { type?: string }) {
+        events.push(`handle:${event.type}`);
+        session["_agentEventQueue"] = session["_agentEventQueue"].then(() =>
+          session["_processAgentEvent"](event),
+        );
+        session["_agentEventQueue"].catch(() => {});
+      },
+      agent: {
+        beforeToolCall: vi.fn(async () => {
+          events.push("hook");
+        }),
+      },
+    };
+
+    installSessionEventWriteLock({
+      session,
+      withSessionWriteLock: async (run) => {
+        events.push("event-lock");
+        return await run();
+      },
+    });
+    installSessionExternalHookWriteLock({
+      session,
+      withSessionWriteLock: async (run) => {
+        events.push("hook-lock");
+        return await run();
+      },
+    });
+
+    const result = session["_handleAgentEvent"]({
+      type: "tool_call",
+    }) as unknown as Promise<unknown>;
+    const completion = await Promise.race([
+      result.then(() => "done"),
+      new Promise<string>((resolve) => {
+        setTimeout(() => resolve("timeout"), 25);
+      }),
+    ]);
+
+    expect(completion).toBe("done");
+    expect(events).toEqual([
+      "event-lock",
+      "handle:tool_call",
+      "process:tool_call",
+      "hook-lock",
+      "hook",
+      "process:end",
+    ]);
+  });
+
+  it("drains queued session events for async hook work spawned after event processing returns", async () => {
+    const events: string[] = [];
+    let releaseQueue!: () => void;
+    let spawnedHook!: Promise<void>;
+    const session = {
+      _agentEventQueue: Promise.resolve(),
+      _extensionRunner: {
+        hasHandlers: vi.fn((eventType: string) => eventType === "tool_call"),
+      },
+      _processAgentEvent: vi.fn(async (event: { type?: string }) => {
+        events.push(`process:${event.type}`);
+        spawnedHook = new Promise<void>((resolve, reject) => {
+          setTimeout(() => {
+            session.agent.beforeToolCall().then(resolve, reject);
+          }, 0);
+        });
+        session["_agentEventQueue"] = new Promise<void>((resolve) => {
+          releaseQueue = resolve;
+        }).then(() => {
+          events.push("queue-drained");
+        });
+        events.push("process:end");
+      }),
+      _handleAgentEvent(event: { type?: string }) {
+        events.push(`handle:${event.type}`);
+        session["_agentEventQueue"] = session["_agentEventQueue"].then(() =>
+          session["_processAgentEvent"](event),
+        );
+        session["_agentEventQueue"].catch(() => {});
+      },
+      agent: {
+        beforeToolCall: vi.fn(async () => {
+          events.push("hook");
+        }),
+      },
+    };
+
+    installSessionEventWriteLock({
+      session,
+      withSessionWriteLock: async (run) => {
+        events.push("event-lock");
+        return await run();
+      },
+    });
+    installSessionExternalHookWriteLock({
+      session,
+      withSessionWriteLock: async (run) => {
+        events.push("hook-lock");
+        return await run();
+      },
+    });
+
+    session["_handleAgentEvent"]({ type: "tool_call" });
+    await new Promise<void>((resolve) => setTimeout(resolve, 5));
+
+    const beforeQueueRelease = await Promise.race([
+      spawnedHook.then(() => "done"),
+      new Promise<string>((resolve) => {
+        setTimeout(() => resolve("waiting"), 25);
+      }),
+    ]);
+
+    expect(beforeQueueRelease).toBe("waiting");
+    expect(events).toEqual(["event-lock", "handle:tool_call", "process:tool_call", "process:end"]);
+
+    releaseQueue();
+    await spawnedHook;
+
+    expect(events).toEqual([
+      "event-lock",
+      "handle:tool_call",
+      "process:tool_call",
+      "process:end",
+      "queue-drained",
+      "hook-lock",
+      "hook",
+    ]);
+  });
+
   it("locks OpenClaw extension hooks that can mutate the session outside agent events", async () => {
     const locked: string[] = [];
     const called: string[] = [];
diff --git a/src/agents/embedded-agent-runner/run/attempt.session-lock.ts b/src/agents/embedded-agent-runner/run/attempt.session-lock.ts
@@ -74,6 +74,32 @@ type LockableFunction = ((...args: unknown[]) => unknown) & {
   __openclawSessionWriteLockInstalled?: boolean;
 };
 
+type SessionEventHookContext = {
+  session: unknown;
+  active: boolean;
+};
+
+const sessionEventHookContext = new AsyncLocalStorage<SessionEventHookContext>();
+
+function isProcessingAgentEventInCurrentChain(session: unknown): boolean {
+  const context = sessionEventHookContext.getStore();
+  return context?.active === true && context.session === session;
+}
+
+async function withProcessingAgentEvent<T>(
+  session: unknown,
+  run: () => Promise<T> | T,
+): Promise<T> {
+  const context: SessionEventHookContext = { session, active: true };
+  return await sessionEventHookContext.run(context, async () => {
+    try {
+      return await run();
+    } finally {
+      context.active = false;
+    }
+  });
+}
+
 function sessionHasExtensionHandlers(session: SessionEventProcessor, eventType: string): boolean {
   const extensionRunner = session["_extensionRunner"];
   const hasHandlers = extensionRunner?.hasHandlers;
@@ -662,6 +688,15 @@ async function waitForSessionEventQueue(session: unknown): Promise<void> {
   }
 }
 
+async function waitForSessionEventQueueBeforeHook(session: unknown): Promise<void> {
+  // Hook calls made by _handleAgentEvent are already inside the current queue
+  // entry. Draining there waits on itself; detached/external hook work still drains.
+  if (isProcessingAgentEventInCurrentChain(session)) {
+    return;
+  }
+  await waitForSessionEventQueue(session);
+}
+
 function installAwaitableSessionEventQueue(session: unknown): void {
   const owner = session as SessionEventQueueBridge;
   const original = owner["_handleAgentEvent"];
@@ -731,10 +766,14 @@ export function installSessionEventWriteLock(params: {
     event: unknown,
     signal?: unknown,
   ) {
-    if (!eventMayReachTranscriptWriters(session, event)) {
-      return await original.call(this, event, signal);
-    }
-    return await params.withSessionWriteLock(async () => await original.call(this, event, signal));
+    return await withProcessingAgentEvent(session, async () => {
+      if (!eventMayReachTranscriptWriters(session, event)) {
+        return await original.call(this, event, signal);
+      }
+      return await params.withSessionWriteLock(
+        async () => await original.call(this, event, signal),
+      );
+    });
   };
   wrapped["__openclawSessionEventQueueAwaitInstalled"] =
     original["__openclawSessionEventQueueAwaitInstalled"];
@@ -756,28 +795,28 @@ export function installSessionExternalHookWriteLock(params: {
       owner: agent as Record<string, unknown>,
       key: "beforeToolCall",
       shouldLock: () => true,
-      waitBeforeLock: () => waitForSessionEventQueue(session),
+      waitBeforeLock: () => waitForSessionEventQueueBeforeHook(session),
       withSessionWriteLock: params.withSessionWriteLock,
     });
     installLockableFunction({
       owner: agent as Record<string, unknown>,
       key: "afterToolCall",
       shouldLock: () => sessionHasExtensionHandlers(session, "tool_result"),
-      waitBeforeLock: () => waitForSessionEventQueue(session),
+      waitBeforeLock: () => waitForSessionEventQueueBeforeHook(session),
       withSessionWriteLock: params.withSessionWriteLock,
     });
     installLockableFunction({
       owner: agent as Record<string, unknown>,
       key: "onPayload",
       shouldLock: () => sessionHasExtensionHandlers(session, "before_provider_request"),
-      waitBeforeLock: () => waitForSessionEventQueue(session),
+      waitBeforeLock: () => waitForSessionEventQueueBeforeHook(session),
       withSessionWriteLock: params.withSessionWriteLock,
     });
     installLockableFunction({
       owner: agent as Record<string, unknown>,
       key: "onResponse",
       shouldLock: () => sessionHasExtensionHandlers(session, "after_provider_response"),
-      waitBeforeLock: () => waitForSessionEventQueue(session),
+      waitBeforeLock: () => waitForSessionEventQueueBeforeHook(session),
       withSessionWriteLock: params.withSessionWriteLock,
     });
   }
