From e01229347a9184848a836bb5d616fb5cf89d9dd5 Mon Sep 17 00:00:00 2001
From: Joey Krug <joeykrug@gmail.com>
Date: Sun, 15 Mar 2026 00:07:47 -0400
Subject: [PATCH] test(agents): add embedded cleanup regressions

---
 .../run/attempt.spawn-workspace.test.ts       | 93 +++++++++++++++++++
 1 file changed, 93 insertions(+)

diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts
index 9bc424f35fd..ca69e168ad6 100644
--- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test.ts
@@ -573,10 +573,103 @@ describe("runEmbeddedAttempt cleanup", () => {
       expect.any(Object),
       "agent:main:test-cleanup",
     );
+    expect(hoisted.clearActiveEmbeddedRunMock.mock.invocationCallOrder[0]).toBeLessThan(
+      hoisted.flushPendingToolResultsAfterIdleMock.mock.invocationCallOrder[0] ??
+        Number.POSITIVE_INFINITY,
+    );
     expect(disposeMock).toHaveBeenCalledTimes(1);
     expect(hoisted.releaseWsSessionMock).toHaveBeenCalledWith("embedded-session");
     expect(hoisted.sessionLockReleaseMock).toHaveBeenCalledTimes(1);
   });
+
+  it("skips active-run clear without masking the original error when subscribe fails before registration", async () => {
+    const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-cleanup-workspace-"));
+    const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-cleanup-agent-"));
+    const sessionFile = path.join(workspaceDir, "session.jsonl");
+    tempPaths.push(workspaceDir, agentDir);
+    await fs.writeFile(sessionFile, "", "utf8");
+
+    const disposeMock = vi.fn();
+    const subscribeError = new Error("subscribe failed");
+    hoisted.subscribeEmbeddedPiSessionMock.mockReset().mockImplementation(() => {
+      throw subscribeError;
+    });
+    hoisted.createAgentSessionMock.mockImplementation(async () => ({
+      session: {
+        ...createDefaultEmbeddedSession(),
+        dispose: disposeMock,
+      },
+    }));
+
+    await expect(
+      runEmbeddedAttempt({
+        sessionId: "embedded-session",
+        sessionKey: "agent:main:test-subscribe-failure",
+        sessionFile,
+        workspaceDir,
+        agentDir,
+        config: {},
+        prompt: "hello",
+        timeoutMs: 10_000,
+        runId: "run-cleanup-subscribe-failure",
+        provider: "openai",
+        modelId: "gpt-test",
+        model: testModel,
+        authStorage: {} as AuthStorage,
+        modelRegistry: {} as ModelRegistry,
+        thinkLevel: "off",
+        senderIsOwner: true,
+        disableMessageTool: true,
+      }),
+    ).rejects.toThrow(subscribeError);
+
+    expect(hoisted.setActiveEmbeddedRunMock).not.toHaveBeenCalled();
+    expect(hoisted.clearActiveEmbeddedRunMock).not.toHaveBeenCalled();
+    expect(hoisted.flushPendingToolResultsAfterIdleMock).toHaveBeenCalledTimes(1);
+    expect(disposeMock).toHaveBeenCalledTimes(1);
+    expect(hoisted.releaseWsSessionMock).toHaveBeenCalledWith("embedded-session");
+    expect(hoisted.sessionLockReleaseMock).toHaveBeenCalledTimes(1);
+  });
+
+  it("clears the active run before waiting for the idle flush", async () => {
+    const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-cleanup-workspace-"));
+    const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-cleanup-agent-"));
+    const sessionFile = path.join(workspaceDir, "session.jsonl");
+    tempPaths.push(workspaceDir, agentDir);
+    await fs.writeFile(sessionFile, "", "utf8");
+
+    hoisted.createAgentSessionMock.mockImplementation(async () => ({
+      session: createDefaultEmbeddedSession(),
+    }));
+
+    const result = await runEmbeddedAttempt({
+      sessionId: "embedded-session",
+      sessionKey: "agent:main:test-cleanup-order",
+      sessionFile,
+      workspaceDir,
+      agentDir,
+      config: {},
+      prompt: "hello",
+      timeoutMs: 10_000,
+      runId: "run-cleanup-order",
+      provider: "openai",
+      modelId: "gpt-test",
+      model: testModel,
+      authStorage: {} as AuthStorage,
+      modelRegistry: {} as ModelRegistry,
+      thinkLevel: "off",
+      senderIsOwner: true,
+      disableMessageTool: true,
+    });
+
+    expect(result.promptError).toBeNull();
+    expect(hoisted.clearActiveEmbeddedRunMock).toHaveBeenCalledTimes(1);
+    expect(hoisted.flushPendingToolResultsAfterIdleMock).toHaveBeenCalledTimes(1);
+    expect(hoisted.clearActiveEmbeddedRunMock.mock.invocationCallOrder[0]).toBeLessThan(
+      hoisted.flushPendingToolResultsAfterIdleMock.mock.invocationCallOrder[0] ??
+        Number.POSITIVE_INFINITY,
+    );
+  });
 });
 
 describe("runEmbeddedAttempt bootstrap warning prompt assembly", () => {