From 2f811a48cb1f1f40347f70f4a2142ff31ec1c760 Mon Sep 17 00:00:00 2001 From: Riad Benguella Date: Fri, 23 Oct 2026 18:11:12 +0000 Subject: [PATCH 1/3] Studio Code: Continue a turn whose tool call leaked into the reply as DSML markup --- apps/cli/ai/runtimes/pi/index.ts | 79 ++++++++------- .../ai/runtimes/pi/leaked-tool-call-markup.ts | 79 +++++++++++++++ .../pi/tests/leaked-tool-call-markup.test.ts | 98 +++++++++++++++++++ apps/cli/ai/tests/pi-runtime.test.ts | 33 +++++++ 4 files changed, 256 insertions(+), 33 deletions(-) create mode 100644 apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts create mode 100644 apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts diff --git a/apps/cli/ai/runtimes/pi/index.ts b/apps/cli/ai/runtimes/pi/index.ts index 1bcce65786..01e809b3da 100644 --- a/apps/cli/ai/runtimes/pi/index.ts +++ b/apps/cli/ai/runtimes/pi/index.ts @@ -52,6 +52,10 @@ import { createPresentDesignOptionsTool } from 'cli/ai/tools/present-design-opti import { createSkillTool } from 'cli/ai/tools/skill'; import { STUDIO_SITES_ROOT } from 'cli/lib/site-paths'; import { getFileToolPrompt } from './file-tool-prompts'; +import { + endedWithLeakedToolCall, + withLeakedToolCallMarkupRemoved, +} from './leaked-tool-call-markup'; import { getPendingWork, type PendingWork } from './pending-work'; import { stripStaleImagesFromContext } from './strip-stale-images'; import { @@ -78,12 +82,16 @@ type ProviderConfigInput = Parameters< ModelRuntime[ 'registerProvider' ] >[ 1 ] const STUDIO_WPCOM_PROVIDER = 'studio-wpcom'; const STUDIO_AGENT_DIR = STUDIO_SITES_ROOT; // A reply with nothing in it ends the turn without the user seeing why (see -// `isEmptyStop`). Resume it with a message the user doesn't see, up to this -// many times per turn. pi leaves a reply that has only reasoning out of the -// next request, so the message carries that reasoning; otherwise the model -// starts the same reasoning over and runs out again. +// `isEmptyStop`), and so does one whose tool call leaked into its text (see +// `endedWithLeakedToolCall`). Resume it with a message the user doesn't see, +// up to this many times per turn. pi leaves a reply that has only reasoning out +// of the next request, so the message carries that reasoning; otherwise the +// model starts the same reasoning over and runs out again. const MAX_EMPTY_STOP_CONTINUES = 2; const EMPTY_STOP_CONTINUE_TYPE = 'studio.empty_stop_continue'; +const LEAKED_TOOL_CALL_CONTINUE_TYPE = 'studio.leaked_tool_call_continue'; +const LEAKED_TOOL_CALL_CONTINUE_MESSAGE = + 'Your last tool call was written as text, so it did not run. Make the tool call again.'; const STUDIO_COMPACTION_SETTINGS = { enabled: true, reserveTokens: 16_384, @@ -249,8 +257,7 @@ async function runAgentSessionTurn( session = await createStudioAgentSession( config, family, resolved.creds, payloadGuardState ); setActiveSession( session ); let emptyStopContinues = 0; - let continueAfterEmptyStop = false; - let cutOffReasoning = ''; + let pendingContinue: { customType: string; content: string } | null = null; unsubscribe = session.subscribe( ( event ) => { updateStudioToolPayloadGuardState( event, payloadGuardState ); if ( @@ -259,13 +266,25 @@ async function runAgentSessionTurn( emptyStopContinues < MAX_EMPTY_STOP_CONTINUES ) { const lastAssistant = findLastAssistant( event.messages ); - if ( isEmptyStop( lastAssistant ) ) { - // Not final: consumers wait for the turn's real end. - continueAfterEmptyStop = true; - cutOffReasoning = ( lastAssistant?.content ?? [] ) + if ( endedWithLeakedToolCall( lastAssistant ) ) { + pendingContinue = { + customType: LEAKED_TOOL_CALL_CONTINUE_TYPE, + content: LEAKED_TOOL_CALL_CONTINUE_MESSAGE, + }; + } else if ( isEmptyStop( lastAssistant ) ) { + const cutOffReasoning = ( lastAssistant?.content ?? [] ) .map( ( block ) => ( block.type === 'thinking' ? block.thinking : '' ) ) .join( '' ) .trim(); + pendingContinue = { + customType: EMPTY_STOP_CONTINUE_TYPE, + content: cutOffReasoning + ? `Your reasoning was cut off before you replied. Continue from where it stopped:\n\n${ cutOffReasoning }` + : 'continue', + }; + } + if ( pendingContinue ) { + // Not final: consumers wait for the turn's real end. config.onEvent( { ...event, willRetry: true } ); return; } @@ -288,23 +307,15 @@ async function runAgentSessionTurn( mimeType: image.mimeType, } ) ), } ); - while ( continueAfterEmptyStop ) { - continueAfterEmptyStop = false; + while ( pendingContinue ) { + const message: { customType: string; content: string } = pendingContinue; + pendingContinue = null; if ( controller.signal.aborted ) { config.onEvent( syntheticErrorAgentEnd( 'aborted', '' ) ); break; } emptyStopContinues++; - await session.sendCustomMessage( - { - customType: EMPTY_STOP_CONTINUE_TYPE, - content: cutOffReasoning - ? `Your reasoning was cut off before you replied. Continue from where it stopped:\n\n${ cutOffReasoning }` - : 'continue', - display: false, - }, - { triggerTurn: true } - ); + await session.sendCustomMessage( { ...message, display: false }, { triggerTurn: true } ); } await getPendingWork( config.session ).settle(); } catch ( error ) { @@ -577,17 +588,19 @@ function createWpcomProviderConfig( headers: creds.extraHeaders, streamSimple: ( m, ctx, options?: SimpleStreamOptions ) => withUsageCapErrorRewrite( - stream( m, stripStaleImagesFromContext( ctx ), { - ...options, - ...( apiModelId - ? { - onPayload: async ( params: unknown, payloadModel ) => { - const next = ( await options?.onPayload?.( params, payloadModel ) ) ?? params; - return { ...( next as Record< string, unknown > ), model: apiModelId }; - }, - } - : {} ), - } ) + withLeakedToolCallMarkupRemoved( + stream( m, stripStaleImagesFromContext( ctx ), { + ...options, + ...( apiModelId + ? { + onPayload: async ( params: unknown, payloadModel ) => { + const next = ( await options?.onPayload?.( params, payloadModel ) ) ?? params; + return { ...( next as Record< string, unknown > ), model: apiModelId }; + }, + } + : {} ), + } ) + ) ), models: [ { diff --git a/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts b/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts new file mode 100644 index 0000000000..b05145f9c9 --- /dev/null +++ b/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts @@ -0,0 +1,79 @@ +import { createAssistantMessageEventStream } from '@earendil-works/pi-ai'; +import type { AssistantMessage, AssistantMessageEventStream } from '@earendil-works/pi-ai'; + +/** + * DeepSeek writes tool calls in its own DSML markup (`<｜DSML｜invoke …>`, + * with fullwidth bars). The upstream that serves it normally turns that markup + * into structured tool calls, but sometimes some of it comes back as reply + * text instead. The tool doesn't run, and because the reply has text and no + * tool call, the turn ends as if the model were done. + */ +const DSML_TAG = /<\/?[｜|]DSML[｜|][^>]*>/g; + +const repliesWithLeakedToolCalls = new WeakSet< AssistantMessage >(); + +/** Removes the span from the first DSML tag to the end of the last one. */ +export function stripLeakedToolCallMarkup( text: string ): string { + const tags = [ ...text.matchAll( DSML_TAG ) ]; + if ( tags.length === 0 ) { + return text; + } + const first = tags[ 0 ]; + const last = tags[ tags.length - 1 ]; + const start = first.index ?? 0; + const end = ( last.index ?? 0 ) + last[ 0 ].length; + return ( text.slice( 0, start ) + text.slice( end ) ).trim(); +} + +function cleanReply( message: AssistantMessage ): void { + let leaked = false; + const content: AssistantMessage[ 'content' ] = []; + for ( const block of message.content ) { + if ( block.type !== 'text' ) { + content.push( block ); + continue; + } + const text = stripLeakedToolCallMarkup( block.text ); + if ( text === block.text ) { + content.push( block ); + continue; + } + leaked = true; + if ( text ) { + content.push( { ...block, text } ); + } + } + if ( ! leaked ) { + return; + } + message.content = content; + if ( ! content.some( ( block ) => block.type === 'toolCall' ) ) { + repliesWithLeakedToolCalls.add( message ); + } +} + +/** + * Strips leaked tool-call markup from the final reply, so it never reaches the + * transcript, the saved session or the next request. A reply that lost a tool + * call this way is remembered for `endedWithLeakedToolCall`. + */ +export function withLeakedToolCallMarkupRemoved( + source: AssistantMessageEventStream +): AssistantMessageEventStream { + const cleaned = createAssistantMessageEventStream(); + void ( async () => { + for await ( const event of source ) { + if ( event.type === 'done' ) { + cleanReply( event.message ); + } + cleaned.push( event ); + } + cleaned.end(); + } )(); + return cleaned; +} + +/** A reply that ended the turn because its only tool call leaked into the text. */ +export function endedWithLeakedToolCall( message: AssistantMessage | undefined ): boolean { + return Boolean( message && repliesWithLeakedToolCalls.has( message ) ); +} diff --git a/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts b/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts new file mode 100644 index 0000000000..1598573bf2 --- /dev/null +++ b/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts @@ -0,0 +1,98 @@ +import { createAssistantMessageEventStream } from '@earendil-works/pi-ai'; +import { describe, expect, it } from 'vitest'; +import { + endedWithLeakedToolCall, + stripLeakedToolCallMarkup, + withLeakedToolCallMarkupRemoved, +} from '../leaked-tool-call-markup'; +import type { AssistantMessage } from '@earendil-works/pi-ai'; + +const LEAKED_CLOSING_TAGS = ' '; +const LEAKED_FULL_CALL = + '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="read">\n<｜DSML｜parameter name="path" string="true">inc/auction-admin.php\n\n'; + +function reply( content: AssistantMessage[ 'content' ] ): AssistantMessage { + return { + role: 'assistant', + content, + api: 'openai-completions', + provider: 'studio-wpcom', + model: 'fast', + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: 'stop', + timestamp: 0, + }; +} + +async function finish( message: AssistantMessage ): Promise< AssistantMessage > { + const source = createAssistantMessageEventStream(); + const cleaned = withLeakedToolCallMarkupRemoved( source ); + source.push( { type: 'done', reason: 'stop', message } ); + source.end(); + return cleaned.result(); +} + +describe( 'stripLeakedToolCallMarkup', () => { + it( 'removes leaked closing tags and keeps the text before them', () => { + expect( + stripLeakedToolCallMarkup( `Let me check the file.\n\n${ LEAKED_CLOSING_TAGS }` ) + ).toBe( 'Let me check the file.' ); + } ); + + it( 'removes a whole leaked tool call, arguments included', () => { + expect( stripLeakedToolCallMarkup( `Reading it now.\n${ LEAKED_FULL_CALL }` ) ).toBe( + 'Reading it now.' + ); + } ); + + it( 'leaves text without DSML tags unchanged, pipes and angle brackets included', () => { + const text = '| a | b |\n|---|---|\n|
| `

` |'; + expect( stripLeakedToolCallMarkup( text ) ).toBe( text ); + } ); +} ); + +describe( 'withLeakedToolCallMarkupRemoved', () => { + it( 'cleans a reply whose tool call leaked and flags it', async () => { + const result = await finish( + reply( [ { type: 'text', text: `Let me check the file.\n\n${ LEAKED_CLOSING_TAGS }` } ] ) + ); + + expect( result.content ).toEqual( [ { type: 'text', text: 'Let me check the file.' } ] ); + expect( endedWithLeakedToolCall( result ) ).toBe( true ); + } ); + + it( 'drops a text block that held only leaked markup', async () => { + const result = await finish( reply( [ { type: 'text', text: LEAKED_CLOSING_TAGS } ] ) ); + + expect( result.content ).toEqual( [] ); + expect( endedWithLeakedToolCall( result ) ).toBe( true ); + } ); + + it( 'cleans but does not flag a reply that still has a tool call', async () => { + const result = await finish( + reply( [ + { type: 'text', text: LEAKED_CLOSING_TAGS }, + { type: 'toolCall', id: 'call-1', name: 'read', arguments: { path: 'style.css' } }, + ] ) + ); + + expect( result.content.map( ( block ) => block.type ) ).toEqual( [ 'toolCall' ] ); + expect( endedWithLeakedToolCall( result ) ).toBe( false ); + } ); + + it( 'leaves an ordinary reply alone', async () => { + const message = reply( [ { type: 'text', text: 'Done.' } ] ); + const result = await finish( message ); + + expect( result ).toBe( message ); + expect( result.content ).toEqual( [ { type: 'text', text: 'Done.' } ] ); + expect( endedWithLeakedToolCall( result ) ).toBe( false ); + } ); +} ); diff --git a/apps/cli/ai/tests/pi-runtime.test.ts b/apps/cli/ai/tests/pi-runtime.test.ts index dfbb6a94ab..72277d0f97 100644 --- a/apps/cli/ai/tests/pi-runtime.test.ts +++ b/apps/cli/ai/tests/pi-runtime.test.ts @@ -1,8 +1,11 @@ +import { createAssistantMessageEventStream } from '@earendil-works/pi-ai'; import { ANTHROPIC_MODELS } from '@earendil-works/pi-ai/providers/anthropic.models'; import { ModelRegistry, SessionManager } from '@earendil-works/pi-coding-agent'; import { AI_MODELS } from '@studio/common/ai/models'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { runStudioAgentTurn, type StudioAgentTurnConfig } from 'cli/ai/runtimes/pi'; +import { withLeakedToolCallMarkupRemoved } from 'cli/ai/runtimes/pi/leaked-tool-call-markup'; +import type { AssistantMessage } from '@earendil-works/pi-ai'; import type { AgentSessionEvent, CreateAgentSessionOptions } from '@earendil-works/pi-coding-agent'; import type { AiModelId } from '@studio/common/ai/models'; @@ -601,6 +604,36 @@ describe( 'pi runtime', () => { expect( mocks.customMessages ).toHaveLength( 0 ); } ); + it( 'continues a turn whose tool call leaked into the text, asking for the call again', async () => { + const leaked = assistantMessage( [ + { type: 'text', text: 'Let me check the file.\n\n ' }, + ] ); + const source = createAssistantMessageEventStream(); + const cleaned = withLeakedToolCallMarkupRemoved( source ); + source.push( { type: 'done', reason: 'stop', message: leaked.message as AssistantMessage } ); + source.end(); + await cleaned.result(); + mocks.nextEvents = runEndingWith( leaked ); + mocks.continueEvents = [ + runEndingWith( assistantMessage( [ { type: 'text', text: 'Done.' } ] ) ), + ]; + + const events = await runFast(); + + expect( mocks.customMessages ).toEqual( [ + { + message: { + customType: 'studio.leaked_tool_call_continue', + content: + 'Your last tool call was written as text, so it did not run. Make the tool call again.', + display: false, + }, + options: { triggerTurn: true }, + }, + ] ); + expect( agentEnds( events ).map( ( event ) => event.willRetry ) ).toEqual( [ true, false ] ); + } ); + it( 'ends the turn as interrupted when the user stops it first', async () => { mocks.nextEvents = runEndingWith( emptyReply() ); const events: AgentSessionEvent[] = []; From 2d42e149241d229b06bc0388fcdc543d370244a6 Mon Sep 17 00:00:00 2001 From: Riad Benguella Date: Fri, 23 Oct 2026 19:33:25 +0100 Subject: [PATCH 2/3] Studio Code: Continue a leaked tool call through the empty-stop continue instead of a second path Co-Authored-By: Claude Opus 5.5 --- apps/cli/ai/runtimes/pi/index.ts | 62 +++++----- .../ai/runtimes/pi/leaked-tool-call-markup.ts | 67 +++-------- .../pi/tests/leaked-tool-call-markup.test.ts | 109 +++++------------- apps/cli/ai/tests/pi-runtime.test.ts | 33 ------ 4 files changed, 68 insertions(+), 203 deletions(-) diff --git a/apps/cli/ai/runtimes/pi/index.ts b/apps/cli/ai/runtimes/pi/index.ts index 01e809b3da..e752fb9fba 100644 --- a/apps/cli/ai/runtimes/pi/index.ts +++ b/apps/cli/ai/runtimes/pi/index.ts @@ -52,10 +52,7 @@ import { createPresentDesignOptionsTool } from 'cli/ai/tools/present-design-opti import { createSkillTool } from 'cli/ai/tools/skill'; import { STUDIO_SITES_ROOT } from 'cli/lib/site-paths'; import { getFileToolPrompt } from './file-tool-prompts'; -import { - endedWithLeakedToolCall, - withLeakedToolCallMarkupRemoved, -} from './leaked-tool-call-markup'; +import { withLeakedToolCallMarkupRemoved } from './leaked-tool-call-markup'; import { getPendingWork, type PendingWork } from './pending-work'; import { stripStaleImagesFromContext } from './strip-stale-images'; import { @@ -82,16 +79,12 @@ type ProviderConfigInput = Parameters< ModelRuntime[ 'registerProvider' ] >[ 1 ] const STUDIO_WPCOM_PROVIDER = 'studio-wpcom'; const STUDIO_AGENT_DIR = STUDIO_SITES_ROOT; // A reply with nothing in it ends the turn without the user seeing why (see -// `isEmptyStop`), and so does one whose tool call leaked into its text (see -// `endedWithLeakedToolCall`). Resume it with a message the user doesn't see, -// up to this many times per turn. pi leaves a reply that has only reasoning out -// of the next request, so the message carries that reasoning; otherwise the -// model starts the same reasoning over and runs out again. +// `isEmptyStop`). Resume it with a message the user doesn't see, up to this +// many times per turn. pi leaves a reply that has only reasoning out of the +// next request, so the message carries that reasoning; otherwise the model +// starts the same reasoning over and runs out again. const MAX_EMPTY_STOP_CONTINUES = 2; const EMPTY_STOP_CONTINUE_TYPE = 'studio.empty_stop_continue'; -const LEAKED_TOOL_CALL_CONTINUE_TYPE = 'studio.leaked_tool_call_continue'; -const LEAKED_TOOL_CALL_CONTINUE_MESSAGE = - 'Your last tool call was written as text, so it did not run. Make the tool call again.'; const STUDIO_COMPACTION_SETTINGS = { enabled: true, reserveTokens: 16_384, @@ -257,7 +250,8 @@ async function runAgentSessionTurn( session = await createStudioAgentSession( config, family, resolved.creds, payloadGuardState ); setActiveSession( session ); let emptyStopContinues = 0; - let pendingContinue: { customType: string; content: string } | null = null; + let continueAfterEmptyStop = false; + let cutOffReasoning = ''; unsubscribe = session.subscribe( ( event ) => { updateStudioToolPayloadGuardState( event, payloadGuardState ); if ( @@ -266,25 +260,13 @@ async function runAgentSessionTurn( emptyStopContinues < MAX_EMPTY_STOP_CONTINUES ) { const lastAssistant = findLastAssistant( event.messages ); - if ( endedWithLeakedToolCall( lastAssistant ) ) { - pendingContinue = { - customType: LEAKED_TOOL_CALL_CONTINUE_TYPE, - content: LEAKED_TOOL_CALL_CONTINUE_MESSAGE, - }; - } else if ( isEmptyStop( lastAssistant ) ) { - const cutOffReasoning = ( lastAssistant?.content ?? [] ) + if ( isEmptyStop( lastAssistant ) ) { + // Not final: consumers wait for the turn's real end. + continueAfterEmptyStop = true; + cutOffReasoning = ( lastAssistant?.content ?? [] ) .map( ( block ) => ( block.type === 'thinking' ? block.thinking : '' ) ) .join( '' ) .trim(); - pendingContinue = { - customType: EMPTY_STOP_CONTINUE_TYPE, - content: cutOffReasoning - ? `Your reasoning was cut off before you replied. Continue from where it stopped:\n\n${ cutOffReasoning }` - : 'continue', - }; - } - if ( pendingContinue ) { - // Not final: consumers wait for the turn's real end. config.onEvent( { ...event, willRetry: true } ); return; } @@ -307,15 +289,23 @@ async function runAgentSessionTurn( mimeType: image.mimeType, } ) ), } ); - while ( pendingContinue ) { - const message: { customType: string; content: string } = pendingContinue; - pendingContinue = null; + while ( continueAfterEmptyStop ) { + continueAfterEmptyStop = false; if ( controller.signal.aborted ) { config.onEvent( syntheticErrorAgentEnd( 'aborted', '' ) ); break; } emptyStopContinues++; - await session.sendCustomMessage( { ...message, display: false }, { triggerTurn: true } ); + await session.sendCustomMessage( + { + customType: EMPTY_STOP_CONTINUE_TYPE, + content: cutOffReasoning + ? `Your reasoning was cut off before you replied. Continue from where it stopped:\n\n${ cutOffReasoning }` + : 'continue', + display: false, + }, + { triggerTurn: true } + ); } await getPendingWork( config.session ).settle(); } catch ( error ) { @@ -562,9 +552,9 @@ function escapePiConfigValue( value: string ): string { } // The wpcom lane only needs pi's stock streaming for each tier's API; the -// custom provider exists to wrap the stream with the usage-cap 429 rewrite -// and to strip stale screenshots, which would otherwise bloat requests past -// the proxy's body limit. +// custom provider exists to wrap the stream with the usage-cap 429 rewrite and +// the leaked tool-call cleanup, and to strip stale screenshots, which would +// otherwise bloat requests past the proxy's body limit. function createWpcomProviderConfig( model: StudioWpcomModel, creds: ResolvedCredentials diff --git a/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts b/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts index b05145f9c9..2a86e4c40c 100644 --- a/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts +++ b/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts @@ -5,58 +5,24 @@ import type { AssistantMessage, AssistantMessageEventStream } from '@earendil-wo * DeepSeek writes tool calls in its own DSML markup (`<｜DSML｜invoke …>`, * with fullwidth bars). The upstream that serves it normally turns that markup * into structured tool calls, but sometimes some of it comes back as reply - * text instead. The tool doesn't run, and because the reply has text and no - * tool call, the turn ends as if the model were done. + * text instead, and the tool never runs. Matches from the first tag to the + * last one. */ -const DSML_TAG = /<\/?[｜|]DSML[｜|][^>]*>/g; +const LEAKED_TOOL_CALL = /<\/?[｜|]DSML[｜|](?:[\s\S]*<\/?[｜|]DSML[｜|])?[^>]*>/; -const repliesWithLeakedToolCalls = new WeakSet< AssistantMessage >(); - -/** Removes the span from the first DSML tag to the end of the last one. */ -export function stripLeakedToolCallMarkup( text: string ): string { - const tags = [ ...text.matchAll( DSML_TAG ) ]; - if ( tags.length === 0 ) { - return text; - } - const first = tags[ 0 ]; - const last = tags[ tags.length - 1 ]; - const start = first.index ?? 0; - const end = ( last.index ?? 0 ) + last[ 0 ].length; - return ( text.slice( 0, start ) + text.slice( end ) ).trim(); -} - -function cleanReply( message: AssistantMessage ): void { - let leaked = false; - const content: AssistantMessage[ 'content' ] = []; - for ( const block of message.content ) { - if ( block.type !== 'text' ) { - content.push( block ); - continue; - } - const text = stripLeakedToolCallMarkup( block.text ); - if ( text === block.text ) { - content.push( block ); - continue; - } - leaked = true; - if ( text ) { - content.push( { ...block, text } ); +// Without a tool call left, the text that carried the leaked one goes too: the +// reply is then an empty stop, which the runtime continues. +function removeLeakedToolCall( message: AssistantMessage ): void { + const hasToolCall = message.content.some( ( block ) => block.type === 'toolCall' ); + message.content = message.content.flatMap( ( block ) => { + if ( block.type !== 'text' || ! LEAKED_TOOL_CALL.test( block.text ) ) { + return [ block ]; } - } - if ( ! leaked ) { - return; - } - message.content = content; - if ( ! content.some( ( block ) => block.type === 'toolCall' ) ) { - repliesWithLeakedToolCalls.add( message ); - } + const text = block.text.replace( LEAKED_TOOL_CALL, '' ).trim(); + return hasToolCall && text ? [ { ...block, text } ] : []; + } ); } -/** - * Strips leaked tool-call markup from the final reply, so it never reaches the - * transcript, the saved session or the next request. A reply that lost a tool - * call this way is remembered for `endedWithLeakedToolCall`. - */ export function withLeakedToolCallMarkupRemoved( source: AssistantMessageEventStream ): AssistantMessageEventStream { @@ -64,7 +30,7 @@ export function withLeakedToolCallMarkupRemoved( void ( async () => { for await ( const event of source ) { if ( event.type === 'done' ) { - cleanReply( event.message ); + removeLeakedToolCall( event.message ); } cleaned.push( event ); } @@ -72,8 +38,3 @@ export function withLeakedToolCallMarkupRemoved( } )(); return cleaned; } - -/** A reply that ended the turn because its only tool call leaked into the text. */ -export function endedWithLeakedToolCall( message: AssistantMessage | undefined ): boolean { - return Boolean( message && repliesWithLeakedToolCalls.has( message ) ); -} diff --git a/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts b/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts index 1598573bf2..34cb11ee50 100644 --- a/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts +++ b/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts @@ -1,98 +1,45 @@ import { createAssistantMessageEventStream } from '@earendil-works/pi-ai'; +import { isEmptyStop } from '@studio/common/ai/session-events'; import { describe, expect, it } from 'vitest'; -import { - endedWithLeakedToolCall, - stripLeakedToolCallMarkup, - withLeakedToolCallMarkupRemoved, -} from '../leaked-tool-call-markup'; +import { withLeakedToolCallMarkupRemoved } from '../leaked-tool-call-markup'; import type { AssistantMessage } from '@earendil-works/pi-ai'; -const LEAKED_CLOSING_TAGS = ' '; -const LEAKED_FULL_CALL = - '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="read">\n<｜DSML｜parameter name="path" string="true">inc/auction-admin.php\n\n'; +const LEAKED_CALL = + '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="read">\n<｜DSML｜parameter name="path" string="true">style.css\n\n'; -function reply( content: AssistantMessage[ 'content' ] ): AssistantMessage { - return { - role: 'assistant', - content, - api: 'openai-completions', - provider: 'studio-wpcom', - model: 'fast', - usage: { - input: 0, - output: 0, - cacheRead: 0, - cacheWrite: 0, - totalTokens: 0, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }, - stopReason: 'stop', - timestamp: 0, - }; -} - -async function finish( message: AssistantMessage ): Promise< AssistantMessage > { +async function finish( content: AssistantMessage[ 'content' ] ): Promise< AssistantMessage > { const source = createAssistantMessageEventStream(); const cleaned = withLeakedToolCallMarkupRemoved( source ); + const message = { role: 'assistant', content, stopReason: 'stop' } as AssistantMessage; source.push( { type: 'done', reason: 'stop', message } ); source.end(); return cleaned.result(); } -describe( 'stripLeakedToolCallMarkup', () => { - it( 'removes leaked closing tags and keeps the text before them', () => { - expect( - stripLeakedToolCallMarkup( `Let me check the file.\n\n${ LEAKED_CLOSING_TAGS }` ) - ).toBe( 'Let me check the file.' ); - } ); - - it( 'removes a whole leaked tool call, arguments included', () => { - expect( stripLeakedToolCallMarkup( `Reading it now.\n${ LEAKED_FULL_CALL }` ) ).toBe( - 'Reading it now.' - ); - } ); - - it( 'leaves text without DSML tags unchanged, pipes and angle brackets included', () => { - const text = '| a | b |\n|---|---|\n|
| `

` |'; - expect( stripLeakedToolCallMarkup( text ) ).toBe( text ); - } ); -} ); - describe( 'withLeakedToolCallMarkupRemoved', () => { - it( 'cleans a reply whose tool call leaked and flags it', async () => { - const result = await finish( - reply( [ { type: 'text', text: `Let me check the file.\n\n${ LEAKED_CLOSING_TAGS }` } ] ) - ); - - expect( result.content ).toEqual( [ { type: 'text', text: 'Let me check the file.' } ] ); - expect( endedWithLeakedToolCall( result ) ).toBe( true ); - } ); - - it( 'drops a text block that held only leaked markup', async () => { - const result = await finish( reply( [ { type: 'text', text: LEAKED_CLOSING_TAGS } ] ) ); - - expect( result.content ).toEqual( [] ); - expect( endedWithLeakedToolCall( result ) ).toBe( true ); + it( 'turns a reply whose only tool call leaked into an empty stop', async () => { + const result = await finish( [ + { type: 'thinking', thinking: 'Read the stylesheet.' }, + { type: 'text', text: `Let me check the file.\n\n${ LEAKED_CALL }` }, + ] ); + + expect( result.content ).toEqual( [ { type: 'thinking', thinking: 'Read the stylesheet.' } ] ); + expect( isEmptyStop( result ) ).toBe( true ); } ); - it( 'cleans but does not flag a reply that still has a tool call', async () => { - const result = await finish( - reply( [ - { type: 'text', text: LEAKED_CLOSING_TAGS }, - { type: 'toolCall', id: 'call-1', name: 'read', arguments: { path: 'style.css' } }, - ] ) - ); - - expect( result.content.map( ( block ) => block.type ) ).toEqual( [ 'toolCall' ] ); - expect( endedWithLeakedToolCall( result ) ).toBe( false ); - } ); - - it( 'leaves an ordinary reply alone', async () => { - const message = reply( [ { type: 'text', text: 'Done.' } ] ); - const result = await finish( message ); - - expect( result ).toBe( message ); - expect( result.content ).toEqual( [ { type: 'text', text: 'Done.' } ] ); - expect( endedWithLeakedToolCall( result ) ).toBe( false ); + it( 'keeps the text and tool calls around leaked markup, and other markup', async () => { + const table = '| a | b |\n|---|---|\n|
| `

` |'; + const toolCall = { type: 'toolCall', id: 'call-1', name: 'read', arguments: {} } as const; + const result = await finish( [ + { type: 'text', text: table }, + { type: 'text', text: 'Reading it now. ' }, + toolCall, + ] ); + + expect( result.content ).toEqual( [ + { type: 'text', text: table }, + { type: 'text', text: 'Reading it now.' }, + toolCall, + ] ); } ); } ); diff --git a/apps/cli/ai/tests/pi-runtime.test.ts b/apps/cli/ai/tests/pi-runtime.test.ts index 72277d0f97..dfbb6a94ab 100644 --- a/apps/cli/ai/tests/pi-runtime.test.ts +++ b/apps/cli/ai/tests/pi-runtime.test.ts @@ -1,11 +1,8 @@ -import { createAssistantMessageEventStream } from '@earendil-works/pi-ai'; import { ANTHROPIC_MODELS } from '@earendil-works/pi-ai/providers/anthropic.models'; import { ModelRegistry, SessionManager } from '@earendil-works/pi-coding-agent'; import { AI_MODELS } from '@studio/common/ai/models'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { runStudioAgentTurn, type StudioAgentTurnConfig } from 'cli/ai/runtimes/pi'; -import { withLeakedToolCallMarkupRemoved } from 'cli/ai/runtimes/pi/leaked-tool-call-markup'; -import type { AssistantMessage } from '@earendil-works/pi-ai'; import type { AgentSessionEvent, CreateAgentSessionOptions } from '@earendil-works/pi-coding-agent'; import type { AiModelId } from '@studio/common/ai/models'; @@ -604,36 +601,6 @@ describe( 'pi runtime', () => { expect( mocks.customMessages ).toHaveLength( 0 ); } ); - it( 'continues a turn whose tool call leaked into the text, asking for the call again', async () => { - const leaked = assistantMessage( [ - { type: 'text', text: 'Let me check the file.\n\n ' }, - ] ); - const source = createAssistantMessageEventStream(); - const cleaned = withLeakedToolCallMarkupRemoved( source ); - source.push( { type: 'done', reason: 'stop', message: leaked.message as AssistantMessage } ); - source.end(); - await cleaned.result(); - mocks.nextEvents = runEndingWith( leaked ); - mocks.continueEvents = [ - runEndingWith( assistantMessage( [ { type: 'text', text: 'Done.' } ] ) ), - ]; - - const events = await runFast(); - - expect( mocks.customMessages ).toEqual( [ - { - message: { - customType: 'studio.leaked_tool_call_continue', - content: - 'Your last tool call was written as text, so it did not run. Make the tool call again.', - display: false, - }, - options: { triggerTurn: true }, - }, - ] ); - expect( agentEnds( events ).map( ( event ) => event.willRetry ) ).toEqual( [ true, false ] ); - } ); - it( 'ends the turn as interrupted when the user stops it first', async () => { mocks.nextEvents = runEndingWith( emptyReply() ); const events: AgentSessionEvent[] = []; From f12fef5f6a11d4b004242ffb557931cc58efa3a3 Mon Sep 17 00:00:00 2001 From: Riad Benguella Date: Fri, 23 Oct 2026 19:54:49 +0100 Subject: [PATCH 3/3] Studio Code: Continue a leaked tool call the upstream still reported as a tool call Co-Authored-By: Claude Opus 5.5 --- .../ai/runtimes/pi/leaked-tool-call-markup.ts | 6 +++++- .../pi/tests/leaked-tool-call-markup.test.ts | 18 ++++++++++++------ 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts b/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts index 2a86e4c40c..e0b9f73ca3 100644 --- a/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts +++ b/apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts @@ -11,9 +11,13 @@ import type { AssistantMessage, AssistantMessageEventStream } from '@earendil-wo const LEAKED_TOOL_CALL = /<\/?[｜|]DSML[｜|](?:[\s\S]*<\/?[｜|]DSML[｜|])?[^>]*>/; // Without a tool call left, the text that carried the leaked one goes too: the -// reply is then an empty stop, which the runtime continues. +// reply is then an empty stop, which the runtime continues. The upstream may +// still report the reply as a tool call. function removeLeakedToolCall( message: AssistantMessage ): void { const hasToolCall = message.content.some( ( block ) => block.type === 'toolCall' ); + if ( ! hasToolCall && message.stopReason === 'toolUse' ) { + message.stopReason = 'stop'; + } message.content = message.content.flatMap( ( block ) => { if ( block.type !== 'text' || ! LEAKED_TOOL_CALL.test( block.text ) ) { return [ block ]; diff --git a/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts b/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts index 34cb11ee50..07f5dedebf 100644 --- a/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts +++ b/apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts @@ -7,10 +7,13 @@ import type { AssistantMessage } from '@earendil-works/pi-ai'; const LEAKED_CALL = '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="read">\n<｜DSML｜parameter name="path" string="true">style.css\n\n'; -async function finish( content: AssistantMessage[ 'content' ] ): Promise< AssistantMessage > { +async function finish( + content: AssistantMessage[ 'content' ], + stopReason: AssistantMessage[ 'stopReason' ] = 'stop' +): Promise< AssistantMessage > { const source = createAssistantMessageEventStream(); const cleaned = withLeakedToolCallMarkupRemoved( source ); - const message = { role: 'assistant', content, stopReason: 'stop' } as AssistantMessage; + const message = { role: 'assistant', content, stopReason } as AssistantMessage; source.push( { type: 'done', reason: 'stop', message } ); source.end(); return cleaned.result(); @@ -18,10 +21,13 @@ async function finish( content: AssistantMessage[ 'content' ] ): Promise< Assist describe( 'withLeakedToolCallMarkupRemoved', () => { it( 'turns a reply whose only tool call leaked into an empty stop', async () => { - const result = await finish( [ - { type: 'thinking', thinking: 'Read the stylesheet.' }, - { type: 'text', text: `Let me check the file.\n\n${ LEAKED_CALL }` }, - ] ); + const result = await finish( + [ + { type: 'thinking', thinking: 'Read the stylesheet.' }, + { type: 'text', text: `Let me check the file.\n\n${ LEAKED_CALL }` }, + ], + 'toolUse' + ); expect( result.content ).toEqual( [ { type: 'thinking', thinking: 'Read the stylesheet.' } ] ); expect( isEmptyStop( result ) ).toBe( true );