Skip to content

Commit f40f8a6

Browse files
vincentkocvincentkoc
authored
fix(agents): harden cli runner hook followups (#70747)
* fix(agents): harden cli runner hook followups * fix(agents): harden cli runner hook followups * fix(agents): harden cli runner hook followups * fix(agents): harden cli runner hook followups * fix(agents): harden cli runner hook followups * fix(agents): harden cli runner hook followups
1 parent e93b3f6 commit f40f8a6

6 files changed

Lines changed: 510 additions & 126 deletions

File tree

‎src/agents/cli-runner.reliability.test.ts‎

Lines changed: 153 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import fs from "node:fs";
22
import os from "node:os";
33
import path from "node:path";
4-
import { CURRENT_SESSION_VERSION, SessionManager } from "@mariozechner/pi-coding-agent";
4+
import { CURRENT_SESSION_VERSION } from "@mariozechner/pi-coding-agent";
55
import { afterEach, describe, expect, it, vi } from "vitest";
66
import {
77
__testing as replyRunTesting,
@@ -18,6 +18,7 @@ import {
1818
} from "./cli-runner.test-support.js";
1919
import { executePreparedCliRun } from "./cli-runner/execute.js";
2020
import { resolveCliNoOutputTimeoutMs } from "./cli-runner/helpers.js";
21+
import { MAX_CLI_SESSION_HISTORY_MESSAGES } from "./cli-runner/session-history.js";
2122
import type { PreparedCliRunContext } from "./cli-runner/types.js";
2223

2324
vi.mock("../plugins/hook-runner-global.js", async () => {
@@ -34,7 +35,9 @@ const mockGetGlobalHookRunner = vi.mocked(getGlobalHookRunner);
3435

3536
function createSessionFile(params?: { history?: Array<{ role: "user"; content: string }> }) {
3637
const dir = fs.mkdtempSync(path.join(os.tmpdir(), "openclaw-cli-hooks-"));
37-
const sessionFile = path.join(dir, "session.jsonl");
38+
vi.stubEnv("OPENCLAW_STATE_DIR", dir);
39+
const sessionFile = path.join(dir, "agents", "main", "sessions", "s1.jsonl");
40+
fs.mkdirSync(path.dirname(sessionFile), { recursive: true });
3841
fs.writeFileSync(
3942
sessionFile,
4043
`${JSON.stringify({
@@ -46,13 +49,22 @@ function createSessionFile(params?: { history?: Array<{ role: "user"; content: s
4649
})}\n`,
4750
"utf-8",
4851
);
49-
const sessionManager = SessionManager.open(sessionFile);
50-
for (const entry of params?.history ?? []) {
51-
sessionManager.appendMessage({
52-
role: entry.role,
53-
content: entry.content,
54-
timestamp: Date.now(),
55-
});
52+
for (const [index, entry] of (params?.history ?? []).entries()) {
53+
fs.appendFileSync(
54+
sessionFile,
55+
`${JSON.stringify({
56+
type: "message",
57+
id: `msg-${index}`,
58+
parentId: index > 0 ? `msg-${index - 1}` : null,
59+
timestamp: new Date(index + 1).toISOString(),
60+
message: {
61+
role: entry.role,
62+
content: entry.content,
63+
timestamp: index + 1,
64+
},
65+
})}\n`,
66+
"utf-8",
67+
);
5668
}
5769
return { dir, sessionFile };
5870
}
@@ -110,6 +122,7 @@ describe("runCliAgent reliability", () => {
110122
afterEach(() => {
111123
replyRunTesting.resetReplyRunRegistry();
112124
mockGetGlobalHookRunner.mockReset();
125+
vi.unstubAllEnvs();
113126
});
114127

115128
it("fails with timeout when no-output watchdog trips", async () => {
@@ -193,6 +206,13 @@ describe("runCliAgent reliability", () => {
193206
});
194207

195208
it("rethrows the retry failure when session-expired recovery retry also fails", async () => {
209+
const hookRunner = {
210+
hasHooks: vi.fn((hookName: string) => ["llm_input", "agent_end"].includes(hookName)),
211+
runLlmInput: vi.fn(async () => undefined),
212+
runLlmOutput: vi.fn(async () => undefined),
213+
runAgentEnd: vi.fn(async () => undefined),
214+
};
215+
mockGetGlobalHookRunner.mockReturnValue(hookRunner as never);
196216
supervisorSpawnMock.mockClear();
197217
supervisorSpawnMock.mockResolvedValueOnce(
198218
createManagedRun({
@@ -218,18 +238,50 @@ describe("runCliAgent reliability", () => {
218238
noOutputTimedOut: false,
219239
}),
220240
);
241+
const { dir, sessionFile } = createSessionFile({
242+
history: [{ role: "user", content: "earlier context" }],
243+
});
221244

222-
await expect(
223-
runPreparedCliAgent(
224-
buildPreparedContext({
225-
sessionKey: "agent:main:subagent:retry",
226-
runId: "run-retry-failure",
227-
cliSessionId: "thread-123",
245+
try {
246+
await expect(
247+
runPreparedCliAgent({
248+
...buildPreparedContext({
249+
sessionKey: "agent:main:subagent:retry",
250+
runId: "run-retry-failure",
251+
cliSessionId: "thread-123",
252+
}),
253+
params: {
254+
...buildPreparedContext({
255+
sessionKey: "agent:main:subagent:retry",
256+
runId: "run-retry-failure",
257+
cliSessionId: "thread-123",
258+
}).params,
259+
agentId: "main",
260+
sessionFile,
261+
workspaceDir: dir,
262+
},
228263
}),
229-
),
230-
).rejects.toThrow("rate limit exceeded");
264+
).rejects.toThrow("rate limit exceeded");
231265

232-
expect(supervisorSpawnMock).toHaveBeenCalledTimes(2);
266+
expect(supervisorSpawnMock).toHaveBeenCalledTimes(2);
267+
await vi.waitFor(() => {
268+
expect(hookRunner.runLlmInput).toHaveBeenCalledTimes(1);
269+
expect(hookRunner.runAgentEnd).toHaveBeenCalledTimes(1);
270+
});
271+
expect(hookRunner.runAgentEnd).toHaveBeenCalledWith(
272+
expect.objectContaining({
273+
success: false,
274+
error: "rate limit exceeded",
275+
messages: [
276+
{ role: "user", content: "earlier context", timestamp: expect.any(Number) },
277+
{ role: "user", content: "hi", timestamp: expect.any(Number) },
278+
],
279+
}),
280+
expect.any(Object),
281+
);
282+
} finally {
283+
fs.rmSync(dir, { recursive: true, force: true });
284+
}
233285
});
234286

235287
it("returns the assembled CLI prompt in meta for raw trace consumers", async () => {
@@ -499,6 +551,89 @@ describe("runCliAgent reliability", () => {
499551
expect.any(Object),
500552
);
501553
});
554+
555+
it("does not emit duplicate llm_input when session-expired recovery succeeds", async () => {
556+
const hookRunner = {
557+
hasHooks: vi.fn((hookName: string) =>
558+
["llm_input", "llm_output", "agent_end"].includes(hookName),
559+
),
560+
runLlmInput: vi.fn(async () => undefined),
561+
runLlmOutput: vi.fn(async () => undefined),
562+
runAgentEnd: vi.fn(async () => undefined),
563+
};
564+
mockGetGlobalHookRunner.mockReturnValue(hookRunner as never);
565+
const { dir, sessionFile } = createSessionFile({
566+
history: Array.from({ length: MAX_CLI_SESSION_HISTORY_MESSAGES + 5 }, (_, index) => ({
567+
role: "user" as const,
568+
content: `history-${index}`,
569+
})),
570+
});
571+
572+
supervisorSpawnMock.mockResolvedValueOnce(
573+
createManagedRun({
574+
reason: "exit",
575+
exitCode: 1,
576+
exitSignal: null,
577+
durationMs: 50,
578+
stdout: "",
579+
stderr: "session expired",
580+
timedOut: false,
581+
noOutputTimedOut: false,
582+
}),
583+
);
584+
supervisorSpawnMock.mockResolvedValueOnce(
585+
createManagedRun({
586+
reason: "exit",
587+
exitCode: 0,
588+
exitSignal: null,
589+
durationMs: 50,
590+
stdout: "recovered output",
591+
stderr: "",
592+
timedOut: false,
593+
noOutputTimedOut: false,
594+
}),
595+
);
596+
597+
try {
598+
await expect(
599+
runPreparedCliAgent({
600+
...buildPreparedContext({
601+
sessionKey: "agent:main:main",
602+
runId: "run-retry-success",
603+
cliSessionId: "thread-123",
604+
}),
605+
params: {
606+
...buildPreparedContext({
607+
sessionKey: "agent:main:main",
608+
runId: "run-retry-success",
609+
cliSessionId: "thread-123",
610+
}).params,
611+
agentId: "main",
612+
sessionFile,
613+
workspaceDir: dir,
614+
},
615+
}),
616+
).resolves.toMatchObject({
617+
payloads: [{ text: "recovered output" }],
618+
});
619+
620+
await vi.waitFor(() => {
621+
expect(hookRunner.runLlmInput).toHaveBeenCalledTimes(1);
622+
expect(hookRunner.runLlmOutput).toHaveBeenCalledTimes(1);
623+
expect(hookRunner.runAgentEnd).toHaveBeenCalledTimes(1);
624+
});
625+
const llmInputCalls = hookRunner.runLlmInput.mock.calls as unknown as Array<Array<unknown>>;
626+
const llmInputEvent = llmInputCalls[0]?.[0] as { historyMessages: unknown[] } | undefined;
627+
expect(llmInputEvent).toBeDefined();
628+
expect(llmInputEvent?.historyMessages).toHaveLength(MAX_CLI_SESSION_HISTORY_MESSAGES);
629+
expect(llmInputEvent?.historyMessages[0]).toMatchObject({
630+
role: "user",
631+
content: `history-5`,
632+
});
633+
} finally {
634+
fs.rmSync(dir, { recursive: true, force: true });
635+
}
636+
});
502637
});
503638

504639
describe("resolveCliNoOutputTimeoutMs", () => {

0 commit comments

Comments
 (0)