Skip to content

Commit ee51169

Browse files
steipetesteipete
authored
perf: reduce session and auth cache hotpath work (#86678)
Move immutable session-store snapshot cloning/freezing off the write path and rebuild snapshots lazily on read. Resolve runtime external auth profiles once per auth-profile save instead of once per OAuth profile. Proof: oxfmt targeted files; pnpm tsgo:core; pnpm check:test-types; node scripts/run-vitest.mjs src/config/sessions.cache.test.ts src/agents/auth-profiles.store.save.test.ts src/agents/auth-profiles/external-oauth.test.ts; autoreview clean.
1 parent 9e93431 commit ee51169

9 files changed

Lines changed: 106 additions & 14 deletions

‎src/agents/auth-profiles.store.save.test.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,12 @@ import {
1414
} from "./auth-profiles/store.js";
1515
import type { AuthProfileStore } from "./auth-profiles/types.js";
1616

17+
const listRuntimeExternalAuthProfilesMock = vi.hoisted(() =>
18+
vi.fn<(_params: unknown) => unknown[]>(() => []),
19+
);
20+
1721
vi.mock("./auth-profiles/external-auth.js", () => ({
22+
listRuntimeExternalAuthProfiles: listRuntimeExternalAuthProfilesMock,
1823
overlayExternalAuthProfiles: <T>(store: T) => store,
1924
shouldPersistExternalAuthProfile: () => true,
2025
syncPersistedExternalCliAuthProfiles: <T>(store: T) => store,
@@ -35,6 +40,48 @@ function expectProfileFields(profile: unknown, expected: Record<string, unknown>
3540
}
3641

3742
describe("saveAuthProfileStore", () => {
43+
it("resolves external auth profiles once per save", async () => {
44+
const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-auth-save-once-"));
45+
const store: AuthProfileStore = {
46+
version: 1,
47+
profiles: {
48+
"openai-codex:one": {
49+
type: "oauth",
50+
provider: "openai-codex",
51+
access: "access-one",
52+
refresh: "refresh-one",
53+
expires: Date.now() + 60_000,
54+
},
55+
"openai-codex:two": {
56+
type: "oauth",
57+
provider: "openai-codex",
58+
access: "access-two",
59+
refresh: "refresh-two",
60+
expires: Date.now() + 60_000,
61+
},
62+
"openai:default": {
63+
type: "api_key",
64+
provider: "openai",
65+
key: "sk-openai",
66+
},
67+
},
68+
};
69+
70+
try {
71+
listRuntimeExternalAuthProfilesMock.mockClear();
72+
73+
saveAuthProfileStore(store, agentDir);
74+
75+
expect(listRuntimeExternalAuthProfilesMock).toHaveBeenCalledTimes(1);
76+
expect(listRuntimeExternalAuthProfilesMock.mock.calls[0]?.[0]).toMatchObject({
77+
store,
78+
agentDir,
79+
});
80+
} finally {
81+
await fs.rm(agentDir, { recursive: true, force: true });
82+
}
83+
});
84+
3885
it("strips plaintext when keyRef/tokenRef are present", async () => {
3986
const structuredCloneSpy = vi.spyOn(globalThis, "structuredClone");
4087
const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-auth-save-"));

‎src/agents/auth-profiles/external-auth.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ function resolveExternalAuthProfileMap(params: {
8989
return resolved;
9090
}
9191

92-
function listRuntimeExternalAuthProfiles(params: {
92+
export function listRuntimeExternalAuthProfiles(params: {
9393
store: AuthProfileStore;
9494
agentDir?: string;
9595
env?: NodeJS.ProcessEnv;

‎src/agents/auth-profiles/oauth-external-auth-passthrough.test-support.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { vi } from "vitest";
22

33
vi.mock("./external-auth.js", () => ({
4+
listRuntimeExternalAuthProfiles: () => [],
45
overlayExternalAuthProfiles: <T>(store: T) => store,
56
shouldPersistExternalAuthProfile: () => true,
67
}));

‎src/agents/auth-profiles/order.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ vi.mock("../../plugins/manifest-registry.js", () => ({
2323
}));
2424

2525
vi.mock("./external-auth.js", () => ({
26+
listRuntimeExternalAuthProfiles: () => [],
2627
overlayExternalAuthProfiles: <T>(store: T) => store,
2728
shouldPersistExternalAuthProfile: () => true,
2829
}));

‎src/agents/auth-profiles/store.ts‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,16 @@ import { loadJsonFile, saveJsonFile } from "../../infra/json-file.js";
77
import { cloneAuthProfileStore } from "./clone.js";
88
import { AUTH_STORE_LOCK_OPTIONS, AUTH_STORE_VERSION, log } from "./constants.js";
99
import {
10+
listRuntimeExternalAuthProfiles,
1011
overlayExternalAuthProfiles,
11-
shouldPersistExternalAuthProfile,
1212
syncPersistedExternalCliAuthProfiles,
1313
} from "./external-auth.js";
1414
import type { ExternalCliAuthDiscovery } from "./external-cli-discovery.js";
15-
import { isSafeToAdoptMainStoreOAuthIdentity } from "./oauth-shared.js";
15+
import {
16+
isSafeToAdoptMainStoreOAuthIdentity,
17+
shouldPersistRuntimeExternalOAuthProfile,
18+
type RuntimeExternalOAuthProfile,
19+
} from "./oauth-shared.js";
1620
import {
1721
ensureAuthStoreFile,
1822
resolveAuthStatePath,
@@ -364,6 +368,7 @@ function shouldKeepProfileInLocalStore(params: {
364368
credential: AuthProfileStore["profiles"][string];
365369
agentDir?: string;
366370
options?: SaveAuthProfileStoreOptions;
371+
externalProfiles: () => RuntimeExternalOAuthProfile[];
367372
}): boolean {
368373
if (params.credential.type !== "oauth") {
369374
return true;
@@ -380,11 +385,10 @@ function shouldKeepProfileInLocalStore(params: {
380385
if (params.options?.filterExternalAuthProfiles === false) {
381386
return true;
382387
}
383-
return shouldPersistExternalAuthProfile({
384-
store: params.store,
388+
return shouldPersistRuntimeExternalOAuthProfile({
385389
profileId: params.profileId,
386390
credential: params.credential,
387-
agentDir: params.agentDir,
391+
profiles: params.externalProfiles(),
388392
});
389393
}
390394

@@ -394,6 +398,12 @@ function buildLocalAuthProfileStoreForSave(params: {
394398
options?: SaveAuthProfileStoreOptions;
395399
}): AuthProfileStore {
396400
const localStore = cloneAuthProfileStore(params.store);
401+
let externalProfiles: RuntimeExternalOAuthProfile[] | undefined;
402+
const getExternalProfiles = (): RuntimeExternalOAuthProfile[] =>
403+
(externalProfiles ??= listRuntimeExternalAuthProfiles({
404+
store: params.store,
405+
agentDir: params.agentDir,
406+
}));
397407
localStore.profiles = Object.fromEntries(
398408
Object.entries(localStore.profiles).filter(([profileId, credential]) =>
399409
shouldKeepProfileInLocalStore({
@@ -402,6 +412,7 @@ function buildLocalAuthProfileStoreForSave(params: {
402412
credential,
403413
agentDir: params.agentDir,
404414
options: params.options,
415+
externalProfiles: getExternalProfiles,
405416
}),
406417
),
407418
);

‎src/agents/pi-auth-json.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { saveAuthProfileStore } from "./auth-profiles/store.js";
66
import { ensurePiAuthJsonFromAuthProfiles } from "./pi-auth-json.js";
77

88
vi.mock("./auth-profiles/external-auth.js", () => ({
9+
listRuntimeExternalAuthProfiles: () => [],
910
overlayExternalAuthProfiles: <T>(store: T) => store,
1011
shouldPersistExternalAuthProfile: () => true,
1112
syncPersistedExternalCliAuthProfiles: <T>(store: T) => store,

‎src/config/sessions.cache.test.ts‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { createSuiteTempRootTracker } from "../test-helpers/temp-dir.js";
55
import {
66
getSerializedSessionStore,
77
getSerializedSessionStoreCacheStatsForTest,
8+
getSessionStoreSnapshotCacheStatsForTest,
89
getSessionStoreStringInternStatsForTest,
910
readSessionStoreCache,
1011
setSerializedSessionStore,
@@ -554,6 +555,35 @@ describe("Session Store Cache", () => {
554555
expect(after["session:1"].displayName).toBe("Updated Session");
555556
});
556557

558+
it("builds immutable session snapshots lazily after writes", async () => {
559+
await saveSessionStore(storePath, createSingleSessionStore());
560+
561+
expect(getSessionStoreSnapshotCacheStatsForTest().entries).toBe(0);
562+
563+
const first = readSessionStoreSnapshot(storePath);
564+
const statsAfterRead = getSessionStoreSnapshotCacheStatsForTest();
565+
const second = readSessionStoreSnapshot(storePath);
566+
567+
expect(first).toBe(second);
568+
expect(Object.isFrozen(first)).toBe(true);
569+
expect(statsAfterRead.entries).toBe(1);
570+
571+
await updateSessionStore(
572+
storePath,
573+
(store) => {
574+
store["session:1"] = {
575+
...store["session:1"],
576+
displayName: "Updated lazily",
577+
updatedAt: Date.now() + 1,
578+
};
579+
},
580+
{ skipMaintenance: true },
581+
);
582+
583+
expect(getSessionStoreSnapshotCacheStatsForTest().entries).toBe(0);
584+
expect(readSessionStoreSnapshot(storePath)["session:1"].displayName).toBe("Updated lazily");
585+
});
586+
557587
it("should refresh cache when store file changes on disk", async () => {
558588
const testStore = createSingleSessionStore();
559589

‎src/config/sessions/store-cache.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,14 @@ export function getSerializedSessionStoreCacheStatsForTest(): {
156156
};
157157
}
158158

159+
export function getSessionStoreSnapshotCacheStatsForTest(): {
160+
entries: number;
161+
} {
162+
return {
163+
entries: SESSION_STORE_SNAPSHOT_CACHE.size(),
164+
};
165+
}
166+
159167
function deepFreeze<T>(value: T, seen = new WeakSet<object>()): DeepReadonly<T> {
160168
if (!value || typeof value !== "object") {
161169
return value as DeepReadonly<T>;

‎src/config/sessions/store.ts‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@ import {
2626
setSerializedSessionStore,
2727
takeMutableSessionStoreCache,
2828
writeSessionStoreCache,
29-
writeSessionStoreSnapshotCache,
3029
} from "./store-cache.js";
3130
import { normalizeStoreSessionKey, resolveSessionStoreEntry } from "./store-entry.js";
3231
import {
@@ -214,13 +213,7 @@ function updateSessionStoreWriteCaches(params: {
214213
sizeBytes: fileStat?.sizeBytes,
215214
serialized: params.serialized,
216215
});
217-
writeSessionStoreSnapshotCache({
218-
storePath: params.storePath,
219-
store: params.store,
220-
mtimeMs: fileStat?.mtimeMs,
221-
sizeBytes: fileStat?.sizeBytes,
222-
serialized: params.serialized,
223-
});
216+
dropSessionStoreSnapshotCache(params.storePath);
224217
}
225218

226219
function loadMutableSessionStoreForWriter(storePath: string): Record<string, SessionEntry> {

0 commit comments

Comments
 (0)