Skip to content

Commit f0a985f

Browse files
jalehmanjalehman
committed
refactor: scope embedded attempt quota reads
1 parent 32d6026 commit f0a985f

21 files changed

Lines changed: 601 additions & 113 deletions

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,12 @@ const legacyWriterNames = new Set([
3030
]);
3131

3232
export const migratedSessionAccessorFiles = new Set([
33+
"src/agents/embedded-agent-runner/run/attempt.ts",
34+
"src/auto-reply/reply/agent-runner-helpers.ts",
35+
"src/auto-reply/reply/agent-runner.ts",
36+
"src/auto-reply/reply/commands-subagents/action-info.ts",
37+
"src/auto-reply/reply/followup-runner.ts",
38+
"src/auto-reply/reply/queue/drain.ts",
3339
"src/commands/export-trajectory.ts",
3440
"src/commands/health.ts",
3541
"src/commands/sandbox-explain.ts",
@@ -54,6 +60,7 @@ export const migratedBundledPluginSessionAccessorFiles = new Set([
5460

5561
export const migratedSessionAccessorWriteFiles = new Set([
5662
"src/agents/command/attempt-execution.shared.ts",
63+
"src/agents/command/session-store.ts",
5764
"src/agents/embedded-agent-runner/run.ts",
5865
"src/agents/embedded-agent-runner/run/attempt.ts",
5966
"src/auto-reply/reply/abort-cutoff.runtime.ts",
@@ -198,6 +205,8 @@ export async function main() {
198205
const readSourceRoots = resolveSourceRoots(repoRoot, [
199206
"extensions/discord/src/monitor",
200207
"extensions/telegram/src",
208+
"src/agents",
209+
"src/auto-reply",
201210
"src/commands",
202211
"src/config/sessions",
203212
"src/cron",

‎src/agents/command/session-store.test.ts‎

Lines changed: 168 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -183,11 +183,21 @@ describe("updateSessionStoreAfterAgentRun", () => {
183183
} as OpenClawConfig;
184184
const sessionKey = "agent:main:explicit:test-maintenance-config";
185185
const sessionId = "test-maintenance-config-session";
186+
const now = Date.now();
186187
const sessionStore: Record<string, SessionEntry> = {
187188
[sessionKey]: {
188189
sessionId,
189-
updatedAt: 1,
190-
},
190+
updatedAt: now,
191+
},
192+
...Object.fromEntries(
193+
Array.from({ length: 45 }, (_, index) => [
194+
`agent:main:stale:${index}`,
195+
{
196+
sessionId: `stale-${index}`,
197+
updatedAt: now - index - 1,
198+
} satisfies SessionEntry,
199+
]),
200+
),
191201
};
192202
await fs.writeFile(storePath, JSON.stringify(sessionStore, null, 2));
193203
const result: EmbeddedAgentRunResult = {
@@ -212,29 +222,10 @@ describe("updateSessionStoreAfterAgentRun", () => {
212222
result,
213223
});
214224

215-
// The gateway write takes cache ownership and supplies single-entry
216-
// persistence so large stores are not rewritten unnecessarily.
217-
const updateOptions = sessionStoreMocks.updateSessionStore.mock.calls.at(-1)?.[2];
218-
expect(updateOptions).toMatchObject({
219-
takeCacheOwnership: true,
220-
maintenanceConfig: {
221-
mode: "enforce",
222-
maxEntries: 42,
223-
},
224-
});
225-
expect(typeof updateOptions?.resolveSingleEntryPersistence).toBe("function");
226-
expect(
227-
updateOptions?.resolveSingleEntryPersistence?.({
228-
sessionId,
229-
updatedAt: 2,
230-
} as SessionEntry),
231-
).toEqual({
232-
sessionKey,
233-
entry: {
234-
sessionId,
235-
updatedAt: 2,
236-
},
237-
});
225+
const persisted = loadSessionStore(storePath, { skipCache: true });
226+
expect(Object.keys(persisted)).toHaveLength(42);
227+
expect(persisted[sessionKey]?.sessionId).toBe(sessionId);
228+
expect(persisted["agent:main:stale:44"]).toBeUndefined();
238229
});
239230
});
240231

@@ -1858,6 +1849,52 @@ describe("updateSessionStoreAfterAgentRun", () => {
18581849
});
18591850
});
18601851

1852+
it("does not recreate a missing persisted row while preserving user-facing state", async () => {
1853+
await withTempSessionStore(async ({ storePath }) => {
1854+
const cfg = {} as OpenClawConfig;
1855+
const sessionKey = "agent:main:explicit:missing-visible-row";
1856+
const sessionId = "missing-visible-row-session";
1857+
const sessionStore: Record<string, SessionEntry> = {
1858+
[sessionKey]: {
1859+
sessionId,
1860+
updatedAt: 1,
1861+
modelProvider: "openai",
1862+
model: "gpt-5.5",
1863+
},
1864+
};
1865+
await fs.writeFile(storePath, JSON.stringify({}, null, 2), "utf8");
1866+
1867+
await updateSessionStoreAfterAgentRun({
1868+
cfg,
1869+
sessionId,
1870+
sessionKey,
1871+
storePath,
1872+
sessionStore,
1873+
defaultProvider: "claude-cli",
1874+
defaultModel: "claude-sonnet-4-6",
1875+
preserveUserFacingSessionModelState: true,
1876+
result: {
1877+
meta: {
1878+
durationMs: 1,
1879+
agentMeta: {
1880+
sessionId,
1881+
provider: "claude-cli",
1882+
model: "claude-sonnet-4-6",
1883+
},
1884+
},
1885+
},
1886+
});
1887+
1888+
expect(sessionStore[sessionKey]).toEqual({
1889+
sessionId,
1890+
updatedAt: 1,
1891+
modelProvider: "openai",
1892+
model: "gpt-5.5",
1893+
});
1894+
expect(loadSessionStore(storePath, { skipCache: true })[sessionKey]).toBeUndefined();
1895+
});
1896+
});
1897+
18611898
it("leaves contextTokens unset when entry has prior model but no contextTokens (heartbeat bleed guard)", async () => {
18621899
await withTempSessionStore(async ({ storePath }) => {
18631900
const cfg = {} as OpenClawConfig;
@@ -2212,6 +2249,58 @@ describe("recordCliCompactionInStore", () => {
22122249
expect(persisted[sessionKey]?.sessionFile).toBe(nextSessionFile);
22132250
});
22142251
});
2252+
2253+
it("recreates a complete persisted row when the caller snapshot survived a missing store row", async () => {
2254+
await withTempSessionStore(async ({ storePath }) => {
2255+
const sessionKey = "agent:main:explicit:test-record-cli-compaction-missing-row";
2256+
const sessionId = "test-record-cli-compaction-missing-row-session";
2257+
const sessionStore: Record<string, SessionEntry> = {
2258+
[sessionKey]: {
2259+
sessionId,
2260+
updatedAt: 1,
2261+
modelProvider: "openai",
2262+
model: "gpt-5.5",
2263+
totalTokens: 12_000,
2264+
totalTokensFresh: true,
2265+
inputTokens: 9_000,
2266+
outputTokens: 100,
2267+
cacheRead: 2_900,
2268+
cacheWrite: 0,
2269+
cliSessionBindings: {
2270+
codex: {
2271+
sessionId: "stale-cli-session",
2272+
},
2273+
},
2274+
cliSessionIds: {
2275+
codex: "stale-cli-session",
2276+
},
2277+
},
2278+
};
2279+
await fs.writeFile(storePath, JSON.stringify({}, null, 2), "utf8");
2280+
2281+
await recordCliCompactionInStore({
2282+
provider: "codex",
2283+
sessionKey,
2284+
sessionStore,
2285+
storePath,
2286+
tokensAfter: 42,
2287+
});
2288+
2289+
const persisted = loadSessionStore(storePath, { skipCache: true })[sessionKey];
2290+
expect(sessionStore[sessionKey]?.sessionId).toBe(sessionId);
2291+
expect(sessionStore[sessionKey]?.modelProvider).toBe("openai");
2292+
expect(sessionStore[sessionKey]?.model).toBe("gpt-5.5");
2293+
expect(sessionStore[sessionKey]?.compactionCount).toBe(1);
2294+
expect(sessionStore[sessionKey]?.totalTokens).toBe(42);
2295+
expect(sessionStore[sessionKey]?.cliSessionBindings?.codex).toBeUndefined();
2296+
expect(persisted?.sessionId).toBe(sessionId);
2297+
expect(persisted?.modelProvider).toBe("openai");
2298+
expect(persisted?.model).toBe("gpt-5.5");
2299+
expect(persisted?.compactionCount).toBe(1);
2300+
expect(persisted?.totalTokens).toBe(42);
2301+
expect(persisted?.cliSessionBindings?.codex).toBeUndefined();
2302+
});
2303+
});
22152304
});
22162305

22172306
describe("clearCliSessionInStore", () => {
@@ -2292,4 +2381,58 @@ describe("clearCliSessionInStore", () => {
22922381
).toBe("claude-session-1");
22932382
});
22942383
});
2384+
2385+
it("clears the caller snapshot and recreates a complete persisted row when the store row is missing", async () => {
2386+
await withTempSessionStore(async ({ storePath }) => {
2387+
const sessionKey = "agent:main:explicit:test-clear-cli-missing-row";
2388+
const entry: SessionEntry = {
2389+
sessionId: "openclaw-session-1",
2390+
updatedAt: 1,
2391+
modelProvider: "anthropic",
2392+
model: "claude-opus-4-6",
2393+
cliSessionBindings: {
2394+
"claude-cli": {
2395+
sessionId: "claude-session-1",
2396+
authEpoch: "epoch-1",
2397+
},
2398+
"codex-cli": {
2399+
sessionId: "codex-session-1",
2400+
},
2401+
},
2402+
cliSessionIds: {
2403+
"claude-cli": "claude-session-1",
2404+
"codex-cli": "codex-session-1",
2405+
},
2406+
claudeCliSessionId: "claude-session-1",
2407+
};
2408+
const sessionStore: Record<string, SessionEntry> = { [sessionKey]: entry };
2409+
await fs.writeFile(storePath, JSON.stringify({}, null, 2), "utf8");
2410+
2411+
const cleared = await clearCliSessionInStore({
2412+
provider: "claude-cli",
2413+
sessionKey,
2414+
sessionStore,
2415+
storePath,
2416+
});
2417+
2418+
const persisted = loadSessionStore(storePath, { skipCache: true })[sessionKey];
2419+
expect(cleared?.sessionId).toBe("openclaw-session-1");
2420+
expect(cleared?.modelProvider).toBe("anthropic");
2421+
expect(cleared?.model).toBe("claude-opus-4-6");
2422+
expect(cleared?.cliSessionBindings?.["claude-cli"]).toBeUndefined();
2423+
expect(cleared?.cliSessionBindings?.["codex-cli"]).toEqual({
2424+
sessionId: "codex-session-1",
2425+
});
2426+
expect(cleared?.claudeCliSessionId).toBeUndefined();
2427+
expect(sessionStore[sessionKey]).toEqual(cleared);
2428+
expect(persisted?.sessionId).toBe("openclaw-session-1");
2429+
expect(persisted?.modelProvider).toBe("anthropic");
2430+
expect(persisted?.model).toBe("claude-opus-4-6");
2431+
expect(persisted?.cliSessionBindings?.["claude-cli"]).toBeUndefined();
2432+
expect(persisted?.cliSessionBindings?.["codex-cli"]).toEqual({
2433+
sessionId: "codex-session-1",
2434+
});
2435+
expect(persisted?.claudeCliSessionId).toBeUndefined();
2436+
});
2437+
});
22952438
});

‎src/agents/command/session-store.ts‎

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,13 @@ import path from "node:path";
55
import { normalizeOptionalString } from "@openclaw/normalization-core/string-coerce";
66
import {
77
canonicalizeAbsoluteSessionFilePath,
8-
mergeSessionEntry,
98
resolveSessionFilePath,
109
resolveSessionFilePathOptions,
1110
setSessionRuntimeModel,
1211
type SessionEntry,
13-
updateSessionStore,
1412
rewriteSessionFileForNewSessionId,
1513
} from "../../config/sessions.js";
16-
import { updateSessionEntry } from "../../config/sessions/session-accessor.js";
14+
import { patchSessionEntry } from "../../config/sessions/session-accessor.js";
1715
import { resolveMaintenanceConfigFromInput } from "../../config/sessions/store-maintenance.js";
1816
import type { OpenClawConfig } from "../../config/types.openclaw.js";
1917
import { resolveAgentIdFromSessionKey } from "../../routing/session-key.js";
@@ -295,21 +293,20 @@ export async function updateSessionStoreAfterAgentRun(params: {
295293
}
296294
: removeLifecycleStateFromMetadataPatch(next);
297295
const maintenanceConfig = resolveMaintenanceConfigFromInput(cfg.session?.maintenance);
298-
const persisted = await updateSessionStore(
299-
storePath,
300-
(store) => {
301-
if (preserveUserFacingRunState && !store[sessionKey]) {
302-
return undefined;
296+
const persisted = await patchSessionEntry(
297+
{
298+
storePath,
299+
sessionKey,
300+
},
301+
(_currentEntry, context) => {
302+
if (preserveUserFacingRunState && !context.existingEntry) {
303+
return null;
303304
}
304-
const merged = mergeSessionEntry(store[sessionKey], metadataPatch);
305-
store[sessionKey] = merged;
306-
return merged;
305+
return metadataPatch;
307306
},
308307
{
309-
takeCacheOwnership: true,
308+
...(preserveUserFacingRunState ? {} : { fallbackEntry: entry }),
310309
maintenanceConfig,
311-
resolveSingleEntryPersistence: (entryLocal) =>
312-
entryLocal ? { sessionKey, entry: entryLocal } : undefined,
313310
},
314311
);
315312
if (persisted) {
@@ -334,7 +331,14 @@ export async function clearCliSessionInStore(params: {
334331
clearCliSession(next, provider);
335332
next.updatedAt = Date.now();
336333

337-
const persisted = await updateSessionEntry({ storePath, sessionKey }, () => next);
334+
const persisted = await patchSessionEntry(
335+
{
336+
storePath,
337+
sessionKey,
338+
},
339+
() => next,
340+
{ fallbackEntry: entry },
341+
);
338342
if (persisted) {
339343
sessionStore[sessionKey] = persisted;
340344
}
@@ -401,7 +405,14 @@ export async function recordCliCompactionInStore(params: {
401405
next.cacheWrite = undefined;
402406
}
403407

404-
const persisted = await updateSessionEntry({ storePath, sessionKey }, () => next);
408+
const persisted = await patchSessionEntry(
409+
{
410+
storePath,
411+
sessionKey,
412+
},
413+
() => next,
414+
{ fallbackEntry: entry },
415+
);
405416
if (persisted) {
406417
sessionStore[sessionKey] = persisted;
407418
}

0 commit comments

Comments
 (0)