diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index e43ea871e..869f29b1b 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/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", diff --git a/skyvern/forge/sdk/services/org_auth_service.py b/skyvern/forge/sdk/services/org_auth_service.py index e9e3a030c..d2d1b38e8 100644 --- a/skyvern/forge/sdk/services/org_auth_service.py +++ b/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) diff --git a/tests/unit/test_org_auth_service.py b/tests/unit/test_org_auth_service.py index 82997c872..eb5cdbb46 100644 --- a/tests/unit/test_org_auth_service.py +++ b/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")