Skip to content

Commit 1e05925

Browse files
steipetesteipete
committed
fix: normalize model override auth handling
1 parent fb47f1c commit 1e05925

15 files changed

Lines changed: 497 additions & 219 deletions

‎src/agents/auth-profiles/session-override.test.ts‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import fs from "node:fs/promises";
2+
import os from "node:os";
3+
import path from "node:path";
4+
import { describe, expect, it } from "vitest";
5+
6+
import type { ClawdbotConfig } from "../../config/config.js";
7+
import type { SessionEntry } from "../../config/sessions.js";
8+
import { resolveSessionAuthProfileOverride } from "./session-override.js";
9+
10+
async function writeAuthStore(agentDir: string) {
11+
const authPath = path.join(agentDir, "auth-profiles.json");
12+
const payload = {
13+
version: 1,
14+
profiles: {
15+
"zai:work": { type: "api_key", provider: "zai", key: "sk-test" },
16+
},
17+
order: {
18+
zai: ["zai:work"],
19+
},
20+
};
21+
await fs.writeFile(authPath, JSON.stringify(payload), "utf-8");
22+
}
23+
24+
describe("resolveSessionAuthProfileOverride", () => {
25+
it("keeps user override when provider alias differs", async () => {
26+
const tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-auth-"));
27+
const prevStateDir = process.env.CLAWDBOT_STATE_DIR;
28+
process.env.CLAWDBOT_STATE_DIR = tmpDir;
29+
try {
30+
const agentDir = path.join(tmpDir, "agent");
31+
await fs.mkdir(agentDir, { recursive: true });
32+
await writeAuthStore(agentDir);
33+
34+
const sessionEntry: SessionEntry = {
35+
sessionId: "s1",
36+
updatedAt: Date.now(),
37+
authProfileOverride: "zai:work",
38+
authProfileOverrideSource: "user",
39+
};
40+
const sessionStore = { "agent:main:main": sessionEntry };
41+
42+
const resolved = await resolveSessionAuthProfileOverride({
43+
cfg: {} as ClawdbotConfig,
44+
provider: "z.ai",
45+
agentDir,
46+
sessionEntry,
47+
sessionStore,
48+
sessionKey: "agent:main:main",
49+
storePath: undefined,
50+
isNewSession: false,
51+
});
52+
53+
expect(resolved).toBe("zai:work");
54+
expect(sessionEntry.authProfileOverride).toBe("zai:work");
55+
} finally {
56+
if (prevStateDir === undefined) delete process.env.CLAWDBOT_STATE_DIR;
57+
else process.env.CLAWDBOT_STATE_DIR = prevStateDir;
58+
await fs.rm(tmpDir, { recursive: true, force: true });
59+
}
60+
});
61+
});

‎src/agents/auth-profiles/session-override.ts‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import type { ClawdbotConfig } from "../../config/config.js";
2+
import { updateSessionStore, type SessionEntry } from "../../config/sessions.js";
3+
import { normalizeProviderId } from "../model-selection.js";
4+
import {
5+
ensureAuthProfileStore,
6+
isProfileInCooldown,
7+
resolveAuthProfileOrder,
8+
} from "../auth-profiles.js";
9+
10+
function isProfileForProvider(params: {
11+
provider: string;
12+
profileId: string;
13+
store: ReturnType<typeof ensureAuthProfileStore>;
14+
}): boolean {
15+
const entry = params.store.profiles[params.profileId];
16+
if (!entry?.provider) return false;
17+
return normalizeProviderId(entry.provider) === normalizeProviderId(params.provider);
18+
}
19+
20+
export async function clearSessionAuthProfileOverride(params: {
21+
sessionEntry: SessionEntry;
22+
sessionStore: Record<string, SessionEntry>;
23+
sessionKey: string;
24+
storePath?: string;
25+
}) {
26+
const { sessionEntry, sessionStore, sessionKey, storePath } = params;
27+
delete sessionEntry.authProfileOverride;
28+
delete sessionEntry.authProfileOverrideSource;
29+
delete sessionEntry.authProfileOverrideCompactionCount;
30+
sessionEntry.updatedAt = Date.now();
31+
sessionStore[sessionKey] = sessionEntry;
32+
if (storePath) {
33+
await updateSessionStore(storePath, (store) => {
34+
store[sessionKey] = sessionEntry;
35+
});
36+
}
37+
}
38+
39+
export async function resolveSessionAuthProfileOverride(params: {
40+
cfg: ClawdbotConfig;
41+
provider: string;
42+
agentDir: string;
43+
sessionEntry?: SessionEntry;
44+
sessionStore?: Record<string, SessionEntry>;
45+
sessionKey?: string;
46+
storePath?: string;
47+
isNewSession: boolean;
48+
}): Promise<string | undefined> {
49+
const {
50+
cfg,
51+
provider,
52+
agentDir,
53+
sessionEntry,
54+
sessionStore,
55+
sessionKey,
56+
storePath,
57+
isNewSession,
58+
} = params;
59+
if (!sessionEntry || !sessionStore || !sessionKey) return sessionEntry?.authProfileOverride;
60+
61+
const store = ensureAuthProfileStore(agentDir, { allowKeychainPrompt: false });
62+
const order = resolveAuthProfileOrder({ cfg, store, provider });
63+
let current = sessionEntry.authProfileOverride?.trim();
64+
65+
if (current && !store.profiles[current]) {
66+
await clearSessionAuthProfileOverride({ sessionEntry, sessionStore, sessionKey, storePath });
67+
current = undefined;
68+
}
69+
70+
if (current && !isProfileForProvider({ provider, profileId: current, store })) {
71+
await clearSessionAuthProfileOverride({ sessionEntry, sessionStore, sessionKey, storePath });
72+
current = undefined;
73+
}
74+
75+
if (current && order.length > 0 && !order.includes(current)) {
76+
await clearSessionAuthProfileOverride({ sessionEntry, sessionStore, sessionKey, storePath });
77+
current = undefined;
78+
}
79+
80+
if (order.length === 0) return undefined;
81+
82+
const pickFirstAvailable = () =>
83+
order.find((profileId) => !isProfileInCooldown(store, profileId)) ?? order[0];
84+
const pickNextAvailable = (active: string) => {
85+
const startIndex = order.indexOf(active);
86+
if (startIndex < 0) return pickFirstAvailable();
87+
for (let offset = 1; offset <= order.length; offset += 1) {
88+
const candidate = order[(startIndex + offset) % order.length];
89+
if (!isProfileInCooldown(store, candidate)) return candidate;
90+
}
91+
return order[startIndex] ?? order[0];
92+
};
93+
94+
const compactionCount = sessionEntry.compactionCount ?? 0;
95+
const storedCompaction =
96+
typeof sessionEntry.authProfileOverrideCompactionCount === "number"
97+
? sessionEntry.authProfileOverrideCompactionCount
98+
: compactionCount;
99+
100+
const source =
101+
sessionEntry.authProfileOverrideSource ??
102+
(typeof sessionEntry.authProfileOverrideCompactionCount === "number"
103+
? "auto"
104+
: current
105+
? "user"
106+
: undefined);
107+
if (source === "user" && current && !isNewSession) {
108+
return current;
109+
}
110+
111+
let next = current;
112+
if (isNewSession) {
113+
next = current ? pickNextAvailable(current) : pickFirstAvailable();
114+
} else if (current && compactionCount > storedCompaction) {
115+
next = pickNextAvailable(current);
116+
} else if (!current || isProfileInCooldown(store, current)) {
117+
next = pickFirstAvailable();
118+
}
119+
120+
if (!next) return current;
121+
const shouldPersist =
122+
next !== sessionEntry.authProfileOverride ||
123+
sessionEntry.authProfileOverrideSource !== "auto" ||
124+
sessionEntry.authProfileOverrideCompactionCount !== compactionCount;
125+
if (shouldPersist) {
126+
sessionEntry.authProfileOverride = next;
127+
sessionEntry.authProfileOverrideSource = "auto";
128+
sessionEntry.authProfileOverrideCompactionCount = compactionCount;
129+
sessionEntry.updatedAt = Date.now();
130+
sessionStore[sessionKey] = sessionEntry;
131+
if (storePath) {
132+
await updateSessionStore(storePath, (store) => {
133+
store[sessionKey] = sessionEntry;
134+
});
135+
}
136+
}
137+
138+
return next;
139+
}

‎src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.test.ts‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,13 +92,16 @@ const makeConfig = (): ClawdbotConfig =>
9292
},
9393
}) satisfies ClawdbotConfig;
9494

95-
const writeAuthStore = async (agentDir: string) => {
95+
const writeAuthStore = async (agentDir: string, opts?: { includeAnthropic?: boolean }) => {
9696
const authPath = path.join(agentDir, "auth-profiles.json");
9797
const payload = {
9898
version: 1,
9999
profiles: {
100100
"openai:p1": { type: "api_key", provider: "openai", key: "sk-one" },
101101
"openai:p2": { type: "api_key", provider: "openai", key: "sk-two" },
102+
...(opts?.includeAnthropic
103+
? { "anthropic:default": { type: "api_key", provider: "anthropic", key: "sk-anth" } }
104+
: {}),
102105
},
103106
usageStats: {
104107
"openai:p1": { lastUsed: 1 },
@@ -206,4 +209,43 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
206209
await fs.rm(workspaceDir, { recursive: true, force: true });
207210
}
208211
});
212+
213+
it("ignores user-locked profile when provider mismatches", async () => {
214+
const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-agent-"));
215+
const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "clawdbot-workspace-"));
216+
try {
217+
await writeAuthStore(agentDir, { includeAnthropic: true });
218+
219+
runEmbeddedAttemptMock.mockResolvedValueOnce(
220+
makeAttempt({
221+
assistantTexts: ["ok"],
222+
lastAssistant: buildAssistant({
223+
stopReason: "stop",
224+
content: [{ type: "text", text: "ok" }],
225+
}),
226+
}),
227+
);
228+
229+
await runEmbeddedPiAgent({
230+
sessionId: "session:test",
231+
sessionKey: "agent:test:mismatch",
232+
sessionFile: path.join(workspaceDir, "session.jsonl"),
233+
workspaceDir,
234+
agentDir,
235+
config: makeConfig(),
236+
prompt: "hello",
237+
provider: "openai",
238+
model: "mock-1",
239+
authProfileId: "anthropic:default",
240+
authProfileIdSource: "user",
241+
timeoutMs: 5_000,
242+
runId: "run:mismatch",
243+
});
244+
245+
expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1);
246+
} finally {
247+
await fs.rm(agentDir, { recursive: true, force: true });
248+
await fs.rm(workspaceDir, { recursive: true, force: true });
249+
}
250+
});
209251
});

‎src/agents/pi-embedded-runner/run.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import {
2323
resolveAuthProfileOrder,
2424
type ResolvedProviderAuth,
2525
} from "../model-auth.js";
26+
import { normalizeProviderId } from "../model-selection.js";
2627
import { ensureClawdbotModelsJson } from "../models-config.js";
2728
import {
2829
classifyFailoverReason,
@@ -116,8 +117,16 @@ export async function runEmbeddedPiAgent(
116117

117118
const authStore = ensureAuthProfileStore(agentDir, { allowKeychainPrompt: false });
118119
const preferredProfileId = params.authProfileId?.trim();
119-
const lockedProfileId =
120-
params.authProfileIdSource === "user" ? preferredProfileId : undefined;
120+
let lockedProfileId = params.authProfileIdSource === "user" ? preferredProfileId : undefined;
121+
if (lockedProfileId) {
122+
const lockedProfile = authStore.profiles[lockedProfileId];
123+
if (
124+
!lockedProfile ||
125+
normalizeProviderId(lockedProfile.provider) !== normalizeProviderId(provider)
126+
) {
127+
lockedProfileId = undefined;
128+
}
129+
}
121130
const profileOrder = resolveAuthProfileOrder({
122131
cfg: params.config,
123132
store: authStore,

‎src/agents/tools/session-status-tool.ts‎

Lines changed: 25 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import {
3636
DEFAULT_AGENT_ID,
3737
resolveAgentIdFromSessionKey,
3838
} from "../../routing/session-key.js";
39+
import { applyModelOverrideToSessionEntry } from "../../sessions/model-overrides.js";
3940
import type { AnyAgentTool } from "./common.js";
4041
import { readStringParam } from "./common.js";
4142
import { resolveInternalSessionKey, resolveMainSessionAlias } from "./sessions-helpers.js";
@@ -240,6 +241,7 @@ export function createSessionStatusTool(opts?: {
240241
throw new Error(`Unknown sessionKey: ${requestedKeyRaw}`);
241242
}
242243

244+
const configured = resolveDefaultModelForAgent({ cfg, agentId });
243245
const modelRaw = readStringParam(params, "model");
244246
let changedModel = false;
245247
if (typeof modelRaw === "string") {
@@ -249,33 +251,33 @@ export function createSessionStatusTool(opts?: {
249251
sessionEntry: resolved.entry,
250252
agentId,
251253
});
252-
const nextEntry: SessionEntry = {
253-
...resolved.entry,
254-
updatedAt: Date.now(),
255-
};
256-
if (selection.kind === "reset" || selection.isDefault) {
257-
delete nextEntry.providerOverride;
258-
delete nextEntry.modelOverride;
259-
delete nextEntry.authProfileOverride;
260-
delete nextEntry.authProfileOverrideSource;
261-
delete nextEntry.authProfileOverrideCompactionCount;
262-
} else {
263-
nextEntry.providerOverride = selection.provider;
264-
nextEntry.modelOverride = selection.model;
265-
delete nextEntry.authProfileOverride;
266-
delete nextEntry.authProfileOverrideSource;
267-
delete nextEntry.authProfileOverrideCompactionCount;
268-
}
269-
store[resolved.key] = nextEntry;
270-
await updateSessionStore(storePath, (nextStore) => {
271-
nextStore[resolved.key] = nextEntry;
254+
const nextEntry: SessionEntry = { ...resolved.entry };
255+
const applied = applyModelOverrideToSessionEntry({
256+
entry: nextEntry,
257+
selection:
258+
selection.kind === "reset"
259+
? {
260+
provider: configured.provider,
261+
model: configured.model,
262+
isDefault: true,
263+
}
264+
: {
265+
provider: selection.provider,
266+
model: selection.model,
267+
isDefault: selection.isDefault,
268+
},
272269
});
273-
resolved.entry = nextEntry;
274-
changedModel = true;
270+
if (applied.updated) {
271+
store[resolved.key] = nextEntry;
272+
await updateSessionStore(storePath, (nextStore) => {
273+
nextStore[resolved.key] = nextEntry;
274+
});
275+
resolved.entry = nextEntry;
276+
changedModel = true;
277+
}
275278
}
276279

277280
const agentDir = resolveAgentDir(cfg, agentId);
278-
const configured = resolveDefaultModelForAgent({ cfg, agentId });
279281
const providerForCard = resolved.entry.providerOverride?.trim() || configured.provider;
280282
const usageProvider = resolveUsageProviderId(providerForCard);
281283
let usageLine: string | undefined;

‎src/auto-reply/reply/directive-handling.impl.ts‎

Lines changed: 6 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { type SessionEntry, updateSessionStore } from "../../config/sessions.js"
1010
import type { ExecAsk, ExecHost, ExecSecurity } from "../../infra/exec-approvals.js";
1111
import { enqueueSystemEvent } from "../../infra/system-events.js";
1212
import { applyVerboseOverride } from "../../sessions/level-overrides.js";
13+
import { applyModelOverrideToSessionEntry } from "../../sessions/model-overrides.js";
1314
import { formatThinkingLevels, formatXHighModelHint, supportsXHighThinking } from "../thinking.js";
1415
import type { ReplyPayload } from "../types.js";
1516
import {
@@ -340,22 +341,11 @@ export async function handleDirectiveOnly(params: {
340341
}
341342
}
342343
if (modelSelection) {
343-
if (modelSelection.isDefault) {
344-
delete sessionEntry.providerOverride;
345-
delete sessionEntry.modelOverride;
346-
} else {
347-
sessionEntry.providerOverride = modelSelection.provider;
348-
sessionEntry.modelOverride = modelSelection.model;
349-
}
350-
if (profileOverride) {
351-
sessionEntry.authProfileOverride = profileOverride;
352-
sessionEntry.authProfileOverrideSource = "user";
353-
delete sessionEntry.authProfileOverrideCompactionCount;
354-
} else if (directives.hasModelDirective) {
355-
delete sessionEntry.authProfileOverride;
356-
delete sessionEntry.authProfileOverrideSource;
357-
delete sessionEntry.authProfileOverrideCompactionCount;
358-
}
344+
applyModelOverrideToSessionEntry({
345+
entry: sessionEntry,
346+
selection: modelSelection,
347+
profileOverride,
348+
});
359349
}
360350
if (directives.hasQueueDirective && directives.queueReset) {
361351
delete sessionEntry.queueMode;

0 commit comments

Comments
 (0)