Skip to content

Commit 961a0f6

Browse files
vincentkocvincentkoc
committed
feat(agents): add harness runtime prewarm hook
1 parent a5eee8f commit 961a0f6

6 files changed

Lines changed: 237 additions & 1 deletion

File tree

‎extensions/codex/harness.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,40 @@ export function createCodexAppServerAgentHarness(options?: {
5050
reason: `provider is not one of: ${[...providerIds].toSorted().join(", ")}`,
5151
};
5252
},
53+
prewarm: async (params) => {
54+
const [
55+
{ resolveCodexAppServerRuntimeOptions },
56+
{ getSharedCodexAppServerClient },
57+
{ resolveCodexAppServerAuthProfileIdForAgent },
58+
{ readCodexAppServerBinding },
59+
] = await Promise.all([
60+
import("./src/app-server/config.js"),
61+
import("./src/app-server/shared-client.js"),
62+
import("./src/app-server/auth-bridge.js"),
63+
import("./src/app-server/session-binding.js"),
64+
]);
65+
const pluginConfig = options?.resolvePluginConfig?.() ?? options?.pluginConfig;
66+
const configuredAppServer = resolveCodexAppServerRuntimeOptions({ pluginConfig });
67+
const binding = params.sessionFile
68+
? await readCodexAppServerBinding(params.sessionFile, {
69+
agentDir: params.agentDir,
70+
config: params.cfg,
71+
})
72+
: undefined;
73+
const authProfileId = resolveCodexAppServerAuthProfileIdForAgent({
74+
authProfileId: params.authProfileId ?? binding?.authProfileId,
75+
agentDir: params.agentDir,
76+
config: params.cfg,
77+
});
78+
await getSharedCodexAppServerClient({
79+
startOptions: configuredAppServer.start,
80+
timeoutMs: 60_000,
81+
agentDir: params.agentDir,
82+
authProfileId,
83+
config: params.cfg,
84+
});
85+
return { warmed: true };
86+
},
5387
runAttempt: async (params) => {
5488
const { runCodexAppServerAttempt } = await import("./src/app-server/run-attempt.js");
5589
return runCodexAppServerAttempt(params, {

‎extensions/codex/index.test.ts‎

Lines changed: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,38 @@
11
import fs from "node:fs";
22
import { createTestPluginApi } from "openclaw/plugin-sdk/plugin-test-api";
3-
import { describe, expect, it, vi } from "vitest";
3+
import { beforeEach, describe, expect, it, vi } from "vitest";
44
import { createCodexAppServerAgentHarness } from "./harness.js";
55
import plugin from "./index.js";
66

77
const runCodexAppServerAttemptMock = vi.hoisted(() => vi.fn());
88
const runCodexAppServerSideQuestionMock = vi.hoisted(() => vi.fn());
9+
const getSharedCodexAppServerClientMock = vi.hoisted(() => vi.fn());
10+
const readCodexAppServerBindingMock = vi.hoisted(() => vi.fn());
11+
const resolveCodexAppServerAuthProfileIdForAgentMock = vi.hoisted(() =>
12+
vi.fn((params: { authProfileId?: string }) => params.authProfileId),
13+
);
14+
const resolveCodexAppServerRuntimeOptionsMock = vi.hoisted(() =>
15+
vi.fn(() => ({ start: { command: "codex", args: ["app-server"] } })),
16+
);
917

1018
vi.mock("./src/app-server/run-attempt.js", () => ({
1119
runCodexAppServerAttempt: runCodexAppServerAttemptMock,
1220
}));
1321
vi.mock("./src/app-server/side-question.js", () => ({
1422
runCodexAppServerSideQuestion: runCodexAppServerSideQuestionMock,
1523
}));
24+
vi.mock("./src/app-server/shared-client.js", () => ({
25+
getSharedCodexAppServerClient: getSharedCodexAppServerClientMock,
26+
}));
27+
vi.mock("./src/app-server/config.js", () => ({
28+
resolveCodexAppServerRuntimeOptions: resolveCodexAppServerRuntimeOptionsMock,
29+
}));
30+
vi.mock("./src/app-server/auth-bridge.js", () => ({
31+
resolveCodexAppServerAuthProfileIdForAgent: resolveCodexAppServerAuthProfileIdForAgentMock,
32+
}));
33+
vi.mock("./src/app-server/session-binding.js", () => ({
34+
readCodexAppServerBinding: readCodexAppServerBindingMock,
35+
}));
1636

1737
function mockCall(mock: { mock: { calls: unknown[][] } }, index = 0) {
1838
return mock.mock.calls.at(index);
@@ -23,6 +43,10 @@ function mockCallArg(mock: { mock: { calls: unknown[][] } }, index = 0, argIndex
2343
}
2444

2545
describe("codex plugin", () => {
46+
beforeEach(() => {
47+
vi.clearAllMocks();
48+
});
49+
2650
it("is opt-in by default", () => {
2751
const manifest = JSON.parse(
2852
fs.readFileSync(new URL("./openclaw.plugin.json", import.meta.url), "utf8"),
@@ -75,6 +99,7 @@ describe("codex plugin", () => {
7599
expect(agentHarnessRegistration.deliveryDefaults).toEqual({
76100
sourceVisibleReplies: "message_tool",
77101
});
102+
expect(typeof agentHarnessRegistration.prewarm).toBe("function");
78103
expect(typeof agentHarnessRegistration.dispose).toBe("function");
79104
expect(mediaProviderRegistration?.id).toBe("codex");
80105
expect(mediaProviderRegistration?.capabilities).toEqual(["image"]);
@@ -96,6 +121,45 @@ describe("codex plugin", () => {
96121
expect(typeof bindingResolvedRegistration?.[0]).toBe("function");
97122
});
98123

124+
it("prewarms existing sessions with the bound app-server auth profile", async () => {
125+
readCodexAppServerBindingMock.mockResolvedValueOnce({ authProfileId: "openai-codex:work" });
126+
getSharedCodexAppServerClientMock.mockResolvedValueOnce({});
127+
const harness = createCodexAppServerAgentHarness({ pluginConfig: { appServer: {} } });
128+
129+
await expect(
130+
harness.prewarm?.({
131+
cfg: { agents: { list: [] } } as never,
132+
agentDir: "/tmp/openclaw-agent",
133+
provider: "codex",
134+
modelId: "gpt-5.5",
135+
sessionKey: "agent:main:main",
136+
sessionFile: "/tmp/openclaw-session.jsonl",
137+
reason: "tui-startup",
138+
}),
139+
).resolves.toEqual({ warmed: true });
140+
141+
expect(resolveCodexAppServerRuntimeOptionsMock).toHaveBeenCalledWith({
142+
pluginConfig: { appServer: {} },
143+
});
144+
expect(readCodexAppServerBindingMock).toHaveBeenCalledWith("/tmp/openclaw-session.jsonl", {
145+
agentDir: "/tmp/openclaw-agent",
146+
config: { agents: { list: [] } },
147+
});
148+
expect(resolveCodexAppServerAuthProfileIdForAgentMock).toHaveBeenCalledWith({
149+
authProfileId: "openai-codex:work",
150+
agentDir: "/tmp/openclaw-agent",
151+
config: { agents: { list: [] } },
152+
});
153+
expect(getSharedCodexAppServerClientMock).toHaveBeenCalledWith({
154+
startOptions: { command: "codex", args: ["app-server"] },
155+
timeoutMs: 60_000,
156+
agentDir: "/tmp/openclaw-agent",
157+
authProfileId: "openai-codex:work",
158+
config: { agents: { list: [] } },
159+
});
160+
expect(runCodexAppServerAttemptMock).not.toHaveBeenCalled();
161+
});
162+
99163
it("registers with capture APIs that do not expose conversation binding hooks yet", () => {
100164
const registerProvider = vi.fn();
101165
const api = createTestPluginApi({

‎src/agents/harness/prewarm.test.ts‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
2+
import {
3+
clearAgentHarnesses,
4+
registerAgentHarness,
5+
restoreRegisteredAgentHarnesses,
6+
listRegisteredAgentHarnesses,
7+
} from "./registry.js";
8+
import { prewarmAgentHarnessRuntime } from "./prewarm.js";
9+
import type { AgentHarness } from "./types.js";
10+
11+
describe("prewarmAgentHarnessRuntime", () => {
12+
let previousHarnesses: ReturnType<typeof listRegisteredAgentHarnesses>;
13+
14+
beforeEach(() => {
15+
previousHarnesses = listRegisteredAgentHarnesses();
16+
clearAgentHarnesses();
17+
});
18+
19+
afterEach(() => {
20+
restoreRegisteredAgentHarnesses(previousHarnesses);
21+
});
22+
23+
it("prewarms a selected plugin harness without starting a turn", async () => {
24+
const prewarm = vi.fn<NonNullable<AgentHarness["prewarm"]>>();
25+
registerAgentHarness({
26+
id: "test-harness",
27+
label: "Test harness",
28+
supports: () => ({ supported: true, priority: 1 }),
29+
prewarm,
30+
runAttempt: vi.fn(),
31+
});
32+
33+
await expect(
34+
prewarmAgentHarnessRuntime({
35+
provider: "test",
36+
modelId: "model-a",
37+
sessionKey: "agent:main:main",
38+
reason: "tui-startup",
39+
}),
40+
).resolves.toEqual({ status: "warmed", harnessId: "test-harness" });
41+
expect(prewarm).toHaveBeenCalledWith(
42+
expect.objectContaining({
43+
provider: "test",
44+
modelId: "model-a",
45+
sessionKey: "agent:main:main",
46+
reason: "tui-startup",
47+
}),
48+
);
49+
});
50+
51+
});

‎src/agents/harness/prewarm.ts‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
import { createSubsystemLogger } from "../../logging/subsystem.js";
2+
import { ensureSelectedAgentHarnessPlugin } from "./runtime-plugin.js";
3+
import { selectAgentHarness } from "./selection.js";
4+
import type { AgentHarnessPrewarmParams } from "./types.js";
5+
6+
const log = createSubsystemLogger("agents/harness");
7+
8+
export type AgentHarnessPrewarmStatus =
9+
| {
10+
status: "pi" | "unsupported";
11+
harnessId: string;
12+
}
13+
| {
14+
status: "warmed";
15+
harnessId: string;
16+
}
17+
| {
18+
status: "failed";
19+
harnessId?: string;
20+
error: string;
21+
};
22+
23+
export async function prewarmAgentHarnessRuntime(
24+
params: AgentHarnessPrewarmParams,
25+
): Promise<AgentHarnessPrewarmStatus> {
26+
try {
27+
if (params.workspaceDir) {
28+
await ensureSelectedAgentHarnessPlugin({
29+
config: params.cfg,
30+
provider: params.provider,
31+
modelId: params.modelId ?? "",
32+
agentId: params.agentId,
33+
sessionKey: params.sessionKey,
34+
workspaceDir: params.workspaceDir,
35+
});
36+
}
37+
const harness = selectAgentHarness({
38+
provider: params.provider,
39+
modelId: params.modelId,
40+
config: params.cfg,
41+
agentId: params.agentId,
42+
sessionKey: params.sessionKey,
43+
});
44+
if (harness.id === "pi") {
45+
return { status: "pi", harnessId: harness.id };
46+
}
47+
if (!harness.prewarm) {
48+
return { status: "unsupported", harnessId: harness.id };
49+
}
50+
await harness.prewarm(params);
51+
return { status: "warmed", harnessId: harness.id };
52+
} catch (error) {
53+
log.warn("agent harness prewarm failed", {
54+
provider: params.provider,
55+
modelId: params.modelId,
56+
sessionKey: params.sessionKey,
57+
agentId: params.agentId,
58+
error,
59+
});
60+
return {
61+
status: "failed",
62+
error: error instanceof Error ? error.message : String(error),
63+
};
64+
}
65+
}

‎src/agents/harness/types.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,24 @@ export type AgentHarnessSideQuestionParams = {
4242
export type AgentHarnessSideQuestionResult = {
4343
text: string;
4444
};
45+
export type AgentHarnessPrewarmReason = "tui-startup" | "gateway-startup" | "manual" | "unknown";
46+
export type AgentHarnessPrewarmParams = {
47+
cfg?: import("../../config/types.openclaw.js").OpenClawConfig;
48+
agentDir?: string;
49+
provider: string;
50+
modelId?: string;
51+
sessionKey?: string;
52+
sessionId?: string;
53+
sessionFile?: string;
54+
agentId?: string;
55+
workspaceDir?: string;
56+
authProfileId?: string;
57+
authProfileIdSource?: "auto" | "user";
58+
reason: AgentHarnessPrewarmReason;
59+
};
60+
export type AgentHarnessPrewarmResult = {
61+
warmed?: boolean;
62+
};
4563
export type AgentHarnessCompactParams =
4664
import("../pi-embedded-runner/compact.types.js").CompactEmbeddedPiSessionParams;
4765
export type AgentHarnessCompactResult =
@@ -77,6 +95,7 @@ export type AgentHarness = {
7795
contextEngineHostCapabilities?: readonly import("../../context-engine/types.js").ContextEngineHostCapability[];
7896
deliveryDefaults?: AgentHarnessDeliveryDefaults;
7997
supports(ctx: AgentHarnessSupportContext): AgentHarnessSupport;
98+
prewarm?(params: AgentHarnessPrewarmParams): Promise<AgentHarnessPrewarmResult | void> | void;
8099
runAttempt(params: AgentHarnessAttemptParams): Promise<AgentHarnessAttemptResult>;
81100
runSideQuestion?(params: AgentHarnessSideQuestionParams): Promise<AgentHarnessSideQuestionResult>;
82101
classify?(

‎src/plugin-sdk/agent-harness-runtime.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ export type {
2929
AgentHarnessCompactParams,
3030
AgentHarnessCompactResult,
3131
AgentHarnessDeliveryDefaults,
32+
AgentHarnessPrewarmParams,
33+
AgentHarnessPrewarmReason,
34+
AgentHarnessPrewarmResult,
3235
AgentHarnessResultClassification,
3336
AgentHarnessSideQuestionParams,
3437
AgentHarnessSideQuestionResult,

0 commit comments

Comments
 (0)