mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-07 08:26:51 +00:00
fix: add finish_reason processing to xai.ts provider (#9929)
This commit is contained in:
parent
3356267aa0
commit
2efebf5da9
2 changed files with 93 additions and 0 deletions
|
|
@ -495,5 +495,87 @@ describe("XAIHandler", () => {
|
|||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should yield tool_call_end events when finish_reason is tool_calls", async () => {
|
||||
// Import NativeToolCallParser to set up state
|
||||
const { NativeToolCallParser } = await import("../../../core/assistant-message/NativeToolCallParser")
|
||||
|
||||
// Clear any previous state
|
||||
NativeToolCallParser.clearRawChunkState()
|
||||
|
||||
const handlerWithTools = new XAIHandler({ apiModelId: "grok-3" })
|
||||
|
||||
mockCreate.mockImplementationOnce(() => {
|
||||
return {
|
||||
[Symbol.asyncIterator]: () => ({
|
||||
next: vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce({
|
||||
done: false,
|
||||
value: {
|
||||
choices: [
|
||||
{
|
||||
delta: {
|
||||
tool_calls: [
|
||||
{
|
||||
index: 0,
|
||||
id: "call_xai_test",
|
||||
function: {
|
||||
name: "test_tool",
|
||||
arguments: '{"arg1":"value"}',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
done: false,
|
||||
value: {
|
||||
choices: [
|
||||
{
|
||||
delta: {},
|
||||
finish_reason: "tool_calls",
|
||||
},
|
||||
],
|
||||
usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
|
||||
},
|
||||
})
|
||||
.mockResolvedValueOnce({ done: true }),
|
||||
}),
|
||||
}
|
||||
})
|
||||
|
||||
const stream = handlerWithTools.createMessage("test prompt", [], {
|
||||
taskId: "test-task-id",
|
||||
tools: testTools,
|
||||
toolProtocol: "native",
|
||||
})
|
||||
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
// Simulate what Task.ts does: when we receive tool_call_partial,
|
||||
// process it through NativeToolCallParser to populate rawChunkTracker
|
||||
if (chunk.type === "tool_call_partial") {
|
||||
NativeToolCallParser.processRawChunk({
|
||||
index: chunk.index,
|
||||
id: chunk.id,
|
||||
name: chunk.name,
|
||||
arguments: chunk.arguments,
|
||||
})
|
||||
}
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Should have tool_call_partial and tool_call_end
|
||||
const partialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
|
||||
const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
|
||||
|
||||
expect(partialChunks).toHaveLength(1)
|
||||
expect(endChunks).toHaveLength(1)
|
||||
expect(endChunks[0].id).toBe("call_xai_test")
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import OpenAI from "openai"
|
|||
|
||||
import { type XAIModelId, xaiDefaultModelId, xaiModels } from "@roo-code/types"
|
||||
|
||||
import { NativeToolCallParser } from "../../core/assistant-message/NativeToolCallParser"
|
||||
import type { ApiHandlerOptions } from "../../shared/api"
|
||||
|
||||
import { ApiStream } from "../transform/stream"
|
||||
|
|
@ -83,6 +84,7 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
|
|||
|
||||
for await (const chunk of stream) {
|
||||
const delta = chunk.choices[0]?.delta
|
||||
const finishReason = chunk.choices[0]?.finish_reason
|
||||
|
||||
if (delta?.content) {
|
||||
yield {
|
||||
|
|
@ -111,6 +113,15 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
|
|||
}
|
||||
}
|
||||
|
||||
// Process finish_reason to emit tool_call_end events
|
||||
// This ensures tool calls are finalized even if the stream doesn't properly close
|
||||
if (finishReason) {
|
||||
const endEvents = NativeToolCallParser.processFinishReason(finishReason)
|
||||
for (const event of endEvents) {
|
||||
yield event
|
||||
}
|
||||
}
|
||||
|
||||
if (chunk.usage) {
|
||||
// Extract detailed token information if available
|
||||
// First check for prompt_tokens_details structure (real API response)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue