Skip to content

Commit 9b2b46f

Browse files
sercadasercada
committed
fix(agents): harden assistant text reconciliation
1 parent 1227486 commit 9b2b46f

3 files changed

Lines changed: 74 additions & 10 deletions

File tree

‎src/agents/pi-embedded-runner/run/assistant-text-persistence.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ function persistedTextCoversAssistantText(persisted: string, candidate: string):
3030
if (candidateNormalized.length <= SHORT_SUBSTRING_COVERAGE_MAX_LENGTH) {
3131
return normalizedTextSegments(persisted).includes(candidateNormalized);
3232
}
33-
return candidateNormalized.length >= 20 && persistedNormalized.includes(candidateNormalized);
33+
return persistedNormalized.includes(candidateNormalized);
3434
}
3535

3636
function toPersistableAssistantText(text: string): string | undefined {

‎src/agents/pi-embedded-runner/run/attempt.test.ts‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
import { streamSimple } from "@earendil-works/pi-ai";
1+
import type { AgentMessage } from "@earendil-works/pi-agent-core";
2+
import { streamSimple, type AssistantMessage } from "@earendil-works/pi-ai";
3+
import { SessionManager } from "@earendil-works/pi-coding-agent";
24
import { describe, expect, it, vi } from "vitest";
35

46
vi.mock("../context-engine-capabilities.js", () => ({
@@ -131,6 +133,48 @@ describe("buildEmbeddedAttemptToolRunContext", () => {
131133
});
132134
});
133135

136+
describe("reconcileAssistantTranscriptAndPreserveLastAssistant", () => {
137+
it("keeps real assistant usage while appending a zero-usage transcript mirror", () => {
138+
const sessionManager = SessionManager.inMemory();
139+
const realLastAssistant = {
140+
role: "assistant",
141+
api: "openai-responses",
142+
provider: "openai-codex",
143+
model: "gpt-5.4",
144+
content: [{ type: "toolCall", id: "call_1", name: "exec", arguments: {} }],
145+
usage: { input: 100, output: 25, totalTokens: 125 },
146+
stopReason: "toolUse",
147+
timestamp: 2,
148+
} as AssistantMessage;
149+
150+
sessionManager.appendMessage({ role: "user", content: "question", timestamp: 1 });
151+
sessionManager.appendMessage(realLastAssistant);
152+
const messagesSnapshot = sessionManager
153+
.getEntries()
154+
.filter((entry) => entry.type === "message")
155+
.map((entry) => (entry as { message: AgentMessage }).message);
156+
157+
const preserved = attemptTesting.reconcileAssistantTranscriptAndPreserveLastAssistant({
158+
sessionManager,
159+
mutableMessagesSnapshot: messagesSnapshot,
160+
prePromptMessageCount: 0,
161+
assistantTexts: ["Final answer delivered after the tool call."],
162+
api: "openai-responses",
163+
provider: "openai-codex",
164+
modelId: "gpt-5.4",
165+
lastAssistant: realLastAssistant,
166+
});
167+
168+
const mirror = messagesSnapshot.at(-1) as AssistantMessage;
169+
expect(preserved).toBe(realLastAssistant);
170+
expect(preserved?.usage).toEqual(realLastAssistant.usage);
171+
expect((preserved?.usage as { totalTokens?: number }).totalTokens).toBe(125);
172+
expect(mirror).not.toBe(realLastAssistant);
173+
expect(JSON.stringify(mirror.content)).toContain("Final answer delivered");
174+
expect((mirror.usage as { totalTokens?: number }).totalTokens).toBe(0);
175+
});
176+
});
177+
134178
describe("buildCallableToolNamesForEmptyAllowlistCheck", () => {
135179
it("ignores auto-added Tool Search controls so bad allowlists still fail", () => {
136180
expect(

‎src/agents/pi-embedded-runner/run/attempt.ts‎

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -784,10 +784,33 @@ async function cancelQueuedSteeringMessage(
784784

785785
export const testing = {
786786
cancelQueuedSteeringMessage,
787+
reconcileAssistantTranscriptAndPreserveLastAssistant,
787788
resolveAttemptStreamAuthProfileId,
788789
steerAndWaitForTranscriptCommit,
789790
};
790791

792+
function reconcileAssistantTranscriptAndPreserveLastAssistant(params: {
793+
sessionManager: Pick<SessionManager, "appendMessage">;
794+
mutableMessagesSnapshot: AgentMessage[];
795+
prePromptMessageCount: number;
796+
assistantTexts: readonly string[];
797+
api: AssistantMessage["api"];
798+
provider: string;
799+
modelId: string;
800+
lastAssistant: AssistantMessage | undefined;
801+
}): AssistantMessage | undefined {
802+
reconcileAssistantTextsWithTranscript({
803+
sessionManager: params.sessionManager,
804+
mutableMessagesSnapshot: params.mutableMessagesSnapshot,
805+
prePromptMessageCount: params.prePromptMessageCount,
806+
assistantTexts: params.assistantTexts,
807+
api: params.api,
808+
provider: params.provider,
809+
modelId: params.modelId,
810+
});
811+
return params.lastAssistant;
812+
}
813+
791814
function resolveAttemptStreamAuthProfileId(
792815
params: Pick<EmbeddedRunAttemptParams, "authProfileId" | "runtimePlan">,
793816
): string | undefined {
@@ -3190,9 +3213,8 @@ export async function runEmbeddedAttempt(
31903213
prompt: string,
31913214
options?: Parameters<typeof activeSession.prompt>[1],
31923215
): Promise<void> =>
3193-
withOwnedSessionTranscriptWrites(
3194-
ownedTranscriptWriteContext,
3195-
async () => abortable(trackPromptSettlePromise(activeSession.prompt(prompt, options))),
3216+
withOwnedSessionTranscriptWrites(ownedTranscriptWriteContext, async () =>
3217+
abortable(trackPromptSettlePromise(activeSession.prompt(prompt, options))),
31963218
);
31973219
const onBlockReply = params.onBlockReply
31983220
? bindOwnedSessionTranscriptWrites(ownedTranscriptWriteContext, params.onBlockReply)
@@ -4322,18 +4344,16 @@ export async function runEmbeddedAttempt(
43224344
!timedOutDuringCompaction &&
43234345
!compactionOccurredThisAttempt
43244346
) {
4325-
const reconciledAssistant = reconcileAssistantTextsWithTranscript({
4326-
sessionManager,
4347+
lastAssistant = reconcileAssistantTranscriptAndPreserveLastAssistant({
4348+
sessionManager: activeSessionManager,
43274349
mutableMessagesSnapshot: messagesSnapshot,
43284350
prePromptMessageCount,
43294351
assistantTexts,
43304352
api: params.model.api,
43314353
provider: params.provider,
43324354
modelId: params.modelId,
4355+
lastAssistant,
43334356
});
4334-
if (reconciledAssistant) {
4335-
lastAssistant = reconciledAssistant;
4336-
}
43374357
}
43384358
cacheBreak = cacheObservabilityEnabled
43394359
? completePromptCacheObservation({

0 commit comments

Comments
 (0)