From d8dc22eac34759db6578e5edc2a524a492414202 Mon Sep 17 00:00:00 2001 From: tobwen <1864057+tobwen@users.noreply.github.com> Date: Tue, 14 Jul 2026 01:21:30 +0000 Subject: [PATCH 1/3] feat(proxy): subagent prompt-cache TTL and session lifecycle Add subagent prompt-cache affinity TTL setting with NULL default (No Cache). Subagent sessions are detected via x-parent-session-id header. - NULL/0 TTL: bridge session and stream lease released immediately after stream end; no sticky mapping created - Positive TTL: bridge session retained for configured duration, then closed; sticky mapping created and cleaned after TTL Unanchored parallel forks (internal_unanchored_parallel) are closed immediately after stream end, releasing their stream lease. This prevents parallel fork sessions from consuming stream capacity for the parent prompt-cache idle TTL (3600s). The is_subagent flag propagates from header detection through account selection to sticky session persistence. The cleanup scheduler respects subagent TTL independently from parent sessions. Includes dashboard settings control, frontend display, Alembic migrations, and OpenSpec specifications. --- .cortexkit/.gitignore | 3 + ...13_080000_add_subagent_prompt_cache_ttl.py | 65 ++++++ app/db/models.py | 12 ++ .../proxy/_service/http_bridge/helpers.py | 44 ++++ .../proxy/_service/http_bridge/mixin.py | 46 ++-- .../proxy/_service/http_bridge/streaming.py | 47 +++++ app/modules/proxy/_service/support.py | 2 + app/modules/proxy/load_balancer.py | 27 ++- app/modules/proxy/service.py | 3 + app/modules/proxy/sticky_repository.py | 57 ++++- app/modules/settings/api.py | 6 + app/modules/settings/repository.py | 9 + app/modules/settings/schemas.py | 2 + app/modules/settings/service.py | 5 + app/modules/sticky_sessions/api.py | 1 + .../sticky_sessions/cleanup_scheduler.py | 20 +- app/modules/sticky_sessions/schemas.py | 1 + app/modules/sticky_sessions/service.py | 21 +- .../settings/components/routing-settings.tsx | 57 +++++ frontend/src/features/settings/payload.ts | 2 + frontend/src/features/settings/schemas.ts | 8 + .../components/sticky-sessions-section.tsx | 6 +- .../hooks/use-sticky-sessions.test.ts | 7 +- .../features/sticky-sessions/schemas.test.ts | 1 + .../src/features/sticky-sessions/schemas.ts | 1 + frontend/src/i18n/locales/en.json | 5 + .../.openspec.yaml | 2 + .../add-subagent-prompt-cache-ttl/design.md | 27 +++ .../add-subagent-prompt-cache-ttl/notes.md | 29 +++ .../add-subagent-prompt-cache-ttl/proposal.md | 29 +++ .../specs/frontend-architecture/spec.md | 20 ++ .../specs/sticky-session-operations/spec.md | 90 ++++++++ .../add-subagent-prompt-cache-ttl/tasks.md | 42 ++++ .../sticky-session-operations/context.md | 3 + scripts/check_proxy_architecture.py | 8 +- tests/integration/test_sticky_sessions_api.py | 35 ++++ tests/unit/test_otel.py | 7 + tests/unit/test_proxy_http_bridge.py | 198 ++++++++++++++++++ .../unit/test_proxy_load_balancer_refresh.py | 6 +- tests/unit/test_select_with_stickiness.py | 117 ++++++++--- .../test_sticky_session_cleanup_scheduler.py | 11 +- 41 files changed, 1007 insertions(+), 75 deletions(-) create mode 100644 .cortexkit/.gitignore create mode 100644 app/db/alembic/versions/20260713_080000_add_subagent_prompt_cache_ttl.py create mode 100644 openspec/changes/add-subagent-prompt-cache-ttl/.openspec.yaml create mode 100644 openspec/changes/add-subagent-prompt-cache-ttl/design.md create mode 100644 openspec/changes/add-subagent-prompt-cache-ttl/notes.md create mode 100644 openspec/changes/add-subagent-prompt-cache-ttl/proposal.md create mode 100644 openspec/changes/add-subagent-prompt-cache-ttl/specs/frontend-architecture/spec.md create mode 100644 openspec/changes/add-subagent-prompt-cache-ttl/specs/sticky-session-operations/spec.md create mode 100644 openspec/changes/add-subagent-prompt-cache-ttl/tasks.md 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..47d81fefaa 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -1851,6 +1851,48 @@ 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 + + +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: + await repositories.sticky_sessions.delete( + affinity.key, + kind=StickySessionKind.PROMPT_CACHE, + is_subagent=True, + ) + except Exception: + logger.warning("Failed to delete completed subagent sticky mapping", exc_info=True) + + def _http_bridge_should_attempt_local_previous_response_recovery(exc: ProxyResponseError) -> bool: payload = exc.payload if not isinstance(payload, dict): @@ -2182,6 +2224,8 @@ 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", + "_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..c3709647c0 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, @@ -252,20 +254,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( @@ -435,6 +427,7 @@ async def _get_or_create_http_bridge_session( session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, ) -> "_HTTPBridgeSession | _HTTPBridgeOwnerForward": settings = _service_get_settings() + dashboard_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) @@ -475,6 +468,7 @@ async def _get_or_create_http_bridge_session( ), ) effective_idle_ttl_seconds = idle_ttl_seconds + is_subagent_session, subagent_prompt_cache_ttl_seconds = _detect_subagent_session(headers, dashboard_settings) forwarded_affinity = ( _forwarded_http_bridge_session_key( headers, @@ -1386,6 +1380,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 +1400,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 +1600,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, @@ -1841,6 +1844,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 +1870,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 +2103,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..4e9f11c6b5 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -76,6 +76,7 @@ _http_bridge_request_needs_unanchored_handoff, _http_bridge_request_stage, _http_bridge_runtime_config, + _http_bridge_session_has_visible_requests, _http_bridge_should_attempt_local_bootstrap_rebind, _http_bridge_should_attempt_local_previous_response_recovery, _http_bridge_should_attempt_soft_affinity_reroute, @@ -472,7 +473,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, *, @@ -2268,3 +2292,26 @@ 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: + ttl = session.subagent_prompt_cache_ttl_seconds + if ttl is None: + await self._close_http_bridge_session(session) + else: + scheduled_last_used_at = session.last_used_at + await self._refresh_completed_subagent_sticky_mapping(session) + + async def _delayed_subagent_close() -> None: + await asyncio.sleep(ttl) + if session.closed or session.last_used_at != scheduled_last_used_at: + return + if _http_bridge_session_has_visible_requests(session): + return + await self._close_http_bridge_session(session) + + self._schedule_cancel_safe_cleanup( + _delayed_subagent_close(), + action="http_bridge_session_close", + request_id=session.key.affinity_key, + ) 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..2236fba95a 100644 --- a/app/modules/proxy/sticky_repository.py +++ b/app/modules/proxy/sticky_repository.py @@ -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() @@ -190,20 +210,35 @@ async def count_entries( 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) -> int: + 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_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 +246,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(), }, ) diff --git a/app/modules/settings/api.py b/app/modules/settings/api.py index b217b45c87..96f92cd80b 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 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..9afe29b0a2 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,19 @@ 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, + ) 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..7fe0bcb452 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,6 +71,7 @@ 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) normalized_account_query = account_query.strip() if account_query else None normalized_key_query = key_query.strip() if key_query else None @@ -98,7 +100,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, @@ -161,12 +165,22 @@ async def purge_entries(self) -> int: cutoff = utcnow() - timedelta(seconds=settings.openai_cache_affinity_max_age_seconds) return await self._repository.purge_prompt_cache_before(cutoff) - 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 +190,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( 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/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..bd294683cd --- /dev/null +++ b/openspec/changes/add-subagent-prompt-cache-ttl/notes.md @@ -0,0 +1,29 @@ +# 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..b3748b3ad8 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_603 +MAX_LOAD_BALANCER_LINES = 3_279 +MAX_HTTP_BRIDGE_MIXIN_LINES = 2_414 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_sticky_sessions_api.py b/tests/integration/test_sticky_sessions_api.py index 3694d75a0c..d7753a90cf 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,6 +62,40 @@ 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, 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..9fe03dff81 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -44,6 +44,192 @@ 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)), + 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_delayed_close_skips_reused_session(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)), + 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() + scheduled: list[Any] = [] + 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( + service, + "_schedule_cancel_safe_cleanup", + lambda coro, **_kwargs: scheduled.append(coro), + ) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: SimpleNamespace(sse_keepalive_interval_seconds=0), + ) + monkeypatch.setattr(http_bridge_streaming_module.asyncio, "sleep", AsyncMock()) + + 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) + assert len(scheduled) == 1 + session.last_used_at += 1.0 + await scheduled[0] + + close_mock.assert_not_awaited() + + +@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.fixture(autouse=True) def _share_proxy_dashboard_settings(monkeypatch: pytest.MonkeyPatch) -> None: class _SettingsCache: @@ -9710,6 +9896,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 +9909,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 +9985,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 +9998,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 +10075,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 +10088,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 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..69926cf025 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,9 @@ 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 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 +99,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 +149,7 @@ 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 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 +160,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 +206,7 @@ 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 bridge_repo.purge_closed_before.assert_called_once() bridge_repo.purge_abandoned_before.assert_called_once() ring_service.purge_stale_before.assert_called_once() From d85007eb75740a19d65fadbf3bbcd927d0f63d78 Mon Sep 17 00:00:00 2001 From: Darafei Praliaskouski Date: Sat, 18 Jul 2026 11:10:14 +0400 Subject: [PATCH 2/3] fix(proxy): isolate subagent prompt-cache sessions --- .../proxy/_service/http_bridge/helpers.py | 34 +++ .../proxy/_service/http_bridge/mixin.py | 16 +- app/modules/proxy/sticky_repository.py | 9 + app/modules/sticky_sessions/service.py | 148 ++++++++++-- frontend/src/i18n/locales/ko.json | 5 + frontend/src/i18n/locales/zh-CN.json | 5 + tests/integration/test_sticky_sessions_api.py | 50 +++- tests/unit/test_proxy_http_bridge.py | 218 ++++++++++++++++++ 8 files changed, 464 insertions(+), 21 deletions(-) diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index 47d81fefaa..95202e1456 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -1869,6 +1869,39 @@ def _detect_subagent_session( return is_subagent, ttl +def _subagent_prompt_cache_bridge_key( + key: _HTTPBridgeSessionKey, + affinity: _AffinityPolicy, + *, + is_subagent: bool, + idle_ttl_seconds: float, + request_scope_id: str, + subagent_prompt_cache_ttl_seconds: int | None, +) -> tuple[_HTTPBridgeSessionKey, _AffinityPolicy, float]: + if not is_subagent or key.affinity_kind != "prompt_cache" or affinity.kind is not StickySessionKind.PROMPT_CACHE: + return key, affinity, idle_ttl_seconds + + if subagent_prompt_cache_ttl_seconds is None: + subagent_affinity_key = f"subagent:{request_scope_id}:{key.affinity_key}" + else: + subagent_affinity_key = f"subagent:{key.affinity_key}" + return ( + _HTTPBridgeSessionKey( + key.affinity_kind, + subagent_affinity_key, + key.api_key_id, + strength=key.strength, + ), + _AffinityPolicy( + key=subagent_affinity_key, + kind=affinity.kind, + reallocate_sticky=affinity.reallocate_sticky, + max_age_seconds=subagent_prompt_cache_ttl_seconds, + ), + idle_ttl_seconds if subagent_prompt_cache_ttl_seconds is None else float(subagent_prompt_cache_ttl_seconds), + ) + + async def _delete_completed_subagent_sticky_mapping( repo_factory: Callable[..., Any], session: Any, @@ -2225,6 +2258,7 @@ def _wrapper(*args: Any, **kwargs: Any) -> Any: "_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", diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index c3709647c0..30548bc98b 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -126,6 +126,7 @@ _renew_durable_http_bridge_lease, _require_http_bridge_bound_account_not_excluded, _reserve_http_bridge_unanchored_handoff, + _subagent_prompt_cache_bridge_key, _track_alias_registration, ) from app.modules.proxy._service.http_bridge.owner_forwarding import _HTTPBridgeOwnerForwardingMixin @@ -367,7 +368,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, @@ -436,11 +436,19 @@ async def _get_or_create_http_bridge_session( 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) + key, affinity, effective_idle_ttl_seconds = _subagent_prompt_cache_bridge_key( + key, + affinity, + is_subagent=is_subagent_session, + idle_ttl_seconds=idle_ttl_seconds, + request_scope_id=request_scope_id, + subagent_prompt_cache_ttl_seconds=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: @@ -467,8 +475,6 @@ async def _get_or_create_http_bridge_session( error_type="server_error", ), ) - effective_idle_ttl_seconds = idle_ttl_seconds - is_subagent_session, subagent_prompt_cache_ttl_seconds = _detect_subagent_session(headers, dashboard_settings) forwarded_affinity = ( _forwarded_http_bridge_session_key( headers, diff --git a/app/modules/proxy/sticky_repository.py b/app/modules/proxy/sticky_repository.py index 2236fba95a..d917257469 100644 --- a/app/modules/proxy/sticky_repository.py +++ b/app/modules/proxy/sticky_repository.py @@ -137,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]]: @@ -145,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, ) @@ -163,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", @@ -176,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, ) @@ -197,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: @@ -204,6 +209,7 @@ 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, ) @@ -267,6 +273,7 @@ def _apply_filters( *, kind: StickySessionKind | None, updated_before: datetime | None, + is_subagent: bool | None, account_query: str | None, key_query: str | None, ): @@ -274,6 +281,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/sticky_sessions/service.py b/app/modules/sticky_sessions/service.py index 7fe0bcb452..4d0f338a70 100644 --- a/app/modules/sticky_sessions/service.py +++ b/app/modules/sticky_sessions/service.py @@ -73,9 +73,14 @@ async def list_entries( 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, + ) if stale_only and kind not in (None, StickySessionKind.PROMPT_CACHE): return StickySessionListData( entries=[], @@ -83,16 +88,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, @@ -146,17 +170,27 @@ 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) @@ -198,10 +232,96 @@ async def _count_stale_prompt_cache_entries( *, kind: StickySessionKind | None, stale_cutoff: datetime, + subagent_stale_cutoff: datetime, ) -> 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, + ) + subagent_count = await self._repository.count_entries( + kind=StickySessionKind.PROMPT_CACHE, + updated_before=subagent_stale_cutoff, + is_subagent=True, ) + 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]: + 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=None, + ) + 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=None, + ) + rows = list(parent_rows) + list(subagent_rows) + rows.sort(key=self._stale_sort_key(sort_by), reverse=sort_dir == "desc") + return rows[offset : offset + limit], len(rows) + + 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/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/tests/integration/test_sticky_sessions_api.py b/tests/integration/test_sticky_sessions_api.py index d7753a90cf..c76e1e7f6d 100644 --- a/tests/integration/test_sticky_sessions_api.py +++ b/tests/integration/test_sticky_sessions_api.py @@ -102,20 +102,22 @@ async def _insert_sticky_session( 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, }, ) @@ -342,6 +344,50 @@ 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 + + @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_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index 9fe03dff81..2fbef29f6c 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -10136,6 +10136,224 @@ 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("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, + 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"}, + 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_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_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_recovers_from_previous_response_id_mapping( monkeypatch: pytest.MonkeyPatch, From 4cce15b8fa988a0045ee9f8468a3198ff3691416 Mon Sep 17 00:00:00 2001 From: Darafei Praliaskouski Date: Sat, 18 Jul 2026 11:51:55 +0400 Subject: [PATCH 3/3] fix(proxy): isolate subagent bridge lifecycle --- .../proxy/_service/http_bridge/helpers.py | 72 ++- .../proxy/_service/http_bridge/mixin.py | 41 +- .../proxy/_service/http_bridge/streaming.py | 58 ++- app/modules/proxy/_service/response_create.py | 1 + app/modules/proxy/sticky_repository.py | 40 +- app/modules/settings/api.py | 1 + .../sticky_sessions/cleanup_scheduler.py | 1 + app/modules/sticky_sessions/service.py | 29 +- .../add-subagent-prompt-cache-ttl/notes.md | 2 - scripts/check_proxy_architecture.py | 4 +- .../test_settings_audit_changed_fields.py | 5 + tests/integration/test_sticky_sessions_api.py | 11 + tests/unit/test_load_balancer_contract.py | 3 +- tests/unit/test_proxy_http_bridge.py | 441 +++++++++++++++++- .../test_sticky_session_cleanup_scheduler.py | 3 + 15 files changed, 625 insertions(+), 87 deletions(-) diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index 95202e1456..7cd45f05bf 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -1872,33 +1872,49 @@ def _detect_subagent_session( 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]: - if not is_subagent or key.affinity_kind != "prompt_cache" or affinity.kind is not StickySessionKind.PROMPT_CACHE: - return key, affinity, idle_ttl_seconds - - if subagent_prompt_cache_ttl_seconds is None: - subagent_affinity_key = f"subagent:{request_scope_id}:{key.affinity_key}" +) -> 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:{key.affinity_key}" + 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( - key.affinity_kind, + subagent_affinity_kind, subagent_affinity_key, key.api_key_id, - strength=key.strength, + 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=subagent_prompt_cache_ttl_seconds, + max_age_seconds=effective_max_age_seconds, ), - idle_ttl_seconds if subagent_prompt_cache_ttl_seconds is None else float(subagent_prompt_cache_ttl_seconds), + effective_idle_ttl_seconds, + effective_subagent_prompt_cache_ttl_seconds, ) @@ -1917,15 +1933,37 @@ async def _delete_completed_subagent_sticky_mapping( return try: async with repo_factory() as repositories: - await repositories.sticky_sessions.delete( - affinity.key, - kind=StickySessionKind.PROMPT_CACHE, - is_subagent=True, - ) + 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): diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index 30548bc98b..b104bc3fab 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -127,6 +127,7 @@ _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 @@ -426,24 +427,36 @@ 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() - dashboard_settings = await _service_get_settings_cache().get() + 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) - key, affinity, effective_idle_ttl_seconds = _subagent_prompt_cache_bridge_key( + subagent_key_result = _subagent_prompt_cache_bridge_key( key, affinity, - is_subagent=is_subagent_session, - idle_ttl_seconds=idle_ttl_seconds, - request_scope_id=request_scope_id, - subagent_prompt_cache_ttl_seconds=subagent_prompt_cache_ttl_seconds, + 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) @@ -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: @@ -1826,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( diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index 4e9f11c6b5..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, @@ -76,7 +77,6 @@ _http_bridge_request_needs_unanchored_handoff, _http_bridge_request_stage, _http_bridge_runtime_config, - _http_bridge_session_has_visible_requests, _http_bridge_should_attempt_local_bootstrap_rebind, _http_bridge_should_attempt_local_previous_response_recovery, _http_bridge_should_attempt_soft_affinity_reroute, @@ -91,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 ( @@ -791,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, @@ -800,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, @@ -813,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, @@ -2295,23 +2326,6 @@ async def _stream_http_bridge_session_events( if _http_bridge_session_is_unanchored_parallel_fork(session): await self._close_http_bridge_session(session) elif session.is_subagent: - ttl = session.subagent_prompt_cache_ttl_seconds - if ttl is None: - await self._close_http_bridge_session(session) - else: - scheduled_last_used_at = session.last_used_at + if session.subagent_prompt_cache_ttl_seconds is not None: await self._refresh_completed_subagent_sticky_mapping(session) - - async def _delayed_subagent_close() -> None: - await asyncio.sleep(ttl) - if session.closed or session.last_used_at != scheduled_last_used_at: - return - if _http_bridge_session_has_visible_requests(session): - return - await self._close_http_bridge_session(session) - - self._schedule_cancel_safe_cleanup( - _delayed_subagent_close(), - action="http_bridge_session_close", - request_id=session.key.affinity_key, - ) + 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/sticky_repository.py b/app/modules/proxy/sticky_repository.py index d917257469..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 @@ -216,9 +216,45 @@ async def count_entries( result = await self._session.execute(statement) return int(result.scalar_one()) - async def purge_prompt_cache_before(self, cutoff: datetime, *, is_subagent: bool | None = None) -> int: + 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_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, diff --git a/app/modules/settings/api.py b/app/modules/settings/api.py index 96f92cd80b..8ecb9ecb96 100644 --- a/app/modules/settings/api.py +++ b/app/modules/settings/api.py @@ -846,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/sticky_sessions/cleanup_scheduler.py b/app/modules/sticky_sessions/cleanup_scheduler.py index 9afe29b0a2..bf16df11b7 100644 --- a/app/modules/sticky_sessions/cleanup_scheduler.py +++ b/app/modules/sticky_sessions/cleanup_scheduler.py @@ -122,6 +122,7 @@ async def _cleanup_as_leader(self) -> None: 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) diff --git a/app/modules/sticky_sessions/service.py b/app/modules/sticky_sessions/service.py index 4d0f338a70..77f5131eee 100644 --- a/app/modules/sticky_sessions/service.py +++ b/app/modules/sticky_sessions/service.py @@ -80,6 +80,8 @@ async def list_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( @@ -197,7 +199,12 @@ async def delete_filtered_entries( 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, @@ -233,6 +240,8 @@ 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 @@ -240,11 +249,15 @@ async def _count_stale_prompt_cache_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 @@ -266,6 +279,7 @@ async def _list_stale_prompt_cache_rows( 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, @@ -275,7 +289,7 @@ async def _list_stale_prompt_cache_rows( sort_by=sort_by, sort_dir=sort_dir, offset=0, - limit=None, + limit=fetch_limit, ) subagent_rows = await self._repository.list_entries( kind=StickySessionKind.PROMPT_CACHE, @@ -286,11 +300,18 @@ async def _list_stale_prompt_cache_rows( sort_by=sort_by, sort_dir=sort_dir, offset=0, - limit=None, + limit=fetch_limit, ) rows = list(parent_rows) + list(subagent_rows) rows.sort(key=self._stale_sort_key(sort_by), reverse=sort_dir == "desc") - return rows[offset : offset + limit], len(rows) + 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, diff --git a/openspec/changes/add-subagent-prompt-cache-ttl/notes.md b/openspec/changes/add-subagent-prompt-cache-ttl/notes.md index bd294683cd..6130071e66 100644 --- a/openspec/changes/add-subagent-prompt-cache-ttl/notes.md +++ b/openspec/changes/add-subagent-prompt-cache-ttl/notes.md @@ -25,5 +25,3 @@ Subagent bridge sessions are closed at response-stream completion. An optional p ## Visibility Gap Active bridge-session visibility is outside this change. - - diff --git a/scripts/check_proxy_architecture.py b/scripts/check_proxy_architecture.py index b3748b3ad8..973b6a5028 100644 --- a/scripts/check_proxy_architecture.py +++ b/scripts/check_proxy_architecture.py @@ -21,9 +21,9 @@ HTTP_BRIDGE_MIXIN_PATH = PROXY_DIR / "_service" / "http_bridge" / "mixin.py" STREAMING_MIXIN_PATH = PROXY_DIR / "_service" / "streaming" / "mixin.py" -MAX_SERVICE_LINES = 2_603 +MAX_SERVICE_LINES = 2_607 MAX_LOAD_BALANCER_LINES = 3_279 -MAX_HTTP_BRIDGE_MIXIN_LINES = 2_414 +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 = 701 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 c76e1e7f6d..df929d97ab 100644 --- a/tests/integration/test_sticky_sessions_api.py +++ b/tests/integration/test_sticky_sessions_api.py @@ -387,6 +387,17 @@ async def test_sticky_sessions_api_stale_only_uses_subagent_ttl(async_client): 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): 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_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index 2fbef29f6c..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 @@ -58,7 +58,7 @@ async def test_fork_session_closed_after_stream_end(monkeypatch: pytest.MonkeyPa headers={}, affinity=proxy_service._AffinityPolicy(key="fork-key"), request_model="gpt-5.4", - account=cast(Any, SimpleNamespace(id="acc-1", status=AccountStatus.ACTIVE)), + 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(), @@ -109,7 +109,7 @@ async def test_fork_session_closed_after_stream_end(monkeypatch: pytest.MonkeyPa @pytest.mark.asyncio -async def test_subagent_delayed_close_skips_reused_session(monkeypatch: pytest.MonkeyPatch) -> None: +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( @@ -120,7 +120,7 @@ async def test_subagent_delayed_close_skips_reused_session(monkeypatch: pytest.M 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-1", status=AccountStatus.ACTIVE, plan_type="plus")), upstream=cast(UpstreamResponsesWebSocket, SimpleNamespace(close=AsyncMock())), upstream_control=proxy_service._WebSocketUpstreamControl(), pending_requests=deque(), @@ -148,22 +148,15 @@ async def test_subagent_delayed_close_skips_reused_session(monkeypatch: pytest.M close_mock = AsyncMock() refresh_mock = AsyncMock() - scheduled: list[Any] = [] 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( - service, - "_schedule_cancel_safe_cleanup", - lambda coro, **_kwargs: scheduled.append(coro), - ) monkeypatch.setattr( proxy_service, "get_settings", lambda: SimpleNamespace(sse_keepalive_interval_seconds=0), ) - monkeypatch.setattr(http_bridge_streaming_module.asyncio, "sleep", AsyncMock()) gen = service._stream_http_bridge_session_events( session=session, @@ -178,11 +171,7 @@ async def test_subagent_delayed_close_skips_reused_session(monkeypatch: pytest.M pass refresh_mock.assert_awaited_once_with(session) - assert len(scheduled) == 1 - session.last_used_at += 1.0 - await scheduled[0] - - close_mock.assert_not_awaited() + close_mock.assert_awaited_once_with(session) @pytest.mark.asyncio @@ -230,6 +219,60 @@ async def __aexit__(self, *_args: object) -> None: ) +@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: @@ -295,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) @@ -4220,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 @@ -4315,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 @@ -5954,6 +6044,143 @@ async def fake_get_or_create( assert captured["preferred_account_id"] is None +@pytest.mark.asyncio +async def test_stream_via_http_bridge_isolates_subagent_key_before_durable_lookup( + 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", + "prompt_cache_key": "parent-cache", + }, + ) + request_state = proxy_service._WebSocketRequestState( + request_id="req-subagent-durable-lookup", + 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) + + 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 + return request_state, '{"type":"response.create"}' + + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey( + "prompt_cache", + "subagent:scope-subagent-durable-lookup:parent-cache", + None, + ), + headers={}, + affinity=proxy_service._AffinityPolicy( + 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-subagent", 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, + ) + 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, + "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, + http_responses_session_bridge_subagent_prompt_cache_ttl_seconds=0, + ) + ) + ), + ), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + 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) + 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()) + + 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 == [] + 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_stream_via_http_bridge_prefers_durable_account_for_soft_prompt_cache_follow_up_recovery( monkeypatch: pytest.MonkeyPatch, @@ -10141,7 +10368,7 @@ async def test_get_or_create_http_bridge_session_namespaces_subagent_prompt_cach monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) - requested_key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "parent-cache", None) + requested_key = proxy_service._HTTPBridgeSessionKey("session_header", "parent-session", None) created_session = proxy_service._HTTPBridgeSession( key=requested_key, headers={}, @@ -10225,7 +10452,10 @@ async def fake_create_http_bridge_session( resolved = await service._get_or_create_http_bridge_session( requested_key, - headers={"x-parent-session-id": "parent-session-1"}, + 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, @@ -10239,6 +10469,7 @@ async def fake_create_http_bridge_session( 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) @@ -10248,6 +10479,76 @@ async def fake_create_http_bridge_session( 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, @@ -10354,6 +10655,110 @@ async def fake_create_http_bridge_session( 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_sticky_session_cleanup_scheduler.py b/tests/unit/test_sticky_session_cleanup_scheduler.py index 69926cf025..e7da88ba9b 100644 --- a/tests/unit/test_sticky_session_cleanup_scheduler.py +++ b/tests/unit/test_sticky_session_cleanup_scheduler.py @@ -88,6 +88,7 @@ async def __aexit__(self, *args): 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() @@ -150,6 +151,7 @@ async def __aexit__(self, *args): await scheduler._cleanup_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() @@ -207,6 +209,7 @@ async def __aexit__(self, *args): await scheduler._cleanup_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()