diff --git a/.cortexkit/.gitignore b/.cortexkit/.gitignore new file mode 100644 index 0000000000..ee3fa486df --- /dev/null +++ b/.cortexkit/.gitignore @@ -0,0 +1,3 @@ +# >>> cortexkit:magic-context +magic-context/ +# <<< cortexkit:magic-context diff --git a/app/db/alembic/versions/20260713_080000_add_subagent_prompt_cache_ttl.py b/app/db/alembic/versions/20260713_080000_add_subagent_prompt_cache_ttl.py new file mode 100644 index 0000000000..0df55ac8e4 --- /dev/null +++ b/app/db/alembic/versions/20260713_080000_add_subagent_prompt_cache_ttl.py @@ -0,0 +1,65 @@ +"""add subagent prompt-cache TTL and sticky is_subagent flag + +Revision ID: 20260713_080000_add_subagent_prompt_cache_ttl +Revises: 20260717_000000_optimize_dashboard_hot_path_indexes +Create Date: 2026-07-13 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "20260713_080000_add_subagent_prompt_cache_ttl" +down_revision = "20260717_000000_optimize_dashboard_hot_path_indexes" +branch_labels = None +depends_on = None + + +def _columns(bind, table_name: str) -> set[str]: + inspector = sa.inspect(bind) + if not inspector.has_table(table_name): + return set() + return {column["name"] for column in inspector.get_columns(table_name)} + + +def upgrade() -> None: + bind = op.get_bind() + + settings_columns = _columns(bind, "dashboard_settings") + if settings_columns and "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds" not in settings_columns: + with op.batch_alter_table("dashboard_settings") as batch_op: + batch_op.add_column( + sa.Column( + "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds", + sa.Integer(), + nullable=True, + server_default=None, + ) + ) + + sticky_columns = _columns(bind, "sticky_sessions") + if sticky_columns and "is_subagent" not in sticky_columns: + with op.batch_alter_table("sticky_sessions") as batch_op: + batch_op.add_column( + sa.Column( + "is_subagent", + sa.Boolean(), + nullable=False, + server_default=sa.false(), + ) + ) + + +def downgrade() -> None: + bind = op.get_bind() + + settings_columns = _columns(bind, "dashboard_settings") + if settings_columns and "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds" in settings_columns: + with op.batch_alter_table("dashboard_settings") as batch_op: + batch_op.drop_column("http_responses_session_bridge_subagent_prompt_cache_ttl_seconds") + + sticky_columns = _columns(bind, "sticky_sessions") + if sticky_columns and "is_subagent" in sticky_columns: + with op.batch_alter_table("sticky_sessions") as batch_op: + batch_op.drop_column("is_subagent") diff --git a/app/db/models.py b/app/db/models.py index a1ec94065d..798da2d8e5 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -598,6 +598,12 @@ class StickySession(Base): onupdate=func.now(), nullable=False, ) + is_subagent: Mapped[bool] = mapped_column( + Boolean, + default=False, + server_default=false(), + nullable=False, + ) class DashboardSettings(Base): @@ -733,6 +739,12 @@ class DashboardSettings(Base): server_default=text("3600"), nullable=False, ) + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds: Mapped[int | None] = mapped_column( + Integer, + default=None, + server_default=None, + nullable=True, + ) http_responses_session_bridge_gateway_safe_mode: Mapped[bool] = mapped_column( Boolean, default=False, diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index 3cc36ef234..7cd45f05bf 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -1851,6 +1851,119 @@ def _http_bridge_previous_response_owner_unavailable_error() -> ProxyResponseErr ) +_SUBAGENT_HEADER_NAMES = ("x-parent-session-id", "x-openai-subagent", "x-codex-parent-thread-id") + + +def _detect_subagent_session( + headers: Mapping[str, str], + dashboard_settings: Any, +) -> tuple[bool, int | None]: + ttl = getattr( + dashboard_settings, + "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds", + None, + ) + if ttl is not None and ttl <= 0: + ttl = None + is_subagent = any(headers.get(name, "").strip() for name in _SUBAGENT_HEADER_NAMES) + return is_subagent, ttl + + +def _subagent_prompt_cache_bridge_key( + key: _HTTPBridgeSessionKey, + affinity: _AffinityPolicy, + is_subagent: bool, + forwarded_request: bool, + idle_ttl_seconds: float, + request_scope_id: str, + subagent_prompt_cache_ttl_seconds: int | None, +) -> tuple[_HTTPBridgeSessionKey, _AffinityPolicy, float, int | None]: + prompt_cache_affinity_key = affinity.key if affinity.kind is StickySessionKind.PROMPT_CACHE else None + is_prompt_cache = bool(prompt_cache_affinity_key) + effective_subagent_prompt_cache_ttl_seconds = subagent_prompt_cache_ttl_seconds if is_prompt_cache else None + if ( + not is_subagent + or forwarded_request + or key.affinity_key.startswith("subagent:") + or (affinity.key is not None and affinity.key.startswith("subagent:")) + ): + return key, affinity, idle_ttl_seconds, effective_subagent_prompt_cache_ttl_seconds + + if is_prompt_cache and subagent_prompt_cache_ttl_seconds is not None: + assert prompt_cache_affinity_key is not None + subagent_affinity_key = f"subagent:{prompt_cache_affinity_key}" + subagent_affinity_kind = StickySessionKind.PROMPT_CACHE.value + effective_idle_ttl_seconds = float(subagent_prompt_cache_ttl_seconds) + effective_max_age_seconds = subagent_prompt_cache_ttl_seconds + else: + subagent_affinity_key = f"subagent:{request_scope_id}:{prompt_cache_affinity_key or key.affinity_key}" + subagent_affinity_kind = key.affinity_kind + effective_idle_ttl_seconds = idle_ttl_seconds + effective_max_age_seconds = affinity.max_age_seconds + return ( + _HTTPBridgeSessionKey( + subagent_affinity_kind, + subagent_affinity_key, + key.api_key_id, + strength=None if subagent_affinity_kind == StickySessionKind.PROMPT_CACHE.value else key.strength, + ), + _AffinityPolicy( + key=subagent_affinity_key, + kind=affinity.kind, + reallocate_sticky=affinity.reallocate_sticky, + max_age_seconds=effective_max_age_seconds, + ), + effective_idle_ttl_seconds, + effective_subagent_prompt_cache_ttl_seconds, + ) + + +async def _delete_completed_subagent_sticky_mapping( + repo_factory: Callable[..., Any], + session: Any, +) -> None: + if not ( + getattr(session, "is_subagent", False) + and getattr(session, "subagent_prompt_cache_ttl_seconds", None) is None + and getattr(getattr(session, "affinity", None), "kind", None) is StickySessionKind.PROMPT_CACHE + ): + return + affinity = getattr(session, "affinity", None) + if affinity is None: + return + try: + async with repo_factory() as repositories: + keys_to_delete = {affinity.key} + if isinstance(affinity.key, str) and affinity.key.startswith("subagent:"): + parts = affinity.key.split(":", 2) + if len(parts) == 3 and parts[2]: + keys_to_delete.add(f"subagent:{parts[2]}") + for key in keys_to_delete: + await repositories.sticky_sessions.delete( + key, + kind=StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ) + except Exception: + logger.warning("Failed to delete completed subagent sticky mapping", exc_info=True) + + +def _sync_reused_http_bridge_session( + session: Any, + metadata: tuple[Any, str | None, str | None, float, bool, int | None], +) -> None: + api_key, request_model, request_service_tier, idle_ttl_seconds, is_subagent, subagent_ttl = metadata + session.api_key = api_key + session.request_model = request_model + session.request_service_tier = request_service_tier + session.last_used_at = _service_time().monotonic() + if not (getattr(session, "is_subagent", False) or is_subagent): + return + session.idle_ttl_seconds = idle_ttl_seconds + session.is_subagent = True + session.subagent_prompt_cache_ttl_seconds = subagent_ttl + + def _http_bridge_should_attempt_local_previous_response_recovery(exc: ProxyResponseError) -> bool: payload = exc.payload if not isinstance(payload, dict): @@ -2182,6 +2295,9 @@ def _wrapper(*args: Any, **kwargs: Any) -> Any: "_http_bridge_previous_response_error_envelope", "_http_bridge_continuity_lost_error_envelope", "_http_bridge_owner_lookup_unavailable_error_envelope", + "_detect_subagent_session", + "_subagent_prompt_cache_bridge_key", + "_delete_completed_subagent_sticky_mapping", "_http_bridge_should_attempt_local_previous_response_recovery", "_http_bridge_is_previous_response_owner_unavailable", "_http_bridge_should_attempt_soft_affinity_reroute", diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index a1f2c125bb..b104bc3fab 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -71,6 +71,8 @@ _HTTP_BRIDGE_INFLIGHT_STARTED_AT_ATTR, _active_http_bridge_instance_ring, _close_http_bridge_session_bounded, + _delete_completed_subagent_sticky_mapping, + _detect_subagent_session, _durable_bridge_lookup_active_owner, _durable_bridge_lookup_allows_local_reuse, _forwarded_http_bridge_session_key, @@ -124,6 +126,8 @@ _renew_durable_http_bridge_lease, _require_http_bridge_bound_account_not_excluded, _reserve_http_bridge_unanchored_handoff, + _subagent_prompt_cache_bridge_key, + _sync_reused_http_bridge_session, _track_alias_registration, ) from app.modules.proxy._service.http_bridge.owner_forwarding import _HTTPBridgeOwnerForwardingMixin @@ -252,20 +256,10 @@ class _HTTPBridgeMixin( _HTTPBridgeUpstreamEventsMixin, _HTTPBridgeServiceProtocol, ): - async def _close_http_bridge_session_bounded( - self, - session: "_HTTPBridgeSession", - *, - reason: str, - ) -> None: + async def _close_http_bridge_session_bounded(self, session: "_HTTPBridgeSession", *, reason: str) -> None: await _close_http_bridge_session_bounded(self, session, reason=reason) - def _schedule_http_bridge_session_closes( - self, - sessions: list["_HTTPBridgeSession"], - *, - reason: str, - ) -> None: + def _schedule_http_bridge_session_closes(self, sessions: list["_HTTPBridgeSession"], *, reason: str) -> None: for session in sessions: if len(self._background_cleanup_tasks) >= _HTTP_BRIDGE_BACKGROUND_CLEANUP_WARN_THRESHOLD: logger.warning( @@ -375,7 +369,6 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, ) -> "_HTTPBridgeSession": ... - @overload async def _get_or_create_http_bridge_session( self, @@ -434,20 +427,41 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, ) -> "_HTTPBridgeSession | _HTTPBridgeOwnerForward": - settings = _service_get_settings() + settings, dashboard_settings = _service_get_settings(), await _service_get_settings_cache().get() request_scope_id = ensure_request_scope_id() api_key_id = api_key.id if api_key is not None else None incoming_turn_state = _sticky_key_from_turn_state_header(headers) incoming_session_key = _sticky_key_from_session_header(headers) initial_session_key = session_header_fallback_key or (key if key.affinity_kind == "session_header" else None) + original_key, original_affinity_key = key, affinity.key original_request_unanchored = _http_bridge_request_needs_unanchored_handoff( key, incoming_turn_state, previous_response_id, forwarded_request, forwarded_original_request_unanchored ) + is_subagent_session, subagent_prompt_cache_ttl_seconds = _detect_subagent_session(headers, dashboard_settings) + subagent_key_result = _subagent_prompt_cache_bridge_key( + key, + affinity, + is_subagent_session, + forwarded_request, + idle_ttl_seconds, + request_scope_id, + subagent_prompt_cache_ttl_seconds, + ) + key, affinity, effective_idle_ttl_seconds, subagent_prompt_cache_ttl_seconds = subagent_key_result + if key != original_key or affinity.key != original_affinity_key: + incoming_turn_state = incoming_session_key = initial_session_key = None + reuse_metadata = ( + api_key, + request_model, + request_service_tier, + effective_idle_ttl_seconds, + is_subagent_session, + subagent_prompt_cache_ttl_seconds, + ) model_transition_rebind = bool( durable_lookup is not None and not _http_bridge_models_compatible(durable_lookup.model, request_model) ) - if model_transition_rebind: - durable_lookup = None + durable_lookup = None if model_transition_rebind else durable_lookup if await _http_bridge_should_wait_for_registration(self, key, settings): skip_registration_gate = False async with self._http_bridge_lock: @@ -474,7 +488,6 @@ async def _get_or_create_http_bridge_session( error_type="server_error", ), ) - effective_idle_ttl_seconds = idle_ttl_seconds forwarded_affinity = ( _forwarded_http_bridge_session_key( headers, @@ -652,10 +665,7 @@ async def _get_or_create_http_bridge_session( assert existing is not None current_instance = settings.http_responses_session_bridge_instance_id if _durable_bridge_lookup_allows_local_reuse(durable_lookup, current_instance=current_instance): - existing.api_key = api_key - existing.request_model = request_model - existing.request_service_tier = request_service_tier - existing.last_used_at = _service_time().monotonic() + _sync_reused_http_bridge_session(existing, reuse_metadata) await _refresh_reused_http_bridge_session_with_handoff( self, existing, @@ -1285,7 +1295,6 @@ async def _get_or_create_http_bridge_session( except Exception: pass continue - if inflight_future is not None and not owns_creation: wait_timeout_seconds = _proxy_admission_wait_timeout_seconds(settings) try: @@ -1355,10 +1364,7 @@ async def _get_or_create_http_bridge_session( ): current_instance = settings.http_responses_session_bridge_instance_id if _durable_bridge_lookup_allows_local_reuse(durable_lookup, current_instance=current_instance): - session.api_key = api_key - session.request_model = request_model - session.request_service_tier = request_service_tier - session.last_used_at = _service_time().monotonic() + _sync_reused_http_bridge_session(session, reuse_metadata) return session if not session.closed and session.account.status == AccountStatus.ACTIVE: old_account_id = session.account.id @@ -1374,7 +1380,6 @@ async def _get_or_create_http_bridge_session( if detached is not None and not retiring_with_visible_requests: self._schedule_http_bridge_session_closes([detached], reason="registry_detach") continue - created_session: _HTTPBridgeSession | None = None session_registered = False try: @@ -1386,6 +1391,8 @@ async def _get_or_create_http_bridge_session( "request_model": request_model, "request_service_tier": request_service_tier, "idle_ttl_seconds": effective_idle_ttl_seconds, + "is_subagent": is_subagent_session, + "subagent_prompt_cache_ttl_seconds": subagent_prompt_cache_ttl_seconds, "request_stage": request_stage, "preferred_account_id": preferred_account_id, "require_preferred_account": require_preferred_account, @@ -1404,7 +1411,13 @@ async def _get_or_create_http_bridge_session( for parameter in create_signature.parameters.values() ) if create_signature is not None and not create_accepts_var_keyword: - for optional_kwarg in ("request_service_tier", "request_usage_budget", "request_deadline"): + for optional_kwarg in ( + "request_service_tier", + "request_usage_budget", + "request_deadline", + "is_subagent", + "subagent_prompt_cache_ttl_seconds", + ): if optional_kwarg not in create_signature.parameters: create_kwargs.pop(optional_kwarg, None) created_session = await create_session(key, **create_kwargs) @@ -1598,6 +1611,7 @@ async def _close_http_bridge_session( api_key=None, response_create_gate=response_create_gate, ) + await _delete_completed_subagent_sticky_mapping(self._repo_factory, session) _log_http_bridge_event( "close", session.key, @@ -1817,12 +1831,10 @@ async def _claim_durable_http_bridge_session( async def _refresh_durable_http_bridge_session(self, session: "_HTTPBridgeSession") -> None: """Renew the durable lease; callers must hold ``self._http_bridge_lock``.""" - await _renew_durable_http_bridge_lease(self, session) async def reconcile_durable_http_bridge_ownership(self) -> int: """Close local sessions whose durable row is owned by another instance/epoch.""" - return await _reconcile_durable_http_bridge_ownership(self) async def _create_http_bridge_session( @@ -1841,6 +1853,8 @@ async def _create_http_bridge_session( fallback_on_preferred_account_unavailable: bool = True, request_usage_budget: ApiKeyRequestUsageBudget | None = None, request_deadline: float | None = None, + is_subagent: bool = False, + subagent_prompt_cache_ttl_seconds: int | None = None, ) -> "_HTTPBridgeSession": request_state = _WebSocketRequestState( request_id=f"http_bridge_connect_{uuid4().hex}", @@ -1865,12 +1879,23 @@ async def _create_http_bridge_session( preferred_candidate_id = preferred_account_id selected_account_lease: AccountLease | None = None while True: + subagent_no_cache = is_subagent and subagent_prompt_cache_ttl_seconds is None select_kwargs = { "request_id": request_state.request_log_id or request_state.request_id, "kind": "http_bridge", "request_stage": request_stage, "api_key": api_key, - "affinity_policy": affinity, + "sticky_key": None if subagent_no_cache else affinity.selection_key, + "sticky_kind": None if subagent_no_cache else affinity.kind, + "reallocate_sticky": affinity.reallocate_sticky, + "sticky_source": None if subagent_no_cache else affinity.codex_session_source, + "legacy_sticky_key": None if subagent_no_cache else affinity.legacy_selection_key, + "spill_bare_session_on_account_cap": False if subagent_no_cache else affinity.spill_on_account_cap, + "require_unambiguous_account": False if subagent_no_cache else affinity.require_unambiguous_account, + "sticky_max_age_seconds": subagent_prompt_cache_ttl_seconds + if is_subagent + else affinity.max_age_seconds, + "sticky_is_subagent": is_subagent, "prefer_earlier_reset_accounts": settings.prefer_earlier_reset_accounts, "prefer_earlier_reset_window": _prefer_earlier_reset_window(settings), "routing_strategy": _routing_strategy(settings), @@ -2087,6 +2112,8 @@ async def _create_http_bridge_session( downstream_turn_state=None, account_lease=selected_account_lease, catalog_omission_quota_admission=selection.catalog_omission_quota_admission, + is_subagent=is_subagent, + subagent_prompt_cache_ttl_seconds=subagent_prompt_cache_ttl_seconds, ) _copy_websocket_route_metadata_to_session(session, request_state) session.upstream_reader = asyncio.create_task(self._relay_http_bridge_upstream_messages(session)) diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index f503f89e5b..ea1e73367e 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -65,6 +65,7 @@ _sticky_key_from_compact_payload as _sticky_key_from_compact_payload, ) from app.modules.proxy._service.http_bridge.helpers import ( + _detect_subagent_session, _effective_http_bridge_idle_ttl_seconds, _http_bridge_durable_lookup_allows_turn_state_takeover, _http_bridge_is_context_overflow_error, @@ -90,6 +91,7 @@ _release_http_bridge_unanchored_handoff, _release_http_bridge_unanchored_handoffs_for_request, _reserve_http_bridge_unanchored_handoff, + _subagent_prompt_cache_bridge_key, _trim_http_bridge_previous_response_input_items, ) from app.modules.proxy._service.http_bridge.service_stubs import ( @@ -472,7 +474,30 @@ async def _registered_turn_state_anchor_lookup( return lookup +def _http_bridge_session_is_unanchored_parallel_fork(session: Any) -> bool: + return session.key.affinity_kind == "internal_unanchored_parallel" + + class _HTTPBridgeStreamingMixin: + async def _refresh_completed_subagent_sticky_mapping(self, session: _HTTPBridgeSession) -> None: + if not ( + session.is_subagent + and session.subagent_prompt_cache_ttl_seconds is not None + and session.affinity.kind is StickySessionKind.PROMPT_CACHE + and session.affinity.key + ): + return + try: + async with cast(Any, self)._repo_factory() as repositories: + await repositories.sticky_sessions.upsert( + session.affinity.key, + session.account.id, + kind=StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ) + except Exception: + logger.warning("Failed to refresh completed subagent sticky mapping", exc_info=True) + async def validate_http_bridge_legacy_forward_anchor( self: Any, *, @@ -767,6 +792,34 @@ def prepare_bridge_request( forwarded_affinity_kind=forwarded_affinity_kind, forwarded_affinity_key=forwarded_affinity_key, ) + request_scope_id = ensure_request_scope_id() + is_subagent_session, subagent_prompt_cache_ttl_seconds = _detect_subagent_session( + headers, + dashboard_settings, + ) + original_bridge_session_key = bridge_session_key + original_affinity_key = affinity.key + ( + bridge_session_key, + affinity, + _effective_lookup_idle_ttl_seconds, + _effective_subagent_prompt_cache_ttl_seconds, + ) = _subagent_prompt_cache_bridge_key( + bridge_session_key, + affinity, + is_subagent_session, + forwarded_request, + idle_ttl_seconds, + request_scope_id, + subagent_prompt_cache_ttl_seconds, + ) + del _effective_lookup_idle_ttl_seconds, _effective_subagent_prompt_cache_ttl_seconds + subagent_bridge_key_isolated = ( + bridge_session_key != original_bridge_session_key or affinity.key != original_affinity_key + ) + if subagent_bridge_key_isolated: + incoming_turn_state_header = None + incoming_session_header = None session_header_fallback_key = ( _make_http_bridge_session_header_fallback_key( headers=headers, @@ -776,10 +829,12 @@ def prepare_bridge_request( if not forwarded_request else None ) + if subagent_bridge_key_isolated: + session_header_fallback_key = None legacy_anchor_lookup = await _legacy_forward_anchor_lookup( durable_bridge=self._durable_bridge, bridge_session_key=bridge_session_key, - turn_state=_sticky_key_from_turn_state_header(headers), + turn_state=incoming_turn_state_header, api_key=api_key, previous_response_id=payload.previous_response_id, forwarded_request=forwarded_request, @@ -789,7 +844,7 @@ def prepare_bridge_request( incoming_turn_state_header = _sticky_key_from_turn_state_header(headers) original_request_unanchored = _http_bridge_request_needs_unanchored_handoff( bridge_session_key, - _sticky_key_from_turn_state_header(headers), + incoming_turn_state_header, payload.previous_response_id, forwarded_request, forwarded_original_request_unanchored, @@ -2268,3 +2323,9 @@ async def _stream_http_bridge_session_events( with anyio.CancelScope(shield=True): await self._detach_http_bridge_request(session, request_state=request_state) session.last_used_at = _service_time().monotonic() + if _http_bridge_session_is_unanchored_parallel_fork(session): + await self._close_http_bridge_session(session) + elif session.is_subagent: + if session.subagent_prompt_cache_ttl_seconds is not None: + await self._refresh_completed_subagent_sticky_mapping(session) + await self._close_http_bridge_session(session) diff --git a/app/modules/proxy/_service/response_create.py b/app/modules/proxy/_service/response_create.py index 8795115449..9686cd454b 100644 --- a/app/modules/proxy/_service/response_create.py +++ b/app/modules/proxy/_service/response_create.py @@ -49,6 +49,7 @@ _OVERSIZED_RESPONSE_CREATE_DUMP_DIR: Path | None = None _RESPONSE_CREATE_COMPATIBILITY_METADATA_HEADERS = ( "x-codex-turn-metadata", + "x-parent-session-id", "x-openai-subagent", "x-codex-parent-thread-id", "x-codex-window-id", diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index a019789f91..44c04059c5 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -886,6 +886,8 @@ class _HTTPBridgeSession: upstream_proxy_endpoint_id: str | None = None upstream_proxy_fallback_used: bool | None = None upstream_proxy_fail_closed_reason: str | None = None + is_subagent: bool = False + subagent_prompt_cache_ttl_seconds: int | None = None def _http_bridge_session_supports_service_tier( diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index 2980d1987c..e3d7d71871 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -364,6 +364,7 @@ async def select_account( spill_bare_session_on_account_cap: bool = False, require_unambiguous_account: bool = False, sticky_max_age_seconds: int | None = None, + sticky_is_subagent: bool = False, prefer_earlier_reset_accounts: bool = False, prefer_earlier_reset_window: ResetPreferenceWindow = "secondary", routing_strategy: RoutingStrategy = "capacity_weighted", @@ -898,6 +899,7 @@ async def load_selection_inputs() -> _SelectionInputs: sticky_kind=sticky_kind, reallocate_sticky=reallocate_sticky, sticky_max_age_seconds=sticky_max_age_seconds, + sticky_is_subagent=sticky_is_subagent, budget_threshold_pct=budget_threshold_pct, secondary_budget_threshold_pct=secondary_budget_threshold_pct, prefer_earlier_reset_accounts=prefer_earlier_reset_accounts, @@ -1559,6 +1561,7 @@ async def _select_with_stickiness( sticky_kind: StickySessionKind | None, reallocate_sticky: bool, sticky_max_age_seconds: int | None, + sticky_is_subagent: bool = False, budget_threshold_pct: float = 95.0, secondary_budget_threshold_pct: float = 100.0, prefer_earlier_reset_accounts: bool, @@ -1588,6 +1591,10 @@ async def _select_with_stickiness( ) if sticky_kind is None: raise ValueError("sticky_kind is required when sticky_key is provided") + sticky_upsert_kwargs = { + "kind": sticky_kind, + "is_subagent": sticky_is_subagent, + } if sticky_existing_account_id is _STICKY_EXISTING_UNSET: existing = await sticky_repo.get_account_id( @@ -1671,7 +1678,11 @@ async def _select_with_stickiness( ) if pinned_result.account is not None: if sticky_max_age_seconds is not None: - await sticky_repo.upsert(sticky_key, pinned.account_id, kind=sticky_kind) + await sticky_repo.upsert( + sticky_key, + pinned.account_id, + **sticky_upsert_kwargs, + ) return pinned_result else: # Reallocate only when a burn-first target exists and can @@ -1725,7 +1736,7 @@ async def _select_with_stickiness( await sticky_repo.upsert( sticky_key, pinned.account_id, - kind=sticky_kind, + **sticky_upsert_kwargs, ) return pinned_result reallocate_sticky = True @@ -1752,7 +1763,11 @@ async def _select_with_stickiness( ) if grace_result.account is not None: if sticky_max_age_seconds is not None: - await sticky_repo.upsert(sticky_key, pinned.account_id, kind=sticky_kind) + await sticky_repo.upsert( + sticky_key, + pinned.account_id, + **sticky_upsert_kwargs, + ) return grace_result if reallocate_sticky: await sticky_repo.delete(sticky_key, kind=sticky_kind) @@ -1788,7 +1803,11 @@ async def _select_with_stickiness( routing_costs_by_account_id=routing_costs_by_account_id, ) if persist_fallback and chosen.account is not None and chosen.account.account_id in account_map: - await sticky_repo.upsert(sticky_key, chosen.account.account_id, kind=sticky_kind) + await sticky_repo.upsert( + sticky_key, + chosen.account.account_id, + **sticky_upsert_kwargs, + ) elif preserve_existing_mapping_on_fallback and chosen.account is not None and existing is not None: # Spillover is deliberately request-local. The alternate may create # its own hard response/file/bridge owner, but local cap pressure diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index 88be05119d..b5c8089495 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -1701,6 +1701,7 @@ async def _select_account_with_budget( spill_bare_session_on_account_cap: bool = False, require_unambiguous_account: bool = False, sticky_max_age_seconds: int | None = None, + sticky_is_subagent: bool = False, prefer_earlier_reset_accounts: bool = False, prefer_earlier_reset_window: ResetPreferenceWindow = "secondary", routing_strategy: RoutingStrategy = "capacity_weighted", @@ -1832,6 +1833,7 @@ async def _select_account_with_budget( sticky_max_age_seconds=preferred_sticky_inputs[3], sticky_source=preferred_sticky_inputs[4], legacy_sticky_key=preferred_sticky_inputs[5], + sticky_is_subagent=sticky_is_subagent, prefer_earlier_reset_accounts=prefer_earlier_reset_accounts, prefer_earlier_reset_window=prefer_earlier_reset_window, routing_strategy=routing_strategy, @@ -1886,6 +1888,7 @@ async def _select_account_with_budget( ), require_unambiguous_account=require_unambiguous_account, sticky_max_age_seconds=sticky_max_age_seconds, + sticky_is_subagent=sticky_is_subagent, prefer_earlier_reset_accounts=prefer_earlier_reset_accounts, prefer_earlier_reset_window=prefer_earlier_reset_window, routing_strategy=routing_strategy, diff --git a/app/modules/proxy/sticky_repository.py b/app/modules/proxy/sticky_repository.py index 83080fbba3..1253d64fe9 100644 --- a/app/modules/proxy/sticky_repository.py +++ b/app/modules/proxy/sticky_repository.py @@ -11,7 +11,7 @@ from sqlalchemy.sql import Insert from app.core.utils.time import to_utc_naive, utcnow -from app.db.models import Account, StickySession, StickySessionKind +from app.db.models import Account, HttpBridgeSessionRecord, HttpBridgeSessionState, StickySession, StickySessionKind from app.db.session import sqlite_writer_section from app.modules.sticky_sessions.schemas import StickySessionSortBy, StickySessionSortDir @@ -64,13 +64,25 @@ async def get_entry(self, key: str, *, kind: StickySessionKind) -> StickySession result = await self._session.execute(statement) return result.scalar_one_or_none() - async def upsert(self, key: str, account_id: str, *, kind: StickySessionKind) -> StickySession: + async def upsert( + self, + key: str, + account_id: str, + *, + kind: StickySessionKind, + is_subagent: bool = False, + ) -> StickySession: # RETURNING collapses the previous upsert + re-select + refresh # (4 round trips) into one statement; this runs inline before the # first upstream byte on sticky requests, so round trips are TTFT. # populate_existing forces the returned row to overwrite any stale # identity-map instance the session may already hold for this key. - statement = self._build_upsert_statement(key, account_id, kind).returning(StickySession) + statement = self._build_upsert_statement( + key, + account_id, + kind, + is_subagent=is_subagent, + ).returning(StickySession) async with sqlite_writer_section(): result = await self._session.execute(statement, execution_options={"populate_existing": True}) row = result.scalar_one_or_none() @@ -79,13 +91,21 @@ async def upsert(self, key: str, account_id: str, *, kind: StickySessionKind) -> raise RuntimeError(f"StickySession upsert failed for key={key!r} kind={kind.value!r}") return row - async def delete(self, key: str, *, kind: StickySessionKind) -> bool: + async def delete( + self, + key: str, + *, + kind: StickySessionKind, + is_subagent: bool | None = None, + ) -> bool: if not key: return False statement = delete(StickySession).where( StickySession.key == key, StickySession.kind == kind, ) + if is_subagent is not None: + statement = statement.where(StickySession.is_subagent == is_subagent) async with sqlite_writer_section(): result = await self._session.execute(statement.returning(StickySession.key)) await self._session.commit() @@ -117,6 +137,7 @@ async def list_entry_identifiers( *, kind: StickySessionKind | None = None, updated_before: datetime | None = None, + is_subagent: bool | None = None, account_query: str | None = None, key_query: str | None = None, ) -> list[tuple[str, StickySessionKind]]: @@ -125,6 +146,7 @@ async def list_entry_identifiers( select(StickySession.key, StickySession.kind), kind=kind, updated_before=updated_before, + is_subagent=is_subagent, account_query=account_query, key_query=key_query, ) @@ -143,6 +165,7 @@ async def list_entries( *, kind: StickySessionKind | None = None, updated_before: datetime | None = None, + is_subagent: bool | None = None, account_query: str | None = None, key_query: str | None = None, sort_by: StickySessionSortBy = "updated_at", @@ -156,6 +179,7 @@ async def list_entries( select(StickySession, Account.email), kind=kind, updated_before=updated_before, + is_subagent=is_subagent, account_query=account_query, key_query=key_query, ) @@ -177,6 +201,7 @@ async def count_entries( *, kind: StickySessionKind | None = None, updated_before: datetime | None = None, + is_subagent: bool | None = None, account_query: str | None = None, key_query: str | None = None, ) -> int: @@ -184,26 +209,78 @@ async def count_entries( select(func.count()).select_from(StickySession).join(Account, Account.id == StickySession.account_id), kind=kind, updated_before=updated_before, + is_subagent=is_subagent, account_query=account_query, key_query=key_query, ) result = await self._session.execute(statement) return int(result.scalar_one()) - async def purge_prompt_cache_before(self, cutoff: datetime) -> int: - return await self.purge_before(cutoff, kind=StickySessionKind.PROMPT_CACHE) + async def purge_prompt_cache_before( + self, + cutoff: datetime, + *, + is_subagent: bool | None = None, + protect_active_bridge_mappings: bool = False, + ) -> int: + if protect_active_bridge_mappings: + return await self._purge_prompt_cache_before_excluding_active_bridge(cutoff, is_subagent=is_subagent) + return await self.purge_before(cutoff, kind=StickySessionKind.PROMPT_CACHE, is_subagent=is_subagent) - async def purge_before(self, cutoff: datetime, *, kind: StickySessionKind | None = None) -> int: + async def _purge_prompt_cache_before_excluding_active_bridge( + self, + cutoff: datetime, + *, + is_subagent: bool | None, + ) -> int: + active_bridge_mapping = ( + select(HttpBridgeSessionRecord.id) + .where( + HttpBridgeSessionRecord.session_key_kind == StickySessionKind.PROMPT_CACHE.value, + HttpBridgeSessionRecord.session_key_value == StickySession.key, + HttpBridgeSessionRecord.state.in_((HttpBridgeSessionState.ACTIVE, HttpBridgeSessionState.DRAINING)), + ) + .exists() + ) + stmt = delete(StickySession).where( + StickySession.updated_at < to_utc_naive(cutoff), + StickySession.kind == StickySessionKind.PROMPT_CACHE, + ~active_bridge_mapping, + ) + if is_subagent is not None: + stmt = stmt.where(StickySession.is_subagent == is_subagent) + async with sqlite_writer_section(): + result = await self._session.execute(stmt.returning(StickySession.key)) + deleted = len(result.scalars().all()) + await self._session.commit() + return deleted + + async def purge_before( + self, + cutoff: datetime, + *, + kind: StickySessionKind | None = None, + is_subagent: bool | None = None, + ) -> int: stmt = delete(StickySession).where(StickySession.updated_at < to_utc_naive(cutoff)) if kind is not None: stmt = stmt.where(StickySession.kind == kind) + if is_subagent is not None: + stmt = stmt.where(StickySession.is_subagent == is_subagent) async with sqlite_writer_section(): result = await self._session.execute(stmt.returning(StickySession.key)) deleted = len(result.scalars().all()) await self._session.commit() return deleted - def _build_upsert_statement(self, key: str, account_id: str, kind: StickySessionKind) -> Insert: + def _build_upsert_statement( + self, + key: str, + account_id: str, + kind: StickySessionKind, + *, + is_subagent: bool, + ) -> Insert: dialect = self._session.get_bind().dialect.name if dialect == "postgresql": insert_fn = pg_insert @@ -211,11 +288,17 @@ def _build_upsert_statement(self, key: str, account_id: str, kind: StickySession insert_fn = sqlite_insert else: raise RuntimeError(f"StickySession upsert unsupported for dialect={dialect!r}") - statement = insert_fn(StickySession).values(key=key, account_id=account_id, kind=kind) + statement = insert_fn(StickySession).values( + key=key, + account_id=account_id, + kind=kind, + is_subagent=is_subagent, + ) return statement.on_conflict_do_update( index_elements=[StickySession.key, StickySession.kind], set_={ "account_id": account_id, + "is_subagent": is_subagent, "updated_at": func.now(), }, ) @@ -226,6 +309,7 @@ def _apply_filters( *, kind: StickySessionKind | None, updated_before: datetime | None, + is_subagent: bool | None, account_query: str | None, key_query: str | None, ): @@ -233,6 +317,8 @@ def _apply_filters( statement = statement.where(StickySession.kind == kind) if updated_before is not None: statement = statement.where(StickySession.updated_at < to_utc_naive(updated_before)) + if is_subagent is not None: + statement = statement.where(StickySession.is_subagent == is_subagent) if account_query: statement = statement.where(func.lower(Account.email).contains(account_query.lower())) if key_query: diff --git a/app/modules/settings/api.py b/app/modules/settings/api.py index b217b45c87..8ecb9ecb96 100644 --- a/app/modules/settings/api.py +++ b/app/modules/settings/api.py @@ -143,6 +143,7 @@ def _dashboard_settings_response(settings) -> DashboardSettingsResponse: openai_cache_affinity_max_age_seconds=settings.openai_cache_affinity_max_age_seconds, dashboard_session_ttl_seconds=settings.dashboard_session_ttl_seconds, http_responses_session_bridge_prompt_cache_idle_ttl_seconds=settings.http_responses_session_bridge_prompt_cache_idle_ttl_seconds, + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=settings.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds, http_responses_session_bridge_gateway_safe_mode=settings.http_responses_session_bridge_gateway_safe_mode, sticky_reallocation_budget_threshold_pct=settings.sticky_reallocation_budget_threshold_pct, sticky_reallocation_primary_budget_threshold_pct=settings.sticky_reallocation_primary_budget_threshold_pct, @@ -700,6 +701,11 @@ async def update_settings( if payload.http_responses_session_bridge_prompt_cache_idle_ttl_seconds is not None else current.http_responses_session_bridge_prompt_cache_idle_ttl_seconds ), + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=( + payload.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds + if "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds" in payload.model_fields_set + else current.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds + ), http_responses_session_bridge_gateway_safe_mode=( payload.http_responses_session_bridge_gateway_safe_mode if payload.http_responses_session_bridge_gateway_safe_mode is not None @@ -840,6 +846,7 @@ async def update_settings( "openai_cache_affinity_max_age_seconds", "dashboard_session_ttl_seconds", "http_responses_session_bridge_prompt_cache_idle_ttl_seconds", + "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds", "http_responses_session_bridge_gateway_safe_mode", "sticky_reallocation_budget_threshold_pct", "sticky_reallocation_primary_budget_threshold_pct", diff --git a/app/modules/settings/repository.py b/app/modules/settings/repository.py index d2dc08786c..6a944e78d8 100644 --- a/app/modules/settings/repository.py +++ b/app/modules/settings/repository.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import cast + from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified @@ -11,6 +13,7 @@ from app.db.models import DashboardSettings _SETTINGS_ID = 1 +_UNSET = object() class SettingsRepository: @@ -63,6 +66,7 @@ async def get_or_create(self) -> DashboardSettings: limit_warmup_windows="both", limit_warmup_model="auto", limit_warmup_prompt="Say OK.", + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=None, limit_warmup_cooldown_seconds=3600, limit_warmup_exhausted_threshold_percent=99.0, limit_warmup_idle_threshold_percent=1.0, @@ -109,6 +113,7 @@ async def update( openai_cache_affinity_max_age_seconds: int | None = None, dashboard_session_ttl_seconds: int | None = None, http_responses_session_bridge_prompt_cache_idle_ttl_seconds: int | None = None, + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds: int | None | object = _UNSET, http_responses_session_bridge_gateway_safe_mode: bool | None = None, sticky_reallocation_budget_threshold_pct: float | None = None, sticky_reallocation_primary_budget_threshold_pct: float | None = None, @@ -189,6 +194,10 @@ async def update( settings.http_responses_session_bridge_prompt_cache_idle_ttl_seconds = ( http_responses_session_bridge_prompt_cache_idle_ttl_seconds ) + if http_responses_session_bridge_subagent_prompt_cache_ttl_seconds is not _UNSET: + settings.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds = cast( + int | None, http_responses_session_bridge_subagent_prompt_cache_ttl_seconds + ) if http_responses_session_bridge_gateway_safe_mode is not None: settings.http_responses_session_bridge_gateway_safe_mode = http_responses_session_bridge_gateway_safe_mode if sticky_reallocation_budget_threshold_pct is not None: diff --git a/app/modules/settings/schemas.py b/app/modules/settings/schemas.py index 52ed53fbf1..a776d0eefb 100644 --- a/app/modules/settings/schemas.py +++ b/app/modules/settings/schemas.py @@ -55,6 +55,7 @@ class DashboardSettingsResponse(DashboardModel): openai_cache_affinity_max_age_seconds: int = Field(gt=0) dashboard_session_ttl_seconds: int = Field(ge=3600) http_responses_session_bridge_prompt_cache_idle_ttl_seconds: int = Field(gt=0) + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds: int | None = Field(default=None, ge=0) http_responses_session_bridge_gateway_safe_mode: bool sticky_reallocation_budget_threshold_pct: float = Field(ge=0.0, le=100.0) sticky_reallocation_primary_budget_threshold_pct: float = Field(ge=0.0, le=100.0) @@ -119,6 +120,7 @@ class DashboardSettingsUpdateRequest(DashboardModel): openai_cache_affinity_max_age_seconds: int | None = Field(default=None, gt=0) dashboard_session_ttl_seconds: int | None = Field(default=None, ge=3600) http_responses_session_bridge_prompt_cache_idle_ttl_seconds: int | None = Field(default=None, gt=0) + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds: int | None = Field(default=None, ge=0) http_responses_session_bridge_gateway_safe_mode: bool | None = None sticky_reallocation_budget_threshold_pct: float | None = Field(default=None, ge=0.0, le=100.0) sticky_reallocation_primary_budget_threshold_pct: float | None = Field(default=None, ge=0.0, le=100.0) diff --git a/app/modules/settings/service.py b/app/modules/settings/service.py index cb75820462..4b33bb9a4c 100644 --- a/app/modules/settings/service.py +++ b/app/modules/settings/service.py @@ -34,6 +34,7 @@ class DashboardSettingsData: openai_cache_affinity_max_age_seconds: int dashboard_session_ttl_seconds: int http_responses_session_bridge_prompt_cache_idle_ttl_seconds: int + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds: int | None http_responses_session_bridge_gateway_safe_mode: bool sticky_reallocation_budget_threshold_pct: float sticky_reallocation_primary_budget_threshold_pct: float @@ -88,6 +89,7 @@ class DashboardSettingsUpdateData: openai_cache_affinity_max_age_seconds: int dashboard_session_ttl_seconds: int http_responses_session_bridge_prompt_cache_idle_ttl_seconds: int + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds: int | None http_responses_session_bridge_gateway_safe_mode: bool sticky_reallocation_budget_threshold_pct: float sticky_reallocation_primary_budget_threshold_pct: float @@ -152,6 +154,7 @@ async def get_settings(self) -> DashboardSettingsData: http_responses_session_bridge_prompt_cache_idle_ttl_seconds=( row.http_responses_session_bridge_prompt_cache_idle_ttl_seconds ), + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=row.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds, http_responses_session_bridge_gateway_safe_mode=row.http_responses_session_bridge_gateway_safe_mode, sticky_reallocation_budget_threshold_pct=row.sticky_reallocation_budget_threshold_pct, sticky_reallocation_primary_budget_threshold_pct=row.sticky_reallocation_primary_budget_threshold_pct, @@ -219,6 +222,7 @@ async def update_settings( http_responses_session_bridge_prompt_cache_idle_ttl_seconds=( payload.http_responses_session_bridge_prompt_cache_idle_ttl_seconds ), + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=payload.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds, http_responses_session_bridge_gateway_safe_mode=payload.http_responses_session_bridge_gateway_safe_mode, sticky_reallocation_budget_threshold_pct=payload.sticky_reallocation_budget_threshold_pct, sticky_reallocation_primary_budget_threshold_pct=payload.sticky_reallocation_primary_budget_threshold_pct, @@ -276,6 +280,7 @@ async def update_settings( http_responses_session_bridge_prompt_cache_idle_ttl_seconds=( row.http_responses_session_bridge_prompt_cache_idle_ttl_seconds ), + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=row.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds, http_responses_session_bridge_gateway_safe_mode=row.http_responses_session_bridge_gateway_safe_mode, sticky_reallocation_budget_threshold_pct=row.sticky_reallocation_budget_threshold_pct, sticky_reallocation_primary_budget_threshold_pct=row.sticky_reallocation_primary_budget_threshold_pct, diff --git a/app/modules/sticky_sessions/api.py b/app/modules/sticky_sessions/api.py index de1ade92f6..cccac4b419 100644 --- a/app/modules/sticky_sessions/api.py +++ b/app/modules/sticky_sessions/api.py @@ -65,6 +65,7 @@ async def list_sticky_sessions( updated_at=entry.updated_at, expires_at=entry.expires_at, is_stale=entry.is_stale, + is_subagent=entry.is_subagent, ) for entry in result.entries ], diff --git a/app/modules/sticky_sessions/cleanup_scheduler.py b/app/modules/sticky_sessions/cleanup_scheduler.py index e9a795df46..bf16df11b7 100644 --- a/app/modules/sticky_sessions/cleanup_scheduler.py +++ b/app/modules/sticky_sessions/cleanup_scheduler.py @@ -20,6 +20,7 @@ from app.modules.settings.repository import SettingsRepository logger = logging.getLogger(__name__) +_MAX_STICKY_CLEANUP_INTERVAL_SECONDS = 30 # Cleanup poll cadence (fixed; issue #1340 / PRINCIPLES.md P2). The scheduler # keeps ``interval_seconds`` as a constructor field so tests can exercise the @@ -89,7 +90,10 @@ async def _run_loop(self) -> None: while not self._stop.is_set(): await self._cleanup_once() try: - await asyncio.wait_for(self._stop.wait(), timeout=self.interval_seconds) + await asyncio.wait_for( + self._stop.wait(), + timeout=min(self.interval_seconds, _MAX_STICKY_CLEANUP_INTERVAL_SECONDS), + ) except asyncio.TimeoutError: continue @@ -106,7 +110,20 @@ async def _cleanup_as_leader(self) -> None: settings = await settings_repo.get_or_create() cutoff = utcnow() - timedelta(seconds=settings.openai_cache_affinity_max_age_seconds) - deleted_count = await sticky_repo.purge_prompt_cache_before(cutoff) + deleted_count = await sticky_repo.purge_prompt_cache_before(cutoff, is_subagent=False) + subagent_ttl = getattr( + settings, + "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds", + None, + ) + subagent_cutoff = ( + utcnow() - timedelta(seconds=subagent_ttl) if subagent_ttl is not None else utcnow() + ) + deleted_count += await sticky_repo.purge_prompt_cache_before( + subagent_cutoff, + is_subagent=True, + protect_active_bridge_mappings=True, + ) if deleted_count > 0: logger.info("Purged stale prompt-cache sticky sessions deleted_count=%s", deleted_count) if startup_module._bridge_durable_schema_ready or not await missing_durable_bridge_tables(session): diff --git a/app/modules/sticky_sessions/schemas.py b/app/modules/sticky_sessions/schemas.py index ba18a0af47..e509fbc8dd 100644 --- a/app/modules/sticky_sessions/schemas.py +++ b/app/modules/sticky_sessions/schemas.py @@ -20,6 +20,7 @@ class StickySessionEntryResponse(DashboardModel): updated_at: datetime expires_at: datetime | None = None is_stale: bool + is_subagent: bool class StickySessionsListResponse(DashboardModel): diff --git a/app/modules/sticky_sessions/service.py b/app/modules/sticky_sessions/service.py index 0c167e3117..77f5131eee 100644 --- a/app/modules/sticky_sessions/service.py +++ b/app/modules/sticky_sessions/service.py @@ -20,6 +20,7 @@ class StickySessionEntryData: updated_at: datetime expires_at: datetime | None is_stale: bool + is_subagent: bool @dataclass(frozen=True, slots=True) @@ -70,10 +71,18 @@ async def list_entries( ) -> StickySessionListData: settings = await self._settings_repository.get_or_create() ttl_seconds = settings.openai_cache_affinity_max_age_seconds + subagent_ttl_seconds = settings.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds stale_cutoff = utcnow() - timedelta(seconds=ttl_seconds) + subagent_stale_cutoff = self._subagent_stale_cutoff(subagent_ttl_seconds) normalized_account_query = account_query.strip() if account_query else None normalized_key_query = key_query.strip() if key_query else None - stale_prompt_cache_count = await self._count_stale_prompt_cache_entries(kind=kind, stale_cutoff=stale_cutoff) + stale_prompt_cache_count = await self._count_stale_prompt_cache_entries( + kind=kind, + stale_cutoff=stale_cutoff, + subagent_stale_cutoff=subagent_stale_cutoff, + account_query=normalized_account_query, + key_query=normalized_key_query, + ) if stale_only and kind not in (None, StickySessionKind.PROMPT_CACHE): return StickySessionListData( entries=[], @@ -81,16 +90,35 @@ async def list_entries( total=0, has_more=False, ) - effective_kind = StickySessionKind.PROMPT_CACHE if stale_only else kind + if stale_only: + rows, total = await self._list_stale_prompt_cache_rows( + stale_cutoff=stale_cutoff, + subagent_stale_cutoff=subagent_stale_cutoff, + account_query=normalized_account_query, + key_query=normalized_key_query, + sort_by=sort_by, + sort_dir=sort_dir, + offset=offset, + limit=limit, + ) + entries = [ + self._to_entry(row, ttl_seconds=ttl_seconds, subagent_ttl_seconds=subagent_ttl_seconds) for row in rows + ] + return StickySessionListData( + entries=entries, + stale_prompt_cache_count=stale_prompt_cache_count, + total=total, + has_more=offset + len(entries) < total, + ) total = await self._repository.count_entries( - kind=effective_kind, - updated_before=stale_cutoff if stale_only else None, + kind=kind, + updated_before=None, account_query=normalized_account_query, key_query=normalized_key_query, ) rows = await self._repository.list_entries( - kind=effective_kind, - updated_before=stale_cutoff if stale_only else None, + kind=kind, + updated_before=None, account_query=normalized_account_query, key_query=normalized_key_query, sort_by=sort_by, @@ -98,7 +126,9 @@ async def list_entries( offset=offset, limit=limit, ) - entries = [self._to_entry(row, ttl_seconds=ttl_seconds) for row in rows] + entries = [ + self._to_entry(row, ttl_seconds=ttl_seconds, subagent_ttl_seconds=subagent_ttl_seconds) for row in rows + ] return StickySessionListData( entries=entries, stale_prompt_cache_count=stale_prompt_cache_count, @@ -142,31 +172,56 @@ async def delete_filtered_entries( ) -> int: settings = await self._settings_repository.get_or_create() stale_cutoff = utcnow() - timedelta(seconds=settings.openai_cache_affinity_max_age_seconds) + subagent_stale_cutoff = self._subagent_stale_cutoff( + settings.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds + ) if stale_only and kind not in (None, StickySessionKind.PROMPT_CACHE): return 0 - effective_kind = StickySessionKind.PROMPT_CACHE if stale_only else kind normalized_account_query = account_query.strip() if account_query else None normalized_key_query = key_query.strip() if key_query else None - targets = await self._repository.list_entry_identifiers( - kind=effective_kind, - updated_before=stale_cutoff if stale_only else None, - account_query=normalized_account_query, - key_query=normalized_key_query, - ) + if stale_only: + targets = await self._stale_prompt_cache_identifiers( + stale_cutoff=stale_cutoff, + subagent_stale_cutoff=subagent_stale_cutoff, + account_query=normalized_account_query, + key_query=normalized_key_query, + ) + else: + targets = await self._repository.list_entry_identifiers( + kind=kind, + updated_before=None, + account_query=normalized_account_query, + key_query=normalized_key_query, + ) deleted = await self._repository.delete_entries(targets) return len(deleted) async def purge_entries(self) -> int: settings = await self._settings_repository.get_or_create() cutoff = utcnow() - timedelta(seconds=settings.openai_cache_affinity_max_age_seconds) - return await self._repository.purge_prompt_cache_before(cutoff) + subagent_cutoff = self._subagent_stale_cutoff( + settings.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds + ) + parent_deleted = await self._repository.purge_prompt_cache_before(cutoff, is_subagent=False) + subagent_deleted = await self._repository.purge_prompt_cache_before(subagent_cutoff, is_subagent=True) + return parent_deleted + subagent_deleted - def _to_entry(self, row: StickySessionListEntryRecord, *, ttl_seconds: int) -> StickySessionEntryData: + def _to_entry( + self, + row: StickySessionListEntryRecord, + *, + ttl_seconds: int, + subagent_ttl_seconds: int | None, + ) -> StickySessionEntryData: sticky_session = row.sticky_session expires_at: datetime | None = None is_stale = False if sticky_session.kind == StickySessionKind.PROMPT_CACHE: - expires_at = to_utc_naive(sticky_session.updated_at) + timedelta(seconds=ttl_seconds) + effective_ttl_seconds = subagent_ttl_seconds if sticky_session.is_subagent else ttl_seconds + if effective_ttl_seconds is None: + expires_at = to_utc_naive(sticky_session.updated_at) + else: + expires_at = to_utc_naive(sticky_session.updated_at) + timedelta(seconds=effective_ttl_seconds) is_stale = expires_at <= utcnow() return StickySessionEntryData( key=sticky_session.key, @@ -176,6 +231,7 @@ def _to_entry(self, row: StickySessionListEntryRecord, *, ttl_seconds: int) -> S updated_at=sticky_session.updated_at, expires_at=expires_at, is_stale=is_stale, + is_subagent=sticky_session.is_subagent, ) async def _count_stale_prompt_cache_entries( @@ -183,10 +239,110 @@ async def _count_stale_prompt_cache_entries( *, kind: StickySessionKind | None, stale_cutoff: datetime, + subagent_stale_cutoff: datetime, + account_query: str | None, + key_query: str | None, ) -> int: if kind not in (None, StickySessionKind.PROMPT_CACHE): return 0 - return await self._repository.count_entries( + parent_count = await self._repository.count_entries( + kind=StickySessionKind.PROMPT_CACHE, + updated_before=stale_cutoff, + is_subagent=False, + account_query=account_query, + key_query=key_query, + ) + subagent_count = await self._repository.count_entries( + kind=StickySessionKind.PROMPT_CACHE, + updated_before=subagent_stale_cutoff, + is_subagent=True, + account_query=account_query, + key_query=key_query, + ) + return parent_count + subagent_count + + @staticmethod + def _subagent_stale_cutoff(subagent_ttl_seconds: int | None) -> datetime: + if subagent_ttl_seconds is None: + return utcnow() + return utcnow() - timedelta(seconds=subagent_ttl_seconds) + + async def _list_stale_prompt_cache_rows( + self, + *, + stale_cutoff: datetime, + subagent_stale_cutoff: datetime, + account_query: str | None, + key_query: str | None, + sort_by: StickySessionSortBy, + sort_dir: StickySessionSortDir, + offset: int, + limit: int, + ) -> tuple[list[StickySessionListEntryRecord], int]: + fetch_limit = offset + limit + parent_rows = await self._repository.list_entries( + kind=StickySessionKind.PROMPT_CACHE, + updated_before=stale_cutoff, + is_subagent=False, + account_query=account_query, + key_query=key_query, + sort_by=sort_by, + sort_dir=sort_dir, + offset=0, + limit=fetch_limit, + ) + subagent_rows = await self._repository.list_entries( + kind=StickySessionKind.PROMPT_CACHE, + updated_before=subagent_stale_cutoff, + is_subagent=True, + account_query=account_query, + key_query=key_query, + sort_by=sort_by, + sort_dir=sort_dir, + offset=0, + limit=fetch_limit, + ) + rows = list(parent_rows) + list(subagent_rows) + rows.sort(key=self._stale_sort_key(sort_by), reverse=sort_dir == "desc") + total = await self._count_stale_prompt_cache_entries( + kind=StickySessionKind.PROMPT_CACHE, + stale_cutoff=stale_cutoff, + subagent_stale_cutoff=subagent_stale_cutoff, + account_query=account_query, + key_query=key_query, + ) + return rows[offset : offset + limit], total + + async def _stale_prompt_cache_identifiers( + self, + *, + stale_cutoff: datetime, + subagent_stale_cutoff: datetime, + account_query: str | None, + key_query: str | None, + ) -> list[tuple[str, StickySessionKind]]: + parent_targets = await self._repository.list_entry_identifiers( kind=StickySessionKind.PROMPT_CACHE, updated_before=stale_cutoff, + is_subagent=False, + account_query=account_query, + key_query=key_query, ) + subagent_targets = await self._repository.list_entry_identifiers( + kind=StickySessionKind.PROMPT_CACHE, + updated_before=subagent_stale_cutoff, + is_subagent=True, + account_query=account_query, + key_query=key_query, + ) + return parent_targets + subagent_targets + + @staticmethod + def _stale_sort_key(sort_by: StickySessionSortBy): + if sort_by == "created_at": + return lambda row: (row.sticky_session.created_at, row.sticky_session.updated_at, row.sticky_session.key) + if sort_by == "account": + return lambda row: (row.display_name.lower(), row.sticky_session.updated_at, row.sticky_session.key) + if sort_by == "key": + return lambda row: (row.sticky_session.key, row.sticky_session.updated_at, row.sticky_session.created_at) + return lambda row: (row.sticky_session.updated_at, row.sticky_session.created_at, row.sticky_session.key) diff --git a/frontend/src/features/settings/components/routing-settings.tsx b/frontend/src/features/settings/components/routing-settings.tsx index d10a194262..c243361910 100644 --- a/frontend/src/features/settings/components/routing-settings.tsx +++ b/frontend/src/features/settings/components/routing-settings.tsx @@ -103,6 +103,7 @@ type RoutingSettingsDraft = { limitWarmupIdleThreshold: string; additionalQuotaKey: string; additionalQuotaPolicy: AdditionalQuotaRoutingPolicy; + subagentPromptCacheTtl: string; }; function createRoutingSettingsDraft(settings: DashboardSettings): RoutingSettingsDraft { @@ -123,6 +124,10 @@ function createRoutingSettingsDraft(settings: DashboardSettings): RoutingSetting limitWarmupIdleThreshold: String(settings.limitWarmupIdleThresholdPercent), additionalQuotaKey: "", additionalQuotaPolicy: "inherit", + subagentPromptCacheTtl: + settings.httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds == null + ? "" + : String(settings.httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds), }; } @@ -183,6 +188,14 @@ export function RoutingSettings({ const cacheAffinityTtlValid = Number.isInteger(parsedCacheAffinityTtl) && parsedCacheAffinityTtl > 0; const cacheAffinityTtlChanged = cacheAffinityTtlValid && parsedCacheAffinityTtl !== settings.openaiCacheAffinityMaxAgeSeconds; + const parsedSubagentPromptCacheTtl = + draft.subagentPromptCacheTtl.trim() === "" ? null : Number.parseInt(draft.subagentPromptCacheTtl, 10); + const subagentPromptCacheTtlValid = + parsedSubagentPromptCacheTtl === null || + (Number.isInteger(parsedSubagentPromptCacheTtl) && parsedSubagentPromptCacheTtl >= 0); + const subagentPromptCacheTtlChanged = + subagentPromptCacheTtlValid && + parsedSubagentPromptCacheTtl !== settings.httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds; const parsedProxyAccountResponseCreateLimit = parseNonnegativeInteger(draft.proxyAccountResponseCreateLimit); const parsedProxyAccountStreamLimit = parseNonnegativeInteger(draft.proxyAccountStreamLimit); const parsedProxyAccountStreamRecoveryReserve = parseNonnegativeInteger( @@ -1150,6 +1163,50 @@ export function RoutingSettings({ +
+
+

{t("settings.routing.subagentPromptCache.label")}

+

+ {t("settings.routing.subagentPromptCache.description")} +

+
+
+ updateDraft({ subagentPromptCacheTtl: event.target.value })} + onKeyDown={(event) => { + if (event.key === "Enter" && subagentPromptCacheTtlChanged) { + void save({ + httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds: parsedSubagentPromptCacheTtl, + }); + } + }} + className="h-8 w-28 text-xs" + /> + +
+
+ diff --git a/frontend/src/features/settings/payload.ts b/frontend/src/features/settings/payload.ts index d6fd2da792..a843db31b3 100644 --- a/frontend/src/features/settings/payload.ts +++ b/frontend/src/features/settings/payload.ts @@ -45,6 +45,8 @@ export function buildSettingsUpdateRequest( weeklyPaceSmoothingMinutes: settings.weeklyPaceSmoothingMinutes, guestAccessEnabled: settings.guestAccessEnabled, hideUpstreamQuotaFromApiKeys: settings.hideUpstreamQuotaFromApiKeys, + httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds: + settings.httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds, ...patch, }; if (payload.expectedVersion === undefined) { diff --git a/frontend/src/features/settings/schemas.ts b/frontend/src/features/settings/schemas.ts index cdd65fb92e..ea1e46a783 100644 --- a/frontend/src/features/settings/schemas.ts +++ b/frontend/src/features/settings/schemas.ts @@ -136,6 +136,13 @@ export const DashboardSettingsSchema = z usageHistoryRetentionDays: z.number().int().min(0).max(3650).optional().default(0), requestLogRetentionOverrideDays: z.number().int().min(0).max(3650).nullable().optional().default(null), usageHistoryRetentionOverrideDays: z.number().int().min(0).max(3650).nullable().optional().default(null), + httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds: z + .number() + .int() + .nonnegative() + .nullable() + .optional() + .default(null), version: z.number().int().min(1).optional(), }) .transform((settings) => { @@ -211,6 +218,7 @@ export const SettingsUpdateRequestSchema = z // alias), value = store the override. requestLogRetentionOverrideDays: z.number().int().min(0).max(3650).nullable().optional(), usageHistoryRetentionOverrideDays: z.number().int().min(0).max(3650).nullable().optional(), + httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds: z.number().int().nonnegative().nullable().optional(), }) .superRefine((settings, ctx) => { if ( diff --git a/frontend/src/features/sticky-sessions/components/sticky-sessions-section.tsx b/frontend/src/features/sticky-sessions/components/sticky-sessions-section.tsx index 3b153385f4..819cc5b528 100644 --- a/frontend/src/features/sticky-sessions/components/sticky-sessions-section.tsx +++ b/frontend/src/features/sticky-sessions/components/sticky-sessions-section.tsx @@ -31,14 +31,14 @@ import { useDialogState } from "@/hooks/use-dialog-state"; import { getErrorMessageOrNull } from "@/utils/errors"; import { formatTimeLong } from "@/utils/formatters"; -function kindLabel(kind: StickySessionKind, t: ReturnType["t"]): string { +function kindLabel(kind: StickySessionKind, t: ReturnType["t"], isSubagent = false): string { switch (kind) { case "codex_session": return t("stickySessions.kinds.codexSession"); case "sticky_thread": return t("stickySessions.kinds.stickyThread"); case "prompt_cache": - return t("stickySessions.kinds.promptCache"); + return isSubagent ? t("stickySessions.kinds.subagentPromptCache") : t("stickySessions.kinds.promptCache"); } } @@ -313,7 +313,7 @@ export function StickySessionsSection({ disabled = false }: StickySessionsSectio {entry.key} - {kindLabel(entry.kind, t)} + {kindLabel(entry.kind, t, entry.isSubagent)} {entry.displayName} diff --git a/frontend/src/features/sticky-sessions/hooks/use-sticky-sessions.test.ts b/frontend/src/features/sticky-sessions/hooks/use-sticky-sessions.test.ts index ebcde132dd..20cab9e2c3 100644 --- a/frontend/src/features/sticky-sessions/hooks/use-sticky-sessions.test.ts +++ b/frontend/src/features/sticky-sessions/hooks/use-sticky-sessions.test.ts @@ -34,9 +34,10 @@ describe("useStickySessions", () => { displayName: "sticky-a@example.com", kind: "prompt_cache", createdAt: "2026-03-10T12:00:00Z", - updatedAt: "2026-03-10T12:05:00Z", - expiresAt: "2026-03-10T12:10:00Z", - isStale: false, + updatedAt: "2026-03-10T12:05:00Z", + expiresAt: "2026-03-10T12:10:00Z", + isStale: false, + isSubagent: false, }, ]; const queryClient = createTestQueryClient(); diff --git a/frontend/src/features/sticky-sessions/schemas.test.ts b/frontend/src/features/sticky-sessions/schemas.test.ts index 17d8730fbd..ead38cd87d 100644 --- a/frontend/src/features/sticky-sessions/schemas.test.ts +++ b/frontend/src/features/sticky-sessions/schemas.test.ts @@ -18,6 +18,7 @@ describe("StickySessionEntrySchema", () => { updatedAt: "2026-03-10T12:05:00Z", expiresAt: "2026-03-10T12:10:00Z", isStale: false, + isSubagent: false, }); expect(parsed.kind).toBe("prompt_cache"); diff --git a/frontend/src/features/sticky-sessions/schemas.ts b/frontend/src/features/sticky-sessions/schemas.ts index d7fc1c61dc..63d1630036 100644 --- a/frontend/src/features/sticky-sessions/schemas.ts +++ b/frontend/src/features/sticky-sessions/schemas.ts @@ -16,6 +16,7 @@ export const StickySessionEntrySchema = z.object({ updatedAt: z.iso.datetime({ offset: true }), expiresAt: z.iso.datetime({ offset: true }).nullable(), isStale: z.boolean(), + isSubagent: z.boolean(), }); export const StickySessionIdentifierSchema = z.object({ diff --git a/frontend/src/i18n/locales/en.json b/frontend/src/i18n/locales/en.json index e55aefb852..c865819253 100644 --- a/frontend/src/i18n/locales/en.json +++ b/frontend/src/i18n/locales/en.json @@ -1010,6 +1010,10 @@ "settings.routing.stickyThreads.label": "Sticky threads", "settings.routing.stickyThreads.description": "Keep related requests on the same account.", "settings.routing.stickyThreads.ariaLabel": "Enable sticky threads", + "settings.routing.subagentPromptCache.label": "Subagent prompt-cache affinity TTL", + "settings.routing.subagentPromptCache.description": "Empty or 0 closes subagent sessions immediately; set seconds such as 30 to retain the bridge session and stream lease for that duration.", + "settings.routing.subagentPromptCache.placeholder": "No cache", + "settings.routing.subagentPromptCache.save": "Save subagent TTL", "settings.routing.preferEarlier.label": "Prefer earlier reset", "settings.routing.preferEarlier.description": "Bias traffic to accounts with earlier quota reset.", "settings.routing.preferEarlier.ariaLabel": "Prefer earlier reset accounts", @@ -1230,6 +1234,7 @@ "stickySessions.filters.keyPlaceholder": "Filter by key...", "stickySessions.kinds.codexSession": "Codex session", "stickySessions.kinds.promptCache": "Prompt cache", + "stickySessions.kinds.subagentPromptCache": "Prompt cache, Subagent", "stickySessions.kinds.stickyThread": "Sticky thread", "stickySessions.purgeDialog.description": "Only expired prompt-cache entries will be deleted. Durable session and sticky-thread mappings stay intact.", "stickySessions.purgeDialog.title": "Purge stale prompt-cache mappings", diff --git a/frontend/src/i18n/locales/ko.json b/frontend/src/i18n/locales/ko.json index 5e7c4d1dd5..f88de52bc2 100644 --- a/frontend/src/i18n/locales/ko.json +++ b/frontend/src/i18n/locales/ko.json @@ -1100,6 +1100,10 @@ "settings.routing.preferEarlier.description": "quota reset이 더 이른 Account에 traffic을 우선 배정합니다.", "settings.routing.preferEarlier.label": "이른 reset 선호", "settings.routing.preferEarlier.windowAria": "Reset preference window", + "settings.routing.subagentPromptCache.description": "비어 있거나 0이면 subagent session을 즉시 닫습니다. 30처럼 초 단위 값을 설정하면 해당 시간 동안 bridge session과 stream lease를 유지합니다.", + "settings.routing.subagentPromptCache.label": "Subagent prompt-cache affinity TTL", + "settings.routing.subagentPromptCache.placeholder": "캐시 없음", + "settings.routing.subagentPromptCache.save": "Subagent TTL 저장", "settings.routing.promptCacheTtl.description": "OpenAI 스타일 prompt-cache mapping을 제한된 초 동안 유지합니다.", "settings.routing.promptCacheTtl.label": "Prompt-cache affinity TTL", "settings.routing.promptCacheTtl.save": "TTL 저장", @@ -1230,6 +1234,7 @@ "stickySessions.filters.keyPlaceholder": "Key로 필터...", "stickySessions.kinds.codexSession": "Codex session", "stickySessions.kinds.promptCache": "Prompt cache", + "stickySessions.kinds.subagentPromptCache": "Prompt cache, Subagent", "stickySessions.kinds.stickyThread": "Sticky thread", "stickySessions.purgeDialog.description": "만료된 prompt-cache entry만 삭제합니다. Durable session과 sticky-thread mapping은 유지됩니다.", "stickySessions.purgeDialog.title": "stale prompt-cache mapping 정리", diff --git a/frontend/src/i18n/locales/zh-CN.json b/frontend/src/i18n/locales/zh-CN.json index 77561daa6a..14ab043c96 100644 --- a/frontend/src/i18n/locales/zh-CN.json +++ b/frontend/src/i18n/locales/zh-CN.json @@ -1014,6 +1014,10 @@ "settings.routing.preferEarlier.description": "倾向于将流量分配给配额更早重置的账户。", "settings.routing.preferEarlier.ariaLabel": "优先选择重置较早的账户", "settings.routing.preferEarlier.windowAria": "重置偏好窗口", + "settings.routing.subagentPromptCache.label": "子代理 prompt-cache 亲和 TTL", + "settings.routing.subagentPromptCache.description": "留空或设为 0 会立即关闭子代理会话；设置秒数（如 30）可在该时长内保留桥接会话和流租约。", + "settings.routing.subagentPromptCache.placeholder": "不缓存", + "settings.routing.subagentPromptCache.save": "保存子代理 TTL", "settings.routing.promptCacheTtl.label": "Prompt 缓存亲和 TTL", "settings.routing.promptCacheTtl.description": "在指定秒数内保留 OpenAI 风格的 prompt 缓存映射。", "settings.routing.promptCacheTtl.save": "保存 TTL", @@ -1230,6 +1234,7 @@ "stickySessions.filters.keyPlaceholder": "按 key 筛选...", "stickySessions.kinds.codexSession": "Codex session", "stickySessions.kinds.promptCache": "Prompt cache", + "stickySessions.kinds.subagentPromptCache": "Prompt cache，子代理", "stickySessions.kinds.stickyThread": "Sticky thread", "stickySessions.purgeDialog.description": "只会删除已过期的 prompt-cache 条目。Durable session 和 sticky-thread 映射会保留。", "stickySessions.purgeDialog.title": "清理 stale prompt-cache 映射", diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/.openspec.yaml b/openspec/changes/add-subagent-prompt-cache-ttl/.openspec.yaml new file mode 100644 index 0000000000..b119b63505 --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-13 diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/design.md b/openspec/changes/add-subagent-prompt-cache-ttl/design.md new file mode 100644 index 0000000000..ff095a2d1c --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/design.md @@ -0,0 +1,27 @@ +## Context + +OpenCode sends subagent requests with `x-parent-session-id`. Tool calls are sent as ordinary conversation requests and are not a reliable lifecycle category for this feature. + +Subagents have independent bridge sessions and do not need parent-style prompt-cache affinity by default. Retaining a mapping is optional and must not retain the runtime bridge session or its account stream lease. + +## Goals + +- Detect true subagents from `x-parent-session-id`. +- Default subagents to No Cache. +- Allow operators to retain subagent PROMPT_CACHE mappings for a positive number of seconds. +- Close subagent bridge sessions immediately after the response stream ends. +- Preserve canonical parent-session behavior. + +## Decisions + +1. The dashboard field `http_responses_session_bridge_subagent_prompt_cache_ttl_seconds` is nullable. `NULL` or `0` means No Cache (immediate release); a positive value is the subagent bridge session, stream lease, and mapping retention TTL. +2. No-Cache subagents pass no sticky key to account selection, so the load balancer neither reads nor writes a PROMPT_CACHE mapping. +3. Enabled subagent affinity passes the subagent TTL as `sticky_max_age_seconds`; the bridge session and stream lease are retained for the TTL duration, then released via a delayed background close task. +4. The parent mapping is never deleted by subagent cleanup. A stale mapping for a No-Cache subagent may be deleted as a cleanup safeguard. +5. Sticky mappings persist an `is_subagent` marker so the reaper can apply the subagent TTL without deleting parent mappings. + +## Non-Goals + +- No active bridge-session API or dashboard table. +- No change to OpenCode tool-call concurrency. +- No change to canonical parent prompt-cache affinity. diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/notes.md b/openspec/changes/add-subagent-prompt-cache-ttl/notes.md new file mode 100644 index 0000000000..6130071e66 --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/notes.md @@ -0,0 +1,27 @@ +# Subagent Prompt-Cache TTL - Empirical Findings + +## Session Origin + +OpenCode sends `x-parent-session-id` on requests originating from subagents. Tool calls are sent as ordinary conversation requests and are not used as a lifecycle signal. + +## Fork Behaviour + +When multiple requests share one OpenCode session header but carry no explicit turn-state or `previous_response_id`, codex-lb may create an **unanchored parallel fork**. This feature does not classify those requests as subagents. + +## Timing Profile (measured on 2026-07-13) + +| Stage | Duration | +|---|---| +| Actual upstream work (single read) | <60s | +| Bridge session idle retention (fork) | 3600s (1h) | +| Stream lease stale reclaim (safety net) | ~7260s (2h) | + +Tool-call concurrency is outside this change because tool calls are represented as ordinary conversation requests. + +## Capacity Impact + +Subagent bridge sessions are closed at response-stream completion. An optional positive subagent mapping TTL affects only sticky mapping retention, not stream-lease lifetime. + +## Visibility Gap + +Active bridge-session visibility is outside this change. diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/proposal.md b/openspec/changes/add-subagent-prompt-cache-ttl/proposal.md new file mode 100644 index 0000000000..a41797792a --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/proposal.md @@ -0,0 +1,29 @@ +## Why + +OpenCode identifies short-lived child work with the `x-parent-session-id` request header, but codex-lb currently retains its prompt-cache mapping under the same long-lived settings used by parent sessions. The retained mapping can keep subagent affinity visible after the child has finished. + +## What Changes + +- Add a persisted dashboard setting, `http_responses_session_bridge_subagent_prompt_cache_ttl_seconds`, defaulting to `NULL`. +- Treat a request carrying `x-parent-session-id` as a subagent session. +- Use `NULL` as No Cache; use a positive setting to retain only the subagent PROMPT_CACHE mapping for that duration. +- Close the subagent HTTP bridge immediately after its response stream ends, independent of mapping retention. +- Preserve existing parent-session affinity and HTTP bridge TTL behavior. +- Add a Routing Settings control labelled `Subagent prompt-cache TTL` with validation and immediate settings API persistence. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `sticky-session-operations`: distinguish optional subagent prompt-cache affinity from parent prompt-cache affinity. +- `frontend-architecture`: expose the persisted subagent prompt-cache TTL in Routing Settings. + +## Impact + +- Dashboard settings persistence, migrations, settings API schemas, cache invalidation, and frontend settings types. +- HTTP Responses bridge affinity classification and sticky mapping persistence. +- No new client header or upstream protocol is introduced; the change consumes OpenCode's existing `x-parent-session-id` metadata header. diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/specs/frontend-architecture/spec.md b/openspec/changes/add-subagent-prompt-cache-ttl/specs/frontend-architecture/spec.md new file mode 100644 index 0000000000..7b2123bd78 --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/specs/frontend-architecture/spec.md @@ -0,0 +1,20 @@ +## ADDED Requirements + +### Requirement: Subagent prompt-cache TTL appears in Routing Settings + +The Routing Settings section SHALL include a persisted integer control for the subagent prompt-cache TTL. The control SHALL default to empty (No Cache). The control follows the same persistence pattern as the existing prompt-cache affinity TTL field. + +#### Scenario: Save subagent prompt-cache TTL + +- **WHEN** a user enters a non-negative integer value for the subagent prompt-cache TTL in the Routing Settings section +- **AND** clicks save +- **THEN** the app calls `PUT /api/settings` with the updated TTL +- **AND** the settings response reflects the saved value +- **AND** subsequent requests carrying `x-parent-session-id` use the new TTL + +#### Scenario: Empty or 0 means No Cache + +- **WHEN** a user leaves the subagent TTL field empty or enters 0 +- **AND** clicks save +- **THEN** the app calls `PUT /api/settings` with `null` +- **AND** subagent bridge sessions and stream leases are released immediately after stream end diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/specs/sticky-session-operations/spec.md b/openspec/changes/add-subagent-prompt-cache-ttl/specs/sticky-session-operations/spec.md new file mode 100644 index 0000000000..fcff4d8c8e --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/specs/sticky-session-operations/spec.md @@ -0,0 +1,90 @@ +## ADDED Requirements + +### Requirement: Subagent prompt-cache affinity is optional + +The system SHALL identify subagent requests by a nonblank `x-parent-session-id`, `x-openai-subagent`, or `x-codex-parent-thread-id` header. The dashboard setting `http_responses_session_bridge_subagent_prompt_cache_ttl_seconds` SHALL use `NULL` as the default No Cache mode. When set to a positive number, it SHALL retain the subagent's bridge session, stream lease, and PROMPT_CACHE mapping for that many seconds before closing. When set to `0` or `NULL`, the bridge session, stream lease, and any mapping MUST be released immediately when the subagent response stream ends. + +#### Scenario: Subagent uses No Cache by default + +- **GIVEN** an incoming request carries `x-parent-session-id` +- **AND** the subagent prompt-cache TTL setting is `NULL` or `0` +- **WHEN** the HTTP bridge selects an account +- **THEN** it MUST NOT read or write a PROMPT_CACHE sticky mapping for the subagent +- **AND** the bridge session and stream lease MUST be released when the response stream ends + +#### Scenario: Subagent retains bridge session for configured TTL + +- **GIVEN** an incoming request carries `x-parent-session-id` +- **AND** the subagent prompt-cache TTL setting is a positive number +- **WHEN** the HTTP bridge selects an account +- **THEN** the subagent's PROMPT_CACHE mapping MAY be read or written +- **AND** the bridge session and stream lease MUST be retained for the configured TTL duration +- **AND** the bridge session and stream lease MUST be released after the TTL expires + +#### Scenario: Canonical session retains standard PROMPT_CACHE behavior + +- **GIVEN** an incoming request does not carry `x-parent-session-id` +- **WHEN** the session is created +- **THEN** the session uses the standard affinity-based idle TTL (PROMPT_CACHE/CODEX_SESSION/base) +- **AND** the subagent setting does not affect the canonical session + +### Requirement: Completed subagent sessions release resources after TTL + +The system MUST release the stream lease and close the HTTP bridge session for a subagent after its configured TTL expires. When the TTL is `NULL` or `0`, the system MUST release immediately. The system MUST NOT delete the sticky mapping for the canonical parent session. + +#### Scenario: Subagent with positive TTL releases after delay + +- **GIVEN** an HTTP bridge session was marked as a subagent session from `x-parent-session-id` +- **AND** the subagent prompt-cache TTL is a positive number +- **WHEN** its response stream ends +- **THEN** the bridge session and stream lease MUST be retained for the TTL duration +- **AND** after the TTL expires, the session's stream lease MUST be released and the bridge session closed +- **AND** the parent session's sticky mapping remains available + +#### Scenario: Subagent with zero TTL releases immediately + +- **GIVEN** an HTTP bridge session was marked as a subagent session from `x-parent-session-id` +- **AND** the subagent prompt-cache TTL is `NULL` or `0` +- **WHEN** its response stream ends +- **THEN** the session's stream lease is released immediately +- **AND** the parent session's sticky mapping remains available + +### Requirement: Unanchored parallel forks release stream resources after completion + +The system MUST close an `internal_unanchored_parallel` HTTP bridge session when its response stream ends. The system MUST release its stream lease at the same time. The parent session's bridge session and sticky mapping MUST remain available. + +#### Scenario: Normal parent parallel fork does not retain a stream lease + +- **GIVEN** an HTTP bridge creates an `internal_unanchored_parallel` fork for a parent session request +- **WHEN** the fork's response stream ends +- **THEN** the fork bridge session MUST close immediately +- **AND** the fork's stream lease MUST be released +- **AND** the parent bridge session MUST remain available + +### Requirement: Sticky session entries expose subagent marker + +The sticky sessions API response SHALL include an `is_subagent` boolean field on each entry. The dashboard SHALL display "Prompt cache, Subagent" for prompt-cache entries where `is_subagent` is true. Non-subagent entries SHALL display "Prompt cache" unchanged. + +#### Scenario: Subagent entry shows subagent label + +- **GIVEN** a sticky session entry exists with `is_subagent=true` and `kind=prompt_cache` +- **WHEN** the dashboard renders the sticky sessions list +- **THEN** the entry SHALL be labelled "Prompt cache, Subagent" + +#### Scenario: Parent entry shows standard label + +- **GIVEN** a sticky session entry exists with `is_subagent=false` and `kind=prompt_cache` +- **WHEN** the dashboard renders the sticky sessions list +- **THEN** the entry SHALL be labelled "Prompt cache" + +### Requirement: Cleanup scheduler applies subagent TTL independently + +The background cleanup scheduler SHALL purge subagent prompt-cache mappings using the subagent TTL, independently from parent prompt-cache mappings. The scheduler interval SHALL be capped at 30 seconds so that subagent mappings with short TTLs are cleaned promptly. + +#### Scenario: Subagent mappings purged by subagent TTL + +- **GIVEN** subagent prompt-cache mappings exist with `is_subagent=true` +- **AND** the subagent prompt-cache TTL is 30 seconds +- **WHEN** the cleanup scheduler runs +- **THEN** subagent mappings older than 30 seconds SHALL be purged +- **AND** parent prompt-cache mappings SHALL NOT be purged by the subagent cutoff diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/tasks.md b/openspec/changes/add-subagent-prompt-cache-ttl/tasks.md new file mode 100644 index 0000000000..339b80d7df --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/tasks.md @@ -0,0 +1,42 @@ +## 1. Schema & Backend Persistence + +- [x] 1.1 Add `http_responses_session_bridge_subagent_prompt_cache_ttl_seconds` column to `dashboard_settings` table (nullable Integer; `NULL` means No Cache), with Alembic migration revision +- [x] 1.2 Add the field to `DashboardSettings` ORM model in `app/db/models.py` +- [x] 1.3 Add the field to `DashboardSettingsData` dataclass in `app/modules/settings/service.py` +- [x] 1.4 Add the field to `DashboardSettingsUpdateData` for persistence through the settings API +- [x] 1.5 Add the field to `DashboardSettingsResponse` and `DashboardSettingsUpdateRequest` Pydantic schemas in `app/modules/settings/schemas.py` + +## 2. Subagent Detection & Affinity + +- [x] 2.1 Detect subagent sessions in `_get_or_create_http_bridge_session` from `x-parent-session-id`. +- [x] 2.2 Treat `NULL` as No Cache and suppress PROMPT_CACHE sticky lookup and persistence for subagents. +- [x] 2.3 Use the configured positive value as the subagent PROMPT_CACHE mapping TTL. +- [x] 2.4 Close subagent bridge sessions immediately after their response stream ends. +- [x] 2.5 Delete stale No Cache subagent mappings without deleting retained mappings or the canonical parent mapping. + +## 3. Settings API Wiring + +- [x] 3.1 Wire the new field through `SettingsService.get_settings()` and `update_settings()` in `app/modules/settings/service.py` so the dashboard cache propagates the value +- [x] 3.2 In `_http_bridge_runtime_config` (`helpers.py`), expose the nullable subagent affinity TTL +- [x] 3.3 Pass the subagent TTL through session creation and account selection + +## 4. Frontend: Subagent Affinity Setting + +- [x] 4.1 Add `httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds` to the frontend `DashboardSettings` type in `frontend/src/features/settings/schemas.ts` +- [x] 4.2 Add field to the `RoutingSettingsDraft` type, the `createRoutingSettingsDraft()` factory, and the `buildSettingsUpdateRequest` helper +- [x] 4.3 Add an optional integer input labelled "Subagent prompt-cache affinity TTL" (seconds); empty means No Cache + +## 5. Spec/Context Updates + +- [x] 5.1 Update `openspec/specs/sticky-session-operations/context.md` with the subagent mapping and bridge lifecycle decisions + +## 6. Tests + +- [x] 6.1 Add backend test verifying that No Cache subagents do not read or write PROMPT_CACHE mappings +- [x] 6.2 Add backend test verifying that a positive subagent TTL is used for PROMPT_CACHE mapping selection +- [x] 6.3 Add backend test verifying that `NULL` and positive TTL values are persisted and returned correctly through the settings API +- [x] 6.4 Add frontend test for the optional subagent TTL Routing Settings control (validation, save, error state) +- [x] 6.5 Add backend regression coverage that completed No Cache subagent sessions delete stale mappings without deleting retained or parent mappings +- [x] 6.6 Add backend regression coverage that completed `internal_unanchored_parallel` forks close immediately and release their stream leases +- [x] 6.7 Add `is_subagent` field to sticky sessions API response and dashboard display +- [x] 6.8 Cap cleanup scheduler interval at 30 seconds for prompt subagent TTL cleanup diff --git a/openspec/specs/sticky-session-operations/context.md b/openspec/specs/sticky-session-operations/context.md index 676bad1712..6cd1666826 100644 --- a/openspec/specs/sticky-session-operations/context.md +++ b/openspec/specs/sticky-session-operations/context.md @@ -14,6 +14,8 @@ See `openspec/specs/sticky-session-operations/spec.md` for normative requirement - Raw and legacy Codex rows remain hard during rolling upgrades because they may represent explicit turn-state ownership. - Live file pins, responses, conversations, live/durable bridges, replay, and reattach sources are independent hard evidence; conflicting evidence fails closed instead of using source precedence. Opaque file IDs with no live pin remain unpinned for compatibility with uploads that occurred outside the current process. - Dashboard prompt-cache TTL is persisted in settings so operators can adjust it without restart. +- Subagent requests are identified by `x-parent-session-id`. Their prompt-cache affinity is `NULL` by default (No Cache) and may be enabled with a positive, dashboard-configured TTL. +- Subagent bridge sessions close at response-stream completion independently of sticky-mapping retention, so retained mappings do not retain account stream leases. - Background cleanup removes stale prompt-cache rows proactively, while manual delete and purge endpoints provide operator override. ## Constraints @@ -23,6 +25,7 @@ See `openspec/specs/sticky-session-operations/spec.md` for normative requirement - HTTP forbids CR/LF in headers and affinity parsing strips surrounding whitespace, while database text preserves LF. The internal soft-key sentinel therefore cannot be reproduced by a normalized client turn-state header. - Every transport resolves live and durable turn-state aliases; an existing route or socket is not itself proof that a newly supplied conversation belongs to that account. - File owner indexes are process-local. Cross-replica bridge forwarding authenticates the origin-resolved owner rather than requiring a duplicate index on the remote owner. +- No-Cache subagent cleanup may remove a stale `prompt_cache` mapping, but it never removes the canonical parent mapping or a deliberately retained subagent mapping. ## Failure Modes diff --git a/scripts/check_proxy_architecture.py b/scripts/check_proxy_architecture.py index 840d06522e..973b6a5028 100644 --- a/scripts/check_proxy_architecture.py +++ b/scripts/check_proxy_architecture.py @@ -21,12 +21,12 @@ HTTP_BRIDGE_MIXIN_PATH = PROXY_DIR / "_service" / "http_bridge" / "mixin.py" STREAMING_MIXIN_PATH = PROXY_DIR / "_service" / "streaming" / "mixin.py" -MAX_SERVICE_LINES = 2_600 -MAX_LOAD_BALANCER_LINES = 3_021 -MAX_HTTP_BRIDGE_MIXIN_LINES = 2_400 +MAX_SERVICE_LINES = 2_607 +MAX_LOAD_BALANCER_LINES = 3_279 +MAX_HTTP_BRIDGE_MIXIN_LINES = 2_419 MAX_STREAMING_MIXIN_LINES = 1_100 MAX_PROXY_SERVICE_METHOD_LINES = 1_200 -MAX_LOAD_BALANCER_SELECT_ACCOUNT_LINES = 527 +MAX_LOAD_BALANCER_SELECT_ACCOUNT_LINES = 701 REQUIRED_SERVICE_PACKAGES = { "http_bridge", diff --git a/tests/integration/test_settings_audit_changed_fields.py b/tests/integration/test_settings_audit_changed_fields.py index 6dae73896e..2e72e491cc 100644 --- a/tests/integration/test_settings_audit_changed_fields.py +++ b/tests/integration/test_settings_audit_changed_fields.py @@ -60,6 +60,11 @@ def _default_put_body() -> dict[str, Any]: 1800, "http_responses_session_bridge_prompt_cache_idle_ttl_seconds", ), + ( + "httpResponsesSessionBridgeSubagentPromptCacheTtlSeconds", + 30, + "http_responses_session_bridge_subagent_prompt_cache_ttl_seconds", + ), ( "httpResponsesSessionBridgeGatewaySafeMode", True, diff --git a/tests/integration/test_sticky_sessions_api.py b/tests/integration/test_sticky_sessions_api.py index 3694d75a0c..df929d97ab 100644 --- a/tests/integration/test_sticky_sessions_api.py +++ b/tests/integration/test_sticky_sessions_api.py @@ -13,6 +13,7 @@ from app.modules.accounts.repository import AccountsRepository from app.modules.proxy.durable_bridge_coordinator import DurableBridgeSessionCoordinator from app.modules.proxy.durable_bridge_repository import DurableBridgeRepository +from app.modules.proxy.sticky_repository import StickySessionsRepository from app.modules.settings.repository import SettingsRepository from app.modules.sticky_sessions.cleanup_scheduler import StickySessionCleanupScheduler @@ -61,26 +62,62 @@ async def _set_affinity_ttl(seconds: int) -> None: await session.commit() +@pytest.mark.asyncio +async def test_subagent_prompt_cache_purge_does_not_delete_parent_mapping(db_setup): + del db_setup + accounts = await _create_accounts() + async with SessionLocal() as session: + repository = StickySessionsRepository(session) + await repository.upsert( + "subagent-cache", + accounts[0].id, + kind=StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ) + await repository.upsert( + "parent-cache", + accounts[0].id, + kind=StickySessionKind.PROMPT_CACHE, + ) + stale_at = utcnow() - timedelta(seconds=31) + await session.execute( + text("UPDATE sticky_sessions SET updated_at = :stale_at WHERE key IN ('subagent-cache', 'parent-cache')"), + {"stale_at": stale_at}, + ) + await session.commit() + + deleted = await repository.purge_prompt_cache_before( + utcnow() - timedelta(seconds=30), + is_subagent=True, + ) + + assert deleted == 1 + assert await repository.get_entry("subagent-cache", kind=StickySessionKind.PROMPT_CACHE) is None + assert await repository.get_entry("parent-cache", kind=StickySessionKind.PROMPT_CACHE) is not None + + async def _insert_sticky_session( *, key: str, account_id: str, kind: StickySessionKind, updated_at_offset_seconds: int, + is_subagent: bool = False, ) -> None: timestamp = utcnow() - timedelta(seconds=updated_at_offset_seconds) async with SessionLocal() as session: await session.execute( text( """ - INSERT INTO sticky_sessions (key, account_id, kind, created_at, updated_at) - VALUES (:key, :account_id, :kind, :timestamp, :timestamp) + INSERT INTO sticky_sessions (key, account_id, kind, is_subagent, created_at, updated_at) + VALUES (:key, :account_id, :kind, :is_subagent, :timestamp, :timestamp) """ ), { "key": key, "account_id": account_id, "kind": kind.value, + "is_subagent": is_subagent, "timestamp": timestamp, }, ) @@ -307,6 +344,61 @@ async def test_sticky_sessions_api_lists_metadata_and_purges_stale(async_client) assert remaining_keys == {"prompt-cache-fresh", "codex-session-old"} +@pytest.mark.asyncio +async def test_sticky_sessions_api_stale_only_uses_subagent_ttl(async_client): + accounts = await _create_accounts() + async with SessionLocal() as session: + settings = await SettingsRepository(session).get_or_create() + settings.openai_cache_affinity_max_age_seconds = 600 + settings.http_responses_session_bridge_subagent_prompt_cache_ttl_seconds = 30 + await session.commit() + + await _insert_sticky_session( + key="parent-cache-fresh", + account_id=accounts[0].id, + kind=StickySessionKind.PROMPT_CACHE, + updated_at_offset_seconds=60, + ) + await _insert_sticky_session( + key="subagent-cache-stale", + account_id=accounts[0].id, + kind=StickySessionKind.PROMPT_CACHE, + updated_at_offset_seconds=60, + is_subagent=True, + ) + + response = await async_client.get("/api/sticky-sessions", params={"staleOnly": "true"}) + assert response.status_code == 200 + payload = response.json() + + assert payload["stalePromptCacheCount"] == 1 + assert payload["total"] == 1 + assert [(entry["key"], entry["isSubagent"], entry["isStale"]) for entry in payload["entries"]] == [ + ("subagent-cache-stale", True, True) + ] + + response = await async_client.post("/api/sticky-sessions/delete-filtered", json={"staleOnly": True}) + assert response.status_code == 200 + assert response.json()["deletedCount"] == 1 + + response = await async_client.get("/api/sticky-sessions") + assert response.status_code == 200 + remaining = {entry["key"]: entry for entry in response.json()["entries"]} + assert set(remaining) == {"parent-cache-fresh"} + assert remaining["parent-cache-fresh"]["isStale"] is False + + await _insert_sticky_session( + key="subagent-cache-stale-purge", + account_id=accounts[0].id, + kind=StickySessionKind.PROMPT_CACHE, + updated_at_offset_seconds=60, + is_subagent=True, + ) + response = await async_client.post("/api/sticky-sessions/purge", json={"staleOnly": True}) + assert response.status_code == 200 + assert response.json()["deletedCount"] == 1 + + @pytest.mark.asyncio async def test_sticky_sessions_api_filters_by_account_and_key(async_client): accounts = await _create_accounts() diff --git a/tests/unit/test_load_balancer_contract.py b/tests/unit/test_load_balancer_contract.py index 2650edc7ba..a9d6f31328 100644 --- a/tests/unit/test_load_balancer_contract.py +++ b/tests/unit/test_load_balancer_contract.py @@ -145,9 +145,10 @@ async def upsert( account_id: str, *, kind: StickySessionKind, + is_subagent: bool = False, ) -> StickySession: self.account_id = account_id - return StickySession(key=key, account_id=account_id, kind=kind) + return StickySession(key=key, account_id=account_id, kind=kind, is_subagent=is_subagent) async def delete(self, *args: Any, **kwargs: Any) -> bool: del args, kwargs diff --git a/tests/unit/test_otel.py b/tests/unit/test_otel.py index 3685a32bc1..300cff5aa5 100644 --- a/tests/unit/test_otel.py +++ b/tests/unit/test_otel.py @@ -288,6 +288,7 @@ def _init_background_db() -> None: monkeypatch.setattr(main, "get_settings", lambda: settings) monkeypatch.setattr(main, "get_settings_cache", lambda: settings_cache) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: settings_cache) monkeypatch.setattr(main, "ensure_auto_bootstrap_token", AsyncMock(return_value=None)) monkeypatch.setattr(main, "get_rate_limit_headers_cache", lambda: rate_limit_cache) monkeypatch.setattr(main, "reload_additional_quota_registry", lambda: None) @@ -372,6 +373,7 @@ async def _register(instance_id: str, *, endpoint_base_url: str | None = None) - monkeypatch.setattr(main, "get_settings", lambda: settings) monkeypatch.setattr(main, "get_settings_cache", lambda: settings_cache) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: settings_cache) monkeypatch.setattr(main, "ensure_auto_bootstrap_token", AsyncMock(return_value=None)) monkeypatch.setattr(main, "get_rate_limit_headers_cache", lambda: rate_limit_cache) monkeypatch.setattr(main, "reload_additional_quota_registry", lambda: None) @@ -464,6 +466,7 @@ async def start(self) -> None: monkeypatch.setattr(main, "get_settings", lambda: settings) monkeypatch.setattr(main, "get_settings_cache", lambda: settings_cache) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: settings_cache) monkeypatch.setattr(main, "ensure_auto_bootstrap_token", AsyncMock(return_value=None)) monkeypatch.setattr(main, "get_rate_limit_headers_cache", lambda: rate_limit_cache) monkeypatch.setattr(main, "reload_additional_quota_registry", lambda: None) @@ -638,6 +641,7 @@ async def _register(instance_id: str, *, endpoint_base_url: str | None = None) - monkeypatch.setattr(main, "get_settings", lambda: settings) monkeypatch.setattr(main, "get_settings_cache", lambda: settings_cache) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: settings_cache) monkeypatch.setattr(main, "ensure_auto_bootstrap_token", AsyncMock(return_value=None)) monkeypatch.setattr(main, "get_rate_limit_headers_cache", lambda: rate_limit_cache) monkeypatch.setattr(main, "reload_additional_quota_registry", lambda: None) @@ -713,6 +717,7 @@ async def test_lifespan_registers_bridge_without_waiting_for_advertise_self_prob monkeypatch.setattr(main, "get_settings", lambda: settings) monkeypatch.setattr(main, "get_settings_cache", lambda: settings_cache) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: settings_cache) monkeypatch.setattr(main, "ensure_auto_bootstrap_token", AsyncMock(return_value=None)) monkeypatch.setattr(main, "get_rate_limit_headers_cache", lambda: rate_limit_cache) monkeypatch.setattr(main, "reload_additional_quota_registry", lambda: None) @@ -786,6 +791,7 @@ async def test_lifespan_fails_fast_when_bridge_durable_schema_is_missing(monkeyp monkeypatch.setattr(main, "get_settings", lambda: settings) monkeypatch.setattr(main, "get_settings_cache", lambda: settings_cache) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: settings_cache) monkeypatch.setattr(main, "ensure_auto_bootstrap_token", AsyncMock(return_value=None)) monkeypatch.setattr(main, "get_rate_limit_headers_cache", lambda: rate_limit_cache) monkeypatch.setattr(main, "reload_additional_quota_registry", lambda: None) @@ -839,6 +845,7 @@ async def test_lifespan_allows_missing_bridge_schema_when_fail_fast_disabled(mon monkeypatch.setattr(main, "get_settings", lambda: settings) monkeypatch.setattr(main, "get_settings_cache", lambda: settings_cache) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: settings_cache) monkeypatch.setattr(main, "ensure_auto_bootstrap_token", AsyncMock(return_value=None)) monkeypatch.setattr(main, "get_rate_limit_headers_cache", lambda: rate_limit_cache) monkeypatch.setattr(main, "reload_additional_quota_registry", lambda: None) diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index a83fc51428..68b21455e9 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -10,7 +10,7 @@ from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any, cast -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, Mock, call import aiohttp import anyio @@ -44,6 +44,235 @@ pytestmark = pytest.mark.unit +@pytest.mark.asyncio +async def test_fork_session_closed_after_stream_end(monkeypatch: pytest.MonkeyPatch) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + + key = proxy_service._HTTPBridgeSessionKey( + "internal_unanchored_parallel", + "fork-key", + None, + ) + session = proxy_service._HTTPBridgeSession( + key=key, + headers={}, + affinity=proxy_service._AffinityPolicy(key="fork-key"), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE, plan_type="plus")), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=1.0, + idle_ttl_seconds=120.0, + ) + + event_queue: asyncio.Queue[str | None] = asyncio.Queue() + await event_queue.put(None) + + request_state = proxy_service._WebSocketRequestState( + request_id="req_fork_close", + response_id=None, + model="gpt-5.4", + service_tier="auto", + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + ) + request_state.event_queue = event_queue + + close_mock = AsyncMock() + monkeypatch.setattr(service, "_close_http_bridge_session", close_mock) + monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_submit_http_bridge_request", AsyncMock()) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: SimpleNamespace(sse_keepalive_interval_seconds=0), + ) + + gen = service._stream_http_bridge_session_events( + session=session, + request_state=request_state, + text_data="{}", + queue_limit=10, + propagate_http_errors=False, + downstream_turn_state=None, + ) + + async for _ in gen: + pass + + close_mock.assert_awaited_once_with(session) + + +@pytest.mark.asyncio +async def test_subagent_prompt_cache_ttl_closes_session_after_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "subagent-cache", None), + headers={}, + affinity=proxy_service._AffinityPolicy( + key="subagent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE, plan_type="plus")), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=1.0, + idle_ttl_seconds=120.0, + is_subagent=True, + subagent_prompt_cache_ttl_seconds=30, + ) + + event_queue: asyncio.Queue[str | None] = asyncio.Queue() + await event_queue.put(None) + request_state = proxy_service._WebSocketRequestState( + request_id="req_subagent_reuse", + response_id=None, + model="gpt-5.4", + service_tier="auto", + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + ) + request_state.event_queue = event_queue + + close_mock = AsyncMock() + refresh_mock = AsyncMock() + monkeypatch.setattr(service, "_close_http_bridge_session", close_mock) + monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_submit_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_refresh_completed_subagent_sticky_mapping", refresh_mock) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: SimpleNamespace(sse_keepalive_interval_seconds=0), + ) + + gen = service._stream_http_bridge_session_events( + session=session, + request_state=request_state, + text_data="{}", + queue_limit=10, + propagate_http_errors=False, + downstream_turn_state=None, + ) + + async for _ in gen: + pass + + refresh_mock.assert_awaited_once_with(session) + close_mock.assert_awaited_once_with(session) + + +@pytest.mark.asyncio +async def test_refresh_completed_subagent_sticky_mapping_touches_prompt_cache_entry() -> None: + sticky_repo = SimpleNamespace(upsert=AsyncMock()) + + class _Repos: + sticky_sessions = sticky_repo + + async def __aenter__(self) -> "_Repos": + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + service = proxy_service.ProxyService(cast(Any, lambda: _Repos())) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "subagent-cache", None), + headers={}, + affinity=proxy_service._AffinityPolicy( + key="subagent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE)), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=1.0, + idle_ttl_seconds=120.0, + is_subagent=True, + subagent_prompt_cache_ttl_seconds=30, + ) + + await service._refresh_completed_subagent_sticky_mapping(session) + + sticky_repo.upsert.assert_awaited_once_with( + "subagent-cache", + "acc-1", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ) + + +@pytest.mark.asyncio +async def test_delete_completed_subagent_sticky_mapping_removes_retained_no_cache_key() -> None: + sticky_repo = SimpleNamespace(delete=AsyncMock()) + + class _Repos: + sticky_sessions = sticky_repo + + async def __aenter__(self) -> "_Repos": + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("session_header", "subagent:scope-no-cache:parent-cache", None), + headers={}, + affinity=proxy_service._AffinityPolicy( + key="subagent:scope-no-cache:parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE)), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=1.0, + idle_ttl_seconds=120.0, + is_subagent=True, + subagent_prompt_cache_ttl_seconds=None, + ) + + await http_bridge_helpers_module._delete_completed_subagent_sticky_mapping(lambda: _Repos(), session) + + sticky_repo.delete.assert_has_awaits( + [ + call( + "subagent:parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ), + call( + "subagent:scope-no-cache:parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ), + ], + any_order=True, + ) + assert sticky_repo.delete.await_count == 2 + + @pytest.fixture(autouse=True) def _share_proxy_dashboard_settings(monkeypatch: pytest.MonkeyPatch) -> None: class _SettingsCache: @@ -109,6 +338,51 @@ def test_forwarded_fork_keeps_authenticated_original_unanchored_state() -> None: ) +def test_subagent_bridge_key_namespaces_non_prompt_cache_keys() -> None: + key, affinity, idle_ttl, subagent_ttl = http_bridge_helpers_module._subagent_prompt_cache_bridge_key( + proxy_service._HTTPBridgeSessionKey("session_header", "parent-session", None), + proxy_service._AffinityPolicy( + key="parent-session", + kind=proxy_service.StickySessionKind.CODEX_SESSION, + max_age_seconds=600, + ), + is_subagent=True, + forwarded_request=False, + idle_ttl_seconds=120.0, + request_scope_id="request-scope-1", + subagent_prompt_cache_ttl_seconds=30, + ) + + assert key.affinity_key == "subagent:request-scope-1:parent-session" + assert affinity.key == "subagent:request-scope-1:parent-session" + assert affinity.kind is proxy_service.StickySessionKind.CODEX_SESSION + assert affinity.max_age_seconds == 600 + assert idle_ttl == 120.0 + assert subagent_ttl is None + + +def test_subagent_bridge_key_preserves_forwarded_or_namespaced_prompt_cache_key() -> None: + forwarded_key, forwarded_affinity, forwarded_idle_ttl, forwarded_subagent_ttl = ( + http_bridge_helpers_module._subagent_prompt_cache_bridge_key( + proxy_service._HTTPBridgeSessionKey("prompt_cache", "subagent:parent-cache", None), + proxy_service._AffinityPolicy( + key="subagent:parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + is_subagent=True, + forwarded_request=True, + idle_ttl_seconds=120.0, + request_scope_id="request-scope-1", + subagent_prompt_cache_ttl_seconds=30, + ) + ) + + assert forwarded_key.affinity_key == "subagent:parent-cache" + assert forwarded_affinity.key == "subagent:parent-cache" + assert forwarded_idle_ttl == 120.0 + assert forwarded_subagent_ttl == 30 + + @pytest.mark.asyncio async def test_legacy_forward_anchor_lookup_accepts_registered_turn_state_alias() -> None: key = proxy_service._HTTPBridgeSessionKey("session_header", "sid-123", None) @@ -4034,6 +4308,7 @@ async def fake_relay(_session: proxy_service._HTTPBridgeSession) -> None: assert "transfer-encoding" not in forwarded assert "upgrade" not in forwarded assert "x-codex-turn-metadata" not in forwarded + assert "x-parent-session-id" not in forwarded assert "x-openai-subagent" not in forwarded assert "x-codex-parent-thread-id" not in forwarded assert "x-codex-window-id" not in forwarded @@ -4129,6 +4404,7 @@ async def open_upstream(_account: object, headers: dict[str, str], **_: object) assert "transfer-encoding" not in forwarded assert "upgrade" not in forwarded assert "x-codex-turn-metadata" not in forwarded + assert "x-parent-session-id" not in forwarded assert "x-openai-subagent" not in forwarded assert "x-codex-parent-thread-id" not in forwarded assert "x-codex-window-id" not in forwarded @@ -5769,7 +6045,7 @@ async def fake_get_or_create( @pytest.mark.asyncio -async def test_stream_via_http_bridge_prefers_durable_account_for_soft_prompt_cache_follow_up_recovery( +async def test_stream_via_http_bridge_isolates_subagent_key_before_durable_lookup( monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) @@ -5777,12 +6053,12 @@ async def test_stream_via_http_bridge_prefers_durable_account_for_soft_prompt_ca { "model": "gpt-5.4", "instructions": "hi", - "input": "hello again", - "prompt_cache_key": "thread-soft-follow-up", + "input": "hello", + "prompt_cache_key": "parent-cache", }, ) request_state = proxy_service._WebSocketRequestState( - request_id="req-soft-prompt-cache-follow-up", + request_id="req-subagent-durable-lookup", model="gpt-5.4", service_tier=None, reasoning_effort=None, @@ -5794,10 +6070,9 @@ async def test_stream_via_http_bridge_prefers_durable_account_for_soft_prompt_ca event_queue = request_state.event_queue assert event_queue is not None await event_queue.put(None) - captured: dict[str, object] = {} def fake_prepare( - prepared_payload: proxy_service.ResponsesRequest, + _prepared_payload: proxy_service.ResponsesRequest, _headers: dict[str, str] | Any, *, api_key: proxy_service.ApiKeyData | None, @@ -5806,18 +6081,21 @@ def fake_prepare( client_ip: str | None = None, ) -> tuple[proxy_service._WebSocketRequestState, str]: del api_key, api_key_reservation, request_id, client_ip - captured["previous_response_id"] = prepared_payload.previous_response_id return request_state, '{"type":"response.create"}' session = proxy_service._HTTPBridgeSession( - key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "thread-soft-follow-up", None), + key=proxy_service._HTTPBridgeSessionKey( + "prompt_cache", + "subagent:scope-subagent-durable-lookup:parent-cache", + None, + ), headers={}, affinity=proxy_service._AffinityPolicy( - key="thread-soft-follow-up", + key="subagent:scope-subagent-durable-lookup:parent-cache", kind=proxy_service.StickySessionKind.PROMPT_CACHE, ), request_model="gpt-5.4", - account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE)), + account=cast(Any, SimpleNamespace(id="acc-subagent", status=AccountStatus.ACTIVE)), upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), upstream_control=proxy_service._WebSocketUpstreamControl(), pending_requests=deque(), @@ -5827,6 +6105,12 @@ def fake_prepare( last_used_at=1.0, idle_ttl_seconds=120.0, ) + lookup_request_targets = AsyncMock(return_value=None) + legacy_anchor_calls: list[dict[str, Any]] = [] + + async def fake_legacy_forward_anchor_lookup(**kwargs: Any) -> None: + legacy_anchor_calls.append(kwargs) + return None monkeypatch.setattr( proxy_service, @@ -5840,6 +6124,7 @@ def fake_prepare( openai_cache_affinity_max_age_seconds=1800, http_responses_session_bridge_prompt_cache_idle_ttl_seconds=3600, http_responses_session_bridge_gateway_safe_mode=False, + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=0, ) ) ), @@ -5847,64 +6132,192 @@ def fake_prepare( ) monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) monkeypatch.setattr( - service._durable_bridge, - "lookup_request_targets", - AsyncMock( - return_value=proxy_service.DurableBridgeLookup( - session_id="sess-soft-follow-up", - canonical_kind="prompt_cache", - canonical_key="thread-soft-follow-up", - api_key_scope="__anonymous__", - account_id="acc-1", - owner_instance_id="instance-a", - owner_epoch=1, - lease_expires_at=datetime.now(timezone.utc), - state=HttpBridgeSessionState.ACTIVE, - latest_turn_state="http_turn_soft_follow_up", - latest_response_id="resp_latest", - ) - ), + http_bridge_streaming_module, + "_legacy_forward_anchor_lookup", + fake_legacy_forward_anchor_lookup, ) + monkeypatch.setattr(service._durable_bridge, "lookup_request_targets", lookup_request_targets) monkeypatch.setattr(service, "_prepare_http_bridge_request", fake_prepare) - - async def fake_get_or_create( - *args: object, - **kwargs: object, - ) -> proxy_service._HTTPBridgeSession: - captured["preferred_account_id"] = kwargs.get("preferred_account_id") - captured["request_stage"] = kwargs.get("request_stage") - return session - - monkeypatch.setattr(service, "_get_or_create_http_bridge_session", fake_get_or_create) + monkeypatch.setattr(service, "_get_or_create_http_bridge_session", AsyncMock(return_value=session)) monkeypatch.setattr(service, "_submit_http_bridge_request", AsyncMock()) monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) - chunks = [ - chunk - async for chunk in service._stream_via_http_bridge( - payload, - headers={"x-codex-turn-state": "http_turn_soft_follow_up"}, - codex_session_affinity=True, - propagate_http_errors=False, - openai_cache_affinity=True, - api_key=None, - api_key_reservation=None, - suppress_text_done_events=False, - idle_ttl_seconds=120.0, - codex_idle_ttl_seconds=1800.0, - max_sessions=8, - queue_limit=4, - ) - ] + request_scope_token = set_request_scope_id("scope-subagent-durable-lookup") + try: + chunks = [ + chunk + async for chunk in service._stream_via_http_bridge( + payload, + headers={ + "x-parent-session-id": "parent-session", + "x-codex-session-id": "parent-session", + "x-codex-turn-state": "parent-turn-state", + }, + codex_session_affinity=False, + propagate_http_errors=False, + openai_cache_affinity=True, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + idle_ttl_seconds=120.0, + codex_idle_ttl_seconds=1800.0, + max_sessions=8, + queue_limit=4, + ) + ] + finally: + reset_request_scope_id(request_scope_token) assert chunks == [] - assert captured["previous_response_id"] is None - assert captured["request_stage"] == "follow_up" - assert captured["preferred_account_id"] == "acc-1" + lookup_request_targets.assert_awaited_once_with( + session_key_kind="turn_state_header", + session_key_value="subagent:scope-subagent-durable-lookup:parent-turn-state", + api_key_id=None, + turn_state=None, + session_header=None, + previous_response_id=None, + ) + assert len(legacy_anchor_calls) == 1 + assert legacy_anchor_calls[0]["turn_state"] is None @pytest.mark.asyncio -async def test_close_http_bridge_session_fails_pending_downstream_requests() -> None: +async def test_stream_via_http_bridge_prefers_durable_account_for_soft_prompt_cache_follow_up_recovery( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + payload = proxy_service.ResponsesRequest.model_validate( + { + "model": "gpt-5.4", + "instructions": "hi", + "input": "hello again", + "prompt_cache_key": "thread-soft-follow-up", + }, + ) + request_state = proxy_service._WebSocketRequestState( + request_id="req-soft-prompt-cache-follow-up", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=1.0, + event_queue=asyncio.Queue(), + transport="http", + ) + event_queue = request_state.event_queue + assert event_queue is not None + await event_queue.put(None) + captured: dict[str, object] = {} + + def fake_prepare( + prepared_payload: proxy_service.ResponsesRequest, + _headers: dict[str, str] | Any, + *, + api_key: proxy_service.ApiKeyData | None, + api_key_reservation: proxy_service.ApiKeyUsageReservationData | None, + request_id: str, + client_ip: str | None = None, + ) -> tuple[proxy_service._WebSocketRequestState, str]: + del api_key, api_key_reservation, request_id, client_ip + captured["previous_response_id"] = prepared_payload.previous_response_id + return request_state, '{"type":"response.create"}' + + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "thread-soft-follow-up", None), + headers={}, + affinity=proxy_service._AffinityPolicy( + key="thread-soft-follow-up", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE)), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=1.0, + idle_ttl_seconds=120.0, + ) + + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: cast( + Any, + SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + sticky_threads_enabled=False, + openai_cache_affinity_max_age_seconds=1800, + http_responses_session_bridge_prompt_cache_idle_ttl_seconds=3600, + http_responses_session_bridge_gateway_safe_mode=False, + ) + ) + ), + ), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + service._durable_bridge, + "lookup_request_targets", + AsyncMock( + return_value=proxy_service.DurableBridgeLookup( + session_id="sess-soft-follow-up", + canonical_kind="prompt_cache", + canonical_key="thread-soft-follow-up", + api_key_scope="__anonymous__", + account_id="acc-1", + owner_instance_id="instance-a", + owner_epoch=1, + lease_expires_at=datetime.now(timezone.utc), + state=HttpBridgeSessionState.ACTIVE, + latest_turn_state="http_turn_soft_follow_up", + latest_response_id="resp_latest", + ) + ), + ) + monkeypatch.setattr(service, "_prepare_http_bridge_request", fake_prepare) + + async def fake_get_or_create( + *args: object, + **kwargs: object, + ) -> proxy_service._HTTPBridgeSession: + captured["preferred_account_id"] = kwargs.get("preferred_account_id") + captured["request_stage"] = kwargs.get("request_stage") + return session + + monkeypatch.setattr(service, "_get_or_create_http_bridge_session", fake_get_or_create) + monkeypatch.setattr(service, "_submit_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) + + chunks = [ + chunk + async for chunk in service._stream_via_http_bridge( + payload, + headers={"x-codex-turn-state": "http_turn_soft_follow_up"}, + codex_session_affinity=True, + propagate_http_errors=False, + openai_cache_affinity=True, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + idle_ttl_seconds=120.0, + codex_idle_ttl_seconds=1800.0, + max_sessions=8, + queue_limit=4, + ) + ] + + assert chunks == [] + assert captured["previous_response_id"] is None + assert captured["request_stage"] == "follow_up" + assert captured["preferred_account_id"] == "acc-1" + + +@pytest.mark.asyncio +async def test_close_http_bridge_session_fails_pending_downstream_requests() -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) event_queue: asyncio.Queue[str | None] = asyncio.Queue() request_state = proxy_service._WebSocketRequestState( @@ -9710,6 +10123,8 @@ async def fake_create_http_bridge_session( preferred_account_id: str | None = None, require_preferred_account: bool = False, fallback_on_preferred_account_unavailable: bool = True, + is_subagent: bool = False, + subagent_prompt_cache_ttl_seconds: int | None = None, ) -> proxy_service._HTTPBridgeSession: del ( headers, @@ -9721,6 +10136,8 @@ async def fake_create_http_bridge_session( preferred_account_id, require_preferred_account, fallback_on_preferred_account_unavailable, + is_subagent, + subagent_prompt_cache_ttl_seconds, ) captured["key"] = create_key return created_session @@ -9795,6 +10212,8 @@ async def fake_create_http_bridge_session( preferred_account_id: str | None = None, require_preferred_account: bool = False, fallback_on_preferred_account_unavailable: bool = True, + is_subagent: bool = False, + subagent_prompt_cache_ttl_seconds: int | None = None, ) -> proxy_service._HTTPBridgeSession: del ( headers, @@ -9806,6 +10225,8 @@ async def fake_create_http_bridge_session( preferred_account_id, require_preferred_account, fallback_on_preferred_account_unavailable, + is_subagent, + subagent_prompt_cache_ttl_seconds, ) captured["key"] = create_key return created_session @@ -9881,6 +10302,8 @@ async def fake_create_http_bridge_session( preferred_account_id: str | None = None, require_preferred_account: bool = False, fallback_on_preferred_account_unavailable: bool = True, + is_subagent: bool = False, + subagent_prompt_cache_ttl_seconds: int | None = None, ) -> proxy_service._HTTPBridgeSession: del ( headers, @@ -9892,6 +10315,8 @@ async def fake_create_http_bridge_session( preferred_account_id, require_preferred_account, fallback_on_preferred_account_unavailable, + is_subagent, + subagent_prompt_cache_ttl_seconds, ) captured["key"] = create_key return created_session @@ -9938,6 +10363,402 @@ async def fake_create_http_bridge_session( assert captured["key"] == requested_key +@pytest.mark.asyncio +async def test_get_or_create_http_bridge_session_namespaces_subagent_prompt_cache_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + requested_key = proxy_service._HTTPBridgeSessionKey("session_header", "parent-session", None) + created_session = proxy_service._HTTPBridgeSession( + key=requested_key, + headers={}, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE)), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=2.0, + idle_ttl_seconds=120.0, + ) + captured: dict[str, object] = {} + + async def fake_create_http_bridge_session( + create_key: proxy_service._HTTPBridgeSessionKey, + *, + headers: dict[str, str], + affinity: proxy_service._AffinityPolicy, + api_key: proxy_service.ApiKeyData | None, + request_model: str | None, + idle_ttl_seconds: float, + request_stage: str = "first_turn", + preferred_account_id: str | None = None, + require_preferred_account: bool = False, + fallback_on_preferred_account_unavailable: bool = True, + is_subagent: bool = False, + subagent_prompt_cache_ttl_seconds: int | None = None, + ) -> proxy_service._HTTPBridgeSession: + del ( + headers, + api_key, + request_model, + request_stage, + preferred_account_id, + require_preferred_account, + fallback_on_preferred_account_unavailable, + ) + captured.update( + { + "key": create_key, + "affinity": affinity, + "idle_ttl_seconds": idle_ttl_seconds, + "is_subagent": is_subagent, + "subagent_prompt_cache_ttl_seconds": subagent_prompt_cache_ttl_seconds, + } + ) + return created_session + + monkeypatch.setattr(service, "_prune_http_bridge_sessions_locked", Mock(return_value=[])) + monkeypatch.setattr(service, "_create_http_bridge_session", fake_create_http_bridge_session) + monkeypatch.setattr(service, "_claim_durable_http_bridge_session", AsyncMock()) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: _make_app_settings(http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=30), + ) + monkeypatch.setattr( + http_bridge_mixin_module, + "_service_get_settings_cache", + lambda: SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=30, + ) + ) + ), + ) + monkeypatch.setattr(proxy_service, "_http_bridge_owner_instance", AsyncMock(return_value="instance-a")) + monkeypatch.setattr( + proxy_service, + "_active_http_bridge_instance_ring", + AsyncMock(return_value=("instance-a", ["instance-a", "instance-b"])), + ) + + resolved = await service._get_or_create_http_bridge_session( + requested_key, + headers={ + "x-parent-session-id": "parent-session-1", + "x-codex-session-id": "parent-session", + }, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + max_sessions=8, + ) + + assert resolved is created_session + subagent_key = captured["key"] + assert isinstance(subagent_key, proxy_service._HTTPBridgeSessionKey) + assert subagent_key.affinity_kind == "prompt_cache" + assert subagent_key.affinity_key == "subagent:parent-cache" + subagent_affinity = captured["affinity"] + assert isinstance(subagent_affinity, proxy_service._AffinityPolicy) + assert subagent_affinity.key == "subagent:parent-cache" + assert captured["idle_ttl_seconds"] == 30.0 + assert captured["is_subagent"] is True + assert captured["subagent_prompt_cache_ttl_seconds"] == 30 + + +@pytest.mark.asyncio +async def test_get_or_create_http_bridge_session_refreshes_reused_subagent_ttl( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + requested_key = proxy_service._HTTPBridgeSessionKey("session_header", "parent-session", None) + stable_subagent_key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "subagent:parent-cache", None) + session = proxy_service._HTTPBridgeSession( + key=stable_subagent_key, + headers={}, + affinity=proxy_service._AffinityPolicy( + key="subagent:parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + max_age_seconds=30, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE, plan_type="plus")), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=2.0, + idle_ttl_seconds=30.0, + is_subagent=True, + subagent_prompt_cache_ttl_seconds=30, + ) + service._http_bridge_sessions[stable_subagent_key] = session + + monkeypatch.setattr(service, "_prune_http_bridge_sessions_locked", Mock(return_value=[])) + monkeypatch.setattr(service, "_create_http_bridge_session", AsyncMock()) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: _make_app_settings(http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=45), + ) + monkeypatch.setattr( + http_bridge_mixin_module, + "_service_get_settings_cache", + lambda: SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=45, + ) + ) + ), + ) + + resolved = await service._get_or_create_http_bridge_session( + requested_key, + headers={ + "x-parent-session-id": "parent-session-1", + "x-codex-session-id": "parent-session", + }, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + max_sessions=8, + ) + + assert resolved is session + assert session.idle_ttl_seconds == 45.0 + assert session.subagent_prompt_cache_ttl_seconds == 45 + + +@pytest.mark.asyncio +async def test_get_or_create_http_bridge_session_no_cache_subagent_uses_request_scoped_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + requested_key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "parent-cache", None) + created_session = proxy_service._HTTPBridgeSession( + key=requested_key, + headers={}, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE)), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=2.0, + idle_ttl_seconds=120.0, + ) + captured: dict[str, object] = {} + + async def fake_create_http_bridge_session( + create_key: proxy_service._HTTPBridgeSessionKey, + *, + headers: dict[str, str], + affinity: proxy_service._AffinityPolicy, + api_key: proxy_service.ApiKeyData | None, + request_model: str | None, + idle_ttl_seconds: float, + request_stage: str = "first_turn", + preferred_account_id: str | None = None, + require_preferred_account: bool = False, + fallback_on_preferred_account_unavailable: bool = True, + is_subagent: bool = False, + subagent_prompt_cache_ttl_seconds: int | None = None, + ) -> proxy_service._HTTPBridgeSession: + del ( + headers, + affinity, + api_key, + request_model, + idle_ttl_seconds, + request_stage, + preferred_account_id, + require_preferred_account, + fallback_on_preferred_account_unavailable, + is_subagent, + subagent_prompt_cache_ttl_seconds, + ) + captured["key"] = create_key + return created_session + + monkeypatch.setattr(service, "_prune_http_bridge_sessions_locked", Mock(return_value=[])) + monkeypatch.setattr(service, "_create_http_bridge_session", fake_create_http_bridge_session) + monkeypatch.setattr(service, "_claim_durable_http_bridge_session", AsyncMock()) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: _make_app_settings(http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=0), + ) + monkeypatch.setattr( + http_bridge_mixin_module, + "_service_get_settings_cache", + lambda: SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=0, + ) + ) + ), + ) + monkeypatch.setattr(proxy_service, "_http_bridge_owner_instance", AsyncMock(return_value="instance-a")) + monkeypatch.setattr( + proxy_service, + "_active_http_bridge_instance_ring", + AsyncMock(return_value=("instance-a", ["instance-a", "instance-b"])), + ) + + request_scope_token = set_request_scope_id("scope-subagent-no-cache") + try: + resolved = await service._get_or_create_http_bridge_session( + requested_key, + headers={"x-parent-session-id": "parent-session-1"}, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + max_sessions=8, + ) + finally: + reset_request_scope_id(request_scope_token) + + assert resolved is created_session + subagent_key = captured["key"] + assert isinstance(subagent_key, proxy_service._HTTPBridgeSessionKey) + assert subagent_key.affinity_key == "subagent:scope-subagent-no-cache:parent-cache" + + +@pytest.mark.asyncio +async def test_get_or_create_http_bridge_session_subagent_ignores_parent_turn_state_alias( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + parent_key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "parent-cache", None) + parent_session = proxy_service._HTTPBridgeSession( + key=parent_key, + headers={"x-codex-turn-state": "parent-turn-state"}, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-parent", status=AccountStatus.ACTIVE)), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=2.0, + idle_ttl_seconds=120.0, + ) + service._http_bridge_sessions[parent_key] = parent_session + service._http_bridge_turn_state_index[ + proxy_service._http_bridge_turn_state_alias_key("parent-turn-state", None) + ] = parent_key + requested_key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "parent-cache", None) + created_session = proxy_service._HTTPBridgeSession( + key=requested_key, + headers={}, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + request_model="gpt-5.4", + account=cast(Any, SimpleNamespace(id="acc-child", status=AccountStatus.ACTIVE)), + upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=2.0, + idle_ttl_seconds=120.0, + ) + create_session = AsyncMock(return_value=created_session) + + monkeypatch.setattr(service, "_prune_http_bridge_sessions_locked", Mock(return_value=[])) + monkeypatch.setattr(service, "_create_http_bridge_session", create_session) + monkeypatch.setattr(service, "_claim_durable_http_bridge_session", AsyncMock()) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: _make_app_settings(http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=0), + ) + monkeypatch.setattr( + http_bridge_mixin_module, + "_service_get_settings_cache", + lambda: SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=0, + ) + ) + ), + ) + monkeypatch.setattr(proxy_service, "_http_bridge_owner_instance", AsyncMock(return_value="instance-a")) + monkeypatch.setattr( + proxy_service, + "_active_http_bridge_instance_ring", + AsyncMock(return_value=("instance-a", ["instance-a", "instance-b"])), + ) + + request_scope_token = set_request_scope_id("scope-subagent-no-alias") + try: + resolved = await service._get_or_create_http_bridge_session( + requested_key, + headers={ + "x-parent-session-id": "parent-session", + "x-codex-turn-state": "parent-turn-state", + }, + affinity=proxy_service._AffinityPolicy( + key="parent-cache", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + ), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + max_sessions=8, + ) + finally: + reset_request_scope_id(request_scope_token) + + assert resolved is created_session + assert resolved is not parent_session + create_call = create_session.await_args + assert create_call is not None + create_key = create_call.args[0] + assert isinstance(create_key, proxy_service._HTTPBridgeSessionKey) + assert create_key.affinity_key == "subagent:scope-subagent-no-alias:parent-cache" + + @pytest.mark.asyncio async def test_get_or_create_http_bridge_session_recovers_from_previous_response_id_mapping( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_proxy_load_balancer_refresh.py b/tests/unit/test_proxy_load_balancer_refresh.py index 556373cd1d..4d53729b48 100644 --- a/tests/unit/test_proxy_load_balancer_refresh.py +++ b/tests/unit/test_proxy_load_balancer_refresh.py @@ -194,12 +194,14 @@ async def get_account_id( ) -> str | None: return None - async def upsert(self, key: str, account_id: str, *, kind: StickySessionKind) -> StickySession: + async def upsert( + self, key: str, account_id: str, *, kind: StickySessionKind, is_subagent: bool = False + ) -> StickySession: row = self._build_row(key, account_id, kind) self.upserts.append(row) return row - async def delete(self, key: str, *, kind: StickySessionKind | None = None) -> bool: + async def delete(self, key: str, *, kind: StickySessionKind | None = None, is_subagent: bool | None = None) -> bool: self.deletes.append((key, kind)) return False diff --git a/tests/unit/test_select_with_stickiness.py b/tests/unit/test_select_with_stickiness.py index b7d6f00bb6..f0229a0abc 100644 --- a/tests/unit/test_select_with_stickiness.py +++ b/tests/unit/test_select_with_stickiness.py @@ -65,12 +65,13 @@ def _make_sticky_repo(existing_account_id: str | None = None) -> AsyncMock: async def _invoke_stickiness( states: list[AccountState], - sticky_key: str, + sticky_key: str | None, sticky_repo: AsyncMock, *, - sticky_kind: StickySessionKind = StickySessionKind.PROMPT_CACHE, + sticky_kind: StickySessionKind | None = StickySessionKind.PROMPT_CACHE, reallocate_sticky: bool = False, sticky_max_age_seconds: int | None = 600, + sticky_is_subagent: bool = False, budget_threshold_pct: float = 95.0, secondary_budget_threshold_pct: float = 100.0, routing_strategy: RoutingStrategy = "usage_weighted", @@ -98,6 +99,7 @@ async def mock_repo_factory(): sticky_kind=sticky_kind, reallocate_sticky=reallocate_sticky, sticky_max_age_seconds=sticky_max_age_seconds, + sticky_is_subagent=sticky_is_subagent, budget_threshold_pct=budget_threshold_pct, secondary_budget_threshold_pct=secondary_budget_threshold_pct, prefer_earlier_reset_accounts=False, @@ -110,6 +112,51 @@ async def mock_repo_factory(): ) +@pytest.mark.asyncio +async def test_no_cache_subagent_skips_sticky_mapping() -> None: + repo = _make_sticky_repo() + + result = await _invoke_stickiness( + [_active("a")], + None, + repo, + sticky_kind=None, + sticky_max_age_seconds=None, + ) + + assert result.account is not None + assert result.account.account_id == "a" + repo.get_account_id.assert_not_called() + repo.upsert.assert_not_called() + repo.delete.assert_not_called() + + +@pytest.mark.asyncio +async def test_positive_subagent_ttl_uses_sticky_mapping() -> None: + repo = _make_sticky_repo() + + result = await _invoke_stickiness( + [_active("a")], + "subagent-cache-key", + repo, + sticky_max_age_seconds=120, + sticky_is_subagent=True, + ) + + assert result.account is not None + repo.get_account_id.assert_awaited_once_with( + "subagent-cache-key", + kind=StickySessionKind.PROMPT_CACHE, + max_age_seconds=120, + ) + repo.upsert.assert_awaited_once_with( + "subagent-cache-key", + "a", + kind=StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ) + + # --------------------------------------------------------------------------- # Fix 1+3: sticky session is preserved when pinned account is temporarily down # --------------------------------------------------------------------------- @@ -178,7 +225,7 @@ async def test_fallback_overwrites_sticky_when_reallocate_sticky_true(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once() - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.STICKY_THREAD) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.STICKY_THREAD, is_subagent=False) @pytest.mark.asyncio @@ -234,7 +281,7 @@ async def test_sticky_deleted_when_pinned_account_removed_from_pool(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key1", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) # --------------------------------------------------------------------------- @@ -259,7 +306,7 @@ async def test_all_accounts_exhausted_keeps_pinned_no_thrashing(): assert result.account is not None assert result.account.account_id == "a" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -279,7 +326,7 @@ async def test_pool_exhausted_but_better_candidate_exists_reallocates(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key1", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -301,7 +348,7 @@ async def test_round_robin_pool_health_check_prefers_budget_safe_candidate(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key-round-robin", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key-round-robin", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key-round-robin", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -364,7 +411,9 @@ async def test_capacity_weighted_pool_health_check_prefers_budget_safe_candidate assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key-capacity-weighted", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key-capacity-weighted", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with( + "key-capacity-weighted", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False + ) @pytest.mark.asyncio @@ -382,7 +431,7 @@ async def test_pool_exhausted_single_account_keeps_pinned(): assert result.account is not None assert result.account.account_id == "a" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -402,7 +451,7 @@ async def test_pool_exhausted_with_custom_threshold(): assert result.account is not None assert result.account.account_id == "a" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -421,7 +470,7 @@ async def test_pool_exhausted_candidate_with_none_usage_triggers_reallocation(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key1", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -440,7 +489,7 @@ async def test_first_request_creates_sticky_mapping(): assert result.account is not None assert result.account.account_id == "a" - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -477,7 +526,9 @@ async def test_first_sticky_request_honors_relative_availability_tuning(): assert result.account is not None assert result.account.account_id == "a" - repo.upsert.assert_called_once_with("key-relative-availability", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with( + "key-relative-availability", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False + ) # --------------------------------------------------------------------------- @@ -503,7 +554,7 @@ async def test_grace_period_returns_pinned_when_reset_imminent(): assert result.account is not None assert result.account.account_id == "a" - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -523,7 +574,7 @@ async def test_grace_period_keeps_rate_limited_pinned_account_even_when_usage_is assert result.account is not None assert result.account.account_id == "a" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -640,7 +691,7 @@ async def test_paused_pinned_account_persists_fallback(): assert result.account is not None assert result.account.account_id == "b" - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -659,7 +710,7 @@ async def test_reauth_required_pinned_account_persists_fallback(): assert result.account is not None assert result.account.account_id == "b" - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) # --------------------------------------------------------------------------- @@ -712,7 +763,7 @@ async def test_codex_session_persists_fallback_during_outage(): assert result.account is not None assert result.account.account_id == "b" - repo.upsert.assert_called_once_with("session_123", "b", kind=StickySessionKind.CODEX_SESSION) + repo.upsert.assert_called_once_with("session_123", "b", kind=StickySessionKind.CODEX_SESSION, is_subagent=False) @pytest.mark.asyncio @@ -734,7 +785,7 @@ async def test_rate_limit_far_away_does_not_reallocate_codex_session_affinity(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("session_123", "b", kind=StickySessionKind.CODEX_SESSION) + repo.upsert.assert_called_once_with("session_123", "b", kind=StickySessionKind.CODEX_SESSION, is_subagent=False) @pytest.mark.asyncio @@ -775,7 +826,7 @@ async def test_budget_exhaustion_triggers_reallocation(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key1", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -795,7 +846,7 @@ async def test_budget_threshold_80_triggers_at_85_percent(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key1", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -815,7 +866,7 @@ async def test_budget_threshold_95_no_reallocation_at_85_percent(): assert result.account is not None assert result.account.account_id == "a" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -837,7 +888,9 @@ async def test_budget_threshold_reallocates_codex_session_affinity(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("codex-session-123", kind=StickySessionKind.CODEX_SESSION) - repo.upsert.assert_called_once_with("codex-session-123", "b", kind=StickySessionKind.CODEX_SESSION) + repo.upsert.assert_called_once_with( + "codex-session-123", "b", kind=StickySessionKind.CODEX_SESSION, is_subagent=False + ) @pytest.mark.asyncio @@ -863,7 +916,7 @@ async def test_budget_threshold_reallocates_sticky_thread_affinity(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("thread-X", kind=StickySessionKind.STICKY_THREAD) - repo.upsert.assert_called_once_with("thread-X", "b", kind=StickySessionKind.STICKY_THREAD) + repo.upsert.assert_called_once_with("thread-X", "b", kind=StickySessionKind.STICKY_THREAD, is_subagent=False) @pytest.mark.asyncio @@ -885,7 +938,9 @@ async def test_budget_threshold_reallocates_to_primary_safe_secondary_pressured_ assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("codex-session-123", kind=StickySessionKind.CODEX_SESSION) - repo.upsert.assert_called_once_with("codex-session-123", "b", kind=StickySessionKind.CODEX_SESSION) + repo.upsert.assert_called_once_with( + "codex-session-123", "b", kind=StickySessionKind.CODEX_SESSION, is_subagent=False + ) @pytest.mark.asyncio @@ -988,7 +1043,9 @@ async def test_fresh_sticky_mapping_uses_normal_budget_gate(): assert result.account is not None assert result.account.account_id == "a" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("new-codex-session", "a", kind=StickySessionKind.CODEX_SESSION) + repo.upsert.assert_called_once_with( + "new-codex-session", "a", kind=StickySessionKind.CODEX_SESSION, is_subagent=False + ) @pytest.mark.asyncio @@ -1011,7 +1068,9 @@ async def test_secondary_budget_threshold_controls_sticky_reallocation(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("codex-session-123", kind=StickySessionKind.CODEX_SESSION) - repo.upsert.assert_called_once_with("codex-session-123", "b", kind=StickySessionKind.CODEX_SESSION) + repo.upsert.assert_called_once_with( + "codex-session-123", "b", kind=StickySessionKind.CODEX_SESSION, is_subagent=False + ) @pytest.mark.asyncio @@ -1033,7 +1092,7 @@ async def test_rate_limit_far_away_triggers_reallocation(): assert result.account is not None assert result.account.account_id == "b" repo.delete.assert_called_once_with("key1", kind=StickySessionKind.PROMPT_CACHE) - repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "b", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) @pytest.mark.asyncio @@ -1057,4 +1116,4 @@ async def test_burn_first_reallocation_only_when_burn_first_is_selectable(): assert result.account is not None assert result.account.account_id == "a" repo.delete.assert_not_called() - repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE) + repo.upsert.assert_called_once_with("key1", "a", kind=StickySessionKind.PROMPT_CACHE, is_subagent=False) diff --git a/tests/unit/test_sticky_session_cleanup_scheduler.py b/tests/unit/test_sticky_session_cleanup_scheduler.py index 51d8ec870c..e7da88ba9b 100644 --- a/tests/unit/test_sticky_session_cleanup_scheduler.py +++ b/tests/unit/test_sticky_session_cleanup_scheduler.py @@ -39,6 +39,7 @@ async def test_cleanup_once_purges_prompt_cache_only(monkeypatch) -> None: dashboard_settings = SimpleNamespace( openai_cache_affinity_max_age_seconds=600, http_responses_session_bridge_prompt_cache_idle_ttl_seconds=600, + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=30, ) settings_repo = AsyncMock() @@ -84,7 +85,10 @@ async def __aexit__(self, *args): ): await scheduler._cleanup_once() - sticky_repo.purge_prompt_cache_before.assert_called_once() + assert sticky_repo.purge_prompt_cache_before.await_count == 2 + assert sticky_repo.purge_prompt_cache_before.await_args_list[0].kwargs["is_subagent"] is False + assert sticky_repo.purge_prompt_cache_before.await_args_list[1].kwargs["is_subagent"] is True + assert sticky_repo.purge_prompt_cache_before.await_args_list[1].kwargs["protect_active_bridge_mappings"] is True sticky_repo.purge_before.assert_not_called() bridge_repo.purge_closed_before.assert_called_once() bridge_repo.purge_abandoned_before.assert_called_once() @@ -96,6 +100,7 @@ async def test_cleanup_once_skips_bridge_purge_when_schema_is_not_ready(monkeypa dashboard_settings = SimpleNamespace( openai_cache_affinity_max_age_seconds=600, http_responses_session_bridge_prompt_cache_idle_ttl_seconds=600, + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=30, ) settings_repo = AsyncMock() @@ -145,7 +150,8 @@ async def __aexit__(self, *args): ): await scheduler._cleanup_once() - sticky_repo.purge_prompt_cache_before.assert_called_once() + assert sticky_repo.purge_prompt_cache_before.await_count == 2 + assert sticky_repo.purge_prompt_cache_before.await_args_list[1].kwargs["protect_active_bridge_mappings"] is True bridge_repo.purge_closed_before.assert_not_called() bridge_repo.purge_abandoned_before.assert_not_called() ring_service.purge_stale_before.assert_called_once() @@ -156,6 +162,7 @@ async def test_cleanup_once_purges_bridge_when_schema_exists_after_startup_flag_ dashboard_settings = SimpleNamespace( openai_cache_affinity_max_age_seconds=600, http_responses_session_bridge_prompt_cache_idle_ttl_seconds=600, + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=30, ) settings_repo = AsyncMock() @@ -201,7 +208,8 @@ async def __aexit__(self, *args): ): await scheduler._cleanup_once() - sticky_repo.purge_prompt_cache_before.assert_called_once() + assert sticky_repo.purge_prompt_cache_before.await_count == 2 + assert sticky_repo.purge_prompt_cache_before.await_args_list[1].kwargs["protect_active_bridge_mappings"] is True bridge_repo.purge_closed_before.assert_called_once() bridge_repo.purge_abandoned_before.assert_called_once() ring_service.purge_stale_before.assert_called_once()