Skip to content

feat: add caching for Token Vault connection token exchanges - #139

Open
kishore7snehil wants to merge 1 commit into
feat/m2m-client-credentialsfrom
feat/token-vault-caching
Open

kishore7snehil wants to merge 1 commit into
feat/m2m-client-credentialsfrom
feat/token-vault-caching

Conversation

@kishore7snehil

Copy link
Copy Markdown
Contributor

📋 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

  • Token Vault Caching: get_access_token_for_connection() caches exchanged tokens when token_store is set. Entries are keyed by tenant, client, caller sub and connection, so one caller or connection never gets another's token.
  • Fail-Open Store Handling: A store read or write failure is logged and the exchange proceeds normally. A malformed or expired entry counts as a cache miss.
  • No Caching Without sub: If the verified token has no usable sub claim, the cache is skipped with a warning.

🔧 API Changes

  • Added an optional keyword argument verified (VerifiedToken, default None) to get_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 raises VerifyAccessTokenError when a store is configured and the token fails verification, or when verified does not match the access token being exchanged.
  • The result now always includes expires_in alongside access_token, expires_at and scope.

📖 Documentation

  • Updated EXAMPLES.md with a Token Vault section covering a basic call and caching
  • Updated README.md and docs/TokenStorage.md to mention Token Vault caching

🧪 Testing

  • This change adds test coverage
  • This change has been tested on the latest version of the platform/language

Contributor Checklist

🤖 Generated with Claude Code

@kishore7snehil
kishore7snehil marked this pull request as ready for review October 6, 2026 06:55
@kishore7snehil
kishore7snehil requested a review from a team as a code owner October 6, 2026 06:55
@kishore7snehil
kishore7snehil force-pushed the feat/token-vault-caching branch from 64a62ec to bdbcef3 Compare October 7, 2026 14:54
@kishore7snehil
kishore7snehil changed the base branch from feat/obo-token-storage to feat/m2m-client-credentials October 7, 2026 14:54
"expires_in": cached["expires_at"] - int(time.time()),
"expires_at": cached["expires_at"],
}
granted = cached.get("granted_scopes")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread EXAMPLES.md
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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Comment thread tests/test_api_client.py


@pytest.mark.asyncio
async def test_get_access_token_for_connection_cache_hit(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@kishore7snehil
kishore7snehil force-pushed the feat/token-vault-caching branch from bdbcef3 to 2abdc32 Compare October 8, 2026 04:07
@kishore7snehil
kishore7snehil force-pushed the feat/m2m-client-credentials branch from fcb1832 to 6576276 Compare October 8, 2026 04:07
@kishore7snehil
kishore7snehil force-pushed the feat/m2m-client-credentials branch from 6576276 to a9668aa Compare October 8, 2026 04:22
@kishore7snehil
kishore7snehil force-pushed the feat/token-vault-caching branch from 2abdc32 to eee148a Compare October 8, 2026 04:22

This branch has not been deployed

No deployments
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.

2 participants