Adapt login to WorkOS Connect: issuer config, sid-aware revocation, branded callback page

internal/cli/env.go

@@ -22,6 +22,8 @@ type envVarSpec struct {
 const (
 	envAPIKey        = "EXTEND_API_KEY"
 	envBaseURL       = "EXTEND_BASE_URL"
+	envOAuthIssuer   = "EXTEND_OAUTH_ISSUER"
+	envOAuthClientID = "EXTEND_OAUTH_CLIENT_ID"
 	envRegion        = "EXTEND_REGION"
 	envWorkspaceID   = "EXTEND_WORKSPACE_ID"
 	envAPIVersion    = "EXTEND_API_VERSION"
@@ -53,6 +55,8 @@ const defaultAPIVersion = "2026-02-09"
 var envVars = []envVarSpec{
 	{Name: envAPIKey, Required: true, Description: "API key (sk_...). Required for API commands unless signed in via 'extend login'."},
 	{Name: envBaseURL, Description: "Override base URL. Wins over EXTEND_REGION."},
+	{Name: envOAuthIssuer, Description: "Authorization server (WorkOS AuthKit domain) for 'extend login', e.g. https://id.extend.ai."},
+	{Name: envOAuthClientID, Description: "OAuth client id for 'extend login'. Defaults to the built-in first-party client."},
 	{Name: envRegion, Description: "Region: us|eu. Selects the regional API endpoint."},
 	{Name: envWorkspaceID, Description: "Workspace ID for org-scoped API keys (sent as X-Extend-Workspace-Id)."},
 	{Name: envAPIVersion, Description: "Pin the API version sent with each request. Defaults to " + defaultAPIVersion + "."},

internal/cli/login.go

@@ -189,11 +189,22 @@ func runLogin(ctx context.Context, app *App, opts loginOptions) error {
 	}
 	pal := paletteFor(app.IO)
 
-	// The redirect policy pins every OAuth call (and /me) to the API
-	// base's origin so request bodies carrying the code, verifier, or
-	// tokens can never be replayed to another host.
+	// The authorization server is the WorkOS AuthKit issuer, not the API:
+	// discovery, the browser authorize URL, and the token endpoint all
+	// live on the issuer domain, while the RFC 8707 resource parameter
+	// and every bearer-authenticated API call point at the API base.
+	issuer, clientID, err := resolveAuthServer(os.Getenv)
+	if err != nil {
+		return err
+	}
+
+	// Two pinned clients, one per origin: OAuth calls (discovery, code
+	// exchange, refresh) may only talk to the issuer, API calls (/me)
+	// only to the API base, so request bodies carrying the code,
+	// verifier, or tokens can never be replayed to another host.
+	authHTTP := oauth.NewHTTPClient(issuer)
 	httpc := oauth.NewHTTPClient(base)
-	eps, err := oauth.Discover(ctx, httpc, base)
+	eps, err := oauth.Discover(ctx, authHTTP, issuer)
 	if err != nil {
 		return err
 	}
@@ -214,7 +225,7 @@ func runLogin(ctx context.Context, app *App, opts loginOptions) error {
 	defer lb.Close()
 
 	authURL := oauth.AuthorizeURL(eps.Authorization, oauth.AuthorizeParams{
-		ClientID:    oauth.DefaultClientID,
+		ClientID:    clientID,
 		RedirectURI: lb.RedirectURI(),
 		State:       state,
 		Challenge:   oauth.ChallengeS256(verifier),
@@ -244,9 +255,9 @@ func runLogin(ctx context.Context, app *App, opts loginOptions) error {
 
 	sp.Update("Completing sign-in...")
 	tokenClient := &oauth.Client{
-		HTTPClient: httpc,
+		HTTPClient: authHTTP,
 		Endpoints:  eps,
-		ClientID:   oauth.DefaultClientID,
+		ClientID:   clientID,
 		Resource:   base,
 	}
 	tr, err := tokenClient.Exchange(ctx, code, verifier, lb.RedirectURI())
@@ -261,21 +272,30 @@ func runLogin(ctx context.Context, app *App, opts loginOptions) error {
 	prev, _ := opts.store.Get(base)
 
 	rec := oauth.Record{
-		AccessToken:        tr.AccessToken,
-		RefreshToken:       tr.RefreshToken,
-		ExpiresAt:          tr.Expiry(time.Now()),
-		TokenEndpoint:      eps.Token,
-		RevocationEndpoint: eps.Revocation,
-		ClientID:           oauth.DefaultClientID,
-		Resource:           base,
+		AccessToken:   tr.AccessToken,
+		RefreshToken:  tr.RefreshToken,
+		ExpiresAt:     tr.Expiry(time.Now()),
+		TokenEndpoint: eps.Token,
+		// No RevocationEndpoint: WorkOS Connect has no RFC 7009
+		// endpoint; logout revokes via the API's /oauth/revoke-current.
+		ClientID: clientID,
+		Resource: base,
 	}
 	if err := opts.store.Set(base, rec); err != nil {
 		sp.Stop("")
 		return fmt.Errorf("store login: %w", err)
 	}
 	if prev != nil && prev.RefreshToken != rec.RefreshToken {
-		// Best-effort: the new login already works either way.
-		_ = revokeRecord(ctx, base, prev)
+		// Revocation is sid-keyed and the sid is the WorkOS consent id:
+		// while the consent stays on file, a re-login issues tokens with
+		// the SAME sid, so revoking the replaced grant would instantly
+		// invalidate the new one too. Only revoke when the sids differ;
+		// a superseded same-sid token dies with the consent at logout.
+		// Best-effort either way: the new login already works.
+		prevSID := oauth.TokenSID(prev.AccessToken)
+		if prevSID == "" || prevSID != oauth.TokenSID(rec.AccessToken) {
+			_ = revokeRecord(ctx, base, prev)
+		}
 	}
 
 	// Personalize the success line from GET /me. Best-effort: the
@@ -323,27 +343,68 @@ func runLogout(ctx context.Context, app *App, store oauth.Store) error {
 	return nil
 }
 
-// revokeRecord revokes a stored login's refresh token server-side,
-// killing its whole grant family per RFC 7009.
+// revokeRecord kills a stored login server-side. WorkOS Connect has no
+// RFC 7009 revocation endpoint; instead POST /oauth/revoke-current on
+// the API denylists the session id inside the presented access token,
+// which invalidates every token of the login instantly (including any
+// still mintable via the refresh token). An expired access token is
+// refreshed first so the revoke call can authenticate.
 func revokeRecord(ctx context.Context, base string, rec *oauth.Record) error {
-	if rec == nil || rec.RefreshToken == "" {
+	if rec == nil || rec.AccessToken == "" {
 		return nil
 	}
-	eps := oauth.DefaultEndpoints(base)
-	if rec.RevocationEndpoint != "" {
-		eps.Revocation = rec.RevocationEndpoint
+	token := rec.AccessToken
+	stale := !rec.ExpiresAt.IsZero() && time.Now().After(rec.ExpiresAt.Add(-30*time.Second))
+	if stale && rec.RefreshToken != "" && rec.TokenEndpoint != "" {
+		clientID := rec.ClientID
+		if clientID == "" {
+			clientID = oauth.DefaultClientID
+		}
+		c := &oauth.Client{
+			HTTPClient: oauth.NewHTTPClient(rec.TokenEndpoint),
+			Endpoints:  oauth.Endpoints{Token: rec.TokenEndpoint},
+			ClientID:   clientID,
+			Resource:   base,
+		}
+		if tr, err := c.Refresh(ctx, rec.RefreshToken); err == nil {
+			token = tr.AccessToken
+		}
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
+		oauth.NormalizeBase(base)+"/oauth/revoke-current", nil)
+	if err != nil {
+		return err
 	}
-	clientID := rec.ClientID
+	req.Header.Set("Authorization", "Bearer "+token)
+	req.Header.Set("User-Agent", userAgent())
+	resp, err := oauth.NewHTTPClient(base).Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return fmt.Errorf("POST /oauth/revoke-current: http %d", resp.StatusCode)
+	}
+	return nil
+}
+
+// resolveAuthServer returns the authorization-server issuer (the WorkOS
+// AuthKit domain) and OAuth client id for 'extend login'. The issuer has
+// no default: it differs per Extend environment and a login attempt
+// without it cannot succeed, so fail with instructions instead.
+func resolveAuthServer(getenv func(string) string) (issuer, clientID string, err error) {
+	issuer = oauth.NormalizeBase(getenv(envOAuthIssuer))
+	if issuer == "" {
+		return "", "", fmt.Errorf("%s is not set; set it to this environment's sign-in domain (e.g. https://id.extend.ai) to use 'extend login'", envOAuthIssuer)
+	}
+	if err := oauth.ValidateBaseURL(issuer); err != nil {
+		return "", "", err
+	}
+	clientID = getenv(envOAuthClientID)
 	if clientID == "" {
 		clientID = oauth.DefaultClientID
 	}
-	c := &oauth.Client{
-		HTTPClient: oauth.NewHTTPClient(base),
-		Endpoints:  eps,
-		ClientID:   clientID,
-		Resource:   base,
-	}
-	return c.Revoke(ctx, rec.RefreshToken)
+	return issuer, clientID, nil
 }
 
 // effectiveBaseURL resolves the API base URL the CLI is pointed at:

internal/cli/login_test.go

@@ -2,6 +2,7 @@ package cli
 
 import (
 	"context"
+	"encoding/base64"
 	"fmt"
 	"io"
 	"net/http"
@@ -24,6 +25,10 @@ func loginTestEnv(t *testing.T, baseURL string) {
 	t.Setenv("XDG_CONFIG_HOME", t.TempDir())
 	t.Setenv(oauth.EnvNoKeyring, "1")
 	t.Setenv("EXTEND_BASE_URL", baseURL)
+	// The fake server plays both the authorization server (issuer) and
+	// the API; in production these are different origins.
+	t.Setenv("EXTEND_OAUTH_ISSUER", baseURL)
+	t.Setenv("EXTEND_OAUTH_CLIENT_ID", "")
 	t.Setenv("EXTEND_API_KEY", "")
 	t.Setenv("EXTEND_REGION", "")
 	t.Setenv("EXTEND_WORKSPACE_ID", "")
@@ -41,7 +46,9 @@ type fakeAuthServer struct {
 	redirectURI string
 	// exchanged records the code redeemed at the token endpoint.
 	exchanged string
-	// revoked records tokens sent to the revoke endpoint.
+	// accessToken overrides the token endpoint's access_token when set.
+	accessToken string
+	// revoked records the bearer tokens presented to /oauth/revoke-current.
 	revoked []string
 	// meHandler, when set, serves GET /me; unset answers 404 so tests
 	// exercise the generic-success fallback by default.
@@ -70,11 +77,14 @@ func newFakeAuthServer(t *testing.T) *fakeAuthServer {
 			t.Errorf("code_verifier does not match the challenge sent to authorize")
 		}
 		f.exchanged = r.PostForm.Get("code")
-		fmt.Fprint(w, `{"access_token":"eoat_test","refresh_token":"eort_test","token_type":"Bearer","expires_in":3600}`)
+		accessToken := f.accessToken
+		if accessToken == "" {
+			accessToken = "eoat_test"
+		}
+		fmt.Fprintf(w, `{"access_token":%q,"refresh_token":"eort_test","token_type":"Bearer","expires_in":3600}`, accessToken)
 	})
-	mux.HandleFunc("/oauth2/revoke", func(w http.ResponseWriter, r *http.Request) {
-		r.ParseForm()
-		f.revoked = append(f.revoked, r.PostForm.Get("token"))
+	mux.HandleFunc("/oauth/revoke-current", func(w http.ResponseWriter, r *http.Request) {
+		f.revoked = append(f.revoked, strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer "))
 		w.WriteHeader(http.StatusOK)
 	})
 	mux.HandleFunc("/me", func(w http.ResponseWriter, r *http.Request) {
@@ -412,12 +422,55 @@ func TestRunLoginRevokesReplacedGrant(t *testing.T) {
 
 	store := oauth.DefaultStore()
 	if err := store.Set(f.srv.URL, oauth.Record{
-		AccessToken:        "eoat_old",
-		RefreshToken:       "eort_old",
-		ExpiresAt:          time.Now().Add(time.Hour),
-		RevocationEndpoint: f.srv.URL + "/oauth2/revoke",
-		ClientID:           oauth.DefaultClientID,
-		Resource:           f.srv.URL,
+		AccessToken:  "eoat_old",
+		RefreshToken: "eort_old",
+		ExpiresAt:    time.Now().Add(time.Hour),
+		ClientID:     oauth.DefaultClientID,
+		Resource:     f.srv.URL,
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	err := runLogin(context.Background(), app, loginOptions{
+		openBrowser: f.browserFor(t, "code-xyz", ""),
+		store:       store,
+	})
+	if err != nil {
+		t.Fatalf("runLogin: %v", err)
+	}
+	if len(f.revoked) != 1 || f.revoked[0] != "eoat_old" {
+		t.Errorf("revoked = %v, want the replaced login's access token", f.revoked)
+	}
+	if rec, _ := store.Get(f.srv.URL); rec == nil || rec.RefreshToken != "eort_test" {
+		t.Errorf("stored record = %+v, want the new login", rec)
+	}
+}
+
+// jwtWithSID builds an unsigned JWT-shaped token whose payload carries
+// the given sid claim, mimicking a WorkOS Connect access token.
+func jwtWithSID(sid string) string {
+	header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
+	payload := base64.RawURLEncoding.EncodeToString([]byte(fmt.Sprintf(`{"sid":%q}`, sid)))
+	return header + "." + payload + ".sig"
+}
+
+// A re-login while the WorkOS consent is still on file issues tokens
+// with the SAME sid as the replaced grant. Revocation is sid-keyed, so
+// revoking the replaced grant would kill the new session too; the login
+// must skip it.
+func TestRunLoginSkipsRevokingSameConsentGrant(t *testing.T) {
+	f := newFakeAuthServer(t)
+	f.accessToken = jwtWithSID("app_consent_shared")
+	loginTestEnv(t, f.srv.URL)
+	app, _ := testAppForLogin()
+
+	store := oauth.DefaultStore()
+	if err := store.Set(f.srv.URL, oauth.Record{
+		AccessToken:  jwtWithSID("app_consent_shared"),
+		RefreshToken: "eort_old",
+		ExpiresAt:    time.Now().Add(time.Hour),
+		ClientID:     oauth.DefaultClientID,
+		Resource:     f.srv.URL,
 	}); err != nil {
 		t.Fatal(err)
 	}
@@ -429,14 +482,46 @@ func TestRunLoginRevokesReplacedGrant(t *testing.T) {
 	if err != nil {
 		t.Fatalf("runLogin: %v", err)
 	}
-	if len(f.revoked) != 1 || f.revoked[0] != "eort_old" {
-		t.Errorf("revoked = %v, want the replaced grant's refresh token", f.revoked)
+	if len(f.revoked) != 0 {
+		t.Errorf("revoked = %v, want none: revoking a same-sid grant would invalidate the new login", f.revoked)
 	}
 	if rec, _ := store.Get(f.srv.URL); rec == nil || rec.RefreshToken != "eort_test" {
 		t.Errorf("stored record = %+v, want the new login", rec)
 	}
 }
 
+// A replaced grant from a different consent (different sid) must still
+// be revoked so it does not linger server-side.
+func TestRunLoginRevokesDifferentConsentGrant(t *testing.T) {
+	f := newFakeAuthServer(t)
+	f.accessToken = jwtWithSID("app_consent_new")
+	loginTestEnv(t, f.srv.URL)
+	app, _ := testAppForLogin()
+
+	oldToken := jwtWithSID("app_consent_old")
+	store := oauth.DefaultStore()
+	if err := store.Set(f.srv.URL, oauth.Record{
+		AccessToken:  oldToken,
+		RefreshToken: "eort_old",
+		ExpiresAt:    time.Now().Add(time.Hour),
+		ClientID:     oauth.DefaultClientID,
+		Resource:     f.srv.URL,
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	err := runLogin(context.Background(), app, loginOptions{
+		openBrowser: f.browserFor(t, "code-xyz", ""),
+		store:       store,
+	})
+	if err != nil {
+		t.Fatalf("runLogin: %v", err)
+	}
+	if len(f.revoked) != 1 || f.revoked[0] != oldToken {
+		t.Errorf("revoked = %v, want the replaced login's access token", f.revoked)
+	}
+}
+
 func TestRunWhoami(t *testing.T) {
 	f := newFakeAuthServer(t)
 	f.meHandler = func(w http.ResponseWriter, r *http.Request) {
@@ -503,21 +588,20 @@ func TestRunLogoutRevokesAndClears(t *testing.T) {
 
 	store := oauth.DefaultStore()
 	if err := store.Set(f.srv.URL, oauth.Record{
-		AccessToken:        "eoat_x",
-		RefreshToken:       "eort_x",
-		ExpiresAt:          time.Now().Add(time.Hour),
-		RevocationEndpoint: f.srv.URL + "/oauth2/revoke",
-		ClientID:           oauth.DefaultClientID,
-		Resource:           f.srv.URL,
+		AccessToken:  "eoat_x",
+		RefreshToken: "eort_x",
+		ExpiresAt:    time.Now().Add(time.Hour),
+		ClientID:     oauth.DefaultClientID,
+		Resource:     f.srv.URL,
 	}); err != nil {
 		t.Fatal(err)
 	}
 
 	if err := runLogout(context.Background(), app, store); err != nil {
 		t.Fatalf("runLogout: %v", err)
 	}
-	if len(f.revoked) != 1 || f.revoked[0] != "eort_x" {
-		t.Errorf("revoked = %v, want the refresh token", f.revoked)
+	if len(f.revoked) != 1 || f.revoked[0] != "eoat_x" {
+		t.Errorf("revoked = %v, want the login's access token", f.revoked)
 	}
 	if rec, _ := store.Get(f.srv.URL); rec != nil {
 		t.Errorf("record should be cleared, got %+v", rec)
@@ -550,8 +634,9 @@ func TestRunLogoutClearsLocallyWhenRevokeFails(t *testing.T) {
 
 	store := oauth.DefaultStore()
 	if err := store.Set(srv.URL, oauth.Record{
-		RefreshToken:       "eort_x",
-		RevocationEndpoint: srv.URL + "/oauth2/revoke",
+		AccessToken:  "eoat_x",
+		RefreshToken: "eort_x",
+		ExpiresAt:    time.Now().Add(time.Hour),
 	}); err != nil {
 		t.Fatal(err)
 	}