Skip to content

Commit 499df39

Browse files
committed
refactor(core): use canonical provider identifiers
1 parent e479b7e commit 499df39

6 files changed

Lines changed: 120 additions & 9 deletions

File tree

src/core/task/Task.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ import {
5454
ConsecutiveMistakeError,
5555
MAX_MCP_TOOLS_THRESHOLD,
5656
countEnabledMcpTools,
57+
providerIdentifiers,
5758
} from "@roo-code/types"
5859
import { TelemetryService } from "@roo-code/telemetry"
5960
import { CloudService } from "@roo-code/cloud"
@@ -4220,7 +4221,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
42204221
// but uses allowedFunctionNames to restrict which tools can be called.
42214222
// Other providers (Anthropic, OpenAI, etc.) don't support this feature yet,
42224223
// so they continue to receive only the filtered tools for the current mode.
4223-
const supportsAllowedFunctionNames = apiConfiguration?.apiProvider === "gemini"
4224+
const supportsAllowedFunctionNames = apiConfiguration?.apiProvider === providerIdentifiers.gemini
42244225

42254226
{
42264227
const provider = this.providerRef.deref()

src/core/task/__tests__/Task.spec.ts

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import * as path from "path"
66
import * as vscode from "vscode"
77
import { Anthropic } from "@anthropic-ai/sdk"
88

9-
import type { GlobalState, ProviderSettings, ModelInfo } from "@roo-code/types"
9+
import { providerIdentifiers, type GlobalState, type ProviderSettings, type ModelInfo } from "@roo-code/types"
1010
import { TelemetryService } from "@roo-code/telemetry"
1111

1212
import { Task } from "../Task"
@@ -1902,6 +1902,52 @@ describe("Cline", () => {
19021902
expect(metadata!.abortSignal).toBe(task.currentRequestAbortController!.signal)
19031903
})
19041904

1905+
it("enables allowed function names through the canonical Gemini provider identifier", async () => {
1906+
const identifiers = providerIdentifiers as Record<string, string>
1907+
const originalIdentifier = identifiers.gemini
1908+
1909+
try {
1910+
identifiers.gemini = "canonical-gemini"
1911+
const apiConfiguration = {
1912+
...mockApiConfig,
1913+
apiProvider: identifiers.gemini,
1914+
} as ProviderSettings
1915+
const task = new Task({
1916+
provider: mockProvider,
1917+
apiConfiguration,
1918+
task: "test task",
1919+
startTask: false,
1920+
})
1921+
1922+
vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt")
1923+
vi.spyOn(task.api, "getModel").mockReturnValue({
1924+
id: mockApiConfig.apiModelId!,
1925+
info: { contextWindow: 200000, maxTokens: 4096 } as ModelInfo,
1926+
})
1927+
const providerState = await mockProvider.getState()
1928+
vi.spyOn(mockProvider, "getState").mockResolvedValue({
1929+
...providerState,
1930+
apiConfiguration,
1931+
autoApprovalEnabled: true,
1932+
requestDelaySeconds: 0,
1933+
})
1934+
const mockStream = (async function* () {
1935+
yield { type: "text", text: "response" } as ApiStreamChunk
1936+
})()
1937+
const createMessageSpy = vi.spyOn(task.api, "createMessage").mockReturnValue(mockStream)
1938+
task.apiConversationHistory = [
1939+
{ role: "user", content: [{ type: "text", text: "test message" }], ts: Date.now() },
1940+
] as any
1941+
1942+
await task.attemptApiRequest(0).next()
1943+
1944+
const [, , metadata] = createMessageSpy.mock.calls[0]!
1945+
expect(metadata?.allowedFunctionNames).toEqual(expect.any(Array))
1946+
} finally {
1947+
identifiers.gemini = originalIdentifier
1948+
}
1949+
})
1950+
19051951
it("should invoke abort on currentRequestAbortController during first-chunk wait", async () => {
19061952
const task = new Task({
19071953
provider: mockProvider,

src/core/webview/ClineProvider.ts

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ import {
5050
DEFAULT_CHECKPOINT_TIMEOUT_SECONDS,
5151
getModelId,
5252
isRetiredProvider,
53+
providerIdentifiers,
5354
} from "@roo-code/types"
5455
import { RateLimitClock, createRateLimitClock } from "../task/RateLimitClock"
5556
import { aggregateTaskCostsRecursive, type AggregatedCosts } from "./aggregateTaskCosts"
@@ -478,7 +479,7 @@ export class ClineProvider
478479
async performPreparationTasks(cline: Task) {
479480
// LMStudio: We need to force model loading in order to read its context
480481
// size; we do it now since we're starting a task with that model selected.
481-
if (cline.apiConfiguration && cline.apiConfiguration.apiProvider === "lmstudio") {
482+
if (cline.apiConfiguration && cline.apiConfiguration.apiProvider === providerIdentifiers.lmstudio) {
482483
try {
483484
if (!hasLoadedFullDetails(cline.apiConfiguration.lmStudioModelId!)) {
484485
await forceFullModelDetailsLoad(
@@ -1003,7 +1004,7 @@ export class ClineProvider
10031004
// Using .find() would miss stale tokens in duplicate/renamed profiles since handleZooCodeCallback
10041005
// uses .filter() and updates all of them — the early-return guard must match.
10051006
const allProfiles = await this.providerSettingsManager.listConfig()
1006-
const zooGatewayProfiles = allProfiles.filter((p) => p.apiProvider === "zoo-gateway")
1007+
const zooGatewayProfiles = allProfiles.filter((p) => p.apiProvider === providerIdentifiers.zooGateway)
10071008

10081009
if (zooGatewayProfiles.length === 0) {
10091010
this.log("[ensureZooGatewayProfileSeeded] No zoo-gateway profile found, creating one")
@@ -1864,14 +1865,14 @@ export class ClineProvider
18641865

18651866
// Check if Zoo Gateway is the currently active profile by apiProvider identity,
18661867
// not by profile name (profile names are user-renameable).
1867-
const isZooGatewayActive = currentSettings.apiProvider === "zoo-gateway"
1868+
const isZooGatewayActive = currentSettings.apiProvider === providerIdentifiers.zooGateway
18681869

18691870
// Always scan ALL profiles and update every zoo-gateway profile with the new token.
18701871
// This ensures renamed profiles, duplicate profiles, and inactive profiles all stay
18711872
// in sync. The model lookup in requestRouterModels uses .find() which returns the
18721873
// first zoo-gateway profile it finds — if that profile has a stale token, requests fail.
18731874
const allProfiles = await this.providerSettingsManager.listConfig()
1874-
const zooProfiles = allProfiles.filter((p) => p.apiProvider === "zoo-gateway")
1875+
const zooProfiles = allProfiles.filter((p) => p.apiProvider === providerIdentifiers.zooGateway)
18751876

18761877
if (zooProfiles.length === 0) {
18771878
// No existing zoo-gateway profile — create the canonical default.

src/core/webview/__tests__/ClineProvider.spec.ts

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import {
1515
DEFAULT_CHECKPOINT_TIMEOUT_SECONDS,
1616
DEFAULT_DIFF_FUZZY_THRESHOLD,
1717
DEFAULT_WRITE_DELAY_MS,
18+
providerIdentifiers,
1819
} from "@roo-code/types"
1920
import { TelemetryService } from "@roo-code/telemetry"
2021

@@ -28,6 +29,7 @@ import { safeWriteJson } from "../../../utils/safeWriteJson"
2829
import { ClineProvider } from "../ClineProvider"
2930
import { Terminal } from "../../../integrations/terminal/Terminal"
3031
import { MessageManager } from "../../message-manager"
32+
import { forceFullModelDetailsLoad } from "../../../api/providers/fetchers/lmstudio"
3133

3234
// Mock setup must come before imports.
3335
vi.mock("../../prompts/sections/custom-instructions")
@@ -257,6 +259,11 @@ vi.mock("../../../api/providers/fetchers/modelCache", () => ({
257259
getModelsFromCache: vi.fn().mockReturnValue(undefined),
258260
}))
259261

262+
vi.mock("../../../api/providers/fetchers/lmstudio", () => ({
263+
hasLoadedFullDetails: vi.fn().mockReturnValue(false),
264+
forceFullModelDetailsLoad: vi.fn().mockResolvedValue(undefined),
265+
}))
266+
260267
vi.mock("../../../services/zoo-code-auth", () => ({
261268
getZooCodeBaseUrl: vi.fn(() => "https://www.zoocode.dev"),
262269
getCachedZooCodeToken: vi.fn(),
@@ -520,6 +527,27 @@ describe("ClineProvider", () => {
520527
expect(ClineProvider.getVisibleInstance()).toBe(provider)
521528
})
522529

530+
test("prepares LM Studio tasks through the canonical provider identifier", async () => {
531+
const identifiers = providerIdentifiers as Record<string, string>
532+
const originalIdentifier = identifiers.lmstudio
533+
534+
try {
535+
identifiers.lmstudio = "canonical-lmstudio"
536+
537+
await provider.performPreparationTasks({
538+
apiConfiguration: {
539+
apiProvider: identifiers.lmstudio,
540+
lmStudioBaseUrl: "http://localhost:1234",
541+
lmStudioModelId: "test-model",
542+
},
543+
} as Task)
544+
545+
expect(forceFullModelDetailsLoad).toHaveBeenCalledWith("http://localhost:1234", "test-model")
546+
} finally {
547+
identifiers.lmstudio = originalIdentifier
548+
}
549+
})
550+
523551
test("resolveWebviewView hydrates the saved terminalProfile into the process-wide Terminal state", async () => {
524552
const setTerminalProfileSpy = vi.spyOn(Terminal, "setTerminalProfile").mockImplementation(() => {})
525553
// Seed the persisted setting so the real getState() returns it during hydration.

src/core/webview/__tests__/webviewMessageHandler.spec.ts

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ vi.mock("../rulesMessageHandler", () => ({
5151
handleOpenRulesDirectory: vi.fn(),
5252
}))
5353

54-
import type { ModelRecord } from "@roo-code/types"
54+
import { providerIdentifiers, type ModelRecord } from "@roo-code/types"
5555

5656
import { webviewMessageHandler } from "../webviewMessageHandler"
5757
import type { ClineProvider } from "../ClineProvider"
@@ -1540,6 +1540,40 @@ describe("zooCodeSignOut", () => {
15401540
expect(mockClineProvider.postStateToWebview).toHaveBeenCalled()
15411541
})
15421542

1543+
it("clears tokens using the canonical Zoo Gateway provider identifier", async () => {
1544+
const identifiers = providerIdentifiers as Record<string, string>
1545+
const originalIdentifier = identifiers.zooGateway
1546+
const upsertProviderProfile = vi.fn().mockResolvedValue(undefined)
1547+
1548+
try {
1549+
identifiers.zooGateway = "canonical-zoo-gateway"
1550+
;(mockClineProvider as any).contextProxy = {
1551+
...mockClineProvider.contextProxy,
1552+
getProviderSettings: vi.fn().mockReturnValue({ apiProvider: identifiers.zooGateway }),
1553+
getValues: vi.fn().mockReturnValue({ currentApiConfigName: "Zoo Gateway" }),
1554+
}
1555+
;(mockClineProvider as any).providerSettingsManager = {
1556+
listConfig: vi.fn().mockResolvedValue([{ name: "Zoo Gateway", apiProvider: identifiers.zooGateway }]),
1557+
getProfile: vi.fn().mockResolvedValue({
1558+
apiProvider: identifiers.zooGateway,
1559+
zooSessionToken: "token-active",
1560+
}),
1561+
saveConfig: vi.fn(),
1562+
}
1563+
;(mockClineProvider as any).upsertProviderProfile = upsertProviderProfile
1564+
1565+
await webviewMessageHandler(mockClineProvider, { type: "zooCodeSignOut" })
1566+
1567+
expect(upsertProviderProfile).toHaveBeenCalledWith(
1568+
"Zoo Gateway",
1569+
expect.not.objectContaining({ zooSessionToken: expect.anything() }),
1570+
true,
1571+
)
1572+
} finally {
1573+
identifiers.zooGateway = originalIdentifier
1574+
}
1575+
})
1576+
15431577
it("still clears the in-memory handler when the active profile token is already empty on disk", async () => {
15441578
const upsertProviderProfile = vi.fn().mockResolvedValue(undefined)
15451579

src/core/webview/webviewMessageHandler.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import {
2222
checkoutDiffPayloadSchema,
2323
checkoutRestorePayloadSchema,
2424
getCompletionCheckpoint,
25+
providerIdentifiers,
2526
} from "@roo-code/types"
2627
import { customToolRegistry } from "@roo-code/core"
2728
import { CloudService } from "@roo-code/cloud"
@@ -2773,11 +2774,11 @@ export const webviewMessageHandler = async (
27732774
const allProfiles = await provider.providerSettingsManager.listConfig()
27742775
// Check if Zoo Gateway is the currently active profile by apiProvider identity
27752776
const currentSettings = provider.contextProxy.getProviderSettings()
2776-
const isZooGatewayActive = currentSettings.apiProvider === "zoo-gateway"
2777+
const isZooGatewayActive = currentSettings.apiProvider === providerIdentifiers.zooGateway
27772778
const currentApiConfigName = provider.contextProxy.getValues().currentApiConfigName
27782779

27792780
for (const entry of allProfiles) {
2780-
if (entry.apiProvider !== "zoo-gateway") {
2781+
if (entry.apiProvider !== providerIdentifiers.zooGateway) {
27812782
continue
27822783
}
27832784

0 commit comments

Comments
 (0)