Fix task execution in VS Code LM API

src/api/providers/vscode-lm.test.ts

@@ -0,0 +1,372 @@
+import * as vscode from "vscode"
+import { VsCodeLmHandler } from "../vscode-lm"
+import { ApiHandlerOptions } from "../../../shared/api"
+import { Anthropic } from "@anthropic-ai/sdk"
+
+// Mock vscode namespace
+jest.mock("vscode", () => {
+	class MockLanguageModelTextPart {
+		type = "text"
+		constructor(public value: string) {}
+	}
+
+	class MockLanguageModelToolCallPart {
+		type = "tool_call"
+		constructor(
+			public callId: string,
+			public name: string,
+			public input: any,
+		) {}
+	}
+
+	return {
+		workspace: {
+			onDidChangeConfiguration: jest.fn((callback) => ({
+				dispose: jest.fn(),
+			})),
+		},
+		CancellationTokenSource: jest.fn(() => ({
+			token: {
+				isCancellationRequested: false,
+				onCancellationRequested: jest.fn(),
+			},
+			cancel: jest.fn(),
+			dispose: jest.fn(),
+		})),
+		CancellationError: class CancellationError extends Error {
+			constructor() {
+				super("Operation cancelled")
+				this.name = "CancellationError"
+			}
+		},
+		LanguageModelChatMessage: {
+			Assistant: jest.fn((content) => ({
+				role: "assistant",
+				content: Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)],
+			})),
+			User: jest.fn((content) => ({
+				role: "user",
+				content: Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)],
+			})),
+		},
+		LanguageModelTextPart: MockLanguageModelTextPart,
+		LanguageModelToolCallPart: MockLanguageModelToolCallPart,
+		lm: {
+			selectChatModels: jest.fn(),
+		},
+	}
+})
+
+const mockLanguageModelChat = {
+	id: "test-model",
+	name: "Test Model",
+	vendor: "test-vendor",
+	family: "test-family",
+	version: "1.0",
+	maxInputTokens: 4096,
+	sendRequest: jest.fn(),
+	countTokens: jest.fn(),
+}
+
+describe("VsCodeLmHandler", () => {
+	let handler: VsCodeLmHandler
+	const defaultOptions: ApiHandlerOptions = {
+		vsCodeLmModelSelector: {
+			vendor: "test-vendor",
+			family: "test-family",
+		},
+	}
+
+	beforeEach(() => {
+		jest.clearAllMocks()
+		handler = new VsCodeLmHandler(defaultOptions)
+	})
+
+	afterEach(() => {
+		handler.dispose()
+	})
+
+	describe("constructor", () => {
+		it("should initialize with provided options", () => {
+			expect(handler).toBeDefined()
+			expect(vscode.workspace.onDidChangeConfiguration).toHaveBeenCalled()
+		})
+
+		it("should handle configuration changes", () => {
+			const callback = (vscode.workspace.onDidChangeConfiguration as jest.Mock).mock.calls[0][0]
+			callback({ affectsConfiguration: () => true })
+			// Should reset client when config changes
+			expect(handler["client"]).toBeNull()
+		})
+	})
+
+	describe("createClient", () => {
+		it("should create client with selector", async () => {
+			const mockModel = { ...mockLanguageModelChat }
+			;(vscode.lm.selectChatModels as jest.Mock).mockResolvedValueOnce([mockModel])
+
+			const client = await handler["createClient"]({
+				vendor: "test-vendor",
+				family: "test-family",
+			})
+
+			expect(client).toBeDefined()
+			expect(client.id).toBe("test-model")
+			expect(vscode.lm.selectChatModels).toHaveBeenCalledWith({
+				vendor: "test-vendor",
+				family: "test-family",
+			})
+		})
+
+		it("should return default client when no models available", async () => {
+			;(vscode.lm.selectChatModels as jest.Mock).mockResolvedValueOnce([])
+
+			const client = await handler["createClient"]({})
+
+			expect(client).toBeDefined()
+			expect(client.id).toBe("default-lm")
+			expect(client.vendor).toBe("vscode")
+		})
+	})
+
+	describe("createMessage", () => {
+		beforeEach(() => {
+			const mockModel = { ...mockLanguageModelChat }
+			;(vscode.lm.selectChatModels as jest.Mock).mockResolvedValueOnce([mockModel])
+			mockLanguageModelChat.countTokens.mockResolvedValue(10)
+		})
+
+		it("should stream text responses", async () => {
+			const systemPrompt = "You are a helpful assistant"
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user" as const,
+					content: "Hello",
+				},
+			]
+
+			const responseText = "Hello! How can I help you?"
+			mockLanguageModelChat.sendRequest.mockResolvedValueOnce({
+				stream: (async function* () {
+					yield new vscode.LanguageModelTextPart(responseText)
+					return
+				})(),
+				text: (async function* () {
+					yield responseText
+					return
+				})(),
+			})
+
+			const stream = handler.createMessage(systemPrompt, messages)
+			const chunks = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			expect(chunks).toHaveLength(2) // Text chunk + usage chunk
+			expect(chunks[0]).toEqual({
+				type: "text",
+				text: responseText,
+			})
+			expect(chunks[1]).toMatchObject({
+				type: "usage",
+				inputTokens: expect.any(Number),
+				outputTokens: expect.any(Number),
+			})
+		})
+
+		it("should handle tool calls", async () => {
+			const systemPrompt = "You are a helpful assistant"
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user" as const,
+					content: "Calculate 2+2",
+				},
+			]
+
+			const toolCallData = {
+				name: "calculator",
+				arguments: { operation: "add", numbers: [2, 2] },
+				callId: "call-1",
+			}
+
+			mockLanguageModelChat.sendRequest.mockResolvedValueOnce({
+				stream: (async function* () {
+					yield new vscode.LanguageModelToolCallPart(
+						toolCallData.callId,
+						toolCallData.name,
+						toolCallData.arguments,
+					)
+					return
+				})(),
+				text: (async function* () {
+					yield JSON.stringify({ type: "tool_call", ...toolCallData })
+					return
+				})(),
+			})
+
+			const stream = handler.createMessage(systemPrompt, messages)
+			const chunks = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			expect(chunks).toHaveLength(2) // Tool call chunk + usage chunk
+			expect(chunks[0]).toEqual({
+				type: "text",
+				text: JSON.stringify({ type: "tool_call", ...toolCallData }),
+			})
+		})
+
+		it("should handle errors", async () => {
+			const systemPrompt = "You are a helpful assistant"
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user" as const,
+					content: "Hello",
+				},
+			]
+
+			mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new Error("API Error"))
+
+			await expect(async () => {
+				const stream = handler.createMessage(systemPrompt, messages)
+				for await (const _ of stream) {
+					// consume stream
+				}
+			}).rejects.toThrow("API Error")
+		})
+
+		it("should execute tasks from tool calls", async () => {
+			const systemPrompt = "You are a helpful assistant"
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user" as const,
+					content: "Execute task",
+				},
+			]
+
+			const toolCallData = {
+				name: "taskExecutor",
+				arguments: { task: "exampleTask" },
+				callId: "call-2",
+			}
+
+			mockLanguageModelChat.sendRequest.mockResolvedValueOnce({
+				stream: (async function* () {
+					yield new vscode.LanguageModelToolCallPart(
+						toolCallData.callId,
+						toolCallData.name,
+						toolCallData.arguments,
+					)
+					return
+				})(),
+				text: (async function* () {
+					yield JSON.stringify({ type: "tool_call", ...toolCallData })
+					return
+				})(),
+			})
+
+			const stream = handler.createMessage(systemPrompt, messages)
+			const chunks = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			expect(chunks).toHaveLength(2) // Tool call chunk + usage chunk
+			expect(chunks[0]).toEqual({
+				type: "text",
+				text: JSON.stringify({ type: "tool_call", ...toolCallData }),
+			})
+		})
+	})
+
+	describe("getModel", () => {
+		it("should return model info when client exists", async () => {
+			const mockModel = { ...mockLanguageModelChat }
+			;(vscode.lm.selectChatModels as jest.Mock).mockResolvedValueOnce([mockModel])
+
+			// Initialize client
+			await handler["getClient"]()
+
+			const model = handler.getModel()
+			expect(model.id).toBe("test-model")
+			expect(model.info).toBeDefined()
+			expect(model.info.contextWindow).toBe(4096)
+		})
+
+		it("should return fallback model info when no client exists", () => {
+			const model = handler.getModel()
+			expect(model.id).toBe("test-vendor/test-family")
+			expect(model.info).toBeDefined()
+		})
+	})
+
+	describe("completePrompt", () => {
+		it("should complete single prompt", async () => {
+			const mockModel = { ...mockLanguageModelChat }
+			;(vscode.lm.selectChatModels as jest.Mock).mockResolvedValueOnce([mockModel])
+
+			const responseText = "Completed text"
+			mockLanguageModelChat.sendRequest.mockResolvedValueOnce({
+				stream: (async function* () {
+					yield new vscode.LanguageModelTextPart(responseText)
+					return
+				})(),
+				text: (async function* () {
+					yield responseText
+					return
+				})(),
+			})
+
+			const result = await handler.completePrompt("Test prompt")
+			expect(result).toBe(responseText)
+			expect(mockLanguageModelChat.sendRequest).toHaveBeenCalled()
+		})
+
+		it("should handle errors during completion", async () => {
+			const mockModel = { ...mockLanguageModelChat }
+			;(vscode.lm.selectChatModels as jest.Mock).mockResolvedValueOnce([mockModel])
+
+			mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new Error("Completion failed"))
+
+			await expect(handler.completePrompt("Test prompt")).rejects.toThrow(
+				"VSCode LM completion error: Completion failed",
+			)
+		})
+
+		it("should execute tasks during completion", async () => {
+			const mockModel = { ...mockLanguageModelChat }
+			;(vscode.lm.selectChatModels as jest.Mock).mockResolvedValueOnce([mockModel])
+
+			const responseText = "Completed text"
+			const toolCallData = {
+				name: "taskExecutor",
+				arguments: { task: "exampleTask" },
+				callId: "call-3",
+			}
+
+			mockLanguageModelChat.sendRequest.mockResolvedValueOnce({
+				stream: (async function* () {
+					yield new vscode.LanguageModelTextPart(responseText)
+					yield new vscode.LanguageModelToolCallPart(
+						toolCallData.callId,
+						toolCallData.name,
+						toolCallData.arguments,
+					)
+					return
+				})(),
+				text: (async function* () {
+					yield responseText
+					yield JSON.stringify({ type: "tool_call", ...toolCallData })
+					return
+				})(),
+			})
+
+			const result = await handler.completePrompt("Test prompt")
+			expect(result).toContain(responseText)
+			expect(result).toContain(JSON.stringify({ type: "tool_call", ...toolCallData }))
+			expect(mockLanguageModelChat.sendRequest).toHaveBeenCalled()
+		})
+	})
+})