changeset

.changeset/little-toys-protect.md

@@ -0,0 +1,5 @@
+---
+"claude-dev": minor
+---
+
+Migrate authCallback to protobus

proto/account.proto

@@ -16,4 +16,7 @@ service AccountService {
     // Handles the user clicking the logout button in the UI.
     // Clears API keys and user state.
     rpc accountLogoutClicked(EmptyRequest) returns (Empty);
+
+    // Subscribe to auth callback events (when authentication tokens are received)
+    rpc subscribeToAuthCallback(EmptyRequest) returns (stream String);
 }

src/core/controller/account/subscribeToAuthCallback.ts

@@ -0,0 +1,59 @@
+import { Controller } from "../index"
+import { EmptyRequest } from "../../../shared/proto/common"
+import { String as ProtoString } from "../../../shared/proto/common"
+import { StreamingResponseHandler, getRequestRegistry } from "../grpc-handler"
+
+// Keep track of active authCallback subscriptions
+const activeAuthCallbackSubscriptions = new Set<StreamingResponseHandler>()
+
+/**
+ * Subscribe to authCallback events
+ * @param controller The controller instance
+ * @param request The empty request
+ * @param responseStream The streaming response handler
+ * @param requestId The ID of the request (passed by the gRPC handler)
+ */
+export async function subscribeToAuthCallback(
+	controller: Controller,
+	request: EmptyRequest,
+	responseStream: StreamingResponseHandler,
+	requestId?: string,
+): Promise<void> {
+	// Add this subscription to the active subscriptions
+	activeAuthCallbackSubscriptions.add(responseStream)
+
+	// Register cleanup when the connection is closed
+	const cleanup = () => {
+		activeAuthCallbackSubscriptions.delete(responseStream)
+	}
+
+	// Register the cleanup function with the request registry if we have a requestId
+	if (requestId) {
+		getRequestRegistry().registerRequest(requestId, cleanup, { type: "authCallback_subscription" }, responseStream)
+	}
+}
+
+/**
+ * Send an authCallback event to all active subscribers
+ * @param customToken The custom token for authentication
+ */
+export async function sendAuthCallbackEvent(customToken: string): Promise<void> {
+	// Send the event to all active subscribers
+	const promises = Array.from(activeAuthCallbackSubscriptions).map(async (responseStream) => {
+		try {
+			const event: ProtoString = {
+				value: customToken,
+			}
+			await responseStream(
+				event,
+				false, // Not the last message
+			)
+		} catch (error) {
+			console.error("Error sending authCallback event:", error)
+			// Remove the subscription if there was an error
+			activeAuthCallbackSubscriptions.delete(responseStream)
+		}
+	})
+
+	await Promise.all(promises)
+}

src/core/controller/index.ts

@@ -52,6 +52,7 @@ import {
 import { Task, cwd } from "../task"
 import { ClineRulesToggles } from "@shared/cline-rules"
 import { sendStateUpdate } from "./state/subscribeToState"
+import { sendAuthCallbackEvent } from "./account/subscribeToAuthCallback"
 import { refreshClineRulesToggles } from "@core/context/instructions/user-instructions/cline-rules"
 import { refreshExternalRulesToggles } from "@core/context/instructions/user-instructions/external-rules"
 import { refreshWorkflowToggles } from "@core/context/instructions/user-instructions/workflows"
@@ -697,10 +698,7 @@ export class Controller {
 			await storeSecret(this.context, "clineApiKey", apiKey)
 
 			// Send custom token to webview for Firebase auth
-			await this.postMessageToWebview({
-				type: "authCallback",
-				customToken,
-			})
+			await sendAuthCallbackEvent(customToken)
 
 			const clineProvider: ApiProvider = "cline"
 			await updateGlobalState(this.context, "apiProvider", clineProvider)

src/shared/ExtensionMessage.ts

@@ -28,7 +28,6 @@ export interface ExtensionMessage {
 		| "requestyModels"
 		| "mcpServers"
 		| "relinquishControl"
-		| "authCallback"
 		| "mcpMarketplaceCatalog"
 		| "mcpDownloadDetails"
 		| "commitSearchResults"

src/shared/WebviewMessage.ts

@@ -16,7 +16,6 @@ export interface WebviewMessage {
 		| "showChatView"
 		| "requestVsCodeLmModels"
 		| "authStateChanged"
-		| "authCallback"
 		| "fetchMcpMarketplace"
 		| "searchCommits"
 		| "fetchLatestMcpServersFromHub"

webview-ui/src/context/FirebaseAuthContext.tsx

@@ -2,6 +2,7 @@ import { User, getAuth, signInWithCustomToken, signOut } from "firebase/auth"
 import { initializeApp } from "firebase/app"
 import React, { createContext, useCallback, useContext, useEffect, useState } from "react"
 import { vscode } from "@/utils/vscode"
+import { AccountServiceClient } from "@/services/grpc-client"
 
 // Firebase configuration from extension
 const firebaseConfig = {
@@ -37,8 +38,6 @@ export const FirebaseAuthProvider: React.FC<{ children: React.ReactNode }> = ({
 			setUser(user)
 			setIsInitialized(true)
 
-			console.log("onAuthStateChanged user", user)
-
 			if (!user) {
 				// when opening the extension in a new webview (ie if you logged in to sidebar webview but then open a popout tab webview) this effect will trigger without the original webview's session, resulting in us clearing out the user info object.
 				// we rely on this object to determine if the user is logged in, so we only want to clear it when the user logs out, rather than whenever a webview without a session is opened.
@@ -73,17 +72,24 @@ export const FirebaseAuthProvider: React.FC<{ children: React.ReactNode }> = ({
 		[auth],
 	)
 
-	// Listen for auth callback from extension
+	// Set up authCallback subscription
 	useEffect(() => {
-		const handleMessage = (event: MessageEvent) => {
-			const message = event.data
-			if (message.type === "authCallback" && message.customToken) {
-				signInWithToken(message.customToken)
-			}
-		}
+		const cleanup = AccountServiceClient.subscribeToAuthCallback(
+			{},
+			{
+				onResponse: (event) => {
+					if (event.value) {
+						signInWithToken(event.value)
+					}
+				},
+				onError: (error) => {
+					console.error("Error in authCallback subscription:", error)
+				},
+				onComplete: () => {},
+			},
+		)
 
-		window.addEventListener("message", handleMessage)
-		return () => window.removeEventListener("message", handleMessage)
+		return cleanup
 	}, [signInWithToken])
 
 	const handleSignOut = useCallback(async () => {