From 4791b40b269b19a3c40daae5ed412ebc5d1b6d42 Mon Sep 17 00:00:00 2001 From: Tianyi Cui <53024+tianyicui@users.noreply.github.com> Date: Sun, 19 Jul 2026 14:18:13 +0800 Subject: [PATCH] fix(llm-pi-ai): classify usage-based context overflow Pass each resolved catalog model capacity into pi-ai stream conversion so the upstream full-message classifier can recognize provider-specific, silent, and length-stop overflow signals. Retain the harness text fallback for legacy provider wording and cover the catalog-resolution path with a mock-provider regression. --- packages/llm/llm-pi-ai/README.md | 2 +- packages/llm/llm-pi-ai/src/adapter.ts | 2 +- packages/llm/llm-pi-ai/src/stream.ts | 34 +++++++++++++++----- packages/llm/llm-pi-ai/tests/adapter.spec.ts | 26 ++++++++++++++- packages/llm/llm-pi-ai/tests/convert.spec.ts | 28 ++++++++++++++++ 5 files changed, 81 insertions(+), 11 deletions(-) diff --git a/packages/llm/llm-pi-ai/README.md b/packages/llm/llm-pi-ai/README.md index e2d56e28a5..1deae8f039 100644 --- a/packages/llm/llm-pi-ai/README.md +++ b/packages/llm/llm-pi-ai/README.md @@ -43,7 +43,7 @@ If a listener rewrites assembled assistant content, the loop drops replay state ## Vocabulary differences - pi-ai tool-call arguments are parsed objects; the harness stores raw JSON strings. The adapter parses input and re-stringifies output. -- pi-ai reports failures as in-stream error events; these map to `finish {kind:'error'|'aborted'}` chunks, with recognized context overflow normalized to `CONTEXT_WINDOW_EXCEEDED`. +- pi-ai reports failures as in-stream error events; these map to `finish {kind:'error'|'aborted'}` chunks. Provider-specific error text and usage signals evaluated against the resolved model's context window normalize overflow to `CONTEXT_WINDOW_EXCEEDED`. - pi-ai folds reasoning tokens into output usage; there is no separate reasoning count to map. - `GenerateOptions.stop` is rejected with `UNSUPPORTED_OPTION` because pi-ai's common streaming surface cannot guarantee it across providers. diff --git a/packages/llm/llm-pi-ai/src/adapter.ts b/packages/llm/llm-pi-ai/src/adapter.ts index 26e28ecd1f..7f40c67da3 100644 --- a/packages/llm/llm-pi-ai/src/adapter.ts +++ b/packages/llm/llm-pi-ai/src/adapter.ts @@ -115,7 +115,7 @@ export class PiAiAdapter extends LlmAdapter { // Harness-owned and therefore win collisions. headers: requestHeaders(profile.headers), }) - yield* toStreamChunks(events) + yield* toStreamChunks(events, model.contextWindow) } finally { options.signal?.removeEventListener('abort', onCallerAbort) controller.abort('consumer stopped streaming') diff --git a/packages/llm/llm-pi-ai/src/stream.ts b/packages/llm/llm-pi-ai/src/stream.ts index f7a0d6ee7f..8bc41cded2 100644 --- a/packages/llm/llm-pi-ai/src/stream.ts +++ b/packages/llm/llm-pi-ai/src/stream.ts @@ -10,6 +10,7 @@ import { CallId, CONTEXT_WINDOW_EXCEEDED_CODE, isContextWindowExceededError, LlmError } from '@deepseek-ai/dsh-llm' import type { FinishReason, StreamChunk, TokenUsage } from '@deepseek-ai/dsh-llm' +import { isContextOverflow } from '@earendil-works/pi-ai' import type { AssistantMessage, AssistantMessageEvent, Usage as PiUsage } from '@earendil-works/pi-ai' import { toPiReplayState } from './replay.ts' @@ -30,10 +31,6 @@ export function mapUsage(usage: PiUsage): TokenUsage { function classifyPiAiError(message: string): string { if (/\b(?:401|403)\b/.test(message)) return 'AUTH' if (/\b429\b|rate.?limit/i.test(message)) return 'RATE_LIMIT' - // TODO: Classify the full message with pi-ai's isContextOverflow() and the - // resolved model's contextWindow so provider-specific and usage-based overflows - // reach automatic compaction. - if (isContextWindowExceededError(message)) return CONTEXT_WINDOW_EXCEEDED_CODE if (/\b400\b|invalid.?request/i.test(message)) return 'INVALID_REQUEST' if (/\b5\d\d\b/.test(message)) return 'SERVER' return 'PI_AI_ERROR' @@ -42,9 +39,22 @@ function classifyPiAiError(message: string): string { /** * Map a terminal pi-ai event to the harness finish reason. * @param message - the assistant message carried by the `done` or `error` event. + * @param contextWindow - resolved catalog capacity for usage-based overflow detection. * @returns the harness reason; `error` yields `{kind: 'error'}` with a code classified from the error text. */ -export function mapStopReason(message: AssistantMessage): FinishReason { +export function mapStopReason(message: AssistantMessage, contextWindow?: number): FinishReason { + const piAiOverflow = isContextOverflow(message, contextWindow) + const harnessOverflow = message.stopReason === 'error' + && message.errorMessage !== undefined + && isContextWindowExceededError(message.errorMessage) + if (piAiOverflow || harnessOverflow) { + return { + kind: 'error', + message: message.errorMessage ?? `pi-ai detected context overflow for model "${message.model}"`, + code: CONTEXT_WINDOW_EXCEEDED_CODE, + } + } + switch (message.stopReason) { case 'stop': return { kind: 'stop' } case 'length': return { kind: 'max-tokens' } @@ -62,10 +72,14 @@ export function mapStopReason(message: AssistantMessage): FinishReason { * mid-stream — failures arrive as `error` events, which become error/aborted * `finish` chunks (the harness protocol's other error-delivery style). * @param events - one assistant turn's pi-ai event stream. + * @param contextWindow - resolved catalog capacity for usage-based overflow detection. * @returns the harness chunks, ending with `usage` then `finish`; throws * `LlmError` (`STREAM_CLOSED`) if the source ends without a terminal event. */ -export async function* toStreamChunks(events: AsyncIterable): AsyncGenerator { +export async function* toStreamChunks( + events: AsyncIterable, + contextWindow?: number, +): AsyncGenerator { // pi-ai contentIndex ↔ our block index map 1:1 (both count blocks from 0 // in stream order), but we track ids per index for tool calls. const toolIds = new Map() @@ -128,13 +142,17 @@ export async function* toStreamChunks(events: AsyncIterable { const result = await assemble(ctx, { model: 'deepseek-v4-flash', messages: [] }) expect(result.finish).toMatchObject({ kind: 'error', code }) }) + + it('uses the resolved catalog context window for usage-based overflow detection', async () => { + const model = getModels('deepseek').find(candidate => candidate.id === 'deepseek-v4-flash') + if (model === undefined) throw new Error('deepseek-v4-flash missing from pi-ai test catalog') + const events = [ + '{"choices":[{"delta":{"role":"assistant","content":""},"index":0,"finish_reason":null}]}', + JSON.stringify({ + choices: [{ delta: {}, index: 0, finish_reason: 'stop' }], + usage: { prompt_tokens: model.contextWindow + 1, completion_tokens: 0 }, + }), + '[DONE]', + ] + const server = await mockServer([{ events }]) + const ctx = await harness(server.url) + + const result = await assemble(ctx, { model: model.id, messages: [] }) + + expect(result.finish).toEqual({ + kind: 'error', + message: `pi-ai detected context overflow for model "${model.id}"`, + code: CONTEXT_WINDOW_EXCEEDED_CODE, + }) + }) }) describe('provider profile lifecycle', () => { diff --git a/packages/llm/llm-pi-ai/tests/convert.spec.ts b/packages/llm/llm-pi-ai/tests/convert.spec.ts index 6c58f67bb7..a2f37ce511 100644 --- a/packages/llm/llm-pi-ai/tests/convert.spec.ts +++ b/packages/llm/llm-pi-ai/tests/convert.spec.ts @@ -537,6 +537,34 @@ describe('mapStopReason / mapUsage', () => { }))).toMatchObject({ kind: 'error', code: 'INVALID_REQUEST' }) }) + it('uses pi-ai provider-specific overflow classification without losing rate-limit exclusions', () => { + expect(mapStopReason(assistant({ + stopReason: 'error', + errorMessage: 'prompt is too long: 213462 tokens > 200000 maximum', + }))).toMatchObject({ kind: 'error', code: CONTEXT_WINDOW_EXCEEDED_CODE }) + expect(mapStopReason(assistant({ + stopReason: 'error', + errorMessage: 'ThrottlingException: Too many tokens, rate limit reached', + }))).toMatchObject({ kind: 'error', code: 'RATE_LIMIT' }) + }) + + it('uses the resolved context window for silent and length-stop overflows', () => { + const silent = assistant({ stopReason: 'stop', usage: usage(101, 0) }) + expect(mapStopReason(silent)).toEqual({ kind: 'stop' }) + expect(mapStopReason(silent, 100)).toEqual({ + kind: 'error', + message: 'pi-ai detected context overflow for model "deepseek-v4-flash"', + code: CONTEXT_WINDOW_EXCEEDED_CODE, + }) + + const truncated = assistant({ stopReason: 'length', usage: usage(80, 0, 19) }) + expect(mapStopReason(truncated)).toEqual({ kind: 'max-tokens' }) + expect(mapStopReason(truncated, 100)).toMatchObject({ + kind: 'error', + code: CONTEXT_WINDOW_EXCEEDED_CODE, + }) + }) + it('maps cache fields only when nonzero', () => { expect(mapUsage(usage(10, 5, 8, 2))).toEqual({ inputTokens: 10,