mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
fix(bedrock): emit tool_call_end events for native tool streaming
Fixes ROO-311: Write to File truncates filenames at special characters before file content The Bedrock provider was not emitting tool_call_end events after tool use streaming, unlike other providers (e.g., OpenAI, Anthropic). The NativeToolCallParser requires these events to finalize tool calls and flush complete arguments. Without them, the parser would continue accumulating arguments in a buffer that never got finalized, causing filenames with special characters like parentheses to get truncated. Changes: - Add contentBlockStop event type to StreamEvent interface - Track tool call IDs when contentBlockStart receives a tool use block - Emit tool_call_end when contentBlockStop is received for a tracked tool block - Emit tool_call_end as fallback on messageStop with stopReason: tool_use Tests added for: - tool_call_end emission when contentBlockStop is received - tool_call_end emission via messageStop fallback - Multiple tool_call_end events for parallel tool calls - No tool_call_end when stopReason is not tool_use
This commit is contained in:
daniel-lxs
parent
2eebf3c2c4
commit
a4331cc489
2 changed files with 374 additions and 2 deletions
|
|
@ -1244,4 +1244,339 @@ describe("AwsBedrockHandler", () => {
|
|||
expect(mockCaptureException).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("native tool streaming support", () => {
|
||||
beforeEach(() => {
|
||||
mockConverseStreamCommand.mockReset()
|
||||
})
|
||||
|
||||
it("should emit tool_call_end when contentBlockStop is received for a tool use block", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
awsAccessKey: "test",
|
||||
awsSecretKey: "test",
|
||||
awsRegion: "us-east-1",
|
||||
})
|
||||
|
||||
// Get the mock send from the handler instance
|
||||
const clientInstance =
|
||||
vi.mocked(BedrockRuntimeClient).mock.results[vi.mocked(BedrockRuntimeClient).mock.results.length - 1]
|
||||
?.value
|
||||
const mockSendFn = clientInstance?.send as ReturnType<typeof vi.fn>
|
||||
|
||||
// Mock the stream response with tool use events
|
||||
const mockStream = (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
contentBlockIndex: 0,
|
||||
start: {
|
||||
toolUse: {
|
||||
toolUseId: "toolu_123",
|
||||
name: "write_to_file",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
contentBlockIndex: 0,
|
||||
delta: {
|
||||
toolUse: {
|
||||
input: '{"path": "test(1).ts", "content": "hello"}',
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockStop: {
|
||||
contentBlockIndex: 0,
|
||||
},
|
||||
}
|
||||
yield { messageStop: { stopReason: "tool_use" } }
|
||||
})()
|
||||
|
||||
mockSendFn.mockResolvedValueOnce({ stream: mockStream })
|
||||
|
||||
const messages: Anthropic.Messages.MessageParam[] = [
|
||||
{
|
||||
role: "user",
|
||||
content: "Create a file test(1).ts",
|
||||
},
|
||||
]
|
||||
|
||||
const generator = handler.createMessage("You are a helpful assistant", messages, {
|
||||
taskId: "test-task-1",
|
||||
tools: [
|
||||
{
|
||||
type: "function" as const,
|
||||
function: {
|
||||
name: "write_to_file",
|
||||
description: "Write content to a file",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: {
|
||||
path: { type: "string" },
|
||||
content: { type: "string" },
|
||||
},
|
||||
required: ["path", "content"],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of generator) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Should have tool_call_partial and tool_call_end events
|
||||
const partialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
|
||||
const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
|
||||
|
||||
expect(partialChunks.length).toBeGreaterThan(0)
|
||||
expect(endChunks).toHaveLength(1)
|
||||
expect(endChunks[0]).toEqual({ type: "tool_call_end", id: "toolu_123" })
|
||||
})
|
||||
|
||||
it("should emit tool_call_end via messageStop fallback when contentBlockStop is not received", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
awsAccessKey: "test",
|
||||
awsSecretKey: "test",
|
||||
awsRegion: "us-east-1",
|
||||
})
|
||||
|
||||
// Get the mock send from the handler instance
|
||||
const clientInstance =
|
||||
vi.mocked(BedrockRuntimeClient).mock.results[vi.mocked(BedrockRuntimeClient).mock.results.length - 1]
|
||||
?.value
|
||||
const mockSendFn = clientInstance?.send as ReturnType<typeof vi.fn>
|
||||
|
||||
// Mock the stream response WITHOUT contentBlockStop (fallback case)
|
||||
const mockStream = (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
contentBlockIndex: 0,
|
||||
start: {
|
||||
toolUse: {
|
||||
toolUseId: "toolu_456",
|
||||
name: "read_file",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
contentBlockIndex: 0,
|
||||
delta: {
|
||||
toolUse: {
|
||||
input: '{"path": "test.ts"}',
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// No contentBlockStop - directly to messageStop
|
||||
yield { messageStop: { stopReason: "tool_use" } }
|
||||
})()
|
||||
|
||||
mockSendFn.mockResolvedValueOnce({ stream: mockStream })
|
||||
|
||||
const messages: Anthropic.Messages.MessageParam[] = [
|
||||
{
|
||||
role: "user",
|
||||
content: "Read a file",
|
||||
},
|
||||
]
|
||||
|
||||
const generator = handler.createMessage("You are a helpful assistant", messages, {
|
||||
taskId: "test-task-2",
|
||||
tools: [
|
||||
{
|
||||
type: "function" as const,
|
||||
function: {
|
||||
name: "read_file",
|
||||
description: "Read a file",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: {
|
||||
path: { type: "string" },
|
||||
},
|
||||
required: ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of generator) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Should emit tool_call_end via messageStop fallback
|
||||
const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
|
||||
expect(endChunks).toHaveLength(1)
|
||||
expect(endChunks[0]).toEqual({ type: "tool_call_end", id: "toolu_456" })
|
||||
})
|
||||
|
||||
it("should emit multiple tool_call_end events for parallel tool calls", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
awsAccessKey: "test",
|
||||
awsSecretKey: "test",
|
||||
awsRegion: "us-east-1",
|
||||
})
|
||||
|
||||
// Get the mock send from the handler instance
|
||||
const clientInstance =
|
||||
vi.mocked(BedrockRuntimeClient).mock.results[vi.mocked(BedrockRuntimeClient).mock.results.length - 1]
|
||||
?.value
|
||||
const mockSendFn = clientInstance?.send as ReturnType<typeof vi.fn>
|
||||
|
||||
// Mock the stream response with multiple tool use blocks
|
||||
const mockStream = (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
// First tool
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
contentBlockIndex: 0,
|
||||
start: {
|
||||
toolUse: {
|
||||
toolUseId: "toolu_first",
|
||||
name: "read_file",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
contentBlockIndex: 0,
|
||||
delta: {
|
||||
toolUse: { input: '{"path": "file1.ts"}' },
|
||||
},
|
||||
},
|
||||
}
|
||||
yield { contentBlockStop: { contentBlockIndex: 0 } }
|
||||
// Second tool
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
contentBlockIndex: 1,
|
||||
start: {
|
||||
toolUse: {
|
||||
toolUseId: "toolu_second",
|
||||
name: "read_file",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
contentBlockIndex: 1,
|
||||
delta: {
|
||||
toolUse: { input: '{"path": "file2.ts"}' },
|
||||
},
|
||||
},
|
||||
}
|
||||
yield { contentBlockStop: { contentBlockIndex: 1 } }
|
||||
yield { messageStop: { stopReason: "tool_use" } }
|
||||
})()
|
||||
|
||||
mockSendFn.mockResolvedValueOnce({ stream: mockStream })
|
||||
|
||||
const messages: Anthropic.Messages.MessageParam[] = [
|
||||
{
|
||||
role: "user",
|
||||
content: "Read two files",
|
||||
},
|
||||
]
|
||||
|
||||
const generator = handler.createMessage("You are a helpful assistant", messages, {
|
||||
taskId: "test-task-3",
|
||||
tools: [
|
||||
{
|
||||
type: "function" as const,
|
||||
function: {
|
||||
name: "read_file",
|
||||
description: "Read a file",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: { path: { type: "string" } },
|
||||
required: ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of generator) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Should emit tool_call_end for both tools
|
||||
const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
|
||||
expect(endChunks).toHaveLength(2)
|
||||
expect(endChunks).toEqual([
|
||||
{ type: "tool_call_end", id: "toolu_first" },
|
||||
{ type: "tool_call_end", id: "toolu_second" },
|
||||
])
|
||||
})
|
||||
|
||||
it("should not emit tool_call_end when stopReason is not tool_use", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
awsAccessKey: "test",
|
||||
awsSecretKey: "test",
|
||||
awsRegion: "us-east-1",
|
||||
})
|
||||
|
||||
// Get the mock send from the handler instance
|
||||
const clientInstance =
|
||||
vi.mocked(BedrockRuntimeClient).mock.results[vi.mocked(BedrockRuntimeClient).mock.results.length - 1]
|
||||
?.value
|
||||
const mockSendFn = clientInstance?.send as ReturnType<typeof vi.fn>
|
||||
|
||||
// Mock a text response (no tool use)
|
||||
const mockStream = (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
contentBlockIndex: 0,
|
||||
start: { text: "" },
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
contentBlockIndex: 0,
|
||||
delta: { text: "Hello, how can I help you?" },
|
||||
},
|
||||
}
|
||||
yield { contentBlockStop: { contentBlockIndex: 0 } }
|
||||
yield { messageStop: { stopReason: "end_turn" } }
|
||||
})()
|
||||
|
||||
mockSendFn.mockResolvedValueOnce({ stream: mockStream })
|
||||
|
||||
const messages: Anthropic.Messages.MessageParam[] = [
|
||||
{
|
||||
role: "user",
|
||||
content: "Hello",
|
||||
},
|
||||
]
|
||||
|
||||
const generator = handler.createMessage("You are a helpful assistant", messages)
|
||||
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of generator) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Should not have any tool_call_end events
|
||||
const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
|
||||
expect(endChunks).toHaveLength(0)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -147,6 +147,9 @@ export interface StreamEvent {
|
|||
}
|
||||
contentBlockStart?: ContentBlockStartEvent
|
||||
contentBlockDelta?: ContentBlockDeltaEvent
|
||||
contentBlockStop?: {
|
||||
contentBlockIndex?: number
|
||||
}
|
||||
metadata?: {
|
||||
usage?: {
|
||||
inputTokens: number
|
||||
|
|
@ -505,6 +508,10 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
throw new Error("No stream available in the response")
|
||||
}
|
||||
|
||||
// Track active tool calls by content block index
|
||||
// Maps contentBlockIndex -> toolUseId for emitting tool_call_end events
|
||||
const activeToolCalls = new Map<number, string>()
|
||||
|
||||
for await (const chunk of response.stream) {
|
||||
// Parse the chunk as JSON if it's a string (for tests)
|
||||
let streamEvent: StreamEvent
|
||||
|
|
@ -618,9 +625,14 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
else if (cbStart.start?.toolUse || cbStart.contentBlock?.toolUse) {
|
||||
const toolUse = cbStart.start?.toolUse || cbStart.contentBlock?.toolUse
|
||||
if (toolUse) {
|
||||
const blockIndex = cbStart.contentBlockIndex ?? 0
|
||||
// Track this tool call for emitting tool_call_end later
|
||||
if (toolUse.toolUseId) {
|
||||
activeToolCalls.set(blockIndex, toolUse.toolUseId)
|
||||
}
|
||||
yield {
|
||||
type: "tool_call_partial",
|
||||
index: cbStart.contentBlockIndex ?? 0,
|
||||
index: blockIndex,
|
||||
id: toolUse.toolUseId,
|
||||
name: toolUse.name,
|
||||
arguments: undefined,
|
||||
|
|
@ -683,8 +695,33 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
}
|
||||
continue
|
||||
}
|
||||
// Handle message stop
|
||||
|
||||
// Handle content block stop - emit tool_call_end for completed tool use blocks
|
||||
if (streamEvent.contentBlockStop) {
|
||||
const blockIndex = streamEvent.contentBlockStop.contentBlockIndex
|
||||
if (blockIndex !== undefined && activeToolCalls.has(blockIndex)) {
|
||||
const toolUseId = activeToolCalls.get(blockIndex)!
|
||||
yield {
|
||||
type: "tool_call_end",
|
||||
id: toolUseId,
|
||||
}
|
||||
activeToolCalls.delete(blockIndex)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle message stop - emit tool_call_end for any remaining active tool calls
|
||||
// This is a fallback in case contentBlockStop wasn't received
|
||||
if (streamEvent.messageStop) {
|
||||
if (streamEvent.messageStop.stopReason === "tool_use") {
|
||||
for (const [, toolUseId] of activeToolCalls) {
|
||||
yield {
|
||||
type: "tool_call_end",
|
||||
id: toolUseId,
|
||||
}
|
||||
}
|
||||
activeToolCalls.clear()
|
||||
}
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue