Skip to content

cloud_api_client: Key the cached LlmApiToken by organization - #58915

Merged
maxdeviant merged 1 commit into
mainfrom
tomhoule-zpwtwnynmqxo
Jun 9, 2026
Merged

cloud_api_client: Key the cached LlmApiToken by organization#58915
maxdeviant merged 1 commit into
mainfrom
tomhoule-zpwtwnynmqxo

Conversation

@tomhoule

@tomhoule tomhoule commented Jun 9, 2026

Copy link
Copy Markdown
Contributor

The tokens are organization-specific (the organization id in the JWT claims). Before this commit, the LlmApiToken cache was an Option<String>, and cached() ignored the organization_id argument on cache hit. Correctness relied on RefreshLlmTokenListener clearing and refreshing the token when it observed an OrganizationChanged event.

That leaves two windows where a token minted for the wrong organization could be served from the cache:

  • the refresh runs in a spawned task, so between the organization switch and the task acquiring the write lock, any request calling cached() would get the previous organization's token.
  • Client::authenticated_llm_request snapshots the organization id at call time; if the server demanded a token refresh after the user had switched organizations mid-request, the retry path would mint a fresh token for the old organization and write it into the shared cache, poisoning it for all subsequent requests until the next refresh event.

The fix is to store the organization id corresponding to the token next to it in the cache, and to check that the cached token is for the correct organization in cached().

Release Notes:

  • Fixed a race where LLM and edit prediction requests made immediately after switching organizations could be attributed to the previously selected organization.

@cla-bot cla-bot Bot added the cla-signed The user has signed the Contributor License Agreement label Jun 9, 2026
@zed-community-bot zed-community-bot Bot added the staff Pull requests authored by a current member of Zed staff label Jun 9, 2026
@tomhoule
tomhoule force-pushed the tomhoule-zpwtwnynmqxo branch from ce800ae to 0830065 Compare June 9, 2026 11:46
@tomhoule
tomhoule changed the base branch from main to push-kuyxlrurmvrv June 9, 2026 11:48
@tomhoule tomhoule changed the title cloud_api_client: Key the cached LLM token by organization cloud_api_client: Key the cached LlmApiToken by organization Jun 9, 2026
Base automatically changed from push-kuyxlrurmvrv to main June 9, 2026 12:09
@tomhoule
tomhoule force-pushed the tomhoule-zpwtwnynmqxo branch from 0830065 to 74800d6 Compare June 9, 2026 12:19
@tomhoule
tomhoule marked this pull request as ready for review June 9, 2026 12:33
Comment thread crates/cloud_api_client/src/llm_token.rs Outdated
@tomhoule
tomhoule force-pushed the tomhoule-zpwtwnynmqxo branch 2 times, most recently from b1fda56 to 2bf96d3 Compare June 9, 2026 12:59
LLM tokens are organization-specific: the organization id in the JWT claims. Before this commit, the `LlmApiToken` cache was an `Option<String>`, and `cached()` ignored the `organization_id` argument on cache hit. Correctness relied on `RefreshLlmTokenListener` clearing and refreshing the token when it observed an `OrganizationChanged` event.

That leaves two windows where a token minted for the wrong organization could be served from the cache:

- the refresh runs in a spawned task, so between the organization switch and the task acquiring the write lock, any request calling `cached()` would get the previous organization's token.
- `Client::authenticated_llm_request` snapshots the organization id at call time; if the server demanded a token refresh after the user had switched organizations mid-request, the retry path would mint a fresh token for the old organization and write it into the shared cache, poisoning it for all subsequent requests until the next refresh event.

The fix is to store the organization id corresponding to the token next to it in the cache, and to check that the cached token is for the correct organization in `cached()`.

Release Notes:

- Fixed a race where LLM requests made immediately after switching organizations could be attributed to the previously selected organization.
@tomhoule
tomhoule force-pushed the tomhoule-zpwtwnynmqxo branch from 2bf96d3 to e37c5b9 Compare June 9, 2026 13:02
token: String,
}

impl fmt::Debug for CachedLlmApiToken {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Adding this to avoid tokens ending up in logs, but obviously that's orthogonal to this PR's purpose.

@maxdeviant
maxdeviant added this pull request to the merge queue Jun 9, 2026
Merged via the queue into main with commit 881b4fb Jun 9, 2026
32 checks passed
@maxdeviant
maxdeviant deleted the tomhoule-zpwtwnynmqxo branch June 9, 2026 19:15
This was referenced Jun 18, 2026
jonx pushed a commit to jonx/zed-aros that referenced this pull request Jul 17, 2026
…ustries#58915)

The tokens are organization-specific (the organization id in the JWT
claims). Before this commit, the `LlmApiToken` cache was an
`Option<String>`, and `cached()` ignored the `organization_id` argument
on cache hit. Correctness relied on `RefreshLlmTokenListener` clearing
and refreshing the token when it observed an `OrganizationChanged`
event.

That leaves two windows where a token minted for the wrong organization
could be served from the cache:

- the refresh runs in a spawned task, so between the organization switch
and the task acquiring the write lock, any request calling `cached()`
would get the previous organization's token.
- `Client::authenticated_llm_request` snapshots the organization id at
call time; if the server demanded a token refresh after the user had
switched organizations mid-request, the retry path would mint a fresh
token for the old organization and write it into the shared cache,
poisoning it for all subsequent requests until the next refresh event.

The fix is to store the organization id corresponding to the token next
to it in the cache, and to check that the cached token is for the
correct organization in `cached()`.

Release Notes:

- Fixed a race where LLM and edit prediction requests made immediately
after switching organizations could be attributed to the previously
selected organization.
jolutz pushed a commit to jolutz/zed that referenced this pull request Aug 8, 2026
…ustries#58915)

The tokens are organization-specific (the organization id in the JWT
claims). Before this commit, the `LlmApiToken` cache was an
`Option<String>`, and `cached()` ignored the `organization_id` argument
on cache hit. Correctness relied on `RefreshLlmTokenListener` clearing
and refreshing the token when it observed an `OrganizationChanged`
event.

That leaves two windows where a token minted for the wrong organization
could be served from the cache:

- the refresh runs in a spawned task, so between the organization switch
and the task acquiring the write lock, any request calling `cached()`
would get the previous organization's token.
- `Client::authenticated_llm_request` snapshots the organization id at
call time; if the server demanded a token refresh after the user had
switched organizations mid-request, the retry path would mint a fresh
token for the old organization and write it into the shared cache,
poisoning it for all subsequent requests until the next refresh event.

The fix is to store the organization id corresponding to the token next
to it in the cache, and to check that the cached token is for the
correct organization in `cached()`.

Release Notes:

- Fixed a race where LLM and edit prediction requests made immediately
after switching organizations could be attributed to the previously
selected organization.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla-signed The user has signed the Contributor License Agreement staff Pull requests authored by a current member of Zed staff

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants