diff --git a/apps/cli/src/ai/agent.ts b/apps/cli/src/ai/agent.ts index 8f17e200..6a1f3724 100644 --- a/apps/cli/src/ai/agent.ts +++ b/apps/cli/src/ai/agent.ts @@ -253,6 +253,8 @@ export async function runServerAgent(options: RunAgentOptions): Promise { maxToolRounds: DEFAULT_MAX_TOOL_ROUNDS, abortSignal, providerSessionId: aiChatId, + locale, + logger: aiLogger ?? undefined, steerMessage, thinkingLevel: thinkingLevel as import('@openchatlab/core').ThinkingLevel | undefined, onConvertToLlm: (filteredMessages) => { @@ -268,12 +270,14 @@ export async function runServerAgent(options: RunAgentOptions): Promise { }, }) - if (result.error) { + if (result.error && !abortSignal?.aborted) { const friendlyMessage = formatAIError(result.error) onEvent({ type: 'error', error: { name: 'AgentError', message: friendlyMessage } }) } - handler.emitStatus('completed', cachedMessages, { force: true }) + handler.emitStatus(abortSignal?.aborted ? 'aborted' : result.error ? 'error' : 'completed', cachedMessages, { + force: true, + }) onEvent({ type: 'done', isFinished: true, usage: result.usage }) } catch (error) { const friendlyMessage = formatAIError(error) diff --git a/apps/desktop/main/ai/agent/index.ts b/apps/desktop/main/ai/agent/index.ts index 8d948d96..d3752a7f 100644 --- a/apps/desktop/main/ai/agent/index.ts +++ b/apps/desktop/main/ai/agent/index.ts @@ -232,6 +232,8 @@ export class Agent { maxToolRounds, abortSignal: this.abortSignal, providerSessionId: this.context.aiChatId, + locale: this.locale, + logger: aiLogger, steerMessage: answerWithoutToolsPrompt, thinkingLevel: this.config.thinkingLevel, streamFn: errorCapturingStreamFn, @@ -319,7 +321,8 @@ export class Agent { const apiPath = typeof apiType === 'string' ? pathMap[apiType] : undefined ctx.url = apiPath ? baseUrl.replace(/\/+$/, '') + apiPath : baseUrl } - if (lastRequestPayload) { + // Output truncation is explained by completion metadata, not the private request body. + if (lastRequestPayload && result.stopReason !== 'length') { try { ctx.requestBody = JSON.stringify(lastRequestPayload, null, 2) } catch { diff --git a/packages/node-runtime/src/ai/__tests__/llm-builder.test.ts b/packages/node-runtime/src/ai/__tests__/llm-builder.test.ts index 884e84c3..8300615b 100644 --- a/packages/node-runtime/src/ai/__tests__/llm-builder.test.ts +++ b/packages/node-runtime/src/ai/__tests__/llm-builder.test.ts @@ -29,6 +29,18 @@ describe('normalizeOpenAICompatibleBaseUrl', () => { }) describe('buildPiModel', () => { + it('reserves output space for reasoning while preserving explicit limits', () => { + for (const config of [ + { provider: 'deepseek', model: 'deepseek-v4-flash' }, + { provider: 'openai', model: 'o3' }, + { provider: 'openai-compatible', model: 'qwen3:8b' }, + ]) { + assert.equal(buildPiModel(config).maxTokens, 16_384) + assert.equal(buildPiModel({ ...config, maxTokens: 4096 }).maxTokens, 4096) + } + assert.equal(buildPiModel({ provider: 'openai', model: 'gpt-4o' }).maxTokens, 4096) + }) + it('builds a Google Generative AI model', () => { const model = buildPiModel({ provider: 'gemini', diff --git a/packages/node-runtime/src/ai/agent/__tests__/core.test.ts b/packages/node-runtime/src/ai/agent/__tests__/core.test.ts index 18490044..7d0beea1 100644 --- a/packages/node-runtime/src/ai/agent/__tests__/core.test.ts +++ b/packages/node-runtime/src/ai/agent/__tests__/core.test.ts @@ -146,6 +146,54 @@ function createOptions(overrides: Partial = {}): AgentCoreOpti } describe('runAgentCore runtime contract', () => { + for (const partialText of ['', 'The main conclusion is']) { + it(`reports output truncation without losing ${partialText ? 'partial text' : 'thinking-only output'}`, async () => { + const events: AgentCoreEvent[] = [] + const content: AssistantMessage['content'] = [{ type: 'thinking', thinking: 'Comparing the evidence' }] + if (partialText) content.push({ type: 'text', text: partialText }) + const result = await runAgentCore( + createOptions({ + locale: 'en-US', + streamFn: createScriptedStream([assistantMessage(content, { stopReason: 'length' })]), + onEvent: (event) => events.push(event), + }) + ) + + assert.match(result.error ?? '', /output limit/i) + assert.equal(result.stopReason, 'length') + assert.deepEqual(result.finalMessages.at(-1)?.content, content) + assert.deepEqual( + events.filter((event) => event.type === 'content').map((event) => event.content), + partialText ? [partialText] : [] + ) + assert.ok(events.some((event) => event.type === 'thinking_delta')) + }) + } + + it('stops on truncated tool arguments instead of spending more tokens on an automatic retry', async () => { + let toolExecuted = false + const result = await runAgentCore( + createOptions({ + tools: [ + createTool('lookup', async () => { + toolExecuted = true + return { content: [{ type: 'text', text: 'Should not run' }], details: {} } + }), + ], + locale: 'en-US', + streamFn: createScriptedStream([ + assistantMessage([{ type: 'toolCall', id: 'partial-call', name: 'lookup', arguments: {} }], { + stopReason: 'length', + }), + ]), + }) + ) + + assert.equal(toolExecuted, false) + assert.equal(result.stopReason, 'length') + assert.match(result.error ?? '', /output limit/i) + }) + it('streams thinking, content, and usage into ChatLab events', async () => { const events: AgentCoreEvent[] = [] const result = await runAgentCore( @@ -182,6 +230,8 @@ describe('runAgentCore runtime contract', () => { cacheReadTokens: 0, cacheWriteTokens: 0, }) + assert.equal(result.error, undefined) + assert.equal(result.stopReason, 'stop') }) it('executes a tool and continues to the final response', async () => { diff --git a/packages/node-runtime/src/ai/agent/core.ts b/packages/node-runtime/src/ai/agent/core.ts index 131a676d..212b2f85 100644 --- a/packages/node-runtime/src/ai/agent/core.ts +++ b/packages/node-runtime/src/ai/agent/core.ts @@ -7,12 +7,18 @@ import { Agent as PiAgentCore } from '@earendil-works/pi-agent-core' import type { AgentEvent as PiAgentEvent, AgentMessage as PiAgentMessage } from '@earendil-works/pi-agent-core' -import { type Message as PiMessage, type Usage as PiUsage, clampThinkingLevel } from '@earendil-works/pi-ai' +import { + type AssistantMessage, + type Message as PiMessage, + type Usage as PiUsage, + clampThinkingLevel, +} from '@earendil-works/pi-ai' import { StreamingThinkTagParser, needsStreamingThinkParsing } from '@openchatlab/core' import type { ToolProgress } from '@openchatlab/shared-types' import type { AgentCoreOptions, AgentCoreResult, AgentTokenUsage } from './types' import { initTokenizer } from '../tokenizer' +import { createAiTranslate } from '../i18n' import { streamSimple as defaultStreamSimple } from '../pi-runtime' import { DEFAULT_MAX_TOOL_ROUNDS } from './constants' import { toPiHistoryMessages, type ReplayOptions } from './history' @@ -51,6 +57,7 @@ export async function runAgentCore(options: AgentCoreOptions): Promise { if (!usage) return @@ -62,7 +69,7 @@ export async function runAgentCore(options: AgentCoreOptions): Promise (result as { isError?: boolean }).isError === true ? { isError: true } : undefined, - shouldStopAfterTurn: async () => hasCompletedFinalAnswerTurn, + // A cut-off tool call must not trigger an automatic paid retry either. + shouldStopAfterTurn: async ({ message }) => message.stopReason === 'length' || hasCompletedFinalAnswerTurn, prepareNextTurnWithContext: ({ context }) => hasReachedToolRoundLimit ? { @@ -238,6 +246,7 @@ export async function runAgentCore(options: AgentCoreOptions): Promise steerMessage?: string /** Override the thinking level for this request. Clamped to what the model supports. */ thinkingLevel?: ThinkingLevel @@ -68,6 +71,7 @@ export interface AgentCoreOptions { export interface AgentCoreResult { usage: AgentTokenUsage error?: string + stopReason?: StopReason finalMessages: Message[] toolsUsed: string[] toolRounds: number diff --git a/packages/node-runtime/src/ai/cross-chat-agent.ts b/packages/node-runtime/src/ai/cross-chat-agent.ts index e29b24e4..964596d6 100644 --- a/packages/node-runtime/src/ai/cross-chat-agent.ts +++ b/packages/node-runtime/src/ai/cross-chat-agent.ts @@ -175,6 +175,8 @@ export async function runCrossChatAgent(options: RunCrossChatAgentOptions): Prom maxToolRounds: DEFAULT_MAX_TOOL_ROUNDS, abortSignal, providerSessionId: aiChatId, + locale, + logger, steerMessage: createAiTranslate(locale)('ai.agent.answerWithoutTools'), thinkingLevel, onConvertToLlm: (messages) => { @@ -196,9 +198,10 @@ export async function runCrossChatAgent(options: RunCrossChatAgentOptions): Prom if (result.error) { onEvent({ type: 'error', error: { name: 'AgentError', message: formatAIError(result.error) } }) } - handler.emitStatus('completed', cachedMessages, { force: true }) - logger?.info('CrossChatAgent', 'Cross-chat agent execution completed', { + handler.emitStatus(result.error ? 'error' : 'completed', cachedMessages, { force: true }) + logger?.info('CrossChatAgent', 'Cross-chat agent execution ended', { aiChatId, + stopReason: result.stopReason, toolRounds: result.toolRounds, toolsUsed: result.toolsUsed.length, }) diff --git a/packages/node-runtime/src/ai/i18n/locales/en-US.ts b/packages/node-runtime/src/ai/i18n/locales/en-US.ts index 4ff4e44f..4ea50864 100644 --- a/packages/node-runtime/src/ai/i18n/locales/en-US.ts +++ b/packages/node-runtime/src/ai/i18n/locales/en-US.ts @@ -259,6 +259,8 @@ Returned summaries are brief descriptions of each segment, helping quickly locat // ===== AI Agent system prompts ===== agent: { answerWithoutTools: 'Please answer based on the information already retrieved, do not call any more tools.', + outputLimitReached: + 'The response reached the model output limit (including reasoning) and is incomplete. Generated content has been kept. Send "continue" or narrow your question and try again.', toolError: 'Error: {{error}}', currentDateIs: 'Current date is', chatContext: { diff --git a/packages/node-runtime/src/ai/i18n/locales/ja-JP.ts b/packages/node-runtime/src/ai/i18n/locales/ja-JP.ts index f758029a..34fb38ff 100644 --- a/packages/node-runtime/src/ai/i18n/locales/ja-JP.ts +++ b/packages/node-runtime/src/ai/i18n/locales/ja-JP.ts @@ -258,6 +258,8 @@ export default { agent: { answerWithoutTools: '取得済みの情報に基づいて回答してください。これ以上ツールを呼び出さないでください。', + outputLimitReached: + 'モデルの出力上限（思考過程を含む）に達したため、回答は未完了です。生成済みの内容は保持されています。「続けて」と送信するか、質問の範囲を絞って再試行してください。', toolError: 'エラー: {{error}}', currentDateIs: '現在の日付は', chatContext: { diff --git a/packages/node-runtime/src/ai/i18n/locales/zh-CN.ts b/packages/node-runtime/src/ai/i18n/locales/zh-CN.ts index d2394131..caffa166 100644 --- a/packages/node-runtime/src/ai/i18n/locales/zh-CN.ts +++ b/packages/node-runtime/src/ai/i18n/locales/zh-CN.ts @@ -249,6 +249,8 @@ export default { agent: { answerWithoutTools: '请根据已获取的信息给出回答，不要再调用工具。', + outputLimitReached: + '本次回复达到模型输出上限（包含思考过程），回答未完成。已保留生成的内容，你可以发送“继续”或缩小问题范围后重试。', toolError: '错误: {{error}}', currentDateIs: '当前日期是', chatContext: { diff --git a/packages/node-runtime/src/ai/i18n/locales/zh-TW.ts b/packages/node-runtime/src/ai/i18n/locales/zh-TW.ts index dee8025b..b1ba0aea 100644 --- a/packages/node-runtime/src/ai/i18n/locales/zh-TW.ts +++ b/packages/node-runtime/src/ai/i18n/locales/zh-TW.ts @@ -249,6 +249,8 @@ export default { agent: { answerWithoutTools: '請根據已取得的資訊給出回答，不要再呼叫工具。', + outputLimitReached: + '本次回覆達到模型輸出上限（包含思考過程），回答尚未完成。已保留產生的內容，你可以傳送「繼續」或縮小問題範圍後重試。', toolError: '錯誤: {{error}}', currentDateIs: '目前日期是', chatContext: { diff --git a/packages/node-runtime/src/ai/llm-builder.ts b/packages/node-runtime/src/ai/llm-builder.ts index 6e2fdd41..89bf4d7f 100644 --- a/packages/node-runtime/src/ai/llm-builder.ts +++ b/packages/node-runtime/src/ai/llm-builder.ts @@ -169,7 +169,8 @@ export function buildPiModel(config: PiModelConfig, options?: BuildPiModelOption input: ['text'], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow, - maxTokens: config.maxTokens ?? 4096, + // Reasoning and visible text share the output budget on compatible APIs. + maxTokens: config.maxTokens ?? (reasoning ? 16_384 : 4096), compat, } } diff --git a/tests/ai-agent-output-limit.test.mts b/tests/ai-agent-output-limit.test.mts new file mode 100644 index 00000000..a15520db --- /dev/null +++ b/tests/ai-agent-output-limit.test.mts @@ -0,0 +1,102 @@ +import assert from 'node:assert/strict' +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { test } from 'node:test' +import { runServerAgent } from '../apps/cli/src/ai/agent' +import { AIChatManager } from '../packages/node-runtime/src/ai/chats' +import { runCrossChatAgent } from '../packages/node-runtime/src/ai/cross-chat-agent' +import type { AgentStreamChunk } from '../packages/node-runtime/src/ai/agent/event-handler' +import { buildPiModel } from '../packages/node-runtime/src/ai/llm-builder' + +// Exercise the real Pi SSE parser and both Node entry points, without a real provider or user data. +test('session and global agents surface truncated output instead of reporting successful completion', async (t) => { + const dir = mkdtempSync(join(tmpdir(), 'chatlab-agent-output-limit-')) + const manager = new AIChatManager(dir, { nativeBinding: process.env.CHATLAB_TEST_SQLITE_NATIVE_BINDING }) + const config = { + provider: 'deepseek' as const, + model: 'deepseek-v4-flash', + baseUrl: 'https://provider.invalid/v1', + apiKey: 'test-placeholder', + } + try { + for (const kind of ['session', 'global'] as const) { + for (const scenario of [ + { name: 'thinking only', text: '', finishReason: 'length' }, + { name: 'partial answer', text: 'The conclusion is', finishReason: 'length' }, + { name: 'complete answer', text: 'The conclusion is complete.', finishReason: 'stop' }, + ]) { + await t.test(`${kind}: ${scenario.name}`, async (t) => { + t.mock.method(globalThis, 'fetch', async () => { + const chunks = [ + { delta: { role: 'assistant', reasoning_content: 'Checking the evidence.' }, finish_reason: null }, + { delta: { content: scenario.text }, finish_reason: null }, + { delta: {}, finish_reason: scenario.finishReason }, + ].map((choice) => ({ + id: 'test-completion', + object: 'chat.completion.chunk', + created: 1, + model: config.model, + choices: [{ index: 0, ...choice }], + })) + return new Response( + chunks.map((chunk) => `data: ${JSON.stringify(chunk)}\n\n`).join('') + 'data: [DONE]\n\n', + { + headers: { 'Content-Type': 'text/event-stream' }, + } + ) + }) + const chat = + kind === 'global' + ? manager.createGlobalAIChat('Output limit', 'general_en') + : manager.createAIChat('session-1', 'Output limit', 'general_en') + const events: AgentStreamChunk[] = [] + const common = { + userMessage: 'Explain what a chat assistant is. Do not query any data.', + aiChatId: chat.id, + historyLeafMessageId: null, + locale: 'en-US', + aiChatManager: manager, + onEvent: (event: AgentStreamChunk) => events.push(event), + } + if (kind === 'global') { + await runCrossChatAgent({ + ...common, + piModel: buildPiModel(config), + apiKey: config.apiKey, + tools: [], + memoryService: { list: () => [] }, + }) + } else { + await runServerAgent({ ...common, llmConfig: config }) + } + + assert.equal( + events + .filter((event) => event.type === 'content') + .map((event) => event.content) + .join(''), + scenario.text + ) + assert.ok(events.some((event) => event.type === 'think' && event.content === 'Checking the evidence.')) + const errors = events.filter((event) => event.type === 'error') + const truncated = scenario.finishReason === 'length' + assert.equal(errors.length, truncated ? 1 : 0) + if (truncated) assert.match((errors[0].error as { message: string }).message, /output limit/i) + assert.deepEqual( + events + .filter((event) => event.type === 'status') + .map((event) => event.status?.phase) + .filter((phase) => phase === 'completed' || phase === 'error' || phase === 'aborted'), + [truncated ? 'error' : 'completed'] + ) + assert.equal(events.filter((event) => event.type === 'done').length, 1) + assert.equal(events.at(-1)?.type, 'done') + }) + } + } + } finally { + manager.close() + rmSync(dir, { recursive: true, force: true }) + } +})