feat(codex): 基于用户 API key 生成稳定的 prompt_cache_key,实现跨 provider key 的 prompt 缓存复用

- prepare_context 传递 user_api_key_id 到 CodexRequestContext
- wrap_request 中调用 patch_openai_cli_request_for_codex 注入 prompt_cache_key
- 客户端已提供 prompt_cache_key 时不覆盖
- 补充对应单元测试
This commit is contained in:
fawney19
2026-03-16 17:25:52 +08:00
parent c070e5a9f6
commit c97c9332eb
9 changed files with 156 additions and 8 deletions

View File

@@ -745,6 +745,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
envelope_tls_profile = envelope.prepare_context(
provider_config=getattr(provider, "config", None),
key_id=str(getattr(key, "id", "") or ""),
user_api_key_id=str(getattr(self.api_key, "id", "") or ""),
is_stream=upstream_is_stream,
provider_id=str(getattr(provider, "id", "") or ""),
key=key,

View File

@@ -346,6 +346,7 @@ class CliStreamMixin:
envelope_tls_profile = envelope.prepare_context(
provider_config=getattr(provider, "config", None),
key_id=str(getattr(key, "id", "") or ""),
user_api_key_id=str(getattr(self.api_key, "id", "") or ""),
is_stream=upstream_is_stream,
provider_id=str(getattr(provider, "id", "") or ""),
key=key,

View File

@@ -167,6 +167,7 @@ class CliSyncMixin:
envelope_tls_profile = envelope.prepare_context(
provider_config=getattr(provider, "config", None),
key_id=str(getattr(key, "id", "") or ""),
user_api_key_id=str(getattr(self.api_key, "id", "") or ""),
is_stream=upstream_is_stream,
provider_id=str(getattr(provider, "id", "") or ""),
key=key,

View File

@@ -629,6 +629,7 @@ class ClaudeCodeEnvelope:
*,
provider_config: Any,
key_id: str,
user_api_key_id: str | None = None, # noqa: ARG002
is_stream: bool,
provider_id: str | None = None,
key: Any = None,

View File

@@ -20,6 +20,7 @@ class CodexRequestContext:
"""
account_id: str | None = None
user_api_key_id: str | None = None
is_compact: bool = False

View File

@@ -21,6 +21,7 @@ from src.services.provider.adapters.codex.context import (
get_codex_request_context,
set_codex_request_context,
)
from src.services.provider.adapters.codex.request_patching import patch_openai_cli_request_for_codex
from src.services.provider.request_context import get_selected_base_url
@@ -35,6 +36,26 @@ class CodexOAuthEnvelope:
# synthetic CLI identity headers here.
return None
def prepare_context(
self,
*,
provider_config: Any, # noqa: ARG002
key_id: str, # noqa: ARG002
user_api_key_id: str | None = None,
is_stream: bool, # noqa: ARG002
provider_id: str | None = None, # noqa: ARG002
key: Any = None, # noqa: ARG002
) -> str | None:
existing_ctx = get_codex_request_context()
set_codex_request_context(
CodexRequestContext(
account_id=existing_ctx.account_id if existing_ctx else None,
user_api_key_id=str(user_api_key_id or "").strip() or None,
is_compact=existing_ctx.is_compact if existing_ctx else False,
)
)
return None
def wrap_request(
self,
request_body: dict[str, Any],
@@ -48,19 +69,25 @@ class CodexOAuthEnvelope:
# Compact sentinel may have been popped earlier by finalize_provider_request;
# prefer the pre-set context var (set by adapter), fall back to request body.
existing_ctx = get_codex_request_context()
user_api_key_id = existing_ctx.user_api_key_id if existing_ctx else None
is_compact = (existing_ctx.is_compact if existing_ctx else False) or bool(
request_body.pop("_aether_compact", False)
request_body.get("_aether_compact", False)
)
patched_request_body = patch_openai_cli_request_for_codex(
request_body,
user_api_key_id=user_api_key_id,
)
set_codex_request_context(
CodexRequestContext(
account_id=str(account_id) if account_id else None,
user_api_key_id=user_api_key_id,
is_compact=is_compact,
)
)
# Context 不需要手动清理: FastAPI 每个请求运行在独立的 asyncio Task 中,
# contextvars 天然隔离, Task 结束后自动回收。
# No wire envelope for Codex; keep request body as-is.
return request_body, url_model
return patched_request_body, url_model
def unwrap_response(self, data: Any) -> Any:
# No response envelope for Codex.

View File

@@ -1,18 +1,33 @@
"""Codex provider request patching helpers (passthrough path).
"""Codex provider request patching helpers.
Codex requests are now treated as passthrough:
- Do not mutate client payload fields.
- Only strip internal sentinel fields that must never reach upstream.
Codex requests are mostly passthrough:
- Do not mutate client payload fields unless Codex-specific compatibility requires it.
- Strip internal sentinel fields that must never reach upstream.
- When the caller's user API key is known and the request did not provide one,
synthesize a stable ``prompt_cache_key`` so prompt caching can be reused.
"""
from __future__ import annotations
import uuid
from typing import Any
from src.core.provider_types import ProviderType
def patch_openai_cli_request_for_codex(request_body: dict[str, Any]) -> dict[str, Any]:
def build_stable_codex_prompt_cache_key(user_api_key_id: str | None) -> str | None:
"""Build a deterministic Codex prompt cache key from the caller's user API key id."""
normalized = str(user_api_key_id or "").strip()
if not normalized:
return None
return str(uuid.uuid5(uuid.NAMESPACE_OID, f"aether:codex:prompt-cache:user:{normalized}"))
def patch_openai_cli_request_for_codex(
request_body: dict[str, Any],
*,
user_api_key_id: str | None = None,
) -> dict[str, Any]:
"""
Patch an OpenAI CLI (Responses API style) request body for Codex gateways.
@@ -21,6 +36,11 @@ def patch_openai_cli_request_for_codex(request_body: dict[str, Any]) -> dict[str
out: dict[str, Any] = dict(request_body)
# Internal routing marker; never send upstream.
out.pop("_aether_compact", None)
prompt_cache_key = str(out.get("prompt_cache_key") or "").strip()
if not prompt_cache_key:
stable_key = build_stable_codex_prompt_cache_key(user_api_key_id)
if stable_key:
out["prompt_cache_key"] = stable_key
return out
@@ -29,6 +49,7 @@ def maybe_patch_request_for_codex(
provider_type: str | None,
provider_api_format: str | None,
request_body: Any,
user_api_key_id: str | None = None,
) -> Any:
"""
Conditionally patch request body for Codex gateways.
@@ -44,10 +65,11 @@ def maybe_patch_request_for_codex(
return request_body
if not isinstance(request_body, dict):
return request_body
return patch_openai_cli_request_for_codex(request_body)
return patch_openai_cli_request_for_codex(request_body, user_api_key_id=user_api_key_id)
__all__ = [
"build_stable_codex_prompt_cache_key",
"maybe_patch_request_for_codex",
"patch_openai_cli_request_for_codex",
]

View File

@@ -61,6 +61,7 @@ class ProviderEnvelope(Protocol):
*,
provider_config: Any,
key_id: str,
user_api_key_id: str | None = None,
is_stream: bool,
provider_id: str | None = None,
key: Any = None,