Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions src/api/providers/vscode-lm.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,11 @@ import { SELECTOR_SEPARATOR, stringifyVsCodeLmModelSelector } from "../../shared
import { normalizeToolSchema } from "../../utils/json-schema"

import { ApiStream } from "../transform/stream"
import { convertToVsCodeLmMessages, extractTextCountFromMessage } from "../transform/vscode-lm-format"
import {
convertToVsCodeLmMessages,
extractTextCountFromMessage,
validateAndRepairToolSequence,
} from "../transform/vscode-lm-format"

import { BaseProvider } from "./base-provider"
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
Expand Down Expand Up @@ -367,9 +371,16 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
}))

// Convert Anthropic messages to VS Code LM messages
const convertedMessages = convertToVsCodeLmMessages(cleanedMessages)

// Validate and repair tool call/result sequences to ensure VSCode LM API requirements are met
// The API requires that assistant messages with ToolCallParts are immediately followed by
// user messages with matching ToolResultParts
const validatedMessages = validateAndRepairToolSequence(convertedMessages)

const vsCodeLmMessages: vscode.LanguageModelChatMessage[] = [
vscode.LanguageModelChatMessage.Assistant(systemPrompt),
...convertToVsCodeLmMessages(cleanedMessages),
...validatedMessages,
]

// Initialize cancellation token for the request
Expand Down
171 changes: 170 additions & 1 deletion src/api/transform/__tests__/vscode-lm-format.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,12 @@
import { Anthropic } from "@anthropic-ai/sdk"
import * as vscode from "vscode"

import { convertToVsCodeLmMessages, convertToAnthropicRole, extractTextCountFromMessage } from "../vscode-lm-format"
import {
convertToVsCodeLmMessages,
convertToAnthropicRole,
extractTextCountFromMessage,
validateAndRepairToolSequence,
} from "../vscode-lm-format"

// Mock crypto using Vitest
vitest.stubGlobal("crypto", {
Expand Down Expand Up @@ -346,3 +351,167 @@ describe("extractTextCountFromMessage", () => {
expect(result).toBe("result-id")
})
})

describe("validateAndRepairToolSequence", () => {
it("should return empty array for empty input", () => {
const result = validateAndRepairToolSequence([])
expect(result).toEqual([])
})

it("should return messages unchanged when no tool calls present", () => {
const mockTextPart = new (vitest.mocked(vscode).LanguageModelTextPart)("Hello")
const messages = [
{
role: vitest.mocked(vscode).LanguageModelChatMessageRole.User,
content: [mockTextPart],
},
{
role: vitest.mocked(vscode).LanguageModelChatMessageRole.Assistant,
content: [mockTextPart],
},
] as vscode.LanguageModelChatMessage[]

const result = validateAndRepairToolSequence(messages)
expect(result).toHaveLength(2)
expect(result[0]).toBe(messages[0])
expect(result[1]).toBe(messages[1])
})

it("should return messages unchanged when tool call is followed by matching tool result", () => {
const mockTextPart = new (vitest.mocked(vscode).LanguageModelTextPart)("I'll help")
const mockToolCallPart = new (vitest.mocked(vscode).LanguageModelToolCallPart)("call-1", "test-tool", {})
const mockToolResultPart = new (vitest.mocked(vscode).LanguageModelToolResultPart)("call-1", [
new (vitest.mocked(vscode).LanguageModelTextPart)("result"),
])

const assistantMessage = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.Assistant,
content: [mockTextPart, mockToolCallPart],
} as vscode.LanguageModelChatMessage

const userMessage = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.User,
content: [mockToolResultPart],
} as vscode.LanguageModelChatMessage

const messages = [assistantMessage, userMessage]

const result = validateAndRepairToolSequence(messages)
expect(result).toHaveLength(2)
expect(result[0]).toBe(assistantMessage)
expect(result[1]).toBe(userMessage)
})

it("should add placeholder tool result when assistant tool call is not followed by user message", () => {
const mockTextPart = new (vitest.mocked(vscode).LanguageModelTextPart)("I'll help")
const mockToolCallPart = new (vitest.mocked(vscode).LanguageModelToolCallPart)("call-1", "test-tool", {})

const assistantMessage = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.Assistant,
content: [mockTextPart, mockToolCallPart],
} as vscode.LanguageModelChatMessage

const messages = [assistantMessage]

const result = validateAndRepairToolSequence(messages)

// Should have original assistant message plus new user message with placeholder
expect(result).toHaveLength(2)
expect(result[0]).toBe(assistantMessage)
expect(result[1].role).toBe(vitest.mocked(vscode).LanguageModelChatMessageRole.User)
// The new user message should have a tool result part
expect(result[1].content).toHaveLength(1)
expect((result[1].content[0] as MockLanguageModelToolResultPart).callId).toBe("call-1")
})

it("should add placeholder tool result when tool call is followed by user message without matching result", () => {
const mockTextPart = new (vitest.mocked(vscode).LanguageModelTextPart)("I'll help")
const mockToolCallPart = new (vitest.mocked(vscode).LanguageModelToolCallPart)("call-1", "test-tool", {})
const mockUserTextPart = new (vitest.mocked(vscode).LanguageModelTextPart)("Continue please")

const assistantMessage = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.Assistant,
content: [mockTextPart, mockToolCallPart],
} as vscode.LanguageModelChatMessage

const userMessage = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.User,
content: [mockUserTextPart],
} as vscode.LanguageModelChatMessage

const messages = [assistantMessage, userMessage]

const result = validateAndRepairToolSequence(messages)

expect(result).toHaveLength(2)
expect(result[0]).toBe(assistantMessage)
// The repaired user message should have placeholder result prepended
expect(result[1].content).toHaveLength(2)
expect((result[1].content[0] as MockLanguageModelToolResultPart).callId).toBe("call-1")
expect((result[1].content[1] as MockLanguageModelTextPart).value).toBe("Continue please")
})

it("should handle multiple tool calls with partial matching results", () => {
const mockTextPart = new (vitest.mocked(vscode).LanguageModelTextPart)("I'll help")
const mockToolCallPart1 = new (vitest.mocked(vscode).LanguageModelToolCallPart)("call-1", "tool-1", {})
const mockToolCallPart2 = new (vitest.mocked(vscode).LanguageModelToolCallPart)("call-2", "tool-2", {})
const mockToolResultPart1 = new (vitest.mocked(vscode).LanguageModelToolResultPart)("call-1", [
new (vitest.mocked(vscode).LanguageModelTextPart)("result-1"),
])

const assistantMessage = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.Assistant,
content: [mockTextPart, mockToolCallPart1, mockToolCallPart2],
} as vscode.LanguageModelChatMessage

const userMessage = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.User,
content: [mockToolResultPart1], // Only has result for call-1, missing call-2
} as vscode.LanguageModelChatMessage

const messages = [assistantMessage, userMessage]

const result = validateAndRepairToolSequence(messages)

expect(result).toHaveLength(2)
expect(result[0]).toBe(assistantMessage)
// The repaired user message should have placeholder for call-2 prepended
expect(result[1].content).toHaveLength(2)
// Placeholder for missing call-2 should be first
expect((result[1].content[0] as MockLanguageModelToolResultPart).callId).toBe("call-2")
// Original result for call-1 should still be present
expect((result[1].content[1] as MockLanguageModelToolResultPart).callId).toBe("call-1")
})

it("should handle tool call followed by another assistant message", () => {
const mockTextPart = new (vitest.mocked(vscode).LanguageModelTextPart)("I'll help")
const mockToolCallPart = new (vitest.mocked(vscode).LanguageModelToolCallPart)("call-1", "test-tool", {})
const mockTextPart2 = new (vitest.mocked(vscode).LanguageModelTextPart)("More assistant text")

const assistantMessage1 = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.Assistant,
content: [mockTextPart, mockToolCallPart],
} as vscode.LanguageModelChatMessage

const assistantMessage2 = {
role: vitest.mocked(vscode).LanguageModelChatMessageRole.Assistant,
content: [mockTextPart2],
} as vscode.LanguageModelChatMessage

const messages = [assistantMessage1, assistantMessage2]

const result = validateAndRepairToolSequence(messages)

// Should insert a user message with placeholder between the two assistant messages
expect(result).toHaveLength(3)
expect(result[0]).toBe(assistantMessage1)
expect(result[1].role).toBe(vitest.mocked(vscode).LanguageModelChatMessageRole.User)
expect((result[1].content[0] as MockLanguageModelToolResultPart).callId).toBe("call-1")
expect(result[2]).toBe(assistantMessage2)
})

it("should handle null/undefined input gracefully", () => {
expect(validateAndRepairToolSequence(null as any)).toEqual(null)
expect(validateAndRepairToolSequence(undefined as any)).toEqual(undefined)
})
})
140 changes: 140 additions & 0 deletions src/api/transform/vscode-lm-format.ts
Original file line number Diff line number Diff line change
Expand Up @@ -194,3 +194,143 @@ export function extractTextCountFromMessage(message: vscode.LanguageModelChatMes
}
return text
}

/**
* Helper function to extract tool call IDs from a message's content.
* Returns an array of call IDs from LanguageModelToolCallPart instances.
*/
function getToolCallIds(message: vscode.LanguageModelChatMessage): string[] {
const callIds: string[] = []
if (Array.isArray(message.content)) {
for (const part of message.content) {
if (part instanceof vscode.LanguageModelToolCallPart) {
callIds.push(part.callId)
}
}
}
return callIds
}

/**
* Helper function to extract tool result IDs from a message's content.
* Returns an array of call IDs from LanguageModelToolResultPart instances.
*/
function getToolResultIds(message: vscode.LanguageModelChatMessage): string[] {
const callIds: string[] = []
if (Array.isArray(message.content)) {
for (const part of message.content) {
if (part instanceof vscode.LanguageModelToolResultPart) {
callIds.push(part.callId)
}
}
}
return callIds
}

/**
* Validates and repairs the tool call/result sequence in VSCode LM messages.
*
* The VSCode Language Model API requires that:
* - When an Assistant message contains LanguageModelToolCallPart(s)
* - The immediately next message MUST be a User message with LanguageModelToolResultPart(s)
* with matching callIds
*
* This function:
* 1. Identifies assistant messages with tool calls
* 2. Checks that the next message is a user message with matching tool results
* 3. If mismatches are found:
* - Missing tool results: Adds placeholder ToolResultParts to the user message
* - Orphaned tool results (no preceding tool call): Removes them
* - No following user message: Creates one with the necessary tool results
*
* @param messages - Array of VSCode LanguageModelChatMessage to validate and repair
* @returns The validated and repaired array of messages
*/
export function validateAndRepairToolSequence(
messages: vscode.LanguageModelChatMessage[],
): vscode.LanguageModelChatMessage[] {
if (!messages || messages.length === 0) {
return messages
}

const repairedMessages: vscode.LanguageModelChatMessage[] = []

for (let i = 0; i < messages.length; i++) {
const currentMessage = messages[i]
const isAssistant = currentMessage.role === vscode.LanguageModelChatMessageRole.Assistant
const toolCallIds = getToolCallIds(currentMessage)

// If this is an assistant message with tool calls
if (isAssistant && toolCallIds.length > 0) {
repairedMessages.push(currentMessage)

const nextMessage = messages[i + 1]
const isNextUser = nextMessage?.role === vscode.LanguageModelChatMessageRole.User

if (isNextUser) {
// Check if the next user message has matching tool results
const toolResultIds = getToolResultIds(nextMessage)
const missingResultIds = toolCallIds.filter((id) => !toolResultIds.includes(id))

if (missingResultIds.length > 0) {
// Need to add placeholder tool results for missing IDs
// Filter existing content to only include valid user message parts (TextPart and ToolResultPart)
const existingContent: (vscode.LanguageModelTextPart | vscode.LanguageModelToolResultPart)[] = []
if (Array.isArray(nextMessage.content)) {
for (const part of nextMessage.content) {
if (
part instanceof vscode.LanguageModelTextPart ||
part instanceof vscode.LanguageModelToolResultPart
) {
existingContent.push(part)
}
}
}

// Add placeholder results at the beginning (tool results come first)
const placeholderResults = missingResultIds.map(
(callId) =>
new vscode.LanguageModelToolResultPart(callId, [
new vscode.LanguageModelTextPart("[Tool result not available]"),
]),
)

// Create new user message with placeholders prepended
const repairedUserMessage = vscode.LanguageModelChatMessage.User([
...placeholderResults,
...existingContent,
])
repairedMessages.push(repairedUserMessage)
i++ // Skip the next message since we've processed it

console.warn(
`Roo Code <Language Model API>: Added ${missingResultIds.length} placeholder tool result(s) for missing call IDs: ${missingResultIds.join(", ")}`,
)
} else {
// All tool results are present, keep the message as-is
repairedMessages.push(nextMessage)
i++ // Skip the next message since we've processed it
}
} else {
// No user message follows - create one with placeholder tool results
const placeholderResults = toolCallIds.map(
(callId) =>
new vscode.LanguageModelToolResultPart(callId, [
new vscode.LanguageModelTextPart("[Tool result not available]"),
]),
)
const newUserMessage = vscode.LanguageModelChatMessage.User(placeholderResults)
repairedMessages.push(newUserMessage)

console.warn(
`Roo Code <Language Model API>: Created user message with ${toolCallIds.length} placeholder tool result(s) for call IDs: ${toolCallIds.join(", ")}`,
)
}
} else {
// For non-tool-call messages, just add them directly
repairedMessages.push(currentMessage)
}
}

return repairedMessages
}
Loading