fix(proxy): complete cache-locality fix for prompt cache hit rate restoration (#273)

* Implement cache locality fix (Tasks 1-5): model-class prefix, budget reallocation, settings columns, bridge TTL wiring, observability logs

* test(migration): verify new dashboard_settings columns

* feat(settings): expose cache-locality fields in dashboard API

* refactor(test): eliminate stickiness test replica, import production code

* feat(observability): populate bridge event cache_key_family and model_class

- Enrich all 15 bridge event call-sites with cache_key_family and model_class kwargs
- cache_key_family sourced from key.affinity_kind (always available)
- model_class derived from _extract_model_class() with None fallback
- No new log events created, only enrichment of existing ones
- Updated test fixture with new settings fields
- All tests passing (17 stickiness tests, bridge integration tests)

* fix(lb): use strict > for budget_threshold_pct to match intent (>95% used)

* fix(test): update idle_ttl_seconds assertion to >= after prompt-cache TTL wiring

* feat(lb): wire configurable budget_threshold_pct through load_balancer

* feat(observability): log reallocation-orphan bridge event

* test(integration): add regression suite for cache-locality completion

* fix(review): address final verification wave rejections

- Fix A: Rename test wrapper _select_with_stickiness to _invoke_stickiness (22 call sites)
- Fix B: Move AccountsRepository and StickySessionsRepository to TYPE_CHECKING block; inline get_settings import
- Fix C: Remove explicit generic type from asyncio.Future annotation in bridge test
- Fix D: Collapse multi-line getattr to single line for grep verification

All verification tests pass.

* fix(ci): resolve ruff lint and ty type-check errors

* fix(proxy): preserve sticky grace path and selection budget

* fix(types): align proxy and test typing contracts

* fix(proxy): scope sticky reallocation and release retry lock

* fix(proxy): restore entropy for empty cache-key fallback

* fix(test): align proxy fixtures with settings wiring

* fix(proxy): scope far-reset reallocation to prompt cache

* fix(proxy): preserve bridge ttl and codex model bucketing

* fix(test): speed up postgres fixture resets

* fix(test): preserve postgres fixture semantics

Author: Soju06

app/core/clients/model_fetcher.py

@@ -48,13 +48,23 @@ def _list_raw(data: dict[str, JsonValue], key: str) -> list[JsonValue]:
     return []
 
 
+def _parse_reasoning_level(value: JsonValue) -> ReasoningLevel | None:
+    if not isinstance(value, dict):
+        return None
+    effort = value.get("effort")
+    description = value.get("description")
+    if not isinstance(effort, str) or not isinstance(description, str):
+        return None
+    return ReasoningLevel(effort=effort, description=description)
+
+
 def _parse_upstream_model(data: dict[str, JsonValue]) -> UpstreamModel:
     raw = {k: v for k, v in data.items() if k not in _FILTERED_FIELDS}
 
     reasoning_levels = tuple(
-        ReasoningLevel(effort=rl.get("effort", ""), description=rl.get("description", ""))
+        parsed_level
         for rl in _list_raw(data, "supported_reasoning_levels")
-        if isinstance(rl, dict)
+        if (parsed_level := _parse_reasoning_level(rl)) is not None
     )
 
     available_in_plans = frozenset(p for p in _list_raw(data, "available_in_plans") if isinstance(p, str))

app/core/clients/proxy.py

@@ -12,19 +12,26 @@
 import socket
 import time
 from dataclasses import dataclass
-from typing import Any, AsyncContextManager, AsyncIterator, Awaitable, Mapping, Protocol, TypeAlias, cast
+from typing import Any, AsyncContextManager, AsyncIterator, Awaitable, Callable, Mapping, Protocol, TypeAlias, cast
 from urllib.parse import ParseResult, urlparse, urlunparse
 
 import aiohttp
 from aiohttp import hdrs
 from aiohttp.client_ws import DEFAULT_WS_CLIENT_TIMEOUT
+from aiohttp.http_websocket import WS_KEY, WebSocketReader, WebSocketWriter
 from multidict import CIMultiDict
 
 from app.core.clients.http import get_http_client
 from app.core.config.settings import get_settings
-from app.core.errors import OpenAIErrorEnvelope, ResponseFailedEvent, openai_error, response_failed_event
+from app.core.errors import (
+    OpenAIErrorDetail,
+    OpenAIErrorEnvelope,
+    ResponseFailedEvent,
+    openai_error,
+    response_failed_event,
+)
 from app.core.openai.model_registry import get_model_registry
-from app.core.openai.models import CompactResponsePayload
+from app.core.openai.models import CompactResponsePayload, OpenAIError
 from app.core.openai.parsing import (
     parse_compact_response_payload,
     parse_error_payload,
@@ -408,7 +415,7 @@ def _error_payload_from_websocket_handshake_error(exc: aiohttp.WSServerHandshake
     if extracted is not None:
         error = parse_error_payload(extracted)
         if error is not None:
-            return {"error": error.model_dump(exclude_none=True)}
+            return {"error": _openai_error_detail(error)}
 
     code = _infer_websocket_handshake_error_code(exc.status, message)
     if code == "invalid_api_key":
@@ -694,13 +701,32 @@ async def _error_payload_from_response(resp: ErrorResponse) -> OpenAIErrorEnvelo
     if isinstance(data, dict):
         error = parse_error_payload(data)
         if error:
-            return {"error": error.model_dump(exclude_none=True)}
+            return {"error": _openai_error_detail(error)}
         message = _extract_upstream_message(data)
         if message:
             return openai_error("upstream_error", message)
     return openai_error("upstream_error", fallback_message)
 
 
+def _openai_error_detail(error: OpenAIError) -> OpenAIErrorDetail:
+    detail: OpenAIErrorDetail = {}
+    if error.message is not None:
+        detail["message"] = error.message
+    if error.type is not None:
+        detail["type"] = error.type
+    if error.code is not None:
+        detail["code"] = error.code
+    if error.param is not None:
+        detail["param"] = error.param
+    if error.plan_type is not None:
+        detail["plan_type"] = error.plan_type
+    if error.resets_at is not None:
+        detail["resets_at"] = error.resets_at
+    if error.resets_in_seconds is not None:
+        detail["resets_in_seconds"] = error.resets_in_seconds
+    return detail
+
+
 def _extract_upstream_message(data: Mapping[str, object]) -> str | None:
     for key in ("message", "detail", "error"):
         value = data.get(key)
@@ -856,8 +882,8 @@ async def _open_upstream_websocket(
     connect_timeout_seconds: float,
     max_msg_size: int,
 ) -> tuple[AsyncContextManager[aiohttp.ClientWebSocketResponse], aiohttp.ClientWebSocketResponse]:
-    request = getattr(session, "request", None)
-    if not callable(request):
+    request_obj = getattr(session, "request", None)
+    if not callable(request_obj):
         websocket_cm = session.ws_connect(
             url,
             headers=headers,
@@ -868,6 +894,7 @@ async def _open_upstream_websocket(
         )
         websocket = await asyncio.wait_for(websocket_cm.__aenter__(), timeout=connect_timeout_seconds)
         return websocket_cm, websocket
+    request = cast(Callable[..., Awaitable[aiohttp.ClientResponse]], request_obj)
 
     request_headers = CIMultiDict(headers)
     request_headers.setdefault(hdrs.UPGRADE, "websocket")
@@ -910,7 +937,7 @@ async def _raise_handshake_error(message: str) -> None:
             await _raise_handshake_error("Invalid connection header")
 
         response_key = resp.headers.get(hdrs.SEC_WEBSOCKET_ACCEPT, "")
-        expected_key = base64.b64encode(hashlib.sha1(sec_key.encode() + aiohttp.client.WS_KEY).digest()).decode()
+        expected_key = base64.b64encode(hashlib.sha1(sec_key.encode() + WS_KEY).digest()).decode()
         if response_key != expected_key:
             await _raise_handshake_error("Invalid challenge response")
 
@@ -922,9 +949,10 @@ async def _raise_handshake_error(message: str) -> None:
 
         transport = conn.transport
         assert transport is not None
-        reader = aiohttp.client.WebSocketDataQueue(conn_proto, 2**16, loop=session._loop)
-        conn_proto.set_parser(aiohttp.client.WebSocketReader(reader, max_msg_size), reader)
-        writer = aiohttp.client.WebSocketWriter(conn_proto, transport, use_mask=True, compress=0, notakeover=False)
+        web_socket_data_queue = cast(Callable[..., Any], getattr(aiohttp.client_ws, "WebSocketDataQueue"))
+        reader = web_socket_data_queue(conn_proto, 2**16, loop=session._loop)
+        conn_proto.set_parser(WebSocketReader(reader, max_msg_size), reader)
+        writer = WebSocketWriter(conn_proto, transport, use_mask=True, compress=0, notakeover=False)
     except BaseException:
         resp.close()
         raise

app/core/clients/proxy_websocket.py

@@ -20,7 +20,8 @@
 
 from app.core.clients.proxy import ProxyResponseError, filter_inbound_headers
 from app.core.config.settings import get_settings
-from app.core.errors import OpenAIErrorEnvelope, openai_error
+from app.core.errors import OpenAIErrorDetail, OpenAIErrorEnvelope, openai_error
+from app.core.openai.models import OpenAIError
 from app.core.openai.parsing import parse_error_payload
 from app.core.utils.request_id import get_request_id
 
@@ -262,4 +263,23 @@ def _try_parse_handshake_error_payload(
     error = parse_error_payload(payload)
     if error is None:
         return None
-    return {"error": error.model_dump(exclude_none=True)}
+    return {"error": _openai_error_detail(error)}
+
+
+def _openai_error_detail(error: OpenAIError) -> OpenAIErrorDetail:
+    detail: OpenAIErrorDetail = {}
+    if error.message is not None:
+        detail["message"] = error.message
+    if error.type is not None:
+        detail["type"] = error.type
+    if error.code is not None:
+        detail["code"] = error.code
+    if error.param is not None:
+        detail["param"] = error.param
+    if error.plan_type is not None:
+        detail["plan_type"] = error.plan_type
+    if error.resets_at is not None:
+        detail["resets_at"] = error.resets_at
+    if error.resets_in_seconds is not None:
+        detail["resets_in_seconds"] = error.resets_in_seconds
+    return detail

app/core/middleware/request_decompression.py

@@ -111,7 +111,7 @@ def _replace_request_body(request: Request, body: bytes) -> None:
     headers.append((b"content-length", str(len(body)).encode("ascii")))
     request.scope["headers"] = headers
     # Ensure subsequent request.headers reflects the updated scope headers.
-    request._headers = None
+    request.__dict__.pop("_headers", None)
 
 
 def add_request_decompression_middleware(app: FastAPI) -> None:

app/core/openai/chat_requests.py

@@ -265,11 +265,13 @@ def _text_format_from_parsed(parsed: ChatResponseFormat) -> ResponsesTextFormat:
         json_schema = parsed.json_schema
         if json_schema is None:
             raise ValueError("'response_format.json_schema' is required when type is 'json_schema'.")
-        return ResponsesTextFormat(
-            type=parsed.type,
-            schema_=json_schema.schema_,
-            name=json_schema.name,
-            strict=json_schema.strict,
+        return ResponsesTextFormat.model_validate(
+            {
+                "type": parsed.type,
+                "schema": json_schema.schema_,
+                "name": json_schema.name,
+                "strict": json_schema.strict,
+            }
         )
     if parsed.type in ("json_object", "text"):
         return ResponsesTextFormat(type=parsed.type)

app/db/alembic/versions/20260213_000000_base_schema.py

@@ -178,6 +178,18 @@ def upgrade() -> None:
                 nullable=False,
                 server_default=sa.text("CURRENT_TIMESTAMP"),
             ),
+            sa.Column(
+                "http_responses_session_bridge_prompt_cache_idle_ttl_seconds",
+                sa.Integer(),
+                nullable=False,
+                server_default=sa.text("3600"),
+            ),
+            sa.Column(
+                "sticky_reallocation_budget_threshold_pct",
+                sa.Float(),
+                nullable=False,
+                server_default=sa.text("95.0"),
+            ),
         )
 
     created_api_keys = not _table_exists(bind, "api_keys")

app/db/alembic/versions/20260330_000000_add_cache_locality_settings.py

@@ -0,0 +1,69 @@
+"""add cache-locality settings fields to dashboard_settings
+
+Revision ID: 20260330_000000_add_cache_locality_settings
+Revises: 20260325_000000_add_request_log_cost
+Create Date: 2026-03-30
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.engine import Connection
+
+# revision identifiers, used by Alembic.
+revision = "20260330_000000_add_cache_locality_settings"
+down_revision = "20260325_000000_add_request_log_cost"
+branch_labels = None
+depends_on = None
+
+
+def _columns(connection: Connection, table_name: str) -> set[str]:
+    inspector = sa.inspect(connection)
+    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()
+    columns = _columns(bind, "dashboard_settings")
+    if not columns:
+        return
+
+    if "http_responses_session_bridge_prompt_cache_idle_ttl_seconds" not in columns:
+        with op.batch_alter_table("dashboard_settings") as batch_op:
+            batch_op.add_column(
+                sa.Column(
+                    "http_responses_session_bridge_prompt_cache_idle_ttl_seconds",
+                    sa.Integer(),
+                    nullable=False,
+                    server_default="3600",
+                )
+            )
+
+    if "sticky_reallocation_budget_threshold_pct" not in columns:
+        with op.batch_alter_table("dashboard_settings") as batch_op:
+            batch_op.add_column(
+                sa.Column(
+                    "sticky_reallocation_budget_threshold_pct",
+                    sa.Float(),
+                    nullable=False,
+                    server_default="95.0",
+                )
+            )
+
+
+def downgrade() -> None:
+    bind = op.get_bind()
+    columns = _columns(bind, "dashboard_settings")
+    if not columns:
+        return
+
+    if "http_responses_session_bridge_prompt_cache_idle_ttl_seconds" in columns:
+        with op.batch_alter_table("dashboard_settings") as batch_op:
+            batch_op.drop_column("http_responses_session_bridge_prompt_cache_idle_ttl_seconds")
+
+    if "sticky_reallocation_budget_threshold_pct" in columns:
+        with op.batch_alter_table("dashboard_settings") as batch_op:
+            batch_op.drop_column("sticky_reallocation_budget_threshold_pct")

app/db/models.py

@@ -201,6 +201,18 @@ class DashboardSettings(Base):
     )
     totp_secret_encrypted: Mapped[bytes | None] = mapped_column(LargeBinary, nullable=True)
     totp_last_verified_step: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    http_responses_session_bridge_prompt_cache_idle_ttl_seconds: Mapped[int] = mapped_column(
+        Integer,
+        default=3600,
+        server_default=text("3600"),
+        nullable=False,
+    )
+    sticky_reallocation_budget_threshold_pct: Mapped[float] = mapped_column(
+        Float,
+        default=95.0,
+        server_default=text("95.0"),
+        nullable=False,
+    )
     created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), nullable=False)
     updated_at: Mapped[datetime] = mapped_column(
         DateTime,

app/modules/firewall/repository.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -15,7 +17,7 @@ class FirewallRepository:
     def __init__(self, session: AsyncSession) -> None:
         self._session = session
 
-    async def list_entries(self) -> list[ApiFirewallAllowlist]:
+    async def list_entries(self) -> Sequence[ApiFirewallAllowlist]:
         result = await self._session.execute(
             select(ApiFirewallAllowlist).order_by(ApiFirewallAllowlist.created_at, ApiFirewallAllowlist.ip_address)
         )