Fix OAuth login hangs after browser callback

cmd/login.go

@@ -64,7 +64,7 @@ func runLogin(cmd *cobra.Command, args []string) error {
 		// Handle common error cases with helpful messages
 		if ctx.Err() == context.Canceled {
 			pterm.Info.Println("Authentication cancelled by user")
-			return nil
+			return ctx.Err()
 		}
 
 		return fmt.Errorf("authentication failed: %w", err)

pkg/auth/oauth.go

@@ -22,14 +22,14 @@ import (
 )
 
 //go:embed success.html
-var successHTMLTemplate string
+var authorizationReceivedHTMLTemplate string
 
 //go:embed favicon.svg
 var faviconSVG string
 
 // Inlined as a data URI because the callback server shuts down before the browser could fetch a served icon.
-var successHTML = strings.Replace(
-	successHTMLTemplate,
+var authorizationReceivedHTML = strings.Replace(
+	authorizationReceivedHTMLTemplate,
 	"__KERNEL_FAVICON__",
 	"data:image/svg+xml;base64,"+base64.StdEncoding.EncodeToString([]byte(faviconSVG)),
 	1,
@@ -45,6 +45,8 @@ const (
 
 	// OAuth scopes - openid for the MCP server flow
 	DefaultScope = "openid email"
+
+	defaultTokenExchangeTimeout = 30 * time.Second
 )
 
 // OAuthConfig represents the OAuth2 configuration
@@ -250,7 +252,7 @@ func (oc *OAuthConfig) StartOAuthFlow(ctx context.Context) (*TokenStorage, error
 		// Success page
 		w.Header().Set("Content-Type", "text/html")
 		w.WriteHeader(http.StatusOK)
-		w.Write([]byte(successHTML))
+		w.Write([]byte(authorizationReceivedHTML))
 
 		// Pass both code and org_id to the channel using JSON encoding
 		result := AuthResult{
@@ -298,21 +300,29 @@ func (oc *OAuthConfig) StartOAuthFlow(ctx context.Context) (*TokenStorage, error
 		return nil, ctx.Err()
 	}
 
+	pterm.Info.Println("Browser authorization received; completing authentication...")
+
 	// Exchange authorization code for tokens
-	return oc.exchangeCodeForTokens(ctx, authCode, orgID)
+	return oc.exchangeCodeForTokens(ctx, authCode, orgID, defaultTokenExchangeTimeout)
 }
 
 // exchangeCodeForTokens exchanges the authorization code for access and refresh tokens
-func (oc *OAuthConfig) exchangeCodeForTokens(ctx context.Context, code, orgID string) (*TokenStorage, error) {
+func (oc *OAuthConfig) exchangeCodeForTokens(ctx context.Context, code, orgID string, timeout time.Duration) (*TokenStorage, error) {
 	// Use PKCE verifier in token exchange, and include org_id if available
 	var opts []oauth2.AuthCodeOption
 	opts = append(opts, oauth2.SetAuthURLParam("code_verifier", oc.Verifier))
 	if orgID != "" {
 		opts = append(opts, oauth2.SetAuthURLParam("org_id", orgID))
 	}
 
-	token, err := oc.Config.Exchange(ctx, code, opts...)
+	exchangeCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	token, err := oc.Config.Exchange(exchangeCtx, code, opts...)
 	if err != nil {
+		if exchangeCtx.Err() == context.DeadlineExceeded {
+			return nil, fmt.Errorf("token exchange timed out after %s: %w", timeout, context.DeadlineExceeded)
+		}
 		return nil, fmt.Errorf("failed to exchange code for token: %w", err)
 	}
 

pkg/auth/oauth_test.go

@@ -1,6 +1,58 @@
 package auth
 
-import "testing"
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestCallbackPageDoesNotClaimAuthenticationComplete(t *testing.T) {
+	t.Parallel()
+
+	if strings.Contains(authorizationReceivedHTML, "authentication successful") {
+		t.Fatal("callback page must not claim authentication succeeded before token exchange completes")
+	}
+	if !strings.Contains(authorizationReceivedHTML, "authorization received") {
+		t.Fatal("callback page should tell the user browser authorization was received")
+	}
+}
+
+func TestTokenExchangeTimesOut(t *testing.T) {
+	release := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case <-r.Context().Done():
+		case <-release:
+		}
+	}))
+	defer server.Close()
+	defer close(release)
+
+	cfg, err := NewOAuthConfig()
+	if err != nil {
+		t.Fatalf("NewOAuthConfig() error = %v", err)
+	}
+	cfg.Config.Endpoint.TokenURL = server.URL
+
+	done := make(chan error, 1)
+	go func() {
+		_, err := cfg.exchangeCodeForTokens(context.Background(), "code", "org", 25*time.Millisecond)
+		done <- err
+	}()
+
+	select {
+	case err := <-done:
+		if err == nil || !strings.Contains(err.Error(), "token exchange timed out") || !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("exchange error = %v, want token exchange timeout", err)
+		}
+	case <-time.After(250 * time.Millisecond):
+		t.Fatal("token exchange did not honor its timeout")
+	}
+}
 
 func TestNewOAuthConfigUsesAuthOverrides(t *testing.T) {
 	t.Setenv("KERNEL_AUTH_BASE_URL", "https://auth.dev.onkernel.com/")

pkg/auth/success.html

@@ -1,7 +1,7 @@
 <!DOCTYPE html>
 <html>
 <head>
-    <title>kernel cli - authenticated</title>
+    <title>kernel cli - authorization received</title>
     <link rel="icon" type="image/svg+xml" href="__KERNEL_FAVICON__">
     <link href="https://fonts.googleapis.com/css2?family=Inter:wght@100..900&display=swap" rel="stylesheet">
     <style>
@@ -88,8 +88,8 @@
                     <path d="M285.036 0C289.434 0 293 3.56173 293 7.95537C293 12.3491 289.434 15.9107 285.036 15.9107C280.638 15.9107 277.072 12.349 277.072 7.95537C277.072 3.56178 280.638 8.45042e-05 285.036 0Z" fill="#1c2024"/>
                 </svg>
             </div>
-            <p class="title">authentication successful</p>
-            <p class="subtitle">you can close this window and return to your terminal.</p>
+            <p class="title">authorization received</p>
+            <p class="subtitle">return to your terminal while the cli completes authentication.</p>
             <div class="check-container">
                 <svg fill="none" stroke="currentColor" viewBox="0 0 24 24">
                     <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M5 13l4 4L19 7"></path>