Skip to content

Commit 198c248

Browse files
steipetesteipete
committed
refactor: share gateway session store migration
1 parent 6464149 commit 198c248

File tree

6 files changed

+46
-85
lines changed

6 files changed

+46
-85
lines changed

‎src/gateway/server-methods/agent.ts‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,7 @@ import { performGatewaySessionReset } from "../session-reset-service.js";
5050
import {
5151
canonicalizeSpawnedByForAgent,
5252
loadSessionEntry,
53-
pruneLegacyStoreKeys,
54-
resolveGatewaySessionStoreTarget,
53+
migrateAndPruneGatewaySessionStoreKey,
5554
} from "../session-utils.js";
5655
import { formatForLog } from "../ws-log.js";
5756
import { waitForAgentJob } from "./agent-job.js";
@@ -425,18 +424,13 @@ export const agentHandlers: GatewayRequestHandlers = {
425424
const mainSessionKey = resolveAgentMainSessionKey({ cfg, agentId });
426425
if (storePath) {
427426
const persisted = await updateSessionStore(storePath, (store) => {
428-
const target = resolveGatewaySessionStoreTarget({
427+
const { primaryKey } = migrateAndPruneGatewaySessionStoreKey({
429428
cfg,
430429
key: requestedSessionKey,
431430
store,
432431
});
433-
pruneLegacyStoreKeys({
434-
store,
435-
canonicalKey: target.canonicalKey,
436-
candidates: target.storeKeys,
437-
});
438-
const merged = mergeSessionEntry(store[canonicalSessionKey], nextEntryPatch);
439-
store[canonicalSessionKey] = merged;
432+
const merged = mergeSessionEntry(store[primaryKey], nextEntryPatch);
433+
store[primaryKey] = merged;
440434
return merged;
441435
});
442436
sessionEntry = persisted;

‎src/gateway/server-methods/sessions.ts‎

Lines changed: 4 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ import {
3131
listSessionsFromStore,
3232
loadCombinedSessionStoreForGateway,
3333
loadSessionEntry,
34-
pruneLegacyStoreKeys,
34+
migrateAndPruneGatewaySessionStoreKey,
3535
readSessionPreviewItemsFromTranscript,
3636
resolveGatewaySessionStoreTarget,
3737
resolveSessionModelRef,
@@ -92,31 +92,6 @@ function rejectWebchatSessionMutation(params: {
9292
return true;
9393
}
9494

95-
function migrateAndPruneSessionStoreKey(params: {
96-
cfg: ReturnType<typeof loadConfig>;
97-
key: string;
98-
store: Record<string, SessionEntry>;
99-
}) {
100-
const target = resolveGatewaySessionStoreTarget({
101-
cfg: params.cfg,
102-
key: params.key,
103-
store: params.store,
104-
});
105-
const primaryKey = target.canonicalKey;
106-
if (!params.store[primaryKey]) {
107-
const existingKey = target.storeKeys.find((candidate) => Boolean(params.store[candidate]));
108-
if (existingKey) {
109-
params.store[primaryKey] = params.store[existingKey];
110-
}
111-
}
112-
pruneLegacyStoreKeys({
113-
store: params.store,
114-
canonicalKey: primaryKey,
115-
candidates: target.storeKeys,
116-
});
117-
return { target, primaryKey, entry: params.store[primaryKey] };
118-
}
119-
12095
export const sessionsHandlers: GatewayRequestHandlers = {
12196
"sessions.list": ({ params, respond }) => {
12297
if (!assertValidParams(params, validateSessionsListParams, "sessions.list", respond)) {
@@ -224,7 +199,7 @@ export const sessionsHandlers: GatewayRequestHandlers = {
224199

225200
const { cfg, target, storePath } = resolveGatewaySessionTargetFromKey(key);
226201
const applied = await updateSessionStore(storePath, async (store) => {
227-
const { primaryKey } = migrateAndPruneSessionStoreKey({ cfg, key, store });
202+
const { primaryKey } = migrateAndPruneGatewaySessionStoreKey({ cfg, key, store });
228203
return await applySessionsPatchToStore({
229204
cfg,
230205
store,
@@ -316,7 +291,7 @@ export const sessionsHandlers: GatewayRequestHandlers = {
316291
}
317292
const sessionId = entry?.sessionId;
318293
const deleted = await updateSessionStore(storePath, (store) => {
319-
const { primaryKey } = migrateAndPruneSessionStoreKey({ cfg, key, store });
294+
const { primaryKey } = migrateAndPruneGatewaySessionStoreKey({ cfg, key, store });
320295
const hadEntry = Boolean(store[primaryKey]);
321296
if (hadEntry) {
322297
delete store[primaryKey];
@@ -385,7 +360,7 @@ export const sessionsHandlers: GatewayRequestHandlers = {
385360
const { cfg, target, storePath } = resolveGatewaySessionTargetFromKey(key);
386361
// Lock + read in a short critical section; transcript work happens outside.
387362
const compactTarget = await updateSessionStore(storePath, (store) => {
388-
const { entry, primaryKey } = migrateAndPruneSessionStoreKey({ cfg, key, store });
363+
const { entry, primaryKey } = migrateAndPruneGatewaySessionStoreKey({ cfg, key, store });
389364
return { entry, primaryKey };
390365
});
391366
const entry = compactTarget.entry;

‎src/gateway/server-node-events.ts‎

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,7 @@ import { defaultRuntime } from "../runtime.js";
1616
import { parseMessageWithAttachments } from "./chat-attachments.js";
1717
import { normalizeRpcAttachmentsToChatAttachments } from "./server-methods/attachment-normalize.js";
1818
import type { NodeEvent, NodeEventContext } from "./server-node-events-types.js";
19-
import {
20-
loadSessionEntry,
21-
pruneLegacyStoreKeys,
22-
resolveGatewaySessionStoreTarget,
23-
} from "./session-utils.js";
19+
import { loadSessionEntry, migrateAndPruneGatewaySessionStoreKey } from "./session-utils.js";
2420
import { formatForLog } from "./ws-log.js";
2521

2622
const MAX_EXEC_EVENT_OUTPUT_CHARS = 180;
@@ -152,17 +148,12 @@ async function touchSessionStore(params: {
152148
return;
153149
}
154150
await updateSessionStore(storePath, (store) => {
155-
const target = resolveGatewaySessionStoreTarget({
151+
const { primaryKey } = migrateAndPruneGatewaySessionStoreKey({
156152
cfg: params.cfg,
157153
key: params.sessionKey,
158154
store,
159155
});
160-
pruneLegacyStoreKeys({
161-
store,
162-
canonicalKey: target.canonicalKey,
163-
candidates: target.storeKeys,
164-
});
165-
store[params.canonicalKey] = {
156+
store[primaryKey] = {
166157
sessionId: params.sessionId,
167158
updatedAt: params.now,
168159
thinkingLevel: params.entry?.thinkingLevel,

‎src/gateway/session-reset-service.ts‎

Lines changed: 6 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -25,38 +25,13 @@ import { ErrorCodes, errorShape } from "./protocol/index.js";
2525
import {
2626
archiveSessionTranscripts,
2727
loadSessionEntry,
28-
pruneLegacyStoreKeys,
28+
migrateAndPruneGatewaySessionStoreKey,
2929
resolveGatewaySessionStoreTarget,
3030
resolveSessionModelRef,
3131
} from "./session-utils.js";
3232

3333
const ACP_RUNTIME_CLEANUP_TIMEOUT_MS = 15_000;
3434

35-
function migrateAndPruneSessionStoreKey(params: {
36-
cfg: ReturnType<typeof loadConfig>;
37-
key: string;
38-
store: Record<string, SessionEntry>;
39-
}) {
40-
const target = resolveGatewaySessionStoreTarget({
41-
cfg: params.cfg,
42-
key: params.key,
43-
store: params.store,
44-
});
45-
const primaryKey = target.canonicalKey;
46-
if (!params.store[primaryKey]) {
47-
const existingKey = target.storeKeys.find((candidate) => Boolean(params.store[candidate]));
48-
if (existingKey) {
49-
params.store[primaryKey] = params.store[existingKey];
50-
}
51-
}
52-
pruneLegacyStoreKeys({
53-
store: params.store,
54-
canonicalKey: primaryKey,
55-
candidates: target.storeKeys,
56-
});
57-
return { target, primaryKey, entry: params.store[primaryKey] };
58-
}
59-
6035
function stripRuntimeModelState(entry?: SessionEntry): SessionEntry | undefined {
6136
if (!entry) {
6237
return entry;
@@ -311,7 +286,11 @@ export async function performGatewaySessionReset(params: {
311286
let oldSessionId: string | undefined;
312287
let oldSessionFile: string | undefined;
313288
const next = await updateSessionStore(storePath, (store) => {
314-
const { primaryKey } = migrateAndPruneSessionStoreKey({ cfg, key: params.key, store });
289+
const { primaryKey } = migrateAndPruneGatewaySessionStoreKey({
290+
cfg,
291+
key: params.key,
292+
store,
293+
});
315294
const currentEntry = store[primaryKey];
316295
const resetEntry = stripRuntimeModelState(currentEntry);
317296
const parsed = parseAgentSessionKey(primaryKey);

‎src/gateway/session-utils.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,31 @@ export function pruneLegacyStoreKeys(params: {
263263
}
264264
}
265265

266+
export function migrateAndPruneGatewaySessionStoreKey(params: {
267+
cfg: ReturnType<typeof loadConfig>;
268+
key: string;
269+
store: Record<string, SessionEntry>;
270+
}) {
271+
const target = resolveGatewaySessionStoreTarget({
272+
cfg: params.cfg,
273+
key: params.key,
274+
store: params.store,
275+
});
276+
const primaryKey = target.canonicalKey;
277+
if (!params.store[primaryKey]) {
278+
const existingKey = target.storeKeys.find((candidate) => Boolean(params.store[candidate]));
279+
if (existingKey) {
280+
params.store[primaryKey] = params.store[existingKey];
281+
}
282+
}
283+
pruneLegacyStoreKeys({
284+
store: params.store,
285+
canonicalKey: primaryKey,
286+
candidates: target.storeKeys,
287+
});
288+
return { target, primaryKey, entry: params.store[primaryKey] };
289+
}
290+
266291
export function classifySessionKey(key: string, entry?: SessionEntry): GatewaySessionRow["kind"] {
267292
if (key === "global") {
268293
return "global";

‎src/gateway/sessions-resolve.ts‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import {
1010
import {
1111
listSessionsFromStore,
1212
loadCombinedSessionStoreForGateway,
13-
pruneLegacyStoreKeys,
13+
migrateAndPruneGatewaySessionStoreKey,
1414
resolveGatewaySessionStoreTarget,
1515
} from "./session-utils.js";
1616

@@ -58,13 +58,10 @@ export async function resolveSessionKeyFromResolveParams(params: {
5858
};
5959
}
6060
await updateSessionStore(target.storePath, (s) => {
61-
const liveTarget = resolveGatewaySessionStoreTarget({ cfg, key, store: s });
62-
const canonicalKey = liveTarget.canonicalKey;
63-
// Migrate the first legacy entry to the canonical key.
64-
if (!s[canonicalKey] && s[legacyKey]) {
65-
s[canonicalKey] = s[legacyKey];
61+
const { primaryKey } = migrateAndPruneGatewaySessionStoreKey({ cfg, key, store: s });
62+
if (!s[primaryKey] && s[legacyKey]) {
63+
s[primaryKey] = s[legacyKey];
6664
}
67-
pruneLegacyStoreKeys({ store: s, canonicalKey, candidates: liveTarget.storeKeys });
6865
});
6966
return { ok: true, key: target.canonicalKey };
7067
}

0 commit comments

Comments
 (0)