diff --git a/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts b/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts index da6c49548..765f63a96 100644 --- a/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts +++ b/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts @@ -195,6 +195,101 @@ describe('createChatOrchestratorRuntime', () => { expect(providerUserMessage).not.toHaveProperty('tools') }) + // ROOT CAUSE: + // + // xsAI kept the assistant tool call and tool result in its private message copy. + // AIRI stored only UI slices, then removed those slices from the next provider request. + // + // We fixed this by storing the provider transcript on the finalized UI message. + // The next request expands that transcript back into chronological provider messages. + it('includes completed tool rounds in the next provider request', async () => { + const harness = createHarness() + + harness.stream.mockImplementationOnce(async (_model, _chatProvider, messages, options) => { + await options?.onStreamEvent?.({ + type: 'tool-call', + toolCallId: 'call-weather', + toolName: 'weather', + args: '{}', + } as StreamEvent) + await options?.onStreamEvent?.({ + type: 'tool-result', + toolCallId: 'call-weather', + result: 'sunny', + } as StreamEvent) + await options?.onStreamEvent?.({ type: 'text-delta', text: 'The weather is sunny.' }) + + await (options as StreamOptions & { onMessages?: (messages: Message[]) => void })?.onMessages?.([ + ...messages, + { + role: 'assistant', + content: '', + tool_calls: [ + { + id: 'call-weather', + type: 'function', + function: { + name: 'weather', + arguments: '{}', + }, + }, + ], + }, + { + role: 'tool', + tool_call_id: 'call-weather', + content: 'sunny', + }, + { + role: 'assistant', + content: 'The weather is sunny.', + }, + ]) + }) + + await harness.runtime.ingest('What is the weather?', { + model: 'gpt-test', + chatProvider: provider, + }) + await harness.runtime.ingest('Can you repeat that?', { + model: 'gpt-test', + chatProvider: provider, + }) + + const messages = harness.stream.mock.calls[1]?.[2] + + expect(messages?.map(message => message.role)).toEqual([ + 'system', + 'user', + 'assistant', + 'tool', + 'assistant', + 'user', + ]) + expect(messages?.[2]).toMatchObject({ + role: 'assistant', + tool_calls: [ + { + id: 'call-weather', + type: 'function', + function: { + name: 'weather', + arguments: '{}', + }, + }, + ], + }) + expect(messages?.[3]).toEqual({ + role: 'tool', + tool_call_id: 'call-weather', + content: 'sunny', + }) + expect(messages?.[4]).toEqual({ + role: 'assistant', + content: 'The weather is sunny.', + }) + }) + /** * @example * Hook order and prompt composition stay compatible with the stage-ui facade. diff --git a/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts b/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts index b528ba1a6..2a71115f2 100644 --- a/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts +++ b/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts @@ -3,7 +3,7 @@ import type { CommonContentPart, Message, ToolMessage } from '@xsai/shared-chat' import type { AgentContextPort } from '../contracts/context-port' import type { AgentForegroundStreamPort } from '../contracts/stream-port' -import type { ChatAssistantMessage, ChatHistoryItem, ChatSlices, ChatStreamEventContext, ChatToolReference, ContextMessage, StreamingAssistantMessage } from '../types/chat' +import type { ChatAssistantMessage, ChatHistoryItem, ChatSlices, ChatStreamEventContext, ChatToolReference, ContextMessage, ErrorMessage, StreamingAssistantMessage } from '../types/chat' import type { LlmUsage, StreamEvent, StreamOptions } from '../types/llm' import { createQueue } from '@proj-airi/stream-kit' @@ -396,23 +396,33 @@ export function createChatOrchestratorRuntime(deps: ChatOrchestratorRuntimeDeps) return fallbackCreatedAt } - function buildProviderMessages(sessionMessagesForSend: ChatHistoryItem[]) { + function buildProviderMessages(sessionMessagesForSend: ChatHistoryItem[]): Array { const nowTs = now() - return sessionMessagesForSend.map((msg) => { + return sessionMessagesForSend.flatMap((msg) => { const { context: _context, id: _id, createdAt: _createdAt, tools: _tools, ...withoutContext } = msg const rawMessage = unwrapMessage(withoutContext) if (rawMessage.role === 'user') { - return prependTextToContent(rawMessage, formatTimePrefix(getStablePromptTimestamp(msg, nowTs))) + return [prependTextToContent(rawMessage, formatTimePrefix(getStablePromptTimestamp(msg, nowTs)))] } if (rawMessage.role === 'assistant') { - const { slices: _slices, tool_results: _toolResults, categorization: _categorization, ...rest } = rawMessage as ChatAssistantMessage - return unwrapMessage(rest) + const { + slices: _slices, + tool_results: _toolResults, + providerTranscript, + categorization: _categorization, + ...rest + } = rawMessage as ChatAssistantMessage + + if (providerTranscript?.length) + return providerTranscript.map(message => unwrapMessage(message)) + + return [unwrapMessage(rest)] } - return rawMessage + return [rawMessage] }) } @@ -711,6 +721,8 @@ export function createChatOrchestratorRuntime(deps: ChatOrchestratorRuntimeDeps) const llmRequestStartedAt = monotonicNow() let llmFirstTokenEmitted = false let generationUsage: LlmUsage = { source: 'unavailable' } + let providerTranscript: Message[] | undefined + const providerInputMessageCount = newMessages.length deps.onLlmRequestStarted?.({ ...correlation, model: options.model, @@ -726,6 +738,16 @@ export function createChatOrchestratorRuntime(deps: ChatOrchestratorRuntimeDeps) }, tools: options.tools, waitForTools: true, + onMessages: (messages) => { + const currentTurnMessages = messages.slice(providerInputMessageCount) + const hasToolRound = currentTurnMessages.some(message => + message.role === 'tool' + || (message.role === 'assistant' && Boolean(message.tool_calls?.length)), + ) + + if (hasToolRound) + providerTranscript = structuredClone(currentTurnMessages) + }, onUsage: (usage) => { generationUsage = usage deps.onLlmGeneration?.({ @@ -802,6 +824,7 @@ export function createChatOrchestratorRuntime(deps: ChatOrchestratorRuntimeDeps) }) await parser.end() + buildingMessage.providerTranscript = providerTranscript deps.onAssistantResponseRendered?.({ ...correlation, model: options.model, diff --git a/packages/core-agent/src/runtime/llm-service.test.ts b/packages/core-agent/src/runtime/llm-service.test.ts index 397de1f9f..54f4a1e60 100644 --- a/packages/core-agent/src/runtime/llm-service.test.ts +++ b/packages/core-agent/src/runtime/llm-service.test.ts @@ -30,10 +30,11 @@ const provider = { function createMockStreamResult( steps: Promise = Promise.resolve([]), totalUsage: Promise<{ inputTokens: number, outputTokens: number, totalTokens: number } | undefined> = Promise.resolve(undefined), + messages: Promise = Promise.resolve([]), ) { return { steps, - messages: Promise.resolve([]), + messages, usage: Promise.resolve(undefined), totalUsage, } @@ -43,6 +44,76 @@ describe('streamFrom tool errors', () => { beforeEach(() => { streamTextMock.mockReset() }) + + it('emits the final xsAI messages after all tool rounds finish', async () => { + const onMessages = vi.fn() + const finalMessages: Message[] = [ + { role: 'user', content: 'Check the weather.' }, + { + role: 'assistant', + content: '', + tool_calls: [ + { + id: 'call-weather', + type: 'function', + function: { name: 'weather', arguments: '{}' }, + }, + ], + }, + { role: 'tool', tool_call_id: 'call-weather', content: 'sunny' }, + { role: 'assistant', content: 'The weather is sunny.' }, + ] + streamTextMock.mockReturnValueOnce(createMockStreamResult( + Promise.resolve([]), + Promise.resolve(undefined), + Promise.resolve(finalMessages), + )) + + await streamFrom({ + model: 'model-a', + chatProvider: provider, + messages: finalMessages.slice(0, 1), + options: { onMessages }, + }) + + expect(onMessages).toHaveBeenCalledTimes(1) + expect(onMessages).toHaveBeenCalledWith(finalMessages) + }) + + it('ignores provider errors after steps resolve while final messages are pending', async () => { + let onEvent: ((event: unknown) => Promise) | undefined + let resolveMessages: ((messages: Message[]) => void) | undefined + const messages = new Promise((resolve) => { + resolveMessages = resolve + }) + + streamTextMock.mockImplementationOnce((options: { onEvent: (event: unknown) => Promise }) => { + onEvent = options.onEvent + return createMockStreamResult(Promise.resolve([]), Promise.resolve(undefined), messages) + }) + + // ROOT CAUSE: + // + // Final message persistence used to delay the steps-settled marker. A late + // provider error could then reject a stream whose authoritative steps + // promise had already resolved. + // + // We mark steps settled before awaiting the final transcript, while still + // treating transcript persistence failures as real stream failures. + const pending = streamFrom({ + model: 'model-a', + chatProvider: provider, + messages: [{ role: 'user', content: 'hello' }] as Message[], + }) + + await vi.waitFor(() => expect(onEvent).toBeTypeOf('function')) + await Promise.resolve() + await onEvent!({ type: 'error', message: 'stream failed', cause: new Error('stream failed') }) + resolveMessages?.([]) + + await expect(pending).resolves.toBeUndefined() + }) + it('requests final streaming usage and emits the reported token totals once', async () => { const onUsage = vi.fn() streamTextMock.mockReturnValueOnce(createMockStreamResult( diff --git a/packages/core-agent/src/runtime/llm-service.ts b/packages/core-agent/src/runtime/llm-service.ts index 85ae89c9c..b26a6c6e5 100644 --- a/packages/core-agent/src/runtime/llm-service.ts +++ b/packages/core-agent/src/runtime/llm-service.ts @@ -222,6 +222,19 @@ export async function streamFrom({ // Ignore any late provider error event emitted after xsAI has already // resolved the authoritative full-step lifecycle. stepsSettled = true + try { + const finalMessages = await streamResult.messages + await options?.onMessages?.(finalMessages) + } + catch (error) { + // Transcript persistence is part of the completed response contract, + // unlike late provider events and optional usage observation. + if (!settled) { + settled = true + reject(error) + } + return + } try { await options?.onStreamEvent?.({ type: 'finish' } as const) } @@ -266,6 +279,8 @@ export async function streamFrom({ rejectOnce(error) console.error('Stream steps error:', error) }) + // `steps` can reject before the success path awaits `messages`. + // Keep this rejection sink so xsAI cannot create an unhandled rejection. void streamResult.messages.catch(error => console.error('Stream messages error:', error)) void streamResult.usage.catch(error => console.error('Stream usage error:', error)) // `steps` and `totalUsage` reject independently when xsAI fails a diff --git a/packages/core-agent/src/types/chat.ts b/packages/core-agent/src/types/chat.ts index 710286517..17cbcd13e 100644 --- a/packages/core-agent/src/types/chat.ts +++ b/packages/core-agent/src/types/chat.ts @@ -27,6 +27,14 @@ export interface ChatAssistantMessage extends AssistantMessage { isError?: boolean result?: string | CommonContentPart[] }[] + /** + * Exact provider messages that xsAI added for this assistant turn. + * + * The chat UI keeps one aggregated assistant message. Tool loops can contain + * multiple assistant and tool messages, so this transcript preserves their + * protocol order for the next provider request. + */ + providerTranscript?: Message[] categorization?: { speech: string reasoning: string diff --git a/packages/core-agent/src/types/llm.ts b/packages/core-agent/src/types/llm.ts index 7bc452fd2..19ff76502 100644 --- a/packages/core-agent/src/types/llm.ts +++ b/packages/core-agent/src/types/llm.ts @@ -25,6 +25,8 @@ export interface StreamOptions { abortSignal?: AbortSignal headers?: Record onStreamEvent?: (event: StreamEvent) => void | Promise + /** Called once with the final xsAI message list after all tool rounds finish. */ + onMessages?: (messages: Message[]) => void | Promise /** Called once after the full stream, including tool rounds, has settled. */ onUsage?: (usage: LlmUsage) => void | Promise /** Internal correlation kept out of the provider request body. */ diff --git a/packages/stage-ui/src/stores/tool-call-rerun.test.ts b/packages/stage-ui/src/stores/tool-call-rerun.test.ts index a42ea1572..9436c8115 100644 --- a/packages/stage-ui/src/stores/tool-call-rerun.test.ts +++ b/packages/stage-ui/src/stores/tool-call-rerun.test.ts @@ -107,6 +107,44 @@ describe('replaceToolCallResult', () => { }, ]) }) + + it('replaces the matching tool message in the provider transcript', () => { + const message = assistantMessage({ + providerTranscript: [ + { + role: 'assistant', + content: '', + tool_calls: [ + { + id: 'call-weather', + type: 'function', + function: { name: 'weather', arguments: '{}' }, + }, + ], + }, + { + role: 'tool', + tool_call_id: 'call-weather', + content: 'old weather', + }, + { + role: 'assistant', + content: 'The old result was returned.', + }, + ], + }) + + const next = replaceToolCallResult(message, { + id: 'call-weather', + result: 'new weather', + }) + + expect(next.providerTranscript?.[1]).toEqual({ + role: 'tool', + tool_call_id: 'call-weather', + content: 'new weather', + }) + }) }) describe('executeToolCallRerun', () => { diff --git a/packages/stage-ui/src/stores/tool-call-rerun.ts b/packages/stage-ui/src/stores/tool-call-rerun.ts index 4956bc4d1..93ae5d037 100644 --- a/packages/stage-ui/src/stores/tool-call-rerun.ts +++ b/packages/stage-ui/src/stores/tool-call-rerun.ts @@ -45,6 +45,16 @@ export function replaceToolCallResult(message: ChatAssistantMessage, result: Too return { ...message, + providerTranscript: message.providerTranscript?.map((providerMessage) => { + if (providerMessage.role === 'tool' && providerMessage.tool_call_id === result.id) { + return { + ...providerMessage, + content: result.result ?? '', + } + } + + return providerMessage + }), slices: message.slices.map((slice) => { if (slice.type === 'tool-call-result' && slice.id === result.id) return resultSlice