fix: user doesn't open webview but trigger completion (#127)

Enhances user authentication by adding logic to manage authentication callbacks without opening the webview when a valid token is present.

Author: xiaojingming

src/activate/handleUri.ts

@@ -2,13 +2,35 @@ import * as vscode from "vscode"
 
 import { ClineProvider } from "../core/webview/ClineProvider"
 
+const handleAuthWithoutVisibleProvider = async (provider: ClineProvider | undefined, params: URLSearchParams) => {
+	if (!provider) {
+		return
+	}
+
+	const code = params.get("code")
+	const state = params.get("state")
+	const token = params.get("token")
+
+	if ((code && state) || token) {
+		await provider?.handleZgsmAuthCallback(code, state, token, false)
+	}
+}
+
 export const handleUri = async (uri: vscode.Uri) => {
 	const path = uri.path
 	const query = new URLSearchParams(uri.query.replace(/\+/g, "%2B"))
 	const visibleProvider = ClineProvider.getVisibleInstance()
+
+	// not open webview to re-login
 	if (!visibleProvider) {
+		if (path === "/callback") {
+			const cacheProvider = ClineProvider.getCacheInstances()
+			handleAuthWithoutVisibleProvider(cacheProvider, query)
+		}
+
 		return
 	}
+
 	switch (path) {
 		case "/glama": {
 			const code = query.get("code")

src/core/config/ContextProxy.ts

@@ -17,6 +17,7 @@ import {
 } from "../../schemas"
 import { logger } from "../../utils/logging"
 import { telemetryService } from "../../services/telemetry/TelemetryService"
+import { ClineProvider } from "../webview/ClineProvider"
 
 type GlobalStateKey = keyof GlobalState
 type SecretStateKey = keyof SecretState
@@ -256,4 +257,14 @@ export class ContextProxy {
 
 		await this.initialize()
 	}
+
+	// get state from originContext.globalState
+	public async getOriginGlobalState(key: string) {
+		return await this?.originalContext?.globalState.get(key)
+	}
+
+	// get state from originContext.secrets
+	public async getOriginSecrets(key: string) {
+		return await this?.originalContext?.secrets.get(key)
+	}
 }

src/core/webview/ClineProvider.ts

@@ -58,6 +58,7 @@ import { WebviewMessage } from "../../shared/WebviewMessage"
 import { afterZgsmPostLogin, getZgsmAccessToken } from "../../zgsmAuth/zgsmAuthHandler"
 import { defaultZgsmAuthConfig } from "../../zgsmAuth/config"
 import { CompletionStatusBar } from "../../../zgsm/src/codeCompletion/completionStatusBar"
+import { CompletionClient } from "../../../zgsm/src/codeCompletion/completionClient"
 
 /**
  * https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default/weather-webview/src/providers/WeatherViewProvider.ts
@@ -1015,11 +1016,11 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 
 	// Zgsm
 
-	async handleZgsmAuthCallback(code: string | null, state: string | null, token: string | null) {
+	async handleZgsmAuthCallback(code: string | null, state: string | null, token: string | null, needVisible = true) {
 		let { apiConfiguration, currentApiConfigName } = await this.getState()
 		const visibleProvider = await ClineProvider.getInstance()
 
-		if (!visibleProvider) {
+		if (!visibleProvider && needVisible) {
 			return
 		}
 
@@ -1522,6 +1523,11 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 		return getWorkspacePath()
 	}
 
+	// has view
+	get hasView() {
+		return !!this.view
+	}
+
 	// dev
 
 	async resetState() {

zgsm/src/codeCompletion/completionClient.ts

@@ -13,7 +13,15 @@ import { Logger } from "../common/log-util"
 import { workspace } from "vscode"
 import { AxiosError } from "axios"
 import { createAuthenticatedHeaders } from "../common/api"
-import { configCompletion, settings, OPENAI_CLIENT_NOT_INITIALIZED } from "../common/constant"
+import {
+	configCompletion,
+	settings,
+	OPENAI_CLIENT_NOT_INITIALIZED,
+	NOT_PROVIDERED,
+	ZGSM_API_KEY,
+	ZGSM_BASE_URL,
+	ZGSM_COMPLETION_URL,
+} from "../common/constant"
 import { CompletionPoint } from "./completionPoint"
 import { CompletionScores } from "./completionScore"
 import { CompletionTrace } from "./completionTrace"
@@ -32,6 +40,28 @@ export class CompletionClient {
 	private reqs: Map<string, any> = new Map<string, any>()
 	private betaMode?: any
 
+	private async getApiConfig(hasView: boolean, apiConfiguration: any) {
+		if (hasView) {
+			return {
+				baseUrl: apiConfiguration.zgsmBaseUrl || defaultZgsmAuthConfig.baseUrl,
+				completionUrl: apiConfiguration.zgsmCompletionUrl || defaultZgsmAuthConfig.completionUrl,
+				apiKey: apiConfiguration.zgsmApiKey || NOT_PROVIDERED,
+			}
+		}
+
+		const context = CompletionClient.providerRef.deref()?.contextProxy
+
+		const contextApiKey = await context?.getOriginSecrets(ZGSM_API_KEY)
+		const contextBaseUrl = await context?.getGlobalState(ZGSM_BASE_URL)
+		const contextCompletionUrl = await context?.getGlobalState(ZGSM_COMPLETION_URL)
+
+		return {
+			baseUrl: contextBaseUrl || defaultZgsmAuthConfig.baseUrl,
+			completionUrl: contextCompletionUrl || defaultZgsmAuthConfig.completionUrl,
+			apiKey: contextApiKey || NOT_PROVIDERED,
+		}
+	}
+
 	public static async setProvider(provider: ClineProvider) {
 		CompletionClient.providerRef = new WeakRef(provider)
 
@@ -119,14 +149,24 @@ export class CompletionClient {
 		const provider = CompletionClient.providerRef.deref()
 
 		const { apiConfiguration } = await provider!.getState()
-		if (!apiConfiguration?.zgsmApiKey) {
+
+		const hasView = !!provider?.hasView
+
+		if (!apiConfiguration?.zgsmApiKey && hasView) {
 			Logger.error("Failed to get login information. Please log in again to use the completion service")
 			return false
 		}
-		const completionUrl = `${apiConfiguration.zgsmBaseUrl || defaultZgsmAuthConfig.baseUrl}${apiConfiguration.zgsmCompletionUrl || defaultZgsmAuthConfig.completionUrl}`
+
+		const config = await this.getApiConfig(hasView, apiConfiguration)
+		const fullUrl = `${config.baseUrl}${config.completionUrl}`
+
+		if (config.apiKey === NOT_PROVIDERED) {
+			return false
+		}
+
 		this.openai = new OpenAI({
-			baseURL: completionUrl,
-			apiKey: apiConfiguration.zgsmApiKey || "not-provided",
+			baseURL: fullUrl,
+			apiKey: config.apiKey,
 		})
 		if (!this.openai) {
 			// Logger.error("Completion: Configuration error: configuration:", configuration, "openai: ", this.openai);
@@ -136,7 +176,7 @@ export class CompletionClient {
 		this.stopWords = workspace.getConfiguration(configCompletion).get("inlineCompletion") ? ["\n", "\r"] : []
 		this.betaMode = workspace.getConfiguration(configCompletion).get("betaMode")
 		Logger.info(
-			`Completion: Create OpenAIApi client, URL: ${completionUrl}, betaMode: ${this.betaMode}, stopWords: ${this.stopWords}`,
+			`Completion: Create OpenAIApi client, URL: ${fullUrl}, betaMode: ${this.betaMode}, stopWords: ${this.stopWords}`,
 		)
 		return true
 	}
@@ -209,8 +249,12 @@ export class CompletionClient {
 		Logger.log(`Completion [${cp.id}]: Sending API request`)
 		const headers = createAuthenticatedHeaders()
 		const repo = workspace?.name?.split(" ")[0] ?? ""
-		this.openai.baseURL = `${apiConfiguration.zgsmBaseUrl || defaultZgsmAuthConfig.baseUrl}${apiConfiguration.zgsmCompletionUrl || defaultZgsmAuthConfig.completionUrl}`
-		this.openai.apiKey = apiConfiguration.zgsmApiKey || "not-provided"
+
+		const config = await this.getApiConfig(!!provider?.hasView, apiConfiguration)
+
+		this.openai.baseURL = `${config.baseUrl}${config.completionUrl}`
+		this.openai.apiKey = config.apiKey
+
 		return this.openai.completions.create(
 			{
 				// no use

zgsm/src/codeCompletion/completionCommands.ts

@@ -6,53 +6,53 @@
  * The above copyright notice and this permission notice shall be included in all
  * copies or substantial portions of the Software.
  */
-import * as vscode from "vscode";
-import { Logger } from "../common/log-util";
-import { setupLangSwitchs } from "../common/services";
+import * as vscode from "vscode"
+import { Logger } from "../common/log-util"
+import { handleStatusBarClick } from "../common/services"
 
-export type Command = { command: string, callback: (...args: any[]) => any, thisArg?: any; };
+export type Command = { command: string; callback: (...args: any[]) => any; thisArg?: any }
 
 /**
  * Add a command to the completion status bar item.
  */
 export const statusBarCommand: Command = {
-    command: "zgsm-statusBar.showInformationMessage",
-    callback: () => setupLangSwitchs()
-};
+	command: "zgsm-statusBar.showInformationMessage",
+	callback: () => handleStatusBarClick(),
+}
 
 /**
  * Set the enabled status of the completion extension.
  */
 function setExtensionStatus(enabled: boolean) {
-    const config = vscode.workspace.getConfiguration();
-    const target = vscode.ConfigurationTarget.Global;
-    Logger.info("Set the status of 'zgsm-completion.enabled' to", enabled);
-    config.update('zgsm-completion.enabled', enabled, target, false).then(console.error);
+	const config = vscode.workspace.getConfiguration()
+	const target = vscode.ConfigurationTarget.Global
+	Logger.info("Set the status of 'zgsm-completion.enabled' to", enabled)
+	config.update("zgsm-completion.enabled", enabled, target, false).then(console.error)
 }
 
 /**
  * Command to enable completion.
  */
 export const turnOnCompletion: Command = {
-    command: "zgsm-completion.enable",
-    callback: () => setExtensionStatus(true)
-};
+	command: "zgsm-completion.enable",
+	callback: () => setExtensionStatus(true),
+}
 
 /**
  * Command to disable completion.
  */
 export const turnOffCompletion: Command = {
-    command: "zgsm-completion.disable",
-    callback: () => setExtensionStatus(false)
-};
+	command: "zgsm-completion.disable",
+	callback: () => setExtensionStatus(false),
+}
 
 /**
  * Command for completion shortcut instructions.
  */
 export const shortKeyCut: Command = {
-    command: "zgsm-completion.shortKeyCut",
-    callback: (context: vscode.ExtensionContext) => {
-        context.workspaceState.update('shortCutKeys', true);
-        vscode.commands.executeCommand('editor.action.inlineSuggest.trigger');
-    }
-};
+	command: "zgsm-completion.shortKeyCut",
+	callback: (context: vscode.ExtensionContext) => {
+		context.workspaceState.update("shortCutKeys", true)
+		vscode.commands.executeCommand("editor.action.inlineSuggest.trigger")
+	},
+}

zgsm/src/common/constant.ts

@@ -204,3 +204,9 @@ export const ACCESS_TOKEN_KEY = `${AUTH_TYPE}.accessToken`
 // OpenAI Client
 export const OPENAI_CLIENT_NOT_INITIALIZED = "OpenAI client not initialized"
 export const OPENAI_REQUEST_ABORTED = "Request was aborted"
+
+export const NOT_PROVIDERED = "not-provided"
+
+export const ZGSM_API_KEY = "zgsmApiKey"
+export const ZGSM_BASE_URL = "zgsmBaseUrl"
+export const ZGSM_COMPLETION_URL = "zgsmCompletionUrl"