diff --git a/src/agents/session-tool-result-guard.ts b/src/agents/session-tool-result-guard.ts index ad69300ced1b..2b0b0e8c2828 100644 --- a/src/agents/session-tool-result-guard.ts +++ b/src/agents/session-tool-result-guard.ts @@ -60,6 +60,7 @@ function resolveMaxToolResultChars(opts?: { maxToolResultChars?: number }): numb type UserAgentMessage = Extract; type CompactionAppendValidator = (entryId: string, appendedText: string) => boolean; +type AppendMessageOptions = Parameters[1]; function isUserAgentMessage(message: AgentMessage): message is UserAgentMessage { return message.role === "user"; @@ -643,6 +644,7 @@ export function installSessionToolResultGuard( const allowSyntheticToolResults = opts?.allowSyntheticToolResults ?? true; const missingToolResultText = opts?.missingToolResultText; const beforeWrite = opts?.beforeMessageWriteHook; + const toolResultTransformerMayMutate = opts?.transformToolResultForPersistence !== undefined; const redactionConfig = opts?.redactLoggingConfig; const maxToolResultChars = resolveMaxToolResultChars(opts); const transcriptSeqByEntryId: TranscriptSeqByEntryId = new Map(); @@ -653,9 +655,10 @@ export function installSessionToolResultGuard( const appendMessageAndCacheTranscriptSeq = ( message: AgentMessage, + options?: AppendMessageOptions, ): { entryId: string; messageSeq?: number; sessionFile?: string | null } => { const parentEntryId = sessionManager.getLeafId(); - const entryId = originalAppend(message as never); + const entryId = originalAppend(message as never, options); void opts?.onMessagePersisted?.(message); const sessionFile = getSessionFile(); if (!sessionFile) { @@ -686,18 +689,20 @@ export function installSessionToolResultGuard( * Run the before_message_write hook. Returns the (possibly modified) message, * or null if the message should be blocked. */ - const applyBeforeWriteHook = (msg: AgentMessage): AgentMessage | null => { + const applyBeforeWriteHook = ( + msg: AgentMessage, + ): { message: AgentMessage; changed: boolean } | null => { if (!beforeWrite) { - return msg; + return { message: msg, changed: false }; } const result = beforeWrite({ message: msg }); if (result?.block) { return null; } if (result?.message) { - return result.message; + return { message: result.message, changed: true }; } - return msg; + return { message: msg, changed: false }; }; const flushPendingToolResults = () => { @@ -711,16 +716,20 @@ export function installSessionToolResultGuard( toolName: name, text: missingToolResultText, }); - const flushed = applyBeforeWriteHook( - persistToolResult(persistMessage(synthetic), { - toolCallId: id, - toolName: name, - isSynthetic: true, - }), - ); + const persistedSynthetic = persistMessage(synthetic); + const transformed = persistToolResult(persistedSynthetic, { + toolCallId: id, + toolName: name, + isSynthetic: true, + }); + const flushed = applyBeforeWriteHook(transformed); if (flushed) { appendMessageAndCacheTranscriptSeq( - capToolResultForPersistence(flushed, maxToolResultChars, redactionConfig), + capToolResultForPersistence(flushed.message, maxToolResultChars, redactionConfig), + { + invalidateSerializedPrefixCache: + persistedSynthetic !== synthetic || toolResultTransformerMayMutate || flushed.changed, + }, ); } } @@ -732,7 +741,8 @@ export function installSessionToolResultGuard( pendingState.clear(); }; - const guardedAppend = (message: AgentMessage) => { + const guardedAppend = (message: AgentMessage, callerOptions?: AppendMessageOptions) => { + const callerInvalidatesCache = callerOptions?.invalidateSerializedPrefixCache === true; let nextMessage = message; const role = (message as { role?: unknown }).role; if (role === "assistant") { @@ -758,23 +768,30 @@ export function installSessionToolResultGuard( const normalizedToolResult = normalizePersistedToolResultName(nextMessage, toolName); // Apply hard size cap before persistence to prevent oversized tool results // from consuming the entire context window on subsequent LLM calls. + const persistedToolResult = persistMessage(normalizedToolResult); const capped = capToolResultForPersistence( - persistMessage(normalizedToolResult), + persistedToolResult, maxToolResultChars, redactionConfig, ); - const persisted = applyBeforeWriteHook( - persistToolResult(capped, { - toolCallId: id ?? undefined, - toolName, - isSynthetic: false, - }), - ); + const transformed = persistToolResult(capped, { + toolCallId: id ?? undefined, + toolName, + isSynthetic: false, + }); + const persisted = applyBeforeWriteHook(transformed); if (!persisted) { return undefined; } return appendMessageAndCacheTranscriptSeq( - capToolResultForPersistence(persisted, maxToolResultChars, redactionConfig), + capToolResultForPersistence(persisted.message, maxToolResultChars, redactionConfig), + { + invalidateSerializedPrefixCache: + callerInvalidatesCache || + persistedToolResult !== normalizedToolResult || + toolResultTransformerMayMutate || + persisted.changed, + }, ).entryId; } @@ -818,10 +835,12 @@ export function installSessionToolResultGuard( flushPendingToolResults(); } - const finalMessage = applyBeforeWriteHook(persistMessage(nextMessage)); - if (!finalMessage) { + const transformedMessage = persistMessage(nextMessage); + const finalWrite = applyBeforeWriteHook(transformedMessage); + if (!finalWrite) { return undefined; } + const finalMessage = finalWrite.message; const finalRole = (finalMessage as { role?: unknown }).role; if ( finalRole === "assistant" && @@ -845,7 +864,10 @@ export function installSessionToolResultGuard( entryId: result, messageSeq, sessionFile, - } = appendMessageAndCacheTranscriptSeq(finalMessage); + } = appendMessageAndCacheTranscriptSeq(finalMessage, { + invalidateSerializedPrefixCache: + callerInvalidatesCache || transformedMessage !== nextMessage || finalWrite.changed, + }); if (sessionFile) { emitSessionTranscriptUpdate({ sessionFile, diff --git a/src/agents/sessions/agent-session.ts b/src/agents/sessions/agent-session.ts index 45524c6a413c..44caa0a0cb79 100644 --- a/src/agents/sessions/agent-session.ts +++ b/src/agents/sessions/agent-session.ts @@ -332,6 +332,7 @@ export class AgentSession { // Branch summarization state private branchSummaryAbortController: AbortController | undefined = undefined; + private extensionModifiedToolResultIds = new Set(); // Retry state private retryAbortController: AbortController | undefined = undefined; @@ -515,6 +516,7 @@ export class AgentSession { if (!hookResult) { return undefined; } + this.extensionModifiedToolResultIds.add(toolCall.id); return { content: hookResult.content, @@ -579,7 +581,7 @@ export class AgentSession { } // Emit to extensions first - await this.emitExtensionEvent(event); + const messageChangedByExtension = await this.emitExtensionEvent(event); // Notify all listeners this.emit( @@ -605,7 +607,13 @@ export class AgentSession { event.message.role === "toolResult" ) { // Regular LLM message - persist as SessionMessageEntry - this.sessionManager.appendMessage(event.message); + const toolResultChangedByExtension = + event.message.role === "toolResult" && + this.extensionModifiedToolResultIds.delete(event.message.toolCallId); + this.sessionManager.appendMessage(event.message, { + invalidateSerializedPrefixCache: + messageChangedByExtension || toolResultChangedByExtension, + }); } // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere @@ -689,7 +697,7 @@ export class AgentSession { } /** Emit extension events based on agent events */ - private async emitExtensionEvent(event: AgentEvent): Promise { + private async emitExtensionEvent(event: AgentEvent): Promise { if (event.type === "agent_start") { this.turnIndex = 0; await this.currentExtensionRunner.emit({ type: "agent_start" }); @@ -732,6 +740,7 @@ export class AgentSession { const replacement = await this.currentExtensionRunner.emitMessageEnd(extensionEvent); if (replacement) { this.replaceMessageInPlace(event.message, replacement); + return true; } } else if (event.type === "tool_execution_start") { const extensionEvent: ToolExecutionStartEvent = { @@ -760,6 +769,7 @@ export class AgentSession { }; await this.currentExtensionRunner.emit(extensionEvent); } + return false; } /** diff --git a/src/agents/sessions/session-manager.test.ts b/src/agents/sessions/session-manager.test.ts index 321c4f2823ef..0776a9a7915f 100644 --- a/src/agents/sessions/session-manager.test.ts +++ b/src/agents/sessions/session-manager.test.ts @@ -454,6 +454,199 @@ describe("SessionManager.open", () => { } }); + it("does not probe custom entry getters before serialization", async () => { + const dir = await makeTempDir(); + const sessionFile = path.join(dir, "session.jsonl"); + const assistantEntry = { + type: "message", + id: "assistant-1", + parentId: null, + timestamp: "2026-06-04T00:00:01.000Z", + message: buildAssistantMessage("message 1"), + }; + await fs.writeFile( + sessionFile, + `${JSON.stringify(buildSessionHeader(dir))}\n${JSON.stringify(assistantEntry)}\n`, + "utf8", + ); + + let accessCount = 0; + const sessionManager = SessionManager.open(sessionFile, dir, dir); + await withOwnedSessionTranscriptWrites( + { + sessionFile, + canAdvanceSessionEntryCache: () => true, + publishSessionFileSnapshot: () => true, + withSessionWriteLock: async (run) => await run(), + }, + async () => { + sessionManager.appendCustomEntry("getter-data", { + get cursor() { + accessCount += 1; + return `value ${accessCount}`; + }, + }); + }, + ); + + const freshEntry = loadEntriesFromFile(sessionFile).find((entry) => entry.type === "custom"); + expect(accessCount).toBe(1); + expect(freshEntry).toMatchObject({ data: { cursor: "value 1" } }); + }); + + it("invalidates custom function serializers before advancing the cache", async () => { + const dir = await makeTempDir(); + const sessionFile = path.join(dir, "session.jsonl"); + const assistantEntry = { + type: "message", + id: "assistant-1", + parentId: null, + timestamp: "2026-06-04T00:00:01.000Z", + message: buildAssistantMessage("message 1"), + }; + const replacementEntry = { + ...assistantEntry, + message: buildAssistantMessage("changed 1"), + }; + const headerLine = JSON.stringify(buildSessionHeader(dir)); + await fs.writeFile(sessionFile, `${headerLine}\n${JSON.stringify(assistantEntry)}\n`, "utf8"); + + const serializer = Object.assign(function serialize() {}, { + toJSON() { + writeFileSync(sessionFile, `${headerLine}\n${JSON.stringify(replacementEntry)}\n`, "utf8"); + return "persisted"; + }, + }); + + const sessionManager = SessionManager.open(sessionFile, dir, dir); + await withOwnedSessionTranscriptWrites( + { + sessionFile, + canAdvanceSessionEntryCache: () => true, + publishSessionFileSnapshot: () => true, + withSessionWriteLock: async (run) => await run(), + }, + async () => { + sessionManager.appendCustomEntry("function-serializer", { value: serializer }); + }, + ); + + const reopenedPrefix = SessionManager.open(sessionFile, dir, dir) + .getEntries() + .find((entry) => entry.id === "assistant-1"); + expect(reopenedPrefix ? readMessageContent(reopenedPrefix) : undefined).toBe("changed 1"); + }); + + it("validates custom message detail hooks before advancing the cache", async () => { + const dir = await makeTempDir(); + const sessionFile = path.join(dir, "session.jsonl"); + const assistantEntry = { + type: "message", + id: "assistant-1", + parentId: null, + timestamp: "2026-06-04T00:00:01.000Z", + message: buildAssistantMessage("message 1"), + }; + const replacementEntry = { + ...assistantEntry, + message: buildAssistantMessage("changed 1"), + }; + const headerLine = JSON.stringify(buildSessionHeader(dir)); + await fs.writeFile(sessionFile, `${headerLine}\n${JSON.stringify(assistantEntry)}\n`, "utf8"); + + const sessionManager = SessionManager.open(sessionFile, dir, dir); + await withOwnedSessionTranscriptWrites( + { + sessionFile, + canAdvanceSessionEntryCache: () => true, + publishSessionFileSnapshot: () => true, + withSessionWriteLock: async (run) => await run(), + }, + async () => { + sessionManager.appendCustomMessageEntry("details-hook", "visible", false, { + value: { + toJSON() { + writeFileSync( + sessionFile, + `${headerLine}\n${JSON.stringify(replacementEntry)}\n`, + "utf8", + ); + return "persisted"; + }, + }, + }); + }, + ); + + expect( + SessionManager.open(sessionFile, dir, dir) + .getEntries() + .filter((entry) => entry.type === "message") + .map((entry) => readMessageContent(entry)), + ).toEqual(["changed 1"]); + }); + + it("invalidates assistant tool-call hook writes before advancing the cache", async () => { + const dir = await makeTempDir(); + const sessionFile = path.join(dir, "session.jsonl"); + const assistantEntry = { + type: "message", + id: "assistant-1", + parentId: null, + timestamp: "2026-06-04T00:00:01.000Z", + message: buildAssistantMessage("message 1"), + }; + const replacementEntry = { + ...assistantEntry, + message: buildAssistantMessage("changed 1"), + }; + const headerLine = JSON.stringify(buildSessionHeader(dir)); + await fs.writeFile(sessionFile, `${headerLine}\n${JSON.stringify(assistantEntry)}\n`, "utf8"); + + const sessionManager = SessionManager.open(sessionFile, dir, dir); + await withOwnedSessionTranscriptWrites( + { + sessionFile, + canAdvanceSessionEntryCache: () => true, + publishSessionFileSnapshot: () => true, + withSessionWriteLock: async (run) => await run(), + }, + async () => { + sessionManager.appendMessage( + { + ...buildAssistantMessage("unused"), + content: [ + { + type: "toolCall", + id: "call-1", + name: "custom", + arguments: { + value: { + toJSON() { + writeFileSync( + sessionFile, + `${headerLine}\n${JSON.stringify(replacementEntry)}\n`, + "utf8", + ); + return "persisted"; + }, + }, + }, + }, + ], + stopReason: "toolUse", + }, + { invalidateSerializedPrefixCache: true }, + ); + }, + ); + + const reopenedPrefix = SessionManager.open(sessionFile, dir, dir) + .getEntries() + .find((entry) => entry.id === "assistant-1"); + expect(reopenedPrefix ? readMessageContent(reopenedPrefix) : undefined).toBe("changed 1"); + }); + it("invalidates incremental repair when append ownership cannot be proven", async () => { const dir = await makeTempDir(); const sessionFile = path.join(dir, "session.jsonl"); diff --git a/src/agents/sessions/session-manager.ts b/src/agents/sessions/session-manager.ts index dee15b1c0120..5135e0634361 100644 --- a/src/agents/sessions/session-manager.ts +++ b/src/agents/sessions/session-manager.ts @@ -164,6 +164,10 @@ export type SessionEntry = /** Raw file entry (includes header) */ export type FileEntry = SessionHeader | SessionEntry; +type AppendPersistenceOptions = { + invalidateSerializedPrefixCache?: boolean; +}; + /** Tree node for getTree() - defensive copy of session structure */ export interface SessionTreeNode { entry: SessionEntry; @@ -562,7 +566,7 @@ function hasCacheableSessionHeader(entries: FileEntry[]): boolean { function rememberWrittenSessionEntries( filePath: string, expectedContent?: string, -): { snapshot: SessionFileSnapshot | undefined; verifiedWrite: boolean } { +): { snapshot: SessionFileSnapshot | undefined; verifiedWrite: boolean; stableRead: boolean } { const resolvedPath = resolve(filePath); // Full rewrites break append continuity used by the pre-run repair cache, // even when the filesystem preserves the inode and the rewritten file grows. @@ -572,11 +576,11 @@ function rememberWrittenSessionEntries( beforeReadSnapshot = readSessionFileSnapshot(resolvedPath); } catch { sessionEntriesCache.delete(resolvedPath); - return { snapshot: undefined, verifiedWrite: false }; + return { snapshot: undefined, verifiedWrite: false, stableRead: false }; } if (beforeReadSnapshot.size > MAX_CACHED_SESSION_BYTES) { sessionEntriesCache.delete(resolvedPath); - return { snapshot: beforeReadSnapshot, verifiedWrite: false }; + return { snapshot: beforeReadSnapshot, verifiedWrite: false, stableRead: false }; } let content: string; @@ -586,14 +590,14 @@ function rememberWrittenSessionEntries( afterReadSnapshot = readSessionFileSnapshot(resolvedPath); } catch { sessionEntriesCache.delete(resolvedPath); - return { snapshot: undefined, verifiedWrite: false }; + return { snapshot: undefined, verifiedWrite: false, stableRead: false }; } if ( (expectedContent !== undefined && content !== expectedContent) || !isSameSessionFileSnapshot(beforeReadSnapshot, afterReadSnapshot) ) { sessionEntriesCache.delete(resolvedPath); - return { snapshot: afterReadSnapshot, verifiedWrite: false }; + return { snapshot: afterReadSnapshot, verifiedWrite: false, stableRead: false }; } rememberSessionEntries( resolvedPath, @@ -604,6 +608,7 @@ function rememberWrittenSessionEntries( return { snapshot: afterReadSnapshot, verifiedWrite: expectedContent !== undefined, + stableRead: true, }; } @@ -614,6 +619,7 @@ function rememberAppendedSessionEntry( serializedAppend: string, cacheOwnedAppend: boolean, publishOwnedAppend: boolean, + invalidateSerializedPrefixCache: boolean, ): { snapshot: SessionFileSnapshot | undefined; cacheAdvanced: boolean; @@ -656,9 +662,9 @@ function rememberAppendedSessionEntry( const cached = sessionEntriesCache.get(resolvedPath); const snapshot = readSessionFileSnapshotIfExists(resolvedPath); - // Owned transcript writes serialize appenders under the session lock. Full - // rewrites refresh the cache explicitly, so identity and size are sufficient - // to advance this append-only snapshot without reading the whole file. + // Owned transcript writes serialize appenders under the session lock. Plain + // appends can advance by stat identity/size; extension-owned message payloads + // may run JSON hooks that rewrite same-size prefix bytes, so they drop cache. const expectedSize = beforeAppendSnapshot.size + appendedByteLength; if ( !snapshot || @@ -674,6 +680,11 @@ function rememberAppendedSessionEntry( invalidateSessionFileRepairCache(resolvedPath); return { snapshot, cacheAdvanced: false, ownedAppendVerified: false }; } + if (invalidateSerializedPrefixCache) { + sessionEntriesCache.delete(resolvedPath); + invalidateSessionFileRepairCache(resolvedPath); + return { snapshot, cacheAdvanced: false, ownedAppendVerified: true }; + } if (snapshot.size > MAX_CACHED_SESSION_BYTES) { sessionEntriesCache.delete(resolvedPath); return { snapshot, cacheAdvanced: false, ownedAppendVerified: true }; @@ -1326,7 +1337,7 @@ export class SessionManager { return this.sessionFile; } - persist(entry: SessionEntry): void { + persist(entry: SessionEntry, options?: AppendPersistenceOptions): void { if (!this.shouldPersist || !this.sessionFile) { return; } @@ -1350,8 +1361,8 @@ export class SessionManager { } } else { // Serialize before taking the prefix snapshot. Custom toJSON methods are - // user code and can mutate the transcript; the cache must validate the - // prefix state that immediately precedes the exact bytes being appended. + // user code and can mutate the transcript; extension-owned writes opt out + // of warm-cache advancement so same-size prefix rewrites cannot publish. const serializedEntry = serializeJsonlEntry(entry); const beforeAppendSnapshot = readSessionFileSnapshotIfExists(this.sessionFile); const canPublishOwnedAppend = Boolean( @@ -1374,6 +1385,7 @@ export class SessionManager { serializedAppend, cacheOwnedAppend, canPublishOwnedAppend, + options?.invalidateSerializedPrefixCache === true, ); this.sessionFileSnapshot = rememberedAppend.snapshot; if (rememberedAppend.ownedAppendVerified) { @@ -1401,11 +1413,11 @@ export class SessionManager { } } - private appendEntry(entry: SessionEntry): void { + private appendEntry(entry: SessionEntry, options?: AppendPersistenceOptions): void { this.fileEntries.push(entry); this.byId.set(entry.id, entry); this.leafId = entry.id; - this.persist(entry); + this.persist(entry, options); } /** Append a message as child of current leaf, then advance leaf. Returns entry id. @@ -1414,7 +1426,10 @@ export class SessionManager { * so it is easier to find them. * These need to be appended via appendCompaction() and appendBranchSummary() methods. */ - appendMessage(message: Message | CustomMessage | BashExecutionMessage): string { + appendMessage( + message: Message | CustomMessage | BashExecutionMessage, + options?: AppendPersistenceOptions, + ): string { const entry: SessionMessageEntry = { type: "message", id: generateId(this.byId), @@ -1422,7 +1437,7 @@ export class SessionManager { timestamp: new Date().toISOString(), message, }; - this.appendEntry(entry); + this.appendEntry(entry, options); return entry.id; } @@ -1472,7 +1487,9 @@ export class SessionManager { details, fromHook, }; - this.appendEntry(entry); + this.appendEntry(entry, { + invalidateSerializedPrefixCache: fromHook === true || details !== undefined, + }); return entry.id; } @@ -1486,7 +1503,7 @@ export class SessionManager { parentId: this.leafId, timestamp: new Date().toISOString(), }; - this.appendEntry(entry); + this.appendEntry(entry, { invalidateSerializedPrefixCache: true }); return entry.id; } @@ -1541,7 +1558,7 @@ export class SessionManager { parentId: this.leafId, timestamp: new Date().toISOString(), }; - this.appendEntry(entry); + this.appendEntry(entry, { invalidateSerializedPrefixCache: true }); return entry.id; } @@ -1748,7 +1765,9 @@ export class SessionManager { details, fromHook, }; - this.appendEntry(entry); + this.appendEntry(entry, { + invalidateSerializedPrefixCache: fromHook === true || details !== undefined, + }); return entry.id; }