Skip to content

Commit bfa3929

Browse files
roomoteroomote
committed
feat: add external MCP server integration for enhance prompt
- Add enhancePrompt settings to provider-settings schema with useExternalServer and endpoint options - Update MessageEnhancer to check for external server configuration and route requests accordingly - Add comprehensive tests for external MCP server functionality including fallback behavior - Import and use getModelId helper for consistent model ID extraction across providers Fixes #6631
1 parent a88238f commit bfa3929

File tree

3 files changed

+311
-2
lines changed

3 files changed

+311
-2
lines changed

‎packages/types/src/provider-settings.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,14 @@ const baseProviderSettingsSchema = z.object({
7777
reasoningEffort: reasoningEffortsSchema.optional(),
7878
modelMaxTokens: z.number().optional(),
7979
modelMaxThinkingTokens: z.number().optional(),
80+
81+
// External MCP server settings for enhance prompt
82+
enhancePrompt: z
83+
.object({
84+
useExternalServer: z.boolean().optional(),
85+
endpoint: z.string().url().optional(),
86+
})
87+
.optional(),
8088
})
8189

8290
// Several of the providers share common model config properties.

‎src/core/webview/__tests__/messageEnhancer.test.ts‎

Lines changed: 253 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
22
import { MessageEnhancer } from "../messageEnhancer"
3-
import { ProviderSettings, ClineMessage } from "@roo-code/types"
3+
import { ProviderSettings, ClineMessage, getModelId } from "@roo-code/types"
44
import { TelemetryService } from "@roo-code/telemetry"
55
import * as singleCompletionHandlerModule from "../../../utils/single-completion-handler"
66
import { ProviderSettingsManager } from "../../config/ProviderSettingsManager"
@@ -9,6 +9,9 @@ import { ProviderSettingsManager } from "../../config/ProviderSettingsManager"
99
vi.mock("../../../utils/single-completion-handler")
1010
vi.mock("@roo-code/telemetry")
1111

12+
// Mock global fetch
13+
global.fetch = vi.fn()
14+
1215
describe("MessageEnhancer", () => {
1316
let mockProviderSettingsManager: ProviderSettingsManager
1417
let mockSingleCompletionHandler: ReturnType<typeof vi.fn>
@@ -254,6 +257,255 @@ describe("MessageEnhancer", () => {
254257
// Should not include task history section
255258
expect(calledPrompt).not.toContain("previous conversation context")
256259
})
260+
261+
describe("External MCP Server", () => {
262+
beforeEach(() => {
263+
vi.mocked(global.fetch).mockReset()
264+
})
265+
266+
it("should use external MCP server when enabled", async () => {
267+
const mockResponse = {
268+
enhancedPrompt: "Enhanced via external server",
269+
}
270+
vi.mocked(global.fetch).mockResolvedValue({
271+
ok: true,
272+
json: vi.fn().mockResolvedValue(mockResponse),
273+
} as any)
274+
275+
const configWithExternalServer: ProviderSettings = {
276+
...mockApiConfiguration,
277+
enhancePrompt: {
278+
useExternalServer: true,
279+
endpoint: "http://localhost:8000/enhance",
280+
},
281+
}
282+
283+
const result = await MessageEnhancer.enhanceMessage({
284+
text: "Test prompt",
285+
apiConfiguration: configWithExternalServer,
286+
listApiConfigMeta: mockListApiConfigMeta,
287+
providerSettingsManager: mockProviderSettingsManager,
288+
})
289+
290+
expect(result.success).toBe(true)
291+
expect(result.enhancedText).toBe("Enhanced via external server")
292+
expect(global.fetch).toHaveBeenCalledWith("http://localhost:8000/enhance", {
293+
method: "POST",
294+
headers: { "Content-Type": "application/json" },
295+
body: JSON.stringify({
296+
prompt: "Test prompt",
297+
context: [],
298+
model: "gpt-4",
299+
}),
300+
})
301+
// Should not call internal enhancement
302+
expect(mockSingleCompletionHandler).not.toHaveBeenCalled()
303+
})
304+
305+
it("should include context messages when using external server", async () => {
306+
const mockResponse = {
307+
enhancedPrompt: "Enhanced with context",
308+
}
309+
vi.mocked(global.fetch).mockResolvedValue({
310+
ok: true,
311+
json: vi.fn().mockResolvedValue(mockResponse),
312+
} as any)
313+
314+
const configWithExternalServer: ProviderSettings = {
315+
...mockApiConfiguration,
316+
enhancePrompt: {
317+
useExternalServer: true,
318+
endpoint: "http://localhost:8000/enhance",
319+
},
320+
}
321+
322+
const mockClineMessages: ClineMessage[] = [
323+
{ type: "ask", text: "User message", ts: 1000 },
324+
{ type: "say", say: "text", text: "Assistant response", ts: 2000 },
325+
]
326+
327+
await MessageEnhancer.enhanceMessage({
328+
text: "Test prompt",
329+
apiConfiguration: configWithExternalServer,
330+
listApiConfigMeta: mockListApiConfigMeta,
331+
currentClineMessages: mockClineMessages,
332+
providerSettingsManager: mockProviderSettingsManager,
333+
})
334+
335+
expect(global.fetch).toHaveBeenCalledWith("http://localhost:8000/enhance", {
336+
method: "POST",
337+
headers: { "Content-Type": "application/json" },
338+
body: JSON.stringify({
339+
prompt: "Test prompt",
340+
context: [
341+
{ role: "user", content: "User message" },
342+
{ role: "assistant", content: "Assistant response" },
343+
],
344+
model: "gpt-4",
345+
}),
346+
})
347+
})
348+
349+
it("should fall back to internal enhancement when external server fails", async () => {
350+
vi.mocked(global.fetch).mockRejectedValue(new Error("Network error"))
351+
352+
const configWithExternalServer: ProviderSettings = {
353+
...mockApiConfiguration,
354+
enhancePrompt: {
355+
useExternalServer: true,
356+
endpoint: "http://localhost:8000/enhance",
357+
},
358+
}
359+
360+
const result = await MessageEnhancer.enhanceMessage({
361+
text: "Test prompt",
362+
apiConfiguration: configWithExternalServer,
363+
listApiConfigMeta: mockListApiConfigMeta,
364+
providerSettingsManager: mockProviderSettingsManager,
365+
})
366+
367+
expect(result.success).toBe(true)
368+
expect(result.enhancedText).toBe("Enhanced prompt text")
369+
expect(mockSingleCompletionHandler).toHaveBeenCalled()
370+
})
371+
372+
it("should fall back when external server returns non-ok status", async () => {
373+
vi.mocked(global.fetch).mockResolvedValue({
374+
ok: false,
375+
status: 500,
376+
statusText: "Internal Server Error",
377+
} as any)
378+
379+
const configWithExternalServer: ProviderSettings = {
380+
...mockApiConfiguration,
381+
enhancePrompt: {
382+
useExternalServer: true,
383+
endpoint: "http://localhost:8000/enhance",
384+
},
385+
}
386+
387+
const result = await MessageEnhancer.enhanceMessage({
388+
text: "Test prompt",
389+
apiConfiguration: configWithExternalServer,
390+
listApiConfigMeta: mockListApiConfigMeta,
391+
providerSettingsManager: mockProviderSettingsManager,
392+
})
393+
394+
expect(result.success).toBe(true)
395+
expect(result.enhancedText).toBe("Enhanced prompt text")
396+
expect(mockSingleCompletionHandler).toHaveBeenCalled()
397+
})
398+
399+
it("should fall back when external server response is missing enhancedPrompt", async () => {
400+
vi.mocked(global.fetch).mockResolvedValue({
401+
ok: true,
402+
json: vi.fn().mockResolvedValue({ wrongField: "value" }),
403+
} as any)
404+
405+
const configWithExternalServer: ProviderSettings = {
406+
...mockApiConfiguration,
407+
enhancePrompt: {
408+
useExternalServer: true,
409+
endpoint: "http://localhost:8000/enhance",
410+
},
411+
}
412+
413+
const result = await MessageEnhancer.enhanceMessage({
414+
text: "Test prompt",
415+
apiConfiguration: configWithExternalServer,
416+
listApiConfigMeta: mockListApiConfigMeta,
417+
providerSettingsManager: mockProviderSettingsManager,
418+
})
419+
420+
expect(result.success).toBe(true)
421+
expect(result.enhancedText).toBe("Enhanced prompt text")
422+
expect(mockSingleCompletionHandler).toHaveBeenCalled()
423+
})
424+
425+
it("should not use external server when useExternalServer is false", async () => {
426+
const configWithDisabledExternalServer: ProviderSettings = {
427+
...mockApiConfiguration,
428+
enhancePrompt: {
429+
useExternalServer: false,
430+
endpoint: "http://localhost:8000/enhance",
431+
},
432+
}
433+
434+
await MessageEnhancer.enhanceMessage({
435+
text: "Test prompt",
436+
apiConfiguration: configWithDisabledExternalServer,
437+
listApiConfigMeta: mockListApiConfigMeta,
438+
providerSettingsManager: mockProviderSettingsManager,
439+
})
440+
441+
expect(global.fetch).not.toHaveBeenCalled()
442+
expect(mockSingleCompletionHandler).toHaveBeenCalled()
443+
})
444+
445+
it("should not use external server when endpoint is missing", async () => {
446+
const configWithoutEndpoint: ProviderSettings = {
447+
...mockApiConfiguration,
448+
enhancePrompt: {
449+
useExternalServer: true,
450+
},
451+
}
452+
453+
await MessageEnhancer.enhanceMessage({
454+
text: "Test prompt",
455+
apiConfiguration: configWithoutEndpoint,
456+
listApiConfigMeta: mockListApiConfigMeta,
457+
providerSettingsManager: mockProviderSettingsManager,
458+
})
459+
460+
expect(global.fetch).not.toHaveBeenCalled()
461+
expect(mockSingleCompletionHandler).toHaveBeenCalled()
462+
})
463+
464+
it("should handle different model ID fields correctly", async () => {
465+
vi.mocked(global.fetch).mockResolvedValue({
466+
ok: true,
467+
json: vi.fn().mockResolvedValue({ enhancedPrompt: "Enhanced" }),
468+
} as any)
469+
470+
// Test with different provider configurations
471+
const configs = [
472+
{
473+
...mockApiConfiguration,
474+
apiModelId: "model-1",
475+
enhancePrompt: { useExternalServer: true, endpoint: "http://localhost:8000/enhance" },
476+
},
477+
{
478+
apiProvider: "ollama" as const,
479+
ollamaModelId: "llama2",
480+
enhancePrompt: { useExternalServer: true, endpoint: "http://localhost:8000/enhance" },
481+
},
482+
{
483+
apiProvider: "openrouter" as const,
484+
openRouterModelId: "gpt-4",
485+
enhancePrompt: { useExternalServer: true, endpoint: "http://localhost:8000/enhance" },
486+
},
487+
]
488+
489+
for (const config of configs) {
490+
vi.mocked(global.fetch).mockClear()
491+
492+
await MessageEnhancer.enhanceMessage({
493+
text: "Test",
494+
apiConfiguration: config,
495+
listApiConfigMeta: mockListApiConfigMeta,
496+
providerSettingsManager: mockProviderSettingsManager,
497+
})
498+
499+
const expectedModel = getModelId(config) || "unknown"
500+
expect(global.fetch).toHaveBeenCalledWith(
501+
expect.any(String),
502+
expect.objectContaining({
503+
body: expect.stringContaining(`"model":"${expectedModel}"`),
504+
}),
505+
)
506+
}
507+
})
508+
})
257509
})
258510

259511
describe("captureTelemetry", () => {

‎src/core/webview/messageEnhancer.ts‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { ProviderSettings, ClineMessage, GlobalState, TelemetryEventName } from "@roo-code/types"
1+
import { ProviderSettings, ClineMessage, GlobalState, TelemetryEventName, getModelId } from "@roo-code/types"
22
import { TelemetryService } from "@roo-code/telemetry"
33
import { supportPrompt } from "../../shared/support-prompt"
44
import { singleCompletionHandler } from "../../utils/single-completion-handler"
@@ -58,6 +58,55 @@ export class MessageEnhancer {
5858
}
5959
}
6060

61+
// Check if external MCP server is enabled
62+
if (configToUse.enhancePrompt?.useExternalServer && configToUse.enhancePrompt?.endpoint) {
63+
try {
64+
// Prepare context messages for external server
65+
const contextMessages =
66+
currentClineMessages
67+
?.filter((msg) => {
68+
if (msg.type === "ask" && msg.text) return true
69+
if (msg.type === "say" && msg.say === "text" && msg.text) return true
70+
return false
71+
})
72+
.slice(-10)
73+
.map((msg) => ({
74+
role: msg.type === "ask" ? "user" : "assistant",
75+
content: msg.text || "",
76+
})) || []
77+
78+
// Make request to external MCP server
79+
const response = await fetch(configToUse.enhancePrompt.endpoint, {
80+
method: "POST",
81+
headers: { "Content-Type": "application/json" },
82+
body: JSON.stringify({
83+
prompt: text,
84+
context: contextMessages,
85+
model: getModelId(configToUse) || "unknown",
86+
}),
87+
})
88+
89+
if (!response.ok) {
90+
throw new Error(`External server returned ${response.status}: ${response.statusText}`)
91+
}
92+
93+
const result = await response.json()
94+
95+
if (result.enhancedPrompt) {
96+
return {
97+
success: true,
98+
enhancedText: result.enhancedPrompt,
99+
}
100+
} else {
101+
throw new Error("External server response missing 'enhancedPrompt' field")
102+
}
103+
} catch (err) {
104+
console.error("Failed to enhance prompt via external server:", err)
105+
// Fallback to default logic
106+
}
107+
}
108+
109+
// Default internal enhancement logic
61110
// Prepare the prompt to enhance
62111
let promptToEnhance = text
63112

0 commit comments

Comments
 (0)