admin/airi
1
0
Fork 0
mirror of https://github.com/moeru-ai/airi.git synced 2026-08-14 00:48:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(core-agent): preserve tool messages across chat turns (#2257)

Browse Source
This commit is contained in:
Neko
2026-08-12 00:06:04 +08:00
committed by GitHub
parent 1ca59ea5c5
commit 9f5249454b
8 changed files with 270 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts
+95
View File
@@ -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.
packages/core-agent/src/runtime/chat-orchestrator-runtime.ts
+30 -7
View File
@@ -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<Message | ErrorMessage> {
const nowTs = now()
return sessionMessagesForSend.map((msg) => {
return sessionMessagesForSend.flatMap<Message | ErrorMessage>((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,
packages/core-agent/src/runtime/llm-service.test.ts
+72 -1
View File
@@ -30,10 +30,11 @@ const provider = {
function createMockStreamResult(
steps: Promise<unknown[]> = Promise.resolve([]),
totalUsage: Promise<{ inputTokens: number, outputTokens: number, totalTokens: number } | undefined> = Promise.resolve(undefined),
messages: Promise<Message[]> = 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<void>) | undefined
let resolveMessages: ((messages: Message[]) => void) | undefined
const messages = new Promise<Message[]>((resolve) => {
resolveMessages = resolve
})
streamTextMock.mockImplementationOnce((options: { onEvent: (event: unknown) => Promise<void> }) => {
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(
packages/core-agent/src/runtime/llm-service.ts
+15
View File
@@ -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
packages/core-agent/src/types/chat.ts
+8
View File
@@ -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
packages/core-agent/src/types/llm.ts
+2
View File
@@ -25,6 +25,8 @@ export interface StreamOptions {
abortSignal?: AbortSignal
headers?: Record<string, string>
onStreamEvent?: (event: StreamEvent) => void | Promise<void>
/** Called once with the final xsAI message list after all tool rounds finish. */
onMessages?: (messages: Message[]) => void | Promise<void>
/** Called once after the full stream, including tool rounds, has settled. */
onUsage?: (usage: LlmUsage) => void | Promise<void>
/** Internal correlation kept out of the provider request body. */
packages/stage-ui/src/stores/tool-call-rerun.test.ts
+38
View File
@@ -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', () => {
packages/stage-ui/src/stores/tool-call-rerun.ts
+10
View File
@@ -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