mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
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:
@@ -745,6 +745,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
|
|||||||
envelope_tls_profile = envelope.prepare_context(
|
envelope_tls_profile = envelope.prepare_context(
|
||||||
provider_config=getattr(provider, "config", None),
|
provider_config=getattr(provider, "config", None),
|
||||||
key_id=str(getattr(key, "id", "") or ""),
|
key_id=str(getattr(key, "id", "") or ""),
|
||||||
|
user_api_key_id=str(getattr(self.api_key, "id", "") or ""),
|
||||||
is_stream=upstream_is_stream,
|
is_stream=upstream_is_stream,
|
||||||
provider_id=str(getattr(provider, "id", "") or ""),
|
provider_id=str(getattr(provider, "id", "") or ""),
|
||||||
key=key,
|
key=key,
|
||||||
|
|||||||
@@ -346,6 +346,7 @@ class CliStreamMixin:
|
|||||||
envelope_tls_profile = envelope.prepare_context(
|
envelope_tls_profile = envelope.prepare_context(
|
||||||
provider_config=getattr(provider, "config", None),
|
provider_config=getattr(provider, "config", None),
|
||||||
key_id=str(getattr(key, "id", "") or ""),
|
key_id=str(getattr(key, "id", "") or ""),
|
||||||
|
user_api_key_id=str(getattr(self.api_key, "id", "") or ""),
|
||||||
is_stream=upstream_is_stream,
|
is_stream=upstream_is_stream,
|
||||||
provider_id=str(getattr(provider, "id", "") or ""),
|
provider_id=str(getattr(provider, "id", "") or ""),
|
||||||
key=key,
|
key=key,
|
||||||
|
|||||||
@@ -167,6 +167,7 @@ class CliSyncMixin:
|
|||||||
envelope_tls_profile = envelope.prepare_context(
|
envelope_tls_profile = envelope.prepare_context(
|
||||||
provider_config=getattr(provider, "config", None),
|
provider_config=getattr(provider, "config", None),
|
||||||
key_id=str(getattr(key, "id", "") or ""),
|
key_id=str(getattr(key, "id", "") or ""),
|
||||||
|
user_api_key_id=str(getattr(self.api_key, "id", "") or ""),
|
||||||
is_stream=upstream_is_stream,
|
is_stream=upstream_is_stream,
|
||||||
provider_id=str(getattr(provider, "id", "") or ""),
|
provider_id=str(getattr(provider, "id", "") or ""),
|
||||||
key=key,
|
key=key,
|
||||||
|
|||||||
@@ -629,6 +629,7 @@ class ClaudeCodeEnvelope:
|
|||||||
*,
|
*,
|
||||||
provider_config: Any,
|
provider_config: Any,
|
||||||
key_id: str,
|
key_id: str,
|
||||||
|
user_api_key_id: str | None = None, # noqa: ARG002
|
||||||
is_stream: bool,
|
is_stream: bool,
|
||||||
provider_id: str | None = None,
|
provider_id: str | None = None,
|
||||||
key: Any = None,
|
key: Any = None,
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ class CodexRequestContext:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
account_id: str | None = None
|
account_id: str | None = None
|
||||||
|
user_api_key_id: str | None = None
|
||||||
is_compact: bool = False
|
is_compact: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ from src.services.provider.adapters.codex.context import (
|
|||||||
get_codex_request_context,
|
get_codex_request_context,
|
||||||
set_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
|
from src.services.provider.request_context import get_selected_base_url
|
||||||
|
|
||||||
|
|
||||||
@@ -35,6 +36,26 @@ class CodexOAuthEnvelope:
|
|||||||
# synthetic CLI identity headers here.
|
# synthetic CLI identity headers here.
|
||||||
return None
|
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(
|
def wrap_request(
|
||||||
self,
|
self,
|
||||||
request_body: dict[str, Any],
|
request_body: dict[str, Any],
|
||||||
@@ -48,19 +69,25 @@ class CodexOAuthEnvelope:
|
|||||||
# Compact sentinel may have been popped earlier by finalize_provider_request;
|
# 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.
|
# prefer the pre-set context var (set by adapter), fall back to request body.
|
||||||
existing_ctx = get_codex_request_context()
|
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(
|
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(
|
set_codex_request_context(
|
||||||
CodexRequestContext(
|
CodexRequestContext(
|
||||||
account_id=str(account_id) if account_id else None,
|
account_id=str(account_id) if account_id else None,
|
||||||
|
user_api_key_id=user_api_key_id,
|
||||||
is_compact=is_compact,
|
is_compact=is_compact,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
# Context 不需要手动清理: FastAPI 每个请求运行在独立的 asyncio Task 中,
|
# Context 不需要手动清理: FastAPI 每个请求运行在独立的 asyncio Task 中,
|
||||||
# contextvars 天然隔离, Task 结束后自动回收。
|
# contextvars 天然隔离, Task 结束后自动回收。
|
||||||
# No wire envelope for Codex; keep request body as-is.
|
# 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:
|
def unwrap_response(self, data: Any) -> Any:
|
||||||
# No response envelope for Codex.
|
# No response envelope for Codex.
|
||||||
|
|||||||
@@ -1,18 +1,33 @@
|
|||||||
"""Codex provider request patching helpers (passthrough path).
|
"""Codex provider request patching helpers.
|
||||||
|
|
||||||
Codex requests are now treated as passthrough:
|
Codex requests are mostly passthrough:
|
||||||
- Do not mutate client payload fields.
|
- Do not mutate client payload fields unless Codex-specific compatibility requires it.
|
||||||
- Only strip internal sentinel fields that must never reach upstream.
|
- 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
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import uuid
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from src.core.provider_types import ProviderType
|
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.
|
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)
|
out: dict[str, Any] = dict(request_body)
|
||||||
# Internal routing marker; never send upstream.
|
# Internal routing marker; never send upstream.
|
||||||
out.pop("_aether_compact", None)
|
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
|
return out
|
||||||
|
|
||||||
|
|
||||||
@@ -29,6 +49,7 @@ def maybe_patch_request_for_codex(
|
|||||||
provider_type: str | None,
|
provider_type: str | None,
|
||||||
provider_api_format: str | None,
|
provider_api_format: str | None,
|
||||||
request_body: Any,
|
request_body: Any,
|
||||||
|
user_api_key_id: str | None = None,
|
||||||
) -> Any:
|
) -> Any:
|
||||||
"""
|
"""
|
||||||
Conditionally patch request body for Codex gateways.
|
Conditionally patch request body for Codex gateways.
|
||||||
@@ -44,10 +65,11 @@ def maybe_patch_request_for_codex(
|
|||||||
return request_body
|
return request_body
|
||||||
if not isinstance(request_body, dict):
|
if not isinstance(request_body, dict):
|
||||||
return request_body
|
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__ = [
|
__all__ = [
|
||||||
|
"build_stable_codex_prompt_cache_key",
|
||||||
"maybe_patch_request_for_codex",
|
"maybe_patch_request_for_codex",
|
||||||
"patch_openai_cli_request_for_codex",
|
"patch_openai_cli_request_for_codex",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -61,6 +61,7 @@ class ProviderEnvelope(Protocol):
|
|||||||
*,
|
*,
|
||||||
provider_config: Any,
|
provider_config: Any,
|
||||||
key_id: str,
|
key_id: str,
|
||||||
|
user_api_key_id: str | None = None,
|
||||||
is_stream: bool,
|
is_stream: bool,
|
||||||
provider_id: str | None = None,
|
provider_id: str | None = None,
|
||||||
key: Any = None,
|
key: Any = None,
|
||||||
|
|||||||
@@ -6,7 +6,9 @@ import jwt
|
|||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from src.api.handlers.base.request_builder import PassthroughRequestBuilder
|
from src.api.handlers.base.request_builder import PassthroughRequestBuilder
|
||||||
|
from src.services.provider.adapters.codex.context import set_codex_request_context
|
||||||
from src.services.provider.adapters.codex.request_patching import (
|
from src.services.provider.adapters.codex.request_patching import (
|
||||||
|
build_stable_codex_prompt_cache_key,
|
||||||
maybe_patch_request_for_codex,
|
maybe_patch_request_for_codex,
|
||||||
patch_openai_cli_request_for_codex,
|
patch_openai_cli_request_for_codex,
|
||||||
)
|
)
|
||||||
@@ -47,6 +49,37 @@ def test_patch_openai_cli_request_for_codex_is_passthrough_except_internal_senti
|
|||||||
assert out["input"][0]["role"] == "system"
|
assert out["input"][0]["role"] == "system"
|
||||||
|
|
||||||
|
|
||||||
|
def test_patch_openai_cli_request_for_codex_generates_stable_prompt_cache_key_from_user_api_key_id() -> (
|
||||||
|
None
|
||||||
|
):
|
||||||
|
req = {"model": "gpt-test", "input": []}
|
||||||
|
|
||||||
|
out = patch_openai_cli_request_for_codex(req, user_api_key_id="user-key-123")
|
||||||
|
|
||||||
|
assert out is not req
|
||||||
|
assert out["prompt_cache_key"] == build_stable_codex_prompt_cache_key("user-key-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_patch_openai_cli_request_for_codex_preserves_existing_prompt_cache_key() -> None:
|
||||||
|
req = {"model": "gpt-test", "input": [], "prompt_cache_key": "client-cache-key"}
|
||||||
|
|
||||||
|
out = patch_openai_cli_request_for_codex(req, user_api_key_id="user-key-123")
|
||||||
|
|
||||||
|
assert out["prompt_cache_key"] == "client-cache-key"
|
||||||
|
|
||||||
|
|
||||||
|
def test_patch_openai_cli_request_for_codex_ignores_provider_key_and_uses_only_user_api_key_id() -> (
|
||||||
|
None
|
||||||
|
):
|
||||||
|
req = {"model": "gpt-test", "input": []}
|
||||||
|
|
||||||
|
out_a = patch_openai_cli_request_for_codex(req, user_api_key_id="user-key-123")
|
||||||
|
out_b = patch_openai_cli_request_for_codex(req, user_api_key_id="user-key-123")
|
||||||
|
|
||||||
|
assert out_a["prompt_cache_key"] == out_b["prompt_cache_key"]
|
||||||
|
assert out_a["prompt_cache_key"] == build_stable_codex_prompt_cache_key("user-key-123")
|
||||||
|
|
||||||
|
|
||||||
def test_maybe_patch_request_for_codex_is_noop_for_non_codex() -> None:
|
def test_maybe_patch_request_for_codex_is_noop_for_non_codex() -> None:
|
||||||
req = {"model": "gpt-test", "input": []}
|
req = {"model": "gpt-test", "input": []}
|
||||||
out = maybe_patch_request_for_codex(
|
out = maybe_patch_request_for_codex(
|
||||||
@@ -119,6 +152,66 @@ def test_codex_envelope_extra_headers_does_not_inject_synthetic_headers() -> Non
|
|||||||
assert codex_oauth_envelope.extra_headers() is None
|
assert codex_oauth_envelope.extra_headers() is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_codex_envelope_wrap_request_injects_stable_prompt_cache_key_from_user_api_key() -> None:
|
||||||
|
from src.services.provider.adapters.codex.envelope import codex_oauth_envelope
|
||||||
|
|
||||||
|
try:
|
||||||
|
codex_oauth_envelope.prepare_context(
|
||||||
|
provider_config=None,
|
||||||
|
key_id="provider-key-123",
|
||||||
|
user_api_key_id="user-key-123",
|
||||||
|
is_stream=True,
|
||||||
|
)
|
||||||
|
out, url_model = codex_oauth_envelope.wrap_request(
|
||||||
|
{"model": "gpt-test", "input": []},
|
||||||
|
model="gpt-test",
|
||||||
|
url_model=None,
|
||||||
|
decrypted_auth_config=None,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
set_codex_request_context(None)
|
||||||
|
|
||||||
|
assert url_model is None
|
||||||
|
assert out["prompt_cache_key"] == build_stable_codex_prompt_cache_key("user-key-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_codex_envelope_wrap_request_same_user_different_provider_keys_share_prompt_cache_key() -> (
|
||||||
|
None
|
||||||
|
):
|
||||||
|
from src.services.provider.adapters.codex.envelope import codex_oauth_envelope
|
||||||
|
|
||||||
|
try:
|
||||||
|
codex_oauth_envelope.prepare_context(
|
||||||
|
provider_config=None,
|
||||||
|
key_id="provider-key-123",
|
||||||
|
user_api_key_id="user-key-123",
|
||||||
|
is_stream=True,
|
||||||
|
)
|
||||||
|
out_a, _ = codex_oauth_envelope.wrap_request(
|
||||||
|
{"model": "gpt-test", "input": []},
|
||||||
|
model="gpt-test",
|
||||||
|
url_model=None,
|
||||||
|
decrypted_auth_config=None,
|
||||||
|
)
|
||||||
|
codex_oauth_envelope.prepare_context(
|
||||||
|
provider_config=None,
|
||||||
|
key_id="provider-key-456",
|
||||||
|
user_api_key_id="user-key-123",
|
||||||
|
is_stream=True,
|
||||||
|
)
|
||||||
|
out_b, _ = codex_oauth_envelope.wrap_request(
|
||||||
|
{"model": "gpt-test", "input": []},
|
||||||
|
model="gpt-test",
|
||||||
|
url_model=None,
|
||||||
|
decrypted_auth_config=None,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
set_codex_request_context(None)
|
||||||
|
|
||||||
|
assert out_a["prompt_cache_key"] == out_b["prompt_cache_key"]
|
||||||
|
assert out_a["prompt_cache_key"] == build_stable_codex_prompt_cache_key("user-key-123")
|
||||||
|
|
||||||
|
|
||||||
def test_codex_passthrough_builder_preserves_real_codex_headers() -> None:
|
def test_codex_passthrough_builder_preserves_real_codex_headers() -> None:
|
||||||
from src.services.provider.adapters.codex.envelope import codex_oauth_envelope
|
from src.services.provider.adapters.codex.envelope import codex_oauth_envelope
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user