refactor: harden session store updates

Co-authored-by: Tyler Yust <tyler6204@users.noreply.github.com>
This commit is contained in:
Peter Steinberger
2026-01-15 23:06:42 +00:00
parent 35492f8513
commit 688a0ce439
24 changed files with 441 additions and 215 deletions

View File

@@ -6,7 +6,7 @@ import {
resolveAgentIdFromSessionKey,
resolveAgentMainSessionKey,
type SessionEntry,
saveSessionStore,
updateSessionStore,
} from "../../config/sessions.js";
import { registerAgentRunContext } from "../../infra/agent-events.js";
import { resolveOutboundTarget } from "../../infra/outbound/targets.js";
@@ -136,7 +136,7 @@ export const agentHandlers: GatewayRequestHandlers = {
let cfgForAgent: ReturnType<typeof loadConfig> | undefined;
if (requestedSessionKey) {
const { cfg, storePath, store, entry, canonicalKey } = loadSessionEntry(requestedSessionKey);
const { cfg, storePath, entry, canonicalKey } = loadSessionEntry(requestedSessionKey);
cfgForAgent = cfg;
const now = Date.now();
const sessionId = entry?.sessionId ?? randomUUID();
@@ -178,11 +178,10 @@ export const agentHandlers: GatewayRequestHandlers = {
const canonicalSessionKey = canonicalKey;
const agentId = resolveAgentIdFromSessionKey(canonicalSessionKey);
const mainSessionKey = resolveAgentMainSessionKey({ cfg, agentId });
if (store) {
store[canonicalSessionKey] = nextEntry;
if (storePath) {
await saveSessionStore(storePath, store);
}
if (storePath) {
await updateSessionStore(storePath, (store) => {
store[canonicalSessionKey] = nextEntry;
});
}
if (canonicalSessionKey === mainSessionKey || canonicalSessionKey === "global") {
context.addChatRun(idem, {

View File

@@ -9,10 +9,9 @@ import {
} from "../../agents/pi-embedded.js";
import { loadConfig } from "../../config/config.js";
import {
loadSessionStore,
resolveMainSessionKey,
type SessionEntry,
saveSessionStore,
updateSessionStore,
} from "../../config/sessions.js";
import { clearCommandLane } from "../../process/command-queue.js";
import {
@@ -30,6 +29,7 @@ import {
archiveFileOnDisk,
listSessionsFromStore,
loadCombinedSessionStoreForGateway,
loadSessionEntry,
resolveGatewaySessionStoreTarget,
resolveSessionTranscriptCandidates,
type SessionsPatchResult,
@@ -106,26 +106,25 @@ export const sessionsHandlers: GatewayRequestHandlers = {
const cfg = loadConfig();
const target = resolveGatewaySessionStoreTarget({ cfg, key });
const storePath = target.storePath;
const store = loadSessionStore(storePath);
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
const applied = await applySessionsPatchToStore({
cfg,
store,
storeKey: primaryKey,
patch: p,
loadGatewayModelCatalog: context.loadGatewayModelCatalog,
const applied = await updateSessionStore(storePath, async (store) => {
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
return await applySessionsPatchToStore({
cfg,
store,
storeKey: primaryKey,
patch: p,
loadGatewayModelCatalog: context.loadGatewayModelCatalog,
});
});
if (!applied.ok) {
respond(false, undefined, applied.error);
return;
}
await saveSessionStore(storePath, store);
const result: SessionsPatchResult = {
ok: true,
path: storePath,
@@ -156,34 +155,35 @@ export const sessionsHandlers: GatewayRequestHandlers = {
const cfg = loadConfig();
const target = resolveGatewaySessionStoreTarget({ cfg, key });
const storePath = target.storePath;
const store = loadSessionStore(storePath);
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
const entry = store[primaryKey];
const now = Date.now();
const next: SessionEntry = {
sessionId: randomUUID(),
updatedAt: now,
systemSent: false,
abortedLastRun: false,
thinkingLevel: entry?.thinkingLevel,
verboseLevel: entry?.verboseLevel,
reasoningLevel: entry?.reasoningLevel,
responseUsage: entry?.responseUsage,
model: entry?.model,
contextTokens: entry?.contextTokens,
sendPolicy: entry?.sendPolicy,
label: entry?.label,
lastChannel: entry?.lastChannel,
lastTo: entry?.lastTo,
skillsSnapshot: entry?.skillsSnapshot,
};
store[primaryKey] = next;
await saveSessionStore(storePath, store);
const next = await updateSessionStore(storePath, (store) => {
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
const entry = store[primaryKey];
const now = Date.now();
const nextEntry: SessionEntry = {
sessionId: randomUUID(),
updatedAt: now,
systemSent: false,
abortedLastRun: false,
thinkingLevel: entry?.thinkingLevel,
verboseLevel: entry?.verboseLevel,
reasoningLevel: entry?.reasoningLevel,
responseUsage: entry?.responseUsage,
model: entry?.model,
contextTokens: entry?.contextTokens,
sendPolicy: entry?.sendPolicy,
label: entry?.label,
lastChannel: entry?.lastChannel,
lastTo: entry?.lastTo,
skillsSnapshot: entry?.skillsSnapshot,
};
store[primaryKey] = nextEntry;
return nextEntry;
});
respond(true, { ok: true, key: target.canonicalKey, entry: next }, undefined);
},
"sessions.delete": async ({ params, respond }) => {
@@ -220,14 +220,7 @@ export const sessionsHandlers: GatewayRequestHandlers = {
const deleteTranscript = typeof p.deleteTranscript === "boolean" ? p.deleteTranscript : true;
const storePath = target.storePath;
const store = loadSessionStore(storePath);
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
const entry = store[primaryKey];
const { entry } = loadSessionEntry(key);
const sessionId = entry?.sessionId;
const existed = Boolean(entry);
clearCommandLane(resolveEmbeddedSessionLane(target.canonicalKey));
@@ -246,8 +239,15 @@ export const sessionsHandlers: GatewayRequestHandlers = {
return;
}
}
if (existed) delete store[primaryKey];
await saveSessionStore(storePath, store);
await updateSessionStore(storePath, (store) => {
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
if (store[primaryKey]) delete store[primaryKey];
});
const archived: string[] = [];
if (deleteTranscript && sessionId) {
@@ -295,14 +295,17 @@ export const sessionsHandlers: GatewayRequestHandlers = {
const cfg = loadConfig();
const target = resolveGatewaySessionStoreTarget({ cfg, key });
const storePath = target.storePath;
const store = loadSessionStore(storePath);
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
const entry = store[primaryKey];
// Lock + read in a short critical section; transcript work happens outside.
const compactTarget = await updateSessionStore(storePath, (store) => {
const primaryKey = target.storeKeys[0] ?? key;
const existingKey = target.storeKeys.find((candidate) => store[candidate]);
if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
store[primaryKey] = store[existingKey];
delete store[existingKey];
}
return { entry: store[primaryKey], primaryKey };
});
const entry = compactTarget.entry;
const sessionId = entry?.sessionId;
if (!sessionId) {
respond(
@@ -358,13 +361,15 @@ export const sessionsHandlers: GatewayRequestHandlers = {
const keptLines = lines.slice(-maxLines);
fs.writeFileSync(filePath, `${keptLines.join("\n")}\n`, "utf-8");
if (store[primaryKey]) {
delete store[primaryKey].inputTokens;
delete store[primaryKey].outputTokens;
delete store[primaryKey].totalTokens;
store[primaryKey].updatedAt = Date.now();
await saveSessionStore(storePath, store);
}
await updateSessionStore(storePath, (store) => {
const entryKey = compactTarget.primaryKey;
const entryToUpdate = store[entryKey];
if (!entryToUpdate) return;
delete entryToUpdate.inputTokens;
delete entryToUpdate.outputTokens;
delete entryToUpdate.totalTokens;
entryToUpdate.updatedAt = Date.now();
});
respond(
true,