fix: OHE-3276 make managed LLM key rotation concurrency-safe (enterprise#439)

[aivong-openhands]

HUMAN:

• A human has tested these changes.

AGENT:

---

Why

Fixes #439. After upgrading enterprise-server to 1.60.0, conversations
intermittently died on the sandbox's first LLM call with a non-retryable
401 AuthenticationError (Invalid proxy server token... token_not_found_in_db)
even though a valid managed key existed and worked for server-side calls.

Root cause (verified by live reproduction): the managed-key refresh runs on
every conversation start, and rotate_managed_llm_key() did
delete_key_by_alias(<shared per-org alias>) then generate_key(...) with no
serialization. Two overlapping rotations for the same org could interleave so
that one rotation's delete_key_by_alias wiped the key the other rotation had
just minted and handed to a sandbox. That sandbox's key was then absent from
LiteLLM's LiteLLM_VerificationToken table -> 401 token_not_found_in_db.

Summary

• rotate_managed_llm_key() now serializes per org with a
  pg_advisory_xact_lock, so rotations for the same org cannot interleave.
• It re-reads the member key under the lock and, when given only_if_current,
  reuses a key a concurrent rotation already minted instead of rotating again
  (idempotent stale-key refresh). This stops overlapping conversation starts
  from orphaning each other's keys.
• Within the lock it clears the shared alias with delete_key_by_alias and then
  mints the replacement. The by-alias clear is not the Managed-key rotation race on conversation start → intermittent token_not_found_in_db 401 (enterprise-server 1.60.0 regression) #439 bug once it runs
  under the lock: serialization means it can no longer race a concurrent
  rotation's freshly-minted key. Clearing by alias (rather than only the
  specific previous key) also self-heals a divergent alias (e.g. after a
  rotation whose LiteLLM key was minted but whose DB commit failed), which would
  otherwise wedge every future rotation for that org on a unique-alias 400.
• The conversation-start path (_maybe_refresh_managed_llm_key) verifies the
  member key and, when stale, calls rotate(only_if_current=<observed key>).

Issue Number

Fixes #439.

How to Test

Unit tests:

uv run pytest tests/unit/server/routes/test_api_keys.py \
  tests/unit/app_server/test_managed_key_rotation.py \
  tests/unit/app_server/test_live_status_app_conversation_service.py

Live validation was run on a KinD + Helm install against the built image
(ghcr.io/openhands/enterprise-server:sha-435b397) and a real LiteLLM proxy.
Three layers, all green:

1. End-to-end (real _maybe_refresh_managed_llm_key + real /chat/completions):
   with the old unserialized rotate, an overlapping start orphans a sandbox's key
   and its real inference call returns 401 token_not_found; with the fix, two
   concurrent starts converge on one fresh key and both real calls authenticate.
   A divergence scenario (stray key under the alias + stale DB key) self-heals
   with the fix (it would wedge the delete-specific variant forever).
2. Negative control proving the advisory lock is load-bearing: with the lock
   disabled (unique lock key per call), 6 concurrent refreshes error on
   unique-alias 400s and orphan a key; with the lock enabled they all converge on
   one valid key. A cross-pod run (3 replicas, separate processes / Postgres
   connections) confirms pg_advisory_xact_lock serializes across connections.
3. Concurrency soak: 25 iterations x 5 concurrent stale-key refreshes (125 total)
   with 0 orphans, 0 errors; every race converges on one valid key.

Type

• Bug fix
• Feature
• Refactor
• Breaking change
• Docs / chore

Notes

No schema changes. The advisory lock is a per-org transaction lock keyed by the
org id; it only serializes managed-key rotations for the same org.

---

Enterprise server image for this PR:

ghcr.io/openhands/enterprise-server:sha-4f31fc2

[github-actions]

Coverage report

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
openhands/app_server/app_conversation
live_status_app_conversation_service.py
openhands/app_server/settings
settings_router.py					210-211
server/routes
api_keys.py
storage
saas_settings_store.py					1182, 1204, 1208
Project Total

This report was generated by python-coverage-comment-action

[aivong-openhands]

Live verification on a Replicated VM ✅

Reproduced #439 on a control build and confirmed this PR resolves it, on a self-hosted
Replicated embedded-cluster VM (single-node k0s, external Postgres/RDS, single-replica
LiteLLM proxy with in-memory positive-only auth cache).

Builds

• Control (bug present, fix: OHE-3276 make managed LLM key rotation concurrency-safe (enterprise#439) #440 absent): enterprise-server:1.61.0.
• Fix under test: this PR's head 4e765a2 → enterprise-server:sha-4e765a2
  (amd64 sha256:c21606c8…), deployed as an in-cluster override of deploy/openhands.
• Confirmed the running fix image actually carries the change:
  rotate_managed_llm_key(self, *, only_if_current: str | None = None) and the per-org
  pg_advisory_xact_lock.
• Caveat: the PR head sits on current main (ahead of 1.61.0), not a 1.61.0 rebase.
  Control was captured on the VM's real 1.61.0; after-runs on sha-4e765a2 on the same VM.

---

1. Server-side collision — already exists 400 on /key/generate

Probe = 4 concurrent rotate_managed_llm_key for one org (reset to a single token before each run).

Build	rotations OK	key/generate 400	surviving tokens
Control 1.61.0	1 / 4 (3 exceptions)	3	1
sha-4e765a2 (×5 runs)	4 / 4 every run	0	1

---

2. Client-side cache 401 — the #439 symptom, reproduced end-to-end

Minted a managed key K0, verified a completion works, then fired two overlapping refreshes and
had a "sandbox" call the proxy with the key the first refresh handed out. The only difference between
the two rows is the only_if_current argument the refresh path passes
(live_status_app_conversation_service._maybe_refresh_managed_llm_key: this PR's L1530 vs 1.61.0's L1525):

buggy caller  rotate_managed_llm_key()                → refreshA/refreshB keys DIFFER (B deleted A's key)
              sandbox completion(keyA)  → HTTP 401
              "Invalid proxy server token passed... Unable to find token in cache or
               `LiteLLM_VerificationTokenTable`"        # == #439 token_not_found_in_db, verbatim

fixed caller  rotate_managed_llm_key(only_if_current=K0) → refreshA/refreshB keys IDENTICAL (B skipped, reused)
              sandbox completion(keyA)  → HTTP 200

---

3. Why #439 is "intermittent" yet can hit on every conversation

• A single isolated conversation start fired exactly one rotation
  (stale_key_detected → rotate_managed_llm_key:rotated, ~85 ms), no skipped_concurrent — it does
  not fail on its own. Confirms it is a genuine race.

• Overlapping starts (all seeing the same stale key) are where it breaks:

  caller	keys minted for 4 overlapping starts	result
  buggy (1.61.0)	4 distinct	3 of 4 keys orphaned → those sandboxes 401
  fixed (this PR)	1 (3 idempotent skips)	0 orphaned → 0 failures

  N overlapping starts orphan N−1 keys on the old path. Because verify-then-rotate runs from several
  places per start (the live-status path at two call sites + the settings path, all on the same shared
  per-org alias) and the stale-key check fires on ~every start, the overlap is effectively structural —
  which is how an "intermittent" race presents as failing on nearly every conversation.

---

Key finding

The advisory lock alone does not fix the 401: running the buggy caller (only_if_current=None) on the
fix image (lock active) still minted 4 distinct keys and orphaned 3. The only_if_current idempotency
guard is the part that prevents deleting a key already handed to a sandbox. This PR ships both, and they
address two distinct faces of the same race:

symptom	fixed by
already exists 400 on /key/generate	pg_advisory_xact_lock (serializes rotations)
token_not_found_in_db 401 at completion (#439)	only_if_current (idempotent refresh, no orphaning)

Verdict: #439 reproduced on 1.61.0; both faces gone on sha-4e765a2 across all runs.

---

_{Drafted by an AI agent (OpenHands) on behalf of the tester.}

[aivong-openhands]

Mutation review of the tests in this PR

I mutation-tested the suites this PR touches: break a behaviour in the source on
purpose, re-run the tests, and see whether anything fails. Coverage says a line
ran; a mutant that survives says nothing was asserting it.

Baseline — tests/unit/server/routes/test_api_keys.py::TestRefreshManagedLlmApiKey

• tests/unit/app_server/test_managed_key_rotation.py
• tests/unit/app_server/test_live_status_app_conversation_service.py:
  232 passed in ~9s, green before mutating. 14 hand-written mutants: 5 controls
  (reverting the actual fix, which must die) and 9 candidates around it.

Caught — all 5 controls, plus 6 of the 9 candidates ❌

#	Mutant	Result
C1	Delete the only_if_current idempotency skip entirely (always rotate)	❌ caught
C2	Delete the per-org pg_advisory_xact_lock call	❌ caught
C3	Conversation start reverts to rotate_managed_llm_key() with no argument	❌ caught
C4	Settings write reverts to rotate_managed_llm_key() with no argument	❌ caught
C5	Route deletes rotation.old_key again (the duplicate cleanup this PR removed)	❌ caught
M1	Lock key made unique per call, so the lock serializes nothing	❌ caught
M2	Drop session.expire_all() before the re-read under the lock	❌ caught
M3	Generate-before-delete: mint first, then clear the alias	❌ caught
M4	Alias-clear failure aborts the rotation (remove the best-effort try/except)	❌ caught
M7	int.from_bytes(..., signed=False) for the lock key	❌ caught
M9	Invert the idempotency guard (old_key == only_if_current)	❌ caught

The central claim of the PR is pinned, and pinned by behaviour rather than by
call-shape assertions. Three things in particular carry that weight:

test_concurrent_rotations_serialize_without_orphaning_survivor is the one that
kills C2 and M1, and it does so because it drives four genuinely overlapping
rotate_managed_llm_key calls through asyncio.gather against real Postgres and
then asserts on the outcome — handed == {final_key}, one generate_key, one
delete_key_by_alias. Deleting the advisory lock is enough to break it. A test
that had asserted "the SQL contains pg_advisory_xact_lock" would have passed
M1, which keeps the call and neuters the key.

test_rotate_openhands_model_attaches_openhands_metadata kills M3 via the
attach_mock ordering assertion, not via the individual assert_awaited_once
calls — those both still hold when you swap the two operations. Since the whole
reason the order is mandatory is LiteLLM's unique-alias constraint, asserting the
order and not just the calls is the right call.

test_route_continues_when_alias_clear_fails kills M4 by making the mock raise
and still demanding refreshed=True plus a completed mint, which is what makes
the best-effort try/except load-bearing rather than decorative.

Survivors — 3 gaps ✅

#	Mutant	Result
M5	Drop and old_key is not None from the idempotency guard	✅ 232 passed
M6	_org_rotation_advisory_lock_key returns a constant 0 for every org	✅ 232 passed
M8	Lock key derived from a 4-byte digest instead of 8	✅ 232 passed

---

M5 — a member with no stored key returns ROTATED with new_key=None

The guard is three clauses:

if (
    only_if_current is not None
    and old_key is not None          # <-- delete this and nothing fails
    and old_key != only_if_current
):
    return ManagedLlmKeyRotation(
        status=ManagedLlmKeyStatus.ROTATED, old_key=old_key, new_key=old_key, ...
    )

Remove the middle clause and a member whose _llm_api_key is empty (the column
is NOT NULL, so unset is stored as '' and read back as old_key is None)
satisfies None != only_if_current and takes the skip branch. I confirmed what
comes back in that case:

status=rotated new_key=None old_key=None generate_called=False alias_called=False

ManagedLlmKeyStatus.ROTATED with new_key=None. Both callers gate on
rotation.status == ROTATED and rotation.new_key, so today they fall through to
rotation_not_applied and return the stale llm — the member silently never
gets a key, and the refresh endpoint reports refreshed=True for a rotation that
minted nothing. That clause is the only thing standing between "unset key" and
"claimed a successful rotation and did nothing", and no test holds it in place.

This one is reachable: _get_effective_llm_api_key and
get_current_managed_llm_key both branch on not org_member._llm_api_key, so
the empty-member-key state is one the storage layer already expects.

@pytest.mark.asyncio
async def test_rotate_only_if_current_mints_when_member_has_no_key(
    self, async_session_maker, managed_env
):
    """A member with no stored key must be minted a real one even under
    ``only_if_current``: ``None != <observed key>`` must not be read as
    "a concurrent rotation already replaced it" and short-circuit to
    ROTATED with ``new_key=None``.
    """
    user_id, org_id = await self._seed(async_session_maker)
    expected_alias = get_openhands_cloud_key_alias(user_id, str(org_id))

    # ``_llm_api_key = ''`` is how an unset member key is stored (the column
    # is NOT NULL); ``rotate`` reads that as ``old_key is None``.
    async with async_session_maker() as session:
        await session.execute(
            update(OrgMember)
            .where(
                OrgMember.org_id == org_id,
                OrgMember.user_id == uuid.UUID(user_id),
            )
            .values(_llm_api_key='')
        )
        await session.commit()

    with self._patched(async_session_maker) as (
        mock_delete_alias,
        mock_generate,
        mock_delete_token,
    ):
        store = SaasSettingsStore(user_id, effective_org_id=org_id)
        rotation = await store.rotate_managed_llm_key(
            only_if_current='sk-a-key-the-caller-observed'
        )

    assert rotation.status == ManagedLlmKeyStatus.ROTATED
    assert rotation.new_key == 'sk-new-managed-key'
    mock_generate.assert_awaited_once()
    mock_delete_alias.assert_awaited_once_with(key_alias=expected_alias)
    mock_delete_token.assert_not_called()

    member = await self._member_key(async_session_maker, org_id, user_id)
    assert member.llm_api_key.get_secret_value() == 'sk-new-managed-key'

Verified: passes on this branch unmodified, fails with M5 applied. Needs
from sqlalchemy import update and from storage.org_member import OrgMember.

M6 / M8 — nothing asserts the lock key is per-org

_org_rotation_advisory_lock_key is a new pure function and has no test. Two
different degradations of it are invisible to the suite:

• M6 — return 0 regardless of org_id. Every org in the deployment then
  serializes managed-key rotations against one global advisory lock. Correctness
  survives; throughput does not, and it does so silently, on the conversation-start
  path. The PR's concurrency test can't see this, because it uses a single org —
  one lock and a correct per-org lock look identical from inside it.
• M8 — digest_size=4. The lock key drops to ~32 bits of entropy, so two
  unrelated orgs collide with probability ~2⁻³² per pair instead of ~2⁻⁶⁴, and a
  collision means two orgs' rotations block each other for no reason. Also
  invisible.

M7 (signed=False) is caught today, but only indirectly: an out-of-range
bigint makes Postgres reject the pg_advisory_xact_lock call, so the concurrency
test blows up. A direct assertion on the helper states the constraint instead of
discovering it through a driver error.

One test covers all three:

def test_rotation_advisory_lock_key_is_per_org_and_fits_bigint():
    org_a, org_b = str(uuid.uuid4()), str(uuid.uuid4())

    key_a = _org_rotation_advisory_lock_key(org_a)
    assert _org_rotation_advisory_lock_key(org_a) == key_a
    assert _org_rotation_advisory_lock_key(org_b) != key_a

    assert -(2**63) <= key_a < 2**63
    # Derived from the full 64 bits, so distinct orgs collide at ~2^-64 rather
    # than the ~2^-32 a narrower digest would give.
    keys = [_org_rotation_advisory_lock_key(str(uuid.uuid4())) for _ in range(16)]
    assert max(abs(k) for k in keys) > 2**48

Verified: passes on this branch unmodified, and fails with M6, M8, and M7
applied. Needs _org_rotation_advisory_lock_key imported from
storage.saas_settings_store. There is precedent for testing a lock key directly
in tests/unit/app_server/test_secrets_api_race.py
(test_secrets_write_lock_key_uses_user_and_org), which asserts the same
"different scopes must not share a lock" property.

Not a test gap

• The rotate(only_if_current=...) signature is a genuinely better shape than a
  test could enforce: making the precondition an argument means a caller cannot
  forget to re-check it, and C3/C4 show both callers are pinned to passing it.
• rotation.old_key is now returned but no longer consumed anywhere in
  production — the route's cleanup was removed and the store cleans up by alias
  under the lock. Not a test concern, but the field is dead weight on the
  ManagedLlmKeyRotation contract and might be worth dropping while the context
  is fresh.
• Nothing in the suite covers the cross-process behaviour of
  pg_advisory_xact_lock (separate connections / pods); the four-way test runs in
  one process against one engine. The PR body says that was validated live, which
  is the right layer for it — a unit test cannot reasonably prove it.

---

This comment was generated by an AI assistant on behalf of the user.

[aivong-openhands]

Thanks for the mutation review — really useful framing. I've closed the three surviving mutants in fa7984b85:

• M5 — added test_rotate_only_if_current_mints_when_member_has_no_key. It seeds a member whose _llm_api_key is '' (read back as old_key is None), then rotates with only_if_current set and asserts a real key is minted (generate awaited once, alias cleared, member ends up with the new key). Verified it fails with the middle and old_key is not None clause removed and passes on the branch unmodified.
• M6 / M8 / M7 — added test_rotation_advisory_lock_key_is_per_org_and_fits_bigint, asserting _org_rotation_advisory_lock_key is stable per org, distinct across orgs, in signed 64-bit range, and derived from the full-width digest. Verified it fails with return 0, digest_size=4, and signed=False applied, and passes unmodified.

All 64 tests in test_api_keys.py pass; pre-commit (ruff/format/mypy) is green.

On the not-a-test-gap note about ManagedLlmKeyRotation.old_key: leaving it as-is for now to keep this PR scoped to the concurrency fix — it's still asserted on by the idempotency tests as a behavioural signal, and dropping it is a contract change better done on its own. Happy to follow up separately if you'd prefer.

This comment was created by an AI agent (OpenHands) on behalf of the user.

[aivong-openhands]

Note: submitted as a COMMENT review rather than APPROVE because GitHub does not allow approving a pull request authored by the same account. The verdict below is nonetheless approve.

Code Review

🟢 Good — a focused, well-reasoned concurrency fix for the #439 token_not_found_in_db 401. I traced the full flow (_maybe_refresh_managed_llm_key / _maybe_rotate_stale_managed_key → rotate_managed_llm_key → LiteLlmManager delete/generate) against the PR head (fa7984b) and found no blocking issues. The serialization is provably correct, the idempotent re-read is the right shape, and the concurrency test is a genuine regression guard against real Postgres.

[CRITICAL ISSUES]

None.

[WHAT WAS VERIFIED]

• [storage/saas_settings_store.py, L1141-1149] Serialization is correct and leak-free: pg_advisory_xact_lock is acquired inside the session's transaction and auto-releases on commit/rollback. Every early-return path (MISSING_MEMBER, BYOK, the skipped_concurrent idempotency branch) exits the async with a_session_maker() block, so the lock is always released, and no intermediate commit precedes the mint+persist — the critical section (re-read → delete-by-alias → generate → persist) is fully covered by one lock hold.

• [storage/saas_settings_store.py, L1150-1158] The re-read under the lock is load-bearing: session.expire_all() + re-select(User) under READ COMMITTED makes a concurrent winner's committed key visible, so a loser sees old_key != only_if_current and reuses the fresh key instead of rotating (and orphaning) it. Using the pre-lock value here would have reintroduced the race — the code correctly reads after acquiring the lock.

• [storage/saas_settings_store.py, L1177-1194] The only_if_current empty-key edge is handled: the old_key is not None guard mints a real key for a member with no stored key (_llm_api_key = '') rather than short-circuiting to new_key=None. Pinned by test_rotate_only_if_current_mints_when_member_has_no_key.

• [server/routes/api_keys.py, L494-499] Removing the route-level delete_key(old_key) is correct: delete-then-generate ordering is mandatory (LiteLLM enforces unique aliases) and is now safe under the per-org lock; the by-alias clear already removes the previous key, so the route's extra best-effort delete was redundant.

[TESTING GAPS]

None blocking. Credit where due: test_concurrent_rotations_serialize_without_orphaning_survivor runs 4 real concurrent rotate_managed_llm_key() calls against a real Postgres DB (postgres_testdb, NullPool → distinct connections, so the advisory lock actually serializes across connections) and asserts exactly one generate_key + one delete_key_by_alias with all callers converging on one key. It fails on the old unserialized code and passes on the new — exactly the right shape, not mock-wiring. The lock-key helper test also pins per-org distinctness and signed 64-bit bigint fit.

[IMPROVEMENT OPPORTUNITIES]

• [storage/saas_settings_store.py, L1195-1207] Divergent-alias cleanup (non-blocking): in the edge case of a prior key minted under a non-canonical alias, dropping the route's specific-key delete leaves that stray key uncleaned. This only leaks an inert LiteLLM_VerificationToken row — no functional impact, and not the #439 orphan-causing-401, since the by-alias clear self-heals the canonical alias. Worth a follow-up at most; no action needed for this PR.

[RISK ASSESSMENT]

• [Overall PR] ⚠️ Risk Assessment: 🟡 MEDIUM

Touches managed-key rotation on the LLM completion/auth boundary and adds DB-level locking, so it is not trivially LOW. But the change is well-contained, additive, backward-compatible (only_if_current defaults to the prior forced-rotation behaviour), has no schema change and no API contract change, and is backed by real-Postgres concurrency tests plus a documented live reproduction. All current-head CI is green, including Python Tests on Linux (3.12). No human-only gate applies.

VERDICT:

✅ Approve — correct fix, right scope, real regression coverage.

KEY INSIGHT:

The fix pairs a per-org advisory lock (serializes rotations) with an idempotent only_if_current re-read (serialized-but-late refreshes reuse rather than re-rotate). Together they close both the "delete wipes a concurrently-minted key" race and the "second start rotates a key already handed to a sandbox" race — without a grace period that the LiteLLM unique-alias constraint would forbid.

---

Improve this review? If any feedback above seems incorrect or irrelevant to this repository, you can teach the reviewer to do better:

1. Add a .agents/skills/custom-codereview-guide.md file to your branch (or edit it if one already exists) with the /codereview trigger and the context the reviewer is missing (e.g., "Security concerns about X do not apply here because Y"). See the customization docs for the required frontmatter format.
2. Re-request a review - the reviewer reads guidelines from the PR branch, so your changes take effect immediately.
3. When your PR is merged, the guideline file goes through normal code review by repository maintainers.

Resolve with AI? Install the iterate skill in your agent and run /iterate to automatically drive this PR through CI, review, and QA until it's merge-ready.

Was this review helpful? React with 👍 or 👎 to give feedback.

---

This review was generated by an AI agent (OpenHands) on behalf of the requesting user.

[ak684]

Reviewed the advisory-lock scope, refreshed read under the lock, and only_if_current callers. The compare-after-lock path makes overlapping stale-key refreshes reuse the winner instead of deleting its replacement. Focused API-key, settings-rotation, and conversation-service tests: 276 passed. No blocking code findings at this head.

GitHub currently reports merge conflicts. Resolve those against main and rerun these checks before merging; this approval applies to the reviewed commit, not an unreviewed conflict resolution.

[ak684]

Resolved the conflicts with current main in 4f31fc2. Preserved main’s stale organization-key repair and this PR’s conditional/concurrent rotation logic.

Validation: 331 tests passed across managed-key rotation, live conversation status, API-key routes, and the settings store. Required pre-commit checks passed. The published Git tree matches the tested tree. CI is rerunning on the new commit.

[ak684]

Re-reviewed the conflict resolution at this head. Both main’s stale managed-key cleanup and this PR’s compare-after-lock rotation guard are preserved. 331 focused key-rotation, conversation-service, API-key and settings-store tests passed; required pre-commit passed. No blocking code findings. Merge once the remaining Python CI job passes.

[openhands-release-bot]

🚀 Released in 1.64.0.