fix(SKY-9507): invalidate org auth cache on update (#5826)

skyvern/forge/sdk/routes/agent_protocol.py

@@ -3277,7 +3277,7 @@ async def update_organization(
                 "artifact_url_expiry_seconds — pick one"
             ),
         )
-    return await app.DATABASE.organizations.update_organization(
+    updated = await app.DATABASE.organizations.update_organization(
         current_org.organization_id,
         max_steps_per_run=org_update.max_steps_per_run,
         max_retries_per_step=org_update.max_retries_per_step,
@@ -3286,6 +3286,9 @@ async def update_organization(
         clear_artifact_url_expiry_seconds=org_update.clear_artifact_url_expiry_seconds,
     )
 
+    org_auth_service.invalidate_cached_org(current_org.organization_id)
+    return updated
+
 
 @legacy_base_router.get(
     "/organizations",

skyvern/forge/sdk/services/org_auth_service.py

@@ -376,7 +376,10 @@ async def resolve_org_from_api_key(
     )
 
 
-@cached(cache=TTLCache(maxsize=CACHE_SIZE, ttl=AUTHENTICATION_TTL))
+_current_org_cache: TTLCache = TTLCache(maxsize=CACHE_SIZE, ttl=AUTHENTICATION_TTL)
+
+
+@cached(cache=_current_org_cache)
 async def _get_current_org_cached(x_api_key: str, db: AgentDB) -> Organization:
     """Authentication is cached for one hour."""
     validation = await resolve_org_from_api_key(x_api_key, db)
@@ -387,3 +390,16 @@ async def _get_current_org_cached(x_api_key: str, db: AgentDB) -> Organization:
         context.organization_id = validation.organization.organization_id
         context.organization_name = validation.organization.organization_name
     return validation.organization
+
+
+def invalidate_cached_org(organization_id: str) -> None:
+    """
+    Drop every cached ``Organization`` entry whose id matches.
+    """
+    keys_to_remove = [
+        key
+        for key, value in list(_current_org_cache.items())
+        if isinstance(value, Organization) and value.organization_id == organization_id
+    ]
+    for key in keys_to_remove:
+        _current_org_cache.pop(key, None)

tests/unit/test_org_auth_service.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from types import SimpleNamespace
 
 import jwt
@@ -6,6 +7,7 @@ from fastapi import HTTPException
 
 from skyvern.config import settings
 from skyvern.forge.sdk.core.security import create_access_token
+from skyvern.forge.sdk.schemas.organizations import Organization
 from skyvern.forge.sdk.services import org_auth_service
 from skyvern.forge.sdk.services.org_auth_service import (
     _get_api_key_debug_fields,
@@ -175,3 +177,50 @@ async def test_resolve_org_from_api_key_returns_403_when_diagnostic_helper_fails
 
     assert exc_info.value.status_code == 403
     assert warnings["diagnostic_error_type"] == "RuntimeError"
+
+
+def _make_org(organization_id: str, name: str = "test-org") -> Organization:
+    now = datetime.utcnow()
+    return Organization(
+        organization_id=organization_id,
+        organization_name=name,
+        created_at=now,
+        modified_at=now,
+    )
+
+
+def test_invalidate_cached_org_drops_only_matching_entries() -> None:
+    cache = org_auth_service._current_org_cache
+    cache.clear()
+    org_a = _make_org("org-a")
+    org_b = _make_org("org-b")
+    cache[("api-key-a", "db")] = org_a
+    cache[("api-key-a-rotated", "db")] = org_a
+    cache[("api-key-b", "db")] = org_b
+
+    org_auth_service.invalidate_cached_org("org-a")
+
+    assert ("api-key-a", "db") not in cache
+    assert ("api-key-a-rotated", "db") not in cache
+    assert ("api-key-b", "db") in cache
+    cache.clear()
+
+
+def test_invalidate_cached_org_is_noop_when_id_absent() -> None:
+    cache = org_auth_service._current_org_cache
+    cache.clear()
+    org = _make_org("org-a")
+    cache[("api-key", "db")] = org
+
+    org_auth_service.invalidate_cached_org("never-seen")
+
+    assert ("api-key", "db") in cache
+    cache.clear()
+
+
+def test_invalidate_cached_org_handles_empty_cache() -> None:
+    cache = org_auth_service._current_org_cache
+    cache.clear()
+
+    # Should not raise.
+    org_auth_service.invalidate_cached_org("anything")