From f534ea9906d2bb57caba54c113fbcd408e9947d1 Mon Sep 17 00:00:00 2001 From: Mitch McAlister Date: Mon, 2 Mar 2026 19:44:40 +0000 Subject: [PATCH] fix: prevent reasoning text leak through handleMessageEnd fallback When enforceFinalTag is active (Google providers), stripBlockTags correctly returns empty for text without tags. However, the handleMessageEnd fallback recovered raw text, bypassing this protection and leaking internal reasoning (e.g. "**Applying single-bot mention rule**NO_REPLY") to Discord. Guard the fallback with enforceFinalTag check: if the provider is supposed to use tags and none were seen, the text is treated as leaked reasoning and suppressed. Also harden stripSilentToken regex to allow bold markdown (**) as separator before NO_REPLY, matching the pattern Gemini Flash Lite produces. Co-Authored-By: Claude Opus 4.6 --- ...pi-embedded-subscribe.handlers.messages.ts | 2 +- ...uppresses-output-without-start-tag.test.ts | 33 +++++++++++++++++-- src/auto-reply/tokens.test.ts | 6 ++++ src/auto-reply/tokens.ts | 2 +- 4 files changed, 38 insertions(+), 5 deletions(-) diff --git a/src/agents/pi-embedded-subscribe.handlers.messages.ts b/src/agents/pi-embedded-subscribe.handlers.messages.ts index a988ebcc6..d58690814 100644 --- a/src/agents/pi-embedded-subscribe.handlers.messages.ts +++ b/src/agents/pi-embedded-subscribe.handlers.messages.ts @@ -288,7 +288,7 @@ export function handleMessageEnd( let mediaUrls = parsedText?.mediaUrls; let hasMedia = Boolean(mediaUrls && mediaUrls.length > 0); - if (!cleanedText && !hasMedia) { + if (!cleanedText && !hasMedia && !ctx.params.enforceFinalTag) { const rawTrimmed = rawText.trim(); const rawStrippedFinal = rawTrimmed.replace(/<\s*\/?\s*final\s*>/gi, "").trim(); const rawCandidate = rawStrippedFinal || rawTrimmed; diff --git a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.test.ts b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.test.ts index 79a8cf50a..0f66888e3 100644 --- a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.test.ts +++ b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.test.ts @@ -4,7 +4,7 @@ import { createStubSessionHarness, emitAssistantTextDelta, emitMessageStartAndEndForAssistantText, - expectSingleAgentEventText, + extractAgentEventPayloads, } from "./pi-embedded-subscribe.e2e-harness.js"; import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js"; @@ -37,7 +37,7 @@ describe("subscribeEmbeddedPiSession", () => { expect(onPartialReply).not.toHaveBeenCalled(); }); - it("emits agent events on message_end even without tags", () => { + it("suppresses agent events on message_end without tags when enforced", () => { const { session, emit } = createStubSessionHarness(); const onAgentEvent = vi.fn(); @@ -49,7 +49,34 @@ describe("subscribeEmbeddedPiSession", () => { onAgentEvent, }); emitMessageStartAndEndForAssistantText({ emit, text: "Hello world" }); - expectSingleAgentEventText(onAgentEvent.mock.calls, "Hello world"); + // With enforceFinalTag, text without tags is treated as leaked + // reasoning and should NOT be recovered by the message_end fallback. + const payloads = extractAgentEventPayloads(onAgentEvent.mock.calls); + expect(payloads).toHaveLength(0); + }); + it("emits via streaming when tags are present and enforcement is on", () => { + const { session, emit } = createStubSessionHarness(); + + const onPartialReply = vi.fn(); + const onAgentEvent = vi.fn(); + + subscribeEmbeddedPiSession({ + session, + runId: "run", + enforceFinalTag: true, + onPartialReply, + onAgentEvent, + }); + + // With enforceFinalTag, content is emitted via streaming (text_delta path), + // NOT recovered from message_end fallback. extractAssistantText strips + // tags, so message_end would see plain text with no markers + // and correctly suppress it (treated as reasoning leak). + emit({ type: "message_start", message: { role: "assistant" } }); + emitAssistantTextDelta({ emit, delta: "Hello world" }); + + expect(onPartialReply).toHaveBeenCalled(); + expect(onPartialReply.mock.calls[0][0].text).toBe("Hello world"); }); it("does not require when enforcement is off", () => { const { session, emit } = createStubSessionHarness(); diff --git a/src/auto-reply/tokens.test.ts b/src/auto-reply/tokens.test.ts index 6dc51d1b7..78db0cffd 100644 --- a/src/auto-reply/tokens.test.ts +++ b/src/auto-reply/tokens.test.ts @@ -62,6 +62,12 @@ describe("stripSilentToken", () => { expect(stripSilentToken(" NO_REPLY ")).toBe(""); }); + it("strips token preceded by bold markdown formatting", () => { + expect(stripSilentToken("**NO_REPLY")).toBe(""); + expect(stripSilentToken("some text **NO_REPLY")).toBe("some text"); + expect(stripSilentToken("reasoning**NO_REPLY")).toBe("reasoning"); + }); + it("works with custom token", () => { expect(stripSilentToken("done HEARTBEAT_OK", "HEARTBEAT_OK")).toBe("done"); }); diff --git a/src/auto-reply/tokens.ts b/src/auto-reply/tokens.ts index 9be470d64..3a0f18d3d 100644 --- a/src/auto-reply/tokens.ts +++ b/src/auto-reply/tokens.ts @@ -24,7 +24,7 @@ export function isSilentReplyText( */ export function stripSilentToken(text: string, token: string = SILENT_REPLY_TOKEN): string { const escaped = escapeRegExp(token); - return text.replace(new RegExp(`(?:^|\\s+)${escaped}\\s*$`), "").trim(); + return text.replace(new RegExp(`(?:^|\\s+|\\*+)${escaped}\\s*$`), "").trim(); } export function isSilentReplyPrefixText(