Skip to content

Commit 6f2869c

Browse files
jalehmanjalehman
authored
refactor: migrate agent session accessors (#96182)
* refactor: migrate agent session accessor writes * refactor: move subagent orphan lookup to reconciliation * test: align session accessor mocks
1 parent 4f5e25a commit 6f2869c

13 files changed

Lines changed: 309 additions & 177 deletions

‎scripts/check-session-accessor-boundary.mjs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,12 +76,15 @@ export const migratedSessionAccessorFiles = new Set([
7676
"packages/memory-host-sdk/src/host/session-files.ts",
7777
"src/acp/runtime/session-meta.ts",
7878
"src/agents/acp-spawn.ts",
79+
"src/agents/auth-profiles/session-override.ts",
7980
"src/agents/embedded-agent-runner/compaction-successor-transcript.ts",
8081
"src/agents/embedded-agent-runner/run/attempt.ts",
8182
"src/agents/embedded-agent-runner/tool-result-truncation.ts",
8283
"src/agents/embedded-agent-runner/transcript-rewrite.ts",
8384
"src/agents/embedded-agent-runner/transcript-runtime-state.ts",
8485
"src/auto-reply/reply/abort.ts",
86+
"src/agents/subagent-control.ts",
87+
"src/agents/subagent-registry-helpers.ts",
8588
"src/auto-reply/reply/agent-runner-helpers.ts",
8689
"src/auto-reply/reply/agent-runner.ts",
8790
"src/auto-reply/reply/commands-subagents/action-info.ts",
@@ -127,12 +130,16 @@ export const migratedBundledPluginSessionAccessorFiles = new Set([
127130

128131
export const migratedSessionAccessorWriteFiles = new Set([
129132
"src/acp/runtime/session-meta.ts",
133+
"src/agents/auth-profiles/session-override.ts",
130134
"src/agents/command/attempt-execution.shared.ts",
131135
"src/agents/command/session-store.ts",
132136
"src/agents/embedded-agent-runner/run.ts",
133137
"src/agents/embedded-agent-runner/run/attempt.ts",
138+
"src/agents/live-model-switch.ts",
134139
"src/agents/main-session-restart-recovery.ts",
135140
"src/auto-reply/reply/abort.ts",
141+
"src/agents/subagent-control.ts",
142+
"src/agents/subagent-registry-helpers.ts",
136143
"src/auto-reply/reply/abort-cutoff.runtime.ts",
137144
"src/auto-reply/reply/agent-runner-cli-dispatch.ts",
138145
"src/auto-reply/reply/agent-runner-execution.ts",

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

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,14 @@ import {
1414
import { ensureAuthProfileStore, hasAnyAuthProfileStoreSource } from "../auth-profiles/store.js";
1515
import { isProfileInCooldown } from "../auth-profiles/usage.js";
1616

17-
const sessionStoreRuntimeLoader = createLazyImportLoader(
18-
() => import("../../config/sessions/store.runtime.js"),
17+
const sessionAccessorLoader = createLazyImportLoader(
18+
() => import("../../config/sessions/session-accessor.js"),
1919
);
2020

21-
// Session-store writes are lazy-loaded so read-only auth resolution paths do not
22-
// import persistence code unless an override must be updated.
23-
function loadSessionStoreRuntime() {
24-
return sessionStoreRuntimeLoader.load();
21+
// Session accessor writes are lazy-loaded so read-only auth resolution paths do
22+
// not import persistence code unless an override must be updated.
23+
function loadSessionAccessor() {
24+
return sessionAccessorLoader.load();
2525
}
2626

2727
// Current session overrides are only valid when the selected provider can use
@@ -83,10 +83,12 @@ export async function clearSessionAuthProfileOverride(params: {
8383
sessionStore[sessionKey] = sessionEntry;
8484
if (storePath) {
8585
await (
86-
await loadSessionStoreRuntime()
87-
).updateSessionStore(storePath, (store) => {
88-
store[sessionKey] = sessionEntry;
89-
});
86+
await loadSessionAccessor()
87+
).patchSessionEntry(
88+
{ storePath, sessionKey },
89+
() => sessionEntry,
90+
{ fallbackEntry: sessionEntry, replaceEntry: true },
91+
);
9092
}
9193
}
9294

@@ -236,10 +238,12 @@ export async function resolveSessionAuthProfileOverride(params: {
236238
sessionStore[sessionKey] = sessionEntry;
237239
if (storePath) {
238240
await (
239-
await loadSessionStoreRuntime()
240-
).updateSessionStore(storePath, (storeLocal) => {
241-
storeLocal[sessionKey] = sessionEntry;
242-
});
241+
await loadSessionAccessor()
242+
).patchSessionEntry(
243+
{ storePath, sessionKey },
244+
() => sessionEntry,
245+
{ fallbackEntry: sessionEntry, replaceEntry: true },
246+
);
243247
}
244248
}
245249

‎src/agents/live-model-switch.test.ts‎

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,14 @@ vi.mock("../config/sessions/store.js", () => ({
3232
updateSessionStore: (...args: unknown[]) => state.updateSessionStoreMock(...args),
3333
}));
3434

35+
vi.mock("../config/sessions/session-accessor.js", () => ({
36+
loadSessionEntry: (scope: { sessionKey: string }) => {
37+
const store = state.loadSessionStoreMock(scope) as Record<string, unknown> | undefined;
38+
return store?.[scope.sessionKey];
39+
},
40+
patchSessionEntry: (...args: unknown[]) => state.updateSessionStoreMock(...args),
41+
}));
42+
3543
vi.mock("../config/sessions/paths.js", () => ({
3644
resolveStorePath: (...args: unknown[]) => state.resolveStorePathMock(...args),
3745
}));
@@ -126,9 +134,29 @@ describe("live model switch", () => {
126134
state.updateSessionStoreMock
127135
.mockReset()
128136
.mockImplementation(
129-
async (_path: string, updater: (store: Record<string, unknown>) => void) => {
130-
const store: Record<string, unknown> = {};
131-
updater(store);
137+
async (
138+
scope: { sessionKey: string },
139+
updater: (
140+
entry: Record<string, unknown>,
141+
) => Promise<Record<string, unknown> | null> | Record<string, unknown> | null,
142+
) => {
143+
const store = state.loadSessionStoreMock(scope) as Record<
144+
string,
145+
Record<string, unknown>
146+
>;
147+
const entry = store?.[scope.sessionKey];
148+
if (!entry) {
149+
return null;
150+
}
151+
const next = await updater(entry);
152+
if (!next) {
153+
return entry;
154+
}
155+
for (const key of Object.keys(entry)) {
156+
delete entry[key];
157+
}
158+
Object.assign(entry, next);
159+
return entry;
132160
},
133161
);
134162
});
@@ -473,12 +501,6 @@ describe("live model switch", () => {
473501
modelOverride: "claude-opus-4-6",
474502
};
475503
state.loadSessionStoreMock.mockReturnValue({ main: sessionEntry });
476-
state.updateSessionStoreMock.mockImplementation(
477-
async (_path: string, updater: (store: Record<string, unknown>) => void) => {
478-
const store: Record<string, typeof sessionEntry> = { main: sessionEntry };
479-
updater(store);
480-
},
481-
);
482504

483505
const { shouldSwitchToLiveModel } = await loadModule();
484506

@@ -538,12 +560,7 @@ describe("live model switch", () => {
538560

539561
it("deletes liveModelSwitchPending from the session entry", async () => {
540562
const sessionEntry = { liveModelSwitchPending: true, sessionId: "s-1" };
541-
state.updateSessionStoreMock.mockImplementation(
542-
async (_path: string, updater: (store: Record<string, unknown>) => void) => {
543-
const store: Record<string, typeof sessionEntry> = { main: sessionEntry };
544-
updater(store);
545-
},
546-
);
563+
state.loadSessionStoreMock.mockReturnValue({ main: sessionEntry });
547564

548565
const { clearLiveModelSwitchPending } = await loadModule();
549566

‎src/agents/live-model-switch.ts‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ import { normalizeProviderId } from "@openclaw/model-catalog-core/provider-id";
55
import { normalizeOptionalString } from "@openclaw/normalization-core/string-coerce";
66

77
import { resolveStorePath } from "../config/sessions/paths.js";
8-
import { loadSessionStore, updateSessionStore } from "../config/sessions/store.js";
8+
import { patchSessionEntry } from "../config/sessions/session-accessor.js";
9+
import { loadSessionStore } from "../config/sessions/store.js";
910
import {
1011
normalizeStoredOverrideModel,
1112
resolveDefaultModelForAgent,
@@ -211,10 +212,13 @@ export async function clearLiveModelSwitchPending(params: {
211212
if (!storePath) {
212213
return;
213214
}
214-
await updateSessionStore(storePath, (store) => {
215-
const entry = store[sessionKey];
216-
if (entry) {
217-
delete entry.liveModelSwitchPending;
218-
}
219-
});
215+
await patchSessionEntry(
216+
{ storePath, sessionKey },
217+
(entry) => {
218+
const next = { ...entry };
219+
delete next.liveModelSwitchPending;
220+
return next;
221+
},
222+
{ replaceEntry: true },
223+
);
220224
}

‎src/agents/subagent-control.test.ts‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@ import fs from "node:fs";
44
import os from "node:os";
55
import path from "node:path";
66
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
7+
import type {
8+
SessionAccessScope,
9+
SessionEntryPatchContext,
10+
SessionEntryPatchOptions,
11+
} from "../config/sessions/session-accessor.js";
712
import type { SessionEntry } from "../config/sessions/types.js";
813
import type { OpenClawConfig } from "../config/types.openclaw.js";
914
import type { CallGatewayOptions } from "../gateway/call.js";
@@ -120,14 +125,33 @@ function setSubagentControlDepsForTest(
120125
testing.setDepsForTest({
121126
abortEmbeddedAgentRun: () => false,
122127
clearSessionQueues: () => ({ followupCleared: 0, laneCleared: 0, keys: [] }),
123-
updateSessionStore: async <T>(
124-
storePath: string,
125-
mutator: (store: Record<string, SessionEntry>) => Promise<T> | T,
128+
patchSessionEntry: async (
129+
scope: SessionAccessScope,
130+
patcher: (
131+
entry: SessionEntry,
132+
context: SessionEntryPatchContext,
133+
) => Promise<Partial<SessionEntry> | null> | Partial<SessionEntry> | null,
134+
options: SessionEntryPatchOptions = {},
126135
) => {
127-
const store = JSON.parse(fs.readFileSync(storePath, "utf-8")) as Record<string, SessionEntry>;
128-
const result = await mutator(store);
129-
fs.writeFileSync(storePath, JSON.stringify(store, null, 2), "utf-8");
130-
return result;
136+
if (!scope.storePath) {
137+
return null;
138+
}
139+
const store = JSON.parse(fs.readFileSync(scope.storePath, "utf-8")) as Record<
140+
string,
141+
SessionEntry
142+
>;
143+
const entry = store[scope.sessionKey];
144+
if (!entry) {
145+
return null;
146+
}
147+
const patch = await patcher(entry, { existingEntry: { ...entry } });
148+
if (!patch) {
149+
return entry;
150+
}
151+
const next = options.replaceEntry ? (patch as SessionEntry) : { ...entry, ...patch };
152+
store[scope.sessionKey] = next;
153+
fs.writeFileSync(scope.storePath, JSON.stringify(store, null, 2), "utf-8");
154+
return next;
131155
},
132156
...overrides,
133157
});
@@ -639,7 +663,7 @@ describe("killSubagentRunAdmin", () => {
639663
});
640664

641665
setSubagentControlDepsForTest({
642-
updateSessionStore: async () => {
666+
patchSessionEntry: async () => {
643667
throw new Error("session store unavailable");
644668
},
645669
});

‎src/agents/subagent-control.ts‎

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import crypto from "node:crypto";
77
import type { ClearSessionQueueResult } from "../auto-reply/reply/queue.js";
88
import { resolveSubagentLabel, sortSubagentRuns } from "../auto-reply/reply/subagents-utils.js";
99
import { resolveStorePath } from "../config/sessions/paths.js";
10-
import { loadSessionStore, updateSessionStore } from "../config/sessions/store.js";
10+
import { loadSessionEntry, patchSessionEntry } from "../config/sessions/session-accessor.js";
1111
import type { SessionEntry } from "../config/sessions/types.js";
1212
import type { OpenClawConfig } from "../config/types.openclaw.js";
1313
import { callGateway } from "../gateway/call.js";
@@ -50,18 +50,18 @@ const SUBAGENT_REPLY_HISTORY_LIMIT = 50;
5050
const steerRateLimit = new Map<string, number>();
5151

5252
type GatewayCaller = typeof callGateway;
53-
type UpdateSessionStore = typeof updateSessionStore;
53+
type PatchSessionEntry = typeof patchSessionEntry;
5454
type AbortEmbeddedAgentRun = (sessionId: string) => boolean;
5555
type ClearSessionQueues = (keys: Array<string | undefined>) => ClearSessionQueueResult;
5656

5757
const defaultSubagentControlDeps = {
5858
callGateway,
59-
updateSessionStore,
59+
patchSessionEntry,
6060
};
6161

6262
let subagentControlDeps: {
6363
callGateway: GatewayCaller;
64-
updateSessionStore: UpdateSessionStore;
64+
patchSessionEntry: PatchSessionEntry;
6565
abortEmbeddedAgentRun?: AbortEmbeddedAgentRun;
6666
clearSessionQueues?: ClearSessionQueues;
6767
} = defaultSubagentControlDeps;
@@ -187,15 +187,15 @@ async function killSubagentRun(params: {
187187
}
188188
if (resolved.entry) {
189189
try {
190-
await subagentControlDeps.updateSessionStore(resolved.storePath, (store) => {
191-
const current = store[childSessionKey];
192-
if (!current) {
193-
return;
194-
}
195-
current.abortedLastRun = true;
196-
current.updatedAt = Date.now();
197-
store[childSessionKey] = current;
198-
});
190+
await subagentControlDeps.patchSessionEntry(
191+
{ storePath: resolved.storePath, sessionKey: childSessionKey },
192+
(current) => ({
193+
...current,
194+
abortedLastRun: true,
195+
updatedAt: Date.now(),
196+
}),
197+
{ replaceEntry: true },
198+
);
199199
} catch (error) {
200200
logVerbose(
201201
`subagents control kill: failed to persist abortedLastRun for ${childSessionKey}: ${formatErrorMessage(error)}`,
@@ -660,8 +660,11 @@ export async function sendControlledSubagentMessage(params: {
660660
const targetSessionKey = params.entry.childSessionKey;
661661
const parsed = parseAgentSessionKey(targetSessionKey);
662662
const storePath = resolveStorePath(params.cfg.session?.store, { agentId: parsed?.agentId });
663-
const store = loadSessionStore(storePath);
664-
const targetSessionEntry = store[targetSessionKey];
663+
const targetSessionEntry = loadSessionEntry({
664+
storePath,
665+
sessionKey: targetSessionKey,
666+
clone: false,
667+
});
665668
const targetSessionId =
666669
typeof targetSessionEntry?.sessionId === "string" && targetSessionEntry.sessionId.trim()
667670
? targetSessionEntry.sessionId.trim()
@@ -724,7 +727,7 @@ export const testing = {
724727
setDepsForTest(
725728
overrides?: Partial<{
726729
callGateway: GatewayCaller;
727-
updateSessionStore: UpdateSessionStore;
730+
patchSessionEntry: PatchSessionEntry;
728731
abortEmbeddedAgentRun: AbortEmbeddedAgentRun;
729732
clearSessionQueues: ClearSessionQueues;
730733
}>,

0 commit comments

Comments
 (0)