Repository navigation
feat: add caching for Token Vault connection token exchanges - #139
kishore7snehil wants to merge 1 commit into
Conversation
64a62ec to
bdbcef3
Compare
| "expires_in": cached["expires_at"] - int(time.time()), | ||
| "expires_at": cached["expires_at"], | ||
| } | ||
| granted = cached.get("granted_scopes") |
There was a problem hiding this comment.
A fresh call always sets scope (even ""), but this hit path only adds scope when granted_scopes is truthy. A connection whose response carries no scope stores granted_scopes="", so the first call returns scope="" and the cached call omits the key. The method docstring promises scope, so a consumer reading result["scope"] succeeds first and then KeyErrors on a hit.
Should we set scope unconditionally here, something like cached.get("granted_scopes") or ""? Also the hit returns the normalized (sorted and deduped) scope while a fresh call returns the provider's raw string, same set but different string, worth keeping in mind.
| incoming_access_token = "incoming-auth0-access-token" | ||
|
|
||
| # Verify once, then pass the result to avoid a second verification inside the exchange. | ||
| verified = await api_client.verify_access_token(access_token=incoming_access_token) |
There was a problem hiding this comment.
This example does verified = await api_client.verify_access_token(...) and then passes verified=verified, but verify_access_token returns a claims dict while the verified kwarg expects a VerifiedToken.
With a store configured the method reads verified.access_token, so this raises AttributeError: 'dict' object has no attribute 'access_token'. Anyone copy pasting the headline example will hit a runtime crash.
Can we construct verified=VerifiedToken(access_token=incoming_access_token, claims=claims) and reword the prose (README line 125 has the same wording)?
| raise VerifyAccessTokenError( | ||
| "verified token does not match the access token being exchanged" | ||
| ) | ||
| cache_key = self._token_vault_cache.cache_key( |
There was a problem hiding this comment.
The cache key is (tenant, client_id, sub, connection) only, but login_hint is forwarded to the exchange a little below.
For a user with more than one linked account on the same connection, a first call with login_hint=A caches a token that a later call with login_hint=B (same sub and connection) will receive, so the second call gets the wrong account's token.
It is scoped to the same user, not a cross user leak. Can we fold a normalized login_hint into the key, or skip the cache when login_hint is present?
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_get_access_token_for_connection_cache_hit( |
There was a problem hiding this comment.
This cache hit test asserts only access_token and the request count. It never checks that the hit carries expires_in, expires_at and scope, or that the hit and fresh results have equal keys, so the shape mismatch above goes undetected.
Can we assert result1.keys() == result2.keys() and add cases with and without scope in the exchange response?
|
|
||
| def cache_key(self, tenant: str, client_id: str, sub: Optional[str], connection: str) -> Optional[str]: | ||
| """Return the cache key, or None when sub is absent (cache skipped with a warning).""" | ||
| if not sub: |
There was a problem hiding this comment.
This guards only with if not sub. A caller supplied verified with a truthy non string sub passes the guard, then the key builder does a str.join with sub and raises TypeError. This runs outside any try, so it breaks the whole exchange instead of degrading to a cache skip. The OBO path checks isinstance(sub, str).
Can we use if not isinstance(sub, str) or not sub: and skip the cache with the existing warning?
|
|
||
| if cached is None: | ||
| return None | ||
| if "access_token" not in cached or "expires_at" not in cached: |
There was a problem hiding this comment.
This malformed entry check only tests for key presence. A non int expires_at then makes cached["expires_at"] <= int(time.time()) raise TypeError outside the try, instead of being treated as a miss like other malformed entries.
Can we also require isinstance(cached.get("expires_at"), int) here (and a non empty str access_token)?
| Raises: | ||
| GetAccessTokenForConnectionError: If there was an issue requesting the access token. | ||
| ApiError: If the token exchange endpoint returns an error. | ||
| VerifyAccessTokenError: If a store is configured and either verified is omitted and the token fails verification, or verified is supplied but does not match the access token being exchanged. |
There was a problem hiding this comment.
Minor, with a store configured and verified omitted this method calls verify_access_token, which can raise MissingOrganizationError or OrganizationNotAllowedError under organization_policy='required'. The Raises block only lists VerifyAccessTokenError.
Can we add both, like the OBO doc-string does?
bdbcef3 to
2abdc32
Compare
fcb1832 to
6576276
Compare
Co-Authored-By: Claude <noreply@anthropic.com>
6576276 to
a9668aa
Compare
2abdc32 to
eee148a
Compare
📋 Changes
This PR extends the token store added in the previous PR to also cache Token Vault exchanges. When a token store is configured,
get_access_token_for_connection()reuses a previously exchanged connection token instead of hitting the token endpoint on every call.✨ Features
get_access_token_for_connection()caches exchanged tokens whentoken_storeis set. Entries are keyed by tenant, client, callersuband connection, so one caller or connection never gets another's token.sub: If the verified token has no usablesubclaim, the cache is skipped with a warning.🔧 API Changes
verified(VerifiedToken, defaultNone) toget_access_token_for_connection(). Callers that have already verified the token (for example an MCP server) can pass it to avoid a second verification. When omitted and a token store is configured, the token is verified before any cache lookup.get_access_token_for_connection()now raisesVerifyAccessTokenErrorwhen a store is configured and the token fails verification, or whenverifieddoes not match the access token being exchanged.expires_inalongsideaccess_token,expires_atandscope.📖 Documentation
EXAMPLES.mdwith a Token Vault section covering a basic call and cachingREADME.mdanddocs/TokenStorage.mdto mention Token Vault caching🧪 Testing
Contributor Checklist
🤖 Generated with Claude Code