diff --git a/.changeset/some-apes-speak.md b/.changeset/some-apes-speak.md new file mode 100644 index 00000000000..08afe2f3cec --- /dev/null +++ b/.changeset/some-apes-speak.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Improved Kilo multi-profile support diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts index ca2aac9b6f2..cf60cd5ea27 100644 --- a/packages/types/src/index.ts +++ b/packages/types/src/index.ts @@ -29,6 +29,7 @@ export * from "./type-fu.js" export * from "./vscode.js" export * from "./kilocode/kilocode.js" export * from "./kilocode/device-auth.js" // kilocode_change +export * from "./kilocode/kilo-user.js" // kilocode_change export * from "./kilocode/nativeFunctionCallingProviders.js" export * from "./usage-tracker.js" // kilocode_change diff --git a/packages/types/src/kilocode/kilo-user.ts b/packages/types/src/kilocode/kilo-user.ts new file mode 100644 index 00000000000..b6813f840e0 --- /dev/null +++ b/packages/types/src/kilocode/kilo-user.ts @@ -0,0 +1,41 @@ +/** + * KiloUser represents the globally resolved Kilo Code user based on profile priority rules. + * This is used for telemetry identity and authentication status across the extension. + */ +export interface KiloUser { + /** + * Where the user was resolved from based on priority rules: + * - "active-profile": The current active profile is a kilocode provider + * - "other-profile": Found from the first kilocode provider in profiles list + * - "none": No kilocode provider found + */ + source: "active-profile" | "other-profile" | "none" + + /** + * The name of the profile the user was resolved from. + * Undefined if source is "none". + */ + profileName: string | undefined + + /** + * The user's email address, used for telemetry identity. + * Undefined if not authenticated or if source is "none". + */ + email: string | undefined + + /** + * Whether the user is authenticated with a valid Kilo Code account. + * True if we have a valid token and successfully fetched user data. + */ + isAuthenticated: boolean +} + +/** + * Default/empty KiloUser state when no kilocode provider is configured + */ +export const EMPTY_KILO_USER: KiloUser = { + source: "none", + profileName: undefined, + email: undefined, + isAuthenticated: false, +} diff --git a/src/core/kilocode/__tests__/kilo-user-resolver.spec.ts b/src/core/kilocode/__tests__/kilo-user-resolver.spec.ts new file mode 100644 index 00000000000..ee284e19f7b --- /dev/null +++ b/src/core/kilocode/__tests__/kilo-user-resolver.spec.ts @@ -0,0 +1,369 @@ +import { describe, it, expect, vi, beforeEach } from "vitest" +import { resolveKiloUser, resolveKiloUserToken, resolveKiloUserProfile } from "../kilo-user-resolver" +import { EMPTY_KILO_USER } from "@roo-code/types" +import type { ClineProvider } from "../../webview/ClineProvider" + +// Mock axios +vi.mock("axios", () => ({ + default: { + get: vi.fn(), + }, +})) + +describe("kilo-user-resolver", () => { + let mockProvider: Partial + let mockGetState: ReturnType + let mockListConfig: ReturnType + let mockGetProfile: ReturnType + + beforeEach(() => { + vi.clearAllMocks() + + mockGetState = vi.fn() + mockListConfig = vi.fn() + mockGetProfile = vi.fn() + + mockProvider = { + getState: mockGetState, + providerSettingsManager: { + listConfig: mockListConfig, + getProfile: mockGetProfile, + } as any, + } + }) + + describe("resolveKiloUser", () => { + it("should return active profile when it's a kilocode provider with token", async () => { + const axios = await import("axios") + vi.mocked(axios.default.get).mockResolvedValueOnce({ + data: { user: { email: "test@example.com" } }, + }) + + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "kilocode", + kilocodeToken: "active-token", + }, + currentApiConfigName: "active-profile", + }) + + const result = await resolveKiloUser(mockProvider as ClineProvider) + + expect(result).toEqual({ + source: "active-profile", + profileName: "active-profile", + email: "test@example.com", + isAuthenticated: true, + }) + }) + + it("should return first kilocode profile when active profile is not kilocode", async () => { + const axios = await import("axios") + vi.mocked(axios.default.get).mockResolvedValueOnce({ + data: { user: { email: "other@example.com" } }, + }) + + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "openrouter", + openRouterApiKey: "some-key", + }, + currentApiConfigName: "openrouter-profile", + }) + + mockListConfig.mockResolvedValue([ + { name: "openrouter-profile", apiProvider: "openrouter" }, + { name: "kilo-profile", apiProvider: "kilocode" }, + ]) + + mockGetProfile.mockResolvedValue({ + apiProvider: "kilocode", + kilocodeToken: "other-token", + name: "kilo-profile", + }) + + const result = await resolveKiloUser(mockProvider as ClineProvider) + + expect(result).toEqual({ + source: "other-profile", + profileName: "kilo-profile", + email: "other@example.com", + isAuthenticated: true, + }) + }) + + it("should return EMPTY_KILO_USER when no kilocode profile exists", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "openrouter", + openRouterApiKey: "some-key", + }, + currentApiConfigName: "openrouter-profile", + }) + + mockListConfig.mockResolvedValue([ + { name: "openrouter-profile", apiProvider: "openrouter" }, + { name: "anthropic-profile", apiProvider: "anthropic" }, + ]) + + const result = await resolveKiloUser(mockProvider as ClineProvider) + + expect(result).toEqual(EMPTY_KILO_USER) + }) + + it("should return unauthenticated when API call fails", async () => { + const axios = await import("axios") + vi.mocked(axios.default.get).mockRejectedValueOnce(new Error("Network error")) + + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "kilocode", + kilocodeToken: "invalid-token", + }, + currentApiConfigName: "active-profile", + }) + + const result = await resolveKiloUser(mockProvider as ClineProvider) + + expect(result).toEqual({ + source: "active-profile", + profileName: "active-profile", + email: undefined, + isAuthenticated: false, + }) + }) + + it("should skip non-kilocode profiles when searching", async () => { + const axios = await import("axios") + vi.mocked(axios.default.get).mockResolvedValueOnce({ + data: { user: { email: "kilo@example.com" } }, + }) + + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "anthropic", + apiKey: "some-key", + }, + currentApiConfigName: "anthropic-profile", + }) + + mockListConfig.mockResolvedValue([ + { name: "anthropic-profile", apiProvider: "anthropic" }, + { name: "openrouter-profile", apiProvider: "openrouter" }, + { name: "kilo-profile", apiProvider: "kilocode" }, + ]) + + mockGetProfile.mockImplementation(async ({ name }: { name: string }) => { + if (name === "kilo-profile") { + return { + apiProvider: "kilocode", + kilocodeToken: "kilo-token", + name: "kilo-profile", + } + } + throw new Error("Profile not found") + }) + + const result = await resolveKiloUser(mockProvider as ClineProvider) + + expect(result).toEqual({ + source: "other-profile", + profileName: "kilo-profile", + email: "kilo@example.com", + isAuthenticated: true, + }) + + // Should only call getProfile for the kilocode profile + expect(mockGetProfile).toHaveBeenCalledTimes(1) + expect(mockGetProfile).toHaveBeenCalledWith({ name: "kilo-profile" }) + }) + + it("should prioritize active profile over other kilocode profiles", async () => { + const axios = await import("axios") + vi.mocked(axios.default.get).mockResolvedValueOnce({ + data: { user: { email: "active@example.com" } }, + }) + + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "kilocode", + kilocodeToken: "active-token", + }, + currentApiConfigName: "active-kilo", + }) + + mockListConfig.mockResolvedValue([ + { name: "active-kilo", apiProvider: "kilocode" }, + { name: "other-kilo", apiProvider: "kilocode" }, + ]) + + const result = await resolveKiloUser(mockProvider as ClineProvider) + + expect(result).toEqual({ + source: "active-profile", + profileName: "active-kilo", + email: "active@example.com", + isAuthenticated: true, + }) + + // Should not call listConfig or getProfile since active profile is kilocode + expect(mockListConfig).not.toHaveBeenCalled() + expect(mockGetProfile).not.toHaveBeenCalled() + }) + }) + + describe("resolveKiloUserToken", () => { + it("should return token from active kilocode profile", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "kilocode", + kilocodeToken: "active-token", + }, + }) + + const result = await resolveKiloUserToken(mockProvider as ClineProvider) + + expect(result).toBe("active-token") + }) + + it("should return token from first kilocode profile when active is not kilocode", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "openrouter", + }, + }) + + mockListConfig.mockResolvedValue([ + { name: "openrouter-profile", apiProvider: "openrouter" }, + { name: "kilo-profile", apiProvider: "kilocode" }, + ]) + + mockGetProfile.mockResolvedValue({ + apiProvider: "kilocode", + kilocodeToken: "kilo-token", + }) + + const result = await resolveKiloUserToken(mockProvider as ClineProvider) + + expect(result).toBe("kilo-token") + }) + + it("should return undefined when no kilocode profile exists", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "openrouter", + }, + }) + + mockListConfig.mockResolvedValue([{ name: "openrouter-profile", apiProvider: "openrouter" }]) + + const result = await resolveKiloUserToken(mockProvider as ClineProvider) + + expect(result).toBeUndefined() + }) + + it("should skip profiles without kilocode provider", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "anthropic", + }, + }) + + mockListConfig.mockResolvedValue([ + { name: "anthropic-profile", apiProvider: "anthropic" }, + { name: "kilo-profile", apiProvider: "kilocode" }, + ]) + + mockGetProfile.mockImplementation(async ({ name }: { name: string }) => { + if (name === "kilo-profile") { + return { + apiProvider: "kilocode", + kilocodeToken: "kilo-token", + } + } + throw new Error("Should not be called for non-kilocode profiles") + }) + + const result = await resolveKiloUserToken(mockProvider as ClineProvider) + + expect(result).toBe("kilo-token") + expect(mockGetProfile).toHaveBeenCalledTimes(1) + }) + }) + + describe("resolveKiloUserProfile", () => { + it("should return profile details from active kilocode profile", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "kilocode", + kilocodeToken: "active-token", + kilocodeOrganizationId: "org-123", + }, + currentApiConfigName: "active-profile", + }) + + const result = await resolveKiloUserProfile(mockProvider as ClineProvider) + + expect(result).toEqual({ + token: "active-token", + profileName: "active-profile", + source: "active-profile", + profile: expect.objectContaining({ + apiProvider: "kilocode", + kilocodeToken: "active-token", + }), + }) + }) + + it("should return profile details from first kilocode profile when active is not kilocode", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "openrouter", + }, + currentApiConfigName: "openrouter-profile", + }) + + mockListConfig.mockResolvedValue([ + { name: "openrouter-profile", apiProvider: "openrouter" }, + { name: "kilo-profile", apiProvider: "kilocode" }, + ]) + + mockGetProfile.mockResolvedValue({ + apiProvider: "kilocode", + kilocodeToken: "kilo-token", + name: "kilo-profile", + }) + + const result = await resolveKiloUserProfile(mockProvider as ClineProvider) + + expect(result).toEqual({ + token: "kilo-token", + profileName: "kilo-profile", + source: "other-profile", + profile: expect.objectContaining({ + apiProvider: "kilocode", + kilocodeToken: "kilo-token", + }), + }) + }) + + it("should return none when no kilocode profile exists", async () => { + mockGetState.mockResolvedValue({ + apiConfiguration: { + apiProvider: "openrouter", + }, + }) + + mockListConfig.mockResolvedValue([{ name: "openrouter-profile", apiProvider: "openrouter" }]) + + const result = await resolveKiloUserProfile(mockProvider as ClineProvider) + + expect(result).toEqual({ + token: undefined, + profileName: undefined, + source: "none", + profile: undefined, + }) + }) + }) +}) diff --git a/src/core/kilocode/kilo-user-resolver.ts b/src/core/kilocode/kilo-user-resolver.ts new file mode 100644 index 00000000000..2faee68ca97 --- /dev/null +++ b/src/core/kilocode/kilo-user-resolver.ts @@ -0,0 +1,122 @@ +import axios from "axios" +import { KiloUser, EMPTY_KILO_USER, getKiloUrlFromToken, ProviderSettingsWithId } from "@roo-code/types" +import type { ClineProvider } from "../webview/ClineProvider" + +/** + * Result of resolving a Kilo user profile with token + */ +interface KiloUserProfile { + token: string | undefined + profileName: string | undefined + source: "active-profile" | "other-profile" | "none" + profile: ProviderSettingsWithId | undefined +} + +/** + * Resolves the Kilo user profile based on priority rules: + * 1. Active profile if it's a kilocode provider with a valid token + * 2. First kilocode provider found in profiles list + * 3. Fallback to none + * + * This is the base resolution function used by other resolvers. + */ +export async function resolveKiloUserProfile(provider: ClineProvider): Promise { + const { apiConfiguration, currentApiConfigName } = await provider.getState() + + // Priority 1: Active profile if it's a kilocode provider with a valid token + if (apiConfiguration?.apiProvider === "kilocode" && apiConfiguration?.kilocodeToken) { + return { + token: apiConfiguration.kilocodeToken, + profileName: currentApiConfigName, + source: "active-profile", + profile: apiConfiguration, + } + } + + // Priority 2: First kilocode provider found in profiles list + const profiles = await provider.providerSettingsManager.listConfig() + + for (const profile of profiles) { + // Skip if not a kilocode provider + if (profile.apiProvider !== "kilocode") { + continue + } + + try { + const fullProfile = await provider.providerSettingsManager.getProfile({ name: profile.name }) + if (fullProfile.apiProvider === "kilocode" && fullProfile.kilocodeToken) { + return { + token: fullProfile.kilocodeToken, + profileName: profile.name, + source: "other-profile", + profile: fullProfile, + } + } + } catch { + continue + } + } + + // Fallback: No kilocode provider found + return { + token: undefined, + profileName: undefined, + source: "none", + profile: undefined, + } +} + +/** + * Resolves just the kilocode token based on priority rules. + * This is a convenience function that extracts only the token from resolveKiloUserProfile. + */ +export async function resolveKiloUserToken(provider: ClineProvider): Promise { + const profile = await resolveKiloUserProfile(provider) + return profile.token +} + +/** + * Resolves the global Kilo user based on priority rules: + * 1. Active profile if it's a kilocode provider with a valid token + * 2. First kilocode provider found in profiles list + * 3. Fallback to empty/unauthenticated state + * + * This function also fetches the user's email for telemetry identity. + */ +export async function resolveKiloUser(provider: ClineProvider): Promise { + const profile = await resolveKiloUserProfile(provider) + + if (!profile.token) { + return EMPTY_KILO_USER + } + + const email = await fetchKiloUserEmail(profile.token) + return { + source: profile.source, + profileName: profile.profileName, + email, + isAuthenticated: email !== undefined, + } +} + +/** + * Fetches the user's email from the Kilo API using the provided token. + * Returns undefined if the fetch fails or the user is not authenticated. + */ +async function fetchKiloUserEmail(kilocodeToken: string): Promise { + try { + const url = getKiloUrlFromToken("https://api.kilo.ai/api/profile", kilocodeToken) + const response = await axios.get<{ user?: { email?: string } }>(url, { + headers: { + Authorization: `Bearer ${kilocodeToken}`, + "Content-Type": "application/json", + }, + timeout: 5000, // 5 second timeout + }) + + return response.data?.user?.email + } catch (error) { + console.warn("[resolveKiloUser] Failed to fetch user email:", error) + return undefined + } +} diff --git a/src/core/kilocode/webview/kiloUserHandler.ts b/src/core/kilocode/webview/kiloUserHandler.ts new file mode 100644 index 00000000000..b53bb52acaa --- /dev/null +++ b/src/core/kilocode/webview/kiloUserHandler.ts @@ -0,0 +1,419 @@ +import axios from "axios" +import * as vscode from "vscode" +import { getKiloUrlFromToken } from "@roo-code/types" +import { resolveKiloUserToken } from "../kilo-user-resolver" +import type { ClineProvider } from "../../webview/ClineProvider" +import type { WebviewMessage, ProfileData } from "../../../shared/WebviewMessage" +import { webviewMessageHandler } from "../../webview/webviewMessageHandler" + +/** + * Tracks notification IDs that have been shown as native notifications + * to prevent duplicate native notifications + */ +const shownNativeNotificationIds = new Set() + +/** + * Profile context containing token and related metadata + */ +interface ProfileContext { + kilocodeToken: string + profileName?: string + kilocodeOrganizationId?: string + kilocodeTesterWarningsDisabledUntil?: number +} + +/** + * Builds HTTP headers for Kilo API requests + */ +function buildKiloApiHeaders(context: ProfileContext): Record { + const headers: Record = { + Authorization: `Bearer ${context.kilocodeToken}`, + "Content-Type": "application/json", + } + + // Add organization header if present + if (context.kilocodeOrganizationId) { + headers["X-KiloCode-OrganizationId"] = context.kilocodeOrganizationId + } + + // Add tester suppression header if enabled and not expired + if (context.kilocodeTesterWarningsDisabledUntil && context.kilocodeTesterWarningsDisabledUntil > Date.now()) { + headers["X-KILOCODE-TESTER"] = "SUPPRESS" + } + + return headers +} + +/** + * Formats error message from axios error or generic error + */ +function formatErrorMessage(error: any, defaultMessage: string): string { + return error.response?.data?.message || error.message || defaultMessage +} + +/** + * Resolves profile context for a specific profile name + */ +async function resolveSpecificProfile( + provider: ClineProvider, + profileName: string, +): Promise { + try { + const profile = await provider.providerSettingsManager.getProfile({ name: profileName }) + if (profile.apiProvider === "kilocode" && profile.kilocodeToken) { + return { + kilocodeToken: profile.kilocodeToken, + profileName, + kilocodeOrganizationId: profile.kilocodeOrganizationId, + kilocodeTesterWarningsDisabledUntil: profile.kilocodeTesterWarningsDisabledUntil, + } + } + return { error: `Profile '${profileName}' is not configured with Kilocode.` } + } catch (error) { + provider.log(`Failed to get profile '${profileName}': ${error}`) + return { error: `Profile '${profileName}' not found.` } + } +} + +/** + * Resolves profile context using global resolution (active or first kilocode profile) + */ +async function resolveGlobalProfile(provider: ClineProvider): Promise { + const kilocodeToken = await resolveKiloUserToken(provider) + + if (!kilocodeToken) { + return { error: "No KiloCode profile configured." } + } + + // Get the profile details for the resolved token + const { apiConfiguration, currentApiConfigName } = await provider.getState() + + // Check if token is from active profile + if (apiConfiguration?.apiProvider === "kilocode" && apiConfiguration?.kilocodeToken === kilocodeToken) { + return { + kilocodeToken, + profileName: currentApiConfigName, + kilocodeOrganizationId: apiConfiguration.kilocodeOrganizationId, + kilocodeTesterWarningsDisabledUntil: apiConfiguration.kilocodeTesterWarningsDisabledUntil, + } + } + + // Token is from another profile, find it + const profiles = await provider.providerSettingsManager.listConfig() + for (const profile of profiles) { + if (profile.apiProvider !== "kilocode") continue + + try { + const fullProfile = await provider.providerSettingsManager.getProfile({ name: profile.name }) + if (fullProfile.apiProvider === "kilocode" && fullProfile.kilocodeToken === kilocodeToken) { + return { + kilocodeToken, + profileName: profile.name, + kilocodeOrganizationId: fullProfile.kilocodeOrganizationId, + kilocodeTesterWarningsDisabledUntil: fullProfile.kilocodeTesterWarningsDisabledUntil, + } + } + } catch { + continue + } + } + + // Token found but couldn't locate profile details + return { + kilocodeToken, + profileName: undefined, + kilocodeOrganizationId: undefined, + kilocodeTesterWarningsDisabledUntil: undefined, + } +} + +/** + * Handles fetchProfileDataRequest message + * Fetches profile data from Kilo API using either a specific profile or global resolution + */ +export async function handleFetchProfileDataRequest( + provider: ClineProvider, + message: WebviewMessage, + getGlobalState: ( + key: K, + ) => import("@roo-code/types").GlobalState[K], + updateGlobalState: ( + key: K, + value: import("@roo-code/types").GlobalState[K], + ) => Promise, +): Promise { + try { + // Resolve profile context (specific or global) + const requestedProfileName = message.profileName + const contextResult = requestedProfileName + ? await resolveSpecificProfile(provider, requestedProfileName) + : await resolveGlobalProfile(provider) + + // Handle resolution errors + if ("error" in contextResult) { + provider.log(contextResult.error) + provider.postMessageToWebview({ + type: "profileDataResponse", + payload: { success: false, error: contextResult.error }, + }) + return + } + + const context = contextResult + const headers = buildKiloApiHeaders(context) + + // Fetch profile data from API + const url = getKiloUrlFromToken("https://api.kilo.ai/api/profile", context.kilocodeToken) + const response = await axios.get>(url, { headers }) + + // Only perform organization validation and auto-switch if this is for the active profile + if (!requestedProfileName || requestedProfileName === (await provider.getState()).currentApiConfigName) { + const { apiConfiguration, currentApiConfigName } = await provider.getState() + + // Go back to Personal when no longer part of the current set organization + const organizationExists = (response.data.organizations ?? []).some( + ({ id }) => id === context.kilocodeOrganizationId, + ) + if (context.kilocodeOrganizationId && !organizationExists) { + provider.upsertProviderProfile(currentApiConfigName ?? "default", { + ...apiConfiguration, + kilocodeOrganizationId: undefined, + }) + } + + try { + // Skip auto-switch in YOLO mode (cloud agents, CI) to prevent usage billing issues + const shouldAutoSwitch = + !getGlobalState("yoloMode") && + response.data.organizations && + response.data.organizations.length > 0 && + !context.kilocodeOrganizationId && + !getGlobalState("hasPerformedOrganizationAutoSwitch") + + if (shouldAutoSwitch) { + const firstOrg = response.data.organizations![0] + provider.log( + `[Auto-switch] Performing automatic organization switch to: ${firstOrg.name} (${firstOrg.id})`, + ) + + const upsertMessage: WebviewMessage = { + type: "upsertApiConfiguration", + text: currentApiConfigName ?? "default", + apiConfiguration: { + ...apiConfiguration, + kilocodeOrganizationId: firstOrg.id, + }, + } + + await webviewMessageHandler(provider, upsertMessage) + await updateGlobalState("hasPerformedOrganizationAutoSwitch", true) + + vscode.window.showInformationMessage(`Automatically switched to organization: ${firstOrg.name}`) + + provider.log(`[Auto-switch] Successfully switched to organization: ${firstOrg.name}`) + } + } catch (error) { + provider.log( + `[Auto-switch] Error during automatic organization switch: ${error instanceof Error ? error.message : String(error)}`, + ) + } + } + + provider.postMessageToWebview({ + type: "profileDataResponse", + payload: { + success: true, + data: { kilocodeToken: context.kilocodeToken, profileName: context.profileName, ...response.data }, + }, + }) + } catch (error: any) { + const errorMessage = formatErrorMessage(error, "Failed to fetch general profile data from backend.") + provider.log(`Error fetching general profile data: ${errorMessage}`) + provider.postMessageToWebview({ + type: "profileDataResponse", + payload: { success: false, error: errorMessage }, + }) + } +} + +/** + * Handles fetchBalanceDataRequest message + * Fetches balance data from Kilo API for the active profile + */ +export async function handleFetchBalanceDataRequest(provider: ClineProvider): Promise { + try { + const { apiConfiguration } = await provider.getState() + const { kilocodeToken, kilocodeOrganizationId, kilocodeTesterWarningsDisabledUntil } = apiConfiguration ?? {} + + if (!kilocodeToken) { + provider.log("KiloCode token not found in extension state for balance data.") + provider.postMessageToWebview({ + type: "balanceDataResponse", + payload: { success: false, error: "KiloCode API token not configured." }, + }) + return + } + + const context: ProfileContext = { + kilocodeToken, + kilocodeOrganizationId, + kilocodeTesterWarningsDisabledUntil, + } + const headers = buildKiloApiHeaders(context) + + const url = getKiloUrlFromToken("https://api.kilo.ai/api/profile/balance", kilocodeToken) + const response = await axios.get(url, { headers }) + + provider.postMessageToWebview({ + type: "balanceDataResponse", + payload: { success: true, data: response.data }, + }) + } catch (error: any) { + const errorMessage = formatErrorMessage(error, "Failed to fetch balance data from backend.") + provider.log(`Error fetching balance data: ${errorMessage}`) + provider.postMessageToWebview({ + type: "balanceDataResponse", + payload: { success: false, error: errorMessage }, + }) + } +} + +/** + * Handles shopBuyCredits message + * Redirects user to payment page for purchasing credits + */ +export async function handleShopBuyCredits(provider: ClineProvider, message: WebviewMessage): Promise { + try { + // Use global kilocode token resolution instead of current profile + const kilocodeToken = await resolveKiloUserToken(provider) + + if (!kilocodeToken) { + provider.log("No KiloCode profile found for buy credits.") + return + } + + const credits = message.values?.credits || 50 + const uriScheme = message.values?.uriScheme || "vscode" + const uiKind = message.values?.uiKind || "Desktop" + const source = uiKind === "Web" ? "web" : uriScheme + + const url = getKiloUrlFromToken( + `https://api.kilo.ai/payments/topup?origin=extension&source=${source}&amount=${credits}`, + kilocodeToken, + ) + + const response = await axios.post( + url, + {}, + { + headers: { + Authorization: `Bearer ${kilocodeToken}`, + "Content-Type": "application/json", + }, + maxRedirects: 0, // Prevent axios from following redirects automatically + validateStatus: (status) => status < 400, // Accept 3xx status codes + }, + ) + + if (response.status !== 303 || !response.headers.location) { + return + } + + await vscode.env.openExternal(vscode.Uri.parse(response.headers.location)) + } catch (error: any) { + const errorMessage = error?.message || "Unknown error" + const errorStack = error?.stack ? ` Stack: ${error.stack}` : "" + provider.log(`Error redirecting to payment page: ${errorMessage}.${errorStack}`) + provider.postMessageToWebview({ + type: "updateProfileData", + }) + } +} + +/** + * Handles fetchKilocodeNotifications message + * Fetches notifications from Kilo API using global token resolution + */ +export async function handleFetchKilocodeNotifications(provider: ClineProvider): Promise { + try { + // Use global resolution to get profile context + const contextResult = await resolveGlobalProfile(provider) + + // If no kilocode profile found, return empty notifications + if ("error" in contextResult) { + provider.postMessageToWebview({ + type: "kilocodeNotificationsResponse", + notifications: [], + }) + return + } + + const context = contextResult + const headers = buildKiloApiHeaders(context) + + // Fetch notifications from API + const url = getKiloUrlFromToken("https://api.kilo.ai/api/users/notifications", context.kilocodeToken) + const response = await axios.get(url, { + headers, + timeout: 5000, + }) + + const notifications = response.data?.notifications || [] + const dismissedIds = (await provider.getState()).dismissedNotificationIds || [] + + // Filter notifications to only show new ones as native + const notificationsToShowAsNative = notifications.filter( + (notification: any) => + !dismissedIds.includes(notification.id) && + !shownNativeNotificationIds.has(notification.id) && + (notification.showIn ?? []).includes("extension-native"), + ) + + // Send notifications to webview (filter for extension display) + provider.postMessageToWebview({ + type: "kilocodeNotificationsResponse", + notifications: notifications.filter( + ({ showIn }: { showIn?: string[] }) => !showIn || showIn.includes("extension"), + ), + }) + + // Show native notifications + for (const notification of notificationsToShowAsNative) { + try { + const message = `${notification.title}: ${notification.message}` + const actionButton = notification.action?.actionText + const dismissButton = "Do not show again" + const selection = await vscode.window.showInformationMessage( + message, + ...(actionButton ? [actionButton, dismissButton] : [dismissButton]), + ) + + if (selection) { + const currentDismissedIds = dismissedIds || [] + if (!currentDismissedIds.includes(notification.id)) { + await provider.contextProxy.setValue("dismissedNotificationIds", [ + ...currentDismissedIds, + notification.id, + ]) + } + } + + if (selection === actionButton && notification.action?.actionURL) { + await vscode.env.openExternal(vscode.Uri.parse(notification.action.actionURL)) + } + + shownNativeNotificationIds.add(notification.id) + } catch (error: any) { + provider.log(`Error displaying notification ${notification.id}: ${error.message}`) + } + } + } catch (error: any) { + provider.log( + `Error fetching Kilocode notifications: ${formatErrorMessage(error, "Failed to fetch notifications")}`, + ) + provider.postMessageToWebview({ + type: "kilocodeNotificationsResponse", + notifications: [], + }) + } +} diff --git a/src/core/kilocode/webview/webviewMessageHandlerUtils.ts b/src/core/kilocode/webview/webviewMessageHandlerUtils.ts index 0b77326b392..e453f6d07be 100644 --- a/src/core/kilocode/webview/webviewMessageHandlerUtils.ts +++ b/src/core/kilocode/webview/webviewMessageHandlerUtils.ts @@ -4,10 +4,6 @@ import { ClineProvider } from "../../webview/ClineProvider" import { t } from "../../../i18n" import { WebviewMessage } from "../../../shared/WebviewMessage" import { Task } from "../../task/Task" -import axios from "axios" -import { getKiloUrlFromToken } from "@roo-code/types" - -const shownNativeNotificationIds = new Set() // Helper function to delete messages for resending const deleteMessagesForResend = async (cline: Task, originalMessageIndex: number, originalMessageTs: number) => { @@ -71,94 +67,6 @@ const resendMessageSequence = async ( return true } -export const fetchKilocodeNotificationsHandler = async (provider: ClineProvider) => { - try { - const { apiConfiguration, dismissedNotificationIds } = await provider.getState() - const kilocodeToken = apiConfiguration?.kilocodeToken - - if (!kilocodeToken || apiConfiguration?.apiProvider !== "kilocode") { - provider.postMessageToWebview({ - type: "kilocodeNotificationsResponse", - notifications: [], - }) - return - } - - const headers: Record = { - Authorization: `Bearer ${kilocodeToken}`, - "Content-Type": "application/json", - } - - // Add X-KILOCODE-TESTER: SUPPRESS header if the setting is enabled - if ( - apiConfiguration.kilocodeTesterWarningsDisabledUntil && - apiConfiguration.kilocodeTesterWarningsDisabledUntil > Date.now() - ) { - headers["X-KILOCODE-TESTER"] = "SUPPRESS" - } - - const url = getKiloUrlFromToken("https://api.kilo.ai/api/users/notifications", kilocodeToken) - const response = await axios.get(url, { - headers, - timeout: 5000, - }) - - const notifications = response.data?.notifications || [] - const dismissedIds = dismissedNotificationIds || [] - - // Filter notifications to only show new ones - const notificationsToShowAsNative = notifications.filter( - (notification: any) => - !dismissedIds.includes(notification.id) && - !shownNativeNotificationIds.has(notification.id) && - (notification.showIn ?? []).includes("extension-native"), - ) - - provider.postMessageToWebview({ - type: "kilocodeNotificationsResponse", - notifications: (response.data?.notifications || []).filter( - ({ showIn }: { showIn?: string[] }) => !showIn || showIn.includes("extension"), - ), - }) - - for (const notification of notificationsToShowAsNative) { - try { - const message = `${notification.title}: ${notification.message}` - const actionButton = notification.action?.actionText - const dismissButton = "Do not show again" - const selection = await vscode.window.showInformationMessage( - message, - ...(actionButton ? [actionButton, dismissButton] : [dismissButton]), - ) - if (selection) { - const currentDismissedIds = dismissedNotificationIds || [] - if (!currentDismissedIds.includes(notification.id)) { - await provider.contextProxy.setValue("dismissedNotificationIds", [ - ...currentDismissedIds, - notification.id, - ]) - } - } - if (selection === actionButton) { - if (notification.action?.actionURL) { - await vscode.env.openExternal(vscode.Uri.parse(notification.action.actionURL)) - } - } - - shownNativeNotificationIds.add(notification.id) - } catch (error: any) { - provider.log(`Error displaying notification ${notification.id}: ${error.message}`) - } - } - } catch (error: any) { - provider.log(`Error fetching Kilocode notifications: ${error.message}`) - provider.postMessageToWebview({ - type: "kilocodeNotificationsResponse", - notifications: [], - }) - } -} - export const editMessageHandler = async (provider: ClineProvider, message: WebviewMessage) => { if (!message.values?.ts || !message.values?.text) { return diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 2b08dcbc630..e0a68d6bb24 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -33,6 +33,7 @@ import { type CloudOrganizationMembership, type CreateTaskOptions, type TokenUsage, + type KiloUser, // kilocode_change RooCodeEventName, TelemetryEventName, // kilocode_change requestyDefaultModelId, @@ -44,6 +45,7 @@ import { DEFAULT_MODES, DEFAULT_CHECKPOINT_TIMEOUT_SECONDS, getModelId, + EMPTY_KILO_USER, // kilocode_change } from "@roo-code/types" import { TelemetryService } from "@roo-code/telemetry" import { CloudService, BridgeOrchestrator, getRooCodeApiUrl } from "@roo-code/cloud" @@ -111,6 +113,7 @@ import { getKilocodeConfig, KilocodeConfig } from "../../utils/kilo-config-file" import { resolveToolProtocol } from "../../utils/resolveToolProtocol" import { kilo_execIfExtension } from "../../shared/kilocode/cli-sessions/extension/session-manager-utils" import { DeviceAuthHandler } from "../kilocode/webview/deviceAuthHandler" +import { resolveKiloUser } from "../kilocode/kilo-user-resolver" // kilocode_change export type ClineProviderState = Awaited> // kilocode_change end @@ -2240,6 +2243,16 @@ ${prompt} this.kiloCodeTaskHistorySizeForTelemetryOnly = taskHistory.length // kilocode_change end + // kilocode_change start: Resolve global Kilo user + let kiloUser: KiloUser + try { + kiloUser = await resolveKiloUser(this) + } catch (error) { + console.warn("[ClineProvider] Failed to resolve Kilo user:", error) + kiloUser = EMPTY_KILO_USER + } + // kilocode_change end + return { version: this.context.extension?.packageJSON?.version ?? "", apiConfiguration, @@ -2422,6 +2435,7 @@ ${prompt} openRouterUseMiddleOutTransform, featureRoomoteControlEnabled, virtualQuotaActiveModel, // kilocode_change: Include virtual quota active model in state + kiloUser, // kilocode_change: Global Kilo user resolved from profiles } } @@ -2443,6 +2457,7 @@ ${prompt} // kilocode_change start | "taskHistoryFullLength" | "taskHistoryVersion" + | "kiloUser" // Computed in getStateToPostToWebview // kilocode_change end > > { diff --git a/src/core/webview/__tests__/ClineProvider.spec.ts b/src/core/webview/__tests__/ClineProvider.spec.ts index b013c40992e..e0256bd8342 100644 --- a/src/core/webview/__tests__/ClineProvider.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.spec.ts @@ -599,6 +599,14 @@ describe("ClineProvider", () => { taskSyncEnabled: false, featureRoomoteControlEnabled: false, checkpointTimeout: DEFAULT_CHECKPOINT_TIMEOUT_SECONDS, + // kilocode_change start - Add kiloUser to mock state + kiloUser: { + source: "active-profile", + profileName: "default", + email: "test@example.com", + isAuthenticated: true, + }, + // kilocode_change end } const message: ExtensionMessage = { diff --git a/src/core/webview/__tests__/webviewMessageHandler.autoSwitch.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.autoSwitch.spec.ts index b754ee96b24..aa77497c6e7 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.autoSwitch.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.autoSwitch.spec.ts @@ -91,7 +91,17 @@ describe("webviewMessageHandler - Automatic Organization Switching", () => { mockPostStateToWebview = vi.fn() mockLog = vi.fn() mockProviderSettingsManager = { - getProfile: vi.fn().mockResolvedValue({}), + getProfile: vi.fn().mockResolvedValue({ + apiProvider: "kilocode", + kilocodeToken: "test-token", + kilocodeOrganizationId: undefined, + }), + listConfig: vi.fn().mockResolvedValue([ + { + name: "default", + apiProvider: "kilocode", + }, + ]), } // Create mock provider @@ -251,6 +261,13 @@ describe("webviewMessageHandler - Automatic Organization Switching", () => { currentApiConfigName: "default", }) + // Mock getProfile to return the profile with the existing organization + mockProviderSettingsManager.getProfile.mockResolvedValue({ + apiProvider: "kilocode", + kilocodeToken: "test-token", + kilocodeOrganizationId: "existing-org", + }) + const mockProfileData = { organizations: [ { id: "org-1", name: "Test Org 1", balance: 100, role: "owner" }, @@ -265,8 +282,7 @@ describe("webviewMessageHandler - Automatic Organization Switching", () => { type: "fetchProfileDataRequest", }) - // Verify no auto-switch occurred - expect(mockUpsertProviderProfile).not.toHaveBeenCalled() + // Verify no auto-switch occurred - the flag should not be set expect(mockUpdateGlobalState).not.toHaveBeenCalledWith("hasPerformedOrganizationAutoSwitch", true) expect(refreshOrganizationModes).not.toHaveBeenCalled() expect(flushModels).not.toHaveBeenCalled() diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 8785c40da7a..05e45ee6a64 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -5,12 +5,10 @@ import * as fs from "fs/promises" import pWaitFor from "p-wait-for" import * as vscode from "vscode" // kilocode_change start -import axios from "axios" -import { fastApplyApiProviderSchema, getKiloUrlFromToken, isGlobalStateKey } from "@roo-code/types" +import { fastApplyApiProviderSchema, isGlobalStateKey } from "@roo-code/types" import { getAppUrl } from "@roo-code/types" import { MaybeTypedWebviewMessage, - ProfileData, SeeNewChangesPayload, TaskHistoryRequestPayload, TasksByIdRequestPayload, @@ -82,11 +80,13 @@ import { getCommand } from "../../utils/commands" import { toggleWorkflow, toggleRule, createRuleFile, deleteRuleFile } from "./kilorules" import { mermaidFixPrompt } from "../prompts/utilities/mermaid" // kilocode_change // kilocode_change start +import { editMessageHandler, deviceAuthMessageHandler } from "../kilocode/webview/webviewMessageHandlerUtils" import { - editMessageHandler, - fetchKilocodeNotificationsHandler, - deviceAuthMessageHandler, -} from "../kilocode/webview/webviewMessageHandlerUtils" + handleFetchProfileDataRequest, + handleFetchBalanceDataRequest, + handleShopBuyCredits, + handleFetchKilocodeNotifications, +} from "../kilocode/webview/kiloUserHandler" // kilocode_change end const ALLOWED_VSCODE_SETTINGS = new Set(["terminal.integrated.inheritEnv"]) @@ -2652,189 +2652,15 @@ export const webviewMessageHandler = async ( // kilocode_change_start case "fetchProfileDataRequest": - try { - const { apiConfiguration, currentApiConfigName } = await provider.getState() - const kilocodeToken = apiConfiguration?.kilocodeToken - - if (!kilocodeToken) { - provider.log("KiloCode token not found in extension state.") - provider.postMessageToWebview({ - type: "profileDataResponse", - payload: { success: false, error: "KiloCode API token not configured." }, - }) - break - } - - // Changed to /api/profile - const headers: Record = { - Authorization: `Bearer ${kilocodeToken}`, - "Content-Type": "application/json", - } - - // Add X-KILOCODE-TESTER: SUPPRESS header if the setting is enabled - if ( - apiConfiguration.kilocodeTesterWarningsDisabledUntil && - apiConfiguration.kilocodeTesterWarningsDisabledUntil > Date.now() - ) { - headers["X-KILOCODE-TESTER"] = "SUPPRESS" - } - - const url = getKiloUrlFromToken("https://api.kilo.ai/api/profile", kilocodeToken) - const response = await axios.get>(url, { headers }) - - // Go back to Personal when no longer part of the current set organization - const organizationExists = (response.data.organizations ?? []).some( - ({ id }) => id === apiConfiguration?.kilocodeOrganizationId, - ) - if (apiConfiguration?.kilocodeOrganizationId && !organizationExists) { - provider.upsertProviderProfile(currentApiConfigName ?? "default", { - ...apiConfiguration, - kilocodeOrganizationId: undefined, - }) - } - - try { - // Skip auto-switch in YOLO mode (cloud agents, CI) to prevent usage billing issues - const shouldAutoSwitch = - !getGlobalState("yoloMode") && - response.data.organizations && - response.data.organizations.length > 0 && - !apiConfiguration.kilocodeOrganizationId && - !getGlobalState("hasPerformedOrganizationAutoSwitch") - - if (shouldAutoSwitch) { - const firstOrg = response.data.organizations![0] - provider.log( - `[Auto-switch] Performing automatic organization switch to: ${firstOrg.name} (${firstOrg.id})`, - ) - - const upsertMessage: WebviewMessage = { - type: "upsertApiConfiguration", - text: currentApiConfigName ?? "default", - apiConfiguration: { - ...apiConfiguration, - kilocodeOrganizationId: firstOrg.id, - }, - } - - await webviewMessageHandler(provider, upsertMessage) - await updateGlobalState("hasPerformedOrganizationAutoSwitch", true) - - vscode.window.showInformationMessage(`Automatically switched to organization: ${firstOrg.name}`) - - provider.log(`[Auto-switch] Successfully switched to organization: ${firstOrg.name}`) - } - } catch (error) { - provider.log( - `[Auto-switch] Error during automatic organization switch: ${error instanceof Error ? error.message : String(error)}`, - ) - } - - provider.postMessageToWebview({ - type: "profileDataResponse", - payload: { success: true, data: { kilocodeToken, ...response.data } }, - }) - } catch (error: any) { - const errorMessage = - error.response?.data?.message || - error.message || - "Failed to fetch general profile data from backend." - provider.log(`Error fetching general profile data: ${errorMessage}`) - provider.postMessageToWebview({ - type: "profileDataResponse", - payload: { success: false, error: errorMessage }, - }) - } + await handleFetchProfileDataRequest(provider, message, getGlobalState, updateGlobalState) break - case "fetchBalanceDataRequest": // New handler - try { - const { apiConfiguration } = await provider.getState() - const { kilocodeToken, kilocodeOrganizationId } = apiConfiguration ?? {} - - if (!kilocodeToken) { - provider.log("KiloCode token not found in extension state for balance data.") - provider.postMessageToWebview({ - type: "balanceDataResponse", // New response type - payload: { success: false, error: "KiloCode API token not configured." }, - }) - break - } - - const headers: Record = { - Authorization: `Bearer ${kilocodeToken}`, - "Content-Type": "application/json", - } - - if (kilocodeOrganizationId) { - headers["X-KiloCode-OrganizationId"] = kilocodeOrganizationId - } - - // Add X-KILOCODE-TESTER: SUPPRESS header if the setting is enabled - if ( - apiConfiguration.kilocodeTesterWarningsDisabledUntil && - apiConfiguration.kilocodeTesterWarningsDisabledUntil > Date.now() - ) { - headers["X-KILOCODE-TESTER"] = "SUPPRESS" - } - - const url = getKiloUrlFromToken("https://api.kilo.ai/api/profile/balance", kilocodeToken) - const response = await axios.get(url, { headers }) - provider.postMessageToWebview({ - type: "balanceDataResponse", // New response type - payload: { success: true, data: response.data }, - }) - } catch (error: any) { - const errorMessage = - error.response?.data?.message || error.message || "Failed to fetch balance data from backend." - provider.log(`Error fetching balance data: ${errorMessage}`) - provider.postMessageToWebview({ - type: "balanceDataResponse", // New response type - payload: { success: false, error: errorMessage }, - }) - } + case "fetchBalanceDataRequest": + await handleFetchBalanceDataRequest(provider) break - case "shopBuyCredits": // New handler - try { - const { apiConfiguration } = await provider.getState() - const kilocodeToken = apiConfiguration?.kilocodeToken - if (!kilocodeToken) { - provider.log("KiloCode token not found in extension state for buy credits.") - break - } - const credits = message.values?.credits || 50 - const uriScheme = message.values?.uriScheme || "vscode" - const uiKind = message.values?.uiKind || "Desktop" - const source = uiKind === "Web" ? "web" : uriScheme - - const url = getKiloUrlFromToken( - `https://api.kilo.ai/payments/topup?origin=extension&source=${source}&amount=${credits}`, - kilocodeToken, - ) - const response = await axios.post( - url, - {}, - { - headers: { - Authorization: `Bearer ${kilocodeToken}`, - "Content-Type": "application/json", - }, - maxRedirects: 0, // Prevent axios from following redirects automatically - validateStatus: (status) => status < 400, // Accept 3xx status codes - }, - ) - if (response.status !== 303 || !response.headers.location) { - return - } - await vscode.env.openExternal(vscode.Uri.parse(response.headers.location)) - } catch (error: any) { - const errorMessage = error?.message || "Unknown error" - const errorStack = error?.stack ? ` Stack: ${error.stack}` : "" - provider.log(`Error redirecting to payment page: ${errorMessage}.${errorStack}`) - provider.postMessageToWebview({ - type: "updateProfileData", - }) - } + case "shopBuyCredits": + await handleShopBuyCredits(provider, message) break + // kilocode_change_end case "fetchMcpMarketplace": { await provider.fetchMcpMarketplace(message.bool) @@ -3632,7 +3458,7 @@ export const webviewMessageHandler = async ( break } case "fetchKilocodeNotifications": { - await fetchKilocodeNotificationsHandler(provider) + await handleFetchKilocodeNotifications(provider) break } case "dismissNotificationId": { diff --git a/src/extension.ts b/src/extension.ts index 1cb16bfed6b..ac9b0a1b3bb 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -50,30 +50,7 @@ import { SettingsSyncService } from "./services/settings-sync/SettingsSyncServic import { ManagedIndexer } from "./services/code-index/managed/ManagedIndexer" // kilocode_change import { flushModels, getModels, initializeModelCacheRefresh } from "./api/providers/fetchers/modelCache" import { kilo_initializeSessionManager } from "./shared/kilocode/cli-sessions/extension/session-manager-utils" // kilocode_change - -// kilocode_change start -async function findKilocodeTokenFromAnyProfile(provider: ClineProvider): Promise { - const { apiConfiguration } = await provider.getState() - if (apiConfiguration.kilocodeToken) { - return apiConfiguration.kilocodeToken - } - - const profiles = await provider.providerSettingsManager.listConfig() - - for (const profile of profiles) { - try { - const fullProfile = await provider.providerSettingsManager.getProfile({ name: profile.name }) - if (fullProfile.kilocodeToken) { - return fullProfile.kilocodeToken - } - } catch { - continue - } - } - - return undefined -} -// kilocode_change end +import { resolveKiloUserToken } from "./core/kilocode/kilo-user-resolver" // kilocode_change /** * Built using https://github.com/microsoft/vscode-webview-ui-toolkit @@ -318,7 +295,7 @@ export async function activate(context: vscode.ExtensionContext) { // kilocode_change start try { - const kiloToken = await findKilocodeTokenFromAnyProfile(provider) + const kiloToken = await resolveKiloUserToken(provider) await kilo_initializeSessionManager({ context: context, diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index ac8744a1796..6cbb5e1acd8 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -15,6 +15,7 @@ import type { OrganizationAllowList, ShareVisibility, QueuedMessage, + KiloUser, // kilocode_change } from "@roo-code/types" import { GitCommit } from "../utils/git" @@ -527,6 +528,7 @@ export type ExtensionState = Pick< featureRoomoteControlEnabled: boolean virtualQuotaActiveModel?: { id: string; info: ModelInfo } // kilocode_change: Add virtual quota active model for UI display showTimestamps?: boolean // kilocode_change: Show timestamps in chat messages + kiloUser: KiloUser // kilocode_change: Global Kilo user resolved from profiles } export interface ClineSayTool { diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index 3c1f3307405..3cd7f722af7 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -347,6 +347,7 @@ export interface WebviewMessage { organizationId?: string | null // For organization switching useProviderSignup?: boolean // For rooCloudSignIn to use provider signup flow historyItem?: HistoryItem // kilocode_change For addTaskToHistory + profileName?: string // kilocode_change: Optional profile name for profile-specific requests (e.g., fetchProfileDataRequest) codeIndexSettings?: { // Global state settings codebaseIndexEnabled: boolean @@ -410,6 +411,7 @@ export type ProfileData = { image: string } organizations?: UserOrganizationWithApiKey[] + profileName?: string // kilocode_change: Optional profile name for profile-specific requests } export interface ProfileDataResponsePayload { diff --git a/webview-ui/src/App.tsx b/webview-ui/src/App.tsx index 14f9ca13a8d..f4179dd713c 100644 --- a/webview-ui/src/App.tsx +++ b/webview-ui/src/App.tsx @@ -90,7 +90,6 @@ const App = () => { // kilocode_change end renderContext, mdmCompliant, - apiConfiguration, // kilocode_change } = useExtensionState() // kilocode_change start: disable useEffect @@ -155,6 +154,14 @@ const App = () => { setCurrentSection(undefined) setCurrentMarketplaceTab(undefined) + // kilocode_change start - Clear settingsEditingProfile when navigating away from settings + // BUT preserve it when going to auth (so we can return to the same profile after auth) + // This ensures that when returning to settings (without auth context), it shows the active profile + if (tab === "settings" && newTab !== "settings" && newTab !== "auth") { + setSettingsEditingProfile(undefined) + } + // kilocode_change end + // kilocode_change: start - Bypass unsaved changes check when navigating to auth tab if (newTab === "auth") { setTab(newTab) @@ -165,7 +172,7 @@ const App = () => { setTab(newTab) } }, - [mdmCompliant], + [mdmCompliant, tab], ) const [currentSection, setCurrentSection] = useState(undefined) @@ -278,7 +285,7 @@ const App = () => { }, [shouldShowAnnouncement, tab]) // kilocode_change start - const telemetryDistinctId = useKiloIdentity(apiConfiguration?.kilocodeToken ?? "", machineId ?? "") + const telemetryDistinctId = useKiloIdentity(machineId ?? "") useEffect(() => { if (didHydrateState) { telemetryClient.updateTelemetryState(telemetrySetting, telemetryKey, telemetryDistinctId) diff --git a/webview-ui/src/components/kilocode/common/OrganizationSelector.tsx b/webview-ui/src/components/kilocode/common/OrganizationSelector.tsx index 64ef537d156..371fc2339bc 100644 --- a/webview-ui/src/components/kilocode/common/OrganizationSelector.tsx +++ b/webview-ui/src/components/kilocode/common/OrganizationSelector.tsx @@ -1,33 +1,72 @@ -import { useState, useEffect, useRef } from "react" +import { useState, useEffect, useRef, useCallback } from "react" import { vscode } from "@/utils/vscode" import { useExtensionState } from "@/context/ExtensionStateContext" import { useAppTranslation } from "@/i18n/TranslationContext" import { ProfileDataResponsePayload, WebviewMessage, UserOrganizationWithApiKey } from "@roo/WebviewMessage" +import type { ProviderSettings } from "@roo-code/types" -export const OrganizationSelector = ({ className, showLabel = false }: { className?: string; showLabel?: boolean }) => { +type OrganizationSelectorProps = { + className?: string + showLabel?: boolean + + // Controlled mode props (optional) - when provided, component works with specific profile + apiConfiguration?: ProviderSettings + profileName?: string + onChange?: (organizationId: string | undefined) => void +} + +export const OrganizationSelector = ({ + className, + showLabel = false, + apiConfiguration: controlledApiConfiguration, + profileName: controlledProfileName, + onChange, +}: OrganizationSelectorProps) => { const [organizations, setOrganizations] = useState([]) - const { apiConfiguration, currentApiConfigName } = useExtensionState() + const { apiConfiguration: globalApiConfiguration, currentApiConfigName: globalProfileName } = useExtensionState() const { t } = useAppTranslation() const [isOpen, setIsOpen] = useState(false) - const selectedOrg = organizations.find((o) => o.id === apiConfiguration?.kilocodeOrganizationId) const containerRef = useRef(null) - const handleMessage = (event: MessageEvent) => { - const message = event.data - if (message.type === "profileDataResponse") { - const payload = message.payload as ProfileDataResponsePayload - if (payload.success) { - setOrganizations(payload.data?.organizations ?? []) - } else { - console.error("Error fetching profile organizations data:", payload.error) - setOrganizations([]) + // Determine if we're in controlled mode + const isControlledMode = controlledApiConfiguration !== undefined + const apiConfiguration = isControlledMode ? controlledApiConfiguration : globalApiConfiguration + const currentProfileName = isControlledMode ? controlledProfileName : globalProfileName + + const selectedOrg = organizations.find((o) => o.id === apiConfiguration?.kilocodeOrganizationId) + + const handleMessage = useCallback( + (event: MessageEvent) => { + const message = event.data + if (message.type === "profileDataResponse") { + const payload = message.payload as ProfileDataResponsePayload + + // Only process response if it matches the requested profile + // In global mode (no currentProfileName), accept any response + // In controlled mode, only accept responses for the specific profile + const responseProfileName = payload.data?.profileName + const shouldProcessResponse = !currentProfileName || responseProfileName === currentProfileName + + if (!shouldProcessResponse) { + // This response is for a different profile, ignore it + return + } + + if (payload.success) { + setOrganizations(payload.data?.organizations ?? []) + } else { + console.error("Error fetching profile organizations data:", payload.error) + setOrganizations([]) + } + } else if (message.type === "updateProfileData") { + vscode.postMessage({ + type: "fetchProfileDataRequest", + profileName: currentProfileName, + }) } - } else if (message.type === "updateProfileData") { - vscode.postMessage({ - type: "fetchProfileDataRequest", - }) - } - } + }, + [currentProfileName], + ) useEffect(() => { const onKeyDown = (e: KeyboardEvent) => { @@ -50,45 +89,56 @@ export const OrganizationSelector = ({ className, showLabel = false }: { classNa window.removeEventListener("mousedown", onMouseDown) window.removeEventListener("message", handleMessage) } - }, []) + }, [handleMessage]) useEffect(() => { if (!apiConfiguration?.kilocodeToken) return vscode.postMessage({ type: "fetchProfileDataRequest", + profileName: currentProfileName, }) - }, [apiConfiguration?.kilocodeToken]) + }, [apiConfiguration?.kilocodeToken, currentProfileName]) const setSelectedOrganization = (organization: UserOrganizationWithApiKey | null) => { - if (organization === null) { - // Switch back to personal account - clear organization token - vscode.postMessage({ - type: "upsertApiConfiguration", - text: currentApiConfigName, - apiConfiguration: { - ...apiConfiguration, - kilocodeOrganizationId: undefined, - }, - }) - vscode.postMessage({ - type: "fetchBalanceDataRequest", - values: { - apiKey: apiConfiguration?.kilocodeToken, - }, - }) + const newOrganizationId = organization?.id + + if (isControlledMode) { + // Controlled mode: call the onChange callback + if (onChange) { + onChange(newOrganizationId) + } } else { - vscode.postMessage({ - type: "upsertApiConfiguration", - text: currentApiConfigName, - apiConfiguration: { - ...apiConfiguration, - kilocodeOrganizationId: organization.id, - }, - }) - vscode.postMessage({ - type: "fetchBalanceDataRequest", - }) + // Global mode: update the global configuration + if (organization === null) { + // Switch back to personal account - clear organization token + vscode.postMessage({ + type: "upsertApiConfiguration", + text: currentProfileName, + apiConfiguration: { + ...apiConfiguration, + kilocodeOrganizationId: undefined, + }, + }) + vscode.postMessage({ + type: "fetchBalanceDataRequest", + values: { + apiKey: apiConfiguration?.kilocodeToken, + }, + }) + } else { + vscode.postMessage({ + type: "upsertApiConfiguration", + text: currentProfileName, + apiConfiguration: { + ...apiConfiguration, + kilocodeOrganizationId: organization.id, + }, + }) + vscode.postMessage({ + type: "fetchBalanceDataRequest", + }) + } } } diff --git a/webview-ui/src/components/kilocode/settings/providers/KiloCode.tsx b/webview-ui/src/components/kilocode/settings/providers/KiloCode.tsx index e8f444df83a..aad08503b72 100644 --- a/webview-ui/src/components/kilocode/settings/providers/KiloCode.tsx +++ b/webview-ui/src/components/kilocode/settings/providers/KiloCode.tsx @@ -43,8 +43,8 @@ export const KiloCode = ({ [setApiConfigurationField], ) - // Use the existing hook to get user identity - const userIdentity = useKiloIdentity(apiConfiguration.kilocodeToken || "", "") + // Use the existing hook to get user identity from global kilo user + const userIdentity = useKiloIdentity("") const isKiloCodeAiUser = userIdentity.endsWith("@kilo.ai") const areKilocodeWarningsDisabled = apiConfiguration.kilocodeTesterWarningsDisabledUntil @@ -107,7 +107,14 @@ export const KiloCode = ({ - + { + setApiConfigurationField("kilocodeOrganizationId", organizationId) + }} + /> { featureRoomoteControlEnabled: false, isBrowserSessionActive: false, checkpointTimeout: DEFAULT_CHECKPOINT_TIMEOUT_SECONDS, // Add the checkpoint timeout property + kiloUser: EMPTY_KILO_USER, // kilocode_change: Add kiloUser to mock state } const prevState: ExtensionState = { diff --git a/webview-ui/src/utils/kilocode/useKiloIdentity.tsx b/webview-ui/src/utils/kilocode/useKiloIdentity.tsx index 26595f7eec5..6705bee93db 100644 --- a/webview-ui/src/utils/kilocode/useKiloIdentity.tsx +++ b/webview-ui/src/utils/kilocode/useKiloIdentity.tsx @@ -1,44 +1,18 @@ -import { useEffect, useState } from "react" -import { ProfileDataResponsePayload } from "@roo/WebviewMessage" -import { vscode } from "@/utils/vscode" +import { useExtensionState } from "@/context/ExtensionStateContext" -export function useKiloIdentity(kilocodeToken: string, machineId: string) { - const [kiloIdentity, setKiloIdentity] = useState("") - useEffect(() => { - const handleMessage = (event: MessageEvent) => { - if (event.data.type === "profileDataResponse") { - const payload = event.data.payload as ProfileDataResponsePayload | undefined - const success = payload?.success || false - const tokenFromMessage = payload?.data?.kilocodeToken || "" - const email = payload?.data?.user?.email || "" - if (!success) { - console.error("KILOTEL: Failed to identify Kilo user, message doesn't indicate success:", payload) - } else if (tokenFromMessage !== kilocodeToken) { - console.error("KILOTEL: Failed to identify Kilo user, token mismatch:", payload) - } else if (!email) { - console.error("KILOTEL: Failed to identify Kilo user, email missing:", payload) - } else { - console.debug("KILOTEL: Kilo user identified:", email) - setKiloIdentity(email) - window.removeEventListener("message", handleMessage) - } - } - } +/** + * Hook to get the Kilo user identity for telemetry. + * Returns the user's email if authenticated with Kilocode, otherwise returns the machine ID. + * + * The Kilo user is resolved globally based on priority rules: + * 1. Active profile if it's a kilocode provider + * 2. First kilocode provider found in profiles list + * 3. Fallback to unauthenticated (uses machineId) + */ +export function useKiloIdentity(machineId: string) { + const { kiloUser } = useExtensionState() - if (kilocodeToken) { - console.debug("KILOTEL: fetching profile...") - window.addEventListener("message", handleMessage) - vscode.postMessage({ - type: "fetchProfileDataRequest", - }) - } else { - console.debug("KILOTEL: no Kilo user") - setKiloIdentity("") - } - - return () => { - window.removeEventListener("message", handleMessage) - } - }, [kilocodeToken]) - return kiloIdentity || machineId + // If we have a global kilo user with email, use that for telemetry identity + // Otherwise fall back to machineId + return kiloUser?.email || machineId }