Studio Code: Continue a turn whose tool call leaked into the reply as DSML markup

apps/cli/ai/runtimes/pi/index.ts

@@ -52,6 +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 { withLeakedToolCallMarkupRemoved } from './leaked-tool-call-markup';
 import { getPendingWork, type PendingWork } from './pending-work';
 import { stripStaleImagesFromContext } from './strip-stale-images';
 import {
@@ -551,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
@@ -577,17 +578,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: [
 			{

apps/cli/ai/runtimes/pi/leaked-tool-call-markup.ts

@@ -0,0 +1,44 @@
+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, and the tool never runs. Matches from the first tag to the
+ * last one.
+ */
+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. 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 ];
+		}
+		const text = block.text.replace( LEAKED_TOOL_CALL, '' ).trim();
+		return hasToolCall && text ? [ { ...block, text } ] : [];
+	} );
+}
+
+export function withLeakedToolCallMarkupRemoved(
+	source: AssistantMessageEventStream
+): AssistantMessageEventStream {
+	const cleaned = createAssistantMessageEventStream();
+	void ( async () => {
+		for await ( const event of source ) {
+			if ( event.type === 'done' ) {
+				removeLeakedToolCall( event.message );
+			}
+			cleaned.push( event );
+		}
+		cleaned.end();
+	} )();
+	return cleaned;
+}

apps/cli/ai/runtimes/pi/tests/leaked-tool-call-markup.test.ts

@@ -0,0 +1,51 @@
+import { createAssistantMessageEventStream } from '@earendil-works/pi-ai';
+import { isEmptyStop } from '@studio/common/ai/session-events';
+import { describe, expect, it } from 'vitest';
+import { withLeakedToolCallMarkupRemoved } from '../leaked-tool-call-markup';
+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</｜DSML｜parameter>\n</｜DSML｜invoke>\n</｜DSML｜tool_calls>';
+
+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 } as AssistantMessage;
+	source.push( { type: 'done', reason: 'stop', message } );
+	source.end();
+	return cleaned.result();
+}
+
+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 }` },
+			],
+			'toolUse'
+		);
+
+		expect( result.content ).toEqual( [ { type: 'thinking', thinking: 'Read the stylesheet.' } ] );
+		expect( isEmptyStop( result ) ).toBe( true );
+	} );
+
+	it( 'keeps the text and tool calls around leaked markup, and other markup', async () => {
+		const table = '| a | b |\n|---|---|\n| <div> | `</p>` |';
+		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. </｜DSML｜invoke> </｜DSML｜tool_calls>' },
+			toolCall,
+		] );
+
+		expect( result.content ).toEqual( [
+			{ type: 'text', text: table },
+			{ type: 'text', text: 'Reading it now.' },
+			toolCall,
+		] );
+	} );
+} );