refactor: migrate LM Studio provider to Vercel AI SDK (#11347)

* refactor: migrate LM Studio provider to Vercel AI SDK

Migrate LmStudioHandler from raw OpenAI SDK to Vercel AI SDK via
OpenAICompatibleHandler base class.

Changes:
- Extend OpenAICompatibleHandler instead of BaseProvider
- Use createOpenAICompatible from @ai-sdk/openai-compatible
- Use streamText/generateText from ai package
- Add extractReasoningMiddleware for <think> tag extraction parity
- Pass draft_model via providerOptions for speculative decoding
- Remove unused getLmStudioModels function (active version in fetchers/)
- Update all tests to mock AI SDK instead of OpenAI SDK

* fix: wrap completePrompt with handleAiSdkError for consistent error handling

---------

Co-authored-by: Roo Code <roomote@roocode.com>

src/api/providers/__tests__/lm-studio-timeout.spec.ts

@@ -1,62 +1,69 @@
 // npx vitest run api/providers/__tests__/lm-studio-timeout.spec.ts
 
+const { mockCreateOpenAICompatible } = vi.hoisted(() => ({
+	mockCreateOpenAICompatible: vi.fn(() => {
+		return vi.fn(() => ({
+			modelId: "llama2",
+			provider: "lmstudio",
+		}))
+	}),
+}))
+
+vi.mock("@ai-sdk/openai-compatible", () => ({
+	createOpenAICompatible: mockCreateOpenAICompatible,
+}))
+
+vi.mock("ai", async (importOriginal) => {
+	const actual = await importOriginal<typeof import("ai")>()
+	return {
+		...actual,
+		streamText: vi.fn(),
+		generateText: vi.fn(),
+	}
+})
+
 import { LmStudioHandler } from "../lm-studio"
 import { ApiHandlerOptions } from "../../../shared/api"
 
-// Mock the timeout config utility
-vitest.mock("../utils/timeout-config", () => ({
-	getApiRequestTimeout: vitest.fn(),
-}))
-
-import { getApiRequestTimeout } from "../utils/timeout-config"
-
-// Mock OpenAI
-const mockOpenAIConstructor = vitest.fn()
-vitest.mock("openai", () => {
-	return {
-		__esModule: true,
-		default: vitest.fn().mockImplementation((config) => {
-			mockOpenAIConstructor(config)
-			return {
-				chat: {
-					completions: {
-						create: vitest.fn(),
-					},
-				},
-			}
-		}),
-	}
-})
-
-describe("LmStudioHandler timeout configuration", () => {
+describe("LmStudioHandler configuration", () => {
 	beforeEach(() => {
-		vitest.clearAllMocks()
+		vi.clearAllMocks()
 	})
 
-	it("should use default timeout of 600 seconds when no configuration is set", () => {
-		;(getApiRequestTimeout as any).mockReturnValue(600000)
-
+	it("should configure the provider with default base URL", () => {
 		const options: ApiHandlerOptions = {
 			apiModelId: "llama2",
 			lmStudioModelId: "llama2",
-			lmStudioBaseUrl: "http://localhost:1234",
 		}
 
 		new LmStudioHandler(options)
 
-		expect(getApiRequestTimeout).toHaveBeenCalled()
-		expect(mockOpenAIConstructor).toHaveBeenCalledWith(
+		expect(mockCreateOpenAICompatible).toHaveBeenCalledWith(
 			expect.objectContaining({
+				name: "lmstudio",
 				baseURL: "http://localhost:1234/v1",
 				apiKey: "noop",
-				timeout: 600000, // 600 seconds in milliseconds
 			}),
 		)
 	})
 
-	it("should use custom timeout when configuration is set", () => {
-		;(getApiRequestTimeout as any).mockReturnValue(1200000) // 20 minutes
+	it("should configure the provider with custom base URL", () => {
+		const options: ApiHandlerOptions = {
+			apiModelId: "llama2",
+			lmStudioModelId: "llama2",
+			lmStudioBaseUrl: "http://localhost:5678",
+		}
 
+		new LmStudioHandler(options)
+
+		expect(mockCreateOpenAICompatible).toHaveBeenCalledWith(
+			expect.objectContaining({
+				baseURL: "http://localhost:5678/v1",
+			}),
+		)
+	})
+
+	it("should use 'noop' as the API key", () => {
 		const options: ApiHandlerOptions = {
 			apiModelId: "llama2",
 			lmStudioModelId: "llama2",
@@ -65,26 +72,9 @@ describe("LmStudioHandler timeout configuration", () => {
 
 		new LmStudioHandler(options)
 
-		expect(mockOpenAIConstructor).toHaveBeenCalledWith(
+		expect(mockCreateOpenAICompatible).toHaveBeenCalledWith(
 			expect.objectContaining({
-				timeout: 1200000, // 1200 seconds in milliseconds
-			}),
-		)
-	})
-
-	it("should handle zero timeout (no timeout)", () => {
-		;(getApiRequestTimeout as any).mockReturnValue(0)
-
-		const options: ApiHandlerOptions = {
-			apiModelId: "llama2",
-			lmStudioModelId: "llama2",
-		}
-
-		new LmStudioHandler(options)
-
-		expect(mockOpenAIConstructor).toHaveBeenCalledWith(
-			expect.objectContaining({
-				timeout: 0, // No timeout
+				apiKey: "noop",
 			}),
 		)
 	})

src/api/providers/__tests__/lmstudio-native-tools.spec.ts

@@ -1,22 +1,28 @@
 // npx vitest run api/providers/__tests__/lmstudio-native-tools.spec.ts
 
-// Mock OpenAI client - must come before other imports
-const mockCreate = vi.fn()
-vi.mock("openai", () => {
+// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
+const { mockStreamText } = vi.hoisted(() => ({
+	mockStreamText: vi.fn(),
+}))
+
+vi.mock("ai", async (importOriginal) => {
+	const actual = await importOriginal<typeof import("ai")>()
 	return {
-		__esModule: true,
-		default: vi.fn().mockImplementation(() => ({
-			chat: {
-				completions: {
-					create: mockCreate,
-				},
-			},
-		})),
+		...actual,
+		streamText: mockStreamText,
 	}
 })
 
+vi.mock("@ai-sdk/openai-compatible", () => ({
+	createOpenAICompatible: vi.fn(() => {
+		return vi.fn(() => ({
+			modelId: "local-model",
+			provider: "lmstudio",
+		}))
+	}),
+}))
+
 import { LmStudioHandler } from "../lm-studio"
-import { NativeToolCallParser } from "../../../core/assistant-message/NativeToolCallParser"
 import type { ApiHandlerOptions } from "../../../shared/api"
 
 describe("LmStudioHandler Native Tools", () => {
@@ -49,128 +55,76 @@ describe("LmStudioHandler Native Tools", () => {
 			lmStudioBaseUrl: "http://localhost:1234",
 		}
 		handler = new LmStudioHandler(mockOptions)
-
-		// Clear NativeToolCallParser state before each test
-		NativeToolCallParser.clearRawChunkState()
 	})
 
 	describe("Native Tool Calling Support", () => {
-		it("should include tools in request when model supports native tools and tools are provided", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [{ delta: { content: "Test response" } }],
-					}
-				},
-			}))
+		it("should include tools in request when tools are provided", async () => {
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Test response" }
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }),
+			})
 
 			const stream = handler.createMessage("test prompt", [], {
 				taskId: "test-task-id",
 				tools: testTools,
 			})
-			await stream.next()
+			for await (const _chunk of stream) {
+				// consume stream
+			}
 
-			expect(mockCreate).toHaveBeenCalledWith(
-				expect.objectContaining({
-					tools: expect.arrayContaining([
-						expect.objectContaining({
-							type: "function",
-							function: expect.objectContaining({
-								name: "test_tool",
-							}),
-						}),
-					]),
-				}),
-			)
-			// parallel_tool_calls should be true by default when not explicitly set
-			const callArgs = mockCreate.mock.calls[0][0]
-			expect(callArgs).toHaveProperty("parallel_tool_calls", true)
+			const callArgs = mockStreamText.mock.calls[0][0]
+			expect(callArgs.tools).toBeDefined()
 		})
 
-		it("should include tool_choice when provided", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [{ delta: { content: "Test response" } }],
-					}
-				},
-			}))
+		it("should include toolChoice when provided", async () => {
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Test response" }
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }),
+			})
 
 			const stream = handler.createMessage("test prompt", [], {
 				taskId: "test-task-id",
 				tools: testTools,
 				tool_choice: "auto",
 			})
-			await stream.next()
+			for await (const _chunk of stream) {
+				// consume stream
+			}
 
-			expect(mockCreate).toHaveBeenCalledWith(
-				expect.objectContaining({
-					tool_choice: "auto",
-				}),
-			)
+			const callArgs = mockStreamText.mock.calls[0][0]
+			expect(callArgs.toolChoice).toBe("auto")
 		})
 
-		it("should always include tools and tool_choice in request (tools are always present after PR #10841)", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [{ delta: { content: "Test response" } }],
-					}
-				},
-			}))
+		it("should yield tool_call_start, tool_call_delta, and tool_call_end chunks from AI SDK stream", async () => {
+			async function* mockFullStream() {
+				yield {
+					type: "tool-input-start",
+					id: "call_lmstudio_123",
+					toolName: "test_tool",
+				}
+				yield {
+					type: "tool-input-delta",
+					id: "call_lmstudio_123",
+					delta: '{"arg1":"value"}',
+				}
+				yield {
+					type: "tool-input-end",
+					id: "call_lmstudio_123",
+				}
+			}
 
-			const stream = handler.createMessage("test prompt", [], {
-				taskId: "test-task-id",
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }),
 			})
-			await stream.next()
-
-			const callArgs = mockCreate.mock.calls[mockCreate.mock.calls.length - 1][0]
-			// Tools are now always present (minimum 6 from ALWAYS_AVAILABLE_TOOLS)
-			expect(callArgs).toHaveProperty("tools")
-			expect(callArgs).toHaveProperty("tool_choice")
-			// parallel_tool_calls should be true by default when not explicitly set
-			expect(callArgs).toHaveProperty("parallel_tool_calls", true)
-		})
-
-		it("should yield tool_call_partial chunks during streaming", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [
-							{
-								delta: {
-									tool_calls: [
-										{
-											index: 0,
-											id: "call_lmstudio_123",
-											function: {
-												name: "test_tool",
-												arguments: '{"arg1":',
-											},
-										},
-									],
-								},
-							},
-						],
-					}
-					yield {
-						choices: [
-							{
-								delta: {
-									tool_calls: [
-										{
-											index: 0,
-											function: {
-												arguments: '"value"}',
-											},
-										},
-									],
-								},
-							},
-						],
-					}
-				},
-			}))
 
 			const stream = handler.createMessage("test prompt", [], {
 				taskId: "test-task-id",
@@ -182,168 +136,56 @@ describe("LmStudioHandler Native Tools", () => {
 				chunks.push(chunk)
 			}
 
-			expect(chunks).toContainEqual({
-				type: "tool_call_partial",
-				index: 0,
+			const startChunks = chunks.filter((chunk) => chunk.type === "tool_call_start")
+			expect(startChunks).toHaveLength(1)
+			expect(startChunks[0]).toEqual({
+				type: "tool_call_start",
 				id: "call_lmstudio_123",
 				name: "test_tool",
-				arguments: '{"arg1":',
 			})
 
-			expect(chunks).toContainEqual({
-				type: "tool_call_partial",
-				index: 0,
-				id: undefined,
-				name: undefined,
-				arguments: '"value"}',
-			})
-		})
-
-		it("should set parallel_tool_calls based on metadata", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [{ delta: { content: "Test response" } }],
-					}
-				},
-			}))
-
-			const stream = handler.createMessage("test prompt", [], {
-				taskId: "test-task-id",
-				tools: testTools,
-				parallelToolCalls: true,
-			})
-			await stream.next()
-
-			expect(mockCreate).toHaveBeenCalledWith(
-				expect.objectContaining({
-					parallel_tool_calls: true,
-				}),
-			)
-		})
-
-		it("should yield tool_call_end events when finish_reason is tool_calls", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [
-							{
-								delta: {
-									tool_calls: [
-										{
-											index: 0,
-											id: "call_lmstudio_test",
-											function: {
-												name: "test_tool",
-												arguments: '{"arg1":"value"}',
-											},
-										},
-									],
-								},
-							},
-						],
-					}
-					yield {
-						choices: [
-							{
-								delta: {},
-								finish_reason: "tool_calls",
-							},
-						],
-					}
-				},
-			}))
-
-			const stream = handler.createMessage("test prompt", [], {
-				taskId: "test-task-id",
-				tools: testTools,
+			const deltaChunks = chunks.filter((chunk) => chunk.type === "tool_call_delta")
+			expect(deltaChunks).toHaveLength(1)
+			expect(deltaChunks[0]).toEqual({
+				type: "tool_call_delta",
+				id: "call_lmstudio_123",
+				delta: '{"arg1":"value"}',
 			})
 
-			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_lmstudio_test")
-		})
-
-		it("should work with parallel tool calls disabled (sends false)", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [{ delta: { content: "Response" } }],
-					}
-				},
-			}))
-
-			const stream = handler.createMessage("test prompt", [], {
-				taskId: "test-task-id",
-				tools: testTools,
-				parallelToolCalls: false,
+			expect(endChunks[0]).toEqual({
+				type: "tool_call_end",
+				id: "call_lmstudio_123",
 			})
-			await stream.next()
-
-			// When parallelToolCalls is false, the parameter should be sent as false
-			const callArgs = mockCreate.mock.calls[0][0]
-			expect(callArgs).toHaveProperty("parallel_tool_calls", false)
 		})
 
 		it("should handle reasoning content alongside tool calls", async () => {
-			mockCreate.mockImplementationOnce(() => ({
-				[Symbol.asyncIterator]: async function* () {
-					yield {
-						choices: [
-							{
-								delta: {
-									content: "<think>Thinking about this...</think>",
-								},
-							},
-						],
-					}
-					yield {
-						choices: [
-							{
-								delta: {
-									tool_calls: [
-										{
-											index: 0,
-											id: "call_after_think",
-											function: {
-												name: "test_tool",
-												arguments: '{"arg1":"result"}',
-											},
-										},
-									],
-								},
-							},
-						],
-					}
-					yield {
-						choices: [
-							{
-								delta: {},
-								finish_reason: "tool_calls",
-							},
-						],
-					}
-				},
-			}))
+			async function* mockFullStream() {
+				yield {
+					type: "reasoning",
+					text: "Thinking about this...",
+				}
+				yield {
+					type: "tool-input-start",
+					id: "call_after_think",
+					toolName: "test_tool",
+				}
+				yield {
+					type: "tool-input-delta",
+					id: "call_after_think",
+					delta: '{"arg1":"result"}',
+				}
+				yield {
+					type: "tool-input-end",
+					id: "call_after_think",
+				}
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }),
+			})
 
 			const stream = handler.createMessage("test prompt", [], {
 				taskId: "test-task-id",
@@ -352,25 +194,60 @@ describe("LmStudioHandler Native Tools", () => {
 
 			const chunks = []
 			for await (const chunk of stream) {
-				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 reasoning, tool_call_partial, and tool_call_end
 			const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
-			const partialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
+			const startChunks = chunks.filter((chunk) => chunk.type === "tool_call_start")
 			const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
 
 			expect(reasoningChunks).toHaveLength(1)
 			expect(reasoningChunks[0].text).toBe("Thinking about this...")
-			expect(partialChunks).toHaveLength(1)
+			expect(startChunks).toHaveLength(1)
+			expect(endChunks).toHaveLength(1)
+		})
+
+		it("should handle text and tool calls in the same response", async () => {
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Here's the result: " }
+				yield {
+					type: "tool-input-start",
+					id: "call_mixed",
+					toolName: "test_tool",
+				}
+				yield {
+					type: "tool-input-delta",
+					id: "call_mixed",
+					delta: '{"arg1":"mixed"}',
+				}
+				yield {
+					type: "tool-input-end",
+					id: "call_mixed",
+				}
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }),
+			})
+
+			const stream = handler.createMessage("test prompt", [], {
+				taskId: "test-task-id",
+				tools: testTools,
+			})
+
+			const chunks = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			const textChunks = chunks.filter((chunk) => chunk.type === "text")
+			const startChunks = chunks.filter((chunk) => chunk.type === "tool_call_start")
+			const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
+
+			expect(textChunks).toHaveLength(1)
+			expect(textChunks[0].text).toBe("Here's the result: ")
+			expect(startChunks).toHaveLength(1)
 			expect(endChunks).toHaveLength(1)
 		})
 	})

src/api/providers/__tests__/lmstudio.spec.ts

@@ -1,63 +1,29 @@
-// Mock OpenAI client - must come before other imports
-const mockCreate = vi.fn()
-vi.mock("openai", () => {
-	return {
-		__esModule: true,
-		default: vi.fn().mockImplementation(() => ({
-			chat: {
-				completions: {
-					create: mockCreate.mockImplementation(async (options) => {
-						if (!options.stream) {
-							return {
-								id: "test-completion",
-								choices: [
-									{
-										message: { role: "assistant", content: "Test response" },
-										finish_reason: "stop",
-										index: 0,
-									},
-								],
-								usage: {
-									prompt_tokens: 10,
-									completion_tokens: 5,
-									total_tokens: 15,
-								},
-							}
-						}
+// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
+const { mockStreamText, mockGenerateText, mockWrapLanguageModel } = vi.hoisted(() => ({
+	mockStreamText: vi.fn(),
+	mockGenerateText: vi.fn(),
+	mockWrapLanguageModel: vi.fn((opts: any) => opts.model),
+}))
 
-						return {
-							[Symbol.asyncIterator]: async function* () {
-								yield {
-									choices: [
-										{
-											delta: { content: "Test response" },
-											index: 0,
-										},
-									],
-									usage: null,
-								}
-								yield {
-									choices: [
-										{
-											delta: {},
-											index: 0,
-										},
-									],
-									usage: {
-										prompt_tokens: 10,
-										completion_tokens: 5,
-										total_tokens: 15,
-									},
-								}
-							},
-						}
-					}),
-				},
-			},
-		})),
+vi.mock("ai", async (importOriginal) => {
+	const actual = await importOriginal<typeof import("ai")>()
+	return {
+		...actual,
+		streamText: mockStreamText,
+		generateText: mockGenerateText,
+		wrapLanguageModel: mockWrapLanguageModel,
 	}
 })
 
+vi.mock("@ai-sdk/openai-compatible", () => ({
+	createOpenAICompatible: vi.fn(() => {
+		return vi.fn(() => ({
+			modelId: "local-model",
+			provider: "lmstudio",
+		}))
+	}),
+}))
+
 import type { Anthropic } from "@anthropic-ai/sdk"
 
 import { LmStudioHandler } from "../lm-studio"
@@ -74,7 +40,7 @@ describe("LmStudioHandler", () => {
 			lmStudioBaseUrl: "http://localhost:1234",
 		}
 		handler = new LmStudioHandler(mockOptions)
-		mockCreate.mockClear()
+		vi.clearAllMocks()
 	})
 
 	describe("constructor", () => {
@@ -102,6 +68,20 @@ describe("LmStudioHandler", () => {
 		]
 
 		it("should handle streaming responses", async () => {
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Test response" }
+			}
+
+			const mockUsage = Promise.resolve({
+				inputTokens: 10,
+				outputTokens: 5,
+			})
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: mockUsage,
+			})
+
 			const stream = handler.createMessage(systemPrompt, messages)
 			const chunks: any[] = []
 			for await (const chunk of stream) {
@@ -114,8 +94,43 @@ describe("LmStudioHandler", () => {
 			expect(textChunks[0].text).toBe("Test response")
 		})
 
+		it("should include usage information", async () => {
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Test response" }
+			}
+
+			const mockUsage = Promise.resolve({
+				inputTokens: 10,
+				outputTokens: 5,
+			})
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: mockUsage,
+			})
+
+			const stream = handler.createMessage(systemPrompt, messages)
+			const chunks: any[] = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
+			expect(usageChunks.length).toBeGreaterThan(0)
+			expect(usageChunks[0].inputTokens).toBe(10)
+			expect(usageChunks[0].outputTokens).toBe(5)
+		})
+
 		it("should handle API errors", async () => {
-			mockCreate.mockRejectedValueOnce(new Error("API Error"))
+			async function* mockFullStream(): AsyncGenerator<{ type: string; text: string }> {
+				yield { type: "text-delta", text: "" }
+				throw new Error("API Error")
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 0, outputTokens: 0 }),
+			})
 
 			const stream = handler.createMessage(systemPrompt, messages)
 
@@ -123,36 +138,37 @@ describe("LmStudioHandler", () => {
 				for await (const _chunk of stream) {
 					// Should not reach here
 				}
-			}).rejects.toThrow("Please check the LM Studio developer logs to debug what went wrong")
+			}).rejects.toThrow()
 		})
 	})
 
 	describe("completePrompt", () => {
 		it("should complete prompt successfully", async () => {
+			mockGenerateText.mockResolvedValue({
+				text: "Test response",
+			})
+
 			const result = await handler.completePrompt("Test prompt")
 			expect(result).toBe("Test response")
-			expect(mockCreate).toHaveBeenCalledWith({
-				model: mockOptions.lmStudioModelId,
-				messages: [{ role: "user", content: "Test prompt" }],
-				temperature: 0,
-				stream: false,
-			})
-		})
-
-		it("should handle API errors", async () => {
-			mockCreate.mockRejectedValueOnce(new Error("API Error"))
-			await expect(handler.completePrompt("Test prompt")).rejects.toThrow(
-				"Please check the LM Studio developer logs to debug what went wrong",
+			expect(mockGenerateText).toHaveBeenCalledWith(
+				expect.objectContaining({
+					prompt: "Test prompt",
+				}),
 			)
 		})
 
 		it("should handle empty response", async () => {
-			mockCreate.mockResolvedValueOnce({
-				choices: [{ message: { content: "" } }],
+			mockGenerateText.mockResolvedValue({
+				text: "",
 			})
 			const result = await handler.completePrompt("Test prompt")
 			expect(result).toBe("")
 		})
+
+		it("should handle API errors with handleAiSdkError", async () => {
+			mockGenerateText.mockRejectedValueOnce(new Error("Connection refused"))
+			await expect(handler.completePrompt("Test prompt")).rejects.toThrow("LM Studio")
+		})
 	})
 
 	describe("getModel", () => {
@@ -164,4 +180,131 @@ describe("LmStudioHandler", () => {
 			expect(modelInfo.info.contextWindow).toBe(128_000)
 		})
 	})
+
+	describe("speculative decoding", () => {
+		it("should include draft_model in providerOptions when speculative decoding is enabled", async () => {
+			const speculativeHandler = new LmStudioHandler({
+				...mockOptions,
+				lmStudioSpeculativeDecodingEnabled: true,
+				lmStudioDraftModelId: "draft-model-id",
+			})
+
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Test" }
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 5, outputTokens: 3 }),
+			})
+
+			const stream = speculativeHandler.createMessage("test prompt", [])
+			for await (const _chunk of stream) {
+				// consume stream
+			}
+
+			expect(mockStreamText).toHaveBeenCalledWith(
+				expect.objectContaining({
+					providerOptions: {
+						lmstudio: { draft_model: "draft-model-id" },
+					},
+				}),
+			)
+		})
+
+		it("should not include draft_model when speculative decoding is disabled", async () => {
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Test" }
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 5, outputTokens: 3 }),
+			})
+
+			const stream = handler.createMessage("test prompt", [])
+			for await (const _chunk of stream) {
+				// consume stream
+			}
+
+			const callArgs = mockStreamText.mock.calls[0][0]
+			expect(callArgs.providerOptions).toBeUndefined()
+		})
+
+		it("should include draft_model in completePrompt when speculative decoding is enabled", async () => {
+			const speculativeHandler = new LmStudioHandler({
+				...mockOptions,
+				lmStudioSpeculativeDecodingEnabled: true,
+				lmStudioDraftModelId: "draft-model-id",
+			})
+
+			mockGenerateText.mockResolvedValue({ text: "Test" })
+
+			await speculativeHandler.completePrompt("Test prompt")
+
+			expect(mockGenerateText).toHaveBeenCalledWith(
+				expect.objectContaining({
+					providerOptions: {
+						lmstudio: { draft_model: "draft-model-id" },
+					},
+				}),
+			)
+		})
+	})
+
+	describe("reasoning middleware", () => {
+		it("should wrap the language model with extractReasoningMiddleware for <think> tags", async () => {
+			async function* mockFullStream() {
+				yield { type: "text-delta", text: "Test" }
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 5, outputTokens: 3 }),
+			})
+
+			const stream = handler.createMessage("test prompt", [])
+			for await (const _chunk of stream) {
+				// consume stream to trigger getLanguageModel()
+			}
+
+			expect(mockWrapLanguageModel).toHaveBeenCalledWith(
+				expect.objectContaining({
+					middleware: expect.any(Object),
+				}),
+			)
+		})
+
+		it("should handle reasoning-delta chunks from middleware-processed stream", async () => {
+			async function* mockFullStream() {
+				yield { type: "reasoning-delta", text: "Let me think about this..." }
+				yield { type: "text-delta", text: "The answer is 42." }
+			}
+
+			mockStreamText.mockReturnValue({
+				fullStream: mockFullStream(),
+				usage: Promise.resolve({ inputTokens: 10, outputTokens: 8 }),
+			})
+
+			const stream = handler.createMessage("test prompt", [])
+			const chunks: any[] = []
+			for await (const chunk of stream) {
+				chunks.push(chunk)
+			}
+
+			const reasoningChunks = chunks.filter((c) => c.type === "reasoning")
+			const textChunks = chunks.filter((c) => c.type === "text")
+
+			expect(reasoningChunks).toHaveLength(1)
+			expect(reasoningChunks[0].text).toBe("Let me think about this...")
+			expect(textChunks).toHaveLength(1)
+			expect(textChunks[0].text).toBe("The answer is 42.")
+		})
+	})
+
+	describe("isAiSdkProvider", () => {
+		it("should return true", () => {
+			expect(handler.isAiSdkProvider()).toBe(true)
+		})
+	})
 })

src/api/providers/lm-studio.ts

@@ -1,39 +1,49 @@
 import { Anthropic } from "@anthropic-ai/sdk"
-import OpenAI from "openai"
-import axios from "axios"
+import { streamText, generateText, ToolSet, wrapLanguageModel, extractReasoningMiddleware, LanguageModel } from "ai"
 
 import { type ModelInfo, openAiModelInfoSaneDefaults, LMSTUDIO_DEFAULT_TEMPERATURE } from "@roo-code/types"
 
 import type { ApiHandlerOptions } from "../../shared/api"
 
-import { NativeToolCallParser } from "../../core/assistant-message/NativeToolCallParser"
-import { TagMatcher } from "../../utils/tag-matcher"
-
-import { convertToOpenAiMessages } from "../transform/openai-format"
+import {
+	convertToAiSdkMessages,
+	convertToolsForAiSdk,
+	processAiSdkStreamPart,
+	mapToolChoice,
+	handleAiSdkError,
+} from "../transform/ai-sdk"
 import { ApiStream } from "../transform/stream"
 
-import { BaseProvider } from "./base-provider"
+import { OpenAICompatibleHandler, OpenAICompatibleConfig } from "./openai-compatible"
 import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
 import { getModelsFromCache } from "./fetchers/modelCache"
-import { getApiRequestTimeout } from "./utils/timeout-config"
-import { handleOpenAIError } from "./utils/openai-error-handler"
-
-export class LmStudioHandler extends BaseProvider implements SingleCompletionHandler {
-	protected options: ApiHandlerOptions
-	private client: OpenAI
-	private readonly providerName = "LM Studio"
 
+export class LmStudioHandler extends OpenAICompatibleHandler implements SingleCompletionHandler {
 	constructor(options: ApiHandlerOptions) {
-		super()
-		this.options = options
+		const modelId = options.lmStudioModelId || ""
+		const baseURL = (options.lmStudioBaseUrl || "http://localhost:1234") + "/v1"
 
-		// LM Studio uses "noop" as a placeholder API key
-		const apiKey = "noop"
+		const models = getModelsFromCache("lmstudio")
+		const modelInfo = (models && modelId && models[modelId]) || openAiModelInfoSaneDefaults
 
-		this.client = new OpenAI({
-			baseURL: (this.options.lmStudioBaseUrl || "http://localhost:1234") + "/v1",
-			apiKey: apiKey,
-			timeout: getApiRequestTimeout(),
+		const config: OpenAICompatibleConfig = {
+			providerName: "lmstudio",
+			baseURL,
+			apiKey: "noop",
+			modelId,
+			modelInfo,
+			temperature: options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
+			modelMaxTokens: options.modelMaxTokens ?? undefined,
+		}
+
+		super(options, config)
+	}
+
+	protected override getLanguageModel(): LanguageModel {
+		const baseModel = this.provider(this.config.modelId)
+		return wrapLanguageModel({
+			model: baseModel,
+			middleware: extractReasoningMiddleware({ tagName: "think" }),
 		})
 	}
 
@@ -42,189 +52,83 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
 		messages: Anthropic.Messages.MessageParam[],
 		metadata?: ApiHandlerCreateMessageMetadata,
 	): ApiStream {
-		const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
-			{ role: "system", content: systemPrompt },
-			...convertToOpenAiMessages(messages),
-		]
+		const model = this.getModel()
+		const languageModel = this.getLanguageModel()
 
-		// -------------------------
-		// Track token usage
-		// -------------------------
-		const toContentBlocks = (
-			blocks: Anthropic.Messages.MessageParam[] | string,
-		): Anthropic.Messages.ContentBlockParam[] => {
-			if (typeof blocks === "string") {
-				return [{ type: "text", text: blocks }]
-			}
+		const aiSdkMessages = convertToAiSdkMessages(messages)
 
-			const result: Anthropic.Messages.ContentBlockParam[] = []
-			for (const msg of blocks) {
-				if (typeof msg.content === "string") {
-					result.push({ type: "text", text: msg.content })
-				} else if (Array.isArray(msg.content)) {
-					for (const part of msg.content) {
-						if (part.type === "text") {
-							result.push({ type: "text", text: part.text })
-						}
-					}
-				}
-			}
-			return result
+		const openAiTools = this.convertToolsForOpenAI(metadata?.tools)
+		const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined
+
+		const requestOptions: Parameters<typeof streamText>[0] = {
+			model: languageModel,
+			system: systemPrompt,
+			messages: aiSdkMessages,
+			temperature: model.temperature ?? this.config.temperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
+			maxOutputTokens: this.getMaxOutputTokens(),
+			tools: aiSdkTools,
+			toolChoice: mapToolChoice(metadata?.tool_choice),
 		}
 
-		let inputTokens = 0
-		try {
-			inputTokens = await this.countTokens([{ type: "text", text: systemPrompt }, ...toContentBlocks(messages)])
-		} catch (err) {
-			console.error("[LmStudio] Failed to count input tokens:", err)
-			inputTokens = 0
+		if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) {
+			requestOptions.providerOptions = {
+				lmstudio: { draft_model: this.options.lmStudioDraftModelId },
+			}
 		}
 
-		let assistantText = ""
+		const result = streamText(requestOptions)
 
 		try {
-			const params: OpenAI.Chat.ChatCompletionCreateParamsStreaming & { draft_model?: string } = {
-				model: this.getModel().id,
-				messages: openAiMessages,
-				temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
-				stream: true,
-				tools: this.convertToolsForOpenAI(metadata?.tools),
-				tool_choice: metadata?.tool_choice,
-				parallel_tool_calls: metadata?.parallelToolCalls ?? true,
-			}
-
-			if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) {
-				params.draft_model = this.options.lmStudioDraftModelId
-			}
-
-			let results
-			try {
-				results = await this.client.chat.completions.create(params)
-			} catch (error) {
-				throw handleOpenAIError(error, this.providerName)
-			}
-
-			const matcher = new TagMatcher(
-				"think",
-				(chunk) =>
-					({
-						type: chunk.matched ? "reasoning" : "text",
-						text: chunk.data,
-					}) as const,
-			)
-
-			for await (const chunk of results) {
-				const delta = chunk.choices[0]?.delta
-				const finishReason = chunk.choices[0]?.finish_reason
-
-				if (delta?.content) {
-					assistantText += delta.content
-					for (const processedChunk of matcher.update(delta.content)) {
-						yield processedChunk
-					}
-				}
-
-				// Handle tool calls in stream - emit partial chunks for NativeToolCallParser
-				if (delta?.tool_calls) {
-					for (const toolCall of delta.tool_calls) {
-						yield {
-							type: "tool_call_partial",
-							index: toolCall.index,
-							id: toolCall.id,
-							name: toolCall.function?.name,
-							arguments: toolCall.function?.arguments,
-						}
-					}
-				}
-
-				// Process finish_reason to emit tool_call_end events
-				if (finishReason) {
-					const endEvents = NativeToolCallParser.processFinishReason(finishReason)
-					for (const event of endEvents) {
-						yield event
-					}
+			for await (const part of result.fullStream) {
+				for (const chunk of processAiSdkStreamPart(part)) {
+					yield chunk
 				}
 			}
 
-			for (const processedChunk of matcher.final()) {
-				yield processedChunk
+			const usage = await result.usage
+			if (usage) {
+				yield this.processUsageMetrics(usage)
 			}
-
-			let outputTokens = 0
-			try {
-				outputTokens = await this.countTokens([{ type: "text", text: assistantText }])
-			} catch (err) {
-				console.error("[LmStudio] Failed to count output tokens:", err)
-				outputTokens = 0
-			}
-
-			yield {
-				type: "usage",
-				inputTokens,
-				outputTokens,
-			} as const
 		} catch (error) {
-			throw new Error(
-				"Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Roo Code's prompts.",
-			)
+			throw handleAiSdkError(error, "LM Studio")
 		}
 	}
 
-	override getModel(): { id: string; info: ModelInfo } {
+	override getModel(): { id: string; info: ModelInfo; maxTokens?: number; temperature?: number } {
 		const models = getModelsFromCache("lmstudio")
-		if (models && this.options.lmStudioModelId && models[this.options.lmStudioModelId]) {
-			return {
-				id: this.options.lmStudioModelId,
-				info: models[this.options.lmStudioModelId],
-			}
-		} else {
-			return {
-				id: this.options.lmStudioModelId || "",
-				info: openAiModelInfoSaneDefaults,
-			}
+		const modelId = this.options.lmStudioModelId || ""
+
+		const info = (models && modelId && models[modelId]) || openAiModelInfoSaneDefaults
+
+		return {
+			id: modelId,
+			info,
+			temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
+			maxTokens: this.options.modelMaxTokens ?? undefined,
 		}
 	}
 
-	async completePrompt(prompt: string): Promise<string> {
+	override async completePrompt(prompt: string): Promise<string> {
+		const languageModel = this.getLanguageModel()
+
+		const options: Parameters<typeof generateText>[0] = {
+			model: languageModel,
+			prompt,
+			maxOutputTokens: this.getMaxOutputTokens(),
+			temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
+		}
+
+		if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) {
+			options.providerOptions = {
+				lmstudio: { draft_model: this.options.lmStudioDraftModelId },
+			}
+		}
+
 		try {
-			// Create params object with optional draft model
-			const params: any = {
-				model: this.getModel().id,
-				messages: [{ role: "user", content: prompt }],
-				temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
-				stream: false,
-			}
-
-			// Add draft model if speculative decoding is enabled and a draft model is specified
-			if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) {
-				params.draft_model = this.options.lmStudioDraftModelId
-			}
-
-			let response
-			try {
-				response = await this.client.chat.completions.create(params)
-			} catch (error) {
-				throw handleOpenAIError(error, this.providerName)
-			}
-			return response.choices[0]?.message.content || ""
+			const { text } = await generateText(options)
+			return text
 		} catch (error) {
-			throw new Error(
-				"Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Roo Code's prompts.",
-			)
+			throw handleAiSdkError(error, "LM Studio")
 		}
 	}
 }
-
-export async function getLmStudioModels(baseUrl = "http://localhost:1234") {
-	try {
-		if (!URL.canParse(baseUrl)) {
-			return []
-		}
-
-		const response = await axios.get(`${baseUrl}/v1/models`)
-		const modelsArray = response.data?.data?.map((model: any) => model.id) || []
-		return [...new Set<string>(modelsArray)]
-	} catch (error) {
-		return []
-	}
-}