Skip to content

Commit 9dbf8f7

Browse files
vincentkocvincentkoc
committed
fix(sessions): guard append cache after extension serialization
1 parent fd806ad commit 9dbf8f7

4 files changed

Lines changed: 292 additions & 48 deletions

File tree

‎src/agents/session-tool-result-guard.ts‎

Lines changed: 48 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ function resolveMaxToolResultChars(opts?: { maxToolResultChars?: number }): numb
6060

6161
type UserAgentMessage = Extract<AgentMessage, { role: "user" }>;
6262
type CompactionAppendValidator = (entryId: string, appendedText: string) => boolean;
63+
type AppendMessageOptions = Parameters<SessionManager["appendMessage"]>[1];
6364

6465
function isUserAgentMessage(message: AgentMessage): message is UserAgentMessage {
6566
return message.role === "user";
@@ -643,6 +644,7 @@ export function installSessionToolResultGuard(
643644
const allowSyntheticToolResults = opts?.allowSyntheticToolResults ?? true;
644645
const missingToolResultText = opts?.missingToolResultText;
645646
const beforeWrite = opts?.beforeMessageWriteHook;
647+
const toolResultTransformerMayMutate = opts?.transformToolResultForPersistence !== undefined;
646648
const redactionConfig = opts?.redactLoggingConfig;
647649
const maxToolResultChars = resolveMaxToolResultChars(opts);
648650
const transcriptSeqByEntryId: TranscriptSeqByEntryId = new Map();
@@ -653,9 +655,10 @@ export function installSessionToolResultGuard(
653655

654656
const appendMessageAndCacheTranscriptSeq = (
655657
message: AgentMessage,
658+
options?: AppendMessageOptions,
656659
): { entryId: string; messageSeq?: number; sessionFile?: string | null } => {
657660
const parentEntryId = sessionManager.getLeafId();
658-
const entryId = originalAppend(message as never);
661+
const entryId = originalAppend(message as never, options);
659662
void opts?.onMessagePersisted?.(message);
660663
const sessionFile = getSessionFile();
661664
if (!sessionFile) {
@@ -686,18 +689,20 @@ export function installSessionToolResultGuard(
686689
* Run the before_message_write hook. Returns the (possibly modified) message,
687690
* or null if the message should be blocked.
688691
*/
689-
const applyBeforeWriteHook = (msg: AgentMessage): AgentMessage | null => {
692+
const applyBeforeWriteHook = (
693+
msg: AgentMessage,
694+
): { message: AgentMessage; changed: boolean } | null => {
690695
if (!beforeWrite) {
691-
return msg;
696+
return { message: msg, changed: false };
692697
}
693698
const result = beforeWrite({ message: msg });
694699
if (result?.block) {
695700
return null;
696701
}
697702
if (result?.message) {
698-
return result.message;
703+
return { message: result.message, changed: true };
699704
}
700-
return msg;
705+
return { message: msg, changed: false };
701706
};
702707

703708
const flushPendingToolResults = () => {
@@ -711,16 +716,20 @@ export function installSessionToolResultGuard(
711716
toolName: name,
712717
text: missingToolResultText,
713718
});
714-
const flushed = applyBeforeWriteHook(
715-
persistToolResult(persistMessage(synthetic), {
716-
toolCallId: id,
717-
toolName: name,
718-
isSynthetic: true,
719-
}),
720-
);
719+
const persistedSynthetic = persistMessage(synthetic);
720+
const transformed = persistToolResult(persistedSynthetic, {
721+
toolCallId: id,
722+
toolName: name,
723+
isSynthetic: true,
724+
});
725+
const flushed = applyBeforeWriteHook(transformed);
721726
if (flushed) {
722727
appendMessageAndCacheTranscriptSeq(
723-
capToolResultForPersistence(flushed, maxToolResultChars, redactionConfig),
728+
capToolResultForPersistence(flushed.message, maxToolResultChars, redactionConfig),
729+
{
730+
invalidateSerializedPrefixCache:
731+
persistedSynthetic !== synthetic || toolResultTransformerMayMutate || flushed.changed,
732+
},
724733
);
725734
}
726735
}
@@ -732,7 +741,8 @@ export function installSessionToolResultGuard(
732741
pendingState.clear();
733742
};
734743

735-
const guardedAppend = (message: AgentMessage) => {
744+
const guardedAppend = (message: AgentMessage, callerOptions?: AppendMessageOptions) => {
745+
const callerInvalidatesCache = callerOptions?.invalidateSerializedPrefixCache === true;
736746
let nextMessage = message;
737747
const role = (message as { role?: unknown }).role;
738748
if (role === "assistant") {
@@ -758,23 +768,30 @@ export function installSessionToolResultGuard(
758768
const normalizedToolResult = normalizePersistedToolResultName(nextMessage, toolName);
759769
// Apply hard size cap before persistence to prevent oversized tool results
760770
// from consuming the entire context window on subsequent LLM calls.
771+
const persistedToolResult = persistMessage(normalizedToolResult);
761772
const capped = capToolResultForPersistence(
762-
persistMessage(normalizedToolResult),
773+
persistedToolResult,
763774
maxToolResultChars,
764775
redactionConfig,
765776
);
766-
const persisted = applyBeforeWriteHook(
767-
persistToolResult(capped, {
768-
toolCallId: id ?? undefined,
769-
toolName,
770-
isSynthetic: false,
771-
}),
772-
);
777+
const transformed = persistToolResult(capped, {
778+
toolCallId: id ?? undefined,
779+
toolName,
780+
isSynthetic: false,
781+
});
782+
const persisted = applyBeforeWriteHook(transformed);
773783
if (!persisted) {
774784
return undefined;
775785
}
776786
return appendMessageAndCacheTranscriptSeq(
777-
capToolResultForPersistence(persisted, maxToolResultChars, redactionConfig),
787+
capToolResultForPersistence(persisted.message, maxToolResultChars, redactionConfig),
788+
{
789+
invalidateSerializedPrefixCache:
790+
callerInvalidatesCache ||
791+
persistedToolResult !== normalizedToolResult ||
792+
toolResultTransformerMayMutate ||
793+
persisted.changed,
794+
},
778795
).entryId;
779796
}
780797

@@ -818,10 +835,12 @@ export function installSessionToolResultGuard(
818835
flushPendingToolResults();
819836
}
820837

821-
const finalMessage = applyBeforeWriteHook(persistMessage(nextMessage));
822-
if (!finalMessage) {
838+
const transformedMessage = persistMessage(nextMessage);
839+
const finalWrite = applyBeforeWriteHook(transformedMessage);
840+
if (!finalWrite) {
823841
return undefined;
824842
}
843+
const finalMessage = finalWrite.message;
825844
const finalRole = (finalMessage as { role?: unknown }).role;
826845
if (
827846
finalRole === "assistant" &&
@@ -845,7 +864,10 @@ export function installSessionToolResultGuard(
845864
entryId: result,
846865
messageSeq,
847866
sessionFile,
848-
} = appendMessageAndCacheTranscriptSeq(finalMessage);
867+
} = appendMessageAndCacheTranscriptSeq(finalMessage, {
868+
invalidateSerializedPrefixCache:
869+
callerInvalidatesCache || transformedMessage !== nextMessage || finalWrite.changed,
870+
});
849871
if (sessionFile) {
850872
emitSessionTranscriptUpdate({
851873
sessionFile,

‎src/agents/sessions/agent-session.ts‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,7 @@ export class AgentSession {
332332

333333
// Branch summarization state
334334
private branchSummaryAbortController: AbortController | undefined = undefined;
335+
private extensionModifiedToolResultIds = new Set<string>();
335336

336337
// Retry state
337338
private retryAbortController: AbortController | undefined = undefined;
@@ -515,6 +516,7 @@ export class AgentSession {
515516
if (!hookResult) {
516517
return undefined;
517518
}
519+
this.extensionModifiedToolResultIds.add(toolCall.id);
518520

519521
return {
520522
content: hookResult.content,
@@ -579,7 +581,7 @@ export class AgentSession {
579581
}
580582

581583
// Emit to extensions first
582-
await this.emitExtensionEvent(event);
584+
const messageChangedByExtension = await this.emitExtensionEvent(event);
583585

584586
// Notify all listeners
585587
this.emit(
@@ -605,7 +607,13 @@ export class AgentSession {
605607
event.message.role === "toolResult"
606608
) {
607609
// Regular LLM message - persist as SessionMessageEntry
608-
this.sessionManager.appendMessage(event.message);
610+
const toolResultChangedByExtension =
611+
event.message.role === "toolResult" &&
612+
this.extensionModifiedToolResultIds.delete(event.message.toolCallId);
613+
this.sessionManager.appendMessage(event.message, {
614+
invalidateSerializedPrefixCache:
615+
messageChangedByExtension || toolResultChangedByExtension,
616+
});
609617
}
610618
// Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
611619

@@ -689,7 +697,7 @@ export class AgentSession {
689697
}
690698

691699
/** Emit extension events based on agent events */
692-
private async emitExtensionEvent(event: AgentEvent): Promise<void> {
700+
private async emitExtensionEvent(event: AgentEvent): Promise<boolean> {
693701
if (event.type === "agent_start") {
694702
this.turnIndex = 0;
695703
await this.currentExtensionRunner.emit({ type: "agent_start" });
@@ -732,6 +740,7 @@ export class AgentSession {
732740
const replacement = await this.currentExtensionRunner.emitMessageEnd(extensionEvent);
733741
if (replacement) {
734742
this.replaceMessageInPlace(event.message, replacement);
743+
return true;
735744
}
736745
} else if (event.type === "tool_execution_start") {
737746
const extensionEvent: ToolExecutionStartEvent = {
@@ -760,6 +769,7 @@ export class AgentSession {
760769
};
761770
await this.currentExtensionRunner.emit(extensionEvent);
762771
}
772+
return false;
763773
}
764774

765775
/**

0 commit comments

Comments
 (0)