Skip to content

Commit 2bc031c

Browse files
vincentkocvincentkoc
committed
perf(cron): keep auth profile runtime cold
1 parent 6a8704c commit 2bc031c

8 files changed

Lines changed: 209 additions & 76 deletions

File tree

‎src/agents/auth-profiles/runtime-snapshots.ts‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
import { resolveAuthStorePath } from "./paths.js";
2+
import type { AuthProfileStore } from "./types.js";
3+
4+
const runtimeAuthStoreSnapshots = new Map<string, AuthProfileStore>();
5+
6+
function resolveRuntimeStoreKey(agentDir?: string): string {
7+
return resolveAuthStorePath(agentDir);
8+
}
9+
10+
function cloneAuthProfileStore(store: AuthProfileStore): AuthProfileStore {
11+
return structuredClone(store);
12+
}
13+
14+
export function getRuntimeAuthProfileStoreSnapshot(
15+
agentDir?: string,
16+
): AuthProfileStore | undefined {
17+
const store = runtimeAuthStoreSnapshots.get(resolveRuntimeStoreKey(agentDir));
18+
return store ? cloneAuthProfileStore(store) : undefined;
19+
}
20+
21+
export function hasRuntimeAuthProfileStoreSnapshot(agentDir?: string): boolean {
22+
return runtimeAuthStoreSnapshots.has(resolveRuntimeStoreKey(agentDir));
23+
}
24+
25+
export function hasAnyRuntimeAuthProfileStoreSource(agentDir?: string): boolean {
26+
const requestedStore = getRuntimeAuthProfileStoreSnapshot(agentDir);
27+
if (requestedStore && Object.keys(requestedStore.profiles).length > 0) {
28+
return true;
29+
}
30+
if (!agentDir) {
31+
return false;
32+
}
33+
const mainStore = getRuntimeAuthProfileStoreSnapshot();
34+
return Boolean(mainStore && Object.keys(mainStore.profiles).length > 0);
35+
}
36+
37+
export function replaceRuntimeAuthProfileStoreSnapshots(
38+
entries: Array<{ agentDir?: string; store: AuthProfileStore }>,
39+
): void {
40+
runtimeAuthStoreSnapshots.clear();
41+
for (const entry of entries) {
42+
runtimeAuthStoreSnapshots.set(
43+
resolveRuntimeStoreKey(entry.agentDir),
44+
cloneAuthProfileStore(entry.store),
45+
);
46+
}
47+
}
48+
49+
export function clearRuntimeAuthProfileStoreSnapshots(): void {
50+
runtimeAuthStoreSnapshots.clear();
51+
}
52+
53+
export function setRuntimeAuthProfileStoreSnapshot(
54+
store: AuthProfileStore,
55+
agentDir?: string,
56+
): void {
57+
runtimeAuthStoreSnapshots.set(resolveRuntimeStoreKey(agentDir), cloneAuthProfileStore(store));
58+
}

‎src/agents/auth-profiles/source-check.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import fs from "node:fs";
2+
import { resolveAuthStatePath, resolveAuthStorePath, resolveLegacyAuthStorePath } from "./paths.js";
3+
import { hasAnyRuntimeAuthProfileStoreSource } from "./runtime-snapshots.js";
4+
5+
function hasStoredAuthProfileFiles(agentDir?: string): boolean {
6+
return (
7+
fs.existsSync(resolveAuthStorePath(agentDir)) ||
8+
fs.existsSync(resolveAuthStatePath(agentDir)) ||
9+
fs.existsSync(resolveLegacyAuthStorePath(agentDir))
10+
);
11+
}
12+
13+
export function hasAnyAuthProfileStoreSource(agentDir?: string): boolean {
14+
if (hasAnyRuntimeAuthProfileStoreSource(agentDir)) {
15+
return true;
16+
}
17+
if (hasStoredAuthProfileFiles(agentDir)) {
18+
return true;
19+
}
20+
21+
const authPath = resolveAuthStorePath(agentDir);
22+
const mainAuthPath = resolveAuthStorePath();
23+
if (agentDir && authPath !== mainAuthPath && hasStoredAuthProfileFiles(undefined)) {
24+
return true;
25+
}
26+
return false;
27+
}

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

Lines changed: 26 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,13 @@ import {
2323
mergeAuthProfileStores,
2424
mergeOAuthFileIntoStore,
2525
} from "./persisted.js";
26+
import {
27+
clearRuntimeAuthProfileStoreSnapshots as clearRuntimeAuthProfileStoreSnapshotsImpl,
28+
getRuntimeAuthProfileStoreSnapshot,
29+
hasRuntimeAuthProfileStoreSnapshot,
30+
replaceRuntimeAuthProfileStoreSnapshots as replaceRuntimeAuthProfileStoreSnapshotsImpl,
31+
setRuntimeAuthProfileStoreSnapshot,
32+
} from "./runtime-snapshots.js";
2633
import { savePersistedAuthProfileState } from "./state.js";
2734
import type { AuthProfileStore } from "./types.js";
2835

@@ -37,7 +44,6 @@ type SaveAuthProfileStoreOptions = {
3744
syncExternalCli?: boolean;
3845
};
3946

40-
const runtimeAuthStoreSnapshots = new Map<string, AuthProfileStore>();
4147
const loadedAuthStoreCache = new Map<
4248
string,
4349
{
@@ -48,72 +54,36 @@ const loadedAuthStoreCache = new Map<
4854
}
4955
>();
5056

51-
function resolveRuntimeStoreKey(agentDir?: string): string {
52-
return resolveAuthStorePath(agentDir);
53-
}
54-
5557
function cloneAuthProfileStore(store: AuthProfileStore): AuthProfileStore {
5658
return structuredClone(store);
5759
}
5860

5961
function resolveRuntimeAuthProfileStore(agentDir?: string): AuthProfileStore | null {
60-
if (runtimeAuthStoreSnapshots.size === 0) {
61-
return null;
62-
}
63-
64-
const mainKey = resolveRuntimeStoreKey(undefined);
65-
const requestedKey = resolveRuntimeStoreKey(agentDir);
66-
const mainStore = runtimeAuthStoreSnapshots.get(mainKey);
67-
const requestedStore = runtimeAuthStoreSnapshots.get(requestedKey);
62+
const mainKey = resolveAuthStorePath(undefined);
63+
const requestedKey = resolveAuthStorePath(agentDir);
64+
const mainStore = getRuntimeAuthProfileStoreSnapshot(undefined);
65+
const requestedStore = getRuntimeAuthProfileStoreSnapshot(agentDir);
6866

6967
if (!agentDir || requestedKey === mainKey) {
7068
if (!mainStore) {
7169
return null;
7270
}
73-
return cloneAuthProfileStore(mainStore);
71+
return mainStore;
7472
}
7573

7674
if (mainStore && requestedStore) {
77-
return mergeAuthProfileStores(
78-
cloneAuthProfileStore(mainStore),
79-
cloneAuthProfileStore(requestedStore),
80-
);
75+
return mergeAuthProfileStores(mainStore, requestedStore);
8176
}
8277
if (requestedStore) {
83-
return cloneAuthProfileStore(requestedStore);
78+
return requestedStore;
8479
}
8580
if (mainStore) {
86-
return cloneAuthProfileStore(mainStore);
81+
return mainStore;
8782
}
8883

8984
return null;
9085
}
9186

92-
function hasStoredAuthProfileFiles(agentDir?: string): boolean {
93-
return (
94-
fs.existsSync(resolveAuthStorePath(agentDir)) ||
95-
fs.existsSync(resolveAuthStatePath(agentDir)) ||
96-
fs.existsSync(resolveLegacyAuthStorePath(agentDir))
97-
);
98-
}
99-
100-
export function replaceRuntimeAuthProfileStoreSnapshots(
101-
entries: Array<{ agentDir?: string; store: AuthProfileStore }>,
102-
): void {
103-
runtimeAuthStoreSnapshots.clear();
104-
for (const entry of entries) {
105-
runtimeAuthStoreSnapshots.set(
106-
resolveRuntimeStoreKey(entry.agentDir),
107-
cloneAuthProfileStore(entry.store),
108-
);
109-
}
110-
}
111-
112-
export function clearRuntimeAuthProfileStoreSnapshots(): void {
113-
runtimeAuthStoreSnapshots.clear();
114-
loadedAuthStoreCache.clear();
115-
}
116-
11787
function readAuthStoreMtimeMs(authPath: string): number | null {
11888
try {
11989
return fs.statSync(authPath).mtimeMs;
@@ -387,23 +357,17 @@ export function ensureAuthProfileStoreForLocalUpdate(agentDir?: string): AuthPro
387357
return mergeAuthProfileStores(mainStore, store);
388358
}
389359

390-
export function hasAnyAuthProfileStoreSource(agentDir?: string): boolean {
391-
const runtimeStore = resolveRuntimeAuthProfileStore(agentDir);
392-
if (runtimeStore && Object.keys(runtimeStore.profiles).length > 0) {
393-
return true;
394-
}
395-
396-
if (hasStoredAuthProfileFiles(agentDir)) {
397-
return true;
398-
}
360+
export { hasAnyAuthProfileStoreSource } from "./source-check.js";
399361

400-
const authPath = resolveAuthStorePath(agentDir);
401-
const mainAuthPath = resolveAuthStorePath();
402-
if (agentDir && authPath !== mainAuthPath && hasStoredAuthProfileFiles(undefined)) {
403-
return true;
404-
}
362+
export function replaceRuntimeAuthProfileStoreSnapshots(
363+
entries: Array<{ agentDir?: string; store: AuthProfileStore }>,
364+
): void {
365+
replaceRuntimeAuthProfileStoreSnapshotsImpl(entries);
366+
}
405367

406-
return false;
368+
export function clearRuntimeAuthProfileStoreSnapshots(): void {
369+
clearRuntimeAuthProfileStoreSnapshotsImpl();
370+
loadedAuthStoreCache.clear();
407371
}
408372

409373
export function saveAuthProfileStore(
@@ -413,7 +377,6 @@ export function saveAuthProfileStore(
413377
): void {
414378
const authPath = resolveAuthStorePath(agentDir);
415379
const statePath = resolveAuthStatePath(agentDir);
416-
const runtimeKey = resolveRuntimeStoreKey(agentDir);
417380
const payload = buildPersistedAuthProfileSecretsStore(store, ({ profileId, credential }) => {
418381
if (credential.type !== "oauth") {
419382
return true;
@@ -440,7 +403,7 @@ export function saveAuthProfileStore(
440403
stateMtimeMs: readAuthStoreMtimeMs(statePath),
441404
store: runtimeStore,
442405
});
443-
if (runtimeAuthStoreSnapshots.has(runtimeKey)) {
444-
runtimeAuthStoreSnapshots.set(runtimeKey, cloneAuthProfileStore(runtimeStore));
406+
if (hasRuntimeAuthProfileStoreSnapshot(agentDir)) {
407+
setRuntimeAuthProfileStoreSnapshot(runtimeStore, agentDir);
445408
}
446409
}

‎src/cron/isolated-agent/run-auth-profile.runtime.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export { resolveSessionAuthProfileOverride } from "../../agents/auth-profiles/session-override.js";

‎src/cron/isolated-agent/run.auth-profile-cold-path.test.ts‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
2+
3+
const hasAnyAuthProfileStoreSourceMock = vi.fn(() => false);
4+
5+
vi.mock("../../agents/auth-profiles/source-check.js", () => ({
6+
hasAnyAuthProfileStoreSource: hasAnyAuthProfileStoreSourceMock,
7+
}));
8+
9+
import {
10+
clearFastTestEnv,
11+
loadRunCronIsolatedAgentTurn,
12+
resolveSessionAuthProfileOverrideMock,
13+
resetRunCronIsolatedAgentTurnHarness,
14+
restoreFastTestEnv,
15+
} from "./run.test-harness.js";
16+
17+
const runCronIsolatedAgentTurn = await loadRunCronIsolatedAgentTurn();
18+
19+
function makeParams(overrides?: Record<string, unknown>) {
20+
return {
21+
cfg: {},
22+
deps: {} as never,
23+
job: {
24+
id: "cron-auth-cold-path",
25+
name: "Auth Cold Path",
26+
schedule: { kind: "cron", expr: "0 * * * *", tz: "UTC" },
27+
sessionTarget: "isolated",
28+
payload: { kind: "agentTurn", message: "run task" },
29+
},
30+
message: "run task",
31+
sessionKey: "cron:auth-cold-path",
32+
...overrides,
33+
};
34+
}
35+
36+
describe("runCronIsolatedAgentTurn auth-profile cold path", () => {
37+
let previousFastTestEnv: string | undefined;
38+
39+
beforeEach(() => {
40+
previousFastTestEnv = clearFastTestEnv();
41+
resetRunCronIsolatedAgentTurnHarness();
42+
hasAnyAuthProfileStoreSourceMock.mockReset();
43+
hasAnyAuthProfileStoreSourceMock.mockReturnValue(false);
44+
});
45+
46+
afterEach(() => {
47+
restoreFastTestEnv(previousFastTestEnv);
48+
});
49+
50+
it("skips auth-profile override resolution when no sources exist", async () => {
51+
const result = await runCronIsolatedAgentTurn(makeParams());
52+
53+
expect(result.status).toBe("ok");
54+
expect(hasAnyAuthProfileStoreSourceMock).toHaveBeenCalledTimes(1);
55+
expect(resolveSessionAuthProfileOverrideMock).not.toHaveBeenCalled();
56+
});
57+
});

‎src/cron/isolated-agent/run.runtime.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ export {
66
resolveDefaultAgentId,
77
resolveAgentSkillsFilter,
88
} from "../../agents/agent-scope.js";
9-
export { resolveSessionAuthProfileOverride } from "../../agents/auth-profiles/session-override.js";
109
export { setCliSessionId } from "../../agents/cli-session.js";
1110
export { lookupContextTokens } from "../../agents/context.js";
1211
export { resolveCronStyleNow } from "../../agents/current-time.js";

‎src/cron/isolated-agent/run.test-harness.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,6 @@ vi.mock("./run.runtime.js", () => ({
9898
resolveAgentWorkspaceDir: vi.fn().mockReturnValue("/tmp/workspace"),
9999
resolveDefaultAgentId: vi.fn().mockReturnValue("default"),
100100
resolveAgentSkillsFilter: resolveAgentSkillsFilterMock,
101-
resolveSessionAuthProfileOverride: resolveSessionAuthProfileOverrideMock,
102101
lookupContextTokens: lookupContextTokensMock,
103102
resolveCronStyleNow: resolveCronStyleNowMock,
104103
DEFAULT_CONTEXT_TOKENS: 128000,
@@ -153,6 +152,10 @@ vi.mock("./run-execution.runtime.js", () => ({
153152
logWarn: (...args: unknown[]) => logWarnMock(...args),
154153
}));
155154

155+
vi.mock("./run-auth-profile.runtime.js", () => ({
156+
resolveSessionAuthProfileOverride: resolveSessionAuthProfileOverrideMock,
157+
}));
158+
156159
vi.mock("../../agents/cli-runner.runtime.js", () => ({
157160
setCliSessionId: vi.fn(),
158161
}));

0 commit comments

Comments
 (0)