Skip to content

Commit 1210071

Browse files
steipetesteipete
committed
fix: preserve cli sessions across model changes
1 parent 236e041 commit 1210071

22 files changed

Lines changed: 613 additions & 81 deletions

‎src/agents/agent-command.ts‎

Lines changed: 18 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ import { ensureAuthProfileStore } from "./auth-profiles.js";
6565
import { clearSessionAuthProfileOverride } from "./auth-profiles/session-override.js";
6666
import { resolveBootstrapWarningSignaturesSeen } from "./bootstrap-budget.js";
6767
import { runCliAgent } from "./cli-runner.js";
68-
import { getCliSessionId, setCliSessionId } from "./cli-session.js";
68+
import { clearCliSession, getCliSessionBinding, setCliSessionBinding } from "./cli-session.js";
6969
import { deliverAgentCommandResult } from "./command/delivery.js";
7070
import { resolveAgentRunContext } from "./command/run-context.js";
7171
import { updateSessionStoreAfterAgentRun } from "./command/session-store.js";
@@ -82,7 +82,6 @@ import {
8282
isCliProvider,
8383
modelKey,
8484
normalizeModelRef,
85-
normalizeProviderId,
8685
parseModelRef,
8786
resolveConfiguredModelRef,
8887
resolveDefaultModelForAgent,
@@ -386,8 +385,12 @@ function runAgentAttempt(params: {
386385
);
387386
const bootstrapPromptWarningSignature =
388387
bootstrapPromptWarningSignaturesSeen[bootstrapPromptWarningSignaturesSeen.length - 1];
388+
const authProfileId =
389+
params.providerOverride === params.authProfileProvider
390+
? params.sessionEntry?.authProfileOverride
391+
: undefined;
389392
if (isCliProvider(params.providerOverride, params.cfg)) {
390-
const cliSessionId = getCliSessionId(params.sessionEntry, params.providerOverride);
393+
const cliSessionBinding = getCliSessionBinding(params.sessionEntry, params.providerOverride);
391394
const runCliWithSession = (nextCliSessionId: string | undefined) =>
392395
runCliAgent({
393396
sessionId: params.sessionId,
@@ -404,17 +407,20 @@ function runAgentAttempt(params: {
404407
runId: params.runId,
405408
extraSystemPrompt: params.opts.extraSystemPrompt,
406409
cliSessionId: nextCliSessionId,
410+
cliSessionBinding:
411+
nextCliSessionId === cliSessionBinding?.sessionId ? cliSessionBinding : undefined,
412+
authProfileId,
407413
bootstrapPromptWarningSignaturesSeen,
408414
bootstrapPromptWarningSignature,
409415
images: params.isFallbackRetry ? undefined : params.opts.images,
410416
streamParams: params.opts.streamParams,
411417
});
412-
return runCliWithSession(cliSessionId).catch(async (err) => {
418+
return runCliWithSession(cliSessionBinding?.sessionId).catch(async (err) => {
413419
// Handle CLI session expired error
414420
if (
415421
err instanceof FailoverError &&
416422
err.reason === "session_expired" &&
417-
cliSessionId &&
423+
cliSessionBinding?.sessionId &&
418424
params.sessionKey &&
419425
params.sessionStore &&
420426
params.storePath
@@ -427,15 +433,7 @@ function runAgentAttempt(params: {
427433
const entry = params.sessionStore[params.sessionKey];
428434
if (entry) {
429435
const updatedEntry = { ...entry };
430-
if (params.providerOverride === "claude-cli") {
431-
delete updatedEntry.claudeCliSessionId;
432-
}
433-
if (updatedEntry.cliSessionIds) {
434-
const normalizedProvider = normalizeProviderId(params.providerOverride);
435-
const newCliSessionIds = { ...updatedEntry.cliSessionIds };
436-
delete newCliSessionIds[normalizedProvider];
437-
updatedEntry.cliSessionIds = newCliSessionIds;
438-
}
436+
clearCliSession(updatedEntry, params.providerOverride);
439437
updatedEntry.updatedAt = Date.now();
440438

441439
await persistSessionEntry({
@@ -453,18 +451,18 @@ function runAgentAttempt(params: {
453451
return runCliWithSession(undefined).then(async (result) => {
454452
// Update session store with new CLI session ID if available
455453
if (
456-
result.meta.agentMeta?.sessionId &&
454+
result.meta.agentMeta?.cliSessionBinding?.sessionId &&
457455
params.sessionKey &&
458456
params.sessionStore &&
459457
params.storePath
460458
) {
461459
const entry = params.sessionStore[params.sessionKey];
462460
if (entry) {
463461
const updatedEntry = { ...entry };
464-
setCliSessionId(
462+
setCliSessionBinding(
465463
updatedEntry,
466464
params.providerOverride,
467-
result.meta.agentMeta.sessionId,
465+
result.meta.agentMeta.cliSessionBinding,
468466
);
469467
updatedEntry.updatedAt = Date.now();
470468

@@ -483,10 +481,6 @@ function runAgentAttempt(params: {
483481
});
484482
}
485483

486-
const authProfileId =
487-
params.providerOverride === params.authProfileProvider
488-
? params.sessionEntry?.authProfileOverride
489-
: undefined;
490484
return runEmbeddedPiAgent({
491485
sessionId: params.sessionId,
492486
sessionKey: params.sessionKey,
@@ -1008,11 +1002,7 @@ async function agentCommandInternal(
10081002
if (overrideModel) {
10091003
const normalizedOverride = normalizeModelRef(overrideProvider, overrideModel);
10101004
const key = modelKey(normalizedOverride.provider, normalizedOverride.model);
1011-
if (
1012-
!isCliProvider(normalizedOverride.provider, cfg) &&
1013-
!allowAnyModel &&
1014-
!allowedModelKeys.has(key)
1015-
) {
1005+
if (!allowAnyModel && !allowedModelKeys.has(key)) {
10161006
const { updated } = applyModelOverrideToSessionEntry({
10171007
entry,
10181008
selection: { provider: defaultProvider, model: defaultModel, isDefault: true },
@@ -1035,11 +1025,7 @@ async function agentCommandInternal(
10351025
const candidateProvider = storedProviderOverride || defaultProvider;
10361026
const normalizedStored = normalizeModelRef(candidateProvider, storedModelOverride);
10371027
const key = modelKey(normalizedStored.provider, normalizedStored.model);
1038-
if (
1039-
isCliProvider(normalizedStored.provider, cfg) ||
1040-
allowAnyModel ||
1041-
allowedModelKeys.has(key)
1042-
) {
1028+
if (allowAnyModel || allowedModelKeys.has(key)) {
10431029
provider = normalizedStored.provider;
10441030
model = normalizedStored.model;
10451031
}
@@ -1057,11 +1043,7 @@ async function agentCommandInternal(
10571043
throw new Error("Invalid model override.");
10581044
}
10591045
const explicitKey = modelKey(explicitRef.provider, explicitRef.model);
1060-
if (
1061-
!isCliProvider(explicitRef.provider, cfg) &&
1062-
!allowAnyModel &&
1063-
!allowedModelKeys.has(explicitKey)
1064-
) {
1046+
if (!allowAnyModel && !allowedModelKeys.has(explicitKey)) {
10651047
throw new Error(
10661048
`Model override "${sanitizeForLog(explicitRef.provider)}/${sanitizeForLog(explicitRef.model)}" is not allowed for agent "${sessionAgentId}".`,
10671049
);

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import type { ImageContent } from "@mariozechner/pi-ai";
22
import { beforeEach, describe, expect, it, vi } from "vitest";
33
import { MAX_IMAGE_BYTES } from "../media/constants.js";
4-
import { loadPromptRefImages } from "./cli-runner/helpers.js";
4+
import { buildCliArgs, loadPromptRefImages } from "./cli-runner/helpers.js";
55
import * as promptImageUtils from "./pi-embedded-runner/run/images.js";
66
import type { SandboxFsBridge } from "./sandbox/fs-bridge.js";
77
import * as toolImages from "./tool-images.js";
@@ -102,3 +102,19 @@ describe("loadPromptRefImages", () => {
102102
expect(result).toEqual([loadedImage]);
103103
});
104104
});
105+
106+
describe("buildCliArgs", () => {
107+
it("keeps passing model overrides on resumed CLI sessions", () => {
108+
expect(
109+
buildCliArgs({
110+
backend: {
111+
command: "codex",
112+
modelArg: "--model",
113+
},
114+
baseArgs: ["exec", "resume", "thread-123"],
115+
modelId: "gpt-5.4",
116+
useResume: true,
117+
}),
118+
).toEqual(["exec", "resume", "thread-123", "--model", "gpt-5.4"]);
119+
});
120+
});

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

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,13 @@ import fs from "node:fs/promises";
22
import os from "node:os";
33
import path from "node:path";
44
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
5+
import { buildAnthropicCliBackend } from "../../extensions/anthropic/cli-backend.js";
6+
import { buildGoogleGeminiCliBackend } from "../../extensions/google/cli-backend.js";
7+
import { buildOpenAICodexCliBackend } from "../../extensions/openai/cli-backend.js";
58
import type { OpenClawConfig } from "../config/config.js";
69
import { resolvePreferredOpenClawTmpDir } from "../infra/tmp-openclaw-dir.js";
10+
import { createEmptyPluginRegistry } from "../plugins/registry.js";
11+
import { setActivePluginRegistry } from "../plugins/runtime.js";
712
import { resolveCliNoOutputTimeoutMs } from "./cli-runner/helpers.js";
813
import type { EmbeddedContextFile } from "./pi-embedded-helpers.js";
914
import type { WorkspaceBootstrapFile } from "./workspace.js";
@@ -157,6 +162,25 @@ describe("runCliAgent with process supervisor", () => {
157162
});
158163

159164
beforeEach(async () => {
165+
const registry = createEmptyPluginRegistry();
166+
registry.cliBackends = [
167+
{
168+
pluginId: "anthropic",
169+
backend: buildAnthropicCliBackend(),
170+
source: "test",
171+
},
172+
{
173+
pluginId: "openai",
174+
backend: buildOpenAICodexCliBackend(),
175+
source: "test",
176+
},
177+
{
178+
pluginId: "google",
179+
backend: buildGoogleGeminiCliBackend(),
180+
source: "test",
181+
},
182+
];
183+
setActivePluginRegistry(registry);
160184
supervisorSpawnMock.mockClear();
161185
enqueueSystemEventMock.mockClear();
162186
requestHeartbeatNowMock.mockClear();
@@ -211,6 +235,94 @@ describe("runCliAgent with process supervisor", () => {
211235
expect(input.scopeKey).toContain("thread-123");
212236
});
213237

238+
it("keeps resuming the CLI across model changes and passes the new model flag", async () => {
239+
mockSuccessfulCliRun();
240+
241+
await runCliAgent({
242+
sessionId: "s1",
243+
sessionFile: "/tmp/session.jsonl",
244+
workspaceDir: "/tmp",
245+
config: {
246+
agents: {
247+
defaults: {
248+
cliBackends: {
249+
"codex-cli": {
250+
command: "codex",
251+
args: ["exec", "--json"],
252+
resumeArgs: ["exec", "resume", "{sessionId}", "--json"],
253+
output: "text",
254+
modelArg: "--model",
255+
sessionMode: "existing",
256+
},
257+
},
258+
},
259+
},
260+
} satisfies OpenClawConfig,
261+
prompt: "hi",
262+
provider: "codex-cli",
263+
model: "gpt-5.4",
264+
timeoutMs: 1_000,
265+
runId: "run-model-switch",
266+
cliSessionBinding: {
267+
sessionId: "thread-123",
268+
authProfileId: "openai:default",
269+
},
270+
authProfileId: "openai:default",
271+
});
272+
273+
const input = supervisorSpawnMock.mock.calls[0]?.[0] as { argv?: string[] };
274+
expect(input.argv).toEqual([
275+
"codex",
276+
"exec",
277+
"resume",
278+
"thread-123",
279+
"--json",
280+
"--model",
281+
"gpt-5.4",
282+
"hi",
283+
]);
284+
});
285+
286+
it("starts a fresh CLI session when the auth profile changes", async () => {
287+
mockSuccessfulCliRun();
288+
289+
await runCliAgent({
290+
sessionId: "s1",
291+
sessionFile: "/tmp/session.jsonl",
292+
workspaceDir: "/tmp",
293+
config: {
294+
agents: {
295+
defaults: {
296+
cliBackends: {
297+
"codex-cli": {
298+
command: "codex",
299+
args: ["exec", "--json"],
300+
resumeArgs: ["exec", "resume", "{sessionId}", "--json"],
301+
output: "text",
302+
modelArg: "--model",
303+
sessionMode: "existing",
304+
},
305+
},
306+
},
307+
},
308+
} satisfies OpenClawConfig,
309+
prompt: "hi",
310+
provider: "codex-cli",
311+
model: "gpt-5.4",
312+
timeoutMs: 1_000,
313+
runId: "run-auth-change",
314+
cliSessionBinding: {
315+
sessionId: "thread-123",
316+
authProfileId: "openai:work",
317+
},
318+
authProfileId: "openai:personal",
319+
});
320+
321+
const input = supervisorSpawnMock.mock.calls[0]?.[0] as { argv?: string[]; scopeKey?: string };
322+
expect(input.argv).toEqual(["codex", "exec", "--json", "--model", "gpt-5.4", "hi"]);
323+
expect(input.scopeKey).toBeUndefined();
324+
});
325+
214326
it("sanitizes dangerous backend env overrides before spawn", async () => {
215327
vi.stubEnv("PATH", "/usr/bin:/bin");
216328
vi.stubEnv("HOME", "/tmp/trusted-home");

0 commit comments

Comments
 (0)