Skip to content

Add Google access token and ID token auth methods - #30

Open
tomchop wants to merge 1 commit into
mainfrom
feat/token-exchange-auth
Open

tomchop wants to merge 1 commit into
mainfrom
feat/token-exchange-auth

Conversation

@tomchop

@tomchop tomchop commented Oct 2, 2026

Copy link
Copy Markdown
Contributor

Adds client methods for POST /api/v2/auth/google-access-token from yeti-platform/yeti#1403 and for the existing ID token endpoint. It also fixes a refresh loop on rejected credentials, which the new methods would otherwise inherit.

What changed

  • auth_google_access_token(token_provider) exchanges a Google OAuth access token for a session token. It's for clients that can mint access tokens without a browser but can't get ID tokens. The server must run with auth.module = oidc and auth.google_access_token_client_ids set. Otherwise the endpoint answers 404, which surfaces as YetiApiError with status_code == 404.
  • auth_id_token(token_provider) does the same with an ID token at /api/v2/auth/oidc-callback-token. That endpoint had no client method.
  • Token providers, not tokens. Both methods take a function that returns a token. When a request gets a 401, refresh_auth() calls the auth method again with no arguments, and the method calls the provider for a fresh token. Session tokens expire after 30 minutes by default and Google access tokens after about an hour, so a long-running client can't keep reusing the first token. There's no new dependency: with google-auth, the provider is a few lines that refresh the credentials and return credentials.token.
  • Refresh can't recurse. Before, a 401 from an auth endpoint went through the same refresh-and-retry as any other request: refresh_auth() → auth_api_key() → 401 → refresh_auth()…. Once a session existed, a revoked API key recursed until RecursionError, sending a request at each level. The three auth methods now share _start_session, which sends the exchange with retries=0. A rejected credential raises YetiAuthError after one attempt.
  • _set_session_token. Auth methods now hand the session token to this method. It updates the headers of the existing requests.Session, where auth_api_key used to replace the session with a new one. That keeps the session's settings (TLS verification, adapters, hooks). It also gives subclasses that send requests through another transport a single method to override.
  • _auth_function_map is typed as dict[str, Callable[[], None]], matching how refresh_auth calls its values.

Verification

Unit and type checks, run the way CI does after poetry install --no-root:

  • poetry run python -m unittest tests/api.py: 47 passed (8 new). On main's yeti/api.py, test_auth_api_key_revoked errors with RecursionError.
  • poetry run pyrefly check: 0 errors

tests/e2e.py against a local Yeti built from main plus yeti-platform/yeti#1403, using API key auth: 14 passed.

I also ran the client against two local Yeti servers built from that branch, one with the exchange enabled and one without. It used real Google access tokens for my account, fetched through the provider and never printed:

[PASS] auth_google_access_token: user_matches=True minted=1 same_session=True requests=['POST /api/v2/auth/google-access-token 200', 'GET /api/v2/auth/me 200']
[PASS] refresh mints a new token: user_matches=True minted=2 requests=['GET /api/v2/auth/me 401', 'POST /api/v2/auth/google-access-token 200', 'GET /api/v2/auth/me 200']
[PASS] disabled user raises after one renewal: YetiAuthError minted=3 requests=['GET /api/v2/auth/me 401', 'POST /api/v2/auth/google-access-token 401']
[PASS] exchange not enabled: YetiApiError 404
[PASS] auth_api_key with renewal: user_matches=True same_session=True requests=['POST /api/v2/auth/api-token 200', 'GET /api/v2/auth/me 401', 'POST /api/v2/auth/api-token 200', 'GET /api/v2/auth/me 200']
[PASS] revoked API key raises after one renewal: YetiAuthError requests=['GET /api/v2/auth/me 401', 'POST /api/v2/auth/api-token 401']
[PASS] auth_id_token rejects an invalid token: YetiAuthError requests=['POST /api/v2/auth/oidc-callback-token 401']
7/7 passed

The request lists come from a response hook on the client's session, and minted is a running count of provider calls. To expire a session, the check replaced the session token with an invalid one. Between steps, it disabled the user, then the API key, directly in the server's database. The output shows one renewal attempt per expired session, and no retry loop when the server rejects the credential. A successful auth_id_token isn't covered: a Google ID token for the test server's client needs a browser flow. The check only confirms that the endpoint receives the token and that a rejection raises YetiAuthError.

Release

No version bump: 2.4.0, from #28, isn't released yet, and this ships with it. auth_google_access_token needs a server with yeti-platform/yeti#1403. Against older servers it raises YetiApiError (404).

auth_google_access_token exchanges a Google OAuth access token at
/api/v2/auth/google-access-token, for clients that can mint access tokens
without a browser but can't get ID tokens. auth_id_token does the same with an
ID token at /api/v2/auth/oidc-callback-token. Both take a function that returns
a token rather than the token itself, and call it again whenever the session
token expires, so long-running clients outlive the first token.

The three auth methods share _start_session, which sends the exchange without
retries. A 401 from an auth endpoint used to trigger refresh_auth, which called
the auth method again: once a session existed, a revoked API key recursed until
RecursionError.

Session tokens go through _set_session_token, which updates the headers of the
existing requests session instead of replacing the session. Subclasses that send
requests through another transport override it.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant