diff --git a/apps/cli/ai/runtimes/pi/index.ts b/apps/cli/ai/runtimes/pi/index.ts index 1bcce65786..e752fb9fba 100644 --- a/apps/cli/ai/runtimes/pi/index.ts +++ b/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: [ { 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..e0b9f73ca3 --- /dev/null +++ b/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; +} 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..07f5dedebf --- /dev/null +++ b/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|