diff --git a/.env.example b/.env.example index 1c5b896fe..d9e43a942 100644 --- a/.env.example +++ b/.env.example @@ -186,8 +186,11 @@ # # TTL and max entries for the sequential strategy's in-memory session->credential # sticky cache. These do not enable disk persistence. -#SESSION_STICKY_ENTRY_TTL_SECONDS=3600 -#SESSION_STICKY_ENTRY_TTL_SECONDS_GEMINI_CLI=3600 +# Default idle cache-affinity lifetime is 300 seconds. Override per provider +# when its upstream prompt-cache invalidation period differs. Set to 0 to disable +# sticky credential bindings for that provider. +#SESSION_STICKY_ENTRY_TTL_SECONDS=300 +#SESSION_STICKY_ENTRY_TTL_SECONDS_GEMINI_CLI=300 #SESSION_STICKY_MAX_ENTRIES=10000 #SESSION_STICKY_MAX_ENTRIES_GEMINI_CLI=10000 # @@ -195,6 +198,11 @@ # Leave unset unless your clients send stable conversation IDs. #TRUSTED_SESSION_ID_FIELDS="conversation_id,thread_id" +# Persist inferred session lineage across proxy restarts. Disabled by default; +# state is stored as content-free schema-3 hashes in session_stickiness.json. +#SESSION_PERSISTENCE_ENABLED=true +#SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS=5.0 + # --- Priority-Based Concurrency Multipliers --- # Credentials can be assigned to priority tiers (1=highest, 2, 3, etc.). # Each tier can have a concurrency multiplier that increases the effective @@ -265,6 +273,15 @@ # EXHAUSTION_COOLDOWN_THRESHOLD_GEMINI_CLI=300 # EXHAUSTION_COOLDOWN_THRESHOLD=300 # Global fallback for all providers +# No-reset quota exhaustion policy. Some quota APIs report an exhausted bucket +# without a reset timestamp when the account cannot access that model/group. +# Policies: warn_only | cooldown | disable_scope +# Providers default to warn_only unless they explicitly opt into a fallback. +# QUOTA_NO_RESET_EXHAUSTION_POLICY_PROVIDER=cooldown +# QUOTA_NO_RESET_COOLDOWN_SECONDS_PROVIDER=86400 +# QUOTA_NO_RESET_EXHAUSTION_POLICY=warn_only +# QUOTA_NO_RESET_COOLDOWN_SECONDS=0 + # ------------------------------------------------------------------------------ # | [ADVANCED] Custom Caps | # ------------------------------------------------------------------------------ @@ -333,6 +350,98 @@ # Default: false # STREAM_RETRY_ON_REASONING_ONLY=false +# --- Optional Structured Config File --- +# Optional JSON config for structured experimental settings such as fallback +# groups, model routes, advisory pricing, streaming observability, and field +# cache rules. Environment variables override JSON config. Do not put API keys, +# OAuth tokens, bearer tokens, or authorization headers in this file. +# Supported top-level sections include: routing, pricing, streaming, retry, +# responses, field_cache, and providers. The providers section is for non-secret +# provider tuning and config-defined custom providers; credentials must stay in +# env or provider credential files. A custom provider can declare api_base, +# protocol_name (openai_chat, responses, anthropic_messages, or gemini), +# endpoint_paths, auth_mode (bearer, x-api-key, x-goog-api-key, custom, or none), +# auth_header_name, models, default_output_protocol, adapter_names, +# adapter_config, native_streaming_supported, field_cache, and model_quota_groups. +# Endpoint paths stay on api_base and may use {model}, {operation}, or {provider}; +# transport selectors such as ?alt=sse are allowed, but credential-like query +# keys, URL credentials, fragments, and cross-origin endpoint URLs are rejected. +# Custom provider definitions are discovered and snapshotted at process startup. +# Supply credentials separately as _API_KEY variables or through the +# existing provider credential facilities. Never put credentials or OAuth data +# in the JSON file. +# LLM_PROXY_CONFIG_FILE=./config/llm-proxy.json + +# Non-streaming generation routes accept X-Proxy-Output-Protocol to select an +# independent response format. Explicit library arguments take priority, then +# this header, then a provider default, then the input protocol. Cross-protocol +# streaming is rejected until canonical stream conversion is enabled. + +# --- Fallback Groups and Model Routes --- +# Ordered fallback groups try targets in order when the previous target fails +# with a retryable provider/category error. Execution suffixes: +# @auto Let the provider choose custom/native/LiteLLM behavior +# @native Require native protocol execution +# @custom Require provider custom execution +# @litellm_fallback Explicitly use LiteLLM fallback +# FALLBACK_GROUPS=code_chain +# FALLBACK_GROUP_CODE_CHAIN=codex/gpt-5.1-codex@native,openai/gpt-5.1@litellm_fallback +# MODEL_ROUTE_CODE=group:code_chain + +# --- Provider Cooldown Activation --- +# Provider-level cooldown is conservative and only intended for large/global +# retry-after events, not every per-credential quota error. Model-capacity +# errors can start a model-scoped cooldown without blocking unrelated models. +# PROVIDER_COOLDOWN_MIN_SECONDS=10 +# PROVIDER_COOLDOWN_DEFAULT_SECONDS=30 +# PROVIDER_COOLDOWN_ON_QUOTA=false +# Repeated transient provider/model failures can increase bounded backoff. +# PROVIDER_BACKOFF_WINDOW_SECONDS=60 +# PROVIDER_BACKOFF_THRESHOLD=3 +# PROVIDER_BACKOFF_BASE_SECONDS=0 # 0/unset means use provider cooldown default +# PROVIDER_BACKOFF_MAX_SECONDS=300 +# FAILURE_HISTORY_MAX_ENTRIES=200 + +# --- Responses API Store Policy --- +# Responses are stored in-memory by default. Use provider_cache for durable JSON +# storage via the existing provider-cache layer. TTL/max limits are disabled +# when unset or <= 0. Failed stream responses are stored by default; in-progress +# streaming state is disabled unless explicitly enabled. +# RESPONSES_STORE_BACKEND=memory +# RESPONSES_STORE_CACHE_NAME=responses +# RESPONSES_STORE_CACHE_PREFIX=responses +# RESPONSES_STORE_CACHE_DIR= +# RESPONSES_STORE_CACHE_MEMORY_TTL_SECONDS=3600 +# RESPONSES_STORE_CACHE_DISK_TTL_SECONDS=172800 +# RESPONSES_STORE_TTL_SECONDS=0 +# RESPONSES_STORE_MAX_ITEMS=0 +# RESPONSES_STORE_FAILED=true +# RESPONSES_STORE_IN_PROGRESS=false + +# --- Streaming Observability --- +# Stream lifecycle metrics are traced by default. TTFB/stall timeout values are +# active when set to >0; keep at 0 to disable. Heartbeats are SSE comments and +# do not count as visible output. +# STREAM_TRACE_METRICS=true +# STREAM_TTFB_TIMEOUT_SECONDS=0 +# STREAM_STALL_TIMEOUT_SECONDS=0 +# STREAM_HEARTBEAT_INTERVAL_SECONDS=0 +# STREAM_HEARTBEAT_SECONDS=0 # Legacy alias +# STREAM_CANCEL_UPSTREAM_ON_DISCONNECT=true + +# --- Advisory Model Pricing --- +# Per-token advisory prices used only when providers do not report actual cost. +# Precedence: skip-cost provider setting > provider-reported cost/SSE cost event +# > provider explicit pricing > env pricing > JSON pricing > LiteLLM metadata. +# Streaming providers can report actual cost with `: cost {...}` comments or +# `event: cost` frames. +# Env names sanitize provider/model by replacing non-alphanumerics with `_`. +# MODEL_PRICE_OPENAI_GPT_5_1_INPUT=0.000001 +# MODEL_PRICE_OPENAI_GPT_5_1_OUTPUT=0.00001 +# MODEL_PRICE_OPENAI_GPT_5_1_CACHE_READ=0.0000001 +# MODEL_PRICE_OPENAI_GPT_5_1_CACHE_WRITE=0.000001 +# MODEL_PRICE_OPENAI_GPT_5_1_REASONING=0.00001 + # ------------------------------------------------------------------------------ # | [ADVANCED] HTTP Timeout Configuration | # ------------------------------------------------------------------------------ diff --git a/.gitignore b/.gitignore index 7ecc7b5d7..97d2cbd38 100644 --- a/.gitignore +++ b/.gitignore @@ -139,6 +139,7 @@ cache/antigravity/thought_signatures.json !tests/ tests/* !tests/test_classifier_scoped_routing.py +!tests/test_nvidia_provider.py !tests/test_session_tracking.py !tests/test_selection_engine.py docs/ignored/ diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 5ab9e6bbd..1ab5a08e3 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -15,7 +15,7 @@ ## Layers **Proxy Application Layer:** -- Purpose: Expose OpenAI- and Anthropic-compatible HTTP endpoints, handle auth, logging, TUI +- Purpose: Expose OpenAI Chat Completions, OpenAI Responses, and Anthropic-compatible HTTP endpoints, handle auth, logging, TUI - Location: `src/proxy_app/` - Contains: FastAPI app, route handlers, Pydantic request/response models, launcher TUI, quota viewer - Depends on: `rotator_library`, `litellm`, `fastapi`, `uvicorn` @@ -49,6 +49,13 @@ - Depends on: `rotator_library.core` - Used by: Client layer's `AnthropicHandler`, proxy routes `/v1/messages`, `/v1/messages/count_tokens` +**Responses API Layer:** +- Purpose: OpenAI Responses API compatibility — create, store, retrieve, and stream response objects with `previous_response_id` continuation +- Location: `src/rotator_library/responses/` +- Contains: `ResponsesService` (orchestrator + `ResponsesServiceError`), `ResponsesBridge` (Responses ↔ chat-completions translation via `ResponsesProtocol`), `ResponsesStore` protocol with `InMemoryResponsesStore` / `ProviderCacheResponsesStore` backends and `create_configured_responses_store` factory, `ResponsesSSEFormatter` / `ResponsesWebSocketFormatter` / `ResponsesStreamEvent` (streaming), `StoredResponse` / `ResponsesStoreSettings` / `generate_response_id` (types) +- Depends on: `rotator_library.protocols`, `rotator_library.streaming`, `rotator_library.usage` (costs/accounting), `rotator_library.client` via injected `RotatingClient` +- Used by: Proxy routes `/v1/responses`, `/v1/responses/{response_id}`, `/v1/responses/{response_id}/input_items` + **Core Types & Config Layer:** - Purpose: Shared type definitions, constants, error classification, config defaults - Location: `src/rotator_library/core/`, `src/rotator_library/config/` @@ -77,6 +84,17 @@ 4. Response is translated back to Anthropic format — `src/rotator_library/anthropic_compat/translator.py` 5. For streaming, `anthropic_streaming_wrapper` wraps the SSE stream — `src/rotator_library/anthropic_compat/streaming.py` +**Responses API Request:** + +1. Client sends POST to `/v1/responses` (streaming or non-streaming) — `src/proxy_app/main.py` +2. `ResponsesService.create_response()` / `stream_response()` parses the payload via `ResponsesProtocol` — `src/rotator_library/responses/service.py` +3. `previous_response_id` resolves the parent `StoredResponse` and its lineage from the scoped `ResponsesStore` for continuation — `src/rotator_library/responses/service.py`, `src/rotator_library/responses/store.py` +4. `ResponsesBridge.to_chat_kwargs()` translates the Responses payload into chat-completions kwargs and emits session-tracking hints — `src/rotator_library/responses/bridge.py` +5. Execution flows through `RotatingClient.acompletion()`, reusing the standard retry/rotation/session-tracking path — `src/rotator_library/client/rotating_client.py` +6. `ResponsesBridge.from_chat_response()` shapes the chat result back into Responses format; usage is recorded via `extract_usage_record` and cost via `CostCalculator` — `src/rotator_library/responses/bridge.py`, `src/rotator_library/usage/` +7. When `store` is true, a `StoredResponse` is persisted scoped by the session isolation key for later retrieval and continuation — `src/rotator_library/responses/store.py` +8. Streaming requests emit Responses SSE events (`response.created`, `response.output_item.added`, `response.output_text.delta`, `response.output_item.done`, `response.completed`) from the chat SSE stream via `ResponsesSSEFormatter` — `src/rotator_library/responses/streaming.py` + **Provider Discovery:** 1. `providers/__init__.py` scans all `*_provider.py` files in `src/rotator_library/providers/` @@ -89,12 +107,12 @@ 1. `RequestContextBuilder` calls `provider.get_session_tracking_hints()` to collect provider-specific evidence — `src/rotator_library/client/request_builder.py` 2. `SessionTracker.infer_session()` builds scoped anchors from explicit IDs, message content, tool-call IDs, and provider hints — `src/rotator_library/session_tracking.py` 3. System/developer prompts are excluded from continuity anchors to prevent harness-level system prompts from cross-binding independent sessions -4. Compaction probe anchors are built separately (`_build_compaction_probe_anchors()`) from large early messages or explicit summarization markers; these identify lineage parents but are not stored on the new child session +4. Compaction probe anchors are built separately (`_build_compaction_probe_anchors()`) from early user/system/developer messages only (assistant, tool, and function-result history is never probed); these identify lineage parents but are not stored on the new child session 5. Normal anchors suppress compaction probe indexes to avoid double-counting summary text as continuity evidence 6. Anchors are namespaced by scope/provider/model so sticky evidence never leaks between credential pools 7. `_best_match()` scores anchor overlap against live sessions with comprehensive tiebreaker (score, strong matches, medium matches, group diversity, response matches, last_seen, session_id); confidence is `strong` (any strong match), `probable` (diverse medium evidence), `weak`, or `none` -8. `_is_compaction_parent_match()` requires response anchor overlap for size-only probes (noisy) or any score > 0 for explicit marker probes -9. Compaction lineage is tracked via `lineage_parent_session_id` on `SessionInference` but does not force sticky continuation +8. `_evaluate_compaction()` requires structural replacement of more than half the parent's high-water request history (`_retained_history_ratio()` below `_COMPACTION_MAX_RETAINED_HISTORY_RATIO`); unmarked summaries must additionally overlap at least two distinct response events (`_MIN_UNMARKED_RESPONSE_GROUPS`) plus a retained request group, while explicit marker probes qualify on any score > 0 +9. Compaction lineage is tracked via `lineage_parent_session_id` on `SessionInference` but does not force sticky continuation of the parent; exact resends of a validated compacted payload reuse the already-created child session via opaque `compaction_replay` anchors (`_find_compaction_replay()` / `_compaction_replay_anchor()`), and changing non-probe history invalidates the replay key; post-compaction requests that retain the validated compacted base context but extend or alter the tail continue the child via `compaction_context` anchors (`_find_compaction_context()` / `_compaction_context_anchor()`) minted only from probe groups that matched parent response evidence; authoritative identity (trusted explicit IDs or provider affinity, via `_is_authoritative_identity_anchor()`) takes precedence over replay/context bindings and suppresses unrelated compaction lineage 10. Returns `SessionInference` with `session_id`, `affinity_key` (deterministic first-pick hint), confidence, and namespace **Credential Selection:** @@ -132,8 +150,8 @@ - Purpose: TTL-based session inference using scoped, compounding evidence anchors with confidence scoring - Location: `src/rotator_library/session_tracking.py` - Pattern: Evidence accumulator with thread-safe anchor store (`threading.RLock` for state, separate `threading.Lock` for save I/O), namespace isolation, schema-versioned JSON persistence via `ResilientStateWriter` -- Key types: `SessionAnchor` (evidence with strength/source/group), `SessionTrackingHints` (provider evidence), `SessionInference` (result with confidence + affinity + lineage), `_MatchCandidate` (scored candidate with `last_seen` tiebreaker) -- Key methods: `infer_session()`, `record_response()`, `flush()`, `_build_compaction_probe_anchors()`, `_is_compaction_parent_match()`, `_prepare_save_locked()`, `_write_save_job()` +- Key types: `SessionAnchor` (evidence with strength/source/group), `SessionTrackingHints` (provider evidence), `SessionInference` (result with confidence + affinity + lineage), `_MatchCandidate` (scored candidate with `response_groups`/`request_groups`/`matched_probe_groups` tracking and `last_seen` tiebreaker), `_CompactionDecision` (validated parent lineage with retained-history ratio and `context_probe_groups`) +- Key methods: `infer_session()`, `record_response()`, `flush()`, `_build_compaction_probe_anchors()`, `_evaluate_compaction()`, `_find_compaction_replay()`, `_find_compaction_context()`, `_compaction_replay_anchor()`, `_compaction_context_anchor()`, `_is_authoritative_identity_anchor()`, `_compaction_marker_probe_groups()`, `_log_inference_decision()`, `_retained_history_ratio()`, `_prepare_save_locked()`, `_write_save_job()` - Save I/O: Dirty generation counter (`_dirty_generation`) decouples state mutations from disk writes; `_prepare_save_locked()` snapshots under state lock, `_write_save_job()` writes under separate I/O lock to avoid blocking inference ## Entry Points @@ -164,6 +182,7 @@ - Error classification: `src/rotator_library/core/errors.py` — `classify_error()`, `should_rotate_on_error()`, `should_retry_same_key()` - Error handler with cooldown parsing: `src/rotator_library/error_handler.py` — parses retry-after headers, duration strings, sets provider cooldowns +- No-reset quota exhaustion: Authoritative quota APIs that report an exhausted bucket with no reset timestamp (e.g., account lacks model-group entitlement) are handled by `UsageManager._handle_no_reset_quota_exhaustion()` in `src/rotator_library/usage/manager.py` — policy `warn_only` | `cooldown` | `disable_scope` from `ProviderUsageConfig.no_reset_exhaustion_policy` applies a scoped fallback cooldown instead of repeated retries - Streaming errors: `StreamedAPIError` raised mid-stream to trigger credential rotation - Credential reauth: `CredentialNeedsReauthError` triggers background OAuth refresh @@ -173,8 +192,8 @@ **Caching:** Provider instances are singletons via `SingletonABCMeta`. Provider-level HTTP caching via `provider_cache.py`. Model info cached by `ModelInfoService` with async refresh. -**Storage:** JSON file persistence for usage data (`usage/usage_*.json`), OAuth credentials in `oauth_creds/`, transaction logs in `logs/transactions/`. Session state persisted to JSON via `ResilientStateWriter` when disk persistence is enabled. Config via `.env` files and environment variables. +**Storage:** JSON file persistence for usage data (`usage/usage_*.json`), OAuth credentials in `oauth_creds/`, transaction logs in `logs/transactions/` written by `TransactionLogger` (`src/rotator_library/transaction_logger.py`) with per-request directories containing client/provider I/O and a JSON-safe payload converter (`_make_json_safe`) for Pydantic/dataclass/timestamp objects. Session state persisted to JSON via `ResilientStateWriter` when disk persistence is enabled. Config via `.env` files and environment variables. **Background Tasks:** `BackgroundRefresher` manages periodic OAuth token refresh (default 10 min) and provider-specific background jobs (quota refresh, etc.) with independent timers. -**Session Tracking:** Thread-safe with two-lock design (`threading.RLock` for anchor store, `threading.Lock` for save I/O), scoped by usage scope/provider/model. Anchor strength levels: `strong` (tool IDs, provider affinity), `medium` (message content hashes, response anchors), `weak` (first-user text, untrusted explicit IDs). System/developer prompts excluded from continuity anchors. Compaction detection uses separate probe anchors (`_build_compaction_probe_anchors()`) with response-anchor validation for size-only probes; lineage tracked via `lineage_parent_session_id` without forcing sticky continuation. Persistence uses schema-versioned JSON (`_PERSISTENCE_SCHEMA_VERSION`) with generation-based write deduplication. Configurable via `TRUSTED_SESSION_ID_FIELDS` env var and per-provider `SESSION_STICKY_*` env vars. +**Session Tracking:** Thread-safe with two-lock design (`threading.RLock` for anchor store, `threading.Lock` for save I/O), scoped by usage scope/provider/model. Anchor strength levels: `strong` (trusted explicit IDs, provider affinity keys, response global IDs), `medium` (message content hashes, response anchors), `weak` (first-user text, raw tool-call IDs, untrusted explicit IDs). System/developer prompts excluded from continuity anchors. Compaction detection uses separate probe anchors (`_build_compaction_probe_anchors()`) restricted to early user/system/developer messages and requires structural replacement of more than half the parent's high-water request history via `_evaluate_compaction()`; unmarked summaries must additionally overlap at least two distinct response events plus a retained request group; authoritative identity (trusted explicit or provider, via `_is_authoritative_identity_anchor()`) suppresses unrelated compaction lineage; exact resends of a validated compacted payload reuse the child session via opaque `compaction_replay` anchors while changed-tail continuations bind via `compaction_context` anchors (`_find_compaction_context()`) minted only from probe groups that matched parent response evidence. Lineage tracked via `lineage_parent_session_id` without forcing sticky continuation of the parent. Trimming and TTL pruning evict weak/ordinary evidence before replay/context identity (`_anchor_eviction_key()` ranks `compaction_context` and `compaction_replay` sources above ordinary anchors) with deterministic value tie-breaking, and late responses cannot resurrect an expired session; session namespaces are immutable — `_refresh_and_bridge()` rejects namespace drift and `record_response()` normalizes fallback callbacks to the session's original namespace. Streaming response identity is recorded only after an explicit provider completion signal (`_sse_has_completion_signal()`: usage-backed final chunk, `finish_reason` paired with usage, or `[DONE]`). Persistence uses schema-versioned JSON (`_PERSISTENCE_SCHEMA_VERSION`) with generation-based write deduplication, dirty state retained on failed writes for retry, stale delayed generations rejected, and anchor ownership rebuilt on load (rejecting malformed containers, non-finite timestamps, expired sessions, orphan anchors, namespace mismatches, invalid strengths, and unsupported schemas); rebuilt sessions are flagged `loaded_from_persistence` for lineage diagnostics. Configurable via `TRUSTED_SESSION_ID_FIELDS` env var, per-provider `SESSION_STICKY_*` env vars, and `SESSION_PERSISTENCE_ENABLED` / `SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS` env vars on `RotatingClient`; every inference emits a temporary warning-level `Session tracker decision` line (`_log_inference_decision()`) with action, session IDs, namespace, confidence, score, persistence origin, and compaction evidence. diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index 62949c9ff..153f19ae4 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -977,31 +977,120 @@ ROTATION_MODE_OPENAI=balanced # Optional sequential session-stickiness controls SESSION_STICKY_WAIT_SECONDS=15 -SESSION_STICKY_ENTRY_TTL_SECONDS=3600 +SESSION_STICKY_ENTRY_TTL_SECONDS=300 SESSION_STICKY_MAX_ENTRIES=10000 ``` #### Session Stickiness Sequential routing uses scoped session evidence to keep related chat requests on -the same credential where possible. Evidence is always bound to the resolved -credential scope, provider, and model/session scope, so classifier/private scopes -do not share sticky sessions. Explicit request IDs are weak by default because -some clients generate random IDs per request; set `TRUSTED_SESSION_ID_FIELDS` only -when the client-provided fields are known stable. +the same credential where possible. Logical session IDs and generic conversation +evidence cross providers and models inside one caller/credential domain. Public, +named-classifier, and ad hoc private credential-bundle domains never share +anchors, trusted IDs, compaction lineage, replay bindings, usage managers, or +Responses storage. Provider-native anchors remain qualified by provider and the +provider's optional `session_scope`. + +Credential affinity is separate from logical lineage. Sticky entries remain +independent per provider/model and expire after that provider's idle TTL. Activity +on Anthropic therefore cannot refresh a DeepSeek binding in the same global +session. The default is 300 seconds; providers may declare +`default_session_sticky_entry_ttl_seconds`, per-provider environment overrides +win, and a TTL of `0` disables sticky bindings for that provider. + +Explicit request IDs are weak by default because some clients generate random IDs +per request; set `TRUSTED_SESSION_ID_FIELDS` only when the fields are stable. Provider plugins may implement `get_session_tracking_hints()` to contribute -provider-specific anchors or a provider session scope. The hook supplies evidence -only; credential selection remains centralized in the rotator. +provider-qualified anchors, affinity, or a native session scope. Proxy-owned +global hints such as Responses IDs use a separate internal-only channel, so a +provider cannot promote its opaque ID into another provider's identity domain. -Compaction lineage detection is intentionally conservative telemetry. It can use -summary-like or unusually large early messages in any role as temporary parent -lookup evidence, but those probe-only anchors are not stored on the child session -and do not force sticky continuation. +Session inference persistence is opt-in for the proxy process: + +```env +SESSION_PERSISTENCE_ENABLED=true +SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS=5.0 +``` + +Set the flush interval to `0` for immediate writes during manual restart testing. +After restart, a continued lineage warning reports `origin=persisted`; schema-3 +state is written to `session_stickiness.json` under the proxy data directory. +External IDs are hashed, file/session/string/source bounds are enforced on load, +and schema-2 provider/model-scoped state is rejected rather than merged. Only +logical lineage is persisted; credential bindings remain in memory. + +Compaction lineage detection is intentionally conservative. A candidate must be +an early user/system/developer message and must replace more than half of the +matched parent's high-water request history. Replacing a minority of history, +including turns in the middle, remains ordinary continuity. Assistant/tool +history is never treated as a size-only compaction probe. + +Explicit summary markers still require matching parent evidence and structural +replacement. Unmarked summaries additionally require overlap with at least two +distinct completed response events and request-side parent evidence, so output +aggregation prompts cannot become compaction merely by quoting responses. +Probe-only summary anchors are not stored as ordinary continuity. An opaque +replay anchor identifies an exact resend, while context anchors derived only from +probe messages that actually matched response evidence keep changed or extended +tails on the validated child. Shared system/user harness probes therefore cannot +become child bindings by position alone. Strong trusted explicit or provider +identity takes precedence over replay/context bindings. Raw tool-call IDs remain +supporting evidence rather than authoritative identity because clients may reuse +counter-like or deterministic IDs across independent conversations. + +Agentic tool structure compounds conservatively. An assistant call is paired +one-to-one with a later `tool`/`function` result carrying the same +`tool_call_id`. The closed event contributes one medium evidence group keyed by +a content-free hash of the ID, function name, and canonicalized arguments. Raw, +unpaired, duplicated, or nameless calls remain weak; a previously closed event +cannot promote an unpaired copy in a later request. One closed event plus an +independent medium message group, or two distinct closed events, reaches the +normal probable-confidence boundary. Tool events support ordinary continuity +but never become compaction probes. + +Fallback routing preserves the global session domain but clears provider-native +affinity before selecting a different provider/model. Response recording remains +in the original caller domain. + +Scoped Responses storage uses `(session_domain, response_id)` ownership. Creation +returns a response-specific `X-Proxy-Session-Domain` capability; non-public GET, +DELETE, input-items, and `previous_response_id` continuation requests must return +that opaque header. The token contains the lookup domain plus a random nonce, and +only its hash is stored, so a deterministic classifier/domain hash cannot be used +to forge access. Continuation validates the parent capability and every stored +ancestor's domain. Credential/provider routing containers are omitted from stored +requests and transform traces. + +Legacy public ProviderCache entries are migrated on first exact ID/scope match. +Pre-capability private entries intentionally reject their old deterministic domain +header because accepting it would restore forgeable access; those records expire +under the configured Responses store TTL. + +Only completed responses contribute response-derived identity. Streaming text is +recorded after an explicit provider completion signal, not from clean iterator +EOF alone, because some transports surface truncation as normal EOF. +Streamed tool calls reconstruct function-name and argument fragments by provider +choice index and tool index. Incremental fragments and cumulative snapshots are +both accepted before the same structural event hashing is applied. + +During the current session-tracking validation period, every inference emits a +warning-level `Session tracker decision` line. `action` is one of `new`, +`continue`, `compaction_child`, `compaction_continue`, `compaction_replay`, or +`untracked`. The line includes the selected `session_id`, any +accepted `matched_session_id`, rejected best `candidate_session_id`, and +`parent_session_id`, plus tracking namespace, confidence, score, persistence +origin, and compaction evidence. `origin=persisted` confirms +that the selected session was restored from `session_stickiness.json`. These +diagnostics are intentionally warnings for terminal testing and are expected to +move to debug level after validation. When optional session persistence is enabled, `session_stickiness.json` is -schema-versioned. Upgrades that change the schema intentionally ignore older -session-stickiness files and rebuild in-memory state instead of attempting an +schema-versioned. Schema 2 persists hashed high-water history profiles, scoped +anchor provenance, response-event groups, and compaction replay/context bindings; raw +message content is never stored. Loading rebuilds ownership from validated anchor +records and ignores malformed, expired, orphaned, cross-namespace, or unsupported +state. Upgrades intentionally ignore older schemas instead of attempting an unsafe migration. #### Per-Model Quota Tracking diff --git a/README.md b/README.md index d770a2496..7c07ce016 100644 --- a/README.md +++ b/README.md @@ -500,6 +500,8 @@ The proxy includes a powerful text-based UI for configuration and management. | `SESSION_STICKY_ENTRY_TTL_SECONDS[_]` | TTL for in-memory session-to-credential sticky entries | | `SESSION_STICKY_MAX_ENTRIES[_]` | Max in-memory sequential sticky entries before LRU pruning | | `TRUSTED_SESSION_ID_FIELDS` | Comma-separated explicit request ID fields to trust as strong evidence; unset by default | +| `SESSION_PERSISTENCE_ENABLED` | Persist content-free session lineage across proxy restarts (`false` by default) | +| `SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS` | Session-state write throttle; use `0` for immediate manual testing (default: `5`) | | `OVERRIDE_TEMPERATURE_ZERO` | `remove` or `set` to prevent tool hallucination | | `GEMINI_CLI_QUOTA_REFRESH_INTERVAL` | Quota baseline refresh interval in seconds (default: 300) | @@ -577,21 +579,45 @@ ROTATION_MODE_OPENAI=balanced ### Session Stickiness -Sequential mode can keep related chat requests on the same credential when the session tracker has enough scoped evidence. Tracking is provider/model/scope-bound, so classifier/private credential pools do not share sticky sessions. +Sequential mode can keep related chat requests on the same credential when the +session tracker has enough evidence. The logical session ID follows a conversation +across providers and models inside one caller/credential domain. Public traffic, +each named classifier, and every ad hoc private credential bundle remain strictly +isolated. Credential bindings are still independent per provider/model, so one +global session can use different sticky credentials on different providers. +Raw tool-call IDs remain weak because clients may reuse counters such as +`call_1`. A call becomes medium evidence only when the request contains a +matching tool result; the event identity hashes the ID, function name, and +canonical arguments. One closed event plus an independent message group, or two +distinct closed events, can therefore continue a sparse agentic session. ```env # Wait up to 15s for a session-bound credential if it is only busy by concurrency. SESSION_STICKY_WAIT_SECONDS=15 SESSION_STICKY_WAIT_SECONDS_GEMINI_CLI=15 -# Bound the in-memory sticky session cache. -SESSION_STICKY_ENTRY_TTL_SECONDS=3600 +# Expire each provider/model binding after five idle minutes. Provider-specific +# overrides win; set a provider TTL to 0 to disable its sticky bindings. +SESSION_STICKY_ENTRY_TTL_SECONDS=300 +SESSION_STICKY_ENTRY_TTL_SECONDS_DEEPSEEK=300 SESSION_STICKY_MAX_ENTRIES=10000 # Optional and conservative: only set this if your clients send stable IDs. TRUSTED_SESSION_ID_FIELDS="conversation_id,thread_id" + +# Optional restart persistence, disabled by default. This stores content-free +# logical lineage only; credential bindings remain in memory. +SESSION_PERSISTENCE_ENABLED=true +SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS=5 ``` +Scoped Responses API creation returns a response-specific +`X-Proxy-Session-Domain` access capability. Present that header when retrieving, +deleting, listing input items, or creating a `previous_response_id` continuation +for that non-public response. Only the capability hash is stored; knowing a +classifier label or deterministic domain hash is insufficient. Responses storage +and lineage remain keyed internally by both the domain and response ID. + ### Priority Multipliers Paid credentials can handle more concurrent requests. Legacy priority multipliers apply to hard max concurrency; provider-specific optimal multipliers can also raise the soft target where supported: diff --git a/STRUCTURE.md b/STRUCTURE.md index fb1ba473f..f8057df4f 100644 --- a/STRUCTURE.md +++ b/STRUCTURE.md @@ -33,18 +33,25 @@ **`src/rotator_library/`:** - Purpose: Portable resilience library for multi-provider API key rotation -- Contains: Client facade, provider plugins, usage tracking, Anthropic compatibility, credential management, session tracking -- Key files: `__init__.py`, `rotating_client.py` (in `client/`), `provider_interface.py` (in `providers/`), `usage_manager.py`, `session_tracking.py` +- Contains: Client facade, provider plugins, usage tracking, Anthropic compatibility, credential management, session tracking, transaction logging +- Key files: `__init__.py`, `rotating_client.py` (in `client/`), `provider_interface.py` (in `providers/`), `usage_manager.py`, `session_tracking.py`, `transaction_logger.py`, `error_handler.py` **`src/rotator_library/session_tracking.py`:** - Purpose: Evidence-based session inference with scoped anchors, confidence scoring, compaction probe detection, and deterministic affinity routing -- Contains: `SessionTracker`, `SessionAnchor`, `SessionTrackingHints`, `SessionInference`, `_MatchCandidate` (with `last_seen` tiebreaker) +- Contains: `SessionTracker`, `SessionAnchor`, `SessionTrackingHints`, `SessionInference`, `_MatchCandidate` (with `response_groups`/`request_groups`/`matched_probe_groups` tracking and `last_seen` tiebreaker), `_CompactionDecision` (validated parent lineage with retained-history ratio and `context_probe_groups`) - Key data types: `SessionAnchor` (evidence with strength/source/group), `SessionTrackingHints` (provider-supplied evidence), `SessionInference` (result with session_id, affinity_key, confidence, lineage_parent_session_id, namespace) -- Anchor strength levels: `strong` (tool-call IDs, provider affinity keys), `medium` (message content hashes, response anchors), `weak` (first-user text, untrusted explicit IDs) -- Scope isolation: Namespaces are `scope:{scope_key}:provider:{provider}:model:{model}` to prevent credential pool leakage -- Compaction probes: Separate anchor path (`_build_compaction_probe_anchors()`) identifies lineage parents from large early messages or explicit summarization markers; probe indexes are suppressed from normal continuity anchors; system/developer prompts excluded from continuity evidence -- Persistence: Schema-versioned JSON disk storage via `ResilientStateWriter` with generation-based write deduplication (`_dirty_generation` / `_save_io_lock`) and configurable flush interval -- Configuration: `TRUSTED_SESSION_ID_FIELDS` env var for trusted explicit ID fields; `max_anchor_records`, `max_anchors_per_session`, `persistence_flush_interval_seconds` constructor args +- Anchor strength levels: `strong` (trusted explicit IDs, provider affinity keys, response global IDs), `medium` (message content hashes, response anchors), `weak` (first-user text, raw tool-call IDs, untrusted explicit IDs) +- Scope isolation: Namespaces are `scope:{scope_key}:provider:{provider}:model:{model}` to prevent credential pool leakage; namespaces are immutable per session (`_refresh_and_bridge()` rejects drift, `record_response()` normalizes fallback callbacks to the original namespace), and eviction ranks `compaction_context`/`compaction_replay` anchors above ordinary evidence before falling back to deterministic value tie-breaking +- Compaction probes: Separate anchor path (`_build_compaction_probe_anchors()`) probes only early user/system/developer messages (assistant/tool/function-result history excluded) and requires structural replacement of more than half the parent's high-water request history via `_evaluate_compaction()` (`_retained_history_ratio()`); unmarked summaries must additionally overlap at least two distinct response events (`_MIN_UNMARKED_RESPONSE_GROUPS`) plus a retained request group; authoritative identity (`_is_authoritative_identity_anchor()`: trusted explicit or provider) takes precedence and suppresses unrelated compaction lineage; exact resends reuse the validated child session via opaque `compaction_replay` anchors while changed-tail continuations bind via `compaction_context` anchors (`_find_compaction_context()` / `_compaction_context_anchor()`) minted only from probe groups that matched parent response evidence; probe indexes are suppressed from normal continuity anchors; system/developer prompts excluded from continuity evidence +- Persistence: Schema-versioned JSON disk storage (v3) via `ResilientStateWriter` with generation-based write deduplication (`_dirty_generation` / `_save_io_lock`), dirty state retained on failed writes, stale delayed generations rejected, anchor ownership rebuilt on load (rejecting malformed containers, non-finite timestamps, expired sessions, orphan anchors, namespace mismatches, invalid strengths, and unsupported schemas), and configurable flush interval +- Configuration: `TRUSTED_SESSION_ID_FIELDS` env var for trusted explicit ID fields; `SESSION_PERSISTENCE_ENABLED` / `SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS` env vars on `RotatingClient` for restart persistence; `max_anchor_records`, `max_anchors_per_session`, `persistence_flush_interval_seconds` constructor args + +**`src/rotator_library/transaction_logger.py`:** +- Purpose: Unified transaction logging between the OpenAI-compatible client layer and provider implementations; each API transaction gets a unique directory with client-level and provider-level I/O +- Contains: `TransactionLogger` class; `_make_json_safe` recursive converter for Pydantic/dataclass/`Path`/timestamp objects with circular-reference tracking; helpers `assemble_streaming_response()`, `_strip_framework_keys()` +- Output layout: `logs/transactions/MMDD_HHMMSS_{provider}_{model}_{request_id}/` containing `request.json`, `response.json`, `streaming_chunks.jsonl`, `metadata.json`, and an optional `provider/` subdir (`request_payload.json`, `response_stream.log`, `final_response.json`, `error.log`) +- Integration: Instantiated by `RequestContextBuilder` and `AnthropicHandler` when request logging is enabled; threaded through `RequestContext.transaction_logger` into the executor, streaming handler, field cache engine, and adapter base +- Toggle: Enabled via proxy `--enable-request-logging` flag **`src/rotator_library/client/`:** - Purpose: Client-side request execution with retry and rotation @@ -58,11 +65,11 @@ **`src/rotator_library/providers/`:** - Purpose: Provider-specific implementations and plugin discovery - Contains: One file per provider implementing `ProviderInterface`, shared utilities, retired providers -- Key files: `provider_interface.py`, `__init__.py` (auto-discovery), `gemini_cli_provider.py`, `gemini_provider.py`, `gemini_auth_base.py`, `google_oauth_base.py`, `openai_provider.py`, `openai_compatible_provider.py`, `openrouter_provider.py`, `deepseek_provider.py`, `nvidia_provider.py`, `mistral_provider.py`, `cohere_provider.py`, `groq_provider.py`, `chutes_provider.py`, `firmware_provider.py`, `nanogpt_provider.py`, `provider_cache.py`, `example_provider.py` +- Key files: `provider_interface.py`, `__init__.py` (auto-discovery), `gemini_provider.py`, `openai_provider.py`, `openai_compatible_provider.py`, `openrouter_provider.py`, `deepseek_provider.py`, `nvidia_provider.py`, `mistral_provider.py`, `cohere_provider.py`, `groq_provider.py`, `chutes_provider.py`, `firmware_provider.py`, `nanogpt_provider.py`, `provider_cache.py`, `example_provider.py` **`src/rotator_library/providers/utilities/`:** - Purpose: Shared provider utility modules for quota tracking and credential management -- Key files: `gemini_credential_manager.py`, `gemini_cli_quota_tracker.py`, `gemini_tool_handler.py`, `gemini_shared_utils.py`, `base_quota_tracker.py`, `nanogpt_quota_tracker.py`, `firmware_quota_tracker.py`, `chutes_quota_tracker.py` +- Key files: `base_quota_tracker.py`, `nanogpt_quota_tracker.py`, `firmware_quota_tracker.py`, `chutes_quota_tracker.py` **`src/rotator_library/usage/`:** - Purpose: Usage tracking, limit enforcement, and credential selection @@ -70,9 +77,9 @@ - Key files: `__init__.py`, `manager.py`, `config.py`, `types.py` **`src/rotator_library/usage/config.py`:** -- Purpose: Per-provider usage configuration with session sticky settings -- Contains: `ProviderUsageConfig` with session sticky controls: `session_sticky_wait_seconds`, `session_sticky_entry_ttl_seconds`, `session_sticky_max_entries` -- Configuration: Per-provider `SESSION_STICKY_WAIT_SECONDS_{PROVIDER}` or global `SESSION_STICKY_WAIT_SECONDS` env vars; similarly for `SESSION_STICKY_ENTRY_TTL_SECONDS` and `SESSION_STICKY_MAX_ENTRIES` +- Purpose: Per-provider usage configuration with session sticky settings and quota-exhaustion policies +- Contains: `ProviderUsageConfig` with session sticky controls (`session_sticky_wait_seconds`, `session_sticky_entry_ttl_seconds`, `session_sticky_max_entries`) and no-reset exhaustion controls (`no_reset_exhaustion_policy` ∈ {`warn_only`, `cooldown`, `disable_scope`}, `no_reset_exhaustion_cooldown_seconds`) +- Configuration: Per-provider `SESSION_STICKY_WAIT_SECONDS_{PROVIDER}` or global `SESSION_STICKY_WAIT_SECONDS` env vars; similarly for `SESSION_STICKY_ENTRY_TTL_SECONDS` and `SESSION_STICKY_MAX_ENTRIES`. Per-provider `QUOTA_NO_RESET_EXHAUSTION_POLICY_{PROVIDER}` / global `QUOTA_NO_RESET_EXHAUSTION_POLICY`, and `QUOTA_NO_RESET_COOLDOWN_SECONDS_{PROVIDER}` / global `QUOTA_NO_RESET_COOLDOWN_SECONDS`; provider classes may set `default_no_reset_exhaustion_policy` / `default_no_reset_exhaustion_cooldown_seconds` as baseline **`src/rotator_library/usage/tracking/`:** - Purpose: Usage recording engine and window management @@ -106,6 +113,13 @@ - Purpose: Anthropic Messages API ↔ OpenAI Chat Completions API translation - Key files: `translator.py`, `models.py`, `streaming.py` +**`src/rotator_library/responses/`:** +- Purpose: OpenAI Responses API compatibility — object creation, retrieval, deletion, and streaming with `previous_response_id` continuation +- Contains: `ResponsesService` (orchestrator) bridging through the chat-completions executor; `ResponsesBridge` (Responses ↔ chat translation via `ResponsesProtocol`); `ResponsesStore` protocol with `InMemoryResponsesStore` (default) and `ProviderCacheResponsesStore` (durable JSON cache) backends plus `create_configured_responses_store` factory; `ResponsesSSEFormatter` / `ResponsesWebSocketFormatter` / `ResponsesStreamEvent` / `ResponsesStreamState` streaming helpers; `StoredResponse`, `ResponsesStoreSettings`, `generate_response_id` types +- Key files: `service.py`, `bridge.py`, `store.py`, `streaming.py`, `types.py`, `__init__.py` +- Scope isolation: `StoredResponse` records are keyed by session isolation key (`derive_session_isolation_key`) so `previous_response_id` continuation cannot cross credential pools; stores enforce TTL (`ResponsesStoreSettings.ttl_seconds`), bounded memory (`max_items`), and prune expired/overflow entries +- Storage backends: Memory is the default; provider-cache backend reuses the existing JSON `ProviderCache` (SHA-256 scoped keys) for durable storage without a new database — selection via `config.experimental` runtime settings + **`src/rotator_library/core/`:** - Purpose: Shared types, constants, utilities, and error definitions - Key files: `types.py` (`RequestContext` with session tracking fields: `session_affinity_key`, `session_tracker`, `session_possible_compaction`, `session_lineage_parent_id`, `session_tracking_namespace`), `config.py`, `constants.py`, `errors.py`, `utils.py` @@ -133,14 +147,14 @@ **Entry Points:** `src/proxy_app/main.py`: FastAPI server, TUI launcher, credential tool **Configuration:** `src/rotator_library/config/defaults.py`: All tunable defaults (rotation mode, cooldowns, fair cycle, concurrency) **Core Logic:** `src/rotator_library/client/executor.py`: Unified retry/rotation engine (~1500 lines) -**Session Tracking:** `src/rotator_library/session_tracking.py`: Evidence-based session inference with scoped anchors (~900 lines) +**Session Tracking:** `src/rotator_library/session_tracking.py`: Evidence-based session inference with scoped anchors (~1900 lines) **Provider Interface:** `src/rotator_library/providers/provider_interface.py`: ABC for all providers (~800 lines) **Usage Facade:** `src/rotator_library/usage/manager.py`: Usage tracking + credential selection facade (~2200 lines) **Tests:** `tests/`: Root-level for integration tests; `tests/refactor/` for parity tests ## Naming Conventions -**Files:** `snake_case.py` — provider files follow `{provider_name}_provider.py` pattern (e.g., `gemini_cli_provider.py`) +**Files:** `snake_case.py` — provider files follow `{provider_name}_provider.py` pattern (e.g., `openai_provider.py`) **Directories:** `snake_case` — package directories match their Python module purpose **Providers:** Named by stripping `_provider` suffix from filename; `nvidia_provider.py` remapped to key `nvidia_nim` **Tests:** `test_{feature_name}.py` — co-located in `tests/` directory @@ -154,6 +168,7 @@ **New limit checker:** `src/rotator_library/usage/limits/{name}.py` — extend limit engine **New proxy endpoint:** `src/proxy_app/main.py` — add route handler to the FastAPI app **New Anthropic translation:** `src/rotator_library/anthropic_compat/` — add models or translation logic +**New Responses store backend:** Implement the `ResponsesStore` protocol in `src/rotator_library/responses/store.py` and select it in `create_configured_responses_store()` — keep scope-keyed retrieval and TTL/overflow pruning **New shared type:** `src/rotator_library/core/types.py` — for types used across multiple packages **New config default:** `src/rotator_library/config/defaults.py` — export from `config/__init__.py` **New utility:** `src/rotator_library/utils/` — for cross-cutting utilities (paths, IO, detection) diff --git a/docs/experimental/00-master-plan.md b/docs/experimental/00-master-plan.md new file mode 100644 index 000000000..2c85482a2 --- /dev/null +++ b/docs/experimental/00-master-plan.md @@ -0,0 +1,132 @@ +# Experimental Native Protocol Roadmap + +This branch is for a long-running experimental rewrite that makes native protocol support the first-class extension point of `rotator_library`, while preserving the existing credential rotation, quota, fair-cycle, session tracking, and provider plugin strengths. + +## Operating Rules + +- Work only on the `experimental` branch. +- Keep all repository work inside `C:\Projects\test\LLM-API-Key-Proxy` and child paths. +- Treat commits as checkpoints. A phase may contain many commits. +- Commit messages must include a body describing what changed, why, tests run, and follow-up considerations. +- Do not commit phase reports written for the user unless explicitly requested. Planning docs under `docs/experimental/` are committed. +- Before each phase implementation, first produce a fresh exhaustive phase plan in conversation text, based on the current code state. Only after that plan is settled should it be written to `docs/experimental/phase-N-*.md`. +- After each phase implementation, call both `explore` and `explore-heavy` agents to review the work against the phase plan, external reference areas, and current proxy behavior. Fix findings and re-review as needed. +- Keep LiteLLM as a fallback path for protocols/providers that are not natively covered yet. Native protocol support should be preferred when available. + +## Strategic Goal + +The target architecture is: + +```text +client API request + -> protocol parse into unified representation + -> field-cache injection + -> adapter chain + -> provider override hooks + -> provider-native request build + -> provider execution and credential rotation + -> provider-native response/stream parse + -> field-cache extraction + -> adapter chain + -> protocol formatting for the client + -> transaction logging for every transform state +``` + +Providers should be able to declare an existing protocol and only override the parts that are genuinely provider-specific. A custom provider should usually be configurable through protocol choice, adapters, field-cache rules, auth strategy, and model options rather than requiring a large bespoke provider implementation. + +## Priority Order + +1. Native protocol foundations, unified types, transformers, adapters, and field-cache rules. +2. OpenAI Responses API support, including future WebSocket extension points. +3. Provider work following the protocol layer: Claude Code, Codex, Copilot, Antigravity, and Gemini CLI parity review. +4. Routing and fallback groups, with optional target-group selectors later. +5. Retry, provider/model cooldown, and failover cleanup. +6. Protocol-aware quota, usage, and cost normalization. +7. Streaming library hardening: SSE now, WebSocket-ready later. +8. Config polish using `.env` and optional JSON. No SQLite dependency for now. +9. Extensive staged tests and review-agent verification. + +## Non-Goals For This Branch + +- Do not make the proxy a full multi-user admin product yet. +- Do not require SQLite or Postgres for the main feature set. +- Do not remove LiteLLM before native coverage exists. +- Do not replace the existing `UsageManager`, fair-cycle, custom caps, or evidence-based `SessionTracker`. +- Do not port frontend/UI work from the external reference gateway. + +## Current Strengths To Preserve + +- Credential-level rotation and priority-aware selection. +- Fair cycle and custom caps. +- Windowed quota tracking and quota groups. +- Evidence-based session tracking with compaction handling. +- Provider plugin discovery. +- Gemini CLI provider behavior unless a reviewed change is clearly better. +- Resilient file/JSON state writing. +- Dynamic OpenAI-compatible provider discovery. + +## Reference Gateway Ideas To Import Carefully + +- Unified protocol/transformer style. +- Adapter registry and configurable provider/model adapters. +- Target groups and direct routing syntax, adapted into fallback-first routing. +- Responses API transformer and storage concepts. +- Stream TTFB/stall detection concepts, implemented with Python-native async primitives. +- Provider/model cooldown and retry-history concepts. +- Usage/cost normalization and provider-reported cost extraction. +- Broader provider support patterns for Claude Code, Codex, Copilot, and Antigravity. + +## Phase Index + +1. Protocol Core. +2. Transform Pass Logging. +3. Adapter and Field Cache System. +4. Responses API and WebSocket-Ready Transport Shape. +5. Provider Protocol Overhaul. +6. Routing and Fallback Groups. +7. Retry/Cooldown/Failover Cleanup. +8. Streaming Library Upgrade. +9. Usage, Quota, and Cost Accuracy. +10. Config Polish. + +Each phase may be subdivided if implementation scope becomes too large. + +## Completeness Matrix + +This matrix exists so the branch does not lose any requested scope while phases evolve. The phase plans are still refreshed before implementation, but every item below must remain accounted for. + +| Requested area | Planned coverage | +| --- | --- | +| Protocols are priority #1 | Phases 1 and 4 create native protocol foundations and Responses support before provider work. | +| Protocols are bases, not gospel | Phase 1 requires override-friendly protocol methods, subclassing, copy/mutate registration, and provider-specific overrides. | +| Move away from LiteLLM | Phase 1 adds a `litellm_fallback` protocol path; later providers should prefer native protocols and use LiteLLM only for unsupported coverage. | +| Add protocols automatically like providers | Phase 1 adds protocol auto-discovery and registry behavior modeled after provider discovery. | +| Cover current providers and reference providers | Phase 1 protocols must cover shapes used by current providers; Phase 5 covers Claude Code, Codex, Copilot, Antigravity, and Gemini CLI parity. | +| Responses API is very needed | Phase 4 is dedicated to Responses, `previous_response_id`, storage, SSE, and WebSocket-ready transport shape. | +| WebSocket support later | Phases 1, 4, and 8 require transport separation so WebSocket can be added without rewriting protocol logic. | +| Adapters/transformers tied to protocols | Phases 1, 2, and 3 define protocol parse/build plus transform tracing, adapter registry, and field-cache rules. | +| Cache and return provider fields | Phase 3 implements configurable extraction/injection rules for request, response, and stream fields with scope and mode controls. | +| Reasoning content and similar fields | Phase 3 explicitly covers reasoning content, thinking signatures, prompt cache keys, response IDs, and provider session IDs. | +| Return all possible or last user/assistant use | Phase 3 modes include `last`, `all`, `last_user_turn`, `last_assistant_turn`, and `per_tool_call`. | +| Per-model custom provider behavior | Phases 3, 5, and 10 cover provider/model field cache rules, adapters, model options, and optional JSON config. | +| Transaction logging after every transform | Phase 2 adds ordered request, response, and stream transform trace passes and integrates them with transaction logging. | +| Comments, docstrings, and key decisions | All implementation phases require docstrings for public abstractions and comments for non-obvious transform, protocol, and future-extension decisions. | +| Providers are priority #2 | Phase 5 follows protocol foundations with Claude Code, Codex, Copilot, Antigravity, and Gemini CLI parity review. | +| Antigravity comparison | Phase 5 explicitly compares the reference Antigravity behavior against `src/rotator_library/providers/_retired/`. | +| Routing is interesting | Phase 6 implements fallback chains first, with target-group selectors later if useful. | +| Fallback groups preferred over target groups | Phase 6 starts with ordered fallback groups and only adds target-group-style selectors after that base works. | +| Retry/cooldown/failover cleanup | Phase 7 makes provider/model cooldown real, adds retry history, backoff, retry-after precedence, and success reset. | +| Quota/usage/cost improvements | Phase 9 adds protocol-aware normalizers, provider-reported cost extraction, structured cost fields, and checker abstractions while keeping existing usage engines. | +| Streaming as library capability | Phase 8 hardens streaming below the proxy route layer with TTFB, TTFT, stall detection, cancellation, and transport-aware stream events. | +| Config via env/json, no SQLite | Phase 10 adds optional JSON config with env overrides and validation. SQLite remains out of scope. | +| Multi-user proxy later | The branch keeps multi-user/admin features as a future expansion and only preserves extension points where natural. | +| Exhaustive tests in stages | Every phase requires tests alongside implementation and phase-end review by both `explore` and `explore-heavy`. | +| Reports are for the user, not git | `06-phase-workflow.md` says planning docs are committed, but phase reports are not committed by default. | + +## Code Quality Expectations + +- Public protocol, adapter, transport, field-cache, and provider-extension classes must have docstrings that explain intent, override points, and future expansion hooks. +- Non-obvious transformations must have comments explaining why data is changed, preserved, reordered, or intentionally dropped. +- Lossy protocol conversions must be documented at the conversion site. +- Future WebSocket, target-group, and multi-user extension seams should be noted in comments where they affect today's design. +- Tests should prefer golden fixtures for protocol shapes and focused unit tests for transform edge cases. diff --git a/docs/experimental/01-protocol-architecture.md b/docs/experimental/01-protocol-architecture.md new file mode 100644 index 000000000..bf6349d83 --- /dev/null +++ b/docs/experimental/01-protocol-architecture.md @@ -0,0 +1,132 @@ +# Native Protocol Architecture + +Protocols are reusable bases, not rigid gospel. Providers can subclass, wrap, copy, or override protocol behavior when a provider deviates from an otherwise standard protocol. + +## Why Protocols First + +The current code relies heavily on LiteLLM and provider-specific transforms. That works, but it makes new protocols hard to reason about and makes debugging transformations difficult. The experimental goal is to make a provider mostly declarative: + +```text +provider = protocol + auth + adapters + field cache rules + model options + quota behavior +``` + +If a provider needs custom behavior, it should override a narrow protocol method instead of forcing an entirely bespoke request path. + +## Auto-Discovery + +Protocols should follow the provider plugin style: + +- protocol modules live under `src/rotator_library/protocols/`. +- modules register concrete protocol classes by name. +- a registry exposes names such as `openai_chat`, `anthropic_messages`, `gemini`, `responses`, and `litellm_fallback`. +- third-party or local protocol modules can be added later with minimal registry changes. + +## Core Types + +The unified representation should be explicit enough to cover all existing providers and the external reference protocols without losing important data. + +Suggested types: + +- `UnifiedRequest` +- `UnifiedResponse` +- `UnifiedStreamEvent` +- `UnifiedMessage` +- `ContentBlock` +- `ToolDefinition` +- `ToolCall` +- `ToolResult` +- `ReasoningBlock` +- `Usage` +- `CostDetails` +- `ProtocolMetadata` + +These types should retain unknown provider-specific metadata in explicit extension dictionaries instead of dropping it. Robustness matters more than a narrow perfect schema. + +## Protocol Interface + +The base protocol should provide default methods that can be overridden: + +- `parse_request(raw_request, context) -> UnifiedRequest` +- `build_request(unified_request, context) -> raw_provider_request` +- `parse_response(raw_response, context) -> UnifiedResponse` +- `format_response(unified_response, context) -> raw_client_response` +- `parse_stream_event(raw_event, context) -> UnifiedStreamEvent` +- `format_stream_event(unified_event, context) -> raw_stream_payload` +- `extract_usage(raw_or_unified, context) -> Usage | None` +- `supports_transport(transport_name) -> bool` + +Provider-specific overrides should receive context that includes provider name, model, credential identity, source protocol, target protocol, request ID, and session tracking information. + +## Initial Protocols + +### OpenAI Chat + +Must support: + +- chat completions request/response. +- stream chunks. +- tools and tool calls. +- function-call legacy shapes. +- reasoning fields from OpenAI-compatible providers. +- cached token and reasoning token usage details. + +### Anthropic Messages + +Must support: + +- messages request/response. +- system content extraction. +- text, image, thinking, redacted thinking, tool_use, tool_result blocks. +- stream lifecycle events. +- count_tokens path later if needed. + +### Gemini + +Must support: + +- generateContent and streamGenerateContent shapes. +- content parts. +- functionCall/functionResponse. +- thought signatures. +- safety settings passthrough without unsafe auto-injection. +- Google/Gemini usage metadata. + +### Responses + +Must support: + +- OpenAI Responses request/response. +- `previous_response_id`. +- output items. +- event streams. +- storage-friendly response objects. +- future WebSocket transport. + +### LiteLLM Fallback + +Must preserve existing behavior for providers/protocols not yet native. This path should be explicit and transaction-logged as a fallback, not hidden. + +## Transport Separation + +Protocol formatting must not be tied only to HTTP SSE. Define a transport boundary so the same unified stream events can be emitted through: + +- non-streaming HTTP JSON. +- HTTP SSE. +- future WebSocket. + +The Responses phase should leave clear extension points for WebSocket even if WebSocket is implemented later. + +## Error Handling + +Protocols should preserve provider error bodies where safe, but format client-facing errors consistently. Parsing errors should include transform-pass names and request IDs to make transaction logs useful. + +## Docstrings And Comments + +Protocol code should include docstrings explaining: + +- which external API shape it models. +- what fields are intentionally preserved in metadata. +- where provider overrides are expected. +- future expansion hooks. + +Comments should explain non-obvious transformations, especially lossy conversions between protocols. diff --git a/docs/experimental/02-transform-logging.md b/docs/experimental/02-transform-logging.md new file mode 100644 index 000000000..5cf74b1d1 --- /dev/null +++ b/docs/experimental/02-transform-logging.md @@ -0,0 +1,84 @@ +# Transform Pass Logging + +Debuggability is a core requirement. Every request, response, and stream payload must be inspectable after each transformation pass. + +## Existing Baseline + +The current project already has transaction and raw I/O logging, but it does not consistently show every intermediate transformed state. The experimental protocol layer must improve this without making normal operation too noisy. + +## Transform Trace Model + +Each request should have a trace containing ordered pass records. A pass record should include: + +- request ID. +- pass name. +- direction: request, response, stream_in, stream_out, error. +- protocol/provider/model context. +- timestamp. +- payload snapshot, redacted if needed. +- optional notes or warnings. +- exception information if the pass failed. + +## Required Request Passes + +Suggested pass names: + +- `raw_client_request` +- `parsed_unified_request` +- `after_session_inference` +- `after_field_cache_injection` +- `after_request_adapters` +- `after_provider_override` +- `provider_request` +- `litellm_fallback_request` when fallback is used + +## Required Response Passes + +- `raw_provider_response` +- `parsed_unified_response` +- `after_field_cache_extraction` +- `after_response_adapters` +- `after_client_protocol_format` +- `final_client_response` + +## Required Stream Passes + +For streams, every event can be large. Logging should support configurable sampling/full capture, but the architecture must be able to record: + +- `raw_provider_stream_event` +- `parsed_unified_stream_event` +- `after_stream_field_cache_extraction` +- `after_stream_adapters` +- `formatted_client_stream_event` + +The transaction logger should be able to record stream events as JSONL to avoid retaining large streams in memory. + +## Redaction + +Even though full multi-user security is out of scope, transform logging must avoid accidental credential leakage. + +Redaction hooks should cover: + +- `Authorization` headers. +- `x-api-key` and related API key headers. +- provider API keys. +- OAuth access and refresh tokens. +- cookies. +- obvious `api_key`, `access_token`, `refresh_token`, `client_secret`, and `Authorization` fields in JSON payloads. + +Redaction should happen at the logging boundary, not by mutating live request objects. + +## Failure Debugging + +When a transform fails, the trace should identify: + +- failed pass name. +- protocol class. +- provider/model. +- whether the failure occurred before or after provider execution. +- original error type. +- redacted payload snapshot when possible. + +## Future Expansion + +Later admin/debug endpoints can read these traces. For now, file-based transaction logging is sufficient. diff --git a/docs/experimental/03-field-cache-rules.md b/docs/experimental/03-field-cache-rules.md new file mode 100644 index 000000000..a93e13636 --- /dev/null +++ b/docs/experimental/03-field-cache-rules.md @@ -0,0 +1,113 @@ +# Field Cache Rules + +Field caching is required for providers that need values from previous responses or stream events to be returned on later requests. Examples include reasoning content, thought signatures, prompt cache keys, provider session IDs, and response IDs. + +## Goals + +- Let custom providers configure what fields to extract and where to inject them. +- Avoid hardcoding every provider-specific memory behavior. +- Preserve strict scoping so values never leak across provider, model, credential, classifier scope, or session. +- Support both non-streaming responses and streaming events. +- Support rules per provider and per model. + +## Rule Shape + +Illustrative JSON shape: + +```json +{ + "name": "reasoning_content", + "source": "response", + "path": "choices.*.message.reasoning_content", + "scope": "session", + "mode": "last", + "inject": { + "target": "request", + "path": "messages[-1].reasoning_content" + } +} +``` + +This is a design sketch, not the final schema. + +## Sources + +- `request` +- `response` +- `stream_event` +- `unified_request` +- `unified_response` +- `unified_stream_event` + +## Targets + +- raw provider request path. +- unified request field. +- protocol metadata. +- provider-specific extension field. + +## Scopes + +- `provider` +- `model` +- `credential` +- `session` +- `conversation` +- combinations of the above when needed. + +The default for conversation-affecting fields should be at least provider+model+session scoped. + +## Modes + +- `last`: only the latest matching value. +- `all`: all matching values within the scope. +- `last_user_turn`: latest value associated with the last user turn. +- `last_assistant_turn`: latest value associated with the last assistant turn. +- `per_tool_call`: keyed by tool call ID. + +## Backing Store + +Use the existing provider cache infrastructure first. Do not require SQLite. + +Potential cache keys should include: + +```text +provider / model / credential-or-scope / session-id / rule-name +``` + +Private/classifier scoped credentials must not share cached fields with global credentials. + +## Examples + +### DeepSeek Reasoning Content + +Extract reasoning content from assistant responses and inject it into the next provider request when the provider expects continuity. + +### Gemini Thought Signatures + +Extract thought signatures from Gemini response parts and return them with matching future content parts. + +### Responses Previous Response ID + +Store response IDs and output items so `previous_response_id` can load prior context. + +### Prompt Cache Keys + +Carry `prompt_cache_key` or equivalent provider cache routing values forward when a provider benefits from stable cache routing. + +## Tests + +Required test categories: + +- extraction from response. +- extraction from stream event. +- injection into next request. +- `last` versus `all` behavior. +- scope isolation. +- missing path is a no-op. +- malformed path produces useful validation error. +- redaction in transform logs. + +## Key Decision + +Field cache rules are a protocol/provider extension system, not a replacement for `SessionTracker`. Session tracking decides continuity and credential affinity; field cache rules preserve provider-specific protocol state. diff --git a/docs/experimental/04-provider-roadmap.md b/docs/experimental/04-provider-roadmap.md new file mode 100644 index 000000000..923f84d84 --- /dev/null +++ b/docs/experimental/04-provider-roadmap.md @@ -0,0 +1,115 @@ +# Provider Roadmap + +Providers follow protocols. The protocol layer must land first so provider work can be small, testable, and declarative where possible. + +## Provider Declaration Target + +A provider should eventually be expressible as: + +```text +provider name + + protocol(s) + + auth strategy + + model definitions/options + + adapter chain + + field cache rules + + quota checker/parser + + optional protocol/provider overrides +``` + +Providers can still have custom Python code. The point is to make custom code narrow. + +## Priority Providers + +1. Claude Code. +2. Codex. +3. Copilot. +4. Antigravity. +5. Gemini CLI review and parity improvements. + +## Claude Code + +Review the external reference gateway for: + +- OAuth/token handling. +- request/response protocol shape. +- tool filtering or tool proxy behavior. +- quota checks. +- stream behavior. +- Claude Code-specific headers and model naming. + +Expected implementation direction: + +- use Anthropic Messages or Responses where applicable. +- add provider-specific adapters for tool behavior. +- field cache rules for thinking/signatures if needed. + +## Codex + +Review the external reference gateway for: + +- Responses API route usage. +- Codex-specific user-agent/version behavior. +- OAuth/account handling. +- cooldown parsing for Codex usage limits. +- stream events. + +Expected implementation direction: + +- build on the Responses protocol. +- add Codex provider auth and headers. +- include version/user-agent support if needed. + +## Copilot + +Review the external reference gateway for: + +- GitHub Copilot OAuth flows. +- endpoint selection. +- model naming. +- quota checker behavior. +- provider-specific request filtering. + +Expected implementation direction: + +- use protocol adapters where possible. +- add provider-specific auth/token refresh only where necessary. + +## Antigravity + +Required comparison: + +- current retired implementation under `src/rotator_library/providers/_retired/`. +- external reference Antigravity provider/checker/OAuth behavior. + +Expected implementation direction: + +- restore only what is still valid. +- reuse protocol and field cache rules. +- avoid resurrecting obsolete device-profile behavior unless clearly required. + +## Gemini CLI + +The current Gemini CLI provider is already deep. Review the external reference gateway only for missed behavior: + +- quota checker details. +- thought signature handling. +- stream transform differences. +- OAuth edge cases. +- Gemini 3 tool behavior. +- request headers and endpoint details. + +Do not rewrite Gemini CLI just for architectural purity. + +## Provider Tests + +Each provider should have tests for: + +- config/load/registration. +- auth header or token acquisition. +- request translation. +- response translation. +- stream translation. +- quota parser/checker behavior. +- field cache extraction/injection. +- LiteLLM fallback not used when native path should apply. diff --git a/docs/experimental/05-routing-retry-usage-roadmap.md b/docs/experimental/05-routing-retry-usage-roadmap.md new file mode 100644 index 000000000..210ccf13a --- /dev/null +++ b/docs/experimental/05-routing-retry-usage-roadmap.md @@ -0,0 +1,114 @@ +# Routing, Retry, Cooldown, Usage, And Cost Roadmap + +These systems should be layered around the protocol/provider work without replacing the existing credential engine. + +## Routing Direction + +The preferred first routing model is fallback chains, not a full target-group router from the external reference gateway. + +Example: + +```json +{ + "fallback_groups": [ + ["gemini_cli/gemini-2.5-pro", "openrouter/google/gemini-2.5-pro", "openai/gpt-4.1"] + ] +} +``` + +Behavior: + +- If the requested model is in a group, start with the requested model and then continue through the rest of the group. +- Retryable provider/model failures can move to the next candidate. +- Non-retryable request errors stop unless explicitly configured otherwise. +- Each candidate delegates to the current credential rotation engine. +- Session tracking and classifier scopes must remain isolated. + +## Target Groups Later + +Reference target groups are still useful later for selectors: + +- `in_order` +- `random` +- `usage` +- `cost` +- `latency` +- `performance` + +But the first implementation should solve ordered fallback groups and the missing/stale fallback module before adding selector complexity. + +## Retry/Cooldown Direction + +Current provider cooldown is effectively inert because `CooldownManager.start_cooldown()` is not called. The real active cooldowns are per-credential in `UsageManager`. + +Upgrade direction: + +- keep per-credential cooldowns in `UsageManager`. +- add provider/model cooldown for provider-wide or model-wide failures. +- add consecutive failure tracking. +- add exponential backoff. +- `retry_after` should override computed backoff. +- success should clear provider/model failure count. +- credential quota exhaustion should not globally cool down a provider unless evidence says the provider/model is globally exhausted. + +## Retry History + +Add structured attempt records: + +- candidate provider/model. +- credential stable ID or masked identity. +- protocol path used. +- status: success, failed, skipped. +- error type. +- retryable decision. +- cooldown decision. +- timing. + +These records should be available to transaction logging and optionally client-facing debug output later. + +## Usage And Cost Direction + +Keep existing windowing, fair cycle, custom caps, quota groups, and usage persistence. Add protocol-aware normalization before usage is recorded. + +Needed normalizers: + +- OpenAI Chat. +- OpenAI Responses. +- Anthropic Messages. +- Gemini. +- OAuth/custom providers. + +Fields to preserve: + +- input tokens. +- output tokens. +- reasoning/thinking tokens. +- cache read tokens. +- cache write tokens. +- total tokens. +- provider-reported cost. +- estimated cost. +- cost source and metadata. + +## Cost Direction + +Provider-reported cost should win when available. Sources include: + +- `usage.cost_details`. +- provider-specific response fields. +- SSE comment lines such as `: cost { ... }`. + +Estimated cost should remain as fallback. + +## Tests + +Required tests: + +- fallback chain success after first failure. +- non-retryable error stops fallback. +- streaming fallback only before visible output. +- provider/model cooldown scope. +- backoff escalation. +- success reset. +- usage normalization for each protocol. +- provider-reported cost precedence. diff --git a/docs/experimental/06-phase-workflow.md b/docs/experimental/06-phase-workflow.md new file mode 100644 index 000000000..839fc88e8 --- /dev/null +++ b/docs/experimental/06-phase-workflow.md @@ -0,0 +1,89 @@ +# Phase Workflow + +This document describes how each phase should be executed. + +## 1. Refresh Understanding + +Before writing phase docs or code, inspect the current implementation relevant to the phase. The master plan is a guide, not a substitute for current-code analysis. + +## 2. Produce Phase Plan In Conversation + +First produce the phase plan in conversation text. This forces a fresh exhaustive design pass with the current implementation in mind. + +The plan should include: + +- goals. +- non-goals. +- files/modules to inspect or modify. +- data model changes. +- public API changes. +- transaction logging implications. +- docstrings, comments, and future-extension notes required for the phase. +- tests to add. +- commit checkpoints. +- risk/rollback notes. + +## 3. Write Phase Plan To Docs + +After the conversation plan is accepted or clearly settled, write it to: + +```text +docs/experimental/phase-N-*.md +``` + +Planning docs are committed. + +## 4. Implement In Checkpoint Commits + +Commits are not phase-only. Commit whenever a coherent slice is finished and tested. + +Each commit body should include: + +- what changed. +- why it changed. +- tests run. +- known limitations or follow-ups. + +## 5. Test Continuously + +Run the most relevant tests after each meaningful slice. Do not wait until the phase end. + +## 6. Review Agents + +At phase end, call exactly these two review perspectives: + +- `explore`: code/file-level verification. +- `explore-heavy`: deeper architecture/reference verification. + +Review prompts should compare the implementation against: + +- the phase plan. +- external reference areas where relevant. +- current proxy behavior that must be preserved. +- transaction logging expectations. +- tests. + +If either agent fails or runs out of context, restart it with a narrower prompt. + +## 7. Address Findings + +Fix real findings. If a finding is intentionally deferred, document it in the user-facing phase report, not necessarily in git. + +Repeat review if the changes are substantial. + +## 8. Report To User + +Write a phase report in the conversation. Reports are for the user and are not committed by default. + +The report should include: + +- completed work. +- commits made. +- tests run. +- review-agent findings and resolutions. +- known limitations. +- next phase recommendation. + +## 9. Move To Next Phase + +Do not rely blindly on previous plans. Start the next phase by refreshing context and producing the next phase plan in conversation text. diff --git a/docs/experimental/07-detailed-phase-roadmap.md b/docs/experimental/07-detailed-phase-roadmap.md new file mode 100644 index 000000000..45026167e --- /dev/null +++ b/docs/experimental/07-detailed-phase-roadmap.md @@ -0,0 +1,540 @@ +# Detailed Phase Roadmap + +This document expands the 10-phase roadmap before implementation begins. It exists to prevent later work from narrowing to only Phase 1 details and losing the full feature set. + +Each phase still requires a fresh conversation plan immediately before implementation. This document is the durable baseline; phase-specific plans can adapt after current-code inspection. + +## Phase 1: Protocol Core + +Purpose: create native protocol foundations without changing live execution. + +Primary deliverables: + +- `src/rotator_library/protocols/` package. +- Auto-discovered protocol registry modeled after provider discovery. +- Override-friendly `ProtocolAdapter` base class. +- Unified request, response, message, content, tool, reasoning, usage, cost, stream event, and context dataclasses. +- Base protocols for OpenAI Chat, Anthropic Messages, Gemini, OpenAI Responses, and LiteLLM fallback. +- JSON-safe serialization helpers for transaction tracing and fixtures. +- Protocol errors that identify protocol name, pass name, and payload preview. + +Key requirements: + +- Protocols are bases, not rigid implementations. +- Providers can subclass, wrap, copy, or override protocol methods. +- Unknown provider fields must be preserved in `extra`/metadata instead of dropped. +- Runtime behavior should not change yet. +- LiteLLM fallback remains explicit and named. + +Tests: + +- Registry discovery and alias resolution. +- Base preservation behavior. +- Round-trip and parse/format fixtures for OpenAI Chat, Anthropic Messages, Gemini, and Responses. +- Stream event smoke coverage. + +Review focus: + +- Ensure protocol abstractions are not too narrow for current providers or external reference protocols. +- Ensure types are trace-friendly for Phase 2. + +## Phase 2: Transform Pass Logging + +Purpose: make every transformation state inspectable for debugging. + +Primary deliverables: + +- Transform trace model with ordered pass records. +- Transaction logger integration for request, response, stream, and error transform states. +- Redaction-at-log-boundary helpers. +- JSONL stream transform logs. +- Pass names shared by protocols, adapters, provider overrides, field-cache rules, and fallback execution. + +Required pass coverage: + +- `raw_client_request` +- `parsed_unified_request` +- `after_session_inference` +- `after_field_cache_injection` +- `after_request_adapters` +- `after_provider_override` +- `provider_request` +- `litellm_fallback_request` +- `raw_provider_response` +- `parsed_unified_response` +- `after_field_cache_extraction` +- `after_response_adapters` +- `after_client_protocol_format` +- `final_client_response` +- `raw_provider_stream_event` +- `parsed_unified_stream_event` +- `after_stream_field_cache_extraction` +- `after_stream_adapters` +- `formatted_client_stream_event` + +Key requirements: + +- Log snapshots must not mutate live request/response objects. +- Logs must preserve enough context to debug provider-specific behavior. +- Secret redaction should cover API keys, OAuth tokens, auth headers, cookies, and common secret field names. +- Logging must be usable by future admin/debug endpoints but remain file-based for now. + +Tests: + +- Trace pass ordering. +- Redaction behavior. +- Request/response/stream JSON serialization. +- Transform failure logging with pass name and protocol/provider context. + +Review focus: + +- Verify every planned pass is reachable from the architecture. +- Verify log output is useful without leaking credentials. + +## Phase 3: Adapter And Field Cache System + +Purpose: let custom providers configure what to transform, cache, return, and reinject without hardcoding every provider. + +Primary deliverables: + +- Adapter registry. +- Adapter chain execution for request, response, and stream events. +- Field-cache rule schema and validation. +- JSON-path-like extraction/injection helpers. +- Cache scope builder using provider, model, credential or classifier scope, session, conversation, and rule name. +- Provider/model-level rule configuration hooks. +- Initial built-in adapter rules for reasoning/thinking-related fields. + +Field cache capabilities: + +- Sources: request, response, stream event, unified request, unified response, unified stream event. +- Targets: raw provider request path, unified request field, protocol metadata, provider extension field. +- Scopes: provider, model, credential, session, conversation, and combinations. +- Modes: `last`, `all`, `last_user_turn`, `last_assistant_turn`, `per_tool_call`. +- Missing paths are no-ops unless strict validation is enabled. + +Examples that must be supported by design: + +- DeepSeek-style reasoning content. +- Anthropic thinking and redacted thinking signatures. +- Gemini thought signatures. +- Prompt cache keys. +- Provider session IDs. +- Responses `previous_response_id` metadata. + +Key requirements: + +- Field cache rules complement `SessionTracker`; they do not replace it. +- Rules must never leak across provider/model/credential/session boundaries. +- Providers can define default rules, and model config can override them. +- Transform logging must capture before/after extraction and injection passes. + +Tests: + +- Extract from response and stream events. +- Inject into later requests. +- Scope isolation. +- Mode behavior. +- Malformed path validation. +- Redacted transform logs. + +Review focus: + +- Verify custom provider authoring becomes declarative for common stateful fields. +- Verify no cross-session or cross-credential leaks. + +## Phase 4: Responses API And WebSocket-Ready Transport Shape + +Purpose: add high-priority OpenAI Responses API support while designing stream transports for future WebSocket support. + +Primary deliverables: + +- `/v1/responses` route. +- `GET /v1/responses/{response_id}` route. +- `DELETE /v1/responses/{response_id}` route. +- Optional Codex alias route if supported by provider work. +- JSON/file response storage, not SQLite. +- TTL cleanup for stored responses. +- `previous_response_id` loading and session anchor integration. +- SSE Responses streaming. +- Transport interfaces that can later support WebSocket without rewriting protocol logic. + +Transport shape: + +- Non-streaming HTTP JSON transport. +- HTTP SSE transport. +- Future WebSocket transport extension point. +- Unified stream events from Phase 1 should flow through transports. + +Key requirements: + +- `previous_response_id` must become strong session evidence where safe. +- Stored response objects must be scoped and cleaned up. +- Responses usage and output items must be normalizable in Phase 9. +- Transform logging must show Responses parse/build/storage-relevant states. + +Tests: + +- Create response. +- Retrieve response. +- Delete response. +- Previous response continuation. +- SSE event formatting. +- TTL cleanup. +- Transport interface extension smoke tests. + +Review focus: + +- Verify compatibility with OpenAI Responses expectations and external reference Responses behavior. +- Verify WebSocket is not blocked by the design. + +## Phase 5: Provider Protocol Overhaul + +Purpose: make providers use native protocols where practical and add priority providers after protocol foundations exist. + +Primary deliverables: + +- Provider protocol declaration mechanism. +- Provider hooks for protocol selection, adapter rules, field cache rules, and model options. +- Claude Code provider implementation or integration path. +- Codex provider implementation or integration path. +- Copilot provider implementation or integration path. +- Restored Antigravity provider if current/reference behavior supports it. +- Gemini CLI parity review and targeted fixes only where the external reference gateway has real improvements. + +Provider priorities: + +- Claude Code. +- Codex. +- Copilot. +- Antigravity. +- Gemini CLI review. + +Antigravity comparison requirements: + +- Compare the external reference implementation to `src/rotator_library/providers/_retired/antigravity_provider.py`. +- Restore only valid behavior. +- Avoid obsolete device-profile or fragile logic unless required by the current service. + +Key requirements: + +- Providers can override protocol methods when the base is close but not exact. +- New providers should avoid monolithic transform logic where adapter rules suffice. +- Native path must be testable independently of live credentials. +- LiteLLM fallback should be explicit if used. + +Tests: + +- Provider registration. +- Protocol selection. +- Auth header/token behavior via mocks. +- Request/response/stream translation. +- Quota/checker parsing. +- Field cache rules. +- No accidental LiteLLM fallback when native path should apply. + +Review focus: + +- Verify provider logic follows protocol architecture instead of reintroducing bespoke protocol code everywhere. +- Verify Gemini CLI improvements do not regress existing behavior. + +## Phase 6: Routing And Fallback Groups + +Purpose: add ordered model/provider fallback while keeping credential rotation inside each candidate. + +Primary deliverables: + +- Fallback group config parser. +- Ordered candidate planner for provider/model fallback chains. +- Retryable/non-retryable fallback decisions. +- Streaming fallback rules that only fallback before visible output. +- Optional target-group structure after fallback groups work. +- Optional selectors after ordered fallback works. + +Fallback behavior: + +- If requested model is in a group, try it first. +- Continue through remaining candidates only for retryable failures or configured fallback conditions. +- Preserve classifier/private credential scopes. +- Preserve session namespace isolation. +- Each candidate delegates to current credential selection and usage tracking. + +Future target group selectors: + +- `in_order`. +- `random`. +- `usage`. +- `cost`. +- `latency`. +- `performance`. + +Key requirements: + +- Fix or replace stale `fallback_groups` expectations without breaking current model resolution. +- Do not replace `UsageManager` or `SelectionEngine`. +- Transform logging should show chosen candidate and fallback attempt history. + +Tests: + +- Ordered fallback after retryable failure. +- Non-retryable failure stops. +- Requested model promotion. +- Exhausted chain reports useful error. +- Streaming fallback before output only. +- Scope/session isolation. + +Review focus: + +- Verify fallback integrates above credential rotation, not inside it. +- Verify behavior matches user preference for fallback chains over target-group complexity. + +## Phase 7: Retry/Cooldown/Failover Cleanup + +Purpose: streamline retry behavior and make provider/model cooldown real. + +Primary deliverables: + +- Replace or activate the currently inert provider `CooldownManager` path. +- Provider/model cooldown keys. +- Consecutive failure tracking. +- Exponential backoff. +- Retry-after precedence over computed cooldown. +- Success reset behavior. +- Retry history records for logging and future debug surfaces. +- Integration with fallback groups from Phase 6. + +Cooldown layers: + +- Credential-level cooldown remains in `UsageManager`. +- Provider/model cooldown applies only to evidence of provider-wide or model-wide failure. +- Credential quota exhaustion must not automatically cool an entire provider. +- Model cooldown should not block healthy models on the same provider. +- Provider cooldown should be reserved for provider-wide failure evidence. + +Retry history fields: + +- candidate provider/model. +- credential stable ID or masked identity. +- protocol path used. +- attempt number. +- status: success, failed, skipped, cooled_down. +- error category and raw classifier result. +- retryable decision. +- cooldown decision and duration. +- timing and latency. + +Key requirements: + +- Preserve the current strong retry-after parser, especially Google/Gemini compound duration parsing. +- Preserve streaming safety: no retry after visible output unless explicitly safe. +- Make cooldown state observable by transaction logs and future endpoint surfaces. + +Tests: + +- `start_cooldown` or successor has production callers. +- Provider cooldown blocks provider-wide only. +- Model cooldown blocks one model only. +- Credential cooldown does not become provider cooldown. +- Retry-after overrides exponential backoff. +- Backoff escalates with repeated failures. +- Success clears counts. +- Retry history is recorded. +- Fallback respects cooldown skips. + +Review focus: + +- Verify no healthy credential/model is suppressed too broadly. +- Verify retry/fallback/streaming interactions are deterministic. + +## Phase 8: Streaming Library Upgrade + +Purpose: harden streaming as a reusable library capability, not just a proxy route behavior. + +Primary deliverables: + +- Transport-aware stream event pipeline using unified stream events. +- Explicit upstream iterator close/cancel on client disconnect. +- TTFB timeout before first emitted output. +- TTFT metrics. +- Throughput stall detector. +- Stream usage extraction through protocol normalizers. +- Stream transform logging for raw, unified, adapted, and formatted states. +- SSE transport improvements and WebSocket-ready transport boundaries. + +Streaming behavior: + +- Retry is allowed before client-visible output when the error is retryable. +- Retry is not allowed after visible output unless a future protocol explicitly proves safety. +- Partial streams should not record accepted session response anchors. +- Completed streams should record response anchors and normalized usage. +- Disconnects should close upstream resources promptly. + +Metrics: + +- time to first byte. +- time to first token/content. +- tokens per second where token counts are available. +- chunk count. +- stall duration. +- completion status. + +Key requirements: + +- Implement with Python-native async patterns, not runtime-specific socket internals from the external reference gateway. +- Keep current conservative retry safety. +- Preserve current usage recording behavior and then improve it via Phase 9 normalizers. + +Tests: + +- Disconnect closes upstream async iterator. +- TTFB timeout retries before output. +- Stall detection trips after grace period. +- No retry after visible content. +- Retry before first output. +- Stream transform logging writes expected pass records. +- SSE formatting remains compatible. + +Review focus: + +- Verify library-level design is not tied only to FastAPI route wrappers. +- Verify future WebSocket can reuse the same unified event stream. + +## Phase 9: Usage, Quota, And Cost Accuracy + +Purpose: make usage and cost accounting protocol-aware while preserving the current usage engine. + +Primary deliverables: + +- Protocol-aware usage normalizer interface. +- Usage normalizers for OpenAI Chat, OpenAI Responses, Anthropic Messages, Gemini, OAuth/custom providers, and LiteLLM fallback. +- Structured cost details model integrated with protocol usage. +- Provider-reported cost extraction. +- SSE cost comment parsing where providers emit cost metadata. +- Reasoning/thinking token normalization. +- Cache read/write token normalization. +- Meter/checker abstraction for proactive provider quota checks where useful. + +Usage fields: + +- input tokens. +- output tokens. +- total tokens. +- reasoning/thinking tokens. +- cache read tokens. +- cache write tokens. +- provider-reported cost. +- estimated cost. +- cost source. +- raw provider usage metadata. + +Cost precedence: + +- Provider-reported cost wins when present and trusted. +- Structured provider cost fields are preferred over text parsing. +- SSE `: cost` comments can supply provider-reported cost. +- Estimated cost remains fallback. + +Current systems to preserve: + +- `UsageManager` facade. +- windowed tracking. +- fair cycle. +- custom caps. +- quota groups. +- classifier/private scope separation. +- JSON usage persistence. + +Key requirements: + +- Do not replace the usage engine. +- Normalize before recording when native protocol path is used. +- Existing LiteLLM response usage should continue working. +- Cost details should be transaction-loggable and eventually available to APIs/TUI. + +Tests: + +- OpenAI Chat usage details. +- Responses usage details. +- Anthropic cache read/write tokens. +- Gemini usage metadata including thoughts/cache when available. +- Reasoning token extraction. +- Provider-reported cost precedence. +- SSE cost comment extraction. +- Existing usage aggregation tests still pass. + +Review focus: + +- Verify accounting is more accurate without disrupting selection and quota logic. +- Verify provider-specific raw usage remains available for debugging. + +## Phase 10: Config Polish + +Purpose: support powerful protocol/routing/provider configuration without SQLite. + +Primary deliverables: + +- Optional JSON config file support. +- Env var pointing to JSON config path. +- Env overrides JSON. +- Validation with actionable errors. +- Config sections for protocols, adapters, field-cache rules, fallback groups, providers, model overrides, quota checkers, and stream settings. +- Documentation/examples for custom provider setup using existing protocols. + +Config priorities: + +- `.env` remains enough for simple setups. +- JSON supports complex nested config that is painful in env vars. +- Env overrides allow quick local changes. +- No SQLite/Postgres requirement. + +Potential JSON sections: + +- `protocols`. +- `providers`. +- `models`. +- `adapters`. +- `field_cache`. +- `fallback_groups`. +- `quota_checkers`. +- `streaming`. +- `logging`. + +Key requirements: + +- Validation errors must name the config section and field. +- Bad config should fail early when possible. +- Config should support per-provider and per-model overrides. +- Config should not require rewriting existing `.env` usage immediately. +- Config docs should show custom provider examples using OpenAI Chat, Anthropic Messages, Gemini, Responses, and LiteLLM fallback. + +Tests: + +- Env-only config. +- JSON-only config. +- Env overrides JSON. +- Bad config failure messages. +- Provider/model adapter rule config. +- Fallback group config. +- Field-cache rule config. + +Review focus: + +- Verify custom provider setup becomes practical and documented. +- Verify no accidental database dependency is introduced. + +## Cross-Phase Review Contract + +Every phase ends with two review agents: + +- `explore` for code/file-level verification. +- `explore-heavy` for deeper architecture and reference comparison. + +Review prompts must include the relevant phase plan, external reference areas, current proxy behavior to preserve, tests run, and transaction logging expectations. If either agent fails or runs out of context, restart with narrower scope. + +## Cross-Phase Testing Contract + +Every implementation phase must add or update tests before being considered complete. New test files may need `git add -f` because the repository ignores most `tests/*` by default. + +## Cross-Phase Documentation Contract + +Implementation code must include docstrings and comments for public extension points, non-obvious transformations, lossy conversions, future WebSocket seams, provider override points, and config decisions. diff --git a/docs/experimental/phase-1-protocol-core.md b/docs/experimental/phase-1-protocol-core.md new file mode 100644 index 000000000..aacfea76d --- /dev/null +++ b/docs/experimental/phase-1-protocol-core.md @@ -0,0 +1,172 @@ +# Phase 1: Protocol Core + +## Goal + +Introduce the native protocol foundation without changing live request execution yet. This phase creates robust, documented primitives that later phases can wire into `RequestExecutor`, provider declarations, adapters, field-cache rules, Responses, and streaming transports. + +## Non-Goals + +- Do not replace LiteLLM execution yet. +- Do not add `/v1/responses` routes yet. +- Do not migrate providers to native protocols yet. +- Do not implement field-cache persistence yet. +- Do not rewrite current Anthropic compatibility routes yet. +- Do not change existing request behavior unless a test-only import path requires a harmless export. + +## Current Code Context + +- No `src/rotator_library/protocols/` package exists yet. +- Provider auto-discovery in `src/rotator_library/providers/__init__.py` is the model for protocol discovery. +- `RequestContext` currently holds execution/session/logging fields and can later receive protocol metadata, but Phase 1 should avoid mutating it unless necessary. +- `ProviderTransforms` is hardcoded and will remain active until later adapter migration. +- `TransactionLogger` currently logs initial request, transformed request, response, and stream chunks; Phase 2 will add transform-pass tracing, but Phase 1 types should be trace-friendly. +- `anthropic_compat` already has useful conversion knowledge, especially thinking/tool block handling, but Phase 1 should not delete or replace it. + +## Files To Add + +- `src/rotator_library/protocols/__init__.py` +- `src/rotator_library/protocols/types.py` +- `src/rotator_library/protocols/base.py` +- `src/rotator_library/protocols/registry.py` +- `src/rotator_library/protocols/openai_chat.py` +- `src/rotator_library/protocols/anthropic_messages.py` +- `src/rotator_library/protocols/gemini.py` +- `src/rotator_library/protocols/responses.py` +- `src/rotator_library/protocols/litellm_fallback.py` +- `tests/test_protocol_registry.py` +- `tests/test_protocol_openai_chat.py` +- `tests/test_protocol_anthropic_messages.py` +- `tests/test_protocol_gemini.py` +- `tests/test_protocol_responses.py` + +## Possible Files To Touch + +- `src/rotator_library/__init__.py` only if a public lazy export is needed. +- `src/rotator_library/core/__init__.py` only if shared exports are cleaner there. +- Avoid modifying `RequestExecutor` in Phase 1 unless tests reveal a strict import issue. + +## Data Model Plan + +- `ProtocolRole`: role names should remain strings in payloads, but internal dataclasses can use simple `str` fields to avoid over-constraining custom protocols. +- `ContentBlock`: typed block with `type`, optional text/image/source/tool fields, and `extra` dict for provider-specific data. +- `UnifiedMessage`: `role`, `content`, `name`, `tool_call_id`, `tool_calls`, `extra`. +- `ToolDefinition`: protocol-neutral tool schema with `name`, `description`, `input_schema`, `extra`. +- `ToolCall`: `id`, `name`, `arguments`, `type`, `index`, `extra`. +- `ToolResult`: `tool_call_id`, `content`, `is_error`, `extra`. +- `ReasoningBlock`: `type`, `text`, `signature`, `redacted`, `extra`. +- `Usage`: input/output/total tokens, cache read/write, reasoning tokens, raw usage, cost details. +- `CostDetails`: provider reported cost, estimated cost, currency, source, metadata. +- `UnifiedRequest`: model, messages, tools, system, stream flag, generation params, response format, previous response ID, metadata, raw payload, extra. +- `UnifiedResponse`: id, model, messages/output, stop reason, usage, metadata, raw payload, extra. +- `UnifiedStreamEvent`: event type, delta/message/tool/usage/error metadata, raw event, extra. +- `ProtocolContext`: provider, model, source protocol, target protocol, request ID, session ID, credential stable ID, transport, transaction metadata, provider options. +- `ProtocolResult` is probably unnecessary in Phase 1; keep methods direct and simple. + +## Protocol Interface Plan + +- `ProtocolAdapter` abstract/base class with override-friendly methods. +- Required class attributes: `name`, `aliases`, `supported_transports`. +- Methods: + - `parse_request(raw_request, context) -> UnifiedRequest` + - `build_request(unified_request, context) -> dict` + - `parse_response(raw_response, context) -> UnifiedResponse` + - `format_response(unified_response, context) -> dict` + - `parse_stream_event(raw_event, context) -> UnifiedStreamEvent` + - `format_stream_event(unified_event, context) -> Any` + - `extract_usage(raw_or_unified, context) -> Usage | None` + - `supports_transport(transport_name) -> bool` +- Defaults should preserve unknown data in `extra` rather than raising. +- Errors should be `ProtocolError` with protocol name, pass name, and optional payload preview. + +## Registry Plan + +- `PROTOCOL_PLUGINS: dict[str, type[ProtocolAdapter]]` +- `register_protocol(cls)` for explicit class registration. +- `get_protocol(name)` returns an instance or class consistently. +- `list_protocols()`. +- Auto-discover modules under `src/rotator_library/protocols/`, skipping private modules and infrastructure modules. +- Support aliases so `openai`, `chat_completions`, and `openai_chat` can resolve to the same adapter. +- Prevent duplicate names unless the class is identical or explicit replacement is requested later. +- Keep registry import safe and lightweight. + +## Initial Protocol Behavior + +### OpenAI Chat + +- Parse chat completions request messages, system/developer/user/assistant/tool roles, text content, multimodal content arrays, tools, tool calls, tool_choice, response_format, stream, temperature/top_p/max_tokens/stop. +- Parse responses with choices, assistant messages, tool_calls, reasoning/reasoning_content, and usage details. +- Parse stream chunks into `UnifiedStreamEvent` while preserving raw delta. +- Format back to OpenAI Chat without losing unknown fields in `extra`. + +### Anthropic Messages + +- Parse separate `system`, messages with content blocks, thinking/redacted_thinking, tool_use, tool_result, images, documents where currently supported. +- Map Anthropic usage to `Usage`. +- Preserve thinking signatures in `ReasoningBlock.extra`. +- Build/format enough for round-trip tests, not full replacement of `anthropic_compat` yet. + +### Gemini + +- Parse `contents`, roles, parts, text, inline data, file data, functionCall, functionResponse, thought/thoughtSignature where present. +- Parse `generationConfig`, `safetySettings`, `tools`, stream flag metadata. +- Map usage metadata prompt/candidates/thoughts/cache where possible. +- Preserve provider-specific safety and generation fields. + +### Responses + +- Parse `input`, `instructions`, `previous_response_id`, tools, metadata, stream. +- Parse `output` items, message content, reasoning items, function/tool calls. +- Parse common response stream events into unified stream events. +- Do not add storage or routes yet. + +### LiteLLM Fallback + +- Wrap existing OpenAI-compatible dicts into unified request/response with raw preservation. +- Exist mainly as a named explicit protocol path for later logging. + +## Transaction Logging Implications + +- Phase 1 does not integrate runtime transform logs yet. +- All dataclasses need `to_dict()`/`from_dict()` or safe serialization helpers so Phase 2 can log every pass cleanly. +- `ProtocolError` should include pass names that Phase 2 can reuse. +- Avoid mutating raw payloads in protocol parse methods unless explicitly building a new provider request. + +## Docstrings And Comments Required + +- Every public protocol class explains which external API shape it models and which parts are intentionally partial/base behavior. +- `ProtocolAdapter` docstring explains override contract and why protocols are bases. +- Registry comments explain auto-discovery and skip rules. +- Conversion helpers comment on lossy or approximate mappings. +- Future-extension comments note where WebSocket, field-cache rules, provider overrides, and target transports will attach. + +## Tests + +- Registry auto-discovers built-in protocols. +- Aliases resolve correctly. +- Duplicate registration behavior is deterministic. +- Base protocol default methods preserve raw payloads. +- OpenAI Chat request round-trip with system/developer/user/assistant/tool messages, tool definitions, tool calls, reasoning content, and usage with cache/reasoning token details. +- Anthropic Messages request round-trip with system field, text blocks, thinking/redacted_thinking blocks, tool_use/tool_result, and usage cache fields. +- Gemini request round-trip with contents and parts, functionCall/functionResponse, thoughtSignature, generationConfig/safetySettings, and usageMetadata. +- Responses request/response parse with instructions, input string and input message list, previous_response_id, output messages, reasoning items, and function/tool call items. +- Stream event parse smoke tests for each protocol where practical. +- Serialization tests ensure unified types are JSON-serializable. + +## Commit Checkpoints + +1. Add protocol dataclasses, errors, and serialization helpers. +2. Add base adapter and registry auto-discovery. +3. Add LiteLLM fallback and OpenAI Chat protocol with tests. +4. Add Anthropic Messages protocol with tests. +5. Add Gemini protocol with tests. +6. Add Responses protocol with tests. +7. Run focused protocol test set and fix issues. +8. Phase review with `explore` and `explore-heavy`, then fixes if needed. + +## Risk And Rollback + +- Keep Phase 1 isolated so rollback is just removing the new package/tests. +- Avoid touching executor behavior to prevent regressions. +- If auto-discovery causes import cycles, switch to explicit built-in imports inside registry while preserving the public registry API. +- If tests become too large for one file, split into `tests/protocols/` only if ignore rules are handled with force-add when committing. +- Since `.gitignore` ignores most `tests/*`, remember to force-add new test files when committing. diff --git a/docs/experimental/phase-10-config-polish.md b/docs/experimental/phase-10-config-polish.md new file mode 100644 index 000000000..fd0b884ab --- /dev/null +++ b/docs/experimental/phase-10-config-polish.md @@ -0,0 +1,286 @@ +# Phase 10 Plan: Config Polish + +## Goal + +Make the new protocol/routing/field-cache/streaming/pricing features configurable in a consistent, documented, testable way without replacing the current `.env` workflow. Phase 10 should add a small optional JSON configuration layer, keep environment variables as the final override, document all new knobs in `.env.example`, and expose validation helpers so invalid config fails clearly instead of silently changing routing or accounting behavior. + +## Non-Goals + +- Do not introduce SQLite or any database. +- Do not replace `.env` as the primary user workflow. +- Do not replace provider class declarations, `UsageManager`, `SessionTracker`, `SelectionEngine`, or provider quota trackers. +- Do not implement full multi-user/security config. +- Do not move secrets into JSON config. +- Do not require JSON config for existing deployments. +- Do not rewrite every direct `os.getenv()` call in the repo. +- Do not change default routing, pricing, usage, or streaming behavior when no new config is present. + +## Configuration Precedence + +- Built-in defaults and provider declarations are the base. +- Optional JSON config may provide structured routing/pricing/streaming/provider metadata. +- Environment variables override JSON config. +- Request-level explicit routing/provider fields still win for that request where already supported. +- Secrets remain environment/OAuth-file based; JSON config should not contain API keys or bearer tokens. + +## Current Code Context + +- `.env.example` documents many legacy env vars but not all Phase 6-9 additions. +- `routing/config.py` currently reads fallback groups and model routes from env only. +- Phase 9 added `ModelPricing`, `CostCalculator`, and `ProviderInterface.get_model_pricing()` but not env/JSON pricing. +- Phase 8 added stream metrics primitives but runtime TTFB/stall/heartbeat env knobs are not implemented. +- Phase 7 added provider cooldown env parsing in `retry_policy.provider_cooldown_env()`. +- Phase 3 added field-cache rule dataclasses, but no user-facing JSON config parser for cache rules. +- Provider base URLs and native provider declarations are mostly provider methods or direct env vars. +- Existing usage config env parsing is large and should not be rewritten in Phase 10. +- Reports are user-facing only and should remain uncommitted. + +## Files To Add + +- `src/rotator_library/config/experimental.py` +- `tests/test_experimental_config.py` +- `tests/test_config_pricing.py` +- `tests/test_config_routing_json.py` +- `tests/test_config_stream_settings.py` +- `tests/test_env_example_experimental_config.py` +- Maybe `docs/experimental/config-reference.md` if the implementation needs more detail than `.env.example`. + +## Files Likely To Touch + +- `src/rotator_library/routing/config.py` +- `src/rotator_library/usage/costs.py` +- `src/rotator_library/client/streaming.py` +- `src/rotator_library/client/executor.py` only if streaming env knobs need to be passed through. +- `src/rotator_library/field_cache/rules.py` or equivalent only if JSON parsing needs a small helper. +- `src/rotator_library/providers/provider_interface.py` only if pricing declarations need docstring/typing alignment. +- `.env.example` +- `docs/experimental/phase-10-config-polish.md` + +## JSON Config Model + +Add `ExperimentalConfig` dataclass with optional sections: + +- `routing` +- `pricing` +- `streaming` +- `field_cache` +- `providers` + +Add loader: + +- `load_experimental_config(path=None, env=None)` +- If `path` is `None`, read from `LLM_PROXY_CONFIG_FILE` or `PROXY_CONFIG_FILE`. +- Missing path returns an empty config. +- Invalid JSON raises `ExperimentalConfigError`. +- Unknown top-level sections should be preserved in metadata or warned about, not fatal. + +Add helpers: + +- `as_bool()` +- `as_int()` +- `as_float()` +- `env_key(provider, model, suffix)` sanitizing provider/model names consistently. + +JSON should not read or interpolate secrets. + +## Routing JSON + +Support JSON shape: + +- `routing.fallback_groups..targets = ["codex/gpt-5.1-codex@native", "openai/gpt-5.1@litellm_fallback"]` +- `routing.fallback_groups..failover_on = ["rate_limit", "server_error"]` +- `routing.fallback_groups..stop_on = ["authentication", "validation"]` +- `routing.model_routes. = "group:code_chain"` or target spec. + +`routing/config.py` should merge JSON first, then env: + +- env `FALLBACK_GROUPS`, `FALLBACK_GROUP_*`, `MODEL_ROUTE_*` override or add entries. +- env overrides should retain current behavior for existing deployments. + +Validation: + +- empty groups are invalid. +- `group:` model routes must reference known groups after merge. +- invalid target specs raise `RoutingConfigError`. + +Tests: + +- JSON-only routing works. +- env override replaces JSON group targets. +- env model route can reference JSON group. +- invalid JSON group route fails clearly. +- existing env-only routing tests still pass. + +## Pricing Config + +Add env pricing support in `CostCalculator` or a small helper: + +- `MODEL_PRICE_{PROVIDER}_{MODEL}_INPUT` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_OUTPUT` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_CACHE_READ` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_CACHE_WRITE` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_REASONING` + +Add JSON pricing support: + +- `pricing...input` +- `pricing...output` +- `pricing...cache_read` +- `pricing...cache_write` +- `pricing...reasoning` +- optional `currency` + +Precedence: + +- provider explicit `get_model_pricing()` first, because provider code may know exact native pricing. +- env pricing overrides JSON pricing. +- JSON pricing applies before LiteLLM fallback. +- LiteLLM remains last fallback. + +If env parsing is invalid, log warning and ignore that component. + +Tests: + +- JSON pricing calculates buckets. +- env pricing overrides JSON pricing. +- provider explicit pricing remains highest priority. +- missing pricing remains `unavailable`. +- `skip_cost_calculation` still wins over all pricing sources. + +## Streaming Settings + +Add `StreamRuntimeSettings` dataclass: + +- `ttfb_timeout_seconds` +- `stall_timeout_seconds` +- `heartbeat_seconds` +- `trace_metrics` + +Load from JSON `streaming` and env: + +- `STREAM_TTFB_TIMEOUT_SECONDS` +- `STREAM_STALL_TIMEOUT_SECONDS` +- `STREAM_HEARTBEAT_SECONDS` +- `STREAM_TRACE_METRICS` + +Phase 10 should not enforce timeouts by default. + +Minimal runtime integration: + +- `StreamingHandler.wrap_stream()` should read settings and conditionally emit lifecycle metrics only when trace metrics enabled, default true to preserve Phase 8 behavior. +- Do not implement heartbeat injection unless small and obviously safe. If implemented, default disabled and only emit SSE comments during controlled tests. +- Do not abort streams on TTFB/stall by default. If values are configured, trace `stream_stall_detected` when observable, but do not sever client streams unless explicitly `STREAM_STALL_ABORT=true` is introduced and tested. Prefer no abort in Phase 10. + +Tests: + +- JSON/env settings parse. +- env overrides JSON. +- `STREAM_TRACE_METRICS=false` suppresses lifecycle trace passes while SSE output remains unchanged. +- default trace behavior remains current. + +## Field-Cache Config + +If feasible, add parser for JSON field-cache rule declarations without wiring every provider: + +- `field_cache..[]` +- fields: `name`, `source`, `path`, `scope`, `mode`, `target_path`, `max_entries`. + +Keep it as helper-only if integration is too risky: + +- parse into existing `FieldCacheRule` dataclasses. +- providers/Phase 10+ can call it from `get_field_cache_rules()`. + +Tests: + +- valid JSON rule parses. +- invalid paths/rule modes fail clearly. +- provider/model wildcard merge order documented and tested if implemented. + +If the existing field-cache dataclasses do not support all desired fields cleanly, document and defer live provider integration. + +## Provider Metadata Config + +- Support safe non-secret provider metadata only. +- API base URLs can remain existing `*_API_BASE` env vars for now. +- JSON may define provider protocol/adapter names for future use, but Phase 10 should not make untrusted JSON instantiate arbitrary classes. +- Allow only names that already exist in protocol/adapter registries if validation is implemented. +- Do not place API keys, OAuth tokens, or authorization headers in JSON. +- Tests can validate config rejects/ignores secret-looking keys like `api_key`, `authorization`, `access_token`. + +## Config Validation And Diagnostics + +Add a validation result: + +- warnings for unknown sections/keys. +- errors for invalid routing target specs and invalid numeric pricing. +- no credential values in error messages. + +Add a CLI-free helper test; do not add a new CLI unless tiny. + +Add transform trace or startup log only if a config is loaded: + +- log path, loaded sections, warning count. +- no config contents with secrets. + +## `.env.example` Updates + +Add section for Phase 6 routing: + +- `FALLBACK_GROUPS` +- `FALLBACK_GROUP_` +- `MODEL_ROUTE_` +- execution suffixes `@auto`, `@native`, `@custom`, `@litellm_fallback`. + +Add section for Phase 7 provider cooldown: + +- `PROVIDER_COOLDOWN_MIN_SECONDS` +- `PROVIDER_COOLDOWN_DEFAULT_SECONDS` +- `PROVIDER_COOLDOWN_ON_QUOTA` +- transient retry delay/jitter already in defaults but should be documented. + +Add section for Phase 8 streaming: + +- `STREAM_RETRY_ON_REASONING_ONLY` +- `STREAM_TRACE_METRICS` +- `STREAM_TTFB_TIMEOUT_SECONDS` +- `STREAM_STALL_TIMEOUT_SECONDS` +- `STREAM_HEARTBEAT_SECONDS` + +Add section for Phase 9 pricing: + +- `MODEL_PRICE___INPUT` +- `MODEL_PRICE___OUTPUT` +- cache/reasoning variants. + +Add optional JSON config: + +- `LLM_PROXY_CONFIG_FILE=./config/llm-proxy.json` +- note that env overrides JSON and secrets should stay in env/OAuth files. + +Keep docs concise; do not document every old legacy env in Phase 10 unless already present. + +## Test Modernization + +- Do not take on the entire stale broad test suite unless small. +- Phase 10 can add focused tests for config parsing and maintained regression set. +- Existing stale ignored tests should not block completion unless they are tracked and part of maintained runs. + +## Implementation Checkpoints + +1. Add `ExperimentalConfig` loader and validation helpers with tests. +2. Integrate JSON+env routing config merge with existing routing tests. +3. Add JSON/env model pricing support to `CostCalculator` with tests. +4. Add stream runtime settings parsing and `STREAM_TRACE_METRICS` runtime integration with tests. +5. Add field-cache JSON rule parser helper if feasible with tests. +6. Update `.env.example` and optional config reference doc. +7. Run Phase 10 focused tests and Phase 1-9 maintained regressions. +8. Review with `explore` and `explore-heavy`; fix findings; write uncommitted Phase 10 report. + +## Risks And Mitigations + +- Config precedence confusion. Mitigation: tests for defaults, JSON, env override, provider explicit override, and request-level route behavior. +- Accidentally allowing secrets in JSON. Mitigation: validation rejects secret-looking keys in safe config sections. +- Routing merge could alter existing env behavior. Mitigation: env-only tests remain unchanged and env overrides JSON. +- Pricing could become authoritative incorrectly. Mitigation: keep costs advisory and return zero/unavailable when ambiguous. +- Streaming trace toggles could hide useful debug data unexpectedly. Mitigation: default trace metrics remains enabled. +- Over-scoping into a full config framework. Mitigation: only structured config for new experimental features; leave legacy env parsing intact. diff --git a/docs/experimental/phase-10b-config-surface-wiring.md b/docs/experimental/phase-10b-config-surface-wiring.md new file mode 100644 index 000000000..01aef40b4 --- /dev/null +++ b/docs/experimental/phase-10b-config-surface-wiring.md @@ -0,0 +1,55 @@ +# Phase 10b: Config Surface Wiring And Runtime Settings Completion + +## Goal + +Correct the Phase 10 validation finding that the config layer exists but does not wire enough runtime surfaces introduced in corrective phases. The proxy remains `.env`-first, optional JSON config is structured and secret-free, and environment variables override JSON values. + +## Non-Goals + +- Do not create a full application settings framework. +- Do not move secrets into JSON config. +- Do not replace provider credential discovery. +- Do not add security or multi-user config. +- Do not introduce SQLite or durable DB config. +- Do not commit user-facing reports. + +## Implementation Plan + +1. Add retry/cooldown runtime settings. + - Parse provider cooldown, provider backoff, and failure-history settings from JSON and env. + - Preserve existing env var names and defaults. + - Env overrides JSON. + +2. Wire retry/cooldown settings into retry policy. + - Keep imports lazy to avoid startup cycles. + - Preserve existing monkeypatch/env test behavior. + +3. Add Responses store settings config. + - Parse `RESPONSES_STORE_TTL_SECONDS`, `RESPONSES_STORE_MAX_ITEMS`, `RESPONSES_STORE_FAILED`, and `RESPONSES_STORE_IN_PROGRESS`. + - Support JSON under `responses.store`. + - Preserve default behavior. + +4. Wire Responses store settings into proxy startup. + - Construct `ResponsesService(store_settings=get_responses_store_settings())` in the FastAPI app path. + - Direct tests can still inject explicit settings. + +5. Harden field-cache/provider config parser coverage. + - Cover TTL, metadata mode hints, inject insert behavior, invalid rule errors, and secret rejection in new sections. + +6. Update `.env.example`. + - Document active streaming timeout/heartbeat settings. + - Document provider/model cooldown and backoff knobs. + - Document Responses store policy knobs. + - Document provider-reported cost precedence and SSE cost comments. + - Show safe structured-config section names. + +## Acceptance Criteria + +- Optional JSON config can configure retry/cooldown/backoff and Responses storage policies. +- Env vars override JSON for all new settings. +- Existing env-only behavior remains compatible. +- ResponsesService at proxy startup uses configured store settings. +- Streaming/pricing/routing/field-cache config regressions remain passing. +- Secret-like JSON keys are still rejected in all new sections. +- `.env.example` documents Phase 7b-10b runtime knobs accurately. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-10c-config-completion.md b/docs/experimental/phase-10c-config-completion.md new file mode 100644 index 000000000..0d37d731b --- /dev/null +++ b/docs/experimental/phase-10c-config-completion.md @@ -0,0 +1,72 @@ +# Phase 10c: Config Honesty, Startup Secret Safety, And Validation Completion + +## Goal + +Close the third-pass Phase 10 findings by making accepted configuration surfaces either live-wired and validated or rejected clearly. The proxy remains `.env`-first; JSON config stays optional, structured, and secret-free. + +## Scope + +- Stop printing the full `PROXY_API_KEY` at startup. +- Wire safe `providers` JSON fields into provider runtime behavior. +- Validate provider protocol/adapter names against registries. +- Reject unsupported provider JSON keys instead of silently accepting non-live surfaces. +- Keep credentials, auth headers, API keys, OAuth tokens, and endpoint secrets out of JSON. +- Make invalid pricing env values non-fatal. +- Validate direct model-route target specs at config load/startup. +- Reject generic `credential` / `credentials` key names in JSON config. +- Preserve Phase 8c Responses streaming runtime settings behavior. + +## Non-Goals + +- Do not introduce SQLite, a database, or a new full application settings framework. +- Do not move secrets into JSON config. +- Do not instantiate arbitrary classes from JSON; only reference already registered protocol/adapter names. +- Do not replace provider declarations; JSON is an override layer for safe metadata. +- Do not rewrite every existing `os.getenv()` call. + +## Implementation Plan + +1. Startup secret masking. + - Add a display helper in `proxy_app/main.py`. + - Show that a key is set without printing the raw key. + - Reuse the helper in both startup banners. + +2. Provider config schema. + - Add `ProviderRuntimeConfig` in `config.experimental`. + - Support `protocol_name`, `adapter_names`, `adapter_config`, `native_streaming_supported`, `field_cache`, and `model_quota_groups`. + - Reject unsupported provider keys. + - Validate protocol and adapter names via existing registries with lazy imports. + +3. Runtime provider wiring. + - Use JSON protocol/adapters/adapter config from `ProviderInterface` methods. + - Append configured field-cache rules after provider-declared rules. + - Let JSON opt into native streaming only when protocol support exists. + - Merge JSON model quota groups before env `QUOTA_GROUPS_*` overrides. + +4. Pricing tolerance. + - Ignore malformed env pricing components rather than raising during requests. + - Let JSON pricing or LiteLLM fallback still apply when env pricing is invalid. + +5. Routing validation. + - Validate direct non-`group:` model routes by parsing them during `load_routing_config_from_env()`. + - Keep unknown group validation unchanged. + +6. Secret rejection expansion. + - Add `credential` and `credentials` variants to secret-key detection. + - Add tests for common camelCase and plural forms. + +7. Documentation and tests. + - Update `.env.example` safe-provider-config documentation. + - Add focused tests for startup masking, provider config wiring, pricing tolerance, routing validation, and secret rejection. + - Re-run config/routing/pricing/provider/Responses-streaming regression slices. + +## Acceptance Criteria + +- Startup never prints the full `PROXY_API_KEY`. +- The `providers` JSON section is live-wired for supported safe fields and rejects unsupported keys. +- Configured protocol/adapter names are registry-validated. +- JSON quota groups affect provider grouping, with env overrides final. +- Invalid pricing env values do not fail requests. +- Direct model-route targets fail clearly at config load if malformed. +- `credential` / `credentials` keys are rejected as secret-like. +- Focused tests pass and both reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-11-session-tracking-hardening-report.md b/docs/experimental/phase-11-session-tracking-hardening-report.md new file mode 100644 index 000000000..aaf9ab8ca --- /dev/null +++ b/docs/experimental/phase-11-session-tracking-hardening-report.md @@ -0,0 +1,208 @@ +# Phase 11: Session Tracking And Compaction Hardening Report + +## Result + +Phase 11 is complete. The captured false-positive families are now rejected, real compaction requires structural replacement of most known history, minority/middle-only replacement remains ordinary continuity, and optional persistence has a strict restart-safe schema and substantially broader failure/concurrency coverage. + +The core hardening is committed; later global-domain, binding, Responses, and +closed-tool-event refinements remain in the staged/working-tree change set. + +Every request now emits a temporary warning-level lineage decision. The log identifies new roots, continuations, compacted children, exact compaction replays, untracked requests, and sessions restored from persistence. +Accepted continuations use `matched_session_id`; weak evidence rejected while creating a new root is reported separately as `candidate_session_id`. + +## Incident Findings + +- The July Mistral report was caused by a downstream location-classifier prompt quoting one freshly generated response, not by the identical client regeneration immediately before it. +- The July DeepSeek events had the same one-response classifier shape. +- The June 12 DeepSeek events were ordinary growing histories whose early assistant messages were treated as size-only probes. +- The June 25 events combined ordinary history, a context contraction without summary replacement, exact resends, and follow-ups. They repeatedly produced descendants from one parent. +- The tracker logged 21 compacted-descendant events across these families. +- Persistence was disabled in the captured run and did not cause the incidents. + +## Compaction Decision Contract + +- Only early user/system/developer messages can be compaction probes. +- Assistant, tool, and function-result history cannot become size-only probes. +- A marker is evidence, not a decision; matched parent evidence is mandatory. +- Each session tracks a content-free high-water request-history profile. +- Retaining at least half of the parent high-water history is ordinary continuity. +- Replacing more than half is structural compaction only when parent evidence is sufficient. +- Unmarked compaction requires overlap with at least two distinct completed response events and request-side parent evidence. +- Duplicate content from one response remains one response event. +- One quoted response cannot create compaction lineage or ordinary sticky response bridging from a user/system/developer role. +- Exact replay of a validated compacted payload reuses the existing child. +- Changed or extended tails retaining the validated summary continue the same child through an evidence-bearing context anchor. +- Context anchors are minted only from probe groups that actually matched parent response evidence; shared long system/user harnesses cannot become bindings by position alone. +- Strong trusted or provider identity takes precedence over replay/context bindings and suppresses unrelated compaction lineage. +- Raw tool-call IDs are supporting evidence, not authoritative identity, because deterministic IDs can be reused across sessions. +- Each one-to-one closed assistant-call/tool-result event contributes one medium + evidence group keyed by hashed ID, function name, and canonical arguments. +- Closure is request-local: persisted medium evidence cannot upgrade a later + unpaired call, duplicate results count once, and one result closes one call. +- Generic evidence and logical IDs cross providers/models only inside one strict + public, classifier, or ad hoc credential-bundle domain. Provider-native + evidence remains provider/session-scope qualified. + +## Tracking And Persistence Hardening + +- Successful responses receive deterministic response-event provenance. +- Response provenance survives when assistant content returns in later request history. +- Shared content anchors retain their first live owner instead of being stolen by an auxiliary session. +- Per-session/global trimming and TTL pruning maintain bidirectional session-anchor ownership. +- Late responses cannot resurrect an expired session. +- Weak/ordinary evidence is evicted before strong replay/context identity with deterministic tie-breaking. +- Fallback response callbacks stay in the logical domain while source-provider affinity is cleared before cross-provider selection. +- Persistence schema 3 stores hashed high-water history, scoped anchor metadata, response-event groups, and replay bindings without raw content or external IDs. +- Loading rejects malformed containers, invalid/non-finite timestamps, unsupported schemas, expired sessions, orphan anchors, namespace mismatches, invalid strengths/sources, malformed history signatures, oversized files, and excess sessions/strings. +- Session anchor sets are rebuilt from validated anchor records instead of trusting duplicated serialized ownership. +- Loaded state enforces both per-session and global caps without orphaning either side. +- Failed and exceptional writes retain dirty state for retry. +- Stale delayed generations cannot overwrite a newer persisted snapshot. +- Disk writes are serialized outside the main tracker lock. + +## Streaming Evidence Hardening + +- Streaming response identity is recorded only after an explicit provider completion signal. +- Bare iterator EOF may remain a transport success for provider compatibility but cannot establish response identity. +- `[DONE]`, processed usage-backed final chunks, and usage-backed raw SSE final frames establish completion. +- A raw intermediate `finish_reason` without usage is insufficient because some providers populate it on every chunk. +- Event-prefixed SSE and `data:` frames with or without a space are parsed. +- Non-object JSON, malformed choices/deltas, and non-list tool-call payloads are ignored safely. +- Duplicate streamed tool-call IDs are deduplicated. +- Streamed function/argument fragments are reconstructed by provider choice and + tool index; cumulative snapshots replace prefixes instead of being appended. + +## Test Coverage + +Compaction regressions cover: + +- Ordinary long growing history. +- Identical regeneration and resend shapes. +- One-response classifier prompts. +- Full marked replacement. +- Most-but-not-all marked and unmarked replacement. +- Exactly-half retention. +- Majority retention. +- Middle-only replacement. +- Context contraction without summary replacement. +- Two distinct response events versus duplicate response content. +- Trusted explicit/provider identity and raw tool-ID non-authority. +- Sparse tool loops with one closed event plus independent message evidence or + two distinct closed events, including provider switches and restart. +- Unpaired calls, duplicate results, duplicate IDs, changed function/arguments, + nameless calls, malformed arguments, and cross-domain events. +- Trusted/provider precedence over existing replay/context bindings. +- Raw short and entropy-looking tool-ID collision resistance. +- Cross-scope isolation. +- Exact replay plus marked/unmarked changed-tail child continuation. +- Shared system/user harness rejection for context binding. +- Two-response aggregation rejection without request-side evidence. +- Competing parents with equal scores but different response-event diversity. + +Persistence/state tests cover: + +- Schema-3 metadata round trip and explicit schema-2 rejection. +- Compaction and replay across two restarts. +- Malformed, unsupported, expired, orphaned, and cross-namespace state. +- Per-session/global caps during runtime and load. +- TTL equality and late responses. +- Shared-anchor ownership. +- Failed writes, exceptional writes, throttled writes, forced flushes, and stale generations. +- Blocked disk I/O without holding the state lock. +- Concurrent inference, response recording, and flushing. +- Deterministic affinity across tracker instances and restart. +- Context/replay TTL, supported cap pressure, restart, and scope isolation. + +Stateful simulations cover: + +- Six ordered agentic tool rounds switching DeepSeek/Anthropic/OpenAI, with response recording, tool results, two mid-loop persistence restarts, cross-provider compaction, exact replay after restart, child continuation, and parent continuation. +- Eight long ordinary turns rotating providers/models, with repeated response content, two persistence restarts, per-turn high-water growth, and no false compaction. +- Long roleplay with same/cross-provider exact redo, edited regeneration, a cross-provider middle response rewrite while later turns remain, rollback, resumed branching, and no false lineage. +- First-turn reroll ambiguity, shared harness exclusion, trusted global identity, provider-native collision resistance, independent provider binding clocks, and strict public/classifier/private-bundle isolation. + +Responses/isolation tests cover: + +- Scoped routing reaches RequestContextBuilder without storing or tracing credential containers. +- Every `previous_response_id` parent requires its response-specific access + capability, and every ancestor must belong to the requesting domain. +- Storage ownership is composite by `(domain, response_id)`, including IDs that would collide under path sanitization. +- Non-public retrieval/deletion/input-items require the unforgeable access header + returned at creation; deterministic classifier/domain hashes are rejected. +- Legacy durable public entries migrate only on exact ID/scope match. Old private + deterministic headers are intentionally rejected rather than weakening the new + capability boundary. +- Ad hoc credential bundles use distinct usage managers; named classifier identity remains stable across credential rotation. +- External anchor IDs are hashed and raw scope strings cannot impersonate internally derived domain markers. + +Streaming tests cover: + +- Completed response evidence. +- EOF without completion evidence. +- Early raw finish reason without usage. +- `[DONE]` and usage-backed completion. +- Event-framed/no-space SSE. +- Non-object/malformed SSE and tool-call payloads. +- Incremental/cumulative arguments, interleaved calls, shared tool indexes across + choices, and choices arriving in separate frames. + +## Verification + +- `python -m pytest tests/test_session_tracking.py -q`: 119 passed, 18 subtests passed. +- Final session/request/selection/routing/classifier/executor slice: 173 passed, 18 subtests passed. +- Final Responses bridge/service/store/routes/streaming/accounting slice: 90 passed. +- Consolidated changed-surface slice: 279 passed, 18 subtests passed. +- Python compilation passed for all changed runtime modules. +- `git diff --check` passed for all Phase 11 files. +- Explore scenario review: no blocker, high, or medium findings after final re-review. +- Explore-heavy isolation/persistence review: no blocker, high, or medium findings after the final domain-marker, composite-storage, and zero-TTL fixes. +- Explore-heavy closed-tool-event review: no blocker, high, or medium findings + after request-local closure, one-to-one pairing, and stable streamed choice + indexing were verified. + +A previous unrestricted test run was not a clean repository-wide signal. Initial +collection was blocked by 13 unrelated retired/Gemini import errors. After +excluding only those known collections, that run reached 791 passing tests plus +16 subtests, with 68 failures from existing refactor drift, missing async-test +plugins, legacy constructor assumptions, and provider-global test state. The +known classifier test pollution on the changed surface is now isolated; the +279-test consolidated changed-surface run is clean. The unrestricted suite was +not rerun after this final cleanup. + +## Files + +- `src/rotator_library/session_tracking.py` +- `src/rotator_library/client/rotating_client.py` +- `src/rotator_library/client/request_builder.py` +- `src/rotator_library/client/scopes.py` +- `src/rotator_library/client/streaming.py` +- `src/rotator_library/client/executor.py` +- `src/rotator_library/core/types.py` +- `src/rotator_library/providers/provider_interface.py` +- `src/rotator_library/providers/provider_cache.py` +- `src/rotator_library/routing/attempts.py` +- `src/rotator_library/usage/config.py` +- `src/rotator_library/usage/selection/strategies/sequential.py` +- `src/rotator_library/responses/bridge.py` +- `src/rotator_library/responses/service.py` +- `src/rotator_library/responses/store.py` +- `src/proxy_app/main.py` +- `src/proxy_app/detailed_logger.py` +- `tests/test_session_tracking.py` +- `tests/test_request_builder_routing.py` +- `tests/test_selection_engine.py` +- `tests/test_routing_attempts.py` +- `tests/test_classifier_scoped_routing.py` +- `tests/test_responses_bridge.py` +- `tests/test_responses_service.py` +- `tests/test_responses_store.py` +- `tests/test_responses_routes.py` +- `tests/test_responses_streaming.py` +- `.env.example` +- `README.md` +- `DOCUMENTATION.md` +- `docs/experimental/phase-11-session-tracking-hardening.md` +- `docs/experimental/phase-11-session-tracking-hardening-report.md` + +## Review Disposition + +The first reviews found valid stream-evidence parsing/completion gaps and several missing boundary tests. The stateful-scenario reviews then found changed-tail child forking, context precedence, shared-harness binding, namespace migration, and raw tool-ID collision risks. The closed-tool-event review found persisted-strength promotion, many-to-one result pairing, and streamed choice-index collisions. The final staged-diff review removed duplicate continuation/scope policy and dead context fields, then found deterministic scoped-response access, missing durable key deletion, capability log exposure, and an inaccessible initial scoped stream state. These were reproduced, fixed, and covered by dedicated restart/isolation tests. Two reported anchor-ownership defects were independently rejected after direct inspection: shared anchors are skipped before entering a second session set, and global trimming already removes the owning session-set entry. diff --git a/docs/experimental/phase-11-session-tracking-hardening.md b/docs/experimental/phase-11-session-tracking-hardening.md new file mode 100644 index 000000000..f1320ec88 --- /dev/null +++ b/docs/experimental/phase-11-session-tracking-hardening.md @@ -0,0 +1,246 @@ +# Phase 11: Session Tracking And Compaction Hardening + +## Goal + +Make session inference conservative, deterministic, restart-safe, and resistant to false compaction detection. Compaction is a structural replacement of all or most of a known conversation with summary context. Replacing only a minority of history, including a section in the middle, remains ordinary continuity. + +## Evidence + +- `logs/proxy.log` contains 21 compacted-descendant events across four clusters. +- Normal growing histories were classified as compaction because a long early assistant message matched its response-derived anchor. +- Side-channel classifier requests were classified as compaction because one long user instruction quoted a single immediately preceding response. +- The reported Mistral sequence included a genuine client-level regeneration, but the regeneration itself was not classified. The downstream classifier request after each generation was classified. +- One DeepSeek sequence removed recent turns, but it did not replace the conversation with a summary. The detector had already classified earlier full-history requests and continued creating descendants for resends and follow-ups. +- False compaction changes routing behavior: probe indexes are removed from normal continuity, a new session is created, and sequential selection may choose another credential. +- Persistence was disabled for the captured run and did not cause the incidents, but its loader, ownership, failure, TTL, and concurrency contracts need broader coverage. + +## Behavioral Contract + +1. A compaction candidate must appear in the first two messages and have an eligible role. +2. Size alone never proves compaction. +3. Ordinary assistant, tool, and function-result history is not compaction-probe material. +4. An explicit marker is evidence, not a decision; a known parent and structural history replacement are still required. +5. Unmarked compaction requires evidence from at least two distinct prior response events plus request-side parent evidence. +6. The candidate must replace most of the parent session's high-water request history. +7. Retaining at least half of the high-water history is ordinary continuity. +8. Replacing a minority section in the middle is ordinary continuity even if that section contains summary-like text. +9. Exact replay of a validated compacted request reuses its existing child session; changed tails retaining the validated summary continue that child. +10. A failed API attempt or partial stream never contributes response identity. +11. Caller/credential isolation remains mandatory; provider and model changes do not change the logical session ID. +12. Raw tool IDs remain weak, but a structurally closed assistant-call/tool-result + pair contributes one medium evidence group keyed by the call ID, function, + and canonicalized arguments. Pairing is ordered and one-to-one, closure must + exist in the current request, and repeated copies of one event count once. + +## Global Session Domain Extension + +The logical session ID is global across providers and models, but only inside one +strict caller/credential isolation domain. Public traffic, each classifier, and +each request-private credential bundle are separate domains and must never share +anchors, compaction lineage, replay bindings, or trusted IDs. + +Generic conversation evidence is domain-global. Provider-supplied native anchors +and affinity hints remain provider-qualified so equal opaque IDs from different +upstreams cannot collide. Credential stickiness remains a separate binding keyed +by provider, model/family affinity domain, and global session ID. + +Provider bindings use an idle timeout aligned with upstream cache lifetime. The +default is 300 seconds, providers may declare a different default, and +`SESSION_STICKY_ENTRY_TTL_SECONDS_` remains the user override. Activity +on one provider must not refresh another provider's binding. + +Session persistence remains disabled by default. Enabling it persists logical +lineage only, not credential bindings. Schema 3 records the new isolation-domain +semantics; schema-2 state is rejected instead of being merged across old +provider/model namespaces. + +Provider-native IDs and affinities are hashed and provider/session-scope +qualified. Proxy-owned global IDs use a typed internal channel. Ad hoc private +credential bundles use the same hash-derived boundary for session tracking, +usage managers, cooldowns, and sticky maps. Responses storage is additionally +owned by `(domain, response_id)` and validates that domain through every stored +ancestor. Transport retrieval and continuation require a response-specific +random capability whose hash is stored with the response; deterministic domain +hashes are not access credentials. + +## Scope + +- Role-aware and structure-aware compaction decisions. +- High-water request-history profiles on live session state. +- Distinct response-event provenance for response anchors. +- Preservation of response provenance when responses return in request history. +- Replay binding for validated compacted children. +- Evidence-bearing context binding for post-compaction child continuation. +- Session/anchor ownership and pruning invariants. +- Strict, schema-versioned persistence loading and restart behavior. +- Save throttling, failure recovery, stale generations, and lock handoff. +- Complete and partial streaming response-anchor integration. +- Payload-free decision diagnostics. +- Extensive unit, regression, persistence, integration, and concurrency tests. + +## Non-Goals + +- Do not infer compaction semantically with an LLM. +- Do not inspect or persist raw prompt/response content in tracker state. +- Do not make one quoted response sufficient for unmarked compaction. +- Do not bind repeated standalone prompts that lack trustworthy session evidence. +- Do not persist credential bindings with logical lineage. +- Do not broaden provider-native response IDs, cache keys, or affinity hints across providers. +- Do not migrate unsupported historical persistence schemas. + +## Implementation Plan + +1. Request-history profiles. + - Build deterministic per-message signatures from normalized role, content, and tool identity. + - Store the largest observed request profile as each session's high-water history. + - Update equal-sized profiles to follow normal fixed-window conversations. + - Compare candidate requests against the matched parent's high-water profile using multiset overlap. + - Treat removal of more than half the parent profile as structural replacement. + +2. Compaction candidate selection. + - Restrict size-only probes to early user/system/developer messages. + - Keep explicit marker recognition limited to early system/developer context. + - Validate a parent match before setting `possible_compaction` or suppressing continuity anchors. + - Reject partial/middle replacement when at least half of parent history remains. + +3. Response-event provenance. + - Derive a deterministic group from each normalized successful response event. + - Assign all medium response content anchors from one response to that event group. + - Count distinct response groups rather than matching chunks. + - Preserve response source/group metadata when the same content is later observed as request history in the same session. + - Require two distinct response groups for unmarked size-only compaction. + +4. Compaction replay binding. + - Derive an opaque replay anchor from the validated probe payload. + - Store it only on a confirmed compacted child. + - Keep parent identity in anchor metadata. + - Resolve an exact replay before creating another child. + - Mint context bindings only from probe groups that actually matched parent response evidence. + - Continue changed/extended tails that retain the validated summary without binding shared harness probes. + - Give trusted explicit/provider identity precedence over replay and context bindings. + +5. State integrity. + - Preserve the first live owner of shared content anchors instead of letting an auxiliary request steal them. + - Keep pruning and both trim paths bidirectionally consistent. + - Preserve stronger/provenance-rich metadata when refreshing the same anchor in one session. + - Verify every stored anchor has one live owner and every session anchor points back to that owner. + - Keep raw tool-call IDs non-authoritative because clients may reuse deterministic IDs. + - Promote only closed tool events to medium evidence; an assistant call must + have a matching tool-result reference in the same request. + - Canonicalize JSON arguments so key order and insignificant formatting do + not fragment one event, while different functions/arguments remain distinct. + - Reconstruct streamed calls by provider choice plus tool index, accepting + incremental fragments and cumulative argument snapshots. + - Normalize fallback response callbacks to the session's original namespace instead of migrating state across scopes. + +6. Persistence hardening. + - Bump the persistence schema for high-water history metadata. + - Treat the anchor mapping as authoritative and rebuild session anchor sets during load. + - Reject malformed containers, timestamps, strengths, namespaces, orphan anchors, expired records, and unsupported schemas without failing startup. + - Clamp anchor expiry to its owning session and enforce configured limits after load. + - Preserve dirty state after failed writes and reject stale delayed snapshots. + - Keep disk I/O outside the main state lock and serialize it under the save-I/O lock. + - Expose opt-in proxy environment settings for persistence and write throttling so restart behavior can be tested through the live API. + +7. Diagnostics and documentation. + - Log every inference decision temporarily at warning level with the action, selected/accepted-match/rejected-candidate/parent IDs, namespace, confidence, persistence origin, retained-history ratio, and response-event count. + - Use stable actions for new, continued, compacted, replayed, and untracked requests so terminal runs expose the complete lineage. + - Keep rejected candidate detail at debug level and move the per-request decision line to debug after validation. + - Never log prompt text or raw anchor values. + - Update session-stickiness documentation with the structural decision contract. + +## Test Matrix + +### False-Positive Regressions + +- Full ordinary `[user, assistant, user]` continuation with a recorded response. +- Long user plus long assistant history from the June request shape. +- One-message classifier quoting one complete response. +- One-message classifier sharing only response chunks. +- Identical client regeneration before the previous response is appended. +- Context branch contraction without a summary. +- Exact resend of contracted ordinary history. +- Shared large system/developer harness prompts. +- Marker text with no matched parent. +- Oversized assistant/tool messages. +- Sanitized structural fixtures matching every observed incident family. + +### Replacement Boundary + +- Entire history replaced by a marked summary. +- Most history replaced by a marked summary. +- Most history replaced by an unmarked summary spanning multiple response events. +- Exactly half retained remains ordinary continuity. +- More than half retained remains ordinary continuity. +- One middle turn replaced remains ordinary continuity. +- Several middle turns replaced while a majority remains is ordinary continuity. +- Summary-like text outside the first two messages remains ordinary continuity. +- Parent high-water history is not reduced by a shorter normal request. + +### Positive And Replay Cases + +- Marked system and developer summaries identify a parent. +- One-response unmarked summary is rejected conservatively. +- Two-response unmarked summary identifies a parent. +- Two-response aggregation without request-side evidence remains independent. +- Duplicate copies of one response count as one response event. +- Exact compacted-request replay reuses one child. +- Marked and unmarked child follow-ups with changed short or long tails continue the child across restart. +- Shared long system/user harnesses cannot become compaction context bindings without matching response evidence. +- Strong trusted/provider identity survives a compaction candidate. +- Raw short or entropy-looking tool IDs cannot bind independent sessions. +- Namespace, provider, model, and credential scope stay isolated. + +### Persistence And State + +- Schema round trip preserves history profiles and all anchor metadata. +- Restart preserves normal continuation, compaction decisions, and replay binding. +- Missing, invalid, non-object, unsupported, malformed, expired, orphaned, and namespace-mismatched state is ignored safely. +- TTL equality expires state. +- Late response recording after expiry does not resurrect a session. +- Per-session/global trimming preserves ownership invariants. +- Replay/context bindings survive supported anchor caps and restart, then expire at the TTL boundary. +- Fallback response namespace mismatches preserve original-scope continuity without cross-scope migration. +- Shared anchors keep one live owner and never create stale secondary ownership. +- Affinity is deterministic across fresh trackers and restart. +- Flush throttling, forced flush, writer failure/recovery, stale generation ordering, and mutation during a blocked write are deterministic. +- Concurrent infer/record/flush operations preserve invariants and do not hold the main lock during disk I/O. + +### Integration + +- Request building performs one inference per logical request. +- Internal credential retries reuse the existing request context/session. +- Successful non-streaming responses record response provenance. +- Completed streams record response provenance. +- Clean EOF without an explicit provider completion signal does not record response provenance. +- Failed, partial, and disconnected streams do not record response provenance. +- Corrected session and affinity values reach sequential selection. +- Sparse tool-only loops continue after two distinct closed events, or after one + closed event plus an independent medium message group. +- Unpaired, duplicated, and cross-domain tool IDs remain insufficient. + +### Stateful Conversation Simulations + +- Run a six-round agentic tool loop by inferring each growing request, recording each assistant tool-call response, appending tool results, and verifying one session at every turn. + - Restart from schema-3 persistence multiple times inside the tool loop and verify the first post-restart request continues the persisted session. +- Compact the completed tool loop while dropping all tool IDs, verify a new child and parent lineage, restart, replay the exact compacted request, then append a child response and verify later summary-prefixed requests continue the child. +- Run an eight-turn long-form ordinary conversation with long user/assistant content, duplicate response content, multiple persistence restarts, monotonically growing high-water history, and no compaction decisions. +- Run a long roleplay conversation with exact generation redo, edited regeneration, rollback to an older request snapshot, resumed branching, and an assistant response edited in the middle while later turns remain. +- Assert every roleplay operation stays on one session, keeps the high-water baseline, and does not become compaction. +- Switch providers and models during agentic, ordinary, and roleplay histories while preserving one global session ID. +- Exercise retries, response rerolls, latest and middle response edits, rollback, and compaction across provider boundaries. +- Verify provider-native anchors never cross providers and every classifier/private credential domain remains isolated. +- Advance provider binding clocks independently and prove one provider expires without refreshing or changing another. + +## Acceptance Criteria + +- All 21 captured event shapes are rejected as compaction under sanitized replay tests. +- Full/most-history summary replacement is detected when parent evidence is sufficient. +- Middle-only/minority replacement does not create a compacted child. +- One quoted response never proves unmarked compaction. +- Exact replay of a validated compacted request returns the same child. +- Restart preserves all decisions and state invariants. +- Persistence corruption or write failure never fails request handling. +- Main tracker lock is not held during disk I/O. +- Targeted and full test suites pass. +- Independent explore and explore-heavy reviews report no unresolved blocker, high, or medium findings. diff --git a/docs/experimental/phase-1b-protocol-breadth-operation-model.md b/docs/experimental/phase-1b-protocol-breadth-operation-model.md new file mode 100644 index 000000000..19ea02c03 --- /dev/null +++ b/docs/experimental/phase-1b-protocol-breadth-operation-model.md @@ -0,0 +1,157 @@ +# Phase 1b: Protocol Breadth And Operation Model + +## Goal + +Correct the main Phase 1 audit failure. Phase 1 created a good chat/message protocol foundation, but the starting requirements asked for protocols as the highest priority and for broad protocol coverage. The current protocol layer is still too chat-oriented: it has OpenAI Chat, Anthropic Messages, Gemini, Responses, and LiteLLM fallback, but no explicit operation model for embeddings, image generation, audio transcription, speech/TTS, Ollama-native chat/generate, MCP-style tool gateway calls, or count-token operations. + +## Non-Goals + +- Do not wire every new operation into live FastAPI routes in this corrective slice. This is protocol foundation work. +- Do not replace current embeddings or Anthropic count_tokens routes yet. +- Do not replace existing provider execution, `UsageManager`, `SessionTracker`, routing, or LiteLLM fallback. +- Do not add SQLite or persistent admin/config databases. +- Do not implement full MCP transport/proxy behavior yet; define the protocol shape and operation seam so it is not blocked later. +- Do not fake WebSocket runtime support; keep explicit future transport/capability metadata. +- Do not touch uncommitted phase reports or `docs/issues/`. + +## Current Code Context + +- `src/rotator_library/protocols/types.py` has unified chat/message/request/response/stream dataclasses, but `UnifiedRequest` has no `operation` field and no dedicated multimodal operation carrier fields. +- `ProtocolAdapter` has `supported_transports` but no `supported_operations` or `supports_operation()`. +- Existing concrete protocols are chat/message/generate-content focused. +- `ContentBlock` already has generic `type`, `source`, `raw`, and `extra`, which can support multimodal payloads with careful extension. +- Registry auto-discovery exists and should be reused for new protocol modules. +- Tests exist only for the initial protocol modules. + +## Files To Add + +- `src/rotator_library/protocols/operation.py` +- `src/rotator_library/protocols/openai_embeddings.py` +- `src/rotator_library/protocols/openai_images.py` +- `src/rotator_library/protocols/openai_audio.py` +- `src/rotator_library/protocols/ollama.py` +- `src/rotator_library/protocols/mcp.py` +- `tests/test_protocol_operation_model.py` +- `tests/test_protocol_openai_embeddings.py` +- `tests/test_protocol_openai_images_audio.py` +- `tests/test_protocol_ollama_mcp.py` + +## Files To Touch + +- `src/rotator_library/protocols/types.py` +- `src/rotator_library/protocols/base.py` +- `src/rotator_library/protocols/__init__.py` +- Existing protocol modules where they should declare `supported_operations`. +- Existing protocol tests if serialization expectations need operation fields added. + +## Data Model Additions + +- Add operation constants in a small `operation.py` module, not a rigid enum, so custom/local protocols can introduce operation strings without fighting the core. +- Initial standard operation names: `chat`, `messages`, `responses`, `count_tokens`, `embeddings`, `image_generation`, `image_edit`, `image_variation`, `audio_transcription`, `audio_translation`, `speech`, `ollama_chat`, `ollama_generate`, `mcp`, and `unknown`. +- Add `operation` to `UnifiedRequest`. +- Add flexible operation carrier fields to `UnifiedRequest`: `input`, `modalities`, and `files`. +- Add `operation` to `UnifiedResponse`. +- Add flexible response fields to `UnifiedResponse`: `data` and `content_type`. +- Keep existing fields and defaults chat-compatible so existing adapters do not break. +- Update explicit serialization fields so transform logging can see the new values. + +## Protocol Adapter Additions + +- Add class attribute `supported_operations`. +- Add `supports_operation(operation_name)`. +- Existing protocols should declare their primary operations. +- The base adapter defaults to `operation="unknown"` when no operation is identified. +- LiteLLM fallback stays broad and explicit, not a native implementation claim. + +## New Protocol Modules + +### OpenAI Embeddings + +- Protocol name: `openai_embeddings`. +- Aliases: `embeddings`. +- Operation: `embeddings`. +- Parse/build `/v1/embeddings` style fields: `model`, `input`, `encoding_format`, `dimensions`, and `user`. +- Parse `data[].embedding` and `usage` from responses. +- Preserve unknown fields. + +### OpenAI Images + +- Protocol name: `openai_images`. +- Aliases: `images`, `image_generation`. +- Operations: `image_generation`, `image_edit`, and `image_variation`. +- Parse/build standard OpenAI image fields: `prompt`, `model`, `n`, `size`, `quality`, `style`, `response_format`, `image`, `mask`, and `user`. +- Parse response `data` entries with `url`, `b64_json`, and `revised_prompt`. +- Preserve file/source metadata without reading file contents. + +### OpenAI Audio + +- Protocol name: `openai_audio`. +- Aliases: `audio`, `audio_transcription`, and `speech`. +- Operations: `audio_transcription`, `audio_translation`, and `speech`. +- Parse/build transcription/translation fields: `file`, `model`, `language`, `prompt`, `response_format`, `temperature`, and `timestamp_granularities`. +- Parse/build speech fields: `input`, `model`, `voice`, `response_format`, and `speed`. +- Preserve binary or text provider responses in `raw`; map JSON responses into `data` or `output` when possible. + +### Ollama + +- Protocol name: `ollama`. +- Operations: `ollama_chat`, `ollama_generate`, and compatible `embeddings`. +- Parse/build `/api/chat`, `/api/generate`, and `/api/embeddings` shapes. +- Map `messages`, `prompt`, `options`, `system`, `template`, `context`, `keep_alive`, `format`, and `stream`. +- Parse final and stream chunks while preserving raw context/performance fields. + +### MCP + +- Protocol name: `mcp`. +- Operation: `mcp`. +- Define a JSON-RPC carrier for `initialize`, `tools/list`, `tools/call`, `resources/list`, `resources/read`, `prompts/list`, `prompts/get`, and generic method calls. +- Do not claim a full MCP proxy implementation in this slice. +- Preserve `jsonrpc`, `method`, `params`, `id`, `result`, and `error` fields exactly. + +## Count Tokens + +- Add count-token operation support where it naturally fits, especially Anthropic Messages. +- Avoid duplicating current count-token routes in this slice. +- Ensure `ProtocolAdapter.supports_operation("count_tokens")` can represent a protocol's support. + +## Tests + +- Operation model serialization includes `operation`, `input`, `modalities`, `files`, response `data`, and `content_type`. +- `ProtocolAdapter.supports_operation()` works and remains override-friendly. +- Registry discovers all new protocols. +- Embeddings request/response parse/build round-trip and usage extraction. +- Images request/response parse/build, including edit file/mask references. +- Audio transcription and speech request/response parse/build. +- Ollama chat/generate parse/build and stream parsing. +- MCP request/response/error round-trip. +- Existing protocol regression tests remain clean. + +## Transaction Logging Implications + +- This slice only adds trace-friendly fields and protocol adapters. +- Do not wire live transform logs here. +- Preserve `raw` and `extra` so Phase 2b can log every state. + +## Docstrings And Comments Required + +- Public classes must explain what API shape they model and what remains intentionally future-facing. +- Operation constants must document that strings are intentionally extensible. +- Lossy conversions, especially binary/audio/image payload preservation, need comments explaining that file contents are not inspected and provider-specific metadata stays in `extra`. + +## Commit Checkpoints + +1. Plan doc commit: `docs(experimental): plan protocol breadth correction`. +2. Operation model/types/base adapter commit with tests. +3. Embeddings/images/audio protocol commit with tests. +4. Ollama/MCP/count-token capability commit with tests. +5. Focused protocol regression run. +6. `explore` and `explore-heavy` review against this plan and the source requirements. +7. Fix findings and re-review if substantial. +8. Write uncommitted Phase 1b report. + +## Risks + +- Adding fields to dataclasses can break exact dict assertions. Mitigation: update tests only where serialization now intentionally includes new fields. +- A rigid operation enum could block custom providers. Mitigation: use string constants/helpers, not a closed enum. +- New non-chat adapters could overclaim live support. Mitigation: protocol package supports parse/build only; route/provider wiring comes later. +- MCP can sprawl. Mitigation: implement JSON-RPC carrier shape only, not a full proxy. diff --git a/docs/experimental/phase-1c-protocol-output-operation-guardrails.md b/docs/experimental/phase-1c-protocol-output-operation-guardrails.md new file mode 100644 index 000000000..a9327d3fa --- /dev/null +++ b/docs/experimental/phase-1c-protocol-output-operation-guardrails.md @@ -0,0 +1,59 @@ +# Phase 1c: Protocol Output Correctness And Native Operation Guardrails + +## Goal + +Close the Phase 1/1b third-pass findings that still affect protocol correctness and later native provider behavior. + +## Scope + +- Fix OpenAI Chat formatted usage so it emits public OpenAI-compatible usage fields rather than unified internal usage fields. +- Fix Responses formatted usage so it emits Responses-compatible usage fields rather than unified internal usage fields. +- Promote OpenAI legacy `function_call` into unified `ToolCall` while preserving legacy round-trip shape. +- Add Ollama response formatting that respects mutated unified responses instead of stale raw payloads. +- Enforce protocol operation compatibility in native execution before transport calls. + +## Non-Goals + +- Do not solve every Phase 5 provider issue here beyond the protocol guardrail needed by native execution. +- Do not enable native streaming for priority providers. +- Do not rewrite protocol registry discovery unless a focused test exposes a direct failure. +- Do not touch unrelated dirty files or user-facing reports. + +## Implementation Plan + +1. OpenAI Chat usage formatting. + - Add a helper that converts `Usage` to `prompt_tokens`, `completion_tokens`, `total_tokens`, `prompt_tokens_details`, `completion_tokens_details`, and `cost_details`. + - Ensure `input_tokens`, `output_tokens`, `raw`, and `extra` do not leak into formatted OpenAI usage. + +2. Responses usage formatting. + - Add a helper that converts `Usage` to `input_tokens`, `output_tokens`, `total_tokens`, `input_tokens_details`, `output_tokens_details`, and `cost_details`. + - Preserve cache-write information only as a safe extension inside `input_tokens_details` when present. + +3. Legacy OpenAI `function_call`. + - Parse legacy `function_call` into a unified `ToolCall` when modern `tool_calls` are not present. + - Preserve the raw legacy field in `extra` so formatting emits `function_call` instead of incorrectly upgrading to `tool_calls`. + +4. Ollama response formatting. + - Implement `OllamaProtocol.format_response()` for chat, generate, and embeddings shapes. + - Merge non-core `extra` and usage/timing fields while honoring mutated unified messages/output/data. + +5. Native operation enforcement. + - In `NativeProviderExecutor.execute()` and `stream()`, reject unsupported protocol operations before parse/build/transport. + - Keep errors sanitized and explicit so bad provider/protocol pairings fail early. + +## Tests + +- `tests/test_protocol_openai_chat.py` +- `tests/test_protocol_responses.py` +- `tests/test_protocol_ollama_mcp.py` +- `tests/test_native_provider_executor.py` or `tests/test_request_executor_native_routing.py` +- Full protocol suite plus native routing smoke tests before review. + +## Acceptance Criteria + +- OpenAI Chat formatted responses expose OpenAI-compatible usage fields. +- Responses formatted responses expose Responses-compatible usage fields. +- Legacy OpenAI `function_call` is represented in unified tool calls and round-trips as legacy `function_call`. +- Ollama formatted responses reflect mutated unified state. +- Native execution rejects unsupported protocol operations before network transport. +- Focused tests pass and both `explore` and `explore-heavy` reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-2-transform-logging.md b/docs/experimental/phase-2-transform-logging.md new file mode 100644 index 000000000..77cfbaeec --- /dev/null +++ b/docs/experimental/phase-2-transform-logging.md @@ -0,0 +1,203 @@ +# Phase 2 Plan: Transform-Pass Logging + +## Goal + +Make transaction logging capable of recording every meaningful request, response, and stream state after each transformation pass, without yet replacing runtime execution with the native protocol adapters. This phase creates the durable trace format and wires it into existing request/executor/provider logging points so later protocol, adapter, and field-cache phases can add more trace entries without redesigning observability. + +## Non-Goals + +- Do not route execution through native protocols yet. +- Do not implement field-cache extraction or injection yet. +- Do not add Responses routes yet. +- Do not add full multi-user/security work. +- Do not rewrite provider-specific logging systems. +- Do not change request/response behavior. + +## Current Code Context + +- `TransactionLogger` creates one directory per transaction and writes legacy files such as `request.json`, `request_transformed.json`, `response.json`, `streaming_chunks.jsonl`, and `metadata.json`. +- `RequestContextBuilder` creates `TransactionLogger` and calls `log_request(kwargs)`. +- `RequestExecutor` calls `log_transformed_request(kwargs, context.kwargs)` after request preparation. +- `RequestExecutor` logs non-streaming responses through `log_response(response_data)`. +- `RequestExecutor._transaction_logging_stream_wrapper()` logs parsed chunks and assembles/logs the final streamed response. +- `ProviderLogger` writes provider-level request payloads, raw response chunks, final responses, and provider errors. +- Phase 1 unified types already provide JSON-safe serialization through `serialize_value()`. + +## Files To Add + +- `src/rotator_library/transform_trace.py` +- `tests/test_transform_trace.py` +- `tests/test_transaction_logger_transform_trace.py` + +## Files Likely To Touch + +- `src/rotator_library/transaction_logger.py` +- `src/rotator_library/client/request_builder.py` +- `src/rotator_library/client/executor.py` + +## Trace Model + +`TransformTraceEntry` fields: + +- `sequence`: monotonically increasing integer per writer instance. +- `timestamp_utc`. +- `component`: `client` or `provider`. +- `pass_name`: stable machine-readable name. +- `direction`: `request`, `response`, `stream`, `error`, or `metadata`. +- `stage`: `client`, `protocol`, `adapter`, `provider`, or `final`. +- `protocol`: optional protocol name. +- `provider`: optional provider name. +- `model`: optional model name. +- `credential_id`: optional stable or masked credential identifier only. +- `transport`: `http`, `sse`, `websocket`, or a future string. +- `changed_from_previous`: optional bool. +- `metadata`: JSON-safe dict. +- `data`: JSON-safe sanitized payload. + +`TransformTraceWriter` responsibilities: + +- Maintain a local sequence counter. +- Write compact ordered entries to `transform_trace.jsonl`. +- Optionally write individual request/response snapshots under `transforms/`. +- Avoid per-chunk snapshot files for stream chunks. +- Never raise logging failures into request execution. + +## Log Files + +Existing files stay for compatibility: + +- `request.json` +- `request_transformed.json` +- `response.json` +- `streaming_chunks.jsonl` +- `metadata.json` + +New files: + +- `transform_trace.jsonl` +- `transforms/0001_raw_client_request.json` +- `transforms/0002_prepared_provider_request.json` +- `transforms/NNNN_.json` for non-stream snapshots + +Stream chunks use JSONL entries rather than one file per chunk. + +## Required Pass Names + +Request passes: + +- `raw_client_request` +- `prepared_provider_request` +- `provider_request_payload` + +Response passes: + +- `raw_client_response` +- `final_client_response` +- `provider_final_response` + +Stream passes: + +- `raw_stream_chunk` +- `parsed_stream_chunk` +- `assembled_stream_response` +- `provider_raw_stream_chunk` + +Error passes: + +- `provider_error` +- `transform_log_error` + +## Sanitization + +This is pragmatic trace hygiene, not full security work. + +Redact recursively by key name: + +- `api_key` +- `credential_identifier` +- `authorization` +- `x-api-key` +- `x-goog-api-key` +- `access_token` +- `refresh_token` +- `client_secret` +- `password` +- `secret` +- `token` + +Rules: + +- Redacted value is `"[REDACTED]"`. +- Redact by key only, not by text value. +- Do not hide ordinary model text containing words such as token. +- Always serialize through Phase 1 JSON-safe helpers. + +## Integration Behavior + +- Trace logging is additive and backward compatible. +- Logging failures never fail requests. +- `log_request()` records `raw_client_request`. +- `log_transformed_request()` records `prepared_provider_request` even when legacy transformed file is skipped because payloads compare equal. +- `log_response()` records `final_client_response`. +- A dedicated method can record `raw_client_response` before final normalization when available. +- `log_stream_chunk()` records `parsed_stream_chunk`. +- `_transaction_logging_stream_wrapper()` records `raw_stream_chunk` before JSON parsing and `assembled_stream_response` before final `log_response()`. +- `ProviderLogger.log_request()` records `provider_request_payload`. +- `ProviderLogger.log_response_chunk()` records `provider_raw_stream_chunk`. +- `ProviderLogger.log_final_response()` records `provider_final_response`. +- `ProviderLogger.log_error()` records `provider_error`. + +## Context Design + +- `TransactionLogger` owns a client-side `TransformTraceWriter`. +- `TransactionContext` carries enough trace metadata for providers to create provider-side trace entries without sharing mutable logger objects. +- `ProviderLogger` creates its own writer with `component="provider"`. +- Client and provider sequence numbers are local to their writer instances. Entries include component and timestamps, so Phase 2 does not promise one global order across independent writers. + +## Comments And Docstrings + +- Public trace classes must explain they are observability-only and must not affect request behavior. +- Redaction helpers must explain why redaction is key-based rather than value-based. +- TransactionLogger comments should distinguish legacy files from the new trace ledger. +- Stream comments should explain why stream chunks use JSONL rather than per-chunk snapshots. + +## Tests + +`tests/test_transform_trace.py`: + +- JSON-safe serialization handles dataclasses and non-primitives. +- Redaction recurses through nested payloads. +- Redaction does not redact normal text values. +- Sequence increments per writer. +- Snapshot filenames are stable and sanitized. + +`tests/test_transaction_logger_transform_trace.py`: + +- `log_request()` writes legacy `request.json` and trace entry `raw_client_request`. +- `log_transformed_request()` writes trace entry even when legacy transformed file is skipped because payloads are equal. +- `log_response()` writes `final_client_response`. +- `log_stream_chunk()` writes `parsed_stream_chunk`. +- ProviderLogger writes provider request, chunk, final, and error trace entries. +- Logging disabled writes nothing. + +Regression tests: + +- Phase 1 protocol tests. +- `tests/test_session_tracking.py`. +- `tests/test_selection_engine.py`. + +## Commit Checkpoints + +1. Add transform trace dataclasses, writer, redaction, and tests. +2. Integrate client-side `TransactionLogger` trace entries and tests. +3. Integrate provider-side `ProviderLogger` trace entries and tests. +4. Run focused and regression tests. +5. Review with `explore` and `explore-heavy`, fix findings, and write the uncommitted Phase 2 report. + +## Risks And Mitigations + +- Stream logs can grow quickly. Use JSONL only for stream chunks. +- Redaction can hide useful fields if too broad. Redact by sensitive key names only. +- Provider/client entry ordering is not globally sequenced in Phase 2. Include component and timestamps. +- Provider SDK objects may not be JSON-native. Always use `serialize_value()`. +- Tests may import an installed package if local `src` is not first on `sys.path`. Keep `tests/conftest.py`. diff --git a/docs/experimental/phase-2b-transform-trace-coverage.md b/docs/experimental/phase-2b-transform-trace-coverage.md new file mode 100644 index 000000000..034874939 --- /dev/null +++ b/docs/experimental/phase-2b-transform-trace-coverage.md @@ -0,0 +1,184 @@ +# Phase 2b: Complete Live Transform-Pass Trace Coverage + +## Goal + +Correct the Phase 2 audit gap. Phase 2 created the trace writer and added important request, response, stream, provider, and error trace entries, but the validation pass found that it still does not capture every meaningful real transformation boundary. Phase 2b makes transform tracing systematic across the live LiteLLM path, provider transforms, native protocol execution, Responses service/streaming, adapter chains, field-cache passes, and stream wrappers without changing request behavior. + +## Non-Goals + +- Do not replace LiteLLM execution in this phase. +- Do not make non-chat protocols live routes here. +- Do not implement field-cache persistence or new field-cache semantics here; only improve logging of existing cache operations. +- Do not implement WebSocket runtime support here. +- Do not add security/multi-user features beyond existing redaction behavior. +- Do not commit user-facing reports. +- Do not touch unrelated dirty `ARCHITECTURE.md`, `STRUCTURE.md`, `.opencode/`, `docs/issues/`, or old phase reports unless explicitly needed. + +## Current Code State + +- `TransformTraceWriter` exists with JSONL and snapshot files. +- `TransactionLogger.log_transform_pass()` and `log_transform_error()` exist and are additive. +- Existing legacy trace entries include `raw_client_request`, `prepared_provider_request`, `final_client_response`, stream chunk entries, provider request/chunk/final/error entries, routing metadata, retry/cooldown metadata, and usage summaries. +- `ProviderTransforms.apply()` mutates request kwargs through built-in transforms, provider hook transforms, model options, and LiteLLM conversion, but it does not emit per-step trace entries. +- `RequestExecutor._prepare_request_kwargs()` calls `ProviderTransforms.apply()`, then `log_transformed_request()` logs only the final prepared provider kwargs. +- Non-streaming LiteLLM/provider responses are logged only after success as `final_client_response`; there is no explicit `raw_provider_response` or `post_usage_normalization_response` distinction. +- Streaming logs `raw_stream_chunk`, `parsed_stream_chunk`, and `assembled_stream_response`, but the live stream handler has additional normalization/usage/error decision work that is not fully trace-covered. +- `NativeProviderExecutor` traces some native passes, but it skips key states: raw native input, parsed unified request, built provider request before adapters, after response adapters, field-cache extraction details, and final formatting boundaries in a complete request/response sequence. +- `run_adapter_chain()` traces `before_adapter_chain` and `after_adapter`, but not a final `after_adapter_chain` summary and not enough metadata to distinguish built-in vs provider-declared adapter chains in live reviews. +- `FieldCacheEngine` logs per-rule operations, but injection/extraction needs start/end summary trace entries so debugging can see when a cache pass happened but no rule matched. +- Responses service currently logs a service create pass, but stream output and bridge conversions are not fully represented as protocol/bridge/final boundaries. + +## Trace Taxonomy + +### Client Request + +- `raw_client_request` +- `pre_provider_transform_request` +- `after_builtin_provider_transform` +- `after_provider_hook_transform` +- `after_provider_model_options` +- `before_litellm_conversion` +- `after_litellm_conversion` +- `prepared_provider_request` + +### LiteLLM / Provider Execution + +- `provider_execution_request` +- `raw_provider_response` +- `post_usage_normalization_response` +- `final_client_response` + +### Native Protocol Request + +- `raw_native_client_request` +- `native_protocol_selected` +- `parsed_native_unified_request` +- `built_native_provider_request` +- `before_adapter_chain` +- `after_adapter` +- `after_adapter_chain` +- `field_cache_injection_start` +- `field_cache_rule_*` +- `field_cache_injection_complete` +- `native_provider_request` + +### Native Protocol Response + +- `raw_native_provider_response` +- `parsed_native_unified_response` +- `formatted_native_response` +- `after_response_adapter_chain` +- `field_cache_extraction_start` +- `field_cache_rule_*` +- `field_cache_extraction_complete` +- `usage_accounting_summary` +- `final_client_response` + +### Native Stream + +- `native_provider_stream_request` +- `raw_native_provider_stream_chunk` +- `parsed_native_stream_event` +- `after_stream_event_adapter_chain` +- `after_field_cache_stream_extraction` +- `formatted_client_stream_event` + +### Responses API + +- `responses_raw_request` +- `responses_parsed_request` +- `responses_bridge_chat_request` +- `responses_bridge_chat_response` +- `responses_stored_response` +- `responses_final_response` +- stream equivalents for created, delta, completed, bridge chunk, and final states where available. + +### Errors + +- Continue using `transform_log_error`. +- Include failed pass name, component, provider, protocol, transport, and sanitized payload. + +## Implementation Plan + +1. Add small local trace helpers where they avoid repeated fragile `if transaction_logger` blocks. + - Avoid global state. + - Avoid dependency cycles. + - Prefer local helpers when a shared helper would create import loops. + +2. Expand `ProviderTransforms.apply()` tracing. + - Add optional keyword-only `transaction_logger`, `credential_id`, `transport`, and `trace_metadata` parameters. + - Trace `pre_provider_transform_request` before any mutation. + - Trace `after_builtin_provider_transform` after each built-in transform that returns a modification string. + - Trace `after_provider_hook_transform` when provider hooks modify or report modifications. + - Trace provider hook errors as `transform_log_error` without changing failure behavior. + - Trace `after_provider_model_options` when options are applied. + - Trace `before_litellm_conversion` and `after_litellm_conversion` around `convert_for_litellm()`. + +3. Pass trace context from `RequestExecutor._prepare_request_kwargs()`. + - Provide transaction logger, provider/model/session/scope/classifier metadata, transport, and stable credential ID where available. + - Add an optional `credential_id` argument rather than changing existing call semantics. + +4. Add raw provider and normalization response trace in non-streaming executor. + - Trace `provider_execution_request` immediately before the provider call after pre-request callback. + - Trace `after_pre_request_callback` if callback changed kwargs. + - Trace `raw_provider_response` immediately after provider call returns. + - Trace `post_usage_normalization_response` after `_normalize_response_usage()` and before returning. + +5. Expand streaming trace coverage. + - Ensure streaming request path gets the same provider-transform traces as non-streaming. + - Keep `raw_stream_chunk`, `parsed_stream_chunk`, and `assembled_stream_response`. + - Add `stream_error_event` for error SSE payloads produced by executor fallback/error handling. + - Add `stream_done_event` for `[DONE]` boundaries with snapshots disabled. + - Do not change emitted SSE bytes. + +6. Expand native provider trace coverage. + - Add `raw_native_client_request` before protocol parse. + - Trace `parsed_native_unified_request` after `protocol.parse_request()`. + - Trace `built_native_provider_request` after `protocol.build_request()` and before adapters. + - Trace `after_request_adapter_chain` after request adapters. + - Trace `parsed_native_unified_response`, `formatted_native_response`, and `after_response_adapter_chain` on response path. + - Add equivalent stream-event boundaries where applicable. + +7. Expand adapter chain summary traces. + - Keep `before_adapter_chain` and per-adapter `after_adapter`. + - Add `after_adapter_chain` with adapter list, stage, and count even when there are no adapters. + - Include `changed_from_previous` where feasible. + - Disable snapshots for stream-event summaries. + +8. Expand field-cache summary traces. + - Log `field_cache_injection_start` / `field_cache_extraction_start` before the pass. + - Log `field_cache_injection_complete` / `field_cache_extraction_complete` after completion. + - Include rule count, matched count, changed count, source, and operation type. + - Keep per-rule trace entries intact. + +9. Expand Responses API and bridge tracing. + - Trace raw request, parsed unified request, bridge chat request, bridge response, stored response, and final response where those boundaries exist. + - Trace emitted streaming response events with stable pass names and snapshots disabled for per-chunk events. + - Preserve route bodies and SSE event ordering. + +10. Tests. + - Extend `tests/test_transaction_logger_transform_trace.py` for new pass names and redaction stability. + - Add or extend tests for `ProviderTransforms.apply()` tracing each stage without changing payload behavior. + - Add native provider executor trace order tests for non-streaming and streaming. + - Add adapter chain summary trace test. + - Add field-cache start/complete trace tests, including no-rules/no-match cases. + - Add Responses service/stream trace tests if existing fixtures can cover it without large integration cost. + - Run protocol regressions because trace serialization touches protocol objects. + +## Commit Checkpoints + +1. `docs(experimental): plan transform trace coverage correction`. +2. Provider transform trace coverage and tests. +3. Executor request/response/stream trace coverage and tests. +4. Native provider/adapter/field-cache trace summaries and tests. +5. Responses trace coverage and tests. +6. Review fixes after `explore` and `explore-heavy`. +7. User-facing Phase 2b report, uncommitted. + +## Risks And Mitigations + +- Trace logging accidentally mutates payloads. Mitigation: pass deep-copied or read-only snapshots and rely on `sanitize_for_trace()`. +- New trace arguments break callers. Mitigation: make all new args keyword-only optional and keep old call sites valid. +- Stream trace output grows too much. Mitigation: keep per-stream-event snapshots disabled and log compact metadata. +- Import cycles. Mitigation: use private local helpers when needed. +- Brittle tests. Mitigation: assert pass presence and causal relative order only where required. diff --git a/docs/experimental/phase-2c-transform-trace-completion.md b/docs/experimental/phase-2c-transform-trace-completion.md new file mode 100644 index 000000000..0de390162 --- /dev/null +++ b/docs/experimental/phase-2c-transform-trace-completion.md @@ -0,0 +1,73 @@ +# Phase 2c: Transform Trace Completion And Redaction Hardening + +## Goal + +Close the Phase 2/2b third-pass tracing findings: Anthropic compatibility coverage, native stream-event adapter tracing, Responses SSE formatting traces, transform-failure traces, final-response ordering, and camelCase secret redaction. + +## Scope + +- Add transform-pass traces around live Anthropic Messages compatibility request/response conversion. +- Add Anthropic streaming conversion traces and explicit upstream close on disconnect or abnormal exit. +- Ensure final client response traces represent the final post-normalization returned payload. +- Run and trace native stream-event adapter chains. +- Trace Responses final SSE formatting boundaries. +- Harden redaction for camelCase secret-bearing keys. +- Emit transform-failure traces for built-in provider transforms and Responses conversion errors. + +## Non-Goals + +- Do not redesign transaction log storage. +- Do not remove legacy request/response JSON logs. +- Do not enable native streaming for priority providers. +- Do not change public API response shapes. + +## Implementation Plan + +1. Anthropic non-stream traces. + - Trace `anthropic_raw_request`, `anthropic_to_openai_request`, `anthropic_openai_response`, `openai_to_anthropic_response`, and `anthropic_final_response`. + +2. Anthropic streaming traces and close safety. + - Trace source OpenAI chunks and emitted Anthropic SSE frames. + - Close upstream stream via `aclose()` / `close()` on disconnect or abnormal exit. + - Trace stream transform errors before emitting the Anthropic error frames. + +3. Provider transform failure traces. + - Wrap built-in provider transforms and emit `transform_log_error` before re-raising. + - Preserve provider-hook behavior but keep sanitized metadata. + +4. Native stream-event adapter chain. + - Run adapter chain for parsed native stream events. + - Trace `after_stream_event_adapter_chain` before formatting. + +5. Responses SSE formatting trace. + - Trace each formatted `ResponsesStreamEvent` frame and terminal frame emitted by `stream_response()`. + +6. Final response trace ordering. + - Verify or adjust `final_client_response` ordering so it occurs after usage normalization/cost accounting. + +7. Redaction hardening. + - Normalize camelCase keys such as `apiKey`, `accessToken`, `refreshToken`, `clientSecret`, and `idToken`. + +8. Responses transform-failure traces. + - Emit standardized transform error traces for Responses parse/bridge/storage conversion errors without changing raised errors. + +## Tests + +- Anthropic non-stream and stream trace tests. +- Anthropic disconnect upstream close test. +- Built-in provider transform failure trace test. +- Native stream-event adapter mutation/trace test. +- Responses SSE formatting trace test. +- CamelCase redaction tests. +- Final response trace-ordering tests. + +## Acceptance Criteria + +- Anthropic compatibility has transform-pass coverage for request conversion, response conversion, and stream conversion. +- Anthropic stream disconnect closes upstream when possible. +- Built-in provider transform exceptions emit transform-failure traces. +- Native stream-event adapter chains are executed and traced. +- Responses stream SSE formatting boundaries are traced. +- Final client response traces reflect post-normalization final payloads. +- CamelCase secret-bearing keys are redacted. +- Focused tests pass and both `explore` and `explore-heavy` reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-3-adapters-field-cache.md b/docs/experimental/phase-3-adapters-field-cache.md new file mode 100644 index 000000000..d7fbf53cf --- /dev/null +++ b/docs/experimental/phase-3-adapters-field-cache.md @@ -0,0 +1,284 @@ +# Phase 3 Plan: Adapter And Field-Cache System + +## Goal + +Add the configurable adapter and field-cache foundation that lets protocol/provider implementations preserve and re-inject provider-specific state without hardcoding every provider. This phase should make it possible to define rules for values like reasoning content, thought signatures, prompt cache keys, provider session IDs, and response IDs, with strict provider/model/credential/session/classifier scoping and transform trace visibility. + +## Non-Goals + +- Do not migrate all existing providers to native protocols yet. +- Do not add `/v1/responses` routes yet. +- Do not replace `ProviderTransforms` yet. +- Do not implement JSON config loading for these rules yet; Phase 10 owns config polish. +- Do not instantiate SQLite or any DB. +- Do not implement provider-specific Claude/Codex/Copilot/Antigravity behavior yet. +- Do not make field cache a replacement for `SessionTracker`. + +## Current Code Context + +- `ProviderTransforms` is still the active runtime transform path and should remain untouched unless a small hook is needed. +- Phase 1 provides unified request/response/stream dataclasses and protocol adapters. +- Phase 2 provides `log_transform_pass()` and `log_transform_error()` for request/response/stream trace states. +- `ProviderCache` already exists and supports async `store_async()` / `retrieve_async()` with memory+disk TTL behavior, but it stores strings and starts async background tasks on construction. +- `RequestContext` has session, scope, classifier, provider, model, and credential metadata needed for cache isolation. +- Existing runtime behavior should stay stable while the new adapter/cache foundation is built and tested in isolation. + +## Files To Add + +- `src/rotator_library/adapters/__init__.py` +- `src/rotator_library/adapters/base.py` +- `src/rotator_library/adapters/registry.py` +- `src/rotator_library/adapters/builtin.py` +- `src/rotator_library/field_cache/__init__.py` +- `src/rotator_library/field_cache/types.py` +- `src/rotator_library/field_cache/paths.py` +- `src/rotator_library/field_cache/store.py` +- `src/rotator_library/field_cache/engine.py` +- `tests/test_adapter_registry.py` +- `tests/test_field_cache_paths.py` +- `tests/test_field_cache_engine.py` +- `tests/test_field_cache_trace.py` + +## Files Possibly Touched + +- `src/rotator_library/providers/provider_interface.py` to add optional provider declarations for adapter names and field-cache rules. +- `src/rotator_library/client/transforms.py` only if adding a no-op future hook is cleaner. +- `src/rotator_library/__init__.py` only if public lazy exports are useful. +- Avoid `RequestExecutor` runtime wiring in the first adapter/cache commits unless tests prove the hook is harmless and necessary. + +## Adapter System + +- Create a protocol-neutral `Adapter` base class. +- Adapters operate on raw dicts, unified dataclasses, or stream events depending on `supported_stages`. +- Adapter methods are override-friendly: + - `transform_request(payload, context) -> payload` + - `transform_response(payload, context) -> payload` + - `transform_stream_event(payload, context) -> payload` +- Adapter context should include provider, model, protocol, credential_id, session_id, scope_key, classifier, transport, metadata, and optional `transaction_logger`. +- Registry should support: + - explicit registration + - alias resolution + - auto-discovery from `src/rotator_library/adapters/` + - duplicate/collision checks similar to protocol registry + - order-preserving adapter chain resolution + +Built-in adapters: + +- `reasoning_rewrite`: copy/rename reasoning fields between configured paths. +- `reasoning_content`: normalize common reasoning content field names. +- `suppress_developer_role`: convert or remove developer-role messages for providers that reject it. +- `model_override`: replace model in outbound payload. +- `field_rename`: generic configured source-path to target-path copy/move. + +These adapters are bases, not gospel. Providers can subclass/copy/mutate adapters or override provider methods. + +## Field-Cache Rules + +`FieldCacheRule` fields: + +- `name` +- `source`: `request`, `response`, `stream_event`, `unified_request`, `unified_response`, `unified_stream_event` +- `path`: JSON-path-like selector +- `mode`: `last`, `all`, `last_user_turn`, `last_assistant_turn`, `per_tool_call` +- `scope`: list/tuple of scope dimensions +- `inject`: optional `FieldCacheInjection` +- `enabled` +- `ttl_seconds` optional, for future store implementations +- `metadata` for provider/model notes + +`FieldCacheInjection` fields: + +- `target`: `request`, `unified_request`, `metadata`, or provider-specific future target +- `path` +- `when_missing_only` +- `insert` +- `as_list` + +Scope dimensions: + +- `provider` +- `model` +- `credential` +- `session` +- `conversation` +- `classifier` + +Default scope for conversation-affecting rules: + +- provider + model + classifier + session. + +Cache keys must include rule name and selected scope values. Missing optional values should use stable placeholders only where safe; for session-scoped rules with no session, default to no-op unless a rule explicitly allows fallback. + +## Path Engine + +Support dotted paths: + +- `choices.0.message.reasoning_content` +- `choices.*.message.reasoning_content` +- `messages[-1].reasoning_content` +- `candidates.*.content.parts.*.thoughtSignature` + +Behavior: + +- Supports dict keys, list indices, `*`, and `[-1]`. +- Extraction returns all matches in stable traversal order. +- Injection creates missing dict containers when unambiguous. +- Injection does not create containers through wildcard paths. +- Malformed paths raise `FieldCachePathError` with useful messages. +- Missing paths are no-op, not errors. + +## Store Plan + +- `FieldCacheStore` protocol/interface with async `get`, `set`, `append`, and `clear`. +- `InMemoryFieldCacheStore` for tests and simple runtime. +- `ProviderCacheFieldStore` wrapper around an injected `ProviderCache`, JSON serializing values. +- Do not instantiate `ProviderCache` in module import paths because it starts async tasks. +- No SQLite. + +## Engine Plan + +`FieldCacheEngine` responsibilities: + +- hold rules and store +- validate rule names and paths +- `extract(source, payload, context, transaction_logger=None)` +- `inject(target, payload, context, transaction_logger=None)` + +Extraction: + +- select rules matching source +- extract values by path +- apply mode +- store under scoped key +- trace `after_field_cache_extraction` with rule name, match count, mode, scope key, and sanitized values summary + +Injection: + +- select rules with matching injection target +- retrieve scoped values +- apply mode result +- inject into a payload copy by path +- trace `after_field_cache_injection` with rule name, hit/miss, target path, and changed flag + +Engine rules: + +- Do not mutate caller payload unless explicitly requested; default returns a deep copy. +- Failures should call `log_transform_error()` with failed rule/pass and then raise validation errors in tests. Runtime integration later can choose fail-open or fail-closed per provider. + +## Modes + +- `last`: store latest extracted value. +- `all`: append extracted values preserving order. +- `last_user_turn`: use metadata/turn marker when available; for raw messages fallback to latest user-associated match if path includes messages. +- `last_assistant_turn`: same for assistant. +- `per_tool_call`: requires a tool-call ID path or extracted object with obvious `id` / `tool_call_id`; otherwise validation error. + +Phase 3 should implement `last` and `all` fully, and provide validated structure for turn/tool modes with tests for validation and simple cases. If a mode needs runtime conversation indexing not yet available, document it as present but limited. + +## Transform Trace Requirements + +Every adapter invocation should be traceable: + +- `before_adapter_chain` +- `after_adapter` +- `after_adapter_chain` + +Every field-cache injection/extraction should be traceable: + +- `before_field_cache_injection` +- `after_field_cache_injection` +- `before_field_cache_extraction` +- `after_field_cache_extraction` + +Errors use Phase 2 `transform_log_error`. + +Trace metadata should include: + +- adapter name +- rule name +- source/target +- path +- mode +- scope dimensions +- cache key hash or readable scoped key with no secrets +- match count +- changed flag + +Do not log huge cached payloads by default; log summaries and sanitized sample values. + +## Provider Declaration Plan + +Add optional provider methods/class attributes only if needed: + +- `protocol_name` +- `adapter_names` +- `field_cache_rules` +- `get_adapter_config(model)` +- `get_field_cache_rules(model)` + +Keep defaults empty/no-op to avoid provider changes. These declarations are for later provider migration. + +## Tests + +Adapter registry: + +- auto-discovers built-ins +- aliases resolve +- duplicate names/collisions are deterministic +- chain order is preserved +- no-op adapter preserves payload + +Built-in adapters: + +- `model_override` +- `suppress_developer_role` +- `field_rename` +- reasoning field copy/rename + +Path engine: + +- dict path extraction +- list index extraction +- wildcard extraction +- `[-1]` extraction +- injection into simple dict path +- missing path no-op +- wildcard injection rejected +- malformed path error + +Field-cache engine: + +- extract response value and inject into next request +- `last` overwrites +- `all` appends +- scope isolation by provider/model/session/classifier/credential +- missing path no-op +- malformed rule validation +- stream_event extraction +- trace entries emitted for injection/extraction +- transform errors emitted for rule failures + +Regression: + +- Phase 1 protocol tests. +- Phase 2 logging tests. +- core session/selection tests. + +## Commit Checkpoints + +1. Add adapter base/registry and built-in no-op/simple adapters with tests. +2. Add field-cache rule dataclasses and path engine with tests. +3. Add field-cache stores and scoped key builder with tests. +4. Add field-cache engine extraction/injection with trace tests. +5. Add optional provider declaration methods if needed. +6. Run focused and regression tests. +7. Review with `explore` and `explore-heavy`, fix findings, and write the uncommitted Phase 3 report. + +## Risks And Mitigations + +- JSON path scope could grow too large. Keep Phase 3 path syntax intentionally small and explicit. +- Cache leaks across providers/models/sessions would be severe. Scope-key tests must be extensive. +- Storing huge model outputs can bloat caches. Engine should store only matched values and trace summaries. +- ProviderCache lifecycle can be awkward because it starts async tasks. Wrap injected instances; do not instantiate globally. +- Turn/tool modes can be under-specified. Validate and implement safe subsets rather than guessing. +- Runtime integration can destabilize requests. Keep Phase 3 mostly isolated until review confirms the foundation. diff --git a/docs/experimental/phase-3b-field-cache-runtime-semantics.md b/docs/experimental/phase-3b-field-cache-runtime-semantics.md new file mode 100644 index 000000000..d65494933 --- /dev/null +++ b/docs/experimental/phase-3b-field-cache-runtime-semantics.md @@ -0,0 +1,147 @@ +# Phase 3b: Field-Cache Runtime Semantics And Persistence + +## Goal + +Correct the Phase 3 audit gap. Phase 3 built the adapter and field-cache foundation, but the audit found that the system is still partly declarative: `last_user_turn`, `last_assistant_turn`, and `per_tool_call` are validated/declared but not semantically implemented enough; the native provider executor creates a fresh `InMemoryFieldCacheStore` by default so cached values do not persist across native requests; JSON-configured field-cache rules are parsed but not merged into runtime provider declarations. + +## Non-Goals + +- Do not introduce SQLite or any database. +- Do not replace `SessionTracker`, `UsageManager`, or provider retry logic. +- Do not wire non-native LiteLLM paths through field-cache rules in this phase. +- Do not make field cache a substitute for session affinity. It preserves provider protocol state only. +- Do not implement a full external admin UI/config editor. +- Do not commit user-facing reports. +- Do not touch unrelated dirty files (`ARCHITECTURE.md`, `STRUCTURE.md`, `.opencode/`, `docs/issues/`, old phase reports). + +## Current Code State + +- `FieldCacheRule.mode` allows `last`, `all`, `last_user_turn`, `last_assistant_turn`, and `per_tool_call`. +- `FieldCacheEngine._store_values()` currently treats `last_user_turn` and `last_assistant_turn` exactly like `last`. +- `per_tool_call` currently extracts tool IDs from each extracted value, which only works when the selected value itself contains the tool ID. It cannot correlate values with a sibling or parent tool-call ID from the original payload. +- `FieldCacheInjection` has `as_list`, `when_missing_only`, and `insert`, but engine injection does not use `insert` and does not support selecting a per-tool-call value at injection time. +- `NativeProviderExecutor.__init__(field_cache_store=None)` passes `None` into each `FieldCacheEngine`, causing each request to get a new `InMemoryFieldCacheStore`. That means default native field-cache persistence is per-engine/per-request, not process-level. +- `ProviderCacheFieldStore` exists and no SQLite is present. +- `ExperimentalConfig.parse_field_cache_rules()` can parse JSON rules, but `RequestExecutor._native_context()` currently only uses `plugin.get_field_cache_rules(model)`. +- `NativeProviderContext.field_cache_context()` includes provider/model/credential/session/conversation/classifier plus metadata; metadata is the right place to carry current request payload/turn/tool-call hints without expanding core scope dimensions. +- Phase 2b hardened traces so field-cache logs expose shapes and counts, not raw cached values. + +## Implementation Plan + +### 1. Process-Local Default Native Store + +- Change `NativeProviderExecutor` so a missing `field_cache_store` creates one shared `InMemoryFieldCacheStore` in the executor instance, not a new store inside every `FieldCacheEngine`. +- Keep injected stores honored exactly as before. +- Add tests proving two calls through the same executor can extract on response and inject on a later request using the default store. +- Do not persist to disk by default; process-local is enough for runtime continuity and respects the no-SQLite rule. + +### 2. Value Envelopes For Contextual Modes + +- Add a small internal stored representation for contextual modes, likely dicts with fields such as `value`, `role`, `turn_index`, `tool_call_id`, and `source_path`. +- Keep stored values JSON-serializable for `ProviderCacheFieldStore`. +- Do not expose envelopes to injected payloads except where needed; injection should unwrap `value`. +- Preserve current external behavior for `last` and `all` as much as possible. + +### 3. Turn-Aware Extraction + +- Implement real `last_user_turn` and `last_assistant_turn` semantics. +- Minimum supported shapes: + - OpenAI/Responses-like `messages[*]` objects with `role`. + - Anthropic-like content lists nested under message objects with `role`. + - Generic payloads where a rule provides metadata hints. +- Add rule metadata hints: + - `turn_container_path`: path to turn/message list, default inferred from common roots such as `messages`. + - `turn_role_path`: role field inside each turn, default `role`. + - `turn_value_path`: value path relative to each turn when inference is better than global path. +- If no turn context can be inferred, skip with `reason="turn_context_not_found"` rather than silently behaving like `last`. +- Store only values from the latest matching user/assistant turn. +- Add tests for user turn, assistant turn, skip/no-op, and metadata-configured relative extraction. + +### 4. Per-Tool-Call Correlation And Injection + +- Keep requiring `metadata.tool_call_id_path` for `per_tool_call`. +- Support existing value-relative extraction when each selected value contains the tool ID. +- Add payload-relative correlation with metadata: + - `tool_container_path` + - `tool_call_id_path` relative to each tool container + - `tool_value_path` relative to each tool container +- Store a mapping of `tool_call_id -> value`. +- Add injection selection: + - If `metadata.inject_tool_call_id_path` exists, extract current tool IDs from target payload and inject matching values. + - If `FieldCacheContext.metadata["tool_call_id"]` exists, inject that value. + - If `inject.as_list=True`, inject all matching values as a list. + - Otherwise skip ambiguous maps with `reason="tool_call_id_not_found"` or `reason="ambiguous_tool_call_values"`. +- Add tests for sibling tool ID/value extraction and target-specific injection. + +### 5. Honor `FieldCacheInjection.insert` + +- Wire `rule.inject.insert` into `inject_path()`. +- Add tests for list insertion where path targets a list index or append-like position if supported by the path engine. +- If the path engine does not support safe insertion yet, add minimal support or explicitly validate/reject unsafe insert paths with clear errors. + +### 6. TTL Handling Without New Persistence + +- `FieldCacheRule.ttl_seconds` exists but stores ignore it. +- Add TTL support to `InMemoryFieldCacheStore`. +- Keep `ProviderCacheFieldStore` compatible: + - If the injected provider cache supports TTL arguments, pass them. + - If not, wrap values with expiry metadata and enforce expiry on `get()`. +- Update `FieldCacheStore` protocol with optional TTL on `set`/`append` only if safe; otherwise add internal engine helper methods to avoid breaking third-party stores. +- Add tests for memory-store expiry using a fake clock if possible. +- No SQLite. + +### 7. Merge JSON-Configured Rules Into Native Runtime Context + +- In `RequestExecutor._native_context()`, load optional `ExperimentalConfig` and merge `parse_field_cache_rules(config, provider, model)` with provider-declared rules. +- Provider-declared rules should come first; JSON rules append and can add/override by rule name. +- Override policy: + - If a JSON rule name matches a provider rule name, JSON replaces that rule. + - Otherwise append. +- Environment remains primary for config path and JSON secrets remain rejected by the existing loader. +- Add tests for native context rule merge without real network calls or credentials. +- Avoid top-level imports that create cycles; use local imports where needed. + +### 8. Tests + +- Engine semantics tests for contextual modes. +- Native executor test proving default process-local persistence across separate `execute()` calls. +- Config merge test for JSON field-cache rules. +- TTL store tests. +- Existing trace and protocol regression tests where touched. + +Focused suites: + +- `tests/test_field_cache_paths.py` +- `tests/test_field_cache_engine.py` +- `tests/test_field_cache_trace.py` +- `tests/test_native_provider_executor.py` +- `tests/test_experimental_config.py` +- relevant provider declaration tests +- Phase 1b/2b regression subsets if field-cache/protocol trace serialization is touched + +### 9. Documentation And Comments + +- Update docstrings in `FieldCacheEngine`, `FieldCacheRule`, and store classes to explain mode semantics. +- Document when turn/tool modes skip rather than fallback. +- Document process-local default store vs injected production store. +- Keep comments focused on algorithmic rationale: turn context inference, tool-call correlation, TTL handling, and JSON rule merge precedence. + +## Commit Checkpoints + +1. `docs(experimental): plan field cache runtime correction`. +2. Default persistent native store. +3. Turn/tool-call mode semantics. +4. TTL/store behavior. +5. JSON rule merge/runtime wiring. +6. Review fixes after `explore` and `explore-heavy`. +7. User-facing Phase 3b report, uncommitted. + +## Acceptance Criteria + +- Native field cache persists across requests by default within one executor instance. +- `last_user_turn` and `last_assistant_turn` have real turn-aware semantics and skip safely when context is unavailable. +- `per_tool_call` can correlate tool IDs with sibling values and inject the right cached value by current tool ID. +- `insert`, `ttl_seconds`, and JSON-configured rules are implemented or explicitly validated with clear errors if a requested shape is unsafe. +- No raw cached values leak in traces after Phase 2b hardening. +- No SQLite or new database is introduced. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-3c-field-cache-runtime-completion.md b/docs/experimental/phase-3c-field-cache-runtime-completion.md new file mode 100644 index 000000000..4c6bb01db --- /dev/null +++ b/docs/experimental/phase-3c-field-cache-runtime-completion.md @@ -0,0 +1,66 @@ +# Phase 3c: Field-Cache Runtime Completion And Trace Safety + +## Goal + +Close the Phase 3/3b third-pass findings around field-cache runtime coverage, adapter-chain trace safety, and credential-scoped isolation. + +## Scope + +- Execute all declared field-cache sources used by native runtime: `request`, `response`, `stream_event`, `unified_request`, `unified_response`, and `unified_stream_event`. +- Execute declared injection targets used by native runtime: `request`, `unified_request`, and `metadata`. +- Suppress generic adapter-chain payload traces in native execution where field-cache-aware redaction has not yet run, and rely on native executor traces that apply rule-aware redaction. +- Fail closed when a rule includes `credential` scope but the runtime lacks a credential identifier. +- Add focused tests for each corrected behavior. + +## Non-Goals + +- Do not replace the field-cache store or add a database. +- Do not make native streaming inject into stream events; streaming extraction remains expected behavior. +- Do not make Copilot declare field-cache rules when none are required. +- Do not implement custom provider-specific turn inference beyond existing metadata hints. + +## Implementation Plan + +1. Credential scope isolation. + - Update cache-key construction so missing `credential` scope returns `None`. + - Keep existing `allow_missing_session` behavior only for `session` scope. + - Add tests proving credential-scoped extraction/injection skips instead of using a shared missing bucket. + +2. Unified request injection and extraction. + - In native non-streaming and streaming paths, create the field-cache engine before protocol build. + - Extract `unified_request` after parsing client payload. + - Inject `metadata` before building protocol context when metadata rules exist. + - Inject `unified_request` before `protocol.build_request()`. + - Hydrate injected serialized unified-request dictionaries back into the current `UnifiedRequest` for supported common fields. + - Trace `after_unified_request_field_cache_injection` and `after_metadata_field_cache_injection`. + +3. Unified response extraction. + - Extract `unified_response` after protocol response parsing and before formatting. + - Keep existing provider-response extraction after response adapter chain. + - Trace `after_unified_response_field_cache_extraction`. + +4. Unified stream-event extraction. + - Extract `unified_stream_event` after native stream-event adapter chain and before formatting. + - Keep existing provider stream-event extraction from formatted event payload. + - Trace `after_unified_stream_event_field_cache_extraction`. + +5. Native adapter trace safety. + - Disable generic adapter-chain traces inside native request/response paths, matching the stream-event path. + - Keep native executor traces after adapter chains, where configured field-cache paths are redacted. + - Add tests with arbitrary configured provider-state paths proving generic `before_adapter_chain`/`after_adapter` traces are not emitted in native execution and native traces redact configured paths. + +6. Tests. + - Field-cache engine credential-scope fail-closed tests. + - Native executor unified request source/target tests. + - Native executor metadata injection tests. + - Native executor unified response extraction tests. + - Native stream unified event extraction tests. + - Native adapter trace redaction/suppression tests. + +## Acceptance Criteria + +- Native runtime executes all declared field-cache sources/targets that it accepts, except stream-event injection which remains intentionally unsupported. +- Missing credential scope never shares a `_none` cache key. +- Native adapter chain traces cannot leak arbitrary configured provider-state fields before rule-aware redaction. +- Existing field-cache and native provider tests continue to pass. +- Both `explore` and `explore-heavy` reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-4-responses-api.md b/docs/experimental/phase-4-responses-api.md new file mode 100644 index 000000000..901510066 --- /dev/null +++ b/docs/experimental/phase-4-responses-api.md @@ -0,0 +1,285 @@ +# Phase 4 Plan: Responses API + +## Goal + +Add a first-class OpenAI-compatible Responses API surface with durable response storage, `previous_response_id` continuation support, HTTP SSE streaming, and an explicit WebSocket transport seam for future implementation. Phase 4 should use the Phase 1 Responses protocol, Phase 2 transform trace logging, and Phase 3 field-cache/adapter foundations without forcing all providers onto native protocols yet. + +## Non-Goals + +- Do not migrate every provider to native protocol execution in this phase. +- Do not remove or replace `/v1/chat/completions`. +- Do not replace `UsageManager`, `SessionTracker`, retry-after parsing, or current credential rotation. +- Do not implement full multi-user/security controls. +- Do not add SQLite or any database. +- Do not implement full provider-specific Responses APIs for Claude Code/Codex/Copilot yet. +- Do not implement WebSocket runtime handling yet; provide transport interfaces/seams and tests for HTTP SSE behavior. + +## Current Code Context + +- FastAPI routes live in `src/proxy_app/main.py`. +- `RotatingClient.acompletion()` is the stable execution entry point today. +- `ResponsesProtocol` already parses/formats Responses request/response/stream shapes and declares future `websocket` support. +- `TransactionLogger` can record new transform passes and errors. +- `FieldCacheEngine` and `ProviderCacheFieldStore` can preserve response IDs/output items but need targeted key deletion for Responses delete endpoints. +- Existing stream wrappers aggregate chat-completions SSE chunks, but Responses SSE needs its own event formatting instead of chat completion chunk output. +- Runtime execution should stay conservative: Phase 4 can bridge Responses requests to the existing completion path until native provider protocols are wired in later phases. + +## Files To Add + +- `src/rotator_library/responses/__init__.py` +- `src/rotator_library/responses/types.py` +- `src/rotator_library/responses/store.py` +- `src/rotator_library/responses/bridge.py` +- `src/rotator_library/responses/service.py` +- `src/rotator_library/responses/streaming.py` +- `tests/test_responses_store.py` +- `tests/test_responses_bridge.py` +- `tests/test_responses_service.py` +- `tests/test_responses_routes.py` +- `tests/test_responses_streaming.py` + +## Files Likely To Touch + +- `src/proxy_app/main.py` +- `src/rotator_library/client/rotating_client.py` +- `src/rotator_library/field_cache/store.py` if targeted delete is added to the store protocol. +- `src/rotator_library/__init__.py` only if lazy public exports are useful. +- Existing protocol tests only if the Responses adapter needs an additive event/output fix. + +## Route/API Scope + +`POST /v1/responses`: + +- Accept raw Responses JSON. +- Support non-streaming response. +- Support `stream: true` with HTTP SSE. +- Validate required `model`. +- Accept `input`, `instructions`, `tools`, `tool_choice`, `reasoning`, `metadata`, `include`, `store`, `previous_response_id`, and common generation params. + +`GET /v1/responses/{response_id}`: + +- Return stored response object when available. +- Return 404-compatible error if missing or deleted. + +`DELETE /v1/responses/{response_id}`: + +- Delete stored response metadata/output. +- Return `{ "id": ..., "object": "response.deleted", "deleted": true }` or equivalent compatible shape. + +`GET /v1/responses/{response_id}/input_items`: + +- Return stored input items for clients that need continuation inspection. + +Optional if simple: + +- `POST /v1/responses/{response_id}/cancel` for active stream cancellation can be planned but not fully implemented unless active stream tracking lands cleanly. + +## Response Storage + +`StoredResponse` fields: + +- `id` +- `created_at` +- `model` +- `status` +- `request` +- `response` +- `input_items` +- `output_items` +- `usage` +- `metadata` +- `session_id` +- `scope_key` +- `classifier` +- `expires_at` optional + +`ResponsesStore` protocol: + +- `save(stored_response)` +- `get(response_id)` +- `delete(response_id)` +- `list_input_items(response_id)` + +Store implementations: + +- `InMemoryResponsesStore` for tests and runtime default if no persistent cache is configured. +- `ProviderCacheResponsesStore` as an optional wrapper around injected `ProviderCache`. +- JSON serialization through `serialize_value()`. +- No global `ProviderCache` construction at import time. +- No SQLite. + +Targeted deletion: + +- Add `delete(key)` to `FieldCacheStore` only if reused directly. +- Prefer a dedicated Responses store so field cache remains small and focused. + +## Response IDs + +- Generate IDs like `resp_` when upstream response lacks an ID. +- Preserve upstream IDs when present. +- Store every response when request `store` is true or omitted if compatibility expects default storage. +- If `store: false`, return the response but do not persist it; `previous_response_id` cannot refer to it later. +- Include `previous_response_id` in stored metadata for lineage/debugging. + +## `previous_response_id` + +- On request with `previous_response_id`, load the stored parent response. +- If not found, return a 404/400-compatible invalid request error rather than silently ignoring. +- Build continuation context conservatively: + - Keep original current request input. + - Add parent response output items into protocol/service metadata for traceability. + - For bridge execution through chat completions, convert parent response message items into prior assistant messages where safe. + - Do not inject unknown output item types into chat messages unless `ResponsesProtocol` can represent them safely. +- Record transform trace pass: + - `responses_previous_response_loaded` + - Include parent ID, output count, input item count, and whether bridge context was expanded. +- This should line up with Phase 3 field-cache storage but does not need to fully rely on field-cache rules yet. + +## Bridge Execution + +`ResponsesBridge` responsibilities: + +- Convert `UnifiedRequest` from `ResponsesProtocol.parse_request()` into chat-completions kwargs for current `RotatingClient.acompletion()`. +- Map `model` to `model`. +- Map `instructions`/system blocks to a system message. +- Map input messages to chat messages. +- Map tool definitions to OpenAI tool shape when possible. +- Map generation params to existing compatible kwargs. +- Map `stream` to `stream`. +- Preserve unsupported Responses fields in trace metadata and request metadata, not silently discard them. +- Rebuild Responses output from chat completion responses using `ResponsesProtocol.format_response()`. + +This bridge is temporary compatibility, not the final native provider path. Docstrings/comments should explain that native provider execution will replace the bridge for covered providers in later phases. + +## Service + +`ResponsesService` responsibilities: + +- Own `ResponsesProtocol`, `ResponsesBridge`, `ResponsesStore`, and optional transform logging. +- `create_response(raw_request, client, request=None)` +- `stream_response(raw_request, client, request=None)` +- `get_response(response_id)` +- `delete_response(response_id)` +- `list_input_items(response_id)` + +Service transform passes: + +- `raw_responses_request` +- `parsed_unified_request` +- `responses_previous_response_loaded` +- `responses_bridge_chat_request` +- `raw_responses_provider_response` or `raw_chat_bridge_response` +- `parsed_unified_response` +- `stored_responses_response` +- `final_responses_response` + +Errors use `log_transform_error()`. + +## Streaming + +- HTTP SSE first. +- Convert chat-completions stream chunks into Responses SSE events. +- Required event names: + - `response.created` + - `response.output_item.added` + - `response.output_text.delta` + - `response.output_item.done` + - `response.completed` + - `response.failed` on errors + - final `data: [DONE]` only if compatibility requires it; prefer Responses-style events and document the chosen behavior. +- Preserve raw stream chunk trace: + - `raw_chat_bridge_stream_chunk` + - `parsed_unified_stream_event` + - `formatted_responses_stream_event` + - `stored_responses_stream_response` +- Accumulate final output items so streamed responses are retrievable by `GET /v1/responses/{id}` after completion. +- Handle client disconnect without crashing the server. +- Do not add the broader stream stall/backpressure overhaul; that belongs to the later streaming phase. + +## WebSocket Seam + +- Add transport-neutral interfaces: + - `ResponsesTransport` or small formatter abstraction. + - `ResponsesSSEFormatter`. + - `ResponsesWebSocketFormatter` placeholder or protocol with `NotImplementedError`. +- Keep protocol/service logic independent from `StreamingResponse`. +- Add tests that formatter/service API can select `sse` now and recognizes `websocket` as a future transport without implementing a WebSocket route. +- Do not expose a WebSocket endpoint yet unless it is a no-op documented placeholder; prefer no route over a misleading route. + +## FastAPI Integration + +- Add routes near existing OpenAI-compatible routes in `main.py`. +- Use existing `verify_api_key`. +- Use existing raw request logging where applicable. +- Use `JSONResponse` for normal responses. +- Use `StreamingResponse(..., media_type="text/event-stream")` for streams. +- Use OpenAI-compatible error response shape for validation/missing response IDs. +- Store service on app state during lifespan if it needs shared storage, or construct module-level in-memory store carefully if no async lifecycle is needed. Prefer `app.state.responses_service`. + +## Tests + +Store tests: + +- save/get/delete response +- input items list +- missing response returns None +- JSON serialization with provider-cache wrapper if implemented +- no SQLite/import-time cache construction + +Bridge tests: + +- Responses input string/list converts to chat messages +- instructions become system message +- tool definitions preserve function call shape +- unsupported fields are preserved in metadata/trace +- chat completion response formats back to Responses output + +Service tests: + +- non-stream create stores response +- `store: false` does not persist +- `previous_response_id` loads parent +- missing previous response raises useful error +- transform trace passes emitted + +Route tests: + +- `POST /v1/responses` non-stream success with fake client +- `POST /v1/responses` missing model returns 400 +- `GET /v1/responses/{id}` success and 404 +- `DELETE /v1/responses/{id}` success +- `GET /v1/responses/{id}/input_items` + +Streaming tests: + +- stream emits Responses SSE event names in order +- deltas aggregate into stored final response +- stream errors emit `response.failed` +- client disconnect path is safe if testable + +Regression tests: + +- Phase 1 protocol tests +- Phase 2 transform logging tests +- Phase 3 adapter/cache tests +- `test_session_tracking.py` +- `test_selection_engine.py` + +## Commit Checkpoints + +1. Add Responses store and tests. +2. Add Responses bridge and tests. +3. Add Responses service with non-stream create/get/delete/input-items and tests. +4. Add HTTP routes and route tests. +5. Add HTTP SSE formatter/streaming conversion and tests. +6. Add WebSocket transport seam tests/docstrings without runtime route. +7. Run focused and regression tests. +8. Review with `explore` and `explore-heavy`, fix findings, and write the uncommitted Phase 4 report. + +## Risks And Mitigations + +- Responses API is broader than chat completions. Keep the bridge explicit about unsupported fields and preserve them in metadata/trace rather than pretending full native support. +- `previous_response_id` can leak context across users/scopes if storage keys are global. Include request/session/scope/classifier metadata and leave full multi-user isolation for later, but avoid credential secrets in storage keys. +- Streaming conversion can lose item fidelity. Accumulate output text and tool/reasoning items explicitly; preserve raw chunks in trace. +- In-memory store will not survive restarts. This is acceptable for Phase 4 default; provider-cache-backed persistence can be enabled where lifecycle is available. +- WebSocket support can be over-promised. Expose the abstraction seam only, not a functioning route, until the later transport phase. diff --git a/docs/experimental/phase-4b-responses-session-storage-transport.md b/docs/experimental/phase-4b-responses-session-storage-transport.md new file mode 100644 index 000000000..81e352540 --- /dev/null +++ b/docs/experimental/phase-4b-responses-session-storage-transport.md @@ -0,0 +1,94 @@ +# Phase 4b Plan: Responses API Session, Storage, And Transport Corrections + +## Goal + +Correct the Phase 4 audit findings without over-claiming full native Responses provider execution. Phase 4 made `/v1/responses` usable through the chat bridge, but the validation pass found that important Responses semantics are still incomplete: `previous_response_id` is only used to replay parent output, not as session-routing evidence; default storage has only incidental expiry support with no runtime TTL policy or max-size pruning; stream generation is SSE-first inside the service instead of cleanly transport-neutral; and route/service tracing/storage should more clearly represent stored/skipped/current-state behavior. + +## Non-Goals + +- Do not implement a WebSocket FastAPI route in this phase unless the transport abstraction makes it trivial and tests stay small. +- Do not replace the bridge with fully native provider Responses execution here. +- Do not introduce SQLite or a database. +- Do not build multi-user security/admin features. +- Do not change current public SSE event order except where tests show it is wrong. +- Do not commit user-facing reports. +- Do not touch unrelated dirty `ARCHITECTURE.md`, `STRUCTURE.md`, `.opencode/`, `docs/issues/`, or old phase reports. + +## Current Code State + +- `ResponsesService.create_response()` parses `previous_response_id`, loads the parent from the store, and passes parent output to `ResponsesBridge.to_chat_kwargs()`. +- `ResponsesBridge.to_chat_kwargs()` stores `previous_response_id` only in private `_responses_bridge` metadata used for trace context, then `ResponsesService` pops that metadata before calling `client.acompletion()`. +- The actual `RotatingClient` request context cannot see `previous_response_id`, so `SessionTracker` cannot use it as a strong anchor for credential affinity. +- `StoredResponse` has `expires_at`, and stores check expiry on `get()`, but `ResponsesService` never sets a TTL policy and `InMemoryResponsesStore` has no max-size pruning. +- `ProviderCacheResponsesStore` exists but app startup always creates `ResponsesService()` with default in-memory store. +- `ResponsesService.stream_response()` directly instantiates `ResponsesSSEFormatter()` and yields formatted strings. +- `ResponsesWebSocketFormatter` exists as a placeholder with `NotImplementedError`, which is honest but does not prove that service logic is formatter-neutral. +- Responses stream storage stores only the completed payload after stream success and can skip storage when `store=false`. It does not store in-progress current state. +- `response.failed` events are yielded on exceptions, but failed responses are not stored as failed state even when `store=true`. + +## Implementation Plan + +1. Add an internal Responses session-hints carrier. + - Introduce private `_session_tracking_hints` request kwargs consumed by `RequestContextBuilder` and removed before provider execution. + - Merge service-level hints with provider hints before `SessionTracker.infer_session()`. + - Preserve provider hints and explain that these hints are proxy-internal evidence, never provider payload. + +2. Make `previous_response_id` a strong Responses session anchor. + - Attach `_session_tracking_hints` for continuation requests with a strong anchor like `responses_previous_response_id:{id}` and deterministic affinity key. + - Do not fake a strong first-turn anchor from a generated response ID before the request executes. + - Add tests proving continuation requests expose hidden hints to context construction and the hidden field is not sent to providers. + +3. Record generated Responses IDs as response-derived metadata. + - Store response IDs and parent IDs clearly in `StoredResponse.metadata`. + - Add a helper for Responses session hints so future native Responses execution can share it. + +4. Add runtime storage policy. + - Add `ResponsesStoreSettings` with `ttl_seconds`, `max_items`, `store_failed`, and `store_in_progress`. + - Preserve current defaults: no expiry, no max pruning, completed responses stored, in-progress updates disabled. + +5. Implement TTL assignment and max-size pruning. + - Set `StoredResponse.expires_at` in `_stored_response()` from settings. + - Add max-item pruning to `InMemoryResponsesStore.save()`. + - Keep provider-cache expiry via `StoredResponse.expires_at`; document max-size limitations if listing is unavailable. + +6. Store failed stream responses when `store=true`. + - Persist `response.failed` payloads when storage policy says failed responses are stored. + - Keep `store=false` skip behavior. + +7. Add a current-state storage seam. + - Add `store_in_progress` default false. + - When enabled, save created/intermediate/completed stream state using the same store API. + - Keep default behavior unchanged. + +8. Refactor stream generation to transport-neutral events. + - Add `ResponsesStreamEvent` and a formatter interface in `responses/streaming.py`. + - Add `ResponsesService.stream_events()` yielding event objects. + - Make `stream_response()` a thin SSE wrapper over `stream_events()`. + - Keep `ResponsesWebSocketFormatter` honest: either pure JSON event formatting or explicit route-level future support, but service logic must not be SSE-specific. + +9. Tighten trace boundaries around the event pipeline. + - Preserve existing `responses_stream_event_*` pass names. + - Avoid trace-only conversions when no transaction logger is present. + - Add trace pass for current-state storage updates if enabled. + +10. Update FastAPI route wiring minimally. + - Keep current `/v1/responses` route behavior. + - Continue using `service.stream_response()` for SSE. + - Defer provider-cache/durable config wiring to config work unless a safe existing setting exists. + +## Tests + +- Responses service tests for hidden session hints, TTL, max pruning, and metadata. +- Responses streaming tests for `stream_events()` order, SSE wrapper compatibility, failed storage, `store=false`, and WebSocket formatter seam. +- Responses store tests for memory max pruning and provider-cache expiry behavior. +- Request-builder/session tracking tests for `_session_tracking_hints` consumption and provider payload cleanup. +- Route regression tests for current HTTP behavior. + +## Acceptance Criteria + +- `previous_response_id` becomes strong session-routing evidence for continuation requests without leaking internal hints to provider payloads. +- Responses storage has explicit TTL/max-size policy and failed-stream storage behavior. +- Streaming service logic can emit transport-neutral event objects, with SSE as a formatter wrapper and WebSocket support no longer requiring service/protocol rewrite. +- Existing `/v1/responses`, retrieve, delete, input_items, and SSE behavior remain compatible. +- Trace-disabled paths avoid unnecessary trace-only conversions. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-4c-responses-storage-lineage-errors.md b/docs/experimental/phase-4c-responses-storage-lineage-errors.md new file mode 100644 index 000000000..b35b1a329 --- /dev/null +++ b/docs/experimental/phase-4c-responses-storage-lineage-errors.md @@ -0,0 +1,60 @@ +# Phase 4c: Responses Storage, Continuation Lineage, And Error Shape + +## Goal + +Close the Phase 4/4b third-pass Responses findings while keeping the existing bridge architecture and no-SQLite constraint. + +## Scope + +- Wire configurable Responses storage at app startup. +- Keep process-local memory storage as the default. +- Add provider-cache-backed JSON storage as an opt-in durable backend. +- Preserve existing `ResponsesStoreSettings` policy for TTL, maximum items, failed storage, and in-progress storage. +- Extend `previous_response_id` continuation context to replay parent input and output lineage, oldest-to-newest. +- Return top-level OpenAI-compatible `error` bodies from Responses routes. + +## Non-Goals + +- Do not introduce SQLite or a new database. +- Do not replace the chat-completions bridge with native Responses execution. +- Do not add a WebSocket route. +- Do not implement a cancel endpoint. +- Do not perfect every rich Responses item type in the bridge; cover known text/message/tool-call cases and preserve unsupported fields. + +## Implementation Plan + +1. Configurable Responses store backend. + - Add runtime config helpers for `memory` and `provider_cache` backends. + - Use existing provider-cache JSON storage for durable mode. + - Keep memory backend as the default. + +2. Startup wiring. + - Construct `ResponsesService(store=..., store_settings=...)` at proxy startup and fallback service creation. + - Keep direct `ResponsesService()` defaults unchanged for tests and embedding users. + +3. Continuation lineage. + - Load parent chains through stored `request.previous_response_id` with depth and cycle guards. + - Pass oldest-to-newest lineage into the bridge. + +4. Bridge replay. + - Replay parent input items as user messages when convertible. + - Replay parent output message items as assistant messages. + - Keep existing parent-output behavior for external callers using the old argument. + +5. Top-level route errors. + - Return `JSONResponse(status_code=..., content={"error": ...})` for Responses service errors. + - Apply to create validation, retrieve, delete, and input-items routes. + +6. Tests. + - Config/env tests for storage backend selection. + - Store tests for provider-cache-backed persistence. + - Bridge/service tests for parent input + output lineage replay. + - Route tests for top-level error bodies. + +## Acceptance Criteria + +- App startup can use durable provider-cache-backed Responses storage via config/env. +- Existing memory storage remains default. +- Continuations replay parent input and output lineage, not just parent output. +- Responses route errors use top-level OpenAI-compatible `error` bodies. +- Focused tests pass and both reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-5-provider-protocol-overhaul.md b/docs/experimental/phase-5-provider-protocol-overhaul.md new file mode 100644 index 000000000..c05d94ad7 --- /dev/null +++ b/docs/experimental/phase-5-provider-protocol-overhaul.md @@ -0,0 +1,263 @@ +# Phase 5 Plan: Native Provider Protocol Overhaul + +## Goal + +Make provider implementations opt into the native protocol, adapter, field-cache, and Responses foundations added in Phases 1-4, then add or restore the priority providers in the requested order: Claude Code, Codex, Copilot, Antigravity, and Gemini CLI parity review. This phase should create a reusable provider-native execution seam first, then implement providers incrementally behind explicit declarations so LiteLLM remains fallback-only for uncovered cases. + +## Non-Goals + +- Do not replace `UsageManager`, `SessionTracker`, `SelectionEngine`, or retry-after parsing. +- Do not implement routing/fallback groups yet; Phase 6 owns ordered fallback chains. +- Do not implement full multi-user/security isolation. +- Do not add SQLite or any database. +- Do not rewrite the entire executor in one pass. +- Do not migrate every existing provider at once. +- Do not copy obsolete retired-provider behavior blindly. +- Do not remove LiteLLM fallback; make fallback explicit and observable. + +## Current Code Context + +- Phase 1 provides reusable protocol adapters: OpenAI Chat, Anthropic Messages, Gemini, Responses, and LiteLLM fallback. +- Phase 2 provides transaction transform tracing and provider trace entries. +- Phase 3 provides provider declarations, adapter chains, field-cache rules, and scoped cache engine. +- Phase 4 provides `/v1/responses`, response storage, HTTP SSE, and WebSocket seam. +- `ProviderInterface` already has optional `get_protocol_name()`, `get_adapter_names()`, `get_adapter_config()`, and `get_field_cache_rules()`. +- Existing custom providers with `has_custom_logic()` include Gemini CLI, Deepseek, and retired Antigravity/IFlow/Qwen-style providers. +- Gemini CLI is large and custom; Phase 5 should review and target improvements rather than destabilize it. +- Antigravity exists only under `src/rotator_library/providers/_retired/` and must be compared against current behavior before restoration. +- No Claude Code, Codex, or Copilot provider files exist in active providers today. + +## Files To Add + +- `src/rotator_library/native_provider/__init__.py` +- `src/rotator_library/native_provider/context.py` +- `src/rotator_library/native_provider/executor.py` +- `src/rotator_library/native_provider/http.py` +- `src/rotator_library/native_provider/streaming.py` +- `src/rotator_library/providers/claude_code_provider.py` +- `src/rotator_library/providers/codex_provider.py` +- `src/rotator_library/providers/copilot_provider.py` +- Possibly `src/rotator_library/providers/antigravity_provider.py` +- `tests/test_native_provider_executor.py` +- `tests/test_native_provider_streaming.py` +- `tests/test_claude_code_provider.py` +- `tests/test_codex_provider.py` +- `tests/test_copilot_provider.py` +- `tests/test_antigravity_provider_restore.py` if restored in this phase +- `tests/test_gemini_cli_protocol_declarations.py` + +## Files Likely To Touch + +- `src/rotator_library/providers/provider_interface.py` only for small optional native execution methods if needed. +- `src/rotator_library/client/executor.py` only if a minimal seam is needed to route declared native providers without breaking current providers. +- `src/rotator_library/client/models.py` only if model/provider resolution needs explicit provider declarations. +- `src/rotator_library/client/request_builder.py` only if provider declarations need context fields before execution. +- `src/rotator_library/responses/service.py` only if native Responses-capable providers can bypass the chat bridge cleanly. +- Existing provider files only for declaration additions or targeted Gemini CLI fixes. + +## Native Provider Execution Seam + +- Add `NativeProviderContext` with provider, model, credential identifier, headers, request context, protocol name, adapter names, field-cache rules, transport, transaction logger, session fields, scope key, classifier, and provider metadata. +- Add `NativeProviderExecutor` that can: + - Resolve protocol via `provider.get_protocol_name(model)`. + - Build `AdapterContext` and run `before_adapter_chain` / `after_adapter` / `after_adapter_chain`. + - Run `FieldCacheEngine.inject()` before provider request. + - Build native HTTP request payload from the selected protocol. + - Send request through a small HTTP transport wrapper. + - Parse provider response through selected protocol. + - Run `FieldCacheEngine.extract()` after response or stream events. + - Format back to the requested client protocol. + - Emit transform trace passes for every step. +- The executor must be provider-opt-in. If `get_protocol_name()` returns `None` or `"litellm_fallback"`, current behavior stays unchanged. +- It must be independently testable with mocked HTTP clients and fake providers. + +## Provider Interface Additions + +Prefer using existing declarations first. Add optional methods only if needed: + +- `get_native_endpoint(model, operation)` +- `get_native_headers(credential_identifier, model, operation)` +- `build_native_request_options(model, operation)` +- `supports_native_operation(operation, model)` +- `should_use_native_protocol(operation, model)` + +Defaults must preserve current behavior. Docstrings must explain provider overrides are encouraged when a base protocol is close but not exact. + +## HTTP Transport + +- Use injected `httpx.AsyncClient`. +- Support JSON POST first. +- Support SSE response iteration for streaming. +- Preserve raw request/response bodies for transform trace. +- Do not add WebSocket runtime transport yet; keep the seam compatible with Phase 4 WebSocket formatter. +- Convert provider HTTP errors into existing error classification paths where possible. + +## Streaming + +- Native stream parser should call protocol stream parsing where available. +- Emit transform passes: + - `native_provider_request` + - `raw_native_provider_stream_chunk` + - `parsed_native_stream_event` + - `after_field_cache_stream_extraction` + - `formatted_client_stream_event` +- Do not fallback after visible output. Phase 6 owns fallback policy; Phase 5 only makes stream behavior observable and safe. + +## Responses Integration + +- For providers with native Responses support, allow `ResponsesService` to bypass `ResponsesBridge`. +- Add a service-level native executor hook only if it can be done without coupling service to provider internals. +- If not clean in Phase 5, leave `/v1/responses` bridge as default and expose provider-native Responses in provider tests/foundation for Phase 6/8 wiring. +- Preserve `/v1/responses` route behavior from Phase 4. + +## Priority Provider Plan + +### Claude Code + +- Add provider first. +- Determine whether it is Anthropic Messages-compatible, OpenAI Chat-compatible, or a dedicated endpoint. +- Start with declarations and mocked native execution tests. +- Preserve Claude reasoning/thinking fields via field-cache rules if present. +- Suppress/transform unsupported roles through adapters instead of bespoke monolithic code. +- Add tests for auth headers, model list behavior, request transform, response transform, streaming text, and field-cache rule declarations. +- If live API details are uncertain, implement an integration path with explicit env names and mocked endpoint behavior rather than guessing secrets or undocumented flows. + +### Codex + +- Add provider second. +- Treat as likely OpenAI/Responses-compatible until provider-specific evidence says otherwise. +- Prefer native Responses protocol if supported; otherwise OpenAI Chat protocol. +- Add tests for protocol selection, Responses bypass or bridge compatibility, auth headers, model naming, and explicit no-LiteLLM path when native mode is declared. + +### Copilot + +- Add provider third. +- Use native protocol declarations and OAuth/header helpers. +- Keep credential acquisition/refresh minimal and mocked unless existing credential machinery already supports it. +- Add field-cache rules for conversation/session IDs if required by provider behavior. +- Add tests for model list, auth header, protocol selection, request conversion, and streaming. + +### Antigravity + +- Compare `src/rotator_library/providers/_retired/antigravity_provider.py`, auth base, quota tracker, and device profile utilities before restoring. +- Restore only valid/current behavior. +- Avoid fragile or obsolete device-profile behavior unless current service behavior requires it. +- Extract reusable pieces: + - model mapping + - auth headers + - schema cleanup + - thinking/reasoning preservation + - quota parsing + - SSE handling +- Do not resurrect the whole monolith unchanged. +- Add a restored provider only after tests describe the stable subset. +- If the current service cannot be validated safely, write an explicit integration-path provider skeleton and defer live-specific behavior. + +### Gemini CLI Parity Review + +- Review active `gemini_cli_provider.py` against the new protocol/adapter/cache foundation. +- Add provider declarations where safe: + - likely protocol `gemini` + - adapters for model override / field rename / role suppression only if needed + - field-cache rules for thought signatures and provider session fields if they match current behavior. +- Do not rewrite Gemini CLI in this phase. +- Add targeted tests proving declarations do not change current behavior. +- Fix only clear parity gaps found during review. + +## LiteLLM Fallback Policy + +- Providers with native declarations should not silently fall back to LiteLLM on the primary path. +- If fallback is allowed, it must be explicit: + - protocol name `"litellm_fallback"` + - trace pass `native_provider_litellm_fallback` + - metadata reason. +- Tests must assert no accidental LiteLLM fallback for providers declared native. + +## Transform Trace Requirements + +Provider-native calls should log: + +- `native_protocol_selected` +- `before_adapter_chain` +- `after_adapter` +- `after_field_cache_injection` +- `native_provider_request` +- `raw_native_provider_response` +- `parsed_native_provider_response` +- `after_field_cache_extraction` +- `final_client_response` + +Provider-native stream calls should log: + +- `native_protocol_selected` +- `native_provider_stream_request` +- `raw_native_provider_stream_chunk` +- `parsed_native_stream_event` +- `after_field_cache_stream_extraction` +- `formatted_client_stream_event` + +Errors should use `log_transform_error()` with provider/protocol/pass context. + +## Field-Cache Requirements + +- Use Phase 3 rules for: + - reasoning content + - thought signatures + - provider session IDs + - prompt cache keys + - previous/response IDs where provider-specific +- Scope must include at least provider + model + classifier + session for conversation-affecting values. +- Credential scope must be added for values tied to an account/token. +- Tests must prove no cross-provider/model/session/credential leakage for provider rules. + +## Testing Plan + +Native executor tests: + +- protocol selection +- adapter chain order +- field-cache injection/extraction +- non-stream request/response trace passes +- stream trace passes +- explicit fallback behavior +- provider opt-in leaves current fallback path untouched for undeclared providers + +Provider tests: + +- registration/discovery for Claude Code, Codex, Copilot, and restored Antigravity if added. +- env var naming and auth header construction via mocks. +- model list parsing via mocked HTTP. +- request payload build and response parse via mocked HTTP. +- SSE stream conversion via mocked chunks. +- field-cache declarations. +- no live credentials required. + +Regression tests: + +- Phase 1 protocol tests. +- Phase 2 logging tests. +- Phase 3 adapter/field-cache tests. +- Phase 4 Responses tests. +- `test_session_tracking.py`. +- `test_selection_engine.py`. + +## Commit Checkpoints + +1. Add native provider context/executor/HTTP transport with tests. +2. Add native streaming support with tests. +3. Add Claude Code provider or integration skeleton with tests. +4. Add Codex provider or integration skeleton with tests. +5. Add Copilot provider or integration skeleton with tests. +6. Compare and restore the safe Antigravity subset, or write a documented deferral with tests. +7. Add Gemini CLI declaration/parity fixes with tests. +8. Run focused and regression tests. +9. Review with `explore` and `explore-heavy`, fix findings, and write the uncommitted Phase 5 report. + +## Risks And Mitigations + +- Provider APIs may be undocumented or volatile. Use mocked behavior and explicit integration seams rather than guessing hidden flows. +- A native executor could destabilize current traffic. Keep native execution opt-in and leave undeclared providers unchanged. +- Restoring Antigravity could reintroduce brittle device/profile logic. Restore only tested stable subsets. +- Streaming fidelity for tool/reasoning deltas may vary by provider. Preserve raw chunks in trace and keep provider-specific parsers override-friendly. +- Field-cache leakage would be serious. Add scope tests per provider rule. +- LiteLLM fallback could hide native failures. Make fallback explicit, traced, and tested. diff --git a/docs/experimental/phase-5b-provider-native-integrations.md b/docs/experimental/phase-5b-provider-native-integrations.md new file mode 100644 index 000000000..e861a9df8 --- /dev/null +++ b/docs/experimental/phase-5b-provider-native-integrations.md @@ -0,0 +1,101 @@ +# Phase 5b Plan: Priority Providers From Skeletons To Mock-Live Native Integrations + +## Goal + +Correct the Phase 5 audit gap. Phase 5 added native-provider seams and priority provider declarations, but the validation pass found the priority providers are still mostly skeletons. Phase 5b will make Claude Code, Codex, Copilot, and Antigravity usable through the native execution path under mock HTTP tests, while preserving Gemini CLI's existing custom execution path and adding parity declarations where they are safe. + +## Non-Goals + +- Do not remove LiteLLM fallback. +- Do not replace Gemini CLI's existing custom provider implementation with the generic native executor in this phase. +- Do not invent device fingerprinting or brittle environment/device-profile behavior for Antigravity. +- Do not add real credential acquisition flows for Copilot/Codex/Claude Code; this phase consumes credentials supplied by the existing credential system. +- Do not use external files outside the project root. +- Do not introduce SQLite or new persistence. +- Do not touch unrelated dirty `ARCHITECTURE.md`, `STRUCTURE.md`, `.opencode/`, `docs/issues/`, or old phase reports. +- Do not commit user-facing phase reports. + +## Current Code State + +- `NativeProviderExecutor` can parse/build/adapter/cache/HTTP/format native requests and streams. +- `RequestExecutor` can select native execution when routing target execution is `native` or auto-detected by provider protocol declaration. +- `_build_native_provider_context()` currently asks every provider for endpoint and headers with operation `"chat"`, which is too generic. +- Priority provider files exist for Claude Code, Codex, Copilot, and Antigravity, but tests mostly cover declarations and helper methods. +- Provider-prefixed model names can leak into upstream native payloads unless each provider normalizes them. +- Native streaming exists, but provider support flags are conservative and not all priority providers have stream coverage. +- Gemini CLI has substantial existing custom logic and must not be silently bypassed by auto-native routing. + +## Implementation Plan + +1. Add provider-native operation resolution. + - Add default methods to `ProviderInterface`: `get_native_operation()`, `normalize_native_model()`, and optional `prepare_native_request()`. + - Preserve current behavior by default. + - Update `RequestExecutor._build_native_provider_context()` to ask providers for operation, endpoint, headers, normalized model, and prepared request metadata. + +2. Make model normalization explicit and tested. + - Claude Code strips `claude_code/`. + - Codex strips `codex/`. + - Copilot strips `copilot/`. + - Antigravity strips `antigravity/` and maps public aliases to internal upstream names. + - Gemini CLI remains custom-path first. + +3. Make provider endpoints operation-aware. + - Claude Code uses `messages` and `/v1/messages`. + - Codex uses `responses` and `/v1/responses`. + - Copilot uses `chat` and `/chat/completions`. + - Antigravity uses Gemini generate/stream-generate endpoints. + +4. Add minimal provider request preparation. + - Normalize model before protocol parsing. + - Allow provider `prepare_native_request()` to deep-copy and adjust request payloads before protocol parsing. + - Trace this pass without adding credentials. + +5. Strengthen provider auth/header behavior. + - Keep current supplied-credential model. + - Add tested header sets for each provider. + - Do not add secrets to JSON config. + +6. Native streaming support declarations. + - Enable only where tested safe. + - Prove native streaming selection through `RequestExecutor` tests. + - Keep unsupported providers on existing fallback/custom paths. + +7. Mock-live RequestExecutor integration tests. + - Prove Claude Code, Codex, Copilot, and Antigravity use the native executor with the correct protocol, operation, endpoint, headers, model normalization, and response formatting under fake HTTP. + - Cover streaming for providers that opt in. + +8. Provider model discovery hardening. + - Test fallback and successful discovery for each priority provider. + - Avoid duplicate prefixes and invalid aliases. + +9. Gemini CLI parity review. + - Keep custom execution path. + - Verify declarations align with Gemini protocol and field-cache paths. + - Explicitly prevent accidental auto-native routing if required. + +10. Documentation and comments. + - Update provider docstrings away from “skeleton” once behavior is mock-live. + - Explain model normalization, Antigravity safety boundaries, and Gemini CLI custom-path deferral. + +## Tests + +- `tests/test_claude_code_provider.py` +- `tests/test_codex_provider.py` +- `tests/test_copilot_provider.py` +- `tests/test_antigravity_provider_restore.py` +- `tests/test_gemini_cli_protocol_declarations.py` +- `tests/test_provider_protocol_declarations.py` +- `tests/test_native_provider_executor.py` +- `tests/test_native_provider_streaming.py` +- `tests/test_request_executor_native_routing.py` +- Relevant protocol, field-cache, and routing regressions. + +## Acceptance Criteria + +- Priority providers are no longer declaration-only skeletons; each has mock-live native `RequestExecutor` coverage or an explicit tested reason for custom-path deferral. +- Native operation and endpoint selection are provider-aware, not hardcoded to `"chat"`. +- Provider-prefixed model names are normalized before native upstream calls. +- Native streaming is enabled only where tested safe. +- Gemini CLI remains on its existing custom path unless explicitly routed otherwise. +- LiteLLM fallback remains available for uncovered providers/protocols. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-5c-provider-native-correction.md b/docs/experimental/phase-5c-provider-native-correction.md new file mode 100644 index 000000000..f07939758 --- /dev/null +++ b/docs/experimental/phase-5c-provider-native-correction.md @@ -0,0 +1,72 @@ +# Phase 5c: Provider-Native Output Protocol And Contract Correction + +## Goal + +Close the Phase 5/5b third-pass provider-native findings while preserving the current safety stance: native streaming remains disabled for priority providers until each live stream path is proven, and LiteLLM remains fallback only for unsupported cases. + +## Scope + +- Return native provider responses in the originating client protocol instead of the provider protocol. +- Add explicit native endpoint/header/operation-selection hooks to `ProviderInterface`. +- Ensure Claude Code native requests always include required Anthropic `max_tokens`. +- Harden Claude Code API-key header behavior without breaking bearer/OAuth-style credentials. +- Fix Antigravity alias normalization so duplicated aliases do not collapse incorrectly and low/high thinking aliases can still affect request metadata. +- Centralize explicit-native streaming fail-closed behavior. +- Keep priority-provider native streaming disabled by default. +- Add focused tests for all blocker/high/medium findings. + +## Non-Goals + +- Do not enable live native streaming for Claude Code, Codex, Copilot, or Antigravity. +- Do not restore retired Antigravity device-profile/fingerprint behavior. +- Do not replace credential rotation, usage tracking, or routing. +- Do not implement all rich Codex Responses item conversions in this phase. +- Do not commit user-facing reports. + +## Implementation Plan + +1. Client target protocol in native context. + - Add `client_protocol_name` to `NativeProviderContext`. + - Default direct/library native use to provider protocol for backwards compatibility. + - Set `client_protocol_name="openai_chat"` from the chat-completions executor path. + - Format non-streaming native responses using the client protocol after provider protocol parsing. + - Format native stream events using the client protocol when a client protocol is set. + +2. Provider interface native contract. + - Add default `get_native_endpoint()` and `get_native_headers()` methods that raise clear `NotImplementedError`. + - Add `supports_native_operation()` defaulting to operation support through provider declarations. + - Add `should_use_native_protocol()` defaulting to true only when a provider declares a native protocol and operation is supported. + - Update executor checks to call hooks rather than relying on `hasattr()`. + +3. Claude Code hardening. + - `prepare_native_request()` ensures Anthropic `max_tokens` is present, using existing request value or `CLAUDE_CODE_MAX_TOKENS` default. + - Header selection supports bearer credentials and API-key credentials: + - `CLAUDE_CODE_AUTH_HEADER=bearer|x-api-key|auto` + - auto uses `x-api-key` for `sk-ant-*` style keys and bearer otherwise. + - Preserve `anthropic-version` and content type. + +4. Antigravity alias/thinking metadata. + - Fix alias-to-upstream mapping to use public alias map directly, not a lossy reverse map. + - Preserve thinking-level hints for `gemini-3-pro-low` / `gemini-3-pro-high` in request metadata before the upstream model is normalized. + - Keep model discovery output stable and prefixed. + +5. Native streaming fail-closed helper. + - Centralize explicit-native streaming unsupported handling so auto and explicit modes use the same fail-closed decision. + - Keep priority providers non-streaming native only. + +6. Tests. + - Native executor: provider protocol response -> OpenAI chat client response. + - Request executor: routed native Claude/Codex/Antigravity chat completions return OpenAI Chat shape. + - Claude Code: missing `max_tokens` gets defaulted and auth header modes behave correctly. + - Antigravity: low/high aliases normalize safely and preserve request metadata. + - ProviderInterface: native hooks exist and unsupported operation checks are explicit. + - Streaming: explicit native streaming fail-closed path uses centralized helper. + +## Acceptance Criteria + +- `/v1/chat/completions` native routes return OpenAI Chat response shape regardless of provider-native protocol. +- Claude Code native Messages requests include `max_tokens`. +- Antigravity model normalization does not lose low/high alias intent. +- Provider native endpoint/header/operation hooks are explicit on `ProviderInterface`. +- Explicit native streaming fails closed through one helper path. +- Focused tests pass and both reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-6-routing-fallback-groups.md b/docs/experimental/phase-6-routing-fallback-groups.md new file mode 100644 index 000000000..d546e8efc --- /dev/null +++ b/docs/experimental/phase-6-routing-fallback-groups.md @@ -0,0 +1,294 @@ +# Phase 6 Plan: Routing And Fallback Groups + +## Goal + +Add explicit ordered routing/fallback groups and a provider-native routing seam that can choose between the current LiteLLM execution path, provider custom logic, and the Phase 5 `NativeProviderExecutor`. Fallback groups are the priority deliverable: a model can resolve to an ordered chain of concrete targets, and the executor can try the next target on eligible failures without replacing the existing credential rotation, usage tracking, session tracking, or retry-after parsing. + +## Non-Goals + +- Do not replace `UsageManager`, `SelectionEngine`, `SessionTracker`, or retry-after parsing. +- Do not add full target-group selector language before fallback groups work. +- Do not implement multi-user/security isolation. +- Do not implement a complete config-file system; Phase 10 owns config polish. +- Do not route every provider natively by default. +- Do not fallback after visible streamed output has been emitted. +- Do not silently fallback to LiteLLM for native providers unless the fallback chain explicitly includes a LiteLLM target. +- Do not make Claude Code/Codex/Copilot/Antigravity live-production claims beyond mocked/tested routing behavior. + +## Current Code Context + +- `RequestExecutor` is the live retry/rotation path. +- `RequestExecutor` already gets provider plugin instances, usage managers, credentials, request context, transaction logger, and LiteLLM provider params. +- Phase 5 `NativeProviderExecutor` is isolated and opt-in. +- Providers can declare `protocol_name`, adapter names, field-cache rules, native endpoints, and native headers. +- Model resolution currently maps a single requested model to one provider/model pair. +- Existing retry loops rotate credentials inside one provider; Phase 6 must add model/provider target fallback around that without breaking per-provider credential rotation. +- Existing transaction trace can log routing decisions. +- Existing streaming handler can rotate credentials before output, but stream fallback after output must remain disallowed. + +## Files To Add + +- `src/rotator_library/routing/__init__.py` +- `src/rotator_library/routing/types.py` +- `src/rotator_library/routing/config.py` +- `src/rotator_library/routing/resolver.py` +- `src/rotator_library/routing/policy.py` +- `src/rotator_library/routing/executor.py` or `attempts.py` +- `tests/test_routing_config.py` +- `tests/test_fallback_resolver.py` +- `tests/test_fallback_policy.py` +- `tests/test_request_executor_fallback_groups.py` +- `tests/test_request_executor_native_routing.py` +- `tests/test_streaming_fallback_policy.py` + +## Files Likely To Touch + +- `src/rotator_library/core/types.py` to add optional routing/fallback fields to `RequestContext`. +- `src/rotator_library/client/request_builder.py` to resolve fallback groups before execution. +- `src/rotator_library/client/models.py` if model resolution needs a group-aware helper. +- `src/rotator_library/client/executor.py` to wrap existing per-provider retry loops in target attempts and call native executor for eligible targets. +- `src/rotator_library/client/rotating_client.py` only if passing route configuration into components requires it. +- `src/rotator_library/providers/provider_interface.py` only if a small native-support method is needed. +- `src/rotator_library/transaction_logger.py` only if adding route correlation helpers is cleaner. + +## Routing Types + +`RouteTarget`: + +- `name`: stable target identifier. +- `provider`: provider key. +- `model`: concrete model name, with provider prefix optional but normalized. +- `protocol`: optional override; defaults to provider declaration. +- `execution`: `auto`, `native`, `custom`, or `litellm_fallback`. +- `priority`: optional numeric order. +- `weight`: future target-group selector support, ignored for ordered fallback. +- `conditions`: optional metadata for later selectors. +- `metadata`. + +`FallbackGroup`: + +- `name`. +- `targets`: ordered `RouteTarget` list. +- `failover_on`: error categories that permit next-target fallback. +- `stop_on`: error categories that must stop. +- `streaming_policy`: e.g. `pre_output_only`. +- `max_targets`: optional guard. +- `metadata`. + +`RoutingDecision`: + +- `requested_model`. +- `group_name` optional. +- `targets`. +- `selected_target_index`. +- `reason`. + +`RouteAttemptResult`: + +- target, success/failure, error classification, emitted_output flag, usage summary. + +## Configuration Plan + +- Phase 6 supports env-first fallback definitions with a small parser. +- Optional JSON config can be parsed if simple, but deeper config merging belongs to Phase 10. +- Env examples: + - `FALLBACK_GROUPS=sonnet_chain,codex_chain` + - `FALLBACK_GROUP_SONNET_CHAIN=claude_code/claude-sonnet-4-5,copilot/claude-sonnet-4-5,anthropic/claude-3-5-sonnet-latest` + - `FALLBACK_GROUP_CODEX_CHAIN=codex/gpt-5.1-codex,openai/gpt-5.1` + - `MODEL_ROUTE_CLAUDE_SONNET=group:sonnet_chain` + - `MODEL_ROUTE_CODEX=group:codex_chain` +- Also allow programmatic construction in tests. +- Config parser must validate: + - group names are unique. + - targets have provider/model. + - empty chains are invalid. + - cycles through group aliases are rejected. + - provider names are not silently guessed when ambiguous. +- No secrets in route config. + +## Target Resolution + +- If requested model maps to a fallback group, return all group targets in order. +- If requested model is already `provider/model`, return a single target. +- If requested model maps through existing model definitions, preserve existing behavior. +- Target model should be normalized with provider prefix before entering `RequestExecutor`. +- Record transform trace pass `routing_decision` with requested model, group, target count, selected target names, execution modes, and no credentials. + +## Fallback Policy + +Try next target only on eligible failures: + +- rate limit/quota/capacity. +- provider unavailable/server error. +- transient connection errors. +- native provider unsupported operation when explicit fallback target exists. + +Do not fallback on: + +- auth errors for all credentials unless the group target explicitly marks auth fallback safe. +- validation/permanent request errors. +- pre-request callback failures. +- client cancellation. +- streaming errors after visible output. + +Respect existing per-provider credential rotation before moving to the next target: + +- one target attempt should let current `RequestExecutor` exhaust eligible credentials according to existing retry logic. +- after a target is exhausted, the fallback layer can advance to the next target. + +Keep error accumulator context across targets so final failure includes all target errors. + +## Native, Custom, And LiteLLM Execution Choice + +`execution=auto`: + +- if plugin has `has_custom_logic()`, use current plugin `acompletion()`. +- else if provider has native protocol declaration and native endpoint/header methods, use `NativeProviderExecutor`. +- else use current LiteLLM call path. + +`execution=native`: require native support or fail this target with a classified unsupported-operation error. + +`execution=custom`: require `has_custom_logic()`. + +`execution=litellm_fallback`: force current LiteLLM path and emit `native_provider_litellm_fallback` / `routing_litellm_fallback` trace metadata. + +Tests must prove native-declared providers do not silently use LiteLLM unless the target says `litellm_fallback`. + +## RequestExecutor Integration + +- Keep the existing `_execute_non_streaming()` and `_execute_streaming()` credential retry loops as the per-target implementation. +- Add a wrapper method that accepts a `RequestContext` with `routing_targets`. +- For each target: + - clone/update context provider/model/usage manager key/credentials for that target. + - record `routing_target_attempt_started`. + - run the existing target execution path. + - on success record `routing_target_attempt_succeeded` and return. + - on failure classify and ask fallback policy if next target is allowed. + - record `routing_target_attempt_failed`. +- Avoid mutating the original request context in-place. +- If no group exists, execute exactly as today. + +## Native Provider Execution Integration + +- Add a small execution branch inside the target attempt where the selected target resolves to native execution. +- Construct `NativeProviderContext` from: + - provider plugin declarations. + - selected model. + - credential identifier/token. + - request context session/scope/classifier. + - transaction logger. + - endpoint/headers from provider methods. +- Use Phase 5 `NativeProviderExecutor` for mocked/tested providers. +- Preserve usage recording through the existing credential context. If native responses do not provide full usage yet, record normalized available usage only and leave cost normalization for Phase 9. +- For custom providers with `has_custom_logic()`, keep existing plugin `acompletion()` branch. +- For undeclared providers, keep LiteLLM path. + +## Streaming Integration + +- Fallback to next target is allowed only before any visible chunk is yielded. +- Track `emitted_output` in the stream wrapper. +- If a stream errors before output, fallback policy may advance to next target. +- If a stream errors after output, propagate the error; do not switch models mid-stream. +- Record: + - `routing_stream_target_attempt_started` + - `routing_stream_target_attempt_failed` + - `routing_stream_target_attempt_succeeded` + - `routing_stream_fallback_blocked_after_output`. +- Keep current stream retry policy for same-target/same-provider errors before group fallback. + +## Transaction Trace Requirements + +- `routing_decision` +- `routing_target_attempt_started` +- `routing_target_attempt_failed` +- `routing_target_attempt_succeeded` +- `routing_fallback_selected` +- `routing_fallback_exhausted` +- `routing_litellm_fallback` +- `routing_native_execution_selected` +- stream equivalents where applicable. +- Trace metadata includes group, target index, provider, model, execution mode, classification, and reason. It must not include raw credentials. + +## Target Groups As Optional Richer Layer + +- Phase 6 focuses on ordered fallback groups. +- Add type seams for future target groups: + - `TargetSelector` + - `TargetGroup` + - `TargetSelectionPolicy` +- Do not implement weighted/latency/cost selector behavior unless it is trivial and isolated. +- Document that Phase 6 fallback groups are deterministic ordered chains. + +## Tests + +Config tests: + +- parse env fallback group. +- invalid empty group. +- duplicate group names. +- provider/model parsing. +- explicit `litellm_fallback` execution target. + +Resolver tests: + +- requested alias maps to group. +- provider/model remains single target. +- target order is preserved. +- existing model prefix behavior stays intact. + +Policy tests: + +- fallback on rate limit/server/transient. +- stop on auth/permanent/pre-request/cancel. +- streaming fallback allowed before output. +- streaming fallback blocked after output. + +Request executor tests: + +- first target success uses no fallback. +- first target rate limits all credentials then second target succeeds. +- final error includes both target failures. +- no fallback group preserves old behavior. +- transaction trace includes routing passes. + +Native routing tests: + +- native-declared provider target uses `NativeProviderExecutor`. +- native-declared provider does not use LiteLLM unless target execution is `litellm_fallback`. +- custom provider still uses plugin `acompletion()`. +- undeclared provider still uses LiteLLM. + +Streaming tests: + +- pre-output target failure falls through to next target. +- post-output failure does not fall through. +- trace records blocked fallback after output. + +Regression: + +- Phase 1 protocol tests. +- Phase 2 logging tests. +- Phase 3 adapter/cache tests. +- Phase 4 Responses tests. +- Phase 5 provider tests. +- `test_session_tracking.py`. +- `test_selection_engine.py`. + +## Commit Checkpoints + +1. Add routing types, config parser, resolver, and policy with tests. +2. Add request context routing fields and target cloning helpers with tests. +3. Integrate non-streaming fallback group wrapper around existing executor path with tests. +4. Integrate native/custom/LiteLLM execution mode selection with tests. +5. Integrate streaming pre-output fallback policy with tests. +6. Add trace pass coverage and run regressions. +7. Review with `explore` and `explore-heavy`, fix findings, and write the uncommitted Phase 6 report. + +## Risks And Mitigations + +- Fallback could double-spend or corrupt usage. Mitigation: each target uses existing credential context and usage manager; group fallback starts only after target failure. +- Fallback could mask permanent request bugs. Mitigation: policy stops on validation/permanent errors. +- Streaming fallback could corrupt client output. Mitigation: disallow fallback after visible output. +- Native routing could accidentally bypass LiteLLM behavior for existing providers. Mitigation: default no group/no native route is unchanged; native path is explicit/declared. +- Config could become too broad. Mitigation: small env parser now; richer JSON config later in Phase 10. +- Session affinity could leak across provider pools. Mitigation: keep `SessionTracker` namespace behavior and clone target context with provider/model-specific namespace where needed. diff --git a/docs/experimental/phase-6b-routing-fallback-correctness.md b/docs/experimental/phase-6b-routing-fallback-correctness.md new file mode 100644 index 000000000..3d53895bd --- /dev/null +++ b/docs/experimental/phase-6b-routing-fallback-correctness.md @@ -0,0 +1,91 @@ +# Phase 6b: Routing/Fallback Correctness And Structured Error Safety + +## Goal + +Correct the Phase 6 audit findings. Phase 6 built ordered fallback groups and live executor wrappers, but fallback still needs stronger safety around hard-stop errors, structured error responses, streaming policy enforcement, and sanitized target summaries. + +## Non-Goals + +- Do not replace `UsageManager`, `SessionTracker`, retry-after parsing, or credential rotation behavior. +- Do not implement rich target-group selector syntax beyond ordered fallback chains. +- Do not add security or multi-user features. +- Do not make native streaming work; unsupported native streaming remains fail-closed. +- Do not include raw provider messages or credentials in cross-target summaries. +- Do not commit user-facing reports. + +## Current State + +- `FallbackPolicy.should_fallback()` uses group `failover_on` and `stop_on` directly, so a group can currently override auth/permanent errors into fallback eligibility. +- `_route_error_type_from_response()` only recognizes retry-like summaries and a few proxy error types; it does not robustly inspect structured status/code/details fields. +- `FallbackGroup.streaming_policy` exists but live streaming fallback does not use it. +- Streaming fallback correctly tracks visible output, but needs stronger tests for error/control frames and `never` streaming policy. +- Non-streaming fallback summaries are structural, but call sites still pass raw exception strings that should be avoided entirely. + +## Implementation Plan + +1. Add hard-stop route error categories. + - Add a non-overridable hard-stop set for auth, forbidden, invalid request, context-window, credential reauth, pre-request callback, cancellation, and configuration errors. + - Document that these are safety boundaries and group policies cannot opt into cross-target fallback for them. + +2. Normalize routing policy vocabulary. + - Add `normalize_route_error_type()` for aliases such as `auth`, `permission_denied`, `bad_request`, `validation`, `transient`, `network`, and `configuration`. + - Use it in policy, routing runner, retry helper, and executor route classification. + +3. Make hard stops win in `FallbackPolicy`. + - Streaming visible output still blocks fallback first. + - Hard-stop categories return false before group policy evaluation. + - Group stop/failover sets are normalized before matching. + +4. Validate routing group policy. + - Reject configured `failover_on` entries that normalize to hard-stop categories. + - Parse and validate `streaming_policy` from JSON routing config. + - Preserve environment override precedence. + +5. Enforce `FallbackGroup.streaming_policy` in live executor paths. + - `never` prevents streaming fallback even before output. + - `pre_output_only` remains the default. + - Trace metadata should include the policy used. + +6. Harden structured response classification. + - Inspect `error.type`, `error.code`, `error.status`, `error.details.status_code`, and detail classification fields before summaries. + - Hard-stop signals win over retryable text. + - Keep retryable classification for 429/quota/rate-limit, 5xx/server, timeout, and connection signals. + +7. Sanitize target-failure summaries. + - Keep raw messages out of fallback details and traces. + - Remove unnecessary raw string arguments from summary call sites. + - Add tests proving secrets/provider text do not appear. + +8. Improve stream fallback frame handling. + - Verify error/control frames are non-visible output. + - Test `event: error`, `type: error`, `response.failed`, `[DONE]`, comments, and heartbeat-like frames. + - Ensure visible text/tool deltas still block fallback. + +9. Add deterministic exhaustion metadata. + - Include sanitized target summaries in routing exhaustion traces for stream and non-stream paths. + - Avoid raw exception text in trace metadata. + +10. Add execution-mode safety tests. + - Explicit native configuration errors must be hard stops. + - Unsupported operation behavior must be clear and tested. + - LiteLLM fallback remains available for retryable errors. + +## Tests + +- `tests/test_fallback_policy.py` +- `tests/test_retry_policy.py` +- `tests/test_routing_config.py` +- `tests/test_fallback_resolver.py` +- `tests/test_routing_executor.py` +- `tests/test_streaming_fallback_policy.py` +- `tests/test_request_executor_native_routing.py` +- Phase 5b provider/native regression subset if executor routing helpers change. + +## Acceptance Criteria + +- Auth, forbidden, invalid request, context-window, credential reauth, pre-request callback, cancellation, and configuration errors never fallback across targets. +- Structured error responses classify deterministically with hard-stop signals taking precedence. +- Streaming fallback respects `FallbackGroup.streaming_policy`. +- Error/control stream frames before visible output do not accidentally lock routing. +- Cross-target summaries and traces remain sanitized. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-6c-routing-fallback-correction.md b/docs/experimental/phase-6c-routing-fallback-correction.md new file mode 100644 index 000000000..4df2eb01f --- /dev/null +++ b/docs/experimental/phase-6c-routing-fallback-correction.md @@ -0,0 +1,62 @@ +# Phase 6c: Routing/Fallback Correctness And Stale Test Repair + +## Goal + +Close all Phase 6/6b third-pass routing findings while preserving the fallback group architecture and the hard-stop safety from Phase 6b. + +## Scope + +- Replace stale `tests/test_fallback_groups.py` with tests for the current routing package. +- Implement requested-model-in-group promotion. +- Align streaming execution-mode behavior with non-streaming behavior for `@custom`, `@native`, and `@litellm_fallback`. +- Ensure explicit native configuration errors are hard-stop configuration errors, not retryable unsupported-operation errors. +- Expand structured route-error alias normalization. +- Adjust fallback target session namespace per target. +- Add focused tests for every blocker/high/medium. + +## Non-Goals + +- Do not resurrect the old `FallbackGroupManager`. +- Do not reintroduce `RotatingClient(model_fallback_groups=...)`. +- Do not implement routing for embeddings unless it is trivial and needed by tests. +- Do not weaken hard-stop policy. +- Do not commit user-facing reports. + +## Implementation Plan + +1. Replace stale fallback tests. + - Rewrite `tests/test_fallback_groups.py` around `routing.config`, `FallbackResolver`, `FallbackAttemptRunner`, and executor helpers. + - Cover env/JSON group parsing, target promotion, execution suffix parsing, and fallback exhaustion behavior. + +2. Requested-model promotion. + - Promote a requested provider/model target to the first attempt when it appears in a fallback group. + - Preserve the relative order of all other targets. + - Apply this both for explicit `MODEL_ROUTE_* = group:name` routes and provider-prefixed requests that appear in any group. + +3. Streaming execution-mode parity. + - Share execution-mode selection logic between streaming and non-streaming paths. + - `@litellm_fallback` must force LiteLLM. + - `@custom` must require custom provider logic. + - `@native` must fail closed if native streaming is unsupported. + - `auto` remains custom first, then native streaming if explicitly supported, then LiteLLM. + +4. Explicit native configuration errors. + - Missing native declaration, unsupported operation, and missing endpoint/header helpers should raise `RoutingExecutionError(error_type="configuration_error")` for explicit native execution failures. + - Hard-stop policy then blocks fallback. + +5. Structured route-error aliases. + - Add common status-less aliases including `invalid_api_key`, `unauthorized`, `invalid_argument`, `rate_limited`, `too_many_requests`, `resource_exhausted`, `unavailable`, `deadline_exceeded`, and context-window variants. + +6. Session namespace adjustment. + - `clone_context_for_target()` should rewrite standard session-tracking namespaces to the target provider/model instead of preserving the first target namespace. + - Unknown/custom namespace shapes remain unchanged. + +## Acceptance Criteria + +- Broad test collection no longer fails on stale fallback imports. +- Requested provider-prefixed models inside a group are tried first. +- Streaming and non-streaming execution-mode selection are consistent. +- Explicit native config errors are hard-stop configuration errors. +- Common structured aliases normalize correctly. +- Fallback target contexts do not reuse the first target's provider/model session namespace. +- Focused tests pass and both reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-7-retry-cooldown-failover.md b/docs/experimental/phase-7-retry-cooldown-failover.md new file mode 100644 index 000000000..ffcfba144 --- /dev/null +++ b/docs/experimental/phase-7-retry-cooldown-failover.md @@ -0,0 +1,253 @@ +# Phase 7 Plan: Retry/Cooldown/Failover Cleanup + +## Goal + +Harden the retry and failover layer now that ordered routing exists. Phase 7 makes provider-level cooldown functional, reduces fallback policy duplication, preserves the current retry-after parser, and improves final error visibility across target chains without replacing `UsageManager`, `SessionTracker`, credential rotation, or the existing retry classifier. + +## Non-Goals + +- Do not replace `UsageManager`, `SelectionEngine`, `SessionTracker`, or the retry-after parser. +- Do not rewrite the entire executor. +- Do not implement the Phase 8 streaming transport/library overhaul. +- Do not implement native streaming dispatch unless a small safe seam is necessary. +- Do not add SQLite or any database. +- Do not implement full JSON config merging; Phase 10 owns config polish. +- Do not make fallback happen after visible streamed output. +- Do not silently fallback to LiteLLM unless the route target explicitly selects `litellm_fallback`. + +## Current Code Context + +- `CooldownManager` exists and `_wait_for_cooldown()` is called before target attempts, but `start_cooldown()` is not called by the executor, so provider cooldown is effectively dormant. +- Existing retry-after parsing lives in `error_handler.py` and is strong; keep it. +- Existing `RequestExecutor._handle_error_with_context()` handles non-streaming per-credential retries and rotations. +- Streaming has duplicated retry/rotation handling in several exception branches. +- Phase 6 added `FallbackPolicy`, `FallbackAttemptRunner`, `clone_context_for_target()`, and executor inline fallback loops. +- Phase 6 final review found no blockers but noted two useful cleanup targets: executor inline loops do not pass `FallbackGroup` overrides to `FallbackPolicy`, and `FallbackAttemptRunner` is tested but not used by live executor. +- Phase 6 response-based fallback maps exhausted structured errors to retryable categories, but final target-chain failures do not yet summarize all target failures. +- Provider-level cooldown should apply to provider-wide or IP/global throttling, not every small per-key retry-after. + +## Files To Add + +- `src/rotator_library/retry_policy.py` +- `tests/test_retry_policy.py` +- `tests/test_cooldown_activation.py` +- `tests/test_request_executor_fallback_error_summary.py` +- Possibly `tests/test_request_executor_group_policy.py` + +## Files Likely To Touch + +- `src/rotator_library/client/executor.py` +- `src/rotator_library/cooldown_manager.py` +- `src/rotator_library/routing/types.py` +- `src/rotator_library/routing/executor.py` +- `src/rotator_library/core/types.py` +- `src/rotator_library/error_handler.py` only if docstring/type comments need alignment; do not weaken parser behavior. +- Existing Phase 6 tests if group policy wiring changes expected trace order. + +## Retry Policy Foundation + +Add a small `retry_policy.py` module to centralize decisions that are currently duplicated: + +- `classify_route_error(error, provider)` +- `should_provider_cooldown(classified_error, *, small_cooldown_threshold, provider_cooldown_threshold)` +- `provider_cooldown_duration(classified_error, default_duration)` +- `should_retry_same_credential(classified_error, small_cooldown_threshold)` +- `should_rotate_credential(classified_error)` +- `is_target_failover_eligible(error_type, group=None, stream=False, emitted_output=False)` + +This module should call existing `classify_error()`, `should_retry_same_key()`, and `should_rotate_on_error()` rather than reimplementing them. + +It should document why small retry-after values stay on the same credential while large/provider-level values can activate provider cooldown. + +It should not own sleeping or mutation; it only returns decisions. + +## Cooldown Activation + +Add provider cooldown activation in non-streaming error handling: + +- If `classified.retry_after` exists and is above `SMALL_COOLDOWN_RETRY_THRESHOLD`, start provider cooldown for that duration when the error is provider-wide enough. +- Candidate categories: `rate_limit`, `server_error`, and possibly provider-specific `quota_exceeded` only when retry-after suggests global reset rather than per-key quota. +- Be conservative: avoid cooling down an entire provider for every credential quota if the error looks per-credential. + +Add env knobs: + +- `PROVIDER_COOLDOWN_MIN_SECONDS` default perhaps same as small-cooldown threshold or a modest value. +- `PROVIDER_COOLDOWN_DEFAULT_SECONDS` for retryable provider-wide errors without retry-after. +- `PROVIDER_COOLDOWN_ON_QUOTA` default false unless evidence indicates provider-wide quota. + +Existing `_wait_for_cooldown()` stays the wait path. + +`CooldownManager.start_cooldown()` should extend only when the new expiry is later than the current expiry, not shorten an active cooldown. + +Add trace pass: + +- `provider_cooldown_started` +- metadata: provider, duration, error_type, retry_after_present, reason. + +Logging failure to start cooldown should never fail the request. + +## Fallback Runner And Live Loop Cleanup + +Either wire `FallbackAttemptRunner` into live non-streaming and streaming wrappers or keep inline wrappers but pass a resolved group object. + +Minimal preferred approach for Phase 7: + +- Add `routing_group` optional field to `RequestContext` or attach group policy data to context. +- `RequestContextBuilder` should carry the `FallbackGroup` resolved from config when routing is active. +- Executor fallback wrappers pass `group=context.routing_group` to `FallbackPolicy.should_fallback()`. +- Keep inline wrappers if that avoids contorting trace behavior. + +If refactoring to `FallbackAttemptRunner` stays small, do it; otherwise leave runner as tested support and make inline loops group-aware. + +Add tests proving group overrides are honored in live executor wrappers for non-streaming and streaming. + +## Cross-Target Error Summary + +Add a target failure accumulator for fallback groups: + +- target name +- provider +- model +- execution mode +- error type +- message summary +- emitted_output for streams + +If all fallback targets fail, final client error should include `fallback_targets` details under `error.details`, with no credentials. + +Preserve existing per-target credential error summaries from `RequestErrorAccumulator`. + +Do not expose credential secrets or raw request data. + +Add trace pass: + +- `routing_fallback_exhausted` +- include accumulated target failures. + +Tests should cover: + +- all targets fail and final response contains both target failures. +- exception failure on first target and structured proxy error on second target are both summarized. +- no credentials in summary. + +## Provider Cooldown And Fallback Interaction + +- If a target enters provider cooldown and still fails/exhausts, fallback to next target may proceed if policy allows. +- If a provider is already cooling down, `_wait_for_cooldown()` should wait only if there is request budget; otherwise target should fail fast in a way fallback can interpret. +- Consider returning/raising a classified `provider_cooldown_budget_exceeded` or mapping to `rate_limit` for target fallback. +- Keep this minimal: do not replace capacity waiting or usage limits. + +## Streaming Retry/Fallback Cleanup + +- Keep Phase 6 invariant: fallback only before visible output. +- Deduplicate obvious streaming error branches only if the change is small and tests cover it. +- Add cooldown activation in streaming error handling where retry-after is available and no visible output has been emitted. +- Do not implement native streaming execution-mode dispatch in Phase 7 unless reviewers insist; Phase 8 owns streaming transport design. +- Add provider cooldown trace for streaming provider-wide throttles. + +## Error-Classification Alignment + +Add tests tying actual `classify_error()` output to fallback policy decisions for: + +- `rate_limit` +- `quota_exceeded` +- `server_error` +- `api_connection` +- `authentication` +- `forbidden` +- `invalid_request` +- `context_window_exceeded` +- `credential_reauth_needed` +- `pre_request_callback_error` +- `cancelled` +- `unsupported_operation` + +Preserve conservative `unknown` behavior unless a test shows a strong reason to fallback. + +## Stale Fallback Test Cleanup + +There is a stale ignored test file `tests/test_fallback_groups.py` that imports a non-existent older module. + +- Do not delete it unless it is tracked and relevant to current runs. +- If it is tracked or causing collection failures in broader test runs, replace/archive it with current routing tests in a focused commit. +- If ignored/untracked, leave it alone unless it blocks CI. + +## Transform Trace Requirements + +Keep existing Phase 6 traces: + +- `routing_decision` +- `routing_target_attempt_started` +- `routing_target_attempt_failed` +- `routing_target_attempt_succeeded` +- `routing_fallback_selected` +- `routing_fallback_exhausted` +- `routing_litellm_fallback` +- `routing_native_execution_selected` +- stream equivalents. + +Add: + +- `provider_cooldown_started` +- `provider_cooldown_skipped` +- `routing_group_policy_applied` if group overrides influence a decision. + +Trace metadata must not include credentials or secrets. + +## Testing Plan + +Retry policy tests: + +- existing classifier output maps to expected retry/rotate/fallback/cooldown decisions. +- small retry-after chooses same-credential retry, not provider cooldown. +- large retry-after can start provider cooldown. +- `unknown` remains conservative for target fallback. + +Cooldown tests: + +- `CooldownManager.start_cooldown()` extends but does not shorten cooldown. +- non-streaming large retry-after starts provider cooldown. +- small retry-after does not start provider cooldown. +- cooldown trace emitted. +- `_wait_for_cooldown()` respects deadline budget. + +Fallback group policy tests: + +- live non-streaming fallback honors group-specific `failover_on`. +- live non-streaming fallback honors group-specific `stop_on`. +- live streaming fallback honors group policy before output. + +Error summary tests: + +- all target failures are summarized in final error details. +- target summary excludes credentials. + +Regression tests: + +- Phase 1 protocol tests. +- Phase 2 transform logging tests. +- Phase 3 adapter/field-cache tests. +- Phase 4 Responses tests. +- Phase 5 provider/native tests. +- Phase 6 routing tests. +- `test_session_tracking.py`. +- `test_selection_engine.py`. + +## Commit Checkpoints + +1. Add retry policy helper module and classifier-alignment tests. +2. Harden `CooldownManager` extension semantics and tests. +3. Wire provider cooldown activation into non-streaming executor with trace tests. +4. Add group-policy awareness to live fallback wrappers with tests. +5. Add cross-target error summaries with tests. +6. Add streaming cooldown/group-policy cleanup tests if not covered earlier. +7. Run focused and regression tests. +8. Review with `explore` and `explore-heavy`; fix findings; write uncommitted Phase 7 report. + +## Risks And Mitigations + +- Provider cooldown could over-throttle healthy credentials. Mitigation: only activate on large retry-after/provider-wide categories, keep quota cooldown disabled by default unless configured. +- Refactoring fallback loops could break trace order. Mitigation: prefer small group-aware changes unless runner integration stays simple. +- Error summaries could leak credentials. Mitigation: summarize target/provider/model/error only; never include credentials. +- Streaming fallback could corrupt output if changed carelessly. Mitigation: preserve Phase 6 visible-output gate and add tests. +- Retry policy could drift from existing parser. Mitigation: call existing parser/classifier helpers instead of rewriting them. diff --git a/docs/experimental/phase-7b-retry-cooldown-backoff.md b/docs/experimental/phase-7b-retry-cooldown-backoff.md new file mode 100644 index 000000000..5ceb14ff7 --- /dev/null +++ b/docs/experimental/phase-7b-retry-cooldown-backoff.md @@ -0,0 +1,81 @@ +# Phase 7b: Retry/Cooldown Backoff And Failure History + +## Goal + +Correct the Phase 7 audit findings. Phase 7 made provider cooldown activation real and centralized retry/failover decisions, but it still lacks provider/model cooldown scopes, scoped backoff, and structured in-memory failure history. + +## Non-Goals + +- Do not replace `UsageManager`, credential cooldowns, usage windows, `classify_error()`, or retry-after parsing. +- Do not introduce SQLite or a new persistence database. +- Do not make native streaming safe; unsupported native streaming remains fail-closed. +- Do not change Phase 6b fallback hard-stop behavior. +- Do not commit user-facing reports. + +## Current State + +- `CooldownManager` tracks provider cooldowns only. +- `RequestExecutor._wait_for_cooldown()` waits on provider cooldown only. +- `retry_policy.decide_provider_cooldown()` returns provider-level decisions without model scope or failure-history context. +- `MODEL_CAPACITY_EXHAUSTED` is noticed in logs but still behaves like a generic provider server error. +- There is no structured provider/model failure-history ring for future observability and bounded backoff decisions. + +## Implementation Plan + +1. Extend `CooldownManager` with scoped cooldown methods. + - Preserve `start_cooldown()`, `is_cooling_down()`, and `get_remaining_cooldown()`. + - Add provider/model scoped methods and extend-only semantics per scope key. + - Provider cooldown blocks all models; model cooldown blocks only that provider/model. + +2. Update executor cooldown waiting. + - Pass model into cooldown wait paths. + - Wait for the max of provider and model cooldown remaining. + - Trace waits without credentials or raw provider text. + +3. Add model-capacity detection. + - Detect `MODEL_CAPACITY_EXHAUSTED`, model capacity, and capacity-exhausted signals from exceptions and dict payloads. + - Keep `error_type="server_error"` for compatibility, but choose model cooldown scope. + +4. Extend cooldown decisions. + - Add `scope`, `model`, and optional backoff metadata to `ProviderCooldownDecision`. + - Large retry-after rate limits stay provider-scoped. + - Model-capacity failures become model-scoped. + - Quota cooldown remains disabled by default. + +5. Add in-memory failure history. + - Bounded ring with timestamp, provider, model, error type, scope, duration, and reason. + - No disk persistence. + - Executor records successful cooldown starts for future observability and backoff. + +6. Add bounded repeated-transient backoff. + - Track repeated transient `server_error`/`api_connection` failures within a configurable window. + - Escalate cooldown duration conservatively up to a max. + - Do not backoff hard-stop categories. + +7. Wire scoped cooldown start in `RequestExecutor`. + - Use scoped cooldown methods when available and fall back to provider-only fakes in tests. + - Keep existing trace pass names with added scope/model metadata. + +8. Update streaming error decisions. + - Keep `decide_streaming_error_action()` side-effect-free. + - Include cooldown scope/model in the decision. + - Visible output still suppresses provider/model cooldown. + +## Tests + +- `tests/test_cooldown_activation.py` +- `tests/test_retry_policy.py` +- `tests/test_streaming_error_handler.py` +- Executor-focused cooldown/trace tests. +- Phase 6b routing regression subset. + +## Acceptance Criteria + +- Provider cooldown behavior remains backwards-compatible and extend-only. +- Model cooldowns do not block unrelated models on the same provider. +- Provider cooldowns still block all models for that provider. +- Model-capacity errors produce model-scoped cooldown/backoff. +- Repeated transient failures can produce bounded scoped backoff from in-memory history. +- Cooldown starts/waits are traceable and sanitized. +- Streaming cooldown decisions expose scope and respect visible-output blocking. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-7c-retry-cooldown-correction.md b/docs/experimental/phase-7c-retry-cooldown-correction.md new file mode 100644 index 000000000..8fba3b192 --- /dev/null +++ b/docs/experimental/phase-7c-retry-cooldown-correction.md @@ -0,0 +1,57 @@ +# Phase 7c: Cooldown Budget, Transient Backoff, And Attempt History + +## Goal + +Close all Phase 7/7b third-pass retry/cooldown findings while preserving existing retry-after parsing and credential-rotation semantics. + +## Scope + +- Fail fast when an active provider/model cooldown exceeds the remaining request deadline budget. +- Prevent a single generic `server_error` / `api_connection` without `retry_after` from starting provider-wide cooldown. +- Record repeated transient failures into `FailureHistory` even when the cooldown threshold has not yet been reached. +- Clear or reduce failure history on successful provider/model calls. +- Populate structured `routing_attempt_history` for live non-streaming and streaming fallback attempts. +- Add tests for cooldown-over-budget, single transient no-cooldown, repeated transient cooldown, success reset, model-scoped cooldown isolation, and stream parity. + +## Non-Goals + +- Do not replace the retry-after parser. +- Do not replace `UsageManager` or `SessionTracker`. +- Do not introduce durable failure-history storage. +- Do not change small retry-after same-credential behavior. +- Do not weaken fallback hard-stop policy from Phase 6c. + +## Implementation Plan + +1. Cooldown-over-budget fail-fast. + - Raise a retryable routing error instead of returning when cooldown remaining exceeds the request deadline budget. + - Trace `cooldown_wait_exceeds_budget` without exposing credentials or provider payloads. + +2. Generic transient cooldown threshold. + - No-`retry_after` `server_error` and `api_connection` should start cooldown only after `FailureHistory.backoff_for()` crosses the configured threshold. + - Large explicit `retry_after` behavior remains unchanged. + +3. Record skipped transient failures. + - Record sanitized transient entries when cooldown is skipped because the backoff threshold is not yet met. + - Keep history bounded and in-memory only. + +4. Success reset. + - Add a clear/reset helper on `FailureHistory`. + - Clear matching provider/model transient entries after successful non-streaming and completed streaming calls. + +5. Structured attempt history. + - Append sanitized attempt entries for fallback failures and successes into `RequestContext.routing_attempt_history`. + - Include error type, target identity, execution mode, output visibility, status code when available, fallback decision, and timing where cheap. + +6. Model cooldown isolation. + - Add tests proving model-scoped cooldowns block only the matching model while provider-scoped cooldowns block all provider models. + +## Acceptance Criteria + +- Cooldown that exceeds request budget does not silently allow execution. +- A single generic transient without retry-after does not start provider-wide cooldown. +- Repeated transient failures can still start bounded cooldown/backoff. +- Successful calls clear matching failure-history entries. +- Live routing attempt history is populated and sanitized. +- Model-scoped cooldown blocks only its model plus provider-wide cooldown still blocks all provider models. +- Focused tests pass and both reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-8-streaming-library-upgrade.md b/docs/experimental/phase-8-streaming-library-upgrade.md new file mode 100644 index 000000000..86c1c7fbc --- /dev/null +++ b/docs/experimental/phase-8-streaming-library-upgrade.md @@ -0,0 +1,329 @@ +# Phase 8 Plan: Streaming Library Upgrade + +## Goal + +Turn streaming into a reusable, transport-aware library layer instead of scattered executor/provider-specific logic. Phase 8 should preserve the current HTTP SSE behavior, add measured stream lifecycle events, consolidate retry/error handling where safe, expose a WebSocket-ready transport seam, and prepare native provider streaming to share the same policies. It must keep the Phase 6/7 invariants: no fallback after visible output, provider cooldown only before visible output or error-only chunks, and no replacement of `UsageManager`, `SessionTracker`, or retry-after parsing. + +## Non-Goals + +- Do not replace `UsageManager`, `SessionTracker`, `SelectionEngine`, or retry-after parsing. +- Do not implement a live WebSocket route yet unless it is strictly a disabled/placeholder seam. +- Do not migrate every provider to native streaming in one pass. +- Do not alter non-streaming request behavior. +- Do not add SQLite or any database. +- Do not implement full cost/quota overhaul; Phase 9 owns usage/quota/cost. +- Do not change client-visible SSE formats except to fix bugs or add compatible error/end events. +- Do not fallback to another target after visible model output. + +## Current Code Context + +- `RequestExecutor._execute_streaming()` owns the active streaming credential retry loop and still has four duplicated exception branches. +- `RequestExecutor._execute_streaming_with_fallback()` wraps target fallback and already tracks visible output. +- Phase 7 added `_stream_chunk_is_visible_output()` and `_can_start_stream_provider_cooldown()` safety helpers. +- `StreamingHandler.wrap_stream()` converts LiteLLM chunks to chat-completions SSE strings, records usage after completion, records response anchors, and handles client disconnect. +- `stream_retry_policy.py` owns reasoning-only retry safety for post-output stream failures. +- `ResponsesService.stream_response()` converts chat stream chunks into Responses SSE events, stores final streamed responses, and has a `ResponsesWebSocketFormatter` placeholder. +- `NativeProviderExecutor.stream()` can stream mocked native provider JSON-line chunks and trace raw/parsed/formatted events, but live executor streaming does not route through native execution modes yet. +- Transform trace already has stream pass names from Phase 2 and later phases. +- There is no central stream metrics object for TTFB/TTFT/stalls/cancellation. + +## Files To Add + +- `src/rotator_library/streaming/__init__.py` +- `src/rotator_library/streaming/events.py` +- `src/rotator_library/streaming/metrics.py` +- `src/rotator_library/streaming/transport.py` +- `src/rotator_library/streaming/policy.py` +- `src/rotator_library/streaming/errors.py` +- `tests/test_stream_events.py` +- `tests/test_stream_metrics.py` +- `tests/test_stream_transport.py` +- `tests/test_stream_policy.py` +- `tests/test_streaming_error_handler.py` +- `tests/test_request_executor_stream_metrics.py` +- `tests/test_native_streaming_transport_seam.py` + +## Files Likely To Touch + +- `src/rotator_library/client/executor.py` +- `src/rotator_library/client/streaming.py` +- `src/rotator_library/client/stream_retry_policy.py` +- `src/rotator_library/native_provider/executor.py` +- `src/rotator_library/native_provider/streaming.py` +- `src/rotator_library/responses/streaming.py` +- `src/rotator_library/responses/service.py` +- `src/proxy_app/main.py` only if route headers/cancellation handling need a small SSE-compatible fix. + +## Streaming Event Model + +Add `StreamEvent` dataclass: + +- `event_type`: `started`, `raw_chunk`, `parsed_chunk`, `delta`, `reasoning_delta`, `tool_delta`, `usage`, `error`, `completed`, `cancelled`, `heartbeat`, `metadata`. +- `protocol`: `openai_chat`, `responses`, `anthropic_messages`, `gemini`, `native`, `litellm_fallback`, or provider-specific. +- `transport`: `sse`, `websocket`, `jsonl`, or future string. +- `data`: JSON-safe event payload. +- `raw`: optional raw chunk for trace only, sanitized/serialized. +- `metadata`: JSON-safe metadata. +- `visible_output`: bool. +- `timestamp_utc`. + +Add helpers: + +- `stream_event_from_sse_chunk()` +- `stream_event_to_sse()` +- `stream_event_to_websocket_message()` placeholder/seam. + +Keep formatters small and override-friendly. + +## Transport Abstraction + +Add `StreamTransportFormatter` base/protocol: + +- `format_event(event)` +- `format_error(error_event)` +- `format_done()` +- `is_terminal_event(event)` + +Add `SSEStreamFormatter` for existing HTTP SSE output. + +Add `WebSocketStreamFormatter` placeholder that formats JSON messages but is not wired to a route yet. + +Add `JSONLineStreamFormatter` for native provider/internal stream tests if useful. + +Responses API can keep `ResponsesSSEFormatter` but should share the transport interface or adapt to it. + +## Metrics And Lifecycle + +Add `StreamMetrics`: + +- `started_at` +- `first_byte_at` +- `first_visible_output_at` +- `last_chunk_at` +- `completed_at` +- `chunk_count` +- `visible_chunk_count` +- `error_count` +- `cancelled` +- properties: `ttfb_seconds`, `ttft_seconds`, `duration_seconds`, `idle_seconds`. + +Add `StreamMonitor`: + +- records raw chunk, formatted chunk, visible output, errors, completion, cancellation. +- can detect stall if `time_since_last_chunk > stall_timeout`. + +Env knobs: + +- `STREAM_TTFB_TIMEOUT_SECONDS` optional/disabled by default. +- `STREAM_STALL_TIMEOUT_SECONDS` optional/disabled by default. +- `STREAM_HEARTBEAT_SECONDS` optional/disabled by default. + +Phase 8 should add metrics and trace events first. Enforced timeouts can be opt-in to avoid behavior surprises. + +## Stream Policy + +- Move/re-export `can_retry_stream_after_error()` into the new streaming policy package, preserving current behavior. +- Keep `stream_retry_policy.py` as a compatibility wrapper if imports exist. +- Add visible-output detection policy: + - Chat-completions content/tool/function deltas are visible. + - Reasoning-only deltas are not visible for fallback only if the existing env allows reasoning-only retry. + - Error chunks and `[DONE]` are not visible. + - Responses `response.output_text.delta` is visible. + - Responses `response.failed` is not visible by itself. +- Add tests for malformed chunks failing closed. + +## Streaming Error Handling + +Add `StreamingErrorDecision` dataclass: + +- `classified` +- `action`: `retry_same`, `rotate`, `fail`, `fallback_allowed`, `fallback_blocked_after_output` +- `start_provider_cooldown`: bool +- `provider_cooldown_duration` +- `reason` + +Add helper that consumes: + +- exception +- provider +- last_streamed_chunk +- attempt +- max_retries +- deadline +- allow_reasoning_only_retry +- retry/cooldown env settings + +It should use existing `classify_error()`, `should_retry_same_key()`, `should_rotate_on_error()`, and Phase 7 retry policy. + +It should not sleep or mutate credential state; executor still owns those side effects. + +Refactor `_execute_streaming()` gradually: + +- First introduce helper and tests. +- Then replace duplicated branch decision logic where safe. +- Preserve exact output semantics and trace ordering. + +If full deduplication is too risky, use helper for cooldown/visibility/fallback decisions but keep branch structure. + +## Fallback And Cooldown Invariants + +- Fallback to next target is allowed only before visible output. +- Provider cooldown can start only before visible output or after error-only chunks. +- Same-credential retry after reasoning-only chunks remains controlled by `STREAM_RETRY_ON_REASONING_ONLY` behavior from current code. +- If a stream emits visible output then errors: + - no cross-target fallback. + - return/raise the current upstream stream failure behavior. + - emit `routing_stream_fallback_blocked_after_output` when in a fallback group. +- Preserve Phase 7 tests. + +## Native Provider Streaming + +- Add a streaming mode seam to `NativeProviderExecutor` that can return `StreamEvent`s or formatted SSE through the common formatter. +- Live `RequestExecutor` does not need to route native streaming fully unless safe, but the mode should be testable: + - `execution=native` streaming target can call native executor stream when provider declares native streaming support. + - if not supported, fail with `unsupported_operation` before output so fallback can choose next target. +- Add provider interface optional method only if needed: + - `supports_native_streaming(model, operation)` + - default false/no-op. +- This should prepare Claude Code/Codex/Copilot/Antigravity skeletons without claiming live support. + +## Responses Streaming Integration + +- Keep current `/v1/responses` SSE behavior. +- Share metrics/trace helpers where possible. +- Preserve stored final streamed response behavior. +- Add tests: + - Responses stream emits visible-output metrics. + - failed stream records error metrics. + - WebSocket formatter seam can format equivalent event payloads but no route is exposed. +- Do not change stored response shape unless necessary. + +## Transaction Trace Requirements + +Add or standardize stream trace passes: + +- `stream_started` +- `stream_first_byte` +- `stream_first_visible_output` +- `stream_stall_detected` +- `stream_heartbeat_sent` +- `stream_completed` +- `stream_cancelled` +- `stream_error_decision` +- `stream_metrics_final` + +Existing pass names stay: + +- `raw_stream_chunk` +- `parsed_stream_chunk` +- `assembled_stream_response` +- provider/native/responses stream pass names. + +Trace metrics must not include raw credentials or auth headers. + +Raw chunks should continue through existing redaction/serialization. + +## Client Disconnect And Cancellation + +- `StreamingHandler.wrap_stream()` already checks `request.is_disconnected()`. +- Add metrics event for cancellation/disconnect. +- Add trace `stream_cancelled`. +- Do not mark success on cancellation. +- Preserve current behavior for partial streams. + +## Heartbeats And Stall Detection + +Implement monitor primitives and tests first. + +Optional runtime heartbeat support: + +- if `STREAM_HEARTBEAT_SECONDS > 0`, emit SSE comment heartbeat `: keep-alive\n\n` or compatible event only when no provider chunk has arrived. +- default disabled to avoid client behavior changes. + +Optional stall detection: + +- if `STREAM_STALL_TIMEOUT_SECONDS > 0`, classify as transient stream failure before visible output or fail after visible output. +- default disabled. + +Tests can use fake clocks. + +## Testing Plan + +Stream event tests: + +- SSE chunk to event parsing. +- Responses event visibility. +- malformed chunk fails closed. +- event-to-SSE formatting. +- WebSocket formatter seam output shape. + +Metrics tests: + +- TTFB and TTFT calculations. +- chunk counts. +- cancellation. +- stall detection with fake clock. + +Policy tests: + +- current reasoning-only retry behavior preserved. +- visible output detection for text/tool deltas. +- error and done chunks not visible. + +Error handler tests: + +- large retry-after before output starts cooldown decision. +- visible output blocks fallback/cooldown. +- transient errors choose same-key retry before max retry. +- permanent errors fail. + +Executor tests: + +- streaming metrics trace passes emitted. +- cancellation trace emitted. +- pre-output fallback still works. +- post-output fallback remains blocked. +- provider cooldown still starts before output. + +Native stream tests: + +- native streaming emits common stream events/trace. +- unsupported native streaming fails before output and can fallback. + +Responses streaming tests: + +- existing route behavior unchanged. +- metrics final trace emitted. + +Regression tests: + +- Phase 1 protocol tests. +- Phase 2 transform logging tests. +- Phase 3 adapter/cache tests. +- Phase 4 Responses tests. +- Phase 5 provider/native tests. +- Phase 6 routing tests. +- Phase 7 retry/cooldown tests. +- `test_session_tracking.py`. +- `test_selection_engine.py`. + +## Commit Checkpoints + +1. Add streaming event/transport/metrics primitives with tests. +2. Move/re-export stream retry/visibility policy with tests. +3. Add streaming error decision helper with tests. +4. Integrate stream metrics and trace passes into `StreamingHandler` and/or executor with tests. +5. Refactor streaming error branches only as far as tests make safe. +6. Add native streaming seam/support detection tests. +7. Add Responses streaming metrics/transport seam tests. +8. Run focused and regression tests. +9. Review with `explore` and `explore-heavy`; fix findings; write uncommitted Phase 8 report. + +## Risks And Mitigations + +- Stream behavior is client-visible. Mitigation: keep SSE output format unchanged by default and add primitives before enforcement. +- Timeout/stall handling can break long reasoning streams. Mitigation: default stall/TTFB enforcement disabled; only metrics on by default. +- Refactoring executor streaming can regress retry semantics. Mitigation: keep branch structure unless shared helper is proven by tests. +- WebSocket support can be over-promised. Mitigation: formatter seam only, no route unless explicitly implemented later. +- Visible-output detection can be too permissive. Mitigation: fail closed on malformed/ambiguous chunks and preserve current tests. +- Metrics can leak data if raw chunks are logged. Mitigation: trace summaries and use existing redaction/serialization paths. diff --git a/docs/experimental/phase-8b-streaming-hardening.md b/docs/experimental/phase-8b-streaming-hardening.md new file mode 100644 index 000000000..932bec89a --- /dev/null +++ b/docs/experimental/phase-8b-streaming-hardening.md @@ -0,0 +1,84 @@ +# Phase 8b: Streaming Hardening, Cancellation, Heartbeats, And Stall Policy + +## Goal + +Correct the Phase 8 audit findings. Phase 8 added stream events, metrics, formatters, and observability, but stream hardening still needs upstream cancellation, active TTFB/stall policy, heartbeat support, and generic native HTTP streaming support. + +## Non-Goals + +- Do not enable native streaming for priority providers from Phase 5b. +- Do not rewrite the entire streaming executor or change existing chat-completions SSE chunk format by default. +- Do not implement a WebSocket FastAPI route. +- Do not weaken Phase 6b fallback visible-output safety or Phase 7b cooldown/retry latch behavior. +- Do not introduce persistence. +- Do not commit user-facing reports. + +## Current State + +- `StreamMonitor` records TTFB, TTFT, chunk counts, errors, cancellation, and stall status, but no active timeout policy uses it. +- `StreamingHandler.wrap_stream()` detects client disconnect but does not guarantee upstream iterator closure. +- No heartbeat frames are emitted during long waits between chunks. +- `NativeHTTPTransport.stream_json_lines()` requires custom injected clients to expose `stream_json_lines()` and does not support generic `httpx.AsyncClient.stream()`. +- Existing tests cover metrics and fallback, but not upstream cancellation, heartbeats, TTFB timeout, stall timeout, or generic native stream transport. + +## Implementation Plan + +1. Extend stream runtime settings. + - Add `ttfb_timeout_seconds`, `stall_timeout_seconds`, `heartbeat_interval_seconds`, and `cancel_upstream_on_disconnect`. + - Add env overrides: `STREAM_TTFB_TIMEOUT_SECONDS`, `STREAM_STALL_TIMEOUT_SECONDS`, `STREAM_HEARTBEAT_INTERVAL_SECONDS`, `STREAM_CANCEL_UPSTREAM_ON_DISCONNECT`. + - Keep defaults behavior-compatible: no timeout/heartbeat unless configured; upstream cancellation enabled on disconnect. + +2. Add upstream stream close helper. + - Close via `aclose()` when available, otherwise `close()`. + - Use it on client disconnect, cancellation, or abnormal stream exit. + - Log close failures only at debug/trace level. + +3. Add heartbeat formatting. + - Add `format_heartbeat()` to SSE/WebSocket/JSONL formatters. + - SSE heartbeat is a comment frame such as `: heartbeat\n\n`. + - Heartbeats must not count as visible output, session evidence, or usage. + +4. Add heartbeat emission in `StreamingHandler.wrap_stream()`. + - When configured, wait for upstream chunks with heartbeat interval timeout and yield heartbeat comments while waiting. + - Default remains no heartbeat. + +5. Add TTFB timeout policy. + - If configured and no first byte arrives in time, raise `StreamedAPIError` with structured `api_connection` timeout payload. + - This occurs before visible output so existing retry/fallback can apply. + +6. Add stall timeout policy. + - If configured and no chunk arrives for the configured interval after first byte, raise `StreamedAPIError` with structured `api_connection` timeout payload. + - Phase 7b visible-output latch must still suppress retry/fallback/cooldown if output was already visible. + +7. Add native `httpx` stream support. + - Keep custom `stream_json_lines()` support first. + - Otherwise use `client.stream("POST", ...)` with `aiter_lines()` or `aiter_bytes()` fallback. + - Preserve `[DONE]`, ignore empty lines, and parse `data:` JSON when possible. + +8. Add lifecycle trace events. + - `stream_heartbeat` + - `stream_ttfb_timeout` + - `stream_stall_timeout` + - `stream_upstream_cancelled` + - `stream_upstream_close_failed` + - Keep snapshots disabled and metadata sanitized. + +## Tests + +- Formatter heartbeat tests. +- Streaming handler disconnect/upstream close tests. +- Heartbeat interval tests. +- TTFB timeout tests. +- Stall timeout tests before and after visible output. +- Native HTTP transport tests for `httpx`-style streaming. +- Phase 7b retry/cooldown/routing regression subset. + +## Acceptance Criteria + +- Client disconnect closes upstream async streams when possible. +- Heartbeats are supported, disabled by default, and non-visible. +- Configured TTFB/stall timeouts produce structured stream errors. +- Visible-output latch still prevents retry/fallback/cooldown after output. +- Native HTTP transport supports generic `httpx` streaming plus custom test clients. +- Stream traces include heartbeat/timeout/cancel metadata without secrets. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-8c-streaming-runtime-hardening.md b/docs/experimental/phase-8c-streaming-runtime-hardening.md new file mode 100644 index 000000000..fa731ef9a --- /dev/null +++ b/docs/experimental/phase-8c-streaming-runtime-hardening.md @@ -0,0 +1,60 @@ +# Phase 8c: Responses Stream Runtime Hardening And Anthropic Close Safety + +## Goal + +Close all Phase 8/8b third-pass streaming findings while preserving the transport-neutral Responses stream model and the existing Anthropic compatibility wrapper. + +## Scope + +- Apply Phase 8b stream runtime settings directly in `ResponsesService.stream_events()`. +- Add Responses heartbeat formatting support so SSE wrappers can emit non-visible heartbeat frames. +- Enforce Responses TTFB timeout before first upstream chunk. +- Enforce Responses stall timeout after a prior upstream chunk/visible output. +- Close upstream Responses chat streams on client disconnect, timeout, or abnormal exit. +- Keep Responses cost-comment handling for Phase 9c, but avoid making heartbeat/cost metadata visible output. +- Ensure Anthropic compatibility streaming always attempts to close upstream on disconnect and wrapper exit, including generator-style streams that only expose close on the iterator. + +## Non-Goals + +- Do not add WebSocket routes. +- Do not replace the Responses chat bridge with native Responses execution. +- Do not change public Responses event ordering except inserting SSE comment heartbeats when configured. +- Do not make timeout defaults active; all timeout/heartbeat knobs remain opt-in through existing config. +- Do not solve Phase 9 cost-comment propagation here except preserving metadata as non-visible. + +## Implementation Plan + +1. Responses heartbeat formatter. + - Add heartbeat handling to `ResponsesSSEFormatter` and `ResponsesWebSocketFormatter`. + - Represent heartbeat as a transport-neutral non-terminal `ResponsesStreamEvent` with `event_name="heartbeat"`. + - SSE heartbeat output is `": heartbeat\n\n"`. + +2. Responses runtime settings. + - Load `get_stream_runtime_settings()` inside `ResponsesService.stream_events()`. + - Use polling around upstream `__anext__()` to enforce optional TTFB and stall timeouts and emit optional heartbeats. + - Keep all settings disabled by default. + +3. Upstream close helper. + - Close upstream iterator/stream with `aclose()` or `close()` on disconnect, timeout, or abnormal exit. + - Trace close/close-failure events. + +4. Disconnect detection. + - Poll `request.is_disconnected()` while waiting for upstream chunks. + - If disconnected, close upstream when configured and stop without emitting model output. + +5. Anthropic close safety. + - Track both the async iterator and original OpenAI stream as close candidates. + - Ensure disconnect and wrapper exit close whichever object exposes `aclose()` / `close()`. + +6. Tests. + - Responses heartbeat, TTFB timeout, stall timeout after output, disconnect close, and heartbeat SSE comment tests. + - Anthropic iterator-only upstream close test. + - Broader streaming/Responses/Anthropic regression slice. + +## Acceptance Criteria + +- Responses streaming honors `STREAM_TTFB_TIMEOUT_SECONDS`, `STREAM_STALL_TIMEOUT_SECONDS`, `STREAM_HEARTBEAT_INTERVAL_SECONDS`, and `STREAM_CANCEL_UPSTREAM_ON_DISCONNECT`. +- Responses heartbeat frames are non-visible SSE comments. +- Responses upstream chat streams are closed on disconnect/timeout/abnormal exit. +- Anthropic compatibility streaming closes upstream even when close is exposed only on the async iterator. +- Focused tests pass and both reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/phase-9-usage-quota-cost.md b/docs/experimental/phase-9-usage-quota-cost.md new file mode 100644 index 000000000..66cb18273 --- /dev/null +++ b/docs/experimental/phase-9-usage-quota-cost.md @@ -0,0 +1,331 @@ +# Phase 9 Plan: Usage, Quota, Cost + +## Goal + +Make usage, quota, and cost accounting consistent across LiteLLM fallback, native protocol adapters, Responses API, streaming, routing fallback chains, and provider-specific quota metadata while preserving the existing `UsageManager`, `SelectionEngine`, fair-cycle behavior, storage format, and provider quota reset logic. Phase 9 should formalize a normalized usage/cost layer that feeds the current usage engine rather than replacing it. + +## Non-Goals + +- Do not replace `UsageManager`, `TrackingEngine`, `LimitEngine`, `SelectionEngine`, `SessionTracker`, or the retry-after parser. +- Do not introduce SQLite or any database. +- Do not implement the full Phase 10 JSON/env config system. +- Do not change credential selection strategy semantics. +- Do not rewrite provider quota trackers. +- Do not require live provider credentials or live pricing API calls. +- Do not make cost enforcement mandatory; cost tracking is additive unless a provider already has limits. +- Do not change client-visible response usage fields except where normalization already happens today. + +## Current Code Context + +- `CredentialContext.mark_success()` already accepts prompt, completion, thinking, cache read/write tokens, approximate cost, and response headers. +- `TrackingEngine.record_usage()` stores request count, token buckets, output tokens, total tokens, and approximate cost into windows/totals. +- `RequestExecutor._extract_usage_tokens()` extracts usage from LiteLLM-like response objects, subtracting reasoning tokens from completion tokens when the provider includes reasoning inside completion. +- `RequestExecutor._calculate_cost()` uses LiteLLM cost helpers and returns `0.0` for providers that set `skip_cost_calculation`. +- `StreamingHandler.wrap_stream()` records stream usage after final chunk and currently calculates cost with `litellm.get_model_info()`. +- Phase 4 Responses streaming maps chat usage into Responses usage for storage. +- Phase 5 native protocols can produce provider-native usage shapes through protocol parse/format paths. +- Phase 6 routing can try multiple targets, but only the successful target should record usage. +- Phase 7 fallback summaries are client-safe and should not be polluted by raw usage/cost internals. +- Phase 8 stream metrics are timing/count-only and do not alter usage recording. + +## Files To Add + +- `src/rotator_library/usage/accounting.py` +- `src/rotator_library/usage/costs.py` +- `src/rotator_library/usage/quota.py` +- `tests/test_usage_accounting.py` +- `tests/test_usage_costs.py` +- `tests/test_usage_quota_snapshots.py` +- `tests/test_executor_usage_accounting.py` +- `tests/test_streaming_usage_accounting.py` +- `tests/test_responses_usage_accounting.py` +- `tests/test_native_usage_accounting.py` + +## Files Likely To Touch + +- `src/rotator_library/usage/types.py` +- `src/rotator_library/usage/manager.py` +- `src/rotator_library/usage/tracking/engine.py` +- `src/rotator_library/client/executor.py` +- `src/rotator_library/client/streaming.py` +- `src/rotator_library/responses/bridge.py` +- `src/rotator_library/responses/service.py` +- `src/rotator_library/native_provider/executor.py` +- `src/rotator_library/core/utils.py` only if current usage normalization helper needs a small adapter. +- Provider files only for optional cost/quota declarations, not broad rewrites. + +## Normalized Usage Model + +Add `UsageRecord` dataclass: + +- `input_tokens` +- `output_tokens` +- `completion_tokens` +- `reasoning_tokens` +- `cache_read_tokens` +- `cache_write_tokens` +- `total_tokens` +- `raw_total_tokens` +- `request_count` +- `source` +- `provider` +- `model` +- `metadata` + +Token semantics: + +- `input_tokens` means billable non-cache-read prompt tokens when provider separates cache-read tokens. +- `cache_read_tokens` means prompt/cache tokens read from provider cache. +- `cache_write_tokens` means prompt/cache tokens written or created. +- `completion_tokens` means visible/output text/tool tokens, excluding reasoning when the provider reports reasoning separately. +- `reasoning_tokens` means hidden thinking/reasoning tokens. +- `output_tokens = completion_tokens + reasoning_tokens`. +- `total_tokens = input_tokens + cache_read_tokens + cache_write_tokens + completion_tokens + reasoning_tokens`. +- `raw_total_tokens` preserves provider-reported total before normalization for debugging. + +Rationale: current code already avoids double-counting reasoning by subtracting thinking from completion when necessary. Phase 9 makes that rule explicit and reusable. + +## Usage Extraction + +Add `extract_usage_record(response_or_usage, provider=None, model=None, source=None)`. + +Support: + +- LiteLLM/OpenAI object usage attributes. +- dict usage fields. +- OpenAI `prompt_tokens_details.cached_tokens`. +- OpenAI `completion_tokens_details.reasoning_tokens`. +- Anthropic `input_tokens`, `output_tokens`, `cache_creation_input_tokens`, `cache_read_input_tokens`. +- Gemini `usageMetadata`, `promptTokenCount`, `candidatesTokenCount`, `thoughtsTokenCount`, `cachedContentTokenCount`, `totalTokenCount`. +- Responses `input_tokens`, `output_tokens`, `output_tokens_details.reasoning_tokens`, `input_tokens_details.cached_tokens`. +- Existing stream usage dicts after `normalize_usage_for_response()`. + +Unknown usage shapes return an empty `UsageRecord` with source metadata rather than raising in runtime paths. + +Tests must cover dicts, objects, nested details, and double-count prevention. + +## Cost Model + +Add `CostBreakdown` dataclass: + +- `input_cost` +- `cache_read_cost` +- `cache_write_cost` +- `output_cost` +- `reasoning_cost` +- `total_cost` +- `currency` +- `pricing_source` +- `metadata` + +Add `ModelPricing` dataclass: + +- per-token prices for input, cache read, cache write, output, reasoning. +- `currency`. +- `source`. + +Add `CostCalculator`: + +- prefers explicit provider/model pricing declarations. +- falls back to LiteLLM model info/completion_cost where applicable. +- returns zero cost for `skip_cost_calculation`. +- does not call network. + +Minimal Phase 9 provider declarations: + +- Optional `get_model_pricing(model)` on `ProviderInterface`, default `None`. +- Providers can later define native pricing without touching accounting code. + +Env pricing can be minimal and Phase-10-ready: + +- `MODEL_PRICE_{PROVIDER}_{MODEL}_INPUT` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_OUTPUT` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_CACHE_READ` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_CACHE_WRITE` +- `MODEL_PRICE_{PROVIDER}_{MODEL}_REASONING` + +If env parsing is too broad, implement programmatic pricing tests and leave env polish to Phase 10. + +## UsageManager Integration + +- Preserve existing `mark_success()` signature for compatibility. +- Add optional `usage_record` and `cost_breakdown` parameters only if needed, but do not force every call site to change. +- Preferred minimal integration: + - Executor and streaming code use `UsageRecord` internally, then pass existing numeric fields to `mark_success()`. + - `approx_cost` remains a float in current storage. + - Add optional metadata later only if storage impact is small. +- Do not change persisted JSON shape unless additive and backward-compatible. +- Add trace pass: + - `usage_accounting_summary` + - Include normalized token fields, raw total, cost total, pricing source, provider/model, and source. + - No credential secrets. + +## Quota Snapshot Model + +Add `QuotaSnapshot` dataclass: + +- `provider` +- `model` +- `quota_group` +- `credential_id` optional stable/masked identifier. +- `window_name` +- `limit` +- `used` +- `remaining` +- `reset_at` +- `source` +- `metadata` + +Add helper to build snapshots from `UsageManager` state: + +- per model. +- per quota group. +- per credential where available. + +Keep `WindowLimitChecker` behavior unchanged. + +Tests should prove snapshots reflect group windows and model windows without affecting limits. + +Optional `UsageAPI` helper can expose snapshots for future UI/TUI work if small. + +## Routing/Fallback Usage Behavior + +- Only successful target attempts record usage. +- Failed target attempts continue to record failures through existing paths. +- Route fallback summaries should include optional `usage_recorded: false` for failed targets only if already tracked internally; do not expose raw costs. +- Tests: + - first target failure + second target success records usage only on second provider/credential. + - native target success records normalized usage. + - explicit LiteLLM fallback target records normalized usage. + +## Streaming Usage + +- Replace ad-hoc stream usage extraction/cost calculation with `UsageRecord` and `CostCalculator`. +- Preserve existing `mark_success()` numeric values and final `[DONE]` behavior. +- Preserve `skip_cost_calculation`. +- Add tests: + - stream final usage with cache/read/write/reasoning records expected buckets. + - stream with missing usage records zero usage but still marks success as today. + - stream metrics from Phase 8 remain independent from token usage. + +## Responses Usage + +- Use `UsageRecord` when converting/storing Responses usage. +- Ensure `previous_response_id` storage is not affected. +- Add tests: + - non-streaming Responses usage maps to normalized fields. + - streaming Responses usage stores expected input/output totals. + - reasoning details are not double counted. + +## Native Provider Usage + +- In `NativeProviderExecutor.execute()`, protocol-formatted provider responses should expose usage in a shape `extract_usage_record()` can understand. +- Add trace summary but do not require the isolated native executor to mutate `UsageManager`. +- Live routed native execution via `RequestExecutor` should record normalized usage because it receives the native response and passes through the same executor accounting. +- Tests: + - native OpenAI-style response usage records normalized fields. + - native Gemini usage metadata records thought/cached tokens. + - native Responses usage records reasoning/details. + +## Quota/Cost Reporting + +- Keep existing quota viewer/API behavior unchanged unless adding optional fields is safe. +- Add cost totals to existing total/window stats already present as `approx_cost`. +- Do not implement hard cost caps unless trivial and isolated; cost caps can be a Phase 10+ config feature. +- Add clear docs/comments that `approx_cost` is advisory and depends on available pricing. + +## Transform Trace Requirements + +Add: + +- `usage_accounting_summary` +- `usage_cost_calculated` +- `quota_snapshot_built` for explicit snapshot APIs/tests, not every request. + +Metadata: + +- provider, model, source, pricing source, skip-cost flag. + +Data: + +- normalized usage record and cost breakdown only. + +Do not log credentials, raw headers, or raw provider errors. + +## Testing Plan + +Accounting tests: + +- OpenAI/LiteLLM object usage. +- OpenAI dict usage with prompt/completion details. +- Anthropic usage. +- Gemini `usageMetadata`. +- Responses usage. +- reasoning double-count prevention. +- missing/unknown usage shape returns empty record. + +Cost tests: + +- explicit pricing calculates all buckets. +- provider `skip_cost_calculation` returns zero. +- LiteLLM fallback path still returns a float when LiteLLM has model info. +- missing pricing returns zero with `pricing_source="unavailable"`. + +Quota snapshot tests: + +- model window snapshot. +- group window snapshot. +- missing window snapshot is empty/no-op. + +Executor tests: + +- non-stream success uses normalized usage. +- trace emits usage accounting summary. +- fallback second target success records usage once. + +Streaming tests: + +- final stream usage uses normalized accounting. +- stream cost honors skip-cost provider. +- Phase 8 metrics still emitted. + +Responses tests: + +- stored Responses usage shape remains compatible. +- normalized trace emitted when transaction logger exists. + +Native tests: + +- native response usage shapes normalize. + +Regression tests: + +- Phase 1 protocol tests. +- Phase 2 transform logging tests. +- Phase 3 adapter/cache tests. +- Phase 4 Responses tests. +- Phase 5 provider/native tests. +- Phase 6 routing tests. +- Phase 7 retry/cooldown tests. +- Phase 8 streaming tests. +- `test_session_tracking.py`. +- `test_selection_engine.py`. + +## Commit Checkpoints + +1. Add `UsageRecord`, extraction helpers, and tests. +2. Add `CostBreakdown`, pricing helpers/calculator, and tests. +3. Wire executor non-streaming accounting and trace summary with tests. +4. Wire streaming accounting/cost calculation with tests. +5. Add quota snapshot helpers and tests. +6. Add Responses/native usage accounting coverage. +7. Run focused and regression tests. +8. Review with `explore` and `explore-heavy`; fix findings; write uncommitted Phase 9 report. + +## Risks And Mitigations + +- Token normalization could change quota accounting. Mitigation: preserve current numeric `mark_success()` buckets and test current OpenAI/LiteLLM behavior. +- Cost estimates may be wrong for providers with unknown pricing. Mitigation: return zero/unavailable rather than guessing; make `approx_cost` advisory. +- Persisted usage JSON compatibility could break. Mitigation: keep existing storage fields and only add optional helpers unless tests prove serialization is safe. +- Provider-specific usage shapes are broad. Mitigation: implement common native shapes now and leave provider overrides possible. +- Fallback chains could double-record usage. Mitigation: only successful attempt calls `mark_success()`; add tests for failure then success. diff --git a/docs/experimental/phase-9b-usage-cost-corrections.md b/docs/experimental/phase-9b-usage-cost-corrections.md new file mode 100644 index 000000000..8cff0bd49 --- /dev/null +++ b/docs/experimental/phase-9b-usage-cost-corrections.md @@ -0,0 +1,76 @@ +# Phase 9b: Provider-Reported Cost, SSE Cost Events, And Quota Snapshot Hardening + +## Goal + +Correct Phase 9 validation findings. Phase 9 normalized token usage and added advisory cost calculation, but provider-reported costs and streaming cost events are not yet first-class runtime accounting inputs. + +## Non-Goals + +- Do not replace `UsageManager`, JSON persistence, or window/limit engines. +- Do not introduce SQLite or any new database. +- Do not make advisory pricing authoritative when providers report actual cost. +- Do not build an admin quota UI. +- Do not commit user-facing reports. +- Do not alter stream fallback/cooldown safety from Phases 6b-8b. + +## Current State + +- `protocols.types.CostDetails` exists and some protocol adapters populate it. +- `UsageRecord` has token buckets but no first-class provider-reported cost fields. +- `extract_usage_record()` ignores `cost_details`, `cost`, `total_cost`, and protocol `Usage.cost`. +- `CostCalculator.calculate()` prefers advisory pricing and drops actual provider cost values. +- `StreamingHandler` does not parse SSE comment/event cost frames. +- `mark_success(approx_cost=...)` can already store an approximate cost total without changing `UsageManager` shape. + +## Implementation Plan + +1. Extend `UsageRecord` with provider-reported cost fields. + - `provider_reported_cost: Optional[float]` + - `cost_currency: str = "USD"` + - `cost_source: Optional[str]` + - Include them in `to_dict()`. + +2. Extract provider-reported costs from provider shapes. + - OpenAI-like `cost_details.total_cost`, `cost_details.cost`, `cost`, `total_cost`, and `provider_reported_cost`. + - Anthropic/Gemini generic cost fields and `costMetadata`. + - Protocol `Usage.cost` values. + +3. Make provider-reported actual cost win in `CostCalculator`. + - `skip_cost_calculation` remains the highest-priority behavior. + - Add `provider_reported_cost` to `CostBreakdown` and make `total_cost` use it without double counting. + - Preserve advisory pricing paths when no provider-reported cost exists. + +4. Parse SSE cost comment/event frames. + - Support `: cost {json}`, `: cost 0.001`, and `event: cost` frames. + - Keep cost comments non-visible and out of session anchors. + - Final provider usage cost overrides earlier cost comments. + +5. Extend formatted SSE usage extraction. + - Ensure formatted SSE chunks with `usage.cost_details` flow through `UsageRecord` and cost calculation. + +6. Harden native and Responses cost traces. + - Native usage traces should include cost breakdown. + - Responses usage traces should show provider-reported cost when present. + +7. Keep quota snapshots honest. + - Do not invent cost totals if usage state does not store them. + - Document/request-token-window scope explicitly. + +## Tests + +- Provider-reported cost extraction tests in `test_usage_accounting.py`. +- Provider-reported cost precedence tests in `test_usage_costs.py`. +- SSE cost comment/event tests in `test_streaming_usage_accounting.py`. +- Native/Responses trace cost tests. +- Quota snapshot honesty tests. +- Phase 8b stream regression subset. + +## Acceptance Criteria + +- Provider-reported actual cost is preserved in `UsageRecord`. +- `CostCalculator` uses provider-reported cost before advisory pricing, while skip-cost still wins. +- SSE cost comments/events update streaming cost accounting and `mark_success(approx_cost=...)`. +- Cost comments remain non-visible for fallback/retry/session purposes. +- Native and Responses usage traces include provider-reported cost when present. +- Quota snapshots remain read-only request/token window reports and do not invent unsupported cost totals. +- Focused tests and dual-agent review pass. diff --git a/docs/experimental/phase-9c-usage-cost-correction.md b/docs/experimental/phase-9c-usage-cost-correction.md new file mode 100644 index 000000000..ec056965f --- /dev/null +++ b/docs/experimental/phase-9c-usage-cost-correction.md @@ -0,0 +1,61 @@ +# Phase 9c: Usage And Cost Correctness Completion + +## Goal + +Close all Phase 9/9b third-pass usage/cost findings while preserving `UsageManager` as the authoritative persistence/limit engine and keeping provider-reported cost precedence from Phase 9b. + +## Scope + +- Preserve provider-reported top-level response cost even when a `usage` object exists. +- Prevent OpenAI-like cache-write tokens from being double-counted in normalized totals/costs. +- Preserve/sum structured provider cost breakdowns that omit `total_cost`. +- Carry Responses streaming SSE cost comments/events into completed usage/cost records. +- Ensure `event: cost` frames are treated as metadata and never visible output for retry/fallback. +- Add native streaming `usage_accounting_summary` traces and preserve provider-reported stream cost where protocol events/raw chunks expose usage. + +## Non-Goals + +- Do not replace `UsageManager`, quota windows, or persisted usage JSON format. +- Do not invent cost totals when neither provider-reported cost nor configured/advisory pricing exists. +- Do not add billing persistence beyond existing usage state and traces. +- Do not implement admin quota/cost dashboards. +- Do not alter public token fields except to correct double-counting. + +## Implementation Plan + +1. Top-level provider cost preservation. + - Merge sibling top-level cost fields into nested `usage` payloads before normalization. + - Covered keys: `cost`, `total_cost`, `cost_details`, `provider_reported_cost`, `currency`, `costMetadata`. + +2. Structured provider cost breakdowns. + - Sum known provider-reported cost fields when no total is present. + - Preserve currency and source metadata. + +3. Cache-write double-count prevention. + - For OpenAI-like usage, subtract both cache-read and cache-write tokens from regular input tokens. + - Keep cache-write tokens as their own normalized bucket. + +4. Responses streaming SSE cost comments/events. + - Parse `: cost ...` comments and `event: cost` frames. + - Merge cost into in-progress/final Responses usage without emitting model output. + - Let final provider usage cost override earlier cost comments. + +5. Metadata visibility for `event: cost`. + - Add stream policy tests locking cost frames as non-visible and retry-safe. + +6. Native streaming usage/cost trace. + - Track stream usage/cost from unified events and raw chunks. + - Emit `usage_accounting_summary` on stream completion. + +7. Tests. + - Cover usage normalization, cost precedence, Responses streaming costs, stream visibility policy, and native streaming cost trace. + +## Acceptance Criteria + +- Top-level provider-reported cost is preserved when a `usage` object exists. +- Cache-write tokens are not double-counted in OpenAI-like normalized totals/costs. +- Structured provider cost breakdowns without totals are summed and treated as provider-reported cost. +- Responses streaming carries SSE cost comments/events into final usage/cost and treats them as metadata. +- `event: cost` frames do not block retry/fallback as visible output. +- Native streaming emits `usage_accounting_summary` and preserves stream cost when available. +- Focused tests pass and both reviewers report no blockers/highs/mediums. diff --git a/docs/experimental/protocol-interoperability-review.md b/docs/experimental/protocol-interoperability-review.md new file mode 100644 index 000000000..133e989c0 --- /dev/null +++ b/docs/experimental/protocol-interoperability-review.md @@ -0,0 +1,728 @@ +# Protocol Interoperability Review And Implementation State + +**Status:** Active correction phase +**Started:** 2026-07-16 +**Branch:** `experimental` +**Purpose:** Durable review, implementation contract, task breakdown, and live state for the protocol interoperability correction. + +## Product Contract + +The proxy must treat client protocols and provider protocols as independent choices. + +For every supported generative protocol: + +```text +client payload + -> client protocol parser + -> canonical request + -> provider protocol builder + -> provider + -> provider protocol parser + -> canonical response or stream events + -> selected output protocol formatter + -> client payload +``` + +The required first-class generative protocols are: + +- OpenAI Chat Completions +- OpenAI Responses +- Anthropic Messages +- Gemini generateContent / streamGenerateContent + +The input protocol, provider protocol, and output protocol must be independently selectable. The default output protocol is the originating client protocol, but callers and configured integrations may choose another supported output protocol. + +Examples that must be valid: + +- Chat client -> Anthropic provider -> Chat response +- Responses client -> Gemini provider -> Responses response +- Gemini client -> OpenAI-compatible provider -> Gemini response +- Gemini client -> provider-native protocol -> Anthropic response +- Anthropic client -> Responses provider -> Chat response when explicitly selected + +Providers must not need to know which protocol the client used. A provider declares only the protocol and operation it expects, plus provider-specific endpoint, authentication, envelope, model alias, cache, quota, and adapter behavior. + +Custom providers configured at runtime must be able to declare a supported provider protocol. Custom integrations must also be able to select an output protocol rather than being limited to OpenAI-compatible output. + +## Original Review Verdict (Superseded) + +The following verdict and reproduced failures describe the branch before the +Phase A-H correction. They are retained as the audit baseline, not as current +behavior. The current beta-readiness verdict is recorded in the final checkpoint. + +The existing branch contains useful protocol primitives but does not implement the product contract end to end. + +The current live architecture uses OpenAI Chat as the effective internal language: + +- Chat requests enter the completion executor directly. +- Anthropic requests are converted to Chat by the legacy compatibility handler before execution. +- Responses requests are converted to Chat by `ResponsesBridge` before execution. +- No public Gemini client route exists. +- Native request execution asks the provider protocol to parse an already Chat-shaped request. +- Provider preparation hooks perform partial Chat-to-provider reshaping. +- Native responses are always formatted as OpenAI Chat before optional outer wrappers convert them again. + +This works for basic text demonstrations but not for general protocol interoperability. + +## Confirmed Architectural Findings + +### Client protocol identity is not authoritative + +The originating protocol is not carried through the common request context. Native execution currently fixes the client response protocol to OpenAI Chat. + +### Native request conversion uses the wrong source parser + +The provider protocol currently parses the incoming request. Correct conversion requires the client protocol to parse the incoming request and the provider protocol to build the outgoing provider request. + +### Public protocol routes bypass the common protocol core + +Anthropic and Responses use dedicated Chat bridges. Gemini is not available as a client-facing protocol. + +### Provider hooks contain protocol translation + +Provider preparation currently performs partial message conversion for some native providers. This creates provider-specific translation islands and handles mostly plain text. Provider hooks must instead run after provider-protocol formatting and be limited to provider quirks. + +### Raw and extension preservation is unsafe across protocols + +The existing `raw` and `extra` fields are valuable for same-protocol fidelity and tracing. Destination formatters currently replay some source-native fields into different protocols, producing hybrid payloads. + +### The unified model is only partly canonical + +Several fields still preserve source-protocol concepts rather than common meaning: + +- content block type names +- system and developer instruction placement +- generation parameter dictionaries +- tool placement and tool configuration +- stop reasons and statuses +- output item layout +- stream event names and lifecycle +- provider-native extensions + +### Streaming is not protocol-independent + +All active native providers disable native streaming by default. Only the OpenAI Chat protocol has a complete cross-protocol stream formatter. Responses streaming is operationally strong but text-only, and Anthropic streaming still uses the legacy Chat wrapper. + +### Existing tests prove adapters more than interoperability + +Most protocol tests verify same-protocol parse/build or parse/format behavior. Native provider tests verify simple text requests and responses. There is no complete client-protocol x provider-protocol x output-protocol contract suite. + +## Directly Reproduced Conversion Failures + +Read-only runtime probes against the current protocol classes reproduced the following: + +- Chat system messages remained conversation messages when written as Anthropic instead of becoming the Anthropic system field. +- Chat images written as Anthropic lost the image source. +- Chat assistant tool calls disappeared from Anthropic output. +- Chat tool results remained an unsupported `tool` role in Anthropic output. +- Anthropic image, thinking, and tool-use blocks leaked into Chat content while malformed duplicate Chat tool calls were also emitted. +- Chat system messages written as Gemini retained an invalid system content role. +- Chat image and tool objects leaked into Gemini parts instead of being converted. +- Chat tool calls written as Responses stayed embedded in message objects instead of becoming Responses function-call items. +- Anthropic and Gemini stop reasons were returned as raw Chat finish reasons. +- Responses reasoning, message, and function-call output items became separate Chat choices instead of one semantically correct Chat response. + +## Preservation Policy + +### Meaning-changing fields + +The proxy must not silently discard or corrupt fields that alter the meaning or required behavior of a request. + +These include: + +- system and developer instructions +- user and assistant content +- tool definitions +- tool choice requirements +- tool calls and tool results +- required images, files, audio, or other modalities +- structured output and schema requirements +- continuation identity +- required reasoning behavior +- explicit response modalities + +If the selected provider protocol cannot represent a required meaning safely, conversion must fail before contacting the provider with an error formatted for the selected output protocol. The error must identify the unsupported capability without exposing secrets or internal provider state. + +### Optional tuning fields + +Advisory controls may be omitted when no safe equivalent exists, provided the omission is recorded in a concise conversion summary. + +Examples include provider-unsupported sampling hints, log-probability requests, service tiers, or optional vendor tuning fields. A future strict mode may promote selected warnings to errors. + +### Unknown and provider-extension fields + +Unknown fields are protocol-owned extensions, not universal fields. + +- Same input and provider protocol: preserve safe unknown fields and replay them unless explicitly blocked. +- Different input and provider protocol: do not forward unknown fields without an explicit mapping. +- Provider response to the same output protocol: preserve safe provider-native extensions. +- Provider response to a different output protocol: keep unknown provider fields internal unless the output protocol has an explicit extension mapping. +- Secret, credential, cache signature, and provider-state fields are never generic passthrough fields. + +This preserves same-protocol compatibility without creating invalid cross-protocol hybrid payloads. + +## Provider-State Policy + +Opaque provider state includes thought signatures, reasoning signatures, provider response IDs, prompt-cache IDs, and similar values that a client may never receive or return. + +Policy: + +- Extract opaque state from provider responses and completed stream events. +- Cache it even when it is not exposed in the selected output protocol. +- Scope it at least by provider, model family, credential, and logical session when provider validity requires those dimensions. +- Inject it only when returning to a compatible provider context. +- Never move opaque signatures from one provider family to another. +- Never assume reasoning text and opaque reasoning signatures are interchangeable. +- Permit client-visible reasoning text to translate only when the destination protocol supports it and policy allows exposure. +- Completion-gate stream extraction so partial or failed streams do not establish false continuation state. + +The existing field-cache system is the intended mechanism. Its declarations, extraction sources, injection timing, scopes, persistence behavior, and trace redaction must be verified during this correction. + +## Canonical Model Requirements + +### Request + +The canonical generative request must represent: + +- logical operation independent of provider endpoint operation +- model selection +- ordered system and developer instructions +- user, assistant, and tool turns +- text content +- image content with URL/data/file identity, MIME type, and detail metadata +- document/file content +- audio/video content through extensible media descriptors +- tool definitions and full input schemas +- provider-native built-in tool categories where meaning can be normalized +- tool choice and function-calling policy +- assistant tool calls +- tool results with correlation identity and error state +- reasoning policy, effort, budget, summary, and visibility +- output token limit +- temperature, top-p, top-k, penalties, seed, and stop sequences +- structured output requirements and JSON schema +- response modalities +- continuation and cache routing identity +- metadata +- source-protocol extensions isolated by protocol +- original payload for tracing and same-protocol preservation + +### Response + +The canonical generative response must represent: + +- response identity and provider-native identity +- model +- ordered output items +- assistant content +- reasoning items and summaries +- tool calls +- annotations and citations +- canonical completion status +- canonical stop reason plus source-native reason +- usage and cost +- errors and incomplete status +- provider-protocol extensions isolated by protocol +- original payload for tracing and same-protocol preservation + +### Stream + +The canonical stream must model lifecycle rather than source event names: + +- response started +- output item started +- text started, delta, and completed +- reasoning started, delta, and completed +- tool call started +- tool argument delta +- tool call completed +- annotations and media events where supported +- usage +- response completed +- response incomplete +- response failed +- transport heartbeat as transport metadata rather than model output + +Every protocol formatter must be able to construct its own valid stream lifecycle from these events. + +## Runtime Architecture Requirements + +### Protocol selection + +Each execution carries three independent protocol identities: + +- input protocol +- provider protocol +- output protocol + +The output protocol defaults to the input protocol. + +### Operation selection + +Logical operation and provider wire operation are separate. + +Chat Completions, Anthropic Messages, Responses, and Gemini generateContent all fulfill the same logical generative operation while using different wire operation names and endpoints. + +### Provider preparation + +Provider-specific preparation runs after the provider protocol has built a valid provider-native payload. + +Allowed responsibilities: + +- provider envelope +- model aliasing +- provider-required defaults +- provider-specific headers outside the payload +- endpoint selection +- narrow schema quirks + +Disallowed responsibility: + +- generic client-protocol to provider-protocol translation + +### Same-protocol fast path + +A same-protocol path may start from the original payload to preserve unknown fields, then overlay canonical changes. This optimization must be source/target gated and must not bypass credential, routing, cache, usage, error, or trace behavior. + +### LiteLLM + +LiteLLM remains an explicit execution fallback for providers that have not opted into native transport. It must not remain the implicit canonical representation for all public client protocols. + +## Public API Requirements + +### Existing routes + +- `/v1/chat/completions` defaults input/output to OpenAI Chat. +- `/v1/messages` defaults input/output to Anthropic Messages. +- `/v1/responses` defaults input/output to OpenAI Responses and retains storage/continuation behavior. + +### Gemini routes + +Add first-class Gemini request support for generateContent, streamGenerateContent, and countTokens-compatible operations. Route aliases and URL model extraction should match practical Gemini clients while preserving the proxy's provider/model routing convention. + +### Output protocol override + +Library callers and configured integrations must be able to select a supported output protocol independently of input protocol. The public HTTP override mechanism must be explicit, authenticated by the existing proxy boundary, and excluded from provider payloads. + +## Configurable Provider Requirements + +Runtime provider configuration must support: + +- provider protocol name +- endpoint or API base behavior +- authentication/header mode using existing secret handling +- native operation mapping +- streaming capability +- ordered adapters +- field-cache rules +- default output protocol for integrations that require one + +Configuration must validate protocol names, operations, transports, and incompatible combinations before requests reach the provider. + +Provider-specific quota, tier, priority, reset-window, and concurrency data remains defined by provider classes. Protocol selection does not centralize provider-variable policy. + +## Test Contract + +The generative protocol contract suite must cover every input/provider/output combination for the four primary protocols. + +For each meaningful combination, verify: + +- plain text +- ordered system and developer instructions +- multimodal text and image input +- tool definitions +- required/automatic/specific tool choice +- assistant tool calls +- tool results +- reasoning text +- opaque reasoning signatures where applicable +- structured output +- generation controls +- canonical stop/status mapping +- usage details +- safe unknown-field behavior +- errors +- non-streaming output +- streaming text +- streaming reasoning +- streaming tool calls and argument deltas +- completion and failure lifecycle + +Same-protocol fidelity and cross-protocol semantic correctness are separate test classes. + +Tests must assert destination payload validity, not merely the presence of a response. + +## Implementation Phases + +### Phase A: Durable review and contract + +State: **completed** + +- Persist this review and implementation state. +- Establish product and preservation policies. +- Record completion criteria and phased work. + +### Phase B: Canonical model and protocol contract + +State: **completed** + +- Add canonical instruction, media, tool policy, generation, status, output-item, error, and stream lifecycle semantics. +- Preserve non-generative operation compatibility. +- Add source-scoped extension handling. +- Separate logical and wire operations. + +### Phase C: Protocol conversions + +State: **completed** + +- Rework Chat, Anthropic, Responses, and Gemini readers and writers around canonical meaning. +- Implement stop/status and generation-control mappings. +- Add destination capability validation. +- Add complete non-streaming cross-protocol golden tests. + +### Phase D: Runtime wiring + +State: **completed** + +- Carry input/provider/output protocol identities through the common request context. +- Parse with the input protocol and build with the provider protocol. +- Parse with the provider protocol and format with the selected output protocol. +- Move provider preparation after provider-protocol formatting. +- Replace hardcoded Chat output selection. +- Keep legacy routes available only as controlled fallback during migration. + +### Phase E: Public Gemini and configurable providers + +State: **completed** + +- Add Gemini generateContent and countTokens-compatible routes. +- Reserve streamGenerateContent for the canonical streaming checkpoint so the + public route never exposes an unconverted or falsely advertised stream. +- Add explicit output protocol selection. +- Extend custom provider configuration to supported native protocols and output defaults. +- Validate configurations before execution. + +### Phase F: Canonical streaming + +State: **completed** + +- Implement canonical stream lifecycle. +- Implement all four input/output stream formatters. +- Preserve text, reasoning, tool calls, argument deltas, usage, completion, and errors. +- Integrate existing timeout, heartbeat, cancellation, retry visibility, and completion-gated state behavior. + +### Phase G: Provider state and integration validation + +State: **completed** + +- Verify hidden signature and provider continuation caching. +- Verify provider/credential/model/session isolation. +- Verify response protocol output does not affect provider-state extraction. +- Run mocked provider contract tests for every native provider. +- Run full local test suite. + +### Phase H: Independent review and correction + +State: **completed** + +- Run an `explore` mapping review after implementation appears complete. +- Run an `explore-heavy` reasoning review after implementation appears complete. +- Reproduce and fix all confirmed blocker, high, and medium findings. +- Repeat both reviews after substantial fixes. +- Update this document with final evidence and remaining intentional limitations. + +## Commit Checkpoints + +Commits are allowed for this correction and should be made after meaningful verified work. + +Planned checkpoints: + +1. Persist review and implementation contract. +2. Canonical model and base protocol contract. +3. Complete non-streaming protocol conversions and contract matrix. +4. Runtime input/provider/output protocol wiring. +5. Public Gemini and configurable provider protocol support. +6. Canonical streaming and stream contract matrix. +7. Provider-state verification and integration hardening. +8. Independent review corrections and final state update. + +Before each commit: + +- inspect status and diff +- run focused tests for the checkpoint +- stage only intended files +- use a concise repository-style commit message + +## Completion Criteria + +This correction is complete only when: + +- all four primary protocols can be used as input protocols +- all four primary protocols can be used as provider protocols +- all four primary protocols can be selected as output protocols +- the default output matches the input protocol +- providers do not translate client protocols themselves +- custom configured providers can declare supported provider protocols +- required semantics are preserved or rejected explicitly before transport +- same-protocol extensions are preserved safely +- cross-protocol extensions do not leak blindly +- non-streaming tools, results, reasoning, media, instructions, and structured output work across supported pairings +- streaming text, reasoning, tools, usage, completion, and failure work across supported pairings +- hidden provider state remains cached and correctly isolated regardless of selected output protocol +- focused and full local tests pass +- final `explore` and `explore-heavy` reviews have no unresolved blocker, high, or medium findings +- this document records final verification evidence and intentional limitations + +## Live State Log + +### 2026-07-16: Initial protocol readiness review + +- Mapped the protocol registry, canonical types, four primary protocol adapters, native executor, public routes, provider preparation hooks, tests, and the in-repo reference gateway. +- Confirmed the protocol product contract was not implemented end to end. +- Reproduced invalid cross-protocol payloads directly. +- Confirmed simple text is the primary working cross-protocol case. +- Confirmed all current public protocol routes converge on OpenAI Chat before provider execution. +- Confirmed native output is hardcoded to OpenAI Chat. +- Confirmed no public Gemini client route exists. +- Confirmed native streaming remains disabled for all four active native providers. +- Product owner confirmed Gemini must be a first-class client protocol. +- Product owner delegated unsupported-field and extension policy decisions to engineering judgment. +- Product owner confirmed opaque provider state should remain cached even when clients do not receive or return it. +- Selected strict preservation for meaning-changing semantics, warning-based best effort for optional tuning fields, same-protocol extension passthrough, and explicit-only cross-protocol extension mapping. + +### 2026-07-16: Canonical non-streaming conversion checkpoint + +- Added independent input/provider/output protocol identities and logical-operation identity to the protocol context and unified records. +- Added canonical source ownership, conversion warnings, media identity, ordered output items, normalized tool arguments/results, and canonical completion reasons. +- Reworked Chat, Anthropic, Responses, and Gemini request/response readers and writers around canonical instruction, content, reasoning, tool, media, generation, status, and usage semantics. +- Added source-aware passthrough: same-protocol extensions remain available while foreign raw fields and reasoning extensions are not replayed into another protocol. +- Added destination validation for required content, media identity, tool identity, tool choice, response modalities, provider-bound continuation/background controls, and failed-response envelopes. +- Added explicit warning behavior for unsupported optional controls and strengthened schema behavior. +- Added provider-state compatibility gating for hidden reasoning/thought signatures. Protocol equality alone is not sufficient to emit opaque state. +- Added a four-by-four non-streaming request and response semantic matrix with destination-wire assertions and focused regressions for all review findings. +- Corrected tool-result behavior so application JSON containing an `error` field is never reclassified as transport failure. Native Anthropic `is_error` is encoded explicitly for protocols without a native error flag without relying on ambiguous reverse inference. +- Ran four iterative `explore` and `explore-heavy` review passes. Both reviewers signed off with no unresolved blocker, high, or medium findings for the isolated canonical non-streaming checkpoint. +- Verification: 172 focused protocol/native-provider tests passed; `git diff --check` passed except repository line-ending notices. +- A raw `pytest -q tests` collection remains blocked by pre-existing retired Antigravity tests and two stale quota tests importing removed modules. This is recorded for the later full-local-verification phase and was not caused by this checkpoint. +- Runtime execution still uses provider-side parsing of Chat-shaped requests and hardcodes Chat output. This is intentionally not claimed as fixed until Phase D. +- Canonical streaming remains intentionally deferred to Phase F. + +### 2026-07-16: Live non-streaming runtime checkpoint + +- Added `RotatingClient.agenerate()` as the protocol-independent generative entry point while preserving `acompletion()` as the Chat-compatible facade. +- The request builder now retains the untouched client payload, parses it with the selected input protocol, and carries canonical meaning plus independent input/output identities through retries and fallback targets. +- Native execution now follows the required sequence: client protocol reader, provider protocol writer, narrow provider preparation, provider transport, provider protocol reader, selected output protocol writer. +- Removed provider-owned Chat-to-Responses and Chat-to-Gemini translation. Codex and Antigravity preparation hooks now receive already-valid provider-native payloads. +- Native provider session hooks and validators receive provider-native request shapes. Native attempts bypass the legacy Chat/LiteLLM provider-transform pipeline; callbacks alone retain a Chat compatibility view. +- Callback compatibility views are deep copies. Nested message/tool edits are merged back into canonical meaning field by field without mutating the immutable baseline or discarding source-native metadata such as Anthropic cache controls. +- Custom and LiteLLM non-streaming responses now pass through the canonical response reader/writer when output differs from Chat. +- Anthropic Messages and Responses non-streaming services now execute through the common protocol path. Responses retains local storage, capability, continuation expansion, and session behavior without converting through Chat. +- Local Responses lineage disables provider continuation-cache injection, preventing locally expanded history from also receiving an upstream `previous_response_id`. +- Opaque thought signatures are extracted from raw provider responses and cached before client formatting. They are never returned to clients, including same-protocol output, and are injected only from provider/credential/session-scoped cache state. +- Centralized top-level provider error normalization handles object, string, and numeric-status envelopes before success parsing. Structured status survives credential rotation, fallback policy, Responses service handling, and protocol-specific HTTP formatting. +- Fixed the public Chat route so successful native dictionaries are returned as successes instead of being mistaken for generic 429 error dictionaries. Raw response logging supports both dictionaries and SDK models. +- Added a 64-case live runtime matrix covering all input/provider/output combinations, a real `agenerate` builder-to-native-executor handoff, real nested callback behavior, two-credential structured-error exhaustion, provider-native hook/validation shape, continuation suppression, response-adapter ordering, and hidden-signature caching/suppression. +- Required iterative `explore` and `explore-heavy` reviews completed. Every confirmed blocker, high, and medium finding was corrected and both reviewers issued final safe-to-commit verdicts. +- Verification: 269 focused protocol/runtime/provider tests passed; all 68 currently tracked test files passed with 763 tests and 18 subtests; compile checks and `git diff --check` passed apart from repository line-ending notices. +- Canonical streaming, public Gemini endpoints, and configuration-defined provider protocols remain explicitly deferred to Phases E and F. + +### 2026-07-16: Public protocol and configurable-provider checkpoint + +- Added first-class Gemini `generateContent` and `countTokens` client routes for + both `/v1beta/models/...` and `/v1/models/...`. Bare Gemini model IDs route to + the Gemini provider unless an explicit model-route alias exists. +- Added independent non-streaming output selection. Precedence is explicit + library argument, `X-Proxy-Output-Protocol`, configured provider default, then + the input protocol. Aliases are normalized case-insensitively and unknown + values fail as client errors. +- Chat, Anthropic, Responses, and Gemini ingress now share that selector. + Responses always executes and stores a native Responses object before an + optional return-format conversion, preserving retrieval and continuation. +- Added config-defined custom providers with startup-snapshotted API base, + provider protocol, endpoint templates, auth-header mode, model list, adapters, + field-cache rules, native-stream capability, and default output protocol. +- Limited configurable provider and output protocols to the four supported + generative protocols. Broader registered operation adapters such as embeddings, + audio, and images are not falsely advertised as interchangeable generation + protocols. +- Kept credentials outside structured config. API bases reject userinfo, query, + and fragments; operation paths remain on the configured origin, reject secret- + bearing query keys and fragments, and permit protocol selectors such as + Gemini's `?alt=sse`. +- Native HTTP transport now treats only 2xx as success and preserves non-success + status plus structured provider bodies for retry, cooldown, and selected- + protocol error formatting. +- Aggregate credential exhaustion now raises a protocol-formatable error without + regressing ordered fallback: structured failure details determine retryability, + and final errors retain all target and attempt summaries. +- Cross-protocol streams are rejected centrally from both `agenerate()` and + direct `acompletion()` entry paths. Gemini `generateContent` also rejects + `stream=true`; the real `streamGenerateContent` route is intentionally coupled + to Phase F's canonical event conversion. +- Iterative `explore` and fresh `explore-heavy` audits found and drove fixes for + header normalization, stream bypasses, endpoint credential safety, native HTTP + errors, Responses error formatting, fallback exhaustion, config snapshots, + non-generative declarations, redirects, and public test tracking. +- Final verification: the complete git-tracked local suite passes 889 tests and + 18 subtests; the widened protocol/public/provider subset passes 390 tests; + `compileall` and staged `git diff --check` pass apart from repository line- + ending notices. +- Final `explore` and `explore-heavy` indexed-state reviews both issued explicit + safe-to-commit verdicts with no unresolved blocker, high, or medium findings. +- Intentional boundaries: custom-provider transport config is snapshotted at + process startup; Gemini `countTokens` is locally estimated; cross-protocol + streaming and the public `streamGenerateContent` route remain Phase F work. + +### 2026-07-16: Canonical streaming checkpoint + +- Added one stateful canonical stream formatter shared by Chat, Anthropic + Messages, Responses, and Gemini. One canonical provider event may expand into + the destination's required start, block, delta, item, usage, terminal, and + error lifecycle frames without replaying provider-native raw fields. +- Added complete four-by-four semantic stream conversion and native + provider/output matrices for text, reasoning, tool calls, fragmented JSON + arguments, usage, completion, failure, hidden signatures, and destination + lifecycle ordering. The actual LiteLLM executor is also exercised against all + four outputs. +- Kept the existing operational stream handler authoritative for timeout, + heartbeat, disconnect, cancellation, cost/accounting, completion-gated session + identity, and retry mechanics. Provider streams normalize to canonical Chat SSE + for this operational pass, then one persistent destination formatter writes the + independently selected client protocol. +- Preserved one destination lifecycle across same-target credential retries. + Anthropic `message_start` and Responses `response.created` are not duplicated + when a role-only or other nonsemantic frame precedes a retryable failure. Each + route-fallback target retains an independent lifecycle because failed-target + pending frames are not exposed. +- Native, LiteLLM, and custom in-band stream errors now enter the same + classification, cooldown, credential-rotation, and route-fallback pipeline. + Provider errors are never yielded as successful content; post-output failures + close the active selected-protocol lifecycle without rotating or changing the + Responses ID. +- Extended visibility policy to Anthropic content/tool events, Responses + reasoning/function items including zero-argument calls, and Gemini candidate + parts. Heartbeats, usage, lifecycle starts, and role-only frames remain + nonsemantic; Chat reasoning-only retry remains behind its existing feature + flag. +- Added a stateful native SSE decoder that assembles `event:` and multiline + `data:` fields at blank-line boundaries and survives byte fragmentation. Bare + NDJSON remains supported. +- Enabled native streaming for Antigravity, Claude Code, Codex, and Copilot. + Stream-event cache rules capture Gemini thought signatures, Anthropic thinking + signatures, and Responses continuation IDs before client formatting within + provider/model/credential/session scope. +- Added first-class Gemini `streamGenerateContent` routes for `/v1beta` and `/v1`. + Anthropic and Responses streaming now use `agenerate()` directly; their old + translation bridges remain only for minimal external facades without that + entry point. Responses stores its terminal Responses object before optional + output reformatting. +- Evidence-less termination is never upgraded to success. Bare `[DONE]`/EOF does + not fabricate OpenAI `stop`, Anthropic `end_turn`, or Gemini `STOP`; Responses + reports `incomplete`. Incomplete Responses items are marked incomplete. +- Gemini tool-call arguments are buffered until one complete JSON object exists, + emitted exactly once, and fail closed on invalid continuation. Genuine + zero-argument calls flush once at terminal. +- Multiple iterative `explore` and `explore-heavy` audits found and drove fixes + for native and in-band errors, failed Responses conversion, zero-argument tool + visibility, lifecycle identity reuse, multiline SSE, Gemini duplicate calls, + active-lifecycle terminal errors, and completion evidence. Both reviewers gave + final safe-to-commit verdicts with no unresolved blocker, high, or medium + findings. +- Final verification: the complete tracked suite plus the force-added stream + matrix passes 958 tests and 18 subtests; `compileall` and `git diff --check` + pass apart from repository line-ending notices. + +### 2026-07-17: Provider-state isolation checkpoint + +- Streaming and non-streaming extraction counterparts for Antigravity thought + signatures, Claude thinking signatures, and Codex continuation IDs now share + one validated logical cache identity. Mixed-mode turns cannot read stale state + from a source-specific cache. +- Cache identity plus every provider, model, credential, session, classifier, + and conversation component uses a full SHA-256 digest. Adversarial delimiter + forms remain distinct and raw scope values never appear in cache keys. +- JSON-configured provider-state rules must retain provider/model/credential/ + session isolation. A configured replacement may tune extraction only while + preserving logical identity, mode, TTL, bounds, injection, continuation, and + tool/turn correlation behavior. +- Continuation destinations are recognized across snake_case, camelCase, and + hyphenated forms. Local Responses lineage suppresses provider continuation by + semantic destination as well as explicit rule metadata. +- `RotatingClient` binds one immutable startup configuration into the executor + and provider singleton instances. Equivalent clients may share it; an + incompatible singleton rebind fails explicitly instead of mutating live + provider behavior. +- Every cache mode enforces serialized byte limits; accumulating and per-tool + modes also enforce value-count limits. Active signature caches retain at most + 1,024 values / 4 MiB for one hour. The process-local store is capped at 10,000 + LRU entries, and ProviderCache uses one atomic backend update for append. +- Legacy injected stores remain compatible without bypassing bounds or masking + unrelated internal `TypeError`s. +- Real `NativeProviderExecutor` contracts cover every stateful provider across + all four output protocols, non-stream -> stream -> non-stream state, model/ + credential/session isolation, missing scopes, Codex local-lineage suppression, + hidden output, and Copilot's explicit no-state behavior. +- Repeated light/heavy audits found and drove fixes for mixed-mode stale state, + key collisions, scope weakening, continuation aliases, config mutation, + unbounded values, append races, singleton rebinding, and mode-specific bounds. + Final heavy verdict: safe to commit. +- Final verification: the complete tracked suite plus the force-added real- + provider contract passes 1,011 tests and 18 subtests; five known + `datetime.utcnow()` deprecation warnings only. Focused provider-state/native/ + config verification passes 298 tests; `compileall` and `git diff --check` pass + apart from repository line-ending notices. + +### 2026-07-17: Final whole-system interoperability verdict + +- Re-ran the complete client-input x provider-protocol x selected-output + contract for non-streaming and streaming behavior, including native, + LiteLLM, and custom execution; public Chat, Anthropic, Responses, and Gemini + routes; fallback and credential rotation; Responses storage/continuation; + provider state; startup config; and non-generative protocol regressions. +- Proxy-side validation, authentication, capacity, timeout, context-window, + and internal failures now use the selected output protocol on every + generative route. Malformed Responses JSON honors an Anthropic or Gemini + selector, and an invalid selector fails with status 400 in the input + protocol rather than producing an internal error. +- Context-window phrase classification is shared by HTTP and structured native + errors, including the common `Context length exceeded` form, so oversized + requests do not rotate credentials or enter route fallback. +- Both custom and code-backed provider `protocol_name` overrides are restricted + at configuration validation to Chat, Responses, Anthropic Messages, or + Gemini. Non-generative adapters remain independently available for their + dedicated operations but cannot be misdeclared as a generative provider. +- Anthropic authentication now honors open-access mode when no proxy API key is + configured, matching Chat and Gemini behavior. +- Removed the obsolete Chat returned-error-dictionary branch. Structured + runtime errors are raised and formatted once, rather than guessed from a + successful dictionary response. +- The tracked suite plus the relevant local-only transaction, error-handler, + and session-forwarding suites pass 1,046 tests and 18 subtests. Five known + `datetime.utcnow()` deprecation warnings remain. `compileall` and + `git diff --check` pass apart from repository line-ending notices. +- The ignored `tests/refactor/` directory is a historical parity suite: 115 + tests still pass and 14 assert superseded cap encodings, state layouts, + transform defaults, or removed private executor methods. It is not used as + current product evidence; current tracked replacements cover those systems. +- Final repeated `explore` and `explore-heavy` reviews report no unresolved + blocker, high, or medium findings and explicitly approve the protocol work + as beta-ready. +- Intentional low-risk boundaries: raw provider transaction captures remain + complete, local, opt-in diagnostics while field-cache traces and all client + output redact opaque state; Gemini token counting is locally estimated; + streaming media/structured-output combinations do not have dedicated matrix + cells beyond the canonical text/reasoning/tool lifecycle coverage. + +**Verdict: the four-protocol interoperability core is beta-ready.** The original +review failures are corrected end to end: input, provider, and output languages +are independent; required meaning is preserved or rejected before transport; +streaming and errors use the selected client language; provider state remains +isolated; and configured providers declare one supported native language +without learning the client's language. + +### Current next action + +Protocol interoperability Phases A-H are complete and beta-ready. No protocol +implementation work remains in this review packet. Continue the wider +experimental-branch beta review with the next concern (transaction/transform +observability noise) only as a separate packet. diff --git a/docs/experimental/third-pass-audit-findings.md b/docs/experimental/third-pass-audit-findings.md new file mode 100644 index 000000000..29658805c --- /dev/null +++ b/docs/experimental/third-pass-audit-findings.md @@ -0,0 +1,361 @@ +# Third-Pass Audit Findings + +This document preserves the complete third-pass validation findings so the remediation pass cannot lose track of them. + +References used by the reviewers: + +- `C:\Projects\test\new 1.txt` +- `C:\Projects\test\new 2.txt` +- General docs: `00-master-plan.md`, `01-protocol-architecture.md`, `02-transform-logging.md`, `03-field-cache-rules.md`, `04-provider-roadmap.md`, `05-routing-retry-usage-roadmap.md`, `06-phase-workflow.md`, `07-detailed-phase-roadmap.md` +- Phase first-pass plans: `phase-1-protocol-core.md` through `phase-10-config-polish.md` +- Phase second-pass plans: `phase-1b-protocol-breadth-operation-model.md` through `phase-10b-config-surface-wiring.md` + +## Overall Verdict + +The third-pass review is not clean. Each phase received one normal `explore` review and one `explore-heavy` review. The reviewers found remaining blockers, highs, mediums, and low-risk residuals across the implemented work. + +## Severity Summary + +| Phase | Highest Severity | Clean? | +|---|---:|---:| +| Phase 1 / 1b Protocols | High | No | +| Phase 2 / 2b Transform Tracing | High | No | +| Phase 3 / 3b Field Cache | Medium | No | +| Phase 4 / 4b Responses | Medium | No | +| Phase 5 / 5b Providers | Blocker | No | +| Phase 6 / 6b Routing/Fallback | Blocker | No | +| Phase 7 / 7b Retry/Cooldown | High | No | +| Phase 8 / 8b Streaming | Medium | No | +| Phase 9 / 9b Usage/Cost | High | No | +| Phase 10 / 10b Config | High | No | + +## Phase 1 / 1b: Protocols + +### Blockers + +- None reported. + +### High + +- Native OpenAI Chat and Responses formatting returns unified usage shape instead of protocol-native usage shape. + - OpenAI Chat `format_response()` can emit `Usage.to_dict()` fields like `input_tokens` / `output_tokens` instead of `prompt_tokens` / `completion_tokens` / `total_tokens`. + - Responses `format_response()` can do the same instead of Responses-compatible usage fields. + - Native executor returns formatted protocol bodies directly to clients, so this can leak wrong usage schemas into live native responses. + +### Medium + +- OpenAI legacy `function_call` is preserved only in `extra`, not modeled as a unified `ToolCall`. +- Ollama lacks a `format_response()` override, so adapter-mutated unified responses can be ignored when `raw` is present. +- Operation compatibility is not enforced in live native provider execution. + - Native execution resolves protocol and operation but does not call `supports_operation()`. + - Misconfigured provider/protocol/operation combinations can silently fall back or drift. + +### Low / Residual + +- No cross-protocol conversion tests such as Anthropic -> unified -> OpenAI or Gemini -> unified -> OpenAI. +- `format_stream_event()` is not overridden by Anthropic, Gemini, Responses, or Ollama; they rely on base raw-pass-through behavior. +- Ollama declares `jsonl` support without JSONL-specific stream formatting. +- Unified types have `to_dict()` but no `from_dict()` deserialization helpers. +- Duplicate helper functions exist across protocol modules. +- No explicit `format_response()` tests for OpenAI Chat, Anthropic Messages, or Gemini. +- `CostDetails` protocol extraction is present in OpenAI Chat and Responses only. +- Registry discovery imports protocol modules without import-error isolation. +- Missing tests for formatted usage schemas, legacy `function_call`, Ollama mutated response formatting, registry failure isolation, and native unsupported operation rejection. + +## Phase 2 / 2b: Transform Tracing + +### Blockers + +- None reported. + +### High + +- Anthropic compatibility path lacks transform-pass tracing for live conversion boundaries. + - Missing explicit traces for Anthropic raw request, Anthropic -> OpenAI conversion, OpenAI -> Anthropic final response, and per-event stream conversion. + - `AnthropicHandler` and `anthropic_streaming_wrapper` mostly rely on legacy transaction logger calls. + +### Medium + +- `raw_provider_stream_response` logs the stream iterator object rather than provider stream data. +- Streaming handler lifecycle traces are gated by `trace_metrics`; chunk traces remain, but lifecycle diagnostics can disappear. +- Non-streaming final response trace can be emitted before usage normalization, making `final_client_response` not truly final. +- Native streaming does not run or trace the stream-event adapter chain. +- Responses streaming does not trace final SSE formatting boundaries. +- Redaction misses common camelCase secret keys such as `apiKey`, `accessToken`, `refreshToken`, and `clientSecret`. +- Non-streaming Responses errors are not traced with standardized transform-failure shape. +- Built-in provider transform exceptions are not logged as transform failures. + +### Low / Residual + +- `ProviderTransforms.apply_sync()` has no trace support; currently not used by live execution. +- `AnthropicHandler.count_tokens()` has no transaction logger or trace entries. +- Provider logger trace entries cannot override per-entry `credential_id`; they rely on context correlation. +- No dedicated test for Anthropic transform tracing. +- Responses service has redundant `if transaction_logger` guards before `_trace()` calls. +- Native provider streaming does not apply adapter chains to stream events. +- Explicit LiteLLM fallback is traceable by metadata but has no pass named `litellm_fallback_request`. + +## Phase 3 / 3b: Adapters And Field Cache + +### Blockers + +- None reported. + +### High + +- None reported. + +### Medium + +- Runtime does not execute all declared field-cache sources and targets. + - `request`, `unified_request`, `unified_response`, and `unified_stream_event` are accepted, but live native execution only injects `request` and extracts `response` / `stream_event`. +- Adapter-chain traces can leak arbitrary configured provider-state fields. + - Native executor applies rule-aware redaction around its own traces, but `run_adapter_chain()` logs payloads directly. +- Credential-scoped field-cache rules do not fail closed when `credential_id` is missing. + - Missing credential can become a shared hashed `_none` bucket in direct/library native executor usage. + +### Low / Residual + +- Copilot provider has no field-cache rules by design. +- `ProviderCacheFieldStore.clear()` depends on injected cache exposing `clear()`. +- Turn-mode inference defaults to common `messages` shape; custom providers need explicit metadata. +- `per_tool_call` injection requires tool-call ID in context metadata or payload path. +- `FieldCacheEngine` defaults to a fresh in-memory store if used directly; `NativeProviderExecutor` correctly shares one per executor instance. +- Native streaming path extracts from stream events but does not inject into stream events, which is expected. +- `ProviderCacheFieldStore.append()` may serialize already wrapped values twice in edge cases. +- Process-local default store is per `NativeProviderExecutor` instance, not a global singleton. +- JSON parsing accepts non-positive `ttl_seconds`, which effectively disables expiry. +- Provider declarations are still conservative/mock-live for several priority providers. + +## Phase 4 / 4b: Responses API + +### Blockers + +- None reported. + +### High + +- None reported. + +### Medium + +- Responses storage is process-local by default; durable JSON/current-state or provider-cache-backed storage is not wired into app startup. +- `previous_response_id` bridge context is lossy. + - The bridge replays parent output only, not parent input or full lineage. +- Responses route errors are wrapped as FastAPI `detail.error` instead of top-level OpenAI-compatible `error` bodies. + +### Low / Residual + +- `ProviderCacheResponsesStore.delete()` returns `False` if the injected cache lacks `delete_async`. +- WebSocket formatter exists but has no FastAPI route or service-neutral integration test. +- `_record_responses_session_anchor()` no-ops if `session_tracker` or `session_id` is unavailable. +- No periodic TTL cleanup task; expired entries prune on access/write. +- Bridge output conversion handles text/tool calls but not all richer Responses item types. +- `store_in_progress` returns opt-in `in_progress` state without a formal retrieval schema. +- `ResponsesStreamState` uses fixed `msg_0` for bridge streams. +- No explicit concurrent access tests for `InMemoryResponsesStore`. +- No cancel endpoint. +- `validate_stream_request()` and `stream_events()` can load the same parent response twice. +- `stream_response(..., transport=...)` always uses SSE formatting. +- Formatted SSE frames themselves are not traced. +- Tests do not include end-to-end FastAPI + real `RotatingClient` sticky continuation or durable app-level storage wiring. + +## Phase 5 / 5b: Provider-Native Integrations + +### Blockers + +- Native provider responses are formatted back to the provider protocol, not the originating client protocol. + - `/v1/chat/completions` routed to Claude Code, Codex, or Antigravity can return Anthropic Messages, Responses API, or Gemini payloads instead of OpenAI Chat payloads. + - Native context does not carry a client target protocol; native executor uses the provider protocol for `format_response()`. + +### High + +- Claude Code native requests can omit required Anthropic `max_tokens`. +- Antigravity model normalization is unsafe for duplicated aliases and can lose low/high thinking-level behavior. + +### Medium + +- `get_native_headers()` and `get_native_endpoint()` are not declared on `ProviderInterface` even though the executor requires them dynamically. +- Native streaming is disabled for all priority providers. This is safe, but it leaves the native streaming path unexercised by real priority providers. +- Claude Code auth header convention may be wrong for API-key credentials because it always uses `Authorization: Bearer` with `anthropic-version`. +- ProviderInterface native contract lacks `supports_native_operation()` / `should_use_native_protocol()` style hooks, making non-stream native auto-selection all-or-nothing by protocol declaration. +- Native streaming fail-closed behavior is duplicated and only partially fail-closed; the executor duplicate catches fewer exceptions than the safe helper. +- Antigravity quota grouping regresses retired parity by grouping models too broadly. +- Field-cache declarations exist, but thinking/signature reinjection is not proven semantically correct for protocol-native continuation requirements. + +### Low / Residual + +- Copilot has no field-cache rules by design. +- Codex message-to-Responses conversion is minimal and does not handle tool calls, multipart content, images, or rich system formatting. +- Antigravity intentionally restores only a conservative safe subset; many retired features remain absent. +- Antigravity has quota groups but no usage reset configs or tier priorities. +- Priority provider model discovery catches broad exceptions and silently falls back to hardcoded lists. +- Execution-mode routing priority is correct and fail-closed. +- Native streaming fail-closed behavior is tested for explicit native streaming. +- Mock-live tests cover direct native request execution for all four priority providers, but not full end-to-end credential rotation/usage manager execution. +- Some provider docstrings still describe skeleton behavior. + +## Phase 6 / 6b: Routing And Fallback + +### Blockers + +- Stale fallback test file imports a missing module, so broad test collection can fail. + - `tests/test_fallback_groups.py` imports `rotator_library.fallback_groups.FallbackGroupManager`, which no longer exists. + - The same file asserts obsolete `RotatingClient(model_fallback_groups=...)` / `fallback_manager` behavior. + +### High + +- Requested-model-in-group promotion is not implemented. + - If the requested model is inside a fallback group, the router should try it first and preserve remaining group order. +- Streaming execution modes do not match non-streaming execution-mode behavior. + - Streaming can ignore explicit `@litellm_fallback` when a plugin has custom logic and can fall through differently for `@custom`. +- Explicit native non-streaming configuration errors can be treated as retryable fallback errors because some `RoutingExecutionError`s default to `unsupported_operation`. +- Structured route-error classification misses common status-less aliases such as `invalid_api_key`, `unauthorized`, `invalid_argument`, `rate_limited`, `too_many_requests`, `resource_exhausted`, and `unavailable`. + +### Medium + +- Fallback target session namespace is cloned from the first target instead of being re-inferred or adjusted per target, risking response anchors under the wrong provider/model namespace. + +### Low / Residual + +- No end-to-end test for env config -> `RequestContextBuilder` -> `RequestExecutor.execute()` full fallback scenario. +- `build_embedding_context()` does not resolve routing groups. +- No explicit test for `_execute_non_streaming_with_fallback()` when all targets return structured error responses. +- `_route_error_type_from_response()` falls back to text-summary scanning when structured fields are ambiguous. +- Fallback exhaustion summaries are sanitized and appear safe. +- Routing route aliases are lowercase keys only; there is no hyphen/space alias normalization. +- Wrapper-level streaming tests are thinner than helper-level coverage for some event frames. + +## Phase 7 / 7b: Retry, Cooldown, Backoff, Failure History + +### Blockers + +- None reported. + +### High + +- Provider/model cooldown can be bypassed when the remaining cooldown exceeds the request deadline budget. + - `_wait_for_cooldown()` logs and returns, and callers continue into credential acquisition/execution. +- Generic transient errors start provider-wide cooldown too aggressively. + - A single `server_error` or `api_connection` without `retry_after` can start provider cooldown immediately. + +### Medium + +- `decide_streaming_error_action()` is implemented and tested but not used by the live executor. +- Failure history lacks success reset; successes do not clear or reduce provider/model failure history. +- Structured retry attempt history is not populated beyond traces/summaries; `routing_attempt_history` exists but is unused in live executor. + +### Low / Residual + +- `FallbackAttemptRunner` is tested but not used by the live executor; inline wrappers remain. +- `decide_streaming_error_action()` does not accept `failure_history`, which would matter if it is adopted live. +- No `test_request_executor_group_policy.py` file despite the plan mentioning it. +- Streaming error branches in the executor are duplicated across several handlers. +- Model-capacity detection is narrow and misses variants such as `CAPACITY_EXHAUSTED` or `capacity-exhausted`. +- Tests miss cooldown-over-budget fail-fast/fallback, single generic transient non-cooldown, success reset, live model cooldown not blocking unrelated models, and live streaming parity with the decision helper. + +## Phase 8 / 8b: Streaming Hardening + +### Blockers + +- None reported. + +### High + +- None reported. + +### Medium + +- Responses streaming bypasses active Phase 8b hardening at the Responses layer. + - No direct TTFB/stall/heartbeat/disconnect/close policy is applied in `ResponsesService.stream_events()`. + - Responses formatter has no heartbeat formatter. +- Anthropic compatibility streaming can stop on disconnect without explicitly closing the upstream OpenAI stream. + +### Low / Residual + +- No formal `StreamTransportFormatter` protocol/ABC. +- Executor streaming error handling remains duplicated across branches. +- Responses streaming does not emit heartbeats or enforce TTFB/stall policy directly. +- Native streaming executor does not integrate with `StreamingHandler.wrap_stream()`. +- Native `_parse_stream_line()` returns raw text on JSON parse failure. +- No specific test for stall timeout after visible output. +- Normalized stream event parser comment says malformed chunks fail closed for visibility, but parser returns `visible_output=False`; routing safety is protected elsewhere. +- `decide_streaming_error_action()` is not used by the live executor. + +## Phase 9 / 9b: Usage, Quota, Cost + +### Blockers + +- None reported. + +### High + +- Provider-reported top-level response costs are ignored when a `usage` object exists. + - `extract_usage_record()` unwraps full responses to `response["usage"]`, discarding sibling `cost`, `total_cost`, `cost_details`, or `provider_reported_cost`. +- OpenAI-like cache-write tokens can be double-counted in normalized totals/costs. + - `input_tokens = prompt_tokens - cache_read`, while `cache_write_tokens` is stored separately and also added to total. +- Responses streaming drops SSE cost comments/events unless final usage carries cost. + +### Medium + +- `event: cost` frames can block streaming retry/fallback despite being metadata. +- Native streaming executor does not emit `usage_accounting_summary` trace. +- Native streaming executor does not preserve provider-reported cost from raw response. +- Structured provider cost breakdowns without `total_cost` are not preserved or summed. + +### Low / Residual + +- Responses bridge passes cost fields through verbatim rather than structuring them, which is acceptable but should be documented. +- No explicit native streaming cost-trace test. +- Streaming cost calculator does not pass `provider_plugin`, so provider explicit pricing can be missed in streaming mode when no provider-reported cost exists. +- Quota snapshot `used` maps to request count, not token count; metadata says request/token window but field naming is ambiguous. +- No integration-level test for Responses usage with provider-reported cost. +- No quota snapshot test for combined model and group filters. +- Planned trace pass names like `usage_cost_calculated` or `quota_snapshot_built` are not implemented if considered distinct from `usage_accounting_summary`. + +## Phase 10 / 10b: Config + +### Blockers + +- None reported. + +### High + +- Full `PROXY_API_KEY` is printed to console at startup. +- Provider/protocol/adapter/model/quota-checker config surfaces are still not live-wired, despite the `providers` JSON section being accepted/documented. + +### Medium + +- Responses streaming ignores Phase 10 streaming runtime settings. +- Pricing env parsing can fail requests instead of warning/ignoring invalid values. +- Routing config validation is not fully startup-safe; direct model-route target specs can parse later at request time. +- Secret rejection is key-name based and misses generic `credential` / `credentials` key names. + +### Low / Residual + +- `STREAM_RETRY_ON_REASONING_ONLY` is env-only and not part of `StreamRuntimeSettings` JSON config. +- `get_responses_store_settings()` returns `Any` rather than a precise type. +- `provider_cooldown_env()` remains a legacy tuple wrapper over richer retry settings. +- `STREAM_HEARTBEAT_SECONDS` legacy alias resolution is subtle. +- No standalone `config-reference.md` was created. +- `providers` section accepts raw dicts but does not validate protocol/adapter names; this is conservative but incomplete for future wiring. +- `.env.example` drift test does not verify the legacy `STREAM_HEARTBEAT_SECONDS=0` default line. + +## Cross-Phase Remediation Order + +1. Phase 5 blocker: native provider output must be converted back to the originating client protocol. +2. Phase 6 blocker: retire/update stale fallback tests and fix routing execution correctness. +3. Phase 10 high: stop printing secrets and make accepted config surfaces honest/live-wired or reject unsupported shapes. +4. Phase 1 high/medium protocol formatting and operation enforcement. +5. Phase 9 high usage/cost correctness. +6. Phase 7 high cooldown budget and transient cooldown behavior. +7. Phase 2 tracing gaps, especially Anthropic compatibility and redaction. +8. Phase 3 field-cache runtime/redaction gaps. +9. Phase 4 Responses storage/continuation/error-shape gaps. +10. Phase 8 streaming hardening gaps for Responses and Anthropic compatibility. + +## Required Completion Bar + +- Every blocker/high/medium listed above must be fixed or explicitly re-reviewed as no longer applicable. +- Low-risk items should be fixed when small/safe; otherwise they must be carried into phase reports as explicit residuals. +- Each remediation phase must follow the established loop: plan in conversation, write plan doc, implement focused commits, run focused tests, run `explore` and `explore-heavy`, fix findings, re-review, write user-facing report without committing the report. diff --git a/src/proxy_app/detailed_logger.py b/src/proxy_app/detailed_logger.py index 07d346ffb..7ae5ff750 100644 --- a/src/proxy_app/detailed_logger.py +++ b/src/proxy_app/detailed_logger.py @@ -5,7 +5,8 @@ """ Raw I/O Logger for the Proxy Layer. -This logger captures the UNMODIFIED HTTP request and response at the proxy boundary. +This logger captures HTTP request and response data at the proxy boundary while +always redacting authentication headers. It is disabled by default and should only be enabled for debugging the proxy itself. Use this when you need to: @@ -18,7 +19,7 @@ Directory structure: logs/raw_io/{YYYYMMDD_HHMMSS}_{request_id}/ - request.json # Unmodified incoming HTTP request + request.json # Incoming request with auth headers redacted streaming_chunks.jsonl # If streaming mode final_response.json # Unmodified outgoing HTTP response metadata.json # Summary metadata @@ -27,7 +28,7 @@ import json import time import uuid -from datetime import datetime +from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, Optional import logging @@ -52,8 +53,8 @@ class RawIOLogger: """ Logs raw HTTP request/response at the proxy boundary. - This captures the EXACT data as received from and sent to the client, - without any transformations. Useful for debugging the proxy itself. + Payloads remain otherwise unchanged, but bearer credentials and cookies are + never persisted. Useful for debugging the proxy itself. DISABLED by default. Enable with --enable-raw-logging flag. @@ -95,18 +96,38 @@ def log_request(self, headers: Dict[str, Any], body: Dict[str, Any]): self.streaming = body.get("stream", False) request_data = { "request_id": self.request_id, - "timestamp_utc": datetime.utcnow().isoformat(), - "headers": dict(headers), + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "headers": self._redact_headers(headers), "body": body, } self._write_json("request.json", request_data) + @staticmethod + def _redact_headers(headers: Dict[str, Any]) -> Dict[str, Any]: + """Remove transport credentials while preserving header structure.""" + + sensitive = { + "authorization", + "cookie", + "proxy-authorization", + "set-cookie", + "x-api-key", + "x-proxy-session-domain", + } + return { + key: "" if str(key).lower() in sensitive else value + for key, value in headers.items() + } + def log_stream_chunk(self, chunk: Dict[str, Any]): """Logs an individual chunk from a streaming response to a JSON Lines file.""" if not self._dir_available: return - log_entry = {"timestamp_utc": datetime.utcnow().isoformat(), "chunk": chunk} + log_entry = { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "chunk": chunk, + } content = json.dumps(log_entry, ensure_ascii=False) + "\n" safe_log_write(self.log_dir / "streaming_chunks.jsonl", content, logging) @@ -119,10 +140,10 @@ def log_final_response( response_data = { "request_id": self.request_id, - "timestamp_utc": datetime.utcnow().isoformat(), + "timestamp_utc": datetime.now(timezone.utc).isoformat(), "status_code": status_code, "duration_ms": round(duration_ms), - "headers": dict(headers) if headers else None, + "headers": self._redact_headers(headers) if headers else None, "body": body, } self._write_json("final_response.json", response_data) diff --git a/src/proxy_app/launcher_tui.py b/src/proxy_app/launcher_tui.py index f8c1eacaf..ebc27644e 100644 --- a/src/proxy_app/launcher_tui.py +++ b/src/proxy_app/launcher_tui.py @@ -185,9 +185,7 @@ def detect_credentials() -> dict: # Scan for env-based OAuth credentials # Maps provider name to the ENV_PREFIX used by the provider # (duplicated from credential_manager to avoid heavy imports) - env_oauth_providers = { - "gemini_cli": "GEMINI_CLI", - } + env_oauth_providers = {} for provider, env_prefix in env_oauth_providers.items(): oauth_count = 0 @@ -313,7 +311,7 @@ def detect_model_filters() -> dict: @staticmethod def detect_provider_settings() -> dict: - """Detect provider-specific settings (Gemini CLI)""" + """Detect provider-specific settings.""" try: from proxy_app.settings_tool import PROVIDER_SETTINGS_MAP except ImportError: diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 4aa689720..af244b07c 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -2,7 +2,6 @@ # Copyright (c) 2026 Mirrowel import time -import uuid # Phase 1: Minimal imports for arg parsing and TUI import asyncio @@ -61,7 +60,7 @@ # Load all .env files from root folder (main .env first, then any additional *.env files) from dotenv import load_dotenv -from glob import glob +from proxy_app.startup_display import mask_secret_for_display as _mask_secret_for_display # Get the application root directory (EXE dir if frozen, else CWD) # Inlined here to avoid triggering heavy rotator_library imports before loading screen @@ -73,7 +72,7 @@ # Load main .env first load_dotenv(_root_dir / ".env") -# Load any additional .env files (e.g., gemini_cli_all_combined.env) +# Load any additional .env files (e.g., provider_credentials.env) _env_files_found = list(_root_dir.glob("*.env")) for _env_file in sorted(_root_dir.glob("*.env")): if _env_file.name != ".env": # Skip main .env (already loaded) @@ -84,10 +83,11 @@ _env_names = [_ef.name for _ef in _env_files_found] print(f"📁 Loaded {len(_env_files_found)} .env file(s): {', '.join(_env_names)}") + # Get proxy API key for display proxy_api_key = os.getenv("PROXY_API_KEY") if proxy_api_key: - key_display = f"✓ {proxy_api_key}" + key_display = f"✓ {_mask_secret_for_display(proxy_api_key)}" else: key_display = "✗ Not Set (INSECURE - anyone can access!)" @@ -115,7 +115,6 @@ print(" → Loading core dependencies...") with _console.status("[dim]Loading core dependencies...", spinner="dots"): - from dotenv import load_dotenv import colorlog import json from typing import AsyncGenerator, Any, List, Optional, Union @@ -134,12 +133,15 @@ print(" → Initializing proxy core...") with _console.status("[dim]Initializing proxy core...", spinner="dots"): from rotator_library import RotatingClient + from rotator_library.client.protocol_selection import format_client_protocol_error, require_same_protocol_stream, resolve_client_output_protocol from rotator_library.credential_manager import CredentialManager - from rotator_library.background_refresher import BackgroundRefresher from rotator_library.model_info_service import init_model_info_service from proxy_app.request_logger import log_request_to_console from proxy_app.batch_manager import EmbeddingBatcher from proxy_app.detailed_logger import RawIOLogger + from rotator_library.responses import ResponsesService, ResponsesServiceError + from rotator_library.core.errors import StructuredAPIResponseError + from rotator_library.transaction_logger import TransactionLogger print(" → Discovering provider plugins...") # Provider lazy loading happens during import, so time it here @@ -563,10 +565,9 @@ async def process_credential(provider: str, path: str, provider_instance): logging.info("OAuth credential processing complete.") oauth_credentials = final_oauth_credentials - # [NEW] Load provider-specific params - litellm_provider_params = { - "gemini_cli": {"project_id": os.getenv("GEMINI_CLI_PROJECT_ID")} - } + # Provider-specific LiteLLM params. API-key Gemini remains configured through + # normal provider environment keys. + litellm_provider_params = {} # Load global timeout from environment (default 30 seconds) global_timeout = int(os.getenv("GLOBAL_TIMEOUT", "30")) @@ -592,6 +593,13 @@ async def process_credential(provider: str, path: str, provider_instance): # print(f"🔑 Credentials loaded: {_total_summary} (API: {_api_summary} | OAuth: {_oauth_summary})") client.background_refresher.start() # Start the background task app.state.rotating_client = client + # Phase 4 Responses API compatibility service. It currently bridges through + # the existing chat-completions client path; later native providers can reuse + # the same route/storage surface without changing clients. + from rotator_library.config.experimental import get_responses_store_settings + from rotator_library.responses import create_configured_responses_store + + app.state.responses_service = ResponsesService(store=create_configured_responses_store(), store_settings=get_responses_store_settings()) # Warn if no provider credentials are configured if not client.all_credentials: @@ -627,6 +635,9 @@ async def process_credential(provider: str, path: str, provider_instance): await client.background_refresher.stop() # Stop the background task on shutdown if app.state.embedding_batcher: await app.state.embedding_batcher.stop() + responses_service = getattr(app.state, "responses_service", None) + if responses_service: + await responses_service.close() await client.close() # Stop model info service @@ -663,6 +674,19 @@ def get_embedding_batcher(request: Request) -> EmbeddingBatcher: return request.app.state.embedding_batcher +def get_responses_service(request: Request) -> ResponsesService: + """Dependency to get the Responses API service instance from app state.""" + + service = getattr(request.app.state, "responses_service", None) + if service is None: + from rotator_library.config.experimental import get_responses_store_settings + from rotator_library.responses import create_configured_responses_store + + service = ResponsesService(store=create_configured_responses_store(), store_settings=get_responses_store_settings()) + request.app.state.responses_service = service + return service + + async def verify_api_key(auth: str = Depends(api_key_header)): """Dependency to verify the proxy API key.""" # If PROXY_API_KEY is not set or empty, skip verification (open access) @@ -675,6 +699,7 @@ async def verify_api_key(auth: str = Depends(api_key_header)): # --- Anthropic API Key Header --- anthropic_api_key_header = APIKeyHeader(name="x-api-key", auto_error=False) +gemini_api_key_header = APIKeyHeader(name="x-goog-api-key", auto_error=False) async def verify_anthropic_api_key( @@ -685,6 +710,8 @@ async def verify_anthropic_api_key( Dependency to verify API key for Anthropic endpoints. Accepts either x-api-key header (Anthropic style) or Authorization Bearer (OpenAI style). """ + if not PROXY_API_KEY: + return x_api_key or auth # Check x-api-key first (Anthropic style) if x_api_key and x_api_key == PROXY_API_KEY: return x_api_key @@ -694,6 +721,25 @@ async def verify_anthropic_api_key( raise HTTPException(status_code=401, detail="Invalid or missing API Key") +async def verify_gemini_api_key( + request: Request, + x_api_key: str = Depends(gemini_api_key_header), + auth: str = Depends(api_key_header), +): + """Accept Gemini header/query authentication or the shared Bearer form.""" + + if not PROXY_API_KEY: + return x_api_key or auth or request.query_params.get("key") + query_key = request.query_params.get("key") + if ( + x_api_key == PROXY_API_KEY + or query_key == PROXY_API_KEY + or auth == f"Bearer {PROXY_API_KEY}" + ): + return x_api_key or query_key or auth + raise HTTPException(status_code=401, detail="Invalid or missing API Key") + + async def streaming_response_wrapper( request: Request, request_data: dict, @@ -886,12 +932,22 @@ async def chat_completions( """ # Raw I/O logger captures unmodified HTTP data at proxy boundary (disabled by default) raw_logger = RawIOLogger() if ENABLE_RAW_LOGGING else None + request_data: dict[str, Any] = {} try: # Read and parse the request body only once at the beginning. try: request_data = await request.json() except json.JSONDecodeError: - raise HTTPException(status_code=400, detail="Invalid JSON in request body.") + status, content = format_client_protocol_error( + client, + request_data, + input_protocol="openai_chat", + request=request, + error="Invalid JSON in request body.", + error_type="invalid_request", + status_code=400, + ) + return JSONResponse(status_code=status, content=content) # Global temperature=0 override (controlled by .env variable, default: OFF) # Low temperature makes models deterministic and prone to following training data @@ -946,8 +1002,17 @@ async def chat_completions( is_streaming = request_data.get("stream", False) if is_streaming: - response_generator = await client.acompletion( - request=request, **request_data + selected_output = resolve_client_output_protocol( + client, + request_data, + input_protocol="openai_chat", + request=request, + ) + require_same_protocol_stream("openai_chat", selected_output) + response_generator = await client.agenerate( + request_data, + input_protocol="openai_chat", + request=request, ) return StreamingResponse( streaming_response_wrapper( @@ -956,15 +1021,11 @@ async def chat_completions( media_type="text/event-stream", ) else: - response = await client.acompletion(request=request, **request_data) - - if isinstance(response, dict): - if raw_logger: - raw_logger.log_final_response( - status_code=429, headers=None, body=response - ) - error_detail = response.get("error", {}).get("message", str(response)) - raise HTTPException(status_code=429, detail=error_detail) + response = await client.agenerate( + request_data, + input_protocol="openai_chat", + request=request, + ) if raw_logger: response_headers = ( @@ -976,26 +1037,67 @@ async def chat_completions( raw_logger.log_final_response( status_code=status_code, headers=response_headers, - body=response.model_dump(), + body=response if isinstance(response, dict) else response.model_dump(), ) return response + except StructuredAPIResponseError as e: + output_protocol = resolve_client_output_protocol( + client, + request_data if isinstance(request_data, dict) else {}, + input_protocol="openai_chat", + request=request, + ) + return JSONResponse(status_code=e.http_status, content=e.to_protocol_payload(output_protocol)) except ( litellm.InvalidRequestError, ValueError, litellm.ContextWindowExceededError, ) as e: - raise HTTPException(status_code=400, detail=f"Invalid Request: {str(e)}") + status, content = format_client_protocol_error( + client, + request_data if isinstance(request_data, dict) else {}, + input_protocol="openai_chat", + request=request, + error=e, + error_type=( + "context_window_exceeded" + if isinstance(e, litellm.ContextWindowExceededError) + else "invalid_request" + ), + status_code=400, + ) + return JSONResponse(status_code=status, content=content) except litellm.AuthenticationError as e: - raise HTTPException(status_code=401, detail=f"Authentication Error: {str(e)}") + status, content = format_client_protocol_error( + client, request_data, input_protocol="openai_chat", request=request, + error=e, error_type="authentication", status_code=401, + ) + return JSONResponse(status_code=status, content=content) except litellm.RateLimitError as e: - raise HTTPException(status_code=429, detail=f"Rate Limit Exceeded: {str(e)}") + status, content = format_client_protocol_error( + client, request_data, input_protocol="openai_chat", request=request, + error=e, error_type="rate_limit", status_code=429, + ) + return JSONResponse(status_code=status, content=content) except (litellm.ServiceUnavailableError, litellm.APIConnectionError) as e: - raise HTTPException(status_code=503, detail=f"Service Unavailable: {str(e)}") + status, content = format_client_protocol_error( + client, request_data, input_protocol="openai_chat", request=request, + error=e, error_type="server_error", status_code=503, + ) + return JSONResponse(status_code=status, content=content) except litellm.Timeout as e: - raise HTTPException(status_code=504, detail=f"Gateway Timeout: {str(e)}") + status, content = format_client_protocol_error( + client, request_data, input_protocol="openai_chat", request=request, + error=e, error_type="proxy_timeout", status_code=504, + ) + return JSONResponse(status_code=status, content=content) except (litellm.InternalServerError, litellm.OpenAIError) as e: - raise HTTPException(status_code=502, detail=f"Bad Gateway: {str(e)}") + status, content = format_client_protocol_error( + client, request_data, input_protocol="openai_chat", request=request, + error=e, error_type="server_error", status_code=502, + ) + return JSONResponse(status_code=status, content=content) except Exception as e: logging.error(f"Request failed after all retries: {e}") # Optionally log the failed request @@ -1008,7 +1110,198 @@ async def chat_completions( raw_logger.log_final_response( status_code=500, headers=None, body={"error": str(e)} ) - raise HTTPException(status_code=500, detail=str(e)) + status, content = format_client_protocol_error( + client, + request_data if isinstance(request_data, dict) else {}, + input_protocol="openai_chat", + request=request, + error=e, + error_type="internal_error", + status_code=500, + ) + return JSONResponse(status_code=status, content=content) + + +def _responses_error_response( + error: ResponsesServiceError, + protocol: str = "responses", +) -> dict[str, Any]: + """Return a Responses service failure in the selected client protocol.""" + + return error.to_protocol_payload(protocol) + + +@app.post("/v1/responses") +async def responses_create( + request: Request, + client: RotatingClient = Depends(get_rotating_client), + service: ResponsesService = Depends(get_responses_service), + _=Depends(verify_api_key), +): + """Create, store, and optionally reformat an OpenAI Responses object.""" + + logger = RawIOLogger() if ENABLE_RAW_LOGGING else None + selected_output = "responses" + try: + request_data = await request.json() + except json.JSONDecodeError: + status, content = format_client_protocol_error( + client, + {}, + input_protocol="responses", + request=request, + error="Invalid JSON in request body.", + error_type="invalid_request", + status_code=400, + ) + return JSONResponse(status_code=status, content=content) + if logger: + logger.log_request( + headers=dict(request.headers), + body=service.redact_request_for_logging(request_data), + ) + try: + selected_output = resolve_client_output_protocol( + client, + request_data, + input_protocol="responses", + request=request, + ) + except ValueError as e: + status, content = format_client_protocol_error( + client, + request_data, + input_protocol="responses", + request=request, + error=e, + error_type="invalid_request", + status_code=400, + ) + return JSONResponse(status_code=status, content=content) + transaction_logger = TransactionLogger("responses", request_data.get("model", "unknown")) if ENABLE_REQUEST_LOGGING else None + try: + request_scope = service.prepare_request_scope(request_data) + previous_response_access_token = request.headers.get( + "X-Proxy-Session-Domain" + ) + if request_data.get("stream"): + require_same_protocol_stream("responses", selected_output) + await service.validate_stream_request( + request_data, + request_scope=request_scope, + previous_response_access_token=previous_response_access_token, + ) + return StreamingResponse( + service.stream_response( + request_data, + client, + request=request, + transaction_logger=transaction_logger, + request_scope=request_scope, + previous_response_access_token=previous_response_access_token, + ), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + "X-Proxy-Session-Domain": request_scope.access_token, + }, + ) + result = await service.create_response( + request_data, + client, + request=request, + transaction_logger=transaction_logger, + request_scope=request_scope, + previous_response_access_token=previous_response_access_token, + ) + if logger: + logger.log_final_response(status_code=200, headers=None, body=result) + return JSONResponse( + content=result, + headers={"X-Proxy-Session-Domain": request_scope.access_token}, + ) + except ResponsesServiceError as e: + payload = _responses_error_response(e, selected_output) + if logger: + logger.log_final_response(status_code=e.status_code, headers=None, body=payload) + return JSONResponse(status_code=e.status_code, content=payload) + except ValueError as e: + payload = _responses_error_response( + ResponsesServiceError(str(e), status_code=400), + selected_output, + ) + return JSONResponse(status_code=400, content=payload) + except Exception as e: + logging.error(f"Responses endpoint error: {e}") + payload = _responses_error_response( + ResponsesServiceError(str(e), status_code=500, error_type="internal_error"), + selected_output, + ) + if logger: + logger.log_final_response(status_code=500, headers=None, body=payload) + return JSONResponse(status_code=500, content=payload) + + +@app.get("/v1/responses/{response_id}") +async def responses_get( + response_id: str, + request: Request, + service: ResponsesService = Depends(get_responses_service), + _=Depends(verify_api_key), +): + """Retrieve a stored Responses object by ID.""" + + try: + return JSONResponse( + content=await service.get_response_with_access_token( + response_id, + request.headers.get("X-Proxy-Session-Domain", "public"), + ) + ) + except ResponsesServiceError as e: + return JSONResponse(status_code=e.status_code, content=_responses_error_response(e)) + + +@app.delete("/v1/responses/{response_id}") +async def responses_delete( + response_id: str, + request: Request, + service: ResponsesService = Depends(get_responses_service), + _=Depends(verify_api_key), +): + """Delete a stored Responses object by ID.""" + + try: + return JSONResponse( + content=await service.delete_response_with_access_token( + response_id, + request.headers.get("X-Proxy-Session-Domain", "public"), + ) + ) + except ResponsesServiceError as e: + return JSONResponse(status_code=e.status_code, content=_responses_error_response(e)) + + +@app.get("/v1/responses/{response_id}/input_items") +async def responses_input_items( + response_id: str, + request: Request, + service: ResponsesService = Depends(get_responses_service), + _=Depends(verify_api_key), +): + """Return stored input items for a Responses object.""" + + try: + return JSONResponse( + content=await service.list_input_items_with_access_token( + response_id, + request.headers.get("X-Proxy-Session-Domain", "public"), + ) + ) + except ResponsesServiceError as e: + return JSONResponse(status_code=e.status_code, content=_responses_error_response(e)) # --- Anthropic Messages API Endpoint --- @@ -1022,8 +1315,8 @@ async def anthropic_messages( """ Anthropic-compatible Messages API endpoint. - Accepts requests in Anthropic's format and returns responses in Anthropic's format. - Internally translates to OpenAI format for processing via LiteLLM. + Accepts Anthropic requests and returns the selected client protocol through + the shared canonical runtime. This endpoint is compatible with Claude Code and other Anthropic API clients. """ @@ -1050,6 +1343,14 @@ async def anthropic_messages( ) # Use the library method to handle the request + if body.stream: + selected_output = resolve_client_output_protocol( + client, + body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", + request=request, + ) + require_same_protocol_stream("anthropic_messages", selected_output) result = await client.anthropic_messages(body, raw_request=request) if body.stream: @@ -1073,40 +1374,54 @@ async def anthropic_messages( ) return JSONResponse(content=result) + except StructuredAPIResponseError as e: + output_protocol = resolve_client_output_protocol( + client, + body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", + request=request, + ) + return JSONResponse(status_code=e.http_status, content=e.to_protocol_payload(output_protocol)) except ( litellm.InvalidRequestError, ValueError, litellm.ContextWindowExceededError, ) as e: - error_response = { - "type": "error", - "error": {"type": "invalid_request_error", "message": str(e)}, - } - raise HTTPException(status_code=400, detail=error_response) + status, content = format_client_protocol_error( + client, body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", request=request, error=e, + error_type=("context_window_exceeded" if isinstance(e, litellm.ContextWindowExceededError) else "invalid_request"), + status_code=400, + ) + return JSONResponse(status_code=status, content=content) except litellm.AuthenticationError as e: - error_response = { - "type": "error", - "error": {"type": "authentication_error", "message": str(e)}, - } - raise HTTPException(status_code=401, detail=error_response) + status, content = format_client_protocol_error( + client, body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", request=request, error=e, + error_type="authentication", status_code=401, + ) + return JSONResponse(status_code=status, content=content) except litellm.RateLimitError as e: - error_response = { - "type": "error", - "error": {"type": "rate_limit_error", "message": str(e)}, - } - raise HTTPException(status_code=429, detail=error_response) + status, content = format_client_protocol_error( + client, body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", request=request, error=e, + error_type="rate_limit", status_code=429, + ) + return JSONResponse(status_code=status, content=content) except (litellm.ServiceUnavailableError, litellm.APIConnectionError) as e: - error_response = { - "type": "error", - "error": {"type": "api_error", "message": str(e)}, - } - raise HTTPException(status_code=503, detail=error_response) + status, content = format_client_protocol_error( + client, body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", request=request, error=e, + error_type="server_error", status_code=503, + ) + return JSONResponse(status_code=status, content=content) except litellm.Timeout as e: - error_response = { - "type": "error", - "error": {"type": "api_error", "message": f"Request timed out: {str(e)}"}, - } - raise HTTPException(status_code=504, detail=error_response) + status, content = format_client_protocol_error( + client, body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", request=request, error=e, + error_type="proxy_timeout", status_code=504, + ) + return JSONResponse(status_code=status, content=content) except Exception as e: logging.error(f"Anthropic messages endpoint error: {e}") if logger: @@ -1115,11 +1430,12 @@ async def anthropic_messages( headers=None, body={"error": str(e)}, ) - error_response = { - "type": "error", - "error": {"type": "api_error", "message": str(e)}, - } - raise HTTPException(status_code=500, detail=error_response) + status, content = format_client_protocol_error( + client, body.model_dump(exclude_none=True), + input_protocol="anthropic_messages", request=request, error=e, + error_type="internal_error", status_code=500, + ) + return JSONResponse(status_code=status, content=content) # --- Anthropic Count Tokens Endpoint --- @@ -1168,6 +1484,131 @@ async def anthropic_count_tokens( raise HTTPException(status_code=500, detail=error_response) +@app.post("/v1beta/models/{model:path}:generateContent") +@app.post("/v1/models/{model:path}:generateContent") +async def gemini_generate_content( + model: str, + request: Request, + client: RotatingClient = Depends(get_rotating_client), + _=Depends(verify_gemini_api_key), +): + """Gemini-compatible generation endpoint using protocol-native routing.""" + + payload: dict[str, Any] = {} + try: + payload = await request.json() + except json.JSONDecodeError: + status, content = format_client_protocol_error( + client, payload, input_protocol="gemini", request=request, + error="Invalid JSON in request body.", error_type="invalid_request", + status_code=400, + ) + return JSONResponse(status_code=status, content=content) + try: + if payload.get("stream"): + raise ValueError( + "Gemini generateContent does not accept stream=true; use streamGenerateContent" + ) + result = await client.gemini_generate(payload, model=model, raw_request=request) + return JSONResponse(content=result) + except StructuredAPIResponseError as error: + output_protocol = resolve_client_output_protocol( + client, + {**payload, "model": model}, + input_protocol="gemini", + request=request, + ) + return JSONResponse(status_code=error.http_status, content=error.to_protocol_payload(output_protocol)) + except (ValueError, litellm.InvalidRequestError) as error: + status, content = format_client_protocol_error( + client, {**payload, "model": model}, input_protocol="gemini", + request=request, error=error, error_type="invalid_request", + status_code=400, + ) + return JSONResponse(status_code=status, content=content) + except Exception as error: + status, content = format_client_protocol_error( + client, {**payload, "model": model}, input_protocol="gemini", + request=request, error=error, error_type="internal_error", + status_code=500, + ) + return JSONResponse(status_code=status, content=content) + + +@app.post("/v1beta/models/{model:path}:streamGenerateContent") +@app.post("/v1/models/{model:path}:streamGenerateContent") +async def gemini_stream_generate_content( + model: str, + request: Request, + client: RotatingClient = Depends(get_rotating_client), + _=Depends(verify_gemini_api_key), +): + """Gemini-compatible streaming generation through canonical routing.""" + + payload: dict[str, Any] = {} + try: + payload = await request.json() + response_stream = await client.gemini_stream_generate( + payload, + model=model, + raw_request=request, + ) + return StreamingResponse( + streaming_response_wrapper(request, payload, response_stream), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + except (json.JSONDecodeError, ValueError, litellm.InvalidRequestError) as error: + status, content = format_client_protocol_error( + client, {**payload, "model": model}, input_protocol="gemini", + request=request, error=error, error_type="invalid_request", + status_code=400, + ) + return JSONResponse(status_code=status, content=content) + except StructuredAPIResponseError as error: + output_protocol = resolve_client_output_protocol( + client, + {**payload, "model": model}, + input_protocol="gemini", + request=request, + ) + return JSONResponse( + status_code=error.http_status, + content=error.to_protocol_payload(output_protocol), + ) + except Exception as error: + status, content = format_client_protocol_error( + client, {**payload, "model": model}, input_protocol="gemini", + request=request, error=error, error_type="internal_error", + status_code=500, + ) + return JSONResponse(status_code=status, content=content) + + +@app.post("/v1beta/models/{model:path}:countTokens") +@app.post("/v1/models/{model:path}:countTokens") +async def gemini_count_tokens( + model: str, + request: Request, + client: RotatingClient = Depends(get_rotating_client), + _=Depends(verify_gemini_api_key), +): + """Gemini-compatible token counting endpoint.""" + + try: + payload = await request.json() + return JSONResponse(content=client.gemini_count_tokens(payload, model=model)) + except (json.JSONDecodeError, ValueError) as error: + return JSONResponse( + status_code=400, + content={"error": {"code": 400, "message": str(error), "status": "INVALID_ARGUMENT"}}, + ) + + @app.post("/v1/embeddings") async def embeddings( request: Request, @@ -1406,8 +1847,8 @@ async def refresh_quota_stats( { "action": "reload" | "force_refresh", "scope": "all" | "provider" | "credential", - "provider": "gemini_cli", // required if scope != "all" - "credential": "gemini_cli_oauth_1.json" // required if scope == "credential" + "provider": "openai", // required if scope != "all" + "credential": "openai_key_1" // required if scope == "credential" } Actions: diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index e25a1e08e..3ade98141 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -31,15 +31,6 @@ # Sentinel value for distinguishing "no pending change" from "pending change to None" _NOT_FOUND = object() -# Import default OAuth port values from provider modules -# These serve as the source of truth for default port values -try: - from rotator_library.providers.gemini_auth_base import GeminiAuthBase - - GEMINI_CLI_DEFAULT_OAUTH_PORT = GeminiAuthBase.CALLBACK_PORT -except ImportError: - GEMINI_CLI_DEFAULT_OAUTH_PORT = 8085 - def clear_screen(subtitle: str = ""): """ Cross-platform terminal clear with optional header. @@ -493,59 +484,8 @@ def remove_multiplier(self, provider: str, priority: int): # PROVIDER-SPECIFIC SETTINGS DEFINITIONS # ============================================================================= -# Gemini CLI provider environment variables -GEMINI_CLI_SETTINGS = { - "GEMINI_CLI_SIGNATURE_CACHE_TTL": { - "type": "int", - "default": 3600, - "description": "Memory cache TTL for thought signatures (seconds)", - }, - "GEMINI_CLI_SIGNATURE_DISK_TTL": { - "type": "int", - "default": 86400, - "description": "Disk cache TTL for thought signatures (seconds)", - }, - "GEMINI_CLI_PRESERVE_THOUGHT_SIGNATURES": { - "type": "bool", - "default": True, - "description": "Preserve thought signatures in client responses", - }, - "GEMINI_CLI_ENABLE_SIGNATURE_CACHE": { - "type": "bool", - "default": True, - "description": "Enable signature caching for multi-turn conversations", - }, - "GEMINI_CLI_GEMINI3_TOOL_FIX": { - "type": "bool", - "default": True, - "description": "Enable Gemini 3 tool hallucination prevention", - }, - "GEMINI_CLI_GEMINI3_TOOL_PREFIX": { - "type": "str", - "default": "gemini3_", - "description": "Prefix added to tool names for Gemini 3 disambiguation", - }, - "GEMINI_CLI_GEMINI3_DESCRIPTION_PROMPT": { - "type": "str", - "default": "\n\nSTRICT PARAMETERS: {params}.", - "description": "Template for strict parameter hints in tool descriptions", - }, - "GEMINI_CLI_PROJECT_ID": { - "type": "str", - "default": "", - "description": "GCP Project ID for paid tier users (required for paid tiers)", - }, - "GEMINI_CLI_OAUTH_PORT": { - "type": "int", - "default": GEMINI_CLI_DEFAULT_OAUTH_PORT, - "description": "Local port for OAuth callback server during authentication", - }, -} - # Map provider names to their settings definitions -PROVIDER_SETTINGS_MAP = { - "gemini_cli": GEMINI_CLI_SETTINGS, -} +PROVIDER_SETTINGS_MAP = {} class ProviderSettingsManager: @@ -1407,7 +1347,7 @@ def launch_model_filter_gui(self): input("Press Enter to continue...") def manage_provider_settings(self): - """Manage provider-specific settings (Gemini CLI)""" + """Manage provider-specific settings.""" while True: clear_screen() diff --git a/src/proxy_app/startup_display.py b/src/proxy_app/startup_display.py new file mode 100644 index 000000000..683b36ffe --- /dev/null +++ b/src/proxy_app/startup_display.py @@ -0,0 +1,16 @@ +"""Startup display helpers that avoid leaking secrets to the console.""" + +from __future__ import annotations + + +def mask_secret_for_display(value: str) -> str: + """Return a startup-safe representation of a configured secret. + + The proxy only needs to show that a key is configured. Printing the full + value makes terminal scrollback a credential leak, so this helper exposes at + most short edge fragments and the length for operator sanity checks. + """ + + if len(value) <= 8: + return f"Set (redacted, {len(value)} chars)" + return f"Set ({value[:4]}...{value[-4:]}, {len(value)} chars)" diff --git a/src/rotator_library/adapters/__init__.py b/src/rotator_library/adapters/__init__.py new file mode 100644 index 000000000..0c28ed20c --- /dev/null +++ b/src/rotator_library/adapters/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Composable payload adapters used between protocols and providers.""" + +from .base import AdapterContext, PayloadAdapter, run_adapter_chain +from .registry import ( + ADAPTER_ALIASES, + ADAPTER_PLUGINS, + get_adapter, + get_adapter_class, + list_adapters, + register_adapter, + resolve_adapter_name, +) + +__all__ = [ + "ADAPTER_ALIASES", + "ADAPTER_PLUGINS", + "AdapterContext", + "PayloadAdapter", + "get_adapter", + "get_adapter_class", + "list_adapters", + "register_adapter", + "resolve_adapter_name", + "run_adapter_chain", +] diff --git a/src/rotator_library/adapters/base.py b/src/rotator_library/adapters/base.py new file mode 100644 index 000000000..89f4e48a8 --- /dev/null +++ b/src/rotator_library/adapters/base.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Base classes for configurable protocol/provider payload adapters.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any, Iterable, Mapping, Optional + + +@dataclass +class AdapterContext: + """Context available to every adapter pass. + + Adapters are intentionally base implementations, not provider law. Provider + code can subclass adapters, copy and mutate them, or bypass them when a + protocol needs special behavior. The context mirrors transform trace fields + so each adapter pass can be debugged by provider/model/session/scope. + """ + + provider: Optional[str] = None + model: Optional[str] = None + protocol: Optional[str] = None + credential_id: Optional[str] = None + session_id: Optional[str] = None + scope_key: Optional[str] = None + classifier: Optional[str] = None + transport: Optional[str] = None + metadata: dict[str, Any] = field(default_factory=dict) + adapter_config: dict[str, dict[str, Any]] = field(default_factory=dict) + transaction_logger: Optional[Any] = None + + def config_for(self, adapter_name: str) -> dict[str, Any]: + """Return adapter-specific config without mutating the original context.""" + + return dict(self.adapter_config.get(adapter_name, {})) + + +class PayloadAdapter: + """Override-friendly base adapter. + + The default adapter is a no-op. Concrete adapters override only the stages + they support. Methods are async so future provider adapters can consult + caches, metadata services, or remote capability probes without changing the + chain runner API. + """ + + name: str = "" + aliases: tuple[str, ...] = () + supported_stages: tuple[str, ...] = ("request", "response", "stream_event") + + async def transform_request(self, payload: Any, context: AdapterContext) -> Any: + return payload + + async def transform_response(self, payload: Any, context: AdapterContext) -> Any: + return payload + + async def transform_stream_event(self, payload: Any, context: AdapterContext) -> Any: + return payload + + async def transform(self, stage: str, payload: Any, context: AdapterContext) -> Any: + """Dispatch a stage-specific transform with a useful error for typos.""" + + if stage not in self.supported_stages: + return payload + if stage == "request": + return await self.transform_request(payload, context) + if stage == "response": + return await self.transform_response(payload, context) + if stage == "stream_event": + return await self.transform_stream_event(payload, context) + raise ValueError(f"Unknown adapter stage: {stage}") + + +def _trace(context: AdapterContext, pass_name: str, payload: Any, *, stage: str, metadata: Mapping[str, Any]) -> None: + logger = context.transaction_logger + if not logger: + return + logger.log_transform_pass( + pass_name, + payload, + direction="stream" if stage == "stream_event" else stage, + stage="adapter", + protocol=context.protocol, + credential_id=context.credential_id, + transport=context.transport, + metadata=dict(metadata), + snapshot=stage != "stream_event", + ) + + +async def run_adapter_chain( + adapters: Iterable[PayloadAdapter], + payload: Any, + context: AdapterContext, + *, + stage: str, + mutate: bool = False, +) -> Any: + """Run adapters in order and emit trace entries around the chain. + + By default the payload is deep-copied before the first adapter. This keeps + Phase 3 isolated and prevents surprise mutations until runtime integration + explicitly chooses mutating behavior for performance. + """ + + current = payload if mutate else deepcopy(payload) + tracing_enabled = context.transaction_logger is not None + original = deepcopy(current) if tracing_enabled else None + adapter_names = [adapter.name for adapter in adapters] + _trace(context, "before_adapter_chain", current, stage=stage, metadata={"adapters": adapter_names}) + for adapter in adapters: + before = deepcopy(current) if tracing_enabled else None + try: + current = await adapter.transform(stage, current, context) + except Exception as exc: + if context.transaction_logger: + context.transaction_logger.log_transform_error( + f"adapter:{adapter.name}:{stage}", + exc, + payload=before if before is not None else current, + stage="adapter", + protocol=context.protocol, + transport=context.transport, + metadata={"adapter": adapter.name, "adapter_stage": stage}, + ) + raise + _trace( + context, + "after_adapter", + current, + stage=stage, + metadata={"adapter": adapter.name, "adapter_stage": stage, "changed": (current != before) if before is not None else None}, + ) + _trace( + context, + "after_adapter_chain", + current, + stage=stage, + metadata={"adapters": adapter_names, "adapter_stage": stage, "adapter_count": len(adapter_names), "changed": (original != current) if original is not None else None}, + ) + return current diff --git a/src/rotator_library/adapters/builtin.py b/src/rotator_library/adapters/builtin.py new file mode 100644 index 000000000..1475d70a7 --- /dev/null +++ b/src/rotator_library/adapters/builtin.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Built-in base adapters for common provider payload quirks.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any +from uuid import uuid4 + +from ..field_cache.paths import extract_path, inject_path +from .base import AdapterContext, PayloadAdapter + + +class NoOpAdapter(PayloadAdapter): + """Adapter that intentionally leaves payloads unchanged.""" + + name = "noop" + aliases = ("none", "passthrough") + + +class ModelOverrideAdapter(PayloadAdapter): + """Replace the outbound model field from adapter config. + + Config shape: + `{ "model": "provider/native-model-name" }` + """ + + name = "model_override" + aliases = ("override_model",) + supported_stages = ("request",) + + async def transform_request(self, payload: Any, context: AdapterContext) -> Any: + config = context.config_for(self.name) + override = config.get("model") or context.metadata.get("model_override") + if not override or not isinstance(payload, dict): + return payload + updated = deepcopy(payload) + updated["model"] = override + return updated + + +class SuppressDeveloperRoleAdapter(PayloadAdapter): + """Convert or remove developer-role messages for providers that reject them. + + Config shape: + `{ "mode": "system" | "user" | "drop" }` + """ + + name = "suppress_developer_role" + aliases = ("developer_role",) + supported_stages = ("request",) + + async def transform_request(self, payload: Any, context: AdapterContext) -> Any: + if not isinstance(payload, dict) or not isinstance(payload.get("messages"), list): + return payload + mode = context.config_for(self.name).get("mode", "system") + if mode not in {"system", "user", "drop"}: + raise ValueError("suppress_developer_role mode must be system, user, or drop") + updated = deepcopy(payload) + messages = [] + for message in updated.get("messages", []): + if not isinstance(message, dict) or message.get("role") != "developer": + messages.append(message) + continue + if mode == "drop": + continue + converted = dict(message) + converted["role"] = mode + messages.append(converted) + updated["messages"] = messages + return updated + + +class ReasoningContentAdapter(PayloadAdapter): + """Normalize common reasoning fields on assistant messages. + + This base adapter copies `reasoning`, `reasoning_content`, or configured + aliases into the configured output field. It deliberately does not delete + source fields; provider-specific subclasses can choose stricter behavior. + """ + + name = "reasoning_content" + aliases = ("reasoning_rewrite",) + supported_stages = ("response",) + + async def transform_response(self, payload: Any, context: AdapterContext) -> Any: + if not isinstance(payload, dict): + return payload + config = context.config_for(self.name) + output_field = config.get("output_field", "reasoning_content") + source_fields = tuple(config.get("source_fields", ("reasoning_content", "reasoning"))) + updated = deepcopy(payload) + for choice in updated.get("choices", []) if isinstance(updated.get("choices"), list) else []: + message = choice.get("message") if isinstance(choice, dict) else None + if not isinstance(message, dict): + continue + if output_field in message: + continue + for source_field in source_fields: + if source_field in message: + message[output_field] = message[source_field] + break + return updated + + +class FieldRenameAdapter(PayloadAdapter): + """Copy configured values between paths on raw payloads. + + Config shape: + `{ "rules": [{ "source_path": "a.b", "target_path": "c.d", "stage": "request", "move": false }] }` + + This adapter is conservative by design: it copies the last matched value by + default, delegates target ambiguity checks to `inject_path`, and only removes + the source for simple dotted-key paths when `move=true`. + """ + + name = "field_rename" + aliases = ("field_copy",) + + async def transform_request(self, payload: Any, context: AdapterContext) -> Any: + return self._transform_stage(payload, context, "request") + + async def transform_response(self, payload: Any, context: AdapterContext) -> Any: + return self._transform_stage(payload, context, "response") + + async def transform_stream_event(self, payload: Any, context: AdapterContext) -> Any: + return self._transform_stage(payload, context, "stream_event") + + def _transform_stage(self, payload: Any, context: AdapterContext, stage: str) -> Any: + if not isinstance(payload, dict): + return payload + updated = deepcopy(payload) + for rule in context.config_for(self.name).get("rules", []): + if rule.get("stage", stage) != stage: + continue + values = extract_path(updated, rule["source_path"]) + if not values: + continue + value = values if rule.get("as_list") else values[-1] + inject_path( + updated, + rule["target_path"], + value, + when_missing_only=bool(rule.get("when_missing_only", False)), + ) + if rule.get("move"): + _delete_simple_path(updated, rule["source_path"]) + return updated + + +class AntigravityEnvelopeAdapter(PayloadAdapter): + """Wrap Gemini payloads in the Antigravity internal request envelope. + + The active provider restores only stable envelope fields. Device profiles, + fingerprints, and other volatile client-emulation fields are intentionally + not generated here until they are verified against current service behavior. + """ + + name = "antigravity_envelope" + supported_stages = ("request",) + + async def transform_request(self, payload: Any, context: AdapterContext) -> Any: + if not isinstance(payload, dict): + return payload + if _looks_like_antigravity_envelope(payload): + return payload + config = context.config_for(self.name) + model = payload.get("model") or context.model + request_payload = {key: deepcopy(value) for key, value in payload.items() if key != "model"} + envelope = { + "model": model, + "request": request_payload, + "requestType": config.get("request_type", "CHAT_COMPLETION"), + "requestId": str(uuid4()), + "userAgent": config.get("user_agent"), + } + project = config.get("project") + if project: + envelope["project"] = project + return {key: value for key, value in envelope.items() if value is not None} + + +def _looks_like_antigravity_envelope(payload: dict[str, Any]) -> bool: + """Return whether a payload already has the controlled envelope shape.""" + + return "request" in payload and "requestType" in payload and "requestId" in payload + + +def _delete_simple_path(payload: dict[str, Any], path: str) -> None: + """Delete a simple dotted dict path after a conservative move operation.""" + + parts = path.split(".") + if any("[" in part or part == "*" for part in parts): + return + current: Any = payload + for part in parts[:-1]: + if not isinstance(current, dict): + return + current = current.get(part) + if isinstance(current, dict): + current.pop(parts[-1], None) diff --git a/src/rotator_library/adapters/registry.py b/src/rotator_library/adapters/registry.py new file mode 100644 index 000000000..7d702c0be --- /dev/null +++ b/src/rotator_library/adapters/registry.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Auto-discovery registry for payload adapters.""" + +from __future__ import annotations + +import importlib +import inspect +import logging +import pkgutil +from typing import Type + +from .base import PayloadAdapter + +lib_logger = logging.getLogger("rotator_library") + +ADAPTER_PLUGINS: dict[str, Type[PayloadAdapter]] = {} +ADAPTER_ALIASES: dict[str, str] = {} +_ADAPTER_INSTANCES: dict[str, PayloadAdapter] = {} + +_INFRASTRUCTURE_MODULES = {"base", "registry"} + + +def register_adapter(adapter_class: Type[PayloadAdapter], *, replace: bool = False) -> Type[PayloadAdapter]: + """Register an adapter class and its aliases with collision checks.""" + + if not inspect.isclass(adapter_class) or not issubclass(adapter_class, PayloadAdapter): + raise TypeError("adapter_class must inherit PayloadAdapter") + if adapter_class is PayloadAdapter: + raise TypeError("cannot register PayloadAdapter itself") + name = adapter_class.name + if not name: + raise ValueError(f"Adapter {adapter_class.__name__} must define a name") + alias_owner = ADAPTER_ALIASES.get(name) + if alias_owner and alias_owner != name and not replace: + raise ValueError(f"Adapter name conflicts with registered alias: {name}") + existing = ADAPTER_PLUGINS.get(name) + if existing and existing is not adapter_class and not replace: + raise ValueError(f"Adapter name already registered: {name}") + if replace and existing and existing is not adapter_class: + for alias, owner in list(ADAPTER_ALIASES.items()): + if owner == name: + ADAPTER_ALIASES.pop(alias, None) + ADAPTER_PLUGINS[name] = adapter_class + _ADAPTER_INSTANCES.pop(name, None) + for alias in adapter_class.aliases: + existing_name = ADAPTER_ALIASES.get(alias) + if existing_name and existing_name != name and not replace: + raise ValueError(f"Adapter alias already registered: {alias}") + if alias in ADAPTER_PLUGINS and alias != name and not replace: + raise ValueError(f"Adapter alias conflicts with registered name: {alias}") + ADAPTER_ALIASES[alias] = name + lib_logger.debug("Registered adapter: %s", name) + return adapter_class + + +def resolve_adapter_name(name: str) -> str: + if name in ADAPTER_PLUGINS: + return name + if name in ADAPTER_ALIASES: + return ADAPTER_ALIASES[name] + raise KeyError(f"Unknown adapter: {name}") + + +def get_adapter_class(name: str) -> Type[PayloadAdapter]: + return ADAPTER_PLUGINS[resolve_adapter_name(name)] + + +def get_adapter(name: str) -> PayloadAdapter: + canonical = resolve_adapter_name(name) + if canonical not in _ADAPTER_INSTANCES: + _ADAPTER_INSTANCES[canonical] = ADAPTER_PLUGINS[canonical]() + return _ADAPTER_INSTANCES[canonical] + + +def list_adapters() -> list[str]: + return sorted(ADAPTER_PLUGINS) + + +def _register_adapters() -> None: + package = importlib.import_module(__package__ or "rotator_library.adapters") + for _, module_name, _ in pkgutil.iter_modules(package.__path__): + if module_name.startswith("_") or module_name in _INFRASTRUCTURE_MODULES: + continue + module = importlib.import_module(f"{package.__name__}.{module_name}") + for attribute_name in dir(module): + attribute = getattr(module, attribute_name) + if ( + inspect.isclass(attribute) + and issubclass(attribute, PayloadAdapter) + and attribute is not PayloadAdapter + and attribute.__module__ == module.__name__ + ): + register_adapter(attribute) + + +_register_adapters() diff --git a/src/rotator_library/anthropic_compat/streaming.py b/src/rotator_library/anthropic_compat/streaming.py index ecb074baa..126ed70ab 100644 --- a/src/rotator_library/anthropic_compat/streaming.py +++ b/src/rotator_library/anthropic_compat/streaming.py @@ -66,11 +66,76 @@ async def anthropic_streaming_wrapper( accumulated_text = "" # Track accumulated text for logging accumulated_thinking = "" # Track accumulated thinking for logging stop_reason_final = "end_turn" # Track final stop reason for logging + upstream_closed = False + stream_iterator = openai_stream.__aiter__() + + def trace_frame(pass_name: str, frame: str, payload: Any | None = None) -> str: + """Trace one Anthropic stream conversion frame and return it for yield.""" + + if transaction_logger: + transaction_logger.log_transform_pass( + pass_name, + payload if payload is not None else frame, + direction="stream", + stage="adapter", + protocol="anthropic_messages", + transport="sse", + snapshot=False, + ) + return frame + + async def close_upstream(reason: str) -> None: + """Close the wrapped OpenAI stream when Anthropic streaming stops early.""" + + nonlocal upstream_closed + if upstream_closed: + return + upstream_closed = True + for candidate in (stream_iterator, openai_stream): + close = getattr(candidate, "aclose", None) + if callable(close): + await close() + break + sync_close = getattr(candidate, "close", None) + if callable(sync_close): + sync_close() + break + if transaction_logger: + transaction_logger.log_transform_pass( + "anthropic_stream_upstream_closed", + {"reason": reason}, + direction="stream", + stage="adapter", + protocol="anthropic_messages", + transport="sse", + snapshot=False, + ) try: - async for chunk_str in openai_stream: + async for chunk_str in stream_iterator: + if transaction_logger: + transaction_logger.log_transform_pass( + "anthropic_stream_source_chunk", + chunk_str, + direction="stream", + stage="provider", + protocol="openai_chat", + transport="sse", + snapshot=False, + ) # Check for client disconnection if callback provided if is_disconnected is not None and await is_disconnected(): + if transaction_logger: + transaction_logger.log_transform_pass( + "anthropic_stream_disconnected", + {"reason": "client_disconnected"}, + direction="stream", + stage="adapter", + protocol="anthropic_messages", + transport="sse", + snapshot=False, + ) + await close_upstream("client_disconnected") break if not chunk_str.strip() or not chunk_str.startswith("data:"): @@ -103,25 +168,25 @@ async def anthropic_streaming_wrapper( "usage": usage_dict, }, } - yield f"event: message_start\ndata: {json.dumps(message_start)}\n\n" + yield trace_frame("anthropic_stream_message_start", f"event: message_start\ndata: {json.dumps(message_start)}\n\n", message_start) message_started = True # Close any open thinking block if thinking_block_started: - yield f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n' + yield trace_frame("anthropic_stream_content_block_stop", f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n') current_block_index += 1 thinking_block_started = False # Close any open text block if content_block_started: - yield f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n' + yield trace_frame("anthropic_stream_content_block_stop", f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n') current_block_index += 1 content_block_started = False # Close all open tool_use blocks for tc_index in sorted(tool_block_indices.keys()): block_idx = tool_block_indices[tc_index] - yield f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {block_idx}}}\n\n' + yield trace_frame("anthropic_stream_content_block_stop", f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {block_idx}}}\n\n') # Determine stop_reason based on whether we had tool calls stop_reason = "tool_use" if tool_calls_by_index else "end_turn" @@ -134,10 +199,10 @@ async def anthropic_streaming_wrapper( final_usage["cache_creation_input_tokens"] = 0 # Send message_delta with final info - yield f'event: message_delta\ndata: {{"type": "message_delta", "delta": {{"stop_reason": "{stop_reason}", "stop_sequence": null}}, "usage": {json.dumps(final_usage)}}}\n\n' + yield trace_frame("anthropic_stream_message_delta", f'event: message_delta\ndata: {{"type": "message_delta", "delta": {{"stop_reason": "{stop_reason}", "stop_sequence": null}}, "usage": {json.dumps(final_usage)}}}\n\n') # Send message_stop - yield 'event: message_stop\ndata: {"type": "message_stop"}\n\n' + yield trace_frame("anthropic_stream_message_stop", 'event: message_stop\ndata: {"type": "message_stop"}\n\n') # Log final Anthropic response if logger provided if transaction_logger: @@ -242,7 +307,7 @@ async def anthropic_streaming_wrapper( "usage": usage_dict, }, } - yield f"event: message_start\ndata: {json.dumps(message_start)}\n\n" + yield trace_frame("anthropic_stream_message_start", f"event: message_start\ndata: {json.dumps(message_start)}\n\n", message_start) message_started = True choices = chunk.get("choices") or [] @@ -261,7 +326,7 @@ async def anthropic_streaming_wrapper( "index": current_block_index, "content_block": {"type": "thinking", "thinking": ""}, } - yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n" + yield trace_frame("anthropic_stream_content_block_start", f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n", block_start) thinking_block_started = True # Send thinking delta @@ -270,7 +335,7 @@ async def anthropic_streaming_wrapper( "index": current_block_index, "delta": {"type": "thinking_delta", "thinking": reasoning_content}, } - yield f"event: content_block_delta\ndata: {json.dumps(block_delta)}\n\n" + yield trace_frame("anthropic_stream_content_delta", f"event: content_block_delta\ndata: {json.dumps(block_delta)}\n\n", block_delta) # Accumulate thinking for logging accumulated_thinking += reasoning_content @@ -279,7 +344,7 @@ async def anthropic_streaming_wrapper( if content: # If we were in a thinking block, close it first if thinking_block_started and not content_block_started: - yield f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n' + yield trace_frame("anthropic_stream_content_block_stop", f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n') current_block_index += 1 thinking_block_started = False @@ -290,7 +355,7 @@ async def anthropic_streaming_wrapper( "index": current_block_index, "content_block": {"type": "text", "text": ""}, } - yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n" + yield trace_frame("anthropic_stream_content_block_start", f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n", block_start) content_block_started = True # Send content delta @@ -299,7 +364,7 @@ async def anthropic_streaming_wrapper( "index": current_block_index, "delta": {"type": "text_delta", "text": content}, } - yield f"event: content_block_delta\ndata: {json.dumps(block_delta)}\n\n" + yield trace_frame("anthropic_stream_content_delta", f"event: content_block_delta\ndata: {json.dumps(block_delta)}\n\n", block_delta) # Accumulate text for logging accumulated_text += content @@ -312,13 +377,13 @@ async def anthropic_streaming_wrapper( if tc_index not in tool_calls_by_index: # Close previous thinking block if open if thinking_block_started: - yield f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n' + yield trace_frame("anthropic_stream_content_block_stop", f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n') current_block_index += 1 thinking_block_started = False # Close previous text block if open if content_block_started: - yield f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n' + yield trace_frame("anthropic_stream_content_block_stop", f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n') current_block_index += 1 content_block_started = False @@ -341,7 +406,7 @@ async def anthropic_streaming_wrapper( "input": {}, }, } - yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n" + yield trace_frame("anthropic_stream_content_block_start", f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n", block_start) # Increment for the next block current_block_index += 1 @@ -361,7 +426,7 @@ async def anthropic_streaming_wrapper( "partial_json": func["arguments"], }, } - yield f"event: content_block_delta\ndata: {json.dumps(block_delta)}\n\n" + yield trace_frame("anthropic_stream_content_delta", f"event: content_block_delta\ndata: {json.dumps(block_delta)}\n\n", block_delta) # Note: We intentionally ignore finish_reason here. # Block closing is handled when we receive [DONE] to avoid @@ -369,6 +434,15 @@ async def anthropic_streaming_wrapper( except Exception as e: logger.error(f"Error in Anthropic streaming wrapper: {e}") + if transaction_logger: + transaction_logger.log_transform_error( + "anthropic_stream_transform", + e, + payload={"request_id": request_id, "model": original_model}, + stage="adapter", + protocol="anthropic_messages", + transport="sse", + ) # If we haven't sent message_start yet, send it now so the client can display the error # Claude Code and other clients may ignore events that come before message_start @@ -395,7 +469,7 @@ async def anthropic_streaming_wrapper( "usage": usage_dict, }, } - yield f"event: message_start\ndata: {json.dumps(message_start)}\n\n" + yield trace_frame("anthropic_stream_message_start", f"event: message_start\ndata: {json.dumps(message_start)}\n\n", message_start) # Send the error as a text content block so it's visible to the user error_message = f"Error: {str(e)}" @@ -404,16 +478,16 @@ async def anthropic_streaming_wrapper( "index": current_block_index, "content_block": {"type": "text", "text": ""}, } - yield f"event: content_block_start\ndata: {json.dumps(error_block_start)}\n\n" + yield trace_frame("anthropic_stream_content_block_start", f"event: content_block_start\ndata: {json.dumps(error_block_start)}\n\n", error_block_start) error_block_delta = { "type": "content_block_delta", "index": current_block_index, "delta": {"type": "text_delta", "text": error_message}, } - yield f"event: content_block_delta\ndata: {json.dumps(error_block_delta)}\n\n" + yield trace_frame("anthropic_stream_content_delta", f"event: content_block_delta\ndata: {json.dumps(error_block_delta)}\n\n", error_block_delta) - yield f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n' + yield trace_frame("anthropic_stream_content_block_stop", f'event: content_block_stop\ndata: {{"type": "content_block_stop", "index": {current_block_index}}}\n\n') # Build final usage with cached tokens final_usage = {"output_tokens": 0} @@ -422,12 +496,15 @@ async def anthropic_streaming_wrapper( final_usage["cache_creation_input_tokens"] = 0 # Send message_delta and message_stop to properly close the stream - yield f'event: message_delta\ndata: {{"type": "message_delta", "delta": {{"stop_reason": "end_turn", "stop_sequence": null}}, "usage": {json.dumps(final_usage)}}}\n\n' - yield 'event: message_stop\ndata: {"type": "message_stop"}\n\n' + yield trace_frame("anthropic_stream_message_delta", f'event: message_delta\ndata: {{"type": "message_delta", "delta": {{"stop_reason": "end_turn", "stop_sequence": null}}, "usage": {json.dumps(final_usage)}}}\n\n') + yield trace_frame("anthropic_stream_message_stop", 'event: message_stop\ndata: {"type": "message_stop"}\n\n') # Also send the formal error event for clients that handle it error_event = { "type": "error", "error": {"type": "api_error", "message": str(e)}, } - yield f"event: error\ndata: {json.dumps(error_event)}\n\n" + yield trace_frame("anthropic_stream_error", f"event: error\ndata: {json.dumps(error_event)}\n\n", error_event) + finally: + if not upstream_closed: + await close_upstream("wrapper_exit") diff --git a/src/rotator_library/client/anthropic.py b/src/rotator_library/client/anthropic.py index 507e82fb9..2bf0366ef 100644 --- a/src/rotator_library/client/anthropic.py +++ b/src/rotator_library/client/anthropic.py @@ -61,9 +61,9 @@ async def messages( """ Handle Anthropic Messages API requests. - This method accepts requests in Anthropic's format, translates them to - OpenAI format internally, processes them through the existing acompletion - method, and returns responses in Anthropic's format. + Non-streaming requests use the shared canonical protocol runtime. + Streaming temporarily retains the established Chat wrapper until the + canonical stream-lifecycle phase is enabled. Args: request: An AnthropicMessagesRequest object @@ -94,9 +94,68 @@ async def messages( request.model_dump(exclude_none=True), filename="anthropic_request.json", ) + _trace_anthropic( + anthropic_logger, + "anthropic_raw_request", + request.model_dump(exclude_none=True), + direction="request", + stage="client", + ) + + anthropic_request = request.model_dump(exclude_none=True) + + if hasattr(self._client, "agenerate"): + selected_output = ( + self._client.resolve_output_protocol( + anthropic_request, + input_protocol="anthropic_messages", + request=raw_request, + ) + if hasattr(self._client, "resolve_output_protocol") + else "anthropic_messages" + ) + response = await self._client.agenerate( + anthropic_request, + input_protocol="anthropic_messages", + output_protocol=selected_output, + request=raw_request, + pre_request_callback=pre_request_callback, + _parent_log_dir=anthropic_logger.log_dir if anthropic_logger and anthropic_logger.log_dir else None, + ) + if request.stream: + return response + anthropic_response = response.model_dump() if hasattr(response, "model_dump") else dict(response) + if selected_output != "anthropic_messages": + return anthropic_response + anthropic_response["id"] = request_id + _trace_anthropic( + anthropic_logger, + "anthropic_native_protocol_response", + anthropic_response, + direction="response", + stage="final", + ) + _trace_anthropic( + anthropic_logger, + "anthropic_final_response", + anthropic_response, + direction="response", + stage="final", + ) + if anthropic_logger: + anthropic_logger.log_response(anthropic_response, filename="anthropic_response.json") + return anthropic_response - # Translate Anthropic request to OpenAI format + # Compatibility fallback for external facades that have not implemented + # the protocol-aware agenerate entry point. openai_request = translate_anthropic_request(request) + _trace_anthropic( + anthropic_logger, + "anthropic_to_openai_request", + openai_request, + direction="request", + stage="adapter", + ) # Pass parent log directory to acompletion for nested logging if anthropic_logger and anthropic_logger.log_dir: @@ -138,9 +197,23 @@ async def messages( if hasattr(response, "model_dump") else dict(response) ) + _trace_anthropic( + anthropic_logger, + "anthropic_openai_response", + openai_response, + direction="response", + stage="provider", + ) anthropic_response = openai_to_anthropic_response( openai_response, original_model ) + _trace_anthropic( + anthropic_logger, + "openai_to_anthropic_response", + anthropic_response, + direction="response", + stage="adapter", + ) # Override the ID with our request ID anthropic_response["id"] = request_id @@ -151,6 +224,13 @@ async def messages( anthropic_response, filename="anthropic_response.json", ) + _trace_anthropic( + anthropic_logger, + "anthropic_final_response", + anthropic_response, + direction="response", + stage="final", + ) return anthropic_response @@ -201,3 +281,26 @@ async def count_tokens( total_tokens = message_tokens + tool_tokens return {"input_tokens": total_tokens} + + +def _trace_anthropic( + logger: Optional[TransactionLogger], + pass_name: str, + payload: Any, + *, + direction: str, + stage: str, + transport: Optional[str] = None, +) -> None: + """Emit an Anthropic compatibility transform trace when logging is enabled.""" + + if not logger: + return + logger.log_transform_pass( + pass_name, + payload, + direction=direction, + stage=stage, + protocol="anthropic_messages", + transport=transport, + ) diff --git a/src/rotator_library/client/executor.py b/src/rotator_library/client/executor.py index d83c61192..a6eb644ef 100644 --- a/src/rotator_library/client/executor.py +++ b/src/rotator_library/client/executor.py @@ -18,6 +18,7 @@ import os import random import time +from copy import deepcopy from typing import ( Any, AsyncGenerator, @@ -47,6 +48,8 @@ StreamedAPIError, ClassifiedError, RequestErrorAccumulator, + StructuredAPIResponseError, + structured_api_response_error, classify_error, should_rotate_on_error, should_retry_same_key, @@ -62,12 +65,24 @@ from ..request_sanitizer import sanitize_request_payload from ..transaction_logger import TransactionLogger from ..failure_logger import log_failure +from ..retry_policy import FailureHistory, decide_provider_cooldown, provider_cooldown_env +from ..routing import FallbackPolicy, clone_context_for_target +from ..routing.policy import normalize_route_error_type +from ..routing.types import RouteTarget +from ..native_provider import NativeHTTPTransport, NativeProviderContext, NativeProviderExecutor +from ..protocols import ProtocolContext, UnifiedStreamEvent, get_protocol +from ..protocols.streaming import convert_protocol_stream, format_canonical_stream_event +from ..native_provider.streaming import provider_supports_native_streaming as native_provider_supports_streaming +from ..field_cache.paths import FieldCachePathError, PathToken, parse_path +from ..transform_trace import REDACTED +from ..usage.accounting import UsageRecord, extract_usage_record +from ..usage.costs import CostBreakdown, CostCalculator from .types import RetryState, AvailabilityStats from .filters import CredentialFilter from .transforms import ProviderTransforms from .streaming import StreamingHandler -from .stream_retry_policy import can_retry_stream_after_error +from ..streaming.policy import can_retry_stream_after_error, is_stream_heartbeat_or_comment, is_visible_stream_output if TYPE_CHECKING: from ..usage import UsageManager @@ -75,6 +90,14 @@ lib_logger = logging.getLogger("rotator_library") +class RoutingExecutionError(RuntimeError): + """Internal error used when a routed target cannot use its requested mode.""" + + def __init__(self, message: str, error_type: str = "configuration_error") -> None: + super().__init__(message) + self.error_type = error_type + + class RequestExecutor: """ Unified retry/rotation logic for all request types. @@ -100,6 +123,7 @@ def __init__( litellm_provider_params: Optional[Dict[str, Any]] = None, litellm_logger_fn: Optional[Any] = None, provider_instances: Optional[Dict[str, Any]] = None, + experimental_config: Optional[Any] = None, ): """ Initialize RequestExecutor. @@ -134,8 +158,11 @@ def __init__( self._abort_on_callback_error = abort_on_callback_error self._litellm_provider_params = litellm_provider_params or {} self._litellm_logger_fn = litellm_logger_fn + self._experimental_config = experimental_config # StreamingHandler no longer needs usage_manager - we pass cred_context directly self._streaming_handler = StreamingHandler() + self._native_executor = NativeProviderExecutor() + self._failure_history = FailureHistory() def _get_transient_retry_delay(self) -> float: """Small jittered delay used before transient retries and rotations.""" @@ -192,14 +219,17 @@ def _get_plugin_instance(self, provider: str) -> Optional[Any]: self._plugin_instances[provider] = plugin_class else: return None - return self._plugin_instances[provider] + instance = self._plugin_instances[provider] + if self._experimental_config is not None and hasattr(instance, "bind_runtime_config"): + instance.bind_runtime_config(self._experimental_config) + return instance def _has_tier_support(self, provider: str) -> bool: """ Check if provider has tier/priority configuration. Providers with tier support define tier_priorities mapping - (e.g., GeminiCli, NanoGpt). + (e.g., NanoGpt). Args: provider: Provider name @@ -358,6 +388,9 @@ async def _prepare_request_kwargs( model: str, cred: str, context: "RequestContext", + *, + credential_id: Optional[str] = None, + native_execution: bool = False, ) -> Dict[str, Any]: """ Prepare request kwargs with transforms, sanitization, and provider params. @@ -371,6 +404,19 @@ async def _prepare_request_kwargs( Returns: Prepared kwargs dict for the LiteLLM call """ + if native_execution: + kwargs = deepcopy(context.kwargs) + self._log_executor_trace( + context, + "native_callback_compatibility_view", + kwargs, + direction="request", + stage="client", + credential_id=credential_id, + metadata={"provider": provider, "model": model}, + ) + return kwargs + # Apply transforms kwargs = await self._transforms.apply( provider, @@ -378,17 +424,58 @@ async def _prepare_request_kwargs( cred, context.kwargs.copy(), provider_config_override=context.provider_config, + transaction_logger=context.transaction_logger, + credential_id=credential_id, + transport="sse" if context.streaming else "http", + trace_metadata={ + "session_id": context.session_id, + "scope_key": context.usage_manager_key, + "classifier": context.classifier, + }, ) - # Sanitize request payload + # Sanitize request payload. Some provider compatibility fields are + # intentionally removed here, so record it as its own transform pass. + before_sanitize = deepcopy(kwargs) if context.transaction_logger else None kwargs = sanitize_request_payload(kwargs, model) + self._log_executor_trace( + context, + "after_request_sanitization", + kwargs, + direction="request", + stage="client", + credential_id=credential_id, + changed_from_previous=(before_sanitize != kwargs) if before_sanitize is not None else None, + metadata={"provider": provider, "model": model}, + ) # Apply provider-specific LiteLLM params + before_params = deepcopy(kwargs) if context.transaction_logger else None self._apply_litellm_provider_params(provider, kwargs) + self._log_executor_trace( + context, + "after_litellm_provider_params", + kwargs, + direction="request", + stage="client", + credential_id=credential_id, + changed_from_previous=(before_params != kwargs) if before_params is not None else None, + metadata={"provider": provider, "model": model}, + ) # Add transaction context for provider logging if context.transaction_logger: kwargs["transaction_context"] = context.transaction_logger.get_context() + self._log_executor_trace( + context, + "after_transaction_context_attached", + kwargs, + direction="request", + stage="client", + credential_id=credential_id, + changed_from_previous=True, + metadata={"provider": provider, "model": model}, + ) return kwargs @@ -472,12 +559,321 @@ async def _run_pre_request_callback( """ if context.pre_request_callback: try: + before = deepcopy(kwargs) if context.transaction_logger else None await context.pre_request_callback(context.request, kwargs) + if before is not None and before != kwargs: + self._log_executor_trace( + context, + "after_pre_request_callback", + kwargs, + direction="request", + stage="client", + changed_from_previous=True, + snapshot=True, + ) except Exception as e: if self._abort_on_callback_error: raise PreRequestCallbackError(str(e)) from e lib_logger.warning(f"Pre-request callback failed: {e}") + async def _execute_provider_request( + self, + provider: str, + model: str, + plugin: Any, + credential_secret: str, + credential_id: str, + kwargs: Dict[str, Any], + context: RequestContext, + ) -> Any: + """Execute one provider request using routed execution-mode rules.""" + + target = _current_route_target(context) + execution = target.execution if target else "auto" + self._log_executor_trace( + context, + "pre_provider_execution_request", + kwargs, + direction="request", + stage="provider", + credential_id=credential_id, + metadata={"execution": execution, "provider": provider, "model": model}, + ) + if execution == "litellm_fallback": + self._log_routing_trace( + context, + "routing_litellm_fallback", + _target_trace(target) if target else {"provider": provider, "model": model}, + ) + response = await self._execute_litellm_request(kwargs, credential_secret, context=context, credential_id=credential_id) + _raise_for_structured_response_error(response) + return self._format_execution_response(response, "openai_chat", context) + + if execution == "custom" or (execution == "auto" and plugin and plugin.has_custom_logic()): + if not plugin or not plugin.has_custom_logic(): + raise RoutingExecutionError(f"Provider {provider} does not support custom execution") + kwargs["credential_identifier"] = credential_secret + self._log_executor_trace( + context, + "provider_execution_request", + kwargs, + direction="request", + stage="provider", + credential_id=credential_id, + metadata={"execution": "custom", "provider": provider, "model": model}, + ) + response = await plugin.acompletion(self._http_client, **kwargs) + _raise_for_structured_response_error(response) + return self._format_execution_response(response, "openai_chat", context) + + if execution == "native" or (execution == "auto" and _should_use_native_protocol(plugin, model, target, kwargs, stream=False, execution=execution)): + native_context, native_request = self._build_native_provider_context( + provider, + model, + plugin, + credential_secret, + credential_id, + context, + target, + raw_request=kwargs, + return_request=True, + ) + self._log_routing_trace( + context, + "routing_native_execution_selected", + _target_trace(target) if target else {"provider": provider, "model": model}, + metadata={"protocol": native_context.protocol_name, "operation": native_context.operation}, + ) + return await self._get_native_executor().execute(native_request, native_context, NativeHTTPTransport(self._http_client)) + + response = await self._execute_litellm_request(kwargs, credential_secret, context=context, credential_id=credential_id) + _raise_for_structured_response_error(response) + return self._format_execution_response(response, "openai_chat", context) + + def _get_native_executor(self) -> NativeProviderExecutor: + """Return the shared native executor for process-local field-cache state.""" + + native_executor = getattr(self, "_native_executor", None) + if native_executor is None: + native_executor = NativeProviderExecutor() + self._native_executor = native_executor + return native_executor + + async def _execute_litellm_request( + self, + kwargs: Dict[str, Any], + credential_secret: str, + *, + context: Optional[RequestContext] = None, + credential_id: Optional[str] = None, + ) -> Any: + """Execute the existing LiteLLM request path.""" + + kwargs["api_key"] = credential_secret + self._apply_litellm_logger(kwargs) + kwargs.pop("transaction_context", None) + if context: + self._log_executor_trace( + context, + "provider_execution_request", + kwargs, + direction="request", + stage="provider", + credential_id=credential_id, + metadata={"execution": "litellm", "provider": context.provider, "model": context.model}, + ) + return await litellm.acompletion(**kwargs) + + @staticmethod + def _format_execution_response( + response: Any, + source_protocol_name: str, + context: RequestContext, + ) -> Any: + """Format a non-native execution result into the selected client protocol.""" + + output_protocol = get_protocol(context.output_protocol_name) + source_protocol = get_protocol(source_protocol_name) + if output_protocol.name == source_protocol.name: + return response + if isinstance(response, dict): + payload = deepcopy(response) + elif hasattr(response, "model_dump"): + payload = response.model_dump() + elif hasattr(response, "dict"): + payload = response.dict() + else: + payload = dict(response) + protocol_context = ProtocolContext( + provider=context.provider, + model=context.model, + source_protocol=source_protocol.name, + target_protocol=output_protocol.name, + input_protocol=context.input_protocol_name, + provider_protocol=source_protocol.name, + output_protocol=output_protocol.name, + source_provider=context.provider, + session_id=context.session_id, + transport="http", + ) + unified_response = source_protocol.parse_response(payload, protocol_context) + unified_response.model = context.model + return output_protocol.format_response(unified_response, protocol_context) + + def _build_native_provider_context( + self, + provider: str, + model: str, + plugin: Any, + credential_secret: str, + credential_id: str, + context: RequestContext, + target: Optional[RouteTarget], + raw_request: Optional[Dict[str, Any]] = None, + transport: str = "http", + stream: bool = False, + return_request: bool = False, + ) -> NativeProviderContext | tuple[NativeProviderContext, Dict[str, Any]]: + """Build native provider context from provider declarations.""" + + if not plugin: + raise RoutingExecutionError(f"Provider {provider} has no plugin for native execution") + protocol_name = _provider_native_protocol(plugin, model, target) + if not protocol_name: + raise RoutingExecutionError(f"Provider {provider} has no native protocol declaration") + public_model = model + native_model = plugin.normalize_native_model(model) if hasattr(plugin, "normalize_native_model") else _strip_provider_prefix(model) + operation = plugin.get_native_operation(native_model, None, stream=stream) if hasattr(plugin, "get_native_operation") else "chat" + if hasattr(plugin, "supports_native_operation") and not plugin.supports_native_operation(native_model, operation): + raise RoutingExecutionError(f"Provider {provider} does not support native operation {operation}") + try: + endpoint = plugin.get_native_endpoint(model=native_model, operation=operation) + headers = plugin.get_native_headers(credential_secret, model=native_model, operation=operation) + except NotImplementedError as exc: + raise RoutingExecutionError(str(exc)) from exc + native_context = NativeProviderContext( + provider=provider, + model=native_model, + protocol_name=protocol_name, + endpoint=endpoint, + operation=operation, + input_protocol_name=context.input_protocol_name, + # The operational stream handler consumes canonical Chat SSE. The + # selected client protocol is applied once after timeout/retry/usage + # and completion-gated session handling. + output_protocol_name="openai_chat" if stream else context.output_protocol_name, + headers=headers, + credential_id=credential_id, + session_id=context.session_id, + scope_key=context.usage_manager_key, + classifier=context.classifier, + transport=transport, + adapter_names=tuple(plugin.get_adapter_names(native_model) if hasattr(plugin, "get_adapter_names") else ()), + adapter_config=dict(plugin.get_adapter_config(native_model) if hasattr(plugin, "get_adapter_config") else {}), + field_cache_rules=_merged_field_cache_rules( + provider, + public_model, + plugin, + config=getattr(self, "_experimental_config", None), + ), + transaction_logger=context.transaction_logger, + metadata={ + "public_model": public_model, + "input_provider": context.input_provider, + "disable_provider_continuation": context.disable_provider_continuation, + }, + request_preparer=plugin.prepare_native_request if hasattr(plugin, "prepare_native_request") else None, + request_validator=plugin.validate_request if hasattr(plugin, "validate_request") else None, + ) + if return_request: + if context.unified_request is None: + raise RoutingExecutionError("Native execution requires a canonical request") + return native_context, self._canonical_request_for_native(context, raw_request or {}) + return native_context + + @staticmethod + def _canonical_request_for_native( + context: RequestContext, + request_payload: Dict[str, Any], + ) -> Any: + """Overlay protocol-safe attempt mutations onto the canonical request. + + Existing transforms and callbacks operate on the Chat execution view. + Only canonical fields changed relative to that view are overlaid, so + source-native semantics that Chat cannot represent remain intact. + """ + + canonical_request = deepcopy(context.unified_request) + chat_protocol = get_protocol("openai_chat") + baseline = chat_protocol.parse_request(context.kwargs) + attempted = chat_protocol.parse_request(request_payload) + original_instructions = [message for message in canonical_request.messages if message.role in {"system", "developer"}] + original_conversation = [message for message in canonical_request.messages if message.role not in {"system", "developer"}] + baseline_instructions = [message for message in baseline.messages if message.role in {"system", "developer"}] + baseline_conversation = [message for message in baseline.messages if message.role not in {"system", "developer"}] + attempted_instructions = [message for message in attempted.messages if message.role in {"system", "developer"}] + attempted_conversation = [message for message in attempted.messages if message.role not in {"system", "developer"}] + if baseline_conversation != attempted_conversation: + original_conversation = _merge_canonical_sequence( + original_conversation, + baseline_conversation, + attempted_conversation, + _merge_canonical_message, + lambda message: (message.role, message.tool_call_id, message.name), + ) + if baseline_instructions != attempted_instructions: + if canonical_request.system and not original_instructions: + baseline_blocks = [block for message in baseline_instructions for block in message.content] + attempted_blocks = [block for message in attempted_instructions for block in message.content] + canonical_request.system = _merge_canonical_sequence( + canonical_request.system, + baseline_blocks, + attempted_blocks, + _merge_canonical_content_block, + lambda block: block.type, + ) + else: + original_instructions = _merge_canonical_sequence( + original_instructions, + baseline_instructions, + attempted_instructions, + _merge_canonical_message, + lambda message: (message.role, message.tool_call_id, message.name), + ) + canonical_request.messages = original_instructions + original_conversation + if baseline.tools != attempted.tools: + canonical_request.tools = _merge_canonical_sequence( + canonical_request.tools, + baseline.tools, + attempted.tools, + _merge_canonical_tool, + lambda tool: (tool.type, tool.name), + ) + for field_name in ("stream", "modalities"): + if getattr(baseline, field_name) != getattr(attempted, field_name): + setattr(canonical_request, field_name, deepcopy(getattr(attempted, field_name))) + metadata_keys = set(baseline.metadata) | set(attempted.metadata) + for key in metadata_keys: + if baseline.metadata.get(key) == attempted.metadata.get(key): + continue + if key in attempted.metadata: + canonical_request.metadata[key] = deepcopy(attempted.metadata[key]) + else: + canonical_request.metadata.pop(key, None) + if baseline.response_format != attempted.response_format: + canonical_request.response_format = deepcopy(attempted.response_format) + generation_keys = set(baseline.generation_params) | set(attempted.generation_params) + for key in generation_keys: + baseline_value = baseline.generation_params.get(key) + attempted_value = attempted.generation_params.get(key) + if baseline_value == attempted_value: + continue + if key in attempted.generation_params: + canonical_request.generation_params[key] = deepcopy(attempted_value) + else: + canonical_request.generation_params.pop(key, None) + return canonical_request + async def execute( self, context: RequestContext, @@ -493,11 +889,325 @@ async def execute( Returns: Response object or async generator for streaming """ + if context.streaming and context.routing_targets: + return self._execute_streaming_with_fallback(context) if context.streaming: return self._execute_streaming(context) + elif context.routing_targets: + return await self._execute_non_streaming_with_fallback(context) else: return await self._execute_non_streaming(context) + async def _execute_non_streaming_with_fallback(self, context: RequestContext) -> Any: + """Execute an ordered non-streaming fallback target chain. + + The normal single-target path remains `_execute_non_streaming()`. This + wrapper only runs when request building has populated `routing_targets`, + preserving existing behavior for all current requests. + """ + + targets = tuple(context.routing_targets or ()) + if not targets: + return await self._execute_non_streaming(context) + policy = FallbackPolicy() + last_failure: Any = None + target_failures: List[Dict[str, Any]] = [] + self._log_routing_trace( + context, + "routing_decision", + {"requested_model": context.model, "target_count": len(targets)}, + metadata={"group": context.routing_group_name, "targets": [_target_trace(target) for target in targets]}, + ) + for index, target in enumerate(targets): + attempt_started_at = time.monotonic() + target_context = clone_context_for_target( + context, + target, + target_index=index, + credentials=_target_scope_value(target, "credentials", context.credentials), + usage_manager_key=_target_scope_value(target, "usage_manager_key", target.provider), + provider_config=_target_scope_value(target, "provider_config", context.provider_config), + credential_secrets=_target_scope_value(target, "credential_secrets", context.credential_secrets), + ) + self._log_routing_trace( + context, + "routing_target_attempt_started", + _target_trace(target), + metadata={"target_index": index, "group": context.routing_group_name}, + ) + try: + result = await self._execute_non_streaming(target_context) + except Exception as exc: + last_failure = exc + error_type = _route_error_type(exc, target.provider) + self._log_routing_trace( + context, + "routing_target_attempt_failed", + _target_trace(target), + metadata={"target_index": index, "error_type": error_type, "exception": exc.__class__.__name__}, + ) + target_failures.append(_target_failure_summary(target, error_type)) + fallback_allowed = index < len(targets) - 1 and policy.should_fallback(error_type, group=context.routing_group) + _append_routing_attempt_history(context, target, index, success=False, error_type=error_type, fallback_allowed=fallback_allowed, duration_ms=_elapsed_ms(attempt_started_at)) + if not fallback_allowed: + self._log_routing_trace(context, "routing_fallback_exhausted", _target_trace(target), metadata={"error_type": error_type, "fallback_targets": target_failures}) + if isinstance(exc, StructuredAPIResponseError): + exc.response = _with_fallback_summary( + deepcopy(exc.response), + target_failures, + context.routing_attempt_history, + ) + raise + self._log_routing_trace(context, "routing_fallback_selected", _target_trace(targets[index + 1]), metadata={"from_target_index": index, "to_target_index": index + 1, "reason": error_type}) + continue + + error_type = _route_error_type_from_response(result) + if error_type: + last_failure = result + self._log_routing_trace( + context, + "routing_target_attempt_failed", + _target_trace(target), + metadata={"target_index": index, "error_type": error_type}, + ) + target_failures.append(_target_failure_summary(target, error_type, status_code=_route_status_code_from_response(result))) + fallback_allowed = index < len(targets) - 1 and policy.should_fallback(error_type, group=context.routing_group) + _append_routing_attempt_history(context, target, index, success=False, error_type=error_type, status_code=_route_status_code_from_response(result), fallback_allowed=fallback_allowed, duration_ms=_elapsed_ms(attempt_started_at)) + if fallback_allowed: + self._log_routing_trace(context, "routing_fallback_selected", _target_trace(targets[index + 1]), metadata={"from_target_index": index, "to_target_index": index + 1, "reason": error_type}) + continue + self._log_routing_trace(context, "routing_fallback_exhausted", _target_trace(target), metadata={"error_type": error_type, "fallback_targets": target_failures}) + return _with_fallback_summary(result, target_failures, context.routing_attempt_history) + + self._log_routing_trace( + context, + "routing_target_attempt_succeeded", + _target_trace(target), + metadata={"target_index": index}, + ) + _append_routing_attempt_history(context, target, index, success=True, duration_ms=_elapsed_ms(attempt_started_at)) + return result + + if isinstance(last_failure, Exception): + raise last_failure + return last_failure + + async def _execute_streaming_with_fallback(self, context: RequestContext) -> AsyncGenerator[str, None]: + """Execute streaming fallback targets with pre-output-only failover.""" + + targets = tuple(context.routing_targets or ()) + if not targets: + async for chunk in self._execute_streaming(context): + yield chunk + return + policy = FallbackPolicy() + target_failures: List[Dict[str, Any]] = [] + self._log_routing_trace( + context, + "routing_decision", + {"requested_model": context.model, "target_count": len(targets), "stream": True}, + metadata={"group": context.routing_group_name, "streaming_policy": _group_streaming_policy(context.routing_group), "targets": [_target_trace(target) for target in targets]}, + ) + for index, target in enumerate(targets): + attempt_started_at = time.monotonic() + emitted_output = False + pending_chunks: List[str] = [] + terminal_error_type: Optional[str] = None + target_context = clone_context_for_target( + context, + target, + target_index=index, + credentials=_target_scope_value(target, "credentials", context.credentials), + usage_manager_key=_target_scope_value(target, "usage_manager_key", target.provider), + provider_config=_target_scope_value(target, "provider_config", context.provider_config), + credential_secrets=_target_scope_value(target, "credential_secrets", context.credential_secrets), + ) + self._log_routing_trace( + context, + "routing_stream_target_attempt_started", + _target_trace(target), + metadata={"target_index": index, "group": context.routing_group_name}, + ) + try: + async for chunk in self._execute_streaming(target_context): + if is_stream_heartbeat_or_comment(chunk): + yield chunk + continue + chunk_error_type = _stream_chunk_error_type(chunk) + if chunk_error_type and not emitted_output: + terminal_error_type = chunk_error_type + pending_chunks.append(chunk) + continue + if _stream_chunk_is_visible_output(chunk): + for pending in pending_chunks: + yield pending + pending_chunks.clear() + emitted_output = True + yield chunk + continue + pending_chunks.append(chunk) + if terminal_error_type and not emitted_output: + error_type = terminal_error_type + self._log_routing_trace( + context, + "routing_stream_target_attempt_failed", + _target_trace(target), + metadata={"target_index": index, "error_type": error_type, "emitted_output": emitted_output, "terminal_error_frame": True}, + ) + target_failures.append(_target_failure_summary(target, error_type)) + fallback_allowed = index < len(targets) - 1 and _streaming_policy_allows_fallback(context.routing_group) and policy.should_fallback(error_type, group=context.routing_group, stream=True, emitted_output=False) + _append_routing_attempt_history(context, target, index, success=False, error_type=error_type, emitted_output=emitted_output, fallback_allowed=fallback_allowed, duration_ms=_elapsed_ms(attempt_started_at)) + if fallback_allowed: + self._log_routing_trace( + context, + "routing_fallback_selected", + _target_trace(targets[index + 1]), + metadata={"from_target_index": index, "to_target_index": index + 1, "reason": error_type, "stream": True}, + ) + continue + self._log_routing_trace(context, "routing_fallback_exhausted", _target_trace(target), metadata={"error_type": error_type, "stream": True, "streaming_policy": _group_streaming_policy(context.routing_group), "fallback_targets": target_failures}) + for pending in pending_chunks: + yield pending + return + for pending in pending_chunks: + yield pending + self._log_routing_trace( + context, + "routing_stream_target_attempt_succeeded", + _target_trace(target), + metadata={"target_index": index, "emitted_output": emitted_output}, + ) + _append_routing_attempt_history(context, target, index, success=True, emitted_output=emitted_output, duration_ms=_elapsed_ms(attempt_started_at)) + return + except Exception as exc: + error_type = _route_error_type(exc, target.provider) + self._log_routing_trace( + context, + "routing_stream_target_attempt_failed", + _target_trace(target), + metadata={"target_index": index, "error_type": error_type, "emitted_output": emitted_output}, + ) + target_failures.append(_target_failure_summary(target, error_type)) + fallback_allowed = (not emitted_output) and index < len(targets) - 1 and _streaming_policy_allows_fallback(context.routing_group) and policy.should_fallback(error_type, group=context.routing_group, stream=True, emitted_output=False) + _append_routing_attempt_history(context, target, index, success=False, error_type=error_type, emitted_output=emitted_output, fallback_allowed=fallback_allowed, duration_ms=_elapsed_ms(attempt_started_at)) + if emitted_output: + self._log_routing_trace( + context, + "routing_stream_fallback_blocked_after_output", + _target_trace(target), + metadata={"target_index": index, "error_type": error_type}, + ) + raise + if fallback_allowed: + self._log_routing_trace( + context, + "routing_fallback_selected", + _target_trace(targets[index + 1]), + metadata={"from_target_index": index, "to_target_index": index + 1, "reason": error_type, "stream": True}, + ) + continue + self._log_routing_trace(context, "routing_fallback_exhausted", _target_trace(target), metadata={"error_type": error_type, "stream": True, "streaming_policy": _group_streaming_policy(context.routing_group), "fallback_targets": target_failures}) + raise + + @staticmethod + def _log_routing_trace(context: RequestContext, pass_name: str, data: Any, *, metadata: Optional[Dict[str, Any]] = None) -> None: + """Record routing trace entries without affecting request execution.""" + + if not context.transaction_logger: + return + context.transaction_logger.log_transform_pass( + pass_name, + data, + direction="metadata", + stage="routing", + metadata=metadata or {}, + snapshot=False, + ) + + @staticmethod + def _log_executor_trace( + context: RequestContext, + pass_name: str, + data: Any, + *, + direction: str, + stage: str, + credential_id: Optional[str] = None, + changed_from_previous: Optional[bool] = None, + metadata: Optional[Dict[str, Any]] = None, + snapshot: bool = True, + ) -> None: + """Record live executor trace boundaries without affecting requests.""" + + if not context.transaction_logger: + return + context.transaction_logger.log_transform_pass( + pass_name, + data, + direction=direction, + stage=stage, + credential_id=credential_id, + transport="sse" if context.streaming else "http", + changed_from_previous=changed_from_previous, + metadata={ + "provider": context.provider, + "model": context.model, + "session_id": context.session_id, + "scope_key": context.usage_manager_key, + "classifier": context.classifier, + **(metadata or {}), + }, + snapshot=snapshot, + ) + + def _terminal_stream_error_lines( + self, + context: RequestContext, + error_data: Dict[str, Any], + *, + protocol_context: Optional[ProtocolContext] = None, + ) -> Tuple[str, ...]: + """Return executor-created terminal SSE lines and trace them first.""" + + output_protocol = getattr(context, "output_protocol_name", None) or "openai_chat" + input_protocol = getattr(context, "input_protocol_name", None) or "openai_chat" + protocol_context = protocol_context or ProtocolContext( + provider=context.provider, + model=context.model, + source_protocol="openai_chat", + target_protocol=output_protocol, + input_protocol=input_protocol, + output_protocol=output_protocol, + request_id=getattr(context, "request_id", None), + session_id=context.session_id, + transport="sse", + ) + error = error_data.get("error", error_data) + lines = tuple( + format_canonical_stream_event( + UnifiedStreamEvent(type="error", error=error), + output_protocol, + protocol_context, + ) + ) + self._log_executor_trace( + context, + "stream_error_event", + error_data, + direction="stream", + stage="client", + snapshot=False, + ) + self._log_executor_trace( + context, + "stream_done_event", + {"frames": lines}, + direction="stream", + stage="final", + snapshot=False, + ) + return lines + async def _prepare_execution( self, context: RequestContext, @@ -516,8 +1226,9 @@ async def _prepare_execution( quota_group = usage_manager.get_model_quota_group(model) await self._ensure_initialized(usage_manager, context, filter_result) - await self._validate_request(provider, model, context.kwargs) - + plugin = self._get_plugin_instance(provider) + if not _attempt_uses_native_protocol(plugin, model, context): + await self._validate_request(provider, model, context.kwargs) if not credentials: raise NoAvailableKeysError(f"No compatible credentials for model {model}") @@ -569,7 +1280,7 @@ async def _execute_non_streaming( break # Wait for provider cooldown - await self._wait_for_cooldown(provider, deadline) + await self._wait_for_cooldown(provider, deadline, model=model, context=context) # Acquire credential using context manager try: @@ -598,22 +1309,33 @@ async def _execute_non_streaming( self._log_acquired_credential( cred, model, state, quota_group, availability, usage_manager ) + plugin = self._get_plugin_instance(provider) + native_execution = _attempt_uses_native_protocol(plugin, model, context) try: # Prepare request kwargs kwargs = await self._prepare_request_kwargs( - provider, model, cred, context + provider, + model, + cred, + context, + credential_id=cred_context.stable_id, + native_execution=native_execution, ) # Log transformed request if it differs from original if context.transaction_logger: context.transaction_logger.log_transformed_request( - kwargs, context.kwargs + kwargs, + context.kwargs, + credential_id=cred_context.stable_id, + metadata={ + "session_id": context.session_id, + "scope_key": context.usage_manager_key, + "classifier": context.classifier, + }, ) - # Get provider plugin - plugin = self._get_plugin_instance(provider) - # Execute request with retries for attempt in range(self._max_retries): try: @@ -624,31 +1346,42 @@ async def _execute_non_streaming( # Pre-request callback await self._run_pre_request_callback(context, kwargs) - # Make the API call - if plugin and plugin.has_custom_logic(): - kwargs["credential_identifier"] = credential_secret - response = await plugin.acompletion( - self._http_client, **kwargs - ) - else: - # Standard LiteLLM call - kwargs["api_key"] = credential_secret - self._apply_litellm_logger(kwargs) - # Remove internal context before litellm call - kwargs.pop("transaction_context", None) - response = await litellm.acompletion(**kwargs) + response = await self._execute_provider_request( + provider, + model, + plugin, + credential_secret, + cred_context.stable_id, + kwargs, + context, + ) + trace_response = _redact_context_field_cache_paths( + response, + context, + "response", + plugin, + config=getattr(self, "_experimental_config", None), + ) + self._log_executor_trace( + context, + "raw_provider_response", + trace_response, + direction="response", + stage="provider", + credential_id=cred_context.stable_id, + metadata={"provider": provider, "model": model}, + ) # Success! Extract token usage if available - ( - prompt_tokens, - completion_tokens, - prompt_tokens_cached, - prompt_tokens_cache_write, - thinking_tokens, - ) = self._extract_usage_tokens(response) - approx_cost = self._calculate_cost( - provider, model, response + usage_record, cost_breakdown = self._account_for_response_usage( + provider, model, response, context ) + prompt_tokens = usage_record.prompt_tokens_for_mark_success + completion_tokens = usage_record.completion_tokens + prompt_tokens_cached = usage_record.cache_read_tokens + prompt_tokens_cache_write = usage_record.cache_write_tokens + thinking_tokens = usage_record.reasoning_tokens + approx_cost = cost_breakdown.total_cost response_headers = self._extract_response_headers( response ) @@ -663,30 +1396,63 @@ async def _execute_non_streaming( approx_cost=approx_cost, response_headers=response_headers, ) + self._clear_failure_history_on_success(provider, model) self._record_session_response(context, response) lib_logger.info( f"Recorded usage from response object for key {mask_credential(cred)}" ) - # Log response if transaction logging enabled + normalized_response = self._normalize_response_usage(response, model) + trace_normalized_response = _redact_context_field_cache_paths( + normalized_response, + context, + "response", + plugin, + config=getattr(self, "_experimental_config", None), + ) + self._log_executor_trace( + context, + "post_usage_normalization_response", + trace_normalized_response, + direction="response", + stage="final", + credential_id=cred_context.stable_id, + metadata={"provider": provider, "model": model}, + ) + # Legacy response.json and final_client_response must + # reflect the same post-normalization payload returned + # to the caller, not the pre-accounting SDK object. if context.transaction_logger: try: - response_data = ( - response.model_dump() - if hasattr(response, "model_dump") - else response - ) - context.transaction_logger.log_response( - response_data - ) + context.transaction_logger.log_response(trace_normalized_response) except Exception as log_err: lib_logger.debug( f"Failed to log response: {log_err}" ) + return normalized_response - return self._normalize_response_usage(response, model) - + except RoutingExecutionError as e: + if e.error_type == "configuration_error": + raise + last_exception = e + action = await self._handle_error_with_context( + e, + cred_context, + model, + provider, + attempt, + error_accumulator, + retry_state, + request_headers, + context, + ) + if action == ErrorAction.RETRY_SAME: + continue + elif action == ErrorAction.ROTATE: + break + else: + raise except Exception as e: last_exception = e action = await self._handle_error_with_context( @@ -698,6 +1464,7 @@ async def _execute_non_streaming( error_accumulator, retry_state, request_headers, + context, ) if action == ErrorAction.RETRY_SAME: @@ -709,6 +1476,9 @@ async def _execute_non_streaming( except PreRequestCallbackError: raise + except RoutingExecutionError as exc: + if exc.error_type == "configuration_error": + raise except Exception: # Let context manager handle cleanup pass @@ -718,11 +1488,19 @@ async def _execute_non_streaming( # All credentials exhausted error_accumulator.timeout_occurred = time.time() >= deadline + if isinstance(last_exception, StructuredAPIResponseError): + raise last_exception if last_exception and not error_accumulator.has_errors(): raise last_exception - # Return error response - return error_accumulator.build_client_error_response() + error_response = error_accumulator.build_client_error_response() + error_type = str(error_response.get("error", {}).get("type") or "proxy_all_credentials_exhausted") + raise StructuredAPIResponseError( + str(error_response.get("error", {}).get("message") or "All credentials exhausted"), + error_type=error_type, + status_code=504 if error_accumulator.timeout_occurred else 503, + response=error_response, + ) async def _execute_streaming( self, @@ -742,6 +1520,7 @@ async def _execute_streaming( provider = context.provider model = context.model deadline = context.deadline + client_protocol_context: Optional[ProtocolContext] = None try: ( @@ -758,8 +1537,8 @@ async def _execute_streaming( "type": "proxy_error", } } - yield f"data: {json.dumps(error_data)}\n\n" - yield "data: [DONE]\n\n" + for line in self._terminal_stream_error_lines(context, error_data): + yield line return error_accumulator = RequestErrorAccumulator() @@ -768,6 +1547,7 @@ async def _execute_streaming( retry_state = RetryState() last_exception: Optional[Exception] = None + last_stream_error_payload: Optional[Dict[str, Any]] = None try: while time.time() < deadline: @@ -785,7 +1565,7 @@ async def _execute_streaming( remaining = deadline - time.time() if remaining <= 0: break - await self._wait_for_cooldown(provider, deadline) + await self._wait_for_cooldown(provider, deadline, model=model, context=context) # Acquire credential using context manager try: @@ -814,27 +1594,39 @@ async def _execute_streaming( self._log_acquired_credential( cred, model, state, quota_group, availability, usage_manager ) + plugin = self._get_plugin_instance(provider) + native_execution = _attempt_uses_native_protocol(plugin, model, context) try: # Prepare request kwargs kwargs = await self._prepare_request_kwargs( - provider, model, cred, context + provider, + model, + cred, + context, + credential_id=cred_context.stable_id, + native_execution=native_execution, ) # Log transformed request if it differs from original if context.transaction_logger: context.transaction_logger.log_transformed_request( - kwargs, context.kwargs + kwargs, + context.kwargs, + credential_id=cred_context.stable_id, + metadata={ + "session_id": context.session_id, + "scope_key": context.usage_manager_key, + "classifier": context.classifier, + }, ) # Add stream usage metadata for active providers. - if "stream_options" not in kwargs: + if not native_execution and "stream_options" not in kwargs: kwargs["stream_options"] = {} - if "include_usage" not in kwargs["stream_options"]: + if not native_execution and "include_usage" not in kwargs["stream_options"]: kwargs["stream_options"]["include_usage"] = True - # Get provider plugin - plugin = self._get_plugin_instance(provider) skip_cost_calculation = bool( plugin and getattr(plugin, "skip_cost_calculation", False) @@ -843,6 +1635,7 @@ async def _execute_streaming( # Execute request with retries for attempt in range(self._max_retries): last_streamed_chunk: Optional[str] = None + stream_visible_output_emitted = False try: lib_logger.info( @@ -853,21 +1646,95 @@ async def _execute_streaming( await self._run_pre_request_callback( context, kwargs ) - - # Make the API call - if plugin and plugin.has_custom_logic(): + target = _current_route_target(context) + execution = target.execution if target else "auto" + stream_provider_protocol = "openai_chat" + + # Make the API call. Keep execution-mode precedence aligned with + # the non-streaming path: explicit LiteLLM wins, explicit custom + # requires custom logic, explicit native fails closed, and auto + # prefers custom before native streaming. + if execution == "litellm_fallback": + kwargs["api_key"] = credential_secret + kwargs["stream"] = True + self._apply_litellm_logger(kwargs) + kwargs.pop("transaction_context", None) + self._log_executor_trace( + context, + "provider_execution_request", + kwargs, + direction="request", + stage="provider", + credential_id=cred_context.stable_id, + metadata={"execution": "litellm_stream", "provider": provider, "model": model}, + ) + stream = await litellm.acompletion(**kwargs) + elif execution == "custom" or (execution == "auto" and plugin and plugin.has_custom_logic()): + if not plugin or not plugin.has_custom_logic(): + raise RoutingExecutionError(f"Provider {provider} does not support custom execution") kwargs["credential_identifier"] = credential_secret + self._log_executor_trace( + context, + "provider_execution_request", + kwargs, + direction="request", + stage="provider", + credential_id=cred_context.stable_id, + metadata={"execution": "custom_stream", "provider": provider, "model": model}, + ) stream = await plugin.acompletion( self._http_client, **kwargs ) + elif _should_use_native_streaming(plugin, model, target, execution, provider): + native_context, native_request = self._build_native_provider_context( + provider, + model, + plugin, + credential_secret, + cred_context.stable_id, + context, + target, + raw_request=kwargs, + transport="sse", + stream=True, + return_request=True, + ) + self._log_routing_trace( + context, + "routing_native_stream_execution_selected", + _target_trace(target) if target else {"provider": provider, "model": model}, + metadata={"protocol": native_context.protocol_name, "operation": native_context.operation}, + ) + stream = self._get_native_executor().stream(native_request, native_context, NativeHTTPTransport(self._http_client)) + stream_provider_protocol = native_context.protocol_name else: kwargs["api_key"] = credential_secret kwargs["stream"] = True self._apply_litellm_logger(kwargs) # Remove internal context before litellm call kwargs.pop("transaction_context", None) + self._log_executor_trace( + context, + "provider_execution_request", + kwargs, + direction="request", + stage="provider", + credential_id=cred_context.stable_id, + metadata={"execution": "litellm_stream", "provider": provider, "model": model}, + ) stream = await litellm.acompletion(**kwargs) + self._log_executor_trace( + context, + "provider_stream_opened", + {"stream_type": f"{type(stream).__module__}.{type(stream).__name__}"}, + direction="response", + stage="provider", + credential_id=cred_context.stable_id, + metadata={"provider": provider, "model": model}, + snapshot=False, + ) + # Hand off to streaming handler with cred_context # The handler will call mark_success on completion base_stream = self._streaming_handler.wrap_stream( @@ -880,6 +1747,31 @@ async def _execute_streaming( response_callback=lambda response: self._record_session_response( context, response ), + success_callback=lambda: self._clear_failure_history_on_success(provider, model), + transaction_logger=context.transaction_logger, + ) + if client_protocol_context is None: + client_protocol_context = ProtocolContext( + provider=provider, + model=model, + source_protocol="openai_chat", + target_protocol=context.output_protocol_name, + input_protocol=context.input_protocol_name, + provider_protocol=stream_provider_protocol, + output_protocol=context.output_protocol_name, + source_provider=provider, + target_provider=None, + provider_state_compatible=False, + request_id=getattr(context, "request_id", None), + session_id=context.session_id, + credential_stable_id=cred_context.stable_id, + transport="sse", + ) + client_stream = convert_protocol_stream( + base_stream, + source_protocol=get_protocol("openai_chat"), + output_protocol=get_protocol(context.output_protocol_name), + context=client_protocol_context, ) lib_logger.info( @@ -892,15 +1784,23 @@ async def _execute_streaming( async for ( chunk ) in self._transaction_logging_stream_wrapper( - base_stream, + client_stream, context.transaction_logger, context.kwargs, + context=context, + plugin=plugin, ): - last_streamed_chunk = chunk + if not is_stream_heartbeat_or_comment(chunk): + last_streamed_chunk = chunk + if _stream_chunk_is_visible_output(chunk): + stream_visible_output_emitted = True yield chunk else: - async for chunk in base_stream: - last_streamed_chunk = chunk + async for chunk in client_stream: + if not is_stream_heartbeat_or_comment(chunk): + last_streamed_chunk = chunk + if _stream_chunk_is_visible_output(chunk): + stream_visible_output_emitted = True yield chunk return @@ -908,6 +1808,12 @@ async def _execute_streaming( last_exception = e original = getattr(e, "data", e) classified = classify_error(original, provider) + last_stream_error_payload = _streamed_error_payload(e, classified) + if _can_start_stream_provider_cooldown( + last_streamed_chunk, + emitted_output=stream_visible_output_emitted, + ): + await self._maybe_start_provider_cooldown(provider, classified, context=context, model=model, original_error=original) log_failure( api_key=cred, model=model, @@ -934,8 +1840,8 @@ async def _execute_streaming( "type": "quota_exhausted", } } - yield f"data: {json.dumps(error_data)}\n\n" - yield "data: [DONE]\n\n" + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line return else: retry_state.reset_quota_failures() @@ -944,19 +1850,11 @@ async def _execute_streaming( cred_context.mark_failure(classified) raise - if not can_retry_stream_after_error( - last_streamed_chunk, - self._stream_retry_on_reasoning_only_enabled(), - ): + if not _can_retry_stream_after_error(last_streamed_chunk, self._stream_retry_on_reasoning_only_enabled(), emitted_output=stream_visible_output_emitted): cred_context.mark_failure(classified) - error_data = { - "error": { - "message": "Upstream stream failed after output began", - "type": classified.error_type, - } - } - yield f"data: {json.dumps(error_data)}\n\n" - yield "data: [DONE]\n\n" + error_data = _streamed_error_payload(e, classified) + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line return small_cooldown_threshold = int( @@ -993,6 +1891,11 @@ async def _execute_streaming( except (RateLimitError, httpx.HTTPStatusError) as e: last_exception = e classified = classify_error(e, provider) + if _can_start_stream_provider_cooldown( + last_streamed_chunk, + emitted_output=stream_visible_output_emitted, + ): + await self._maybe_start_provider_cooldown(provider, classified, context=context, model=model, original_error=e) log_failure( api_key=cred, model=model, @@ -1019,8 +1922,8 @@ async def _execute_streaming( "type": "quota_exhausted", } } - yield f"data: {json.dumps(error_data)}\n\n" - yield "data: [DONE]\n\n" + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line return else: retry_state.reset_quota_failures() @@ -1029,6 +1932,13 @@ async def _execute_streaming( cred_context.mark_failure(classified) raise + if not _can_retry_stream_after_error(last_streamed_chunk, self._stream_retry_on_reasoning_only_enabled(), emitted_output=stream_visible_output_emitted): + cred_context.mark_failure(classified) + error_data = {"error": {"message": "Upstream stream failed after output began", "type": classified.error_type}} + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line + return + # Check for small cooldown - retry same key instead of rotating small_cooldown_threshold = int( os.environ.get( @@ -1068,6 +1978,11 @@ async def _execute_streaming( ) as e: last_exception = e classified = classify_error(e, provider) + if _can_start_stream_provider_cooldown( + last_streamed_chunk, + emitted_output=stream_visible_output_emitted, + ): + await self._maybe_start_provider_cooldown(provider, classified, context=context, model=model, original_error=e) log_failure( api_key=cred, model=model, @@ -1076,6 +1991,13 @@ async def _execute_streaming( request_headers=request_headers, ) + if not _can_retry_stream_after_error(last_streamed_chunk, self._stream_retry_on_reasoning_only_enabled(), emitted_output=stream_visible_output_emitted): + cred_context.mark_failure(classified) + error_data = {"error": {"message": "Upstream stream failed after output began", "type": classified.error_type}} + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line + return + if attempt >= self._max_retries - 1: error_accumulator.record_error( cred, classified, str(e)[:150] @@ -1101,9 +2023,35 @@ async def _execute_streaming( continue # Retry + except RoutingExecutionError as e: + if e.error_type == "configuration_error": + raise + last_exception = e + classified = classify_error(e, provider) + if _can_start_stream_provider_cooldown(last_streamed_chunk, emitted_output=stream_visible_output_emitted): + await self._maybe_start_provider_cooldown(provider, classified, context=context, model=model, original_error=e) + log_failure(api_key=cred, model=model, attempt=attempt + 1, error=e, request_headers=request_headers) + error_accumulator.record_error(cred, classified, str(e)[:150]) + if not should_rotate_on_error(classified): + cred_context.mark_failure(classified) + raise + if not _can_retry_stream_after_error(last_streamed_chunk, self._stream_retry_on_reasoning_only_enabled(), emitted_output=stream_visible_output_emitted): + cred_context.mark_failure(classified) + error_data = {"error": {"message": "Upstream stream failed after output began", "type": classified.error_type}} + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line + return + cred_context.mark_failure(classified) + break + except Exception as e: last_exception = e classified = classify_error(e, provider) + if _can_start_stream_provider_cooldown( + last_streamed_chunk, + emitted_output=stream_visible_output_emitted, + ): + await self._maybe_start_provider_cooldown(provider, classified, context=context, model=model, original_error=e) log_failure( api_key=cred, model=model, @@ -1119,6 +2067,13 @@ async def _execute_streaming( cred_context.mark_failure(classified) raise + if not _can_retry_stream_after_error(last_streamed_chunk, self._stream_retry_on_reasoning_only_enabled(), emitted_output=stream_visible_output_emitted): + cred_context.mark_failure(classified) + error_data = {"error": {"message": "Upstream stream failed after output began", "type": classified.error_type}} + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line + return + small_cooldown_threshold = int( os.environ.get( "SMALL_COOLDOWN_RETRY_THRESHOLD", @@ -1152,6 +2107,9 @@ async def _execute_streaming( except PreRequestCallbackError: raise + except RoutingExecutionError as exc: + if exc.error_type == "configuration_error": + raise except Exception: # Let context manager handle cleanup pass @@ -1162,20 +2120,23 @@ async def _execute_streaming( # All credentials exhausted or timeout error_accumulator.timeout_occurred = time.time() >= deadline error_data = error_accumulator.build_client_error_response() - yield f"data: {json.dumps(error_data)}\n\n" - yield "data: [DONE]\n\n" + if last_stream_error_payload: + _merge_stream_error_details(error_data, last_stream_error_payload) + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line except NoAvailableKeysError as e: lib_logger.error(f"No keys available: {e}") error_data = {"error": {"message": str(e), "type": "proxy_busy"}} - yield f"data: {json.dumps(error_data)}\n\n" - yield "data: [DONE]\n\n" + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line except Exception as e: lib_logger.error(f"Unhandled exception in streaming: {e}", exc_info=True) - error_data = {"error": {"message": str(e), "type": "proxy_internal_error"}} - yield f"data: {json.dumps(error_data)}\n\n" - yield "data: [DONE]\n\n" + classified = classify_error(e, context.provider) + error_data = {"error": {"message": "Streaming request failed", "type": classified.error_type, "details": {"status_code": classified.status_code}}} + for line in self._terminal_stream_error_lines(context, error_data, protocol_context=client_protocol_context): + yield line def _apply_litellm_provider_params( self, provider: str, kwargs: Dict[str, Any] @@ -1209,9 +2170,12 @@ async def _wait_for_cooldown( self, provider: str, deadline: float, + *, + model: Optional[str] = None, + context: Optional[RequestContext] = None, ) -> None: """ - Wait for provider-level cooldown to end. + Wait for provider/model cooldown to end. Args: provider: Provider name @@ -1220,15 +2184,31 @@ async def _wait_for_cooldown( if not self._cooldown: return - remaining = await self._cooldown.get_remaining_cooldown(provider) + if hasattr(self._cooldown, "get_max_remaining"): + remaining = await self._cooldown.get_max_remaining(provider, model=model) + else: + remaining = await self._cooldown.get_remaining_cooldown(provider) if remaining > 0: budget = deadline - time.time() if remaining > budget: lib_logger.warning( f"Provider {provider} cooldown ({remaining:.1f}s) exceeds budget ({budget:.1f}s)" ) - return # Will fail on no keys available + self._log_provider_cooldown_trace( + context, + "cooldown_wait_exceeds_budget", + provider, + ClassifiedError(error_type="rate_limit", original_exception=RuntimeError("cooldown exceeds budget"), retry_after=int(remaining)), + int(remaining), + "cooldown_exceeds_budget", + model=model, + ) + raise RoutingExecutionError( + f"Provider {provider} cooldown exceeds request budget", + error_type="rate_limit", + ) lib_logger.info(f"Waiting {remaining:.1f}s for {provider} cooldown") + self._log_cooldown_wait_trace(context, provider, model, remaining) await asyncio.sleep(remaining) async def _handle_error_with_context( @@ -1241,6 +2221,7 @@ async def _handle_error_with_context( error_accumulator: RequestErrorAccumulator, retry_state: RetryState, request_headers: Dict[str, Any], + context: Optional[RequestContext] = None, ) -> str: """ Handle an error and determine next action. @@ -1289,6 +2270,8 @@ async def _handle_error_with_context( cred_context.mark_failure(classified) return ErrorAction.FAIL + await self._maybe_start_provider_cooldown(provider, classified, context=context, model=model, original_error=error) + # Check if should retry same key (including small cooldown auto-retry) small_cooldown_threshold = int( os.environ.get( @@ -1334,6 +2317,135 @@ async def _handle_error_with_context( ) return ErrorAction.ROTATE + async def _maybe_start_provider_cooldown( + self, + provider: str, + classified: ClassifiedError, + *, + context: Optional[RequestContext], + model: Optional[str] = None, + original_error: Any = None, + ) -> None: + """Start provider-wide cooldown for large provider-level throttles. + + This is intentionally conservative: small retry-after values stay on the + same credential path, and quota cooldown is disabled by default because + most quota errors are per credential or account. + """ + + if not self._cooldown: + return + small_cooldown_threshold = int( + os.environ.get( + "SMALL_COOLDOWN_RETRY_THRESHOLD", DEFAULT_SMALL_COOLDOWN_RETRY_THRESHOLD + ) + ) + min_seconds, default_seconds, cooldown_on_quota = provider_cooldown_env() + decision = decide_provider_cooldown( + classified, + small_cooldown_threshold=small_cooldown_threshold, + provider_cooldown_min_seconds=min_seconds, + default_duration=default_seconds, + cooldown_on_quota=cooldown_on_quota, + provider=provider, + model=model, + original_error=original_error, + failure_history=getattr(self, "_failure_history", None), + ) + if not decision.should_start: + if decision.reason == "transient_backoff_threshold_not_met" and classified.error_type in {"server_error", "api_connection"}: + history = getattr(self, "_failure_history", None) + if history is not None: + history.record(provider=provider, model=decision.model, error_type=classified.error_type, scope=decision.scope, duration=0, reason=decision.reason) + self._log_provider_cooldown_trace( + context, + "provider_cooldown_skipped", + provider, + classified, + decision.duration, + decision.reason, + scope=decision.scope, + model=decision.model, + backoff_level=decision.backoff_level, + ) + return + try: + if hasattr(self._cooldown, "start_scoped_cooldown"): + await self._cooldown.start_scoped_cooldown(provider, decision.duration, model=decision.model, scope=decision.scope, reason=decision.reason) + else: + await self._cooldown.start_cooldown(provider, decision.duration) + history = getattr(self, "_failure_history", None) + if history is not None: + history.record(provider=provider, model=decision.model, error_type=classified.error_type, scope=decision.scope, duration=decision.duration, reason=decision.reason) + self._log_provider_cooldown_trace( + context, + "provider_cooldown_started", + provider, + classified, + decision.duration, + decision.reason, + scope=decision.scope, + model=decision.model, + backoff_level=decision.backoff_level, + ) + except Exception as exc: + lib_logger.debug("Failed to start provider cooldown for %s: %s", provider, exc) + + def _clear_failure_history_on_success(self, provider: str, model: Optional[str]) -> None: + """Clear transient failure history after a successful provider/model call.""" + + history = getattr(self, "_failure_history", None) + if history is not None and hasattr(history, "clear"): + history.clear(provider=provider, model=model) + + @staticmethod + def _log_provider_cooldown_trace( + context: Optional[RequestContext], + pass_name: str, + provider: str, + classified: ClassifiedError, + duration: int, + reason: str, + *, + scope: str = "provider", + model: Optional[str] = None, + backoff_level: int = 0, + ) -> None: + if not context or not context.transaction_logger: + return + context.transaction_logger.log_transform_pass( + pass_name, + {"provider": provider, "model": model, "scope": scope, "error_type": classified.error_type, "duration": duration}, + direction="metadata", + stage="retry", + metadata={ + "provider": provider, + "duration": duration, + "scope": scope, + "model": model, + "backoff_level": backoff_level, + "error_type": classified.error_type, + "retry_after_present": classified.retry_after is not None, + "reason": reason, + }, + snapshot=False, + ) + + @staticmethod + def _log_cooldown_wait_trace(context: Optional[RequestContext], provider: str, model: Optional[str], remaining: float) -> None: + """Trace cooldown waits without exposing credentials or payloads.""" + + if not context or not context.transaction_logger: + return + context.transaction_logger.log_transform_pass( + "cooldown_wait", + {"provider": provider, "model": model, "remaining": remaining}, + direction="metadata", + stage="retry", + metadata={"provider": provider, "model": model, "remaining": remaining}, + snapshot=False, + ) + def _record_session_response(self, context: RequestContext, response: Any) -> None: """Let the tracker learn anchors emitted by a successful response. @@ -1348,7 +2460,7 @@ def _record_session_response(self, context: RequestContext, response: Any) -> No context.session_id, provider=context.provider, model=context.model, - scope_key=context.usage_manager_key, + scope_key=context.session_isolation_key, tracking_namespace=context.session_tracking_namespace, response=response, ) @@ -1377,8 +2489,10 @@ async def _validate_request( provider: str, model: str, kwargs: Dict[str, Any], + *, + plugin: Any = None, ) -> None: - plugin = self._get_plugin_instance(provider) + plugin = plugin or self._get_plugin_instance(provider) if not plugin or not hasattr(plugin, "validate_request"): return @@ -1390,60 +2504,52 @@ async def _validate_request( if isinstance(result, str): raise ValueError(result) - def _extract_usage_tokens(self, response: Any) -> tuple[int, int, int, int, int]: - prompt_tokens = 0 - completion_tokens = 0 - cached_tokens = 0 - cache_write_tokens = 0 - thinking_tokens = 0 - - if hasattr(response, "usage") and response.usage: - prompt_tokens = getattr(response.usage, "prompt_tokens", 0) or 0 - completion_tokens = getattr(response.usage, "completion_tokens", 0) or 0 - - prompt_details = getattr(response.usage, "prompt_tokens_details", None) - if prompt_details: - if isinstance(prompt_details, dict): - cached_tokens = prompt_details.get("cached_tokens", 0) or 0 - cache_write_tokens = ( - prompt_details.get("cache_creation_tokens", 0) or 0 - ) - else: - cached_tokens = getattr(prompt_details, "cached_tokens", 0) or 0 - cache_write_tokens = ( - getattr(prompt_details, "cache_creation_tokens", 0) or 0 - ) + def _account_for_response_usage( + self, + provider: str, + model: str, + response: Any, + context: RequestContext, + ) -> tuple[UsageRecord, CostBreakdown]: + """Normalize usage and advisory cost for one successful response.""" - completion_details = getattr( - response.usage, "completion_tokens_details", None - ) - if completion_details: - if isinstance(completion_details, dict): - thinking_tokens = completion_details.get("reasoning_tokens", 0) or 0 - else: - thinking_tokens = ( - getattr(completion_details, "reasoning_tokens", 0) or 0 - ) + usage_record = extract_usage_record( + response, + provider=provider, + model=model, + source="executor_response", + ) + plugin = self._get_plugin_instance(provider) + cost_breakdown = CostCalculator(provider_plugin=plugin).calculate( + usage_record, + model=model, + response=response, + ) + self._trace_usage_accounting(context, usage_record, cost_breakdown) + return usage_record, cost_breakdown - cache_read_tokens = getattr(response.usage, "cache_read_tokens", None) - if cache_read_tokens is not None: - cached_tokens = cache_read_tokens or 0 - cache_creation_tokens = getattr( - response.usage, "cache_creation_tokens", None - ) - if cache_creation_tokens is not None: - cache_write_tokens = cache_creation_tokens or 0 - - if thinking_tokens and completion_tokens >= thinking_tokens: - completion_tokens = completion_tokens - thinking_tokens - - uncached_prompt = max(0, prompt_tokens - cached_tokens) - return ( - uncached_prompt, - completion_tokens, - cached_tokens, - cache_write_tokens, - thinking_tokens, + @staticmethod + def _trace_usage_accounting( + context: RequestContext, + usage_record: UsageRecord, + cost_breakdown: CostBreakdown, + ) -> None: + """Record normalized usage/cost trace data without affecting requests.""" + + if not context.transaction_logger: + return + context.transaction_logger.log_transform_pass( + "usage_accounting_summary", + {"usage": usage_record.to_dict(), "cost": cost_breakdown.to_dict()}, + direction="metadata", + stage="final", + metadata={ + "provider": usage_record.provider, + "model": usage_record.model, + "source": usage_record.source, + "pricing_source": cost_breakdown.pricing_source, + }, + snapshot=False, ) @staticmethod @@ -1453,39 +2559,22 @@ def _normalize_response_usage(response: Any, model: str) -> Any: Delegates to normalize_usage_for_response which handles both dicts (streaming) and pydantic objects (non-streaming). - Internal tracking values from _extract_usage_tokens are unaffected. + Internal tracking values from UsageRecord accounting are unaffected. """ - if hasattr(response, "usage") and response.usage: + if isinstance(response, dict): + normalize_usage_for_response(response.get("usage"), model) + elif hasattr(response, "usage") and response.usage: normalize_usage_for_response(response.usage, model) return response - def _calculate_cost(self, provider: str, model: str, response: Any) -> float: - plugin = self._get_plugin_instance(provider) - if plugin and getattr(plugin, "skip_cost_calculation", False): - return 0.0 - - try: - if isinstance(response, litellm.EmbeddingResponse): - model_info = litellm.get_model_info(model) - input_cost = model_info.get("input_cost_per_token") - if input_cost: - return (response.usage.prompt_tokens or 0) * input_cost - return 0.0 - - cost = litellm.completion_cost( - completion_response=response, - model=model, - ) - return float(cost) if cost is not None else 0.0 - except Exception as exc: - lib_logger.debug(f"Cost calculation failed for {model}: {exc}") - return 0.0 - async def _transaction_logging_stream_wrapper( self, stream: AsyncGenerator[str, None], transaction_logger: TransactionLogger, request_kwargs: Dict[str, Any], + *, + context: Optional[RequestContext] = None, + plugin: Any = None, ) -> AsyncGenerator[str, None]: """ Wrap a stream to log chunks and final response to TransactionLogger. @@ -1503,6 +2592,29 @@ async def _transaction_logging_stream_wrapper( chunks = [] async for sse_line in stream: + trace_sse_line = _redact_stream_sse_for_trace( + sse_line, + context, + plugin, + config=getattr(self, "_experimental_config", None), + ) + transaction_logger.log_transform_pass( + "raw_stream_chunk", + trace_sse_line, + direction="stream", + stage="client", + transport="sse", + snapshot=False, + ) + if sse_line.startswith("data: [DONE]"): + transaction_logger.log_transform_pass( + "stream_done_event", + {"raw": trace_sse_line}, + direction="stream", + stage="final", + transport="sse", + snapshot=False, + ) yield sse_line # Parse and accumulate for final logging @@ -1514,7 +2626,23 @@ async def _transaction_logging_stream_wrapper( if content: chunk_data = json.loads(content) chunks.append(chunk_data) - transaction_logger.log_stream_chunk(chunk_data) + trace_chunk_data = _redact_context_field_cache_paths( + chunk_data, + context, + "stream", + plugin, + config=getattr(self, "_experimental_config", None), + ) if context else chunk_data + transaction_logger.log_stream_chunk(trace_chunk_data) + if isinstance(chunk_data, dict) and chunk_data.get("error") is not None: + transaction_logger.log_transform_pass( + "stream_error_event", + trace_chunk_data, + direction="stream", + stage="client", + transport="sse", + snapshot=False, + ) except json.JSONDecodeError: lib_logger.debug( f"Failed to parse chunk for logging: {sse_line[:100]}" @@ -1524,8 +2652,819 @@ async def _transaction_logging_stream_wrapper( if chunks: try: final_response = TransactionLogger.assemble_streaming_response(chunks) - transaction_logger.log_response(final_response) + trace_final_response = _redact_context_field_cache_paths( + final_response, + context, + "stream", + plugin, + config=getattr(self, "_experimental_config", None), + ) if context else final_response + transaction_logger.log_transform_pass( + "assembled_stream_response", + trace_final_response, + direction="response", + stage="client", + transport="sse", + ) + transaction_logger.log_response(trace_final_response) except Exception as e: lib_logger.debug( f"Failed to assemble/log final streaming response: {e}" ) + + +def _target_trace(target: RouteTarget) -> Dict[str, Any]: + """Return non-secret route target metadata for transaction traces.""" + + return { + "name": target.name, + "provider": target.provider, + "model": target.prefixed_model, + "execution": target.execution, + "protocol": target.protocol, + } + + +def _merge_canonical_sequence( + original: list[Any], + baseline: list[Any], + attempted: list[Any], + merge_item: Any, + identity: Any, +) -> list[Any]: + """Apply Chat-view edits while retaining source-native item metadata.""" + + if len(original) == len(baseline) == len(attempted): + return [ + deepcopy(source) if before == after else merge_item(source, before, after) + for source, before, after in zip(original, baseline, attempted) + ] + merged: list[Any] = [] + unused = set(range(min(len(original), len(baseline)))) + for index, after in enumerate(attempted): + exact = next((position for position in unused if baseline[position] == after), None) + if exact is not None: + merged.append(deepcopy(original[exact])) + unused.remove(exact) + continue + if index in unused and identity(baseline[index]) == identity(after): + merged.append(merge_item(original[index], baseline[index], after)) + unused.remove(index) + continue + merged.append(_without_source_artifacts(after)) + return merged + + +def _merge_canonical_message(original: Any, baseline: Any, attempted: Any) -> Any: + if original.role != baseline.role or baseline.role != attempted.role: + return _without_source_artifacts(attempted) + merged = deepcopy(original) + for field_name in ("role", "name", "tool_call_id"): + if getattr(baseline, field_name) != getattr(attempted, field_name): + setattr(merged, field_name, deepcopy(getattr(attempted, field_name))) + if baseline.content != attempted.content: + merged.content = _merge_canonical_sequence( + original.content, + baseline.content, + attempted.content, + _merge_canonical_content_block, + lambda block: block.type, + ) + if baseline.tool_calls != attempted.tool_calls: + merged.tool_calls = _merge_canonical_sequence( + original.tool_calls, + baseline.tool_calls, + attempted.tool_calls, + _merge_canonical_tool_call, + lambda call: (call.type, call.id, call.name), + ) + if baseline.reasoning != attempted.reasoning: + merged.reasoning = deepcopy(attempted.reasoning) + return merged + + +def _merge_canonical_content_block(original: Any, baseline: Any, attempted: Any) -> Any: + if original.type != baseline.type or baseline.type != attempted.type: + return _without_source_artifacts(attempted) + merged = deepcopy(original) + for field_name in ("type", "text", "source"): + if getattr(baseline, field_name) != getattr(attempted, field_name): + setattr(merged, field_name, deepcopy(getattr(attempted, field_name))) + if baseline.tool_call != attempted.tool_call: + merged.tool_call = ( + _merge_canonical_tool_call(original.tool_call, baseline.tool_call, attempted.tool_call) + if original.tool_call and baseline.tool_call and attempted.tool_call + else deepcopy(attempted.tool_call) + ) + if baseline.tool_result != attempted.tool_result: + merged.tool_result = deepcopy(attempted.tool_result) + if baseline.reasoning != attempted.reasoning: + merged.reasoning = deepcopy(attempted.reasoning) + return merged + + +def _merge_canonical_tool_call(original: Any, baseline: Any, attempted: Any) -> Any: + if original is None or baseline is None or attempted is None: + return deepcopy(attempted) + merged = deepcopy(original) + changed = False + for field_name in ("type", "id", "name", "arguments", "index"): + if getattr(baseline, field_name) != getattr(attempted, field_name): + setattr(merged, field_name, deepcopy(getattr(attempted, field_name))) + changed = True + if changed: + merged.raw = None + merged.extra = {} + return merged + + +def _merge_canonical_tool(original: Any, baseline: Any, attempted: Any) -> Any: + if original.type != baseline.type or baseline.type != attempted.type: + return _without_source_artifacts(attempted) + merged = deepcopy(original) + for field_name in ("type", "name", "description", "input_schema"): + if getattr(baseline, field_name) != getattr(attempted, field_name): + setattr(merged, field_name, deepcopy(getattr(attempted, field_name))) + return merged + + +def _without_source_artifacts(value: Any) -> Any: + """Copy a callback-created canonical value without Chat wire fragments.""" + + copied = deepcopy(value) + if hasattr(copied, "raw"): + copied.raw = None + if hasattr(copied, "extra"): + copied.extra = {} + if hasattr(copied, "content") and isinstance(copied.content, list): + copied.content = [_without_source_artifacts(item) for item in copied.content] + if hasattr(copied, "tool_calls") and isinstance(copied.tool_calls, list): + copied.tool_calls = [_without_source_artifacts(item) for item in copied.tool_calls] + return copied + + +def _current_route_target(context: RequestContext) -> Optional[RouteTarget]: + """Return the currently selected route target from context metadata.""" + + targets = tuple(context.routing_targets or ()) + if not targets: + return None + if context.routing_target_index < 0 or context.routing_target_index >= len(targets): + return None + return targets[context.routing_target_index] + + +def _provider_native_protocol(plugin: Any, model: str, target: Optional[RouteTarget]) -> Optional[str]: + """Resolve native protocol from target override or provider declaration.""" + + if target and target.protocol: + return target.protocol + if plugin and hasattr(plugin, "get_protocol_name"): + return plugin.get_protocol_name(model) + return None + + +def _should_use_native_protocol(plugin: Any, model: str, target: Optional[RouteTarget], kwargs: Dict[str, Any], *, stream: bool, execution: str) -> bool: + """Return whether auto routing should use provider-native execution.""" + + protocol_name = _provider_native_protocol(plugin, model, target) + if not plugin or not protocol_name: + return False + native_model = plugin.normalize_native_model(model) if hasattr(plugin, "normalize_native_model") else _strip_provider_prefix(model) + operation = plugin.get_native_operation(native_model, None, stream=stream) if hasattr(plugin, "get_native_operation") else "chat" + hook = getattr(plugin, "should_use_native_protocol", None) + if callable(hook): + return bool(hook(model=native_model, operation=operation, stream=stream, execution=execution)) + support = getattr(plugin, "supports_native_operation", None) + return bool(support(native_model, operation) if callable(support) else True) + + +def _strip_provider_prefix(model: str) -> str: + """Return model without the proxy-facing provider prefix.""" + + return model.split("/", 1)[1] if "/" in model else model + + +def _provider_supports_native_streaming(plugin: Any, model: str) -> bool: + """Return whether a provider declares native streaming support.""" + + support = getattr(plugin, "supports_native_streaming", None) + if not callable(support): + return False + operation = "chat" + resolver = getattr(plugin, "get_native_operation", None) + if callable(resolver): + try: + operation = resolver(model, None, stream=True) + except TypeError: + try: + operation = resolver(model) + except Exception: + return False + except Exception: + return False + return native_provider_supports_streaming(plugin, model=model, operation=operation) + + +def _should_use_native_streaming(plugin: Any, model: str, target: Optional[RouteTarget], execution: str, provider: str) -> bool: + """Return whether streaming may use the native executor. + + Explicit `@native` routing is still constrained by provider capability. + Falling through to native streaming when a provider has not opted in is unsafe + because the generic stream wrapper currently expects LiteLLM-shaped chunks. + """ + + if execution == "native": + if _provider_supports_native_streaming(plugin, model): + return True + raise RoutingExecutionError( + f"Provider {provider} does not support native streaming for {model}", + error_type="configuration_error", + ) + if execution != "auto" or not plugin: + return False + if not _should_use_native_protocol(plugin, model, target, {"model": model, "stream": True}, stream=True, execution=execution): + return False + return _provider_supports_native_streaming(plugin, model) + + +def _redact_context_field_cache_paths( + payload: Any, + context: RequestContext, + direction: str, + plugin: Any, + *, + config: Any = None, +) -> Any: + """Redact configured field-cache paths before executor-level traces. + + Native provider execution already redacts its internal trace passes. The + client executor also logs returned responses, so it must apply the same + rule-aware redaction there without mutating the client-facing response. + """ + + if direction not in {"response", "stream"} or not plugin or not _provider_native_protocol(plugin, context.model, _current_route_target(context)): + return payload + try: + rules = _merged_field_cache_rules( + context.provider, + context.model, + plugin, + config=config, + ) + except RoutingExecutionError: + raise + except Exception: + return payload + if not rules: + return payload + redacted = deepcopy(payload) + for rule in rules: + if direction == "response" and getattr(rule, "source", None) not in {"response", "unified_response"}: + continue + if direction == "stream" and getattr(rule, "source", None) not in {"stream_event", "unified_stream_event", "response", "unified_response"}: + continue + for path in _trace_redaction_paths((rule.path,), direction=direction): + try: + tokens = parse_path(path) + _redact_trace_path(redacted, tokens) + _redact_trace_leaf_key(redacted, tokens) + except (FieldCachePathError, TypeError, ValueError): + continue + return redacted + + +def _trace_redaction_paths(paths: tuple[str, ...] | list[str], *, direction: str) -> list[str]: + """Return configured paths plus raw-stream envelope fallbacks for traces.""" + + expanded: list[str] = [] + for path in paths: + expanded.append(path) + if direction == "stream" and path.startswith("raw."): + expanded.append(path[4:]) + return expanded + + +def _redact_stream_sse_for_trace( + sse_line: str, + context: Optional[RequestContext], + plugin: Any, + *, + config: Any = None, +) -> str: + """Return a trace-only SSE line with configured native cache paths redacted.""" + + if not context or not isinstance(sse_line, str) or not sse_line.startswith("data: ") or sse_line.startswith("data: [DONE]"): + return sse_line + try: + payload = json.loads(sse_line[6:].strip()) + except json.JSONDecodeError: + return sse_line + redacted = _redact_context_field_cache_paths( + payload, + context, + "stream", + plugin, + config=config, + ) + if redacted is payload: + return sse_line + return f"data: {json.dumps(redacted, separators=(',', ':'))}\n\n" + + +def _redact_trace_path(value: Any, tokens: tuple[PathToken, ...]) -> None: + if not tokens: + return + token = tokens[0] + rest = tokens[1:] + if token.kind == "key": + if isinstance(value, dict) and token.value in value: + if rest: + _redact_trace_path(value[token.value], rest) + else: + value[token.value] = REDACTED + return + if token.kind == "index": + if isinstance(value, list) and value: + index = int(token.value) + if -len(value) <= index < len(value): + if rest: + _redact_trace_path(value[index], rest) + else: + value[index] = REDACTED + return + if token.kind == "wildcard": + if isinstance(value, dict): + for key in list(value.keys()): + if rest: + _redact_trace_path(value[key], rest) + else: + value[key] = REDACTED + elif isinstance(value, list): + for index, item in enumerate(value): + if rest: + _redact_trace_path(item, rest) + else: + value[index] = REDACTED + + +def _redact_trace_leaf_key(value: Any, tokens: tuple[PathToken, ...]) -> None: + """Redact terminal configured cache keys across duplicated trace envelopes.""" + + leaf = next((token.value for token in reversed(tokens) if token.kind == "key"), None) + if not leaf: + return + if isinstance(value, dict): + for key, item in list(value.items()): + if key == leaf: + value[key] = REDACTED + else: + _redact_trace_leaf_key(item, tokens) + elif isinstance(value, list): + for item in value: + _redact_trace_leaf_key(item, tokens) + + +def _merged_field_cache_rules( + provider: str, + model: str, + plugin: Any, + *, + config: Any = None, +) -> tuple[Any, ...]: + """Merge provider-declared and JSON-configured field-cache rules. + + Provider declarations are the safe default. Optional JSON config can add or + replace rules by name so operators can tune protocol-state preservation per + provider/model without editing provider code. The import is local to keep the + experimental config layer out of executor module initialization. + """ + + declared = list(plugin.get_field_cache_rules(model) if plugin and hasattr(plugin, "get_field_cache_rules") else ()) + try: + from ..config.experimental import load_experimental_config, parse_field_cache_rules + + configured = list( + parse_field_cache_rules( + config if config is not None else load_experimental_config(), + provider, + model, + ) + ) + except Exception as exc: + raise RoutingExecutionError(f"Invalid field-cache configuration for {provider}/{model}", error_type="configuration_error") from exc + if not configured: + return tuple(declared) + merged: dict[str, Any] = {getattr(rule, "name", str(index)): rule for index, rule in enumerate(declared)} + order = [getattr(rule, "name", str(index)) for index, rule in enumerate(declared)] + for rule in configured: + name = getattr(rule, "name", "") + if name and name not in merged: + order.append(name) + if name: + declared_rule = merged.get(name) + if declared_rule is not None and not _safe_field_cache_override( + declared_rule, + rule, + ): + raise RoutingExecutionError( + f"Configured field-cache rule {name!r} cannot weaken provider state isolation or injection behavior", + error_type="configuration_error", + ) + merged[name] = rule + return tuple(merged[name] for name in order if name in merged) + + +def _safe_field_cache_override(declared: Any, configured: Any) -> bool: + """Allow extraction tuning only when provider-state behavior is unchanged.""" + + behavior_keys = ( + "provider_continuation", + "tool_call_id_path", + "tool_container_path", + "tool_value_path", + "inject_tool_call_id_path", + "turn_container_path", + "turn_role_path", + "turn_value_path", + ) + return ( + getattr(configured, "cache_key", None) == getattr(declared, "cache_key", None) + and getattr(configured, "mode", None) == getattr(declared, "mode", None) + and getattr(configured, "scope", None) == getattr(declared, "scope", None) + and getattr(configured, "ttl_seconds", None) == getattr(declared, "ttl_seconds", None) + and getattr(configured, "allow_missing_session", None) == getattr(declared, "allow_missing_session", None) + and getattr(configured, "max_values", None) == getattr(declared, "max_values", None) + and getattr(configured, "max_bytes", None) == getattr(declared, "max_bytes", None) + and getattr(configured, "inject", None) == getattr(declared, "inject", None) + and all( + (getattr(configured, "metadata", {}) or {}).get(key) + == (getattr(declared, "metadata", {}) or {}).get(key) + for key in behavior_keys + ) + ) + + +def _target_scope_value(target: RouteTarget, key: str, default: Any) -> Any: + """Read request-scope metadata attached by routing resolution.""" + + scope = target.metadata.get("request_scope") if isinstance(target.metadata, dict) else None + if isinstance(scope, dict) and key in scope: + return scope[key] + return default + + +def _route_error_type(error: BaseException, provider: Optional[str] = None) -> str: + """Map an exception to a fallback-policy error type.""" + + if isinstance(error, asyncio.CancelledError): + return "cancelled" + if isinstance(error, StructuredAPIResponseError): + response_type = _route_error_type_from_response(error.response) + if response_type: + return response_type + explicit = getattr(error, "error_type", None) + if explicit: + return normalize_route_error_type(str(explicit)) + classified = classify_error(error, provider) + return normalize_route_error_type(classified.error_type) + + +def _route_error_type_from_response(response: Any) -> Optional[str]: + """Infer retryability from the proxy's structured error response.""" + + if not isinstance(response, dict) or not isinstance(response.get("error"), dict): + return None + error = response["error"] + details = error.get("details") if isinstance(error.get("details"), dict) else {} + candidates = _structured_error_candidates(error, details) + hard_stop = _first_route_error_candidate(candidates, _HARD_STOP_ROUTE_ERRORS) + if hard_stop: + return hard_stop + retryable = _first_route_error_candidate(candidates, _RETRYABLE_ROUTE_ERRORS) + if retryable: + return retryable + normal_summary = str(details.get("normal_error_summary", "")).lower() + if any(token in normal_summary for token in ("authentication", "forbidden", "invalid_request", "context_window", "credential_reauth", "configuration_error")): + return _summary_hard_stop_type(normal_summary) + if any(token in normal_summary for token in ("rate_limit", "quota", "capacity")): + return "quota_exceeded" if "quota" in normal_summary else "rate_limit" + if any(token in normal_summary for token in ("server_error", "api_connection", "transient")): + return "api_connection" if "api_connection" in normal_summary else "server_error" + error_type = normalize_route_error_type(str(error.get("type", ""))) + if error_type in {"proxy_timeout", "proxy_all_credentials_exhausted"}: + return "rate_limit" + return None + + +def _attempt_uses_native_protocol(plugin: Any, model: str, context: RequestContext) -> bool: + """Predict execution mode before attempt-specific request preparation.""" + + target = _current_route_target(context) + execution = target.execution if target else "auto" + if execution in {"custom", "litellm_fallback"}: + return False + if execution == "auto" and plugin and getattr(plugin, "has_custom_logic", lambda: False)(): + return False + if context.streaming: + return _should_use_native_streaming(plugin, model, target, execution, context.provider) + return execution == "native" or _should_use_native_protocol( + plugin, + model, + target, + context.kwargs, + stream=context.streaming, + execution=execution, + ) + + +def _raise_for_structured_response_error(response: Any) -> None: + """Prevent structured errors from being formatted as empty successes.""" + + error = structured_api_response_error(response) + if error: + raise error + + +def _route_status_code_from_response(response: Any) -> Optional[int]: + """Return a structured status code without reading raw provider text.""" + + if not isinstance(response, dict) or not isinstance(response.get("error"), dict): + return None + error = response["error"] + details = error.get("details") if isinstance(error.get("details"), dict) else {} + for candidate in (details.get("status_code"), details.get("status"), error.get("status_code"), error.get("status"), error.get("code")): + try: + return int(candidate) + except (TypeError, ValueError): + continue + return None + + +def _target_failure_summary(target: RouteTarget, error_type: str, *, status_code: Optional[int] = None) -> Dict[str, Any]: + """Return a client-safe fallback target failure summary.""" + + summary = { + "target": target.name, + "provider": target.provider, + "model": target.prefixed_model, + "execution": target.execution, + "error_type": normalize_route_error_type(error_type), + # Provider error text can contain raw upstream payload fragments or + # credential-like identifiers. Keep the cross-target summary structural; + # detailed per-credential errors remain in the existing sanitized error + # accumulator. + "message": "", + } + if status_code is not None: + summary["status_code"] = status_code + return summary + + +def _append_routing_attempt_history( + context: RequestContext, + target: RouteTarget, + target_index: int, + *, + success: bool, + error_type: Optional[str] = None, + status_code: Optional[int] = None, + emitted_output: Optional[bool] = None, + fallback_allowed: Optional[bool] = None, + duration_ms: Optional[int] = None, +) -> None: + """Append sanitized live fallback-attempt metadata to the request context.""" + + entry: Dict[str, Any] = { + "target_index": target_index, + "target": target.name, + "provider": target.provider, + "model": target.prefixed_model, + "execution": target.execution, + "success": success, + } + if error_type is not None: + entry["error_type"] = normalize_route_error_type(error_type) + if status_code is not None: + entry["status_code"] = status_code + if emitted_output is not None: + entry["emitted_output"] = bool(emitted_output) + if fallback_allowed is not None: + entry["fallback_allowed"] = bool(fallback_allowed) + if duration_ms is not None: + entry["duration_ms"] = max(0, duration_ms) + context.routing_attempt_history.append(entry) + + +def _elapsed_ms(started_at: float) -> int: + return int((time.monotonic() - started_at) * 1000) + + +def _group_streaming_policy(group: Any) -> str: + """Return the active streaming fallback policy for trace and decisions.""" + + return str(getattr(group, "streaming_policy", "pre_output_only") or "pre_output_only") + + +def _streaming_policy_allows_fallback(group: Any) -> bool: + """Return whether a group permits pre-output streaming fallback.""" + + return _group_streaming_policy(group) != "never" + + +_HARD_STOP_ROUTE_ERRORS = { + "authentication", + "forbidden", + "invalid_request", + "context_window_exceeded", + "credential_reauth_needed", + "pre_request_callback_error", + "cancelled", + "configuration_error", +} +_RETRYABLE_ROUTE_ERRORS = {"rate_limit", "quota_exceeded", "server_error", "api_connection", "unsupported_operation"} + + +def _structured_error_candidates(error: Dict[str, Any], details: Dict[str, Any]) -> List[str]: + """Return normalized structured error hints before reading free-form text.""" + + values: List[Any] = [ + error.get("type"), + error.get("code"), + error.get("status"), + details.get("classification"), + details.get("error_type"), + details.get("status"), + ] + for abnormal in details.get("abnormal_errors") or []: + if isinstance(abnormal, dict): + values.append(abnormal.get("error_type")) + values.append(abnormal.get("status_code")) + status_code = _route_status_code_from_response({"error": error}) + if status_code == 400: + values.append("invalid_request") + elif status_code == 401: + values.append("authentication") + elif status_code == 403: + values.append("forbidden") + elif status_code == 429: + values.append("rate_limit") + elif status_code is not None and status_code >= 500: + values.append("server_error") + return [normalize_route_error_type(str(value)) for value in values if value not in (None, "")] + + +def _first_route_error_candidate(candidates: List[str], allowed: set[str]) -> Optional[str]: + """Return the first structured candidate in an allowed policy set.""" + + for candidate in candidates: + if candidate in allowed: + return candidate + return None + + +def _summary_hard_stop_type(summary: str) -> str: + """Map legacy summary text to a hard-stop category without using raw text.""" + + if "credential_reauth" in summary: + return "credential_reauth_needed" + if "context_window" in summary: + return "context_window_exceeded" + if "configuration_error" in summary or "config" in summary: + return "configuration_error" + if "forbidden" in summary: + return "forbidden" + if "authentication" in summary: + return "authentication" + return "invalid_request" + + +def _with_fallback_summary(response: Any, target_failures: List[Dict[str, Any]], attempt_history: Optional[List[Dict[str, Any]]] = None) -> Any: + """Attach fallback target summaries to a structured error response.""" + + if not target_failures or not isinstance(response, dict) or not isinstance(response.get("error"), dict): + return response + details = response["error"].setdefault("details", {}) + if isinstance(details, dict): + details["fallback_targets"] = list(target_failures) + if attempt_history: + details["routing_attempt_history"] = list(attempt_history) + return response + + +def _stream_chunk_is_visible_output(chunk: str) -> bool: + """Return whether a stream chunk should block cross-target fallback. + + Only client-visible model output should lock the route. Empty frames, DONE + sentinels, and structured error frames are not considered visible output, so + a provider that fails before producing content can still fall through to the + next ordered target. + """ + + return is_visible_stream_output(chunk) + + +def _stream_chunk_error_type(chunk: str) -> Optional[str]: + """Return a route error type for terminal stream error frames. + + Per-target stream executors can emit a structured error SSE and `[DONE]` + instead of raising. Fallback wrappers must treat those frames as target + failures before visible output, while still forwarding them if no fallback is + available. + """ + + payload = _stream_chunk_payload(chunk) + if not isinstance(payload, dict): + return None + event_type = normalize_route_error_type(str(payload.get("event_type") or payload.get("type") or "")) + if event_type == "error": + error = payload.get("error") if isinstance(payload.get("error"), dict) else payload + return _route_error_type_from_response({"error": error}) or "server_error" + if event_type == "response.failed": + response = payload.get("response") if isinstance(payload.get("response"), dict) else {} + error = payload.get("error") if isinstance(payload.get("error"), dict) else response.get("error") if isinstance(response.get("error"), dict) else {"type": "server_error"} + return _route_error_type_from_response({"error": error}) or "server_error" + if isinstance(payload.get("error"), dict): + return _route_error_type_from_response({"error": payload["error"]}) or "server_error" + return None + + +def _stream_chunk_payload(chunk: str) -> Optional[Dict[str, Any]]: + """Parse a minimal SSE payload for routing decisions only.""" + + text = str(chunk or "").strip() + if not text: + return None + event_type = None + data_lines: List[str] = [] + for line in text.splitlines(): + stripped = line.strip() + if not stripped or stripped.startswith(":"): + continue + if stripped.startswith("event:"): + event_type = stripped[6:].strip() + continue + if stripped.startswith("data:"): + data_lines.append(stripped[5:].strip()) + if not data_lines: + return {"event_type": event_type} if event_type else None + data = "\n".join(data_lines).strip() + if not data or data == "[DONE]": + return None + try: + parsed = json.loads(data) + except json.JSONDecodeError: + return None + if not isinstance(parsed, dict): + return None + if event_type and "event_type" not in parsed: + parsed["event_type"] = event_type + return parsed + + +def _can_start_stream_provider_cooldown(last_streamed_chunk: Optional[str], *, emitted_output: bool = False) -> bool: + """Return whether a streaming failure occurred before visible output.""" + + if emitted_output: + return False + return last_streamed_chunk is None or not _stream_chunk_is_visible_output(last_streamed_chunk) + + +def _can_retry_stream_after_error(last_streamed_chunk: Optional[str], allow_reasoning_only_retry: bool, *, emitted_output: bool = False) -> bool: + """Return whether a stream can retry/rotate without duplicating output.""" + + if emitted_output: + return False + return can_retry_stream_after_error(last_streamed_chunk, allow_reasoning_only_retry) + + +def _streamed_error_payload(error: StreamedAPIError, classified: Any) -> Dict[str, Any]: + """Return a terminal stream error while preserving structured timeout data.""" + + data = getattr(error, "data", None) + if isinstance(data, dict) and isinstance(data.get("error"), dict): + source = data["error"] + return { + "error": { + "message": source.get("message") or "Upstream stream failed after output began", + "type": source.get("type") or classified.error_type, + "details": source.get("details") or {"status_code": classified.status_code}, + } + } + return {"error": {"message": "Upstream stream failed after output began", "type": classified.error_type, "details": {"status_code": classified.status_code}}} + + +def _merge_stream_error_details(error_data: Dict[str, Any], stream_error: Dict[str, Any]) -> None: + """Preserve structured stream timeout metadata in aggregate stream errors.""" + + target = error_data.get("error") if isinstance(error_data, dict) else None + source = stream_error.get("error") if isinstance(stream_error, dict) else None + if not isinstance(target, dict) or not isinstance(source, dict): + return + details = source.get("details") + if not isinstance(details, dict) or "timeout_type" not in details: + return + merged = dict(target.get("details") or {}) if isinstance(target.get("details"), dict) else {} + merged.update(details) + target["details"] = merged + target["type"] = source.get("type") or target.get("type") diff --git a/src/rotator_library/client/gemini.py b/src/rotator_library/client/gemini.py new file mode 100644 index 000000000..0e8707304 --- /dev/null +++ b/src/rotator_library/client/gemini.py @@ -0,0 +1,105 @@ +"""Gemini client-surface handler backed by the shared protocol runtime.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, Optional + +from ..protocols import OPERATION_COUNT_TOKENS, ProtocolContext, get_protocol +from ..routing import load_routing_config_from_env + +if TYPE_CHECKING: + from .rotating_client import RotatingClient + + +class GeminiHandler: + """Expose Gemini wire requests without leaking that format into providers.""" + + def __init__(self, client: "RotatingClient") -> None: + self._client = client + + async def generate( + self, + payload: dict[str, Any], + *, + model: str, + raw_request: Optional[Any] = None, + ) -> Any: + """Execute one Gemini generateContent request through canonical routing.""" + + request_payload = dict(payload) + if request_payload.get("stream"): + raise ValueError( + "Gemini generateContent does not accept stream=true; use streamGenerateContent" + ) + request_payload["model"] = self._routable_model(model) + return await self._client.agenerate( + request_payload, + input_protocol="gemini", + request=raw_request, + ) + + async def stream_generate( + self, + payload: dict[str, Any], + *, + model: str, + raw_request: Optional[Any] = None, + ) -> Any: + """Execute one Gemini streamGenerateContent request canonically.""" + + request_payload = dict(payload) + request_payload["model"] = self._routable_model(model) + request_payload["stream"] = True + return await self._client.agenerate( + request_payload, + input_protocol="gemini", + request=raw_request, + ) + + def count_tokens(self, payload: dict[str, Any], *, model: str) -> dict[str, int]: + """Count a Gemini request locally using its canonical Chat projection.""" + + request_payload = dict(payload) + request_payload["model"] = self._routable_model(model) + gemini = get_protocol("gemini") + unified = gemini.parse_request( + request_payload, + ProtocolContext( + source_protocol="gemini", + target_protocol="gemini", + model=request_payload["model"], + metadata={"operation": OPERATION_COUNT_TOKENS}, + ), + ) + chat_request = get_protocol("openai_chat").build_request( + unified, + ProtocolContext( + source_protocol="gemini", + target_protocol="openai_chat", + input_protocol="gemini", + output_protocol="gemini", + model=request_payload["model"], + ), + ) + total = self._client.token_count( + model=request_payload["model"], + messages=chat_request.get("messages") or [], + ) + tools = chat_request.get("tools") or [] + if tools: + total += self._client.token_count( + model=request_payload["model"], + text=json.dumps(tools, separators=(",", ":")), + ) + return {"totalTokens": total} + + @staticmethod + def _routable_model(model: str) -> str: + """Keep configured aliases, otherwise default Gemini-style IDs to Gemini.""" + + normalized = str(model or "").removeprefix("models/") + if "/" in normalized: + return normalized + routes = load_routing_config_from_env().model_routes + return normalized if normalized.lower() in routes else f"gemini/{normalized}" diff --git a/src/rotator_library/client/protocol_selection.py b/src/rotator_library/client/protocol_selection.py new file mode 100644 index 000000000..df7aa1a7a --- /dev/null +++ b/src/rotator_library/client/protocol_selection.py @@ -0,0 +1,91 @@ +"""Resolve client output protocol choices without coupling providers to routes.""" + +from __future__ import annotations + +from typing import Any, Optional + +from ..protocols import get_protocol +from ..core.errors import protocol_error_payload + + +OUTPUT_PROTOCOL_HEADER = "X-Proxy-Output-Protocol" + + +def request_output_protocol(request: Optional[Any]) -> Optional[str]: + """Return the canonical output protocol requested at the HTTP boundary.""" + + headers = getattr(request, "headers", None) + if not headers: + return None + value = headers.get(OUTPUT_PROTOCOL_HEADER) or headers.get(OUTPUT_PROTOCOL_HEADER.lower()) + if not value: + return None + return canonical_protocol_name(str(value)) + + +def canonical_protocol_name(value: str) -> str: + """Resolve protocol aliases and reject unknown output formats early.""" + + try: + return get_protocol(str(value).strip().lower()).name + except KeyError as exc: + raise ValueError(f"Unsupported output protocol: {value}") from exc + + +def require_same_protocol_stream(input_protocol: str, output_protocol: str) -> None: + """Validate that both stream protocols use canonical generative events.""" + + source = canonical_protocol_name(input_protocol) + target = canonical_protocol_name(output_protocol) + supported = {"openai_chat", "anthropic_messages", "responses", "gemini"} + if source not in supported or target not in supported: + raise ValueError( + f"Streaming conversion from {source} to {target} is not supported" + ) + + +def resolve_client_output_protocol( + client: Any, + payload: dict[str, Any], + *, + input_protocol: str, + request: Optional[Any] = None, +) -> str: + """Resolve output through RotatingClient or a compatible minimal facade.""" + + resolver = getattr(client, "resolve_output_protocol", None) + if callable(resolver): + return resolver(payload, input_protocol=input_protocol, request=request) + return request_output_protocol(request) or canonical_protocol_name(input_protocol) + + +def format_client_protocol_error( + client: Any, + payload: dict[str, Any], + *, + input_protocol: str, + request: Any, + error: BaseException | str, + error_type: str, + status_code: int, +) -> tuple[int, dict[str, Any]]: + """Format a proxy-side failure using the selected client protocol.""" + + try: + output_protocol = resolve_client_output_protocol( + client, + payload, + input_protocol=input_protocol, + request=request, + ) + except ValueError as selection_error: + output_protocol = canonical_protocol_name(input_protocol) + error = selection_error + error_type = "invalid_request" + status_code = 400 + return protocol_error_payload( + error, + output_protocol, + error_type=error_type, + status_code=status_code, + ) diff --git a/src/rotator_library/client/request_builder.py b/src/rotator_library/client/request_builder.py index c9455ec23..e47f4b1c7 100644 --- a/src/rotator_library/client/request_builder.py +++ b/src/rotator_library/client/request_builder.py @@ -3,12 +3,19 @@ """RequestContext construction for RotatingClient public request methods.""" -import time import inspect +import time +from copy import deepcopy from typing import Any, Awaitable, Callable, Dict, Optional from ..core.types import RequestContext +from ..protocols import ProtocolContext, get_protocol +from ..routing import FallbackResolver, RoutingConfigError, load_routing_config_from_env +from ..routing.types import RouteTarget, RoutingDecision +from ..session_tracking import SessionTrackingHints from ..transaction_logger import TransactionLogger +from .scopes import derive_session_isolation_key +from .protocol_selection import require_same_protocol_stream class RequestContextBuilder: @@ -35,13 +42,35 @@ def __init__( self._get_provider_instance = get_provider_instance @staticmethod - def _pop_scope_kwargs(kwargs: Dict[str, Any]) -> tuple[Optional[str], Any, Any, bool]: + def _pop_scope_kwargs(kwargs: Dict[str, Any]) -> tuple[Optional[str], Any, Any, bool, Any]: classifier = kwargs.pop("classifier", None) request_api_keys = kwargs.pop("api_keys", None) request_providers = kwargs.pop("providers", None) private = bool(kwargs.pop("private", False)) + session_tracking_hints = kwargs.pop("_session_tracking_hints", None) kwargs.pop("model_filters", None) - return classifier, request_api_keys, request_providers, private + return classifier, request_api_keys, request_providers, private, session_tracking_hints + + @staticmethod + def _session_isolation_key( + classifier: Optional[str], + request_api_keys: Any, + request_providers: Any, + private: bool, + ) -> str: + """Return a provider-independent caller/credential isolation key. + + Ad hoc credentials and provider overrides are hashed as one bundle so + secrets never enter logs or persistence. Named classifiers intentionally + own one authoritative domain across their providers and credential changes. + """ + + return derive_session_isolation_key( + classifier, + request_api_keys, + request_providers, + private, + ) @staticmethod def _provider_from_model(model: str) -> str: @@ -51,11 +80,53 @@ def _provider_from_model(model: str) -> str: def _raise_no_provider(model: str) -> None: raise ValueError(f"Invalid model format or no credentials for provider: {model}") + def _resolve_routing_decision(self, model: str) -> Optional[RoutingDecision]: + """Resolve env-configured fallback routing, if any applies.""" + + config = load_routing_config_from_env() + if not config.fallback_groups and not config.model_routes: + return None + try: + decision = FallbackResolver(config).resolve(model) + if decision.reason == "direct_provider_model" and model.lower() not in config.model_routes: + return None + return decision + except RoutingConfigError: + if "/" in model: + return None + raise + + @staticmethod + def _with_request_scope(target: RouteTarget, scope: Dict[str, Any]) -> RouteTarget: + """Attach per-provider request scope to a route target without secrets in traces.""" + + metadata = dict(target.metadata) + metadata["request_scope"] = { + "credentials": list(scope["credentials"]), + "usage_manager_key": scope["usage_manager_key"], + "provider_config": scope["provider_config"], + "credential_secrets": dict(scope["credential_secrets"]), + } + return RouteTarget( + provider=target.provider, + model=target.model, + name=target.name, + protocol=target.protocol, + execution=target.execution, + priority=target.priority, + weight=target.weight, + conditions=dict(target.conditions), + metadata=metadata, + ) + async def _get_session_hints( self, provider: str, model: str, kwargs: Dict[str, Any], + *, + unified_request: Any = None, + input_protocol: str = "openai_chat", ) -> Any: """Ask the provider for optional session evidence before routing. @@ -71,7 +142,31 @@ async def _get_session_hints( if not hook: return None try: - result = hook(kwargs, model=model) + provider_request = kwargs + if unified_request is not None: + provider_protocol_name = plugin.get_protocol_name(model) if hasattr(plugin, "get_protocol_name") else "openai_chat" + provider_protocol = get_protocol(provider_protocol_name) + native_model = plugin.normalize_native_model(model) if hasattr(plugin, "normalize_native_model") else model + provider_view = deepcopy(unified_request) + provider_view.model = native_model + provider_request = provider_protocol.build_request( + provider_view, + ProtocolContext( + input_protocol=input_protocol, + provider_protocol=provider_protocol.name, + output_protocol=input_protocol, + source_protocol=input_protocol, + target_protocol=provider_protocol.name, + source_provider=None, + target_provider=provider, + provider=provider, + model=native_model, + ), + ) + operation = plugin.get_native_operation(native_model, None, stream=bool(provider_view.stream)) if hasattr(plugin, "get_native_operation") else "generate" + if hasattr(plugin, "prepare_native_request"): + provider_request = plugin.prepare_native_request(provider_request, native_model, operation) + result = hook(provider_request, model=model) if inspect.isawaitable(result): result = await result return result @@ -88,17 +183,91 @@ async def _get_session_hints( ) return None + @staticmethod + def _merge_session_hints(*hints: Any) -> Any: + """Merge proxy-internal and provider session evidence. + + Internal hints are removed from request kwargs before provider execution. + They let services such as Responses expose stable continuation IDs to the + centralized tracker without adding provider-visible payload fields. + """ + + merged = SessionTrackingHints() + seen = False + for index, hint in enumerate(hints): + if not hint: + continue + allow_global = index == 0 and isinstance(hint, SessionTrackingHints) + if isinstance(hint, dict): + hint = SessionTrackingHints( + strong_anchors=list(hint.get("strong_anchors") or []), + medium_anchors=list(hint.get("medium_anchors") or []), + weak_anchors=list(hint.get("weak_anchors") or []), + affinity_key=hint.get("affinity_key"), + session_scope=hint.get("session_scope"), + ) + if not isinstance(hint, SessionTrackingHints): + continue + seen = True + merged.strong_anchors.extend(hint.strong_anchors) + merged.medium_anchors.extend(hint.medium_anchors) + merged.weak_anchors.extend(hint.weak_anchors) + if allow_global: + merged.global_strong_anchors.extend(hint.global_strong_anchors) + merged.global_medium_anchors.extend(hint.global_medium_anchors) + merged.global_weak_anchors.extend(hint.global_weak_anchors) + if not merged.affinity_key and hint.affinity_key: + merged.affinity_key = hint.affinity_key + if not merged.session_scope and hint.session_scope: + merged.session_scope = hint.session_scope + return merged if seen else None + async def build_completion_context( self, request: Optional[Any], pre_request_callback: Optional[Callable], kwargs: Dict[str, Any], ) -> RequestContext: - classifier, request_api_keys, request_providers, private = self._pop_scope_kwargs( + classifier, request_api_keys, request_providers, private, internal_session_hints = self._pop_scope_kwargs( kwargs ) + parent_log_dir = kwargs.pop("_parent_log_dir", None) + disable_provider_continuation = bool(kwargs.pop("_disable_provider_continuation", False)) + requested_input_protocol = str(kwargs.pop("_input_protocol", "openai_chat") or "openai_chat") + raw_output_protocol = kwargs.pop("_output_protocol", None) + output_protocol_explicit = raw_output_protocol not in (None, "") + requested_output_protocol = str(raw_output_protocol or requested_input_protocol) + input_protocol = get_protocol(requested_input_protocol) + output_protocol = get_protocol(requested_output_protocol) + protocol_request = deepcopy(kwargs) + protocol_context = ProtocolContext( + source_protocol=input_protocol.name, + target_protocol=input_protocol.name, + input_protocol=input_protocol.name, + output_protocol=output_protocol.name, + ) + unified_request = input_protocol.parse_request(protocol_request, protocol_context) + if input_protocol.name != "openai_chat": + kwargs = get_protocol("openai_chat").build_request( + unified_request, + ProtocolContext( + source_protocol=input_protocol.name, + target_protocol="openai_chat", + input_protocol=input_protocol.name, + provider_protocol="openai_chat", + output_protocol=output_protocol.name, + ), + ) + session_isolation_key = self._session_isolation_key( + classifier, + request_api_keys, + request_providers, + private, + ) model = kwargs.get("model", "") - provider = self._provider_from_model(model) + routing_decision = self._resolve_routing_decision(model) + routing_targets = routing_decision.targets if routing_decision else None + provider = routing_targets[0].provider if routing_targets else self._provider_from_model(model) if not provider: self._raise_no_provider(model) @@ -112,9 +281,34 @@ async def build_completion_context( if not scope["credentials"]: self._raise_no_provider(model) - parent_log_dir = kwargs.pop("_parent_log_dir", None) - resolved_model = self._model_resolver.resolve_model_id(model, provider) + if routing_targets: + scoped_targets = [] + for index, target in enumerate(routing_targets): + target_scope = scope if index == 0 else await self._resolve_scope_for_provider( + target.provider, + classifier, + request_api_keys, + request_providers, + private, + ) + if not target_scope["credentials"]: + self._raise_no_provider(target.prefixed_model) + scoped_targets.append(self._with_request_scope(target, target_scope)) + routing_targets = tuple(scoped_targets) + + resolved_model = self._model_resolver.resolve_model_id(routing_targets[0].prefixed_model if routing_targets else model, provider) kwargs["model"] = resolved_model + if not output_protocol_explicit and self._get_provider_instance: + plugin = self._get_provider_instance(provider) + default_output = ( + plugin.get_default_output_protocol(resolved_model) + if plugin and hasattr(plugin, "get_default_output_protocol") + else None + ) + if default_output: + output_protocol = get_protocol(default_output) + if unified_request.stream: + require_same_protocol_stream(input_protocol.name, output_protocol.name) transaction_logger = None if self._get_enable_request_logging(): @@ -130,9 +324,25 @@ async def build_completion_context( kwargs, provider=provider, model=resolved_model, - scope_key=scope["usage_manager_key"], - hints=await self._get_session_hints(provider, resolved_model, kwargs), + scope_key=session_isolation_key, + hints=self._merge_session_hints( + internal_session_hints, + await self._get_session_hints( + provider, + resolved_model, + kwargs, + unified_request=unified_request, + input_protocol=input_protocol.name, + ), + ), + _trusted_isolation_key=True, ) + if transaction_logger: + transaction_logger.set_trace_context( + session_id=session.session_id, + scope_key=scope["usage_manager_key"], + classifier=scope["classifier"], + ) return RequestContext( model=resolved_model, @@ -144,9 +354,8 @@ async def build_completion_context( session_id=session.session_id, session_affinity_key=session.affinity_key, session_tracker=self._session_tracker, - session_possible_compaction=session.possible_compaction, - session_lineage_parent_id=session.lineage_parent_session_id, session_tracking_namespace=session.tracking_namespace, + session_isolation_key=session_isolation_key, request=request, pre_request_callback=pre_request_callback, transaction_logger=transaction_logger, @@ -154,6 +363,15 @@ async def build_completion_context( provider_config=scope["provider_config"], credential_secrets=scope["credential_secrets"], classifier=scope["classifier"], + routing_targets=routing_targets, + routing_group_name=routing_decision.group_name if routing_decision else None, + input_protocol_name=input_protocol.name, + output_protocol_name=output_protocol.name, + protocol_request=protocol_request, + unified_request=unified_request, + input_provider=provider, + disable_provider_continuation=disable_provider_continuation, + routing_group=routing_decision.group if routing_decision else None, ) async def build_embedding_context( @@ -162,9 +380,15 @@ async def build_embedding_context( pre_request_callback: Optional[Callable], kwargs: Dict[str, Any], ) -> RequestContext: - classifier, request_api_keys, request_providers, private = self._pop_scope_kwargs( + classifier, request_api_keys, request_providers, private, internal_session_hints = self._pop_scope_kwargs( kwargs ) + session_isolation_key = self._session_isolation_key( + classifier, + request_api_keys, + request_providers, + private, + ) model = kwargs.get("model", "") provider = self._provider_from_model(model) if not provider: @@ -184,8 +408,12 @@ async def build_embedding_context( kwargs, provider=provider, model=model, - scope_key=scope["usage_manager_key"], - hints=await self._get_session_hints(provider, model, kwargs), + scope_key=session_isolation_key, + hints=self._merge_session_hints( + internal_session_hints, + await self._get_session_hints(provider, model, kwargs), + ), + _trusted_isolation_key=True, ) return RequestContext( @@ -198,9 +426,8 @@ async def build_embedding_context( session_id=session.session_id, session_affinity_key=session.affinity_key, session_tracker=self._session_tracker, - session_possible_compaction=session.possible_compaction, - session_lineage_parent_id=session.lineage_parent_session_id, session_tracking_namespace=session.tracking_namespace, + session_isolation_key=session_isolation_key, request=request, pre_request_callback=pre_request_callback, usage_manager_key=scope["usage_manager_key"], diff --git a/src/rotator_library/client/rotating_client.py b/src/rotator_library/client/rotating_client.py index 7d496ecc9..b5c929014 100644 --- a/src/rotator_library/client/rotating_client.py +++ b/src/rotator_library/client/rotating_client.py @@ -36,11 +36,14 @@ from .transforms import ProviderTransforms from .executor import RequestExecutor from .anthropic import AnthropicHandler -from .scopes import ScopeManager +from .gemini import GeminiHandler +from .scopes import NO_AUTH_CREDENTIAL, ScopeManager from .model_discovery import ModelDiscoveryService from .usage_managers import UsageManagerRegistry from .request_builder import RequestContextBuilder from .quota import QuotaService +from .protocol_selection import canonical_protocol_name, request_output_protocol, require_same_protocol_stream +from ..routing import FallbackResolver, load_routing_config_from_env from ..session_tracking import SessionTracker # Import providers and other dependencies @@ -63,6 +66,50 @@ lib_logger = logging.getLogger("rotator_library") +def _resolve_session_persistence_settings( + enabled: Optional[bool], + flush_interval_seconds: Optional[float], +) -> tuple[bool, float]: + """Resolve explicit session-persistence settings before environment defaults.""" + + if enabled is None: + enabled = os.getenv( + "SESSION_PERSISTENCE_ENABLED", + "false", + ).strip().lower() in {"1", "true", "yes", "on"} + if flush_interval_seconds is None: + raw_flush_interval = os.getenv( + "SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS", + "5.0", + ) + try: + flush_interval_seconds = max(0.0, float(raw_flush_interval)) + except (TypeError, ValueError): + flush_interval_seconds = 5.0 + lib_logger.warning( + "Invalid SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS=%r; using 5.0", + raw_flush_interval, + ) + return enabled, float(flush_interval_seconds) + + +def _add_configured_no_auth_credentials( + credentials: Dict[str, List[str]], + provider_names: set[str], + *, + config: Any = None, +) -> None: + """Give configured no-auth providers one internal rotation/accounting slot.""" + + from ..config.experimental import get_provider_runtime_config, load_experimental_config + + active = config or load_experimental_config() + for provider in provider_names: + runtime = get_provider_runtime_config(provider, config=active) + if runtime.api_base and runtime.auth_mode == "none" and not credentials.get(provider): + credentials[provider] = [NO_AUTH_CREDENTIAL] + + class RotatingClient: """ A client that intelligently rotates and retries API keys using LiteLLM, @@ -99,8 +146,8 @@ def __init__( rotation_tolerance: float = DEFAULT_ROTATION_TOLERANCE, data_dir: Optional[Union[str, Path]] = None, session_stickiness_ttl_seconds: int = 3600, - session_persistence_enabled: bool = False, - session_persistence_flush_interval_seconds: float = 5.0, + session_persistence_enabled: Optional[bool] = None, + session_persistence_flush_interval_seconds: Optional[float] = None, ): """ Initialize the RotatingClient. @@ -109,6 +156,16 @@ def __init__( """ # Resolve data directory self.data_dir = Path(data_dir).resolve() if data_dir else get_default_root() + from ..config.experimental import load_experimental_config + + self._experimental_config = load_experimental_config() + ( + session_persistence_enabled, + session_persistence_flush_interval_seconds, + ) = _resolve_session_persistence_settings( + session_persistence_enabled, + session_persistence_flush_interval_seconds, + ) # Configure logging configure_failure_logger(get_logs_dir(self.data_dir)) @@ -131,11 +188,6 @@ def __init__( api_keys = {p: k for p, k in api_keys.items() if k} oauth_credentials = {p: c for p, c in oauth_credentials.items() if c} - if not api_keys and not oauth_credentials: - lib_logger.warning( - "No provider credentials configured. Client will be unable to make requests." - ) - # Discover OAuth credentials if not provided if oauth_credentials: self.oauth_credentials = oauth_credentials @@ -151,6 +203,15 @@ def __init__( self.all_credentials.setdefault(provider, []).extend(keys) for provider, paths in self.oauth_credentials.items(): self.all_credentials.setdefault(provider, []).extend(paths) + _add_configured_no_auth_credentials( + self.all_credentials, + set(PROVIDER_PLUGINS), + config=self._experimental_config, + ) + if not self.all_credentials: + lib_logger.warning( + "No provider credentials configured. Client will be unable to make requests." + ) self.api_keys = api_keys self.oauth_providers = set(self.oauth_credentials.keys()) @@ -276,6 +337,7 @@ def __init__( litellm_provider_params=self.litellm_provider_params, litellm_logger_fn=self._litellm_logger_fn, provider_instances=self._provider_instances, + experimental_config=self._experimental_config, ) self._model_list_cache: Dict[str, List[str]] = {} @@ -321,6 +383,7 @@ def __init__( # Initialize Anthropic compatibility handler self._anthropic_handler = AnthropicHandler(self) + self._gemini_handler = GeminiHandler(self) @staticmethod def _normalize_mode_concurrency(values: Dict[str, Dict[str, int]]) -> None: @@ -529,15 +592,82 @@ async def list_scope_credentials( classifier, provider=provider, include_secrets=include_secrets ) + async def agenerate( + self, + payload: Dict[str, Any], + *, + input_protocol: str, + output_protocol: Optional[str] = None, + request: Optional[Any] = None, + pre_request_callback: Optional[callable] = None, + **routing_kwargs: Any, + ) -> Union[Any, AsyncGenerator[str, None]]: + """Execute a generative request with independent wire protocols. + + Provider selection is still driven by the request model and existing + routing controls. Providers receive only their declared native format; + the selected output defaults to the client's input format. + """ + + selected_output = self.resolve_output_protocol( + payload, + input_protocol=input_protocol, + request=request, + explicit=output_protocol, + ) + if payload.get("stream"): + require_same_protocol_stream(input_protocol, selected_output) + kwargs = dict(payload) + kwargs.update(routing_kwargs) + kwargs["_input_protocol"] = input_protocol + kwargs["_output_protocol"] = selected_output + return await self.acompletion( + request=request, + pre_request_callback=pre_request_callback, + **kwargs, + ) + + def resolve_output_protocol( + self, + payload: Dict[str, Any], + *, + input_protocol: str, + request: Optional[Any] = None, + explicit: Optional[str] = None, + ) -> str: + """Resolve explicit, HTTP, provider-default, then input output format.""" + + selected = explicit or request_output_protocol(request) + if selected: + return canonical_protocol_name(selected) + model = str(payload.get("model") or "") + provider = model.split("/", 1)[0] if "/" in model else "" + if not provider and model: + routing = load_routing_config_from_env() + if model.lower() in routing.model_routes: + decision = FallbackResolver(routing).resolve(model) + if decision.targets: + provider = decision.targets[0].provider + plugin = self._get_provider_instance(provider) if provider else None + configured = ( + plugin.get_default_output_protocol(model) + if plugin and hasattr(plugin, "get_default_output_protocol") + else None + ) + return canonical_protocol_name(configured or input_protocol) + async def acompletion( self, request: Optional[Any] = None, pre_request_callback: Optional[callable] = None, **kwargs, ) -> Union[Any, AsyncGenerator[str, None]]: + request_context_callback = kwargs.pop("_request_context_callback", None) context = await self._request_builder.build_completion_context( request, pre_request_callback, kwargs ) + if callable(request_context_callback): + request_context_callback(context) return await self._executor.execute(context) async def aembedding( @@ -570,6 +700,37 @@ def token_count(self, **kwargs) -> int: return base_count + async def gemini_generate( + self, + payload: Dict[str, Any], + *, + model: str, + raw_request: Optional[Any] = None, + ) -> Any: + """Execute a native Gemini client request through shared routing.""" + + return await self._gemini_handler.generate(payload, model=model, raw_request=raw_request) + + async def gemini_stream_generate( + self, + payload: Dict[str, Any], + *, + model: str, + raw_request: Optional[Any] = None, + ) -> Any: + """Execute a Gemini streamGenerateContent request through shared routing.""" + + return await self._gemini_handler.stream_generate( + payload, + model=model, + raw_request=raw_request, + ) + + def gemini_count_tokens(self, payload: Dict[str, Any], *, model: str) -> Dict[str, int]: + """Return Gemini-compatible local token usage for a request.""" + + return self._gemini_handler.count_tokens(payload, model=model) + def _model_cache_key( self, provider: str, @@ -674,7 +835,10 @@ def _get_provider_instance(self, provider: str) -> Optional[Any]: if provider not in self._provider_instances: plugin_class = self._provider_plugins.get(provider) if plugin_class: - self._provider_instances[provider] = plugin_class() + instance = plugin_class() + if hasattr(instance, "bind_runtime_config"): + instance.bind_runtime_config(self._experimental_config) + self._provider_instances[provider] = instance else: return None diff --git a/src/rotator_library/client/scopes.py b/src/rotator_library/client/scopes.py index 5eeb4ae54..888f90003 100644 --- a/src/rotator_library/client/scopes.py +++ b/src/rotator_library/client/scopes.py @@ -6,11 +6,56 @@ import asyncio import hashlib import hmac +import json import re from pathlib import Path from typing import Any, Awaitable, Callable, Dict, List, Optional +NO_AUTH_CREDENTIAL = "__proxy_no_auth__" + + +def derive_session_isolation_key( + classifier: Optional[str], + request_api_keys: Any, + request_providers: Any, + private: bool, +) -> str: + """Return a stable provider-independent caller/credential domain key.""" + + # A named classifier is the caller's explicit isolation boundary. Rotating + # credentials or provider overrides inside it must not create hidden domains. + if classifier: + digest = hashlib.sha256(str(classifier).encode("utf-8")).hexdigest()[:24] + return f"classifier:{digest}" + if not request_api_keys and not request_providers: + return "public" + + normalized_api_keys = ScopeManager.normalize_api_key_map(request_api_keys) + normalized_providers = ScopeManager.normalize_provider_map(request_providers) + + def canonical(value: Any) -> Any: + if isinstance(value, dict): + return { + str(key): canonical(item) + for key, item in sorted(value.items(), key=lambda pair: str(pair[0])) + } + if isinstance(value, (list, tuple, set)): + items = [canonical(item) for item in value] + return sorted(items, key=lambda item: json.dumps(item, sort_keys=True, default=str)) + if value is None or isinstance(value, (str, int, float, bool)): + return value + return str(value) + + payload = { + "api_keys": canonical(normalized_api_keys), + "providers": canonical(normalized_providers), + "private": bool(private), + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return "bundle:" + hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + class ScopeManager: """Manage registered classifier state and resolve per-request scopes.""" @@ -149,8 +194,26 @@ async def resolve_scope_for_provider( credential_entry = registered["credentials"][provider] raw_credentials = list(credential_entry.get("keys", [])) credential_private = bool(credential_entry.get("private", True)) + elif self._all_credentials.get(provider) == [NO_AUTH_CREDENTIAL]: + # A no-auth provider still needs one selector/accounting slot in each + # isolated usage pool, but there is no secret to fingerprint or leak. + raw_credentials = [NO_AUTH_CREDENTIAL] + credential_private = False - scope_name = classifier or "default" + session_domain = derive_session_isolation_key( + classifier, + request_api_keys, + request_providers, + private, + ) + # Ad hoc bundles need their own usage manager as well as their own + # session namespace. Otherwise cooldowns, sticky entries, and credential + # registries from unrelated private bundles would share "default". + scope_name = ( + session_domain + if session_domain.startswith("bundle:") + else (classifier or "default") + ) credentials: List[str] = [] credential_secrets: Dict[str, str] = {} if credential_private: diff --git a/src/rotator_library/client/stream_retry_policy.py b/src/rotator_library/client/stream_retry_policy.py index ccf0a593c..c776d5cdd 100644 --- a/src/rotator_library/client/stream_retry_policy.py +++ b/src/rotator_library/client/stream_retry_policy.py @@ -1,78 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-only # Copyright (c) 2026 Mirrowel -"""Retry-safety policy for stream errors after output has started.""" +"""Backward-compatible import path for stream retry policy.""" -import json -from typing import Any, Optional +from ..streaming.policy import can_retry_stream_after_error - -_REASONING_FIELDS = ( - "reasoning", - "reasoning_content", - "thinking", - "thinking_content", -) - - -def can_retry_stream_after_error( - last_streamed_chunk: Optional[str], - allow_reasoning_only_retry: bool, -) -> bool: - """ - Return whether a stream can be retried after an upstream error. - - Retrying is always safe before any chunk has been emitted. After that, it is - only allowed when explicitly enabled and the latest chunk is clearly - reasoning/thinking-only. Ambiguous chunks fail closed. - """ - if last_streamed_chunk is None: - return True - if not allow_reasoning_only_retry: - return False - - payload = last_streamed_chunk.strip() - if not payload.startswith("data:"): - return False - payload = payload[5:].strip() - if not payload or payload == "[DONE]": - return False - try: - data = json.loads(payload) - except json.JSONDecodeError: - return False - - has_reasoning = False - choices = data.get("choices") - if not isinstance(choices, list): - return False - - for choice in choices: - if not isinstance(choice, dict): - return False - for source in (choice, choice.get("delta"), choice.get("message")): - if not isinstance(source, dict): - continue - if ( - _has_text(source.get("content")) - or _has_text(source.get("text")) - or source.get("tool_calls") - or source.get("function_call") - ): - return False - if any(_has_text(source.get(key)) for key in _REASONING_FIELDS): - has_reasoning = True - - return has_reasoning - - -def _has_text(value: Any) -> bool: - if isinstance(value, str): - return bool(value.strip()) - if isinstance(value, list): - for item in value: - if isinstance(item, str) and item.strip(): - return True - if isinstance(item, dict) and _has_text(item.get("text")): - return True - return False +__all__ = ["can_retry_stream_after_error"] diff --git a/src/rotator_library/client/streaming.py b/src/rotator_library/client/streaming.py index 01900b50d..aeb9943f8 100644 --- a/src/rotator_library/client/streaming.py +++ b/src/rotator_library/client/streaming.py @@ -13,23 +13,37 @@ - Client disconnect handling """ +import asyncio import codecs +import contextlib import json import logging import re +import time +from copy import deepcopy +from dataclasses import replace from typing import Any, AsyncGenerator, AsyncIterator, Callable, Dict, List, Optional, TYPE_CHECKING -import litellm - from ..core.errors import StreamedAPIError, CredentialNeedsReauthError from ..core.types import ProcessedChunk from ..core.utils import normalize_usage_for_response +from ..streaming import StreamEvent, StreamMonitor, stream_event_from_sse_chunk +from ..streaming.transport import SSEStreamFormatter +from ..usage.accounting import UsageRecord, extract_usage_record +from ..usage.costs import CostBreakdown, CostCalculator if TYPE_CHECKING: from ..usage.manager import CredentialContext lib_logger = logging.getLogger("rotator_library") +_MIXED_REASONING_CONTENT_MODELS = frozenset( + { + "google/diffusiongemma-26b-a4b-it", + "nvidia_nim/google/diffusiongemma-26b-a4b-it", + } +) + class StreamingHandler: """ @@ -50,6 +64,8 @@ async def wrap_stream( cred_context: Optional["CredentialContext"] = None, skip_cost_calculation: bool = False, response_callback: Optional[Callable[[Dict[str, Any]], None]] = None, + success_callback: Optional[Callable[[], None]] = None, + transaction_logger: Optional[Any] = None, ) -> AsyncGenerator[str, None]: """ Wrap a LiteLLM stream with error handling and usage tracking. @@ -70,6 +86,7 @@ async def wrap_stream( SSE-formatted strings: "data: {...}\\n\\n" """ stream_completed = False + response_identity_complete = False error_buffer = StreamBuffer() # Use StreamBuffer for JSON reassembly accumulated_finish_reason: Optional[str] = None has_tool_calls = False @@ -79,12 +96,61 @@ async def wrap_stream( prompt_tokens_uncached = 0 completion_tokens = 0 thinking_tokens = 0 + usage_record = UsageRecord(source="stream") assistant_parts: List[str] = [] tool_call_ids: List[str] = [] + tool_call_events: Dict[tuple[int, int], Dict[str, str]] = {} + monitor = StreamMonitor(clock=time.monotonic) + from ..config.experimental import get_stream_runtime_settings + + stream_settings = get_stream_runtime_settings() + formatter = SSEStreamFormatter() + split_mixed_reasoning_content = ( + model.lower() in _MIXED_REASONING_CONTENT_MODELS + ) + upstream_closed = False + stream_cancelled = False + last_heartbeat_at = monitor.metrics.started_at + lifecycle_logger = transaction_logger + self._log_stream_lifecycle( + lifecycle_logger, + "stream_started", + monitor, + StreamEvent("started", protocol="openai_chat"), + ) # Use manual iteration to allow continue after partial JSON errors stream_iterator = stream.__aiter__() + async def close_upstream(reason: str, *, force: bool = False) -> None: + """Best-effort close for upstream async streams. + + Client disconnects and timeout failures should not leave provider + HTTP streams running in the background. Close failures are logged as + lifecycle metadata only and never replace the original stream error. + """ + + nonlocal upstream_closed + if upstream_closed or (not force and not stream_settings.cancel_upstream_on_disconnect): + return + upstream_closed = True + for candidate in (stream_iterator, stream): + try: + closer = getattr(candidate, "aclose", None) + if closer: + await closer() + self._log_stream_lifecycle(lifecycle_logger, "stream_upstream_cancelled", monitor, StreamEvent("cancelled", protocol="openai_chat", data={"reason": reason})) + return + closer = getattr(candidate, "close", None) + if closer: + closer() + self._log_stream_lifecycle(lifecycle_logger, "stream_upstream_cancelled", monitor, StreamEvent("cancelled", protocol="openai_chat", data={"reason": reason})) + return + except Exception as exc: + lib_logger.debug("Failed to close upstream stream: %s", exc) + self._log_stream_lifecycle(lifecycle_logger, "stream_upstream_close_failed", monitor, StreamEvent("error", protocol="openai_chat", data={"reason": reason, "error_type": type(exc).__name__})) + return + try: while True: try: @@ -95,23 +161,96 @@ async def wrap_stream( ) break - chunk = await stream_iterator.__anext__() + next_task = asyncio.create_task(stream_iterator.__anext__()) + try: + while True: + wait_seconds = _next_stream_wait_seconds(monitor, stream_settings, last_heartbeat_at) + wait_tasks = {next_task} + disconnect_task = None + if request is not None: + disconnect_task = asyncio.create_task(request.is_disconnected()) + wait_tasks.add(disconnect_task) + done, _ = await asyncio.wait(wait_tasks, timeout=wait_seconds) + if disconnect_task is not None: + if disconnect_task in done and disconnect_task.result(): + stream_cancelled = True + next_task.cancel() + with contextlib.suppress(asyncio.CancelledError, StopAsyncIteration): + await next_task + await close_upstream("client_disconnect") + return + if not disconnect_task.done(): + disconnect_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await disconnect_task + if next_task in done: + chunk = next_task.result() + break + + timeout_error = _stream_timeout_error(monitor, stream_settings) + if timeout_error: + next_task.cancel() + with contextlib.suppress(asyncio.CancelledError, StopAsyncIteration): + await next_task + await close_upstream(timeout_error[0], force=True) + self._log_stream_lifecycle(lifecycle_logger, timeout_error[2], monitor, StreamEvent("error", protocol="openai_chat", data={"error": timeout_error[1]})) + raise StreamedAPIError(timeout_error[1]["message"], data={"error": timeout_error[1]}) + + if _heartbeat_due(monitor, stream_settings, last_heartbeat_at): + heartbeat = formatter.format_heartbeat() + last_heartbeat_at = time.monotonic() + self._log_stream_lifecycle(lifecycle_logger, "stream_heartbeat", monitor, StreamEvent("heartbeat", protocol="openai_chat", visible_output=False)) + yield heartbeat + except Exception: + if not next_task.done(): + next_task.cancel() + with contextlib.suppress(asyncio.CancelledError, StopAsyncIteration): + await next_task + raise + + raw_event = StreamEvent( + "raw_chunk", + protocol="openai_chat", + raw=chunk, + data={"chunk_index": monitor.metrics.chunk_count + 1}, + ) + if monitor.metrics.first_byte_at is None: + self._log_stream_lifecycle( + lifecycle_logger, + "stream_first_byte", + monitor, + raw_event, + ) # Clear error buffer on successful chunk receipt error_buffer.reset() # Process chunk + cost_event_record = _usage_record_from_sse_cost_chunk(chunk, model=model) processed = self._process_chunk( chunk, accumulated_finish_reason, has_tool_calls, model, ) - self._collect_session_response_anchors( - processed.sse_string, - assistant_parts, - tool_call_ids, + if cost_event_record.provider_reported_cost is not None: + usage_record = _merge_usage_cost(usage_record, cost_event_record) + if not processed.sse_string: + stream_completed = True + response_identity_complete = True + break + output_frames = ( + self._split_mixed_reasoning_content_frames( + processed.sse_string + ) + if split_mixed_reasoning_content + else [processed.sse_string] ) + if processed.finish_reason or any( + self._sse_has_completion_signal(frame) + for frame in output_frames + ): + response_identity_complete = True # Update tracking state if processed.has_tool_calls: @@ -121,54 +260,42 @@ async def wrap_stream( # Only update if not already tool_calls (highest priority) accumulated_finish_reason = processed.finish_reason if processed.usage and isinstance(processed.usage, dict): - # Extract token counts from final chunk - prompt_tokens = processed.usage.get("prompt_tokens", 0) - completion_tokens = processed.usage.get("completion_tokens", 0) - prompt_details = processed.usage.get("prompt_tokens_details") - if prompt_details: - if isinstance(prompt_details, dict): - prompt_tokens_cached = ( - prompt_details.get("cached_tokens", 0) or 0 - ) - prompt_tokens_cache_write = ( - prompt_details.get("cache_creation_tokens", 0) or 0 - ) - else: - prompt_tokens_cached = ( - getattr(prompt_details, "cached_tokens", 0) or 0 - ) - prompt_tokens_cache_write = ( - getattr(prompt_details, "cache_creation_tokens", 0) - or 0 - ) - completion_details = processed.usage.get( - "completion_tokens_details" + next_usage_record = extract_usage_record( + processed.usage, + model=model, + source="stream_final_chunk", ) - if completion_details: - if isinstance(completion_details, dict): - thinking_tokens = ( - completion_details.get("reasoning_tokens", 0) or 0 - ) - else: - thinking_tokens = ( - getattr(completion_details, "reasoning_tokens", 0) - or 0 - ) - if processed.usage.get("cache_read_tokens") is not None: - prompt_tokens_cached = ( - processed.usage.get("cache_read_tokens") or 0 - ) - if processed.usage.get("cache_creation_tokens") is not None: - prompt_tokens_cache_write = ( - processed.usage.get("cache_creation_tokens") or 0 - ) - if thinking_tokens and completion_tokens >= thinking_tokens: - completion_tokens = completion_tokens - thinking_tokens - prompt_tokens_uncached = max( - 0, prompt_tokens - prompt_tokens_cached + if next_usage_record.provider_reported_cost is None and usage_record.provider_reported_cost is not None: + next_usage_record = _merge_usage_cost(next_usage_record, usage_record) + usage_record = next_usage_record + prompt_tokens = usage_record.input_tokens + usage_record.cache_read_tokens + completion_tokens = usage_record.completion_tokens + thinking_tokens = usage_record.reasoning_tokens + prompt_tokens_cached = usage_record.cache_read_tokens + prompt_tokens_cache_write = usage_record.cache_write_tokens + prompt_tokens_uncached = usage_record.prompt_tokens_for_mark_success + + for frame in output_frames: + self._collect_session_response_anchors( + frame, + assistant_parts, + tool_call_ids, + tool_call_events, ) - - yield processed.sse_string + event = stream_event_from_sse_chunk(frame) + first_visible = ( + event.visible_output + and monitor.metrics.first_visible_output_at is None + ) + monitor.record_event(event) + if first_visible: + self._log_stream_lifecycle( + lifecycle_logger, + "stream_first_visible_output", + monitor, + event, + ) + yield frame except StopAsyncIteration: # Stream ended normally @@ -221,37 +348,60 @@ async def wrap_stream( ) # Not a JSON-related error, re-raise + monitor.metrics.error_count += 1 + await close_upstream("stream_exception", force=True) raise except StreamedAPIError: # Re-raise for retry loop + await close_upstream("streamed_api_error", force=True) + raise + + except asyncio.CancelledError: + stream_cancelled = True + monitor.cancel() + self._log_stream_lifecycle( + lifecycle_logger, + "stream_cancelled", + monitor, + StreamEvent("cancelled", protocol="openai_chat", data={"reason": "task_cancelled"}), + ) + await close_upstream("task_cancelled", force=True) raise finally: # Record usage if stream completed if stream_completed: if cred_context: - approx_cost = 0.0 - if not skip_cost_calculation: - approx_cost = self._calculate_stream_cost( - model, - prompt_tokens_uncached + prompt_tokens_cached, - completion_tokens + thinking_tokens, - ) + cost_breakdown = self._calculate_stream_cost_breakdown( + model, + usage_record, + skip_cost_calculation=skip_cost_calculation, + ) + self._log_stream_usage_accounting( + transaction_logger, + usage_record, + cost_breakdown, + ) cred_context.mark_success( prompt_tokens=prompt_tokens_uncached, completion_tokens=completion_tokens, thinking_tokens=thinking_tokens, prompt_tokens_cache_read=prompt_tokens_cached, prompt_tokens_cache_write=prompt_tokens_cache_write, - approx_cost=approx_cost, + approx_cost=cost_breakdown.total_cost, ) - - if response_callback and (assistant_parts or tool_call_ids): - # Intentionally only record response anchors after a complete - # stream. Partial/aborted streams can contain text the client - # never accepted, so using them for identity would over-bind - # failed or disconnected sessions. + if success_callback: + success_callback() + + if ( + response_callback + and response_identity_complete + and (assistant_parts or tool_call_ids) + ): + # A clean iterator EOF is not enough identity evidence: some + # transports surface upstream truncation as normal EOF. Require + # a provider completion signal before recording response anchors. response_callback( { "choices": [ @@ -259,23 +409,77 @@ async def wrap_stream( "message": { "role": "assistant", "content": "".join(assistant_parts), - "tool_calls": [ - {"id": call_id} for call_id in tool_call_ids - ], + "tool_calls": self._assembled_tool_calls( + tool_call_ids, + tool_call_events, + ), } } ] } ) + monitor.complete() + self._log_stream_lifecycle( + lifecycle_logger, + "stream_completed", + monitor, + StreamEvent("completed", protocol="openai_chat"), + ) + self._log_stream_lifecycle( + lifecycle_logger, + "stream_metrics_final", + monitor, + StreamEvent("metadata", protocol="openai_chat"), + ) + # Yield [DONE] for completed streams yield "data: [DONE]\n\n" + elif request and await request.is_disconnected(): + stream_cancelled = True + monitor.cancel() + self._log_stream_lifecycle( + lifecycle_logger, + "stream_cancelled", + monitor, + StreamEvent("cancelled", protocol="openai_chat"), + ) + await close_upstream("client_disconnect") + elif stream_cancelled: + await close_upstream("stream_cancelled", force=True) + + @staticmethod + def _log_stream_lifecycle( + transaction_logger: Optional[Any], + pass_name: str, + monitor: StreamMonitor, + event: StreamEvent, + ) -> None: + """Emit stream lifecycle metrics without affecting stream delivery.""" + + if not transaction_logger: + return + try: + transaction_logger.log_transform_pass( + pass_name, + {"event": event.to_dict(), "metrics": monitor.metrics.to_dict()}, + direction="stream", + stage="client", + protocol=event.protocol, + transport="sse", + metadata={"event_type": event.event_type}, + snapshot=False, + ) + except Exception as exc: + lib_logger.debug("Stream lifecycle trace failed: %s", exc) + def _collect_session_response_anchors( self, sse_string: str, assistant_parts: List[str], tool_call_ids: List[str], + tool_call_events: Optional[Dict[tuple[int, int], Dict[str, str]]] = None, ) -> None: """Collect lightweight response evidence for session tracking. @@ -283,24 +487,213 @@ def _collect_session_response_anchors( chunks. We keep a synthetic assistant message so the core tracker can use the same response-anchor path as non-streaming responses. """ - if not sse_string.startswith("data: "): - return - payload = sse_string[6:].strip() - if not payload or payload == "[DONE]": - return + for payload in self._sse_data_payloads(sse_string): + if not payload or payload == "[DONE]": + continue + try: + data = json.loads(payload) + except json.JSONDecodeError: + continue + if not isinstance(data, dict): + continue + choices = data.get("choices") or [] + if not isinstance(choices, list): + continue + for choice_position, choice in enumerate(choices): + if not isinstance(choice, dict): + continue + choice_index = choice.get("index", choice_position) + try: + stable_choice_index = int(choice_index) + except (TypeError, ValueError): + stable_choice_index = choice_position + delta = choice.get("delta") or {} + if not isinstance(delta, dict): + continue + content = delta.get("content") + if isinstance(content, str) and content: + assistant_parts.append(content) + tool_calls = delta.get("tool_calls") or [] + if not isinstance(tool_calls, list): + continue + for position, tool_call in enumerate(tool_calls): + call_id = tool_call.get("id") if isinstance(tool_call, dict) else None + if call_id and str(call_id) not in tool_call_ids: + tool_call_ids.append(str(call_id)) + if tool_call_events is None or not isinstance(tool_call, dict): + continue + index = tool_call.get("index", position) + try: + event_index = int(index) + except (TypeError, ValueError): + event_index = position + event = tool_call_events.setdefault( + (stable_choice_index, event_index), + {}, + ) + if call_id: + event["id"] = str(call_id) + function = tool_call.get("function") + if not isinstance(function, dict): + continue + if function.get("name"): + event["name"] = str(function["name"]) + arguments = function.get("arguments") + if isinstance(arguments, str) and arguments: + event["arguments"] = self._merge_streamed_tool_arguments( + event.get("arguments", ""), + arguments, + ) + + @staticmethod + def _assembled_tool_calls( + tool_call_ids: List[str], + tool_call_events: Dict[tuple[int, int], Dict[str, str]], + ) -> List[Dict[str, Any]]: + """Return reconstructed tool calls for response-derived session evidence.""" + + calls: List[Dict[str, Any]] = [] + emitted_ids: set[str] = set() + for index in sorted(tool_call_events): + event = tool_call_events[index] + call_id = event.get("id") + if not call_id: + continue + call: Dict[str, Any] = {"id": call_id} + if event.get("name"): + call["function"] = { + "name": event["name"], + "arguments": event.get("arguments", ""), + } + calls.append(call) + emitted_ids.add(call_id) + calls.extend( + {"id": call_id} + for call_id in tool_call_ids + if call_id not in emitted_ids + ) + return calls + + @staticmethod + def _merge_streamed_tool_arguments(current: str, incoming: str) -> str: + """Accept both incremental fragments and cumulative argument snapshots.""" + + if not current: + return incoming + if incoming.startswith(current): + return incoming + if current.startswith(incoming): + return current + return current + incoming + + def _sse_has_completion_signal(self, sse_string: str) -> bool: + """Return whether an SSE frame explicitly marks response completion.""" + + for payload in self._sse_data_payloads(sse_string): + if payload == "[DONE]": + return True + try: + data = json.loads(payload) + except json.JSONDecodeError: + continue + if not isinstance(data, dict): + continue + choices = data.get("choices") or [] + if not isinstance(choices, list): + continue + has_finish_reason = any( + isinstance(choice, dict) and choice.get("finish_reason") is not None + for choice in choices + ) + usage = data.get("usage") + has_meaningful_usage = isinstance(usage, dict) and any( + isinstance(value, (int, float)) and value > 0 + for value in usage.values() + ) + if isinstance(usage, dict) and (has_finish_reason or has_meaningful_usage): + return True + return False + + @staticmethod + def _sse_data_payloads(sse_string: str) -> List[str]: + """Extract SSE data payloads from plain or event-prefixed frames.""" + + payloads: List[str] = [] + for frame in re.split(r"\r?\n\r?\n", sse_string): + data_lines = [] + for line in frame.splitlines(): + if line.startswith("data:"): + data_lines.append(line[5:].lstrip()) + if data_lines: + payloads.append("\n".join(data_lines).strip()) + return payloads + + def _split_mixed_reasoning_content_frames( + self, + sse_string: str, + ) -> List[str]: + """Separate reasoning and final text that share one OpenAI delta. + + DiffusionGemma can end reasoning and begin answer text in the same + chunk. Several OpenAI-compatible clients treat a delta as one state or + the other and discard its content when reasoning is also present. + Fast checks keep normal stream chunks out of the JSON parsing path. + """ + if ( + '"reasoning_content"' not in sse_string + or '"content"' not in sse_string + ): + return [sse_string] + if '"content": null' in sse_string or '"content":null' in sse_string: + return [sse_string] + + payloads = self._sse_data_payloads(sse_string) + if len(payloads) != 1: + return [sse_string] + try: - data = json.loads(payload) + payload = json.loads(payloads[0]) except json.JSONDecodeError: - return - for choice in data.get("choices") or []: - delta = choice.get("delta") or {} - content = delta.get("content") - if content: - assistant_parts.append(str(content)) - for tool_call in delta.get("tool_calls") or []: - call_id = tool_call.get("id") if isinstance(tool_call, dict) else None - if call_id: - tool_call_ids.append(str(call_id)) + return [sse_string] + if not isinstance(payload, dict): + return [sse_string] + + choices = payload.get("choices") + if not isinstance(choices, list) or len(choices) != 1: + return [sse_string] + choice = choices[0] + if not isinstance(choice, dict): + return [sse_string] + delta = choice.get("delta") + if not isinstance(delta, dict): + return [sse_string] + + content = delta.get("content") + reasoning_fields = [ + field + for field in ("reasoning_content", "reasoning") + if isinstance(delta.get(field), str) and delta[field] + ] + if not isinstance(content, str) or not content or not reasoning_fields: + return [sse_string] + + reasoning_payload = deepcopy(payload) + reasoning_choice = reasoning_payload["choices"][0] + reasoning_delta = reasoning_choice["delta"] + for field in ("content", "tool_calls", "function_call", "audio", "refusal"): + reasoning_delta.pop(field, None) + reasoning_choice["finish_reason"] = None + reasoning_payload.pop("usage", None) + + content_payload = deepcopy(payload) + content_delta = content_payload["choices"][0]["delta"] + for field in reasoning_fields: + content_delta.pop(field, None) + + return [ + f"data: {json.dumps(reasoning_payload)}\n\n", + f"data: {json.dumps(content_payload)}\n\n", + ] def _process_chunk( self, @@ -325,6 +718,19 @@ def _process_chunk( ProcessedChunk with SSE string and metadata """ # Convert chunk to dict + if isinstance(chunk, str): + stripped = chunk.strip() + if stripped == "[DONE]" or stripped == "data: [DONE]": + return ProcessedChunk(sse_string="", finish_reason="stop") + if stripped.startswith("data:") or stripped.startswith("event:") or stripped.startswith(":"): + error = self._in_band_error_payload(chunk) + if error is not None: + raise StreamedAPIError( + str(error.get("message") or error.get("type") or "Provider stream failed"), + data={"error": error}, + ) + usage = _usage_from_sse_string(chunk) + return ProcessedChunk(sse_string=chunk if chunk.endswith("\n\n") else f"{chunk}\n\n", usage=usage) if hasattr(chunk, "model_dump"): chunk_dict = chunk.model_dump() elif hasattr(chunk, "dict"): @@ -332,8 +738,16 @@ def _process_chunk( else: chunk_dict = chunk - # Extract metadata before modifying - usage = chunk_dict.get("usage") + error = self._in_band_error_payload(chunk_dict) + if error is not None: + raise StreamedAPIError( + str(error.get("message") or error.get("type") or "Provider stream failed"), + data={"error": error}, + ) + + # Extract metadata before modifying. Providers can report cost as a + # sibling of `usage`, so keep those fields attached for normalization. + usage = _usage_from_chunk_dict(chunk_dict) finish_reason = None chunk_has_tool_calls = False @@ -350,7 +764,7 @@ def _process_chunk( # Detect final chunk using multiple signals: # 1. Primary: has usage with any meaningful token count > 0 - # 2. Secondary: has usage (even empty) + source has finish_reason (Fallback case) + # 2. Secondary: has usage (even empty) + source has finish_reason has_meaningful_usage = ( usage and isinstance(usage, dict) @@ -387,7 +801,7 @@ def _process_chunk( # INTERMEDIATE CHUNK: Never emit finish_reason choice["finish_reason"] = None - usage = chunk_dict.get("usage") + usage = _usage_from_chunk_dict(chunk_dict) if isinstance(usage, dict): normalize_usage_for_response(usage, model) @@ -398,6 +812,46 @@ def _process_chunk( has_tool_calls=chunk_has_tool_calls, ) + @classmethod + def _in_band_error_payload(cls, chunk: Any) -> Optional[Dict[str, Any]]: + """Return a structured provider error embedded in a stream frame.""" + + event_name = None + payload = chunk + if isinstance(chunk, str): + event_names = [line[6:].strip() for line in chunk.splitlines() if line.startswith("event:")] + event_name = event_names[-1] if event_names else None + data_payloads = cls._sse_data_payloads(chunk) + if not data_payloads: + return None + try: + payload = json.loads(data_payloads[-1]) + except json.JSONDecodeError: + if event_name == "error": + return {"type": "upstream_error", "message": data_payloads[-1]} + return None + if not isinstance(payload, dict): + return None + error = payload.get("error") + if isinstance(error, dict): + return deepcopy(error) + if error is not None: + return {"type": "upstream_error", "message": str(error)} + event_type = str(event_name or payload.get("event_type") or payload.get("type") or "") + if event_type == "response.failed": + response = payload.get("response") if isinstance(payload.get("response"), dict) else {} + nested = response.get("error") + if isinstance(nested, dict): + return deepcopy(nested) + if event_type in {"error", "response.error"}: + result = deepcopy(payload) + result.pop("event_type", None) + if result.get("type") in {None, "error", "response.error"}: + result["type"] = str(result.get("error_type") or "upstream_error") + result.setdefault("message", "Provider stream failed") + return result + return None + def _try_extract_error( self, exception: Exception, @@ -446,25 +900,214 @@ def _try_extract_error( return None - def _calculate_stream_cost( + def _calculate_stream_cost_breakdown( self, model: str, - prompt_tokens: int, - completion_tokens: int, - ) -> float: + usage_record: UsageRecord, + *, + skip_cost_calculation: bool, + ) -> CostBreakdown: + """Calculate advisory stream cost through the shared cost helper.""" + + if skip_cost_calculation: + return CostBreakdown(pricing_source="skipped") + return CostCalculator().calculate(usage_record, model=model) + + @staticmethod + def _log_stream_usage_accounting( + transaction_logger: Optional[Any], + usage_record: UsageRecord, + cost_breakdown: CostBreakdown, + ) -> None: + """Trace normalized stream usage without affecting stream delivery.""" + + if not transaction_logger: + return + transaction_logger.log_transform_pass( + "usage_accounting_summary", + {"usage": usage_record.to_dict(), "cost": cost_breakdown.to_dict()}, + direction="metadata", + stage="final", + transport="sse", + metadata={ + "source": usage_record.source, + "pricing_source": cost_breakdown.pricing_source, + }, + snapshot=False, + ) + + +def _next_stream_wait_seconds( + monitor: StreamMonitor, + settings: Any, + last_heartbeat_at: float, +) -> Optional[float]: + """Return the next active stream policy deadline. + + A pending upstream `__anext__()` task is not cancelled for heartbeats. The + handler waits until the nearest policy deadline, emits heartbeat metadata if + needed, and keeps waiting on the same upstream task. + """ + + candidates: list[float] = [] + now = time.monotonic() + if settings.heartbeat_seconds: + candidates.append(max(0.0, last_heartbeat_at + settings.heartbeat_seconds - now)) + if settings.ttfb_timeout_seconds and monitor.metrics.first_byte_at is None: + candidates.append(max(0.0, monitor.metrics.started_at + settings.ttfb_timeout_seconds - now)) + if settings.stall_timeout_seconds and monitor.metrics.first_byte_at is not None: + last_chunk_at = monitor.metrics.last_chunk_at or monitor.metrics.first_byte_at + candidates.append(max(0.0, last_chunk_at + settings.stall_timeout_seconds - now)) + return min(candidates) if candidates else None + + +def _heartbeat_due(monitor: StreamMonitor, settings: Any, last_heartbeat_at: float) -> bool: + """Return whether a heartbeat should be emitted while upstream is idle.""" + + if not settings.heartbeat_seconds: + return False + return time.monotonic() - last_heartbeat_at >= settings.heartbeat_seconds + + +def _stream_timeout_error(monitor: StreamMonitor, settings: Any) -> Optional[tuple[str, dict[str, Any], str]]: + """Return a structured stream timeout error when a configured deadline expires.""" + + now = time.monotonic() + if settings.ttfb_timeout_seconds and monitor.metrics.first_byte_at is None: + if now - monitor.metrics.started_at >= settings.ttfb_timeout_seconds: + return ( + "ttfb_timeout", + { + "message": "Stream timed out before first byte", + "type": "api_connection", + "details": {"timeout_type": "ttfb", "timeout_seconds": settings.ttfb_timeout_seconds}, + }, + "stream_ttfb_timeout", + ) + if settings.stall_timeout_seconds and monitor.metrics.first_byte_at is not None: + last_chunk_at = monitor.metrics.last_chunk_at or monitor.metrics.first_byte_at + if now - last_chunk_at >= settings.stall_timeout_seconds: + return ( + "stall_timeout", + { + "message": "Stream stalled while waiting for provider data", + "type": "api_connection", + "details": {"timeout_type": "stall", "timeout_seconds": settings.stall_timeout_seconds}, + }, + "stream_stall_timeout", + ) + return None + + +def _usage_from_sse_string(chunk: str) -> Optional[dict[str, Any]]: + """Extract usage from already formatted SSE chunks when native streams pass through.""" + + text = chunk.strip() + data_lines: list[str] = [] + for line in text.splitlines(): + stripped = line.strip() + if stripped.startswith("data:"): + data_lines.append(stripped[5:].strip()) + if not data_lines: + return None + payload = "\n".join(data_lines).strip() + if not payload or payload == "[DONE]": + return None + try: + data = json.loads(payload) + except json.JSONDecodeError: + return None + if not isinstance(data, dict): + return None + usage = data.get("usage") + if not isinstance(usage, dict): + return None + merged = dict(usage) + for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd", "currency", "costMetadata"): + if key in data and key not in merged: + merged[key] = data[key] + return merged + + +def _usage_from_chunk_dict(chunk: Any) -> Optional[dict[str, Any]]: + """Return dict chunk usage with top-level provider cost siblings.""" + + if not isinstance(chunk, dict): + return None + usage = chunk.get("usage") + if not isinstance(usage, dict): + return None + merged = dict(usage) + for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd", "currency", "costMetadata"): + if key in chunk and key not in merged: + merged[key] = chunk[key] + return merged + + +def _usage_record_from_sse_cost_chunk(chunk: Any, *, model: str) -> UsageRecord: + """Extract provider-reported stream cost from SSE comments/events.""" + + if not isinstance(chunk, str): + return UsageRecord(source="stream_cost_event", model=model) + cost_payload = _sse_cost_payload(chunk) + if cost_payload is None: + return UsageRecord(source="stream_cost_event", model=model) + if isinstance(cost_payload, (int, float, str)): + cost_payload = {"provider_reported_cost": cost_payload, "source": "sse_cost"} + if not isinstance(cost_payload, dict): + return UsageRecord(source="stream_cost_event", model=model) + return extract_usage_record( + {"usage": {"provider_reported_cost": cost_payload.get("provider_reported_cost", cost_payload.get("request_cost_usd", cost_payload.get("total_cost", cost_payload.get("cost", cost_payload.get("estimated_cost"))))), "currency": cost_payload.get("currency", "USD"), "cost_details": cost_payload}}, + model=model, + source="stream_cost_event", + ) + + +def _sse_cost_payload(chunk: str) -> Any: + """Parse `: cost ...` comments and `event: cost` frames.""" + + event_type: Optional[str] = None + data_lines: list[str] = [] + for line in chunk.strip().splitlines(): + stripped = line.strip() + if stripped.startswith(":"): + comment = stripped[1:].strip() + if comment.startswith("cost"): + return _parse_cost_text(comment[4:].strip()) + continue + if stripped.startswith("event:"): + event_type = stripped[6:].strip() + continue + if stripped.startswith("data:"): + data_lines.append(stripped[5:].strip()) + if event_type == "cost" and data_lines: + return _parse_cost_text("\n".join(data_lines).strip()) + return None + + +def _parse_cost_text(text: str) -> Any: + if not text: + return None + try: + return json.loads(text) + except json.JSONDecodeError: try: - model_info = litellm.get_model_info(model) - input_cost = model_info.get("input_cost_per_token") - output_cost = model_info.get("output_cost_per_token") - total_cost = 0.0 - if input_cost: - total_cost += prompt_tokens * input_cost - if output_cost: - total_cost += completion_tokens * output_cost - return total_cost - except Exception as exc: - lib_logger.debug(f"Stream cost calculation failed for {model}: {exc}") - return 0.0 + return float(text) + except ValueError: + return None + + +def _merge_usage_cost(base: UsageRecord, cost_record: UsageRecord) -> UsageRecord: + """Copy provider-reported cost onto an existing normalized usage record.""" + + if cost_record.provider_reported_cost is None: + return base + return replace( + base, + provider_reported_cost=cost_record.provider_reported_cost, + cost_currency=cost_record.cost_currency, + cost_source=cost_record.cost_source, + ) class StreamBuffer: diff --git a/src/rotator_library/client/transforms.py b/src/rotator_library/client/transforms.py index fb2a02d88..78f052ba4 100644 --- a/src/rotator_library/client/transforms.py +++ b/src/rotator_library/client/transforms.py @@ -15,6 +15,7 @@ """ import logging +from copy import deepcopy from typing import Any, Callable, Dict, List, Optional lib_logger = logging.getLogger("rotator_library") @@ -81,6 +82,11 @@ async def apply( credential: str, kwargs: Dict[str, Any], provider_config_override: Optional[Dict[str, Any]] = None, + *, + transaction_logger: Optional[Any] = None, + credential_id: Optional[str] = None, + transport: Optional[str] = None, + trace_metadata: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """ Apply all applicable transforms to request kwargs. @@ -95,43 +101,146 @@ async def apply( Modified kwargs """ modifications: List[str] = [] + trace_metadata = dict(trace_metadata or {}) + _trace_transform_pass( + transaction_logger, + "pre_provider_transform_request", + kwargs, + provider=provider, + model=model, + credential_id=credential_id, + transport=transport, + metadata={"phase": "start", **trace_metadata}, + ) # 1. Apply built-in transforms for transform_provider, transforms in self._transforms.items(): # Check if transform applies (provider match or model contains pattern) if transform_provider == provider or transform_provider in model.lower(): for transform in transforms: - result = transform(kwargs, model, provider) + before = deepcopy(kwargs) if transaction_logger else None + try: + result = transform(kwargs, model, provider) + except Exception as exc: + if transaction_logger: + transaction_logger.log_transform_error( + "builtin_provider_transform", + exc, + payload=before if before is not None else kwargs, + stage="client", + transport=transport, + metadata={ + "provider": provider, + "model": model, + "credential_id": credential_id, + "transform_provider": transform_provider, + "transform_name": getattr(transform, "__name__", repr(transform)), + **trace_metadata, + }, + ) + raise if result: modifications.append(result) + _trace_transform_pass( + transaction_logger, + "after_builtin_provider_transform", + kwargs, + provider=provider, + model=model, + credential_id=credential_id, + transport=transport, + changed_from_previous=(before != kwargs) if before is not None else None, + metadata={ + "transform_provider": transform_provider, + "transform_name": getattr(transform, "__name__", repr(transform)), + "modification": result, + **trace_metadata, + }, + ) # 2. Apply provider hook transforms (async) plugin = self._get_plugin_instance(provider) if plugin and hasattr(plugin, "transform_request"): try: + before = deepcopy(kwargs) if transaction_logger else None hook_result = await plugin.transform_request(kwargs, model, credential) if hook_result: modifications.extend(hook_result) + if hook_result or (before is not None and before != kwargs): + _trace_transform_pass( + transaction_logger, + "after_provider_hook_transform", + kwargs, + provider=provider, + model=model, + credential_id=credential_id, + transport=transport, + changed_from_previous=(before != kwargs) if before is not None else None, + metadata={"modifications": hook_result or [], **trace_metadata}, + ) except Exception as e: lib_logger.debug(f"Provider transform_request hook failed: {e}") + if transaction_logger: + transaction_logger.log_transform_error( + "provider_hook_transform", + e, + payload=kwargs, + stage="client", + transport=transport, + metadata={"provider": provider, "model": model, "credential_id": credential_id, **trace_metadata}, + ) # 3. Apply model-specific options from provider if plugin and hasattr(plugin, "get_model_options"): model_options = plugin.get_model_options(model) if model_options: + before = deepcopy(kwargs) if transaction_logger else None for key, value in model_options.items(): if key == "reasoning_effort": kwargs["reasoning_effort"] = value elif key not in kwargs: kwargs[key] = value modifications.append(f"applied model options for {model}") + _trace_transform_pass( + transaction_logger, + "after_provider_model_options", + kwargs, + provider=provider, + model=model, + credential_id=credential_id, + transport=transport, + changed_from_previous=(before != kwargs) if before is not None else None, + metadata={"model_options": deepcopy(model_options) if transaction_logger else None, **trace_metadata}, + ) # 4. Apply LiteLLM conversion if config available if self._config and hasattr(self._config, "convert_for_litellm"): + before = deepcopy(kwargs) if transaction_logger else None + _trace_transform_pass( + transaction_logger, + "before_litellm_conversion", + kwargs, + provider=provider, + model=model, + credential_id=credential_id, + transport=transport, + metadata={"provider_config_override": bool(provider_config_override), **trace_metadata}, + ) kwargs = self._config.convert_for_litellm( provider_override=provider_config_override, **kwargs, ) + _trace_transform_pass( + transaction_logger, + "after_litellm_conversion", + kwargs, + provider=provider, + model=model, + credential_id=credential_id, + transport=transport, + changed_from_previous=(before != kwargs) if before is not None else None, + metadata={"provider_config_override": bool(provider_config_override), **trace_metadata}, + ) if modifications: lib_logger.debug( @@ -312,3 +421,36 @@ def _transform_dedaluslabs_tool_choice( # caused 400 errors on models that don't support those categories (e.g. Gemma). # See gemini_provider.py for the full removal comment with previous defaults. # ========================================================================= + + +def _trace_transform_pass( + transaction_logger: Optional[Any], + pass_name: str, + payload: Dict[str, Any], + *, + provider: str, + model: str, + credential_id: Optional[str], + transport: Optional[str], + metadata: Dict[str, Any], + changed_from_previous: Optional[bool] = None, +) -> None: + """Record provider-transform states without changing transform behavior. + + Transform tracing is observability-only. This helper centralizes pass + metadata so individual transforms can stay focused on payload mutation while + the transaction trace still shows each live request state. + """ + + if not transaction_logger: + return + transaction_logger.log_transform_pass( + pass_name, + payload, + direction="request", + stage="client", + credential_id=credential_id, + transport=transport, + changed_from_previous=changed_from_previous, + metadata={"provider": provider, "model": model, **metadata}, + ) diff --git a/src/rotator_library/config/experimental.py b/src/rotator_library/config/experimental.py new file mode 100644 index 000000000..ada8084d0 --- /dev/null +++ b/src/rotator_library/config/experimental.py @@ -0,0 +1,828 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Optional structured configuration for experimental native features. + +The proxy remains environment-first: existing `.env` variables keep working and +environment variables override this JSON layer. This module intentionally avoids +secrets. API keys, OAuth tokens, bearer headers, and similar values must remain +in environment variables or provider-managed credential files. +""" + +from __future__ import annotations + +import json +import os +import re +import string +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping, Optional +from urllib.parse import parse_qsl, urlparse + +from ..field_cache import FieldCacheInjection, FieldCacheRule +from ..usage.costs import ModelPricing + +_CONFIG_ENV_KEYS = ("LLM_PROXY_CONFIG_FILE", "PROXY_CONFIG_FILE") +_KNOWN_SECTIONS = {"routing", "pricing", "streaming", "field_cache", "providers", "retry", "responses"} +_SECRET_KEY_PARTS = ("api_key", "apikey", "authorization", "access_token", "accesstoken", "refresh_token", "refreshtoken", "oauth_token", "oauthtoken", "oauth_token_secret", "oauthtokensecret", "id_token", "idtoken", "token_secret", "tokensecret", "client_secret", "clientsecret", "secret_key", "secretkey", "bearer_token", "bearertoken", "credential", "credentials", "password") +_PROVIDER_CONFIG_KEYS = { + "protocol_name", + "api_base", + "endpoint_paths", + "auth_mode", + "auth_header_name", + "models", + "default_output_protocol", + "adapter_names", + "adapter_config", + "native_streaming_supported", + "field_cache", + "model_quota_groups", +} +_HTTP_TOKEN_RE = re.compile(r"^[!#$%&'*+.^_`|~0-9A-Za-z-]+$") +_PROVIDER_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]*$") +_ENDPOINT_TEMPLATE_FIELDS = {"model", "operation", "provider"} + + +class ExperimentalConfigError(ValueError): + """Raised when optional structured config is malformed or unsafe.""" + + +@dataclass(frozen=True) +class ExperimentalConfig: + """Parsed optional JSON config. + + Sections are stored as dictionaries rather than deep custom classes so Phase + 10 can layer config onto existing feature-specific parsers without creating + a second full application configuration system. + """ + + routing: dict[str, Any] = field(default_factory=dict) + pricing: dict[str, Any] = field(default_factory=dict) + streaming: dict[str, Any] = field(default_factory=dict) + field_cache: dict[str, Any] = field(default_factory=dict) + providers: dict[str, Any] = field(default_factory=dict) + retry: dict[str, Any] = field(default_factory=dict) + responses: dict[str, Any] = field(default_factory=dict) + unknown_sections: dict[str, Any] = field(default_factory=dict) + warnings: tuple[str, ...] = () + path: Optional[str] = None + + @property + def is_empty(self) -> bool: + return not (self.routing or self.pricing or self.streaming or self.field_cache or self.providers or self.retry or self.responses or self.unknown_sections) + + +@dataclass(frozen=True) +class StreamRuntimeSettings: + """Runtime stream observability settings. + + Timeout and heartbeat values default to disabled so existing long-running + reasoning streams keep working. Operators can opt into active stream + hardening through env or JSON config without changing provider code. + """ + + ttfb_timeout_seconds: Optional[float] = None + stall_timeout_seconds: Optional[float] = None + heartbeat_seconds: Optional[float] = None + cancel_upstream_on_disconnect: bool = True + trace_metrics: bool = True + + +@dataclass(frozen=True) +class RetryRuntimeSettings: + """Runtime retry/cooldown settings layered from JSON and env.""" + + provider_cooldown_min_seconds: int = 10 + provider_cooldown_default_seconds: int = 30 + provider_cooldown_on_quota: bool = False + provider_backoff_window_seconds: int = 60 + provider_backoff_threshold: int = 3 + provider_backoff_base_seconds: Optional[int] = None + provider_backoff_max_seconds: int = 300 + failure_history_max_entries: int = 200 + + +@dataclass(frozen=True) +class ResponsesStoreRuntimeSettings: + """Runtime backend selection for Responses storage.""" + + backend: str = "memory" + cache_name: str = "responses" + cache_prefix: str = "responses" + cache_dir: Optional[str] = None + cache_memory_ttl_seconds: int = 3600 + cache_disk_ttl_seconds: int = 172800 + + +@dataclass(frozen=True) +class ProviderRuntimeConfig: + """Safe provider metadata loaded from optional JSON config.""" + + protocol_name: Optional[str] = None + api_base: Optional[str] = None + endpoint_paths: dict[str, str] = field(default_factory=dict) + auth_mode: str = "bearer" + auth_header_name: Optional[str] = None + models: tuple[str, ...] = () + default_output_protocol: Optional[str] = None + adapter_names: Optional[tuple[str, ...]] = None + adapter_config: dict[str, dict[str, Any]] = field(default_factory=dict) + native_streaming_supported: Optional[bool] = None + field_cache_rules: tuple[FieldCacheRule, ...] = () + model_quota_groups: Optional[dict[str, list[str]]] = None + + +def load_experimental_config(path: str | os.PathLike[str] | None = None, env: Mapping[str, str] | None = None) -> ExperimentalConfig: + """Load optional JSON config from an explicit path or config env var.""" + + source = env if env is not None else os.environ + resolved = Path(path) if path is not None else _path_from_env(source) + if resolved is None or not resolved.exists(): + return ExperimentalConfig(path=str(resolved) if resolved is not None else None) + try: + data = json.loads(resolved.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ExperimentalConfigError(f"Invalid JSON config at {resolved}: {exc.msg}") from exc + if not isinstance(data, dict): + raise ExperimentalConfigError("JSON config root must be an object") + _reject_secret_keys(data) + _validate_provider_sections(data.get("providers", {})) + warnings = tuple(f"Unknown config section '{key}' ignored by current runtime" for key in data if key not in _KNOWN_SECTIONS) + unknown = {key: value for key, value in data.items() if key not in _KNOWN_SECTIONS} + return ExperimentalConfig( + routing=_dict_section(data, "routing"), + pricing=_dict_section(data, "pricing"), + streaming=_dict_section(data, "streaming"), + field_cache=_dict_section(data, "field_cache"), + providers=_dict_section(data, "providers"), + retry=_dict_section(data, "retry"), + responses=_dict_section(data, "responses"), + unknown_sections=unknown, + warnings=warnings, + path=str(resolved), + ) + + +def load_config_from_mapping(data: Mapping[str, Any]) -> ExperimentalConfig: + """Build config from an in-memory mapping for tests and provider helpers.""" + + _reject_secret_keys(data) + _validate_provider_sections(data.get("providers", {})) + warnings = tuple(f"Unknown config section '{key}' ignored by current runtime" for key in data if key not in _KNOWN_SECTIONS) + return ExperimentalConfig( + routing=_dict_section(data, "routing"), + pricing=_dict_section(data, "pricing"), + streaming=_dict_section(data, "streaming"), + field_cache=_dict_section(data, "field_cache"), + providers=_dict_section(data, "providers"), + retry=_dict_section(data, "retry"), + responses=_dict_section(data, "responses"), + unknown_sections={key: value for key, value in data.items() if key not in _KNOWN_SECTIONS}, + warnings=warnings, + ) + + +def get_configured_model_pricing( + provider: str, + model: str, + *, + config: ExperimentalConfig | None = None, + env: Mapping[str, str] | None = None, +) -> Optional[ModelPricing]: + """Return JSON/env pricing for a provider/model, with env taking priority.""" + + source = env if env is not None else os.environ + env_pricing = _pricing_from_env(provider, model, source) + if env_pricing: + return env_pricing + active = config if config is not None else load_experimental_config(env=source) + pricing_section = active.pricing.get(provider, {}) if isinstance(active.pricing, dict) else {} + raw = pricing_section.get(model) if isinstance(pricing_section, dict) else None + if not isinstance(raw, dict): + return None + return ModelPricing( + input_cost_per_token=as_float(raw.get("input", raw.get("input_cost_per_token", 0.0)), name="pricing.input"), + output_cost_per_token=as_float(raw.get("output", raw.get("output_cost_per_token", 0.0)), name="pricing.output"), + cache_read_cost_per_token=as_float(raw.get("cache_read", raw.get("cache_read_cost_per_token", 0.0)), name="pricing.cache_read"), + cache_write_cost_per_token=as_float(raw.get("cache_write", raw.get("cache_write_cost_per_token", 0.0)), name="pricing.cache_write"), + reasoning_cost_per_token=as_float(raw.get("reasoning", raw.get("reasoning_cost_per_token", 0.0)), name="pricing.reasoning"), + currency=str(raw.get("currency", "USD")), + source="json_config", + ) + + +def get_stream_runtime_settings( + *, + config: ExperimentalConfig | None = None, + env: Mapping[str, str] | None = None, +) -> StreamRuntimeSettings: + """Return stream runtime settings with environment overriding JSON.""" + + source = env if env is not None else os.environ + active = config if config is not None else load_experimental_config(env=source) + streaming = active.streaming if isinstance(active.streaming, dict) else {} + + return StreamRuntimeSettings( + ttfb_timeout_seconds=_optional_positive_float(_env_or_json(source, "STREAM_TTFB_TIMEOUT_SECONDS", streaming, "ttfb_timeout_seconds"), "STREAM_TTFB_TIMEOUT_SECONDS"), + stall_timeout_seconds=_optional_positive_float(_env_or_json(source, "STREAM_STALL_TIMEOUT_SECONDS", streaming, "stall_timeout_seconds"), "STREAM_STALL_TIMEOUT_SECONDS"), + heartbeat_seconds=_optional_positive_float(_env_or_json(source, "STREAM_HEARTBEAT_INTERVAL_SECONDS", streaming, "heartbeat_interval_seconds", default=_env_or_json(source, "STREAM_HEARTBEAT_SECONDS", streaming, "heartbeat_seconds")), "STREAM_HEARTBEAT_INTERVAL_SECONDS"), + cancel_upstream_on_disconnect=as_bool(_env_or_json(source, "STREAM_CANCEL_UPSTREAM_ON_DISCONNECT", streaming, "cancel_upstream_on_disconnect", default=True), name="STREAM_CANCEL_UPSTREAM_ON_DISCONNECT"), + trace_metrics=as_bool(_env_or_json(source, "STREAM_TRACE_METRICS", streaming, "trace_metrics", default=True), name="STREAM_TRACE_METRICS"), + ) + + +def get_retry_runtime_settings( + *, + config: ExperimentalConfig | None = None, + env: Mapping[str, str] | None = None, +) -> RetryRuntimeSettings: + """Return retry/cooldown settings with environment overriding JSON.""" + + source = env if env is not None else os.environ + active = config if config is not None else load_experimental_config(env=source) + retry = active.retry if isinstance(active.retry, dict) else {} + cooldown = retry.get("provider_cooldown", {}) if isinstance(retry.get("provider_cooldown"), dict) else retry + backoff = retry.get("backoff", {}) if isinstance(retry.get("backoff"), dict) else retry + return RetryRuntimeSettings( + provider_cooldown_min_seconds=max(0, _int_setting(source, "PROVIDER_COOLDOWN_MIN_SECONDS", cooldown, "provider_cooldown_min_seconds", 10)), + provider_cooldown_default_seconds=max(0, _int_setting(source, "PROVIDER_COOLDOWN_DEFAULT_SECONDS", cooldown, "provider_cooldown_default_seconds", 30)), + provider_cooldown_on_quota=_bool_setting(source, "PROVIDER_COOLDOWN_ON_QUOTA", cooldown, "provider_cooldown_on_quota", False), + provider_backoff_window_seconds=max(0, _int_setting(source, "PROVIDER_BACKOFF_WINDOW_SECONDS", backoff, "provider_backoff_window_seconds", 60)), + provider_backoff_threshold=max(1, _int_setting(source, "PROVIDER_BACKOFF_THRESHOLD", backoff, "provider_backoff_threshold", 3)), + provider_backoff_base_seconds=_optional_int_setting(source, "PROVIDER_BACKOFF_BASE_SECONDS", backoff, "provider_backoff_base_seconds"), + provider_backoff_max_seconds=max(1, _int_setting(source, "PROVIDER_BACKOFF_MAX_SECONDS", backoff, "provider_backoff_max_seconds", 300)), + failure_history_max_entries=max(1, _int_setting(source, "FAILURE_HISTORY_MAX_ENTRIES", backoff, "failure_history_max_entries", 200)), + ) + + +def get_responses_store_settings( + *, + config: ExperimentalConfig | None = None, + env: Mapping[str, str] | None = None, +) -> Any: + """Return Responses store settings with environment overriding JSON.""" + + from ..responses import ResponsesStoreSettings + + source = env if env is not None else os.environ + active = config if config is not None else load_experimental_config(env=source) + responses = active.responses if isinstance(active.responses, dict) else {} + store = responses.get("store", {}) if isinstance(responses.get("store"), dict) else responses + ttl_seconds = _optional_positive_int(_env_or_json(source, "RESPONSES_STORE_TTL_SECONDS", store, "ttl_seconds"), "RESPONSES_STORE_TTL_SECONDS") + max_items = _optional_positive_int(_env_or_json(source, "RESPONSES_STORE_MAX_ITEMS", store, "max_items"), "RESPONSES_STORE_MAX_ITEMS") + return ResponsesStoreSettings( + ttl_seconds=ttl_seconds, + max_items=max_items, + store_failed=as_bool(_env_or_json(source, "RESPONSES_STORE_FAILED", store, "store_failed", default=True), name="RESPONSES_STORE_FAILED"), + store_in_progress=as_bool(_env_or_json(source, "RESPONSES_STORE_IN_PROGRESS", store, "store_in_progress", default=False), name="RESPONSES_STORE_IN_PROGRESS"), + ) + + +def get_responses_store_runtime_settings( + *, + config: ExperimentalConfig | None = None, + env: Mapping[str, str] | None = None, +) -> ResponsesStoreRuntimeSettings: + """Return Responses store backend settings with env overriding JSON.""" + + source = env if env is not None else os.environ + active = config if config is not None else load_experimental_config(env=source) + responses = active.responses if isinstance(active.responses, dict) else {} + store = responses.get("store", {}) if isinstance(responses.get("store"), dict) else responses + backend = str(_env_or_json(source, "RESPONSES_STORE_BACKEND", store, "backend", default="memory")).strip().lower() + if backend not in {"memory", "provider_cache"}: + raise ExperimentalConfigError("RESPONSES_STORE_BACKEND must be 'memory' or 'provider_cache'") + return ResponsesStoreRuntimeSettings( + backend=backend, + cache_name=str(_env_or_json(source, "RESPONSES_STORE_CACHE_NAME", store, "cache_name", default="responses")), + cache_prefix=str(_env_or_json(source, "RESPONSES_STORE_CACHE_PREFIX", store, "cache_prefix", default="responses")), + cache_dir=_optional_string(_env_or_json(source, "RESPONSES_STORE_CACHE_DIR", store, "cache_dir")), + cache_memory_ttl_seconds=max(1, _int_setting(source, "RESPONSES_STORE_CACHE_MEMORY_TTL_SECONDS", store, "cache_memory_ttl_seconds", 3600)), + cache_disk_ttl_seconds=max(1, _int_setting(source, "RESPONSES_STORE_CACHE_DISK_TTL_SECONDS", store, "cache_disk_ttl_seconds", 172800)), + ) + + +def get_provider_runtime_config( + provider: str, + model: str = "", + *, + config: ExperimentalConfig | None = None, + env: Mapping[str, str] | None = None, +) -> ProviderRuntimeConfig: + """Return safe JSON-configured provider metadata for one provider/model.""" + + source = env if env is not None else os.environ + active = config if config is not None else load_experimental_config(env=source) + providers = active.providers if isinstance(active.providers, dict) else {} + raw = providers.get(provider, {}) + if not isinstance(raw, Mapping) or not raw: + return ProviderRuntimeConfig() + _validate_provider_sections({provider: raw}) + protocol_name = _configured_provider_protocol(raw.get("protocol_name")) + adapter_names = _configured_adapters(raw.get("adapter_names")) + adapter_config = _configured_adapter_config(raw.get("adapter_config", {})) + native_streaming_supported = None + if "native_streaming_supported" in raw: + native_streaming_supported = as_bool(raw.get("native_streaming_supported"), name="providers.native_streaming_supported") + field_cache_rules = _configured_provider_field_cache(provider, model, raw.get("field_cache")) + model_quota_groups = _configured_quota_groups(raw.get("model_quota_groups")) if "model_quota_groups" in raw else None + return ProviderRuntimeConfig( + protocol_name=protocol_name, + api_base=_configured_api_base(raw.get("api_base")), + endpoint_paths=_configured_endpoint_paths(raw.get("endpoint_paths")), + auth_mode=_configured_auth_mode(raw.get("auth_mode")), + auth_header_name=_configured_auth_header_name(raw.get("auth_header_name")), + models=_configured_models(raw.get("models")), + default_output_protocol=_configured_output_protocol(raw.get("default_output_protocol")), + adapter_names=adapter_names, + adapter_config=adapter_config, + native_streaming_supported=native_streaming_supported, + field_cache_rules=field_cache_rules, + model_quota_groups=model_quota_groups, + ) + + +def parse_field_cache_rules(config: ExperimentalConfig, provider: str, model: str) -> tuple[FieldCacheRule, ...]: + """Parse configured field-cache rules for a provider/model. + + Wildcard model rules are returned before exact model rules so providers can + define general preservation behavior and then append model-specific rules. + This helper is intentionally not auto-wired into providers; providers decide + whether external config is appropriate for their protocol state. + """ + + provider_rules = config.field_cache.get(provider, {}) if isinstance(config.field_cache, dict) else {} + if not isinstance(provider_rules, dict): + raise ExperimentalConfigError("field_cache provider section must be an object") + raw_rules: list[Any] = [] + keys = ["*"] + if "/" in model: + keys.append(model.split("/", 1)[1]) + keys.append(model) + for key in dict.fromkeys(keys): + value = provider_rules.get(key, []) + if isinstance(value, list): + raw_rules.extend(value) + elif value not in (None, []): + raise ExperimentalConfigError("field_cache model rules must be a list") + parsed_rules = [] + for rule in raw_rules: + if not isinstance(rule, dict): + raise ExperimentalConfigError("field_cache rule entries must be objects") + parsed_rules.append(_field_cache_rule_from_dict(rule)) + return tuple(parsed_rules) + + +def as_bool(value: Any, *, name: str) -> bool: + """Parse a JSON/env boolean value.""" + + if isinstance(value, bool): + return value + if isinstance(value, (int, float)) and value in (0, 1): + return bool(value) + if isinstance(value, str): + text = value.strip().lower() + if text in {"1", "true", "yes", "on"}: + return True + if text in {"0", "false", "no", "off"}: + return False + raise ExperimentalConfigError(f"Invalid boolean for {name}") + + +def as_float(value: Any, *, name: str) -> float: + """Parse a JSON/env float value with redacted errors.""" + + try: + return float(value) + except (TypeError, ValueError) as exc: + raise ExperimentalConfigError(f"Invalid number for {name}") from exc + + +def as_int(value: Any, *, name: str) -> int: + """Parse a JSON/env integer value with redacted errors.""" + + try: + return int(value) + except (TypeError, ValueError) as exc: + raise ExperimentalConfigError(f"Invalid integer for {name}") from exc + + +def env_price_key(provider: str, model: str, suffix: str) -> str: + """Return normalized model-price environment variable name.""" + + return f"MODEL_PRICE_{_env_part(provider)}_{_env_part(model)}_{_env_part(suffix)}" + + +def _path_from_env(env: Mapping[str, str]) -> Optional[Path]: + for key in _CONFIG_ENV_KEYS: + value = env.get(key) + if value: + return Path(value) + return None + + +def _dict_section(data: Mapping[str, Any], key: str) -> dict[str, Any]: + value = data.get(key, {}) + if value is None: + return {} + if not isinstance(value, dict): + raise ExperimentalConfigError(f"Config section '{key}' must be an object") + return dict(value) + + +def _reject_secret_keys(value: Any, path: str = "config") -> None: + if isinstance(value, Mapping): + for key, nested in value.items(): + if _is_secret_like_key(key): + raise ExperimentalConfigError(f"Unsafe secret-like key in JSON config at {path}.{key}") + _reject_secret_keys(nested, f"{path}.{key}") + elif isinstance(value, list): + for index, nested in enumerate(value): + _reject_secret_keys(nested, f"{path}[{index}]") + + +def _validate_provider_sections(value: Any) -> None: + if value in (None, {}): + return + if not isinstance(value, Mapping): + raise ExperimentalConfigError("providers config section must be an object") + for provider, raw in value.items(): + if not _PROVIDER_NAME_RE.fullmatch(str(provider)): + raise ExperimentalConfigError( + f"providers key {provider!r} must contain only letters, numbers, underscores, or hyphens" + ) + if not isinstance(raw, Mapping): + raise ExperimentalConfigError(f"providers.{provider} must be an object") + unsupported = set(str(key) for key in raw) - _PROVIDER_CONFIG_KEYS + if unsupported: + raise ExperimentalConfigError(f"providers.{provider} contains unsupported keys: {', '.join(sorted(unsupported))}") + _configured_provider_protocol(raw.get("protocol_name")) + _configured_output_protocol(raw.get("default_output_protocol")) + _configured_adapters(raw.get("adapter_names")) + _configured_adapter_config(raw.get("adapter_config", {})) + _configured_api_base(raw.get("api_base")) + _configured_endpoint_paths(raw.get("endpoint_paths")) + auth_mode = _configured_auth_mode(raw.get("auth_mode")) + auth_header_name = _configured_auth_header_name(raw.get("auth_header_name")) + if auth_mode == "custom" and not auth_header_name: + raise ExperimentalConfigError("providers.auth_header_name is required when auth_mode is custom") + if auth_mode == "none" and not raw.get("protocol_name"): + raise ExperimentalConfigError( + "providers.protocol_name is required when auth_mode is none" + ) + _configured_models(raw.get("models")) + if "native_streaming_supported" in raw: + as_bool(raw.get("native_streaming_supported"), name="providers.native_streaming_supported") + if "model_quota_groups" in raw: + _configured_quota_groups(raw.get("model_quota_groups")) + + +def _configured_protocol(value: Any) -> Optional[str]: + if value in (None, ""): + return None + name = str(value).strip().lower() + try: + from ..protocols import get_protocol + + protocol = get_protocol(name) + except Exception as exc: + raise ExperimentalConfigError(f"Unknown provider protocol_name {name!r}") from exc + return protocol.name + + +def _configured_output_protocol(value: Any) -> Optional[str]: + """Return one supported generative client output protocol.""" + + protocol_name = _configured_protocol(value) + if protocol_name and protocol_name not in { + "openai_chat", + "responses", + "anthropic_messages", + "gemini", + }: + raise ExperimentalConfigError( + f"default_output_protocol must be a generative protocol, got {protocol_name!r}" + ) + return protocol_name + + +def _configured_provider_protocol(value: Any) -> Optional[str]: + """Return one protocol supported by the generative provider runtime.""" + + protocol_name = _configured_protocol(value) + if protocol_name and protocol_name not in { + "openai_chat", + "responses", + "anthropic_messages", + "gemini", + }: + raise ExperimentalConfigError( + f"protocol_name must be a supported generative protocol, got {protocol_name!r}" + ) + return protocol_name + + +def _configured_api_base(value: Any) -> Optional[str]: + if value in (None, ""): + return None + url = str(value).strip() + parsed = urlparse(url) + if parsed.scheme not in {"http", "https"} or not parsed.netloc: + raise ExperimentalConfigError("providers.api_base must be an http(s) URL") + if parsed.username or parsed.password or parsed.query or parsed.fragment: + raise ExperimentalConfigError( + "providers.api_base cannot contain credentials, query parameters, or fragments" + ) + return url.rstrip("/") + + +def _configured_endpoint_paths(value: Any) -> dict[str, str]: + if value in (None, {}): + return {} + if not isinstance(value, Mapping): + raise ExperimentalConfigError("providers.endpoint_paths must be an object") + result: dict[str, str] = {} + for operation, path in value.items(): + name = str(operation).strip() + if not isinstance(path, str): + raise ExperimentalConfigError("providers.endpoint_paths values must be strings") + rendered = path.strip() + if not name or not rendered: + raise ExperimentalConfigError("providers.endpoint_paths entries require non-empty operation and path values") + if not _PROVIDER_NAME_RE.fullmatch(name): + raise ExperimentalConfigError("providers.endpoint_paths operation names are invalid") + try: + fields = { + field_name + for _, field_name, _, _ in string.Formatter().parse(rendered) + if field_name is not None + } + except ValueError as exc: + raise ExperimentalConfigError("providers.endpoint_paths contains an invalid template") from exc + unsupported = fields - _ENDPOINT_TEMPLATE_FIELDS + if unsupported: + raise ExperimentalConfigError( + f"providers.endpoint_paths contains unsupported placeholders: {', '.join(sorted(unsupported))}" + ) + if not rendered.startswith("/") or rendered.startswith("//"): + raise ExperimentalConfigError( + "providers.endpoint_paths values must be absolute paths on api_base" + ) + parsed_path = urlparse(rendered) + if parsed_path.fragment: + raise ExperimentalConfigError( + "providers.endpoint_paths cannot contain fragments" + ) + for query_key, _ in parse_qsl(parsed_path.query, keep_blank_values=True): + if _is_secret_like_key(query_key): + raise ExperimentalConfigError( + "providers.endpoint_paths cannot contain secret-bearing query parameters" + ) + result[name] = rendered + return result + + +def _configured_auth_mode(value: Any) -> str: + mode = str(value or "bearer").strip().lower().replace("_", "-") + aliases = {"x-api-key": "x-api-key", "x-goog-api-key": "x-goog-api-key", "none": "none", "custom": "custom", "bearer": "bearer"} + if mode not in aliases: + raise ExperimentalConfigError("providers.auth_mode must be bearer, x-api-key, x-goog-api-key, custom, or none") + return aliases[mode] + + +def _is_secret_like_key(value: Any) -> bool: + """Return whether a config key names credential or authentication data.""" + + key_text = str(value).lower() + compact_key = re.sub(r"[^a-z0-9]+", "", key_text) + underscored_key = re.sub(r"[^a-z0-9]+", "_", key_text) + return any( + part in key_text or part in compact_key or part in underscored_key + for part in _SECRET_KEY_PARTS + ) + + +def _configured_auth_header_name(value: Any) -> Optional[str]: + if value in (None, ""): + return None + name = str(value).strip() + if not _HTTP_TOKEN_RE.fullmatch(name): + raise ExperimentalConfigError("providers.auth_header_name must be a valid HTTP header name") + return name + + +def _configured_models(value: Any) -> tuple[str, ...]: + if value in (None, []): + return () + if not isinstance(value, (list, tuple)): + raise ExperimentalConfigError("providers.models must be a list") + if not all(isinstance(model, str) and model.strip() for model in value): + raise ExperimentalConfigError("providers.models entries must be non-empty strings") + models = tuple(model.strip() for model in value) + return models + + +def _configured_adapters(value: Any) -> Optional[tuple[str, ...]]: + if value in (None, ""): + return None + if isinstance(value, str): + names = tuple(part.strip() for part in value.split(",") if part.strip()) + elif isinstance(value, (list, tuple)): + names = tuple(str(part) for part in value) + else: + raise ExperimentalConfigError("providers.adapter_names must be a string or list") + for name in names: + _validate_adapter_name(name) + return names + + +def _configured_adapter_config(value: Any) -> dict[str, dict[str, Any]]: + if value in (None, {}): + return {} + if not isinstance(value, Mapping): + raise ExperimentalConfigError("providers.adapter_config must be an object") + result: dict[str, dict[str, Any]] = {} + for name, config in value.items(): + adapter_name = str(name) + _validate_adapter_name(adapter_name) + if not isinstance(config, Mapping): + raise ExperimentalConfigError("providers.adapter_config entries must be objects") + result[adapter_name] = dict(config) + return result + + +def _validate_adapter_name(name: str) -> None: + try: + from ..adapters import get_adapter + + get_adapter(name) + except Exception as exc: + raise ExperimentalConfigError(f"Unknown provider adapter {name!r}") from exc + + +def _configured_provider_field_cache(provider: str, model: str, value: Any) -> tuple[FieldCacheRule, ...]: + if value in (None, {}, []): + return () + if isinstance(value, list): + section = {provider: {"*": value}} + elif isinstance(value, Mapping): + section = {provider: dict(value)} + else: + raise ExperimentalConfigError("providers.field_cache must be an object or list") + return parse_field_cache_rules(load_config_from_mapping({"field_cache": section}), provider, model) + + +def _configured_quota_groups(value: Any) -> dict[str, list[str]]: + if not isinstance(value, Mapping): + raise ExperimentalConfigError("providers.model_quota_groups must be an object") + result: dict[str, list[str]] = {} + for group, models in value.items(): + if not isinstance(models, list) or not all(isinstance(model, str) for model in models): + raise ExperimentalConfigError("providers.model_quota_groups values must be string arrays") + result[str(group)] = list(models) + return result + + +def _pricing_from_env(provider: str, model: str, env: Mapping[str, str]) -> Optional[ModelPricing]: + suffixes = { + "input": "INPUT", + "output": "OUTPUT", + "cache_read": "CACHE_READ", + "cache_write": "CACHE_WRITE", + "reasoning": "REASONING", + } + values: dict[str, float] = {} + for field_name, suffix in suffixes.items(): + key = env_price_key(provider, model, suffix) + raw = env.get(key) + if raw not in (None, ""): + try: + values[field_name] = as_float(raw, name=key) + except ExperimentalConfigError: + continue + if not values: + return None + return ModelPricing( + input_cost_per_token=values.get("input", 0.0), + output_cost_per_token=values.get("output", 0.0), + cache_read_cost_per_token=values.get("cache_read", 0.0), + cache_write_cost_per_token=values.get("cache_write", 0.0), + reasoning_cost_per_token=values.get("reasoning", 0.0), + source="env", + ) + + +def _env_or_json(env: Mapping[str, str], env_key: str, data: Mapping[str, Any], json_key: str, default: Any = None) -> Any: + if env_key in env: + return env[env_key] + return data.get(json_key, default) + + +def _int_setting(env: Mapping[str, str], env_key: str, data: Mapping[str, Any], json_key: str, default: int) -> int: + if env_key in env: + try: + return int(env.get(env_key) or default) + except (TypeError, ValueError): + return default + return as_int(data.get(json_key, default), name=env_key) + + +def _optional_int_setting(env: Mapping[str, str], env_key: str, data: Mapping[str, Any], json_key: str) -> Optional[int]: + if env_key in env: + try: + parsed = int(env.get(env_key) or 0) + except (TypeError, ValueError): + return None + return parsed if parsed > 0 else None + return _optional_positive_int(data.get(json_key), env_key) + + +def _bool_setting(env: Mapping[str, str], env_key: str, data: Mapping[str, Any], json_key: str, default: bool) -> bool: + if env_key in env: + try: + return as_bool(env.get(env_key), name=env_key) + except ExperimentalConfigError: + return default + return as_bool(data.get(json_key, default), name=env_key) + + +def _optional_positive_float(value: Any, name: str) -> Optional[float]: + if value in (None, ""): + return None + parsed = as_float(value, name=name) + # Zero and negative values mean "not configured" for timeout-like knobs. + # Runtime enforcement is intentionally disabled by default. + return parsed if parsed > 0 else None + + +def _optional_positive_int(value: Any, name: str) -> Optional[int]: + if value in (None, ""): + return None + parsed = as_int(value, name=name) + return parsed if parsed > 0 else None + + +def _optional_string(value: Any) -> Optional[str]: + if value in (None, ""): + return None + return str(value) + + +def _field_cache_rule_from_dict(data: Mapping[str, Any]) -> FieldCacheRule: + inject_data = data.get("inject") + inject = None + if isinstance(inject_data, Mapping): + inject = FieldCacheInjection( + target=str(inject_data.get("target", "request")), + path=str(inject_data.get("path", data.get("target_path", ""))), + when_missing_only=as_bool(inject_data.get("when_missing_only", False), name="field_cache.inject.when_missing_only"), + insert=as_bool(inject_data.get("insert", False), name="field_cache.inject.insert"), + as_list=as_bool(inject_data.get("as_list", False), name="field_cache.inject.as_list"), + ) + elif inject_data is not None: + raise ExperimentalConfigError("field_cache.inject must be an object") + elif data.get("target_path"): + inject = FieldCacheInjection(target=str(data.get("target", "request")), path=str(data["target_path"])) + scope = data.get("scope", ("provider", "model", "credential", "session")) + if isinstance(scope, str): + scope_values = tuple(part.strip() for part in scope.split(",") if part.strip()) + elif isinstance(scope, (list, tuple)): + scope_values = tuple(str(part) for part in scope) + else: + raise ExperimentalConfigError("field_cache.scope must be a string or list") + if {"provider", "model", "credential", "session"}.difference(scope_values): + raise ExperimentalConfigError( + "field_cache.scope must include provider, model, credential, and session" + ) + try: + return FieldCacheRule( + name=str(data["name"]), + source=str(data["source"]), + path=str(data["path"]), + cache_key=_optional_string(data.get("cache_key")), + mode=str(data.get("mode", "last")), + scope=scope_values, + inject=inject, + enabled=as_bool(data.get("enabled", True), name="field_cache.enabled"), + ttl_seconds=int(data["ttl_seconds"]) if data.get("ttl_seconds") is not None else None, + metadata=_metadata_dict(data.get("metadata", {})), + allow_missing_session=as_bool(data.get("allow_missing_session", False), name="field_cache.allow_missing_session"), + max_values=int(data.get("max_values", 1024)), + max_bytes=int(data.get("max_bytes", 4 * 1024 * 1024)), + ) + except KeyError as exc: + raise ExperimentalConfigError(f"Missing field-cache rule key {exc.args[0]}") from exc + except ValueError as exc: + raise ExperimentalConfigError(str(exc)) from exc + + +def _metadata_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return dict(value) + raise ExperimentalConfigError("field_cache.metadata must be an object") + + +def _env_part(value: str) -> str: + return re.sub(r"[^A-Z0-9]+", "_", value.upper()).strip("_") diff --git a/src/rotator_library/cooldown_manager.py b/src/rotator_library/cooldown_manager.py index 0d1bb63ee..a4dbafefd 100644 --- a/src/rotator_library/cooldown_manager.py +++ b/src/rotator_library/cooldown_manager.py @@ -3,9 +3,21 @@ import asyncio import time +from dataclasses import dataclass from typing import Dict +@dataclass(frozen=True) +class CooldownSnapshot: + """Read-only view of one active cooldown scope for tests/observability.""" + + provider: str + scope: str + model: str | None + remaining: float + reason: str | None = None + + class CooldownManager: """ Manages global cooldown periods for API providers to handle IP-based rate limiting. @@ -15,34 +27,115 @@ class CooldownManager: def __init__(self): self._cooldowns: Dict[str, float] = {} + self._metadata: Dict[str, dict[str, str | None]] = {} self._lock = asyncio.Lock() async def is_cooling_down(self, provider: str) -> bool: """Checks if a provider is currently in a cooldown period.""" - async with self._lock: - return ( - provider in self._cooldowns and time.time() < self._cooldowns[provider] - ) + return await self.is_scoped_cooling_down(provider, scope="provider") async def start_cooldown(self, provider: str, duration: int): """ Initiates or extends a cooldown period for a provider. The cooldown is set to the current time plus the specified duration. + + A shorter new cooldown must not shorten an existing longer cooldown; + provider-wide throttles often arrive concurrently from several requests. """ - async with self._lock: - self._cooldowns[provider] = time.time() + duration + await self.start_scoped_cooldown(provider, duration, scope="provider") async def get_cooldown_remaining(self, provider: str) -> float: """ Returns the remaining cooldown time in seconds for a provider. Returns 0 if the provider is not in a cooldown period. """ - async with self._lock: - if provider in self._cooldowns: - remaining = self._cooldowns[provider] - time.time() - return max(0, remaining) - return 0 + return await self.get_scoped_remaining(provider, scope="provider") async def get_remaining_cooldown(self, provider: str) -> float: """Backward-compatible alias for get_cooldown_remaining.""" return await self.get_cooldown_remaining(provider) + + async def start_scoped_cooldown( + self, + provider: str, + duration: int, + *, + model: str | None = None, + scope: str = "provider", + reason: str | None = None, + ) -> None: + """Start or extend a provider/model cooldown scope. + + Model scopes are intentionally separate from provider scopes because + capacity failures often belong to one model deployment rather than the + entire provider. Provider scopes remain available for provider-wide + throttles and block every model for the provider. + """ + + key = _cooldown_key(provider, scope=scope, model=model) + async with self._lock: + new_expiry = time.time() + max(0, duration) + current_expiry = self._cooldowns.get(key, 0) + if new_expiry > current_expiry: + self._cooldowns[key] = new_expiry + self._metadata[key] = {"provider": provider, "scope": scope, "model": model, "reason": reason} + + async def get_scoped_remaining(self, provider: str, *, model: str | None = None, scope: str = "provider") -> float: + """Return remaining seconds for one cooldown scope.""" + + key = _cooldown_key(provider, scope=scope, model=model) + async with self._lock: + return self._remaining_for_key(key, time.time()) + + async def get_max_remaining(self, provider: str, *, model: str | None = None) -> float: + """Return the max provider/model cooldown remaining for a request.""" + + async with self._lock: + now = time.time() + provider_remaining = self._remaining_for_key(_cooldown_key(provider, scope="provider"), now) + model_remaining = self._remaining_for_key(_cooldown_key(provider, scope="model", model=model), now) if model else 0 + return max(provider_remaining, model_remaining) + + async def is_scoped_cooling_down(self, provider: str, *, model: str | None = None, scope: str = "provider") -> bool: + """Return whether one scoped cooldown is currently active.""" + + return (await self.get_scoped_remaining(provider, model=model, scope=scope)) > 0 + + async def snapshot(self) -> tuple[CooldownSnapshot, ...]: + """Return active cooldown scopes for tests and future observability.""" + + async with self._lock: + now = time.time() + snapshots = [] + for key, expires_at in list(self._cooldowns.items()): + remaining = max(0, expires_at - now) + if remaining <= 0: + continue + metadata = self._metadata.get(key, {}) + snapshots.append( + CooldownSnapshot( + provider=str(metadata.get("provider") or key), + scope=str(metadata.get("scope") or "provider"), + model=metadata.get("model"), + remaining=remaining, + reason=metadata.get("reason"), + ) + ) + return tuple(snapshots) + + def _remaining_for_key(self, key: str, now: float) -> float: + expires_at = self._cooldowns.get(key) + if expires_at is None: + return 0 + remaining = expires_at - now + if remaining <= 0: + self._cooldowns.pop(key, None) + self._metadata.pop(key, None) + return 0 + return remaining + + +def _cooldown_key(provider: str, *, scope: str, model: str | None = None) -> str: + if scope == "model" and model: + return f"provider:{provider}:model:{model}" + return f"provider:{provider}" diff --git a/src/rotator_library/core/constants.py b/src/rotator_library/core/constants.py index 6da64c83a..1a1a4ed05 100644 --- a/src/rotator_library/core/constants.py +++ b/src/rotator_library/core/constants.py @@ -69,7 +69,7 @@ # Provider-specific providers that use request_count instead of success_count # for credential selection (because failed requests also consume quota) -REQUEST_COUNT_PROVIDERS = frozenset({"gemini_cli", "chutes", "nanogpt"}) +REQUEST_COUNT_PROVIDERS = frozenset({"chutes", "nanogpt"}) # Usage manager storage USAGE_FILE_NAME = "usage.json" # New format diff --git a/src/rotator_library/core/errors.py b/src/rotator_library/core/errors.py index 5acd9fc77..1d8b1f436 100644 --- a/src/rotator_library/core/errors.py +++ b/src/rotator_library/core/errors.py @@ -34,6 +34,7 @@ is_rate_limit_error, is_server_error, is_unrecoverable_error, + is_context_window_error_text, # Constants ABNORMAL_ERROR_TYPES, NORMAL_ERROR_TYPES, @@ -62,6 +63,156 @@ def __init__(self, message: str, data=None): self.data = data +class StructuredAPIResponseError(Exception): + """Raise a structured provider error before success-format conversion.""" + + def __init__( + self, + message: str, + *, + error_type: str, + status_code: int | None = None, + response: dict | None = None, + headers: dict | None = None, + ) -> None: + super().__init__(message) + self.error_type = error_type + self.status_code = status_code + self.response = response or {} + self.headers = headers or {} + + @property + def http_status(self) -> int: + if self.status_code and 400 <= self.status_code <= 599: + return self.status_code + return { + "authentication": 401, + "forbidden": 403, + "rate_limit": 429, + "quota_exceeded": 429, + "invalid_request": 400, + "context_window_exceeded": 400, + "server_error": 502, + "proxy_timeout": 504, + "proxy_all_credentials_exhausted": 503, + "not_found": 404, + }.get(self.error_type, 502) + + def to_protocol_payload(self, protocol: str) -> dict: + """Format one terminal provider error in the selected client protocol.""" + + message = str(self) + if protocol == "anthropic_messages": + error_type = { + "authentication": "authentication_error", + "forbidden": "permission_error", + "rate_limit": "rate_limit_error", + "quota_exceeded": "rate_limit_error", + "invalid_request": "invalid_request_error", + "context_window_exceeded": "invalid_request_error", + "not_found": "not_found_error", + }.get(self.error_type, "api_error") + return {"type": "error", "error": {"type": error_type, "message": message}} + if protocol == "gemini": + status = { + "authentication": "UNAUTHENTICATED", + "forbidden": "PERMISSION_DENIED", + "rate_limit": "RESOURCE_EXHAUSTED", + "quota_exceeded": "RESOURCE_EXHAUSTED", + "invalid_request": "INVALID_ARGUMENT", + "context_window_exceeded": "INVALID_ARGUMENT", + "not_found": "NOT_FOUND", + "proxy_timeout": "DEADLINE_EXCEEDED", + "proxy_all_credentials_exhausted": "UNAVAILABLE", + }.get(self.error_type, "INTERNAL") + return {"error": {"code": self.http_status, "message": message, "status": status}} + return { + "error": { + "message": message, + "type": self.error_type, + "code": self.error_type, + } + } + + +def structured_api_response_error( + response, + *, + headers: dict | None = None, +) -> StructuredAPIResponseError | None: + """Normalize top-level provider error envelopes across execution modes.""" + + if not isinstance(response, dict) or "error" not in response or response.get("error") in (None, "", False): + return None + value = response.get("error") + details = value if isinstance(value, dict) else {"message": str(value)} + raw_status = next( + ( + candidate + for candidate in ( + details.get("status_code"), + details.get("code"), + details.get("status"), + response.get("status_code"), + response.get("status"), + ) + if candidate is not None + ), + None, + ) + try: + status_code = int(raw_status) + except (TypeError, ValueError): + status_code = None + descriptor = " ".join( + str(details.get(key) or "") + for key in ("type", "status", "code", "message") + ).lower() + if is_context_window_error_text(descriptor): + error_type = "context_window_exceeded" + elif status_code == 429 or any(token in descriptor for token in ("rate", "quota", "resource_exhausted")): + error_type = "quota_exceeded" if "quota" in descriptor or "resource_exhausted" in descriptor else "rate_limit" + elif status_code == 401 or "auth" in descriptor or "unauthorized" in descriptor or "unauthenticated" in descriptor: + error_type = "authentication" + elif status_code == 403 or "forbidden" in descriptor or "permission_denied" in descriptor: + error_type = "forbidden" + elif (status_code is not None and status_code >= 500) or any(token in descriptor for token in ("server", "unavailable", "internal")): + error_type = "server_error" + else: + error_type = "invalid_request" + message = str(details.get("message") or details.get("status") or value or "Provider returned a structured error response") + return StructuredAPIResponseError( + message, + error_type=error_type, + status_code=status_code, + response=response, + headers=headers, + ) + + +def is_structured_error_payload(response) -> bool: + """Return whether a value is an explicit top-level API error envelope.""" + + return isinstance(response, dict) and "error" in response and response.get("error") not in (None, "", False) + + +def protocol_error_payload( + error: BaseException | str, + protocol: str, + *, + error_type: str, + status_code: int, +) -> tuple[int, dict]: + """Render one proxy-side terminal failure in the selected client protocol.""" + + structured = StructuredAPIResponseError( + str(error), + error_type=error_type, + status_code=status_code, + ) + return structured.http_status, structured.to_protocol_payload(protocol) + + __all__ = [ # Exception classes "NoAvailableKeysError", @@ -70,6 +221,10 @@ def __init__(self, message: str, data=None): "EmptyResponseError", "TransientQuotaError", "StreamedAPIError", + "StructuredAPIResponseError", + "structured_api_response_error", + "is_structured_error_payload", + "protocol_error_payload", # Error classification "ClassifiedError", "RequestErrorAccumulator", diff --git a/src/rotator_library/core/types.py b/src/rotator_library/core/types.py index c502592ec..4194b5572 100644 --- a/src/rotator_library/core/types.py +++ b/src/rotator_library/core/types.py @@ -14,9 +14,7 @@ Callable, Dict, List, - Literal, Optional, - Set, Tuple, Union, ) @@ -79,9 +77,8 @@ class RequestContext: session_id: Optional[str] = None session_affinity_key: Optional[str] = None session_tracker: Optional[Any] = None - session_possible_compaction: bool = False - session_lineage_parent_id: Optional[str] = None session_tracking_namespace: Optional[str] = None + session_isolation_key: Optional[str] = None request: Optional[Any] = None # FastAPI Request object pre_request_callback: Optional[Callable] = None transaction_logger: Optional[Any] = None @@ -89,6 +86,17 @@ class RequestContext: provider_config: Optional[Dict[str, Any]] = None credential_secrets: Dict[str, str] = field(default_factory=dict) classifier: Optional[str] = None + routing_targets: Optional[Any] = None + routing_group_name: Optional[str] = None + routing_group: Optional[Any] = None + routing_target_index: int = 0 + routing_attempt_history: List[Dict[str, Any]] = field(default_factory=list) + input_protocol_name: str = "openai_chat" + output_protocol_name: str = "openai_chat" + protocol_request: Dict[str, Any] = field(default_factory=dict) + unified_request: Optional[Any] = None + input_provider: Optional[str] = None + disable_provider_continuation: bool = False @dataclass diff --git a/src/rotator_library/credential_manager.py b/src/rotator_library/credential_manager.py index 1092e5729..e9b0f9499 100644 --- a/src/rotator_library/credential_manager.py +++ b/src/rotator_library/credential_manager.py @@ -12,17 +12,12 @@ lib_logger = logging.getLogger("rotator_library") -# Standard directories where tools like `gemini login` store credentials. -DEFAULT_OAUTH_DIRS = { - "gemini_cli": Path.home() / ".gemini", - # Add other providers like 'claude' here if they have a standard CLI path -} +# Standard directories where OAuth-enabled CLI tools store credentials. +DEFAULT_OAUTH_DIRS = {} # OAuth providers that support environment variable-based credentials # Maps provider name to the ENV_PREFIX used by the provider -ENV_OAUTH_PROVIDERS = { - "gemini_cli": "GEMINI_CLI", -} +ENV_OAUTH_PROVIDERS = {} class CredentialManager: @@ -61,11 +56,11 @@ def _discover_env_oauth_credentials(self) -> Dict[str, List[str]]: Discover OAuth credentials defined via environment variables. Supports two formats: - 1. Single credential: GEMINI_CLI_ACCESS_TOKEN + GEMINI_CLI_REFRESH_TOKEN - 2. Multiple credentials: GEMINI_CLI_1_ACCESS_TOKEN + GEMINI_CLI_1_REFRESH_TOKEN, etc. + 1. Single credential: PROVIDER_ACCESS_TOKEN + PROVIDER_REFRESH_TOKEN + 2. Multiple credentials: PROVIDER_1_ACCESS_TOKEN + PROVIDER_1_REFRESH_TOKEN, etc. Returns: - Dict mapping provider name to list of virtual paths (e.g., "env://gemini_cli/1") + Dict mapping provider name to list of virtual paths (e.g., "env://provider/1") """ env_credentials: Dict[str, Set[str]] = {} @@ -73,7 +68,6 @@ def _discover_env_oauth_credentials(self) -> Dict[str, List[str]]: found_indices: Set[str] = set() # Check for numbered credentials (PROVIDER_N_ACCESS_TOKEN pattern) - # Pattern: GEMINI_CLI_1_ACCESS_TOKEN, GEMINI_CLI_2_ACCESS_TOKEN, etc. numbered_pattern = re.compile(rf"^{env_prefix}_(\d+)_ACCESS_TOKEN$") for key in self.env_vars.keys(): diff --git a/src/rotator_library/credential_tool.py b/src/rotator_library/credential_tool.py index 8240ba961..9a9afb83d 100644 --- a/src/rotator_library/credential_tool.py +++ b/src/rotator_library/credential_tool.py @@ -26,7 +26,6 @@ get_provider_api_key_var, get_provider_display_name, ) -from .providers.utilities.gemini_shared_utils import format_tier_for_display def _get_oauth_base_dir() -> Path: @@ -61,9 +60,7 @@ def _ensure_providers_loaded(): # OAuth provider display names mapping (no "(OAuth)" suffix - context makes it clear) -OAUTH_FRIENDLY_NAMES = { - "gemini_cli": "Gemini CLI", -} +OAUTH_FRIENDLY_NAMES = {} def _extract_key_number(key_name: str) -> int: @@ -78,8 +75,10 @@ def _extract_key_number(key_name: str) -> int: return int(match.group(1)) if match else 0 -# Note: _normalize_tier_name was replaced with format_tier_for_display -# from providers.utilities.gemini_shared_utils for centralized tier handling +def format_tier_for_display(tier: str) -> str: + """Format a stored provider tier for compact credential summaries.""" + + return tier.replace("_", "-").replace("-tier", "").lower() def _count_tiers(credentials: list) -> dict: @@ -263,10 +262,10 @@ def _get_oauth_credentials_summary() -> dict: Returns: Dict mapping provider names to lists of credential info dicts. - Example: {"gemini_cli": [{"email": "user@example.com", "tier": "free-tier", ...}, ...]} + Example: {"provider": [{"email": "user@example.com", "tier": "free-tier", ...}, ...]} """ provider_factory, _ = _ensure_providers_loaded() - oauth_providers = ["gemini_cli"] + oauth_providers = provider_factory.get_available_providers() oauth_summary = {} for provider_name in oauth_providers: @@ -505,7 +504,7 @@ def _display_provider_credentials(provider_name: str): Display all credentials for a specific OAuth provider. Args: - provider_name: The provider key (e.g., "gemini_cli") + provider_name: The provider key. """ provider_factory, _ = _ensure_providers_loaded() @@ -529,22 +528,10 @@ def _display_provider_credentials(provider_name: str): table.add_column("File", style="yellow") table.add_column("Email/Identifier", style="cyan") - # Add tier/project columns for Google OAuth providers - if provider_name == "gemini_cli": - table.add_column("Tier", style="green") - table.add_column("Project", style="dim") for i, cred in enumerate(credentials, 1): file_name = Path(cred["file_path"]).name email = cred.get("email", "unknown") - - if provider_name == "gemini_cli": - tier = cred.get("tier", "-") - project = cred.get("project_id", "-") - if project and len(project) > 20: - project = project[:17] + "..." - table.add_row(str(i), file_name, email, tier or "-", project or "-") - else: - table.add_row(str(i), file_name, email) + table.add_row(str(i), file_name, email) console.print(table) console.print("") @@ -555,7 +542,7 @@ async def _edit_oauth_credential_email(provider_name: str): Edit the email field of an OAuth credential. Args: - provider_name: The provider key (e.g., "gemini_cli") + provider_name: The provider key. """ provider_factory, _ = _ensure_providers_loaded() @@ -769,24 +756,10 @@ async def _view_oauth_credentials_detail(provider_name: str): table.add_column("File", style="yellow") table.add_column("Email/Identifier", style="cyan") - # Add tier/project columns for Google OAuth providers - if provider_name == "gemini_cli": - table.add_column("Tier", style="green") - table.add_column("Project", style="dim") for i, cred in enumerate(credentials, 1): file_name = Path(cred["file_path"]).name email = cred.get("email", "unknown") - - if provider_name == "gemini_cli": - tier = ( - format_tier_for_display(cred.get("tier")) if cred.get("tier") else "-" - ) - project = cred.get("project_id", "-") - if project and len(project) > 25: - project = project[:22] + "..." - table.add_row(str(i), file_name, email, tier, project or "-") - else: - table.add_row(str(i), file_name, email) + table.add_row(str(i), file_name, email) console.print(table) console.print(f"\n[dim]Total: {len(credentials)} credential(s)[/dim]") @@ -1032,8 +1005,8 @@ async def _edit_oauth_credential_menu(): Panel( Text.from_markup( "[bold yellow]Warning:[/bold yellow] Editing OAuth credentials is generally not recommended.\n" - "For Gemini CLI OAuth credentials, the email is automatically\n" - "retrieved during authentication and changing it may cause confusion." + "OAuth provider metadata is usually retrieved during authentication,\n" + "and changing it manually may cause confusion." ), style="yellow", title="Edit OAuth Credential", @@ -1188,9 +1161,7 @@ async def setup_api_key(): litellm_api_keys.add(api_key_var) # OAuth-only providers to exclude entirely from API key setup - oauth_only_providers = { - "gemini_cli", # OAuth-only - } + oauth_only_providers = set() # Base classes to exclude base_classes = { @@ -1703,10 +1674,7 @@ async def setup_new_credential(provider_name: str): auth_instance = auth_class() # Build display name for better user experience - oauth_friendly_names = { - "gemini_cli": "Gemini CLI (OAuth)", - } - display_name = oauth_friendly_names.get( + display_name = OAUTH_FRIENDLY_NAMES.get( provider_name, provider_name.replace("_", " ").title() ) @@ -1760,101 +1728,6 @@ async def setup_new_credential(provider_name: str): ) -async def export_gemini_cli_to_env(): - """ - Export a Gemini CLI credential JSON file to .env format. - Uses the auth class's build_env_lines() and list_credentials() methods. - """ - clear_screen("Export Gemini CLI Credential") - - # Get auth instance for this provider - provider_factory, _ = _ensure_providers_loaded() - auth_class = provider_factory.get_provider_auth_class("gemini_cli") - auth_instance = auth_class() - - # List available credentials using auth class - credentials = auth_instance.list_credentials(_get_oauth_base_dir()) - - if not credentials: - console.print( - Panel( - "No Gemini CLI credentials found. Please add one first using 'Add OAuth Credential'.", - style="bold red", - title="No Credentials", - ) - ) - return - - # Display available credentials - cred_text = Text() - for i, cred_info in enumerate(credentials): - cred_text.append( - f" {i + 1}. {Path(cred_info['file_path']).name} ({cred_info['email']})\n" - ) - - console.print( - Panel( - cred_text, - title="Available Gemini CLI Credentials", - style="bold blue", - ) - ) - - choice = Prompt.ask( - Text.from_markup( - "[bold]Please select a credential to export or type [red]'b'[/red] to go back[/bold]" - ), - choices=[str(i + 1) for i in range(len(credentials))] + ["b"], - show_choices=False, - ) - - if choice.lower() == "b": - return - - try: - choice_index = int(choice) - 1 - if 0 <= choice_index < len(credentials): - cred_info = credentials[choice_index] - - # Use auth class to export - env_path = auth_instance.export_credential_to_env( - cred_info["file_path"], _get_oauth_base_dir() - ) - - if env_path: - numbered_prefix = f"GEMINI_CLI_{cred_info['number']}" - success_text = Text.from_markup( - f"Successfully exported credential to [bold yellow]'{Path(env_path).name}'[/bold yellow]\n\n" - f"[bold]Environment variable prefix:[/bold] [cyan]{numbered_prefix}_*[/cyan]\n\n" - f"[bold]To use this credential:[/bold]\n" - f"1. Copy the contents to your main .env file, OR\n" - f"2. Source it: [bold cyan]source {Path(env_path).name}[/bold cyan] (Linux/Mac)\n" - f"3. Or on Windows: [bold cyan]Get-Content {Path(env_path).name} | ForEach-Object {{ $_ -replace '^([^#].*)$', 'set $1' }} | cmd[/bold cyan]\n\n" - f"[bold]To combine multiple credentials:[/bold]\n" - f"Copy lines from multiple .env files into one file.\n" - f"Each credential uses a unique number ({numbered_prefix}_*)." - ) - console.print(Panel(success_text, style="bold green", title="Success")) - else: - console.print( - Panel( - "Failed to export credential", style="bold red", title="Error" - ) - ) - else: - console.print("[bold red]Invalid choice. Please try again.[/bold red]") - except ValueError: - console.print( - "[bold red]Invalid input. Please enter a number or 'b'.[/bold red]" - ) - except Exception as e: - console.print( - Panel( - f"An error occurred during export: {e}", style="bold red", title="Error" - ) - ) - - async def export_all_provider_credentials(provider_name: str): """ Export all credentials for a specific provider to individual .env files. @@ -2018,10 +1891,8 @@ async def combine_all_credentials(): """ clear_screen("Combine All Credentials") - # List of providers that support OAuth credentials - oauth_providers = ["gemini_cli"] - provider_factory, _ = _ensure_providers_loaded() + oauth_providers = provider_factory.get_available_providers() combined_lines = [ "# Combined All Provider Credentials", @@ -2116,18 +1987,25 @@ async def export_credentials_submenu(): while True: clear_screen("Export Credentials") + provider_factory, _ = _ensure_providers_loaded() + oauth_providers = provider_factory.get_available_providers() + if not oauth_providers: + console.print( + Panel( + "No active OAuth credential providers are available.", + style="bold yellow", + title="No OAuth Providers", + ) + ) + console.print("\n[dim]Press Enter to go back...[/dim]") + input() + return + console.print( Panel( Text.from_markup( - "[bold]Individual Exports:[/bold]\n" - "1. Export Gemini CLI credential\n" - "\n" - "[bold]Bulk Exports (per provider):[/bold]\n" - "2. Export ALL Gemini CLI credentials\n" - "\n" "[bold]Combine Credentials:[/bold]\n" - "3. Combine all Gemini CLI into one file\n" - "4. Combine ALL providers into one file" + "1. Combine ALL providers into one file" ), title="Choose export option", style="bold blue", @@ -2138,36 +2016,14 @@ async def export_credentials_submenu(): Text.from_markup( "[bold]Please select an option or type [red]'b'[/red] to go back[/bold]" ), - choices=[ - "1", - "2", - "3", - "4", - "b", - ], + choices=["1", "b"], show_choices=False, ) if export_choice.lower() == "b": break - # Individual exports if export_choice == "1": - await export_gemini_cli_to_env() - console.print("\n[dim]Press Enter to return to export menu...[/dim]") - input() - # Bulk exports (all credentials for a provider) - elif export_choice == "2": - await export_all_provider_credentials("gemini_cli") - console.print("\n[dim]Press Enter to return to export menu...[/dim]") - input() - # Combine per provider - elif export_choice == "3": - await combine_provider_credentials("gemini_cli") - console.print("\n[dim]Press Enter to return to export menu...[/dim]") - input() - # Combine all providers - elif export_choice == "4": await combine_all_credentials() console.print("\n[dim]Press Enter to return to export menu...[/dim]") input() @@ -2227,6 +2083,17 @@ async def main(clear_on_start=True): provider_factory, _ = _ensure_providers_loaded() available_providers = provider_factory.get_available_providers() + if not available_providers: + console.print( + Panel( + "No active OAuth credential providers are available.", + style="bold yellow", + title="No OAuth Providers", + ) + ) + console.print("\n[dim]Press Enter to return to main menu...[/dim]") + input() + continue provider_text = Text() for i, provider in enumerate(available_providers): diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index e0130c17c..0359e49dd 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -23,6 +23,23 @@ lib_logger = logging.getLogger("rotator_library") +_CONTEXT_WINDOW_ERROR_PATTERNS = ( + "context_length", + "context length", + "max_tokens", + "token limit", + "context window", + "too many tokens", + "too long", +) + + +def is_context_window_error_text(value: Any) -> bool: + """Return whether provider text describes an oversized context request.""" + + normalized = str(value or "").lower() + return any(pattern in normalized for pattern in _CONTEXT_WINDOW_ERROR_PATTERNS) + def _parse_duration_string(duration_str: str) -> Optional[int]: """ @@ -90,7 +107,7 @@ def extract_retry_after_from_body(error_body: Optional[str]) -> Optional[int]: Extract the retry-after time from an API error response body. Handles various error formats including: - - Gemini CLI: "Your quota will reset after 39s." + - Provider text: "Your quota will reset after 39s." - Google RPC: "quota will reset after 156h14m36s" - Generic: "quota will reset after 120s", "retry after 60s" @@ -167,7 +184,7 @@ class EmptyResponseError(Exception): Treated as a transient server-side issue (503 equivalent). Attributes: - provider: The provider name (e.g., "gemini_cli") + provider: The provider name (e.g., "openai") model: The model that was requested message: Human-readable message about the error """ @@ -193,7 +210,7 @@ class TransientQuotaError(Exception): Treated as a transient server-side issue (503 equivalent), same as EmptyResponseError. Attributes: - provider: The provider name (e.g., "gemini_cli") + provider: The provider name (e.g., "openai") model: The model that was requested message: Human-readable message about the error """ @@ -680,6 +697,25 @@ def get_retry_after(error: Exception) -> Optional[int]: except (ValueError, TypeError): pass + # Structured native errors retain decoded bodies and headers without needing + # to expose transport objects to protocol formatters. + structured_response = getattr(error, "response", None) + if isinstance(structured_response, dict): + try: + result = _extract_retry_from_json_body(json.dumps(structured_response)) + if result is not None: + return result + except (TypeError, ValueError): + pass + structured_headers = getattr(error, "headers", None) + if isinstance(structured_headers, dict): + retry_header = structured_headers.get("retry-after") or structured_headers.get("Retry-After") + if retry_header is not None: + try: + return max(0, int(retry_header)) + except (TypeError, ValueError): + pass + # 1. Try to parse JSON from the error string representation # Some exceptions embed JSON in their string representation error_str = str(error) @@ -759,12 +795,26 @@ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedEr if isinstance(e, dict): payload = e.get("error", e) if isinstance(payload, dict): + details = payload.get("details") if isinstance(payload.get("details"), dict) else {} code = payload.get("code") + status_value = payload.get("status_code") or payload.get("status") or details.get("status_code") or details.get("status") or code status = str(payload.get("status", "")).upper() try: - status_code = int(code) if code is not None else None + status_code = int(status_value) if status_value is not None else None except (TypeError, ValueError): status_code = None + structured_type = _classify_structured_error_text(payload, details) + if status_code == 400: + return ClassifiedError(error_type=structured_type or "invalid_request", original_exception=e, status_code=status_code) + if status_code == 401: + return ClassifiedError(error_type="authentication", original_exception=e, status_code=status_code) + if status_code == 403: + return ClassifiedError(error_type="forbidden", original_exception=e, status_code=status_code) + if status_code == 429: + body = str(payload).lower() + return ClassifiedError(error_type="quota_exceeded" if "quota" in body or "resource_exhausted" in body else "rate_limit", original_exception=e, status_code=status_code) + if structured_type: + return ClassifiedError(error_type=structured_type, original_exception=e, status_code=status_code) if (status_code is not None and status_code >= 500) or status in { "INTERNAL", "UNAVAILABLE", @@ -775,6 +825,15 @@ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedEr status_code=status_code or 503, ) + explicit_error_type = getattr(e, "error_type", None) + if explicit_error_type: + return ClassifiedError( + error_type=str(explicit_error_type).strip().lower().replace("-", "_").replace(" ", "_"), + original_exception=e, + status_code=getattr(e, "status_code", None), + retry_after=get_retry_after(e), + ) + error_text = str(e) error_type_name = type(e).__name__ if ( @@ -904,17 +963,7 @@ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedEr ) if status_code == 400: # Check for context window / token limit errors with more specific patterns - if any( - pattern in error_body - for pattern in [ - "context_length", - "max_tokens", - "token limit", - "context window", - "too many tokens", - "too long", - ] - ): + if is_context_window_error_text(error_body): return ClassifiedError( error_type="context_window_exceeded", original_exception=e, @@ -925,11 +974,6 @@ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedEr original_exception=e, status_code=status_code, ) - return ClassifiedError( - error_type="invalid_request", - original_exception=e, - status_code=status_code, - ) if 400 <= status_code < 500: # Other 4xx errors - generally client errors return ClassifiedError( @@ -1073,6 +1117,13 @@ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedEr status_code=status_code or 503, ) + if "model_capacity_exhausted" in error_text.lower() or "model capacity" in error_text.lower() or "capacity exhausted" in error_text.lower(): + return ClassifiedError( + error_type="server_error", + original_exception=e, + status_code=status_code or 503, + ) + # Fallback for any other unclassified errors return ClassifiedError( error_type="unknown", original_exception=e, status_code=status_code @@ -1084,6 +1135,34 @@ def is_rate_limit_error(e: Exception) -> bool: return isinstance(e, RateLimitError) +def _classify_structured_error_text(payload: dict, details: dict) -> Optional[str]: + """Classify structured error type/code/status text without raw messages.""" + + values = [payload.get("type"), payload.get("code"), payload.get("status"), details.get("error_type"), details.get("classification"), details.get("status")] + normalized = {str(value).strip().lower().replace("-", "_").replace(" ", "_") for value in values if value not in (None, "")} + if normalized & {"authentication", "auth", "unauthorized", "invalid_api_key"}: + return "authentication" + if normalized & {"forbidden", "permission_denied", "access_denied"}: + return "forbidden" + if normalized & {"context_window_exceeded", "context_length", "context_length_exceeded", "too_many_tokens"}: + return "context_window_exceeded" + if normalized & {"invalid_request", "bad_request", "validation", "invalid_argument"}: + return "invalid_request" + if normalized & {"credential_reauth_needed"}: + return "credential_reauth_needed" + if normalized & {"configuration_error", "config", "configuration"}: + return "configuration_error" + if normalized & {"rate_limit", "rate_limited", "too_many_requests"}: + return "rate_limit" + if normalized & {"quota_exceeded", "resource_exhausted", "quota"}: + return "quota_exceeded" + if normalized & {"server_error", "internal", "unavailable"}: + return "server_error" + if normalized & {"api_connection", "network", "connection"}: + return "api_connection" + return None + + def is_server_error(e: Exception) -> bool: """Checks if the exception is a temporary server-side error.""" return isinstance( diff --git a/src/rotator_library/field_cache/__init__.py b/src/rotator_library/field_cache/__init__.py new file mode 100644 index 000000000..ae5d62aa2 --- /dev/null +++ b/src/rotator_library/field_cache/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Configurable provider field-cache rules and helpers.""" + +from .paths import FieldCachePathError, extract_path, inject_path, parse_path +from .engine import FieldCacheEngine, FieldCacheOperation, build_cache_key +from .store import FieldCacheStore, InMemoryFieldCacheStore, ProviderCacheFieldStore +from .types import FieldCacheContext, FieldCacheInjection, FieldCacheRule + +__all__ = [ + "FieldCacheContext", + "FieldCacheEngine", + "FieldCacheInjection", + "FieldCacheOperation", + "FieldCachePathError", + "FieldCacheRule", + "FieldCacheStore", + "InMemoryFieldCacheStore", + "ProviderCacheFieldStore", + "build_cache_key", + "extract_path", + "inject_path", + "parse_path", +] diff --git a/src/rotator_library/field_cache/engine.py b/src/rotator_library/field_cache/engine.py new file mode 100644 index 000000000..776a5437e --- /dev/null +++ b/src/rotator_library/field_cache/engine.py @@ -0,0 +1,587 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Extraction and injection engine for field-cache rules.""" + +from __future__ import annotations + +import hashlib +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any, Iterable, Optional + +from .paths import FieldCachePathError, extract_path, inject_path, parse_path +from .store import ( + FieldCacheStore, + InMemoryFieldCacheStore, + _bounded_append_values, + _bounded_set_value, +) +from .types import FieldCacheContext, FieldCacheRule + + +@dataclass +class FieldCacheOperation: + """Summary of one field-cache rule application.""" + + rule_name: str + cache_key: Optional[str] + matched: int = 0 + changed: bool = False + hit: bool = False + skipped: bool = False + reason: Optional[str] = None + sample_values: list[Any] = field(default_factory=list) + + +def _safe_scope_value(value: str) -> str: + return hashlib.sha256(value.encode("utf-8")).hexdigest() + + +def _unsupported_store_keyword(error: TypeError) -> bool: + message = str(error).lower() + return "unexpected keyword" in message or "keyword argument" in message + + +def _shared_cache_signature(rule: FieldCacheRule) -> tuple[Any, ...]: + """Return settings that must agree for source-specific cache counterparts.""" + + injection = rule.inject + behavior_metadata = tuple( + (key, repr(rule.metadata.get(key))) + for key in ( + "provider_continuation", + "tool_call_id_path", + "tool_container_path", + "tool_value_path", + "inject_tool_call_id_path", + "turn_container_path", + "turn_role_path", + "turn_value_path", + ) + ) + return ( + rule.mode, + rule.scope, + rule.ttl_seconds, + rule.allow_missing_session, + rule.max_values, + rule.max_bytes, + injection.target if injection else None, + injection.path if injection else None, + injection.when_missing_only if injection else None, + injection.insert if injection else None, + injection.as_list if injection else None, + behavior_metadata, + ) + + +def build_cache_key(rule: FieldCacheRule, context: FieldCacheContext) -> Optional[str]: + """Build a scoped cache key or return None when required scope is absent.""" + + parts = [f"rule={_safe_scope_value(rule.cache_key or rule.name)}"] + for scope in rule.scope: + value = context.value_for_scope(scope) + if value is None or value == "": + if scope == "credential": + return None + if scope == "session" and not rule.allow_missing_session: + return None + value = "_none" + safe_value = _safe_scope_value(value) + parts.append(f"{scope}={safe_value}") + return "|".join(parts) + + +class FieldCacheEngine: + """Apply field-cache extraction and injection rules. + + The engine preserves provider protocol state; it is not session tracking. + It defaults to copying payloads before injection so providers can opt into + mutation explicitly. Turn and tool-call modes skip safely when the requested + context cannot be inferred rather than silently falling back to `last`. + """ + + def __init__(self, rules: Iterable[FieldCacheRule], store: Optional[FieldCacheStore] = None) -> None: + self.rules = tuple(rules) + self.store = store or InMemoryFieldCacheStore() + self._validate_rules() + + def _validate_rules(self) -> None: + names: set[str] = set() + shared_keys: dict[str, tuple[Any, ...]] = {} + for rule in self.rules: + if rule.name in names: + raise ValueError(f"Duplicate field-cache rule name: {rule.name}") + names.add(rule.name) + parse_path(rule.path) + if rule.inject: + parse_path(rule.inject.path) + if rule.mode == "per_tool_call" and not rule.metadata.get("tool_call_id_path"): + raise ValueError("per_tool_call field-cache mode requires metadata.tool_call_id_path") + if rule.cache_key: + signature = _shared_cache_signature(rule) + previous = shared_keys.get(rule.cache_key) + if previous is not None and previous != signature: + raise ValueError( + f"Field-cache rules sharing cache_key {rule.cache_key!r} must use identical mode, scope, TTL, injection, and correlation behavior" + ) + shared_keys[rule.cache_key] = signature + + async def extract( + self, + source: str, + payload: Any, + context: FieldCacheContext, + *, + transaction_logger: Optional[Any] = None, + ) -> list[FieldCacheOperation]: + operations: list[FieldCacheOperation] = [] + rules = self._rules_for_source(source) + self._trace_summary(transaction_logger, "field_cache_extraction_start", payload, source=source, target=None, rules=rules, operations=operations) + for rule in rules: + operation = FieldCacheOperation(rule_name=rule.name, cache_key=build_cache_key(rule, context)) + self._trace(transaction_logger, "before_field_cache_extraction", payload, rule, operation, source=source) + if not operation.cache_key: + operation.skipped = True + operation.reason = "missing_required_scope" + operations.append(operation) + self._trace(transaction_logger, "after_field_cache_extraction", payload, rule, operation, source=source) + continue + try: + values = extract_path(payload, rule.path) + operation.matched = len(values) + operation.sample_values = _sample_values(values) + if values or rule.mode in {"last_user_turn", "last_assistant_turn", "per_tool_call"}: + operation.changed = await self._store_values(rule, operation.cache_key, values, payload, operation) + except Exception as exc: + self._log_error(transaction_logger, "field_cache_extract", exc, payload, rule) + raise + operations.append(operation) + self._trace(transaction_logger, "after_field_cache_extraction", payload, rule, operation, source=source) + self._trace_summary(transaction_logger, "field_cache_extraction_complete", payload, source=source, target=None, rules=rules, operations=operations) + return operations + + async def inject( + self, + target: str, + payload: Any, + context: FieldCacheContext, + *, + transaction_logger: Optional[Any] = None, + mutate: bool = False, + ) -> tuple[Any, list[FieldCacheOperation]]: + updated = payload if mutate else deepcopy(payload) + operations: list[FieldCacheOperation] = [] + rules = self._rules_for_injection(target) + self._trace_summary(transaction_logger, "field_cache_injection_start", updated, source=None, target=target, rules=rules, operations=operations) + for rule in rules: + operation = FieldCacheOperation(rule_name=rule.name, cache_key=build_cache_key(rule, context)) + if not rule.inject: + continue + self._trace(transaction_logger, "before_field_cache_injection", updated, rule, operation, target=target) + if not operation.cache_key: + operation.skipped = True + operation.reason = "missing_required_scope" + operations.append(operation) + self._trace(transaction_logger, "after_field_cache_injection", updated, rule, operation, target=target) + continue + try: + cached = await self.store.get(operation.cache_key) + if cached is None: + operation.reason = "cache_miss" + operations.append(operation) + self._trace(transaction_logger, "after_field_cache_injection", updated, rule, operation, target=target) + continue + operation.hit = True + value = self._injection_value(rule, cached, updated, context, operation) + if operation.skipped: + operations.append(operation) + self._trace(transaction_logger, "after_field_cache_injection", updated, rule, operation, target=target) + continue + operation.changed = inject_path( + updated, + rule.inject.path, + value, + when_missing_only=rule.inject.when_missing_only, + insert=rule.inject.insert, + ) + operation.sample_values = _sample_values(value if isinstance(value, list) else [value]) + except Exception as exc: + self._log_error(transaction_logger, "field_cache_inject", exc, updated, rule) + raise + operations.append(operation) + self._trace(transaction_logger, "after_field_cache_injection", updated, rule, operation, target=target) + self._trace_summary(transaction_logger, "field_cache_injection_complete", updated, source=None, target=target, rules=rules, operations=operations) + return updated, operations + + def _rules_for_source(self, source: str) -> list[FieldCacheRule]: + return [rule for rule in self.rules if rule.enabled and rule.source == source] + + def _rules_for_injection(self, target: str) -> list[FieldCacheRule]: + return [rule for rule in self.rules if rule.enabled and rule.inject and rule.inject.target == target] + + async def _store_values(self, rule: FieldCacheRule, cache_key: str, values: list[Any], payload: Any, operation: FieldCacheOperation) -> bool: + if rule.mode == "all": + await self._store_append( + cache_key, + values, + ttl_seconds=rule.ttl_seconds, + max_values=rule.max_values, + max_bytes=rule.max_bytes, + ) + return True + if rule.mode == "last": + await self._store_set( + cache_key, + _wrap_cached_value(values[-1]), + ttl_seconds=rule.ttl_seconds, + max_values=None, + max_bytes=rule.max_bytes, + trim_collections=False, + ) + return True + if rule.mode in {"last_user_turn", "last_assistant_turn"}: + role = "user" if rule.mode == "last_user_turn" else "assistant" + turn_values = _turn_values(rule, payload, role) + if not turn_values: + operation.skipped = True + operation.reason = "turn_context_not_found" + return False + operation.matched = len(turn_values) + operation.sample_values = _sample_values(turn_values) + await self._store_set( + cache_key, + _wrap_cached_value(turn_values[-1]), + ttl_seconds=rule.ttl_seconds, + max_values=None, + max_bytes=rule.max_bytes, + trim_collections=False, + ) + return True + if rule.mode == "per_tool_call": + stored = _tool_call_values(rule, payload, values) + if not stored: + operation.skipped = True + operation.reason = "tool_call_id_not_found" + return False + operation.matched = len(stored) + operation.sample_values = _sample_values(list(stored.values())) + await self._store_set( + cache_key, + stored, + ttl_seconds=rule.ttl_seconds, + max_values=rule.max_values, + max_bytes=rule.max_bytes, + trim_collections=True, + ) + return True + raise ValueError(f"Unsupported field-cache mode: {rule.mode}") + + async def _store_set( + self, + cache_key: str, + value: Any, + *, + ttl_seconds: Optional[int], + max_values: Optional[int], + max_bytes: Optional[int], + trim_collections: bool, + ) -> None: + bounded = _bounded_set_value( + value, + max_values=max_values, + max_bytes=max_bytes, + trim_collections=trim_collections, + ) + try: + await self.store.set(cache_key, bounded, ttl_seconds=ttl_seconds) + except TypeError as exc: + if not _unsupported_store_keyword(exc): + raise + # Preserve compatibility with simple injected stores that implement + # the original set(key, value) shape. TTL is best-effort there. + await self.store.set(cache_key, bounded) + + async def _store_append( + self, + cache_key: str, + values: list[Any], + *, + ttl_seconds: Optional[int], + max_values: Optional[int], + max_bytes: Optional[int], + ) -> None: + try: + await self.store.append( + cache_key, + values, + ttl_seconds=ttl_seconds, + max_values=max_values, + max_bytes=max_bytes, + ) + except TypeError as exc: + if not _unsupported_store_keyword(exc): + raise + current = await self.store.get(cache_key) + if not isinstance(current, list): + current = [] + bounded = _bounded_append_values( + current, + values, + max_values=max_values, + max_bytes=max_bytes, + ) + try: + await self.store.set( + cache_key, + bounded, + ttl_seconds=ttl_seconds, + ) + except TypeError as set_exc: + if not _unsupported_store_keyword(set_exc): + raise + await self.store.set(cache_key, bounded) + + def _injection_value(self, rule: FieldCacheRule, cached: Any, payload: Any, context: FieldCacheContext, operation: FieldCacheOperation) -> Any: + """Select the cached value to inject for the rule's mode. + + Per-tool-call maps require a current tool-call ID so the engine never + injects an arbitrary provider signature into the wrong tool result. + """ + + if rule.mode == "per_tool_call": + if not isinstance(cached, dict): + operation.skipped = True + operation.reason = "invalid_tool_call_cache" + return None + ids = _injection_tool_ids(rule, payload, context) + if not ids: + operation.skipped = True + operation.reason = "tool_call_id_not_found" + return None + matches = [_unwrap_cached_value(cached[str(tool_id)]) for tool_id in ids if str(tool_id) in cached] + if not matches: + operation.skipped = True + operation.reason = "tool_call_cache_miss" + return None + if rule.inject and rule.inject.as_list: + return matches + if len(matches) == 1: + return matches[0] + operation.skipped = True + operation.reason = "ambiguous_tool_call_values" + return None + if rule.mode == "all": + return cached if isinstance(cached, list) else [cached] + if rule.inject and rule.inject.as_list: + unwrapped = _unwrap_cached_value(cached) + return unwrapped if isinstance(unwrapped, list) else [unwrapped] + return _unwrap_cached_value(cached) + + def _trace( + self, + transaction_logger: Optional[Any], + pass_name: str, + payload: Any, + rule: FieldCacheRule, + operation: FieldCacheOperation, + **extra_metadata: Any, + ) -> None: + if not transaction_logger: + return + transaction_logger.log_transform_pass( + pass_name, + _payload_shape(payload), + direction=_trace_direction(pass_name, rule.source, extra_metadata), + stage="adapter", + metadata={ + "rule_name": rule.name, + "source": rule.source, + "path": rule.path, + "mode": rule.mode, + "scope": list(rule.scope), + "cache_key": operation.cache_key, + "matched": operation.matched, + "changed": operation.changed, + "hit": operation.hit, + "skipped": operation.skipped, + "reason": operation.reason, + # Cached fields can include provider signatures or session keys. + # Trace only shape/count metadata; keep raw samples out of logs. + "sample_value_count": len(operation.sample_values), + "sample_value_types": [type(value).__name__ for value in operation.sample_values[:3]], + **extra_metadata, + }, + snapshot=rule.source != "stream_event", + ) + + def _trace_summary( + self, + transaction_logger: Optional[Any], + pass_name: str, + payload: Any, + *, + source: Optional[str], + target: Optional[str], + rules: list[FieldCacheRule], + operations: list[FieldCacheOperation], + ) -> None: + """Record cache-pass boundaries even when no individual rule matches.""" + + if not transaction_logger: + return + transaction_logger.log_transform_pass( + pass_name, + _payload_shape(payload), + direction="request" if target or source == "request" else "response" if source == "response" else "stream" if source == "stream_event" else "metadata", + stage="adapter", + metadata={ + "source": source, + "target": target, + "rule_count": len(rules), + "operation_count": len(operations), + "matched_count": sum(operation.matched for operation in operations), + "changed_count": sum(1 for operation in operations if operation.changed), + "hit_count": sum(1 for operation in operations if operation.hit), + "skipped_count": sum(1 for operation in operations if operation.skipped), + }, + snapshot=(source != "stream_event"), + ) + + def _log_error(self, transaction_logger: Optional[Any], pass_name: str, error: BaseException, payload: Any, rule: FieldCacheRule) -> None: + if not transaction_logger: + return + transaction_logger.log_transform_error( + pass_name, + error, + payload=_payload_shape(payload), + stage="adapter", + metadata={"rule_name": rule.name, "path": rule.path, "mode": rule.mode}, + ) + + +def _last_value(value: Any) -> Any: + if isinstance(value, list): + return value[-1] if value else None + return value + + +def _wrap_cached_value(value: Any) -> dict[str, Any]: + """Wrap one extracted value so list-valued fields stay intact on injection.""" + + return {"__field_cache_value__": True, "value": deepcopy(value)} + + +def _unwrap_cached_value(value: Any) -> Any: + if isinstance(value, dict) and value.get("__field_cache_value__") is True: + return deepcopy(value.get("value")) + return _last_value(value) + + +def _turn_values(rule: FieldCacheRule, payload: Any, role: str) -> list[Any]: + """Return values from the latest turn matching `role`. + + Rules can provide explicit turn paths for provider-specific payloads. The + default handles the common `messages[*]` shape used by OpenAI-compatible and + Responses-like requests. + """ + + container_path = rule.metadata.get("turn_container_path", "messages") + role_path = rule.metadata.get("turn_role_path", "role") + value_path = rule.metadata.get("turn_value_path") or _message_relative_path(rule.path, container_path) + if not value_path: + return [] + turns = extract_path(payload, str(container_path)) + if len(turns) == 1 and isinstance(turns[0], list): + turns = turns[0] + latest: list[Any] = [] + for turn in turns: + roles = extract_path(turn, str(role_path)) + if roles and str(roles[0]) == role: + values = extract_path(turn, str(value_path)) + if values: + latest = values + return latest + + +def _message_relative_path(path: str, container_path: str) -> Optional[str]: + prefixes = (f"{container_path}.*.", f"{container_path}[-1].") + for prefix in prefixes: + if path.startswith(prefix): + return path[len(prefix) :] + return None + + +def _tool_call_values(rule: FieldCacheRule, payload: Any, values: list[Any]) -> dict[str, Any]: + """Correlate cached values to provider tool-call IDs.""" + + container_path = rule.metadata.get("tool_container_path") + tool_id_path = rule.metadata.get("tool_call_id_path") + tool_value_path = rule.metadata.get("tool_value_path") + stored: dict[str, Any] = {} + if container_path and tool_value_path: + containers = extract_path(payload, str(container_path)) + if len(containers) == 1 and isinstance(containers[0], list): + containers = containers[0] + for container in containers: + tool_ids = extract_path(container, str(tool_id_path)) if tool_id_path else [] + tool_values = extract_path(container, str(tool_value_path)) + if tool_ids and tool_values: + stored[str(tool_ids[0])] = _wrap_cached_value(tool_values[-1]) + return stored + for value in values: + tool_ids = extract_path(value, str(tool_id_path)) if tool_id_path else [] + if tool_ids: + stored[str(tool_ids[0])] = _wrap_cached_value(value) + return stored + + +def _injection_tool_ids(rule: FieldCacheRule, payload: Any, context: FieldCacheContext) -> list[str]: + configured = context.metadata.get("tool_call_id") + if configured: + return [str(configured)] + inject_path_value = rule.metadata.get("inject_tool_call_id_path") + if inject_path_value: + return [str(value) for value in extract_path(payload, str(inject_path_value))] + return [] + + +def _trace_direction(pass_name: str, source: str, metadata: dict[str, Any]) -> str: + if "injection" in pass_name: + target = metadata.get("target") + if target in {"stream_event", "unified_stream_event"}: + return "stream" + if target in {"request", "unified_request", "metadata"}: + return "request" + return "response" + if source in {"stream_event", "unified_stream_event"}: + return "stream" + if source in {"request", "unified_request"}: + return "request" + return "response" + + +def _sample_values(values: list[Any], *, max_items: int = 3, max_text: int = 500) -> list[Any]: + samples: list[Any] = [] + for value in values[:max_items]: + if isinstance(value, str) and len(value) > max_text: + samples.append(f"{value[:max_text]}...") + else: + samples.append(value) + return samples + + +def _payload_shape(payload: Any) -> dict[str, Any]: + """Return non-sensitive payload shape metadata for cache traces. + + Cache rules often target provider signatures, session IDs, and other opaque + state. Logging full payloads would expose exactly the fields the cache is + designed to preserve, so field-cache traces record structure only. + """ + + if isinstance(payload, dict): + return {"payload_type": "dict", "keys": sorted(str(key) for key in payload.keys())[:20]} + if isinstance(payload, list): + return {"payload_type": "list", "length": len(payload)} + return {"payload_type": type(payload).__name__} diff --git a/src/rotator_library/field_cache/paths.py b/src/rotator_library/field_cache/paths.py new file mode 100644 index 000000000..1dfc29203 --- /dev/null +++ b/src/rotator_library/field_cache/paths.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Small JSON-path-like selector used by field-cache rules.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + + +class FieldCachePathError(ValueError): + """Raised when a field-cache path is malformed or cannot be injected.""" + + +@dataclass(frozen=True) +class PathToken: + kind: Literal["key", "index", "wildcard"] + value: str | int | None = None + + +def parse_path(path: str) -> tuple[PathToken, ...]: + """Parse a minimal dotted path with indexes and wildcards. + + Supported examples: `choices.0.message.content`, `choices.*.message`, + `messages[-1].reasoning_content`. Escaping is intentionally unsupported in + Phase 3 so malformed provider configs fail early and clearly. + """ + + if not path or path.startswith(".") or path.endswith(".") or ".." in path: + raise FieldCachePathError(f"Malformed field-cache path: {path!r}") + tokens: list[PathToken] = [] + for segment in path.split("."): + if not segment: + raise FieldCachePathError(f"Malformed field-cache path: {path!r}") + _parse_segment(segment, tokens, path) + return tuple(tokens) + + +def _parse_segment(segment: str, tokens: list[PathToken], full_path: str) -> None: + if segment == "*": + tokens.append(PathToken("wildcard")) + return + cursor = 0 + while cursor < len(segment): + bracket = segment.find("[", cursor) + if bracket == -1: + value = segment[cursor:] + if value: + if value.lstrip("-").isdigit(): + tokens.append(PathToken("index", int(value))) + else: + tokens.append(PathToken("key", value)) + return + if bracket > cursor: + tokens.append(PathToken("key", segment[cursor:bracket])) + close = segment.find("]", bracket) + if close == -1: + raise FieldCachePathError(f"Unclosed index in field-cache path: {full_path!r}") + raw_index = segment[bracket + 1 : close] + if raw_index == "*": + tokens.append(PathToken("wildcard")) + else: + try: + tokens.append(PathToken("index", int(raw_index))) + except ValueError as exc: + raise FieldCachePathError(f"Invalid index {raw_index!r} in field-cache path: {full_path!r}") from exc + cursor = close + 1 + + +def extract_path(payload: Any, path: str) -> list[Any]: + """Return every value matching path in stable traversal order.""" + + tokens = parse_path(path) + current = [payload] + for token in tokens: + next_values: list[Any] = [] + for value in current: + next_values.extend(_extract_token(value, token)) + current = next_values + if not current: + break + return current + + +def _extract_token(value: Any, token: PathToken) -> list[Any]: + if token.kind == "key": + if isinstance(value, dict) and token.value in value: + return [value[token.value]] + return [] + if token.kind == "index": + if isinstance(value, list) and value: + index = int(token.value) + if -len(value) <= index < len(value): + return [value[index]] + return [] + if token.kind == "wildcard": + if isinstance(value, dict): + return list(value.values()) + if isinstance(value, list): + return list(value) + return [] + raise FieldCachePathError(f"Unknown path token: {token}") + + +def inject_path(payload: Any, path: str, injected_value: Any, *, when_missing_only: bool = False, insert: bool = False) -> bool: + """Inject a value at a simple path, creating dict containers as needed. + + Wildcard injection is rejected because creating multiple branches can be + ambiguous and provider-specific. List indexes must already exist; this keeps + mutation predictable for message-tail use cases like `messages[-1].field`. + `insert=True` is intentionally limited to final list-index tokens so rules + cannot accidentally create provider-specific list structures. + """ + + tokens = parse_path(path) + if any(token.kind == "wildcard" for token in tokens): + raise FieldCachePathError("Wildcard injection is not supported") + if not isinstance(payload, dict): + raise FieldCachePathError("Field-cache injection root must be a dict") + current = payload + for index, token in enumerate(tokens): + is_last = index == len(tokens) - 1 + if token.kind == "key": + if not isinstance(current, dict): + raise FieldCachePathError(f"Cannot inject key {token.value!r} into non-dict value") + key = str(token.value) + if is_last: + if insert: + raise FieldCachePathError("insert=True requires a final list index token") + if when_missing_only and key in current: + return False + changed = current.get(key) != injected_value + current[key] = injected_value + return changed + if key not in current or current[key] is None: + current[key] = [] if tokens[index + 1].kind == "index" else {} + current = current[key] + continue + if token.kind == "index": + if not isinstance(current, list) or (not current and not (is_last and insert)): + raise FieldCachePathError("Cannot inject into missing or empty list") + list_index = int(token.value) + if is_last and insert: + if list_index < 0: + list_index = max(0, len(current) + list_index) + if not (0 <= list_index <= len(current)): + raise FieldCachePathError(f"List index out of range for field-cache insertion: {list_index}") + if when_missing_only: + return False + current.insert(list_index, injected_value) + return True + if not (-len(current) <= list_index < len(current)): + raise FieldCachePathError(f"List index out of range for field-cache injection: {list_index}") + if is_last: + if when_missing_only and current[list_index] is not None: + return False + changed = current[list_index] != injected_value + current[list_index] = injected_value + return changed + current = current[list_index] + continue + raise FieldCachePathError(f"Unsupported injection token: {token}") + return False diff --git a/src/rotator_library/field_cache/store.py b/src/rotator_library/field_cache/store.py new file mode 100644 index 000000000..5a31c7f31 --- /dev/null +++ b/src/rotator_library/field_cache/store.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Async stores for field-cache values.""" + +from __future__ import annotations + +import asyncio +import json +import time +from copy import deepcopy +from typing import Any, Callable, Protocol + +from ..protocols import serialize_value + +_TTL_ENVELOPE_MARKER = "__llm_proxy_field_cache_ttl_v1__" +_APPEND_LOCKS: dict[int, asyncio.Lock] = {} + + +class FieldCacheStore(Protocol): + """Minimal async store interface used by `FieldCacheEngine`.""" + + async def get(self, key: str) -> Any: ... + + async def set(self, key: str, value: Any, *, ttl_seconds: int | None = None) -> None: ... + + async def append( + self, + key: str, + values: list[Any], + *, + ttl_seconds: int | None = None, + max_values: int | None = None, + max_bytes: int | None = None, + ) -> list[Any]: ... + + async def clear(self) -> None: ... + + +class InMemoryFieldCacheStore: + """Simple process-local store with optional per-key TTL. + + This is the default native runtime store. It intentionally persists only for + the Python process and avoids a database while still preserving protocol + state across requests handled by the same executor instance. + """ + + def __init__( + self, + *, + clock: Callable[[], float] | None = None, + max_entries: int = 10_000, + ) -> None: + self._values: dict[str, Any] = {} + self._expires_at: dict[str, float] = {} + self._last_access: dict[str, float] = {} + self._clock = clock or time.monotonic + self._max_entries = max(1, int(max_entries)) + + async def get(self, key: str) -> Any: + if self._is_expired(key): + self._remove(key) + return None + if key in self._values: + self._last_access[key] = self._clock() + return deepcopy(self._values.get(key)) + + async def set(self, key: str, value: Any, *, ttl_seconds: int | None = None) -> None: + self._prune() + self._values[key] = deepcopy(value) + self._last_access[key] = self._clock() + self._set_expiry(key, ttl_seconds) + self._trim() + + async def append( + self, + key: str, + values: list[Any], + *, + ttl_seconds: int | None = None, + max_values: int | None = None, + max_bytes: int | None = None, + ) -> list[Any]: + current = await self.get(key) + if not isinstance(current, list): + current = [] + current = _bounded_append_values( + deepcopy(current), + deepcopy(values), + max_values=max_values, + max_bytes=max_bytes, + ) + self._values[key] = current + self._last_access[key] = self._clock() + self._set_expiry(key, ttl_seconds) + self._trim() + return deepcopy(current) + + async def clear(self) -> None: + self._values.clear() + self._expires_at.clear() + self._last_access.clear() + + def _set_expiry(self, key: str, ttl_seconds: int | None) -> None: + if ttl_seconds is None or ttl_seconds <= 0: + self._expires_at.pop(key, None) + return + self._expires_at[key] = self._clock() + ttl_seconds + + def _is_expired(self, key: str) -> bool: + expires_at = self._expires_at.get(key) + return expires_at is not None and expires_at <= self._clock() + + def _prune(self) -> None: + for key in tuple(self._expires_at): + if self._is_expired(key): + self._remove(key) + + def _trim(self) -> None: + overflow = len(self._values) - self._max_entries + if overflow <= 0: + return + oldest = sorted(self._values, key=lambda key: (self._last_access.get(key, 0.0), key)) + for key in oldest[:overflow]: + self._remove(key) + + def _remove(self, key: str) -> None: + self._values.pop(key, None) + self._expires_at.pop(key, None) + self._last_access.pop(key, None) + + +class ProviderCacheFieldStore: + """Field-cache store backed by an injected `ProviderCache` instance. + + The wrapper does not create `ProviderCache` itself because that class starts + background async tasks during initialization. Providers or later config code + should own that lifecycle and pass an initialized cache here. + """ + + def __init__(self, provider_cache: Any) -> None: + self._cache = provider_cache + self._append_lock = _shared_append_lock(provider_cache) + + async def get(self, key: str) -> Any: + raw = await self._cache.retrieve_async(key) + return _decode_provider_cache_value(raw) + + async def set(self, key: str, value: Any, *, ttl_seconds: int | None = None) -> None: + await self._cache.store_async( + key, + _encode_provider_cache_value(value, ttl_seconds=ttl_seconds), + ) + + async def append( + self, + key: str, + values: list[Any], + *, + ttl_seconds: int | None = None, + max_values: int | None = None, + max_bytes: int | None = None, + ) -> list[Any]: + if hasattr(self._cache, "update_async"): + result: list[Any] = [] + + def update(raw: str | None) -> str: + current = _decode_provider_cache_value(raw) + if not isinstance(current, list): + current = [] + bounded = _bounded_append_values( + current, + serialize_value(values), + max_values=max_values, + max_bytes=max_bytes, + ) + result[:] = bounded + return _encode_provider_cache_value( + bounded, + ttl_seconds=ttl_seconds, + ) + + await self._cache.update_async(key, update) + return result + async with self._append_lock: + current = await self.get(key) + if not isinstance(current, list): + current = [] + current = _bounded_append_values( + current, + serialize_value(values), + max_values=max_values, + max_bytes=max_bytes, + ) + await self.set(key, current, ttl_seconds=ttl_seconds) + return current + + async def clear(self) -> None: + await self._cache.clear() + + +def _shared_append_lock(provider_cache: Any) -> asyncio.Lock: + lock = getattr(provider_cache, "_field_cache_append_lock", None) + if isinstance(lock, asyncio.Lock): + return lock + lock = asyncio.Lock() + try: + setattr(provider_cache, "_field_cache_append_lock", lock) + except (AttributeError, TypeError): + lock = _APPEND_LOCKS.setdefault(id(provider_cache), lock) + return lock + + +def _decode_provider_cache_value(raw: str | None) -> Any: + if raw is None: + return None + value = json.loads(raw) + if isinstance(value, dict) and value.get(_TTL_ENVELOPE_MARKER) is True: + expires_at = value.get("expires_at") + if isinstance(expires_at, (int, float)) and expires_at <= time.time(): + return None + return value.get("value") + return value + + +def _encode_provider_cache_value(value: Any, *, ttl_seconds: int | None) -> str: + payload = serialize_value(value) + if ttl_seconds is not None and ttl_seconds > 0: + payload = { + _TTL_ENVELOPE_MARKER: True, + "expires_at": time.time() + ttl_seconds, + "value": payload, + } + return json.dumps(payload, ensure_ascii=False) + + +def _bounded_append_values( + current: list[Any], + values: list[Any], + *, + max_values: int | None, + max_bytes: int | None, +) -> list[Any]: + """Append newest state while enforcing provider-rule memory bounds.""" + + combined = current + values + if max_values is not None and len(combined) > max_values: + combined = combined[-max_values:] + if max_bytes is not None: + while combined and len(json.dumps(serialize_value(combined), ensure_ascii=False).encode("utf-8")) > max_bytes: + if len(combined) == 1: + raise ValueError("Field-cache value exceeds max_bytes") + combined.pop(0) + return combined + + +def _bounded_set_value( + value: Any, + *, + max_values: int | None, + max_bytes: int | None, + trim_collections: bool, +) -> Any: + """Bound scalar or correlated-map state before replacing a cache value.""" + + bounded = deepcopy(value) + if trim_collections and max_values is not None: + if isinstance(bounded, dict) and len(bounded) > max_values: + bounded = dict(list(bounded.items())[-max_values:]) + elif isinstance(bounded, list) and len(bounded) > max_values: + bounded = bounded[-max_values:] + if max_bytes is None: + return bounded + while len(json.dumps(serialize_value(bounded), ensure_ascii=False).encode("utf-8")) > max_bytes: + if not trim_collections: + raise ValueError("Field-cache value exceeds max_bytes") + if isinstance(bounded, dict) and len(bounded) > 1: + bounded.pop(next(iter(bounded))) + continue + if isinstance(bounded, list) and len(bounded) > 1: + bounded.pop(0) + continue + raise ValueError("Field-cache value exceeds max_bytes") + return bounded diff --git a/src/rotator_library/field_cache/types.py b/src/rotator_library/field_cache/types.py new file mode 100644 index 000000000..ea325ba10 --- /dev/null +++ b/src/rotator_library/field_cache/types.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Data types for provider field-cache rules.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal, Optional + +FieldCacheSource = Literal[ + "request", + "response", + "stream_event", + "unified_request", + "unified_response", + "unified_stream_event", +] +FieldCacheTarget = Literal["request", "unified_request", "metadata"] +FieldCacheMode = Literal["last", "all", "last_user_turn", "last_assistant_turn", "per_tool_call"] +FieldCacheScope = Literal["provider", "model", "credential", "session", "conversation", "classifier"] + +DEFAULT_SCOPE: tuple[FieldCacheScope, ...] = ("provider", "model", "credential", "session") +REQUIRED_PROVIDER_STATE_SCOPE = frozenset({"provider", "model", "credential", "session"}) +_VALID_SOURCES = {"request", "response", "stream_event", "unified_request", "unified_response", "unified_stream_event"} +_VALID_TARGETS = {"request", "unified_request", "metadata"} +_VALID_SCOPES = {"provider", "model", "credential", "session", "conversation", "classifier"} + + +@dataclass(frozen=True) +class FieldCacheInjection: + """Where and how a cached value should be injected into a later payload.""" + + target: FieldCacheTarget + path: str + when_missing_only: bool = False + insert: bool = False + as_list: bool = False + + +@dataclass(frozen=True) +class FieldCacheRule: + """Declarative rule for extracting and re-injecting provider state. + + Rules are protocol/provider extensions, not session-affinity logic. Session + tracking decides continuity; field-cache rules preserve protocol state such + as reasoning content, thought signatures, prompt cache keys, and response IDs. + """ + + name: str + source: FieldCacheSource + path: str + mode: FieldCacheMode = "last" + scope: tuple[FieldCacheScope, ...] = DEFAULT_SCOPE + inject: Optional[FieldCacheInjection] = None + enabled: bool = True + ttl_seconds: Optional[int] = None + metadata: dict[str, Any] = field(default_factory=dict) + allow_missing_session: bool = False + cache_key: Optional[str] = None + max_values: Optional[int] = 1024 + max_bytes: Optional[int] = 4 * 1024 * 1024 + + def __post_init__(self) -> None: + if not self.name or any(char in self.name for char in "/\\:"): + raise ValueError("FieldCacheRule.name must be non-empty and filesystem-safe") + if self.cache_key is not None and ( + not self.cache_key or any(char in self.cache_key for char in "/\\:") + ): + raise ValueError("FieldCacheRule.cache_key must be non-empty and filesystem-safe") + if self.mode not in {"last", "all", "last_user_turn", "last_assistant_turn", "per_tool_call"}: + raise ValueError(f"Unsupported field-cache mode: {self.mode}") + if self.source not in _VALID_SOURCES: + raise ValueError(f"Unsupported field-cache source: {self.source}") + if self.inject and self.inject.target not in _VALID_TARGETS: + raise ValueError(f"Unsupported field-cache injection target: {self.inject.target}") + if not self.scope: + raise ValueError("FieldCacheRule.scope must contain at least one dimension") + invalid_scopes = [scope for scope in self.scope if scope not in _VALID_SCOPES] + if invalid_scopes: + raise ValueError(f"Unsupported field-cache scope: {invalid_scopes[0]}") + if self.max_values is not None and self.max_values <= 0: + raise ValueError("FieldCacheRule.max_values must be positive") + if self.max_bytes is not None and self.max_bytes <= 0: + raise ValueError("FieldCacheRule.max_bytes must be positive") + if self.inject and is_provider_continuation_path(self.inject.path): + if self.metadata.get("provider_continuation") is not True: + raise ValueError( + "Continuation field-cache injection requires metadata.provider_continuation=true" + ) + + +def is_provider_continuation_path(path: str) -> bool: + normalized = path.replace("[", ".").replace("]", "") + leaf = normalized.rsplit(".", 1)[-1] + leaf = "".join(character for character in leaf.lower() if character.isalnum()) + return leaf in { + "previousresponseid", + "providerresponseid", + "continuationid", + "conversationid", + } + + +@dataclass(frozen=True) +class FieldCacheContext: + """Scope values used to isolate cached provider fields.""" + + provider: Optional[str] = None + model: Optional[str] = None + credential_id: Optional[str] = None + session_id: Optional[str] = None + conversation_id: Optional[str] = None + classifier: Optional[str] = None + metadata: dict[str, Any] = field(default_factory=dict) + + def value_for_scope(self, scope: FieldCacheScope) -> Optional[str]: + if scope == "provider": + return self.provider + if scope == "model": + return self.model + if scope == "credential": + return self.credential_id + if scope == "session": + return self.session_id + if scope == "conversation": + return self.conversation_id + if scope == "classifier": + return self.classifier + raise ValueError(f"Unsupported field-cache scope: {scope}") diff --git a/src/rotator_library/model_info_service.py b/src/rotator_library/model_info_service.py index e9c86b153..721426c63 100644 --- a/src/rotator_library/model_info_service.py +++ b/src/rotator_library/model_info_service.py @@ -81,7 +81,6 @@ # PROVIDER_ALIASES = { "nvidia_nim": ["nvidia"], - "gemini_cli": ["google"], "gemini": ["google"], } @@ -1115,7 +1114,6 @@ def _get_alias_candidates(self, model_id: str) -> List[str]: Examples: nvidia_nim/mistralai/model -> nvidia/mistralai/model - gemini_cli/gemini-2.5-flash -> google/gemini-2.5-flash gemini/gemini-2.5-pro -> google/gemini-2.5-pro """ parts = model_id.split("/") diff --git a/src/rotator_library/native_provider/__init__.py b/src/rotator_library/native_provider/__init__.py new file mode 100644 index 000000000..dc63572ad --- /dev/null +++ b/src/rotator_library/native_provider/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Opt-in native provider execution helpers.""" + +from .context import NativeProviderContext +from .executor import NativeProviderExecutor +from .http import NativeHTTPTransport +from .streaming import stream_event_payload + +__all__ = ["NativeHTTPTransport", "NativeProviderContext", "NativeProviderExecutor", "stream_event_payload"] diff --git a/src/rotator_library/native_provider/context.py b/src/rotator_library/native_provider/context.py new file mode 100644 index 000000000..9b03b4f85 --- /dev/null +++ b/src/rotator_library/native_provider/context.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Context objects for native provider execution.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Optional + +from ..adapters import AdapterContext +from ..field_cache import FieldCacheContext, FieldCacheRule +from ..protocols import ProtocolContext + + +@dataclass +class NativeProviderContext: + """Metadata needed to execute a provider through native protocol helpers. + + This context intentionally mirrors the trace, adapter, protocol, and + field-cache contexts so provider-native execution can remain opt-in and + testable without changing the existing LiteLLM-backed path. + """ + + provider: str + model: str + protocol_name: str + endpoint: str + operation: str = "chat" + input_protocol_name: Optional[str] = None + output_protocol_name: Optional[str] = None + headers: dict[str, str] = field(default_factory=dict) + credential_id: Optional[str] = None + session_id: Optional[str] = None + scope_key: Optional[str] = None + classifier: Optional[str] = None + transport: str = "http" + adapter_names: tuple[str, ...] = () + adapter_config: dict[str, dict[str, Any]] = field(default_factory=dict) + field_cache_rules: tuple[FieldCacheRule, ...] = () + transaction_logger: Optional[Any] = None + metadata: dict[str, Any] = field(default_factory=dict) + request_preparer: Optional[Callable[..., dict[str, Any]]] = None + request_validator: Optional[Callable[..., Any]] = None + + def protocol_context( + self, + *, + source_protocol: Optional[str] = None, + target_protocol: Optional[str] = None, + source_provider: Optional[str] = None, + target_provider: Optional[str] = None, + provider_state_compatible: bool = False, + ) -> ProtocolContext: + """Build a protocol context for parse/build/format passes.""" + + input_protocol = self.input_protocol_name or self.protocol_name + output_protocol = self.output_protocol_name or input_protocol + return ProtocolContext( + provider=self.provider, + model=self.model, + source_protocol=source_protocol or input_protocol, + target_protocol=target_protocol or self.protocol_name, + input_protocol=input_protocol, + provider_protocol=self.protocol_name, + output_protocol=output_protocol, + source_provider=source_provider, + target_provider=target_provider, + provider_state_compatible=provider_state_compatible, + session_id=self.session_id, + credential_stable_id=self.credential_id, + transport=self.transport, + provider_options={"operation": self.operation}, + metadata={"operation": self.operation, **dict(self.metadata)}, + ) + + def adapter_context(self) -> AdapterContext: + """Build an adapter context for provider payload adapters.""" + + return AdapterContext( + provider=self.provider, + model=self.model, + protocol=self.protocol_name, + credential_id=self.credential_id, + session_id=self.session_id, + scope_key=self.scope_key, + classifier=self.classifier, + transport=self.transport, + metadata={"operation": self.operation, **dict(self.metadata)}, + adapter_config=dict(self.adapter_config), + transaction_logger=self.transaction_logger, + ) + + def field_cache_context(self) -> FieldCacheContext: + """Build a field-cache context with provider isolation metadata.""" + + return FieldCacheContext( + provider=self.provider, + model=self.model, + credential_id=self.credential_id, + session_id=self.session_id, + conversation_id=self.scope_key, + classifier=self.classifier, + metadata={"operation": self.operation, **dict(self.metadata)}, + ) diff --git a/src/rotator_library/native_provider/executor.py b/src/rotator_library/native_provider/executor.py new file mode 100644 index 000000000..6602a3542 --- /dev/null +++ b/src/rotator_library/native_provider/executor.py @@ -0,0 +1,626 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Opt-in executor for provider-native protocol calls.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import replace +from typing import Any, AsyncGenerator + +from ..adapters import get_adapter, run_adapter_chain +from ..field_cache import FieldCacheEngine, InMemoryFieldCacheStore +from ..field_cache.types import is_provider_continuation_path +from ..core.errors import StreamedAPIError, structured_api_response_error +from ..field_cache.paths import FieldCachePathError, PathToken, parse_path +from ..protocols import ProtocolError, get_protocol, serialize_value +from ..protocols.types import UnifiedRequest +from ..transform_trace import REDACTED +from ..usage.accounting import extract_usage_record +from ..usage.costs import CostCalculator +from .context import NativeProviderContext +from .http import NativeHTTPTransport +from .streaming import stream_event_payload + + +class NativeProviderExecutor: + """Run one native provider request through protocol/adapter/cache passes. + + The default field-cache store is process-local per executor. That preserves + provider protocol state across native requests without adding a database; + production callers can still inject a persistent store when needed. + """ + + def __init__(self, *, field_cache_store: Any = None) -> None: + self.field_cache_store = field_cache_store or InMemoryFieldCacheStore() + + async def execute(self, raw_request: dict[str, Any] | UnifiedRequest, context: NativeProviderContext, transport: NativeHTTPTransport) -> dict[str, Any]: + """Execute a non-streaming native provider request.""" + + logger = context.transaction_logger + provider_protocol = get_protocol(context.protocol_name) + input_protocol = get_protocol(context.input_protocol_name or context.protocol_name) + output_protocol = get_protocol(context.output_protocol_name or context.input_protocol_name or context.protocol_name) + context = _without_provider_continuation_rules(context) + self._ensure_supported_operation(provider_protocol, context) + self._trace(context, "native_protocol_selected", {"input_protocol": input_protocol.name, "provider_protocol": provider_protocol.name, "output_protocol": output_protocol.name}, direction="metadata", stage="protocol") + try: + self._trace(context, "raw_native_client_request", raw_request, direction="request", stage="client") + cache_engine = FieldCacheEngine(context.field_cache_rules, store=self.field_cache_store) + context = await self._inject_metadata(context, cache_engine) + input_context = context.protocol_context( + source_protocol=input_protocol.name, + target_protocol=provider_protocol.name, + source_provider=context.metadata.get("input_provider"), + target_provider=context.provider, + ) + unified_request = deepcopy(raw_request) if isinstance(raw_request, UnifiedRequest) else input_protocol.parse_request(raw_request, input_context) + unified_request.model = context.model + self._trace(context, "parsed_native_unified_request", unified_request, direction="request", stage="protocol") + await cache_engine.extract("unified_request", serialize_value(unified_request), context.field_cache_context(), transaction_logger=logger) + self._trace(context, "after_unified_request_field_cache_extraction", {"source": "unified_request"}, direction="request", stage="adapter", snapshot=False) + request_before_injection = unified_request + unified_request = await self._inject_unified_request(unified_request, context, cache_engine) + provider_state_compatible = unified_request is not request_before_injection + provider_context = context.protocol_context( + source_protocol=input_protocol.name, + target_protocol=provider_protocol.name, + source_provider=context.metadata.get("input_provider"), + target_provider=context.provider, + provider_state_compatible=provider_state_compatible, + ) + provider_request = provider_protocol.build_request(unified_request, provider_context) + self._trace(context, "built_native_provider_request", provider_request, direction="request", stage="protocol") + provider_request = self._prepare_provider_request(provider_request, context) + self._trace(context, "provider_native_request_prepared", provider_request, direction="request", stage="provider") + adapters = [get_adapter(name) for name in context.adapter_names] + adapter_context = context.adapter_context() + adapter_context.transaction_logger = None + provider_request = await run_adapter_chain(adapters, provider_request, adapter_context, stage="request") + self._trace(context, "after_request_adapter_chain", provider_request, direction="request", stage="adapter") + provider_request, _ = await cache_engine.inject( + "request", + provider_request, + context.field_cache_context(), + transaction_logger=logger, + ) + self._trace(context, "after_field_cache_injection", provider_request, direction="request", stage="adapter") + await cache_engine.extract("request", provider_request, context.field_cache_context(), transaction_logger=logger) + self._trace(context, "after_request_field_cache_extraction", {"source": "request"}, direction="request", stage="adapter", snapshot=False) + await self._validate_provider_request(provider_request, context) + self._trace(context, "native_provider_request", provider_request, direction="request", stage="provider") + raw_response = await transport.post_json(context.endpoint, headers=context.headers, payload=provider_request) + self._trace(context, "raw_native_provider_response", raw_response, direction="response", stage="provider") + structured_error = structured_api_response_error(raw_response) + if structured_error: + raise structured_error + await cache_engine.extract("response", raw_response, context.field_cache_context(), transaction_logger=logger) + self._trace(context, "after_response_field_cache_extraction", {"source": "response", "payload": "raw_provider_response"}, direction="response", stage="adapter", snapshot=False) + response_context = context.protocol_context( + source_protocol=provider_protocol.name, + target_protocol=output_protocol.name, + source_provider=context.provider, + target_provider=None, + provider_state_compatible=False, + ) + unified_response = provider_protocol.parse_response(raw_response, response_context) + unified_response.model = str(context.metadata.get("public_model") or unified_response.model or context.model) + self._trace(context, "parsed_native_unified_response", unified_response, direction="response", stage="protocol") + await cache_engine.extract("unified_response", serialize_value(unified_response), context.field_cache_context(), transaction_logger=logger) + self._trace(context, "after_unified_response_field_cache_extraction", {"source": "unified_response"}, direction="response", stage="adapter", snapshot=False) + self._trace(context, "native_response_protocol_selected", {"protocol": output_protocol.name}, direction="metadata", stage="protocol", snapshot=False) + client_response = output_protocol.format_response(unified_response, response_context) + self._trace(context, "formatted_native_response", client_response, direction="response", stage="protocol") + adapter_context = context.adapter_context() + adapter_context.transaction_logger = None + client_response = await run_adapter_chain(adapters, client_response, adapter_context, stage="response") + self._trace(context, "after_response_adapter_chain", client_response, direction="response", stage="adapter") + usage_record = extract_usage_record( + client_response, + provider=context.provider, + model=context.model, + source="native_provider_response", + ) + raw_usage_record = extract_usage_record( + raw_response, + provider=context.provider, + model=context.model, + source="native_provider_raw_response", + ) + if usage_record.provider_reported_cost is None and raw_usage_record.provider_reported_cost is not None: + usage_record = replace( + usage_record, + provider_reported_cost=raw_usage_record.provider_reported_cost, + cost_currency=raw_usage_record.cost_currency, + cost_source=raw_usage_record.cost_source, + ) + cost_breakdown = CostCalculator().calculate(usage_record, model=context.model, provider=context.provider) + self._trace( + context, + "usage_accounting_summary", + {"usage": usage_record.to_dict(), "cost": cost_breakdown.to_dict()}, + direction="metadata", + stage="final", + snapshot=False, + ) + self._trace(context, "final_client_response", client_response, direction="response", stage="final") + return client_response + except Exception as exc: + if logger: + logger.log_transform_error( + "native_provider_execute", + exc, + payload=raw_request, + stage="provider", + protocol=context.protocol_name, + metadata={"provider": context.provider, "model": context.model}, + ) + raise + + async def stream(self, raw_request: dict[str, Any] | UnifiedRequest, context: NativeProviderContext, transport: NativeHTTPTransport) -> AsyncGenerator[Any, None]: + """Execute a streaming native provider request and yield client events.""" + + logger = context.transaction_logger + protocol = get_protocol(context.protocol_name) + input_protocol = get_protocol(context.input_protocol_name or context.protocol_name) + output_protocol = get_protocol(context.output_protocol_name or context.input_protocol_name or context.protocol_name) + context = _without_provider_continuation_rules(context) + self._ensure_supported_operation(protocol, context) + self._trace(context, "native_protocol_selected", {"input_protocol": input_protocol.name, "provider_protocol": protocol.name, "output_protocol": output_protocol.name}, direction="metadata", stage="protocol") + try: + self._trace(context, "raw_native_client_request", raw_request, direction="request", stage="client") + cache_engine = FieldCacheEngine(context.field_cache_rules, store=self.field_cache_store) + context = await self._inject_metadata(context, cache_engine) + input_context = context.protocol_context( + source_protocol=input_protocol.name, + target_protocol=protocol.name, + source_provider=context.metadata.get("input_provider"), + target_provider=context.provider, + ) + if isinstance(raw_request, UnifiedRequest): + unified_request = deepcopy(raw_request) + else: + request_payload = dict(raw_request) + request_payload["stream"] = True + unified_request = input_protocol.parse_request(request_payload, input_context) + unified_request.model = context.model + unified_request.stream = True + self._trace(context, "parsed_native_unified_request", unified_request, direction="request", stage="protocol") + await cache_engine.extract("unified_request", serialize_value(unified_request), context.field_cache_context(), transaction_logger=logger) + self._trace(context, "after_unified_request_field_cache_extraction", {"source": "unified_request"}, direction="request", stage="adapter", snapshot=False) + request_before_injection = unified_request + unified_request = await self._inject_unified_request(unified_request, context, cache_engine) + provider_context = context.protocol_context( + source_protocol=input_protocol.name, + target_protocol=protocol.name, + source_provider=context.metadata.get("input_provider"), + target_provider=context.provider, + provider_state_compatible=( + unified_request is not request_before_injection + or ( + input_protocol.name == protocol.name + and context.metadata.get("input_provider") == context.provider + ) + ), + ) + provider_request = protocol.build_request(unified_request, provider_context) + self._trace(context, "built_native_provider_request", provider_request, direction="request", stage="protocol") + provider_request = self._prepare_provider_request(provider_request, context) + self._trace(context, "provider_native_request_prepared", provider_request, direction="request", stage="provider") + adapters = [get_adapter(name) for name in context.adapter_names] + adapter_context = context.adapter_context() + adapter_context.transaction_logger = None + provider_request = await run_adapter_chain(adapters, provider_request, adapter_context, stage="request") + self._trace(context, "after_request_adapter_chain", provider_request, direction="request", stage="adapter") + provider_request, _ = await cache_engine.inject( + "request", + provider_request, + context.field_cache_context(), + transaction_logger=logger, + ) + self._trace(context, "after_field_cache_injection", provider_request, direction="request", stage="adapter") + await cache_engine.extract("request", provider_request, context.field_cache_context(), transaction_logger=logger) + self._trace(context, "after_request_field_cache_extraction", {"source": "request"}, direction="request", stage="adapter", snapshot=False) + await self._validate_provider_request(provider_request, context) + self._trace(context, "native_provider_stream_request", provider_request, direction="request", stage="provider") + usage_record = extract_usage_record(None, provider=context.provider, model=context.model, source="native_provider_stream") + response_context = context.protocol_context( + source_protocol=protocol.name, + target_protocol=output_protocol.name, + source_provider=context.provider, + target_provider=None, + provider_state_compatible=False, + ) + async for raw_chunk in transport.stream_json_lines(context.endpoint, headers=context.headers, payload=provider_request): + self._trace(context, "raw_native_provider_stream_chunk", raw_chunk, direction="stream", stage="provider") + event = protocol.parse_stream_event(raw_chunk, response_context) + self._trace(context, "parsed_native_unified_stream_event", event, direction="stream", stage="protocol", snapshot=False) + if event.type == "error" or event.error is not None: + error = event.error if isinstance(event.error, dict) else {"message": str(event.error or "Provider stream failed")} + raise StreamedAPIError( + str(error.get("message") or error.get("type") or "Provider stream failed"), + data={"error": deepcopy(error)}, + ) + usage_record = _merge_stream_usage_records( + usage_record, + extract_usage_record(serialize_value(event), provider=context.provider, model=context.model, source="native_stream_event"), + extract_usage_record(raw_chunk, provider=context.provider, model=context.model, source="native_raw_stream_event"), + ) + if event.type == "done": + event_payload = stream_event_payload(event) + self._trace(context, "parsed_native_stream_event", event_payload, direction="stream", stage="protocol") + formatted = output_protocol.format_stream_event(event, response_context) + for frame in _formatted_stream_frames(formatted): + self._trace(context, "formatted_client_stream_event", frame, direction="stream", stage="final", snapshot=False) + yield frame + break + adapter_context = context.adapter_context() + # Native stream traces apply field-cache path redaction below. + # Suppress generic adapter-chain snapshots here so provider state + # cannot leak before rule-aware redaction runs. + adapter_context.transaction_logger = None + event = await run_adapter_chain(adapters, event, adapter_context, stage="stream_event") + self._trace(context, "after_stream_event_adapter_chain", event, direction="stream", stage="adapter", snapshot=False) + await cache_engine.extract("unified_stream_event", serialize_value(event), context.field_cache_context(), transaction_logger=logger) + self._trace(context, "after_unified_stream_event_field_cache_extraction", {"source": "unified_stream_event"}, direction="stream", stage="adapter", snapshot=False) + event_payload = stream_event_payload(event) + self._trace(context, "parsed_native_stream_event", event_payload, direction="stream", stage="protocol") + await cache_engine.extract("stream_event", event_payload, context.field_cache_context(), transaction_logger=logger) + self._trace( + context, + "after_field_cache_stream_extraction", + {"source": "stream_event"}, + direction="stream", + stage="adapter", + snapshot=False, + ) + formatted = output_protocol.format_stream_event(event, response_context) + for frame in _formatted_stream_frames(formatted): + self._trace(context, "formatted_client_stream_event", frame, direction="stream", stage="final", snapshot=False) + yield frame + cost_breakdown = CostCalculator().calculate(usage_record, model=context.model, provider=context.provider) + self._trace( + context, + "usage_accounting_summary", + {"usage": usage_record.to_dict(), "cost": cost_breakdown.to_dict()}, + direction="metadata", + stage="final", + snapshot=False, + ) + except Exception as exc: + if logger: + logger.log_transform_error( + "native_provider_stream", + exc, + payload=raw_request, + stage="provider", + protocol=context.protocol_name, + transport=context.transport, + metadata={"provider": context.provider, "model": context.model}, + ) + raise + + @staticmethod + def _prepare_provider_request( + provider_request: dict[str, Any], + context: NativeProviderContext, + ) -> dict[str, Any]: + """Apply narrow provider quirks after protocol-native formatting.""" + + prepared = dict(provider_request) + public_model = context.metadata.get("public_model") + if public_model: + prepared["_proxy_model"] = public_model + if context.request_preparer: + prepared = dict( + context.request_preparer( + prepared, + model=context.model, + operation=context.operation, + ) + ) + prepared.pop("_proxy_model", None) + return prepared + + @staticmethod + async def _validate_provider_request(provider_request: dict[str, Any], context: NativeProviderContext) -> None: + """Run provider validation only after the provider payload exists.""" + + validator = context.request_validator + if not callable(validator): + return + result = validator(provider_request, context.model) + if hasattr(result, "__await__"): + result = await result + if result is False: + raise ProtocolError( + f"Request validation failed for {context.provider}/{context.model}", + protocol=context.protocol_name, + pass_name="provider_validation", + ) + if isinstance(result, str): + raise ProtocolError( + result, + protocol=context.protocol_name, + pass_name="provider_validation", + ) + + @staticmethod + def _trace( + context: NativeProviderContext, + pass_name: str, + data: Any, + *, + direction: str, + stage: str, + metadata: dict[str, Any] | None = None, + snapshot: bool = True, + ) -> None: + if not context.transaction_logger: + return + context.transaction_logger.log_transform_pass( + pass_name, + _redact_field_cache_paths(data, context, direction), + direction=direction, + stage=stage, + protocol=context.protocol_name, + credential_id=context.credential_id, + transport=context.transport, + metadata={ + "provider": context.provider, + "model": context.model, + "session_id": context.session_id, + "scope_key": context.scope_key, + "classifier": context.classifier, + **(metadata or {}), + }, + snapshot=snapshot, + ) + + @staticmethod + def _ensure_supported_operation(protocol: Any, context: NativeProviderContext) -> None: + """Fail before transport when provider and protocol operations disagree.""" + + if protocol.supports_operation(context.operation): + return + raise ProtocolError( + f"provider {context.provider} requested unsupported operation {context.operation!r}", + protocol=protocol.name, + pass_name="native_operation_check", + payload={"provider": context.provider, "model": context.model, "operation": context.operation}, + ) + + async def _inject_metadata(self, context: NativeProviderContext, cache_engine: FieldCacheEngine) -> NativeProviderContext: + """Inject cached metadata before protocol/adapter contexts are built.""" + + metadata, operations = await cache_engine.inject( + "metadata", + dict(context.metadata), + context.field_cache_context(), + transaction_logger=context.transaction_logger, + ) + if operations: + self._trace(context, "after_metadata_field_cache_injection", metadata, direction="metadata", stage="adapter", snapshot=False) + if metadata == context.metadata: + return context + return replace(context, metadata=metadata) + + async def _inject_unified_request( + self, + unified_request: UnifiedRequest, + context: NativeProviderContext, + cache_engine: FieldCacheEngine, + ) -> UnifiedRequest: + """Inject cached values into a serialized unified request and hydrate it.""" + + serialized = serialize_value(unified_request) + injected, operations = await cache_engine.inject( + "unified_request", + serialized, + context.field_cache_context(), + transaction_logger=context.transaction_logger, + ) + if operations: + self._trace(context, "after_unified_request_field_cache_injection", injected, direction="request", stage="adapter") + if injected == serialized: + return unified_request + return _hydrate_unified_request(unified_request, injected) + + +def _merge_stream_usage_records(base: Any, event_record: Any, raw_record: Any) -> Any: + """Merge native stream usage, preserving raw provider cost when needed.""" + + selected = event_record if _usage_record_has_token_values(event_record) else base + if not _usage_record_has_token_values(selected) and _usage_record_has_token_values(raw_record): + selected = raw_record + if selected.provider_reported_cost is None and base.provider_reported_cost is not None: + selected = replace( + selected, + provider_reported_cost=base.provider_reported_cost, + cost_currency=base.cost_currency, + cost_source=base.cost_source, + ) + if selected.provider_reported_cost is None and raw_record.provider_reported_cost is not None: + selected = replace( + selected, + provider_reported_cost=raw_record.provider_reported_cost, + cost_currency=raw_record.cost_currency, + cost_source=raw_record.cost_source, + ) + return selected + + +def _without_provider_continuation_rules(context: NativeProviderContext) -> NativeProviderContext: + """Suppress upstream continuation IDs when proxy history was expanded.""" + + if not context.metadata.get("disable_provider_continuation"): + return context + rules = tuple( + rule + for rule in context.field_cache_rules + if not ( + (rule.metadata or {}).get("provider_continuation") + or (rule.inject and is_provider_continuation_path(rule.inject.path)) + ) + ) + return context if rules == context.field_cache_rules else replace(context, field_cache_rules=rules) + + +def _usage_record_has_values(record: Any) -> bool: + return bool( + record.input_tokens + or record.completion_tokens + or record.reasoning_tokens + or record.cache_read_tokens + or record.cache_write_tokens + or record.raw_total_tokens + or record.provider_reported_cost is not None + ) + + +def _usage_record_has_token_values(record: Any) -> bool: + return bool( + record.input_tokens + or record.completion_tokens + or record.reasoning_tokens + or record.cache_read_tokens + or record.cache_write_tokens + or record.raw_total_tokens + ) + + +def _redact_field_cache_paths(data: Any, context: NativeProviderContext, direction: str) -> Any: + """Redact configured cache paths before broad native payload traces. + + Field-cache rules can inject opaque state under arbitrary configured keys, + so key-based trace redaction is not enough. Native traces apply the active + rules' source and injection paths to a copy before handing data to the normal + transaction trace sanitizer. + """ + + if not context.field_cache_rules: + return data + redacted = serialize_value(deepcopy(data)) + for rule in context.field_cache_rules: + paths: list[str] = [] + if direction == "request" and rule.inject: + paths.append(rule.inject.path) + if direction == "metadata" and rule.inject and rule.inject.target == "metadata": + paths.append(rule.inject.path) + if direction in {"response", "stream"}: + paths.append(rule.path) + for path in _trace_redaction_paths(paths, direction=direction): + try: + tokens = parse_path(path) + _redact_path(redacted, tokens) + _redact_leaf_key(redacted, tokens) + except (FieldCachePathError, TypeError, ValueError): + continue + return redacted + + +def _hydrate_unified_request(original: UnifiedRequest, injected: Any) -> UnifiedRequest: + """Hydrate common unified-request fields after serialized cache injection. + + Field-cache path injection operates on JSON-like dictionaries. Protocol + builders still expect `UnifiedRequest`, so this helper copies supported + top-level fields back onto the dataclass while leaving complex message/tool + objects untouched unless providers handle them through provider-payload rules. + """ + + if not isinstance(injected, dict): + return original + safe_fields = { + "operation", + "model", + "stream", + "input", + "modalities", + "files", + "generation_params", + "response_format", + "previous_response_id", + "metadata", + "raw", + "extra", + } + values = {field_name: getattr(original, field_name) for field_name in UnifiedRequest._fields} + for field_name in safe_fields: + if field_name in injected: + values[field_name] = injected[field_name] + return UnifiedRequest(**values) + + +def _trace_redaction_paths(paths: list[str], *, direction: str) -> list[str]: + """Return configured paths plus raw-stream envelope fallbacks for traces.""" + + expanded: list[str] = [] + for path in paths: + expanded.append(path) + if direction == "stream" and path.startswith("raw."): + expanded.append(path[4:]) + return expanded + + +def _redact_path(value: Any, tokens: tuple[PathToken, ...]) -> None: + if not tokens: + return + token = tokens[0] + rest = tokens[1:] + if token.kind == "key": + if isinstance(value, dict) and token.value in value: + if rest: + _redact_path(value[token.value], rest) + else: + value[token.value] = REDACTED + return + if token.kind == "index": + if isinstance(value, list) and value: + index = int(token.value) + if -len(value) <= index < len(value): + if rest: + _redact_path(value[index], rest) + else: + value[index] = REDACTED + return + if token.kind == "wildcard": + if isinstance(value, dict): + for key in list(value.keys()): + if rest: + _redact_path(value[key], rest) + else: + value[key] = REDACTED + elif isinstance(value, list): + for index, item in enumerate(value): + if rest: + _redact_path(item, rest) + else: + value[index] = REDACTED + + +def _redact_leaf_key(value: Any, tokens: tuple[PathToken, ...]) -> None: + """Redact the configured terminal key wherever stream traces duplicate it.""" + + leaf = next((token.value for token in reversed(tokens) if token.kind == "key"), None) + if not leaf: + return + if isinstance(value, dict): + for key, item in list(value.items()): + if key == leaf: + value[key] = REDACTED + else: + _redact_leaf_key(item, tokens) + elif isinstance(value, list): + for item in value: + _redact_leaf_key(item, tokens) + + +def _formatted_stream_frames(formatted: Any) -> list[Any]: + """Normalize one adapter result into destination frames.""" + + if formatted is None: + return [] + if isinstance(formatted, list): + return formatted + return [formatted] diff --git a/src/rotator_library/native_provider/http.py b/src/rotator_library/native_provider/http.py new file mode 100644 index 000000000..16dd7a41f --- /dev/null +++ b/src/rotator_library/native_provider/http.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Small HTTP transport wrapper for native provider calls.""" + +from __future__ import annotations + +import json +from typing import Any, AsyncIterator + +from ..core.errors import StructuredAPIResponseError, structured_api_response_error + + +class NativeHTTPTransport: + """Execute provider-native JSON HTTP requests through an injected client.""" + + def __init__(self, client: Any) -> None: + self.client = client + + async def post_json(self, endpoint: str, *, headers: dict[str, str], payload: dict[str, Any]) -> Any: + """POST JSON and return a decoded response body. + + The wrapper keeps HTTP behavior easy to mock. It does not own retries or + credential rotation; those remain in the existing executor/usage layer. + """ + + response = await self.client.post(endpoint, headers=headers, json=payload) + await _raise_for_http_error(response) + if hasattr(response, "json"): + return response.json() + return response + + async def stream_json_lines(self, endpoint: str, *, headers: dict[str, str], payload: dict[str, Any]) -> AsyncIterator[Any]: + """Yield provider stream chunks from an injected streaming-capable client. + + Provider-specific test clients can still expose `stream_json_lines()`. + When a normal `httpx.AsyncClient`-style object is injected, this method + now uses `client.stream()` directly so native streaming has a real HTTP + seam without enabling any provider that has not opted in safely. + """ + + if hasattr(self.client, "stream_json_lines"): + async for chunk in self.client.stream_json_lines(endpoint, headers=headers, json=payload): + yield chunk + return + if hasattr(self.client, "stream"): + async with self.client.stream("POST", endpoint, headers=headers, json=payload) as response: + await _raise_for_http_error(response, read_stream=True) + if hasattr(response, "aiter_lines"): + decoder = _SSELineDecoder() + async for line in response.aiter_lines(): + for parsed in decoder.feed(line): + yield parsed + for parsed in decoder.flush(): + yield parsed + return + if hasattr(response, "aiter_bytes"): + buffer = "" + decoder = _SSELineDecoder() + async for chunk in response.aiter_bytes(): + text = chunk.decode("utf-8", errors="replace") if isinstance(chunk, (bytes, bytearray)) else str(chunk) + buffer += text + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + for parsed in decoder.feed(line): + yield parsed + if buffer: + for parsed in decoder.feed(buffer): + yield parsed + for parsed in decoder.flush(): + yield parsed + return + raise NotImplementedError("Injected native HTTP client does not expose streaming support") + + +async def _raise_for_http_error(response: Any, *, read_stream: bool = False) -> None: + """Preserve non-2xx provider bodies and status for retry/error formatting.""" + + status = getattr(response, "status_code", None) + if status is None: + if hasattr(response, "raise_for_status"): + response.raise_for_status() + return + try: + status_code = int(status) + except (TypeError, ValueError): + status_code = None + if status_code is None or 200 <= status_code < 300: + if hasattr(response, "raise_for_status"): + response.raise_for_status() + return + if read_stream: + reader = getattr(response, "aread", None) + if callable(reader): + await reader() + payload: Any = None + if hasattr(response, "json"): + try: + payload = response.json() + except Exception: + payload = None + if not isinstance(payload, dict): + text = getattr(response, "text", None) + payload = { + "error": { + "message": str(text or f"Provider returned HTTP {status_code}"), + "code": status_code, + } + } + payload.setdefault("status_code", status_code) + if isinstance(payload.get("error"), dict): + payload["error"].setdefault("status_code", status_code) + headers = dict(getattr(response, "headers", {}) or {}) + error = structured_api_response_error(payload, headers=headers) + if error: + raise error + raise StructuredAPIResponseError( + f"Provider returned HTTP {status_code}", + error_type="server_error" if status_code >= 500 else "invalid_request", + status_code=status_code, + response=payload, + headers=headers, + ) + + +def _parse_stream_line(line: Any) -> Any: + """Parse one HTTP streaming line while preserving provider sentinels.""" + + if line is None: + return None + text = line.decode("utf-8", errors="replace") if isinstance(line, (bytes, bytearray)) else str(line) + text = text.strip() + if not text: + return None + if text.startswith(":"): + return None + if text.startswith("data:"): + text = text[len("data:") :].strip() + if text == "[DONE]": + return "[DONE]" + try: + return json.loads(text) + except json.JSONDecodeError: + return text + + +class _SSELineDecoder: + """Assemble SSE fields until the blank-line event delimiter.""" + + def __init__(self) -> None: + self.event_name: str | None = None + self.data_lines: list[str] = [] + + def feed(self, line: Any) -> list[Any]: + text = line.decode("utf-8", errors="replace") if isinstance(line, (bytes, bytearray)) else str(line) + text = text.rstrip("\r") + if not text: + return self.flush() + if text.startswith(":"): + return [] + if text.startswith("event:"): + self.event_name = text[len("event:") :].strip() + return [] + if text.startswith("data:"): + self.data_lines.append(text[len("data:") :].lstrip()) + return [] + output = self.flush() + parsed = _parse_stream_line(text) + if parsed is not None: + output.append(parsed) + return output + + def flush(self) -> list[Any]: + if not self.data_lines: + self.event_name = None + return [] + text = "\n".join(self.data_lines).strip() + event_name = self.event_name + self.event_name = None + self.data_lines = [] + if text == "[DONE]": + return ["[DONE]"] + try: + parsed: Any = json.loads(text) + except json.JSONDecodeError: + parsed = text + if isinstance(parsed, dict) and event_name and not (parsed.get("type") or parsed.get("event")): + parsed["type"] = event_name + return [parsed] diff --git a/src/rotator_library/native_provider/streaming.py b/src/rotator_library/native_provider/streaming.py new file mode 100644 index 000000000..b043a2338 --- /dev/null +++ b/src/rotator_library/native_provider/streaming.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Streaming helpers for provider-native execution.""" + +from __future__ import annotations + +from typing import Any + +from ..streaming import StreamEvent, stream_event_from_sse_chunk + + +def stream_event_payload(event: Any) -> Any: + """Return a JSON-safe payload for stream field-cache and trace passes.""" + + if hasattr(event, "to_dict"): + return event.to_dict() + return event + + +def provider_supports_native_streaming(provider: Any, *, model: str = "", operation: str = "chat") -> bool: + """Return explicit provider native-streaming support. + + Providers must opt in. Missing methods and exceptions fail closed so routed + streaming can fallback before output rather than accidentally claiming native + streaming support. + """ + + method = getattr(provider, "supports_native_streaming", None) + if not method: + return False + try: + return bool(method(model=model, operation=operation)) + except TypeError: + return bool(method(model, operation)) + except Exception: + return False + + +def native_stream_event_from_formatted(formatted: Any, *, protocol: str = "openai_chat") -> StreamEvent: + """Convert a formatted native stream chunk into the common event seam.""" + + if isinstance(formatted, str): + return stream_event_from_sse_chunk(formatted, protocol=protocol) + return StreamEvent("parsed_chunk", protocol=protocol, data=formatted, raw=formatted) diff --git a/src/rotator_library/protocols/__init__.py b/src/rotator_library/protocols/__init__.py new file mode 100644 index 000000000..e520d0427 --- /dev/null +++ b/src/rotator_library/protocols/__init__.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Native protocol adapters for provider-independent request handling. + +Importing this package auto-discovers built-in protocol adapters, similar to the +provider plugin system. Runtime execution is not changed by Phase 1; the package +only exposes reusable protocol primitives for later phases. +""" + +from .base import ProtocolAdapter +from .operation import ( + OPERATION_AUDIO_TRANSCRIPTION, + OPERATION_AUDIO_TRANSLATION, + OPERATION_CHAT, + OPERATION_COUNT_TOKENS, + OPERATION_EMBEDDINGS, + OPERATION_GENERATE, + OPERATION_IMAGE_EDIT, + OPERATION_IMAGE_GENERATION, + OPERATION_IMAGE_VARIATION, + OPERATION_MCP, + OPERATION_MESSAGES, + OPERATION_OLLAMA_CHAT, + OPERATION_OLLAMA_GENERATE, + OPERATION_RESPONSES, + OPERATION_SPEECH, + OPERATION_UNKNOWN, + normalize_operation, +) +from .registry import ( + PROTOCOL_ALIASES, + PROTOCOL_PLUGINS, + get_protocol, + get_protocol_class, + list_protocols, + register_protocol, + resolve_protocol_name, +) +from .types import ( + ContentBlock, + ConversionWarning, + CostDetails, + MediaSource, + OutputItem, + ProtocolContext, + ProtocolError, + ReasoningBlock, + ToolCall, + ToolDefinition, + ToolResult, + UnifiedMessage, + UnifiedRequest, + UnifiedResponse, + UnifiedStreamEvent, + Usage, + first_text, + serialize_value, + text_blocks, +) + +__all__ = [ + "PROTOCOL_ALIASES", + "PROTOCOL_PLUGINS", + "ContentBlock", + "ConversionWarning", + "CostDetails", + "MediaSource", + "OPERATION_AUDIO_TRANSCRIPTION", + "OPERATION_AUDIO_TRANSLATION", + "OPERATION_CHAT", + "OPERATION_COUNT_TOKENS", + "OPERATION_EMBEDDINGS", + "OPERATION_GENERATE", + "OPERATION_IMAGE_EDIT", + "OPERATION_IMAGE_GENERATION", + "OPERATION_IMAGE_VARIATION", + "OPERATION_MCP", + "OPERATION_MESSAGES", + "OPERATION_OLLAMA_CHAT", + "OPERATION_OLLAMA_GENERATE", + "OPERATION_RESPONSES", + "OPERATION_SPEECH", + "OPERATION_UNKNOWN", + "ProtocolAdapter", + "ProtocolContext", + "ProtocolError", + "ReasoningBlock", + "OutputItem", + "ToolCall", + "ToolDefinition", + "ToolResult", + "UnifiedMessage", + "UnifiedRequest", + "UnifiedResponse", + "UnifiedStreamEvent", + "Usage", + "first_text", + "get_protocol", + "get_protocol_class", + "list_protocols", + "normalize_operation", + "register_protocol", + "resolve_protocol_name", + "serialize_value", + "text_blocks", +] diff --git a/src/rotator_library/protocols/anthropic_messages.py b/src/rotator_library/protocols/anthropic_messages.py new file mode 100644 index 000000000..42b6109bb --- /dev/null +++ b/src/rotator_library/protocols/anthropic_messages.py @@ -0,0 +1,634 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Anthropic Messages protocol adapter. + +This adapter captures the native Messages shape as a reusable base. The existing +compatibility routes remain active; this module gives future provider-native +execution a loss-conscious parser/builder with thinking and tool block support. +""" + +from __future__ import annotations + +import json +from copy import deepcopy +from typing import Any, ClassVar, Iterable + +from .base import ProtocolAdapter +from .canonical import ( + canonical_stop_reason, + canonical_structured_output, + canonical_tool_arguments, + canonical_tool_choice, + coalesce_assistant_message, + conversation_messages, + format_stop_reason, + format_structured_output, + format_tool_choice, + instruction_blocks, + is_same_protocol, + message_reasoning, + message_tool_calls, + message_tool_results, + may_emit_opaque_provider_state, + normalize_tool_result_messages, + ordered_message_blocks, + retain_supported_generation_params, + resolve_tool_result_names, + source_extensions, + tool_arguments_object, + tool_result_text, +) +from .operation import OPERATION_COUNT_TOKENS, OPERATION_GENERATE, OPERATION_MESSAGES, OPERATION_UNKNOWN, normalize_operation +from .validation import validate_generative_request, validate_generative_response +from .types import ( + ContentBlock, + MediaSource, + ProtocolContext, + ReasoningBlock, + ToolCall, + ToolDefinition, + ToolResult, + UnifiedMessage, + UnifiedRequest, + UnifiedResponse, + UnifiedStreamEvent, + Usage, + first_text, + text_blocks, +) + +_GENERATION_PARAMS = { + "max_tokens", + "metadata", + "output_config", + "stop_sequences", + "temperature", + "thinking", + "tool_choice", + "top_k", + "top_p", +} + +_REQUEST_CORE_FIELDS = {"model", "messages", "system", "tools", "stream", *_GENERATION_PARAMS} + + +class AnthropicMessagesProtocol(ProtocolAdapter): + """Adapter for Anthropic Messages requests, responses, and stream events. + + Thinking and redacted-thinking blocks are represented as reasoning blocks so + later field-cache rules can extract signatures without relying on a bespoke + provider implementation. + """ + + name: ClassVar[str] = "anthropic_messages" + aliases: ClassVar[tuple[str, ...]] = ("anthropic", "messages", "claude_messages") + supported_transports: ClassVar[tuple[str, ...]] = ("http", "sse") + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_MESSAGES, OPERATION_COUNT_TOKENS) + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + source_generation = {k: deepcopy(request[k]) for k in _GENERATION_PARAMS if k in request and k != "metadata"} + generation_params = _parse_anthropic_generation_params(source_generation) + if "tool_choice" in generation_params: + generation_params["tool_choice"] = canonical_tool_choice(generation_params["tool_choice"], self.name) + structured_output = canonical_structured_output(request.get("output_config"), self.name) + if structured_output: + generation_params["structured_output"] = structured_output + messages = resolve_tool_result_names( + normalize_tool_result_messages([self._parse_message(message) for message in request.get("messages") or []]) + ) + return UnifiedRequest( + operation=_operation_from_context(context, OPERATION_MESSAGES), + logical_operation=OPERATION_GENERATE, + model=str(request.get("model") or getattr(context, "model", None) or ""), + messages=messages, + system=self._parse_system(request.get("system")), + tools=[self._parse_tool_definition(tool) for tool in request.get("tools") or []], + stream=bool(request.get("stream", False)), + generation_params=generation_params, + response_format=structured_output, + metadata=deepcopy(request.get("metadata") or {}), + source_protocol=self.name, + extensions={self.name: {"generation_params": source_generation}}, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _REQUEST_CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + validate_generative_request(unified_request, self.name, context) + preserve_source = is_same_protocol(context, self.name, unified_request.source_protocol) + payload: dict[str, Any] = { + "model": unified_request.model, + "messages": self._format_messages( + conversation_messages(unified_request), + preserve_source=preserve_source, + emit_opaque_state=may_emit_opaque_provider_state(context, preserve_source=preserve_source), + ), + } + system = self._format_system(instruction_blocks(unified_request), preserve_source=preserve_source) + if system is not None: + payload["system"] = system + if unified_request.tools: + payload["tools"] = [self._format_tool_definition(tool, preserve_source=preserve_source) for tool in unified_request.tools] + if unified_request.stream: + payload["stream"] = True + if unified_request.metadata: + payload["metadata"] = deepcopy(unified_request.metadata) + payload.update(self._format_generation_params(unified_request, preserve_source=preserve_source)) + payload.update(source_extensions(unified_request.extra, context, self.name, unified_request.source_protocol)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + response = _as_dict(raw_response) + operation = _response_operation(response, context) + message = UnifiedMessage( + role=str(response.get("role") or "assistant"), + content=self._parse_content(response.get("content")), + raw=deepcopy(response), + extra={"type": response.get("type")}, + ) + self._promote_message_blocks(message) + return UnifiedResponse( + operation=operation, + logical_operation=OPERATION_GENERATE if operation != OPERATION_COUNT_TOKENS else OPERATION_UNKNOWN, + id=response.get("id"), + model=response.get("model") or getattr(context, "model", None), + messages=[] if operation == OPERATION_COUNT_TOKENS else [message] if response else [], + stop_reason=canonical_stop_reason(response.get("stop_reason")), + usage=self.extract_usage(response, context), + metadata={"stop_sequence": response.get("stop_sequence"), "type": response.get("type"), "native_stop_reason": response.get("stop_reason")}, + source_protocol=self.name, + raw=deepcopy(response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"id", "type", "role", "content", "model", "stop_reason", "stop_sequence", "usage"}}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + if unified_response.operation == OPERATION_COUNT_TOKENS: + usage = unified_response.usage + payload = deepcopy(unified_response.extra) + # Normalized usage wins over raw preserved fields so later adapters + # can correct counts without stale provider keys shadowing them. + payload["input_tokens"] = usage.input_tokens if usage else 0 + return payload + validate_generative_response(unified_response, self.name) + preserve_source = is_same_protocol(context, self.name, unified_response.source_protocol) + message = unified_response.messages[0] if preserve_source and unified_response.messages else coalesce_assistant_message(unified_response.messages) + payload = { + "id": unified_response.id, + "type": unified_response.metadata.get("type", "message"), + "role": message.role, + "content": self._format_assistant_content(message, preserve_source=preserve_source, emit_opaque_state=may_emit_opaque_provider_state(context, preserve_source=preserve_source)), + "model": unified_response.model, + "stop_reason": format_stop_reason(unified_response.stop_reason, self.name), + "stop_sequence": unified_response.metadata.get("stop_sequence"), + "usage": self._format_usage(unified_response.usage), + } + payload.update(source_extensions(unified_response.extra, context, self.name, unified_response.source_protocol)) + return {k: v for k, v in payload.items() if v is not None} + + def parse_stream_event(self, raw_event: Any, context: ProtocolContext | None = None) -> UnifiedStreamEvent: + event = _decode_sse_data(raw_event) + if event == "[DONE]": + return UnifiedStreamEvent(type="done", operation=OPERATION_MESSAGES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type="done", raw=deepcopy(raw_event)) + data = _as_dict(event) + event_type = str(data.get("type") or "chunk") + + if event_type == "error" or data.get("error") is not None: + return UnifiedStreamEvent(type="error", operation=OPERATION_MESSAGES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, error=deepcopy(data.get("error", data)), raw=deepcopy(raw_event), extra={"payload": data}) + if event_type == "message_start": + response = self.parse_response(data.get("message") or {}, context) + return UnifiedStreamEvent(type="message_start", operation=OPERATION_MESSAGES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, message=response.messages[0] if response.messages else None, usage=response.usage, raw=deepcopy(raw_event), extra={"payload": data}) + if event_type == "message_delta": + stop_reason = canonical_stop_reason((data.get("delta") or {}).get("stop_reason")) + return UnifiedStreamEvent(type="message_delta", operation=OPERATION_MESSAGES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, usage=self.extract_usage(data.get("usage") or {}, context), stop_reason=stop_reason, raw=deepcopy(raw_event), extra={"payload": data, "stop_reason": stop_reason}) + if event_type in {"content_block_start", "content_block_delta", "content_block_stop"}: + return self._parse_content_stream_event(data, raw_event) + if event_type == "message_stop": + event_type = "done" + return UnifiedStreamEvent(type=event_type, operation=OPERATION_MESSAGES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=str(data.get("type") or "chunk"), raw=deepcopy(raw_event), extra={"payload": data}) + + def extract_usage(self, raw_or_unified: Any, context: ProtocolContext | None = None) -> Usage | None: + if isinstance(raw_or_unified, (UnifiedResponse, UnifiedStreamEvent)): + return raw_or_unified.usage + payload = _as_dict(raw_or_unified) + usage = payload.get("usage") if isinstance(payload.get("usage"), dict) else payload + if not isinstance(usage, dict) or not any(k.endswith("tokens") for k in usage): + return None + input_tokens = int(usage.get("input_tokens") or 0) + output_tokens = int(usage.get("output_tokens") or 0) + cache_write = int(usage.get("cache_creation_input_tokens") or 0) + cache_read = int(usage.get("cache_read_input_tokens") or 0) + return Usage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=int(usage.get("total_tokens") or input_tokens + output_tokens), + cache_read_tokens=cache_read, + cache_write_tokens=cache_write, + raw=deepcopy(usage), + ) + + def _parse_system(self, system: Any) -> list[ContentBlock]: + if system is None: + return [] + if isinstance(system, str): + return text_blocks(system) + return self._parse_content(system) + + def _format_system(self, blocks: Iterable[ContentBlock], *, preserve_source: bool = True) -> Any: + block_list = list(blocks) + if not block_list: + return None + if preserve_source and all(block.type == "text" and not isinstance(block.raw, dict) and not block.extra for block in block_list): + return first_text(block_list) or "" + return self._format_content(block_list, preserve_source=preserve_source) + + def _parse_message(self, message: dict[str, Any]) -> UnifiedMessage: + payload = dict(message or {}) + unified = UnifiedMessage( + role=str(payload.get("role") or "user"), + content=self._parse_content(payload.get("content")), + raw=deepcopy(message), + extra={k: deepcopy(v) for k, v in payload.items() if k not in {"role", "content"}}, + ) + self._promote_message_blocks(unified) + return unified + + def _format_messages( + self, + messages: Iterable[UnifiedMessage], + *, + preserve_source: bool, + emit_opaque_state: bool, + ) -> list[dict[str, Any]]: + """Format canonical turns and merge adjacent Anthropic roles.""" + + formatted: list[dict[str, Any]] = [] + for message in messages: + role = "assistant" if message.role in {"assistant", "model"} else "user" + content = self._format_assistant_content(message, preserve_source=preserve_source, emit_opaque_state=emit_opaque_state) if role == "assistant" else self._format_user_content(message, preserve_source=preserve_source) + payload = {"role": role, "content": content} + if preserve_source: + payload.update(deepcopy(message.extra)) + if formatted and formatted[-1]["role"] == role: + previous = formatted[-1].get("content") + if not isinstance(previous, list): + previous = [{"type": "text", "text": str(previous or "")}] + previous.extend(content) + formatted[-1]["content"] = previous + else: + formatted.append(payload) + return formatted + + def _format_assistant_content( + self, + message: UnifiedMessage, + *, + preserve_source: bool, + emit_opaque_state: bool = True, + ) -> list[dict[str, Any]]: + """Format reasoning, visible content, and tool calls in Anthropic order.""" + + return self._format_content( + ordered_message_blocks(message), + preserve_source=preserve_source, + emit_opaque_state=emit_opaque_state, + ) + + def _format_user_content(self, message: UnifiedMessage, *, preserve_source: bool) -> list[dict[str, Any]]: + """Format user content and canonical tool results.""" + + return self._format_content(ordered_message_blocks(message), preserve_source=preserve_source) + + def _format_message( + self, + message: UnifiedMessage, + *, + preserve_source: bool = True, + emit_opaque_state: bool = True, + ) -> dict[str, Any]: + role = "assistant" if message.role in {"assistant", "model"} else "user" + content = self._format_assistant_content(message, preserve_source=preserve_source, emit_opaque_state=emit_opaque_state) if role == "assistant" else self._format_user_content(message, preserve_source=preserve_source) + payload = {"role": role, "content": content} + if preserve_source: + payload.update(deepcopy(message.extra)) + return payload + + def _parse_content(self, content: Any) -> list[ContentBlock]: + if content is None: + return [] + if isinstance(content, str): + return text_blocks(content) + if not isinstance(content, list): + return [ContentBlock(type="unknown", raw=deepcopy(content))] + return [self._parse_content_block(block) for block in content] + + def _parse_content_block(self, block: Any) -> ContentBlock: + if isinstance(block, str): + return ContentBlock(type="text", text=block, raw=block) + if not isinstance(block, dict): + return ContentBlock(type="unknown", raw=deepcopy(block)) + block_type = str(block.get("type") or "text") + if block_type == "text": + return ContentBlock(type="text", text=block.get("text", ""), raw=deepcopy(block)) + if block_type in {"image", "document"}: + source = _parse_anthropic_media_source(block.get("source"), kind=block_type) + return ContentBlock(type=block_type, source=source, raw=deepcopy(block), extra=_without(block, {"type", "source"})) + if block_type in {"thinking", "redacted_thinking"}: + reasoning = ReasoningBlock( + type=block_type, + text=block.get("thinking"), + signature=block.get("signature"), + redacted=block_type == "redacted_thinking", + raw=deepcopy(block), + extra=_without(block, {"type", "thinking", "signature"}), + ) + return ContentBlock(type="reasoning", reasoning=reasoning, raw=deepcopy(block)) + if block_type == "tool_use": + return ContentBlock( + type="tool_call", + tool_call=ToolCall(id=block.get("id"), name=block.get("name"), arguments=canonical_tool_arguments(block.get("input")), type="function", raw=deepcopy(block)), + raw=deepcopy(block), + extra=_without(block, {"type", "id", "name", "input"}), + ) + if block_type == "tool_result": + return ContentBlock( + type="tool_result", + tool_result=ToolResult(tool_call_id=block.get("tool_use_id"), content=canonical_tool_arguments(block.get("content")), is_error=block.get("is_error"), raw=deepcopy(block)), + raw=deepcopy(block), + extra=_without(block, {"type", "tool_use_id", "content", "is_error"}), + ) + return ContentBlock(type=block_type, raw=deepcopy(block), extra=_without(block, {"type"})) + + def _format_content( + self, + blocks: Iterable[ContentBlock], + *, + preserve_source: bool = True, + emit_opaque_state: bool = True, + ) -> list[dict[str, Any]]: + formatted = [] + for block in blocks: + if block.type == "text": + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {"type": "text"} + payload["type"] = "text" + payload["text"] = block.text or "" + formatted.append(payload) + elif block.reasoning: + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {"type": "redacted_thinking" if block.reasoning.redacted else "thinking"} + payload["type"] = "redacted_thinking" if block.reasoning.redacted else "thinking" + if block.reasoning.text is not None: + payload["thinking"] = block.reasoning.text + if emit_opaque_state and block.reasoning.signature is not None: + payload["signature"] = block.reasoning.signature + elif not emit_opaque_state: + payload.pop("signature", None) + if preserve_source: + payload.update(deepcopy(block.reasoning.extra)) + formatted.append(payload) + elif block.tool_call: + formatted.append(self._format_tool_call(block.tool_call, preserve_source=preserve_source)) + elif block.tool_result: + formatted.append(self._format_tool_result(block.tool_result, preserve_source=preserve_source)) + elif block.type in {"image", "document"}: + formatted.append(_format_anthropic_media(block, preserve_source=preserve_source)) + elif preserve_source and isinstance(block.raw, dict): + formatted.append(deepcopy(block.raw)) + return formatted + + def _parse_tool_definition(self, tool: dict[str, Any]) -> ToolDefinition: + payload = dict(tool or {}) + return ToolDefinition( + name=str(payload.get("name") or ""), + description=payload.get("description"), + input_schema=deepcopy(payload.get("input_schema") or {}), + type="function", + extra=_without(payload, {"name", "description", "input_schema"}), + ) + + def _format_tool_definition(self, tool: ToolDefinition, *, preserve_source: bool = True) -> dict[str, Any]: + payload = {"name": tool.name, "input_schema": deepcopy(tool.input_schema)} + if tool.description is not None: + payload["description"] = tool.description + if preserve_source: + payload.update(deepcopy(tool.extra)) + return payload + + def _format_tool_call(self, call: ToolCall, *, preserve_source: bool) -> dict[str, Any]: + payload = deepcopy(call.raw) if preserve_source and isinstance(call.raw, dict) else {} + payload.update({"type": "tool_use", "id": call.id or "", "name": call.name or "", "input": tool_arguments_object(call.arguments)}) + return payload + + def _format_tool_result(self, result: ToolResult, *, preserve_source: bool) -> dict[str, Any]: + payload = deepcopy(result.raw) if preserve_source and isinstance(result.raw, dict) else {} + if preserve_source and isinstance(result.content, (str, list)): + content = deepcopy(result.content) + else: + content = tool_result_text(result.content) + payload.update({"type": "tool_result", "tool_use_id": result.tool_call_id or "", "content": content}) + if result.is_error is not None: + payload["is_error"] = result.is_error + return payload + + def _format_generation_params(self, request: UnifiedRequest, *, preserve_source: bool) -> dict[str, Any]: + params = deepcopy(request.generation_params) + original = request.extensions.get(self.name, {}).get("generation_params") if preserve_source else None + payload = deepcopy(original) if isinstance(original, dict) else {} + if "max_output_tokens" in params: + payload["max_tokens"] = params.pop("max_output_tokens") + if "stop_sequences" in params: + payload["stop_sequences"] = params.pop("stop_sequences") + if "tool_choice" in params: + payload["tool_choice"] = format_tool_choice(params.pop("tool_choice"), self.name) + if "structured_output" in params: + payload["output_config"] = format_structured_output(params.pop("structured_output"), self.name) + reasoning = params.pop("reasoning", None) + if isinstance(reasoning, dict): + thinking: dict[str, Any] = {} + if reasoning.get("budget_tokens") is not None: + thinking = {"type": "enabled", "budget_tokens": reasoning["budget_tokens"]} + elif reasoning.get("enabled") is False: + thinking = {"type": "disabled"} + if thinking: + payload["thinking"] = thinking + supported = {"temperature", "top_k", "top_p"} + payload.update( + retain_supported_generation_params( + request, + params, + supported=supported, + target_protocol=self.name, + ) + ) + return payload + + def _format_usage(self, usage: Usage | None) -> dict[str, int] | None: + if usage is None: + return None + payload = {"input_tokens": usage.input_tokens, "output_tokens": usage.output_tokens} + if usage.cache_write_tokens: + payload["cache_creation_input_tokens"] = usage.cache_write_tokens + if usage.cache_read_tokens: + payload["cache_read_input_tokens"] = usage.cache_read_tokens + return payload + + def _promote_message_blocks(self, message: UnifiedMessage) -> None: + for block in message.content: + if block.tool_call: + message.tool_calls.append(block.tool_call) + if block.reasoning: + message.reasoning.append(block.reasoning) + + def _parse_content_stream_event(self, data: dict[str, Any], raw_event: Any) -> UnifiedStreamEvent: + block = data.get("content_block") if isinstance(data.get("content_block"), dict) else None + delta = data.get("delta") if isinstance(data.get("delta"), dict) else None + content_block = None + if block: + content_block = self._parse_content_block(block) + if content_block.tool_call: + content_block.tool_call.index = data.get("index") + if content_block.tool_call.arguments == {}: + content_block.tool_call.arguments = None + elif delta: + delta_type = delta.get("type") + if delta_type == "text_delta": + content_block = ContentBlock(type="text", text=delta.get("text"), raw=deepcopy(delta)) + elif delta_type in {"thinking_delta", "signature_delta"}: + reasoning = ReasoningBlock(type=str(delta_type), text=delta.get("thinking"), signature=delta.get("signature"), extra=_without(delta, {"type", "thinking", "signature"})) + content_block = ContentBlock(type=str(delta_type), reasoning=reasoning, raw=deepcopy(delta)) + elif delta_type == "input_json_delta": + call = ToolCall(index=data.get("index"), arguments=delta.get("partial_json") or "") + content_block = ContentBlock(type="tool_call", tool_call=call, raw=deepcopy(delta)) + message = UnifiedMessage(role="assistant", content=[content_block] if content_block else []) + self._promote_message_blocks(message) + return UnifiedStreamEvent(type=str(data.get("type") or "content_block_delta"), operation=OPERATION_MESSAGES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=str(data.get("type") or "content_block_delta"), delta=message, raw=deepcopy(raw_event), extra={"payload": data, "index": data.get("index")}) + + +def _operation_from_context(context: ProtocolContext | None, default: str) -> str: + supported = {OPERATION_MESSAGES, OPERATION_COUNT_TOKENS} + if context and isinstance(context.provider_options, dict): + operation = normalize_operation(context.provider_options.get("operation")) + if operation in supported: + return operation + if context and isinstance(context.metadata, dict): + operation = normalize_operation(context.metadata.get("operation")) + if operation in supported: + return operation + return default + + +def _response_operation(response: dict[str, Any], context: ProtocolContext | None) -> str: + requested = _operation_from_context(context, OPERATION_MESSAGES) + if requested == OPERATION_COUNT_TOKENS: + return OPERATION_COUNT_TOKENS + if "input_tokens" in response and not response.get("content") and not response.get("id"): + return OPERATION_COUNT_TOKENS + return OPERATION_MESSAGES + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return deepcopy(value) + if hasattr(value, "model_dump"): + return value.model_dump() + if hasattr(value, "dict"): + return value.dict() + return {} + + +def _decode_sse_data(raw_event: Any) -> Any: + from .streaming import decode_sse_data + + return decode_sse_data(raw_event) + + +def _without(payload: dict[str, Any], keys: set[str]) -> dict[str, Any]: + return {k: deepcopy(v) for k, v in payload.items() if k not in keys} + + +def _parse_anthropic_generation_params(source: dict[str, Any]) -> dict[str, Any]: + """Normalize Anthropic controls into canonical names.""" + + params = deepcopy(source) + max_tokens = params.pop("max_tokens", None) + if max_tokens is not None: + params["max_output_tokens"] = max_tokens + stop_sequences = params.pop("stop_sequences", None) + if stop_sequences is not None: + params["stop_sequences"] = deepcopy(stop_sequences) + thinking = params.pop("thinking", None) + if isinstance(thinking, dict): + params["reasoning"] = { + "enabled": thinking.get("type") != "disabled", + "budget_tokens": thinking.get("budget_tokens"), + } + return params + + +def _parse_anthropic_media_source(value: Any, *, kind: str) -> MediaSource: + """Normalize Anthropic image/document source objects.""" + + payload = value if isinstance(value, dict) else {} + source_type = str(payload.get("type") or "") + if source_type == "base64": + source_kind = "base64" + elif source_type in {"url", "text"}: + source_kind = "url" if source_type == "url" else "text" + elif payload.get("file_id"): + source_kind = "file" + else: + source_kind = source_type or kind + return MediaSource( + kind=source_kind, + media_type=payload.get("media_type"), + url=payload.get("url"), + data=payload.get("data") or payload.get("text"), + file_id=payload.get("file_id"), + raw=deepcopy(value), + extra=_without(payload, {"type", "media_type", "url", "data", "text", "file_id"}), + ) + + +def _coerce_media_source(value: Any) -> MediaSource: + """Coerce legacy media dictionaries into canonical form.""" + + if isinstance(value, MediaSource): + return value + if isinstance(value, str): + return MediaSource(kind="url", url=value, raw=value) + payload = value if isinstance(value, dict) else {} + url = payload.get("url") + data = payload.get("data") + file_id = payload.get("file_id") + return MediaSource( + kind="file" if file_id else "base64" if data else "url", + media_type=payload.get("media_type") or payload.get("mime_type"), + url=url, + data=data, + file_id=file_id, + detail=payload.get("detail"), + raw=deepcopy(value), + ) + + +def _format_anthropic_media(block: ContentBlock, *, preserve_source: bool) -> dict[str, Any]: + """Format canonical media as an Anthropic content block.""" + + source = _coerce_media_source(block.source) + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {"type": block.type} + if source.kind == "base64": + payload["source"] = { + "type": "base64", + "media_type": source.media_type or "application/octet-stream", + "data": source.data or "", + } + elif source.file_id: + payload["source"] = {"type": "file", "file_id": source.file_id} + elif source.url: + payload["source"] = {"type": "url", "url": source.url} + elif source.data and block.type == "document": + payload["source"] = {"type": "text", "media_type": source.media_type or "text/plain", "data": source.data} + else: + payload["source"] = {"type": source.kind, "data": source.data or ""} + payload["type"] = block.type + return payload diff --git a/src/rotator_library/protocols/base.py b/src/rotator_library/protocols/base.py new file mode 100644 index 000000000..4b6396486 --- /dev/null +++ b/src/rotator_library/protocols/base.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Base classes for native protocol adapters. + +Protocol adapters are intentionally override-friendly. They provide reusable +defaults for custom providers, but providers can override any method when a +service uses an almost-standard protocol with provider-specific quirks. +""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any, ClassVar + +from .operation import OPERATION_UNKNOWN, normalize_operation +from .types import ( + ProtocolContext, + ProtocolError, + UnifiedRequest, + UnifiedResponse, + UnifiedStreamEvent, + Usage, +) + + +class ProtocolAdapter: + """Base adapter for converting between raw protocol payloads and unified types. + + Subclasses should override only the methods they need. The default behavior + is deliberately conservative: preserve raw payloads and avoid lossy + assumptions. Later phases will layer transform logging, field-cache rules, + and provider override hooks around this interface. + """ + + name: ClassVar[str] = "base" + aliases: ClassVar[tuple[str, ...]] = () + supported_transports: ClassVar[tuple[str, ...]] = ("http", "sse") + future_transports: ClassVar[tuple[str, ...]] = () + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_UNKNOWN,) + + def supports_transport(self, transport_name: str) -> bool: + """Return whether this protocol can format the requested transport.""" + + return transport_name in self.supported_transports + + def is_future_transport(self, transport_name: str) -> bool: + """Return whether this protocol has an intentional future transport seam.""" + + return transport_name in self.future_transports + + def supports_operation(self, operation_name: str) -> bool: + """Return whether this adapter natively models an operation. + + Operation names are string based so custom protocols can add their own + values. The default base adapter only claims ``unknown`` and keeps raw + payloads intact; concrete adapters should list every operation they can + parse/build without relying on LiteLLM. + """ + + return normalize_operation(operation_name) in self.supported_operations + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + """Parse a raw client/provider request into a unified request.""" + + request = dict(raw_request or {}) + return UnifiedRequest( + operation=normalize_operation(request.get("operation")), + model=str(request.get("model") or getattr(context, "model", None) or ""), + stream=bool(request.get("stream", False)), + input=deepcopy(request.get("input")), + modalities=list(request.get("modalities") or []), + files=list(request.get("files") or []), + raw=deepcopy(raw_request), + extra={ + k: deepcopy(v) + for k, v in request.items() + if k not in {"operation", "model", "stream", "input", "modalities", "files"} + }, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + """Build a provider request from a unified request. + + The base implementation returns the original raw dict when present. This + keeps fallback providers safe and gives custom protocol subclasses a + predictable starting point. + """ + + if isinstance(unified_request.raw, dict): + return deepcopy(unified_request.raw) + if not isinstance(unified_request.raw, type(None)): + raise ProtocolError( + "cannot build dict request from non-dict raw payload", + protocol=self.name, + pass_name="build_request", + payload=unified_request.raw, + ) + payload = {"model": unified_request.model, "stream": unified_request.stream} + if unified_request.operation != OPERATION_UNKNOWN: + payload["operation"] = unified_request.operation + if unified_request.input is not None: + payload["input"] = deepcopy(unified_request.input) + if unified_request.modalities: + payload["modalities"] = deepcopy(unified_request.modalities) + if unified_request.files: + payload["files"] = deepcopy(unified_request.files) + payload.update(deepcopy(unified_request.extra)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + """Parse a raw response into a unified response.""" + + response = raw_response if isinstance(raw_response, dict) else {} + return UnifiedResponse( + operation=normalize_operation(response.get("operation") if isinstance(response, dict) else None), + id=response.get("id") if isinstance(response, dict) else None, + model=response.get("model") if isinstance(response, dict) else getattr(context, "model", None), + data=deepcopy(response.get("data") or []) if isinstance(response, dict) else [], + content_type=response.get("content_type") if isinstance(response, dict) else None, + raw=deepcopy(raw_response), + extra=deepcopy(response) if isinstance(response, dict) else {}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + """Format a unified response for a client protocol.""" + + if isinstance(unified_response.raw, dict): + return deepcopy(unified_response.raw) + payload = deepcopy(unified_response.extra) + if unified_response.operation != OPERATION_UNKNOWN: + payload.setdefault("operation", unified_response.operation) + if unified_response.id is not None: + payload.setdefault("id", unified_response.id) + if unified_response.model is not None: + payload.setdefault("model", unified_response.model) + if unified_response.data: + payload.setdefault("data", deepcopy(unified_response.data)) + if unified_response.content_type is not None: + payload.setdefault("content_type", unified_response.content_type) + return payload + + def parse_stream_event(self, raw_event: Any, context: ProtocolContext | None = None) -> UnifiedStreamEvent: + """Parse one raw stream event. + + Subclasses should preserve the original event in ``raw`` because Phase 2 + transform logging needs both provider-native and unified states. + """ + + event_type = "done" if raw_event == "[DONE]" else "chunk" + return UnifiedStreamEvent(type=event_type, raw=deepcopy(raw_event)) + + def format_stream_event(self, unified_event: UnifiedStreamEvent, context: ProtocolContext | None = None) -> Any: + """Format one canonical stream event for this destination protocol.""" + + from .streaming import format_canonical_stream_event + + return format_canonical_stream_event(unified_event, self.name, context) + + def extract_usage(self, raw_or_unified: Any, context: ProtocolContext | None = None) -> Usage | None: + """Extract normalized usage when the protocol can identify it.""" + + if isinstance(raw_or_unified, UnifiedResponse): + return raw_or_unified.usage + if isinstance(raw_or_unified, UnifiedStreamEvent): + return raw_or_unified.usage + if isinstance(raw_or_unified, dict) and isinstance(raw_or_unified.get("usage"), dict): + usage = raw_or_unified["usage"] + return Usage( + input_tokens=int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0), + output_tokens=int(usage.get("completion_tokens") or usage.get("output_tokens") or 0), + total_tokens=int(usage.get("total_tokens") or 0), + raw=deepcopy(usage), + ) + return None diff --git a/src/rotator_library/protocols/canonical.py b/src/rotator_library/protocols/canonical.py new file mode 100644 index 000000000..1bcab10c8 --- /dev/null +++ b/src/rotator_library/protocols/canonical.py @@ -0,0 +1,596 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Shared canonical semantics for cross-protocol generative conversion. + +Protocol modules own wire parsing and formatting. This module owns only meanings +that must be identical across those modules: logical operations, completion +reasons, instruction placement, source-aware passthrough, and warning records. +""" + +from __future__ import annotations + +from copy import deepcopy +import json +from typing import Any, Iterable, Optional + +from .types import ( + ContentBlock, + ConversionWarning, + ProtocolContext, + ReasoningBlock, + ToolCall, + ToolResult, + UnifiedMessage, + UnifiedRequest, + serialize_value, +) + + +STOP_REASON_STOP = "stop" +STOP_REASON_MAX_TOKENS = "max_tokens" +STOP_REASON_TOOL_USE = "tool_use" +STOP_REASON_CONTENT_FILTER = "content_filter" +STOP_REASON_ERROR = "error" +STOP_REASON_INCOMPLETE = "incomplete" +STOP_REASON_UNKNOWN = "unknown" + + +_STOP_REASON_ALIASES = { + "stop": STOP_REASON_STOP, + "end_turn": STOP_REASON_STOP, + "stop_sequence": STOP_REASON_STOP, + "completed": STOP_REASON_STOP, + "length": STOP_REASON_MAX_TOKENS, + "max_tokens": STOP_REASON_MAX_TOKENS, + "max_output_tokens": STOP_REASON_MAX_TOKENS, + "tool_calls": STOP_REASON_TOOL_USE, + "function_call": STOP_REASON_TOOL_USE, + "tool_use": STOP_REASON_TOOL_USE, + "content_filter": STOP_REASON_CONTENT_FILTER, + "safety": STOP_REASON_CONTENT_FILTER, + "blocklist": STOP_REASON_CONTENT_FILTER, + "prohibited_content": STOP_REASON_CONTENT_FILTER, + "recitation": STOP_REASON_CONTENT_FILTER, + "failed": STOP_REASON_ERROR, + "error": STOP_REASON_ERROR, + "incomplete": STOP_REASON_INCOMPLETE, +} + + +_TARGET_STOP_REASONS = { + "openai_chat": { + STOP_REASON_STOP: "stop", + STOP_REASON_MAX_TOKENS: "length", + STOP_REASON_TOOL_USE: "tool_calls", + STOP_REASON_CONTENT_FILTER: "content_filter", + STOP_REASON_ERROR: None, + STOP_REASON_INCOMPLETE: "length", + STOP_REASON_UNKNOWN: None, + }, + "anthropic_messages": { + STOP_REASON_STOP: "end_turn", + STOP_REASON_MAX_TOKENS: "max_tokens", + STOP_REASON_TOOL_USE: "tool_use", + STOP_REASON_CONTENT_FILTER: "refusal", + STOP_REASON_ERROR: None, + STOP_REASON_INCOMPLETE: "max_tokens", + STOP_REASON_UNKNOWN: None, + }, + "gemini": { + STOP_REASON_STOP: "STOP", + STOP_REASON_MAX_TOKENS: "MAX_TOKENS", + STOP_REASON_TOOL_USE: "STOP", + STOP_REASON_CONTENT_FILTER: "SAFETY", + STOP_REASON_ERROR: "OTHER", + STOP_REASON_INCOMPLETE: "MAX_TOKENS", + STOP_REASON_UNKNOWN: "OTHER", + }, + "responses": { + STOP_REASON_STOP: "completed", + STOP_REASON_MAX_TOKENS: "incomplete", + STOP_REASON_TOOL_USE: "completed", + STOP_REASON_CONTENT_FILTER: "incomplete", + STOP_REASON_ERROR: "failed", + STOP_REASON_INCOMPLETE: "incomplete", + STOP_REASON_UNKNOWN: "incomplete", + }, +} + + +def canonical_stop_reason(value: Any) -> Optional[str]: + """Normalize a provider/client completion reason into a stable meaning.""" + + if value is None: + return None + normalized = str(value).strip().lower() + return _STOP_REASON_ALIASES.get(normalized, STOP_REASON_UNKNOWN) + + +def format_stop_reason(value: Optional[str], target_protocol: str) -> Optional[str]: + """Return the target protocol's public completion reason.""" + + if value is None: + return None + return _TARGET_STOP_REASONS.get(target_protocol, {}).get(value, value) + + +def is_same_protocol( + context: ProtocolContext | None, + protocol_name: str, + source_protocol: str | None = None, +) -> bool: + """Return whether source-owned raw fields are safe to replay. + + A missing context is not proof of ownership. Parsed unified objects carry + their source protocol, so direct parse/build calls still preserve native + fields without making manually constructed or foreign objects unsafe. + """ + + effective_source = context.source_protocol if context and context.source_protocol else source_protocol + return effective_source == protocol_name + + +def source_extensions( + extra: dict[str, Any], + context: ProtocolContext | None, + protocol_name: str, + source_protocol: str | None = None, +) -> dict[str, Any]: + """Return source extensions only for a same-protocol destination.""" + + return deepcopy(extra) if is_same_protocol(context, protocol_name, source_protocol) else {} + + +def may_emit_opaque_provider_state( + context: ProtocolContext | None, + *, + preserve_source: bool, +) -> bool: + """Return whether opaque signatures may leave canonical/cache state. + + Opaque state is suppressed unless real execution explicitly proves provider + compatibility through a compatible-domain flag or identical provider IDs. + """ + + if not preserve_source: + return False + if context is None: + return False + if context.provider_state_compatible: + return True + return bool( + context.source_provider + and context.target_provider + and context.source_provider == context.target_provider + ) + + +def instruction_messages(request: UnifiedRequest) -> list[UnifiedMessage]: + """Return ordered canonical system/developer instructions. + + Older parsers store a separate ``system`` block list. That field remains a + compatibility carrier until all non-generative consumers migrate, but it is + promoted only when no explicit system message already exists. + """ + + instructions = [message for message in request.messages if message.role in {"system", "developer"}] + if request.system and not any(message.role == "system" for message in instructions): + instructions.insert(0, UnifiedMessage(role="system", content=deepcopy(request.system))) + return instructions + + +def conversation_messages(request: UnifiedRequest) -> list[UnifiedMessage]: + """Return messages excluding system/developer instruction turns.""" + + return [message for message in request.messages if message.role not in {"system", "developer"}] + + +def instruction_blocks(request: UnifiedRequest) -> list[ContentBlock]: + """Flatten ordered instruction turns for protocols with one system field.""" + + blocks: list[ContentBlock] = [] + for message in instruction_messages(request): + blocks.extend(deepcopy(message.content)) + return blocks + + +def add_conversion_warning( + request: UnifiedRequest, + *, + code: str, + message: str, + field: str | None, + target_protocol: str, +) -> None: + """Record a deliberate omission of an optional conversion hint.""" + + request.warnings.append( + ConversionWarning( + code=code, + message=message, + field=field, + source_protocol=request.source_protocol, + target_protocol=target_protocol, + ) + ) + + +def retain_supported_generation_params( + request: UnifiedRequest, + params: dict[str, Any], + *, + supported: set[str], + target_protocol: str, +) -> dict[str, Any]: + """Return supported optional controls and record every deliberate omission.""" + + kept: dict[str, Any] = {} + for key, value in params.items(): + if key in supported: + kept[key] = value + continue + add_conversion_warning( + request, + code="unsupported_optional_control", + message=f"{target_protocol} has no safe mapping for optional control '{key}'", + field=key, + target_protocol=target_protocol, + ) + return kept + + +def canonical_tool_choice(value: Any, source_protocol: str) -> dict[str, Any] | None: + """Normalize tool-choice spellings used by the four generative APIs.""" + + if value is None: + return None + if isinstance(value, str): + mode = value.lower() + if mode in {"auto", "none", "required"}: + return {"mode": mode} + if mode in {"any", "any_required"}: + return {"mode": "required"} + return {"mode": "named", "name": value} + if not isinstance(value, dict): + return {"mode": "auto"} + + value_type = str(value.get("type") or value.get("mode") or "auto").lower() + if source_protocol == "gemini": + # Gemini normalization normally happens in its protocol module because + # the value is nested in toolConfig; accept the canonical intermediate. + if "allowed_names" in value: + return deepcopy(value) + if value_type in {"auto", "none"}: + return {"mode": value_type} + if value_type in {"required", "any"}: + return {"mode": "required", "allowed_names": deepcopy(value.get("allowed_names") or [])} + if value_type in {"function", "tool", "named"}: + function = value.get("function") if isinstance(value.get("function"), dict) else {} + name = value.get("name") or function.get("name") + return {"mode": "named", "name": name} + return {"mode": "auto"} + + +def format_tool_choice(value: Any, target_protocol: str) -> Any: + """Format canonical tool choice for a destination protocol.""" + + choice = value if isinstance(value, dict) and "mode" in value else canonical_tool_choice(value, target_protocol) + if not choice: + return None + mode = choice.get("mode", "auto") + name = choice.get("name") + allowed_names = deepcopy(choice.get("allowed_names") or []) + if target_protocol == "openai_chat": + if mode == "named": + return {"type": "function", "function": {"name": name or ""}} + return "required" if mode == "required" else mode + if target_protocol == "anthropic_messages": + if mode == "named": + return {"type": "tool", "name": name or ""} + return {"type": "any" if mode == "required" else mode} + if target_protocol == "responses": + if mode == "named": + return {"type": "function", "name": name or ""} + return "required" if mode == "required" else mode + if target_protocol == "gemini": + if mode == "none": + config: dict[str, Any] = {"mode": "NONE"} + elif mode == "required": + config = {"mode": "ANY"} + if allowed_names: + config["allowedFunctionNames"] = allowed_names + elif mode == "named": + config = {"mode": "ANY", "allowedFunctionNames": [name] if name else []} + else: + config = {"mode": "AUTO"} + return {"functionCallingConfig": config} + return deepcopy(value) + + +def canonical_structured_output(value: Any, source_protocol: str) -> dict[str, Any] | None: + """Normalize JSON/object schema controls without source wrappers.""" + + if not isinstance(value, dict): + return None + if source_protocol == "openai_chat": + output_type = value.get("type") + if output_type == "json_schema" and isinstance(value.get("json_schema"), dict): + schema = value["json_schema"] + return { + "type": "json_schema", + "name": schema.get("name"), + "schema": deepcopy(schema.get("schema")), + "strict": schema.get("strict"), + } + if output_type == "json_object": + return {"type": "json_object"} + return deepcopy(value) + if source_protocol == "responses": + return { + "type": value.get("type") or "json_schema", + "name": value.get("name"), + "schema": deepcopy(value.get("schema")), + "strict": value.get("strict"), + } + if source_protocol == "anthropic_messages": + format_value = value.get("format") if isinstance(value.get("format"), dict) else value + return { + "type": format_value.get("type") or "json_schema", + "name": format_value.get("name"), + "schema": deepcopy(format_value.get("schema")), + "strict": format_value.get("strict"), + } + if source_protocol == "gemini": + normalized = deepcopy(value) + if normalized.get("type") == "json_schema" and normalized.get("schema") is not None: + normalized.setdefault("strict", True) + return normalized + return deepcopy(value) + + +def format_structured_output(value: Any, target_protocol: str) -> Any: + """Format a canonical structured-output requirement for a destination.""" + + if not isinstance(value, dict): + return None + output_type = value.get("type") or "json_schema" + if target_protocol == "openai_chat": + if output_type == "json_object": + return {"type": "json_object"} + return { + "type": "json_schema", + "json_schema": { + key: deepcopy(item) + for key, item in { + "name": value.get("name") or "response", + "schema": value.get("schema") or {}, + "strict": value.get("strict"), + }.items() + if item is not None + }, + } + if target_protocol == "responses": + if output_type == "json_object": + return {"type": "json_object"} + return { + key: deepcopy(item) + for key, item in { + "type": "json_schema", + "name": value.get("name") or "response", + "schema": value.get("schema") or {}, + "strict": value.get("strict"), + }.items() + if item is not None + } + if target_protocol == "anthropic_messages": + if output_type == "json_object": + return {"format": {"type": "json_schema", "schema": {"type": "object"}}} + return { + "format": { + key: deepcopy(item) + for key, item in { + "type": "json_schema", + "name": value.get("name"), + "schema": value.get("schema") or {}, + "strict": value.get("strict"), + }.items() + if item is not None + } + } + if target_protocol == "gemini": + return { + "responseMimeType": "application/json", + "responseJsonSchema": deepcopy(value.get("schema")) if output_type != "json_object" else None, + } + return deepcopy(value) + + +def message_tool_calls(message: UnifiedMessage) -> list[ToolCall]: + """Return de-duplicated calls from message and content representations.""" + + calls: list[ToolCall] = [] + seen: set[tuple[Optional[str], Optional[str], str]] = set() + for call in [*message.tool_calls, *[block.tool_call for block in message.content if block.tool_call]]: + if call is None: + continue + key = (call.id, call.name, json.dumps(serialize_value(call.arguments), sort_keys=True, separators=(",", ":"))) + if key not in seen: + calls.append(call) + seen.add(key) + return calls + + +def message_reasoning(message: UnifiedMessage) -> list[ReasoningBlock]: + """Return de-duplicated reasoning from message and content representations.""" + + blocks: list[ReasoningBlock] = [] + seen: set[tuple[Optional[str], Optional[str], bool]] = set() + for reasoning in [*message.reasoning, *[block.reasoning for block in message.content if block.reasoning]]: + if reasoning is None: + continue + key = (reasoning.text, reasoning.signature, reasoning.redacted) + if key not in seen: + blocks.append(reasoning) + seen.add(key) + return blocks + + +def message_tool_results(message: UnifiedMessage) -> list[ToolResult]: + """Return tool results embedded in a canonical message.""" + + return [block.tool_result for block in message.content if block.tool_result is not None] + + +def resolve_tool_result_names(messages: Iterable[UnifiedMessage]) -> list[UnifiedMessage]: + """Enrich result records with function names from preceding calls. + + Chat and Anthropic identify results by call ID, while Gemini requires the + function name on its response part. Keeping both in the canonical record + allows either direction without provider-specific history lookups. + """ + + message_list = list(messages) + names: dict[str, str] = {} + ids_by_name: dict[str, list[str]] = {} + result_index_by_name: dict[str, int] = {} + call_index = 0 + for message in message_list: + for call in message_tool_calls(message): + if not call.id: + call.id = f"call_{call_index}" + call.extra["synthetic_id"] = True + call_index += 1 + if call.id and call.name: + names[call.id] = call.name + ids_by_name.setdefault(call.name, []).append(call.id) + for result in message_tool_results(message): + if not result.name and result.tool_call_id in names: + result.name = names[result.tool_call_id] + if result.name and (not result.tool_call_id or result.tool_call_id == result.name): + candidates = ids_by_name.get(result.name, []) + result_index = result_index_by_name.get(result.name, 0) + if result_index < len(candidates): + result.tool_call_id = candidates[result_index] + result.extra["synthetic_tool_call_id"] = True + result_index_by_name[result.name] = result_index + 1 + return message_list + + +def normalize_tool_result_messages(messages: Iterable[UnifiedMessage]) -> list[UnifiedMessage]: + """Split embedded provider result blocks into canonical tool-role turns.""" + + normalized: list[UnifiedMessage] = [] + for message in messages: + if not any(block.tool_result for block in message.content): + normalized.append(message) + continue + groups: list[tuple[bool, list[ContentBlock]]] = [] + for block in message.content: + is_result = block.tool_result is not None + if groups and groups[-1][0] == is_result: + groups[-1][1].append(block) + else: + groups.append((is_result, [block])) + for group_index, (is_result, blocks) in enumerate(groups): + split = deepcopy(message) + split.content = blocks + split.role = "tool" if is_result else message.role + split.tool_calls = [block.tool_call for block in blocks if block.tool_call] + split.reasoning = [block.reasoning for block in blocks if block.reasoning] + split.tool_call_id = blocks[0].tool_result.tool_call_id if is_result and blocks[0].tool_result else None + if group_index: + split.raw = None + split.extra = {} + normalized.append(split) + return normalized + + +def ordered_message_blocks(message: UnifiedMessage) -> list[ContentBlock]: + """Return one ordered block sequence without duplicate promoted fields.""" + + blocks = deepcopy(message.content) + reasoning_keys = { + (block.reasoning.text, block.reasoning.signature, block.reasoning.redacted) + for block in blocks + if block.reasoning + } + call_keys = { + (block.tool_call.id, block.tool_call.name, json.dumps(serialize_value(block.tool_call.arguments), sort_keys=True)) + for block in blocks + if block.tool_call + } + missing_reasoning: list[ContentBlock] = [] + for reasoning in message.reasoning: + key = (reasoning.text, reasoning.signature, reasoning.redacted) + if key not in reasoning_keys: + missing_reasoning.append(ContentBlock(type="reasoning", reasoning=deepcopy(reasoning))) + reasoning_keys.add(key) + if missing_reasoning: + blocks = [*missing_reasoning, *blocks] + for call in message.tool_calls: + key = (call.id, call.name, json.dumps(serialize_value(call.arguments), sort_keys=True)) + if key not in call_keys: + blocks.append(ContentBlock(type="tool_call", tool_call=deepcopy(call))) + call_keys.add(key) + return blocks + + +def coalesce_assistant_message(messages: Iterable[UnifiedMessage]) -> UnifiedMessage: + """Collapse item-oriented provider output into one assistant message.""" + + merged = UnifiedMessage(role="assistant") + for message in messages: + if message.role not in {"assistant", "model"}: + continue + merged.content.extend(ordered_message_blocks(message)) + merged.reasoning = message_reasoning(merged) + merged.tool_calls = message_tool_calls(merged) + return merged + + +def canonical_tool_arguments(value: Any) -> Any: + """Decode JSON tool arguments while retaining invalid source text.""" + + if not isinstance(value, str): + return deepcopy(value) + try: + return json.loads(value) + except (TypeError, ValueError): + return value + + +def tool_arguments_object(value: Any) -> dict[str, Any]: + """Return object arguments required by Anthropic and Gemini.""" + + normalized = canonical_tool_arguments(value) + if isinstance(normalized, dict): + return normalized + raise ValueError("tool arguments must be a JSON object") + + +def tool_arguments_text(value: Any) -> str: + """Return compact JSON arguments required by Chat and Responses.""" + + if value is None: + return "" + if isinstance(value, str): + return value + return json.dumps(serialize_value(value), separators=(",", ":")) + + +def tool_result_text(value: Any) -> str: + """Return the string result representation required by Chat/Responses.""" + + if isinstance(value, str): + return value + return json.dumps(serialize_value(value), separators=(",", ":")) + + +def tool_result_object(value: Any) -> dict[str, Any]: + """Return the object result representation required by Gemini.""" + + normalized = canonical_tool_arguments(value) + if isinstance(normalized, dict): + return normalized + return {"result": serialize_value(normalized)} diff --git a/src/rotator_library/protocols/gemini.py b/src/rotator_library/protocols/gemini.py new file mode 100644 index 000000000..8c572c7ae --- /dev/null +++ b/src/rotator_library/protocols/gemini.py @@ -0,0 +1,670 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Gemini generateContent protocol adapter. + +The adapter preserves Gemini-native content parts, thought signatures, safety +settings, tools, and generation configuration so later native providers can use +the same base without forcing an OpenAI-compatible intermediate shape. +""" + +from __future__ import annotations + +import json +from copy import deepcopy +from typing import Any, ClassVar, Iterable + +from .base import ProtocolAdapter +from .canonical import ( + add_conversion_warning, + canonical_stop_reason, + canonical_structured_output, + canonical_tool_arguments, + coalesce_assistant_message, + conversation_messages, + format_stop_reason, + format_structured_output, + format_tool_choice, + instruction_blocks, + is_same_protocol, + message_reasoning, + message_tool_calls, + message_tool_results, + may_emit_opaque_provider_state, + normalize_tool_result_messages, + ordered_message_blocks, + retain_supported_generation_params, + resolve_tool_result_names, + source_extensions, + tool_arguments_object, + tool_result_object, +) +from .operation import OPERATION_CHAT, OPERATION_COUNT_TOKENS, OPERATION_GENERATE, OPERATION_UNKNOWN, normalize_operation +from .validation import validate_generative_request, validate_generative_response +from .types import ( + ContentBlock, + MediaSource, + ProtocolContext, + ReasoningBlock, + ToolCall, + ToolDefinition, + ToolResult, + UnifiedMessage, + UnifiedRequest, + UnifiedResponse, + UnifiedStreamEvent, + Usage, +) + +_REQUEST_CORE_FIELDS = { + "model", + "contents", + "systemInstruction", + "system_instruction", + "tools", + "generationConfig", + "generation_config", + "safetySettings", + "safety_settings", + "toolConfig", + "tool_config", + "stream", +} + + +class GeminiProtocol(ProtocolAdapter): + """Adapter for Gemini ``generateContent`` and stream event shapes. + + Gemini parts are richer than simple chat messages. Unknown part fields remain + in ``extra`` and raw payloads are preserved so provider-specific subclasses + can refine behavior without losing data. + """ + + name: ClassVar[str] = "gemini" + aliases: ClassVar[tuple[str, ...]] = ("google_gemini", "generate_content") + supported_transports: ClassVar[tuple[str, ...]] = ("http", "sse") + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_CHAT, OPERATION_COUNT_TOKENS, "generate", "stream_generate") + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + generation_config = deepcopy(request.get("generationConfig") or request.get("generation_config") or {}) + safety_settings = deepcopy(request.get("safetySettings") or request.get("safety_settings") or []) + tool_config = deepcopy(request.get("toolConfig") or request.get("tool_config") or {}) + generation_params = _parse_gemini_generation_params(generation_config, tool_config) + if safety_settings: + generation_params["safety_settings"] = safety_settings + return UnifiedRequest( + operation=_operation_from_context(context, OPERATION_CHAT), + logical_operation=OPERATION_GENERATE, + model=str(request.get("model") or getattr(context, "model", None) or ""), + messages=resolve_tool_result_names( + normalize_tool_result_messages([self._parse_content(content) for content in request.get("contents") or []]) + ), + system=self._parse_system(request.get("systemInstruction") or request.get("system_instruction")), + tools=self._parse_tools(request.get("tools") or []), + stream=bool(request.get("stream", False)), + modalities=[str(value).lower() for value in generation_config.get("responseModalities") or []], + generation_params=generation_params, + response_format=deepcopy(generation_params.get("structured_output")), + source_protocol=self.name, + extensions={self.name: {"generationConfig": generation_config, "safetySettings": safety_settings, "toolConfig": tool_config}}, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _REQUEST_CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + validate_generative_request(unified_request, self.name, context) + preserve_source = is_same_protocol(context, self.name, unified_request.source_protocol) + emit_opaque_state = may_emit_opaque_provider_state(context, preserve_source=preserve_source) + payload: dict[str, Any] = { + "contents": [ + self._format_content(message, preserve_source=preserve_source, emit_opaque_state=emit_opaque_state) + for message in resolve_tool_result_names(deepcopy(conversation_messages(unified_request))) + ], + } + if unified_request.model: + payload["model"] = unified_request.model + instructions = instruction_blocks(unified_request) + if instructions: + payload["systemInstruction"] = { + "parts": self._format_parts( + instructions, + preserve_source=preserve_source, + emit_opaque_state=emit_opaque_state, + ) + } + generation_config, safety_settings, tool_config = self._format_generation_params(unified_request, preserve_source=preserve_source) + if generation_config: + payload["generationConfig"] = deepcopy(generation_config) + if safety_settings: + payload["safetySettings"] = deepcopy(safety_settings) + if tool_config: + payload["toolConfig"] = deepcopy(tool_config) + if unified_request.tools: + payload["tools"] = self._format_tools(unified_request.tools, preserve_source=preserve_source) + if unified_request.stream: + payload["stream"] = True + payload.update(source_extensions(unified_request.extra, context, self.name, unified_request.source_protocol)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + response = _as_dict(raw_response) + messages: list[UnifiedMessage] = [] + stop_reason = None + for candidate in response.get("candidates") or []: + if not isinstance(candidate, dict): + continue + content = candidate.get("content") if isinstance(candidate.get("content"), dict) else {} + message = self._parse_content(content) + message.extra["candidate"] = _without(candidate, {"content"}) + messages.append(message) + if candidate.get("finishReason") is not None: + stop_reason = candidate.get("finishReason") + canonical_reason = canonical_stop_reason(stop_reason) + if canonical_reason == "stop" and any(message_tool_calls(message) for message in messages): + canonical_reason = "tool_use" + return UnifiedResponse( + operation=_response_operation(response, context), + logical_operation=OPERATION_GENERATE, + id=response.get("responseId") or response.get("id"), + model=response.get("modelVersion") or getattr(context, "model", None), + messages=messages, + stop_reason=canonical_reason, + usage=self.extract_usage(response, context), + metadata={"promptFeedback": deepcopy(response.get("promptFeedback")), "modelVersion": response.get("modelVersion"), "native_stop_reason": stop_reason}, + source_protocol=self.name, + raw=deepcopy(response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"responseId", "id", "modelVersion", "candidates", "usageMetadata", "promptFeedback"}}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + if unified_response.operation == OPERATION_COUNT_TOKENS: + usage = unified_response.usage + payload = deepcopy(unified_response.extra) + # Normalized usage wins over raw preserved fields so later adapters + # can correct counts without stale provider keys shadowing them. + payload["totalTokens"] = usage.total_tokens if usage else 0 + return payload + validate_generative_response(unified_response, self.name) + preserve_source = is_same_protocol(context, self.name, unified_response.source_protocol) + emit_opaque_state = may_emit_opaque_provider_state(context, preserve_source=preserve_source) + messages = unified_response.messages if preserve_source else [coalesce_assistant_message(unified_response.messages)] + candidates = [] + for index, message in enumerate(message for message in messages if message.content or message.reasoning or message.tool_calls): + candidate = {"index": index, "content": self._format_content(message, preserve_source=preserve_source, emit_opaque_state=emit_opaque_state)} + if unified_response.stop_reason: + candidate["finishReason"] = format_stop_reason(unified_response.stop_reason, self.name) + if preserve_source: + candidate.update(deepcopy(message.extra.get("candidate") or {})) + candidates.append(candidate) + payload = { + "responseId": unified_response.id, + "modelVersion": unified_response.model, + "candidates": candidates, + "usageMetadata": self._format_usage(unified_response.usage), + "promptFeedback": deepcopy(unified_response.metadata.get("promptFeedback")), + } + payload.update(source_extensions(unified_response.extra, context, self.name, unified_response.source_protocol)) + return {k: v for k, v in payload.items() if v is not None} + + def parse_stream_event(self, raw_event: Any, context: ProtocolContext | None = None) -> UnifiedStreamEvent: + event = _decode_sse_data(raw_event) + if event == "[DONE]": + return UnifiedStreamEvent(type="done", operation=OPERATION_CHAT, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type="done", raw=deepcopy(raw_event)) + data = _as_dict(event) + if data.get("error") is not None: + return UnifiedStreamEvent(type="error", operation=OPERATION_CHAT, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type="error", error=deepcopy(data["error"]), raw=deepcopy(raw_event), extra={"payload": data}) + response = self.parse_response(data, context) + message = response.messages[0] if response.messages else None + return UnifiedStreamEvent( + type="message_delta" if message else "chunk", + operation=response.operation, + logical_operation=OPERATION_GENERATE, + source_protocol=self.name, + native_type="gemini.chunk", + delta=message, + usage=response.usage, + stop_reason=response.stop_reason, + raw=deepcopy(raw_event), + extra={"payload": data, "finish_reason": response.stop_reason}, + ) + + def extract_usage(self, raw_or_unified: Any, context: ProtocolContext | None = None) -> Usage | None: + if isinstance(raw_or_unified, (UnifiedResponse, UnifiedStreamEvent)): + return raw_or_unified.usage + payload = _as_dict(raw_or_unified) + usage = payload.get("usageMetadata") if isinstance(payload.get("usageMetadata"), dict) else payload + if not isinstance(usage, dict) or (not any(key.endswith("TokenCount") for key in usage) and "totalTokens" not in usage): + return None + input_tokens = int(usage.get("promptTokenCount") or 0) + output_tokens = int(usage.get("candidatesTokenCount") or 0) + reasoning_tokens = int(usage.get("thoughtsTokenCount") or 0) + return Usage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=int(usage.get("totalTokenCount") or usage.get("totalTokens") or input_tokens + output_tokens + reasoning_tokens), + cache_read_tokens=int(usage.get("cachedContentTokenCount") or 0), + reasoning_tokens=reasoning_tokens, + raw=deepcopy(usage), + ) + + def _parse_system(self, system: Any) -> list[ContentBlock]: + if system is None: + return [] + if isinstance(system, str): + return [ContentBlock(type="text", text=system, raw=system)] + if isinstance(system, dict): + return self._parse_parts(system.get("parts") or []) + return [] + + def _parse_content(self, content: dict[str, Any]) -> UnifiedMessage: + payload = dict(content or {}) + role = str(payload.get("role") or "model") + # Gemini uses "model" where chat protocols usually say "assistant". + normalized_role = "assistant" if role == "model" else role + message = UnifiedMessage( + role=normalized_role, + content=self._parse_parts(payload.get("parts") or []), + raw=deepcopy(content), + extra={"gemini_role": role, **_without(payload, {"role", "parts"})}, + ) + for block in message.content: + if block.tool_call: + message.tool_calls.append(block.tool_call) + if block.reasoning: + message.reasoning.append(block.reasoning) + return message + + def _format_content( + self, + message: UnifiedMessage, + *, + preserve_source: bool = True, + emit_opaque_state: bool = True, + ) -> dict[str, Any]: + role = message.extra.get("gemini_role") if preserve_source else None + role = role or ("model" if message.role in {"assistant", "model"} else "user") + parts = self._format_parts( + ordered_message_blocks(message), + preserve_source=preserve_source, + emit_opaque_state=emit_opaque_state, + ) + payload = {"role": role, "parts": parts} + if preserve_source: + payload.update({k: deepcopy(v) for k, v in message.extra.items() if k != "gemini_role"}) + return payload + + def _parse_parts(self, parts: Iterable[Any]) -> list[ContentBlock]: + blocks = [] + for part in parts: + blocks.append(self._parse_part(part)) + return blocks + + def _parse_part(self, part: Any) -> ContentBlock: + if isinstance(part, str): + return ContentBlock(type="text", text=part, raw=part) + if not isinstance(part, dict): + return ContentBlock(type="unknown", raw=deepcopy(part)) + if "text" in part: + reasoning = None + if part.get("thought") or part.get("thoughtSignature"): + reasoning = ReasoningBlock(type="reasoning", text=part.get("text"), signature=part.get("thoughtSignature"), raw=deepcopy(part), extra=_without(part, {"text", "thought", "thoughtSignature"})) + return ContentBlock(type="reasoning", text=part.get("text", ""), reasoning=reasoning, raw=deepcopy(part), extra=_without(part, {"text"})) + return ContentBlock(type="text", text=part.get("text", ""), raw=deepcopy(part), extra=_without(part, {"text"})) + if "inlineData" in part or "inline_data" in part: + source = part.get("inlineData") or part.get("inline_data") + media = _parse_gemini_media_source(source, inline=True) + return ContentBlock(type=_media_block_type(media.media_type), source=media, raw=deepcopy(part), extra=_without(part, {"inlineData", "inline_data"})) + if "fileData" in part or "file_data" in part: + source = part.get("fileData") or part.get("file_data") + media = _parse_gemini_media_source(source, inline=False) + return ContentBlock(type=_media_block_type(media.media_type), source=media, raw=deepcopy(part), extra=_without(part, {"fileData", "file_data"})) + if "functionCall" in part or "function_call" in part: + call = part.get("functionCall") or part.get("function_call") or {} + return ContentBlock(type="tool_call", tool_call=ToolCall(id=call.get("id"), name=call.get("name"), arguments=canonical_tool_arguments(call.get("args")), type="function", raw=deepcopy(call)), raw=deepcopy(part), extra=_without(part, {"functionCall", "function_call"})) + if "functionResponse" in part or "function_response" in part: + response = part.get("functionResponse") or part.get("function_response") or {} + result_content = canonical_tool_arguments(response.get("response")) + return ContentBlock(type="tool_result", tool_result=ToolResult(tool_call_id=response.get("id") or response.get("name"), name=response.get("name"), content=result_content, raw=deepcopy(response)), raw=deepcopy(part), extra=_without(part, {"functionResponse", "function_response"})) + return ContentBlock(type="unknown", raw=deepcopy(part), extra=deepcopy(part)) + + def _format_parts( + self, + blocks: Iterable[ContentBlock], + *, + preserve_source: bool = True, + emit_opaque_state: bool = True, + ) -> list[dict[str, Any]]: + parts = [] + for block in blocks: + if block.tool_call: + parts.append(self._format_tool_call(block.tool_call, preserve_source=preserve_source)) + elif block.tool_result: + parts.append(self._format_tool_result(block.tool_result, preserve_source=preserve_source)) + elif block.type in {"image", "audio", "video", "file", "document"}: + parts.append(_format_gemini_media(block, preserve_source=preserve_source)) + elif block.reasoning: + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {} + payload["text"] = block.reasoning.text or "" + payload["thought"] = True + if emit_opaque_state and block.reasoning.signature: + payload["thoughtSignature"] = block.reasoning.signature + elif not emit_opaque_state: + payload.pop("thoughtSignature", None) + parts.append(payload) + else: + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {} + payload["text"] = block.text or "" + if preserve_source: + payload.update(deepcopy(block.extra)) + parts.append(payload) + return parts + + def _parse_tools(self, tools: Iterable[dict[str, Any]]) -> list[ToolDefinition]: + parsed: list[ToolDefinition] = [] + for container_index, tool in enumerate(tools): + payload = dict(tool or {}) + declarations = payload.get("functionDeclarations") or payload.get("function_declarations") or [] + if declarations: + for index, declaration in enumerate(declarations): + if not isinstance(declaration, dict): + continue + parsed.append( + ToolDefinition( + name=str(declaration.get("name") or ""), + description=declaration.get("description"), + input_schema=deepcopy(declaration.get("parameters") or {}), + type="function", + extra={"raw_container": deepcopy(tool), "container_index": container_index, "declaration_index": index}, + ) + ) + continue + parsed.append( + ToolDefinition( + name=str(payload.get("name") or payload.get("type") or "gemini_tool"), + description=payload.get("description"), + input_schema=deepcopy(payload.get("parameters") or {}), + type=str(payload.get("type") or next(iter(payload.keys()), "tool")), + extra={"raw": deepcopy(tool)}, + ) + ) + return parsed + + def _format_tools(self, tools: Iterable[ToolDefinition], *, preserve_source: bool = True) -> list[dict[str, Any]]: + grouped: dict[int, dict[str, Any]] = {} + ungrouped: list[dict[str, Any]] = [] + for tool in tools: + raw_container = tool.extra.get("raw_container") + container_index = tool.extra.get("container_index") + declaration_index = tool.extra.get("declaration_index") + if preserve_source and isinstance(raw_container, dict) and isinstance(container_index, int) and isinstance(declaration_index, int): + container = grouped.setdefault(container_index, deepcopy(raw_container)) + declarations = container.setdefault("functionDeclarations", []) + while len(declarations) <= declaration_index: + declarations.append({}) + declaration = deepcopy(declarations[declaration_index]) if isinstance(declarations[declaration_index], dict) else {} + declaration["name"] = tool.name + if tool.description is not None: + declaration["description"] = tool.description + declaration["parameters"] = deepcopy(tool.input_schema) + declarations[declaration_index] = declaration + continue + ungrouped.append(self._format_tool(tool, preserve_source=preserve_source)) + return [grouped[index] for index in sorted(grouped)] + ungrouped + + def _format_tool(self, tool: ToolDefinition, *, preserve_source: bool = True) -> dict[str, Any]: + raw = tool.extra.get("raw") + if preserve_source and isinstance(raw, dict): + return deepcopy(raw) + return {"functionDeclarations": [{"name": tool.name, "description": tool.description, "parameters": deepcopy(tool.input_schema)}]} + + def _format_tool_call(self, call: ToolCall, *, preserve_source: bool) -> dict[str, Any]: + payload = deepcopy(call.raw) if preserve_source and isinstance(call.raw, dict) else {} + function_call = {"name": call.name or "", "args": tool_arguments_object(call.arguments)} + if call.id and not (preserve_source and call.extra.get("synthetic_id")): + function_call["id"] = call.id + payload["functionCall"] = function_call + return payload + + def _format_tool_result(self, result: ToolResult, *, preserve_source: bool) -> dict[str, Any]: + payload = deepcopy(result.raw) if preserve_source and isinstance(result.raw, dict) else {} + result_content = {"error": result.content} if result.is_error else result.content + response = {"name": result.name or result.tool_call_id or "", "response": tool_result_object(result_content)} + if result.tool_call_id and result.name and not (preserve_source and result.extra.get("synthetic_tool_call_id")): + response["id"] = result.tool_call_id + payload["functionResponse"] = response + return payload + + def _format_generation_params(self, request: UnifiedRequest, *, preserve_source: bool) -> tuple[dict[str, Any], list[Any], dict[str, Any]]: + source = request.extensions.get(self.name, {}) if preserve_source else {} + generation = deepcopy(source.get("generationConfig") or {}) + safety = deepcopy(source.get("safetySettings") or []) + tool_config = deepcopy(source.get("toolConfig") or {}) + params = deepcopy(request.generation_params) + canonical_safety = params.pop("safety_settings", None) + if canonical_safety is not None: + safety = deepcopy(canonical_safety) + if request.modalities: + generation["responseModalities"] = [str(value).upper() for value in request.modalities] + mapping = { + "max_output_tokens": "maxOutputTokens", + "stop_sequences": "stopSequences", + "top_p": "topP", + "top_k": "topK", + "temperature": "temperature", + "candidate_count": "candidateCount", + "seed": "seed", + "frequency_penalty": "frequencyPenalty", + "presence_penalty": "presencePenalty", + } + for canonical, wire in mapping.items(): + if canonical in params: + generation[wire] = params.pop(canonical) + structured = params.pop("structured_output", None) + if isinstance(structured, dict): + if structured.get("strict") is False: + add_conversion_warning( + request, + code="structured_output_strictness_strengthened", + message="Gemini enforces its response schema; explicit strict=false was strengthened", + target_protocol=self.name, + field="structured_output.strict", + ) + generation.update( + { + key: value + for key, value in format_structured_output(structured, self.name).items() + if value is not None + } + ) + reasoning = params.pop("reasoning", None) + if isinstance(reasoning, dict): + thinking: dict[str, Any] = {} + if reasoning.get("budget_tokens") is not None: + thinking["thinkingBudget"] = reasoning["budget_tokens"] + if reasoning.get("include_thoughts") is not None: + thinking["includeThoughts"] = reasoning["include_thoughts"] + if thinking: + generation["thinkingConfig"] = thinking + tool_choice = params.pop("tool_choice", None) + if tool_choice is not None: + tool_config = format_tool_choice(tool_choice, self.name) + retain_supported_generation_params( + request, + params, + supported=set(), + target_protocol=self.name, + ) + return generation, safety, tool_config + + def _format_usage(self, usage: Usage | None) -> dict[str, int] | None: + if usage is None: + return None + payload = { + "promptTokenCount": usage.input_tokens, + "candidatesTokenCount": usage.output_tokens, + "totalTokenCount": usage.total_tokens, + } + if usage.reasoning_tokens: + payload["thoughtsTokenCount"] = usage.reasoning_tokens + if usage.cache_read_tokens: + payload["cachedContentTokenCount"] = usage.cache_read_tokens + return payload + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return deepcopy(value) + if hasattr(value, "model_dump"): + return value.model_dump() + if hasattr(value, "dict"): + return value.dict() + return {} + + +def _operation_from_context(context: ProtocolContext | None, default: str) -> str: + supported = {OPERATION_CHAT, OPERATION_COUNT_TOKENS, "generate", "stream_generate"} + if context and isinstance(context.provider_options, dict): + operation = normalize_operation(context.provider_options.get("operation")) + if operation in supported: + return operation + if context and isinstance(context.metadata, dict): + operation = normalize_operation(context.metadata.get("operation")) + if operation in supported: + return operation + return default + + +def _response_operation(response: dict[str, Any], context: ProtocolContext | None) -> str: + requested = _operation_from_context(context, OPERATION_CHAT) + if requested == OPERATION_COUNT_TOKENS: + return OPERATION_COUNT_TOKENS + if "totalTokens" in response and "candidates" not in response: + return OPERATION_COUNT_TOKENS + return requested if requested in {"generate", "stream_generate"} else OPERATION_CHAT + + +def _decode_sse_data(raw_event: Any) -> Any: + from .streaming import decode_sse_data + + return decode_sse_data(raw_event) + + +def _without(payload: dict[str, Any], keys: set[str]) -> dict[str, Any]: + return {k: deepcopy(v) for k, v in payload.items() if k not in keys} + + +def _parse_gemini_generation_params(generation: dict[str, Any], tool_config: dict[str, Any]) -> dict[str, Any]: + """Normalize Gemini generation and tool controls.""" + + params: dict[str, Any] = {} + mapping = { + "maxOutputTokens": "max_output_tokens", + "stopSequences": "stop_sequences", + "topP": "top_p", + "topK": "top_k", + "temperature": "temperature", + "candidateCount": "candidate_count", + "seed": "seed", + "frequencyPenalty": "frequency_penalty", + "presencePenalty": "presence_penalty", + } + for wire, canonical in mapping.items(): + if wire in generation: + params[canonical] = deepcopy(generation[wire]) + if generation.get("responseMimeType") is not None or generation.get("responseJsonSchema") is not None or generation.get("responseSchema") is not None: + params["structured_output"] = canonical_structured_output( + { + "type": "json_schema" if generation.get("responseJsonSchema") is not None or generation.get("responseSchema") is not None else "json_object", + "schema": deepcopy(generation.get("responseJsonSchema") or generation.get("responseSchema")), + }, + "gemini", + ) + thinking = generation.get("thinkingConfig") + if isinstance(thinking, dict): + params["reasoning"] = { + "budget_tokens": thinking.get("thinkingBudget"), + "include_thoughts": thinking.get("includeThoughts"), + } + if tool_config: + params["tool_choice"] = _parse_gemini_tool_choice(tool_config) + return params + + +def _parse_gemini_tool_choice(tool_config: dict[str, Any]) -> Any: + """Normalize Gemini function-calling mode and allow-list.""" + + config = tool_config.get("functionCallingConfig") or tool_config.get("function_calling_config") or {} + if not isinstance(config, dict): + return deepcopy(tool_config) + mode = str(config.get("mode") or "AUTO").lower() + names = deepcopy(config.get("allowedFunctionNames") or config.get("allowed_function_names") or []) + if mode == "none": + return {"mode": "none"} + if mode == "any" and len(names) == 1: + return {"mode": "named", "name": names[0]} + if mode == "any": + return {"mode": "required", "allowed_names": names} + return {"mode": "auto"} + + +def _parse_gemini_media_source(value: Any, *, inline: bool) -> MediaSource: + """Normalize Gemini inlineData and fileData sources.""" + + payload = value if isinstance(value, dict) else {} + return MediaSource( + kind="base64" if inline else "url", + media_type=payload.get("mimeType") or payload.get("mime_type"), + url=payload.get("fileUri") or payload.get("file_uri"), + data=payload.get("data"), + file_id=payload.get("fileId") or payload.get("file_id"), + raw=deepcopy(value), + extra=_without(payload, {"mimeType", "mime_type", "fileUri", "file_uri", "data", "fileId", "file_id"}), + ) + + +def _media_block_type(media_type: Any) -> str: + """Return the canonical media category for a MIME type.""" + + value = str(media_type or "").lower() + if value.startswith("image/"): + return "image" + if value.startswith("audio/"): + return "audio" + if value.startswith("video/"): + return "video" + return "file" + + +def _coerce_media_source(value: Any) -> MediaSource: + """Coerce legacy dictionaries into a canonical media source.""" + + if isinstance(value, MediaSource): + return value + if isinstance(value, str): + return MediaSource(kind="url", url=value, raw=value) + payload = value if isinstance(value, dict) else {} + return MediaSource( + kind="base64" if payload.get("data") else "url", + media_type=payload.get("mimeType") or payload.get("mime_type") or payload.get("media_type"), + url=payload.get("fileUri") or payload.get("file_uri") or payload.get("url"), + data=payload.get("data"), + file_id=payload.get("fileId") or payload.get("file_id"), + raw=deepcopy(value), + ) + + +def _format_gemini_media(block: ContentBlock, *, preserve_source: bool) -> dict[str, Any]: + """Format canonical media as a Gemini content part.""" + + source = _coerce_media_source(block.source) + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {} + if source.data: + payload["inlineData"] = {"mimeType": source.media_type or "application/octet-stream", "data": source.data} + else: + file_data: dict[str, Any] = { + "mimeType": source.media_type or "application/octet-stream", + "fileUri": source.url or source.file_id or "", + } + payload["fileData"] = file_data + return payload diff --git a/src/rotator_library/protocols/litellm_fallback.py b/src/rotator_library/protocols/litellm_fallback.py new file mode 100644 index 000000000..6c66352c6 --- /dev/null +++ b/src/rotator_library/protocols/litellm_fallback.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Explicit fallback protocol for the existing LiteLLM-shaped path. + +This adapter intentionally does very little beyond preserving raw payloads. It +gives later execution and transform-logging code a named protocol path whenever +native adapters do not yet cover a provider or request shape. +""" + +from __future__ import annotations + +from typing import ClassVar + +from .base import ProtocolAdapter +from .operation import OPERATION_UNKNOWN + + +class LiteLLMFallbackProtocol(ProtocolAdapter): + """Protocol adapter that keeps the current LiteLLM-compatible payload shape. + + Providers should prefer native protocol adapters when available. This class + remains as a safe compatibility base for unsupported provider shapes and as + a clear marker in future transaction transform traces. + """ + + name: ClassVar[str] = "litellm_fallback" + aliases: ClassVar[tuple[str, ...]] = ("litellm", "fallback") + supported_transports: ClassVar[tuple[str, ...]] = ("http", "sse") + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_UNKNOWN,) diff --git a/src/rotator_library/protocols/mcp.py b/src/rotator_library/protocols/mcp.py new file mode 100644 index 000000000..e51a1ec5c --- /dev/null +++ b/src/rotator_library/protocols/mcp.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""MCP JSON-RPC carrier protocol adapter. + +This is not a full MCP proxy implementation. It gives the native protocol layer a +lossless request/response carrier for future MCP gateway work, keeping method, +params, ids, results, errors, and JSON-RPC batch arrays intact for transform +logging and routing. +""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any, ClassVar + +from .base import ProtocolAdapter +from .operation import OPERATION_MCP +from .types import ProtocolContext, UnifiedRequest, UnifiedResponse, UnifiedStreamEvent + +_REQUEST_CORE_FIELDS = {"jsonrpc", "id", "method", "params"} +_RESPONSE_CORE_FIELDS = {"jsonrpc", "id", "result", "error"} + + +class MCPProtocol(ProtocolAdapter): + """Adapter for MCP-style JSON-RPC request and response envelopes.""" + + name: ClassVar[str] = "mcp" + aliases: ClassVar[tuple[str, ...]] = ("model_context_protocol", "jsonrpc_mcp") + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_MCP,) + supported_transports: ClassVar[tuple[str, ...]] = ("http",) + future_transports: ClassVar[tuple[str, ...]] = ("sse", "websocket") + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + if isinstance(raw_request, list): + return UnifiedRequest(operation=OPERATION_MCP, input=deepcopy(raw_request), metadata={"batch": True}, raw=deepcopy(raw_request)) + request = dict(raw_request or {}) + metadata = { + "jsonrpc": request.get("jsonrpc", "2.0"), + "method": request.get("method"), + "has_id": "id" in request, + "has_params": "params" in request, + } + if "id" in request: + metadata["id"] = deepcopy(request.get("id")) + return UnifiedRequest( + operation=OPERATION_MCP, + input=deepcopy(request.get("params")) if "params" in request else None, + metadata=metadata, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _REQUEST_CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> Any: + if unified_request.metadata.get("batch"): + return deepcopy(unified_request.input or []) + payload = { + "jsonrpc": unified_request.metadata.get("jsonrpc", "2.0"), + "method": unified_request.metadata.get("method"), + } + if unified_request.metadata.get("has_params", True): + payload["params"] = deepcopy(unified_request.input) + if unified_request.metadata.get("has_id"): + payload["id"] = deepcopy(unified_request.metadata.get("id")) + payload.update(deepcopy(unified_request.extra)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + if isinstance(raw_response, list): + return UnifiedResponse(operation=OPERATION_MCP, data=deepcopy(raw_response), metadata={"batch": True}, raw=deepcopy(raw_response)) + response = raw_response if isinstance(raw_response, dict) else {} + metadata = {"jsonrpc": response.get("jsonrpc", "2.0"), "has_id": "id" in response} + if "id" in response: + metadata["id"] = deepcopy(response.get("id")) + extra = {k: deepcopy(v) for k, v in response.items() if k not in _RESPONSE_CORE_FIELDS} + if "error" in response: + # JSON-RPC errors are not provider exceptions here; they are protocol + # payloads that must survive transform logging and response rebuilds. + extra["error"] = deepcopy(response["error"]) + return UnifiedResponse( + operation=OPERATION_MCP, + data=[deepcopy(response["result"])] if "result" in response else [], + metadata=metadata, + raw=deepcopy(raw_response), + extra=extra, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> Any: + if unified_response.metadata.get("batch"): + return deepcopy(unified_response.data) + payload = {"jsonrpc": unified_response.metadata.get("jsonrpc", "2.0")} + if unified_response.metadata.get("has_id"): + payload["id"] = deepcopy(unified_response.metadata.get("id")) + if unified_response.extra.get("error") is not None: + payload["error"] = deepcopy(unified_response.extra["error"]) + else: + payload["result"] = deepcopy(unified_response.data[0] if unified_response.data else None) + payload.update({k: deepcopy(v) for k, v in unified_response.extra.items() if k != "error"}) + return payload + + def parse_stream_event(self, raw_event: Any, context: ProtocolContext | None = None) -> UnifiedStreamEvent: + data = raw_event if isinstance(raw_event, dict) else {"event": raw_event} + return UnifiedStreamEvent(type=str(data.get("method") or data.get("type") or "message"), operation=OPERATION_MCP, error=deepcopy(data.get("error")), raw=deepcopy(raw_event), extra=deepcopy(data)) diff --git a/src/rotator_library/protocols/ollama.py b/src/rotator_library/protocols/ollama.py new file mode 100644 index 000000000..748cca1c7 --- /dev/null +++ b/src/rotator_library/protocols/ollama.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Ollama-native chat, generate, and embeddings protocol adapter.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from typing import Any, ClassVar + +from .base import ProtocolAdapter +from .operation import OPERATION_EMBEDDINGS, OPERATION_OLLAMA_CHAT, OPERATION_OLLAMA_GENERATE, normalize_operation +from .types import ContentBlock, ProtocolContext, UnifiedMessage, UnifiedRequest, UnifiedResponse, UnifiedStreamEvent, Usage, first_text + +_OPTION_FIELDS = {"options", "format", "keep_alive", "template", "context", "raw", "suffix"} +_CORE_FIELDS = {"operation", "model", "messages", "prompt", "input", "stream", "system", *_OPTION_FIELDS} + + +class OllamaProtocol(ProtocolAdapter): + """Adapter for Ollama `/api/chat`, `/api/generate`, and embeddings shapes.""" + + name: ClassVar[str] = "ollama" + aliases: ClassVar[tuple[str, ...]] = ("ollama_native",) + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_OLLAMA_CHAT, OPERATION_OLLAMA_GENERATE, OPERATION_EMBEDDINGS) + supported_transports: ClassVar[tuple[str, ...]] = ("http", "jsonl") + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + operation = _ollama_operation(request, context) + input_field = "input" + if operation == OPERATION_EMBEDDINGS and "input" not in request and "prompt" in request: + # Older Ollama embeddings endpoints use `prompt`; newer endpoints use + # `input`. Preserve the original spelling so round-trips are lossless. + input_field = "prompt" + return UnifiedRequest( + operation=operation, + model=str(request.get("model") or getattr(context, "model", None) or ""), + messages=[_message_from_ollama(message) for message in request.get("messages") or []], + system=[ContentBlock(type="text", text=str(request["system"]))] if "system" in request else [], + stream=bool(request.get("stream", False)), + input=deepcopy(request.get(input_field) if operation == OPERATION_EMBEDDINGS else request.get("prompt")), + generation_params={k: deepcopy(request[k]) for k in _OPTION_FIELDS if k in request}, + metadata={**({"embedding_input_field": input_field} if operation == OPERATION_EMBEDDINGS else {}), "has_stream": "stream" in request}, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + payload: dict[str, Any] = {"model": unified_request.model} + if unified_request.stream or unified_request.metadata.get("has_stream"): + payload["stream"] = unified_request.stream + if unified_request.operation == OPERATION_OLLAMA_CHAT: + payload["messages"] = [_message_to_ollama(message) for message in unified_request.messages] + elif unified_request.operation == OPERATION_EMBEDDINGS: + input_field = str(unified_request.metadata.get("embedding_input_field") or "input") + payload[input_field] = deepcopy(unified_request.input) + else: + payload["prompt"] = deepcopy(unified_request.input) + if unified_request.system: + payload["system"] = "".join(block.text or "" for block in unified_request.system if block.type == "text") + payload.update(deepcopy(unified_request.generation_params)) + payload.update(deepcopy(unified_request.extra)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + response = raw_response if isinstance(raw_response, dict) else {} + output = [] + messages: list[UnifiedMessage] = [] + if isinstance(response.get("message"), dict): + message = _message_from_ollama(response["message"]) + messages.append(message) + output.append(message.to_dict()) + elif "response" in response: + output.append(response.get("response")) + return UnifiedResponse( + operation=_ollama_operation(response, context), + model=response.get("model") or getattr(context, "model", None), + messages=messages, + output=output, + data=deepcopy(response.get("embeddings") or response.get("embedding") or []), + usage=_ollama_usage(response), + raw=deepcopy(raw_response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"model", "message", "response", "embeddings", "embedding", "prompt_eval_count", "eval_count", "total_duration", "load_duration", "prompt_eval_duration", "eval_duration"}}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + """Format a unified response back to an Ollama response shape. + + Ollama responses are often mutated by adapters after parsing. Do not + return `raw` wholesale here; rebuild the public fields from unified state + and then merge preserved extras/timing values. + """ + + payload = deepcopy(unified_response.extra) + if unified_response.model: + payload["model"] = unified_response.model + operation = unified_response.operation or _ollama_operation(payload, context) + if operation == OPERATION_OLLAMA_CHAT: + if unified_response.messages: + payload["message"] = _message_to_ollama(unified_response.messages[0]) + elif operation == OPERATION_EMBEDDINGS: + raw = unified_response.raw if isinstance(unified_response.raw, dict) else {} + key = "embedding" if "embedding" in raw and "embeddings" not in raw else "embeddings" + payload[key] = deepcopy(unified_response.data) + else: + payload["response"] = _ollama_response_text(unified_response) + if unified_response.usage and isinstance(unified_response.usage.raw, dict): + for key, value in unified_response.usage.raw.items(): + if key.endswith("duration"): + payload[key] = deepcopy(value) + if unified_response.usage.input_tokens: + payload["prompt_eval_count"] = unified_response.usage.input_tokens + if unified_response.usage.output_tokens: + payload["eval_count"] = unified_response.usage.output_tokens + return {k: v for k, v in payload.items() if v is not None} + + def parse_stream_event(self, raw_event: Any, context: ProtocolContext | None = None) -> UnifiedStreamEvent: + data = _json_event(raw_event) + if not isinstance(data, dict): + return UnifiedStreamEvent(type="metadata", operation=OPERATION_OLLAMA_GENERATE, raw=deepcopy(raw_event), extra={"unparsed": True}) + event_type = "done" if data.get("done") else "message_delta" + delta = None + if isinstance(data.get("message"), dict): + delta = _message_from_ollama(data["message"]) + elif data.get("response"): + delta = UnifiedMessage(role="assistant", content=[ContentBlock(type="text", text=str(data["response"]))]) + return UnifiedStreamEvent(type=event_type, operation=_ollama_operation(data, context), delta=delta, usage=_ollama_usage(data), raw=deepcopy(raw_event), extra=deepcopy(data)) + + +def _ollama_operation(request: dict[str, Any], context: ProtocolContext | None = None) -> str: + explicit = normalize_operation(request.get("operation")) + if explicit in {OPERATION_OLLAMA_CHAT, OPERATION_OLLAMA_GENERATE, OPERATION_EMBEDDINGS}: + return explicit + if context and isinstance(context.provider_options, dict): + operation = normalize_operation(context.provider_options.get("operation")) + if operation in {OPERATION_OLLAMA_CHAT, OPERATION_OLLAMA_GENERATE, OPERATION_EMBEDDINGS}: + return operation + if context and isinstance(context.metadata, dict): + operation = normalize_operation(context.metadata.get("operation")) + if operation in {OPERATION_OLLAMA_CHAT, OPERATION_OLLAMA_GENERATE, OPERATION_EMBEDDINGS}: + return operation + if "messages" in request or "message" in request: + return OPERATION_OLLAMA_CHAT + if "prompt" in request and request.get("endpoint") != "embeddings": + return OPERATION_OLLAMA_GENERATE + if "embeddings" in request or "embedding" in request or "input" in request or request.get("endpoint") == "embeddings": + return OPERATION_EMBEDDINGS + return OPERATION_OLLAMA_GENERATE + + +def _message_from_ollama(message: dict[str, Any]) -> UnifiedMessage: + return UnifiedMessage(role=str(message.get("role") or "assistant"), content=[ContentBlock(type="text", text=str(message.get("content") or ""))], raw=deepcopy(message), extra={k: deepcopy(v) for k, v in message.items() if k not in {"role", "content"}}) + + +def _message_to_ollama(message: UnifiedMessage) -> dict[str, Any]: + payload = {"role": message.role, "content": "".join(block.text or "" for block in message.content if block.type == "text")} + payload.update(deepcopy(message.extra)) + return payload + + +def _ollama_response_text(response: UnifiedResponse) -> str: + if response.output: + return "".join(str(item) for item in response.output if item is not None) + if response.messages: + text = first_text(response.messages[0].content) + if text is not None: + return text + return "" + + +def _json_event(raw_event: Any) -> Any: + if isinstance(raw_event, dict): + return raw_event + if not isinstance(raw_event, str): + return None + text = raw_event.strip() + if text.startswith("data:"): + text = text[5:].strip() + try: + return json.loads(text) + except json.JSONDecodeError: + return None + + +def _ollama_usage(response: dict[str, Any]) -> Usage | None: + prompt_tokens = int(response.get("prompt_eval_count") or 0) + output_tokens = int(response.get("eval_count") or 0) + if not prompt_tokens and not output_tokens: + return None + return Usage(input_tokens=prompt_tokens, output_tokens=output_tokens, raw={k: deepcopy(v) for k, v in response.items() if k.endswith("count") or k.endswith("duration")}) diff --git a/src/rotator_library/protocols/openai_audio.py b/src/rotator_library/protocols/openai_audio.py new file mode 100644 index 000000000..392f5b19e --- /dev/null +++ b/src/rotator_library/protocols/openai_audio.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""OpenAI-compatible audio transcription/translation and speech protocol adapter.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any, ClassVar + +from .base import ProtocolAdapter +from .operation import OPERATION_AUDIO_TRANSCRIPTION, OPERATION_AUDIO_TRANSLATION, OPERATION_SPEECH, normalize_operation +from .types import ProtocolContext, UnifiedRequest, UnifiedResponse + +_AUDIO_OPTION_FIELDS = {"language", "response_format", "temperature", "timestamp_granularities"} +_SPEECH_OPTION_FIELDS = {"voice", "response_format", "speed"} +_CORE_FIELDS = {"operation", "model", "file", "input", "prompt", *_AUDIO_OPTION_FIELDS, *_SPEECH_OPTION_FIELDS} + + +class OpenAIAudioProtocol(ProtocolAdapter): + """Adapter for OpenAI audio transcription/translation and speech requests. + + Audio file and generated-audio bytes are intentionally preserved instead of + interpreted. Transports/providers decide multipart and binary handling; this + protocol only records enough structure for routing, tracing, and tests. + """ + + name: ClassVar[str] = "openai_audio" + aliases: ClassVar[tuple[str, ...]] = ("audio", "audio_transcription", "speech", "tts") + supported_operations: ClassVar[tuple[str, ...]] = ( + OPERATION_AUDIO_TRANSCRIPTION, + OPERATION_AUDIO_TRANSLATION, + OPERATION_SPEECH, + ) + supported_transports: ClassVar[tuple[str, ...]] = ("http",) + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + operation = _audio_operation(request, context) + files = [] + if "file" in request: + files.append({"field": "file", "value": deepcopy(request["file"])}) + options = _SPEECH_OPTION_FIELDS if operation == OPERATION_SPEECH else _AUDIO_OPTION_FIELDS + return UnifiedRequest( + operation=operation, + model=str(request.get("model") or getattr(context, "model", None) or ""), + input=deepcopy(request.get("input") if operation == OPERATION_SPEECH else request.get("prompt")), + files=files, + generation_params={k: deepcopy(request[k]) for k in options if k in request}, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + payload: dict[str, Any] = {"model": unified_request.model} + if unified_request.operation == OPERATION_SPEECH: + payload["input"] = deepcopy(unified_request.input) + elif unified_request.input is not None: + payload["prompt"] = deepcopy(unified_request.input) + for file_entry in unified_request.files: + if isinstance(file_entry, dict) and file_entry.get("field"): + payload[str(file_entry["field"])] = deepcopy(file_entry.get("value")) + payload.update(deepcopy(unified_request.generation_params)) + payload.update(deepcopy(unified_request.extra)) + return payload + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> Any: + """Format audio responses without forcing binary/text payloads into JSON. + + Speech endpoints can return raw audio bytes while transcription endpoints + can return JSON or plain text depending on `response_format`. Returning + the preserved raw body for non-dict responses keeps the protocol adapter + honest until a transport layer decides headers and streaming behavior. + """ + + if unified_response.raw is not None and not isinstance(unified_response.raw, dict): + return deepcopy(unified_response.raw) + payload: dict[str, Any] = {} + if unified_response.output: + payload["text"] = deepcopy(unified_response.output[0]) + if unified_response.data: + payload["data"] = deepcopy(unified_response.data) + if unified_response.content_type: + payload["content_type"] = unified_response.content_type + payload.update(deepcopy(unified_response.extra)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + if isinstance(raw_response, dict): + response = raw_response + operation = normalize_operation(response.get("operation")) + if operation == "unknown": + operation = OPERATION_AUDIO_TRANSCRIPTION + return UnifiedResponse( + operation=_context_operation(context, operation), + model=response.get("model") or getattr(context, "model", None), + output=[deepcopy(response["text"])] if "text" in response else [], + data=deepcopy(response.get("data") or []), + content_type=response.get("content_type") or "application/json", + raw=deepcopy(raw_response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"operation", "model", "text", "data", "content_type"}}, + ) + content_type = "text/plain" if isinstance(raw_response, str) else "application/octet-stream" + # Providers normally return transcription text as JSON/text and speech as + # bytes. Context can still override this heuristic for unusual endpoints. + default_operation = OPERATION_AUDIO_TRANSCRIPTION if isinstance(raw_response, str) else OPERATION_SPEECH + return UnifiedResponse(operation=_context_operation(context, default_operation), content_type=content_type, raw=deepcopy(raw_response), output=[deepcopy(raw_response)] if isinstance(raw_response, str) else []) + + +def _audio_operation(request: dict[str, Any], context: ProtocolContext | None = None) -> str: + explicit = normalize_operation(request.get("operation")) + if explicit in {OPERATION_AUDIO_TRANSCRIPTION, OPERATION_AUDIO_TRANSLATION, OPERATION_SPEECH}: + return explicit + contextual = _context_operation(context, "unknown") + if contextual in {OPERATION_AUDIO_TRANSCRIPTION, OPERATION_AUDIO_TRANSLATION, OPERATION_SPEECH}: + return contextual + if "voice" in request or ("input" in request and "file" not in request): + return OPERATION_SPEECH + return OPERATION_AUDIO_TRANSCRIPTION + + +def _context_operation(context: ProtocolContext | None, default: str) -> str: + if context and isinstance(context.provider_options, dict): + operation = normalize_operation(context.provider_options.get("operation")) + if operation in {OPERATION_AUDIO_TRANSCRIPTION, OPERATION_AUDIO_TRANSLATION, OPERATION_SPEECH}: + return operation + if context and isinstance(context.metadata, dict): + operation = normalize_operation(context.metadata.get("operation")) + if operation in {OPERATION_AUDIO_TRANSCRIPTION, OPERATION_AUDIO_TRANSLATION, OPERATION_SPEECH}: + return operation + return default diff --git a/src/rotator_library/protocols/openai_chat.py b/src/rotator_library/protocols/openai_chat.py new file mode 100644 index 000000000..70f3dddf0 --- /dev/null +++ b/src/rotator_library/protocols/openai_chat.py @@ -0,0 +1,774 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""OpenAI Chat Completions protocol adapter. + +The adapter models the common OpenAI-compatible chat shape used by many current +providers. It is a reusable base, not a final authority: providers can subclass +or override pieces when they need non-standard fields, stricter ordering, or +different stream semantics. +""" + +from __future__ import annotations + +import json +from copy import deepcopy +from typing import Any, ClassVar, Iterable, Optional + +from .base import ProtocolAdapter +from .canonical import ( + canonical_stop_reason, + canonical_structured_output, + canonical_tool_arguments, + coalesce_assistant_message, + canonical_tool_choice, + format_stop_reason, + format_structured_output, + format_tool_choice, + conversation_messages, + instruction_messages, + is_same_protocol, + retain_supported_generation_params, + resolve_tool_result_names, + source_extensions, + tool_arguments_text, +) +from .operation import OPERATION_CHAT, OPERATION_GENERATE +from .validation import validate_generative_request, validate_generative_response +from .types import ( + ContentBlock, + CostDetails, + MediaSource, + ProtocolContext, + ReasoningBlock, + ToolCall, + ToolDefinition, + ToolResult, + UnifiedMessage, + UnifiedRequest, + UnifiedResponse, + UnifiedStreamEvent, + Usage, + first_text, + serialize_value, + text_blocks, +) + +_GENERATION_PARAMS = { + "frequency_penalty", + "logit_bias", + "logprobs", + "max_completion_tokens", + "max_tokens", + "n", + "parallel_tool_calls", + "presence_penalty", + "reasoning_effort", + "seed", + "service_tier", + "stop", + "stream_options", + "temperature", + "tool_choice", + "top_logprobs", + "top_p", + "user", +} + +_REQUEST_CORE_FIELDS = { + "model", + "messages", + "modalities", + "audio", + "tools", + "stream", + "response_format", + "metadata", + *_GENERATION_PARAMS, +} + + +class OpenAIChatProtocol(ProtocolAdapter): + """Adapter for OpenAI Chat Completions request, response, and stream chunks. + + Unknown OpenAI-compatible extension fields are preserved in ``extra`` so a + custom provider can still use them through later adapter or field-cache + phases. Lossy conversions are avoided unless the source shape itself uses a + compact representation, such as string message content. + """ + + name: ClassVar[str] = "openai_chat" + aliases: ClassVar[tuple[str, ...]] = ( + "openai", + "chat", + "chat_completions", + "openai_chat_completions", + ) + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_CHAT,) + supported_transports: ClassVar[tuple[str, ...]] = ("http", "sse") + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + messages = resolve_tool_result_names([self._parse_message(message) for message in request.get("messages") or []]) + tools = [self._parse_tool_definition(tool) for tool in request.get("tools") or []] + source_generation_params = {k: deepcopy(request[k]) for k in _GENERATION_PARAMS if k in request} + generation_params = _parse_openai_generation_params(source_generation_params) + structured_output = canonical_structured_output(request.get("response_format"), self.name) + if structured_output: + generation_params["structured_output"] = structured_output + if "tool_choice" in generation_params: + generation_params["tool_choice"] = canonical_tool_choice(generation_params["tool_choice"], self.name) + if request.get("audio") is not None: + generation_params["audio_output"] = deepcopy(request["audio"]) + extra = {k: deepcopy(v) for k, v in request.items() if k not in _REQUEST_CORE_FIELDS} + + return UnifiedRequest( + operation=OPERATION_CHAT, + logical_operation=OPERATION_GENERATE, + model=str(request.get("model") or getattr(context, "model", None) or ""), + messages=messages, + tools=tools, + stream=bool(request.get("stream", False)), + modalities=[str(value).lower() for value in request.get("modalities") or []], + generation_params=generation_params, + response_format=structured_output, + metadata=deepcopy(request.get("metadata") or {}), + source_protocol=self.name, + extensions={self.name: {"generation_params": source_generation_params, "response_format": deepcopy(request.get("response_format"))}}, + raw=deepcopy(raw_request), + extra=extra, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + validate_generative_request(unified_request, self.name, context) + preserve_source = is_same_protocol(context, self.name, unified_request.source_protocol) + payload: dict[str, Any] = { + "model": unified_request.model, + "messages": self._format_request_messages( + [*instruction_messages(unified_request), *conversation_messages(unified_request)], + preserve_source=preserve_source, + ), + } + if unified_request.tools: + payload["tools"] = [self._format_tool_definition(tool, preserve_source=preserve_source) for tool in unified_request.tools] + if unified_request.stream: + payload["stream"] = True + if unified_request.modalities: + payload["modalities"] = deepcopy(unified_request.modalities) + if unified_request.metadata: + payload["metadata"] = deepcopy(unified_request.metadata) + payload.update(self._format_generation_params(unified_request, preserve_source=preserve_source)) + payload.update(source_extensions(unified_request.extra, context, self.name, unified_request.source_protocol)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + response = _as_dict(raw_response) + messages: list[UnifiedMessage] = [] + stop_reason = None + for choice in response.get("choices") or []: + if not isinstance(choice, dict): + continue + message_payload = choice.get("message") or {} + if message_payload: + messages.append(self._parse_message(message_payload)) + if choice.get("finish_reason") is not None: + stop_reason = choice.get("finish_reason") + + return UnifiedResponse( + operation=OPERATION_CHAT, + logical_operation=OPERATION_GENERATE, + id=response.get("id"), + model=response.get("model") or getattr(context, "model", None), + messages=messages, + stop_reason=canonical_stop_reason(stop_reason), + usage=self.extract_usage(response, context), + metadata={ + "object": response.get("object"), + "created": response.get("created"), + "system_fingerprint": response.get("system_fingerprint"), + "native_stop_reason": stop_reason, + }, + source_protocol=self.name, + raw=deepcopy(response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"id", "object", "created", "model", "choices", "usage", "system_fingerprint"}}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + validate_generative_response(unified_response, self.name) + preserve_source = is_same_protocol(context, self.name, unified_response.source_protocol) + messages = unified_response.messages if preserve_source else [coalesce_assistant_message(unified_response.messages)] + choices = [] + for index, message in enumerate(messages): + choices.append( + { + "index": index, + "message": _format_response_message( + self._format_message(message, preserve_source=preserve_source), + message, + ), + "finish_reason": format_stop_reason(unified_response.stop_reason, self.name), + } + ) + payload = { + "id": unified_response.id, + "object": unified_response.metadata.get("object", "chat.completion"), + "created": unified_response.metadata.get("created"), + "model": unified_response.model, + "choices": choices, + "usage": _format_openai_usage(unified_response.usage), + } + payload.update(source_extensions(unified_response.extra, context, self.name, unified_response.source_protocol)) + return {k: v for k, v in payload.items() if v is not None} + + def parse_stream_event(self, raw_event: Any, context: ProtocolContext | None = None) -> UnifiedStreamEvent: + event = _decode_sse_data(raw_event) + if event == "[DONE]": + return UnifiedStreamEvent(type="done", operation=OPERATION_CHAT, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type="done", raw=deepcopy(raw_event)) + data = _as_dict(event) + if data.get("error") is not None: + return UnifiedStreamEvent(type="error", operation=OPERATION_CHAT, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type="error", error=deepcopy(data["error"]), raw=deepcopy(raw_event), extra={"payload": data}) + + delta_message = None + finish_reason = None + for choice in data.get("choices") or []: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") or {} + if delta: + delta_message = self._parse_message({"role": delta.get("role", "assistant"), **delta}) + finish_reason = choice.get("finish_reason") if choice.get("finish_reason") is not None else finish_reason + break + + usage = self.extract_usage(data, context) + return UnifiedStreamEvent( + type="message_delta" if delta_message else "chunk", + operation=OPERATION_CHAT, + logical_operation=OPERATION_GENERATE, + source_protocol=self.name, + native_type="chat.completion.chunk", + delta=delta_message, + usage=usage, + stop_reason=canonical_stop_reason(finish_reason), + raw=deepcopy(raw_event), + extra={ + "id": data.get("id"), + "model": data.get("model"), + "finish_reason": canonical_stop_reason(finish_reason), + "payload": data, + }, + ) + + def format_stream_event(self, unified_event: UnifiedStreamEvent, context: ProtocolContext | None = None) -> Any: + from .streaming import format_canonical_stream_event + + return format_canonical_stream_event(unified_event, self.name, context) + + def extract_usage(self, raw_or_unified: Any, context: ProtocolContext | None = None) -> Usage | None: + if isinstance(raw_or_unified, (UnifiedResponse, UnifiedStreamEvent)): + return raw_or_unified.usage + payload = _as_dict(raw_or_unified) + usage = payload.get("usage") + if not isinstance(usage, dict): + return None + prompt_details = usage.get("prompt_tokens_details") or {} + completion_details = usage.get("completion_tokens_details") or {} + if not isinstance(prompt_details, dict): + prompt_details = {} + if not isinstance(completion_details, dict): + completion_details = {} + cost = None + cost_details = usage.get("cost_details") + if isinstance(cost_details, dict): + provider_cost = cost_details.get("total_cost") or cost_details.get("request_cost_usd") or cost_details.get("cost") or cost_details.get("estimated_cost") + cost = CostDetails( + provider_reported_cost=float(provider_cost) if provider_cost is not None else None, + currency=str(cost_details.get("currency") or "USD"), + source="usage.cost_details", + metadata={k: deepcopy(v) for k, v in cost_details.items() if k not in {"total_cost", "cost", "currency"}}, + ) + return Usage( + input_tokens=int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0), + output_tokens=int(usage.get("completion_tokens") or usage.get("output_tokens") or 0), + total_tokens=int(usage.get("total_tokens") or 0), + cache_read_tokens=int(prompt_details.get("cached_tokens") or usage.get("cache_read_tokens") or 0), + cache_write_tokens=int(prompt_details.get("cache_creation_tokens") or usage.get("cache_creation_tokens") or 0), + reasoning_tokens=int(completion_details.get("reasoning_tokens") or usage.get("reasoning_tokens") or 0), + cost=cost, + raw=deepcopy(usage), + ) + + def _parse_message(self, message: dict[str, Any]) -> UnifiedMessage: + payload = dict(message or {}) + reasoning = _extract_reasoning(payload) + role = str(payload.get("role") or "assistant") + content = self._parse_content(payload.get("content")) + if role == "tool": + result_content = canonical_tool_arguments(payload.get("content")) + content = [ + ContentBlock( + type="tool_result", + tool_result=ToolResult( + tool_call_id=payload.get("tool_call_id"), + name=payload.get("name"), + content=result_content, + raw=deepcopy(message), + ), + raw=deepcopy(message), + ) + ] + return UnifiedMessage( + role=role, + content=content, + name=payload.get("name"), + tool_call_id=payload.get("tool_call_id"), + tool_calls=self._parse_message_tool_calls(payload), + reasoning=reasoning, + raw=deepcopy(message), + extra={k: deepcopy(v) for k, v in payload.items() if k not in {"role", "content", "name", "tool_call_id", "tool_calls", "reasoning", "reasoning_content"}}, + ) + + def _format_request_messages(self, messages: Iterable[UnifiedMessage], *, preserve_source: bool) -> list[dict[str, Any]]: + """Format messages, expanding protocols that embed tool results in user turns.""" + + formatted: list[dict[str, Any]] = [] + for message in messages: + result_blocks = [block for block in message.content if block.tool_result] + if result_blocks and not (message.role == "tool" and len(result_blocks) == 1): + residual = [block for block in message.content if not block.tool_result] + if residual: + residual_message = deepcopy(message) + residual_message.content = residual + residual_message.tool_call_id = None + formatted.append(self._format_message(residual_message, preserve_source=False)) + for block in result_blocks: + result = block.tool_result + if result is None: + continue + formatted.append( + { + "role": "tool", + "tool_call_id": result.tool_call_id, + "content": _tool_result_text(result.content), + } + ) + continue + formatted.append(self._format_message(message, preserve_source=preserve_source)) + return formatted + + def _format_message(self, message: UnifiedMessage, *, preserve_source: bool = True) -> dict[str, Any]: + payload: dict[str, Any] = {"role": message.role} + if message.name: + payload["name"] = message.name + if message.tool_call_id: + payload["tool_call_id"] = message.tool_call_id + result_blocks = [block.tool_result for block in message.content if block.tool_result] + if message.role == "tool" and result_blocks: + result = result_blocks[0] + payload["tool_call_id"] = result.tool_call_id or message.tool_call_id + if result.name: + payload["name"] = result.name + content = _tool_result_text({"error": result.content} if result.is_error else result.content) + else: + content = self._format_content(message.content, preserve_source=preserve_source) + if content is not None: + payload["content"] = content + extra = deepcopy(message.extra) if preserve_source else {} + legacy_function_call = extra.get("function_call") + tool_calls = _message_tool_calls(message) + if tool_calls and not legacy_function_call: + payload["tool_calls"] = [self._format_tool_call(call, preserve_source=preserve_source) for call in tool_calls] + elif tool_calls and legacy_function_call: + call = tool_calls[0] + extra["function_call"] = {"name": call.name or "", "arguments": tool_arguments_text(call.arguments)} + if message.reasoning: + # OpenAI-compatible providers use multiple names for reasoning text. + # Prefer the common extension field while keeping all blocks in extra. + text = "".join(block.text or "" for block in message.reasoning if block.text) + if text: + payload["reasoning_content"] = text + payload.update(extra) + return payload + + def _parse_message_tool_calls(self, payload: dict[str, Any]) -> list[ToolCall]: + """Return modern and legacy OpenAI function calls as unified tools.""" + + modern_calls = payload.get("tool_calls") or [] + if modern_calls: + return [self._parse_tool_call(call) for call in modern_calls] + legacy_call = payload.get("function_call") + if isinstance(legacy_call, dict): + return [ + ToolCall( + id=None, + name=legacy_call.get("name"), + arguments=legacy_call.get("arguments"), + type="function", + raw=deepcopy(legacy_call), + extra={"legacy_function_call": True}, + ) + ] + return [] + + def _parse_content(self, content: Any) -> list[ContentBlock]: + if content is None: + return [] + if isinstance(content, str): + return text_blocks(content) + if not isinstance(content, list): + return [ContentBlock(type="unknown", raw=deepcopy(content))] + blocks = [] + for block in content: + if isinstance(block, str): + blocks.append(ContentBlock(type="text", text=block, raw=block)) + continue + if not isinstance(block, dict): + blocks.append(ContentBlock(type="unknown", raw=deepcopy(block))) + continue + block_type = block.get("type", "text") + if block_type == "text": + blocks.append(ContentBlock(type="text", text=block.get("text", ""), raw=deepcopy(block), extra=_without(block, {"type", "text"}))) + elif block_type in {"image_url", "input_image"}: + raw_source = deepcopy(block.get("image_url") or block.get("source")) + source = _openai_media_source(raw_source, kind="image") + blocks.append(ContentBlock(type="image", source=source, raw=deepcopy(block), extra=_without(block, {"type", "image_url", "source"}))) + elif block_type in {"input_audio", "audio"}: + raw_source = deepcopy(block.get("input_audio") or block.get("audio") or block.get("source")) + source = _openai_media_source(raw_source, kind="audio") + blocks.append(ContentBlock(type="audio", source=source, raw=deepcopy(block), extra=_without(block, {"type", "input_audio", "audio", "source"}))) + elif block_type in {"file", "input_file"}: + source = _openai_media_source(block, kind="file") + blocks.append(ContentBlock(type="file", source=source, raw=deepcopy(block), extra=_without(block, {"type", "file_id", "file_data", "filename"}))) + else: + blocks.append(ContentBlock(type=str(block_type), raw=deepcopy(block), extra=_without(block, {"type"}))) + return blocks + + def _format_content(self, blocks: Iterable[ContentBlock], *, preserve_source: bool = True) -> Any: + block_list = list(blocks) + if not block_list: + return None + if all(block.type == "text" and (not preserve_source or not isinstance(block.raw, dict)) and (not preserve_source or not block.extra) for block in block_list): + return first_text(block_list) or "" + formatted = [] + for block in block_list: + if block.type == "text": + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {"type": "text"} + payload["type"] = "text" + payload["text"] = block.text or "" + if preserve_source: + payload.update(deepcopy(block.extra)) + formatted.append(payload) + elif block.type == "image": + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {"type": "image_url"} + payload["type"] = "image_url" + payload["image_url"] = _format_openai_image_source(block.source) + if preserve_source: + payload.update(deepcopy(block.extra)) + formatted.append(payload) + elif block.type == "audio": + source = _media_source(block.source) + payload = {"type": "input_audio", "input_audio": {"data": source.data or "", "format": _audio_format(source.media_type)}} + formatted.append(payload) + elif block.type in {"file", "document"}: + source = _media_source(block.source) + file_payload: dict[str, Any] = {"type": "file"} + if source.file_id: + file_payload["file_id"] = source.file_id + elif source.data: + file_payload["file_data"] = source.data + elif source.url: + file_payload["file_url"] = source.url + formatted.append(file_payload) + elif preserve_source and isinstance(block.raw, dict): + formatted.append(deepcopy(block.raw)) + return formatted + + def _parse_tool_definition(self, tool: dict[str, Any]) -> ToolDefinition: + payload = dict(tool or {}) + function = payload.get("function") if isinstance(payload.get("function"), dict) else payload + return ToolDefinition( + name=str(function.get("name") or ""), + description=function.get("description"), + input_schema=deepcopy(function.get("parameters") or function.get("input_schema") or {}), + type=str(payload.get("type") or "function"), + extra={"raw": deepcopy(tool), **_without(payload, {"type", "function"})}, + ) + + def _format_tool_definition(self, tool: ToolDefinition, *, preserve_source: bool = True) -> dict[str, Any]: + raw = tool.extra.get("raw") + payload = deepcopy(raw) if preserve_source and isinstance(raw, dict) else { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description, + "parameters": deepcopy(tool.input_schema), + }, + } + payload["type"] = "function" + function = payload.get("function") if isinstance(payload.get("function"), dict) else {} + function.update({"name": tool.name, "description": tool.description, "parameters": deepcopy(tool.input_schema)}) + payload["function"] = {k: v for k, v in function.items() if v is not None} + return payload + + def _parse_tool_call(self, call: dict[str, Any]) -> ToolCall: + payload = dict(call or {}) + function = payload.get("function") if isinstance(payload.get("function"), dict) else {} + arguments: Any = canonical_tool_arguments(function.get("arguments")) + return ToolCall( + id=payload.get("id"), + name=function.get("name") or payload.get("name"), + arguments=arguments, + type=str(payload.get("type") or "function"), + index=payload.get("index"), + raw=deepcopy(call), + extra={**_without(function, {"name", "arguments"}), **_without(payload, {"id", "function", "type", "index", "name"})}, + ) + + def _format_tool_call(self, call: ToolCall, *, preserve_source: bool = True) -> dict[str, Any]: + payload = deepcopy(call.raw) if preserve_source and isinstance(call.raw, dict) else {} + payload["type"] = "function" + if call.id: + payload["id"] = call.id + if call.index is not None: + payload["index"] = call.index + function = deepcopy(payload.get("function")) if isinstance(payload.get("function"), dict) else {} + function["name"] = call.name or "" + function["arguments"] = tool_arguments_text(call.arguments) + payload["function"] = function + return payload + + def _format_generation_params(self, request: UnifiedRequest, *, preserve_source: bool) -> dict[str, Any]: + """Format canonical controls into OpenAI Chat field names.""" + + params = deepcopy(request.generation_params) + if preserve_source: + original = request.extensions.get(self.name, {}).get("generation_params") + payload = deepcopy(original) if isinstance(original, dict) else {} + else: + payload = {} + if "max_output_tokens" in params: + payload["max_completion_tokens"] = params.pop("max_output_tokens") + if "stop_sequences" in params: + payload["stop"] = params.pop("stop_sequences") + reasoning = params.pop("reasoning", None) + if isinstance(reasoning, dict) and reasoning.get("effort") is not None: + payload["reasoning_effort"] = reasoning["effort"] + if "structured_output" in params: + payload["response_format"] = format_structured_output(params.pop("structured_output"), self.name) + if "tool_choice" in params: + payload["tool_choice"] = format_tool_choice(params.pop("tool_choice"), self.name) + if "audio_output" in params: + payload["audio"] = deepcopy(params.pop("audio_output")) + supported = { + "frequency_penalty", + "logit_bias", + "logprobs", + "n", + "parallel_tool_calls", + "presence_penalty", + "seed", + "service_tier", + "stream_options", + "temperature", + "top_logprobs", + "top_p", + "user", + } + payload.update( + retain_supported_generation_params( + request, + params, + supported=supported, + target_protocol=self.name, + ) + ) + return payload + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return deepcopy(value) + if hasattr(value, "model_dump"): + return value.model_dump() + if hasattr(value, "dict"): + return value.dict() + return {} + + +def _decode_sse_data(raw_event: Any) -> Any: + from .streaming import decode_sse_data + + return decode_sse_data(raw_event) + + +def _extract_reasoning(payload: dict[str, Any]) -> list[ReasoningBlock]: + blocks = [] + for field_name in ("reasoning_content", "reasoning"): + value = payload.get(field_name) + if value: + blocks.append(ReasoningBlock(type=field_name, text=str(value), extra={"source_field": field_name})) + return blocks + + +def _format_openai_usage(usage: Usage | None) -> dict[str, Any] | None: + """Format normalized usage using OpenAI Chat's public field names.""" + + if usage is None: + return None + payload: dict[str, Any] = { + "prompt_tokens": usage.input_tokens, + "completion_tokens": usage.output_tokens, + "total_tokens": usage.total_tokens or (usage.input_tokens + usage.output_tokens), + } + prompt_details: dict[str, Any] = {} + if usage.cache_read_tokens: + prompt_details["cached_tokens"] = usage.cache_read_tokens + if usage.cache_write_tokens: + prompt_details["cache_creation_tokens"] = usage.cache_write_tokens + if prompt_details: + payload["prompt_tokens_details"] = prompt_details + completion_details: dict[str, Any] = {} + if usage.reasoning_tokens: + completion_details["reasoning_tokens"] = usage.reasoning_tokens + if completion_details: + payload["completion_tokens_details"] = completion_details + if usage.cost: + cost_details: dict[str, Any] = dict(usage.cost.metadata) + if usage.cost.provider_reported_cost is not None: + cost_details["total_cost"] = usage.cost.provider_reported_cost + elif usage.cost.estimated_cost is not None: + cost_details["estimated_cost"] = usage.cost.estimated_cost + cost_details["currency"] = usage.cost.currency + if usage.cost.source: + cost_details["source"] = usage.cost.source + payload["cost_details"] = cost_details + return payload + + +def _format_response_message(payload: dict[str, Any], message: UnifiedMessage) -> dict[str, Any]: + """Return Chat Completions response-message shape. + + Request messages may legitimately preserve content-part arrays. Assistant + response messages from non-chat native protocols often arrive as text parts; + Chat Completions clients expect the final message content to be a string in + that common case. + """ + + if payload.get("content") is not None and message.content: + if all(block.type in {"text", "input_text", "output_text"} and not block.extra for block in message.content): + payload = dict(payload) + payload["content"] = _first_response_text(message.content) or "" + return payload + + +def _first_response_text(blocks: Iterable[ContentBlock]) -> Optional[str]: + parts = [block.text for block in blocks if block.type in {"text", "input_text", "output_text"} and block.text] + return "".join(parts) if parts else first_text(blocks) + + +def _parse_openai_generation_params(source: dict[str, Any]) -> dict[str, Any]: + """Normalize OpenAI Chat controls into protocol-independent field names.""" + + params = deepcopy(source) + max_tokens = params.pop("max_completion_tokens", params.pop("max_tokens", None)) + if max_tokens is not None: + params["max_output_tokens"] = max_tokens + stop = params.pop("stop", None) + if stop is not None: + params["stop_sequences"] = [stop] if isinstance(stop, str) else deepcopy(stop) + reasoning_effort = params.pop("reasoning_effort", None) + if reasoning_effort is not None: + params["reasoning"] = {"effort": reasoning_effort} + return params + + +def _openai_media_source(value: Any, *, kind: str) -> MediaSource: + """Parse OpenAI media fields into a canonical media source.""" + + if isinstance(value, str): + if value.startswith("data:") and ";base64," in value: + prefix, data = value.split(",", 1) + return MediaSource(kind="base64", media_type=prefix[5:].split(";", 1)[0], data=data, raw=value) + return MediaSource(kind="url", url=value, raw=value) + payload = value if isinstance(value, dict) else {} + url = payload.get("url") or payload.get("file_url") + data = payload.get("data") or payload.get("file_data") + media_type = payload.get("media_type") or payload.get("mime_type") or payload.get("format") + if isinstance(url, str) and url.startswith("data:") and ";base64," in url: + prefix, encoded = url.split(",", 1) + media_type = media_type or prefix[5:].split(";", 1)[0] + data = encoded + url = None + file_id = payload.get("file_id") + source_kind = "file" if file_id else "base64" if data else "url" + return MediaSource( + kind=source_kind, + media_type=media_type, + url=url, + data=data, + file_id=file_id, + detail=payload.get("detail"), + raw=deepcopy(value), + extra=_without(payload, {"url", "file_url", "data", "file_data", "file_id", "media_type", "mime_type", "format", "detail"}), + ) + + +def _media_source(value: Any) -> MediaSource: + """Coerce legacy dictionary media sources into the canonical type.""" + + if isinstance(value, MediaSource): + return value + return _openai_media_source(value, kind="media") + + +def _format_openai_image_source(value: Any) -> dict[str, Any]: + """Format a canonical image source for Chat Completions.""" + + source = _media_source(value) + if source.url: + url = source.url + elif source.data: + media_type = source.media_type or "application/octet-stream" + url = f"data:{media_type};base64,{source.data}" + elif source.file_id: + url = source.file_id + else: + url = "" + payload: dict[str, Any] = {"url": url} + if source.detail: + payload["detail"] = source.detail + return payload + + +def _audio_format(media_type: Optional[str]) -> str: + """Return OpenAI's compact audio format label from a MIME type.""" + + if not media_type: + return "wav" + return media_type.rsplit("/", 1)[-1].lower() + + +def _message_tool_calls(message: UnifiedMessage) -> list[ToolCall]: + """Return de-duplicated calls from both canonical message representations.""" + + calls = list(message.tool_calls) + seen = {(call.id, call.name, tool_arguments_text(call.arguments)) for call in calls} + for block in message.content: + if not block.tool_call: + continue + key = (block.tool_call.id, block.tool_call.name, tool_arguments_text(block.tool_call.arguments)) + if key not in seen: + calls.append(block.tool_call) + seen.add(key) + return calls + + +def _tool_result_text(value: Any) -> str: + """Serialize a canonical tool result for Chat's tool message content.""" + + if isinstance(value, str): + return value + return json.dumps(serialize_value(value), separators=(",", ":")) + + +def _without(payload: dict[str, Any], keys: set[str]) -> dict[str, Any]: + return {k: deepcopy(v) for k, v in payload.items() if k not in keys} diff --git a/src/rotator_library/protocols/openai_embeddings.py b/src/rotator_library/protocols/openai_embeddings.py new file mode 100644 index 000000000..1ba955b99 --- /dev/null +++ b/src/rotator_library/protocols/openai_embeddings.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""OpenAI-compatible embeddings protocol adapter.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any, ClassVar + +from .base import ProtocolAdapter +from .operation import OPERATION_EMBEDDINGS +from .types import ProtocolContext, UnifiedRequest, UnifiedResponse, Usage + +_REQUEST_CORE_FIELDS = {"model", "input", "encoding_format", "dimensions", "user", "operation"} +_REQUEST_OPTION_FIELDS = {"encoding_format", "dimensions", "user"} + + +class OpenAIEmbeddingsProtocol(ProtocolAdapter): + """Adapter for `/v1/embeddings` style request and response payloads. + + The adapter intentionally treats embedding vectors as opaque data entries. + That keeps it usable for OpenAI-compatible providers with additional index, + metadata, or sparse-vector fields without narrowing the schema too early. + """ + + name: ClassVar[str] = "openai_embeddings" + aliases: ClassVar[tuple[str, ...]] = ("embeddings", "openai_embedding") + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_EMBEDDINGS,) + supported_transports: ClassVar[tuple[str, ...]] = ("http",) + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + return UnifiedRequest( + operation=OPERATION_EMBEDDINGS, + model=str(request.get("model") or getattr(context, "model", None) or ""), + input=deepcopy(request.get("input")), + generation_params={k: deepcopy(request[k]) for k in _REQUEST_OPTION_FIELDS if k in request}, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _REQUEST_CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + payload = {"model": unified_request.model, "input": deepcopy(unified_request.input)} + payload.update(deepcopy(unified_request.generation_params)) + payload.update(deepcopy(unified_request.extra)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + response = raw_response if isinstance(raw_response, dict) else {} + return UnifiedResponse( + operation=OPERATION_EMBEDDINGS, + model=response.get("model") or getattr(context, "model", None), + data=deepcopy(response.get("data") or []), + usage=self.extract_usage(response, context), + raw=deepcopy(raw_response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"model", "data", "usage"}}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + payload = {"object": "list", "data": deepcopy(unified_response.data)} + if unified_response.model: + payload["model"] = unified_response.model + if unified_response.usage: + payload["usage"] = unified_response.usage.raw or unified_response.usage.to_dict() + payload.update(deepcopy(unified_response.extra)) + return payload + + def extract_usage(self, raw_or_unified: Any, context: ProtocolContext | None = None) -> Usage | None: + if isinstance(raw_or_unified, UnifiedResponse): + return raw_or_unified.usage + usage = raw_or_unified.get("usage") if isinstance(raw_or_unified, dict) else None + if not isinstance(usage, dict): + return None + return Usage( + input_tokens=int(usage.get("prompt_tokens") or usage.get("input_tokens") or usage.get("total_tokens") or 0), + output_tokens=int(usage.get("output_tokens") or 0), + total_tokens=int(usage.get("total_tokens") or 0), + raw=deepcopy(usage), + ) diff --git a/src/rotator_library/protocols/openai_images.py b/src/rotator_library/protocols/openai_images.py new file mode 100644 index 000000000..fbf382dbb --- /dev/null +++ b/src/rotator_library/protocols/openai_images.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""OpenAI-compatible image generation/edit/variation protocol adapter.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any, ClassVar + +from .base import ProtocolAdapter +from .operation import OPERATION_IMAGE_EDIT, OPERATION_IMAGE_GENERATION, OPERATION_IMAGE_VARIATION, normalize_operation +from .types import ProtocolContext, UnifiedRequest, UnifiedResponse + +_OPTION_FIELDS = {"n", "size", "quality", "style", "response_format", "user", "background", "moderation"} +_CORE_FIELDS = {"operation", "model", "prompt", "image", "mask", *_OPTION_FIELDS} + + +class OpenAIImagesProtocol(ProtocolAdapter): + """Adapter for image generation, edit, and variation request shapes. + + File references are preserved as metadata in ``UnifiedRequest.files``. The + adapter never reads file contents; multipart assembly belongs to transport or + provider execution code so protocol parsing remains side-effect free. + """ + + name: ClassVar[str] = "openai_images" + aliases: ClassVar[tuple[str, ...]] = ("images", "image_generation", "openai_image") + supported_operations: ClassVar[tuple[str, ...]] = ( + OPERATION_IMAGE_GENERATION, + OPERATION_IMAGE_EDIT, + OPERATION_IMAGE_VARIATION, + ) + supported_transports: ClassVar[tuple[str, ...]] = ("http",) + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + operation = _image_operation(request) + files = [] + if "image" in request: + files.append({"field": "image", "value": deepcopy(request["image"])}) + if "mask" in request: + files.append({"field": "mask", "value": deepcopy(request["mask"])}) + return UnifiedRequest( + operation=operation, + model=str(request.get("model") or getattr(context, "model", None) or ""), + input=deepcopy(request.get("prompt")), + files=files, + generation_params={k: deepcopy(request[k]) for k in _OPTION_FIELDS if k in request}, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + payload: dict[str, Any] = {} + if unified_request.model: + payload["model"] = unified_request.model + if unified_request.input is not None: + payload["prompt"] = deepcopy(unified_request.input) + for file_entry in unified_request.files: + if isinstance(file_entry, dict) and file_entry.get("field"): + payload[str(file_entry["field"])] = deepcopy(file_entry.get("value")) + payload.update(deepcopy(unified_request.generation_params)) + payload.update(deepcopy(unified_request.extra)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + response = raw_response if isinstance(raw_response, dict) else {} + return UnifiedResponse( + operation=_context_operation(context, OPERATION_IMAGE_GENERATION), + model=response.get("model") or getattr(context, "model", None), + data=deepcopy(response.get("data") or []), + raw=deepcopy(raw_response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"model", "data"}}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + payload = {"data": deepcopy(unified_response.data)} + if unified_response.model: + payload["model"] = unified_response.model + payload.update(deepcopy(unified_response.extra)) + return payload + + +def _image_operation(request: dict[str, Any]) -> str: + explicit = normalize_operation(request.get("operation")) + if explicit in {OPERATION_IMAGE_GENERATION, OPERATION_IMAGE_EDIT, OPERATION_IMAGE_VARIATION}: + return explicit + if "image" in request and "prompt" in request: + return OPERATION_IMAGE_EDIT + if "image" in request: + return OPERATION_IMAGE_VARIATION + return OPERATION_IMAGE_GENERATION + + +def _context_operation(context: ProtocolContext | None, default: str) -> str: + if context and isinstance(context.provider_options, dict): + operation = normalize_operation(context.provider_options.get("operation")) + if operation in {OPERATION_IMAGE_GENERATION, OPERATION_IMAGE_EDIT, OPERATION_IMAGE_VARIATION}: + return operation + if context and isinstance(context.metadata, dict): + operation = normalize_operation(context.metadata.get("operation")) + if operation in {OPERATION_IMAGE_GENERATION, OPERATION_IMAGE_EDIT, OPERATION_IMAGE_VARIATION}: + return operation + return default diff --git a/src/rotator_library/protocols/operation.py b/src/rotator_library/protocols/operation.py new file mode 100644 index 000000000..da9541c58 --- /dev/null +++ b/src/rotator_library/protocols/operation.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Shared protocol operation names. + +Operations are deliberately plain strings instead of a closed enum. The native +protocol system is meant to be extended by local/custom providers, so the core +must provide well-known names without blocking new operations that are not known +today. +""" + +from __future__ import annotations + +from typing import Final + +OPERATION_UNKNOWN: Final[str] = "unknown" +# The four conversational wire operations map to this logical operation. Keeping +# it separate from endpoint-specific names allows a Chat request to target a +# Messages, Responses, or Gemini generate endpoint without failing capability +# checks merely because the external operation names differ. +OPERATION_GENERATE: Final[str] = "generate" +OPERATION_CHAT: Final[str] = "chat" +OPERATION_MESSAGES: Final[str] = "messages" +OPERATION_RESPONSES: Final[str] = "responses" +OPERATION_COUNT_TOKENS: Final[str] = "count_tokens" +OPERATION_EMBEDDINGS: Final[str] = "embeddings" +OPERATION_IMAGE_GENERATION: Final[str] = "image_generation" +OPERATION_IMAGE_EDIT: Final[str] = "image_edit" +OPERATION_IMAGE_VARIATION: Final[str] = "image_variation" +OPERATION_AUDIO_TRANSCRIPTION: Final[str] = "audio_transcription" +OPERATION_AUDIO_TRANSLATION: Final[str] = "audio_translation" +OPERATION_SPEECH: Final[str] = "speech" +OPERATION_OLLAMA_CHAT: Final[str] = "ollama_chat" +OPERATION_OLLAMA_GENERATE: Final[str] = "ollama_generate" +OPERATION_MCP: Final[str] = "mcp" + + +def normalize_operation(operation: str | None) -> str: + """Normalize an operation name while preserving custom extensions.""" + + if not operation: + return OPERATION_UNKNOWN + return str(operation).strip().lower() or OPERATION_UNKNOWN diff --git a/src/rotator_library/protocols/registry.py b/src/rotator_library/protocols/registry.py new file mode 100644 index 000000000..a27a192fa --- /dev/null +++ b/src/rotator_library/protocols/registry.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Auto-discovery registry for native protocol adapters.""" + +from __future__ import annotations + +import importlib +import inspect +import logging +import pkgutil +from typing import Type + +from .base import ProtocolAdapter + +lib_logger = logging.getLogger("rotator_library") + +PROTOCOL_PLUGINS: dict[str, Type[ProtocolAdapter]] = {} +PROTOCOL_ALIASES: dict[str, str] = {} +_PROTOCOL_INSTANCES: dict[str, ProtocolAdapter] = {} + +_INFRASTRUCTURE_MODULES = {"base", "operation", "registry", "types"} + + +def register_protocol(protocol_class: Type[ProtocolAdapter], *, replace: bool = False) -> Type[ProtocolAdapter]: + """Register a protocol adapter class and its aliases. + + The registry mirrors the provider plugin system while staying stricter about + duplicate names. This matters for custom protocol modules: accidental alias + collisions should fail early instead of silently changing conversion logic. + """ + + if not inspect.isclass(protocol_class) or not issubclass(protocol_class, ProtocolAdapter): + raise TypeError("protocol_class must inherit ProtocolAdapter") + if protocol_class is ProtocolAdapter: + raise TypeError("cannot register ProtocolAdapter itself") + + name = protocol_class.name + if not name: + raise ValueError(f"Protocol {protocol_class.__name__} must define a name") + alias_owner = PROTOCOL_ALIASES.get(name) + if alias_owner and alias_owner != name and not replace: + raise ValueError(f"Protocol name conflicts with registered alias: {name}") + + existing = PROTOCOL_PLUGINS.get(name) + if existing and existing is not protocol_class and not replace: + raise ValueError(f"Protocol name already registered: {name}") + + if replace and existing and existing is not protocol_class: + for alias, owner in list(PROTOCOL_ALIASES.items()): + if owner == name: + PROTOCOL_ALIASES.pop(alias, None) + PROTOCOL_PLUGINS[name] = protocol_class + _PROTOCOL_INSTANCES.pop(name, None) + + for alias in protocol_class.aliases: + existing_name = PROTOCOL_ALIASES.get(alias) + if existing_name and existing_name != name and not replace: + raise ValueError(f"Protocol alias already registered: {alias}") + if alias in PROTOCOL_PLUGINS and alias != name and not replace: + raise ValueError(f"Protocol alias conflicts with registered name: {alias}") + PROTOCOL_ALIASES[alias] = name + + lib_logger.debug("Registered protocol: %s", name) + return protocol_class + + +def resolve_protocol_name(name: str) -> str: + """Resolve aliases to canonical protocol names.""" + + if name in PROTOCOL_PLUGINS: + return name + if name in PROTOCOL_ALIASES: + return PROTOCOL_ALIASES[name] + raise KeyError(f"Unknown protocol: {name}") + + +def get_protocol_class(name: str) -> Type[ProtocolAdapter]: + """Return a registered protocol adapter class by name or alias.""" + + return PROTOCOL_PLUGINS[resolve_protocol_name(name)] + + +def get_protocol(name: str) -> ProtocolAdapter: + """Return a shared stateless protocol adapter instance by name or alias.""" + + canonical = resolve_protocol_name(name) + if canonical not in _PROTOCOL_INSTANCES: + _PROTOCOL_INSTANCES[canonical] = PROTOCOL_PLUGINS[canonical]() + return _PROTOCOL_INSTANCES[canonical] + + +def list_protocols() -> list[str]: + """Return canonical protocol names in deterministic order.""" + + return sorted(PROTOCOL_PLUGINS) + + +def _register_protocols() -> None: + """Discover protocol modules in this package and register adapter classes. + + Private modules and infrastructure modules are skipped so local experiments + can live next to production protocols without being imported accidentally. + """ + + package = importlib.import_module(__package__ or "rotator_library.protocols") + for _, module_name, _ in pkgutil.iter_modules(package.__path__): + if module_name.startswith("_") or module_name in _INFRASTRUCTURE_MODULES: + continue + module = importlib.import_module(f"{package.__name__}.{module_name}") + for attribute_name in dir(module): + attribute = getattr(module, attribute_name) + if ( + inspect.isclass(attribute) + and issubclass(attribute, ProtocolAdapter) + and attribute is not ProtocolAdapter + and attribute.__module__ == module.__name__ + ): + register_protocol(attribute) + + +_register_protocols() diff --git a/src/rotator_library/protocols/responses.py b/src/rotator_library/protocols/responses.py new file mode 100644 index 000000000..42f97dab6 --- /dev/null +++ b/src/rotator_library/protocols/responses.py @@ -0,0 +1,777 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""OpenAI Responses protocol adapter. + +Responses is important enough to model natively rather than forcing it through a +chat-completions shape. This adapter focuses on loss-conscious parsing and +formatting; storage, routes, and WebSocket transport are later phases. +""" + +from __future__ import annotations + +import json +from copy import deepcopy +from typing import Any, ClassVar, Iterable + +from .base import ProtocolAdapter +from .canonical import ( + add_conversion_warning, + canonical_stop_reason, + canonical_structured_output, + canonical_tool_arguments, + canonical_tool_choice, + coalesce_assistant_message, + conversation_messages, + format_stop_reason, + format_structured_output, + format_tool_choice, + instruction_blocks, + is_same_protocol, + message_reasoning, + message_tool_calls, + message_tool_results, + ordered_message_blocks, + retain_supported_generation_params, + resolve_tool_result_names, + source_extensions, + tool_arguments_text, + tool_result_text, +) +from .operation import OPERATION_GENERATE, OPERATION_RESPONSES +from .validation import validate_generative_request, validate_generative_response +from .types import ( + ContentBlock, + CostDetails, + MediaSource, + OutputItem, + ProtocolContext, + ReasoningBlock, + ToolCall, + ToolDefinition, + ToolResult, + UnifiedMessage, + UnifiedRequest, + UnifiedResponse, + UnifiedStreamEvent, + Usage, + first_text, + text_blocks, +) + +_GENERATION_PARAMS = { + "background", + "conversation", + "include", + "instructions", + "max_output_tokens", + "max_tool_calls", + "parallel_tool_calls", + "prompt", + "prompt_cache_key", + "reasoning", + "safety_identifier", + "service_tier", + "store", + "stream_options", + "temperature", + "text", + "tool_choice", + "top_p", + "top_logprobs", + "truncation", + "user", +} + +_REQUEST_CORE_FIELDS = { + "model", + "input", + "metadata", + "modalities", + "previous_response_id", + "stream", + "tools", + *_GENERATION_PARAMS, +} + + +class ResponsesProtocol(ProtocolAdapter): + """Adapter for OpenAI Responses request, response, and event stream shapes. + + The protocol keeps output items in addition to parsed messages because later + response storage and continuation features need item-level fidelity. + """ + + name: ClassVar[str] = "responses" + aliases: ClassVar[tuple[str, ...]] = ("openai_responses", "response_api") + supported_transports: ClassVar[tuple[str, ...]] = ("http", "sse") + supported_operations: ClassVar[tuple[str, ...]] = (OPERATION_RESPONSES,) + future_transports: ClassVar[tuple[str, ...]] = ("websocket",) + + def parse_request(self, raw_request: dict[str, Any], context: ProtocolContext | None = None) -> UnifiedRequest: + request = dict(raw_request or {}) + source_generation = {k: deepcopy(request[k]) for k in _GENERATION_PARAMS if k in request and k != "instructions"} + generation_params = _parse_responses_generation_params(source_generation) + if "tool_choice" in generation_params: + generation_params["tool_choice"] = canonical_tool_choice(generation_params["tool_choice"], self.name) + return UnifiedRequest( + operation=OPERATION_RESPONSES, + logical_operation=OPERATION_GENERATE, + model=str(request.get("model") or getattr(context, "model", None) or ""), + messages=resolve_tool_result_names(self._parse_input(request.get("input"))), + system=text_blocks(request.get("instructions")) if request.get("instructions") is not None else [], + tools=[self._parse_tool(tool) for tool in request.get("tools") or []], + stream=bool(request.get("stream", False)), + modalities=[str(value).lower() for value in request.get("modalities") or []], + generation_params=generation_params, + response_format=deepcopy(generation_params.get("structured_output")), + previous_response_id=request.get("previous_response_id"), + metadata=deepcopy(request.get("metadata") or {}), + source_protocol=self.name, + extensions={self.name: {"generation_params": source_generation}}, + raw=deepcopy(raw_request), + extra={k: deepcopy(v) for k, v in request.items() if k not in _REQUEST_CORE_FIELDS}, + ) + + def build_request(self, unified_request: UnifiedRequest, context: ProtocolContext | None = None) -> dict[str, Any]: + validate_generative_request(unified_request, self.name, context) + preserve_source = is_same_protocol(context, self.name, unified_request.source_protocol) + payload: dict[str, Any] = { + "model": unified_request.model, + "input": self._format_input(conversation_messages(unified_request), preserve_source=preserve_source), + } + instructions = "\n\n".join(block.text or "" for block in instruction_blocks(unified_request) if block.text) + if instructions: + payload["instructions"] = instructions + if unified_request.previous_response_id: + payload["previous_response_id"] = unified_request.previous_response_id + if unified_request.tools: + payload["tools"] = [self._format_tool(tool, preserve_source=preserve_source) for tool in unified_request.tools] + if unified_request.stream: + payload["stream"] = True + if unified_request.modalities: + payload["modalities"] = deepcopy(unified_request.modalities) + if unified_request.metadata: + payload["metadata"] = deepcopy(unified_request.metadata) + payload.update(self._format_generation_params(unified_request, preserve_source=preserve_source)) + payload.update(source_extensions(unified_request.extra, context, self.name, unified_request.source_protocol)) + return payload + + def parse_response(self, raw_response: Any, context: ProtocolContext | None = None) -> UnifiedResponse: + response = _as_dict(raw_response) + output = deepcopy(response.get("output") or []) + messages: list[UnifiedMessage] = [] + items: list[OutputItem] = [] + for index, item in enumerate(output): + if isinstance(item, dict): + parsed = self._parse_output_item(item) + if parsed: + parsed.extra["_output_index"] = index + messages.append(parsed) + items.append(_output_item_from_message(parsed, item)) + stop_reason = canonical_stop_reason(response.get("status")) + if stop_reason == "stop" and any(message_tool_calls(message) for message in messages): + stop_reason = "tool_use" + return UnifiedResponse( + operation=OPERATION_RESPONSES, + logical_operation=OPERATION_GENERATE, + id=response.get("id"), + model=response.get("model") or getattr(context, "model", None), + messages=messages, + items=items, + output=output, + stop_reason=stop_reason, + usage=self.extract_usage(response, context), + metadata={"object": response.get("object"), "created_at": response.get("created_at"), "native_status": response.get("status"), "incomplete_details": deepcopy(response.get("incomplete_details"))}, + source_protocol=self.name, + raw=deepcopy(response), + extra={k: deepcopy(v) for k, v in response.items() if k not in {"id", "object", "created_at", "model", "output", "usage", "status"}}, + ) + + def format_response(self, unified_response: UnifiedResponse, context: ProtocolContext | None = None) -> dict[str, Any]: + validate_generative_response(unified_response, self.name) + preserve_source = is_same_protocol(context, self.name, unified_response.source_protocol) + if preserve_source and unified_response.output: + output = deepcopy(unified_response.output) + for fallback_index, message in enumerate(unified_response.messages): + output_index = message.extra.get("_output_index", fallback_index) + if isinstance(output_index, int) and 0 <= output_index < len(output): + output[output_index] = self._format_output_message(message, output_index) + else: + output.append(self._format_output_message(message, fallback_index)) + else: + output = self._format_canonical_output(coalesce_assistant_message(unified_response.messages)) + payload = { + "id": unified_response.id, + "object": unified_response.metadata.get("object", "response"), + "created_at": unified_response.metadata.get("created_at"), + "model": unified_response.model, + "status": format_stop_reason(unified_response.stop_reason, self.name), + "output": output, + "usage": _format_responses_usage(unified_response.usage), + } + payload.update(source_extensions(unified_response.extra, context, self.name, unified_response.source_protocol)) + return {k: v for k, v in payload.items() if v is not None} + + def parse_stream_event(self, raw_event: Any, context: ProtocolContext | None = None) -> UnifiedStreamEvent: + event = _decode_sse_data(raw_event) + if event == "[DONE]": + return UnifiedStreamEvent(type="done", operation=OPERATION_RESPONSES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type="done", raw=deepcopy(raw_event)) + data = _as_dict(event) + event_type = str(data.get("type") or data.get("event") or "chunk") + if event_type in {"error", "response.error"} or data.get("error") is not None: + return UnifiedStreamEvent(type="error", operation=OPERATION_RESPONSES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, error=deepcopy(data.get("error", data)), raw=deepcopy(raw_event), extra={"payload": data}) + if event_type in {"response.completed", "response.failed", "response.incomplete"}: + response_payload = data.get("response") if isinstance(data.get("response"), dict) else {} + response = self.parse_response(response_payload, context) + error = None + if event_type == "response.failed": + error = deepcopy(response_payload.get("error") or data.get("error") or {"type": "upstream_error", "message": "Provider response failed"}) + return UnifiedStreamEvent(type=event_type, operation=OPERATION_RESPONSES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, message=response.messages[0] if response.messages else None, usage=response.usage, error=error, stop_reason=response.stop_reason, raw=deepcopy(raw_event), extra={"payload": data}) + if event_type == "response.output_text.delta": + message = UnifiedMessage(role="assistant", content=text_blocks(data.get("delta") or "")) + return UnifiedStreamEvent(type="message_delta", operation=OPERATION_RESPONSES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, delta=message, output_index=data.get("output_index"), content_index=data.get("content_index"), raw=deepcopy(raw_event), extra={"payload": data}) + if event_type == "response.function_call_arguments.delta": + call = ToolCall(id=data.get("call_id") or data.get("item_id"), arguments=data.get("delta") or "", index=data.get("output_index")) + message = UnifiedMessage(role="assistant", content=[ContentBlock(type="tool_call", tool_call=call)], tool_calls=[call]) + return UnifiedStreamEvent(type="tool_call_delta", operation=OPERATION_RESPONSES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, delta=message, tool_call=call, item_id=data.get("item_id"), output_index=data.get("output_index"), raw=deepcopy(raw_event), extra={"payload": data}) + if event_type in {"response.output_item.added", "response.output_item.done"} and isinstance(data.get("item"), dict): + message = self._parse_output_item(data["item"]) + if message and message.tool_calls: + for call in message.tool_calls: + call.index = data.get("output_index") + if call.arguments in ({}, ""): + call.arguments = None + return UnifiedStreamEvent(type=event_type, operation=OPERATION_RESPONSES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, message=message, raw=deepcopy(raw_event), extra={"payload": data}) + return UnifiedStreamEvent(type=event_type, operation=OPERATION_RESPONSES, logical_operation=OPERATION_GENERATE, source_protocol=self.name, native_type=event_type, raw=deepcopy(raw_event), extra={"payload": data}) + + def extract_usage(self, raw_or_unified: Any, context: ProtocolContext | None = None) -> Usage | None: + if isinstance(raw_or_unified, (UnifiedResponse, UnifiedStreamEvent)): + return raw_or_unified.usage + payload = _as_dict(raw_or_unified) + usage = payload.get("usage") if isinstance(payload.get("usage"), dict) else payload + if not isinstance(usage, dict) or not any(key.endswith("tokens") for key in usage): + return None + input_details = usage.get("input_tokens_details") if isinstance(usage.get("input_tokens_details"), dict) else {} + output_details = usage.get("output_tokens_details") if isinstance(usage.get("output_tokens_details"), dict) else {} + cost = None + cost_details = usage.get("cost_details") + if isinstance(cost_details, dict): + provider_cost = cost_details.get("total_cost") or cost_details.get("request_cost_usd") or cost_details.get("cost") or cost_details.get("estimated_cost") + cost = CostDetails( + provider_reported_cost=float(provider_cost) if provider_cost is not None else None, + currency=str(cost_details.get("currency") or "USD"), + source="usage.cost_details", + metadata={k: deepcopy(v) for k, v in cost_details.items() if k not in {"total_cost", "cost", "currency"}}, + ) + return Usage( + input_tokens=int(usage.get("input_tokens") or 0), + output_tokens=int(usage.get("output_tokens") or 0), + total_tokens=int(usage.get("total_tokens") or 0), + cache_read_tokens=int(input_details.get("cached_tokens") or 0), + cache_write_tokens=int(input_details.get("cache_creation_tokens") or usage.get("cache_creation_tokens") or 0), + reasoning_tokens=int(output_details.get("reasoning_tokens") or 0), + cost=cost, + raw=deepcopy(usage), + ) + + def _parse_input(self, input_value: Any) -> list[UnifiedMessage]: + if input_value is None: + return [] + if isinstance(input_value, str): + return [UnifiedMessage(role="user", content=text_blocks(input_value), raw=input_value)] + if not isinstance(input_value, list): + return [UnifiedMessage(role="user", content=[ContentBlock(type="unknown", raw=deepcopy(input_value))], raw=deepcopy(input_value))] + messages = [] + for item in input_value: + if isinstance(item, dict): + messages.append(self._parse_input_item(item)) + else: + messages.append(UnifiedMessage(role="user", content=text_blocks(str(item)), raw=deepcopy(item))) + return messages + + def _parse_input_item(self, item: dict[str, Any]) -> UnifiedMessage: + item_type = item.get("type") + if item_type in {"message", None}: + return UnifiedMessage( + role=str(item.get("role") or "user"), + content=self._parse_content(item.get("content")), + raw=deepcopy(item), + extra={k: deepcopy(v) for k, v in item.items() if k not in {"type", "role", "content"}}, + ) + if item_type == "function_call_output": + result_content = canonical_tool_arguments(item.get("output")) + return UnifiedMessage( + role="tool", + content=[ContentBlock(type="tool_result", tool_result=ToolResult(tool_call_id=item.get("call_id"), content=result_content), raw=deepcopy(item))], + tool_call_id=item.get("call_id"), + raw=deepcopy(item), + ) + if item_type in {"function_call", "custom_tool_call"}: + call = ToolCall( + id=item.get("call_id") or item.get("id"), + name=item.get("name"), + arguments=canonical_tool_arguments(item.get("arguments") or item.get("input")), + type="function" if item_type == "function_call" else str(item_type), + raw=deepcopy(item), + ) + return UnifiedMessage( + role="assistant", + content=[ContentBlock(type="tool_call", tool_call=call, raw=deepcopy(item))], + tool_calls=[call], + raw=deepcopy(item), + ) + if item_type == "reasoning": + reasoning = ReasoningBlock(type="reasoning", text=_reasoning_text(item), raw=deepcopy(item)) + return UnifiedMessage( + role="assistant", + content=[ContentBlock(type="reasoning", reasoning=reasoning, raw=deepcopy(item))], + reasoning=[reasoning], + raw=deepcopy(item), + ) + return UnifiedMessage(role=str(item.get("role") or "user"), content=[ContentBlock(type=str(item_type or "unknown"), raw=deepcopy(item))], raw=deepcopy(item)) + + def _format_input(self, messages: Iterable[UnifiedMessage], *, preserve_source: bool) -> list[dict[str, Any]]: + """Format canonical turns into ordered Responses input items.""" + + items: list[dict[str, Any]] = [] + for message in messages: + visible: list[ContentBlock] = [] + + def flush_visible() -> None: + if not visible: + return + residual_message = deepcopy(message) + residual_message.content = list(visible) + residual_message.tool_calls = [] + residual_message.reasoning = [] + items.append(self._format_input_message(residual_message, preserve_source=preserve_source)) + visible.clear() + + for block in ordered_message_blocks(message): + if block.reasoning: + flush_visible() + if block.reasoning.text: + items.append({"type": "reasoning", "summary": [{"type": "summary_text", "text": block.reasoning.text}]}) + elif block.tool_call: + flush_visible() + items.append(self._format_function_call(block.tool_call, preserve_source=preserve_source)) + elif block.tool_result: + flush_visible() + items.append(self._format_function_result(block.tool_result, preserve_source=preserve_source)) + else: + visible.append(block) + flush_visible() + return items + + def _format_input_message(self, message: UnifiedMessage, *, preserve_source: bool = True) -> dict[str, Any]: + if preserve_source and isinstance(message.raw, dict): + payload = deepcopy(message.raw) + if payload.get("type") == "function_call_output": + payload["call_id"] = message.tool_call_id or payload.get("call_id") + result = message.content[0].tool_result if message.content and message.content[0].tool_result else None + if result: + payload["output"] = deepcopy(result.content) + return payload + payload["role"] = message.role + payload["content"] = self._format_content(message.content, role=message.role, preserve_source=preserve_source) + return payload + role = "assistant" if message.role in {"assistant", "model"} else "user" + return {"type": "message", "role": role, "content": self._format_content(message.content, role=role, preserve_source=preserve_source)} + + def _parse_output_item(self, item: dict[str, Any]) -> UnifiedMessage | None: + item_type = item.get("type") + if item_type == "message": + return UnifiedMessage( + role=str(item.get("role") or "assistant"), + content=self._parse_content(item.get("content")), + raw=deepcopy(item), + extra={k: deepcopy(v) for k, v in item.items() if k not in {"type", "role", "content"}}, + ) + if item_type == "reasoning": + reasoning = ReasoningBlock(type="reasoning", text=_reasoning_text(item), extra={k: deepcopy(v) for k, v in item.items() if k not in {"type", "summary"}}) + reasoning.raw = deepcopy(item) + return UnifiedMessage(role="assistant", content=[ContentBlock(type="reasoning", reasoning=reasoning, raw=deepcopy(item))], reasoning=[reasoning], raw=deepcopy(item)) + if item_type in {"function_call", "custom_tool_call"}: + call = ToolCall(id=item.get("call_id") or item.get("id"), name=item.get("name"), arguments=canonical_tool_arguments(item.get("arguments") or item.get("input")), type="function" if item_type == "function_call" else str(item_type), raw=deepcopy(item)) + return UnifiedMessage(role="assistant", content=[ContentBlock(type="tool_call", tool_call=call, raw=deepcopy(item))], tool_calls=[call], raw=deepcopy(item)) + return None + + def _format_output_message(self, message: UnifiedMessage, index: int) -> dict[str, Any]: + if isinstance(message.raw, dict): + payload = deepcopy(message.raw) + item_type = payload.get("type") + if item_type == "message": + payload["role"] = message.role + payload["content"] = self._format_content(message.content, role="assistant", output=True, preserve_source=True) + return payload + if item_type == "reasoning" and message.reasoning: + payload["summary"] = [{"type": "summary_text", "text": message.reasoning[0].text or ""}] + return payload + if item_type in {"function_call", "custom_tool_call"} and message.tool_calls: + call = message.tool_calls[0] + payload["call_id"] = call.id + payload["name"] = call.name + payload["arguments"] = tool_arguments_text(call.arguments) + return payload + return {"id": f"msg_{index}", "type": "message", "role": message.role, "content": self._format_content(message.content, role=message.role, output=True, preserve_source=False)} + + def _format_canonical_output(self, message: UnifiedMessage) -> list[dict[str, Any]]: + """Build ordered Responses output items from one canonical assistant turn.""" + + output: list[dict[str, Any]] = [] + visible: list[ContentBlock] = [] + item_index = 0 + + def flush_visible() -> None: + nonlocal item_index + if not visible: + return + output.append( + { + "id": f"msg_{item_index}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": self._format_content(visible, role="assistant", output=True, preserve_source=False), + } + ) + visible.clear() + item_index += 1 + + for block in ordered_message_blocks(message): + if block.reasoning: + flush_visible() + if block.reasoning.text: + output.append( + { + "id": f"rs_{item_index}", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": block.reasoning.text}], + "status": "completed", + } + ) + item_index += 1 + elif block.tool_call: + flush_visible() + item = self._format_function_call(block.tool_call, preserve_source=False) + item.setdefault("id", f"fc_{item_index}") + item["status"] = "completed" + output.append(item) + item_index += 1 + elif block.tool_result: + flush_visible() + output.append(self._format_function_result(block.tool_result, preserve_source=False)) + item_index += 1 + else: + visible.append(block) + flush_visible() + return output + + def _parse_content(self, content: Any) -> list[ContentBlock]: + if content is None: + return [] + if isinstance(content, str): + return text_blocks(content) + if not isinstance(content, list): + return [ContentBlock(type="unknown", raw=deepcopy(content))] + blocks = [] + for block in content: + if isinstance(block, str): + blocks.append(ContentBlock(type="input_text", text=block, raw=block)) + continue + if not isinstance(block, dict): + blocks.append(ContentBlock(type="unknown", raw=deepcopy(block))) + continue + block_type = str(block.get("type") or "text") + if block_type in {"input_text", "output_text", "text"}: + blocks.append(ContentBlock(type="text", text=block.get("text", ""), raw=deepcopy(block), extra={"source_type": block_type, **_without(block, {"type", "text"})})) + elif block_type in {"input_image", "image_url"}: + source = _parse_responses_media_source(block) + blocks.append(ContentBlock(type="image", source=source, raw=deepcopy(block), extra={"source_type": block_type, **_without(block, {"type", "image_url", "source"})})) + elif block_type in {"input_file", "file"}: + source = _parse_responses_media_source(block) + blocks.append(ContentBlock(type="file", source=source, raw=deepcopy(block), extra={"source_type": block_type, **_without(block, {"type", "file_id", "file_data", "file_url"})})) + else: + blocks.append(ContentBlock(type=block_type, raw=deepcopy(block), extra=_without(block, {"type"}))) + return blocks + + def _format_content(self, blocks: Iterable[ContentBlock], *, role: str = "user", output: bool = False, preserve_source: bool = True) -> list[dict[str, Any]]: + formatted = [] + for block in blocks: + if block.type == "text": + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {} + payload["type"] = "output_text" if output or role in {"assistant", "model"} else "input_text" + payload["text"] = block.text or "" + if preserve_source: + payload.update({k: deepcopy(v) for k, v in block.extra.items() if k != "source_type"}) + formatted.append(payload) + elif block.type == "image": + payload = deepcopy(block.raw) if preserve_source and isinstance(block.raw, dict) else {"type": "input_image"} + payload["type"] = "input_image" + payload.update(_format_responses_image_source(block.source)) + if preserve_source: + payload.update({k: deepcopy(v) for k, v in block.extra.items() if k != "source_type"}) + formatted.append(payload) + elif block.type in {"file", "document"}: + payload = {"type": "input_file"} + payload.update(_format_responses_file_source(block.source)) + formatted.append(payload) + elif preserve_source and isinstance(block.raw, dict): + formatted.append(deepcopy(block.raw)) + return formatted + + def _parse_tool(self, tool: dict[str, Any]) -> ToolDefinition: + payload = dict(tool or {}) + parameters = payload.get("parameters") or payload.get("input_schema") or {} + return ToolDefinition( + name=str(payload.get("name") or ""), + description=payload.get("description"), + input_schema=deepcopy(parameters), + type=str(payload.get("type") or "function"), + extra={k: deepcopy(v) for k, v in payload.items() if k not in {"type", "name", "description", "parameters", "input_schema"}}, + ) + + def _format_tool(self, tool: ToolDefinition, *, preserve_source: bool = True) -> dict[str, Any]: + payload = {"type": "function" if tool.type == "function" else tool.type, "name": tool.name, "parameters": deepcopy(tool.input_schema)} + if tool.description is not None: + payload["description"] = tool.description + if preserve_source: + payload.update(deepcopy(tool.extra)) + return payload + + def _format_function_call(self, call: ToolCall, *, preserve_source: bool) -> dict[str, Any]: + payload = deepcopy(call.raw) if preserve_source and isinstance(call.raw, dict) else {} + payload.update( + { + "type": "function_call", + "call_id": call.id or "", + "name": call.name or "", + "arguments": tool_arguments_text(call.arguments), + } + ) + return payload + + def _format_function_result(self, result: ToolResult, *, preserve_source: bool) -> dict[str, Any]: + payload = deepcopy(result.raw) if preserve_source and isinstance(result.raw, dict) else {} + result_content = {"error": result.content} if result.is_error else result.content + payload.update({"type": "function_call_output", "call_id": result.tool_call_id or "", "output": tool_result_text(result_content)}) + return payload + + def _format_generation_params(self, request: UnifiedRequest, *, preserve_source: bool) -> dict[str, Any]: + params = deepcopy(request.generation_params) + original = request.extensions.get(self.name, {}).get("generation_params") if preserve_source else None + payload = deepcopy(original) if isinstance(original, dict) else {} + if "max_output_tokens" in params: + payload["max_output_tokens"] = params.pop("max_output_tokens") + if "stop_sequences" in params: + # Responses currently has no universal stop field. Keep it only when + # an explicitly compatible provider extension supplied one. + params.pop("stop_sequences") + add_conversion_warning( + request, + code="unsupported_optional_control", + message="responses has no portable stop-sequence request field", + field="stop_sequences", + target_protocol=self.name, + ) + reasoning = params.pop("reasoning", None) + if isinstance(reasoning, dict): + payload["reasoning"] = deepcopy(reasoning) + structured = params.pop("structured_output", None) + if isinstance(structured, dict): + payload["text"] = {"format": format_structured_output(structured, self.name)} + if "tool_choice" in params: + payload["tool_choice"] = format_tool_choice(params.pop("tool_choice"), self.name) + supported = { + "background", + "conversation", + "include", + "max_tool_calls", + "parallel_tool_calls", + "prompt", + "prompt_cache_key", + "safety_identifier", + "service_tier", + "store", + "stream_options", + "temperature", + "top_logprobs", + "top_p", + "truncation", + "user", + } + payload.update( + retain_supported_generation_params( + request, + params, + supported=supported, + target_protocol=self.name, + ) + ) + return payload + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return deepcopy(value) + if hasattr(value, "model_dump"): + return value.model_dump() + if hasattr(value, "dict"): + return value.dict() + return {} + + +def _decode_sse_data(raw_event: Any) -> Any: + from .streaming import decode_sse_data + + return decode_sse_data(raw_event) + + +def _reasoning_text(item: dict[str, Any]) -> str | None: + summary = item.get("summary") + if isinstance(summary, list): + parts = [] + for part in summary: + if isinstance(part, dict) and part.get("text"): + parts.append(str(part["text"])) + elif isinstance(part, str): + parts.append(part) + return "".join(parts) if parts else None + return str(summary) if summary else None + + +def _without(payload: dict[str, Any], keys: set[str]) -> dict[str, Any]: + return {k: deepcopy(v) for k, v in payload.items() if k not in keys} + + +def _format_responses_usage(usage: Usage | None) -> dict[str, Any] | None: + """Format normalized usage using OpenAI Responses public field names.""" + + if usage is None: + return None + payload: dict[str, Any] = { + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, + "total_tokens": usage.total_tokens or (usage.input_tokens + usage.output_tokens), + } + input_details: dict[str, Any] = {} + if usage.cache_read_tokens: + input_details["cached_tokens"] = usage.cache_read_tokens + if usage.cache_write_tokens: + # OpenAI Responses does not have a universal cache-write field, but this + # extension keeps provider-reported cache creation visible without + # leaking the unified internal `cache_write_tokens` key. + input_details["cache_creation_tokens"] = usage.cache_write_tokens + if input_details: + payload["input_tokens_details"] = input_details + output_details: dict[str, Any] = {} + if usage.reasoning_tokens: + output_details["reasoning_tokens"] = usage.reasoning_tokens + if output_details: + payload["output_tokens_details"] = output_details + if usage.cost: + cost_details: dict[str, Any] = dict(usage.cost.metadata) + if usage.cost.provider_reported_cost is not None: + cost_details["total_cost"] = usage.cost.provider_reported_cost + elif usage.cost.estimated_cost is not None: + cost_details["estimated_cost"] = usage.cost.estimated_cost + cost_details["currency"] = usage.cost.currency + if usage.cost.source: + cost_details["source"] = usage.cost.source + payload["cost_details"] = cost_details + return payload + + +def _parse_responses_generation_params(source: dict[str, Any]) -> dict[str, Any]: + """Normalize Responses controls into canonical names.""" + + params = deepcopy(source) + text = params.pop("text", None) + if isinstance(text, dict) and isinstance(text.get("format"), dict): + params["structured_output"] = canonical_structured_output(text["format"], "responses") + reasoning = params.get("reasoning") + if isinstance(reasoning, dict): + params["reasoning"] = deepcopy(reasoning) + return params + + +def _parse_responses_media_source(block: dict[str, Any]) -> MediaSource: + """Normalize Responses image and file content fields.""" + + value = block.get("image_url") or block.get("file_url") or block.get("source") + if isinstance(value, str): + if value.startswith("data:") and ";base64," in value: + prefix, data = value.split(",", 1) + return MediaSource(kind="base64", media_type=prefix[5:].split(";", 1)[0], data=data, raw=deepcopy(block)) + return MediaSource(kind="url", url=value, detail=block.get("detail"), raw=deepcopy(block)) + return MediaSource( + kind="file" if block.get("file_id") else "base64" if block.get("file_data") else "url", + media_type=block.get("mime_type") or block.get("media_type"), + url=block.get("file_url"), + data=block.get("file_data"), + file_id=block.get("file_id"), + detail=block.get("detail"), + raw=deepcopy(block), + ) + + +def _coerce_media_source(value: Any) -> MediaSource: + """Coerce legacy media dictionaries into canonical form.""" + + if isinstance(value, MediaSource): + return value + if isinstance(value, str): + return MediaSource(kind="url", url=value, raw=value) + payload = value if isinstance(value, dict) else {} + return MediaSource( + kind="file" if payload.get("file_id") else "base64" if payload.get("data") or payload.get("file_data") else "url", + media_type=payload.get("mime_type") or payload.get("media_type"), + url=payload.get("url") or payload.get("file_url"), + data=payload.get("data") or payload.get("file_data"), + file_id=payload.get("file_id"), + detail=payload.get("detail"), + raw=deepcopy(value), + ) + + +def _format_responses_image_source(value: Any) -> dict[str, Any]: + """Format a canonical image source for Responses input content.""" + + source = _coerce_media_source(value) + if source.url: + image_url = source.url + elif source.data: + image_url = f"data:{source.media_type or 'application/octet-stream'};base64,{source.data}" + else: + image_url = source.file_id or "" + payload: dict[str, Any] = {"image_url": image_url} + if source.detail: + payload["detail"] = source.detail + return payload + + +def _format_responses_file_source(value: Any) -> dict[str, Any]: + """Format a canonical file source for Responses input content.""" + + source = _coerce_media_source(value) + if source.file_id: + return {"file_id": source.file_id} + if source.data: + return {"file_data": source.data} + if source.url: + return {"file_url": source.url} + return {"file_data": ""} + + +def _output_item_from_message(message: UnifiedMessage, raw: dict[str, Any]) -> OutputItem: + """Create an ordered canonical output item alongside compatibility messages.""" + + item_type = str(raw.get("type") or "message") + if item_type == "reasoning": + reasoning = message_reasoning(message) + return OutputItem(type="reasoning", id=raw.get("id"), reasoning=reasoning[0] if reasoning else None, status=raw.get("status"), raw=deepcopy(raw)) + if item_type in {"function_call", "custom_tool_call"}: + calls = message_tool_calls(message) + return OutputItem(type="tool_call", id=raw.get("id"), tool_call=calls[0] if calls else None, status=raw.get("status"), raw=deepcopy(raw)) + return OutputItem(type="message", id=raw.get("id"), role=message.role, content=deepcopy(message.content), status=raw.get("status"), raw=deepcopy(raw)) diff --git a/src/rotator_library/protocols/streaming.py b/src/rotator_library/protocols/streaming.py new file mode 100644 index 000000000..cda43f226 --- /dev/null +++ b/src/rotator_library/protocols/streaming.py @@ -0,0 +1,684 @@ +"""Stateful canonical conversion for generative streaming protocols. + +Protocol adapters parse provider frames into :class:`UnifiedStreamEvent`. This +module owns the inverse operation because destination protocols have different +lifecycle requirements: one canonical delta can expand into several SSE frames, +and terminal events must close every destination-owned content block exactly +once. Operational concerns such as timeouts, retries, cancellation, and +heartbeats remain in ``client.streaming``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +import uuid +from typing import Any, AsyncIterator + +from .canonical import ( + STOP_REASON_UNKNOWN, + format_stop_reason, + ordered_message_blocks, + tool_arguments_text, +) +from .types import ( + ContentBlock, + ProtocolContext, + ProtocolError, + ToolCall, + UnifiedMessage, + UnifiedStreamEvent, + Usage, + serialize_value, +) + + +@dataclass +class StreamFormatState: + """Destination-owned lifecycle state for one converted stream.""" + + protocol: str + response_id: str + model: str + started: bool = False + terminal: bool = False + completion_emitted: bool = False + role_emitted: bool = False + stop_reason: str | None = None + usage: Usage | None = None + next_index: int = 0 + open_blocks: dict[str, int] = field(default_factory=dict) + block_order: list[str] = field(default_factory=list) + item_ids: dict[str, str] = field(default_factory=dict) + item_kinds: dict[str, str] = field(default_factory=dict) + text_by_key: dict[str, str] = field(default_factory=dict) + tool_arguments: dict[str, str] = field(default_factory=dict) + tool_names: dict[str, str] = field(default_factory=dict) + tool_ids: dict[str, str] = field(default_factory=dict) + emitted_tools: set[str] = field(default_factory=set) + + +class ProtocolStreamConverter: + """Convert raw source frames to one independently selected output protocol.""" + + def __init__( + self, + source_protocol: Any, + output_protocol: Any, + context: ProtocolContext, + ) -> None: + self.source_protocol = source_protocol + self.output_protocol = output_protocol + self.context = context + self.state = stream_format_state(context, output_protocol.name) + + def convert(self, raw_event: Any) -> list[Any]: + """Parse and format one source frame, expanding destination lifecycle frames.""" + + event = self.source_protocol.parse_stream_event(raw_event, self.context) + return format_canonical_stream_event( + event, + self.output_protocol.name, + self.context, + state=self.state, + ) + + +async def convert_protocol_stream( + stream: AsyncIterator[Any], + *, + source_protocol: Any, + output_protocol: Any, + context: ProtocolContext, +) -> AsyncIterator[Any]: + """Convert a resilient source stream while preserving transport heartbeats.""" + + converter = ProtocolStreamConverter(source_protocol, output_protocol, context) + async for raw_event in stream: + if isinstance(raw_event, str) and raw_event.lstrip().startswith(":"): + yield raw_event + continue + for frame in converter.convert(raw_event): + yield frame + + +def stream_format_state( + context: ProtocolContext | None, + protocol: str, +) -> StreamFormatState: + """Return persistent formatter state stored on the stream protocol context.""" + + metadata = context.metadata if context is not None else {} + states = metadata.setdefault("_stream_format_states", {}) + state = states.get(protocol) + if isinstance(state, StreamFormatState): + return state + model = str((context.model if context else None) or metadata.get("model") or "") + request_id = str((context.request_id if context else None) or metadata.get("request_id") or uuid.uuid4().hex) + prefix = {"openai_chat": "chatcmpl", "anthropic_messages": "msg", "responses": "resp"}.get(protocol, "stream") + state = StreamFormatState(protocol=protocol, response_id=f"{prefix}_{request_id}", model=model) + states[protocol] = state + return state + + +def format_canonical_stream_event( + event: UnifiedStreamEvent, + target_protocol: str, + context: ProtocolContext | None = None, + *, + state: StreamFormatState | None = None, +) -> list[Any]: + """Format one canonical event into zero or more destination wire frames.""" + + state = state or stream_format_state(context, target_protocol) + if state.terminal: + return [] + if event.usage is not None: + state.usage = event.usage + if event.stop_reason: + state.stop_reason = event.stop_reason + elif event.extra.get("stop_reason"): + state.stop_reason = str(event.extra["stop_reason"]) + + if target_protocol == "openai_chat": + return _format_openai(event, state) + if target_protocol == "anthropic_messages": + return _format_anthropic(event, state) + if target_protocol == "responses": + return _format_responses(event, state) + if target_protocol == "gemini": + return _format_gemini(event, state) + raise ProtocolError( + f"Canonical streaming is not supported for {target_protocol}", + protocol=target_protocol, + pass_name="format_stream_event", + ) + + +def _format_openai(event: UnifiedStreamEvent, state: StreamFormatState) -> list[str]: + if event.type == "error" or event.error is not None: + state.terminal = True + return [_data_frame({"error": _error_payload(event.error)}), "data: [DONE]\n\n"] + + frames: list[str] = [] + delta = _openai_delta(event.delta or event.message) + if delta: + if not state.role_emitted: + delta.setdefault("role", "assistant") + state.role_emitted = True + frames.append(_data_frame(_openai_chunk(state, delta=delta, finish_reason=None, usage=event.usage))) + + terminal = _is_terminal(event) + if event.stop_reason or event.extra.get("stop_reason"): + reason = event.stop_reason or event.extra.get("stop_reason") + state.stop_reason = str(reason) + frames.append(_data_frame(_openai_chunk(state, delta={}, finish_reason=format_stop_reason(state.stop_reason, "openai_chat"), usage=event.usage))) + if terminal: + frames.append("data: [DONE]\n\n") + state.terminal = True + elif event.usage is not None and not delta: + frames.append(_data_frame(_openai_chunk(state, delta={}, finish_reason=None, usage=event.usage))) + return frames + + +def _format_anthropic(event: UnifiedStreamEvent, state: StreamFormatState) -> list[str]: + if event.type == "error" or event.error is not None: + state.terminal = True + return [_event_frame("error", {"type": "error", "error": _error_payload(event.error)})] + + frames = _anthropic_start(state) + for block in _event_blocks(event): + key, block_type = _block_key(block, state) + if key not in state.open_blocks: + index = state.next_index + state.next_index += 1 + state.open_blocks[key] = index + state.block_order.append(key) + frames.append(_event_frame("content_block_start", { + "type": "content_block_start", + "index": index, + "content_block": _anthropic_block_start(block, key, state), + })) + index = state.open_blocks[key] + delta = _anthropic_block_delta(block, key, state) + if delta is not None: + frames.append(_event_frame("content_block_delta", { + "type": "content_block_delta", + "index": index, + "delta": delta, + })) + + if _is_terminal(event): + for key in state.block_order: + frames.append(_event_frame("content_block_stop", { + "type": "content_block_stop", + "index": state.open_blocks[key], + })) + reason = format_stop_reason(state.stop_reason, "anthropic_messages") + frames.append(_event_frame("message_delta", { + "type": "message_delta", + "delta": {"stop_reason": reason, "stop_sequence": None}, + "usage": _anthropic_usage(state.usage, output_only=True), + })) + frames.append(_event_frame("message_stop", {"type": "message_stop"})) + state.terminal = True + return frames + + +def _format_responses(event: UnifiedStreamEvent, state: StreamFormatState) -> list[str]: + if event.type == "error" or event.error is not None: + frames = _responses_start(state) + error = _error_payload(event.error) + frames.append(_event_frame("response.failed", { + "type": "response.failed", + "response": _responses_object(state, status="failed", error=error), + })) + frames.append("data: [DONE]\n\n") + state.terminal = True + return frames + + frames = _responses_start(state) + for block in _event_blocks(event): + key, kind = _block_key(block, state) + if key not in state.item_ids: + item_id = _responses_item_id(kind, state.next_index) + state.next_index += 1 + state.item_ids[key] = item_id + state.item_kinds[key] = kind + frames.extend(_responses_item_start(block, key, item_id, state)) + frames.extend(_responses_item_delta(block, key, state.item_ids[key], state)) + + if _is_terminal(event): + status = format_stop_reason(state.stop_reason or STOP_REASON_UNKNOWN, "responses") or "incomplete" + item_status = "completed" if status == "completed" else "incomplete" + for key, item_id in state.item_ids.items(): + frames.extend(_responses_item_done(key, item_id, state, item_status=item_status)) + event_name = "response.failed" if status == "failed" else "response.completed" if status == "completed" else "response.incomplete" + frames.append(_event_frame(event_name, { + "type": event_name, + "response": _responses_object(state, status=status), + })) + frames.append("data: [DONE]\n\n") + state.terminal = True + return frames + + +def _format_gemini(event: UnifiedStreamEvent, state: StreamFormatState) -> list[str]: + if event.type == "error" or event.error is not None: + state.terminal = True + return [_data_frame({"error": _error_payload(event.error)})] + + parts: list[dict[str, Any]] = [] + for block in _event_blocks(event): + if block.tool_call: + key, _ = _block_key(block, state) + call = block.tool_call + state.tool_names[key] = call.name or state.tool_names.get(key, "") + state.tool_ids[key] = call.id or state.tool_ids.get(key, "") + fragment = tool_arguments_text(call.arguments) + if key in state.emitted_tools: + if fragment: + raise ProtocolError( + "Gemini tool-call arguments continued after a complete object was emitted", + protocol="gemini", + pass_name="format_stream_event", + payload={"tool_call": key}, + ) + continue + state.tool_arguments[key] = state.tool_arguments.get(key, "") + fragment + if not state.tool_arguments[key]: + continue + try: + arguments = json.loads(state.tool_arguments[key]) + except json.JSONDecodeError: + continue + parts.append(_gemini_function_call_part(key, arguments, state)) + state.emitted_tools.add(key) + elif block.reasoning: + parts.append({"text": block.reasoning.text or "", "thought": True}) + elif block.type == "text": + parts.append({"text": block.text or ""}) + + if _is_terminal(event): + for key in state.tool_names: + if key in state.emitted_tools: + continue + arguments_text = state.tool_arguments.get(key, "") + if arguments_text and not _is_json(arguments_text): + raise ProtocolError( + "Gemini cannot emit an incomplete streamed tool-call argument object", + protocol="gemini", + pass_name="format_stream_event", + payload={"tool_call": key}, + ) + arguments = json.loads(arguments_text) if arguments_text else {} + parts.append(_gemini_function_call_part(key, arguments, state)) + state.emitted_tools.add(key) + + frames: list[str] = [] + finish_reason = None + if event.stop_reason or event.extra.get("stop_reason"): + state.stop_reason = str(event.stop_reason or event.extra.get("stop_reason")) + finish_reason = format_stop_reason(state.stop_reason, "gemini") + if parts or finish_reason or event.usage is not None: + candidate: dict[str, Any] = {"index": 0} + if parts: + candidate["content"] = {"role": "model", "parts": parts} + if finish_reason: + candidate["finishReason"] = finish_reason + payload: dict[str, Any] = {"candidates": [candidate]} + if state.model: + payload["modelVersion"] = state.model + usage = _gemini_usage(event.usage) + if usage: + payload["usageMetadata"] = usage + frames.append(_data_frame(payload)) + if finish_reason: + state.completion_emitted = True + if _is_terminal(event): + state.terminal = True + return frames + + +def _gemini_function_call_part( + key: str, + arguments: Any, + state: StreamFormatState, +) -> dict[str, Any]: + """Build one single-shot Gemini function call from buffered fragments.""" + + function_call: dict[str, Any] = { + "name": state.tool_names.get(key, ""), + "args": arguments, + } + if state.tool_ids.get(key): + function_call["id"] = state.tool_ids[key] + return {"functionCall": function_call} + + +def _event_blocks(event: UnifiedStreamEvent) -> list[ContentBlock]: + if _is_terminal(event) or event.type.endswith(".done"): + # Provider terminal snapshots repeat content already delivered as deltas. + # They remain available on the event for accounting/storage but are not + # emitted as a second client-visible delta. + return [] + message = event.delta or event.message + if message is None: + return [] + return ordered_message_blocks(message) + + +def _block_key(block: ContentBlock, state: StreamFormatState) -> tuple[str, str]: + if block.tool_call: + call = block.tool_call + identity = call.index if call.index is not None else call.id or call.name or "default" + return f"tool:{identity}", "tool" + if block.reasoning: + return "reasoning:0", "reasoning" + return "text:0", "text" + + +def _openai_delta(message: UnifiedMessage | None) -> dict[str, Any]: + if message is None: + return {} + delta: dict[str, Any] = {} + text = "".join(block.text or "" for block in ordered_message_blocks(message) if block.type == "text" and not block.reasoning) + reasoning = "".join(block.reasoning.text or "" for block in ordered_message_blocks(message) if block.reasoning) + if text: + delta["content"] = text + if reasoning: + delta["reasoning_content"] = reasoning + calls = [block.tool_call for block in ordered_message_blocks(message) if block.tool_call] + if calls: + delta["tool_calls"] = [ + { + "index": call.index if call.index is not None else index, + "id": call.id, + "type": call.type or "function", + "function": { + "name": call.name, + "arguments": tool_arguments_text(call.arguments), + }, + } + for index, call in enumerate(calls) + ] + delta["tool_calls"] = [ + {key: value for key, value in call.items() if value is not None} + for call in delta["tool_calls"] + ] + return delta + + +def _openai_chunk( + state: StreamFormatState, + *, + delta: dict[str, Any], + finish_reason: str | None, + usage: Usage | None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "id": state.response_id, + "object": "chat.completion.chunk", + "model": state.model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + formatted_usage = _openai_usage(usage) + if formatted_usage: + payload["usage"] = formatted_usage + return payload + + +def _anthropic_start(state: StreamFormatState) -> list[str]: + if state.started: + return [] + state.started = True + return [_event_frame("message_start", { + "type": "message_start", + "message": { + "id": state.response_id, + "type": "message", + "role": "assistant", + "content": [], + "model": state.model, + "stop_reason": None, + "stop_sequence": None, + "usage": _anthropic_usage(state.usage), + }, + })] + + +def _anthropic_block_start(block: ContentBlock, key: str, state: StreamFormatState) -> dict[str, Any]: + if block.tool_call: + call = block.tool_call + state.tool_names[key] = call.name or state.tool_names.get(key, "") + state.tool_ids[key] = call.id or state.tool_ids.get(key, f"call_{state.open_blocks[key]}") + return {"type": "tool_use", "id": state.tool_ids[key], "name": state.tool_names[key], "input": {}} + if block.reasoning: + return {"type": "thinking", "thinking": ""} + return {"type": "text", "text": ""} + + +def _anthropic_block_delta(block: ContentBlock, key: str, state: StreamFormatState) -> dict[str, Any] | None: + if block.tool_call: + call = block.tool_call + if call.name: + state.tool_names[key] = call.name + fragment = tool_arguments_text(call.arguments) + state.tool_arguments[key] = state.tool_arguments.get(key, "") + fragment + return {"type": "input_json_delta", "partial_json": fragment} if fragment else None + if block.reasoning: + text = block.reasoning.text or "" + return {"type": "thinking_delta", "thinking": text} if text else None + text = block.text or "" + return {"type": "text_delta", "text": text} if text else None + + +def _responses_start(state: StreamFormatState) -> list[str]: + if state.started: + return [] + state.started = True + return [_event_frame("response.created", { + "type": "response.created", + "response": _responses_object(state, status="in_progress"), + })] + + +def _responses_item_id(kind: str, index: int) -> str: + return f"{'fc' if kind == 'tool' else 'rs' if kind == 'reasoning' else 'msg'}_{index}" + + +def _responses_item_start(block: ContentBlock, key: str, item_id: str, state: StreamFormatState) -> list[str]: + kind = state.item_kinds[key] + if kind == "tool" and block.tool_call: + call = block.tool_call + state.tool_names[key] = call.name or "" + state.tool_ids[key] = call.id or item_id + item = {"id": item_id, "type": "function_call", "call_id": state.tool_ids[key], "name": state.tool_names[key], "arguments": "", "status": "in_progress"} + return [_event_frame("response.output_item.added", {"type": "response.output_item.added", "output_index": state.next_index - 1, "item": item})] + if kind == "reasoning": + item = {"id": item_id, "type": "reasoning", "summary": [], "status": "in_progress"} + return [ + _event_frame("response.output_item.added", {"type": "response.output_item.added", "output_index": state.next_index - 1, "item": item}), + _event_frame("response.reasoning_summary_part.added", {"type": "response.reasoning_summary_part.added", "item_id": item_id, "output_index": state.next_index - 1, "summary_index": 0, "part": {"type": "summary_text", "text": ""}}), + ] + item = {"id": item_id, "type": "message", "role": "assistant", "content": [], "status": "in_progress"} + return [ + _event_frame("response.output_item.added", {"type": "response.output_item.added", "output_index": state.next_index - 1, "item": item}), + _event_frame("response.content_part.added", {"type": "response.content_part.added", "item_id": item_id, "output_index": state.next_index - 1, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}}), + ] + + +def _responses_item_delta(block: ContentBlock, key: str, item_id: str, state: StreamFormatState) -> list[str]: + kind = state.item_kinds[key] + output_index = list(state.item_ids).index(key) + if kind == "tool" and block.tool_call: + fragment = tool_arguments_text(block.tool_call.arguments) + state.tool_arguments[key] = state.tool_arguments.get(key, "") + fragment + return [_event_frame("response.function_call_arguments.delta", {"type": "response.function_call_arguments.delta", "item_id": item_id, "output_index": output_index, "delta": fragment})] if fragment else [] + text = block.reasoning.text if block.reasoning else block.text + if not text: + return [] + state.text_by_key[key] = state.text_by_key.get(key, "") + text + event_name = "response.reasoning_summary_text.delta" if kind == "reasoning" else "response.output_text.delta" + payload = {"type": event_name, "item_id": item_id, "output_index": output_index, "delta": text} + payload["summary_index" if kind == "reasoning" else "content_index"] = 0 + return [_event_frame(event_name, payload)] + + +def _responses_item_done( + key: str, + item_id: str, + state: StreamFormatState, + *, + item_status: str, +) -> list[str]: + kind = state.item_kinds[key] + output_index = list(state.item_ids).index(key) + if kind == "tool": + arguments = state.tool_arguments.get(key, "") + item = {"id": item_id, "type": "function_call", "call_id": state.tool_ids.get(key, item_id), "name": state.tool_names.get(key, ""), "arguments": arguments, "status": item_status} + return [ + _event_frame("response.function_call_arguments.done", {"type": "response.function_call_arguments.done", "item_id": item_id, "output_index": output_index, "arguments": arguments}), + _event_frame("response.output_item.done", {"type": "response.output_item.done", "output_index": output_index, "item": item}), + ] + text = state.text_by_key.get(key, "") + if kind == "reasoning": + item = {"id": item_id, "type": "reasoning", "summary": [{"type": "summary_text", "text": text}], "status": item_status} + return [ + _event_frame("response.reasoning_summary_text.done", {"type": "response.reasoning_summary_text.done", "item_id": item_id, "output_index": output_index, "summary_index": 0, "text": text}), + _event_frame("response.reasoning_summary_part.done", {"type": "response.reasoning_summary_part.done", "item_id": item_id, "output_index": output_index, "summary_index": 0, "part": item["summary"][0]}), + _event_frame("response.output_item.done", {"type": "response.output_item.done", "output_index": output_index, "item": item}), + ] + part = {"type": "output_text", "text": text, "annotations": []} + item = {"id": item_id, "type": "message", "role": "assistant", "content": [part], "status": item_status} + return [ + _event_frame("response.output_text.done", {"type": "response.output_text.done", "item_id": item_id, "output_index": output_index, "content_index": 0, "text": text}), + _event_frame("response.content_part.done", {"type": "response.content_part.done", "item_id": item_id, "output_index": output_index, "content_index": 0, "part": part}), + _event_frame("response.output_item.done", {"type": "response.output_item.done", "output_index": output_index, "item": item}), + ] + + +def _responses_object(state: StreamFormatState, *, status: str, error: Any = None) -> dict[str, Any]: + output = [] + item_status = "completed" if status == "completed" else "in_progress" if status == "in_progress" else "incomplete" + for key, item_id in state.item_ids.items(): + kind = state.item_kinds[key] + if kind == "tool": + output.append({"id": item_id, "type": "function_call", "call_id": state.tool_ids.get(key, item_id), "name": state.tool_names.get(key, ""), "arguments": state.tool_arguments.get(key, ""), "status": item_status}) + elif kind == "reasoning": + output.append({"id": item_id, "type": "reasoning", "summary": [{"type": "summary_text", "text": state.text_by_key.get(key, "")}], "status": item_status}) + else: + output.append({"id": item_id, "type": "message", "role": "assistant", "content": [{"type": "output_text", "text": state.text_by_key.get(key, ""), "annotations": []}], "status": item_status}) + payload: dict[str, Any] = {"id": state.response_id, "object": "response", "status": status, "model": state.model, "output": output} + usage = _responses_usage(state.usage) + if usage: + payload["usage"] = usage + if error is not None: + payload["error"] = error + return payload + + +def _is_terminal(event: UnifiedStreamEvent) -> bool: + return event.type in { + "done", + "message_stop", + "response.completed", + "response.failed", + "response.incomplete", + "completed", + } + + +def _data_frame(payload: dict[str, Any]) -> str: + return f"data: {json.dumps(serialize_value(payload), ensure_ascii=False)}\n\n" + + +def decode_sse_data(raw_event: Any) -> Any: + """Decode one SSE frame while accepting optional event and comment lines.""" + + if not isinstance(raw_event, str): + return raw_event + text = raw_event.strip() + data_lines = [ + line[len("data:") :].strip() + for line in text.splitlines() + if line.strip().startswith("data:") + ] + if data_lines: + text = "\n".join(data_lines).strip() + if text == "[DONE]": + return text + try: + return json.loads(text) + except json.JSONDecodeError: + return raw_event + + +def _event_frame(name: str, payload: dict[str, Any]) -> str: + return f"event: {name}\ndata: {json.dumps(serialize_value(payload), ensure_ascii=False)}\n\n" + + +def _error_payload(error: Any) -> dict[str, Any]: + if isinstance(error, dict): + return serialize_value(error) + return {"type": "server_error", "message": str(error or "Provider stream failed")} + + +def _openai_usage(usage: Usage | None) -> dict[str, Any] | None: + if usage is None: + return None + payload: dict[str, Any] = { + "prompt_tokens": usage.input_tokens + usage.cache_read_tokens + usage.cache_write_tokens, + "completion_tokens": usage.output_tokens, + "total_tokens": usage.total_tokens, + } + if usage.cache_read_tokens or usage.cache_write_tokens: + payload["prompt_tokens_details"] = {"cached_tokens": usage.cache_read_tokens, "cache_creation_tokens": usage.cache_write_tokens} + if usage.reasoning_tokens: + payload["completion_tokens_details"] = {"reasoning_tokens": usage.reasoning_tokens} + return payload + + +def _anthropic_usage(usage: Usage | None, *, output_only: bool = False) -> dict[str, int]: + if usage is None: + return {"output_tokens": 0} if output_only else {"input_tokens": 0, "output_tokens": 0} + payload = {"output_tokens": usage.output_tokens} + if not output_only: + payload["input_tokens"] = usage.input_tokens + if usage.cache_read_tokens: + payload["cache_read_input_tokens"] = usage.cache_read_tokens + if usage.cache_write_tokens: + payload["cache_creation_input_tokens"] = usage.cache_write_tokens + return payload + + +def _responses_usage(usage: Usage | None) -> dict[str, Any] | None: + if usage is None: + return None + return { + "input_tokens": usage.input_tokens + usage.cache_read_tokens + usage.cache_write_tokens, + "input_tokens_details": {"cached_tokens": usage.cache_read_tokens}, + "output_tokens": usage.output_tokens, + "output_tokens_details": {"reasoning_tokens": usage.reasoning_tokens}, + "total_tokens": usage.total_tokens, + } + + +def _gemini_usage(usage: Usage | None) -> dict[str, int] | None: + if usage is None: + return None + return { + "promptTokenCount": usage.input_tokens + usage.cache_read_tokens, + "candidatesTokenCount": usage.output_tokens, + "totalTokenCount": usage.total_tokens, + "cachedContentTokenCount": usage.cache_read_tokens, + "thoughtsTokenCount": usage.reasoning_tokens, + } + + +def _is_json(value: str) -> bool: + try: + json.loads(value) + except json.JSONDecodeError: + return False + return True diff --git a/src/rotator_library/protocols/types.py b/src/rotator_library/protocols/types.py new file mode 100644 index 000000000..e74ab1496 --- /dev/null +++ b/src/rotator_library/protocols/types.py @@ -0,0 +1,551 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Protocol-neutral data structures used by native protocol adapters. + +These types intentionally model common LLM API concepts without pretending the +set is complete. Providers and future protocol implementations can preserve +non-standard fields in ``extra`` or ``raw`` instead of dropping them. That +preservation is important for transform-pass logging, field-cache rules, and +provider-specific overrides added in later experimental phases. +""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import asdict, dataclass, field, is_dataclass +from datetime import date, datetime +from decimal import Decimal +from pathlib import Path +from typing import Any, ClassVar, Iterable, Mapping, Optional + +from .operation import OPERATION_UNKNOWN + + +JsonObject = dict[str, Any] + + +def serialize_value(value: Any) -> Any: + """Return a JSON-friendly copy of a protocol value. + + The transaction logging phase needs reliable snapshots after every protocol + and adapter pass. This helper keeps that concern centralized and avoids + mutating live request/response objects while preparing logs or fixtures. + """ + + if isinstance(value, ProtocolSerializable): + return value.to_dict() + if is_dataclass(value): + return serialize_value(asdict(value)) + if isinstance(value, Mapping): + return {str(k): serialize_value(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [serialize_value(v) for v in value] + if isinstance(value, bytes): + try: + return value.decode("utf-8") + except UnicodeDecodeError: + return value.hex() + if isinstance(value, (datetime, date)): + return value.isoformat() + if isinstance(value, Decimal): + return float(value) + if isinstance(value, Path): + return str(value) + try: + import json + + json.dumps(value) + return deepcopy(value) + except (TypeError, ValueError): + return repr(value) + + +def copy_mapping(value: Optional[Mapping[str, Any]]) -> JsonObject: + """Return a deep-copied dict for extension fields.""" + + return serialize_value(dict(value or {})) + + +class ProtocolSerializable: + """Mixin for dataclasses that need stable dict serialization. + + Concrete protocol types define ``_fields`` so ``to_dict`` stays explicit and + future additions do not accidentally disappear from transform logs. + """ + + _fields: ClassVar[tuple[str, ...]] = () + + def to_dict(self) -> JsonObject: + return {field_name: serialize_value(getattr(self, field_name)) for field_name in self._fields} + + +@dataclass +class CostDetails(ProtocolSerializable): + """Normalized cost metadata from provider-reported or estimated sources.""" + + provider_reported_cost: Optional[float] = None + estimated_cost: Optional[float] = None + currency: str = "USD" + source: Optional[str] = None + metadata: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "provider_reported_cost", + "estimated_cost", + "currency", + "source", + "metadata", + ) + + +@dataclass +class Usage(ProtocolSerializable): + """Protocol-neutral token and cost usage values. + + Existing usage tracking has provider-specific extraction paths. Native + protocols use this shape first, then later phases can normalize it into the + current usage manager without replacing that engine. + """ + + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + cache_read_tokens: int = 0 + cache_write_tokens: int = 0 + reasoning_tokens: int = 0 + cost: Optional[CostDetails] = None + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "input_tokens", + "output_tokens", + "total_tokens", + "cache_read_tokens", + "cache_write_tokens", + "reasoning_tokens", + "cost", + "raw", + "extra", + ) + + def __post_init__(self) -> None: + if self.total_tokens <= 0: + # Most APIs report reasoning/thinking tokens as a detail of output + # tokens. Protocol-specific normalizers can provide a larger total + # when a provider documents reasoning as a separate bucket. + self.total_tokens = self.input_tokens + self.output_tokens + + +@dataclass +class ReasoningBlock(ProtocolSerializable): + """Reasoning/thinking content and signatures preserved across protocols.""" + + type: str = "reasoning" + text: Optional[str] = None + signature: Optional[str] = None + redacted: bool = False + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ("type", "text", "signature", "redacted", "raw", "extra") + + +@dataclass +class MediaSource(ProtocolSerializable): + """Canonical media identity independent of a protocol's wire spelling. + + A media value can be URL-backed, embedded, or provider-file-backed. Protocol + parsers populate the common fields and retain the source object only for + same-protocol fidelity. Cross-protocol writers consume the common fields. + """ + + kind: str = "url" + media_type: Optional[str] = None + url: Optional[str] = None + data: Optional[str] = None + file_id: Optional[str] = None + detail: Optional[str] = None + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "kind", + "media_type", + "url", + "data", + "file_id", + "detail", + "raw", + "extra", + ) + + +@dataclass +class ToolCall(ProtocolSerializable): + """Protocol-neutral tool/function call emitted by an assistant.""" + + id: Optional[str] = None + name: Optional[str] = None + arguments: Any = None + type: str = "function" + index: Optional[int] = None + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ("id", "name", "arguments", "type", "index", "raw", "extra") + + +@dataclass +class ToolResult(ProtocolSerializable): + """Protocol-neutral result associated with a prior tool call.""" + + tool_call_id: Optional[str] = None + name: Optional[str] = None + content: Any = None + is_error: Optional[bool] = None + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ("tool_call_id", "name", "content", "is_error", "raw", "extra") + + +@dataclass +class ToolDefinition(ProtocolSerializable): + """Protocol-neutral tool schema exposed to a model.""" + + name: str = "" + description: Optional[str] = None + input_schema: JsonObject = field(default_factory=dict) + type: str = "function" + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ("name", "description", "input_schema", "type", "extra") + + +@dataclass +class ContentBlock(ProtocolSerializable): + """A single message content block. + + ``type`` is canonical (for example ``text``, ``image``, ``tool_call``, or + ``reasoning``). ``raw`` and ``extra`` retain source-protocol details for + same-protocol fidelity and field-cache extraction; foreign formatters must + not replay them without an explicit mapping. + """ + + type: str = "text" + text: Optional[str] = None + source: Any = None + tool_call: Optional[ToolCall] = None + tool_result: Optional[ToolResult] = None + reasoning: Optional[ReasoningBlock] = None + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "type", + "text", + "source", + "tool_call", + "tool_result", + "reasoning", + "raw", + "extra", + ) + + +@dataclass +class ConversionWarning(ProtocolSerializable): + """A non-fatal, explicit loss of an optional conversion hint.""" + + code: str + message: str + field: Optional[str] = None + source_protocol: Optional[str] = None + target_protocol: Optional[str] = None + + _fields: ClassVar[tuple[str, ...]] = ( + "code", + "message", + "field", + "source_protocol", + "target_protocol", + ) + + +@dataclass +class UnifiedMessage(ProtocolSerializable): + """A protocol-neutral chat/message turn.""" + + role: str + content: list[ContentBlock] = field(default_factory=list) + name: Optional[str] = None + tool_call_id: Optional[str] = None + tool_calls: list[ToolCall] = field(default_factory=list) + reasoning: list[ReasoningBlock] = field(default_factory=list) + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "role", + "content", + "name", + "tool_call_id", + "tool_calls", + "reasoning", + "raw", + "extra", + ) + + +@dataclass +class UnifiedRequest(ProtocolSerializable): + """A request after parsing from a client or provider protocol.""" + + operation: str = OPERATION_UNKNOWN + logical_operation: str = OPERATION_UNKNOWN + model: str = "" + messages: list[UnifiedMessage] = field(default_factory=list) + system: list[ContentBlock] = field(default_factory=list) + tools: list[ToolDefinition] = field(default_factory=list) + stream: bool = False + input: Any = None + modalities: list[str] = field(default_factory=list) + files: list[Any] = field(default_factory=list) + generation_params: JsonObject = field(default_factory=dict) + response_format: Any = None + previous_response_id: Optional[str] = None + metadata: JsonObject = field(default_factory=dict) + source_protocol: Optional[str] = None + extensions: dict[str, JsonObject] = field(default_factory=dict) + warnings: list[ConversionWarning] = field(default_factory=list) + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "operation", + "logical_operation", + "model", + "messages", + "system", + "tools", + "stream", + "input", + "modalities", + "files", + "generation_params", + "response_format", + "previous_response_id", + "metadata", + "source_protocol", + "extensions", + "warnings", + "raw", + "extra", + ) + + +@dataclass +class OutputItem(ProtocolSerializable): + """One ordered semantic output item from a generative response.""" + + type: str + id: Optional[str] = None + role: Optional[str] = None + content: list[ContentBlock] = field(default_factory=list) + tool_call: Optional[ToolCall] = None + reasoning: Optional[ReasoningBlock] = None + status: Optional[str] = None + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "type", + "id", + "role", + "content", + "tool_call", + "reasoning", + "status", + "raw", + "extra", + ) + + +@dataclass +class UnifiedResponse(ProtocolSerializable): + """A complete provider/client response in protocol-neutral form.""" + + operation: str = OPERATION_UNKNOWN + logical_operation: str = OPERATION_UNKNOWN + id: Optional[str] = None + model: Optional[str] = None + messages: list[UnifiedMessage] = field(default_factory=list) + items: list[OutputItem] = field(default_factory=list) + output: list[Any] = field(default_factory=list) + data: list[Any] = field(default_factory=list) + content_type: Optional[str] = None + stop_reason: Optional[str] = None + usage: Optional[Usage] = None + metadata: JsonObject = field(default_factory=dict) + source_protocol: Optional[str] = None + extensions: dict[str, JsonObject] = field(default_factory=dict) + warnings: list[ConversionWarning] = field(default_factory=list) + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "operation", + "logical_operation", + "id", + "model", + "messages", + "items", + "output", + "data", + "content_type", + "stop_reason", + "usage", + "metadata", + "source_protocol", + "extensions", + "warnings", + "raw", + "extra", + ) + + +@dataclass +class UnifiedStreamEvent(ProtocolSerializable): + """A single protocol-neutral stream event. + + Future SSE and WebSocket transports should consume this type instead of raw + provider chunks so transport code can stay independent from protocol parsing. + """ + + type: str + operation: str = OPERATION_UNKNOWN + logical_operation: str = OPERATION_UNKNOWN + source_protocol: Optional[str] = None + native_type: Optional[str] = None + delta: Optional[UnifiedMessage] = None + message: Optional[UnifiedMessage] = None + tool_call: Optional[ToolCall] = None + usage: Optional[Usage] = None + error: Any = None + item_id: Optional[str] = None + output_index: Optional[int] = None + content_index: Optional[int] = None + stop_reason: Optional[str] = None + raw: Any = None + extra: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "type", + "operation", + "logical_operation", + "source_protocol", + "native_type", + "delta", + "message", + "tool_call", + "usage", + "error", + "item_id", + "output_index", + "content_index", + "stop_reason", + "raw", + "extra", + ) + + +@dataclass +class ProtocolContext(ProtocolSerializable): + """Execution context passed through protocol methods. + + Only a small subset is needed in Phase 1, but the fields anticipate later + provider overrides, transaction tracing, field-cache scoping, and transport + selection without forcing those systems to exist yet. + """ + + provider: Optional[str] = None + model: Optional[str] = None + source_protocol: Optional[str] = None + target_protocol: Optional[str] = None + input_protocol: Optional[str] = None + provider_protocol: Optional[str] = None + output_protocol: Optional[str] = None + source_provider: Optional[str] = None + target_provider: Optional[str] = None + provider_state_compatible: bool = False + request_id: Optional[str] = None + session_id: Optional[str] = None + credential_stable_id: Optional[str] = None + transport: str = "http" + provider_options: JsonObject = field(default_factory=dict) + metadata: JsonObject = field(default_factory=dict) + + _fields: ClassVar[tuple[str, ...]] = ( + "provider", + "model", + "source_protocol", + "target_protocol", + "input_protocol", + "provider_protocol", + "output_protocol", + "source_provider", + "target_provider", + "provider_state_compatible", + "request_id", + "session_id", + "credential_stable_id", + "transport", + "provider_options", + "metadata", + ) + + +class ProtocolError(ValueError): + """Error raised by protocol parsing/building passes.""" + + def __init__( + self, + message: str, + *, + protocol: str, + pass_name: str, + payload: Any = None, + ): + self.protocol = protocol + self.pass_name = pass_name + self.payload_preview = _payload_preview(payload) + details = f"{protocol}.{pass_name}: {message}" + if self.payload_preview is not None: + details = f"{details} | payload={self.payload_preview}" + super().__init__(details) + + +def _payload_preview(payload: Any, limit: int = 500) -> Optional[str]: + if payload is None: + return None + text = repr(serialize_value(payload)) + if len(text) > limit: + return text[: limit - 3] + "..." + return text + + +def text_blocks(text: Optional[str]) -> list[ContentBlock]: + """Return a single text block list for simple string content.""" + + if text is None: + return [] + return [ContentBlock(type="text", text=str(text))] + + +def first_text(blocks: Iterable[ContentBlock]) -> Optional[str]: + """Return concatenated text from content blocks, or ``None`` if absent.""" + + parts = [block.text for block in blocks if block.type == "text" and block.text] + return "".join(parts) if parts else None diff --git a/src/rotator_library/protocols/validation.py b/src/rotator_library/protocols/validation.py new file mode 100644 index 000000000..9e85d8697 --- /dev/null +++ b/src/rotator_library/protocols/validation.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Destination capability validation for generative protocol conversion.""" + +from __future__ import annotations + +from typing import Any + +from .canonical import ( + STOP_REASON_ERROR, + canonical_tool_arguments, + is_same_protocol, + message_tool_calls, + message_tool_results, +) +from .types import MediaSource, ProtocolContext, ProtocolError, UnifiedRequest, UnifiedResponse + + +_CONTENT_CAPABILITIES: dict[str, set[str]] = { + "openai_chat": {"text", "image", "audio", "file", "document", "reasoning", "tool_call", "tool_result"}, + "anthropic_messages": {"text", "image", "file", "document", "reasoning", "tool_call", "tool_result"}, + "responses": {"text", "image", "file", "document", "reasoning", "tool_call", "tool_result"}, + "gemini": {"text", "image", "audio", "video", "file", "document", "reasoning", "tool_call", "tool_result"}, +} + +_RESPONSE_MODALITIES: dict[str, set[str]] = { + "openai_chat": {"text", "audio"}, + "anthropic_messages": {"text"}, + "responses": {"text", "audio"}, + "gemini": {"text", "audio", "image"}, +} + + +def validate_generative_request( + request: UnifiedRequest, + target_protocol: str, + context: ProtocolContext | None, +) -> None: + """Reject meaning-changing cross-protocol losses before provider transport. + + Same-protocol requests may retain future native content through their raw + payload. Cross-protocol requests must use a known canonical meaning so a + source-native object can never be emitted as a malformed foreign object. + """ + + if is_same_protocol(context, target_protocol, request.source_protocol): + return + if request.previous_response_id and target_protocol != "responses": + raise ProtocolError( + "A provider-bound previous_response_id cannot be translated safely", + protocol=target_protocol, + pass_name="validate_request", + payload={"field": "previous_response_id"}, + ) + provider_bound_responses_fields = [ + field + for field in ("background", "conversation", "prompt") + if request.generation_params.get(field) not in (None, False) + ] + if provider_bound_responses_fields and target_protocol != "responses": + raise ProtocolError( + "Provider-bound Responses controls cannot be translated safely", + protocol=target_protocol, + pass_name="validate_request", + payload={"fields": provider_bound_responses_fields}, + ) + if request.generation_params.get("safety_settings") and target_protocol != "gemini": + raise ProtocolError( + "Gemini safety settings have no equivalent in the selected provider protocol", + protocol=target_protocol, + pass_name="validate_request", + payload={"field": "safety_settings"}, + ) + unsupported_modalities = set(request.modalities) - _RESPONSE_MODALITIES.get(target_protocol, {"text"}) + if unsupported_modalities: + raise ProtocolError( + f"{target_protocol} cannot produce required response modalities: {sorted(unsupported_modalities)}", + protocol=target_protocol, + pass_name="validate_request", + payload={"field": "modalities", "unsupported": sorted(unsupported_modalities)}, + ) + supported = _CONTENT_CAPABILITIES.get(target_protocol, set()) + block_groups = [("system", request.system)] + [ + (f"message:{message_index}", message.content) + for message_index, message in enumerate(request.messages) + ] + for group_name, blocks in block_groups: + for block_index, block in enumerate(blocks): + if block.type in supported: + if block.type in {"image", "audio", "video", "file", "document"} and not _has_media_identity(block.source): + raise ProtocolError( + f"Cannot represent {block.type} content without URL, data, or file identity", + protocol=target_protocol, + pass_name="validate_request", + payload={"group": group_name, "content_index": block_index, "content_type": block.type}, + ) + continue + raise ProtocolError( + f"Cannot represent required content type '{block.type}' in {target_protocol}", + protocol=target_protocol, + pass_name="validate_request", + payload={"group": group_name, "content_index": block_index, "content_type": block.type}, + ) + for tool_index, tool in enumerate(request.tools): + if tool.type != "function": + raise ProtocolError( + f"Cannot safely translate tool type '{tool.type}' into {target_protocol}", + protocol=target_protocol, + pass_name="validate_request", + payload={"tool_index": tool_index, "tool_type": tool.type, "tool_name": tool.name}, + ) + if not tool.name or not isinstance(tool.input_schema, dict): + raise ProtocolError( + "Function tools require a name and object input schema", + protocol=target_protocol, + pass_name="validate_request", + payload={"tool_index": tool_index, "tool_name": tool.name}, + ) + for message_index, message in enumerate(request.messages): + for call_index, call in enumerate(message_tool_calls(message)): + arguments = canonical_tool_arguments(call.arguments) + if not call.name or (target_protocol != "gemini" and not call.id): + raise ProtocolError( + "Tool calls require a name and correlation ID", + protocol=target_protocol, + pass_name="validate_request", + payload={"message_index": message_index, "call_index": call_index}, + ) + if target_protocol in {"anthropic_messages", "gemini"} and not isinstance(arguments, dict): + raise ProtocolError( + f"{target_protocol} requires tool arguments to be a JSON object", + protocol=target_protocol, + pass_name="validate_request", + payload={"message_index": message_index, "call_index": call_index}, + ) + for result_index, result in enumerate(message_tool_results(message)): + if not result.tool_call_id: + raise ProtocolError( + "Tool results require correlation identity", + protocol=target_protocol, + pass_name="validate_request", + payload={"message_index": message_index, "result_index": result_index}, + ) + if target_protocol == "gemini" and not result.name: + raise ProtocolError( + "Gemini tool results require the originating function name", + protocol=target_protocol, + pass_name="validate_request", + payload={"message_index": message_index, "result_index": result_index}, + ) + _validate_tool_choice(request, target_protocol) + + +def validate_generative_response(response: UnifiedResponse, target_protocol: str) -> None: + """Reject failed provider responses that a target success envelope cannot express.""" + + if response.stop_reason == STOP_REASON_ERROR and target_protocol != "responses": + raise ProtocolError( + f"{target_protocol} cannot represent a failed provider response as a successful completion", + protocol=target_protocol, + pass_name="validate_response", + payload={"stop_reason": response.stop_reason}, + ) + + +def _validate_tool_choice(request: UnifiedRequest, target_protocol: str) -> None: + """Require named and allow-listed choices to reference declared tools.""" + + choice: Any = request.generation_params.get("tool_choice") + if not isinstance(choice, dict): + return + available = {tool.name for tool in request.tools if tool.name} + if choice.get("mode") == "named": + name = str(choice.get("name") or "").strip() + if not name: + raise ProtocolError( + "Named tool choice requires a tool name", + protocol=target_protocol, + pass_name="validate_request", + payload={"field": "tool_choice"}, + ) + if name not in available: + raise ProtocolError( + f"Named tool choice references unavailable tool {name!r}", + protocol=target_protocol, + pass_name="validate_request", + payload={"field": "tool_choice", "name": name}, + ) + allowed_names = {str(name) for name in choice.get("allowed_names") or []} + missing = allowed_names - available + if missing: + raise ProtocolError( + f"Tool choice references unavailable tools: {sorted(missing)}", + protocol=target_protocol, + pass_name="validate_request", + payload={"field": "tool_choice", "names": sorted(missing)}, + ) + + +def _has_media_identity(source: object) -> bool: + """Return whether a canonical or legacy media source is transportable.""" + + if isinstance(source, MediaSource): + return bool(source.url or source.data or source.file_id) + if isinstance(source, str): + return bool(source) + if isinstance(source, dict): + return bool( + source.get("url") + or source.get("data") + or source.get("file_id") + or source.get("fileUri") + or source.get("file_uri") + ) + return False diff --git a/src/rotator_library/provider_factory.py b/src/rotator_library/provider_factory.py index 3bb95bd54..1176d61b2 100644 --- a/src/rotator_library/provider_factory.py +++ b/src/rotator_library/provider_factory.py @@ -3,11 +3,9 @@ # src/rotator_library/provider_factory.py -from .providers.gemini_auth_base import GeminiAuthBase - -PROVIDER_MAP = { - "gemini_cli": GeminiAuthBase, -} +# OAuth providers are intentionally opt-in here. Retired OAuth providers live in +# providers/_retired and must not be imported by active startup paths. +PROVIDER_MAP = {} def get_provider_auth_class(provider_name: str): """ diff --git a/src/rotator_library/providers/__init__.py b/src/rotator_library/providers/__init__.py index f719eb209..5cb5c5a51 100644 --- a/src/rotator_library/providers/__init__.py +++ b/src/rotator_library/providers/__init__.py @@ -4,7 +4,7 @@ import importlib import pkgutil import os -from typing import Dict, Type +from typing import Any, Dict, Optional, Type from .provider_interface import ProviderInterface # --- Provider Plugin System --- @@ -13,34 +13,56 @@ PROVIDER_PLUGINS: Dict[str, Type[ProviderInterface]] = {} -class DynamicOpenAICompatibleProvider: - """ - Dynamic provider class for custom OpenAI-compatible providers. - Created at runtime for providers with _API_BASE environment variables - that are NOT known LiteLLM providers. +def _create_dynamic_plugin_class( + name: str, + config_snapshot: Any = None, +) -> Type[ProviderInterface]: + """Create one ProviderInterface implementation bound to a config name.""" + + class DynamicPlugin(DynamicOpenAICompatibleProvider, ProviderInterface): + provider_env_name = name + + def __init__(self): + DynamicOpenAICompatibleProvider.__init__( + self, + name, + config_snapshot=config_snapshot, + ) + + DynamicPlugin.__name__ = f"{''.join(part.title() for part in name.split('_'))}DynamicProvider" + return DynamicPlugin - Environment variable pattern: - _API_BASE - The API base URL (required) - _API_KEY - The API key - Example: - MYSERVER_API_BASE=http://localhost:8000/v1 - MYSERVER_API_KEY=sk-xxx +class DynamicOpenAICompatibleProvider: + """ + Dynamic provider for safe config or ``*_API_BASE`` declarations. - Note: For known providers (openai, anthropic, etc.), setting _API_BASE - will override their default endpoint without creating a custom provider. + Environment-only declarations retain the existing OpenAI-compatible + LiteLLM path. Structured config may additionally opt the provider into any + registered native protocol without storing credentials in the config file. """ # Class attribute - no need to instantiate skip_cost_calculation: bool = True - def __init__(self, provider_name: str): + def __init__(self, provider_name: str, *, config_snapshot: Any = None): self.provider_name = provider_name - # Get API base URL from environment (using _API_BASE pattern) - self.api_base = os.getenv(f"{provider_name.upper()}_API_BASE") + self.provider_env_name = provider_name + from ..config.experimental import get_provider_runtime_config, load_experimental_config + + self._config_snapshot = ( + config_snapshot + if config_snapshot is not None + else load_experimental_config() + ) + runtime = get_provider_runtime_config( + provider_name, + config=self._config_snapshot, + ) + self.api_base = runtime.api_base or os.getenv(f"{provider_name.upper()}_API_BASE") if not self.api_base: raise ValueError( - f"Environment variable {provider_name.upper()}_API_BASE is required for custom OpenAI-compatible provider" + f"API base URL is required for dynamic provider {provider_name!r}" ) # Import model definitions @@ -48,15 +70,47 @@ def __init__(self, provider_name: str): self.model_definitions = ModelDefinitions() - def get_models(self, api_key: str, client): - """Delegate to OpenAI-compatible provider implementation.""" - from .openai_compatible_provider import OpenAICompatibleProvider - - # Create temporary instance to reuse logic - temp_provider = OpenAICompatibleProvider(self.provider_name) - return temp_provider.get_models(api_key, client) - - def get_model_options(self, model_name: str) -> Dict[str, any]: + def _runtime_config(self, model: str = ""): + from ..config.experimental import get_provider_runtime_config + + return get_provider_runtime_config( + self.provider_name, + model, + config=self._config_snapshot, + ) + + def _get_runtime_config(self, model: str = ""): + """Keep custom provider transport identity immutable after startup.""" + + return self._runtime_config(model) + + def get_api_base(self) -> str: + return str(self._runtime_config().api_base or self.api_base).rstrip("/") + + async def get_models(self, api_key: str, client): + """Return configured models or discover common ``/models`` shapes.""" + configured = self._runtime_config().models + if configured: + return [ + model if model.startswith(f"{self.provider_name}/") else f"{self.provider_name}/{model}" + for model in configured + ] + response = await client.get( + f"{self.get_api_base()}/models", + headers=self.get_native_headers(api_key, operation="models"), + ) + response.raise_for_status() + payload = response.json() + entries = payload.get("data") or payload.get("models") or [] if isinstance(payload, dict) else [] + models: list[str] = [] + for entry in entries: + raw_id = entry.get("id") or entry.get("name") if isinstance(entry, dict) else entry + model_id = str(raw_id or "").removeprefix("models/") + if model_id: + models.append(f"{self.provider_name}/{model_id}") + return models + + def get_model_options(self, model_name: str) -> Dict[str, Any]: """Get model options from static definitions.""" # Extract model name without provider prefix if present if "/" in model_name: @@ -69,8 +123,80 @@ def has_custom_logic(self) -> bool: return False def get_auth_header(self, credential_identifier: str) -> Dict[str, str]: - """Returns the standard Bearer token header.""" - return {"Authorization": f"Bearer {credential_identifier}"} + """Return the configured credential header.""" + return self.get_native_headers(credential_identifier) + + def get_native_operation( + self, + model: str = "", + request: Optional[Dict[str, Any]] = None, + stream: bool = False, + ) -> str: + protocol = self.get_protocol_name(model) + if protocol == "responses": + return "responses" + if protocol == "anthropic_messages": + return "messages" + if protocol == "gemini": + return "stream_generate" if stream else "generate" + return "chat" + + def normalize_native_model(self, model: str = "") -> str: + prefix = f"{self.provider_name}/" + return model[len(prefix):] if model.startswith(prefix) else model + + def get_native_endpoint(self, model: str = "", operation: str = "chat") -> str: + runtime = self._runtime_config(model) + protocol = self.get_protocol_name(model) or "openai_chat" + defaults = { + "openai_chat": {"chat": "/chat/completions"}, + "responses": {"responses": "/responses"}, + "anthropic_messages": {"messages": "/messages"}, + "gemini": { + "generate": "/models/{model}:generateContent", + "stream_generate": "/models/{model}:streamGenerateContent?alt=sse", + "count_tokens": "/models/{model}:countTokens", + }, + } + path = runtime.endpoint_paths.get(operation) or defaults.get(protocol, {}).get(operation) + if not path: + raise NotImplementedError( + f"Dynamic provider {self.provider_name} has no endpoint for {protocol}/{operation}" + ) + rendered = path.format( + model=self.normalize_native_model(model), + operation=operation, + provider=self.provider_name, + ) + if rendered.startswith(("http://", "https://")): + return rendered + return f"{self.get_api_base()}/{rendered.lstrip('/')}" + + def get_native_headers( + self, + credential_identifier: str, + model: str = "", + operation: str = "chat", + ) -> Dict[str, str]: + runtime = self._runtime_config(model) + headers: Dict[str, str] = {"Content-Type": "application/json"} + if operation == "stream_generate": + headers["Accept"] = "text/event-stream" + if runtime.auth_mode == "none": + return headers + if runtime.auth_mode == "x-api-key": + headers["x-api-key"] = credential_identifier + elif runtime.auth_mode == "x-goog-api-key": + headers["x-goog-api-key"] = credential_identifier + elif runtime.auth_mode == "custom": + if not runtime.auth_header_name: + raise ValueError( + f"Dynamic provider {self.provider_name} requires auth_header_name for custom auth" + ) + headers[runtime.auth_header_name] = credential_identifier + else: + headers["Authorization"] = f"Bearer {credential_identifier}" + return headers def _register_providers(): @@ -80,6 +206,9 @@ def _register_providers(): """ package_path = __path__ package_name = __name__ + from ..config.experimental import load_experimental_config + + config_snapshot = load_experimental_config() # First, register file-based providers. Archive/private modules are skipped # so retired providers can remain in-tree without becoming accessible. @@ -101,7 +230,7 @@ def _register_providers(): and issubclass(attribute, ProviderInterface) and attribute is not ProviderInterface ): - # Derives 'gemini_cli' from 'gemini_cli_provider.py' + # Derives 'openrouter' from 'openrouter_provider.py' # Remap 'nvidia' to 'nvidia_nim' to align with litellm's provider name provider_name = module_name.replace("_provider", "") if provider_name == "nvidia": @@ -132,16 +261,10 @@ def _register_providers(): if provider_name in PROVIDER_PLUGINS: continue - # Create a dynamic plugin class - def create_plugin_class(name): - class DynamicPlugin(DynamicOpenAICompatibleProvider): - def __init__(self): - super().__init__(name) - - return DynamicPlugin - - # Create and register the plugin class - plugin_class = create_plugin_class(provider_name) + plugin_class = _create_dynamic_plugin_class( + provider_name, + config_snapshot=config_snapshot, + ) PROVIDER_PLUGINS[provider_name] = plugin_class import logging @@ -149,6 +272,56 @@ def __init__(self): f"Registered dynamic provider: {provider_name}" ) + # Structured config can define custom providers without a parallel API_BASE + # environment variable. Credentials remain in the existing secret stores. + configured = config_snapshot.providers + for raw_name, raw in configured.items(): + provider_name = str(raw_name).lower() + if provider_name in PROVIDER_PLUGINS: + transport_keys = { + "api_base", + "endpoint_paths", + "auth_mode", + "auth_header_name", + "models", + } & set(raw if isinstance(raw, dict) else {}) + if transport_keys: + raise ValueError( + f"Provider {provider_name!r} is implemented in code; custom transport keys are not applied: " + f"{', '.join(sorted(transport_keys))}" + ) + continue + if not isinstance(raw, dict): + continue + configured_protocol = raw.get("protocol_name") + if configured_protocol: + from ..protocols import get_protocol + + protocol_name = get_protocol(str(configured_protocol).strip().lower()).name + if protocol_name not in { + "openai_chat", + "responses", + "anthropic_messages", + "gemini", + }: + raise ValueError( + f"Configured custom provider {provider_name!r} requires a supported generative protocol, got {protocol_name!r}" + ) + api_base = str(raw.get("api_base") or os.getenv(f"{provider_name.upper()}_API_BASE") or "").strip() + if not api_base: + raise ValueError( + f"Configured custom provider {provider_name!r} requires providers.{provider_name}.api_base" + ) + PROVIDER_PLUGINS[provider_name] = _create_dynamic_plugin_class( + provider_name, + config_snapshot=config_snapshot, + ) + import logging + + logging.getLogger("rotator_library").debug( + f"Registered config-defined provider: {provider_name}" + ) + # Discover and register providers when the package is imported _register_providers() diff --git a/src/rotator_library/providers/_retired/README.md b/src/rotator_library/providers/_retired/README.md index b8977ab7d..0149163e9 100644 --- a/src/rotator_library/providers/_retired/README.md +++ b/src/rotator_library/providers/_retired/README.md @@ -6,3 +6,7 @@ registered, discovered, or exposed by the proxy. Archived modules are kept for historical reference only. Do not import them from active code or add them back to provider factory/discovery wiring unless the provider is explicitly un-retired. + +Retired OAuth-backed Google/CLI providers and their helper modules are archived +here together so active startup paths do not import OAuth credential flows. The +API-key Gemini provider remains active outside this folder. diff --git a/src/rotator_library/providers/_retired/antigravity_auth_base.py b/src/rotator_library/providers/_retired/antigravity_auth_base.py index e22102669..1bcff5e24 100644 --- a/src/rotator_library/providers/_retired/antigravity_auth_base.py +++ b/src/rotator_library/providers/_retired/antigravity_auth_base.py @@ -24,7 +24,7 @@ # Import tier utilities from shared module # These are re-exported here for backwards compatibility with existing imports -from .utilities.gemini_shared_utils import ( +from .gemini_shared_utils import ( # Tier constants TIER_ULTRA, TIER_PRO, @@ -195,7 +195,7 @@ async def _call_load_code_assist( Returns: Tuple of (response_data, successful_endpoint) or (None, None) on failure """ - from .utilities.gemini_shared_utils import ANTIGRAVITY_LOAD_ENDPOINT_ORDER + from .gemini_shared_utils import ANTIGRAVITY_LOAD_ENDPOINT_ORDER core_client_metadata = { "ideType": "IDE_UNSPECIFIED", @@ -257,7 +257,7 @@ async def _call_onboard_user( Returns: Response data dict or None on failure """ - from .utilities.gemini_shared_utils import ANTIGRAVITY_ENDPOINT_FALLBACKS + from .gemini_shared_utils import ANTIGRAVITY_ENDPOINT_FALLBACKS last_error = None for endpoint in ANTIGRAVITY_ENDPOINT_FALLBACKS: diff --git a/src/rotator_library/providers/_retired/antigravity_provider.py b/src/rotator_library/providers/_retired/antigravity_provider.py index ee569b5c3..b0f9531eb 100644 --- a/src/rotator_library/providers/_retired/antigravity_provider.py +++ b/src/rotator_library/providers/_retired/antigravity_provider.py @@ -52,11 +52,11 @@ import httpx import litellm -from .provider_interface import ProviderInterface, UsageResetConfigDef, QuotaGroupMap +from ..provider_interface import ProviderInterface, UsageResetConfigDef, QuotaGroupMap from .antigravity_auth_base import AntigravityAuthBase -from .provider_cache import ProviderCache -from .utilities.antigravity_quota_tracker import AntigravityQuotaTracker -from .utilities.gemini_shared_utils import ( +from ..provider_cache import ProviderCache +from .antigravity_quota_tracker import AntigravityQuotaTracker +from .gemini_shared_utils import ( env_bool, env_int, inline_schema_refs, @@ -70,16 +70,16 @@ TIER_PRIORITIES, DEFAULT_TIER_PRIORITY, ) -from ..transaction_logger import AntigravityProviderLogger -from .utilities.gemini_tool_handler import GeminiToolHandler -from .utilities.gemini_credential_manager import GeminiCredentialManager -from ..model_definitions import ModelDefinitions -from ..timeout_config import TimeoutConfig -from ..error_handler import EmptyResponseError, TransientQuotaError -from ..utils.paths import get_logs_dir, get_cache_dir +from ...transaction_logger import AntigravityProviderLogger +from .gemini_tool_handler import GeminiToolHandler +from .gemini_credential_manager import GeminiCredentialManager +from ...model_definitions import ModelDefinitions +from ...timeout_config import TimeoutConfig +from ...error_handler import EmptyResponseError, TransientQuotaError +from ...utils.paths import get_logs_dir, get_cache_dir if TYPE_CHECKING: - from ..usage import UsageManager + from ...usage import UsageManager # ============================================================================= @@ -1613,7 +1613,7 @@ def _get_antigravity_headers( email = self._get_credential_email(credential_path) if email: try: - from .utilities.device_profile import ( + from .device_profile import ( get_or_create_fingerprint, build_fingerprint_headers, ) diff --git a/src/rotator_library/providers/_retired/antigravity_quota_tracker.py b/src/rotator_library/providers/_retired/antigravity_quota_tracker.py index dfecc3edf..9faa3dc2a 100644 --- a/src/rotator_library/providers/_retired/antigravity_quota_tracker.py +++ b/src/rotator_library/providers/_retired/antigravity_quota_tracker.py @@ -36,7 +36,7 @@ import httpx -from .base_quota_tracker import BaseQuotaTracker, QUOTA_DISCOVERY_DELAY_SECONDS +from ..utilities.base_quota_tracker import BaseQuotaTracker, QUOTA_DISCOVERY_DELAY_SECONDS if TYPE_CHECKING: from ...usage import UsageManager diff --git a/src/rotator_library/providers/gemini_auth_base.py b/src/rotator_library/providers/_retired/gemini_auth_base.py similarity index 99% rename from src/rotator_library/providers/gemini_auth_base.py rename to src/rotator_library/providers/_retired/gemini_auth_base.py index e0b037260..c170deaf5 100644 --- a/src/rotator_library/providers/gemini_auth_base.py +++ b/src/rotator_library/providers/_retired/gemini_auth_base.py @@ -13,7 +13,7 @@ import httpx from .google_oauth_base import GoogleOAuthBase -from .utilities.gemini_shared_utils import ( +from .gemini_shared_utils import ( CODE_ASSIST_ENDPOINT, GEMINI_CLI_UA_VERSION, GEMINI_CLI_NODE_CLIENT_VERSION, diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/_retired/gemini_cli_provider.py similarity index 97% rename from src/rotator_library/providers/gemini_cli_provider.py rename to src/rotator_library/providers/_retired/gemini_cli_provider.py index fe0080ca3..2e28098cd 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/_retired/gemini_cli_provider.py @@ -10,11 +10,12 @@ import time import asyncio from typing import List, Dict, Any, AsyncGenerator, Union, Optional, Tuple -from .provider_interface import ProviderInterface, QuotaGroupMap, UsageResetConfigDef +from ..provider_interface import ProviderInterface, QuotaGroupMap, UsageResetConfigDef from .gemini_auth_base import GeminiAuthBase -from .provider_cache import ProviderCache -from .utilities.gemini_cli_quota_tracker import GeminiCliQuotaTracker -from .utilities.gemini_shared_utils import ( +from ..provider_cache import ProviderCache +from ...field_cache import FieldCacheInjection, FieldCacheRule +from .gemini_cli_quota_tracker import GeminiCliQuotaTracker +from .gemini_shared_utils import ( env_bool, env_int, inline_schema_refs, @@ -33,16 +34,16 @@ TIER_PRIORITIES, DEFAULT_TIER_PRIORITY, ) -from ..transaction_logger import ProviderLogger -from .utilities.gemini_tool_handler import GeminiToolHandler -from .utilities.gemini_credential_manager import GeminiCredentialManager -from ..model_definitions import ModelDefinitions -from ..timeout_config import TimeoutConfig -from ..utils.paths import get_cache_dir -from ..session_tracking import SessionTrackingHints +from ...transaction_logger import ProviderLogger +from .gemini_tool_handler import GeminiToolHandler +from .gemini_credential_manager import GeminiCredentialManager +from ...model_definitions import ModelDefinitions +from ...timeout_config import TimeoutConfig +from ...utils.paths import get_cache_dir +from ...session_tracking import SessionTrackingHints import litellm from litellm.exceptions import RateLimitError -from ..error_handler import extract_retry_after_from_body +from ...error_handler import extract_retry_after_from_body import os from pathlib import Path import uuid @@ -156,6 +157,24 @@ class GeminiCliProvider( # Provider name for env var lookups (QUOTA_GROUPS_GEMINI_CLI_*) provider_env_name: str = "gemini_cli" + # Native protocol declarations for the Phase 5 opt-in execution seam. The + # existing custom Gemini CLI execution path remains active; these metadata + # hooks let future native routing reuse the Gemini protocol and field-cache + # engine without rewriting this provider first. + protocol_name: str = "gemini" + adapter_names: tuple[str, ...] = () + field_cache_rules = ( + FieldCacheRule( + name="gemini_cli_thought_signature", + source="response", + path="candidates.*.content.parts.*.thoughtSignature", + mode="all", + scope=("provider", "model", "credential", "session"), + inject=FieldCacheInjection(target="request", path="metadata.thoughtSignatures", as_list=True), + metadata={"purpose": "mirror existing Gemini CLI thoughtSignature preservation in the native cache seam"}, + ), + ) + # Tier name -> priority mapping (from centralized tier utilities) # Lower numbers = higher priority (ULTRA=1 > PRO=2 > FREE=3) tier_priorities = TIER_PRIORITIES @@ -175,6 +194,12 @@ class GeminiCliProvider( ), } + # retrieveUserQuota can report an exhausted bucket without resetTime when a + # project/account has no entitlement for that model group. Treat that as a + # scoped fallback cooldown instead of repeatedly retrying and logging errors. + default_no_reset_exhaustion_policy = "cooldown" + default_no_reset_exhaustion_cooldown_seconds = 24 * 60 * 60 + # Model quota groups - models that share quota/cooldown timing # Verified 2026-01-07 via quota verification tests # Can be overridden via env: QUOTA_GROUPS_GEMINI_CLI_{GROUP}="model1,model2" diff --git a/src/rotator_library/providers/utilities/gemini_cli_quota_tracker.py b/src/rotator_library/providers/_retired/gemini_cli_quota_tracker.py similarity index 99% rename from src/rotator_library/providers/utilities/gemini_cli_quota_tracker.py rename to src/rotator_library/providers/_retired/gemini_cli_quota_tracker.py index 8b0261e6a..e7a40738b 100644 --- a/src/rotator_library/providers/utilities/gemini_cli_quota_tracker.py +++ b/src/rotator_library/providers/_retired/gemini_cli_quota_tracker.py @@ -35,7 +35,7 @@ import httpx -from .base_quota_tracker import BaseQuotaTracker +from ..utilities.base_quota_tracker import BaseQuotaTracker from .gemini_shared_utils import ( CODE_ASSIST_ENDPOINT, GEMINI_CLI_UA_VERSION, @@ -484,6 +484,10 @@ async def retrieve_user_quota( # Reset time parsing failed; leave reset_timestamp as None pass + exhaustion_reason = None + if is_exhausted and reset_timestamp is None: + exhaustion_reason = "no_reset_time" + buckets_data.append( { "model_id": bucket.get("modelId"), @@ -493,6 +497,7 @@ async def retrieve_user_quota( "reset_timestamp": reset_timestamp, "token_type": bucket.get("tokenType"), "is_exhausted": is_exhausted, + "exhaustion_reason": exhaustion_reason, } ) diff --git a/src/rotator_library/providers/utilities/gemini_credential_manager.py b/src/rotator_library/providers/_retired/gemini_credential_manager.py similarity index 100% rename from src/rotator_library/providers/utilities/gemini_credential_manager.py rename to src/rotator_library/providers/_retired/gemini_credential_manager.py diff --git a/src/rotator_library/providers/utilities/gemini_shared_utils.py b/src/rotator_library/providers/_retired/gemini_shared_utils.py similarity index 100% rename from src/rotator_library/providers/utilities/gemini_shared_utils.py rename to src/rotator_library/providers/_retired/gemini_shared_utils.py diff --git a/src/rotator_library/providers/utilities/gemini_tool_handler.py b/src/rotator_library/providers/_retired/gemini_tool_handler.py similarity index 100% rename from src/rotator_library/providers/utilities/gemini_tool_handler.py rename to src/rotator_library/providers/_retired/gemini_tool_handler.py diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/_retired/google_oauth_base.py similarity index 99% rename from src/rotator_library/providers/google_oauth_base.py rename to src/rotator_library/providers/_retired/google_oauth_base.py index c72a72c91..3ea948db0 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/_retired/google_oauth_base.py @@ -22,10 +22,10 @@ from rich.text import Text from rich.markup import escape as rich_escape -from ..utils.headless_detection import is_headless_environment -from ..utils.reauth_coordinator import get_reauth_coordinator -from ..utils.resilient_io import safe_write_json -from ..error_handler import CredentialNeedsReauthError +from ...utils.headless_detection import is_headless_environment +from ...utils.reauth_coordinator import get_reauth_coordinator +from ...utils.resilient_io import safe_write_json +from ...error_handler import CredentialNeedsReauthError lib_logger = logging.getLogger("rotator_library") diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py new file mode 100644 index 000000000..16dec6ec2 --- /dev/null +++ b/src/rotator_library/providers/antigravity_provider.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Antigravity provider integration restored from safe retired pieces.""" + +from __future__ import annotations + +import os +from typing import Any, List, Optional + +import httpx + +from ..field_cache import FieldCacheInjection, FieldCacheRule +from .provider_interface import ProviderInterface + +BASE_URLS = [ + "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal", + "https://daily-cloudcode-pa.googleapis.com/v1internal", + "https://cloudcode-pa.googleapis.com/v1internal", +] +ANTIGRAVITY_HEADERS = { + "User-Agent": "antigravity/1.15.8 windows/amd64", + "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1", + "Client-Metadata": '{"ideType":"IDE_UNSPECIFIED","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}', +} +AVAILABLE_MODELS = [ + "gemini-2.5-flash", + "gemini-2.5-flash-lite", + "gemini-3-pro-preview", + "gemini-3-flash", + "claude-sonnet-4.5", + "claude-opus-4.5", + "claude-opus-4.6", +] +MODEL_ALIAS_MAP = { + "rev19-uic3-1p": "gemini-2.5-computer-use-preview-10-2025", + "gemini-3-pro-image": "gemini-3-pro-image-preview", + "gemini-3-pro-low": "gemini-3-pro-preview", + "gemini-3-pro-high": "gemini-3-pro-preview", + "claude-sonnet-4-5": "claude-sonnet-4.5", + "claude-opus-4-5": "claude-opus-4.5", + "claude-opus-4-6": "claude-opus-4.6", +} +_AMBIGUOUS_REVERSE_ALIASES = {"gemini-3-pro-preview"} +MODEL_ALIAS_REVERSE = {public: internal for internal, public in MODEL_ALIAS_MAP.items() if public not in _AMBIGUOUS_REVERSE_ALIASES} +EXCLUDED_MODELS = {"chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-2.5-pro"} +FIELD_CACHE_TTL_SECONDS = 3600 +FIELD_CACHE_MAX_VALUES = 1024 +FIELD_CACHE_MAX_BYTES = 4 * 1024 * 1024 + + +class AntigravityProvider(ProviderInterface): + """Safe restored Antigravity integration path. + + The retired provider contained valuable model/header/endpoint knowledge mixed + with fragile device-profile and monolithic transform logic. This active + skeleton restores only stable declarations and helpers so native provider + work can proceed behind tests before any live routing is enabled. + """ + + provider_env_name = "antigravity" + protocol_name = "gemini" + adapter_names: tuple[str, ...] = ("antigravity_envelope",) + field_cache_rules = ( + FieldCacheRule( + name="antigravity_thought_signature", + cache_key="antigravity_thought_signatures", + source="response", + path="candidates.*.content.parts.*.thoughtSignature", + mode="all", + scope=("provider", "model", "credential", "session"), + ttl_seconds=FIELD_CACHE_TTL_SECONDS, + max_values=FIELD_CACHE_MAX_VALUES, + max_bytes=FIELD_CACHE_MAX_BYTES, + inject=FieldCacheInjection(target="request", path="request.metadata.thoughtSignatures", as_list=True), + metadata={"purpose": "preserve Gemini thought signatures across Antigravity turns"}, + ), + FieldCacheRule( + name="antigravity_stream_thought_signature", + cache_key="antigravity_thought_signatures", + source="stream_event", + path="raw.candidates.*.content.parts.*.thoughtSignature", + mode="all", + scope=("provider", "model", "credential", "session"), + ttl_seconds=FIELD_CACHE_TTL_SECONDS, + max_values=FIELD_CACHE_MAX_VALUES, + max_bytes=FIELD_CACHE_MAX_BYTES, + inject=FieldCacheInjection(target="request", path="request.metadata.thoughtSignatures", as_list=True), + metadata={"purpose": "preserve streamed Gemini thought signatures across Antigravity turns"}, + ), + ) + native_streaming_supported = True + model_quota_groups = { + "gemini_3_pro": ["gemini-3-pro-preview", "gemini-3-pro-low", "gemini-3-pro-high"], + "gemini_3_flash": ["gemini-3-flash"], + "gemini_2_5_flash": ["gemini-2.5-flash"], + "gemini_2_5_flash_lite": ["gemini-2.5-flash-lite"], + "claude_sonnet_4_5": ["claude-sonnet-4.5"], + "claude_opus_4_5": ["claude-opus-4.5"], + "claude_opus_4_6": ["claude-opus-4.6"], + } + default_rotation_mode = "sequential" + + async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: + """Fetch available Antigravity models or return the restored safe list.""" + + try: + response = await client.post(self.get_native_endpoint(operation="models"), headers=self.get_native_headers(api_key), json={}, timeout=30) + response.raise_for_status() + models = self._models_from_response(response.json()) + if models: + return [self._with_prefix(model) for model in models] + except Exception: + return [self._with_prefix(model) for model in AVAILABLE_MODELS] + return [self._with_prefix(model) for model in AVAILABLE_MODELS] + + def get_api_base(self) -> str: + """Return the first configured Antigravity base URL.""" + + configured = os.getenv("ANTIGRAVITY_API_BASE") + return configured.rstrip("/") if configured else BASE_URLS[0] + + def get_native_headers(self, credential_identifier: str, model: str = "", operation: str = "generate") -> dict[str, str]: + """Return static Antigravity headers plus bearer auth. + + Retired per-device fingerprint headers are intentionally not restored; + they are brittle and should only return if a current service requirement + is verified with tests. + """ + + headers = { + "Authorization": f"Bearer {credential_identifier}", + "Content-Type": "application/json", + **ANTIGRAVITY_HEADERS, + } + if operation == "stream_generate": + headers["Accept"] = "text/event-stream" + return headers + + def get_native_operation(self, model: str = "", request: dict[str, Any] | None = None, stream: bool = False) -> str: + """Return the Gemini generate operation used by Antigravity endpoints.""" + + return "stream_generate" if stream else "generate" + + def normalize_native_model(self, model: str) -> str: + """Strip the proxy prefix and map public aliases to upstream names.""" + + clean = model.split("/", 1)[1] if model.startswith("antigravity/") else model + return self._alias_to_internal(clean) + + def prepare_native_request(self, request: dict[str, Any], model: str = "", operation: str = "") -> dict[str, Any]: + """Apply Antigravity model and thinking details to Gemini-native input. + + The provider intentionally keeps this to model alias and message-shape + handling only. Device profile and fingerprint behavior stays out of the + active integration until it is verified against current service behavior. + """ + + prepared = dict(request) + public_model = str(request.get("_proxy_model") or request.get("model") or "") + thinking_level = _thinking_level_from_model(public_model) + if model: + prepared["model"] = _model_with_thinking_variant(model, thinking_level) + prepared.pop("_proxy_model", None) + if thinking_level: + generation_config = prepared.setdefault("generationConfig", {}) + generation_config.setdefault("thinkingConfig", {})["thinkingLevel"] = thinking_level + prepared.setdefault("metadata", {})["thinking_level"] = thinking_level + return prepared + + def supports_native_streaming(self, model: str = "", operation: str = "generate") -> bool: + """Return whether the Gemini SSE endpoint is enabled for this model.""" + + if self._get_runtime_config(model).native_streaming_supported is not None: + return super().supports_native_streaming(model, operation) + return operation == "stream_generate" + + def get_native_endpoint(self, model: str = "", operation: str = "generate") -> str: + """Return Antigravity internal operation endpoints.""" + + if operation == "models": + return f"{self.get_api_base()}:fetchAvailableModels" + if operation == "stream_generate": + return f"{self.get_api_base()}:streamGenerateContent?alt=sse" + return f"{self.get_api_base()}:generateContent" + + def get_adapter_config(self, model: str = "") -> dict[str, dict[str, Any]]: + """Configure the safe Antigravity internal request envelope.""" + + config = { + "antigravity_envelope": { + "project": os.getenv("ANTIGRAVITY_PROJECT", ""), + "user_agent": ANTIGRAVITY_HEADERS["User-Agent"], + "request_type": os.getenv("ANTIGRAVITY_REQUEST_TYPE", "CHAT_COMPLETION"), + } + } + for adapter, override in super().get_adapter_config(model).items(): + config.setdefault(adapter, {}).update(override) + return config + + def get_model_tier_requirement(self, model: str) -> Optional[int]: + """Antigravity exposes no restored model-tier restriction.""" + + return None + + def normalize_model_for_tracking(self, model: str) -> str: + """Normalize internal names to public aliases while preserving prefix.""" + + if "/" in model: + provider, clean_model = model.split("/", 1) + return f"{provider}/{self._api_to_user_model(clean_model)}" + return self._api_to_user_model(model) + + @staticmethod + def _with_prefix(model: str) -> str: + if model.startswith("antigravity/"): + return model + return f"antigravity/{model}" + + @staticmethod + def _alias_to_internal(alias: str) -> str: + if alias in {"rev19-uic3-1p", "gemini-3-pro-image", "gemini-3-pro-low", "gemini-3-pro-high"}: + return MODEL_ALIAS_MAP.get(alias, alias) + return MODEL_ALIAS_REVERSE.get(alias, MODEL_ALIAS_MAP.get(alias, alias)) + + @staticmethod + def _api_to_user_model(internal: str) -> str: + return MODEL_ALIAS_MAP.get(internal, internal) + + def _models_from_response(self, payload: dict[str, Any]) -> list[str]: + raw_models: list[str] = [] + if isinstance(payload.get("models"), dict): + raw_models.extend(payload["models"].keys()) + if isinstance(payload.get("data"), list): + raw_models.extend(item.get("id") for item in payload["data"] if isinstance(item, dict) and item.get("id")) + result = [] + for model in raw_models: + public = self._api_to_user_model(str(model)) + if public not in EXCLUDED_MODELS and public in AVAILABLE_MODELS and public not in result: + result.append(public) + return result + + +def _thinking_level_from_model(model: str) -> Optional[str]: + clean = model.split("/", 1)[1] if model.startswith("antigravity/") else model + if clean.endswith("-low"): + return "low" + if clean.endswith("-high"): + return "high" + return None + + +def _model_with_thinking_variant(model: str, thinking_level: Optional[str]) -> str: + if model == "gemini-3-pro-preview" and thinking_level in {"low", "high"}: + return f"gemini-3-pro-{thinking_level}" + return model diff --git a/src/rotator_library/providers/claude_code_provider.py b/src/rotator_library/providers/claude_code_provider.py new file mode 100644 index 000000000..64d128e8a --- /dev/null +++ b/src/rotator_library/providers/claude_code_provider.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Claude Code provider integration for native Anthropic Messages execution.""" + +from __future__ import annotations + +import os +from typing import Any, List + +import httpx + +from ..field_cache import FieldCacheInjection, FieldCacheRule +from .provider_interface import ProviderInterface + +DEFAULT_API_BASE = "https://api.anthropic.com" +FALLBACK_MODELS = ["claude_code/claude-sonnet-4-5", "claude_code/claude-opus-4-5"] +FIELD_CACHE_TTL_SECONDS = 3600 +FIELD_CACHE_MAX_VALUES = 1024 +FIELD_CACHE_MAX_BYTES = 4 * 1024 * 1024 + + +class ClaudeCodeProvider(ProviderInterface): + """Provider declaration for Claude Code style native requests. + + The provider starts as an explicit integration path rather than a guessed live + implementation. It declares protocol/adapters/cache rules and exposes mocked + auth/model helpers so later native wiring can use it without reintroducing a + monolithic provider transform. + """ + + provider_env_name = "claude_code" + protocol_name = "anthropic_messages" + adapter_names = ("suppress_developer_role",) + native_streaming_supported = True + field_cache_rules = ( + FieldCacheRule( + name="claude_code_thinking_signature", + cache_key="claude_code_thinking_signatures", + source="response", + path="content.*.signature", + mode="all", + scope=("provider", "model", "credential", "session"), + ttl_seconds=FIELD_CACHE_TTL_SECONDS, + max_values=FIELD_CACHE_MAX_VALUES, + max_bytes=FIELD_CACHE_MAX_BYTES, + inject=FieldCacheInjection(target="request", path="metadata.thinking_signatures", as_list=True), + metadata={"purpose": "preserve Claude thinking signatures for follow-up requests"}, + ), + FieldCacheRule( + name="claude_code_stream_thinking_signature", + cache_key="claude_code_thinking_signatures", + source="stream_event", + path="raw.delta.signature", + mode="all", + scope=("provider", "model", "credential", "session"), + ttl_seconds=FIELD_CACHE_TTL_SECONDS, + max_values=FIELD_CACHE_MAX_VALUES, + max_bytes=FIELD_CACHE_MAX_BYTES, + inject=FieldCacheInjection(target="request", path="metadata.thinking_signatures", as_list=True), + metadata={"purpose": "preserve streamed Claude thinking signatures for follow-up requests"}, + ), + ) + default_rotation_mode = "sequential" + + async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: + """Fetch provider models with a conservative fallback list. + + Model discovery is intentionally mock-friendly. If the configured service + does not expose a standard `/v1/models` response, provider work can later + override this without changing protocol declarations. + """ + + try: + response = await client.get(f"{self.get_api_base().rstrip('/')}/v1/models", headers=self.get_native_headers(api_key), timeout=30) + response.raise_for_status() + models = [item.get("id") for item in response.json().get("data", []) if isinstance(item, dict) and item.get("id")] + if models: + return [self._with_prefix(model) for model in models] + except Exception: + return list(FALLBACK_MODELS) + return list(FALLBACK_MODELS) + + def get_api_base(self) -> str: + """Return the configured Claude Code API base URL.""" + + return os.getenv("CLAUDE_CODE_API_BASE", DEFAULT_API_BASE) + + def get_native_headers(self, credential_identifier: str, model: str = "", operation: str = "messages") -> dict[str, str]: + """Return headers for native mocked HTTP requests.""" + + mode = os.getenv("CLAUDE_CODE_AUTH_HEADER", "auto").strip().lower() + use_api_key = mode == "x-api-key" or (mode == "auto" and credential_identifier.startswith("sk-ant-")) + headers = { + "anthropic-version": os.getenv("CLAUDE_CODE_ANTHROPIC_VERSION", "2023-06-01"), + "content-type": "application/json", + } + if use_api_key: + headers["x-api-key"] = credential_identifier + else: + headers["Authorization"] = f"Bearer {credential_identifier}" + return headers + + def get_native_operation(self, model: str = "", request: dict[str, Any] | None = None, stream: bool = False) -> str: + """Claude Code uses the Anthropic Messages operation for completions.""" + + return "messages" + + def normalize_native_model(self, model: str) -> str: + """Strip the proxy provider prefix before sending upstream.""" + + return model.split("/", 1)[1] if model.startswith("claude_code/") else model + + def supports_native_streaming(self, model: str = "", operation: str = "messages") -> bool: + """Return whether Anthropic Messages SSE is enabled for this model.""" + + if self._get_runtime_config(model).native_streaming_supported is not None: + return super().supports_native_streaming(model, operation) + return operation == "messages" + + def prepare_native_request(self, request: dict[str, Any], model: str = "", operation: str = "messages") -> dict[str, Any]: + """Ensure Anthropic Messages-required fields are present.""" + + prepared = dict(request) + prepared.setdefault("max_tokens", int(os.getenv("CLAUDE_CODE_MAX_TOKENS", "4096"))) + return prepared + + def get_native_endpoint(self, model: str = "", operation: str = "messages") -> str: + """Return the provider endpoint for a native operation.""" + + if operation == "models": + return f"{self.get_api_base().rstrip('/')}/v1/models" + return f"{self.get_api_base().rstrip('/')}/v1/messages" + + def get_adapter_config(self, model: str = "") -> dict[str, dict[str, Any]]: + """Configure adapters without hardcoding provider transforms.""" + + config = {"suppress_developer_role": {"mode": "user"}} + for adapter, override in super().get_adapter_config(model).items(): + config.setdefault(adapter, {}).update(override) + return config + + @staticmethod + def _with_prefix(model: str) -> str: + if model.startswith("claude_code/"): + return model + return f"claude_code/{model}" diff --git a/src/rotator_library/providers/codex_provider.py b/src/rotator_library/providers/codex_provider.py new file mode 100644 index 000000000..70b2242f5 --- /dev/null +++ b/src/rotator_library/providers/codex_provider.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Codex provider integration for native Responses execution.""" + +from __future__ import annotations + +import os +from typing import Any, List + +import httpx + +from ..field_cache import FieldCacheInjection, FieldCacheRule +from .provider_interface import ProviderInterface + +DEFAULT_API_BASE = "https://api.openai.com" +FALLBACK_MODELS = ["codex/codex-mini-latest", "codex/gpt-5.1-codex"] +FIELD_CACHE_TTL_SECONDS = 3600 + + +class CodexProvider(ProviderInterface): + """Provider declaration for Codex-style native Responses requests.""" + + provider_env_name = "codex" + protocol_name = "responses" + adapter_names: tuple[str, ...] = () + native_streaming_supported = True + field_cache_rules = ( + FieldCacheRule( + name="codex_previous_response_id", + cache_key="codex_previous_response_id", + source="response", + path="id", + scope=("provider", "model", "credential", "session"), + ttl_seconds=FIELD_CACHE_TTL_SECONDS, + inject=FieldCacheInjection(target="request", path="previous_response_id", when_missing_only=True), + metadata={ + "purpose": "preserve Responses continuation IDs for Codex sessions", + "provider_continuation": True, + }, + ), + FieldCacheRule( + name="codex_stream_previous_response_id", + cache_key="codex_previous_response_id", + source="stream_event", + path="raw.response.id", + scope=("provider", "model", "credential", "session"), + ttl_seconds=FIELD_CACHE_TTL_SECONDS, + inject=FieldCacheInjection(target="request", path="previous_response_id", when_missing_only=True), + metadata={ + "purpose": "preserve streamed Responses continuation IDs for Codex sessions", + "provider_continuation": True, + }, + ), + ) + default_rotation_mode = "sequential" + + async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: + """Fetch Codex-visible models with fallback names for offline tests.""" + + try: + response = await client.get(self.get_native_endpoint(operation="models"), headers=self.get_native_headers(api_key), timeout=30) + response.raise_for_status() + models = [item.get("id") for item in response.json().get("data", []) if isinstance(item, dict) and item.get("id")] + codex_models = [model for model in models if "codex" in model.lower()] + if codex_models: + return [self._with_prefix(model) for model in codex_models] + except Exception: + return list(FALLBACK_MODELS) + return list(FALLBACK_MODELS) + + def get_api_base(self) -> str: + """Return the configured Codex API base URL.""" + + return os.getenv("CODEX_API_BASE", DEFAULT_API_BASE) + + def get_native_headers(self, credential_identifier: str, model: str = "", operation: str = "responses") -> dict[str, str]: + """Return headers for Codex native HTTP calls.""" + + return {"Authorization": f"Bearer {credential_identifier}", "content-type": "application/json"} + + def get_native_operation(self, model: str = "", request: dict | None = None, stream: bool = False) -> str: + """Codex native calls use the Responses operation.""" + + return "responses" + + def normalize_native_model(self, model: str) -> str: + """Strip the proxy provider prefix before sending upstream.""" + + return model.split("/", 1)[1] if model.startswith("codex/") else model + + def supports_native_streaming(self, model: str = "", operation: str = "responses") -> bool: + """Return whether Responses SSE is enabled for this model.""" + + if self._get_runtime_config(model).native_streaming_supported is not None: + return super().supports_native_streaming(model, operation) + return operation == "responses" + + def prepare_native_request(self, request: dict[str, Any], model: str = "", operation: str = "") -> dict[str, Any]: + """Apply the normalized upstream model to a Responses-native payload.""" + + prepared = dict(request) + if model: + prepared["model"] = model + return prepared + + def get_native_endpoint(self, model: str = "", operation: str = "responses") -> str: + """Return the native Codex endpoint for an operation.""" + + suffix = "/v1/models" if operation == "models" else "/v1/responses" + return f"{self.get_api_base().rstrip('/')}{suffix}" + + @staticmethod + def _with_prefix(model: str) -> str: + if model.startswith("codex/"): + return model + return f"codex/{model}" diff --git a/src/rotator_library/providers/copilot_provider.py b/src/rotator_library/providers/copilot_provider.py new file mode 100644 index 000000000..57dec62a2 --- /dev/null +++ b/src/rotator_library/providers/copilot_provider.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Copilot provider integration for native OpenAI Chat execution.""" + +from __future__ import annotations + +import os +from typing import List + +import httpx + +from .provider_interface import ProviderInterface + +DEFAULT_API_BASE = "https://api.githubcopilot.com" +FALLBACK_MODELS = ["copilot/gpt-4.1", "copilot/claude-sonnet-4-5"] + + +class CopilotProvider(ProviderInterface): + """Provider declaration for Copilot-style OpenAI-compatible chat calls. + + The skeleton intentionally avoids inventing field-cache rules until a stable + provider session/conversation field is identified. This keeps Copilot native + support explicit and testable without guessing hidden behavior. + """ + + provider_env_name = "copilot" + protocol_name = "openai_chat" + adapter_names = ("suppress_developer_role",) + field_cache_rules: tuple = () + default_rotation_mode = "sequential" + native_streaming_supported = True + + async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: + """Fetch Copilot-visible models with a safe fallback list.""" + + try: + response = await client.get(self.get_native_endpoint(operation="models"), headers=self.get_native_headers(api_key), timeout=30) + response.raise_for_status() + models = [item.get("id") for item in response.json().get("data", []) if isinstance(item, dict) and item.get("id")] + if models: + return [self._with_prefix(model) for model in models] + except Exception: + return list(FALLBACK_MODELS) + return list(FALLBACK_MODELS) + + def get_api_base(self) -> str: + """Return the configured Copilot API base URL.""" + + return os.getenv("COPILOT_API_BASE", DEFAULT_API_BASE) + + def get_native_headers(self, credential_identifier: str, model: str = "", operation: str = "chat") -> dict[str, str]: + """Return headers for Copilot native HTTP calls.""" + + return { + "Authorization": f"Bearer {credential_identifier}", + "content-type": "application/json", + "Copilot-Integration-Id": os.getenv("COPILOT_INTEGRATION_ID", "llm-api-key-proxy"), + } + + def get_native_operation(self, model: str = "", request: dict | None = None, stream: bool = False) -> str: + """Copilot exposes an OpenAI-compatible chat operation.""" + + return "chat" + + def normalize_native_model(self, model: str) -> str: + """Strip the proxy provider prefix before sending upstream.""" + + return model.split("/", 1)[1] if model.startswith("copilot/") else model + + def supports_native_streaming(self, model: str = "", operation: str = "chat") -> bool: + """Return whether Chat Completions SSE is enabled for this model.""" + + if self._get_runtime_config(model).native_streaming_supported is not None: + return super().supports_native_streaming(model, operation) + return operation == "chat" + + def get_native_endpoint(self, model: str = "", operation: str = "chat") -> str: + """Return the Copilot endpoint for a native operation.""" + + suffix = "/models" if operation == "models" else "/chat/completions" + return f"{self.get_api_base().rstrip('/')}{suffix}" + + def get_adapter_config(self, model: str = "") -> dict[str, dict[str, str]]: + """Configure role suppression declaratively for OpenAI-compatible chat.""" + + config = {"suppress_developer_role": {"mode": "system"}} + for adapter, override in super().get_adapter_config(model).items(): + config.setdefault(adapter, {}).update(override) + return config + + @staticmethod + def _with_prefix(model: str) -> str: + if model.startswith("copilot/"): + return model + return f"copilot/{model}" diff --git a/src/rotator_library/providers/nvidia_provider.py b/src/rotator_library/providers/nvidia_provider.py index 25f4ef464..649db729c 100644 --- a/src/rotator_library/providers/nvidia_provider.py +++ b/src/rotator_library/providers/nvidia_provider.py @@ -4,7 +4,6 @@ import httpx import logging from typing import List, Dict, Any -import litellm from .provider_interface import ProviderInterface lib_logger = logging.getLogger("rotator_library") @@ -54,6 +53,9 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str] KIMI_K2_MODEL_PATTERNS = [ "kimi-k2.", ] + DIFFUSION_GEMMA_MODEL_EXACT = [ + "google/diffusiongemma-26b-a4b-it", + ] V4_EFFORT_MAP = { "low": "high", @@ -61,7 +63,7 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str] "high": "max", "max": "max", } - DISABLE_VALUES = {"none", "disable", "off"} + DISABLE_VALUES = {"none", "disable", "disabled", "off"} def _is_v3_deepseek(self, model_name: str) -> bool: if model_name in self.V3_MODEL_EXACT: @@ -77,14 +79,18 @@ def _is_mistral_reasoning(self, model_name: str) -> bool: def _is_kimi_k2(self, model_name: str) -> bool: return any(p in model_name for p in self.KIMI_K2_MODEL_PATTERNS) + def _is_diffusion_gemma(self, model_name: str) -> bool: + return model_name in self.DIFFUSION_GEMMA_MODEL_EXACT + def handle_thinking_parameter(self, payload: Dict[str, Any], model: str): """ - Configures thinking and reasoning_effort for DeepSeek and Mistral models on NVIDIA. + Configures thinking and reasoning_effort for supported NVIDIA models. DeepSeek V3.x: only thinking=True/False in chat_template_kwargs. DeepSeek V4: thinking + mapped reasoning_effort (high/max). Mistral: reasoning_effort="high" via extra_body (LiteLLM drops unsupported top-level params). - Incoming reasoning_effort of none/disable/off disables thinking/effort. + DiffusionGemma: enable_thinking=True/False in top-level chat_template_kwargs. + Incoming reasoning_effort of none/disable/disabled/off disables thinking/effort. """ model_name = model.split("/", 1)[1] if "/" in model else model @@ -92,8 +98,15 @@ def handle_thinking_parameter(self, payload: Dict[str, Any], model: str): is_v4 = self._is_v4_deepseek(model_name) is_mistral = self._is_mistral_reasoning(model_name) is_kimi = self._is_kimi_k2(model_name) - - if not is_v3 and not is_v4 and not is_mistral and not is_kimi: + is_diffusion_gemma = self._is_diffusion_gemma(model_name) + + if ( + not is_v3 + and not is_v4 + and not is_mistral + and not is_kimi + and not is_diffusion_gemma + ): return reasoning_effort = payload.get("reasoning_effort") @@ -103,6 +116,19 @@ def handle_thinking_parameter(self, payload: Dict[str, Any], model: str): and reasoning_effort.lower() in self.DISABLE_VALUES ) + if is_diffusion_gemma: + payload.pop("reasoning_effort", None) + template_kwargs = payload.get("chat_template_kwargs") + if not isinstance(template_kwargs, dict): + template_kwargs = {} + payload["chat_template_kwargs"] = template_kwargs + template_kwargs["enable_thinking"] = not is_disabled + lib_logger.info( + f"NVIDIA: DiffusionGemma '{model_name}' — " + f"enable_thinking={not is_disabled}" + ) + return + if is_kimi: payload.pop("reasoning_effort", None) thinking = not is_disabled @@ -135,6 +161,7 @@ def handle_thinking_parameter(self, payload: Dict[str, Any], model: str): kwargs = payload["extra_body"]["chat_template_kwargs"] if is_disabled: + payload.pop("reasoning_effort", None) kwargs["thinking"] = False lib_logger.info( f"NVIDIA: DeepSeek '{model_name}' — thinking DISABLED " diff --git a/src/rotator_library/providers/provider_cache.py b/src/rotator_library/providers/provider_cache.py index a1b539675..8627990cd 100644 --- a/src/rotator_library/providers/provider_cache.py +++ b/src/rotator_library/providers/provider_cache.py @@ -25,7 +25,7 @@ import os import time from pathlib import Path -from typing import Any, Dict, Optional, Tuple +from typing import Any, Callable, Dict, Optional, Tuple from ..utils.resilient_io import safe_write_json @@ -107,6 +107,8 @@ def __init__( else _env_bool(f"{env_prefix}_ENABLE", True) ) self._dirty = False + self._deleted_keys: set[str] = set() + self._clear_disk_on_save = False self._write_interval = write_interval or _env_int( f"{env_prefix}_WRITE_INTERVAL", 60 ) @@ -232,11 +234,18 @@ async def _save_to_disk(self) -> bool: # Step 2: Filter existing disk entries by disk_ttl (not memory_ttl) # This preserves entries that expired from memory but are still valid on disk - valid_disk_entries = { - k: v - for k, v in existing_entries.items() - if now - v.get("timestamp", 0) <= self._disk_ttl - } + deleted_keys = set(self._deleted_keys) + clear_disk = self._clear_disk_on_save + valid_disk_entries = ( + {} + if clear_disk + else { + k: v + for k, v in existing_entries.items() + if k not in deleted_keys + and now - v.get("timestamp", 0) <= self._disk_ttl + } + ) # Step 3: Merge - memory entries take precedence (fresher timestamps) merged_entries = valid_disk_entries.copy() @@ -269,6 +278,9 @@ async def _save_to_disk(self) -> bool: ): self._stats["writes"] += 1 self._disk_available = True + self._deleted_keys.difference_update(deleted_keys) + if clear_disk: + self._clear_disk_on_save = False # Log merge info only when we preserved disk-only entries (infrequent) if preserved_from_disk > 0: lib_logger.debug( @@ -360,6 +372,7 @@ async def _async_store(self, key: str, value: str) -> None: """Async implementation of store.""" async with self._lock: self._cache[key] = (value, time.time()) + self._deleted_keys.discard(key) self._dirty = True async def store_async(self, key: str, value: str) -> None: @@ -370,6 +383,23 @@ async def store_async(self, key: str, value: str) -> None: """ await self._async_store(key, value) + async def update_async( + self, + key: str, + update: Callable[[Optional[str]], str], + ) -> str: + """Atomically update one value within this cache backend instance.""" + + await self.retrieve_async(key) + async with self._lock: + current = self._cache.get(key) + raw = current[0] if current is not None else None + updated = update(raw) + self._cache[key] = (updated, time.time()) + self._deleted_keys.discard(key) + self._dirty = True + return updated + def retrieve(self, key: str) -> Optional[str]: """ Retrieve a value by key (synchronous, with optional async disk fallback). @@ -422,6 +452,20 @@ async def retrieve_async(self, key: str) -> Optional[str]: self._stats["misses"] += 1 return None + async def delete_async(self, key: str) -> bool: + """Delete one key from memory and durable storage.""" + + existed = key in self._cache + if not existed and self._enable_disk: + existed = await self.retrieve_async(key) is not None + async with self._lock: + self._cache.pop(key, None) + self._deleted_keys.add(key) + self._dirty = True + if self._enable_disk: + await self._save_to_disk() + return existed + async def _check_disk_fallback(self, key: str) -> None: """Check disk for key and load into memory if found (background).""" try: @@ -507,6 +551,8 @@ async def clear(self) -> None: """Clear all cached data.""" async with self._lock: self._cache.clear() + self._deleted_keys.clear() + self._clear_disk_on_save = True self._dirty = True if self._enable_disk: await self._save_to_disk() diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index fb5811bf4..45433d8e9 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -135,7 +135,7 @@ class ProviderInterface(ABC, metaclass=SingletonABCMeta): # TIER CONFIGURATION - Override in subclass # ========================================================================= - # Provider name for env var lookups (e.g., "gemini_cli") + # Provider name for env var lookups (e.g., "openai") # Used for: QUOTA_GROUPS_{provider_env_name}_{GROUP} provider_env_name: str = "" @@ -270,6 +270,15 @@ class ProviderInterface(ABC, metaclass=SingletonABCMeta): Union[int, Tuple[int, ...], str], Dict[str, Dict[str, Any]] ] = {} + # Native protocol/adapter declarations introduced for the experimental + # protocol stack. Defaults are intentionally no-op so existing providers keep + # the LiteLLM-backed execution path until they opt into native protocols. + protocol_name: Optional[str] = None + default_output_protocol: Optional[str] = None + adapter_names: Tuple[str, ...] = () + field_cache_rules: Tuple[Any, ...] = () + native_streaming_supported: bool = False + @abstractmethod async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: """ @@ -301,10 +310,12 @@ def get_session_tracking_hints( """Return provider-specific evidence for core session tracking. Providers can expose stable native markers, custom request structure, or - a preferred affinity key before credential selection. The return value is - intentionally evidence-only: providers must not select credentials or - mutate sticky state. Core routing merges these hints with generic anchors - and applies the same confidence policy for every provider. + a preferred affinity key before credential selection. Native evidence is + provider-qualified and may use ``SessionTrackingHints.session_scope`` for + model/family partitioning without fragmenting the global logical session. + The return value is intentionally evidence-only: providers must not select + credentials or mutate sticky state. Core routing merges these hints with + generic anchors and applies the same confidence policy for every provider. Expansion path: providers that know their upstream cache/session scope can return a richer ``SessionTrackingHints`` object from @@ -313,6 +324,180 @@ def get_session_tracking_hints( """ return None + def get_protocol_name(self, model: str = "") -> Optional[str]: + """Return the native protocol adapter name this provider prefers. + + Providers may override this method when protocol choice varies by model. + Returning ``None`` keeps the current fallback execution behavior. Later + provider phases will use this as the bridge from provider declarations to + native protocol parsing/building. + """ + + configured = self._get_runtime_config(model).protocol_name + return configured or self.protocol_name + + def get_default_output_protocol(self, model: str = "") -> Optional[str]: + """Return the integration output protocol selected for this provider.""" + + configured = self._get_runtime_config(model).default_output_protocol + return configured or self.default_output_protocol + + def get_adapter_names(self, model: str = "") -> Tuple[str, ...]: + """Return ordered adapter names for this provider/model. + + The order is significant and is preserved by the adapter chain runner. + Providers can override this for model-specific quirks without mutating the + global adapter registry. + """ + + configured = self._get_runtime_config(model).adapter_names + return tuple(configured) if configured is not None else tuple(self.adapter_names) + + def get_adapter_config(self, model: str = "") -> Dict[str, Dict[str, Any]]: + """Return adapter-specific config keyed by adapter name. + + Config is intentionally a plain dict so custom providers can define it in + env/JSON later without importing adapter classes. Phase 10 will add formal + config loading and validation. + """ + + return dict(self._get_runtime_config(model).adapter_config) + + def get_field_cache_rules(self, model: str = "") -> Tuple[Any, ...]: + """Return field-cache rules for provider-specific protocol state. + + Rules preserve provider state such as reasoning content, thought + signatures, prompt cache keys, provider session IDs, and response IDs. + They are not a replacement for ``SessionTracker``; session tracking still + decides continuity and credential affinity. + """ + + configured = self._get_runtime_config(model).field_cache_rules + return tuple(self.field_cache_rules) + tuple(configured) + + def supports_native_streaming(self, model: str = "", operation: str = "chat") -> bool: + """Return whether this provider explicitly supports native streaming. + + The default is intentionally false. Phase 8 keeps live streaming + conservative: providers must opt in before routed streaming can use the + native stream executor instead of current custom/LiteLLM behavior. + """ + + configured = self._get_runtime_config(model).native_streaming_supported + if configured is None: + return self.native_streaming_supported + return bool(configured and self.supports_native_operation(model, operation)) + + def _get_runtime_config(self, model: str = "") -> Any: + """Return optional JSON runtime config for this provider. + + The helper keeps config loading lazy so provider imports do not depend on + the experimental config layer during startup discovery. + """ + + from ..config.experimental import get_provider_runtime_config + + return get_provider_runtime_config( + self._provider_config_key(), + model, + config=getattr(self, "_runtime_config_snapshot", None), + ) + + def bind_runtime_config(self, config: Any) -> None: + """Bind the immutable process-start configuration used by this instance.""" + + existing = getattr(self, "_runtime_config_snapshot", None) + if existing is not None and existing != config: + raise RuntimeError( + f"Provider singleton {self._provider_config_key()!r} is already bound to a different startup configuration" + ) + self._runtime_config_snapshot = config + + def _provider_config_key(self) -> str: + """Return the JSON providers-section key for this provider.""" + + if self.provider_env_name: + return self.provider_env_name.lower() + name = self.__class__.__name__ + if name.endswith("Provider"): + name = name[: -len("Provider")] + return name.lower() + + def supports_native_operation(self, model: str = "", operation: str = "chat") -> bool: + """Return whether this provider supports a native operation.""" + + protocol_name = self.get_protocol_name(model) + if not protocol_name: + return False + try: + from ..protocols import get_protocol + + return get_protocol(protocol_name).supports_operation(operation) + except Exception: + return False + + def should_use_native_protocol(self, model: str = "", operation: str = "chat", *, stream: bool = False, execution: str = "auto") -> bool: + """Return whether routing should use this provider's native protocol.""" + + if stream and not self.supports_native_streaming(model, operation): + return False + return bool(self.get_protocol_name(model) and self.supports_native_operation(model, operation)) + + def get_native_operation(self, model: str = "", request: Optional[Dict[str, Any]] = None, stream: bool = False) -> str: + """Return the provider-native operation for a request. + + Providers that expose native protocols often use operation names that are + not simply ``chat``: Anthropic-compatible providers use ``messages``, + Responses providers use ``responses``, and Gemini-style providers use a + generate operation. The default stays ``chat`` so existing OpenAI-chat + providers remain compatible unless they opt in to something richer. + """ + + return "chat" + + def get_native_endpoint(self, model: str = "", operation: str = "chat") -> str: + """Return the upstream endpoint for a native operation.""" + + raise NotImplementedError(f"{self.__class__.__name__} does not define a native endpoint") + + def get_native_headers(self, credential_identifier: str, model: str = "", operation: str = "chat") -> Dict[str, str]: + """Return non-payload HTTP headers for native requests.""" + + raise NotImplementedError(f"{self.__class__.__name__} does not define native headers") + + def normalize_native_model(self, model: str) -> str: + """Return the upstream model name for native provider calls. + + The proxy-facing model commonly includes a provider prefix such as + ``provider/model``. Native upstream APIs usually expect only ``model``. + Providers may override this for aliases, but stripping the first prefix + is the safe default for native execution. + """ + + return model.split("/", 1)[1] if "/" in model else model + + def prepare_native_request(self, request: Dict[str, Any], model: str = "", operation: str = "") -> Dict[str, Any]: + """Return a provider-adjusted native request payload. + + The declared provider protocol has already built a valid native payload + before this hook runs. Implementations may add provider envelopes, + aliases, or required defaults, but must never translate a client protocol. + Credentials remain in ``get_native_headers()`` so payload traces never + mix request data with secrets. + """ + + return dict(request) + + def get_model_pricing(self, model: str = "") -> Optional[Any]: + """Return optional local pricing metadata for advisory cost tracking. + + Providers can return `usage.costs.ModelPricing` or a compatible dict. + The default is `None`, which lets cost accounting safely fall back to + LiteLLM model metadata or report pricing as unavailable. + """ + + return None + async def acompletion( self, client: httpx.AsyncClient, **kwargs ) -> Union[ @@ -411,9 +596,8 @@ def get_model_tier_requirement(self, model: str) -> Optional[int]: Minimum required priority level or None if no restrictions Example: - For Gemini CLI: - - gemini-3-*: requires priority 1 (paid tier only) - - gemini-2.5-*: no restriction (None) + A provider may restrict high-capability models to priority 1 + credentials while allowing lower-cost models for all priorities. """ return None @@ -459,7 +643,7 @@ def get_rotation_mode(cls, provider_name: str) -> str: then falls back to the class's default_rotation_mode. Args: - provider_name: The provider name (e.g., "gemini_cli") + provider_name: The provider name (e.g., "openai") Returns: "balanced" or "sequential" @@ -623,12 +807,16 @@ def _get_effective_quota_groups(self) -> QuotaGroupMap: Env format: QUOTA_GROUPS_{PROVIDER}_{GROUP}="model1,model2" Set empty string to disable a default group. """ - if not self.provider_env_name or not self.model_quota_groups: - return self.model_quota_groups + configured_groups = self._get_runtime_config().model_quota_groups or {} + base_groups: QuotaGroupMap = {group: list(models) for group, models in self.model_quota_groups.items()} + for group, models in configured_groups.items(): + base_groups[group] = list(models) + if not self.provider_env_name: + return base_groups result: QuotaGroupMap = {} - for group_name, default_models in self.model_quota_groups.items(): + for group_name, default_models in base_groups.items(): env_key = ( f"QUOTA_GROUPS_{self.provider_env_name.upper()}_{group_name.upper()}" ) diff --git a/src/rotator_library/providers/utilities/__init__.py b/src/rotator_library/providers/utilities/__init__.py index 6f30a7ceb..1cb57a358 100644 --- a/src/rotator_library/providers/utilities/__init__.py +++ b/src/rotator_library/providers/utilities/__init__.py @@ -3,53 +3,13 @@ # Utilities for provider implementations from .base_quota_tracker import BaseQuotaTracker -from .gemini_cli_quota_tracker import GeminiCliQuotaTracker - -# Shared utilities for Gemini-based providers -from .gemini_shared_utils import ( - env_bool, - env_int, - inline_schema_refs, - normalize_type_arrays, - clean_gemini_schema, - recursively_parse_json_strings, - GEMINI3_TOOL_RENAMES, - GEMINI3_TOOL_RENAMES_REVERSE, - FINISH_REASON_MAP, - DEFAULT_SAFETY_SETTINGS, -) -from .gemini_tool_handler import GeminiToolHandler -from .gemini_credential_manager import GeminiCredentialManager # Re-export loggers from transaction_logger for backward compatibility from ...transaction_logger import ProviderLogger -# Deprecated aliases for backward compatibility with external consumers -# These map old class names to their new equivalents -GeminiFileLogger = ProviderLogger -GeminiCliFileLogger = ProviderLogger - __all__ = [ # Quota trackers "BaseQuotaTracker", - "GeminiCliQuotaTracker", - # Shared utilities - "env_bool", - "env_int", - "inline_schema_refs", - "normalize_type_arrays", - "clean_gemini_schema", - "recursively_parse_json_strings", - "GEMINI3_TOOL_RENAMES", - "GEMINI3_TOOL_RENAMES_REVERSE", - "FINISH_REASON_MAP", - "DEFAULT_SAFETY_SETTINGS", # Loggers (from transaction_logger) "ProviderLogger", - # Deprecated logger aliases (for backward compatibility) - "GeminiFileLogger", - "GeminiCliFileLogger", - # Mixins - "GeminiToolHandler", - "GeminiCredentialManager", ] diff --git a/src/rotator_library/providers/utilities/base_quota_tracker.py b/src/rotator_library/providers/utilities/base_quota_tracker.py index 3bda8979e..0545438e4 100644 --- a/src/rotator_library/providers/utilities/base_quota_tracker.py +++ b/src/rotator_library/providers/utilities/base_quota_tracker.py @@ -5,7 +5,7 @@ Base Quota Tracking Mixin Provides shared quota tracking infrastructure for providers that use OAuth -credentials with quota-based rate limiting (e.g., Gemini CLI). +credentials with quota-based rate limiting. This base class handles: - Learned costs management (load/save/lookup) @@ -17,7 +17,7 @@ Subclasses must implement: - _fetch_quota_for_credential() - Provider-specific quota API call - _extract_model_quota_from_response() - Parse provider-specific response format -- _get_provider_prefix() - Return provider prefix for model names (e.g., "gemini_cli") +- _get_provider_prefix() - Return provider prefix for model names (e.g., "nanogpt") Required from provider (via mixin inheritance): - self.project_id_cache: Dict[str, str] @@ -66,8 +66,8 @@ class BaseQuotaTracker: Base mixin class providing shared quota tracking functionality. Subclasses provide: - - provider_env_prefix: str (e.g., "GEMINI_CLI") - - cache_subdir: str (e.g., "gemini_cli") + - provider_env_prefix: str (e.g., "NANOGPT") + - cache_subdir: str (e.g., "nanogpt") - default_quota_costs: Dict[str, Dict[str, float]] - tier -> model -> cost% - default_quota_cost_unknown: float - fallback cost for unknown models - user_to_api_model_map: Dict[str, str] - optional model name mappings @@ -84,7 +84,7 @@ class BaseQuotaTracker: # ========================================================================= # Environment variable prefix for credential discovery - # e.g., "GEMINI_CLI" looks for GEMINI_CLI_1_ACCESS_TOKEN, etc. + # e.g., "NANOGPT" looks for NANOGPT_1_ACCESS_TOKEN, etc. provider_env_prefix: str = "" # Cache subdirectory name for learned costs file @@ -166,7 +166,7 @@ def _get_provider_prefix(self) -> str: Get the provider prefix for model names. Returns: - Provider prefix (e.g., "gemini_cli") + Provider prefix (e.g., "nanogpt") """ pass @@ -378,7 +378,7 @@ def discover_all_credentials( # 2. Env-based credentials # Check for {PREFIX}_1_ACCESS_TOKEN, {PREFIX}_2_ACCESS_TOKEN, etc. env_prefix = self.provider_env_prefix - provider_name = self.cache_subdir # e.g., "gemini_cli" + provider_name = self.cache_subdir # e.g., "nanogpt" for i in range(1, ENV_CREDENTIAL_DISCOVERY_LIMIT): # Upper limit if os.getenv(f"{env_prefix}_{i}_ACCESS_TOKEN"): @@ -559,6 +559,7 @@ async def _store_baselines_to_usage_manager( bucket = self._find_bucket_for_model(quota_data, user_model) reset_timestamp = bucket.get("reset_timestamp") if bucket else None valid_reset_ts = reset_timestamp if remaining < 1.0 else None + exhaustion_reason = bucket.get("exhaustion_reason") if bucket else None # DEFAULT: Always apply exhaustion if remaining == 0.0 exactly # (API is authoritative for most providers) @@ -573,6 +574,7 @@ async def _store_baselines_to_usage_manager( quota_group=quota_group, force=force, apply_exhaustion=apply_exhaustion, + exhaustion_reason=exhaustion_reason, ) stored_count += 1 diff --git a/src/rotator_library/responses/__init__.py b/src/rotator_library/responses/__init__.py new file mode 100644 index 000000000..0ba921406 --- /dev/null +++ b/src/rotator_library/responses/__init__.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Responses API service, storage, and streaming helpers.""" + +from .bridge import ResponsesBridge +from .service import ResponsesService, ResponsesServiceError +from .store import InMemoryResponsesStore, ProviderCacheResponsesStore, ResponsesStore, create_configured_responses_store +from .streaming import ResponsesSSEFormatter, ResponsesStreamEvent, ResponsesWebSocketFormatter +from .types import ResponsesStoreSettings, StoredResponse, generate_response_id + +__all__ = [ + "InMemoryResponsesStore", + "ProviderCacheResponsesStore", + "ResponsesBridge", + "ResponsesService", + "ResponsesServiceError", + "ResponsesStoreSettings", + "ResponsesSSEFormatter", + "ResponsesStreamEvent", + "ResponsesStore", + "ResponsesWebSocketFormatter", + "StoredResponse", + "create_configured_responses_store", + "generate_response_id", +] diff --git a/src/rotator_library/responses/bridge.py b/src/rotator_library/responses/bridge.py new file mode 100644 index 000000000..8c5aa2c8f --- /dev/null +++ b/src/rotator_library/responses/bridge.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Bridge between Responses requests and the current chat-completions executor.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any, Optional + +from ..protocols import ContentBlock, ToolDefinition, UnifiedMessage, UnifiedRequest, serialize_value +from ..protocols.responses import ResponsesProtocol +from ..session_tracking import SessionTrackingHints +from .types import generate_response_id + +_CHAT_GENERATION_KEYS = { + "frequency_penalty", + "logit_bias", + "logprobs", + "max_completion_tokens", + "max_tokens", + "max_output_tokens", + "metadata", + "n", + "parallel_tool_calls", + "presence_penalty", + "reasoning", + "reasoning_effort", + "response_format", + "seed", + "stop", + "stream_options", + "temperature", + "tool_choice", + "top_logprobs", + "top_p", + "user", +} +PROXY_ROUTING_KEYS = {"classifier", "api_keys", "providers", "private", "model_filters"} + + +class ResponsesBridge: + """Temporary compatibility bridge from Responses to chat completions. + + Later provider phases should call native Responses-capable providers directly. + Until then, this bridge makes `/v1/responses` useful while preserving enough + metadata and trace detail to debug fields that cannot be represented in chat. + """ + + def __init__(self, protocol: Optional[ResponsesProtocol] = None) -> None: + self.protocol = protocol or ResponsesProtocol() + + def to_chat_kwargs( + self, + unified: UnifiedRequest, + *, + parent_response: Optional[dict[str, Any]] = None, + parent_responses: Optional[list[dict[str, Any]]] = None, + ) -> dict[str, Any]: + """Convert a parsed Responses request to chat-completions kwargs.""" + + messages: list[dict[str, Any]] = [] + system_text = _blocks_to_text(unified.system) + if system_text: + messages.append({"role": "system", "content": system_text}) + if parent_responses is not None: + for parent in parent_responses: + messages.extend(_parent_request_to_messages(parent.get("request") or {})) + messages.extend(_parent_output_to_messages(parent.get("output") or parent.get("response", {}).get("output") or [])) + elif parent_response: + messages.extend(_parent_output_to_messages(parent_response.get("output") or [])) + messages.extend(_message_to_chat(message) for message in unified.messages) + kwargs: dict[str, Any] = { + "model": unified.model, + "messages": messages, + "stream": unified.stream, + } + if unified.tools: + kwargs["tools"] = [_tool_to_chat(tool) for tool in unified.tools] + for key, value in unified.generation_params.items(): + if key in _CHAT_GENERATION_KEYS: + kwargs[_chat_generation_key(key)] = deepcopy(value) + if unified.metadata: + kwargs.setdefault("metadata", deepcopy(unified.metadata)) + unsupported = { + "previous_response_id": unified.previous_response_id, + "extra": { + key: deepcopy(value) + for key, value in unified.extra.items() + if key not in PROXY_ROUTING_KEYS + }, + } + kwargs["_responses_bridge"] = {k: v for k, v in unsupported.items() if v} + hints = responses_session_hints(unified.previous_response_id) + if hints: + kwargs["_session_tracking_hints"] = hints + return kwargs + + def from_chat_response( + self, + chat_response: Any, + unified_request: UnifiedRequest, + *, + response_id: Optional[str] = None, + ) -> dict[str, Any]: + """Convert a chat-completions response into a Responses object.""" + + response = _as_dict(chat_response) + output = [] + for index, choice in enumerate(response.get("choices") or []): + if not isinstance(choice, dict): + continue + message = choice.get("message") or {} + output.extend(_chat_message_to_output_items(message, index)) + responses_payload = { + "id": response_id or response.get("id") or generate_response_id(), + "object": "response", + "created_at": response.get("created") or response.get("created_at"), + "model": response.get("model") or unified_request.model, + "status": _status_from_chat(response), + "output": output, + "usage": _usage_to_responses(response), + } + if unified_request.metadata: + responses_payload["metadata"] = deepcopy(unified_request.metadata) + return self.protocol.format_response(self.protocol.parse_response(responses_payload)) + + +def _chat_generation_key(key: str) -> str: + if key == "max_output_tokens": + return "max_tokens" + return key + + +def responses_session_hints( + previous_response_id: Optional[str], +) -> SessionTrackingHints | None: + """Return proxy-internal sticky routing evidence for Responses continuations.""" + + if not previous_response_id: + return None + anchor = f"responses_previous_response_id:{previous_response_id}" + return SessionTrackingHints( + global_strong_anchors=[anchor], + affinity_key=anchor, + ) + + +def _message_to_chat(message: UnifiedMessage) -> dict[str, Any]: + if message.role == "tool": + result = _tool_result_from_message(message) + payload = {"role": "tool", "content": result["content"]} + if result.get("tool_call_id"): + payload["tool_call_id"] = result["tool_call_id"] + return payload + payload = {"role": message.role, "content": _blocks_to_chat_content(message.content)} + if message.name: + payload["name"] = message.name + if message.tool_call_id: + payload["tool_call_id"] = message.tool_call_id + if message.tool_calls: + payload["tool_calls"] = [call.to_dict() for call in message.tool_calls] + return payload + + +def _tool_result_from_message(message: UnifiedMessage) -> dict[str, Any]: + for block in message.content: + if block.tool_result is not None: + content = block.tool_result.content + return { + "tool_call_id": message.tool_call_id or block.tool_result.tool_call_id, + "content": content if isinstance(content, str) else serialize_value(content), + } + return {"tool_call_id": message.tool_call_id, "content": _blocks_to_chat_content(message.content)} + + +def _blocks_to_chat_content(blocks: list[ContentBlock]) -> Any: + if not blocks: + return "" + if len(blocks) == 1 and blocks[0].text is not None: + return blocks[0].text + content = [] + for block in blocks: + if block.text is not None: + content.append({"type": "text", "text": block.text}) + elif block.source is not None: + content.append({"type": block.type, "source": deepcopy(block.source)}) + elif block.raw is not None: + content.append(deepcopy(block.raw)) + return content or "" + + +def _blocks_to_text(blocks: list[ContentBlock]) -> str: + return "".join(block.text or "" for block in blocks if block.text) + + +def _tool_to_chat(tool: ToolDefinition) -> dict[str, Any]: + if tool.type == "function": + return { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description or "", + "parameters": deepcopy(tool.input_schema), + }, + } + payload = {"type": tool.type, "name": tool.name} + payload.update(deepcopy(tool.extra)) + return payload + + +def _parent_output_to_messages(output: list[Any]) -> list[dict[str, Any]]: + messages = [] + for item in output: + if not isinstance(item, dict): + continue + if item.get("type") == "message": + text = _responses_content_to_text(item.get("content") or []) + if text: + messages.append({"role": item.get("role") or "assistant", "content": text}) + elif item.get("type") in {"function_call", "custom_tool_call"}: + tool_call = _responses_tool_call_to_chat(item) + if tool_call: + messages.append({"role": "assistant", "content": None, "tool_calls": [tool_call]}) + return messages + + +def _responses_tool_call_to_chat(item: dict[str, Any]) -> dict[str, Any] | None: + """Convert a stored Responses tool-call item to OpenAI chat shape.""" + + call_id = item.get("call_id") or item.get("id") + name = item.get("name") or item.get("tool_name") or item.get("type") + arguments = item.get("arguments", item.get("input", "")) + if call_id is None or name is None: + return None + return { + "id": call_id, + "type": "function", + "function": {"name": str(name), "arguments": arguments if isinstance(arguments, str) else serialize_value(arguments)}, + } + + +def _parent_request_to_messages(request: dict[str, Any]) -> list[dict[str, Any]]: + """Replay convertible parent Responses input items as chat messages.""" + + if not isinstance(request, dict): + return [] + protocol = ResponsesProtocol() + try: + unified = protocol.parse_request(request) + except Exception: + return _raw_input_to_messages(request.get("input")) + return [_message_to_chat(message) for message in unified.messages] + + +def _raw_input_to_messages(value: Any) -> list[dict[str, Any]]: + if value in (None, ""): + return [] + if isinstance(value, str): + return [{"role": "user", "content": value}] + if isinstance(value, list): + messages: list[dict[str, Any]] = [] + for item in value: + if isinstance(item, str): + messages.append({"role": "user", "content": item}) + elif isinstance(item, dict) and item.get("role") and item.get("content") is not None: + messages.append({"role": item.get("role"), "content": item.get("content")}) + return messages + return [] + + +def _chat_message_to_output_items(message: dict[str, Any], index: int) -> list[dict[str, Any]]: + items = [] + content = message.get("content") + if content is not None: + items.append( + { + "id": f"msg_{index}", + "type": "message", + "role": message.get("role") or "assistant", + "content": [{"type": "output_text", "text": content if isinstance(content, str) else serialize_value(content)}], + } + ) + for tool_call in message.get("tool_calls") or []: + function = tool_call.get("function") if isinstance(tool_call, dict) else None + if isinstance(function, dict): + items.append( + { + "id": tool_call.get("id"), + "type": "function_call", + "call_id": tool_call.get("id"), + "name": function.get("name"), + "arguments": function.get("arguments"), + } + ) + return items + + +def _responses_content_to_text(content: list[Any]) -> str: + parts = [] + for block in content: + if isinstance(block, dict): + text = block.get("text") + if text: + parts.append(str(text)) + elif isinstance(block, str): + parts.append(block) + return "".join(parts) + + +def _status_from_chat(response: dict[str, Any]) -> str: + for choice in response.get("choices") or []: + finish_reason = choice.get("finish_reason") if isinstance(choice, dict) else None + if finish_reason in {"length", "content_filter"}: + return "incomplete" + return "completed" + + +def _usage_to_responses(usage: Any) -> Any: + if isinstance(usage, dict) and isinstance(usage.get("usage"), dict): + nested = dict(usage["usage"]) + for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd", "currency", "costMetadata"): + if key in usage and key not in nested: + nested[key] = deepcopy(usage[key]) + usage = nested + if not isinstance(usage, dict): + return usage + result = { + "input_tokens": usage.get("prompt_tokens", usage.get("input_tokens", 0)), + "output_tokens": usage.get("completion_tokens", usage.get("output_tokens", 0)), + "total_tokens": usage.get("total_tokens", 0), + } + prompt_details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") + if isinstance(prompt_details, dict): + result["input_tokens_details"] = {"cached_tokens": prompt_details.get("cached_tokens", 0)} + completion_details = usage.get("completion_tokens_details") or usage.get("output_tokens_details") + if isinstance(completion_details, dict): + result["output_tokens_details"] = {"reasoning_tokens": completion_details.get("reasoning_tokens", 0)} + for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd", "currency", "costMetadata"): + if key in usage: + result[key] = deepcopy(usage[key]) + if "request_cost_usd" in result and "cost_details" not in result: + result["cost_details"] = {"request_cost_usd": result["request_cost_usd"]} + if "estimated_cost" in result and "cost_details" not in result: + result["cost_details"] = {"estimated_cost": result["estimated_cost"]} + return result + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return deepcopy(value) + if hasattr(value, "model_dump"): + return value.model_dump() + if hasattr(value, "dict"): + return value.dict() + return serialize_value(value) if isinstance(serialize_value(value), dict) else {} diff --git a/src/rotator_library/responses/service.py b/src/rotator_library/responses/service.py new file mode 100644 index 000000000..87f8d2c71 --- /dev/null +++ b/src/rotator_library/responses/service.py @@ -0,0 +1,1577 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Service layer for the OpenAI-compatible Responses API.""" + +from __future__ import annotations + +import asyncio +import hashlib +import hmac +import json +import re +import secrets +import time +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, AsyncGenerator, NoReturn, Optional + +from ..protocols import ProtocolContext, get_protocol +from ..protocols.streaming import ProtocolStreamConverter +from ..streaming import StreamEvent, StreamMonitor +from ..config.experimental import get_stream_runtime_settings +from ..client.scopes import derive_session_isolation_key +from ..core.errors import StructuredAPIResponseError +from ..usage.accounting import extract_usage_record +from ..usage.costs import CostCalculator +from ..protocols.responses import ResponsesProtocol +from .bridge import PROXY_ROUTING_KEYS, ResponsesBridge, responses_session_hints +from .store import InMemoryResponsesStore, ResponsesStore +from .streaming import ( + ResponsesSSEFormatter, + ResponsesStreamEvent, + ResponsesStreamState, + output_item_added_payload, + output_item_done_payload, + output_text_delta_payload, + parse_chat_sse_chunk, + response_completed_payload, + response_created_payload, + response_failed_payload, +) +from .types import ResponsesStoreSettings, StoredResponse +from .types import generate_response_id + + +_STORED_REQUEST_FIELDS = { + "include", + "input", + "instructions", + "max_output_tokens", + "metadata", + "model", + "parallel_tool_calls", + "previous_response_id", + "reasoning", + "store", + "stream", + "temperature", + "text", + "tool_choice", + "tools", + "top_p", + "truncation", + "user", +} +_SCOPE_ACCESS_TOKEN_RE = re.compile( + r"^(classifier:[0-9a-f]{24}|bundle:[0-9a-f]{64})\.([A-Za-z0-9_-]{32,64})$" +) + + +@dataclass(frozen=True) +class ResponsesRequestScope: + """Internal scope plus the unforgeable capability returned to HTTP clients.""" + + key: str + access_token: str + + @property + def access_token_hash(self) -> str: + return hashlib.sha256(self.access_token.encode("utf-8")).hexdigest() + + +class ResponsesServiceError(ValueError): + """Error with an HTTP-compatible status code for proxy routes.""" + + def __init__(self, message: str, *, status_code: int = 400, error_type: str = "invalid_request_error") -> None: + self.status_code = status_code + self.error_type = error_type + super().__init__(message) + + def to_protocol_payload(self, protocol: str) -> dict[str, Any]: + """Format this service failure in an independently selected protocol.""" + + if protocol == "responses": + return { + "error": { + "message": str(self), + "type": self.error_type, + "code": self.status_code, + } + } + normalized = { + "authentication_error": "authentication", + "permission_error": "forbidden", + "rate_limit_error": "rate_limit", + "invalid_request_error": "invalid_request", + "not_found_error": "not_found", + }.get(self.error_type, self.error_type) + return StructuredAPIResponseError( + str(self), + error_type=normalized, + status_code=self.status_code, + ).to_protocol_payload(protocol) + + +class ResponsesService: + """Create, store, retrieve, and delete Responses API objects. + + Phase 4 deliberately bridges through the existing chat-completions execution + path. Native Responses-capable provider execution will replace the bridge for + covered providers in later phases without changing the route/storage surface. + """ + + def __init__( + self, + *, + protocol: Optional[ResponsesProtocol] = None, + bridge: Optional[ResponsesBridge] = None, + store: Optional[ResponsesStore] = None, + store_settings: Optional[ResponsesStoreSettings] = None, + ) -> None: + self.store_settings = store_settings or ResponsesStoreSettings() + self.protocol = protocol or ResponsesProtocol() + self.bridge = bridge or ResponsesBridge(self.protocol) + self.store = store or InMemoryResponsesStore(max_items=self.store_settings.max_items) + + @staticmethod + def request_scope_key(raw_request: dict[str, Any]) -> str: + """Return the opaque caller/credential domain for a Responses request.""" + + return _request_isolation_key(raw_request) + + @staticmethod + def redact_request_for_logging(raw_request: dict[str, Any]) -> dict[str, Any]: + """Return a recursive credential-free copy for transport-level logs.""" + + redacted = _redact_sensitive_fields(raw_request) + return redacted if isinstance(redacted, dict) else {} + + def prepare_request_scope( + self, + raw_request: dict[str, Any], + ) -> ResponsesRequestScope: + """Create the scope capability used by transport-facing retrieval APIs.""" + + scope_key = self.request_scope_key(raw_request) + access_token = ( + "public" + if scope_key == "public" + else f"{scope_key}.{secrets.token_urlsafe(32)}" + ) + return ResponsesRequestScope(scope_key, access_token) + + def _resolve_request_scope( + self, + raw_request: dict[str, Any], + request_scope: Optional[ResponsesRequestScope], + ) -> ResponsesRequestScope: + expected_key = self.request_scope_key(raw_request) + resolved = request_scope or self.prepare_request_scope(raw_request) + if resolved.key != expected_key: + raise ResponsesServiceError( + "Responses request scope does not match routing credentials", + status_code=400, + ) + return resolved + + async def close(self) -> None: + """Close the configured response store and its owned background tasks.""" + + close = getattr(self.store, "close", None) + if close: + await close() + + async def create_response( + self, + raw_request: dict[str, Any], + client: Any, + *, + request: Optional[Any] = None, + transaction_logger: Optional[Any] = None, + request_scope: Optional[ResponsesRequestScope] = None, + previous_response_access_token: Optional[str] = None, + ) -> dict[str, Any]: + """Create a non-streaming Responses object through the chat bridge.""" + + if not raw_request.get("model"): + raise ResponsesServiceError("'model' is required", status_code=400) + if raw_request.get("stream"): + raise ResponsesServiceError("Use stream_response for streaming requests", status_code=400) + + resolved_scope = self._resolve_request_scope(raw_request, request_scope) + selected_output = ( + client.resolve_output_protocol( + raw_request, + input_protocol="responses", + request=request, + ) + if hasattr(client, "resolve_output_protocol") + else "responses" + ) + isolation_key = resolved_scope.key + safe_request = _safe_stored_request(raw_request) + self._trace(transaction_logger, "responses_raw_request", safe_request, direction="request", stage="client") + try: + unified = self.protocol.parse_request(raw_request, ProtocolContext(source_protocol="responses")) + except Exception as exc: + self._log_transform_error(transaction_logger, "responses_parse_request", exc, safe_request) + raise + if transaction_logger: + self._trace(transaction_logger, "responses_parsed_request", _redact_sensitive_fields(unified.to_dict()), direction="request", stage="protocol") + + parent = await self._load_previous_response( + unified.previous_response_id, + transaction_logger, + expected_scope_key=isolation_key, + access_token=previous_response_access_token, + ) + try: + parent_lineage = await self._load_response_lineage( + parent, + expected_scope_key=isolation_key, + ) + except Exception as exc: + self._log_transform_error( + transaction_logger, + "responses_bridge_chat_request", + exc, + _redact_sensitive_fields(unified.to_dict()), + ) + raise + session_hints = responses_session_hints(unified.previous_response_id) + session_info: dict[str, Any] = { + "scope_access_hash": resolved_scope.access_token_hash, + } + if hasattr(client, "agenerate"): + native_request = _expanded_responses_request(raw_request, parent_lineage) + internal_kwargs = _internal_client_kwargs(client, session_hints, session_info) + self._trace( + transaction_logger, + "responses_native_protocol_request", + _redact_sensitive_fields(native_request), + direction="request", + stage="protocol", + metadata={"lineage_depth": len(parent_lineage), "has_session_hints": bool(session_hints)}, + ) + try: + response = await client.agenerate( + native_request, + input_protocol="responses", + output_protocol="responses", + request=request, + _disable_provider_continuation=bool(parent_lineage), + **_routing_kwargs(raw_request), + **internal_kwargs, + ) + except StructuredAPIResponseError as exc: + raise ResponsesServiceError( + str(exc), + error_type=exc.error_type, + status_code=exc.http_status, + ) from exc + response_payload = self._response_to_dict(response) + self._trace(transaction_logger, "responses_native_protocol_response", response_payload, direction="response", stage="provider") + else: + chat_kwargs = self.bridge.to_chat_kwargs( + unified, + parent_responses=[stored.to_dict() for stored in parent_lineage] if parent_lineage else None, + ) + bridge_metadata = chat_kwargs.pop("_responses_bridge", {}) + chat_kwargs.pop("_session_tracking_hints", None) + chat_kwargs.update(_routing_kwargs(raw_request)) + chat_kwargs.update(_internal_client_kwargs(client, session_hints, session_info)) + self._trace( + transaction_logger, + "responses_bridge_chat_request", + _without_internal_kwargs(chat_kwargs), + direction="request", + stage="adapter", + metadata={"bridge_metadata": {"extra_keys": sorted((bridge_metadata.get("extra") or {}).keys()), "has_session_hints": bool(session_hints)}}, + ) + chat_kwargs["_output_protocol"] = "openai_chat" + chat_response = await client.acompletion(request=request, **chat_kwargs) + if transaction_logger: + self._trace(transaction_logger, "responses_bridge_chat_response", self._response_to_dict(chat_response), direction="response", stage="provider") + try: + response_payload = self.bridge.from_chat_response(chat_response, unified) + except Exception as exc: + self._log_transform_error(transaction_logger, "responses_bridge_chat_response", exc, self._response_to_dict(chat_response)) + raise + _record_responses_session_anchor(session_info, response_payload) + self._trace(transaction_logger, "responses_parsed_response", response_payload, direction="response", stage="protocol") + self._trace_responses_usage(transaction_logger, response_payload, unified.model, source="responses_response") + + if raw_request.get("store", True): + stored = self._stored_response(raw_request, response_payload, parent, session_info=session_info) + try: + await self.store.save(stored) + except Exception as exc: + self._log_transform_error(transaction_logger, "responses_store_response", exc, stored.to_dict()) + raise + self._trace(transaction_logger, "responses_stored_response", stored.to_dict(), direction="metadata", stage="final") + + self._trace(transaction_logger, "responses_final_response", response_payload, direction="response", stage="final") + if selected_output == "responses": + return response_payload + unified_response = self.protocol.parse_response( + response_payload, + ProtocolContext(source_protocol="responses", target_protocol=selected_output), + ) + return get_protocol(selected_output).format_response( + unified_response, + ProtocolContext( + source_protocol="responses", + target_protocol=selected_output, + input_protocol="responses", + output_protocol=selected_output, + provider_state_compatible=False, + ), + ) + + async def stream_response( + self, + raw_request: dict[str, Any], + client: Any, + *, + request: Optional[Any] = None, + transaction_logger: Optional[Any] = None, + transport: str = "sse", + request_scope: Optional[ResponsesRequestScope] = None, + previous_response_access_token: Optional[str] = None, + ) -> AsyncGenerator[str, None]: + """Stream a Responses API request as HTTP SSE events.""" + + if hasattr(client, "agenerate"): + async for frame in self._stream_native_response( + raw_request, + client, + request=request, + transaction_logger=transaction_logger, + transport=transport, + request_scope=request_scope, + previous_response_access_token=previous_response_access_token, + ): + yield frame + return + + formatter = ResponsesSSEFormatter() + async for event in self.stream_events( + raw_request, + client, + request=request, + transaction_logger=transaction_logger, + transport=transport, + request_scope=request_scope, + previous_response_access_token=previous_response_access_token, + ): + formatted = formatter.format_stream_event(event) + self._trace( + transaction_logger, + "responses_sse_formatted_event", + formatted, + direction="stream", + stage="final", + metadata={"event_name": event.event_name, "terminal": event.terminal, "transport": transport}, + scrub_strings=True, + ) + if event.heartbeat: + self._trace( + transaction_logger, + "responses_sse_formatted_heartbeat", + formatted, + direction="stream", + stage="final", + metadata={"transport": transport}, + scrub_strings=True, + ) + yield formatted + + async def _stream_native_response( + self, + raw_request: dict[str, Any], + client: Any, + *, + request: Optional[Any], + transaction_logger: Optional[Any], + transport: str, + request_scope: Optional[ResponsesRequestScope], + previous_response_access_token: Optional[str], + ) -> AsyncGenerator[str, None]: + """Stream through the canonical runtime while retaining Responses storage.""" + + stream_request = dict(raw_request) + stream_request["stream"] = True + resolved_scope = self._resolve_request_scope(stream_request, request_scope) + unified = self.protocol.parse_request( + stream_request, + ProtocolContext(source_protocol="responses", transport=transport), + ) + parent = await self._load_previous_response( + unified.previous_response_id, + transaction_logger, + expected_scope_key=resolved_scope.key, + access_token=previous_response_access_token, + ) + parent_lineage = await self._load_response_lineage( + parent, + expected_scope_key=resolved_scope.key, + ) + native_request = _expanded_responses_request(stream_request, parent_lineage) + session_hints = responses_session_hints(unified.previous_response_id) + session_info: dict[str, Any] = { + "scope_access_hash": resolved_scope.access_token_hash, + } + selected_output = ( + client.resolve_output_protocol( + stream_request, + input_protocol="responses", + request=request, + ) + if hasattr(client, "resolve_output_protocol") + else "responses" + ) + self._trace( + transaction_logger, + "responses_native_protocol_stream_request", + _redact_sensitive_fields(native_request), + direction="request", + stage="protocol", + metadata={"lineage_depth": len(parent_lineage), "selected_output": selected_output}, + ) + response_stream = await client.agenerate( + native_request, + input_protocol="responses", + output_protocol="responses", + request=request, + _disable_provider_continuation=bool(parent_lineage), + **_routing_kwargs(raw_request), + **_internal_client_kwargs(client, session_hints, session_info), + ) + response_context = ProtocolContext( + model=unified.model, + source_protocol="responses", + target_protocol=selected_output, + input_protocol="responses", + provider_protocol="responses", + output_protocol=selected_output, + transport=transport, + provider_state_compatible=False, + ) + converter = ( + ProtocolStreamConverter( + self.protocol, + get_protocol(selected_output), + response_context, + ) + if selected_output != "responses" + else None + ) + completed = False + async for raw_frame in response_stream: + if isinstance(raw_frame, str) and raw_frame.lstrip().startswith(":"): + yield raw_frame + continue + event = self.protocol.parse_stream_event(raw_frame, response_context) + payload = event.extra.get("payload") if isinstance(event.extra, dict) else None + response_payload = payload.get("response") if isinstance(payload, dict) and isinstance(payload.get("response"), dict) else None + if response_payload and event.type in {"response.completed", "response.failed", "response.incomplete"}: + completed = True + _record_responses_session_anchor(session_info, response_payload) + self._trace_responses_usage(transaction_logger, response_payload, unified.model, source="responses_stream") + stored = await self._store_stream_response( + stream_request, + response_payload, + parent, + failed=event.type == "response.failed", + transaction_logger=transaction_logger, + session_info=session_info, + ) + self._trace( + transaction_logger, + "responses_stored_stream_response" if stored else "responses_store_skipped", + response_payload if stored else {"response_id": response_payload.get("id")}, + direction="metadata", + stage="final", + ) + if converter is None: + yield raw_frame + else: + for formatted in converter.convert(raw_frame): + yield formatted + if not completed: + raise ResponsesServiceError( + "Responses stream ended without a terminal response event", + status_code=502, + error_type="upstream_error", + ) + + async def validate_stream_request( + self, + raw_request: dict[str, Any], + *, + request_scope: Optional[ResponsesRequestScope] = None, + previous_response_access_token: Optional[str] = None, + ) -> None: + """Validate stream-only preconditions before an HTTP response starts.""" + + if not raw_request.get("model"): + raise ResponsesServiceError("'model' is required", status_code=400) + previous_response_id = raw_request.get("previous_response_id") + if previous_response_id: + resolved_scope = self._resolve_request_scope(raw_request, request_scope) + await self._load_previous_response( + str(previous_response_id), + None, + expected_scope_key=resolved_scope.key, + access_token=previous_response_access_token, + ) + + async def stream_events( + self, + raw_request: dict[str, Any], + client: Any, + *, + request: Optional[Any] = None, + transaction_logger: Optional[Any] = None, + transport: str = "sse", + request_scope: Optional[ResponsesRequestScope] = None, + previous_response_access_token: Optional[str] = None, + ) -> AsyncGenerator[ResponsesStreamEvent, None]: + """Yield transport-neutral Responses events for streaming transports.""" + + if not raw_request.get("model"): + raise ResponsesServiceError("'model' is required", status_code=400) + stream_request = dict(raw_request) + stream_request["stream"] = True + resolved_scope = self._resolve_request_scope(stream_request, request_scope) + isolation_key = resolved_scope.key + safe_stream_request = _safe_stored_request(stream_request) + self._trace(transaction_logger, "responses_raw_request", safe_stream_request, direction="request", stage="client") + try: + unified = self.protocol.parse_request(stream_request, ProtocolContext(source_protocol="responses", transport=transport)) + except Exception as exc: + self._log_transform_error(transaction_logger, "responses_parse_request", exc, safe_stream_request) + raise + if transaction_logger: + self._trace(transaction_logger, "responses_parsed_request", _redact_sensitive_fields(unified.to_dict()), direction="request", stage="protocol") + parent = await self._load_previous_response( + unified.previous_response_id, + transaction_logger, + expected_scope_key=isolation_key, + access_token=previous_response_access_token, + ) + try: + parent_lineage = await self._load_response_lineage( + parent, + expected_scope_key=isolation_key, + ) + chat_kwargs = self.bridge.to_chat_kwargs(unified, parent_responses=[stored.to_dict() for stored in parent_lineage] if parent_lineage else None) + except Exception as exc: + self._log_transform_error( + transaction_logger, + "responses_bridge_chat_request", + exc, + _redact_sensitive_fields(unified.to_dict()), + ) + raise + bridge_metadata = chat_kwargs.pop("_responses_bridge", {}) + session_hints = chat_kwargs.pop("_session_tracking_hints", None) + session_info: dict[str, Any] = { + "scope_access_hash": resolved_scope.access_token_hash, + } + chat_kwargs.update(_routing_kwargs(stream_request)) + chat_kwargs.update(_internal_client_kwargs(client, session_hints, session_info)) + chat_kwargs["stream"] = True + trace_chat_kwargs = _without_internal_kwargs(chat_kwargs) + self._trace( + transaction_logger, + "responses_bridge_chat_request", + trace_chat_kwargs, + direction="request", + stage="adapter", + metadata={ + "bridge_metadata": { + "extra_keys": sorted((bridge_metadata.get("extra") or {}).keys()), + "has_session_hints": bool(session_hints), + }, + "transport": transport, + }, + ) + + response_id = generate_response_id() + state = ResponsesStreamState(response_id=response_id, model=unified.model) + usage = None + item_started = False + monitor = StreamMonitor(clock=time.monotonic) + stream_settings = get_stream_runtime_settings() + chat_stream = None + stream_iterator = None + upstream_closed = False + pending_next_task = None + pending_next_started_at = None + pending_next_last_heartbeat_at = None + acquire_task = None + acquire_started_at = None + acquire_last_heartbeat_at = None + ttfb_started_at = time.monotonic() + + async def cancel_task(task: Any) -> None: + """Cancel and await an in-flight stream task before closing its source.""" + + if task is None or task.done(): + return + task.cancel() + try: + await task + except (asyncio.CancelledError, StopAsyncIteration): + return + except Exception: + return + + async def close_upstream(reason: str) -> None: + """Best-effort close for upstream Responses bridge streams.""" + + nonlocal upstream_closed + if upstream_closed: + return + attempted = False + for candidate in (stream_iterator, chat_stream): + if candidate is None: + continue + attempted = True + try: + closer = getattr(candidate, "aclose", None) + if callable(closer): + await closer() + upstream_closed = True + self._trace(transaction_logger, "responses_stream_upstream_closed", {"reason": reason}, direction="stream", stage="provider", metadata={"transport": transport}) + return + closer = getattr(candidate, "close", None) + if callable(closer): + closer() + upstream_closed = True + self._trace(transaction_logger, "responses_stream_upstream_closed", {"reason": reason}, direction="stream", stage="provider", metadata={"transport": transport}) + return + except Exception as exc: + self._trace(transaction_logger, "responses_stream_upstream_close_failed", {"reason": reason, "error_type": type(exc).__name__}, direction="stream", stage="provider", metadata={"transport": transport}) + continue + if attempted: + self._trace(transaction_logger, "responses_stream_upstream_close_failed", {"reason": reason, "error_type": "no_close_method"}, direction="stream", stage="provider", metadata={"transport": transport}) + + async def next_upstream_chunk(*, first: bool) -> tuple[str, Any]: + """Return the next upstream chunk or a control marker.""" + + timeout = stream_settings.ttfb_timeout_seconds if first else stream_settings.stall_timeout_seconds + heartbeat = stream_settings.heartbeat_seconds + nonlocal pending_next_task, pending_next_started_at, pending_next_last_heartbeat_at + if pending_next_task is None: + pending_next_task = asyncio.create_task(stream_iterator.__anext__()) + pending_next_started_at = time.monotonic() + pending_next_last_heartbeat_at = pending_next_started_at + next_task = pending_next_task + started_at = ttfb_started_at if first else (pending_next_started_at or time.monotonic()) + while True: + if next_task.done(): + pending_next_task = None + pending_next_started_at = None + pending_next_last_heartbeat_at = None + return "chunk", next_task.result() + if request is not None and await request.is_disconnected(): + self._trace(transaction_logger, "responses_stream_disconnected", {"reason": "client_disconnected"}, direction="stream", stage="client", metadata={"transport": transport}) + if stream_settings.cancel_upstream_on_disconnect: + await cancel_task(next_task) + pending_next_task = None + await close_upstream("client_disconnected") + return "disconnect", None + elapsed = time.monotonic() - started_at + waits = [] + if timeout is not None: + remaining_timeout = timeout - elapsed + if remaining_timeout <= 0: + await cancel_task(next_task) + pending_next_task = None + await close_upstream("ttfb_timeout" if first else "stall_timeout") + raise ResponsesServiceError( + f"Responses stream {'TTFB' if first else 'stall'} timeout", + status_code=504, + error_type="api_connection", + ) + waits.append(remaining_timeout) + if heartbeat is not None: + last_heartbeat_at = pending_next_last_heartbeat_at or started_at + remaining_heartbeat = heartbeat - (time.monotonic() - last_heartbeat_at) + if remaining_heartbeat <= 0: + pending_next_last_heartbeat_at = time.monotonic() + return "heartbeat", None + waits.append(remaining_heartbeat) + wait_timeout = min(waits) if waits else None + if wait_timeout is None: + chunk = await next_task + pending_next_task = None + pending_next_started_at = None + pending_next_last_heartbeat_at = None + return "chunk", chunk + done, _ = await asyncio.wait({next_task}, timeout=wait_timeout) + if done: + pending_next_task = None + pending_next_started_at = None + pending_next_last_heartbeat_at = None + return "chunk", next_task.result() + last_heartbeat_at = pending_next_last_heartbeat_at or started_at + if heartbeat is not None and time.monotonic() - last_heartbeat_at >= heartbeat: + pending_next_last_heartbeat_at = time.monotonic() + return "heartbeat", None + + async def acquire_upstream_stream() -> tuple[str, Any]: + """Acquire the upstream stream under the same TTFB/disconnect policy.""" + + nonlocal acquire_task, acquire_started_at, acquire_last_heartbeat_at + if acquire_task is None: + acquire_task = asyncio.create_task(client.acompletion(request=request, **chat_kwargs)) + acquire_started_at = ttfb_started_at + acquire_last_heartbeat_at = acquire_started_at + task = acquire_task + started_at = acquire_started_at or time.monotonic() + timeout = stream_settings.ttfb_timeout_seconds + heartbeat = stream_settings.heartbeat_seconds + while True: + if task.done(): + acquire_task = None + acquire_started_at = None + acquire_last_heartbeat_at = None + return "stream", task.result() + if request is not None and await request.is_disconnected(): + self._trace(transaction_logger, "responses_stream_disconnected", {"reason": "client_disconnected", "phase": "acquire"}, direction="stream", stage="client", metadata={"transport": transport}) + await cancel_task(task) + acquire_task = None + acquire_started_at = None + acquire_last_heartbeat_at = None + return "disconnect", None + waits = [] + elapsed = time.monotonic() - started_at + if timeout is not None: + remaining_timeout = timeout - elapsed + if remaining_timeout <= 0: + await cancel_task(task) + acquire_task = None + acquire_started_at = None + acquire_last_heartbeat_at = None + raise ResponsesServiceError("Responses stream TTFB timeout", status_code=504, error_type="api_connection") + waits.append(remaining_timeout) + if heartbeat is not None: + last_heartbeat_at = acquire_last_heartbeat_at or started_at + remaining_heartbeat = heartbeat - (time.monotonic() - last_heartbeat_at) + if remaining_heartbeat <= 0: + acquire_last_heartbeat_at = time.monotonic() + return "heartbeat", None + waits.append(remaining_heartbeat) + wait_timeout = min(waits) if waits else None + if wait_timeout is None: + stream = await task + acquire_task = None + acquire_started_at = None + acquire_last_heartbeat_at = None + return "stream", stream + done, _ = await asyncio.wait({task}, timeout=wait_timeout) + if done: + acquire_task = None + acquire_started_at = None + acquire_last_heartbeat_at = None + return "stream", task.result() + last_heartbeat_at = acquire_last_heartbeat_at or started_at + if heartbeat is not None and time.monotonic() - last_heartbeat_at >= heartbeat: + acquire_last_heartbeat_at = time.monotonic() + return "heartbeat", None + + if transaction_logger: + self._trace( + transaction_logger, + "stream_started", + {"event": StreamEvent("started", protocol="responses").to_dict(), "metrics": monitor.metrics.to_dict()}, + direction="stream", + stage="client", + metadata={"transport": transport}, + ) + created = response_created_payload(response_id, unified.model) + self._trace(transaction_logger, "responses_stream_event_created", created, direction="stream", stage="final", metadata={"transport": transport}) + await self._store_stream_current_state( + stream_request, + created, + parent, + transaction_logger=transaction_logger, + session_info=session_info, + ) + yield ResponsesStreamEvent("response.created", created) + try: + while chat_stream is None: + marker, acquired = await acquire_upstream_stream() + if marker == "disconnect": + return + if marker == "heartbeat": + self._trace(transaction_logger, "responses_stream_heartbeat", {"comment": "heartbeat", "phase": "acquire"}, direction="stream", stage="transport", metadata={"transport": transport}) + yield ResponsesStreamEvent("heartbeat", {"comment": "heartbeat"}) + continue + chat_stream = acquired + stream_iterator = chat_stream.__aiter__() + first_chunk = True + while True: + try: + marker, raw_chunk = await next_upstream_chunk(first=first_chunk) + except StopAsyncIteration: + break + if marker == "disconnect": + return + if marker == "heartbeat": + self._trace(transaction_logger, "responses_stream_heartbeat", {"comment": "heartbeat"}, direction="stream", stage="transport", metadata={"transport": transport}) + yield ResponsesStreamEvent("heartbeat", {"comment": "heartbeat"}) + continue + first_chunk = False + if monitor.metrics.first_byte_at is None: + monitor.record_event(StreamEvent("raw_chunk", protocol="responses", raw=raw_chunk)) + if transaction_logger: + self._trace( + transaction_logger, + "stream_first_byte", + {"metrics": monitor.metrics.to_dict()}, + direction="stream", + stage="provider", + metadata={"transport": transport}, + ) + self._trace(transaction_logger, "raw_chat_bridge_stream_chunk", raw_chunk, direction="stream", stage="provider") + cost_usage = _responses_sse_cost_usage(raw_chunk) + if cost_usage is not None: + usage = _merge_responses_stream_usage(usage, cost_usage) + self._trace(transaction_logger, "responses_stream_cost_event", cost_usage, direction="stream", stage="metadata", metadata={"transport": transport}) + continue + chunk = parse_chat_sse_chunk(raw_chunk) + if not chunk or chunk.get("type") == "done": + continue + self._trace(transaction_logger, "parsed_unified_stream_event", chunk, direction="stream", stage="protocol") + if chunk.get("error") is not None or chunk.get("type") == "error": + raise ResponsesServiceError(_stream_error_message(chunk), status_code=502, error_type="upstream_error") + if chunk.get("usage"): + usage = _merge_responses_stream_usage(_responses_chunk_usage(chunk), usage) + delta = _chunk_text_delta(chunk) + if not delta: + continue + if not item_started: + item_started = True + added = output_item_added_payload(state) + self._trace(transaction_logger, "responses_stream_event_output_item_added", added, direction="stream", stage="final", metadata={"transport": transport}) + yield ResponsesStreamEvent("response.output_item.added", added) + state = ResponsesStreamState( + response_id=state.response_id, + model=state.model, + output_text=state.output_text + delta, + output_item_id=state.output_item_id, + ) + event = output_text_delta_payload(state, delta) + first_visible = monitor.metrics.first_visible_output_at is None + monitor.record_event( + StreamEvent( + "delta", + protocol="responses", + data=event, + visible_output=True, + ) + ) + if first_visible: + if transaction_logger: + self._trace( + transaction_logger, + "stream_first_visible_output", + {"event": event, "metrics": monitor.metrics.to_dict()}, + direction="stream", + stage="final", + metadata={"transport": transport}, + ) + self._trace(transaction_logger, "responses_stream_event_output_text_delta", event, direction="stream", stage="final", metadata={"transport": transport}) + await self._store_stream_current_state(stream_request, _current_stream_payload(state), parent, transaction_logger=transaction_logger, session_info=session_info) + yield ResponsesStreamEvent("response.output_text.delta", event) + + if not item_started: + added = output_item_added_payload(state) + self._trace(transaction_logger, "responses_stream_event_output_item_added", added, direction="stream", stage="final", metadata={"transport": transport}) + yield ResponsesStreamEvent("response.output_item.added", added) + done_item = output_item_done_payload(state) + self._trace(transaction_logger, "responses_stream_event_output_item_done", done_item, direction="stream", stage="final", metadata={"transport": transport}) + yield ResponsesStreamEvent("response.output_item.done", done_item) + completed = response_completed_payload(state, _usage_to_responses_stream(usage)) + _record_responses_session_anchor(session_info, completed) + self._trace_responses_usage(transaction_logger, completed, unified.model, source="responses_stream") + stored = await self._store_stream_response(stream_request, completed, parent, transaction_logger=transaction_logger, session_info=session_info) + if stored: + self._trace(transaction_logger, "responses_stored_stream_response", completed, direction="metadata", stage="final") + else: + self._trace(transaction_logger, "responses_store_skipped", {"response_id": completed.get("id")}, direction="metadata", stage="final") + monitor.complete() + if transaction_logger: + self._trace( + transaction_logger, + "stream_completed", + {"metrics": monitor.metrics.to_dict()}, + direction="stream", + stage="final", + metadata={"transport": transport}, + ) + self._trace( + transaction_logger, + "stream_metrics_final", + {"metrics": monitor.metrics.to_dict()}, + direction="stream", + stage="final", + metadata={"transport": transport}, + ) + self._trace(transaction_logger, "responses_stream_event_completed", completed, direction="stream", stage="final", metadata={"transport": transport}) + yield ResponsesStreamEvent("response.completed", completed) + self._trace(transaction_logger, "stream_done_event", {"raw": "done"}, direction="stream", stage="final", metadata={"transport": transport}) + yield ResponsesStreamEvent("done", {}, terminal=True) + except Exception as exc: + monitor.record_event(StreamEvent("error", protocol="responses", data={"error_type": exc.__class__.__name__})) + failed = response_failed_payload(response_id, unified.model, _stream_failure_error(exc)) + if state.output_text: + failed["output"] = [output_item_done_payload(state)["item"]] + self._log_transform_error(transaction_logger, "responses_stream", exc, stream_request) + stored = await self._store_stream_response(stream_request, failed, parent, failed=True, transaction_logger=transaction_logger, session_info=session_info) + if stored: + self._trace(transaction_logger, "responses_stored_failed_stream_response", {"response_id": failed.get("id"), "status": "failed"}, direction="metadata", stage="final") + self._trace(transaction_logger, "responses_stream_event_failed", failed, direction="stream", stage="final", metadata={"transport": transport}, scrub_strings=True) + if transaction_logger: + self._trace( + transaction_logger, + "stream_metrics_final", + {"metrics": monitor.metrics.to_dict()}, + direction="stream", + stage="final", + metadata={"transport": transport, "failed": True}, + ) + yield ResponsesStreamEvent("response.failed", failed) + self._trace(transaction_logger, "stream_done_event", {"raw": "done"}, direction="stream", stage="final", metadata={"transport": transport, "failed": True}) + yield ResponsesStreamEvent("done", {}, terminal=True) + finally: + if chat_stream is None and acquire_task is not None and acquire_task.done() and not acquire_task.cancelled(): + try: + chat_stream = acquire_task.result() + stream_iterator = chat_stream.__aiter__() + except Exception: + chat_stream = None + await cancel_task(pending_next_task) + await cancel_task(acquire_task) + if chat_stream is not None and not upstream_closed: + await close_upstream("wrapper_exit") + + async def get_response( + self, + response_id: str, + *, + scope_key: str = "public", + ) -> dict[str, Any]: + """Return a stored response payload or raise a 404-compatible error.""" + + stored = await self._stored_or_not_found(response_id, scope_key) + return deepcopy(stored.response) + + async def get_response_with_access_token( + self, + response_id: str, + access_token: str = "public", + ) -> dict[str, Any]: + """Return a response only when the transport capability is valid.""" + + stored = await self._stored_for_access(response_id, access_token) + return deepcopy(stored.response) + + async def delete_response( + self, + response_id: str, + *, + scope_key: str = "public", + ) -> dict[str, Any]: + """Delete a stored response and return a compatible deletion object.""" + + await self._stored_or_not_found(response_id, scope_key) + deleted = await self.store.delete(response_id, scope_key) + if not deleted: + raise ResponsesServiceError(f"Response not found: {response_id}", status_code=404, error_type="not_found_error") + return {"id": response_id, "object": "response.deleted", "deleted": True} + + async def delete_response_with_access_token( + self, + response_id: str, + access_token: str = "public", + ) -> dict[str, Any]: + """Delete a response only when the transport capability is valid.""" + + stored = await self._stored_for_access(response_id, access_token) + deleted = await self.store.delete(response_id, stored.scope_key or "public") + if not deleted: + self._raise_response_not_found(response_id) + return {"id": response_id, "object": "response.deleted", "deleted": True} + + async def list_input_items( + self, + response_id: str, + *, + scope_key: str = "public", + ) -> dict[str, Any]: + """Return stored input items for a response continuation.""" + + stored = await self._stored_or_not_found(response_id, scope_key) + return {"object": "list", "data": deepcopy(stored.input_items)} + + async def list_input_items_with_access_token( + self, + response_id: str, + access_token: str = "public", + ) -> dict[str, Any]: + """Return input items only when the transport capability is valid.""" + + stored = await self._stored_for_access(response_id, access_token) + return {"object": "list", "data": deepcopy(stored.input_items)} + + async def _stored_for_access( + self, + response_id: str, + access_token: str, + ) -> StoredResponse: + if access_token == "public": + return await self._stored_or_not_found(response_id, "public") + match = _SCOPE_ACCESS_TOKEN_RE.fullmatch(str(access_token)) + if not match: + self._raise_response_not_found(response_id) + scope_key = match.group(1) + stored = await self._stored_or_not_found(response_id, scope_key) + expected_hash = str(stored.metadata.get("scope_access_hash") or "") + actual_hash = hashlib.sha256(str(access_token).encode("utf-8")).hexdigest() + if not expected_hash or not hmac.compare_digest(expected_hash, actual_hash): + self._raise_response_not_found(response_id) + return stored + + async def _stored_or_not_found( + self, + response_id: str, + scope_key: str, + ) -> StoredResponse: + stored = await self.store.get(response_id, scope_key) + if stored is None or stored.scope_key != scope_key: + self._raise_response_not_found(response_id) + return stored + + @staticmethod + def _raise_response_not_found(response_id: str) -> NoReturn: + raise ResponsesServiceError( + f"Response not found: {response_id}", + status_code=404, + error_type="not_found_error", + ) + + async def _load_response_lineage( + self, + parent: Optional[StoredResponse], + *, + expected_scope_key: str, + max_depth: int = 20, + ) -> list[StoredResponse]: + """Return parent continuation lineage from oldest to newest.""" + + if parent is None: + return [] + lineage: list[StoredResponse] = [] + seen: set[str] = set() + current: Optional[StoredResponse] = parent + while current is not None and current.id not in seen and len(lineage) < max_depth: + if current.scope_key != expected_scope_key: + raise ResponsesServiceError( + f"Previous response not found: {current.id}", + status_code=404, + error_type="not_found_error", + ) + seen.add(current.id) + lineage.append(current) + previous_id = current.request.get("previous_response_id") if isinstance(current.request, dict) else None + if not previous_id: + break + current = await self.store.get(str(previous_id), expected_scope_key) + return list(reversed(lineage)) + + async def _load_previous_response( + self, + response_id: Optional[str], + transaction_logger: Optional[Any], + *, + expected_scope_key: str, + access_token: Optional[str] = None, + ) -> Optional[StoredResponse]: + if not response_id: + return None + if expected_scope_key == "public": + parent = await self._stored_or_not_found(response_id, "public") + else: + parent = await self._stored_for_access(response_id, access_token or "") + if parent.scope_key != expected_scope_key: + self._raise_response_not_found(response_id) + if transaction_logger: + self._trace( + transaction_logger, + "responses_previous_response_loaded", + parent.to_dict(), + direction="metadata", + stage="adapter", + metadata={ + "previous_response_id": response_id, + "output_count": len(parent.output_items), + "input_item_count": len(parent.input_items), + "bridge_context_expanded": True, + }, + ) + return parent + + def _stored_response( + self, + raw_request: dict[str, Any], + response_payload: dict[str, Any], + parent: Optional[StoredResponse], + *, + session_info: Optional[dict[str, Any]] = None, + ) -> StoredResponse: + session_info = session_info or {} + return StoredResponse( + id=str(response_payload["id"]), + created_at=float(response_payload.get("created_at") or time.time()), + model=str(response_payload.get("model") or raw_request.get("model") or ""), + status=str(response_payload.get("status") or "completed"), + request=_safe_stored_request(raw_request), + response=deepcopy(response_payload), + input_items=_input_items(raw_request), + output_items=deepcopy(response_payload.get("output") or []), + usage=deepcopy(response_payload.get("usage")) if isinstance(response_payload.get("usage"), dict) else None, + metadata={ + "previous_response_id": parent.id if parent else raw_request.get("previous_response_id"), + "response_id": response_payload.get("id"), + "scope_access_hash": session_info.get("scope_access_hash"), + }, + session_id=session_info.get("session_id") or (parent.session_id if parent else None), + scope_key=( + session_info.get("scope_key") + or (parent.scope_key if parent else None) + or self.request_scope_key(raw_request) + ), + classifier=session_info.get("classifier") or (parent.classifier if parent else None), + expires_at=_expires_at(self.store_settings), + ) + + @staticmethod + def _response_to_dict(response: Any) -> Any: + if isinstance(response, dict): + return deepcopy(response) + if hasattr(response, "model_dump"): + return response.model_dump() + if hasattr(response, "dict"): + return response.dict() + return repr(response) + + @staticmethod + def _trace( + transaction_logger: Optional[Any], + pass_name: str, + data: Any, + *, + direction: str, + stage: str, + metadata: Optional[dict[str, Any]] = None, + scrub_strings: bool = False, + ) -> None: + if not transaction_logger: + return + transaction_logger.log_transform_pass( + pass_name, + data, + direction=direction, + stage=stage, + protocol="responses", + metadata=metadata or {}, + scrub_strings=scrub_strings, + ) + + @staticmethod + def _log_transform_error(transaction_logger: Optional[Any], pass_name: str, error: BaseException, payload: Any) -> None: + if transaction_logger: + transaction_logger.log_transform_error(pass_name, error, payload=payload, stage="adapter", protocol="responses") + + def _trace_responses_usage( + self, + transaction_logger: Optional[Any], + response_payload: dict[str, Any], + model: str, + *, + source: str, + ) -> None: + """Trace normalized Responses usage without changing stored payloads.""" + + if not transaction_logger: + return + usage = response_payload.get("usage") if isinstance(response_payload, dict) else None + if not usage: + return + record = extract_usage_record(usage, provider="responses", model=model, source=source) + cost_breakdown = CostCalculator().calculate(record, model=model, provider="responses") + self._trace( + transaction_logger, + "usage_accounting_summary", + {"usage": record.to_dict(), "cost": cost_breakdown.to_dict()}, + direction="metadata", + stage="final", + metadata={"source": source, "pricing_source": cost_breakdown.pricing_source}, + ) + + async def _store_stream_response( + self, + raw_request: dict[str, Any], + response_payload: dict[str, Any], + parent: Optional[StoredResponse], + *, + failed: bool = False, + transaction_logger: Optional[Any] = None, + session_info: Optional[dict[str, Any]] = None, + ) -> bool: + if not raw_request.get("store", True): + return False + if failed and not self.store_settings.store_failed: + return False + stored = self._stored_response(raw_request, response_payload, parent, session_info=session_info) + try: + await self.store.save(stored) + except Exception as exc: + self._log_transform_error(transaction_logger, "responses_store_stream_response", exc, stored.to_dict()) + raise + return True + + async def _store_stream_current_state( + self, + raw_request: dict[str, Any], + response_payload: dict[str, Any], + parent: Optional[StoredResponse], + *, + transaction_logger: Optional[Any], + session_info: Optional[dict[str, Any]] = None, + ) -> bool: + """Optionally persist in-progress stream state for retrieval surfaces.""" + + if not self.store_settings.store_in_progress or not raw_request.get("store", True): + return False + stored = self._stored_response(raw_request, response_payload, parent, session_info=session_info) + try: + await self.store.save(stored) + except Exception as exc: + self._log_transform_error(transaction_logger, "responses_store_stream_current_state", exc, stored.to_dict()) + raise + self._trace( + transaction_logger, + "responses_stored_stream_current_state", + {"response_id": response_payload.get("id"), "status": response_payload.get("status")}, + direction="metadata", + stage="final", + ) + return True + + +def _input_items(raw_request: dict[str, Any]) -> list[Any]: + value = raw_request.get("input") + if value is None: + return [] + return deepcopy(value if isinstance(value, list) else [value]) + + +def _expanded_responses_request( + raw_request: dict[str, Any], + lineage: list[StoredResponse], +) -> dict[str, Any]: + """Expand proxy-owned continuation history into native Responses input.""" + + expanded = { + key: deepcopy(value) + for key, value in raw_request.items() + if key not in PROXY_ROUTING_KEYS + } + input_items: list[Any] = [] + for stored in lineage: + input_items.extend(deepcopy(stored.input_items)) + input_items.extend(deepcopy(stored.output_items)) + input_items.extend(_input_items(raw_request)) + expanded["input"] = input_items + expanded.pop("previous_response_id", None) + return expanded + + +def _routing_kwargs(raw_request: dict[str, Any]) -> dict[str, Any]: + """Carry proxy routing controls to RequestContextBuilder without tracing them.""" + + return { + key: deepcopy(raw_request[key]) + for key in PROXY_ROUTING_KEYS + if key in raw_request + } + + +def _request_isolation_key(raw_request: dict[str, Any]) -> str: + return derive_session_isolation_key( + raw_request.get("classifier"), + raw_request.get("api_keys"), + raw_request.get("providers"), + bool(raw_request.get("private", False)), + ) + + +def _safe_stored_request(raw_request: dict[str, Any]) -> dict[str, Any]: + """Persist only standard continuation fields, never routing credentials.""" + + return { + key: deepcopy(raw_request[key]) + for key in _STORED_REQUEST_FIELDS + if key in raw_request + } + + +def _redact_sensitive_fields(value: Any) -> Any: + """Remove routing/auth containers recursively from diagnostic payloads.""" + + if isinstance(value, dict): + return { + key: _redact_sensitive_fields(item) + for key, item in value.items() + if str(key).lower() + not in { + "api_key", + "api_keys", + "authorization", + "credential_secrets", + "provider_config", + "providers", + "x-api-key", + } + } + if isinstance(value, list): + return [_redact_sensitive_fields(item) for item in value] + if isinstance(value, tuple): + return tuple(_redact_sensitive_fields(item) for item in value) + return deepcopy(value) + + +def _current_stream_payload(state: ResponsesStreamState) -> dict[str, Any]: + """Return a retrievable in-progress Responses object for stream state.""" + + payload = response_completed_payload(state) + payload["status"] = "in_progress" + return payload + + +def _expires_at(settings: ResponsesStoreSettings) -> Optional[float]: + """Return the expiration timestamp for a new stored response, if enabled.""" + + ttl = settings.ttl_seconds + if ttl is None or ttl <= 0: + return None + return time.time() + ttl + + +def _internal_client_kwargs(client: Any, session_hints: Any, session_info: dict[str, Any]) -> dict[str, Any]: + """Return hidden kwargs only for the internal RotatingClient path.""" + + if not _supports_internal_context_kwargs(client): + return {} + kwargs: dict[str, Any] = {"_request_context_callback": _capture_request_context(session_info)} + if session_hints: + kwargs["_session_tracking_hints"] = session_hints + return kwargs + + +def _supports_internal_context_kwargs(client: Any) -> bool: + """Return whether a client is the proxy's internal rotating client.""" + + return hasattr(client, "_request_builder") and hasattr(client, "_executor") + + +def _capture_request_context(session_info: dict[str, Any]): + """Build a callback that records non-secret request context metadata.""" + + def capture(context: Any) -> None: + session_info["session_id"] = getattr(context, "session_id", None) + session_info["scope_key"] = getattr(context, "session_isolation_key", None) + session_info["classifier"] = getattr(context, "classifier", None) + session_info["session_tracker"] = getattr(context, "session_tracker", None) + session_info["provider"] = getattr(context, "provider", None) + session_info["model"] = getattr(context, "model", None) + session_info["tracking_namespace"] = getattr(context, "session_tracking_namespace", None) + + return capture + + +def _without_internal_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + """Return trace/provider-visible kwargs without proxy-internal controls.""" + + return _redact_sensitive_fields( + { + key: deepcopy(value) + for key, value in kwargs.items() + if not key.startswith("_") and key not in PROXY_ROUTING_KEYS + } + ) + + +def _record_responses_session_anchor(session_info: dict[str, Any], response_payload: dict[str, Any]) -> None: + """Record emitted Responses IDs as response-derived session evidence.""" + + tracker = session_info.get("session_tracker") + session_id = session_info.get("session_id") + if not tracker or not session_id or not response_payload.get("id"): + return + tracker.record_response( + session_id, + provider=session_info.get("provider"), + model=session_info.get("model"), + scope_key=session_info.get("scope_key"), + tracking_namespace=session_info.get("tracking_namespace"), + response={"id": response_payload.get("id"), "object": "response"}, + ) + + +def _stream_error_message(chunk: dict[str, Any]) -> str: + """Return a compact, client-safe message for upstream stream error chunks.""" + + error = chunk.get("error") + if isinstance(error, dict): + message = error.get("message") or error.get("type") + if message: + return str(message) + message = chunk.get("message") + return str(message) if message else "Upstream stream error" + + +def _stream_failure_error(exc: Exception) -> dict[str, Any]: + """Return a client-safe Responses stream failure object.""" + + error_type = getattr(exc, "error_type", None) or exc.__class__.__name__ + result = {"message": str(exc), "type": str(error_type)} + if isinstance(exc, ResponsesServiceError) and error_type == "api_connection": + text = str(exc).lower() + if "ttfb" in text: + result["timeout_type"] = "ttfb" + elif "stall" in text: + result["timeout_type"] = "stall" + return result + + +def _responses_sse_cost_usage(chunk: Any) -> Optional[dict[str, Any]]: + """Extract Responses stream cost metadata from SSE comments/events.""" + + if not isinstance(chunk, str): + return None + payload = _responses_sse_cost_payload(chunk) + if payload is None: + return None + if isinstance(payload, (int, float, str)): + payload = {"provider_reported_cost": payload, "source": "responses_sse_cost"} + if not isinstance(payload, dict): + return None + cost = payload.get("provider_reported_cost", payload.get("request_cost_usd", payload.get("total_cost", payload.get("cost", payload.get("estimated_cost"))))) + if cost is None: + return None + return { + "provider_reported_cost": cost, + "currency": payload.get("currency", "USD"), + "cost_details": payload, + } + + +def _responses_chunk_usage(chunk: dict[str, Any]) -> Any: + """Return stream chunk usage with sibling cost metadata preserved.""" + + usage = chunk.get("usage") + if not isinstance(usage, dict): + return usage + merged = dict(usage) + for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd", "currency", "costMetadata"): + if key in chunk and key not in merged: + merged[key] = deepcopy(chunk[key]) + return merged + + +def _responses_sse_cost_payload(chunk: str) -> Any: + event_type = None + data_lines: list[str] = [] + for line in chunk.strip().splitlines(): + stripped = line.strip() + if stripped.startswith(":"): + comment = stripped[1:].strip() + if comment.startswith("cost"): + return _parse_cost_text(comment[4:].strip()) + continue + if stripped.startswith("event:"): + event_type = stripped[6:].strip() + continue + if stripped.startswith("data:"): + data_lines.append(stripped[5:].strip()) + if event_type == "cost" and data_lines: + return _parse_cost_text("\n".join(data_lines).strip()) + return None + + +def _parse_cost_text(text: str) -> Any: + if not text: + return None + try: + return json.loads(text) + except json.JSONDecodeError: + try: + return float(text) + except ValueError: + return None + + +def _merge_responses_stream_usage(primary: Any, fallback_cost: Any) -> Any: + """Merge earlier stream cost metadata into later token usage when needed.""" + + if not isinstance(primary, dict): + return fallback_cost if primary is None else primary + if not isinstance(fallback_cost, dict): + return primary + merged = deepcopy(primary) + has_cost = any(key in merged for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd")) + if has_cost: + return merged + for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd", "currency"): + if key in fallback_cost: + merged[key] = deepcopy(fallback_cost[key]) + return merged + + +def _chunk_text_delta(chunk: dict[str, Any]) -> str: + choices = chunk.get("choices") if isinstance(chunk.get("choices"), list) else [] + if not choices: + return "" + delta = choices[0].get("delta") if isinstance(choices[0], dict) else None + if not isinstance(delta, dict): + return "" + content = delta.get("content") + return content if isinstance(content, str) else "" + + +def _usage_to_responses_stream(usage: Any) -> Any: + if not isinstance(usage, dict): + return usage + result = { + "input_tokens": usage.get("prompt_tokens", usage.get("input_tokens", 0)), + "output_tokens": usage.get("completion_tokens", usage.get("output_tokens", 0)), + "total_tokens": usage.get("total_tokens", 0), + } + prompt_details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") + if isinstance(prompt_details, dict): + result["input_tokens_details"] = {"cached_tokens": prompt_details.get("cached_tokens", 0)} + completion_details = usage.get("completion_tokens_details") or usage.get("output_tokens_details") + if isinstance(completion_details, dict): + result["output_tokens_details"] = {"reasoning_tokens": completion_details.get("reasoning_tokens", 0)} + for key in ("cost_details", "cost", "total_cost", "estimated_cost", "provider_reported_cost", "request_cost_usd", "currency"): + if key in usage: + result[key] = deepcopy(usage[key]) + return result diff --git a/src/rotator_library/responses/store.py b/src/rotator_library/responses/store.py new file mode 100644 index 000000000..05ee09c6b --- /dev/null +++ b/src/rotator_library/responses/store.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Storage backends for Responses API objects.""" + +from __future__ import annotations + +import hashlib +import json +from copy import deepcopy +from pathlib import Path +from typing import Any, Optional, Protocol + +from .types import StoredResponse + + +class ResponsesStore(Protocol): + """Minimal async store for response retrieval and continuation.""" + + async def save(self, response: StoredResponse) -> None: ... + + async def get( + self, + response_id: str, + scope_key: str = "public", + ) -> Optional[StoredResponse]: ... + + async def delete(self, response_id: str, scope_key: str = "public") -> bool: ... + + async def list_input_items( + self, + response_id: str, + scope_key: str = "public", + ) -> Optional[list[Any]]: ... + + async def close(self) -> None: ... + + +class InMemoryResponsesStore: + """Process-local Responses store. + + This is the Phase 4 default because it has no async lifecycle and avoids a + new persistence dependency. A provider-cache-backed store can be injected by + later configuration code when disk persistence is desired. + """ + + def __init__(self, *, max_items: int | None = None) -> None: + self._responses: dict[tuple[str, str], StoredResponse] = {} + self.max_items = max_items if max_items and max_items > 0 else None + + async def save(self, response: StoredResponse) -> None: + self._prune_expired() + key = (response.scope_key or "public", response.id) + self._responses[key] = StoredResponse.from_dict(response.to_dict()) + self._prune_overflow() + + async def get( + self, + response_id: str, + scope_key: str = "public", + ) -> Optional[StoredResponse]: + key = (scope_key, response_id) + response = self._responses.get(key) + if response is None: + return None + if response.is_expired(): + self._responses.pop(key, None) + return None + return StoredResponse.from_dict(response.to_dict()) + + async def delete(self, response_id: str, scope_key: str = "public") -> bool: + return self._responses.pop((scope_key, response_id), None) is not None + + async def list_input_items( + self, + response_id: str, + scope_key: str = "public", + ) -> Optional[list[Any]]: + response = await self.get(response_id, scope_key) + if response is None: + return None + return deepcopy(response.input_items) + + async def close(self) -> None: + """Memory storage owns no external resources.""" + + def _prune_expired(self) -> None: + for key, response in list(self._responses.items()): + if response.is_expired(): + self._responses.pop(key, None) + + def _prune_overflow(self) -> None: + if not self.max_items: + return + while len(self._responses) > self.max_items: + oldest_key = min( + self._responses, + key=lambda key: self._responses[key].created_at, + ) + self._responses.pop(oldest_key, None) + + +class ProviderCacheResponsesStore: + """Responses store backed by an injected `ProviderCache` instance. + + The wrapper does not instantiate `ProviderCache`; the factory injects it and + `close()` delegates lifecycle shutdown for its background tasks. + """ + + def __init__(self, provider_cache: Any, *, prefix: str = "responses") -> None: + self._cache = provider_cache + self._prefix = prefix + + async def save(self, response: StoredResponse) -> None: + await self._cache.store_async( + self._key(response.id, response.scope_key or "public"), + json.dumps(response.to_dict(), ensure_ascii=False), + ) + flush = getattr(self._cache, "_save_to_disk", None) + if callable(flush): + await flush() + + async def get( + self, + response_id: str, + scope_key: str = "public", + ) -> Optional[StoredResponse]: + raw = await self._cache.retrieve_async(self._key(response_id, scope_key)) + legacy_key = None + if raw is None: + legacy_key = self._legacy_key(response_id) + raw = await self._cache.retrieve_async(legacy_key) + if raw is None: + return None + response = StoredResponse.from_dict(json.loads(raw)) + response_scope = response.scope_key or "public" + if response.id != response_id or response_scope != scope_key: + return None + if legacy_key is not None: + response.scope_key = response_scope + await self.save(response) + delete = getattr(self._cache, "delete_async", None) + if delete: + await delete(legacy_key) + if response.is_expired(): + await self.delete(response_id, scope_key) + return None + return response + + async def delete(self, response_id: str, scope_key: str = "public") -> bool: + delete = getattr(self._cache, "delete_async", None) + if delete: + return bool(await delete(self._key(response_id, scope_key))) + # Third-party cache adapters may not support key-level deletion. Avoid + # clearing unrelated provider state when that capability is absent. + return False + + async def list_input_items( + self, + response_id: str, + scope_key: str = "public", + ) -> Optional[list[Any]]: + response = await self.get(response_id, scope_key) + if response is None: + return None + return deepcopy(response.input_items) + + async def close(self) -> None: + """Stop the injected cache's writer and cleanup tasks.""" + + shutdown = getattr(self._cache, "shutdown", None) + if shutdown: + await shutdown() + + def _key(self, response_id: str, scope_key: str) -> str: + digest = hashlib.sha256( + f"{scope_key}\x00{response_id}".encode("utf-8") + ).hexdigest() + return f"{self._prefix}:{digest}" + + def _legacy_key(self, response_id: str) -> str: + safe_id = response_id.replace("/", "_").replace("\\", "_").replace(":", "_") + return f"{self._prefix}:{safe_id}" + + +def create_configured_responses_store(*, config: Any = None, env: Any = None) -> ResponsesStore: + """Create the configured Responses store backend. + + Memory remains the default. The provider-cache backend uses the existing JSON + cache implementation so durable storage does not require a new database. + """ + + from ..config.experimental import get_responses_store_runtime_settings, get_responses_store_settings + from ..providers.provider_cache import create_provider_cache + + runtime = get_responses_store_runtime_settings(config=config, env=env) + settings = get_responses_store_settings(config=config, env=env) + if runtime.backend == "memory": + return InMemoryResponsesStore(max_items=settings.max_items) + cache_dir = Path(runtime.cache_dir) if runtime.cache_dir else None + provider_cache = create_provider_cache( + runtime.cache_name, + cache_dir=cache_dir, + memory_ttl_seconds=runtime.cache_memory_ttl_seconds, + disk_ttl_seconds=runtime.cache_disk_ttl_seconds, + env_prefix=f"{runtime.cache_name.upper().replace('-', '_')}_CACHE", + ) + return ProviderCacheResponsesStore(provider_cache, prefix=runtime.cache_prefix) diff --git a/src/rotator_library/responses/streaming.py b/src/rotator_library/responses/streaming.py new file mode 100644 index 000000000..f76111134 --- /dev/null +++ b/src/rotator_library/responses/streaming.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""HTTP SSE formatting for the Responses API compatibility layer.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + +from ..protocols import serialize_value + + +@dataclass(frozen=True) +class ResponsesStreamState: + """Mutable-by-replacement state accumulated from chat stream chunks.""" + + response_id: str + model: str + output_text: str = "" + output_item_id: str = "msg_0" + + +@dataclass(frozen=True) +class ResponsesStreamEvent: + """Transport-neutral Responses stream event. + + Service code yields these events. Formatters decide how to serialize them for + SSE, WebSocket, or future transports without duplicating protocol logic. + """ + + event_name: str + payload: dict[str, Any] + terminal: bool = False + + @property + def heartbeat(self) -> bool: + """Return whether this event is transport metadata, not model output.""" + + return self.event_name == "heartbeat" + + +class ResponsesSSEFormatter: + """Format Responses API events as HTTP Server-Sent Events.""" + + transport = "sse" + + def format_event(self, event_name: str, payload: dict[str, Any]) -> str: + return f"event: {event_name}\ndata: {json.dumps(serialize_value(payload), ensure_ascii=False)}\n\n" + + def format_stream_event(self, event: ResponsesStreamEvent) -> str: + """Format one transport-neutral event for HTTP SSE.""" + + if event.heartbeat: + return self.format_heartbeat(str(event.payload.get("comment") or "heartbeat")) + if event.terminal: + return self.done() + return self.format_event(event.event_name, event.payload) + + def format_heartbeat(self, comment: str = "heartbeat") -> str: + """Return a non-visible SSE comment heartbeat frame.""" + + safe_comment = comment.replace("\r", " ").replace("\n", " ") + return f": {safe_comment}\n\n" + + def done(self) -> str: + """Return the final compatibility sentinel used by many SSE clients.""" + + return "data: [DONE]\n\n" + + +class ResponsesWebSocketFormatter: + """Placeholder transport seam for future WebSocket Responses support.""" + + transport = "websocket" + future_supported = True + + def format_event(self, event_name: str, payload: dict[str, Any]) -> str: + return json.dumps({"event": event_name, "data": serialize_value(payload)}, ensure_ascii=False) + + def format_stream_event(self, event: ResponsesStreamEvent) -> str: + """Format one transport-neutral event as a WebSocket message payload.""" + + if event.heartbeat: + return json.dumps({"event": "heartbeat", "data": {"visible_output": False}}, ensure_ascii=False) + if event.terminal: + return json.dumps({"event": "done", "data": {}}, ensure_ascii=False) + return self.format_event(event.event_name, event.payload) + + +def parse_chat_sse_chunk(chunk: Any) -> dict[str, Any] | None: + """Decode a chat-completions stream chunk into a dict if possible.""" + + if isinstance(chunk, dict): + return chunk + if not isinstance(chunk, str): + return None + text = chunk.strip() + if not text: + return None + event_name = None + data_lines: list[str] = [] + for line in text.splitlines(): + if line.startswith("event:"): + event_name = line[len("event:") :].strip() + elif line.startswith("data:"): + data_lines.append(line[len("data:") :].strip()) + if data_lines: + text = "\n".join(data_lines).strip() + if text == "[DONE]": + return {"type": "done"} + try: + payload = json.loads(text) + except json.JSONDecodeError: + return None + if isinstance(payload, dict) and event_name and payload.get("type") is None: + payload["type"] = event_name + return payload if isinstance(payload, dict) else None + + +def response_created_payload(response_id: str, model: str) -> dict[str, Any]: + return {"id": response_id, "object": "response", "status": "in_progress", "model": model, "output": []} + + +def output_item_added_payload(state: ResponsesStreamState) -> dict[str, Any]: + return { + "response_id": state.response_id, + "output_index": 0, + "item": {"id": state.output_item_id, "type": "message", "role": "assistant", "content": []}, + } + + +def output_text_delta_payload(state: ResponsesStreamState, delta: str) -> dict[str, Any]: + return { + "response_id": state.response_id, + "item_id": state.output_item_id, + "output_index": 0, + "content_index": 0, + "delta": delta, + } + + +def output_item_done_payload(state: ResponsesStreamState) -> dict[str, Any]: + return { + "response_id": state.response_id, + "output_index": 0, + "item": { + "id": state.output_item_id, + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": state.output_text}], + }, + } + + +def response_completed_payload(state: ResponsesStreamState, usage: Any = None) -> dict[str, Any]: + payload = { + "id": state.response_id, + "object": "response", + "status": "completed", + "model": state.model, + "output": [output_item_done_payload(state)["item"]], + } + if usage is not None: + payload["usage"] = usage + return payload + + +def response_failed_payload(response_id: str, model: str, error: Any) -> dict[str, Any]: + return {"id": response_id, "object": "response", "status": "failed", "model": model, "error": serialize_value(error)} diff --git a/src/rotator_library/responses/types.py b/src/rotator_library/responses/types.py new file mode 100644 index 000000000..9b2a91e29 --- /dev/null +++ b/src/rotator_library/responses/types.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Data types for the Responses API compatibility layer.""" + +from __future__ import annotations + +import secrets +import time +from dataclasses import dataclass, field +from typing import Any, Optional + +from ..protocols import serialize_value + + +def generate_response_id() -> str: + """Return a local Responses-compatible identifier. + + Upstream IDs are preserved when providers return them. This helper is only + used by the bridge path when the current chat-completions backend has no + native Responses ID to expose. + """ + + return f"resp_{secrets.token_urlsafe(18).replace('-', '').replace('_', '')[:24]}" + + +@dataclass +class ResponsesStoreSettings: + """Runtime policy for Responses object storage. + + The defaults preserve Phase 4 behavior. Operators can opt into TTL, bounded + memory, failed-response persistence, or in-progress updates without changing + the store interface or introducing a database. + """ + + ttl_seconds: Optional[int] = None + max_items: Optional[int] = None + store_failed: bool = True + store_in_progress: bool = False + + +@dataclass +class StoredResponse: + """Persisted response object used for retrieval and continuation. + + The shape stores both client-facing response data and enough request/session + metadata for `previous_response_id` debugging. It deliberately avoids storing + credential secrets; callers should pass only stable identifiers if they need + credential correlation later. + """ + + id: str + model: str + status: str + response: dict[str, Any] + request: dict[str, Any] = field(default_factory=dict) + input_items: list[Any] = field(default_factory=list) + output_items: list[Any] = field(default_factory=list) + usage: Optional[dict[str, Any]] = None + metadata: dict[str, Any] = field(default_factory=dict) + created_at: float = field(default_factory=time.time) + session_id: Optional[str] = None + scope_key: Optional[str] = None + classifier: Optional[str] = None + expires_at: Optional[float] = None + + def is_expired(self, now: Optional[float] = None) -> bool: + """Return whether the response should be treated as unavailable.""" + + return self.expires_at is not None and (now if now is not None else time.time()) >= self.expires_at + + def to_dict(self) -> dict[str, Any]: + """Serialize with JSON-safe values for disk/cache persistence.""" + + return serialize_value( + { + "id": self.id, + "created_at": self.created_at, + "model": self.model, + "status": self.status, + "request": self.request, + "response": self.response, + "input_items": self.input_items, + "output_items": self.output_items, + "usage": self.usage, + "metadata": self.metadata, + "session_id": self.session_id, + "scope_key": self.scope_key, + "classifier": self.classifier, + "expires_at": self.expires_at, + } + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "StoredResponse": + """Rehydrate a stored response from a JSON-compatible dict.""" + + return cls( + id=str(data["id"]), + created_at=float(data.get("created_at") or time.time()), + model=str(data.get("model") or ""), + status=str(data.get("status") or "completed"), + request=dict(data.get("request") or {}), + response=dict(data.get("response") or {}), + input_items=list(data.get("input_items") or []), + output_items=list(data.get("output_items") or []), + usage=data.get("usage") if isinstance(data.get("usage"), dict) else None, + metadata=dict(data.get("metadata") or {}), + session_id=data.get("session_id"), + scope_key=data.get("scope_key"), + classifier=data.get("classifier"), + expires_at=data.get("expires_at"), + ) diff --git a/src/rotator_library/retry_policy.py b/src/rotator_library/retry_policy.py new file mode 100644 index 000000000..313c35c53 --- /dev/null +++ b/src/rotator_library/retry_policy.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Retry, cooldown, and target-failover policy helpers. + +This module centralizes decisions that sit above the existing error classifier. +It deliberately delegates parsing and classification to `error_handler.py` so the +proxy keeps its current retry-after parser and credential-rotation semantics. +""" + +from __future__ import annotations + +import asyncio +import os +import time +from collections import deque +from dataclasses import dataclass +from typing import Any, Optional + +from .error_handler import ClassifiedError, classify_error, should_retry_same_key, should_rotate_on_error +from .routing import FallbackPolicy +from .routing.policy import normalize_route_error_type +from .routing.types import FallbackGroup + +DEFAULT_PROVIDER_COOLDOWN_DEFAULT_SECONDS = 30 + + +@dataclass(frozen=True) +class ProviderCooldownDecision: + """Decision describing whether a provider-level cooldown should start.""" + + should_start: bool + duration: int = 0 + reason: str = "not_applicable" + scope: str = "provider" + model: Optional[str] = None + backoff_level: int = 0 + + +@dataclass(frozen=True) +class FailureHistoryEntry: + """One sanitized provider/model failure event kept in memory only.""" + + timestamp: float + provider: str + model: Optional[str] + error_type: str + scope: str + duration: int + reason: str + + +def classify_route_error(error: BaseException, provider: Optional[str] = None) -> str: + """Map an exception into the vocabulary consumed by fallback policy.""" + + if isinstance(error, asyncio.CancelledError): + return "cancelled" + explicit = getattr(error, "error_type", None) + if explicit: + return normalize_route_error_type(str(explicit)) + return normalize_route_error_type(classify_error(error, provider).error_type) + + +def should_retry_same_credential(classified_error: ClassifiedError, small_cooldown_threshold: int) -> bool: + """Return whether the current credential should be retried before rotation.""" + + return should_retry_same_key(classified_error, small_cooldown_threshold) + + +def should_rotate_credential(classified_error: ClassifiedError) -> bool: + """Return whether a classified failure should rotate to another credential.""" + + return should_rotate_on_error(classified_error) + + +def decide_provider_cooldown( + classified_error: ClassifiedError, + *, + small_cooldown_threshold: int, + provider_cooldown_min_seconds: int, + default_duration: int = DEFAULT_PROVIDER_COOLDOWN_DEFAULT_SECONDS, + cooldown_on_quota: bool = False, + provider: Optional[str] = None, + model: Optional[str] = None, + original_error: Any = None, + failure_history: "FailureHistory | None" = None, +) -> ProviderCooldownDecision: + """Return whether a provider-wide cooldown should be activated. + + Small retry-after values are intentionally left to same-credential retry to + preserve cache/session locality. Larger retry-after values can indicate a + provider-wide or IP-level throttle and are therefore safe to coordinate via + the provider cooldown manager. + """ + + error_type = classified_error.error_type + if error_type == "quota_exceeded" and not cooldown_on_quota: + return ProviderCooldownDecision(False, reason="quota_cooldown_disabled") + if error_type not in {"rate_limit", "server_error", "api_connection", "quota_exceeded"}: + return ProviderCooldownDecision(False, reason="non_provider_cooldown_error") + scope = "model" if model and is_model_capacity_error(original_error or classified_error.original_exception) else "provider" + + retry_after = classified_error.retry_after + if retry_after is not None: + if retry_after <= 0: + return ProviderCooldownDecision(False, reason="non_positive_retry_after") + if retry_after < small_cooldown_threshold: + return ProviderCooldownDecision(False, reason="small_retry_after") + if retry_after < provider_cooldown_min_seconds: + return ProviderCooldownDecision(False, reason="below_provider_cooldown_minimum") + return ProviderCooldownDecision(True, duration=int(retry_after), reason="retry_after", scope=scope, model=model if scope == "model" else None) + + if error_type in {"server_error", "api_connection"} and default_duration >= provider_cooldown_min_seconds: + backoff_level = 0 + duration = int(default_duration) + if scope == "model": + return ProviderCooldownDecision(True, duration=duration, reason="model_capacity_cooldown", scope=scope, model=model, backoff_level=backoff_level) + if failure_history is None: + return ProviderCooldownDecision(False, reason="missing_failure_history", scope=scope, model=model if scope == "model" else None) + backoff = failure_history.backoff_for(provider=provider, error_type=error_type, scope=scope, model=model if scope == "model" else None, default_duration=duration) + duration = backoff.duration + backoff_level = backoff.level + if backoff_level <= 0: + return ProviderCooldownDecision(False, reason="transient_backoff_threshold_not_met", scope=scope, model=model if scope == "model" else None) + return ProviderCooldownDecision(True, duration=duration, reason="model_capacity_cooldown" if scope == "model" else "default_transient_cooldown", scope=scope, model=model if scope == "model" else None, backoff_level=backoff_level) + return ProviderCooldownDecision(False, reason="missing_retry_after") + + +@dataclass(frozen=True) +class BackoffDecision: + """Bounded backoff duration derived from recent transient failures.""" + + duration: int + level: int = 0 + + +class FailureHistory: + """Bounded in-memory provider/model failure history. + + This is intentionally process-local. It provides enough recent context for + conservative cooldown backoff and future observability without introducing a + persistence layer or changing credential usage accounting. + """ + + def __init__(self, *, max_entries: int | None = None, clock: Any = None) -> None: + settings = _retry_settings() + self.max_entries = max(1, max_entries if max_entries is not None else settings.failure_history_max_entries) + self._entries: deque[FailureHistoryEntry] = deque(maxlen=self.max_entries) + self._clock = clock or time.time + + def record(self, *, provider: str, model: Optional[str], error_type: str, scope: str, duration: int, reason: str) -> None: + """Record one sanitized cooldown/failure event.""" + + self._entries.append( + FailureHistoryEntry( + timestamp=float(self._clock()), + provider=provider, + model=model, + error_type=error_type, + scope=scope, + duration=duration, + reason=reason, + ) + ) + + def snapshot(self) -> tuple[FailureHistoryEntry, ...]: + """Return recent entries for tests and future read-only reporting.""" + + return tuple(self._entries) + + def clear(self, *, provider: str, model: Optional[str] = None, scope: Optional[str] = None, error_type: Optional[str] = None) -> None: + """Clear matching failure entries after a successful provider/model call.""" + + kept = [ + entry + for entry in self._entries + if not ( + entry.provider == provider + and (scope is None or entry.scope == scope) + and (error_type is None or entry.error_type == error_type) + and (entry.scope != "model" or model is None or entry.model == model) + ) + ] + self._entries.clear() + self._entries.extend(kept) + + def backoff_for(self, *, provider: Optional[str], error_type: str, scope: str, model: Optional[str], default_duration: int) -> BackoffDecision: + """Return bounded backoff for repeated transient failures.""" + + settings = _retry_settings() + window = settings.provider_backoff_window_seconds + threshold = max(1, settings.provider_backoff_threshold) + base = max(1, settings.provider_backoff_base_seconds or default_duration) + max_seconds = max(base, settings.provider_backoff_max_seconds) + now = float(self._clock()) + recent = [ + entry + for entry in self._entries + if now - entry.timestamp <= window + and entry.provider == provider + and entry.error_type == error_type + and entry.scope == scope + and (scope != "model" or entry.model == model) + ] + if len(recent) + 1 < threshold: + return BackoffDecision(default_duration, level=0) + level = len(recent) + 1 - threshold + 1 + return BackoffDecision(min(max_seconds, base * (2 ** (level - 1))), level=level) + + +def is_model_capacity_error(error: Any) -> bool: + """Return whether an error indicates model/deployment capacity exhaustion.""" + + if error is None: + return False + if isinstance(error, dict): + text = str(error).lower() + else: + parts = [str(error)] + response = getattr(error, "response", None) + if response is not None: + parts.append(str(getattr(response, "text", ""))) + body = getattr(error, "body", None) + if body is not None: + parts.append(str(body)) + text = " ".join(parts).lower() + return "model_capacity_exhausted" in text or "model capacity" in text or "capacity exhausted" in text + + +def provider_cooldown_env() -> tuple[int, int, bool]: + """Read provider-cooldown env controls with conservative defaults.""" + + settings = _retry_settings() + return settings.provider_cooldown_min_seconds, settings.provider_cooldown_default_seconds, settings.provider_cooldown_on_quota + + +def is_target_failover_eligible( + error_type: str, + *, + group: FallbackGroup | None = None, + stream: bool = False, + emitted_output: bool = False, +) -> bool: + """Return whether a target failure may advance to the next target.""" + + return FallbackPolicy().should_fallback(error_type, group=group, stream=stream, emitted_output=emitted_output) + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + + +def _retry_settings() -> Any: + """Load retry settings lazily to avoid config import cycles at startup.""" + + from .config.experimental import get_retry_runtime_settings + + return get_retry_runtime_settings() diff --git a/src/rotator_library/routing/__init__.py b/src/rotator_library/routing/__init__.py new file mode 100644 index 000000000..9892616bc --- /dev/null +++ b/src/rotator_library/routing/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Routing and fallback group primitives.""" + +from .config import RoutingConfigError, load_routing_config_from_env, parse_route_target +from .attempts import clone_context_for_target +from .executor import FallbackAttemptRunner, FallbackExhaustedError +from .policy import FallbackPolicy +from .resolver import FallbackResolver +from .types import FallbackGroup, RouteTarget, RoutingConfig, RoutingDecision, TargetGroup, TargetSelector + +__all__ = [ + "FallbackGroup", + "FallbackAttemptRunner", + "FallbackExhaustedError", + "FallbackPolicy", + "FallbackResolver", + "RouteTarget", + "RoutingConfig", + "RoutingConfigError", + "RoutingDecision", + "TargetGroup", + "TargetSelector", + "clone_context_for_target", + "load_routing_config_from_env", + "parse_route_target", +] diff --git a/src/rotator_library/routing/attempts.py b/src/rotator_library/routing/attempts.py new file mode 100644 index 000000000..a767c6ee6 --- /dev/null +++ b/src/rotator_library/routing/attempts.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Helpers for creating per-target request contexts.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Any, Sequence + +from ..core.types import RequestContext +from .types import RouteTarget + + +def clone_context_for_target( + context: RequestContext, + target: RouteTarget, + *, + credentials: Sequence[str] | None = None, + usage_manager_key: str | None = None, + provider_config: dict[str, Any] | None = None, + credential_secrets: dict[str, str] | None = None, + target_index: int = 0, +) -> RequestContext: + """Return a target-specific request context without mutating the original. + + Fallback routing must preserve the original request context for traceability + and for safe retry decisions. This helper copies the request kwargs, updates + the selected model/provider, and carries routing metadata for executor traces. + """ + + kwargs: dict[str, Any] = dict(context.kwargs) + kwargs["model"] = target.prefixed_model + next_usage_key = usage_manager_key if usage_manager_key is not None else target.provider + return replace( + context, + model=target.prefixed_model, + provider=target.provider, + kwargs=kwargs, + credentials=list(credentials) if credentials is not None else list(context.credentials), + usage_manager_key=next_usage_key, + provider_config=provider_config if provider_config is not None else context.provider_config, + credential_secrets=dict(credential_secrets) if credential_secrets is not None else dict(context.credential_secrets), + routing_target_index=target_index, + session_affinity_key=( + context.session_affinity_key + if target.provider == context.provider and target.prefixed_model == context.model + else None + ), + ) diff --git a/src/rotator_library/routing/config.py b/src/rotator_library/routing/config.py new file mode 100644 index 000000000..473bdc23f --- /dev/null +++ b/src/rotator_library/routing/config.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Environment parser for Phase 6 fallback routing configuration.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping + +from .policy import normalize_route_error_set +from .types import DEFAULT_FAILOVER_ON, DEFAULT_STOP_ON, HARD_STOP_ON, FallbackGroup, RouteTarget, RoutingConfig + + +class RoutingConfigError(ValueError): + """Raised when fallback routing configuration is invalid.""" + + +def parse_route_target(spec: str) -> RouteTarget: + """Parse `provider/model` with optional `@execution` suffix. + + The suffix is intentionally small because Phase 6 prioritizes ordered + fallback groups. Rich selector syntax belongs to the config polish phase. + """ + + text = spec.strip() + if not text: + raise RoutingConfigError("route target cannot be empty") + target_text, _, execution = text.partition("@") + if "/" not in target_text: + raise RoutingConfigError(f"route target requires provider/model: {spec}") + provider, model = target_text.split("/", 1) + if not provider or not model: + raise RoutingConfigError(f"route target requires provider/model: {spec}") + return RouteTarget(provider=provider.strip(), model=model.strip(), execution=(execution.strip() or "auto")) + + +def load_routing_config_from_env(env: Mapping[str, str] | None = None, config: object | None = None) -> RoutingConfig: + """Load fallback groups and model-route aliases from JSON then environment. + + Environment variables intentionally remain the final override layer so the + existing `.env` deployment model keeps working exactly as before. The + optional JSON object is a convenience for structured routing plans. + """ + + source = env if env is not None else os.environ + if config is None: + from ..config.experimental import load_experimental_config + + config = load_experimental_config(env=source) + groups: dict[str, FallbackGroup] = _groups_from_json_config(config) + + group_names = _csv(source.get("FALLBACK_GROUPS", "")) + if len(group_names) != len(set(group_names)): + raise RoutingConfigError("fallback group names must be unique") + for name in group_names: + key = f"FALLBACK_GROUP_{_env_key(name)}" + target_specs = _csv(source.get(key, "")) + if not target_specs: + raise RoutingConfigError(f"fallback group {name} has no targets") + failover_on = _policy_set(source.get(f"{key}_FAILOVER_ON"), DEFAULT_FAILOVER_ON, name) + stop_on = _policy_set(source.get(f"{key}_STOP_ON"), DEFAULT_STOP_ON, name, allow_hard_stop=True) + streaming_policy = _streaming_policy(source.get(f"{key}_STREAMING_POLICY", "pre_output_only"), name) + groups[name] = FallbackGroup(name=name, targets=tuple(parse_route_target(spec) for spec in target_specs), failover_on=failover_on, stop_on=stop_on, streaming_policy=streaming_policy) + + model_routes: dict[str, str] = _model_routes_from_json_config(config) + for key, value in source.items(): + if not key.startswith("MODEL_ROUTE_"): + continue + model_alias = key[len("MODEL_ROUTE_") :].lower() + route = value.strip() + model_routes[model_alias] = route + _validate_model_routes(model_routes, groups) + return RoutingConfig(fallback_groups=groups, model_routes=model_routes) + + +def _groups_from_json_config(config: object) -> dict[str, FallbackGroup]: + routing = getattr(config, "routing", {}) + if not isinstance(routing, Mapping): + return {} + raw_groups = routing.get("fallback_groups", {}) + if not isinstance(raw_groups, Mapping): + raise RoutingConfigError("routing.fallback_groups must be an object") + groups: dict[str, FallbackGroup] = {} + for name, raw_group in raw_groups.items(): + if not isinstance(raw_group, Mapping): + raise RoutingConfigError(f"fallback group {name} must be an object") + raw_targets = raw_group.get("targets", []) + if not isinstance(raw_targets, list) or not raw_targets: + raise RoutingConfigError(f"fallback group {name} has no targets") + groups[str(name)] = FallbackGroup( + name=str(name), + targets=tuple(parse_route_target(str(spec)) for spec in raw_targets), + failover_on=_policy_set(raw_group.get("failover_on"), DEFAULT_FAILOVER_ON, str(name)), + stop_on=_policy_set(raw_group.get("stop_on"), DEFAULT_STOP_ON, str(name), allow_hard_stop=True), + streaming_policy=_streaming_policy(raw_group.get("streaming_policy", "pre_output_only"), str(name)), + max_targets=int(raw_group["max_targets"]) if raw_group.get("max_targets") is not None else None, + metadata=dict(raw_group.get("metadata", {})) if isinstance(raw_group.get("metadata", {}), Mapping) else {}, + ) + return groups + + +def _model_routes_from_json_config(config: object) -> dict[str, str]: + routing = getattr(config, "routing", {}) + if not isinstance(routing, Mapping): + return {} + raw_routes = routing.get("model_routes", {}) + if not isinstance(raw_routes, Mapping): + raise RoutingConfigError("routing.model_routes must be an object") + return {str(alias).lower(): str(route).strip() for alias, route in raw_routes.items()} + + +def _validate_model_routes(model_routes: Mapping[str, str], groups: Mapping[str, FallbackGroup]) -> None: + for key, route in model_routes.items(): + if route.startswith("group:") and route[len("group:") :] not in groups: + raise RoutingConfigError(f"model route {key} references unknown fallback group {route}") + if not route.startswith("group:"): + parse_route_target(route) + + +def _string_set(value: object, default: frozenset[str]) -> frozenset[str]: + if value is None: + return default + if isinstance(value, str): + return frozenset(_csv(value)) + if isinstance(value, list): + return frozenset(str(item) for item in value) + raise RoutingConfigError("routing policy lists must be strings or arrays") + + +def _policy_set(value: object, default: frozenset[str], group_name: str, *, allow_hard_stop: bool = False) -> frozenset[str]: + values = normalize_route_error_set(_string_set(value, default)) + unsafe = values & HARD_STOP_ON + if unsafe and not allow_hard_stop: + raise RoutingConfigError(f"fallback group {group_name} failover_on cannot include hard-stop errors: {', '.join(sorted(unsafe))}") + return values + + +def _streaming_policy(value: object, group_name: str): + policy = str(value or "pre_output_only").strip().lower() + if policy not in {"pre_output_only", "never"}: + raise RoutingConfigError(f"fallback group {group_name} has unsupported streaming_policy {value!r}") + return policy + + +def _csv(value: str) -> list[str]: + return [part.strip() for part in value.split(",") if part.strip()] + + +def _env_key(value: str) -> str: + return value.upper().replace("-", "_") diff --git a/src/rotator_library/routing/executor.py b/src/rotator_library/routing/executor.py new file mode 100644 index 000000000..dce897e9c --- /dev/null +++ b/src/rotator_library/routing/executor.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Ordered fallback attempt runner.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any + +from .policy import FallbackPolicy, normalize_route_error_type +from .types import FallbackGroup, RouteAttemptResult, RouteTarget, RoutingDecision + +AttemptCallback = Callable[[RouteTarget, int], Awaitable[Any]] + + +class FallbackExhaustedError(Exception): + """Raised when every eligible fallback target fails.""" + + def __init__(self, decision: RoutingDecision, attempts: tuple[RouteAttemptResult, ...]) -> None: + self.decision = decision + self.attempts = attempts + summary = ", ".join(f"{attempt.target.name}:{attempt.error_type}" for attempt in attempts) + super().__init__(f"fallback group exhausted for {decision.requested_model}: {summary}") + + +class FallbackAttemptRunner: + """Run ordered route targets using an injected per-target attempt callback. + + The runner is independent from `RequestExecutor`. Phase 6 can unit-test + fallback control flow here, then wire the callback to existing per-target + credential retry logic without duplicating policy decisions. + """ + + def __init__(self, policy: FallbackPolicy | None = None) -> None: + self.policy = policy or FallbackPolicy() + + async def run(self, decision: RoutingDecision, attempt: AttemptCallback, *, stream: bool = False) -> Any: + """Try targets in order until one succeeds or fallback is exhausted.""" + + return await self.run_group(decision, decision.group, attempt, stream=stream) + + async def run_group( + self, + decision: RoutingDecision, + group: FallbackGroup | None, + attempt: AttemptCallback, + *, + stream: bool = False, + ) -> Any: + """Try targets while honoring optional group-specific policy overrides.""" + + attempts: list[RouteAttemptResult] = [] + for index, target in enumerate(decision.targets): + try: + return await attempt(target, index) + except Exception as exc: + error_type = _error_type(exc) + emitted_output = bool(getattr(exc, "emitted_output", False)) + attempts.append(RouteAttemptResult(target=target, success=False, error_type=error_type, emitted_output=emitted_output)) + has_next = index < len(decision.targets) - 1 + if not has_next or (stream and getattr(group, "streaming_policy", "pre_output_only") == "never") or not self.policy.should_fallback(error_type, group=group, emitted_output=emitted_output, stream=stream): + raise FallbackExhaustedError(decision, tuple(attempts)) from exc + raise FallbackExhaustedError(decision, tuple(attempts)) + + +def _error_type(error: BaseException) -> str: + return normalize_route_error_type(str(getattr(error, "error_type", error.__class__.__name__))) diff --git a/src/rotator_library/routing/policy.py b/src/rotator_library/routing/policy.py new file mode 100644 index 000000000..549ee114a --- /dev/null +++ b/src/rotator_library/routing/policy.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Fallback eligibility policy for ordered route chains.""" + +from __future__ import annotations + +from .types import DEFAULT_FAILOVER_ON, DEFAULT_STOP_ON, HARD_STOP_ON, FallbackGroup + + +_ALIASES = { + "auth": "authentication", + "invalid_api_key": "authentication", + "invalid_key": "authentication", + "unauthorized": "authentication", + "permission": "forbidden", + "permission_denied": "forbidden", + "access_denied": "forbidden", + "bad_request": "invalid_request", + "invalid_argument": "invalid_request", + "bad_argument": "invalid_request", + "validation": "invalid_request", + "permanent": "invalid_request", + "context_length": "context_window_exceeded", + "context_length_exceeded": "context_window_exceeded", + "context_window": "context_window_exceeded", + "too_many_tokens": "context_window_exceeded", + "max_tokens_exceeded": "context_window_exceeded", + "token_limit_exceeded": "context_window_exceeded", + "pre_request_callback": "pre_request_callback_error", + "configuration": "configuration_error", + "config": "configuration_error", + "quota": "quota_exceeded", + "resource_exhausted": "quota_exceeded", + "capacity": "rate_limit", + "rate_limited": "rate_limit", + "too_many_requests": "rate_limit", + "transient": "server_error", + "unavailable": "server_error", + "network": "api_connection", + "connection": "api_connection", + "deadline_exceeded": "api_connection", + "timeout": "api_connection", + "400": "invalid_request", + "401": "authentication", + "403": "forbidden", + "429": "rate_limit", + "500": "server_error", + "502": "server_error", + "503": "server_error", + "504": "server_error", +} + + +def normalize_route_error_type(error_type: str | None) -> str: + """Return the policy vocabulary for a raw classifier or config value. + + Routing config is user-facing, while the executor consumes classifier output. + Keeping normalization in one place prevents aliases such as `auth` from + bypassing non-overridable hard-stop categories. + """ + + normalized = str(error_type or "").strip().lower().replace("-", "_").replace(" ", "_") + return _ALIASES.get(normalized, normalized) + + +def normalize_route_error_set(values: frozenset[str]) -> frozenset[str]: + """Normalize a policy set while preserving unknown custom categories.""" + + return frozenset(normalize_route_error_type(value) for value in values) + + +class FallbackPolicy: + """Decide whether a failed target may advance to the next target.""" + + def should_fallback( + self, + error_type: str, + *, + group: FallbackGroup | None = None, + emitted_output: bool = False, + stream: bool = False, + ) -> bool: + """Return whether fallback is allowed for a classified failure.""" + + normalized = normalize_route_error_type(error_type) + if stream and emitted_output: + return False + if normalized in HARD_STOP_ON: + return False + active_stop = normalize_route_error_set(group.stop_on if group else DEFAULT_STOP_ON) + active_failover = normalize_route_error_set(group.failover_on if group else DEFAULT_FAILOVER_ON) + if normalized in active_stop: + return False + return normalized in active_failover diff --git a/src/rotator_library/routing/resolver.py b/src/rotator_library/routing/resolver.py new file mode 100644 index 000000000..53b66d58b --- /dev/null +++ b/src/rotator_library/routing/resolver.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Resolve requested models to direct targets or fallback groups.""" + +from __future__ import annotations + +from .config import RoutingConfigError, parse_route_target +from .types import FallbackGroup, RouteTarget, RoutingConfig, RoutingDecision + + +class FallbackResolver: + """Resolve a model name using deterministic fallback group rules.""" + + def __init__(self, config: RoutingConfig) -> None: + self.config = config + + def resolve(self, requested_model: str) -> RoutingDecision: + """Return the ordered targets for a requested model.""" + + route = self.config.model_routes.get(requested_model.lower()) + if route and route.startswith("group:"): + group_name = route[len("group:") :] + group = self.config.fallback_groups.get(group_name) + if not group: + raise RoutingConfigError(f"unknown fallback group {group_name}") + targets = _promote_requested_target(group, requested_model) + reason = "model_route_group_promoted" if targets != group.targets else "model_route_group" + return RoutingDecision(requested_model=requested_model, group_name=group.name, group=group, targets=targets, reason=reason) + if route: + return RoutingDecision(requested_model=requested_model, targets=(parse_route_target(route),), reason="model_route_target") + for group in self.config.fallback_groups.values(): + targets = _promote_requested_target(group, requested_model) + if targets != group.targets or any(_same_target(target, requested_model) for target in group.targets): + return RoutingDecision(requested_model=requested_model, group_name=group.name, group=group, targets=targets, reason="provider_model_group_promoted") + if "/" in requested_model: + return RoutingDecision(requested_model=requested_model, targets=(parse_route_target(requested_model),), reason="direct_provider_model") + raise RoutingConfigError(f"model {requested_model!r} is not provider-prefixed and has no route") + + +def _promote_requested_target(group: FallbackGroup, requested_model: str) -> tuple[RouteTarget, ...]: + """Return group targets with the requested provider/model attempted first.""" + + matching = [target for target in group.targets if _same_target(target, requested_model)] + if not matching: + return group.targets + selected = matching[0] + return (selected, *(target for target in group.targets if target is not selected)) + + +def _same_target(target: RouteTarget, requested_model: str) -> bool: + return target.prefixed_model.lower() == requested_model.lower() diff --git a/src/rotator_library/routing/types.py b/src/rotator_library/routing/types.py new file mode 100644 index 000000000..b0598e00b --- /dev/null +++ b/src/rotator_library/routing/types.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Typed route targets and fallback groups.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal, Protocol + +ExecutionMode = Literal["auto", "native", "custom", "litellm_fallback"] +StreamingFallbackPolicy = Literal["pre_output_only", "never"] + +DEFAULT_FAILOVER_ON = frozenset( + { + "rate_limit", + "quota_exceeded", + "server_error", + "api_connection", + "unsupported_operation", + # Human-friendly aliases for config files; classifier output uses the + # names above, but config authors should not need to know every internal + # error string. + "quota", + "capacity", + "transient", + } +) +DEFAULT_STOP_ON = frozenset( + { + "authentication", + "forbidden", + "invalid_request", + "context_window_exceeded", + "credential_reauth_needed", + "pre_request_callback_error", + "cancelled", + # Config aliases retained for readability. + "auth", + "validation", + "permanent", + "pre_request_callback", + } +) +HARD_STOP_ON = frozenset( + { + "authentication", + "forbidden", + "invalid_request", + "context_window_exceeded", + "credential_reauth_needed", + "pre_request_callback_error", + "cancelled", + "configuration_error", + } +) + + +@dataclass(frozen=True) +class RouteTarget: + """One concrete provider/model execution target in a fallback chain.""" + + provider: str + model: str + name: str = "" + protocol: str | None = None + execution: ExecutionMode = "auto" + priority: int | None = None + weight: float | None = None + conditions: dict[str, Any] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.provider or not self.model: + raise ValueError("route targets require provider and model") + if self.execution not in {"auto", "native", "custom", "litellm_fallback"}: + raise ValueError(f"unsupported execution mode: {self.execution}") + if not self.name: + object.__setattr__(self, "name", f"{self.provider}/{self.model}") + + @property + def prefixed_model(self) -> str: + """Return `provider/model` without double-prefixing an already-prefixed model.""" + + return self.model if self.model.startswith(f"{self.provider}/") else f"{self.provider}/{self.model}" + + +@dataclass(frozen=True) +class FallbackGroup: + """Deterministic ordered chain of route targets.""" + + name: str + targets: tuple[RouteTarget, ...] + failover_on: frozenset[str] = DEFAULT_FAILOVER_ON + stop_on: frozenset[str] = DEFAULT_STOP_ON + streaming_policy: StreamingFallbackPolicy = "pre_output_only" + max_targets: int | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("fallback group name is required") + if not self.targets: + raise ValueError("fallback groups require at least one target") + if self.max_targets is not None and self.max_targets <= 0: + raise ValueError("max_targets must be positive") + if self.max_targets is not None and len(self.targets) > self.max_targets: + raise ValueError("fallback group target count exceeds max_targets") + + +@dataclass(frozen=True) +class RoutingConfig: + """Routing configuration loaded from env or tests.""" + + fallback_groups: dict[str, FallbackGroup] = field(default_factory=dict) + model_routes: dict[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class RoutingDecision: + """Resolved routing plan for a requested model.""" + + requested_model: str + targets: tuple[RouteTarget, ...] + group_name: str | None = None + group: FallbackGroup | None = None + selected_target_index: int = 0 + reason: str = "direct" + + +@dataclass(frozen=True) +class RouteAttemptResult: + """Result summary for one attempted route target.""" + + target: RouteTarget + success: bool + error_type: str | None = None + emitted_output: bool = False + usage: dict[str, Any] = field(default_factory=dict) + + +class TargetSelector(Protocol): + """Future target-group selector seam; Phase 6 keeps ordered fallback only.""" + + def select(self, targets: tuple[RouteTarget, ...]) -> RouteTarget: + """Return one target from a richer target group.""" + + +@dataclass(frozen=True) +class TargetGroup: + """Future richer target group; not used for Phase 6 ordered fallback.""" + + name: str + targets: tuple[RouteTarget, ...] + selector: str = "ordered" diff --git a/src/rotator_library/session_tracking.py b/src/rotator_library/session_tracking.py index c30b3b738..c4c8c5a64 100644 --- a/src/rotator_library/session_tracking.py +++ b/src/rotator_library/session_tracking.py @@ -25,11 +25,13 @@ import hashlib import json import logging +import math import os import re import threading import time import uuid +from collections import Counter from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, Iterable, List, Optional @@ -59,14 +61,19 @@ class SessionTrackingHints: """Provider-supplied tracking evidence. Providers should return evidence, not routing decisions. Core routing still - owns credential selection. This lets providers expose stable native markers - or provider-specific request structure while keeping sticky policy centralized. + owns credential selection. Provider anchors and affinity are automatically + qualified by provider and optional ``session_scope``; global anchors are + reserved for proxy-owned/client-global identity such as Responses IDs. """ strong_anchors: List[str] = field(default_factory=list) medium_anchors: List[str] = field(default_factory=list) weak_anchors: List[str] = field(default_factory=list) + global_strong_anchors: List[str] = field(default_factory=list) + global_medium_anchors: List[str] = field(default_factory=list) + global_weak_anchors: List[str] = field(default_factory=list) affinity_key: Optional[str] = None + # Partitions provider-native anchors only; it never changes logical identity. session_scope: Optional[str] = None @@ -96,6 +103,8 @@ class _SessionState: affinity_key: Optional[str] = None anchors: set[str] = field(default_factory=set) last_seen: float = 0.0 + history_signatures: tuple[str, ...] = field(default_factory=tuple) + loaded_from_persistence: bool = False @dataclass @@ -119,6 +128,10 @@ class _MatchCandidate: medium_groups: set[str] = field(default_factory=set) provider_matches: int = 0 response_matches: int = 0 + response_groups: set[str] = field(default_factory=set) + request_groups: set[str] = field(default_factory=set) + matched_probe_groups: set[str] = field(default_factory=set) + response_probe_groups: set[str] = field(default_factory=set) last_seen: float = 0.0 @property @@ -145,6 +158,21 @@ def is_sticky_match(self) -> bool: return self.confidence in {"strong", "probable"} +@dataclass(frozen=True) +class _CompactionDecision: + """Validated structural replacement of a known parent history.""" + + parent_session_id: Optional[str] = None + marker_compaction: bool = False + retained_history_ratio: Optional[float] = None + response_group_count: int = 0 + context_probe_groups: frozenset[str] = frozenset() + + @property + def possible_compaction(self) -> bool: + return self.parent_session_id is not None + + class SessionTracker: """TTL-based session inference with scoped, compounding evidence anchors. @@ -158,7 +186,27 @@ class SessionTracker: _STRONG_SCORE = 100 _MEDIUM_SCORE = 35 _WEAK_SCORE = 5 - _PERSISTENCE_SCHEMA_VERSION = 1 + _PERSISTENCE_SCHEMA_VERSION = 3 + _COMPACTION_MAX_RETAINED_HISTORY_RATIO = 0.5 + _MIN_UNMARKED_RESPONSE_GROUPS = 2 + _COMPACTION_PROBE_ROLES = {"user", "system", "developer"} + _MAX_PERSISTED_HISTORY_SIGNATURES = 4096 + _MAX_PERSISTED_FILE_BYTES = 16 * 1024 * 1024 + _MAX_PERSISTED_SESSIONS = 10000 + _MAX_PERSISTED_STRING_LENGTH = 1024 + _PERSISTED_ANCHOR_SOURCES = { + "compaction_context", + "compaction_replay", + "explicit", + "first_user", + "global_hint", + "message", + "provider", + "response", + "tool", + "tool_event", + "window", + } def __init__( self, @@ -209,6 +257,7 @@ def infer_session( model: Optional[str] = None, scope_key: Optional[str] = None, hints: Optional[Any] = None, + _trusted_isolation_key: bool = False, ) -> SessionInference: """Infer live session and deterministic affinity from a request payload.""" with self._lock: @@ -218,6 +267,7 @@ def infer_session( model=model, scope_key=scope_key, hints=hints, + trusted_isolation_key=_trusted_isolation_key, ) save_job = self._prepare_save_locked() self._write_save_job(save_job) @@ -231,6 +281,7 @@ def _infer_session_locked( model: Optional[str], scope_key: Optional[str], hints: Optional[Any], + trusted_isolation_key: bool, ) -> SessionInference: now = time.time() self._prune(now) @@ -238,32 +289,205 @@ def _infer_session_locked( hints = self._coerce_hints(hints) namespace = self._namespace( provider, - model, scope_key=scope_key, - session_scope=hints.session_scope if hints else None, + trusted_isolation_key=trusted_isolation_key, + ) + history_signatures = self._request_history_signatures(request_data) + probe_indexes = self._compaction_probe_indexes(request_data) + compaction_probe_anchors = self._build_compaction_probe_anchors( + request_data, + namespace, + probe_indexes=probe_indexes, + ) + unsuppressed_normal_anchors = self._build_anchors( + request_data, + namespace, + hints, + provider=provider, ) - compaction_probe_anchors = self._build_compaction_probe_anchors(request_data, namespace) compaction_match = ( self._best_match(compaction_probe_anchors, namespace, now) if compaction_probe_anchors else None ) - marker_compaction = self._looks_like_compaction(request_data) - possible_compaction = marker_compaction or self._is_compaction_parent_match( + marker_probe_groups = self._compaction_marker_probe_groups(request_data) + marker_compaction = bool(marker_probe_groups) + compaction = self._evaluate_compaction( compaction_match, marker_compaction=marker_compaction, + marker_probe_groups=marker_probe_groups, + history_signatures=history_signatures, ) - normal_anchors = self._build_anchors( - request_data, - namespace, - hints, - suppressed_continuity_indexes=( - self._compaction_probe_indexes(request_data) if possible_compaction else None - ), + authoritative_anchors = [ + anchor + for anchor in unsuppressed_normal_anchors + if self._is_authoritative_identity_anchor(anchor) + ] + authoritative_match = ( + self._best_match(authoritative_anchors, namespace, now) + if authoritative_anchors + else None + ) + if authoritative_match and authoritative_match.is_sticky_match: + possible_compaction = ( + compaction.possible_compaction + and compaction.parent_session_id == authoritative_match.session_id + ) + normal_anchors = ( + self._build_anchors( + request_data, + namespace, + hints, + provider=provider, + suppressed_continuity_indexes=probe_indexes, + ) + if possible_compaction + else unsuppressed_normal_anchors + ) + state = self._refresh_and_bridge( + authoritative_match.session_id, + namespace, + normal_anchors, + now, + affinity_key=self._affinity_from_anchors(normal_anchors, namespace), + history_signatures=history_signatures, + ) + return self._log_inference_decision( + SessionInference( + session_id=state.session_id, + affinity_key=self._effective_affinity(state, hints, provider), + confidence=authoritative_match.confidence, + match_score=authoritative_match.score, + possible_compaction=possible_compaction, + tracking_namespace=namespace, + ), + action="compaction_continue" if possible_compaction else "continue", + matched_session_id=state.session_id, + compaction=compaction if possible_compaction else None, + provider=provider, + model=model, + ) + + has_authoritative_identity = bool(authoritative_anchors) + context_binding = ( + None + if has_authoritative_identity + else self._find_compaction_context( + compaction_probe_anchors, + namespace, + now, + ) + ) + replay_record = ( + None + if has_authoritative_identity + else self._find_compaction_replay( + compaction_probe_anchors, + history_signatures, + namespace, + now, + ) + ) + if replay_record: + replay_anchor = self._compaction_replay_anchor( + compaction_probe_anchors, + history_signatures, + namespace, + parent_session_id=replay_record.group, + ) + normal_anchors = self._build_anchors( + request_data, + namespace, + hints, + provider=provider, + suppressed_continuity_indexes=probe_indexes, + ) + context_anchors = ( + [context_binding[1]] + if context_binding and context_binding[0].session_id == replay_record.session_id + else [] + ) + normal_anchors = self._dedupe_anchors( + [*normal_anchors, *context_anchors, replay_anchor] + ) + state = self._refresh_and_bridge( + replay_record.session_id, + namespace, + normal_anchors, + now, + affinity_key=self._affinity_from_anchors(normal_anchors, namespace), + history_signatures=history_signatures, + ) + return self._log_inference_decision( + SessionInference( + session_id=state.session_id, + affinity_key=self._effective_affinity(state, hints, provider), + confidence="strong", + match_score=self._STRONG_SCORE, + possible_compaction=True, + lineage_parent_session_id=replay_record.group, + tracking_namespace=namespace, + ), + action="compaction_replay", + matched_session_id=state.session_id, + provider=provider, + model=model, + ) + + if context_binding: + context_record, context_anchor = context_binding + normal_anchors = self._build_anchors( + request_data, + namespace, + hints, + provider=provider, + suppressed_continuity_indexes=probe_indexes, + ) + normal_anchors = self._dedupe_anchors([*normal_anchors, context_anchor]) + state = self._refresh_and_bridge( + context_record.session_id, + namespace, + normal_anchors, + now, + affinity_key=self._affinity_from_anchors(normal_anchors, namespace), + history_signatures=history_signatures, + ) + return self._log_inference_decision( + SessionInference( + session_id=state.session_id, + affinity_key=self._effective_affinity(state, hints, provider), + confidence="strong", + match_score=self._STRONG_SCORE, + possible_compaction=False, + lineage_parent_session_id=context_record.group, + tracking_namespace=namespace, + ), + action="compaction_continue", + matched_session_id=state.session_id, + provider=provider, + model=model, + ) + + possible_compaction = compaction.possible_compaction + normal_anchors = ( + self._build_anchors( + request_data, + namespace, + hints, + provider=provider, + suppressed_continuity_indexes=probe_indexes, + ) + if possible_compaction + else unsuppressed_normal_anchors ) if not normal_anchors and not compaction_probe_anchors: - return SessionInference(session_id=None, tracking_namespace=namespace) + return self._log_inference_decision( + SessionInference(session_id=None, tracking_namespace=namespace), + action="untracked", + provider=provider, + model=model, + ) match = self._best_match(normal_anchors, namespace, now) if normal_anchors else None @@ -278,24 +502,50 @@ def _infer_session_locked( normal_anchors, now, affinity_key=self._affinity_from_anchors(normal_anchors, namespace), + history_signatures=history_signatures, ) - return SessionInference( - session_id=state.session_id, - affinity_key=state.affinity_key, - confidence=match.confidence, - match_score=match.score, - possible_compaction=possible_compaction, - tracking_namespace=namespace, + return self._log_inference_decision( + SessionInference( + session_id=state.session_id, + affinity_key=self._effective_affinity(state, hints, provider), + confidence=match.confidence, + match_score=match.score, + possible_compaction=possible_compaction, + lineage_parent_session_id=( + compaction.parent_session_id + if possible_compaction + and compaction.parent_session_id != state.session_id + else None + ), + tracking_namespace=namespace, + ), + action="compaction_continue" if possible_compaction else "continue", + matched_session_id=match.session_id, + compaction=compaction, + provider=provider, + model=model, ) - parent_id = ( - compaction_match.session_id - if self._is_compaction_parent_match( - compaction_match, - marker_compaction=marker_compaction, + parent_id = compaction.parent_session_id + if possible_compaction: + context_anchors = self._compaction_context_anchors( + compaction_probe_anchors, + compaction.context_probe_groups, + namespace, + parent_session_id=parent_id, + ) + normal_anchors = self._dedupe_anchors( + [ + *normal_anchors, + *context_anchors, + self._compaction_replay_anchor( + compaction_probe_anchors, + history_signatures, + namespace, + parent_session_id=parent_id, + ), + ] ) - else (match.session_id if match and possible_compaction else None) - ) session_id = str(uuid.uuid4()) state = self._create_session( session_id, @@ -303,24 +553,74 @@ def _infer_session_locked( normal_anchors, now, affinity_key=self._affinity_from_anchors(normal_anchors, namespace), + history_signatures=history_signatures, ) - if parent_id: - lib_logger.info( - "Session tracker: possible compacted descendant %s -> %s for %s", - parent_id, - session_id, - namespace, - ) - return SessionInference( - session_id=state.session_id, - affinity_key=state.affinity_key, - confidence="weak" if match else "none", - match_score=match.score if match else 0, - possible_compaction=possible_compaction, - lineage_parent_session_id=parent_id, - tracking_namespace=namespace, + return self._log_inference_decision( + SessionInference( + session_id=state.session_id, + affinity_key=self._effective_affinity(state, hints, provider), + confidence="weak" if match else "none", + match_score=match.score if match else 0, + possible_compaction=possible_compaction, + lineage_parent_session_id=parent_id, + tracking_namespace=namespace, + ), + action="compaction_child" if parent_id else "new", + candidate_session_id=match.session_id if match else None, + compaction=compaction, + provider=provider, + model=model, ) + def _log_inference_decision( + self, + inference: SessionInference, + *, + action: str, + matched_session_id: Optional[str] = None, + candidate_session_id: Optional[str] = None, + compaction: Optional[_CompactionDecision] = None, + provider: Optional[str] = None, + model: Optional[str] = None, + ) -> SessionInference: + """Emit temporary warning-level lineage diagnostics for every request.""" + + state = self._sessions.get(inference.session_id) if inference.session_id else None + origin = ( + "persisted" + if state and state.loaded_from_persistence + else ("memory" if state else "none") + ) + retained_ratio = ( + f"{compaction.retained_history_ratio:.3f}" + if compaction and compaction.retained_history_ratio is not None + else "-" + ) + response_events = compaction.response_group_count if compaction else 0 + marker = compaction.marker_compaction if compaction else False + lib_logger.warning( + "Session tracker decision: action=%s session_id=%s matched_session_id=%s " + "candidate_session_id=%s parent_session_id=%s namespace=%s " + "provider=%s model=%s confidence=%s score=%d origin=%s " + "possible_compaction=%s marker=%s retained_history=%s response_events=%d", + action, + inference.session_id or "-", + matched_session_id or "-", + candidate_session_id or "-", + inference.lineage_parent_session_id or "-", + inference.tracking_namespace or "-", + provider or "-", + model or "-", + inference.confidence, + inference.match_score, + origin, + inference.possible_compaction, + marker, + retained_ratio, + response_events, + ) + return inference + def record_response( self, session_id: Optional[str], @@ -339,16 +639,25 @@ def record_response( """ save_job = None with self._lock: - if not session_id or response is None or session_id not in self._sessions: - return now = time.time() - state = self._sessions[session_id] - namespace = tracking_namespace or state.namespace or self._namespace( - provider, model, scope_key=scope_key - ) - anchors = self._anchors_from_response(response, namespace) - if anchors: - self._refresh_and_bridge(session_id, namespace, anchors, now) + self._prune(now) + if session_id and response is not None and session_id in self._sessions: + state = self._sessions[session_id] + namespace = tracking_namespace or state.namespace or self._namespace( + provider, scope_key=scope_key + ) + if namespace != state.namespace: + lib_logger.warning( + "Session tracker normalized response namespace mismatch: " + "session_id=%s expected_namespace=%s received_namespace=%s", + session_id, + state.namespace, + namespace, + ) + namespace = state.namespace + anchors = self._anchors_from_response(response, namespace) + if anchors: + self._refresh_and_bridge(session_id, namespace, anchors, now) save_job = self._prepare_save_locked() self._write_save_job(save_job) @@ -366,6 +675,7 @@ def _create_session( now: float, *, affinity_key: Optional[str], + history_signatures: tuple[str, ...] = (), ) -> _SessionState: state = _SessionState( session_id=session_id, @@ -375,7 +685,14 @@ def _create_session( last_seen=now, ) self._sessions[session_id] = state - self._refresh_and_bridge(session_id, namespace, anchors, now, affinity_key=affinity_key) + self._refresh_and_bridge( + session_id, + namespace, + anchors, + now, + affinity_key=affinity_key, + history_signatures=history_signatures, + ) return state def _refresh_and_bridge( @@ -386,26 +703,51 @@ def _refresh_and_bridge( now: float, *, affinity_key: Optional[str] = None, + history_signatures: Optional[tuple[str, ...]] = None, ) -> _SessionState: expires_at = now + self.ttl_seconds state = self._sessions.setdefault( session_id, _SessionState(session_id=session_id, namespace=namespace, expires_at=expires_at), ) - state.namespace = namespace + if state.namespace != namespace: + raise ValueError( + f"Session {session_id} belongs to {state.namespace}, not {namespace}" + ) state.expires_at = expires_at state.last_seen = now if affinity_key and not state.affinity_key: state.affinity_key = affinity_key + if history_signatures and len(history_signatures) >= len(state.history_signatures): + # Keep the largest observed request as the structural baseline. Equal + # fixed-window histories advance to the newest normal conversation. + state.history_signatures = history_signatures for anchor in anchors: + existing = self._anchors.get(anchor.value) + if existing and existing.session_id != session_id: + # A shared prompt/chunk is not proof that a newer independent + # session owns it. Preserve the first live owner until TTL expiry. + continue + strength = ( + self._strongest(anchor.strength, existing.strength) + if existing + else anchor.strength + ) + source = anchor.source + group = anchor.group + if existing and existing.source == "response" and anchor.source != "response": + # Responses commonly return as ordinary assistant history. Keep + # their response-event provenance for future lineage decisions. + source = existing.source + group = existing.group state.anchors.add(anchor.value) self._anchors[anchor.value] = _AnchorRecord( session_id=session_id, namespace=namespace, - strength=anchor.strength, - source=anchor.source, - group=anchor.group, + strength=strength, + source=source, + group=group, expires_at=expires_at, last_seen=now, ) @@ -428,7 +770,16 @@ def _best_match( continue candidate = candidates.setdefault(record.session_id, _MatchCandidate(record.session_id)) candidate.last_seen = max(candidate.last_seen, record.last_seen) - strength = self._strongest(anchor.strength, record.strength) + if anchor.source == "compaction_probe" and anchor.group: + candidate.matched_probe_groups.add(anchor.group) + # Closure belongs to the current request. A previously closed event + # must not upgrade a later unpaired copy merely because the stored + # anchor retained medium strength. + strength = ( + anchor.strength + if anchor.source == "tool_event" + else self._strongest(anchor.strength, record.strength) + ) if strength == "strong": candidate.score += self._STRONG_SCORE candidate.strong_matches += 1 @@ -440,8 +791,18 @@ def _best_match( candidate.medium_groups.add(group) if anchor.source == "provider" or record.source == "provider": candidate.provider_matches += 1 - if anchor.source == "response" or record.source == "response": + response_overlap = anchor.source == "response" or record.source == "response" + if response_overlap and self._allows_response_bridge(anchor, record): candidate.response_matches += 1 + response_group = ( + anchor.group if anchor.source == "response" else record.group + ) + if response_group: + candidate.response_groups.add(response_group) + if anchor.source == "compaction_probe" and anchor.group: + candidate.response_probe_groups.add(anchor.group) + if record.source == "message" and record.group: + candidate.request_groups.add(record.group) else: candidate.score += self._WEAK_SCORE candidate.weak_matches += 1 @@ -455,6 +816,7 @@ def _best_match( item.strong_matches, item.medium_matches, len(item.medium_groups), + len(item.response_groups), item.response_matches, item.provider_matches, item.last_seen, @@ -472,6 +834,9 @@ def _coerce_hints(self, hints: Optional[Any]) -> Optional[SessionTrackingHints]: strong_anchors=list(hints.get("strong_anchors") or []), medium_anchors=list(hints.get("medium_anchors") or []), weak_anchors=list(hints.get("weak_anchors") or []), + global_strong_anchors=list(hints.get("global_strong_anchors") or []), + global_medium_anchors=list(hints.get("global_medium_anchors") or []), + global_weak_anchors=list(hints.get("global_weak_anchors") or []), affinity_key=hints.get("affinity_key"), session_scope=hints.get("session_scope"), ) @@ -483,11 +848,12 @@ def _build_anchors( namespace: str, hints: Optional[Any], *, + provider: Optional[str] = None, allow_system_continuity: bool = False, suppressed_continuity_indexes: Optional[set[int]] = None, ) -> List[SessionAnchor]: anchors: List[SessionAnchor] = [] - anchors.extend(self._anchors_from_provider_hints(hints, namespace)) + anchors.extend(self._anchors_from_provider_hints(hints, namespace, provider)) anchors.extend(self._anchors_from_explicit_ids(request_data, namespace)) messages = request_data.get("messages") or [] @@ -507,27 +873,31 @@ def _build_compaction_probe_anchors( self, request_data: Dict[str, Any], namespace: str, + *, + probe_indexes: Optional[set[int]] = None, ) -> List[SessionAnchor]: """Build temporary anchors for compaction lineage lookup only. Compaction summaries often replace prior user/assistant history and may - be sent as system, developer, user, or assistant messages. These anchors - are compared against existing response/message anchors to identify a - likely parent, but they are not stored on the newly-created child session. + be sent as system, developer, or user messages. Assistant/tool history is + deliberately excluded because replaying it is ordinary continuity. These + anchors are compared against existing response/message anchors to identify + a likely parent, but they are not stored on the newly-created child session. """ messages = request_data.get("messages") or [] if not isinstance(messages, list) or not messages: return [] anchors: List[SessionAnchor] = [] - probe_indexes = self._compaction_probe_indexes(request_data) + if probe_indexes is None: + probe_indexes = self._compaction_probe_indexes(request_data) for index, message in enumerate(messages[:2]): if index not in probe_indexes: continue if not isinstance(message, dict): continue text = self._normalize_text(self._extract_text(message.get("content"))) - role = str(message.get("role", "")) + role = str(message.get("role", "")).lower() anchors.append( SessionAnchor( self._scoped(namespace, f"message:{role}:{self._hash_text(text)}"), @@ -555,58 +925,285 @@ def _compaction_probe_indexes(self, request_data: Dict[str, Any]) -> set[int]: for index, message in enumerate(messages[:2]): if not isinstance(message, dict): continue + role = str(message.get("role", "")).lower() + if role not in self._COMPACTION_PROBE_ROLES: + continue text = self._normalize_text(self._extract_text(message.get("content"))) if self._is_compaction_probe_text(text): indexes.add(index) return indexes - def _is_compaction_parent_match( + def _evaluate_compaction( self, match: Optional[_MatchCandidate], *, marker_compaction: bool, - ) -> bool: + marker_probe_groups: set[str], + history_signatures: tuple[str, ...], + ) -> _CompactionDecision: + """Require parent evidence and replacement of most known history. + + One response overlap is intentionally insufficient for unmarked input: + side-channel classifiers often quote the latest response verbatim. A real + unmarked summary must span distinct response events and remove more than + half of the parent's high-water request history. + """ + if not match: - return False - # Explicit summary markers are allowed to produce weaker lineage because - # the result is telemetry and still starts a new live sticky session. - if marker_compaction: - return match.score > 0 - # Size-only probes are much noisier. Require overlap with prior response - # anchors so ordinary long repeated prompts/history are not classified as - # compaction descendants from normal message overlap alone. - return match.response_matches > 0 + return _CompactionDecision() + parent = self._sessions.get(match.session_id) + if not parent or not parent.history_signatures: + return _CompactionDecision() + retained_ratio = self._retained_history_ratio( + history_signatures, + parent.history_signatures, + ) + if retained_ratio >= self._COMPACTION_MAX_RETAINED_HISTORY_RATIO: + lib_logger.debug( + "Session tracker rejected compaction candidate for %s: " + "retained_history=%.3f", + match.session_id, + retained_ratio, + ) + return _CompactionDecision() + + response_group_count = len(match.response_groups) + qualifies = ( + match.score > 0 + if marker_compaction + else ( + response_group_count >= self._MIN_UNMARKED_RESPONSE_GROUPS + and bool(match.request_groups) + ) + ) + if not qualifies: + lib_logger.debug( + "Session tracker rejected compaction candidate for %s: " + "marker=%s, response_events=%d, request_groups=%d", + match.session_id, + marker_compaction, + response_group_count, + len(match.request_groups), + ) + return _CompactionDecision() + marker_context_groups = match.matched_probe_groups.intersection(marker_probe_groups) + return _CompactionDecision( + parent_session_id=match.session_id, + marker_compaction=marker_compaction, + retained_history_ratio=retained_ratio, + response_group_count=response_group_count, + context_probe_groups=frozenset( + (marker_context_groups or match.response_probe_groups) + if marker_compaction + else match.response_probe_groups + ), + ) + + def _compaction_marker_probe_groups(self, request_data: Dict[str, Any]) -> set[str]: + """Return early system/developer probe groups carrying explicit markers.""" + + messages = request_data.get("messages") or [] + groups: set[str] = set() + if not isinstance(messages, list): + return groups + for index, message in enumerate(messages[:2]): + if not isinstance(message, dict): + continue + role = str(message.get("role", "")).lower() + text = self._extract_text(message.get("content")) + if role in {"system", "developer"} and self._has_compaction_marker(text): + groups.add(f"compaction_probe:{index}") + return groups + + def _compaction_replay_anchor( + self, + probe_anchors: List[SessionAnchor], + history_signatures: tuple[str, ...], + namespace: str, + *, + parent_session_id: Optional[str], + ) -> SessionAnchor: + """Build an opaque exact-replay key for a validated compacted payload.""" + + digest = self._hash_json( + { + "probe_anchors": sorted(anchor.value for anchor in probe_anchors), + "history_signatures": list(history_signatures), + } + ) + return SessionAnchor( + self._scoped(namespace, f"compaction_replay:{digest}"), + "strong", + source="compaction_replay", + group=parent_session_id, + ) + + def _compaction_context_anchor( + self, + probe_anchors: List[SessionAnchor], + namespace: str, + *, + parent_session_id: Optional[str], + ) -> SessionAnchor: + """Bind later requests that retain one validated compacted base context.""" + + digest = self._hash_json(sorted(anchor.value for anchor in probe_anchors)) + return SessionAnchor( + self._scoped(namespace, f"compaction_context:{digest}"), + "strong", + source="compaction_context", + group=parent_session_id, + ) + + def _compaction_context_anchors( + self, + probe_anchors: List[SessionAnchor], + probe_groups: Iterable[str], + namespace: str, + *, + parent_session_id: Optional[str], + ) -> List[SessionAnchor]: + """Build context keys only for probes that actually matched the parent.""" + + anchors: List[SessionAnchor] = [] + for group in sorted(set(probe_groups)): + grouped = [anchor for anchor in probe_anchors if anchor.group == group] + if grouped: + anchors.append( + self._compaction_context_anchor( + grouped, + namespace, + parent_session_id=parent_session_id, + ) + ) + return anchors + + @staticmethod + def _is_authoritative_identity_anchor(anchor: SessionAnchor) -> bool: + """Return whether an anchor must take precedence over replay bindings.""" + + return anchor.strength == "strong" and anchor.source in { + "explicit", + "global_hint", + "provider", + } + + def _find_compaction_replay( + self, + probe_anchors: List[SessionAnchor], + history_signatures: tuple[str, ...], + namespace: str, + now: float, + ) -> Optional[_AnchorRecord]: + if not probe_anchors: + return None + anchor = self._compaction_replay_anchor( + probe_anchors, + history_signatures, + namespace, + parent_session_id=None, + ) + record = self._anchors.get(anchor.value) + if ( + not record + or record.source != "compaction_replay" + or record.namespace != namespace + or record.expires_at <= now + or record.session_id not in self._sessions + or not record.group + ): + return None + return record + + def _find_compaction_context( + self, + probe_anchors: List[SessionAnchor], + namespace: str, + now: float, + ) -> Optional[tuple[_AnchorRecord, SessionAnchor]]: + probe_groups = sorted({anchor.group for anchor in probe_anchors if anchor.group}) + for probe_group in probe_groups: + grouped = [anchor for anchor in probe_anchors if anchor.group == probe_group] + anchor = self._compaction_context_anchor( + grouped, + namespace, + parent_session_id=None, + ) + record = self._anchors.get(anchor.value) + if ( + record + and record.source == "compaction_context" + and record.namespace == namespace + and record.expires_at > now + and record.session_id in self._sessions + and record.group + ): + return ( + record, + SessionAnchor( + anchor.value, + "strong", + source="compaction_context", + group=record.group, + ), + ) + return None def _anchors_from_provider_hints( self, hints: Optional[SessionTrackingHints], namespace: str, + provider: Optional[str], ) -> List[SessionAnchor]: if not hints: return [] anchors: List[SessionAnchor] = [] + # A provider may partition only its native evidence/affinity domain. This + # never fragments the global logical session or caller isolation domain. + native_scope = hints.session_scope or "provider" + provider_key = self._hash_text(f"{provider or 'unknown'}:{native_scope}") for strength, attr in ( ("strong", "strong_anchors"), ("medium", "medium_anchors"), ("weak", "weak_anchors"), ): for value in getattr(hints, attr, []) or []: + value_hash = self._hash_text(str(value)) anchors.append( SessionAnchor( - self._scoped(namespace, f"provider:{value}"), + self._scoped(namespace, f"provider:{provider_key}:{value_hash}"), strength, source="provider", - group=f"provider:{value}", + group=f"provider:{provider_key}:{value_hash}", + ) + ) + for strength, attr in ( + ("strong", "global_strong_anchors"), + ("medium", "global_medium_anchors"), + ("weak", "global_weak_anchors"), + ): + for value in getattr(hints, attr, []) or []: + value_hash = self._hash_text(str(value)) + anchors.append( + SessionAnchor( + self._scoped(namespace, f"global:{value_hash}"), + strength, + source="global_hint", + group=f"global:{value_hash}", ) ) affinity_key = getattr(hints, "affinity_key", None) if affinity_key: + affinity_hash = self._hash_text(str(affinity_key)) anchors.append( SessionAnchor( - self._scoped(namespace, f"provider_affinity:{affinity_key}"), + self._scoped( + namespace, + f"provider:{provider_key}:affinity:{affinity_hash}", + ), "strong", source="provider", - group="provider_affinity", + group=f"provider:{provider_key}:affinity", ) ) return anchors @@ -630,10 +1227,11 @@ def _anchors_from_explicit_ids( ): value = request_data.get(key) if value: + value_hash = self._hash_text(str(value)) strength = "strong" if key in self.trusted_explicit_fields else "weak" anchors.append( SessionAnchor( - self._scoped(namespace, f"explicit:{key}:{value}"), + self._scoped(namespace, f"explicit:{key}:{value_hash}"), strength, source="explicit", group=f"explicit:{key}", @@ -647,6 +1245,7 @@ def _anchors_from_messages( namespace: str, *, source: str = "message", + evidence_group: Optional[str] = None, allow_system_continuity: bool = False, suppressed_continuity_indexes: Optional[set[int]] = None, ) -> List[SessionAnchor]: @@ -654,6 +1253,7 @@ def _anchors_from_messages( normalized_messages: List[Dict[str, Any]] = [] tool_ids: List[str] = [] first_user_text: Optional[str] = None + closed_tool_calls = self._closed_tool_call_positions(messages) for index, message in enumerate(messages): if not isinstance(message, dict): @@ -666,36 +1266,56 @@ def _anchors_from_messages( tool_call_id = message.get("tool_call_id") if tool_call_id: tool_id = str(tool_call_id) + tool_id_hash = self._hash_text(tool_id) tool_ids.append(tool_id) normalized["tool_call_id"] = tool_id anchors.append( SessionAnchor( - self._scoped(namespace, f"tool:{tool_id}"), - "strong", + self._scoped(namespace, f"tool:{tool_id_hash}"), + "weak", source="tool", - group=f"tool:{tool_id}", + group=f"tool:{tool_id_hash}", ) ) tool_calls = message.get("tool_calls") or [] if isinstance(tool_calls, list) and tool_calls: call_ids: List[str] = [] - for tool_call in tool_calls: + for tool_call_index, tool_call in enumerate(tool_calls): if not isinstance(tool_call, dict): continue call_id = tool_call.get("id") if call_id: call_id = str(call_id) + call_id_hash = self._hash_text(call_id) call_ids.append(call_id) tool_ids.append(call_id) anchors.append( SessionAnchor( - self._scoped(namespace, f"tool:{call_id}"), - "strong", + self._scoped(namespace, f"tool:{call_id_hash}"), + "weak", source="tool", - group=f"tool:{call_id}", + group=f"tool:{call_id_hash}", ) ) + event = self._tool_event_descriptor(tool_call) + if role.lower() == "assistant" and event is not None: + event_hash = self._hash_json(event) + anchors.append( + SessionAnchor( + self._scoped( + namespace, + f"tool_event:{event_hash}", + ), + ( + "medium" + if (index, tool_call_index) in closed_tool_calls + else "weak" + ), + source="tool_event", + group=f"tool_event:{event_hash}", + ) + ) if call_ids: normalized["tool_calls"] = call_ids @@ -719,7 +1339,7 @@ def _anchors_from_messages( self._scoped(namespace, f"message:{role}:{self._hash_text(normalized_text)}"), "medium", source=source, - group=f"{source}:{index}", + group=evidence_group or f"{source}:{index}:{role.lower()}", ) ) for chunk_hash in self._content_chunk_hashes(normalized_text): @@ -728,7 +1348,7 @@ def _anchors_from_messages( self._scoped(namespace, f"chunk:{chunk_hash}"), "medium", source=source, - group=f"{source}:{index}", + group=evidence_group or f"{source}:{index}:{role.lower()}", ) ) @@ -741,7 +1361,7 @@ def _anchors_from_messages( anchors.append( SessionAnchor( self._scoped(namespace, "tool_group:" + self._hash_json(sorted(tool_ids))), - "strong", + "weak", source="tool", group="tool_group", ) @@ -773,6 +1393,27 @@ def _anchors_from_response(self, response: Any, namespace: str) -> List[SessionA data = response.model_dump() if hasattr(response, "model_dump") else response if not isinstance(data, dict): return [] + anchors: List[SessionAnchor] = [] + response_id = data.get("id") + if response_id: + response_id_hash = self._hash_text( + f"responses_previous_response_id:{response_id}" + ) + # Responses API continuations identify their parent with + # previous_response_id. Record the emitted response id as strong + # response evidence so the next request can route back to the exact + # session/credential that produced it. + anchors.append( + SessionAnchor( + self._scoped( + namespace, + f"global:{response_id_hash}", + ), + "strong", + source="response", + group="responses_previous_response_id", + ) + ) messages: List[Dict[str, Any]] = [] for choice in data.get("choices") or []: if not isinstance(choice, dict): @@ -782,17 +1423,102 @@ def _anchors_from_response(self, response: Any, namespace: str) -> List[SessionA response_message = dict(message) response_message.setdefault("role", "assistant") messages.append(response_message) - return self._anchors_from_messages(messages, namespace, source="response") if messages else [] + if messages: + response_group = "response_event:" + self._hash_json( + [self._message_signature(message) for message in messages] + ) + anchors.extend( + self._anchors_from_messages( + messages, + namespace, + source="response", + evidence_group=response_group, + ) + ) + return anchors + + @staticmethod + def _closed_tool_call_positions( + messages: List[Dict[str, Any]], + ) -> set[tuple[int, int]]: + """Pair each tool result with one earlier assistant call of the same ID.""" + + pending: Dict[str, List[tuple[int, int]]] = {} + closed: set[tuple[int, int]] = set() + for message_index, message in enumerate(messages): + if not isinstance(message, dict): + continue + role = str(message.get("role", "")).lower() + if role == "assistant": + tool_calls = message.get("tool_calls") or [] + if not isinstance(tool_calls, list): + continue + for tool_call_index, tool_call in enumerate(tool_calls): + if not isinstance(tool_call, dict) or not tool_call.get("id"): + continue + pending.setdefault(str(tool_call["id"]), []).append( + (message_index, tool_call_index) + ) + continue + if role not in {"tool", "function"} or not message.get("tool_call_id"): + continue + candidates = pending.get(str(message["tool_call_id"])) + if candidates: + closed.add(candidates.pop(0)) + return closed + + def _tool_event_descriptor( + self, + tool_call: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + """Return stable call structure used by pending and closed tool events.""" + + call_id = tool_call.get("id") + function = tool_call.get("function") + if isinstance(function, dict): + name = function.get("name") or tool_call.get("name") + arguments = function.get("arguments") + else: + name = tool_call.get("name") + arguments = tool_call.get("arguments") + if not call_id or not name: + return None + return { + "id": str(call_id), + "name": str(name), + "arguments": self._canonical_tool_arguments(arguments), + } + + def _canonical_tool_arguments(self, arguments: Any) -> Any: + """Normalize JSON arguments without retaining their plaintext in anchors.""" + + if not isinstance(arguments, str): + return arguments + stripped = arguments.strip() + if not stripped: + return "" + try: + return json.loads(stripped) + except (TypeError, ValueError, json.JSONDecodeError): + return self._normalize_text(stripped) def _affinity_from_anchors( self, anchors: List[SessionAnchor], namespace: str, ) -> Optional[str]: - strong = sorted(anchor.value for anchor in anchors if anchor.strength == "strong") + strong = sorted( + anchor.value + for anchor in anchors + if anchor.strength == "strong" and anchor.source != "provider" + ) if strong: return self._scoped(namespace, "affinity:" + self._hash_json(strong[:4])) - medium_anchors = [anchor for anchor in anchors if anchor.strength == "medium"] + medium_anchors = [ + anchor + for anchor in anchors + if anchor.strength == "medium" and anchor.source != "provider" + ] medium_groups = { anchor.group for anchor in medium_anchors @@ -806,20 +1532,25 @@ def _affinity_from_anchors( return self._scoped(namespace, "affinity:" + self._hash_json(medium[:8])) return None - def _looks_like_compaction(self, request_data: Dict[str, Any]) -> bool: - messages = request_data.get("messages") or [] - if not isinstance(messages, list) or not messages: - return False - summary_texts: List[str] = [] - for message in messages[:2]: - if not isinstance(message, dict): - continue - role = str(message.get("role", "")).lower() - if role not in {"system", "developer"}: - continue - summary_texts.append(self._extract_text(message.get("content"))) - joined = "\n".join(summary_texts) - return self._has_compaction_marker(joined) + @staticmethod + def _effective_affinity( + state: _SessionState, + hints: Optional[SessionTrackingHints], + provider: Optional[str], + ) -> Optional[str]: + """Use provider affinity only for the provider handling this request.""" + + if hints: + if hints.affinity_key: + return hints.affinity_key + native = sorted(str(value) for value in hints.strong_anchors if value) + if native: + payload = json.dumps( + [provider or "unknown", hints.session_scope or "provider", native], + separators=(",", ":"), + ) + return "provider-affinity:" + hashlib.sha256(payload.encode("utf-8")).hexdigest() + return state.affinity_key def _is_compaction_probe_text(self, text: str) -> bool: if not text: @@ -848,7 +1579,10 @@ def _trim_session_anchors(self, state: _SessionState) -> None: return sorted_anchors = sorted( state.anchors, - key=lambda value: self._anchors.get(value, _AnchorRecord("", "", "weak", "", None, 0, 0)).last_seen, + key=lambda value: ( + self._anchor_eviction_key(self._anchors.get(value)), + value, + ), ) for value in sorted_anchors[: len(state.anchors) - self.max_anchors_per_session]: state.anchors.discard(value) @@ -860,28 +1594,49 @@ def _trim_global_anchors(self) -> None: if len(self._anchors) <= self.max_anchor_records: return overage = len(self._anchors) - self.max_anchor_records - for value, record in sorted(self._anchors.items(), key=lambda item: item[1].last_seen)[:overage]: + for value, record in sorted( + self._anchors.items(), + key=lambda item: (self._anchor_eviction_key(item[1]), item[0]), + )[:overage]: state = self._sessions.get(record.session_id) if state: state.anchors.discard(value) del self._anchors[value] def _prune(self, now: Optional[float] = None) -> None: - now = now or time.time() + if now is None: + now = time.time() + changed = False expired_sessions = [key for key, state in self._sessions.items() if state.expires_at <= now] for session_id in expired_sessions: state = self._sessions.pop(session_id) + changed = True for anchor in list(state.anchors): record = self._anchors.get(anchor) if record and record.session_id == session_id: del self._anchors[anchor] expired_anchors = [key for key, record in self._anchors.items() if record.expires_at <= now] for anchor in expired_anchors: - del self._anchors[anchor] + record = self._anchors.pop(anchor) + changed = True + state = self._sessions.get(record.session_id) + if state: + state.anchors.discard(anchor) + if changed: + self._mark_dirty() def _load(self) -> None: if not self.persistence_path: return + try: + if ( + self.persistence_path.exists() + and self.persistence_path.stat().st_size > self._MAX_PERSISTED_FILE_BYTES + ): + lib_logger.warning("Ignoring oversized session_stickiness.json") + return + except OSError: + return data = safe_read_json(self.persistence_path, lib_logger) if not isinstance(data, dict): return @@ -891,43 +1646,105 @@ def _load(self) -> None: "Ignoring unsupported session_stickiness.json format; session persistence will rebuild in memory." ) return - sessions = data.get("sessions", {}) - anchors = data.get("anchors", {}) - for session_id, payload in sessions.items(): + sessions = data.get("sessions") + anchors = data.get("anchors") + if not isinstance(sessions, dict) or not isinstance(anchors, dict): + lib_logger.info( + "Ignoring malformed session_stickiness.json containers; " + "session persistence will rebuild in memory." + ) + return + def persisted_last_seen(item: tuple[Any, Any]) -> float: + payload = item[1] if not isinstance(payload, dict): + return 0.0 + return self._finite_float(payload.get("last_seen")) or 0.0 + + ordered_sessions = sorted( + sessions.items(), + key=persisted_last_seen, + reverse=True, + )[: self._MAX_PERSISTED_SESSIONS] + for session_id, payload in ordered_sessions: + if not isinstance(session_id, str) or not session_id or not isinstance(payload, dict): + continue + if len(session_id) > self._MAX_PERSISTED_STRING_LENGTH: continue - expires_at = float(payload.get("expires_at", 0.0)) - if expires_at <= now: + expires_at = self._finite_float(payload.get("expires_at")) + if expires_at is None or expires_at <= now: continue + namespace = payload.get("namespace") + if not isinstance(namespace, str) or not namespace.startswith("session-domain:"): + continue + if len(namespace) > self._MAX_PERSISTED_STRING_LENGTH: + continue + last_seen = self._finite_float(payload.get("last_seen")) + history_payload = payload.get("history_signatures") or [] + history_signatures = ( + tuple( + value + for value in history_payload[: self._MAX_PERSISTED_HISTORY_SIGNATURES] + if isinstance(value, str) and re.fullmatch(r"[0-9a-f]{64}", value) + ) + if isinstance(history_payload, list) + else () + ) + affinity_key = payload.get("affinity_key") + if affinity_key is not None and not isinstance(affinity_key, str): + affinity_key = None + if isinstance(affinity_key, str) and len(affinity_key) > self._MAX_PERSISTED_STRING_LENGTH: + affinity_key = None self._sessions[session_id] = _SessionState( session_id=session_id, - namespace=str(payload.get("namespace") or "global"), + namespace=namespace, expires_at=expires_at, - affinity_key=payload.get("affinity_key"), - anchors=set(payload.get("anchors") or []), - last_seen=float(payload.get("last_seen", now)), + affinity_key=affinity_key, + anchors=set(), + last_seen=last_seen if last_seen is not None else now, + history_signatures=history_signatures, + loaded_from_persistence=True, ) for value, payload in anchors.items(): if not isinstance(payload, dict): continue session_id = payload.get("session_id") - expires_at = float(payload.get("expires_at", 0.0)) - if not session_id or expires_at <= now: + state = self._sessions.get(session_id) if isinstance(session_id, str) else None + expires_at = self._finite_float(payload.get("expires_at")) + if not state or expires_at is None or expires_at <= now: + continue + namespace = payload.get("namespace") + if namespace != state.namespace: + continue + if not isinstance(value, str) or not value.startswith(f"{namespace}:"): + continue + if len(value) > self._MAX_PERSISTED_STRING_LENGTH: + continue + strength = payload.get("strength") + source = payload.get("source") + group = payload.get("group") + if strength not in {"weak", "medium", "strong"}: continue - namespace = str(payload.get("namespace") or "global") + if source not in self._PERSISTED_ANCHOR_SOURCES: + continue + if group is not None and not isinstance(group, str): + continue + if isinstance(group, str) and len(group) > self._MAX_PERSISTED_STRING_LENGTH: + continue + last_seen = self._finite_float(payload.get("last_seen")) self._anchors[value] = _AnchorRecord( session_id=session_id, namespace=namespace, - strength=str(payload.get("strength") or "medium"), - source=str(payload.get("source") or "generic"), - group=payload.get("group"), - expires_at=expires_at, - last_seen=float(payload.get("last_seen", now)), + strength=strength, + source=source, + group=group, + expires_at=min(expires_at, state.expires_at), + last_seen=last_seen if last_seen is not None else now, ) - self._sessions.setdefault( - session_id, - _SessionState(session_id=session_id, namespace=namespace, expires_at=expires_at), - ).anchors.add(value) + state.anchors.add(value) + + for state in self._sessions.values(): + self._trim_session_anchors(state) + self._trim_global_anchors() def _prepare_save_locked( self, @@ -947,8 +1764,8 @@ def _prepare_save_locked( "namespace": state.namespace, "expires_at": state.expires_at, "affinity_key": state.affinity_key, - "anchors": sorted(state.anchors), "last_seen": state.last_seen, + "history_signatures": list(state.history_signatures), } for session_id, state in self._sessions.items() }, @@ -983,7 +1800,11 @@ def _write_save_job( with self._save_io_lock: if generation < self._last_persisted_generation: return - success = writer.write(payload) + try: + success = writer.write(payload) + except Exception as exc: + lib_logger.warning("Failed to persist session tracking state: %s", exc) + success = False if success: self._last_persisted_generation = generation if not success: @@ -999,17 +1820,38 @@ def _mark_dirty(self) -> None: def _namespace( self, provider: Optional[str], - model: Optional[str], *, scope_key: Optional[str] = None, - session_scope: Optional[str] = None, + trusted_isolation_key: bool = False, ) -> str: - # The resolved usage/classifier scope is part of the namespace so sticky - # evidence never leaks between private/classifier-scoped credential pools. - allowed_scope = scope_key or "default" - provider_key = provider or "global" - model_key = session_scope or model or "default" - return f"scope:{allowed_scope}:provider:{provider_key}:model:{model_key}" + # Logical sessions cross providers and models, but never caller/credential + # isolation domains. Provider-native evidence is qualified separately. + allowed_scope = self._normalize_isolation_key( + scope_key, + provider, + trusted=trusted_isolation_key, + ) + return f"session-domain:{allowed_scope}" + + @staticmethod + def _normalize_isolation_key( + scope_key: Optional[str], + provider: Optional[str], + *, + trusted: bool, + ) -> str: + """Accept internal domain markers only from RequestContextBuilder.""" + + scope = str(scope_key or provider or "public") + provider_key = (provider or "").strip().lower() + if not scope_key or scope.strip().lower() in {"public", provider_key}: + return "public" + if trusted and ( + re.fullmatch(r"classifier:[0-9a-f]{24}", scope) + or re.fullmatch(r"bundle:[0-9a-f]{64}", scope) + ): + return scope + return "scope:" + hashlib.sha256(scope.encode("utf-8")).hexdigest() def _trusted_fields_from_env(self) -> List[str]: raw = os.getenv("TRUSTED_SESSION_ID_FIELDS", "") @@ -1022,9 +1864,102 @@ def _strongest(self, left: str, right: str) -> str: order = {"weak": 0, "medium": 1, "strong": 2} return left if order.get(left, 0) >= order.get(right, 0) else right + def _anchor_eviction_key( + self, + record: Optional[_AnchorRecord], + ) -> tuple[int, int, float]: + """Evict weak/ordinary evidence before strong replay/context identity.""" + + if record is None: + return (-1, 0, 0.0) + strength_rank = {"weak": 0, "medium": 1, "strong": 2}.get(record.strength, 0) + replay_rank = 1 if record.source in {"compaction_context", "compaction_replay"} else 0 + return (strength_rank, replay_rank, record.last_seen) + + def _allows_response_bridge( + self, + anchor: SessionAnchor, + record: _AnchorRecord, + ) -> bool: + """Only assistant-role request history may bridge from one response. + + Compaction probes are evaluated separately and may summarize responses in + a user/system message. Ordinary user prompts that merely quote a response + must not become sticky to the producer session. + """ + + message_group = None + if anchor.source == "message": + message_group = anchor.group + elif record.source == "message": + message_group = record.group + return message_group is None or message_group.endswith(":assistant") + + def _finite_float(self, value: Any) -> Optional[float]: + """Parse persisted numeric metadata without allowing NaN or infinity.""" + + try: + parsed = float(value) + except (TypeError, ValueError): + return None + return parsed if math.isfinite(parsed) else None + def _is_substantial_text(self, text: str) -> bool: return len(text) >= 24 and len(text.split()) >= 4 + def _request_history_signatures(self, request_data: Dict[str, Any]) -> tuple[str, ...]: + """Return content-free signatures for structural history comparison.""" + + messages = request_data.get("messages") or [] + if not isinstance(messages, list): + return () + return tuple( + signature + for message in messages + if isinstance(message, dict) + for signature in [self._message_signature(message)] + if signature + ) + + def _message_signature(self, message: Dict[str, Any]) -> str: + """Hash the normalized message shape without retaining raw content.""" + + payload: Dict[str, Any] = { + "role": str(message.get("role", "")).lower(), + "content": self._normalize_content(message.get("content")), + } + tool_call_id = message.get("tool_call_id") + if tool_call_id: + payload["tool_call_id"] = str(tool_call_id) + tool_calls = message.get("tool_calls") or [] + if isinstance(tool_calls, list): + normalized_calls = [] + for call in tool_calls: + if not isinstance(call, dict): + continue + normalized_calls.append( + { + key: call.get(key) + for key in ("id", "type", "function", "name") + if call.get(key) is not None + } + ) + if normalized_calls: + payload["tool_calls"] = normalized_calls + return self._hash_json(payload) + + def _retained_history_ratio( + self, + current: tuple[str, ...], + previous: tuple[str, ...], + ) -> float: + """Measure how much of the parent's high-water history remains.""" + + if not previous: + return 1.0 + retained = sum((Counter(current) & Counter(previous)).values()) + return retained / len(previous) + def _content_chunk_hashes(self, text: str) -> List[str]: words = text.split() if len(words) < 8: diff --git a/src/rotator_library/streaming/__init__.py b/src/rotator_library/streaming/__init__.py new file mode 100644 index 000000000..209c068a8 --- /dev/null +++ b/src/rotator_library/streaming/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Reusable streaming primitives shared by protocol and provider layers.""" + +from .events import StreamEvent, stream_event_from_sse_chunk +from .errors import StreamingErrorDecision, decide_streaming_error_action +from .metrics import StreamMetrics, StreamMonitor +from .policy import can_retry_stream_after_error, is_visible_stream_output +from .transport import JSONLineStreamFormatter, SSEStreamFormatter, WebSocketStreamFormatter + +__all__ = [ + "JSONLineStreamFormatter", + "SSEStreamFormatter", + "StreamEvent", + "StreamingErrorDecision", + "StreamMetrics", + "StreamMonitor", + "WebSocketStreamFormatter", + "can_retry_stream_after_error", + "decide_streaming_error_action", + "is_visible_stream_output", + "stream_event_from_sse_chunk", +] diff --git a/src/rotator_library/streaming/errors.py b/src/rotator_library/streaming/errors.py new file mode 100644 index 000000000..94467e228 --- /dev/null +++ b/src/rotator_library/streaming/errors.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Side-effect-free decisions for streaming failures.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from ..error_handler import ClassifiedError, classify_error, should_retry_same_key, should_rotate_on_error +from ..retry_policy import ProviderCooldownDecision, decide_provider_cooldown +from .policy import can_retry_stream_after_error, is_visible_stream_output + + +@dataclass(frozen=True) +class StreamingErrorDecision: + """Decision returned by streaming error policy before executor side effects.""" + + classified: ClassifiedError + action: str + start_provider_cooldown: bool = False + provider_cooldown_duration: int = 0 + provider_cooldown_scope: str = "provider" + provider_cooldown_model: str | None = None + reason: str = "" + + +def decide_streaming_error_action( + error: Exception, + *, + provider: str, + last_streamed_chunk: str | None, + attempt: int, + max_retries: int, + small_cooldown_threshold: int, + provider_cooldown_min_seconds: int, + provider_cooldown_default_seconds: int, + cooldown_on_quota: bool = False, + allow_reasoning_only_retry: bool = False, + model: str | None = None, + emitted_output: bool = False, +) -> StreamingErrorDecision: + """Classify a stream failure without sleeping or mutating state. + + The executor remains responsible for logging, credential state, sleeping, + provider cooldown mutation, and fallback. This helper only makes the same + decision consistently across streaming exception branches. + """ + + classified = classify_error(error, provider) + cooldown = _cooldown_decision( + classified, + small_cooldown_threshold=small_cooldown_threshold, + provider_cooldown_min_seconds=provider_cooldown_min_seconds, + provider_cooldown_default_seconds=provider_cooldown_default_seconds, + cooldown_on_quota=cooldown_on_quota, + last_streamed_chunk=last_streamed_chunk, + emitted_output=emitted_output, + provider=provider, + model=model, + original_error=error, + ) + if not should_rotate_on_error(classified): + return _decision(classified, "fail", cooldown, "non_rotatable") + if emitted_output or not can_retry_stream_after_error(last_streamed_chunk, allow_reasoning_only_retry): + return _decision(classified, "fallback_blocked_after_output", cooldown, "visible_output") + if should_retry_same_key(classified, small_cooldown_threshold) and attempt < max_retries - 1: + return _decision(classified, "retry_same", cooldown, "retry_same_credential") + return _decision(classified, "rotate", cooldown, "rotate_credential") + + +def _cooldown_decision( + classified: ClassifiedError, + *, + small_cooldown_threshold: int, + provider_cooldown_min_seconds: int, + provider_cooldown_default_seconds: int, + cooldown_on_quota: bool, + last_streamed_chunk: str | None, + emitted_output: bool, + provider: str, + model: str | None, + original_error: Exception, +) -> ProviderCooldownDecision: + if emitted_output or is_visible_stream_output(last_streamed_chunk): + return ProviderCooldownDecision(False, reason="visible_output") + return decide_provider_cooldown( + classified, + small_cooldown_threshold=small_cooldown_threshold, + provider_cooldown_min_seconds=provider_cooldown_min_seconds, + default_duration=provider_cooldown_default_seconds, + cooldown_on_quota=cooldown_on_quota, + provider=provider, + model=model, + original_error=original_error, + ) + + +def _decision( + classified: ClassifiedError, + action: str, + cooldown: ProviderCooldownDecision, + reason: str, +) -> StreamingErrorDecision: + return StreamingErrorDecision( + classified=classified, + action=action, + start_provider_cooldown=cooldown.should_start, + provider_cooldown_duration=cooldown.duration, + provider_cooldown_scope=cooldown.scope, + provider_cooldown_model=cooldown.model, + reason=reason if not cooldown.should_start else f"{reason};cooldown:{cooldown.reason}", + ) diff --git a/src/rotator_library/streaming/events.py b/src/rotator_library/streaming/events.py new file mode 100644 index 000000000..294b6b267 --- /dev/null +++ b/src/rotator_library/streaming/events.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Transport-neutral stream event model.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime +import json +from typing import Any, Optional + +from ..protocols import serialize_value + + +@dataclass(frozen=True) +class StreamEvent: + """A normalized stream event before transport-specific formatting. + + The event model is intentionally broad so providers can attach native data + without inventing a new stream lifecycle object. Raw data is for trace/debug + use only and should still pass through transaction-log redaction when logged. + """ + + event_type: str + data: Any = field(default_factory=dict) + protocol: Optional[str] = None + transport: str = "sse" + raw: Any = None + metadata: dict[str, Any] = field(default_factory=dict) + visible_output: bool = False + timestamp_utc: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe representation for traces and transports.""" + + return { + "event_type": self.event_type, + "protocol": self.protocol, + "transport": self.transport, + "data": serialize_value(self.data), + "raw": serialize_value(self.raw), + "metadata": serialize_value(self.metadata), + "visible_output": self.visible_output, + "timestamp_utc": self.timestamp_utc, + } + + +def stream_event_from_sse_chunk(chunk: Any, *, protocol: str = "openai_chat") -> StreamEvent: + """Parse an SSE `data:` chunk into a normalized stream event. + + Malformed or non-JSON chunks are treated as metadata and fail closed for + visibility. This avoids accidental fallback after ambiguous client output. + """ + + if isinstance(chunk, dict): + data = chunk + raw = chunk + elif isinstance(chunk, str): + raw = chunk + text = chunk.strip() + if text.startswith("data:"): + text = text[len("data:") :].strip() + if text == "[DONE]": + return StreamEvent("completed", protocol=protocol, raw=chunk, data={"done": True}) + try: + data = json.loads(text) + except json.JSONDecodeError: + return StreamEvent("metadata", protocol=protocol, raw=chunk, data={"malformed": True}) + else: + return StreamEvent("metadata", protocol=protocol, raw=chunk, data={"unsupported_chunk": True}) + + if isinstance(data, dict) and data.get("error"): + return StreamEvent("error", protocol=protocol, raw=raw, data=data, visible_output=False) + visible = _openai_chat_visible(data) if protocol == "openai_chat" else False + event_type = "delta" if visible else "parsed_chunk" + return StreamEvent(event_type, protocol=protocol, raw=raw, data=data, visible_output=visible) + + +def _openai_chat_visible(data: Any) -> bool: + if not isinstance(data, dict): + return False + choices = data.get("choices") + if not isinstance(choices, list): + return False + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") or {} + message = choice.get("message") or {} + for source in (delta, message): + if not isinstance(source, dict): + continue + if _has_visible_text(source.get("content")) or _has_visible_text(source.get("text")): + return True + if source.get("tool_calls") or source.get("function_call"): + return True + return False + + +def _has_visible_text(value: Any) -> bool: + if isinstance(value, str): + return bool(value.strip()) + if isinstance(value, list): + for item in value: + if isinstance(item, str) and item.strip(): + return True + if isinstance(item, dict) and _has_visible_text(item.get("text")): + return True + return False diff --git a/src/rotator_library/streaming/metrics.py b/src/rotator_library/streaming/metrics.py new file mode 100644 index 000000000..64470d700 --- /dev/null +++ b/src/rotator_library/streaming/metrics.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Stream lifecycle metrics and stall monitoring.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, Optional + +from .events import StreamEvent + + +Clock = Callable[[], float] + + +@dataclass +class StreamMetrics: + """Timing and lifecycle counters for one stream.""" + + started_at: float + first_byte_at: Optional[float] = None + first_visible_output_at: Optional[float] = None + last_chunk_at: Optional[float] = None + completed_at: Optional[float] = None + chunk_count: int = 0 + visible_chunk_count: int = 0 + error_count: int = 0 + cancelled: bool = False + + @property + def ttfb_seconds(self) -> Optional[float]: + return None if self.first_byte_at is None else self.first_byte_at - self.started_at + + @property + def ttft_seconds(self) -> Optional[float]: + return None if self.first_visible_output_at is None else self.first_visible_output_at - self.started_at + + @property + def duration_seconds(self) -> Optional[float]: + return None if self.completed_at is None else self.completed_at - self.started_at + + @property + def idle_seconds(self) -> Optional[float]: + if self.last_chunk_at is None or self.completed_at is None: + return None + return self.completed_at - self.last_chunk_at + + def to_dict(self) -> dict: + return { + "started_at": self.started_at, + "first_byte_at": self.first_byte_at, + "first_visible_output_at": self.first_visible_output_at, + "last_chunk_at": self.last_chunk_at, + "completed_at": self.completed_at, + "chunk_count": self.chunk_count, + "visible_chunk_count": self.visible_chunk_count, + "error_count": self.error_count, + "cancelled": self.cancelled, + "ttfb_seconds": self.ttfb_seconds, + "ttft_seconds": self.ttft_seconds, + "duration_seconds": self.duration_seconds, + "idle_seconds": self.idle_seconds, + } + + +class StreamMonitor: + """Record stream lifecycle events without changing stream behavior.""" + + def __init__(self, *, clock: Clock) -> None: + self._clock = clock + self.metrics = StreamMetrics(started_at=clock()) + + def record_event(self, event: StreamEvent) -> None: + now = self._clock() + if self.metrics.first_byte_at is None: + self.metrics.first_byte_at = now + self.metrics.last_chunk_at = now + self.metrics.chunk_count += 1 + if event.visible_output: + self.metrics.visible_chunk_count += 1 + if self.metrics.first_visible_output_at is None: + self.metrics.first_visible_output_at = now + if event.event_type == "error": + self.metrics.error_count += 1 + + def complete(self) -> None: + self.metrics.completed_at = self._clock() + + def cancel(self) -> None: + self.metrics.cancelled = True + self.metrics.completed_at = self._clock() + + def is_stalled(self, timeout_seconds: float) -> bool: + if timeout_seconds <= 0 or self.metrics.last_chunk_at is None: + return False + return self._clock() - self.metrics.last_chunk_at > timeout_seconds diff --git a/src/rotator_library/streaming/policy.py b/src/rotator_library/streaming/policy.py new file mode 100644 index 000000000..f2244c2b8 --- /dev/null +++ b/src/rotator_library/streaming/policy.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Stream retry and visible-output policy.""" + +from __future__ import annotations + +import json +from typing import Any, Optional + +_REASONING_FIELDS = ( + "reasoning", + "reasoning_content", + "thinking", + "thinking_content", +) + + +def can_retry_stream_after_error(last_streamed_chunk: Optional[str], allow_reasoning_only_retry: bool) -> bool: + """Return whether an upstream stream can be retried after an error.""" + + if last_streamed_chunk is None: + return True + if is_stream_heartbeat_or_comment(last_streamed_chunk): + return True + metadata = _sse_json(last_streamed_chunk, malformed_is_visible=False) + if isinstance(metadata, dict) and (metadata.get("event_type") or metadata.get("type")) == "cost": + return True + if metadata is None: + return False + if is_visible_stream_output(last_streamed_chunk): + return False + + has_reasoning = False + choices = metadata.get("choices") + if not isinstance(choices, list): + return True + for choice in choices: + if not isinstance(choice, dict): + return False + for source in (choice, choice.get("delta"), choice.get("message")): + if not isinstance(source, dict): + continue + if _has_visible_text(source.get("content")) or _has_visible_text(source.get("text")): + return False + if source.get("tool_calls") or source.get("function_call"): + return False + if any(_has_visible_text(source.get(key)) for key in _REASONING_FIELDS): + has_reasoning = True + return allow_reasoning_only_retry if has_reasoning else True + + +def is_visible_stream_output(chunk: Optional[str], *, protocol: str = "openai_chat") -> bool: + """Return whether a formatted stream chunk should block fallback. + + Malformed or ambiguous chunks fail closed by counting as visible output. This + preserves the existing safety rule that route fallback must not happen after + a client may have received model output. + """ + + if chunk is None: + return False + data = _sse_json(chunk, malformed_is_visible=True) + if data is _MALFORMED_VISIBLE: + return True + if data is None: + return False + if data.get("error"): + return False + event_type = data.get("event_type") or data.get("type") + if isinstance(event_type, str) and event_type.startswith("response."): + return _responses_visible(data) + if isinstance(event_type, str) and event_type.startswith("content_block_"): + return _anthropic_visible(data) + if isinstance(data.get("candidates"), list): + return _gemini_visible(data) + if protocol == "responses": + return _responses_visible(data) + return _openai_chat_visible(data) + + +def is_stream_heartbeat_or_comment(chunk: Optional[str]) -> bool: + """Return true for SSE comment-only frames that must not affect retry state.""" + + if chunk is None: + return False + payload = chunk.strip() + return bool(payload) and all(line.startswith(":") for line in payload.splitlines() if line.strip()) + + +_MALFORMED_VISIBLE = object() + + +def _sse_json(chunk: str, *, malformed_is_visible: bool) -> dict[str, Any] | object | None: + payload = chunk.strip() + if not payload: + return None + if all(line.startswith(":") for line in payload.splitlines() if line.strip()): + return None + event_type = None + data_lines: list[str] = [] + for line in payload.splitlines(): + stripped = line.strip() + if not stripped or stripped.startswith(":"): + continue + if stripped.startswith("event:"): + event_type = stripped[6:].strip() + continue + if stripped.startswith("data:"): + data_lines.append(stripped[5:].strip()) + continue + return _MALFORMED_VISIBLE if malformed_is_visible else None + if not data_lines: + if event_type in {"error", "response.failed"}: + return {"event_type": event_type} + return None + payload = "\n".join(data_lines).strip() + if not payload or payload == "[DONE]": + return None + try: + parsed = json.loads(payload) + except json.JSONDecodeError: + return _MALFORMED_VISIBLE if malformed_is_visible else None + if not isinstance(parsed, dict): + if event_type == "cost": + return {"event_type": "cost", "value": parsed} + return _MALFORMED_VISIBLE if malformed_is_visible else None + if event_type and "event_type" not in parsed: + parsed["event_type"] = event_type + return parsed + + +def _openai_chat_visible(data: dict[str, Any]) -> bool: + choices = data.get("choices") + if not isinstance(choices, list): + return False + for choice in choices: + if not isinstance(choice, dict): + continue + for source in (choice.get("delta"), choice.get("message")): + if not isinstance(source, dict): + continue + if _has_visible_text(source.get("content")) or _has_visible_text(source.get("text")): + return True + if source.get("tool_calls") or source.get("function_call"): + return True + return False + + +def _responses_visible(data: dict[str, Any]) -> bool: + event_type = data.get("event_type") or data.get("type") + if event_type == "response.output_text.delta": + return bool(str(data.get("delta", "")).strip()) + if event_type == "response.output_item.added": + item = data.get("item") + if isinstance(item, dict) and item.get("type") in {"function_call", "custom_tool_call"}: + return bool(item.get("call_id") or item.get("id") or item.get("name")) + if isinstance(event_type, str) and ("function_call" in event_type or "tool_call" in event_type): + return _has_visible_text(data.get("delta")) or _has_visible_text(data.get("arguments")) or _has_visible_text(data.get("item")) + if isinstance(event_type, str) and "reasoning" in event_type and event_type.endswith(".delta"): + return _has_visible_text(data.get("delta")) + if event_type == "response.failed": + return False + return False + + +def _anthropic_visible(data: dict[str, Any]) -> bool: + """Return whether an Anthropic content event exposes model output.""" + + event_type = data.get("event_type") or data.get("type") + if event_type == "content_block_start": + block = data.get("content_block") + return isinstance(block, dict) and block.get("type") == "tool_use" + if event_type != "content_block_delta": + return False + delta = data.get("delta") + if not isinstance(delta, dict): + return False + return any( + _has_visible_text(delta.get(key)) + for key in ("text", "thinking", "partial_json") + ) + + +def _gemini_visible(data: dict[str, Any]) -> bool: + """Return whether a Gemini candidate frame exposes content or a tool call.""" + + for candidate in data.get("candidates") or []: + if not isinstance(candidate, dict): + continue + content = candidate.get("content") + if not isinstance(content, dict): + continue + for part in content.get("parts") or []: + if not isinstance(part, dict): + continue + if _has_visible_text(part.get("text")) or part.get("functionCall") or part.get("function_call"): + return True + return False + + +def _has_visible_text(value: Any) -> bool: + if isinstance(value, str): + return bool(value.strip()) + if isinstance(value, list): + for item in value: + if isinstance(item, str) and item.strip(): + return True + if isinstance(item, dict) and _has_visible_text(item.get("text")): + return True + return False diff --git a/src/rotator_library/streaming/transport.py b/src/rotator_library/streaming/transport.py new file mode 100644 index 000000000..b7ebd1b43 --- /dev/null +++ b/src/rotator_library/streaming/transport.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Transport formatters for normalized stream events.""" + +from __future__ import annotations + +import json + +from .events import StreamEvent + + +class SSEStreamFormatter: + """Format stream events as HTTP Server-Sent Events.""" + + transport = "sse" + + def format_event(self, event: StreamEvent) -> str: + return f"event: {event.event_type}\ndata: {json.dumps(event.to_dict(), ensure_ascii=False)}\n\n" + + def format_error(self, event: StreamEvent) -> str: + return self.format_event(event) + + def format_done(self) -> str: + return "data: [DONE]\n\n" + + def format_heartbeat(self, comment: str = "heartbeat") -> str: + """Return an SSE comment heartbeat frame. + + Comment frames keep HTTP connections active without becoming model + output, so routing/session code must continue treating them as + non-visible stream data. + """ + + safe_comment = comment.replace("\r", " ").replace("\n", " ") + return f": {safe_comment}\n\n" + + def is_terminal_event(self, event: StreamEvent) -> bool: + return event.event_type in {"completed", "cancelled", "error"} + + +class WebSocketStreamFormatter: + """Format stream events as WebSocket JSON messages without exposing a route.""" + + transport = "websocket" + future_supported = True + + def format_event(self, event: StreamEvent) -> dict: + return {"type": event.event_type, "payload": event.to_dict()} + + def format_error(self, event: StreamEvent) -> dict: + return self.format_event(event) + + def format_done(self) -> dict: + return {"type": "completed", "payload": {"done": True}} + + def format_heartbeat(self) -> dict: + """Return the future WebSocket heartbeat message shape.""" + + return {"type": "heartbeat", "payload": {"visible_output": False}} + + def is_terminal_event(self, event: StreamEvent) -> bool: + return event.event_type in {"completed", "cancelled", "error"} + + +class JSONLineStreamFormatter: + """Format stream events as newline-delimited JSON for provider tests.""" + + transport = "jsonl" + + def format_event(self, event: StreamEvent) -> str: + return json.dumps(event.to_dict(), ensure_ascii=False) + "\n" + + def format_error(self, event: StreamEvent) -> str: + return self.format_event(event) + + def format_done(self) -> str: + return json.dumps({"event_type": "completed", "done": True}) + "\n" + + def format_heartbeat(self) -> str: + """Return a JSONL heartbeat record for test transports.""" + + return json.dumps({"event_type": "heartbeat", "visible_output": False}) + "\n" + + def is_terminal_event(self, event: StreamEvent) -> bool: + return event.event_type in {"completed", "cancelled", "error"} diff --git a/src/rotator_library/transaction_logger.py b/src/rotator_library/transaction_logger.py index 1a5fa6df8..b9b5a857a 100644 --- a/src/rotator_library/transaction_logger.py +++ b/src/rotator_library/transaction_logger.py @@ -28,11 +28,12 @@ import logging import time import uuid -from dataclasses import dataclass -from datetime import datetime +from dataclasses import dataclass, fields, is_dataclass +from datetime import UTC, date, datetime from pathlib import Path from typing import Any, Dict, Optional, Union +from .transform_trace import TransformTraceWriter, provider_snapshot_namespace from .utils.paths import get_logs_dir lib_logger = logging.getLogger("rotator_library") @@ -50,6 +51,77 @@ def _strip_framework_keys(data: Dict[str, Any]) -> Dict[str, Any]: return {k: v for k, v in data.items() if k not in FRAMEWORK_KEYS} +def _make_json_safe(value: Any, _seen: Optional[set[int]] = None) -> Any: + """Return a JSON-serializable copy of provider/client logging payloads. + + Provider implementations may hand transaction logging LiteLLM/Pydantic + models, dataclasses, paths, timestamps, bytes, or provider-specific helper + objects. Logging must never fail the request path, so this function converts + common structured objects to their dictionary form and falls back to strings + only at unknown leaves. + """ + + if value is None or isinstance(value, (str, int, float, bool)): + return value + + if isinstance(value, (datetime, date)): + return value.isoformat() + if isinstance(value, Path): + return str(value) + if isinstance(value, bytes): + try: + return value.decode("utf-8") + except UnicodeDecodeError: + return value.hex() + + seen = _seen if _seen is not None else set() + object_id = id(value) + if object_id in seen: + return "[CIRCULAR]" + + if isinstance(value, dict): + seen.add(object_id) + try: + return {str(k): _make_json_safe(v, seen) for k, v in value.items()} + finally: + seen.discard(object_id) + + if isinstance(value, (list, tuple, set, frozenset)): + seen.add(object_id) + try: + return [_make_json_safe(item, seen) for item in value] + finally: + seen.discard(object_id) + + if is_dataclass(value) and not isinstance(value, type): + seen.add(object_id) + try: + return { + field.name: _make_json_safe(getattr(value, field.name), seen) + for field in fields(value) + } + finally: + seen.discard(object_id) + + for method_name in ("model_dump", "dict", "to_dict"): + method = getattr(value, method_name, None) + if not callable(method): + continue + seen.add(object_id) + try: + try: + converted = method() + except TypeError: + converted = method(exclude_none=False) + return _make_json_safe(converted, seen) + except Exception: + continue + finally: + seen.discard(object_id) + + return str(value) + + def _get_transactions_dir() -> Path: """Get the transactions log directory, creating it if needed.""" logs_dir = get_logs_dir() @@ -58,6 +130,12 @@ def _get_transactions_dir() -> Path: return transactions_dir +def _utc_timestamp() -> str: + """Return an ISO-8601 UTC timestamp for log records.""" + + return datetime.now(UTC).isoformat() + + def _sanitize_name(name: str) -> str: """Sanitize a name for use in directory/file names.""" # Replace problematic characters with underscores @@ -85,11 +163,26 @@ class TransactionContext: """Whether logging is enabled.""" provider: str - """Provider name (e.g., 'gemini_cli', 'openai').""" + """Provider name (e.g., 'openai', 'anthropic').""" model: str """Model name (sanitized for filesystem use).""" + trace_model: Optional[str] = None + """Exact model name used in transform trace entries.""" + + session_id: Optional[str] = None + """Inferred session id for trace correlation, when available.""" + + scope_key: Optional[str] = None + """Usage scope key for trace correlation, when available.""" + + classifier: Optional[str] = None + """Classifier/private routing label for trace correlation, when available.""" + + trace_enabled: bool = False + """Whether provider loggers should append transform trace entries.""" + class TransactionLogger: """ @@ -112,10 +205,15 @@ class TransactionLogger: "request_id", "provider", "model", + "trace_model", + "session_id", + "scope_key", + "classifier", "streaming", "api_format", "_dir_available", "_context", + "_trace_writer", ) def __init__( @@ -130,7 +228,7 @@ def __init__( Initialize transaction logger. Args: - provider: Provider name (e.g., 'gemini_cli', 'openai') + provider: Provider name (e.g., 'openai', 'anthropic') model: Model name (will be sanitized for filesystem) enabled: Whether logging is enabled api_format: API format prefix ('oai' for OpenAI, 'ant' for Anthropic) @@ -140,10 +238,14 @@ def __init__( self.start_time = time.time() self.request_id = str(uuid.uuid4())[:8] # 8-char short ID self.provider = provider + self.trace_model = model + self.session_id: Optional[str] = None + self.scope_key: Optional[str] = None + self.classifier: Optional[str] = None self.api_format = api_format # Strip provider prefix from model if present - # e.g., "gemini_cli/gemini-2.5-pro" -> "gemini-2.5-pro" + # e.g., "openai/gpt-4.1" -> "gpt-4.1" model_name = model if "/" in model_name and model_name.split("/")[0] == provider: model_name = model_name.split("/", 1)[1] @@ -153,6 +255,7 @@ def __init__( self.log_dir: Optional[Path] = None self._dir_available = False self._context: Optional[TransactionContext] = None + self._trace_writer: Optional[TransformTraceWriter] = None if not enabled: return @@ -174,6 +277,14 @@ def __init__( try: self.log_dir.mkdir(parents=True, exist_ok=True) self._dir_available = True + self._trace_writer = TransformTraceWriter( + self.log_dir, + component="client", + provider=provider, + model=self.trace_model, + request_id=self.request_id, + enabled=True, + ) except Exception as e: lib_logger.error(f"TransactionLogger: Failed to create directory: {e}") self.enabled = False @@ -192,9 +303,104 @@ def get_context(self) -> TransactionContext: enabled=self.enabled, provider=self.provider, model=self.model, + trace_model=self.trace_model, + session_id=self.session_id, + scope_key=self.scope_key, + classifier=self.classifier, + trace_enabled=bool(self._trace_writer), ) return self._context + def set_trace_context( + self, + *, + session_id: Optional[str] = None, + scope_key: Optional[str] = None, + classifier: Optional[str] = None, + ) -> None: + """Attach routing/session metadata discovered after logger creation.""" + + if session_id is not None: + self.session_id = session_id + if scope_key is not None: + self.scope_key = scope_key + if classifier is not None: + self.classifier = classifier + if self._trace_writer: + self._trace_writer.update_context( + session_id=self.session_id, + scope_key=self.scope_key, + classifier=self.classifier, + ) + if self._context: + self._context.session_id = self.session_id + self._context.scope_key = self.scope_key + self._context.classifier = self.classifier + + def log_transform_pass( + self, + pass_name: str, + data: Any, + *, + direction: str, + stage: str, + protocol: Optional[str] = None, + credential_id: Optional[str] = None, + transport: Optional[str] = None, + changed_from_previous: Optional[bool] = None, + metadata: Optional[Dict[str, Any]] = None, + scrub_strings: bool = False, + snapshot: bool = True, + ) -> None: + """Record an additive transform trace entry if tracing is available.""" + + if not self.enabled or not self._dir_available or not self._trace_writer: + return + self._trace_writer.record( + pass_name, + data, + direction=direction, + stage=stage, + protocol=protocol, + credential_id=credential_id, + transport=transport, + changed_from_previous=changed_from_previous, + metadata=metadata, + scrub_strings=scrub_strings, + snapshot=snapshot, + ) + + def log_transform_error( + self, + failed_pass_name: str, + error: BaseException, + *, + payload: Any = None, + stage: str = "client", + protocol: Optional[str] = None, + transport: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: + """Record a standardized transform/logging failure without raising.""" + + error_data = { + "failed_pass_name": failed_pass_name, + "error_type": type(error).__name__, + "message": str(error), + "payload": payload, + } + self.log_transform_pass( + "transform_log_error", + error_data, + direction="error", + stage=stage, + protocol=protocol, + transport=transport, + metadata=metadata, + scrub_strings=True, + snapshot=False, + ) + def log_request( self, request_data: Dict[str, Any], filename: str = "request.json" ) -> None: @@ -212,15 +418,25 @@ def log_request( data = { "request_id": self.request_id, - "timestamp_utc": datetime.utcnow().isoformat(), + "timestamp_utc": _utc_timestamp(), "data": request_data, } + self.log_transform_pass( + "raw_client_request", + request_data, + direction="request", + stage="client", + transport="sse" if self.streaming else "http", + ) self._write_json(filename, data) def log_transformed_request( self, transformed_data: Dict[str, Any], original_data: Dict[str, Any], + *, + credential_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, ) -> None: """ Log the transformed request if it differs from the original. @@ -239,18 +455,32 @@ def log_transformed_request( stripped_transformed = _strip_framework_keys(transformed_data) stripped_original = _strip_framework_keys(original_data) + changed_from_previous: Optional[bool] = None try: - if json.dumps(stripped_transformed, sort_keys=True, default=str) == json.dumps( + changed_from_previous = json.dumps(stripped_transformed, sort_keys=True, default=str) != json.dumps( stripped_original, sort_keys=True, default=str - ): - return + ) except (TypeError, ValueError): - pass + changed_from_previous = None + + self.log_transform_pass( + "prepared_provider_request", + transformed_data, + direction="request", + stage="client", + credential_id=credential_id, + transport="sse" if transformed_data.get("stream") else "http", + changed_from_previous=changed_from_previous, + metadata=metadata, + ) + + if changed_from_previous is False: + return logged = _strip_framework_keys(transformed_data) data = { "request_id": self.request_id, - "timestamp_utc": datetime.utcnow().isoformat(), + "timestamp_utc": _utc_timestamp(), "data": logged, } self._write_json("request_transformed.json", data) @@ -266,9 +496,17 @@ def log_stream_chunk(self, chunk: Dict[str, Any]) -> None: return log_entry = { - "timestamp_utc": datetime.utcnow().isoformat(), - "chunk": chunk, + "timestamp_utc": _utc_timestamp(), + "chunk": _make_json_safe(chunk), } + self.log_transform_pass( + "parsed_stream_chunk", + chunk, + direction="stream", + stage="client", + transport="sse", + snapshot=False, + ) content = json.dumps(log_entry, ensure_ascii=False) + "\n" self._append_text("streaming_chunks.jsonl", content) @@ -294,29 +532,44 @@ def log_response( end_time = time.time() duration_ms = (end_time - self.start_time) * 1000 + safe_response = _make_json_safe(response_data) data = { "request_id": self.request_id, - "timestamp_utc": datetime.utcnow().isoformat(), + "timestamp_utc": _utc_timestamp(), "status_code": status_code, "duration_ms": round(duration_ms), - "headers": dict(headers) if headers else None, - "data": response_data, + "headers": _make_json_safe(dict(headers)) if headers else None, + "data": safe_response, } + self.log_transform_pass( + "final_client_response", + response_data, + direction="response", + stage="final", + transport="sse" if self.streaming else "http", + metadata={"status_code": status_code, "headers": dict(headers) if headers else None}, + ) self._write_json(filename, data) # Also write metadata - self._log_metadata(response_data, status_code, duration_ms) + self._log_metadata(safe_response, status_code, duration_ms) def _log_metadata( self, response_data: Dict[str, Any], status_code: int, duration_ms: float ) -> None: """Log transaction metadata summary.""" + if not isinstance(response_data, dict): + response_data = {"value": response_data} usage = response_data.get("usage") or {} + if not isinstance(usage, dict): + usage = {} model = response_data.get("model", self.model) finish_reason = "N/A" if "choices" in response_data and response_data["choices"]: - finish_reason = response_data["choices"][0].get("finish_reason", "N/A") + first_choice = response_data["choices"][0] + if isinstance(first_choice, dict): + finish_reason = first_choice.get("finish_reason", "N/A") # Check for provider subdirectory has_provider_logs = False @@ -331,7 +584,7 @@ def _log_metadata( metadata = { "request_id": self.request_id, - "timestamp_utc": datetime.utcnow().isoformat(), + "timestamp_utc": _utc_timestamp(), "duration_ms": round(duration_ms), "status_code": status_code, "provider": self.provider, @@ -365,7 +618,12 @@ def _extract_reasoning(self, response_data: Dict[str, Any]) -> Optional[str]: return response_data["reasoning"] if "choices" in response_data and response_data["choices"]: - message = response_data["choices"][0].get("message", {}) + first_choice = response_data["choices"][0] + if not isinstance(first_choice, dict): + return None + message = first_choice.get("message", {}) + if not isinstance(message, dict): + return None if "reasoning" in message: return message["reasoning"] if "reasoning_content" in message: @@ -379,7 +637,7 @@ def _write_json(self, filename: str, data: Dict[str, Any]) -> None: return try: with open(self.log_dir / filename, "w", encoding="utf-8") as f: - json.dump(data, f, indent=2, ensure_ascii=False) + json.dump(_make_json_safe(data), f, indent=2, ensure_ascii=False) except Exception as e: lib_logger.error(f"TransactionLogger: Failed to write {filename}: {e}") @@ -543,7 +801,7 @@ class ProviderLogger: or add custom methods for provider-specific logging needs. """ - __slots__ = ("enabled", "log_dir") + __slots__ = ("enabled", "log_dir", "_trace_writer") def __init__(self, context: Optional[TransactionContext]): """ @@ -554,6 +812,7 @@ def __init__(self, context: Optional[TransactionContext]): """ self.enabled = False self.log_dir: Optional[Path] = None + self._trace_writer: Optional[TransformTraceWriter] = None if context is None or not context.enabled: return @@ -563,10 +822,46 @@ def __init__(self, context: Optional[TransactionContext]): try: self.log_dir.mkdir(parents=True, exist_ok=True) + if getattr(context, "trace_enabled", False): + self._trace_writer = TransformTraceWriter( + context.log_dir, + component="provider", + provider=context.provider, + model=context.trace_model or context.model, + request_id=context.request_id, + session_id=context.session_id, + scope_key=context.scope_key, + classifier=context.classifier, + snapshot_namespace=provider_snapshot_namespace(), + enabled=True, + ) except Exception as e: lib_logger.error(f"ProviderLogger: Failed to create directory: {e}") self.enabled = False + def _log_transform_pass( + self, + pass_name: str, + data: Any, + *, + direction: str, + stage: str = "provider", + transport: Optional[str] = None, + scrub_strings: bool = False, + snapshot: bool = True, + ) -> None: + if not self.enabled or not self._trace_writer: + return + self._trace_writer.record( + pass_name, + data, + direction=direction, + stage=stage, + transport=transport, + scrub_strings=scrub_strings, + snapshot=snapshot, + ) + def log_request(self, payload: Dict[str, Any]) -> None: """ Log the request payload sent to the provider API. @@ -574,6 +869,12 @@ def log_request(self, payload: Dict[str, Any]) -> None: Args: payload: The transformed request payload """ + self._log_transform_pass( + "provider_request_payload", + payload, + direction="request", + transport="http", + ) self._write_json("request_payload.json", payload) def log_response_chunk(self, chunk: str) -> None: @@ -583,6 +884,13 @@ def log_response_chunk(self, chunk: str) -> None: Args: chunk: Raw chunk string from the stream """ + self._log_transform_pass( + "provider_raw_stream_chunk", + chunk, + direction="stream", + transport="sse", + snapshot=False, + ) self._append_text("response_stream.log", chunk + "\n") def log_final_response(self, response_data: Dict[str, Any]) -> None: @@ -592,6 +900,11 @@ def log_final_response(self, response_data: Dict[str, Any]) -> None: Args: response_data: The complete response data """ + self._log_transform_pass( + "provider_final_response", + response_data, + direction="response", + ) self._write_json("final_response.json", response_data) def log_error(self, error_message: str) -> None: @@ -601,7 +914,14 @@ def log_error(self, error_message: str) -> None: Args: error_message: The error message to log """ - timestamp = datetime.utcnow().isoformat() + timestamp = _utc_timestamp() + self._log_transform_pass( + "provider_error", + {"timestamp_utc": timestamp, "message": error_message}, + direction="error", + scrub_strings=True, + snapshot=False, + ) self._append_text("error.log", f"[{timestamp}] {error_message}\n") def log_extra(self, filename: str, data: Union[Dict[str, Any], str]) -> None: @@ -625,7 +945,7 @@ def _write_json(self, filename: str, data: Dict[str, Any]) -> None: return try: with open(self.log_dir / filename, "w", encoding="utf-8") as f: - json.dump(data, f, indent=2, ensure_ascii=False) + json.dump(_make_json_safe(data), f, indent=2, ensure_ascii=False) except Exception as e: lib_logger.error(f"ProviderLogger: Failed to write {filename}: {e}") diff --git a/src/rotator_library/transform_trace.py b/src/rotator_library/transform_trace.py new file mode 100644 index 000000000..513e37a56 --- /dev/null +++ b/src/rotator_library/transform_trace.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Transform-pass trace primitives for transaction logging. + +The trace layer is observability-only: failures here must never change request +execution. Transaction logging uses this module to snapshot each meaningful +request, response, and stream state as later protocol/adapters mutate payloads. +""" + +from __future__ import annotations + +import json +import logging +import re +import uuid +from dataclasses import dataclass, field +from dataclasses import is_dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Mapping, Optional + +from .protocols import serialize_value + +lib_logger = logging.getLogger("rotator_library") + +REDACTED = "[REDACTED]" + +_SENSITIVE_KEYS = frozenset( + { + "api-key", + "credential-identifier", + "authorization", + "proxy-authorization", + "cookie", + "set-cookie", + "x-api-key", + "x-goog-api-key", + "openai-api-key", + "api-key", + "api-secret", + "access-token", + "refresh-token", + "id-token", + "client-secret", + "password", + "secret", + "token", + # Provider protocol-state fields can contain opaque signatures or + # reasoning continuations that field-cache rules preserve and re-inject. + "reasoning-content", + "reasoningcontent", + "thought-signature", + "thoughtsignature", + "signature", + "state", + "provider-state", + "provider-session-id", + "providersessionid", + "prompt-cache-key", + "promptcachekey", + "cache-key", + "cachekey", + "thinking-signatures", + "thinkingsignatures", + "thought-signatures", + "thoughtsignatures", + } +) + +_SENSITIVE_TEXT_RE = re.compile( + r"(?im)\b(authorization|proxy-authorization|x-api-key|x-goog-api-key|openai-api-key|api[_-]?key|access[_-]?token|refresh[_-]?token|client[_-]?secret|cookie|set-cookie)\b(['\"]?\s*[:=]\s*['\"]?)([^'\"\r\n,}]+)" +) + + +def _normalise_key(key: Any) -> str: + camel_split = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "-", str(key).strip()) + return camel_split.lower().replace("_", "-") + + +def scrub_sensitive_text(value: str) -> str: + """Scrub obvious credential-bearing header/query fragments from text. + + General trace payloads use key-based redaction to avoid hiding model text. + Error strings and object reprs are different: providers often embed HTTP + headers or query strings in exception text, so this targeted scrub only + applies when callers opt into string scrubbing. + """ + + return _SENSITIVE_TEXT_RE.sub(lambda match: f"{match.group(1)}{match.group(2)}{REDACTED}", value) + + +def _object_mapping(value: Any) -> Optional[dict[str, Any]]: + """Extract structured data from common SDK objects before using repr().""" + + for method_name in ("model_dump", "dict"): + method = getattr(value, method_name, None) + if callable(method): + try: + result = method() + if isinstance(result, Mapping): + return dict(result) + except Exception: + pass + structured: dict[str, Any] = {"type": f"{type(value).__module__}.{type(value).__name__}"} + found = False + for attr in ("status_code", "headers", "url", "method", "text", "content"): + if hasattr(value, attr): + try: + structured[attr] = getattr(value, attr) + found = True + except Exception: + pass + if found: + return structured + if hasattr(value, "__dict__") and not is_dataclass(value): + try: + return {"type": structured["type"], "attributes": vars(value)} + except Exception: + return None + return None + + +def sanitize_for_trace(value: Any, *, scrub_strings: bool = False) -> Any: + """Return a JSON-safe, recursively redacted value for trace files. + + Redaction is intentionally key-based rather than value-based. Model text may + legitimately mention tokens, passwords, or secrets; hiding by value would make + debugging transformations unreliable. Sensitive framework fields are redacted + only when their key name is known to carry credentials. + """ + + if isinstance(value, Mapping): + sanitized = {} + block_type = str(value.get("type", "")).lower() + source_field = str(value.get("source_field", "")).lower() + if not source_field and isinstance(value.get("extra"), Mapping): + source_field = str(value["extra"].get("source_field", "")).lower() + is_reasoning_block = "reasoning" in block_type or source_field in {"reasoning_content", "thoughtsignature", "thought_signature", "signature"} + for key, item in value.items(): + if _normalise_key(key) in _SENSITIVE_KEYS or (is_reasoning_block and _normalise_key(key) == "text"): + sanitized[str(key)] = REDACTED + else: + sanitized[str(key)] = sanitize_for_trace(item, scrub_strings=scrub_strings) + return sanitized + if isinstance(value, (list, tuple, set, frozenset)): + return [sanitize_for_trace(item, scrub_strings=scrub_strings) for item in value] + if isinstance(value, str): + return scrub_sensitive_text(value) if scrub_strings else value + + object_mapping = _object_mapping(value) + if object_mapping is not None: + return sanitize_for_trace(object_mapping, scrub_strings=scrub_strings) + + serialized = serialize_value(value) + if isinstance(serialized, dict): + return sanitize_for_trace(serialized, scrub_strings=scrub_strings) + if isinstance(serialized, list): + return [sanitize_for_trace(item, scrub_strings=scrub_strings) for item in serialized] + if isinstance(serialized, str): + return scrub_sensitive_text(serialized) if scrub_strings else serialized + return serialized + + +def sanitize_filename(value: str) -> str: + """Return a stable, filesystem-safe name component for snapshot files.""" + + safe = value.strip().lower().replace(" ", "_") or "trace" + for char in '/\\:*?"<>|': + safe = safe.replace(char, "_") + return "".join(char if char.isalnum() or char in "._-" else "_" for char in safe) + + +@dataclass +class TransformTraceEntry: + """A single transform-pass observation. + + Entries are deliberately protocol-neutral. Later phases can add protocol, + adapter, field-cache, routing, and transport passes without changing the file + format used by this phase. + """ + + sequence: int + component: str + pass_name: str + direction: str + stage: str + request_id: Optional[str] = None + timestamp_utc: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) + protocol: Optional[str] = None + provider: Optional[str] = None + model: Optional[str] = None + credential_id: Optional[str] = None + transport: Optional[str] = None + changed_from_previous: Optional[bool] = None + session_id: Optional[str] = None + scope_key: Optional[str] = None + classifier: Optional[str] = None + metadata: dict[str, Any] = field(default_factory=dict) + data: Any = None + scrub_strings: bool = False + + def to_dict(self) -> dict[str, Any]: + return { + "sequence": self.sequence, + "component": self.component, + "pass_name": self.pass_name, + "direction": self.direction, + "stage": self.stage, + "request_id": self.request_id, + "timestamp_utc": self.timestamp_utc, + "protocol": self.protocol, + "provider": self.provider, + "model": self.model, + "credential_id": self.credential_id, + "transport": self.transport, + "changed_from_previous": self.changed_from_previous, + "session_id": self.session_id, + "scope_key": self.scope_key, + "classifier": self.classifier, + "metadata": sanitize_for_trace(self.metadata, scrub_strings=self.scrub_strings), + "data": sanitize_for_trace(self.data, scrub_strings=self.scrub_strings), + } + + +class TransformTraceWriter: + """Append-only writer for transform trace entries and snapshots. + + One writer owns one local sequence counter. Phase 2 intentionally does not + promise a global ordering across client and provider writers; entries include + component and timestamps so interleaved logs remain understandable. + """ + + def __init__( + self, + log_dir: Path, + *, + component: str, + provider: Optional[str] = None, + model: Optional[str] = None, + request_id: Optional[str] = None, + session_id: Optional[str] = None, + scope_key: Optional[str] = None, + classifier: Optional[str] = None, + snapshot_namespace: Optional[str] = None, + enabled: bool = True, + ) -> None: + self.log_dir = log_dir + self.component = component + self.provider = provider + self.model = model + self.request_id = request_id + self.session_id = session_id + self.scope_key = scope_key + self.classifier = classifier + self.snapshot_namespace = snapshot_namespace + self.enabled = enabled + self._sequence = 0 + self.trace_file = log_dir / "transform_trace.jsonl" + self.snapshot_dir = log_dir / "transforms" + + def update_context( + self, + *, + session_id: Optional[str] = None, + scope_key: Optional[str] = None, + classifier: Optional[str] = None, + ) -> None: + """Update immutable-ish correlation fields discovered after creation.""" + + if session_id is not None: + self.session_id = session_id + if scope_key is not None: + self.scope_key = scope_key + if classifier is not None: + self.classifier = classifier + + def record( + self, + pass_name: str, + data: Any, + *, + direction: str, + stage: str, + protocol: Optional[str] = None, + credential_id: Optional[str] = None, + transport: Optional[str] = None, + changed_from_previous: Optional[bool] = None, + session_id: Optional[str] = None, + scope_key: Optional[str] = None, + classifier: Optional[str] = None, + metadata: Optional[dict[str, Any]] = None, + scrub_strings: bool = False, + snapshot: bool = True, + ) -> Optional[TransformTraceEntry]: + """Record a transform pass, swallowing logging failures.""" + + if not self.enabled: + return None + self._sequence += 1 + entry = TransformTraceEntry( + sequence=self._sequence, + component=self.component, + pass_name=pass_name, + direction=direction, + stage=stage, + request_id=self.request_id, + protocol=protocol, + provider=self.provider, + model=self.model, + credential_id=credential_id, + transport=transport, + changed_from_previous=changed_from_previous, + session_id=session_id if session_id is not None else self.session_id, + scope_key=scope_key if scope_key is not None else self.scope_key, + classifier=classifier if classifier is not None else self.classifier, + metadata=metadata or {}, + data=data, + scrub_strings=scrub_strings, + ) + try: + self.log_dir.mkdir(parents=True, exist_ok=True) + with open(self.trace_file, "a", encoding="utf-8") as handle: + handle.write(json.dumps(entry.to_dict(), ensure_ascii=False) + "\n") + if snapshot and direction != "stream": + self.snapshot_dir.mkdir(parents=True, exist_ok=True) + namespace = f"{sanitize_filename(self.snapshot_namespace)}_" if self.snapshot_namespace else "" + snapshot_name = f"{entry.sequence:04d}_{namespace}{sanitize_filename(pass_name)}.json" + with open(self.snapshot_dir / snapshot_name, "w", encoding="utf-8") as handle: + json.dump(entry.to_dict(), handle, indent=2, ensure_ascii=False) + except Exception as exc: + lib_logger.debug("Transform trace write failed for %s: %s", pass_name, exc) + return entry + + +def provider_snapshot_namespace() -> str: + """Return a short namespace that prevents provider snapshot collisions.""" + + return f"provider_{uuid.uuid4().hex[:8]}" diff --git a/src/rotator_library/usage/__init__.py b/src/rotator_library/usage/__init__.py index f8ae947fe..2ab7f5bd8 100644 --- a/src/rotator_library/usage/__init__.py +++ b/src/rotator_library/usage/__init__.py @@ -56,11 +56,21 @@ # Main facade (imports components above) from .manager import UsageManager, CredentialContext +from .accounting import UsageRecord, extract_usage_record +from .costs import CostBreakdown, CostCalculator, ModelPricing +from .quota import QuotaSnapshot, build_quota_snapshots __all__ = [ # Main public API "UsageManager", "CredentialContext", + "UsageRecord", + "CostBreakdown", + "CostCalculator", + "ModelPricing", + "QuotaSnapshot", + "build_quota_snapshots", + "extract_usage_record", # Types "WindowStats", "TotalStats", diff --git a/src/rotator_library/usage/accounting.py b/src/rotator_library/usage/accounting.py new file mode 100644 index 000000000..59e6a727b --- /dev/null +++ b/src/rotator_library/usage/accounting.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Normalized usage accounting across provider protocols. + +The existing `UsageManager` remains the source of truth for persistence, +windows, and credential selection. This module only converts provider-specific +usage payloads into the numeric buckets that `CredentialContext.mark_success()` +already accepts. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Optional + +from ..protocols import serialize_value + + +@dataclass(frozen=True) +class UsageRecord: + """Provider-neutral token usage buckets. + + `completion_tokens` excludes `reasoning_tokens` when providers report hidden + reasoning separately. This preserves the current no-double-count behavior and + gives Phase 9+ cost logic clear input/output/reasoning buckets. + """ + + input_tokens: int = 0 + completion_tokens: int = 0 + reasoning_tokens: int = 0 + cache_read_tokens: int = 0 + cache_write_tokens: int = 0 + raw_total_tokens: int = 0 + request_count: int = 1 + source: str = "unknown" + provider: Optional[str] = None + model: Optional[str] = None + provider_reported_cost: Optional[float] = None + cost_currency: str = "USD" + cost_source: Optional[str] = None + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def output_tokens(self) -> int: + return self.completion_tokens + self.reasoning_tokens + + @property + def total_tokens(self) -> int: + return ( + self.input_tokens + + self.cache_read_tokens + + self.cache_write_tokens + + self.completion_tokens + + self.reasoning_tokens + ) + + @property + def prompt_tokens_for_mark_success(self) -> int: + """Return non-cache-read prompt tokens for existing usage storage.""" + + return self.input_tokens + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe summary for transform traces and tests.""" + + return { + "input_tokens": self.input_tokens, + "completion_tokens": self.completion_tokens, + "reasoning_tokens": self.reasoning_tokens, + "cache_read_tokens": self.cache_read_tokens, + "cache_write_tokens": self.cache_write_tokens, + "output_tokens": self.output_tokens, + "total_tokens": self.total_tokens, + "raw_total_tokens": self.raw_total_tokens, + "request_count": self.request_count, + "source": self.source, + "provider": self.provider, + "model": self.model, + "provider_reported_cost": self.provider_reported_cost, + "cost_currency": self.cost_currency, + "cost_source": self.cost_source, + "metadata": serialize_value(self.metadata), + } + + +def extract_usage_record( + response_or_usage: Any, + *, + provider: Optional[str] = None, + model: Optional[str] = None, + source: str = "response", +) -> UsageRecord: + """Extract normalized usage from a response object or usage payload.""" + + usage = _unwrap_usage(response_or_usage) + if usage is None: + return UsageRecord(provider=provider, model=model, source=source) + data = _as_dict(usage) + if not data: + return UsageRecord(provider=provider, model=model, source=source) + + if "usageMetadata" in data and isinstance(data["usageMetadata"], dict): + data = data["usageMetadata"] + + if _looks_like_gemini(data): + return _from_gemini_usage(data, provider=provider, model=model, source=source) + if _looks_like_anthropic(data): + return _from_anthropic_usage(data, provider=provider, model=model, source=source) + return _from_openai_like_usage(data, provider=provider, model=model, source=source) + + +def _unwrap_usage(value: Any) -> Any: + if value is None: + return None + if isinstance(value, dict): + if "usage" in value: + usage = _as_dict(value.get("usage")) + return _merge_top_level_cost_fields(usage, value) + return value + data = _as_dict(value) + if "usage" in data: + usage = _as_dict(data.get("usage")) + return _merge_top_level_cost_fields(usage, data) + return getattr(value, "usage", value) + + +def _merge_top_level_cost_fields(usage: dict[str, Any], response: dict[str, Any]) -> dict[str, Any]: + """Copy sibling cost metadata into a nested usage payload.""" + + merged = dict(usage) + for key in ("cost", "total_cost", "estimated_cost", "cost_details", "provider_reported_cost", "request_cost_usd", "currency", "costMetadata"): + if key in response and key not in merged: + merged[key] = response[key] + return merged + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + if hasattr(value, "to_dict"): + dumped = value.to_dict() + return dumped if isinstance(dumped, dict) else {} + if hasattr(value, "model_dump"): + dumped = value.model_dump() + return dumped if isinstance(dumped, dict) else {} + if hasattr(value, "dict"): + dumped = value.dict() + return dumped if isinstance(dumped, dict) else {} + result: dict[str, Any] = {} + for key in ( + "prompt_tokens", + "usage", + "completion_tokens", + "total_tokens", + "prompt_tokens_details", + "completion_tokens_details", + "cache_read_tokens", + "cache_creation_tokens", + "input_tokens", + "output_tokens", + "input_tokens_details", + "output_tokens_details", + "cache_creation_input_tokens", + "cache_read_input_tokens", + "cached_tokens", + "cache_creation_tokens", + "cache_write_tokens", + "reasoning_tokens", + "thinking_tokens", + "cost", + "cost_details", + "costMetadata", + "provider_reported_cost", + "request_cost_usd", + "estimated_cost", + "total_cost", + "currency", + ): + if hasattr(value, key): + result[key] = getattr(value, key) + return result + + +def _from_openai_like_usage(data: dict[str, Any], *, provider: Optional[str], model: Optional[str], source: str) -> UsageRecord: + prompt_tokens = _int(data.get("prompt_tokens", data.get("input_tokens", 0))) + completion_tokens = _int(data.get("completion_tokens", data.get("output_tokens", 0))) + raw_total = _int(data.get("total_tokens", data.get("raw_total_tokens", 0))) + + prompt_details = _as_dict(data.get("prompt_tokens_details") or data.get("input_tokens_details") or {}) + completion_details = _as_dict(data.get("completion_tokens_details") or data.get("output_tokens_details") or {}) + cache_read = _int( + data.get( + "cache_read_tokens", + data.get("cached_tokens", prompt_details.get("cached_tokens", 0)), + ) + ) + cache_write = _int( + data.get( + "cache_creation_tokens", + data.get("cache_write_tokens", prompt_details.get("cache_creation_tokens", 0)), + ) + ) + reasoning = _int( + data.get( + "reasoning_tokens", + completion_details.get("reasoning_tokens", completion_details.get("thinking_tokens", 0)), + ) + ) + if reasoning and completion_tokens >= reasoning: + completion_tokens -= reasoning + input_tokens = max(0, prompt_tokens - cache_read - cache_write) + cost = _extract_cost(data) + return UsageRecord( + input_tokens=input_tokens, + completion_tokens=completion_tokens, + reasoning_tokens=reasoning, + cache_read_tokens=cache_read, + cache_write_tokens=cache_write, + raw_total_tokens=raw_total, + source=source, + provider=provider, + model=model, + provider_reported_cost=cost[0], + cost_currency=cost[1], + cost_source=cost[2], + metadata={"shape": "openai_like"}, + ) + + +def _from_anthropic_usage(data: dict[str, Any], *, provider: Optional[str], model: Optional[str], source: str) -> UsageRecord: + cache_read = _int(data.get("cache_read_input_tokens", data.get("cache_read_tokens", 0))) + cache_write = _int(data.get("cache_creation_input_tokens", data.get("cache_creation_tokens", 0))) + input_tokens = max(0, _int(data.get("input_tokens", 0)) - cache_read - cache_write) + output_tokens = _int(data.get("output_tokens", data.get("completion_tokens", 0))) + reasoning = _int(data.get("reasoning_tokens", data.get("thinking_tokens", 0))) + if reasoning and output_tokens >= reasoning: + output_tokens -= reasoning + cost = _extract_cost(data) + return UsageRecord( + input_tokens=input_tokens, + completion_tokens=output_tokens, + reasoning_tokens=reasoning, + cache_read_tokens=cache_read, + cache_write_tokens=cache_write, + raw_total_tokens=_int(data.get("total_tokens", 0)), + source=source, + provider=provider, + model=model, + provider_reported_cost=cost[0], + cost_currency=cost[1], + cost_source=cost[2], + metadata={"shape": "anthropic"}, + ) + + +def _from_gemini_usage(data: dict[str, Any], *, provider: Optional[str], model: Optional[str], source: str) -> UsageRecord: + cache_read = _int(data.get("cachedContentTokenCount", data.get("cache_read_tokens", 0))) + prompt_tokens = _int(data.get("promptTokenCount", data.get("prompt_tokens", 0))) + reasoning = _int(data.get("thoughtsTokenCount", data.get("reasoning_tokens", 0))) + completion = _int(data.get("candidatesTokenCount", data.get("completion_tokens", 0))) + if reasoning and completion >= reasoning: + completion -= reasoning + cost = _extract_cost(data) + return UsageRecord( + input_tokens=max(0, prompt_tokens - cache_read), + completion_tokens=completion, + reasoning_tokens=reasoning, + cache_read_tokens=cache_read, + raw_total_tokens=_int(data.get("totalTokenCount", data.get("total_tokens", 0))), + source=source, + provider=provider, + model=model, + provider_reported_cost=cost[0], + cost_currency=cost[1], + cost_source=cost[2], + metadata={"shape": "gemini"}, + ) + + +def _extract_cost(data: dict[str, Any]) -> tuple[Optional[float], str, Optional[str]]: + """Extract actual provider-reported cost without guessing prices. + + Advisory pricing belongs in `usage.costs`. This helper only preserves cost + values explicitly reported by a provider or protocol adapter. + """ + + cost_payload = _as_dict(data.get("cost") or data.get("cost_details") or data.get("costMetadata") or {}) + raw_cost = ( + cost_payload.get("provider_reported_cost", cost_payload.get("request_cost_usd")) + if cost_payload + else None + ) + breakdown_used = False + if raw_cost is None and cost_payload: + raw_cost = cost_payload.get("total_cost", cost_payload.get("total")) + if raw_cost is None and cost_payload: + raw_cost = cost_payload.get("cost") + if raw_cost is None and cost_payload: + breakdown_total = _sum_cost_breakdown(cost_payload) + if breakdown_total is not None: + raw_cost = breakdown_total + breakdown_used = True + if raw_cost is None: + raw_cost = data.get("provider_reported_cost", data.get("request_cost_usd", data.get("total_cost", data.get("cost", data.get("estimated_cost"))))) + cost_value = _float_or_none(raw_cost) + currency = str(cost_payload.get("currency") or data.get("currency") or "USD") + source = cost_payload.get("source") or ("provider_reported_breakdown" if breakdown_used else ("provider_reported" if cost_value is not None else None)) + return cost_value, currency, str(source) if source else None + + +def _sum_cost_breakdown(payload: dict[str, Any]) -> Optional[float]: + """Sum structured provider cost buckets when no total is reported.""" + + total = 0.0 + found = False + upstream_total = _float_or_none(payload.get("upstream_inference_cost")) + if upstream_total is not None: + total += upstream_total + found = True + else: + for key in ( + "upstream_inference_prompt_cost", + "upstream_inference_completions_cost", + "upstream_inference_input_cost", + "upstream_inference_output_cost", + ): + value = _float_or_none(payload.get(key)) + if value is not None: + total += value + found = True + for key in ( + "input_cost", + "prompt_cost", + "cache_read_cost", + "cache_write_cost", + "cached_input_cost", + "cache_write_input_cost", + "output_cost", + "completion_cost", + "completions_cost", + "reasoning_cost", + "thinking_cost", + "image_input_cost", + "image_output_cost", + "audio_input_cost", + "audio_output_cost", + "data_storage_cost", + "estimated_cost", + "request_cost", + "web_search_cost", + "search_cost", + ): + value = _float_or_none(payload.get(key)) + if value is not None: + total += value + found = True + ticks = _float_or_none(payload.get("cost_in_usd_ticks")) + if ticks is not None: + total += ticks / 10_000_000_000 + found = True + return total if found else None + + +def _looks_like_gemini(data: dict[str, Any]) -> bool: + return any(key in data for key in ("promptTokenCount", "candidatesTokenCount", "thoughtsTokenCount", "cachedContentTokenCount")) + + +def _looks_like_anthropic(data: dict[str, Any]) -> bool: + return any(key in data for key in ("cache_creation_input_tokens", "cache_read_input_tokens")) and "input_tokens" in data + + +def _int(value: Any) -> int: + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 + + +def _float_or_none(value: Any) -> Optional[float]: + if value is None: + return None + try: + return max(0.0, float(value)) + except (TypeError, ValueError): + return None diff --git a/src/rotator_library/usage/config.py b/src/rotator_library/usage/config.py index b9cd57a18..dd8f7c9c1 100644 --- a/src/rotator_library/usage/config.py +++ b/src/rotator_library/usage/config.py @@ -510,6 +510,12 @@ class ProviderUsageConfig: # Exhaustion threshold (cooldown must exceed this to count as "exhausted") exhaustion_cooldown_threshold: int = DEFAULT_EXHAUSTION_COOLDOWN_THRESHOLD + # Policy for authoritative quota APIs that report exhaustion without a reset + # timestamp. Defaults to legacy warn-only behavior; providers with known + # entitlement-style semantics can opt into a fallback cooldown. + no_reset_exhaustion_policy: str = "warn_only" + no_reset_exhaustion_cooldown_seconds: int = 0 + # Window limits blocking (if True, block credentials when window quota exhausted locally) # Default False: only API errors (cooldowns) should block, not local tracking window_limits_enabled: bool = False @@ -521,7 +527,9 @@ class ProviderUsageConfig: # is only blocked by concurrency. Configuring this keeps cache-locality # deployments from forcing the same latency tradeoff on low-latency ones. session_sticky_wait_seconds: float = 15.0 - session_sticky_entry_ttl_seconds: int = 3600 + # Five minutes matches the common idle lifetime of upstream prompt caches. + # Providers and per-provider environment settings can override this default. + session_sticky_entry_ttl_seconds: int = 300 session_sticky_max_entries: int = 10000 def get_effective_multiplier(self, priority: int) -> int: @@ -689,13 +697,27 @@ def load_provider_usage_config( if hasattr(plugin_class, "default_session_sticky_entry_ttl_seconds"): sticky_ttl = plugin_class.default_session_sticky_entry_ttl_seconds if sticky_ttl is not None: - config.session_sticky_entry_ttl_seconds = max(1, int(sticky_ttl)) + config.session_sticky_entry_ttl_seconds = max(0, int(sticky_ttl)) if hasattr(plugin_class, "default_session_sticky_max_entries"): sticky_max = plugin_class.default_session_sticky_max_entries if sticky_max is not None: config.session_sticky_max_entries = max(100, int(sticky_max)) + if hasattr(plugin_class, "default_no_reset_exhaustion_policy"): + config.no_reset_exhaustion_policy = str( + plugin_class.default_no_reset_exhaustion_policy + ).lower() + + if hasattr(plugin_class, "default_no_reset_exhaustion_cooldown_seconds"): + try: + config.no_reset_exhaustion_cooldown_seconds = max( + 0, + int(plugin_class.default_no_reset_exhaustion_cooldown_seconds), + ) + except (TypeError, ValueError): + pass + # Fair cycle if hasattr(plugin_class, "default_fair_cycle_config"): fc_config = plugin_class.default_fair_cycle_config @@ -818,7 +840,7 @@ def _parse_concurrency_env(var_name: str) -> Optional[int]: env_sticky_ttl = os.getenv("SESSION_STICKY_ENTRY_TTL_SECONDS") if env_sticky_ttl: try: - config.session_sticky_entry_ttl_seconds = max(1, int(env_sticky_ttl)) + config.session_sticky_entry_ttl_seconds = max(0, int(env_sticky_ttl)) except ValueError: lib_logger.warning( f"Invalid session sticky entry TTL value '{env_sticky_ttl}'. Keeping default." @@ -915,6 +937,33 @@ def _parse_concurrency_env(var_name: str) -> Optional[int]: except ValueError: pass + env_no_reset_policy = os.getenv(f"QUOTA_NO_RESET_EXHAUSTION_POLICY_{provider_upper}") + if env_no_reset_policy is None: + env_no_reset_policy = os.getenv("QUOTA_NO_RESET_EXHAUSTION_POLICY") + if env_no_reset_policy: + policy = env_no_reset_policy.strip().lower() + if policy in {"warn_only", "cooldown", "disable_scope"}: + config.no_reset_exhaustion_policy = policy + else: + lib_logger.warning( + f"Invalid QUOTA_NO_RESET_EXHAUSTION_POLICY_{provider_upper}='{env_no_reset_policy}'. " + f"Keeping default '{config.no_reset_exhaustion_policy}'." + ) + + env_no_reset_cooldown = os.getenv(f"QUOTA_NO_RESET_COOLDOWN_SECONDS_{provider_upper}") + if env_no_reset_cooldown is None: + env_no_reset_cooldown = os.getenv("QUOTA_NO_RESET_COOLDOWN_SECONDS") + if env_no_reset_cooldown: + try: + config.no_reset_exhaustion_cooldown_seconds = max( + 0, int(env_no_reset_cooldown) + ) + except ValueError: + lib_logger.warning( + f"Invalid QUOTA_NO_RESET_COOLDOWN_SECONDS_{provider_upper}='{env_no_reset_cooldown}'. " + "Keeping default." + ) + # Priority multipliers from env # Format: CONCURRENCY_MULTIPLIER_{PROVIDER}_PRIORITY_{N}=value # Format: CONCURRENCY_MULTIPLIER_{PROVIDER}_PRIORITY_{N}_{MODE}=value diff --git a/src/rotator_library/usage/costs.py b/src/rotator_library/usage/costs.py new file mode 100644 index 000000000..6a37aeb82 --- /dev/null +++ b/src/rotator_library/usage/costs.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Advisory cost calculation for normalized usage records.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Optional + +import litellm + +from ..protocols import serialize_value +from .accounting import UsageRecord + + +@dataclass(frozen=True) +class ModelPricing: + """Per-token pricing for one provider/model. + + Prices are advisory and local-only; this module never calls a network pricing + endpoint. Providers can return this object from `get_model_pricing()` later. + """ + + input_cost_per_token: float = 0.0 + cache_read_cost_per_token: float = 0.0 + cache_write_cost_per_token: float = 0.0 + output_cost_per_token: float = 0.0 + reasoning_cost_per_token: float = 0.0 + currency: str = "USD" + source: str = "explicit" + + +@dataclass(frozen=True) +class CostBreakdown: + """Advisory request cost split by normalized usage bucket.""" + + input_cost: float = 0.0 + cache_read_cost: float = 0.0 + cache_write_cost: float = 0.0 + output_cost: float = 0.0 + reasoning_cost: float = 0.0 + provider_reported_cost: float = 0.0 + currency: str = "USD" + pricing_source: str = "unavailable" + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def total_cost(self) -> float: + if self.provider_reported_cost: + return self.provider_reported_cost + return self.input_cost + self.cache_read_cost + self.cache_write_cost + self.output_cost + self.reasoning_cost + + def to_dict(self) -> dict[str, Any]: + return { + "input_cost": self.input_cost, + "cache_read_cost": self.cache_read_cost, + "cache_write_cost": self.cache_write_cost, + "output_cost": self.output_cost, + "reasoning_cost": self.reasoning_cost, + "provider_reported_cost": self.provider_reported_cost, + "total_cost": self.total_cost, + "currency": self.currency, + "pricing_source": self.pricing_source, + "metadata": serialize_value(self.metadata), + } + + +class CostCalculator: + """Calculate advisory costs without replacing usage tracking.""" + + def __init__(self, *, provider_plugin: Any = None, use_litellm_fallback: bool = True, config: Any = None, env: Any = None) -> None: + self.provider_plugin = provider_plugin + self.use_litellm_fallback = use_litellm_fallback + self.config = config + self.env = env + + def calculate(self, usage: UsageRecord, *, model: str, response: Any = None, provider: str | None = None) -> CostBreakdown: + """Return an advisory cost breakdown for a normalized usage record.""" + + if self.provider_plugin and getattr(self.provider_plugin, "skip_cost_calculation", False): + return CostBreakdown(pricing_source="skipped", metadata={"reason": "provider_skip_cost_calculation"}) + if usage.provider_reported_cost is not None: + return CostBreakdown( + provider_reported_cost=usage.provider_reported_cost, + currency=usage.cost_currency, + pricing_source=usage.cost_source or "provider_reported", + metadata={"actual_provider_reported": True}, + ) + pricing = self._provider_pricing(model) + if pricing: + return _calculate_from_pricing(usage, pricing) + pricing = self._configured_pricing(provider or usage.provider or _provider_from_model(model), model) + if pricing: + return _calculate_from_pricing(usage, pricing) + if self.use_litellm_fallback: + lite = self._litellm_cost(usage, model=model, response=response) + if lite: + return lite + return CostBreakdown(pricing_source="unavailable") + + def _provider_pricing(self, model: str) -> Optional[ModelPricing]: + if not self.provider_plugin: + return None + method = getattr(self.provider_plugin, "get_model_pricing", None) + if not method: + return None + pricing = method(model) + if isinstance(pricing, ModelPricing): + return pricing + if isinstance(pricing, dict): + return ModelPricing(**pricing) + return None + + def _configured_pricing(self, provider: str | None, model: str) -> Optional[ModelPricing]: + if not provider: + return None + from ..config.experimental import get_configured_model_pricing + + return get_configured_model_pricing(provider, _model_without_provider(provider, model), config=self.config, env=self.env) + + @staticmethod + def _litellm_cost(usage: UsageRecord, *, model: str, response: Any = None) -> Optional[CostBreakdown]: + if response is not None: + try: + cost = litellm.completion_cost(completion_response=response, model=model) + if cost is not None: + return CostBreakdown(output_cost=float(cost), pricing_source="litellm_completion_cost") + except Exception: + pass + try: + model_info = litellm.get_model_info(model) + except Exception: + return None + input_price = float(model_info.get("input_cost_per_token") or 0.0) + output_price = float(model_info.get("output_cost_per_token") or 0.0) + if input_price == 0.0 and output_price == 0.0: + return None + pricing = ModelPricing( + input_cost_per_token=input_price, + cache_read_cost_per_token=float(model_info.get("cache_read_input_token_cost") or input_price), + cache_write_cost_per_token=float(model_info.get("cache_creation_input_token_cost") or input_price), + output_cost_per_token=output_price, + reasoning_cost_per_token=output_price, + source="litellm_model_info", + ) + return _calculate_from_pricing(usage, pricing) + + +def _calculate_from_pricing(usage: UsageRecord, pricing: ModelPricing) -> CostBreakdown: + return CostBreakdown( + input_cost=usage.input_tokens * pricing.input_cost_per_token, + cache_read_cost=usage.cache_read_tokens * pricing.cache_read_cost_per_token, + cache_write_cost=usage.cache_write_tokens * pricing.cache_write_cost_per_token, + output_cost=usage.completion_tokens * pricing.output_cost_per_token, + reasoning_cost=usage.reasoning_tokens * (pricing.reasoning_cost_per_token or pricing.output_cost_per_token), + currency=pricing.currency, + pricing_source=pricing.source, + ) + + +def _provider_from_model(model: str) -> Optional[str]: + return model.split("/", 1)[0] if "/" in model else None + + +def _model_without_provider(provider: str, model: str) -> str: + prefix = f"{provider}/" + return model[len(prefix) :] if model.startswith(prefix) else model diff --git a/src/rotator_library/usage/integration/hooks.py b/src/rotator_library/usage/integration/hooks.py index d980f6f2c..b9feb1838 100644 --- a/src/rotator_library/usage/integration/hooks.py +++ b/src/rotator_library/usage/integration/hooks.py @@ -212,7 +212,7 @@ async def dispatch_request_complete( behavior for request counting, cooldowns, or exhaustion marking. Args: - provider: Provider name (e.g., "gemini_cli", "openai") + provider: Provider name (e.g., "openai", "anthropic") credential: Credential accessor (file path or API key) model: Model that was called (with provider prefix) success: Whether the request succeeded diff --git a/src/rotator_library/usage/manager.py b/src/rotator_library/usage/manager.py index e1937b1e7..bda028488 100644 --- a/src/rotator_library/usage/manager.py +++ b/src/rotator_library/usage/manager.py @@ -1385,7 +1385,7 @@ def _get_grouped_models(self, group: str) -> List[str]: Returns: List of normalized, deduplicated model names with provider prefix - (e.g., ["gemini_cli/gemini-2.5-pro", "gemini_cli/gemini-3-pro-preview"]) + (e.g., ["provider/model-a", "provider/model-b"]) """ plugin_instance = self._get_provider_plugin_instance() @@ -1499,6 +1499,7 @@ async def update_quota_baseline( quota_group: Optional[str] = None, force: bool = False, apply_exhaustion: bool = False, + exhaustion_reason: Optional[str] = None, ) -> Optional[Dict[str, Any]]: """ Update quota baseline from provider API response. @@ -1521,6 +1522,8 @@ async def update_quota_baseline( Provider controls when this is set based on its semantics (e.g., providers can choose initial-only or always-on refresh when remaining == 0). + exhaustion_reason: Provider-supplied reason for exhausted quota, + such as no_reset_time for entitlement/no-access style buckets. Returns: Cooldown info dict if cooldown was applied, None otherwise @@ -1599,17 +1602,87 @@ async def update_quota_baseline( "cooldown_hours": max(0.0, (quota_reset_ts - time.time()) / 3600), } else: - # ERROR: Provider says exhausted but no reset timestamp! - lib_logger.error( - f"Quota exhausted for {cooldown_target} on " - f"{mask_credential(accessor, style='full')} but no reset_timestamp " - f"provided by API - cannot apply cooldown" + result = await self._handle_no_reset_quota_exhaustion( + state=state, + accessor=accessor, + model=model, + cooldown_target=cooldown_target, + exhaustion_reason=exhaustion_reason, ) + if result: + await self._save_if_needed() + return result await self._save_if_needed() return None + async def _handle_no_reset_quota_exhaustion( + self, + state: CredentialState, + accessor: str, + model: str, + cooldown_target: str, + exhaustion_reason: Optional[str], + ) -> Optional[Dict[str, Any]]: + """Handle authoritative quota exhaustion when the API omits reset time. + + Some providers use a zero/null quota bucket with no reset timestamp to + indicate that the current account cannot access a model group at all. + Provider usage config decides whether that remains warn-only or becomes + a scoped fallback cooldown. + """ + + policy = (self._config.no_reset_exhaustion_policy or "warn_only").lower() + configured_duration = max( + 0, int(self._config.no_reset_exhaustion_cooldown_seconds or 0) + ) + + if policy == "warn_only" or configured_duration <= 0 and policy == "cooldown": + lib_logger.warning( + f"Quota exhausted for {cooldown_target} on " + f"{mask_credential(accessor, style='full')} but no reset_timestamp " + f"provided by API; no fallback cooldown configured" + ) + return None + + if policy not in {"cooldown", "disable_scope"}: + lib_logger.warning( + f"Quota exhausted for {cooldown_target} on " + f"{mask_credential(accessor, style='full')} but no reset_timestamp " + f"provided by API; unknown no-reset policy '{policy}'" + ) + return None + + duration = configured_duration + if policy == "disable_scope" and duration <= 0: + duration = 365 * 24 * 60 * 60 + + await self._tracking.apply_cooldown( + state=state, + reason="quota_no_reset_exhausted", + duration=duration, + model_or_group=cooldown_target, + source="api_quota_no_reset", + ) + + cooldown_until = time.time() + duration + lib_logger.warning( + f"Quota exhausted for {cooldown_target} on " + f"{mask_credential(accessor, style='full')} with no reset_timestamp " + f"provided by API; applying {policy} fallback cooldown for " + f"{round(duration / 3600, 2)}h" + ) + + return { + "cooldown_until": cooldown_until, + "reason": "quota_no_reset_exhausted", + "exhaustion_reason": exhaustion_reason, + "model": model, + "cooldown_hours": max(0.0, duration / 3600), + "policy": policy, + } + # ========================================================================= # WINDOW CLEANUP # ========================================================================= diff --git a/src/rotator_library/usage/quota.py b/src/rotator_library/usage/quota.py new file mode 100644 index 000000000..7d490c2a0 --- /dev/null +++ b/src/rotator_library/usage/quota.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Read-only quota snapshot helpers built from existing usage state.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping, Optional + +from ..error_handler import mask_credential +from ..protocols import serialize_value +from .types import CredentialState, WindowStats + + +@dataclass(frozen=True) +class QuotaSnapshot: + """Client-safe view of one usage window. + + Snapshots are reporting-only. They intentionally do not participate in limit + checks, selection, or persistence so existing quota behavior remains owned by + `UsageManager`, `TrackingEngine`, and `LimitEngine`. + """ + + provider: str + model: Optional[str] + quota_group: Optional[str] + credential_id: Optional[str] + window_name: str + limit: Optional[int] + used: int + remaining: Optional[int] + reset_at: Optional[float] + source: str + metadata: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "provider": self.provider, + "model": self.model, + "quota_group": self.quota_group, + "credential_id": self.credential_id, + "window_name": self.window_name, + "limit": self.limit, + "used": self.used, + "remaining": self.remaining, + "reset_at": self.reset_at, + "source": self.source, + "metadata": serialize_value(self.metadata), + } + + +def build_quota_snapshots( + *, + provider: str, + states: Mapping[str, CredentialState], + model: Optional[str] = None, + quota_group: Optional[str] = None, + include_credentials: bool = True, +) -> list[QuotaSnapshot]: + """Build read-only request/token quota snapshots from credential states. + + The current usage state stores request/token windows, not a reliable + provider-cost ledger. Snapshots therefore avoid inventing cost totals; cost + reporting can be added later only if the underlying state owns that data. + """ + + snapshots: list[QuotaSnapshot] = [] + for stable_id, state in states.items(): + credential_id = mask_credential(stable_id) if include_credentials else None + if model: + model_stats = state.get_model_stats(model, create=False) + if model_stats: + snapshots.extend( + _snapshots_for_windows( + provider=provider, + model=model, + quota_group=None, + credential_id=credential_id, + windows=model_stats.windows, + source="model", + ) + ) + if quota_group: + group_stats = state.get_group_stats(quota_group, create=False) + if group_stats: + snapshots.extend( + _snapshots_for_windows( + provider=provider, + model=model, + quota_group=quota_group, + credential_id=credential_id, + windows=group_stats.windows, + source="group", + ) + ) + return snapshots + + +def _snapshots_for_windows( + *, + provider: str, + model: Optional[str], + quota_group: Optional[str], + credential_id: Optional[str], + windows: Mapping[str, WindowStats], + source: str, +) -> list[QuotaSnapshot]: + return [ + QuotaSnapshot( + provider=provider, + model=model, + quota_group=quota_group, + credential_id=credential_id, + window_name=window.name, + limit=window.limit, + used=window.request_count, + remaining=window.remaining, + reset_at=window.reset_at, + source=source, + metadata={"scope": "request_token_window"}, + ) + for window in windows.values() + ] diff --git a/src/rotator_library/usage/selection/strategies/sequential.py b/src/rotator_library/usage/selection/strategies/sequential.py index 101af1819..e59b645a8 100644 --- a/src/rotator_library/usage/selection/strategies/sequential.py +++ b/src/rotator_library/usage/selection/strategies/sequential.py @@ -41,7 +41,7 @@ class SequentialStrategy: def __init__( self, fallback_multiplier: int = 1, - sticky_entry_ttl_seconds: int = 3600, + sticky_entry_ttl_seconds: int = 300, max_sticky_entries: int = 10000, ): """ @@ -52,7 +52,7 @@ def __init__( when not explicitly configured """ self.fallback_multiplier = fallback_multiplier - self.sticky_entry_ttl_seconds = max(1, sticky_entry_ttl_seconds) + self.sticky_entry_ttl_seconds = max(0, sticky_entry_ttl_seconds) self.max_sticky_entries = max(100, max_sticky_entries) # Track current "sticky" credential per model session or model-group fallback. self._current: Dict[tuple, _StickyEntry] = {} @@ -126,7 +126,7 @@ def _select_locked( ) # Make it sticky - if selected: + if selected and self.sticky_entry_ttl_seconds > 0: self._current[key] = _StickyEntry(selected, now) self._trim_sticky() masked = ( @@ -204,10 +204,13 @@ def get_current( """ with self._lock: key = (provider, model_or_group, session_id or "__default__") + if self.sticky_entry_ttl_seconds <= 0: + self._current.pop(key, None) + return None entry = self._current.get(key) if not entry: return None - if time.time() - entry.last_seen > self.sticky_entry_ttl_seconds: + if time.time() - entry.last_seen >= self.sticky_entry_ttl_seconds: del self._current[key] return None return entry.credential @@ -215,10 +218,13 @@ def get_current( def _prune_sticky(self, now: float) -> None: # Locking contract: callers must hold self._lock. This helper is kept # private so select/get_current can batch prune + selection atomically. + if self.sticky_entry_ttl_seconds <= 0: + self._current.clear() + return expired = [ key for key, entry in self._current.items() - if now - entry.last_seen > self.sticky_entry_ttl_seconds + if now - entry.last_seen >= self.sticky_entry_ttl_seconds ] for key in expired: del self._current[key] diff --git a/src/rotator_library/utils/paths.py b/src/rotator_library/utils/paths.py index 3f77ad82f..e9b125856 100644 --- a/src/rotator_library/utils/paths.py +++ b/src/rotator_library/utils/paths.py @@ -58,7 +58,7 @@ def get_cache_dir( Args: root: Optional root directory. If None, uses get_default_root(). - subdir: Optional subdirectory name (e.g., "gemini_cli") + subdir: Optional subdirectory name (e.g., "provider_cache") Returns: Path to the cache directory (or subdirectory) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..df150dc59 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +SRC = ROOT / "src" + +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) diff --git a/tests/test_adapter_registry.py b/tests/test_adapter_registry.py new file mode 100644 index 000000000..cef410c85 --- /dev/null +++ b/tests/test_adapter_registry.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.adapters import ( + AdapterContext, + PayloadAdapter, + get_adapter, + get_adapter_class, + list_adapters, + register_adapter, + resolve_adapter_name, + run_adapter_chain, +) +from rotator_library.transaction_logger import TransactionLogger + + +def _trace_entries(log_dir): + return [json.loads(line) for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +def test_adapter_registry_auto_discovers_builtins_and_aliases() -> None: + adapters = list_adapters() + + assert "noop" in adapters + assert "model_override" in adapters + assert "field_rename" in adapters + assert resolve_adapter_name("passthrough") == "noop" + assert resolve_adapter_name("field_copy") == "field_rename" + assert get_adapter_class("override_model").name == "model_override" + assert get_adapter("none") is get_adapter("noop") + + +def test_adapter_registry_rejects_duplicate_names_and_alias_collisions() -> None: + class DuplicateNoop(PayloadAdapter): + name = "noop" + + class AliasCollision(PayloadAdapter): + name = "custom_alias_collision" + aliases = ("noop",) + + with pytest.raises(ValueError): + register_adapter(DuplicateNoop) + with pytest.raises(ValueError): + register_adapter(AliasCollision) + + +@pytest.mark.asyncio +async def test_noop_adapter_preserves_payload_without_mutating_original() -> None: + payload = {"model": "original", "messages": []} + + result = await run_adapter_chain([get_adapter("noop")], payload, AdapterContext(), stage="request") + + assert result == payload + assert result is not payload + + +@pytest.mark.asyncio +async def test_model_override_adapter_changes_model_from_config() -> None: + payload = {"model": "public-name", "messages": []} + context = AdapterContext(adapter_config={"model_override": {"model": "native-name"}}) + + result = await run_adapter_chain([get_adapter("model_override")], payload, context, stage="request") + + assert result["model"] == "native-name" + assert payload["model"] == "public-name" + + +@pytest.mark.asyncio +async def test_suppress_developer_role_converts_or_drops_messages() -> None: + payload = {"messages": [{"role": "developer", "content": "rules"}, {"role": "user", "content": "hi"}]} + + system_result = await run_adapter_chain([get_adapter("suppress_developer_role")], payload, AdapterContext(), stage="request") + drop_result = await run_adapter_chain( + [get_adapter("suppress_developer_role")], + payload, + AdapterContext(adapter_config={"suppress_developer_role": {"mode": "drop"}}), + stage="request", + ) + + assert system_result["messages"][0]["role"] == "system" + assert [message["role"] for message in drop_result["messages"]] == ["user"] + + +@pytest.mark.asyncio +async def test_reasoning_content_adapter_copies_common_reasoning_field() -> None: + payload = {"choices": [{"message": {"role": "assistant", "reasoning": "hidden"}}]} + + result = await run_adapter_chain([get_adapter("reasoning_content")], payload, AdapterContext(), stage="response") + + assert result["choices"][0]["message"]["reasoning_content"] == "hidden" + + +@pytest.mark.asyncio +async def test_field_rename_adapter_copies_and_moves_configured_fields() -> None: + payload = {"old": {"field": "value"}, "messages": [{}]} + context = AdapterContext( + adapter_config={ + "field_rename": { + "rules": [ + {"source_path": "old.field", "target_path": "messages[-1].new_field", "move": True} + ] + } + } + ) + + result = await run_adapter_chain([get_adapter("field_rename")], payload, context, stage="request") + + assert result["messages"][-1]["new_field"] == "value" + assert "field" not in result["old"] + assert payload["old"]["field"] == "value" + + +@pytest.mark.asyncio +async def test_antigravity_envelope_wraps_user_request_key() -> None: + adapter = get_adapter("antigravity_envelope") + + result = await adapter.transform_request( + {"model": "gemini-3-flash", "request": {"user_supplied": True}, "contents": [{"parts": [{"text": "hi"}]}]}, + AdapterContext(adapter_config={"antigravity_envelope": {"request_type": "CHAT_COMPLETION", "user_agent": "test-agent"}}), + ) + + assert result["requestType"] == "CHAT_COMPLETION" + assert result["requestId"] + assert result["request"]["request"] == {"user_supplied": True} + assert result["request"]["contents"][0]["parts"][0]["text"] == "hi" + + +@pytest.mark.asyncio +async def test_antigravity_envelope_is_idempotent_for_controlled_envelope() -> None: + adapter = get_adapter("antigravity_envelope") + payload = {"model": "gemini-3-flash", "request": {"contents": []}, "requestType": "CHAT_COMPLETION", "requestId": "id"} + + result = await adapter.transform_request(payload, AdapterContext()) + + assert result == payload + + +@pytest.mark.asyncio +async def test_adapter_chain_order_is_preserved() -> None: + payload = {"model": "public", "messages": [{"role": "developer", "content": "rules"}]} + context = AdapterContext(adapter_config={"model_override": {"model": "native"}}) + + result = await run_adapter_chain( + [get_adapter("model_override"), get_adapter("suppress_developer_role")], + payload, + context, + stage="request", + ) + + assert result["model"] == "native" + assert result["messages"][0]["role"] == "system" + + +@pytest.mark.asyncio +async def test_adapter_chain_traces_final_summary(tmp_path) -> None: + logger = TransactionLogger("native", "native/test", parent_dir=tmp_path) + payload = {"model": "public", "messages": []} + context = AdapterContext( + adapter_config={"model_override": {"model": "native"}}, + transaction_logger=logger, + protocol="openai_chat", + ) + + result = await run_adapter_chain([get_adapter("model_override")], payload, context, stage="request") + + assert result["model"] == "native" + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert pass_names == ["before_adapter_chain", "after_adapter", "after_adapter_chain"] + assert entries[-1]["metadata"]["adapter_count"] == 1 + assert entries[-1]["metadata"]["changed"] is True diff --git a/tests/test_anthropic_transform_tracing.py b/tests/test_anthropic_transform_tracing.py new file mode 100644 index 000000000..45a8ba8ac --- /dev/null +++ b/tests/test_anthropic_transform_tracing.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import json + +import pytest + +import rotator_library.client.anthropic as anthropic_client_module +from rotator_library.anthropic_compat import AnthropicMessagesRequest +from rotator_library.anthropic_compat.streaming import anthropic_streaming_wrapper +from rotator_library.client.anthropic import AnthropicHandler +from rotator_library.transaction_logger import TransactionLogger + + +def _trace_entries(log_dir): + return [json.loads(line) for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +class FakeRotatingClient: + enable_request_logging = True + + async def acompletion(self, **kwargs): + return { + "id": "chat_1", + "model": kwargs["model"], + "choices": [{"message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + +class FakeNativeProtocolClient: + enable_request_logging = False + + def __init__(self) -> None: + self.call = None + + async def agenerate(self, payload, *, input_protocol, output_protocol, **kwargs): + self.call = { + "payload": payload, + "input_protocol": input_protocol, + "output_protocol": output_protocol, + } + if payload.get("stream"): + async def stream(): + yield 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_stream","type":"message","role":"assistant","content":[],"model":"claude-test","usage":{"input_tokens":1,"output_tokens":0}}}\n\n' + yield 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n' + yield 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"native stream"}}\n\n' + yield 'event: message_stop\ndata: {"type":"message_stop"}\n\n' + return stream() + return { + "id": "msg_provider", + "type": "message", + "role": "assistant", + "model": payload["model"], + "content": [{"type": "text", "text": "native answer"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + + +@pytest.mark.asyncio +async def test_anthropic_handler_traces_conversion_boundaries(tmp_path, monkeypatch) -> None: + created: list[TransactionLogger] = [] + + def logger_factory(provider, model, enabled=True, api_format="ant", parent_dir=None): + logger = TransactionLogger(provider, model, enabled=enabled, api_format=api_format, parent_dir=tmp_path) + created.append(logger) + return logger + + monkeypatch.setattr(anthropic_client_module, "TransactionLogger", logger_factory) + request = AnthropicMessagesRequest(model="openai/gpt-test", max_tokens=16, messages=[{"role": "user", "content": "hi"}]) + + response = await AnthropicHandler(FakeRotatingClient()).messages(request) + + assert response["content"][0]["text"] == "ok" + pass_names = [entry["pass_name"] for entry in _trace_entries(created[0].log_dir)] + assert "anthropic_raw_request" in pass_names + assert "anthropic_to_openai_request" in pass_names + assert "anthropic_openai_response" in pass_names + assert "openai_to_anthropic_response" in pass_names + assert "anthropic_final_response" in pass_names + + +@pytest.mark.asyncio +async def test_anthropic_handler_uses_protocol_native_runtime_when_available() -> None: + client = FakeNativeProtocolClient() + request = AnthropicMessagesRequest( + model="claude_code/claude-test", + max_tokens=16, + system="rule", + messages=[{"role": "user", "content": "hi"}], + ) + + response = await AnthropicHandler(client).messages(request) + + assert client.call["input_protocol"] == "anthropic_messages" + assert client.call["output_protocol"] == "anthropic_messages" + assert client.call["payload"]["system"] == "rule" + assert response["content"][0]["text"] == "native answer" + assert response["id"].startswith("msg_") + + +@pytest.mark.asyncio +async def test_anthropic_handler_uses_protocol_native_runtime_for_streams() -> None: + client = FakeNativeProtocolClient() + request = AnthropicMessagesRequest( + model="claude_code/claude-test", + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + stream=True, + ) + + stream = await AnthropicHandler(client).messages(request) + output = "".join([chunk async for chunk in stream]) + + assert client.call["input_protocol"] == "anthropic_messages" + assert client.call["output_protocol"] == "anthropic_messages" + assert "native stream" in output + + +class ClosingOpenAIStream: + def __init__(self) -> None: + self.closed = False + self._chunks = iter(['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n']) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._chunks) + except StopIteration: + raise StopAsyncIteration + + async def aclose(self) -> None: + self.closed = True + + +class IteratorOnlyCloseStream: + def __init__(self) -> None: + self.iterator = IteratorOnlyCloseStreamIterator() + + def __aiter__(self): + return self.iterator + + +class IteratorOnlyCloseStreamIterator: + def __init__(self) -> None: + self.closed = False + self._chunks = iter(['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n']) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._chunks) + except StopIteration: + raise StopAsyncIteration + + async def aclose(self) -> None: + self.closed = True + + +@pytest.mark.asyncio +async def test_anthropic_stream_traces_and_closes_on_disconnect(tmp_path) -> None: + logger = TransactionLogger("anthropic", "claude-test", parent_dir=tmp_path) + stream = ClosingOpenAIStream() + + async def disconnected() -> bool: + return True + + chunks = [chunk async for chunk in anthropic_streaming_wrapper(stream, "claude-test", is_disconnected=disconnected, transaction_logger=logger)] + + assert chunks == [] + assert stream.closed is True + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "anthropic_stream_source_chunk" in pass_names + assert "anthropic_stream_disconnected" in pass_names + assert "anthropic_stream_upstream_closed" in pass_names + + +@pytest.mark.asyncio +async def test_anthropic_stream_closes_iterator_only_upstream() -> None: + stream = IteratorOnlyCloseStream() + + async def disconnected() -> bool: + return True + + chunks = [chunk async for chunk in anthropic_streaming_wrapper(stream, "claude-test", is_disconnected=disconnected)] + + assert chunks == [] + assert stream.iterator.closed is True + + +@pytest.mark.asyncio +async def test_anthropic_stream_traces_emitted_frames(tmp_path) -> None: + logger = TransactionLogger("anthropic", "claude-test", parent_dir=tmp_path) + + async def stream(): + yield 'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n' + yield "data: [DONE]\n\n" + + chunks = [chunk async for chunk in anthropic_streaming_wrapper(stream(), "claude-test", transaction_logger=logger)] + + assert any("event: message_start" in chunk for chunk in chunks) + assert any("event: message_stop" in chunk for chunk in chunks) + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "anthropic_stream_message_start" in pass_names + assert "anthropic_stream_content_delta" in pass_names + assert "anthropic_stream_message_delta" in pass_names + assert "anthropic_stream_message_stop" in pass_names diff --git a/tests/test_antigravity_provider_restore.py b/tests/test_antigravity_provider_restore.py new file mode 100644 index 000000000..f652b2aa9 --- /dev/null +++ b/tests/test_antigravity_provider_restore.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import pytest + +from rotator_library.providers import PROVIDER_PLUGINS +from rotator_library.providers.antigravity_provider import AntigravityProvider + + +class FakeResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + +class FakeClient: + def __init__(self, payload): + self.payload = payload + self.calls = [] + + async def post(self, url, *, headers, json, timeout): + self.calls.append({"url": url, "headers": headers, "json": json, "timeout": timeout}) + return FakeResponse(self.payload) + + +def test_antigravity_provider_is_discovered() -> None: + assert "antigravity" in PROVIDER_PLUGINS + + +def test_antigravity_provider_restores_safe_declarations() -> None: + provider = AntigravityProvider() + + assert provider.get_protocol_name("gemini-3-flash") == "gemini" + assert provider.get_adapter_names("gemini-3-flash") == ("antigravity_envelope",) + assert provider.get_adapter_config("gemini-3-flash")["antigravity_envelope"]["request_type"] == "CHAT_COMPLETION" + assert provider.get_model_tier_requirement("antigravity/gemini-3-flash") is None + rules = provider.get_field_cache_rules("gemini-3-flash") + assert rules[0].name == "antigravity_thought_signature" + assert rules[0].scope == ("provider", "model", "credential", "session") + + +def test_antigravity_quota_groups_are_model_family_scoped() -> None: + groups = AntigravityProvider.model_quota_groups + + assert groups["gemini_3_pro"] == ["gemini-3-pro-preview", "gemini-3-pro-low", "gemini-3-pro-high"] + assert groups["gemini_3_flash"] == ["gemini-3-flash"] + assert groups["gemini_2_5_flash"] == ["gemini-2.5-flash"] + assert groups["gemini_2_5_flash_lite"] == ["gemini-2.5-flash-lite"] + assert groups["claude_sonnet_4_5"] == ["claude-sonnet-4.5"] + + +def test_antigravity_provider_builds_static_headers_without_device_profile(monkeypatch) -> None: + monkeypatch.setenv("ANTIGRAVITY_API_BASE", "https://antigravity.test/v1internal") + provider = AntigravityProvider() + + headers = provider.get_native_headers("token") + + assert provider.get_native_endpoint(operation="generate") == "https://antigravity.test/v1internal:generateContent" + assert provider.get_native_endpoint(operation="stream_generate") == "https://antigravity.test/v1internal:streamGenerateContent?alt=sse" + assert provider.get_native_endpoint(operation="models") == "https://antigravity.test/v1internal:fetchAvailableModels" + assert headers["Authorization"] == "Bearer token" + assert headers["X-Goog-Api-Client"] == "google-cloud-sdk vscode_cloudshelleditor/0.1" + assert "Accept" not in headers + assert provider.get_native_headers("token", operation="stream_generate")["Accept"] == "text/event-stream" + assert "X-Client-Device-Id" not in headers + + +def test_antigravity_model_aliases_and_tracking_normalization() -> None: + provider = AntigravityProvider() + + assert provider._alias_to_internal("claude-sonnet-4.5") == "claude-sonnet-4-5" + assert provider.normalize_native_model("antigravity/claude-sonnet-4.5") == "claude-sonnet-4-5" + assert provider.normalize_native_model("antigravity/gemini-3-pro-preview") == "gemini-3-pro-preview" + assert provider.normalize_model_for_tracking("antigravity/claude-sonnet-4-5") == "antigravity/claude-sonnet-4.5" + + +def test_antigravity_native_operation_model_and_stream_support() -> None: + provider = AntigravityProvider() + + assert provider.get_native_operation("gemini-3-flash", {}, stream=False) == "generate" + assert provider.get_native_operation("gemini-3-flash", {}, stream=True) == "stream_generate" + assert provider.supports_native_streaming("gemini-3-flash", operation="stream_generate") is True + assert provider.supports_native_streaming("gemini-3-flash", operation="generate") is False + prepared = provider.prepare_native_request({"model": "antigravity/gemini-3-pro-low"}, model="gemini-3-pro-preview", operation="generate") + assert prepared["model"] == "gemini-3-pro-low" + assert prepared["generationConfig"]["thinkingConfig"]["thinkingLevel"] == "low" + + +def test_antigravity_prepare_native_request_only_adjusts_gemini_native_payload() -> None: + provider = AntigravityProvider() + + prepared = provider.prepare_native_request( + {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + model="gemini-3-flash", + operation="generate", + ) + + assert prepared["model"] == "gemini-3-flash" + assert prepared["contents"] == [{"role": "user", "parts": [{"text": "hello"}]}] + assert "messages" not in prepared + + +@pytest.mark.asyncio +async def test_antigravity_get_models_filters_and_aliases_mocked_response(monkeypatch) -> None: + monkeypatch.setenv("ANTIGRAVITY_API_BASE", "https://antigravity.test/v1internal") + client = FakeClient({"models": {"gemini-3-pro-low": {}, "chat_20706": {}, "claude-opus-4-6": {}}}) + provider = AntigravityProvider() + + models = await provider.get_models("token", client) + + assert models == ["antigravity/gemini-3-pro-preview", "antigravity/claude-opus-4.6"] + assert client.calls[0]["url"] == "https://antigravity.test/v1internal:fetchAvailableModels" + assert client.calls[0]["headers"]["Authorization"] == "Bearer token" + + +@pytest.mark.asyncio +async def test_antigravity_get_models_falls_back_on_errors() -> None: + class BrokenClient: + async def post(self, *args, **kwargs): + raise RuntimeError("offline") + + models = await AntigravityProvider().get_models("token", BrokenClient()) + + assert "antigravity/gemini-3-flash" in models + assert "antigravity/claude-opus-4.6" in models diff --git a/tests/test_classifier_scoped_routing.py b/tests/test_classifier_scoped_routing.py index f71e8c6c7..f51b487ac 100644 --- a/tests/test_classifier_scoped_routing.py +++ b/tests/test_classifier_scoped_routing.py @@ -608,6 +608,37 @@ async def run_test(): run_async(run_test()) +def test_ad_hoc_private_bundles_never_share_usage_manager_state(tmp_path): + client = _make_client(tmp_path, api_keys={"openai": ["global-openai-key"]}) + + async def run_test(): + try: + first = await client._resolve_scope_for_provider( + "openai", + None, + {"openai": ["private-secret-a"]}, + None, + True, + ) + second = await client._resolve_scope_for_provider( + "openai", + None, + {"openai": ["private-secret-b"]}, + None, + True, + ) + + assert first["usage_manager_key"] != second["usage_manager_key"] + assert first["classifier"].startswith("bundle:") + assert second["classifier"].startswith("bundle:") + assert first["usage_manager_key"] in client.usage_managers + assert second["usage_manager_key"] in client.usage_managers + finally: + await _close(client) + + run_async(run_test()) + + def test_model_discovery_cache_key_changes_with_scoped_credentials(tmp_path): seen = [] @@ -675,7 +706,8 @@ async def run_test(): run_async(run_test()) -def test_provider_config_override_routes_without_global_mutation(tmp_path): +def test_provider_config_override_routes_without_global_mutation(tmp_path, monkeypatch): + monkeypatch.delenv("LOGFARE_API_BASE", raising=False) client = _make_client(tmp_path, api_keys={"openai": ["global-openai-key"]}) try: diff --git a/tests/test_claude_code_provider.py b/tests/test_claude_code_provider.py new file mode 100644 index 000000000..8e1ffcada --- /dev/null +++ b/tests/test_claude_code_provider.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import pytest + +from rotator_library.adapters import AdapterContext, get_adapter, run_adapter_chain +from rotator_library.providers import PROVIDER_PLUGINS +from rotator_library.providers.claude_code_provider import ClaudeCodeProvider + + +class FakeResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + +class FakeClient: + def __init__(self, payload): + self.payload = payload + self.calls = [] + + async def get(self, url, *, headers, timeout): + self.calls.append({"url": url, "headers": headers, "timeout": timeout}) + return FakeResponse(self.payload) + + +def test_claude_code_provider_is_discovered() -> None: + assert "claude_code" in PROVIDER_PLUGINS + + +def test_claude_code_provider_declares_native_protocol_adapters_and_cache_rules() -> None: + provider = ClaudeCodeProvider() + + assert provider.get_protocol_name("claude-sonnet-4-5") == "anthropic_messages" + assert provider.get_adapter_names("claude-sonnet-4-5") == ("suppress_developer_role",) + assert provider.get_adapter_config("claude-sonnet-4-5") == {"suppress_developer_role": {"mode": "user"}} + rules = provider.get_field_cache_rules("claude-sonnet-4-5") + assert rules[0].name == "claude_code_thinking_signature" + assert rules[0].scope == ("provider", "model", "credential", "session") + + +@pytest.mark.asyncio +async def test_claude_code_adapter_config_converts_developer_role_to_user() -> None: + provider = ClaudeCodeProvider() + + result = await run_adapter_chain( + [get_adapter("suppress_developer_role")], + {"messages": [{"role": "developer", "content": "rules"}]}, + AdapterContext(adapter_config=provider.get_adapter_config("claude-sonnet-4-5")), + stage="request", + ) + + assert result["messages"][0]["role"] == "user" + + +def test_claude_code_provider_builds_native_headers_and_endpoint(monkeypatch) -> None: + monkeypatch.setenv("CLAUDE_CODE_API_BASE", "https://claude-code.test") + monkeypatch.setenv("CLAUDE_CODE_ANTHROPIC_VERSION", "2099-01-01") + provider = ClaudeCodeProvider() + + assert provider.get_native_endpoint(operation="messages") == "https://claude-code.test/v1/messages" + assert provider.get_native_headers("token") == { + "Authorization": "Bearer token", + "anthropic-version": "2099-01-01", + "content-type": "application/json", + } + + +def test_claude_code_native_operation_model_and_stream_support() -> None: + provider = ClaudeCodeProvider() + + assert provider.get_native_operation("claude-sonnet-4-5", {}, stream=False) == "messages" + assert provider.normalize_native_model("claude_code/claude-sonnet-4-5") == "claude-sonnet-4-5" + assert provider.supports_native_streaming("claude-sonnet-4-5", operation="messages") is True + assert provider.supports_native_streaming("claude-sonnet-4-5", operation="chat") is False + + +@pytest.mark.asyncio +async def test_claude_code_provider_get_models_uses_mocked_models_endpoint(monkeypatch) -> None: + monkeypatch.setenv("CLAUDE_CODE_API_BASE", "https://claude-code.test") + client = FakeClient({"data": [{"id": "claude-sonnet-test"}]}) + provider = ClaudeCodeProvider() + + models = await provider.get_models("token", client) + + assert models == ["claude_code/claude-sonnet-test"] + assert client.calls[0]["url"] == "https://claude-code.test/v1/models" + assert client.calls[0]["headers"]["Authorization"] == "Bearer token" + + +@pytest.mark.asyncio +async def test_claude_code_provider_get_models_falls_back_on_errors() -> None: + class BrokenClient: + async def get(self, *args, **kwargs): + raise RuntimeError("offline") + + assert await ClaudeCodeProvider().get_models("token", BrokenClient()) == [ + "claude_code/claude-sonnet-4-5", + "claude_code/claude-opus-4-5", + ] diff --git a/tests/test_codex_provider.py b/tests/test_codex_provider.py new file mode 100644 index 000000000..1fcee16e2 --- /dev/null +++ b/tests/test_codex_provider.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import pytest + +from rotator_library.providers import PROVIDER_PLUGINS +from rotator_library.providers.codex_provider import CodexProvider + + +class FakeResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + +class FakeClient: + def __init__(self, payload): + self.payload = payload + self.calls = [] + + async def get(self, url, *, headers, timeout): + self.calls.append({"url": url, "headers": headers, "timeout": timeout}) + return FakeResponse(self.payload) + + +def test_codex_provider_is_discovered() -> None: + assert "codex" in PROVIDER_PLUGINS + + +def test_codex_provider_declares_responses_protocol_and_cache_rule() -> None: + provider = CodexProvider() + + assert provider.get_protocol_name("codex-mini-latest") == "responses" + assert provider.get_adapter_names("codex-mini-latest") == () + rules = provider.get_field_cache_rules("codex-mini-latest") + assert rules[0].name == "codex_previous_response_id" + assert rules[0].inject.path == "previous_response_id" + + +def test_codex_provider_builds_native_headers_and_endpoint(monkeypatch) -> None: + monkeypatch.setenv("CODEX_API_BASE", "https://codex.test") + provider = CodexProvider() + + assert provider.get_native_endpoint(operation="responses") == "https://codex.test/v1/responses" + assert provider.get_native_endpoint(operation="models") == "https://codex.test/v1/models" + assert provider.get_native_headers("token") == {"Authorization": "Bearer token", "content-type": "application/json"} + + +def test_codex_native_operation_model_and_stream_support() -> None: + provider = CodexProvider() + + assert provider.get_native_operation("gpt-5.1-codex", {}, stream=False) == "responses" + assert provider.normalize_native_model("codex/gpt-5.1-codex") == "gpt-5.1-codex" + assert provider.supports_native_streaming("gpt-5.1-codex", operation="responses") is True + assert provider.supports_native_streaming("gpt-5.1-codex", operation="chat") is False + + +def test_codex_prepare_native_request_only_adjusts_responses_native_payload() -> None: + provider = CodexProvider() + + prepared = provider.prepare_native_request( + {"model": "codex/gpt-5.1-codex", "input": [{"role": "user", "content": "hello"}]}, + model="gpt-5.1-codex", + operation="responses", + ) + + assert prepared["model"] == "gpt-5.1-codex" + assert prepared["input"] == [{"role": "user", "content": "hello"}] + assert "messages" not in prepared + + +@pytest.mark.asyncio +async def test_codex_provider_get_models_filters_codex_models(monkeypatch) -> None: + monkeypatch.setenv("CODEX_API_BASE", "https://codex.test") + client = FakeClient({"data": [{"id": "gpt-5.1-codex"}, {"id": "gpt-5.1"}]}) + provider = CodexProvider() + + models = await provider.get_models("token", client) + + assert models == ["codex/gpt-5.1-codex"] + assert client.calls[0]["url"] == "https://codex.test/v1/models" + + +@pytest.mark.asyncio +async def test_codex_provider_get_models_falls_back_on_errors() -> None: + class BrokenClient: + async def get(self, *args, **kwargs): + raise RuntimeError("offline") + + assert await CodexProvider().get_models("token", BrokenClient()) == ["codex/codex-mini-latest", "codex/gpt-5.1-codex"] diff --git a/tests/test_config_pricing.py b/tests/test_config_pricing.py new file mode 100644 index 000000000..67dd39437 --- /dev/null +++ b/tests/test_config_pricing.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from rotator_library.config.experimental import env_price_key, load_config_from_mapping +from rotator_library.usage.accounting import UsageRecord +from rotator_library.usage.costs import CostCalculator, ModelPricing + + +class ExplicitPricingProvider: + def get_model_pricing(self, model: str) -> ModelPricing: + return ModelPricing(input_cost_per_token=9.0, source="provider") + + +class SkipCostProvider: + skip_cost_calculation = True + + +def _usage() -> UsageRecord: + return UsageRecord(input_tokens=10, completion_tokens=5, reasoning_tokens=2, cache_read_tokens=3, cache_write_tokens=4, provider="openai", model="gpt-test") + + +def test_json_pricing_calculates_all_buckets() -> None: + config = load_config_from_mapping( + {"pricing": {"openai": {"gpt-test": {"input": 1.0, "output": 2.0, "reasoning": 3.0, "cache_read": 0.5, "cache_write": 0.75}}}} + ) + + cost = CostCalculator(config=config, use_litellm_fallback=False).calculate(_usage(), model="openai/gpt-test") + + assert cost.pricing_source == "json_config" + assert cost.input_cost == 10.0 + assert cost.output_cost == 10.0 + assert cost.reasoning_cost == 6.0 + assert cost.cache_read_cost == 1.5 + assert cost.cache_write_cost == 3.0 + + +def test_env_pricing_overrides_json_pricing() -> None: + config = load_config_from_mapping({"pricing": {"openai": {"gpt-test": {"input": 1.0}}}}) + env = {env_price_key("openai", "gpt-test", "input"): "4.0"} + + cost = CostCalculator(config=config, env=env, use_litellm_fallback=False).calculate(_usage(), model="openai/gpt-test") + + assert cost.pricing_source == "env" + assert cost.input_cost == 40.0 + + +def test_invalid_env_pricing_falls_back_to_json_pricing() -> None: + config = load_config_from_mapping({"pricing": {"openai": {"gpt-test": {"input": 1.0}}}}) + env = {env_price_key("openai", "gpt-test", "input"): "not-a-number"} + + cost = CostCalculator(config=config, env=env, use_litellm_fallback=False).calculate(_usage(), model="openai/gpt-test") + + assert cost.pricing_source == "json_config" + assert cost.input_cost == 10.0 + + +def test_mixed_invalid_env_pricing_uses_valid_components() -> None: + env = { + env_price_key("openai", "gpt-test", "input"): "not-a-number", + env_price_key("openai", "gpt-test", "output"): "2.0", + } + + cost = CostCalculator(env=env, use_litellm_fallback=False).calculate(_usage(), model="openai/gpt-test") + + assert cost.pricing_source == "env" + assert cost.input_cost == 0.0 + assert cost.output_cost == 10.0 + + +def test_provider_pricing_beats_env_and_json_pricing() -> None: + config = load_config_from_mapping({"pricing": {"openai": {"gpt-test": {"input": 1.0}}}}) + env = {env_price_key("openai", "gpt-test", "input"): "4.0"} + + cost = CostCalculator(provider_plugin=ExplicitPricingProvider(), config=config, env=env, use_litellm_fallback=False).calculate(_usage(), model="openai/gpt-test") + + assert cost.pricing_source == "provider" + assert cost.input_cost == 90.0 + + +def test_skip_cost_provider_beats_all_config_pricing() -> None: + config = load_config_from_mapping({"pricing": {"openai": {"gpt-test": {"input": 1.0}}}}) + + cost = CostCalculator(provider_plugin=SkipCostProvider(), config=config, use_litellm_fallback=False).calculate(_usage(), model="openai/gpt-test") + + assert cost.pricing_source == "skipped" + assert cost.total_cost == 0.0 + + +def test_missing_pricing_remains_unavailable() -> None: + cost = CostCalculator(use_litellm_fallback=False).calculate(_usage(), model="openai/gpt-test") + + assert cost.pricing_source == "unavailable" diff --git a/tests/test_config_routing_json.py b/tests/test_config_routing_json.py new file mode 100644 index 000000000..dd50f6eb6 --- /dev/null +++ b/tests/test_config_routing_json.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import pytest + +from rotator_library.config.experimental import load_config_from_mapping +from rotator_library.routing import RoutingConfigError, load_routing_config_from_env + + +def test_json_routing_config_loads_fallback_group_and_route() -> None: + experimental = load_config_from_mapping( + { + "routing": { + "fallback_groups": { + "code_chain": { + "targets": ["codex/gpt-5.1@native", "openai/gpt-5.1@litellm_fallback"], + "failover_on": ["rate_limit"], + "stop_on": ["authentication"], + } + }, + "model_routes": {"code": "group:code_chain"}, + } + } + ) + + config = load_routing_config_from_env({}, config=experimental) + + assert config.fallback_groups["code_chain"].targets[0].execution == "native" + assert config.fallback_groups["code_chain"].failover_on == frozenset({"rate_limit"}) + assert config.model_routes["code"] == "group:code_chain" + + +def test_env_group_overrides_json_group_targets() -> None: + experimental = load_config_from_mapping( + {"routing": {"fallback_groups": {"chain": {"targets": ["codex/gpt"]}}, "model_routes": {"code": "group:chain"}}} + ) + + config = load_routing_config_from_env( + {"FALLBACK_GROUPS": "chain", "FALLBACK_GROUP_CHAIN": "openai/gpt-5.1", "MODEL_ROUTE_CODE": "group:chain"}, + config=experimental, + ) + + assert config.fallback_groups["chain"].targets[0].provider == "openai" + + +def test_env_route_can_reference_json_group() -> None: + experimental = load_config_from_mapping({"routing": {"fallback_groups": {"chain": {"targets": ["codex/gpt"]}}}}) + + config = load_routing_config_from_env({"MODEL_ROUTE_CODE": "group:chain"}, config=experimental) + + assert config.model_routes["code"] == "group:chain" + + +def test_json_route_rejects_unknown_group() -> None: + experimental = load_config_from_mapping({"routing": {"model_routes": {"code": "group:missing"}}}) + + with pytest.raises(RoutingConfigError): + load_routing_config_from_env({}, config=experimental) + + +def test_json_route_rejects_invalid_direct_target() -> None: + experimental = load_config_from_mapping({"routing": {"model_routes": {"code": "not-a-target"}}}) + + with pytest.raises(RoutingConfigError): + load_routing_config_from_env({}, config=experimental) + + +def test_env_route_rejects_invalid_direct_target() -> None: + with pytest.raises(RoutingConfigError): + load_routing_config_from_env({"MODEL_ROUTE_CODE": "not-a-target"}, config=load_config_from_mapping({})) diff --git a/tests/test_config_stream_settings.py b/tests/test_config_stream_settings.py new file mode 100644 index 000000000..5da2c5141 --- /dev/null +++ b/tests/test_config_stream_settings.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import pytest + +from rotator_library.config.experimental import ExperimentalConfigError, get_stream_runtime_settings, load_config_from_mapping + + +def test_stream_settings_parse_json_values() -> None: + config = load_config_from_mapping( + {"streaming": {"ttfb_timeout_seconds": 5, "stall_timeout_seconds": 30, "heartbeat_interval_seconds": 10, "cancel_upstream_on_disconnect": False, "trace_metrics": False}} + ) + + settings = get_stream_runtime_settings(config=config, env={}) + + assert settings.ttfb_timeout_seconds == 5 + assert settings.stall_timeout_seconds == 30 + assert settings.heartbeat_seconds == 10 + assert settings.cancel_upstream_on_disconnect is False + assert settings.trace_metrics is False + + +def test_stream_settings_env_overrides_json_values() -> None: + config = load_config_from_mapping({"streaming": {"trace_metrics": True, "heartbeat_seconds": 10}}) + + settings = get_stream_runtime_settings(config=config, env={"STREAM_TRACE_METRICS": "false", "STREAM_HEARTBEAT_INTERVAL_SECONDS": "2"}) + + assert settings.trace_metrics is False + assert settings.heartbeat_seconds == 2 + + +def test_stream_settings_accept_legacy_heartbeat_env_name() -> None: + settings = get_stream_runtime_settings(env={"STREAM_HEARTBEAT_SECONDS": "3"}) + + assert settings.heartbeat_seconds == 3 + + +def test_stream_settings_invalid_boolean_fails_clearly() -> None: + with pytest.raises(ExperimentalConfigError): + get_stream_runtime_settings(env={"STREAM_TRACE_METRICS": "maybe"}) diff --git a/tests/test_cooldown_activation.py b/tests/test_cooldown_activation.py new file mode 100644 index 000000000..b7f0490a3 --- /dev/null +++ b/tests/test_cooldown_activation.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from rotator_library.client.executor import RequestExecutor, RoutingExecutionError, _can_start_stream_provider_cooldown +from rotator_library.cooldown_manager import CooldownManager +from rotator_library.core.types import RequestContext +from rotator_library.error_handler import ClassifiedError +from rotator_library.transaction_logger import TransactionLogger + + +class FakeCooldown: + def __init__(self) -> None: + self.started = [] + self.scoped_started = [] + self.waits = [] + + async def start_cooldown(self, provider, duration): + self.started.append((provider, duration)) + + async def start_scoped_cooldown(self, provider, duration, *, model=None, scope="provider", reason=None): + self.scoped_started.append((provider, duration, model, scope, reason)) + + async def get_max_remaining(self, provider, *, model=None): + self.waits.append((provider, model)) + return 0 + + +class BudgetCooldown(FakeCooldown): + def __init__(self, remaining: float) -> None: + super().__init__() + self.remaining = remaining + + async def get_max_remaining(self, provider, *, model=None): + self.waits.append((provider, model)) + return self.remaining + + +@pytest.mark.asyncio +async def test_start_cooldown_extends_but_does_not_shorten() -> None: + manager = CooldownManager() + + await manager.start_cooldown("provider", 30) + initial = await manager.get_remaining_cooldown("provider") + await asyncio.sleep(0.01) + await manager.start_cooldown("provider", 1) + after_shorter = await manager.get_remaining_cooldown("provider") + await manager.start_cooldown("provider", 60) + after_longer = await manager.get_remaining_cooldown("provider") + + assert after_shorter > 25 + assert after_shorter <= initial + assert after_longer > after_shorter + + +@pytest.mark.asyncio +async def test_model_cooldown_is_independent_from_provider_cooldown() -> None: + manager = CooldownManager() + + await manager.start_scoped_cooldown("provider", 30, model="model-a", scope="model", reason="capacity") + + assert await manager.is_scoped_cooling_down("provider", model="model-a", scope="model") is True + assert await manager.is_scoped_cooling_down("provider", model="model-b", scope="model") is False + assert await manager.is_cooling_down("provider") is False + + +@pytest.mark.asyncio +async def test_max_remaining_uses_provider_or_model_scope() -> None: + manager = CooldownManager() + + await manager.start_cooldown("provider", 1) + await manager.start_scoped_cooldown("provider", 30, model="model-a", scope="model") + + assert await manager.get_max_remaining("provider", model="model-a") > 20 + assert 0 < await manager.get_max_remaining("provider", model="model-b") <= 1 + + +@pytest.mark.asyncio +async def test_cooldown_snapshot_reports_scopes() -> None: + manager = CooldownManager() + + await manager.start_scoped_cooldown("provider", 30, model="model-a", scope="model", reason="capacity") + snapshot = await manager.snapshot() + + assert snapshot[0].provider == "provider" + assert snapshot[0].scope == "model" + assert snapshot[0].model == "model-a" + assert snapshot[0].reason == "capacity" + + +def _classified(error_type: str, retry_after=None) -> ClassifiedError: + return ClassifiedError(error_type, original_exception=Exception(error_type), retry_after=retry_after) + + +def _executor(cooldown) -> RequestExecutor: + executor = RequestExecutor.__new__(RequestExecutor) + executor._cooldown = cooldown + from rotator_library.retry_policy import FailureHistory + + executor._failure_history = FailureHistory() + return executor + + +def _context(logger) -> RequestContext: + return RequestContext( + model="openai/gpt-test", + provider="openai", + kwargs={"model": "openai/gpt-test"}, + streaming=False, + credentials=["cred"], + deadline=9999999999.0, + transaction_logger=logger, + ) + + +@pytest.mark.asyncio +async def test_large_retry_after_starts_provider_cooldown_and_traces(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("SMALL_COOLDOWN_RETRY_THRESHOLD", "10") + monkeypatch.setenv("PROVIDER_COOLDOWN_MIN_SECONDS", "10") + cooldown = FakeCooldown() + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + + await _executor(cooldown)._maybe_start_provider_cooldown( + "openai", + _classified("rate_limit", retry_after=60), + context=_context(logger), + ) + + assert cooldown.scoped_started == [("openai", 60, None, "provider", "retry_after")] + trace = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "provider_cooldown_started" in trace + + +@pytest.mark.asyncio +async def test_small_retry_after_skips_provider_cooldown(monkeypatch) -> None: + monkeypatch.setenv("SMALL_COOLDOWN_RETRY_THRESHOLD", "10") + cooldown = FakeCooldown() + + await _executor(cooldown)._maybe_start_provider_cooldown( + "openai", + _classified("rate_limit", retry_after=3), + context=None, + ) + + assert cooldown.started == [] + assert cooldown.scoped_started == [] + + +@pytest.mark.asyncio +async def test_model_capacity_starts_model_scoped_cooldown(monkeypatch) -> None: + monkeypatch.setenv("PROVIDER_COOLDOWN_DEFAULT_SECONDS", "30") + monkeypatch.setenv("PROVIDER_COOLDOWN_MIN_SECONDS", "10") + cooldown = FakeCooldown() + executor = _executor(cooldown) + + await executor._maybe_start_provider_cooldown( + "openai", + _classified("server_error"), + context=None, + model="gpt-5", + original_error=Exception("MODEL_CAPACITY_EXHAUSTED"), + ) + + assert cooldown.scoped_started == [("openai", 30, "gpt-5", "model", "model_capacity_cooldown")] + assert executor._failure_history.snapshot()[0].scope == "model" + + +@pytest.mark.asyncio +async def test_wait_for_cooldown_uses_model_scope_when_available() -> None: + cooldown = FakeCooldown() + + await _executor(cooldown)._wait_for_cooldown("openai", 9999999999.0, model="gpt-5") + + assert cooldown.waits == [("openai", "gpt-5")] + + +@pytest.mark.asyncio +async def test_wait_for_cooldown_exceeding_budget_fails_fast() -> None: + cooldown = BudgetCooldown(remaining=60) + + with pytest.raises(RoutingExecutionError) as exc: + await _executor(cooldown)._wait_for_cooldown("openai", 1.0, model="gpt-5") + + assert exc.value.error_type == "rate_limit" + assert cooldown.waits == [("openai", "gpt-5")] + + +@pytest.mark.asyncio +async def test_generic_transient_records_history_before_starting_cooldown(monkeypatch) -> None: + monkeypatch.setenv("PROVIDER_BACKOFF_THRESHOLD", "2") + monkeypatch.setenv("PROVIDER_COOLDOWN_DEFAULT_SECONDS", "10") + monkeypatch.setenv("PROVIDER_COOLDOWN_MIN_SECONDS", "10") + cooldown = FakeCooldown() + executor = _executor(cooldown) + + await executor._maybe_start_provider_cooldown("openai", _classified("server_error"), context=None, model="gpt-5") + assert cooldown.scoped_started == [] + assert executor._failure_history.snapshot()[0].reason == "transient_backoff_threshold_not_met" + + await executor._maybe_start_provider_cooldown("openai", _classified("server_error"), context=None, model="gpt-5") + assert cooldown.scoped_started == [("openai", 10, None, "provider", "default_transient_cooldown")] + + +def test_streaming_provider_cooldown_gate_allows_only_pre_output_failures() -> None: + assert _can_start_stream_provider_cooldown(None) is True + assert _can_start_stream_provider_cooldown('data: {"error":{"type":"rate_limit"}}\n\n') is True + assert _can_start_stream_provider_cooldown('data: {"choices":[{"delta":{"content":"visible"}}]}\n\n') is False + assert _can_start_stream_provider_cooldown('data: {"usage":{"total_tokens":1}}\n\n', emitted_output=True) is False diff --git a/tests/test_copilot_provider.py b/tests/test_copilot_provider.py new file mode 100644 index 000000000..9a5f252d5 --- /dev/null +++ b/tests/test_copilot_provider.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import pytest + +from rotator_library.adapters import AdapterContext, get_adapter, run_adapter_chain +from rotator_library.providers import PROVIDER_PLUGINS +from rotator_library.providers.copilot_provider import CopilotProvider + + +class FakeResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + +class FakeClient: + def __init__(self, payload): + self.payload = payload + self.calls = [] + + async def get(self, url, *, headers, timeout): + self.calls.append({"url": url, "headers": headers, "timeout": timeout}) + return FakeResponse(self.payload) + + +def test_copilot_provider_is_discovered() -> None: + assert "copilot" in PROVIDER_PLUGINS + + +def test_copilot_provider_declares_openai_chat_protocol_and_adapter() -> None: + provider = CopilotProvider() + + assert provider.get_protocol_name("gpt-4.1") == "openai_chat" + assert provider.get_adapter_names("gpt-4.1") == ("suppress_developer_role",) + assert provider.get_adapter_config("gpt-4.1") == {"suppress_developer_role": {"mode": "system"}} + assert provider.get_field_cache_rules("gpt-4.1") == () + + +@pytest.mark.asyncio +async def test_copilot_adapter_config_converts_developer_role_to_system() -> None: + provider = CopilotProvider() + + result = await run_adapter_chain( + [get_adapter("suppress_developer_role")], + {"messages": [{"role": "developer", "content": "rules"}]}, + AdapterContext(adapter_config=provider.get_adapter_config("gpt-4.1")), + stage="request", + ) + + assert result["messages"][0]["role"] == "system" + + +def test_copilot_provider_builds_native_headers_and_endpoint(monkeypatch) -> None: + monkeypatch.setenv("COPILOT_API_BASE", "https://copilot.test") + monkeypatch.setenv("COPILOT_INTEGRATION_ID", "proxy-test") + provider = CopilotProvider() + + assert provider.get_native_endpoint(operation="chat") == "https://copilot.test/chat/completions" + assert provider.get_native_endpoint(operation="models") == "https://copilot.test/models" + assert provider.get_native_headers("token") == { + "Authorization": "Bearer token", + "content-type": "application/json", + "Copilot-Integration-Id": "proxy-test", + } + + +def test_copilot_native_operation_model_and_stream_support() -> None: + provider = CopilotProvider() + + assert provider.get_native_operation("gpt-4.1", {}, stream=False) == "chat" + assert provider.normalize_native_model("copilot/gpt-4.1") == "gpt-4.1" + assert provider.supports_native_streaming("gpt-4.1", operation="chat") is True + assert provider.supports_native_streaming("gpt-4.1", operation="responses") is False + + +@pytest.mark.asyncio +async def test_copilot_provider_get_models_uses_mocked_endpoint(monkeypatch) -> None: + monkeypatch.setenv("COPILOT_API_BASE", "https://copilot.test") + client = FakeClient({"data": [{"id": "gpt-4.1"}]}) + provider = CopilotProvider() + + models = await provider.get_models("token", client) + + assert models == ["copilot/gpt-4.1"] + assert client.calls[0]["url"] == "https://copilot.test/models" + + +@pytest.mark.asyncio +async def test_copilot_provider_get_models_falls_back_on_errors() -> None: + class BrokenClient: + async def get(self, *args, **kwargs): + raise RuntimeError("offline") + + assert await CopilotProvider().get_models("token", BrokenClient()) == ["copilot/gpt-4.1", "copilot/claude-sonnet-4-5"] diff --git a/tests/test_env_example_experimental_config.py b/tests/test_env_example_experimental_config.py new file mode 100644 index 000000000..b8b915989 --- /dev/null +++ b/tests/test_env_example_experimental_config.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from pathlib import Path + + +def test_env_example_documents_experimental_config_knobs() -> None: + text = Path(".env.example").read_text(encoding="utf-8") + + for key in ( + "LLM_PROXY_CONFIG_FILE", + "FALLBACK_GROUPS", + "FALLBACK_GROUP_CODE_CHAIN", + "MODEL_ROUTE_CODE", + "PROVIDER_COOLDOWN_MIN_SECONDS", + "PROVIDER_COOLDOWN_DEFAULT_SECONDS", + "PROVIDER_COOLDOWN_ON_QUOTA", + "PROVIDER_BACKOFF_WINDOW_SECONDS", + "PROVIDER_BACKOFF_THRESHOLD", + "PROVIDER_BACKOFF_BASE_SECONDS", + "PROVIDER_BACKOFF_MAX_SECONDS", + "FAILURE_HISTORY_MAX_ENTRIES", + "RESPONSES_STORE_TTL_SECONDS", + "RESPONSES_STORE_MAX_ITEMS", + "RESPONSES_STORE_FAILED", + "RESPONSES_STORE_IN_PROGRESS", + "STREAM_TRACE_METRICS", + "STREAM_TTFB_TIMEOUT_SECONDS", + "STREAM_STALL_TIMEOUT_SECONDS", + "STREAM_HEARTBEAT_INTERVAL_SECONDS", + "STREAM_HEARTBEAT_SECONDS", + "STREAM_CANCEL_UPSTREAM_ON_DISCONNECT", + "MODEL_PRICE_OPENAI_GPT_5_1_INPUT", + "MODEL_PRICE_OPENAI_GPT_5_1_REASONING", + ): + assert key in text + + assert "Do not put API keys" in text + assert "protocol_name (openai_chat, responses, anthropic_messages, or gemini)" in text + assert "default_output_protocol" in text + assert "model_quota_groups" in text + assert "credentials must stay in" in text + for default_line in ( + "PROVIDER_COOLDOWN_MIN_SECONDS=10", + "PROVIDER_COOLDOWN_DEFAULT_SECONDS=30", + "PROVIDER_COOLDOWN_ON_QUOTA=false", + "PROVIDER_BACKOFF_WINDOW_SECONDS=60", + "PROVIDER_BACKOFF_THRESHOLD=3", + "PROVIDER_BACKOFF_BASE_SECONDS=0", + "PROVIDER_BACKOFF_MAX_SECONDS=300", + "FAILURE_HISTORY_MAX_ENTRIES=200", + "RESPONSES_STORE_TTL_SECONDS=0", + "RESPONSES_STORE_MAX_ITEMS=0", + "RESPONSES_STORE_FAILED=true", + "RESPONSES_STORE_IN_PROGRESS=false", + "STREAM_TRACE_METRICS=true", + "STREAM_TTFB_TIMEOUT_SECONDS=0", + "STREAM_STALL_TIMEOUT_SECONDS=0", + "STREAM_HEARTBEAT_INTERVAL_SECONDS=0", + "STREAM_CANCEL_UPSTREAM_ON_DISCONNECT=true", + ): + assert default_line in text diff --git a/tests/test_executor_usage_accounting.py b/tests/test_executor_usage_accounting.py new file mode 100644 index 000000000..1f28cc537 --- /dev/null +++ b/tests/test_executor_usage_accounting.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import json +from types import SimpleNamespace + +from rotator_library.client.executor import RequestExecutor +from rotator_library.core.types import RequestContext +from rotator_library.transaction_logger import TransactionLogger + + +def _executor() -> RequestExecutor: + return RequestExecutor({}, None, None, None, {}, None) + + +def test_executor_accounts_for_non_streaming_usage_and_cost_trace(tmp_path, monkeypatch) -> None: + executor = _executor() + response = SimpleNamespace( + usage=SimpleNamespace( + prompt_tokens=100, + completion_tokens=30, + prompt_tokens_details={"cached_tokens": 40}, + completion_tokens_details={"reasoning_tokens": 10}, + ) + ) + logger = TransactionLogger("openai", "gpt-test", parent_dir=tmp_path) + context = RequestContext( + provider="openai", + model="gpt-test", + kwargs={}, + streaming=False, + credentials=[], + deadline=0, + transaction_logger=logger, + ) + monkeypatch.setattr( + "rotator_library.usage.costs.litellm.get_model_info", + lambda model: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + + usage, cost = executor._account_for_response_usage("openai", "gpt-test", response, context) + + assert usage.input_tokens == 60 + assert usage.cache_read_tokens == 40 + assert usage.completion_tokens == 20 + assert usage.reasoning_tokens == 10 + assert cost.total_cost > 0 + entries = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + assert entries[-1]["pass_name"] == "usage_accounting_summary" + assert entries[-1]["data"]["usage"]["total_tokens"] == usage.total_tokens + + +def test_executor_accounting_uses_configured_env_pricing(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("MODEL_PRICE_OPENAI_GPT_TEST_INPUT", "2.0") + executor = _executor() + response = SimpleNamespace(usage=SimpleNamespace(prompt_tokens=3, completion_tokens=0)) + context = RequestContext( + provider="openai", + model="gpt-test", + kwargs={}, + streaming=False, + credentials=[], + deadline=0, + transaction_logger=TransactionLogger("openai", "gpt-test", parent_dir=tmp_path), + ) + + _, cost = executor._account_for_response_usage("openai", "gpt-test", response, context) + + assert cost.pricing_source == "env" + assert cost.input_cost == 6.0 + + +def test_normalize_response_usage_handles_dict_responses() -> None: + response = { + "usage": { + "prompt_tokens": 4, + "completion_tokens": 1, + "completion_tokens_details": {"reasoning_tokens": 3}, + } + } + + result = RequestExecutor._normalize_response_usage(response, "gpt-test") + + assert result is response + assert response["usage"]["completion_tokens"] == 4 + assert response["usage"]["total_tokens"] == 8 diff --git a/tests/test_experimental_config.py b/tests/test_experimental_config.py new file mode 100644 index 000000000..e5099c274 --- /dev/null +++ b/tests/test_experimental_config.py @@ -0,0 +1,338 @@ +from __future__ import annotations + +import pytest + +from rotator_library.config.experimental import ( + ExperimentalConfigError, + as_int, + env_price_key, + get_responses_store_settings, + get_responses_store_runtime_settings, + get_retry_runtime_settings, + get_stream_runtime_settings, + load_config_from_mapping, + load_experimental_config, + parse_field_cache_rules, +) + + +def test_missing_config_file_returns_empty(tmp_path) -> None: + config = load_experimental_config(tmp_path / "missing.json", env={}) + + assert config.is_empty + + +def test_loads_config_from_env_path(tmp_path) -> None: + path = tmp_path / "config.json" + path.write_text('{"routing":{"model_routes":{"code":"group:chain"}},"extra":{}}', encoding="utf-8") + + config = load_experimental_config(env={"LLM_PROXY_CONFIG_FILE": str(path)}) + + assert config.routing["model_routes"]["code"] == "group:chain" + assert config.unknown_sections == {"extra": {}} + assert config.warnings + + +def test_rejects_secret_like_json_keys() -> None: + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"providers": {"openai": {"api_key": "hidden"}}}) + + +def test_invalid_json_raises(tmp_path) -> None: + path = tmp_path / "config.json" + path.write_text("{", encoding="utf-8") + + with pytest.raises(ExperimentalConfigError): + load_experimental_config(path, env={}) + + +def test_stream_runtime_settings_env_overrides_json() -> None: + config = load_config_from_mapping({"streaming": {"trace_metrics": True, "stall_timeout_seconds": 30}}) + + settings = get_stream_runtime_settings(config=config, env={"STREAM_TRACE_METRICS": "false"}) + + assert settings.trace_metrics is False + assert settings.stall_timeout_seconds == 30 + + +def test_retry_runtime_settings_env_overrides_json() -> None: + config = load_config_from_mapping( + { + "retry": { + "provider_cooldown": {"provider_cooldown_min_seconds": 20, "provider_cooldown_on_quota": True}, + "backoff": {"provider_backoff_threshold": 4, "failure_history_max_entries": 50}, + } + } + ) + + settings = get_retry_runtime_settings(config=config, env={"PROVIDER_COOLDOWN_MIN_SECONDS": "5"}) + + assert settings.provider_cooldown_min_seconds == 5 + assert settings.provider_cooldown_on_quota is True + assert settings.provider_backoff_threshold == 4 + assert settings.failure_history_max_entries == 50 + + +def test_responses_store_settings_env_overrides_json() -> None: + config = load_config_from_mapping( + { + "responses": { + "store": { + "ttl_seconds": 60, + "max_items": 10, + "store_failed": False, + "store_in_progress": True, + } + } + } + ) + + settings = get_responses_store_settings(config=config, env={"RESPONSES_STORE_MAX_ITEMS": "5", "RESPONSES_STORE_FAILED": "true"}) + + assert settings.ttl_seconds == 60 + assert settings.max_items == 5 + assert settings.store_failed is True + assert settings.store_in_progress is True + + +def test_responses_store_runtime_settings_env_overrides_json(tmp_path) -> None: + config = load_config_from_mapping( + { + "responses": { + "store": { + "backend": "memory", + "cache_name": "json_responses", + "cache_prefix": "json_prefix", + "cache_dir": str(tmp_path / "json-cache"), + "cache_memory_ttl_seconds": 10, + "cache_disk_ttl_seconds": 20, + } + } + } + ) + + settings = get_responses_store_runtime_settings( + config=config, + env={"RESPONSES_STORE_BACKEND": "provider_cache", "RESPONSES_STORE_CACHE_MEMORY_TTL_SECONDS": "30"}, + ) + + assert settings.backend == "provider_cache" + assert settings.cache_name == "json_responses" + assert settings.cache_prefix == "json_prefix" + assert settings.cache_dir == str(tmp_path / "json-cache") + assert settings.cache_memory_ttl_seconds == 30 + assert settings.cache_disk_ttl_seconds == 20 + + +def test_responses_store_runtime_rejects_unknown_backend() -> None: + with pytest.raises(ExperimentalConfigError): + get_responses_store_runtime_settings(config=load_config_from_mapping({"responses": {"store": {"backend": "sqlite"}}}), env={}) + + +def test_new_config_sections_still_reject_secret_like_keys() -> None: + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"responses": {"store": {"authorization": "hidden"}}}) + + +@pytest.mark.parametrize("secret_key", ["secret_key", "secret-key", "apiKey", "client-secret", "oauth_token", "oauthToken", "oauth-token", "id_token", "oauth_token_secret", "credential", "credentials", "providerCredential", "oauthCredentials"]) +def test_secret_key_variants_are_rejected(secret_key: str) -> None: + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"retry": {secret_key: "hidden"}}) + + +def test_provider_config_rejects_unsupported_keys() -> None: + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"providers": {"openai": {"unknown_setting": True}}}) + + +def test_provider_config_validates_protocol_and_adapter_names() -> None: + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"providers": {"openai": {"protocol_name": "missing_protocol"}}}) + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"providers": {"openai": {"adapter_names": ["missing_adapter"]}}}) + + +@pytest.mark.parametrize( + "provider_config", + [ + {"api_base": "file:///tmp/provider"}, + {"api_base": "https://user:secret@example.test/v1"}, + {"api_base": "https://example.test/v1?token=secret"}, + {"auth_mode": "query-secret"}, + {"auth_mode": "custom"}, + {"api_base": "https://example.test", "auth_mode": "none"}, + {"auth_mode": "custom", "auth_header_name": "Bad Header"}, + {"models": "model-a"}, + {"models": [{"id": "model-a"}]}, + {"endpoint_paths": ["/chat/completions"]}, + {"endpoint_paths": {"chat": "chat/completions"}}, + {"endpoint_paths": {"chat": "https://other.example/chat"}}, + {"endpoint_paths": {"chat": "/chat?api_key=secret"}}, + {"endpoint_paths": {"chat": "/chat#token=secret"}}, + {"endpoint_paths": {"chat": "/models/{unknown}:generate"}}, + {"endpoint_paths": {"chat": 42}}, + {"default_output_protocol": "missing_protocol"}, + {"default_output_protocol": "embeddings"}, + ], +) +def test_provider_config_validates_dynamic_transport_settings(provider_config) -> None: + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"providers": {"configured": provider_config}}) + + +def test_provider_config_rejects_invalid_provider_names() -> None: + with pytest.raises(ExperimentalConfigError): + load_config_from_mapping({"providers": {"bad/provider": {"api_base": "https://example.test"}}}) + + +def test_retry_runtime_settings_malformed_env_preserves_defaults() -> None: + config = load_config_from_mapping({"retry": {"provider_cooldown_default_seconds": 45}}) + + settings = get_retry_runtime_settings( + config=config, + env={ + "PROVIDER_COOLDOWN_DEFAULT_SECONDS": "not-an-int", + "PROVIDER_COOLDOWN_ON_QUOTA": "not-a-bool", + "PROVIDER_BACKOFF_BASE_SECONDS": "not-an-int", + }, + ) + + assert settings.provider_cooldown_default_seconds == 30 + assert settings.provider_cooldown_on_quota is False + assert settings.provider_backoff_base_seconds is None + + +def test_field_cache_rules_parse_wildcard_then_model_specific() -> None: + config = load_config_from_mapping( + { + "field_cache": { + "gemini_cli": { + "*": [{"name": "thought", "source": "response", "path": "$.thought", "target_path": "$.cached_thought"}], + "gemini-3": [{"name": "signature", "source": "response", "path": "$.sig", "scope": ["provider", "model", "credential", "session"]}], + } + } + } + ) + + rules = parse_field_cache_rules(config, "gemini_cli", "gemini-3") + + assert [rule.name for rule in rules] == ["thought", "signature"] + assert rules[0].inject is not None + + +def test_field_cache_rules_match_unprefixed_model_alias() -> None: + config = load_config_from_mapping( + { + "field_cache": { + "gemini_cli": { + "gemini-3": [{"name": "signature", "source": "response", "path": "sig"}], + } + } + } + ) + + rules = parse_field_cache_rules(config, "gemini_cli", "gemini_cli/gemini-3") + + assert [rule.name for rule in rules] == ["signature"] + + +def test_field_cache_rule_parses_ttl_metadata_and_insert_injection() -> None: + config = load_config_from_mapping( + { + "field_cache": { + "provider": { + "*": [ + { + "name": "tool_state", + "cache_key": "shared_tool_state", + "source": "stream_event", + "path": "raw.tool.state", + "mode": "per_tool_call", + "ttl_seconds": 120, + "metadata": {"tool_container_path": "tools"}, + "inject": {"target": "request", "path": "metadata.tool_state", "insert": True}, + } + ] + } + } + } + ) + + rule = parse_field_cache_rules(config, "provider", "model")[0] + + assert rule.ttl_seconds == 120 + assert rule.cache_key == "shared_tool_state" + assert rule.metadata == {"tool_container_path": "tools"} + assert rule.inject is not None + assert rule.inject.insert is True + + +def test_field_cache_rule_rejects_invalid_config_values() -> None: + config = load_config_from_mapping( + { + "field_cache": { + "provider": { + "*": [ + {"name": "bad_source", "source": "not_a_source", "path": "x"}, + ] + } + } + } + ) + + with pytest.raises(ValueError, match="Unsupported field-cache source"): + parse_field_cache_rules(config, "provider", "model") + + +def test_field_cache_rules_reject_malformed_shapes() -> None: + config = load_config_from_mapping({"field_cache": {"provider": {"*": ["not-a-rule"]}}}) + + with pytest.raises(ExperimentalConfigError, match="rule entries"): + parse_field_cache_rules(config, "provider", "model") + + +def test_field_cache_config_rejects_weakened_provider_state_scope() -> None: + config = load_config_from_mapping({ + "field_cache": { + "provider": { + "*": [{ + "name": "unsafe", + "source": "response", + "path": "state", + "scope": ["provider", "model", "session"], + }] + } + } + }) + + with pytest.raises(ExperimentalConfigError, match="credential"): + parse_field_cache_rules(config, "provider", "model") + + +def test_field_cache_rules_reject_non_list_model_rules() -> None: + config = load_config_from_mapping({"field_cache": {"provider": {"*": {"name": "bad"}}}}) + + with pytest.raises(ExperimentalConfigError, match="model rules"): + parse_field_cache_rules(config, "provider", "model") + + +def test_field_cache_rules_reject_malformed_nested_shapes() -> None: + bad_inject = load_config_from_mapping({"field_cache": {"provider": {"*": [{"name": "bad", "source": "response", "path": "x", "inject": "not-object"}]}}}) + + with pytest.raises(ExperimentalConfigError, match="inject"): + parse_field_cache_rules(bad_inject, "provider", "model") + + bad_metadata = load_config_from_mapping({"field_cache": {"provider": {"*": [{"name": "bad", "source": "response", "path": "x", "metadata": "not-object"}]}}}) + + with pytest.raises(ExperimentalConfigError, match="metadata"): + parse_field_cache_rules(bad_metadata, "provider", "model") + + +def test_env_price_key_sanitizes_provider_and_model() -> None: + assert env_price_key("openai", "gpt-5.1-mini", "cache_read") == "MODEL_PRICE_OPENAI_GPT_5_1_MINI_CACHE_READ" + + +def test_as_int_parses_integers_with_redacted_errors() -> None: + assert as_int("5", name="TEST_INT") == 5 + with pytest.raises(ExperimentalConfigError, match="TEST_INT"): + as_int("not-secret-value", name="TEST_INT") diff --git a/tests/test_fallback_attempt_runner.py b/tests/test_fallback_attempt_runner.py new file mode 100644 index 000000000..167e5698d --- /dev/null +++ b/tests/test_fallback_attempt_runner.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import pytest + +from rotator_library.routing import FallbackAttemptRunner, FallbackExhaustedError, RoutingDecision, parse_route_target +from rotator_library.routing.types import FallbackGroup + + +class ClassifiedFailure(Exception): + def __init__(self, error_type: str, *, emitted_output: bool = False) -> None: + super().__init__(error_type) + self.error_type = error_type + self.emitted_output = emitted_output + + +def _decision() -> RoutingDecision: + targets = (parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1")) + return RoutingDecision( + requested_model="code", + group_name="code_chain", + targets=targets, + group=FallbackGroup(name="code_chain", targets=targets), + reason="model_route_group", + ) + + +@pytest.mark.asyncio +async def test_attempt_runner_returns_first_success_without_fallback() -> None: + calls = [] + + async def attempt(target, index): + calls.append(target.prefixed_model) + return {"target": target.prefixed_model} + + result = await FallbackAttemptRunner().run(_decision(), attempt) + + assert result == {"target": "codex/gpt-5.1-codex"} + assert calls == ["codex/gpt-5.1-codex"] + + +@pytest.mark.asyncio +async def test_attempt_runner_falls_back_on_retryable_error() -> None: + calls = [] + + async def attempt(target, index): + calls.append(target.prefixed_model) + if index == 0: + raise ClassifiedFailure("rate_limit") + return {"target": target.prefixed_model} + + result = await FallbackAttemptRunner().run(_decision(), attempt) + + assert result == {"target": "openai/gpt-5.1"} + assert calls == ["codex/gpt-5.1-codex", "openai/gpt-5.1"] + + +@pytest.mark.asyncio +async def test_attempt_runner_stops_on_permanent_error() -> None: + async def attempt(target, index): + raise ClassifiedFailure("validation") + + with pytest.raises(FallbackExhaustedError) as exc: + await FallbackAttemptRunner().run(_decision(), attempt) + + assert len(exc.value.attempts) == 1 + assert exc.value.attempts[0].error_type == "invalid_request" + + +@pytest.mark.asyncio +async def test_attempt_runner_blocks_stream_fallback_after_output() -> None: + async def attempt(target, index): + raise ClassifiedFailure("rate_limit", emitted_output=True) + + with pytest.raises(FallbackExhaustedError) as exc: + await FallbackAttemptRunner().run(_decision(), attempt, stream=True) + + assert len(exc.value.attempts) == 1 + assert exc.value.attempts[0].emitted_output is True + + +@pytest.mark.asyncio +async def test_attempt_runner_hard_stops_group_policy_overrides() -> None: + group = FallbackGroup( + name="custom", + targets=_decision().targets, + failover_on=frozenset({"authentication"}), + stop_on=frozenset({"validation"}), + ) + calls = [] + + async def attempt(target, index): + calls.append(index) + if index == 0: + raise ClassifiedFailure("authentication") + return {"target": target.prefixed_model} + + with pytest.raises(FallbackExhaustedError): + await FallbackAttemptRunner().run_group(_decision(), group, attempt) + + assert calls == [0] + + +@pytest.mark.asyncio +async def test_attempt_runner_respects_never_streaming_policy() -> None: + group = FallbackGroup( + name="custom", + targets=_decision().targets, + failover_on=frozenset({"rate_limit"}), + streaming_policy="never", + ) + calls = [] + + async def attempt(target, index): + calls.append(index) + raise ClassifiedFailure("rate_limit") + + with pytest.raises(FallbackExhaustedError): + await FallbackAttemptRunner().run_group(_decision(), group, attempt, stream=True) + + assert calls == [0] + + +@pytest.mark.asyncio +async def test_attempt_runner_run_uses_decision_group_streaming_policy() -> None: + decision = _decision() + never_group = FallbackGroup(name="code_chain", targets=decision.targets, failover_on=frozenset({"rate_limit"}), streaming_policy="never") + decision = RoutingDecision(requested_model=decision.requested_model, group_name=decision.group_name, targets=decision.targets, group=never_group) + calls = [] + + async def attempt(target, index): + calls.append(index) + raise ClassifiedFailure("rate_limit") + + with pytest.raises(FallbackExhaustedError): + await FallbackAttemptRunner().run(decision, attempt, stream=True) + + assert calls == [0] diff --git a/tests/test_fallback_groups.py b/tests/test_fallback_groups.py new file mode 100644 index 000000000..66f24bdb0 --- /dev/null +++ b/tests/test_fallback_groups.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import pytest + +from rotator_library.routing import FallbackResolver, FallbackGroup, RouteTarget, RoutingConfig +from rotator_library.routing.config import RoutingConfigError, load_routing_config_from_env, parse_route_target +from rotator_library.routing.executor import FallbackAttemptRunner + + +def test_parse_route_target_supports_execution_suffix() -> None: + target = parse_route_target("openai/gpt-5@litellm_fallback") + + assert target.provider == "openai" + assert target.model == "gpt-5" + assert target.execution == "litellm_fallback" + + +def test_env_fallback_group_config_loads_current_routing_types() -> None: + config = load_routing_config_from_env( + { + "FALLBACK_GROUPS": "main", + "FALLBACK_GROUP_MAIN": "openai/gpt-5@litellm_fallback,anthropic/claude@native", + "MODEL_ROUTE_GPT5": "group:main", + }, + config=RoutingConfig(), + ) + + group = config.fallback_groups["main"] + assert [target.prefixed_model for target in group.targets] == ["openai/gpt-5", "anthropic/claude"] + assert group.targets[0].execution == "litellm_fallback" + assert group.targets[1].execution == "native" + assert config.model_routes["gpt5"] == "group:main" + + +def test_resolver_promotes_requested_provider_model_inside_group() -> None: + group = FallbackGroup( + name="main", + targets=( + RouteTarget("openai", "gpt-5"), + RouteTarget("anthropic", "claude"), + RouteTarget("google", "gemini"), + ), + ) + decision = FallbackResolver(RoutingConfig(fallback_groups={"main": group})).resolve("anthropic/claude") + + assert decision.reason == "provider_model_group_promoted" + assert [target.prefixed_model for target in decision.targets] == ["anthropic/claude", "openai/gpt-5", "google/gemini"] + + +def test_resolver_promotes_requested_model_for_group_route_alias() -> None: + group = FallbackGroup( + name="main", + targets=(RouteTarget("openai", "gpt-5"), RouteTarget("anthropic", "claude")), + ) + config = RoutingConfig(fallback_groups={"main": group}, model_routes={"anthropic/claude": "group:main"}) + + decision = FallbackResolver(config).resolve("anthropic/claude") + + assert decision.reason == "model_route_group_promoted" + assert [target.prefixed_model for target in decision.targets] == ["anthropic/claude", "openai/gpt-5"] + + +def test_resolver_rejects_missing_group_route() -> None: + with pytest.raises(RoutingConfigError): + FallbackResolver(RoutingConfig(model_routes={"alias": "group:missing"})).resolve("alias") + + +@pytest.mark.asyncio +async def test_fallback_runner_uses_decision_group_policy() -> None: + group = FallbackGroup(name="main", targets=(RouteTarget("a", "one"), RouteTarget("b", "two")), failover_on=frozenset({"rate_limit"})) + decision = FallbackResolver(RoutingConfig(fallback_groups={"main": group}, model_routes={"alias": "group:main"})).resolve("alias") + attempts: list[str] = [] + + class RateLimitError(RuntimeError): + error_type = "rate_limit" + + async def attempt(target, index): + attempts.append(target.prefixed_model) + if target.provider == "a": + raise RateLimitError("rate limit") + return "ok" + + result = await FallbackAttemptRunner().run(decision, attempt) + + assert result == "ok" + assert attempts == ["a/one", "b/two"] diff --git a/tests/test_fallback_policy.py b/tests/test_fallback_policy.py new file mode 100644 index 000000000..74adaeaab --- /dev/null +++ b/tests/test_fallback_policy.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from rotator_library.routing import FallbackPolicy, parse_route_target +from rotator_library.routing.policy import normalize_route_error_type +from rotator_library.routing.types import FallbackGroup + + +def test_policy_falls_back_on_retryable_categories() -> None: + policy = FallbackPolicy() + + assert policy.should_fallback("rate_limit") is True + assert policy.should_fallback("quota_exceeded") is True + assert policy.should_fallback("server_error") is True + assert policy.should_fallback("api_connection") is True + + +def test_policy_stops_on_permanent_categories() -> None: + policy = FallbackPolicy() + + assert policy.should_fallback("authentication") is False + assert policy.should_fallback("forbidden") is False + assert policy.should_fallback("invalid_request") is False + assert policy.should_fallback("context_window_exceeded") is False + assert policy.should_fallback("credential_reauth_needed") is False + assert policy.should_fallback("pre_request_callback_error") is False + assert policy.should_fallback("cancelled") is False + + +def test_policy_blocks_stream_fallback_after_visible_output() -> None: + assert FallbackPolicy().should_fallback("rate_limit", stream=True, emitted_output=True) is False + + +def test_policy_allows_stream_fallback_before_visible_output() -> None: + assert FallbackPolicy().should_fallback("rate_limit", stream=True, emitted_output=False) is True + + +def test_policy_respects_safe_group_overrides() -> None: + group = FallbackGroup( + name="auth_safe", + targets=(parse_route_target("a/model"), parse_route_target("b/model")), + failover_on=frozenset({"network"}), + stop_on=frozenset({"validation"}), + ) + + assert FallbackPolicy().should_fallback("api_connection", group=group) is True + assert FallbackPolicy().should_fallback("validation", group=group) is False + + +def test_policy_hard_stops_cannot_be_overridden_by_group_failover() -> None: + group = FallbackGroup( + name="unsafe", + targets=(parse_route_target("a/model"), parse_route_target("b/model")), + failover_on=frozenset({"auth", "configuration"}), + stop_on=frozenset(), + ) + + assert FallbackPolicy().should_fallback("authentication", group=group) is False + assert FallbackPolicy().should_fallback("configuration_error", group=group) is False + + +def test_policy_normalizes_user_facing_aliases() -> None: + assert normalize_route_error_type("auth") == "authentication" + assert normalize_route_error_type("permission-denied") == "forbidden" + assert normalize_route_error_type("bad request") == "invalid_request" + assert normalize_route_error_type("context_length_exceeded") == "context_window_exceeded" + assert FallbackPolicy().should_fallback("network") is True + assert FallbackPolicy().should_fallback("validation") is False + + +def test_policy_normalizes_common_structured_provider_aliases() -> None: + assert normalize_route_error_type("invalid_api_key") == "authentication" + assert normalize_route_error_type("unauthorized") == "authentication" + assert normalize_route_error_type("invalid_argument") == "invalid_request" + assert normalize_route_error_type("max_tokens_exceeded") == "context_window_exceeded" + assert normalize_route_error_type("rate_limited") == "rate_limit" + assert normalize_route_error_type("too_many_requests") == "rate_limit" + assert normalize_route_error_type("resource_exhausted") == "quota_exceeded" + assert normalize_route_error_type("unavailable") == "server_error" + assert normalize_route_error_type("deadline_exceeded") == "api_connection" diff --git a/tests/test_fallback_resolver.py b/tests/test_fallback_resolver.py new file mode 100644 index 000000000..674a5d559 --- /dev/null +++ b/tests/test_fallback_resolver.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import pytest + +from rotator_library.routing import FallbackResolver, RoutingConfigError, load_routing_config_from_env + + +def test_resolver_maps_alias_to_fallback_group_in_order() -> None: + config = load_routing_config_from_env( + { + "FALLBACK_GROUPS": "code_chain", + "FALLBACK_GROUP_CODE_CHAIN": "codex/gpt-5.1-codex,openai/gpt-5.1", + "MODEL_ROUTE_CODEX": "group:code_chain", + } + ) + + decision = FallbackResolver(config).resolve("codex") + + assert decision.group_name == "code_chain" + assert [target.prefixed_model for target in decision.targets] == ["codex/gpt-5.1-codex", "openai/gpt-5.1"] + assert decision.reason == "model_route_group" + + +def test_resolver_keeps_provider_prefixed_model_as_direct_target() -> None: + decision = FallbackResolver(load_routing_config_from_env({})).resolve("openai/gpt-5.1") + + assert decision.group_name is None + assert decision.targets[0].provider == "openai" + assert decision.targets[0].model == "gpt-5.1" + assert decision.reason == "direct_provider_model" + + +def test_resolver_maps_alias_to_single_target() -> None: + config = load_routing_config_from_env({"MODEL_ROUTE_FAST": "openai/gpt-5.1"}) + + decision = FallbackResolver(config).resolve("fast") + + assert decision.targets[0].prefixed_model == "openai/gpt-5.1" + assert decision.reason == "model_route_target" + + +def test_resolver_rejects_unprefixed_model_without_route() -> None: + with pytest.raises(RoutingConfigError): + FallbackResolver(load_routing_config_from_env({})).resolve("gpt-5.1") diff --git a/tests/test_field_cache_engine.py b/tests/test_field_cache_engine.py new file mode 100644 index 000000000..0760dda88 --- /dev/null +++ b/tests/test_field_cache_engine.py @@ -0,0 +1,862 @@ +from __future__ import annotations + +import asyncio +import json + +import pytest + +from rotator_library.field_cache import ( + FieldCacheContext, + FieldCacheEngine, + FieldCacheInjection, + FieldCacheRule, + InMemoryFieldCacheStore, + ProviderCacheFieldStore, + build_cache_key, +) +from rotator_library.field_cache.types import is_provider_continuation_path +from rotator_library.transaction_logger import TransactionLogger + + +def _trace_entries(log_dir): + return [json.loads(line) for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +def _reasoning_rule(mode: str = "last", scope=("provider", "model", "credential", "session")) -> FieldCacheRule: + return FieldCacheRule( + name="reasoning_content", + source="response", + path="choices.*.message.reasoning_content", + mode=mode, + scope=scope, + inject=FieldCacheInjection(target="request", path="messages[-1].reasoning_content"), + ) + + +def _context(**overrides) -> FieldCacheContext: + values = {"provider": "openai", "model": "gpt-test", "credential_id": "credential-a", "session_id": "session-a", "classifier": "global"} + values.update(overrides) + return FieldCacheContext(**values) + + +@pytest.mark.asyncio +async def test_extract_response_value_and_inject_into_next_request() -> None: + engine = FieldCacheEngine([_reasoning_rule()]) + response = {"choices": [{"message": {"reasoning_content": "hidden"}}]} + request = {"messages": [{"role": "user", "content": "hi"}]} + + operations = await engine.extract("response", response, _context()) + updated, injection_operations = await engine.inject("request", request, _context()) + + assert operations[0].matched == 1 + assert injection_operations[0].hit is True + assert updated["messages"][-1]["reasoning_content"] == "hidden" + assert "reasoning_content" not in request["messages"][-1] + + +@pytest.mark.asyncio +async def test_last_mode_overwrites_prior_value() -> None: + engine = FieldCacheEngine([_reasoning_rule()]) + + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "first"}}]}, _context()) + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "second"}}]}, _context()) + updated, _ = await engine.inject("request", {"messages": [{}]}, _context()) + + assert updated["messages"][-1]["reasoning_content"] == "second" + + +@pytest.mark.asyncio +async def test_all_mode_appends_values() -> None: + engine = FieldCacheEngine([_reasoning_rule(mode="all")]) + + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "first"}}]}, _context()) + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "second"}}]}, _context()) + updated, _ = await engine.inject("request", {"messages": [{}]}, _context()) + + assert updated["messages"][-1]["reasoning_content"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_source_counterparts_share_one_logical_cache() -> None: + injection = FieldCacheInjection(target="request", path="metadata.signature") + response_rule = FieldCacheRule( + name="response_signature", + cache_key="provider_signature", + source="response", + path="signature", + scope=("provider", "model", "session"), + inject=injection, + ) + stream_rule = FieldCacheRule( + name="stream_signature", + cache_key="provider_signature", + source="stream_event", + path="raw.signature", + scope=("provider", "model", "session"), + inject=injection, + ) + engine = FieldCacheEngine([response_rule, stream_rule]) + + await engine.extract("response", {"signature": "non-stream"}, _context()) + first, operations = await engine.inject("request", {"metadata": {}}, _context()) + await engine.extract("stream_event", {"raw": {"signature": "stream"}}, _context()) + second, _ = await engine.inject("request", {"metadata": {}}, _context()) + + assert build_cache_key(response_rule, _context()) == build_cache_key(stream_rule, _context()) + assert first["metadata"]["signature"] == "non-stream" + assert second["metadata"]["signature"] == "stream" + assert [operation.hit for operation in operations] == [True, True] + + +def test_shared_cache_key_rejects_incompatible_counterparts() -> None: + response_rule = FieldCacheRule( + name="response_signature", + cache_key="provider_signature", + source="response", + path="signature", + mode="last", + ) + stream_rule = FieldCacheRule( + name="stream_signature", + cache_key="provider_signature", + source="stream_event", + path="raw.signature", + mode="all", + ) + + with pytest.raises(ValueError, match="identical mode, scope, TTL, injection, and correlation"): + FieldCacheEngine([response_rule, stream_rule]) + + +def test_shared_cache_key_rejects_continuation_filter_mismatch() -> None: + common = { + "cache_key": "provider_continuation", + "path": "id", + "inject": FieldCacheInjection( + target="request", + path="previous_response_id", + when_missing_only=True, + ), + } + response_rule = FieldCacheRule( + name="response_id", + source="response", + metadata={"provider_continuation": True}, + **common, + ) + assert response_rule.metadata["provider_continuation"] is True + with pytest.raises(ValueError, match="metadata.provider_continuation"): + FieldCacheRule( + name="stream_id", + source="stream_event", + metadata={"provider_continuation": False}, + **common, + ) + + +@pytest.mark.asyncio +async def test_scope_isolation_by_session_and_classifier() -> None: + rule = _reasoning_rule(scope=("provider", "model", "session", "classifier")) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "a"}}]}, _context(session_id="session-a")) + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "b"}}]}, _context(session_id="session-b")) + + updated, _ = await engine.inject("request", {"messages": [{}]}, _context(session_id="session-a")) + + assert updated["messages"][-1]["reasoning_content"] == "a" + assert build_cache_key(rule, _context(session_id="session-a")) != build_cache_key(rule, _context(session_id="session-b")) + + +@pytest.mark.asyncio +async def test_scope_isolation_by_credential_and_provider() -> None: + rule = _reasoning_rule(scope=("provider", "model", "credential")) + engine = FieldCacheEngine([rule]) + + await engine.extract( + "response", + {"choices": [{"message": {"reasoning_content": "cred-a"}}]}, + _context(provider="openai", credential_id="credential-a"), + ) + await engine.extract( + "response", + {"choices": [{"message": {"reasoning_content": "cred-b"}}]}, + _context(provider="openai", credential_id="credential-b"), + ) + + updated, _ = await engine.inject("request", {"messages": [{}]}, _context(provider="openai", credential_id="credential-a")) + + assert updated["messages"][-1]["reasoning_content"] == "cred-a" + assert build_cache_key(rule, _context(provider="openai", credential_id="credential-a")) != build_cache_key( + rule, _context(provider="other", credential_id="credential-a") + ) + + +def test_scope_key_hashes_adversarial_provider_and_model_delimiters() -> None: + rule = _reasoning_rule() + keys = { + build_cache_key(rule, _context(model=value)) + for value in ("a:b", "a/b", "a_b", "a|b", "a=b") + } + + assert None not in keys + assert len(keys) == 5 + assert all("a:b" not in key and "a/b" not in key and "a|b" not in key for key in keys) + assert all( + len(component.split("=", 1)[1]) == 64 + for key in keys + for component in key.split("|") + ) + + +@pytest.mark.parametrize( + "path", + ( + "previous_response_id", + "previousResponseId", + "previous-response-id", + "metadata.conversationId", + "request[continuation_id]", + ), +) +def test_continuation_aliases_require_semantic_registration(path: str) -> None: + assert is_provider_continuation_path(path) is True + with pytest.raises(ValueError, match="provider_continuation"): + FieldCacheRule( + name="continuation", + source="response", + path="id", + inject=FieldCacheInjection(target="request", path=path), + ) + + +@pytest.mark.asyncio +async def test_missing_session_scope_skips_by_default() -> None: + engine = FieldCacheEngine([_reasoning_rule()]) + + operations = await engine.extract("response", {"choices": [{"message": {"reasoning_content": "x"}}]}, _context(session_id=None)) + + assert operations[0].skipped is True + assert operations[0].reason == "missing_required_scope" + + +@pytest.mark.asyncio +async def test_missing_credential_scope_skips_instead_of_sharing_none_bucket() -> None: + rule = _reasoning_rule(scope=("provider", "model", "credential")) + engine = FieldCacheEngine([rule]) + + operations = await engine.extract( + "response", + {"choices": [{"message": {"reasoning_content": "x"}}]}, + _context(credential_id=None), + ) + updated, injection_operations = await engine.inject("request", {"messages": [{}]}, _context(credential_id=None)) + + assert operations[0].skipped is True + assert operations[0].reason == "missing_required_scope" + assert injection_operations[0].skipped is True + assert injection_operations[0].reason == "missing_required_scope" + assert updated == {"messages": [{}]} + assert build_cache_key(rule, _context(credential_id=None)) is None + + +@pytest.mark.asyncio +async def test_missing_path_is_noop() -> None: + engine = FieldCacheEngine([_reasoning_rule()]) + + operations = await engine.extract("response", {"choices": []}, _context()) + + assert operations[0].matched == 0 + assert operations[0].changed is False + + +@pytest.mark.asyncio +async def test_stream_event_extraction() -> None: + rule = FieldCacheRule( + name="provider_session_id", + source="stream_event", + path="metadata.provider_session_id", + scope=("provider", "model", "session"), + inject=FieldCacheInjection(target="request", path="metadata.provider_session_id"), + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("stream_event", {"metadata": {"provider_session_id": "sid_1"}}, _context()) + updated, _ = await engine.inject("request", {"metadata": {}}, _context()) + + assert updated["metadata"]["provider_session_id"] == "sid_1" + + +@pytest.mark.asyncio +async def test_trace_sample_values_are_truncated() -> None: + rule = _reasoning_rule() + engine = FieldCacheEngine([rule]) + long_value = "x" * 700 + + operations = await engine.extract( + "response", + {"choices": [{"message": {"reasoning_content": long_value}}]}, + _context(), + ) + + assert operations[0].sample_values[0].endswith("...") + + +@pytest.mark.asyncio +async def test_field_cache_trace_omits_raw_sample_values(tmp_path) -> None: + logger = TransactionLogger("openai", "gpt-test", parent_dir=tmp_path) + rule = _reasoning_rule() + engine = FieldCacheEngine([rule]) + + await engine.extract( + "response", + {"choices": [{"message": {"reasoning_content": "provider-signature-secret"}}]}, + _context(), + transaction_logger=logger, + ) + + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "provider-signature-secret" not in trace_text + entries = _trace_entries(logger.log_dir) + after_entry = next(entry for entry in entries if entry["pass_name"] == "after_field_cache_extraction") + assert after_entry["metadata"]["sample_value_count"] == 1 + assert after_entry["metadata"]["sample_value_types"] == ["str"] + + +@pytest.mark.asyncio +async def test_field_cache_error_trace_omits_raw_payload_values(tmp_path) -> None: + class FailingStore(InMemoryFieldCacheStore): + async def set(self, key, value, *, ttl_seconds=None): + raise RuntimeError("store failed") + + logger = TransactionLogger("openai", "gpt-test", parent_dir=tmp_path) + engine = FieldCacheEngine([_reasoning_rule()], store=FailingStore()) + + with pytest.raises(RuntimeError): + await engine.extract( + "response", + {"choices": [{"message": {"reasoning_content": "provider-signature-secret"}}]}, + _context(), + transaction_logger=logger, + ) + + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "provider-signature-secret" not in trace_text + assert "payload_type" in trace_text + + +@pytest.mark.asyncio +async def test_field_cache_traces_start_and_complete_even_without_matching_rules(tmp_path) -> None: + logger = TransactionLogger("openai", "gpt-test", parent_dir=tmp_path) + engine = FieldCacheEngine([]) + + operations = await engine.extract("response", {"choices": []}, _context(), transaction_logger=logger) + updated, injection_operations = await engine.inject("request", {"messages": []}, _context(), transaction_logger=logger) + + assert operations == [] + assert injection_operations == [] + assert updated == {"messages": []} + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert pass_names == [ + "field_cache_extraction_start", + "field_cache_extraction_complete", + "field_cache_injection_start", + "field_cache_injection_complete", + ] + assert entries[1]["metadata"]["rule_count"] == 0 + assert entries[-1]["metadata"]["operation_count"] == 0 + + +def test_per_tool_call_requires_tool_call_id_path() -> None: + with pytest.raises(ValueError): + FieldCacheEngine([ + FieldCacheRule(name="tool_state", source="response", path="tool_calls.*", mode="per_tool_call") + ]) + + +@pytest.mark.asyncio +async def test_provider_cache_field_store_wraps_json_string_cache() -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.values = {} + + async def retrieve_async(self, key: str): + return self.values.get(key) + + async def store_async(self, key: str, value: str) -> None: + json.loads(value) + self.values[key] = value + + async def clear(self) -> None: + self.values.clear() + + store = ProviderCacheFieldStore(FakeProviderCache()) + + await store.set("key", {"value": 1}) + assert await store.get("key") == {"value": 1} + await store.append("key", [{"value": 2}]) + assert await store.get("key") == [{"value": 2}] + + +@pytest.mark.asyncio +async def test_in_memory_store_expires_ttl_values() -> None: + now = 100.0 + store = InMemoryFieldCacheStore(clock=lambda: now) + + await store.set("key", "value", ttl_seconds=5) + assert await store.get("key") == "value" + now = 106.0 + assert await store.get("key") is None + + +@pytest.mark.asyncio +async def test_provider_cache_field_store_expires_ttl_values(monkeypatch) -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.values = {} + + async def retrieve_async(self, key: str): + return self.values.get(key) + + async def store_async(self, key: str, value: str) -> None: + json.loads(value) + self.values[key] = value + + async def clear(self) -> None: + self.values.clear() + + now = 100.0 + monkeypatch.setattr("rotator_library.field_cache.store.time.time", lambda: now) + store = ProviderCacheFieldStore(FakeProviderCache()) + + await store.set("key", "value", ttl_seconds=5) + assert await store.get("key") == "value" + now = 106.0 + assert await store.get("key") is None + + +@pytest.mark.asyncio +async def test_in_memory_store_returns_deep_copies() -> None: + store = InMemoryFieldCacheStore() + await store.set("key", {"nested": []}) + value = await store.get("key") + value["nested"].append("mutated") + + assert await store.get("key") == {"nested": []} + + +@pytest.mark.asyncio +async def test_in_memory_store_prunes_expired_and_lru_entries() -> None: + now = [0.0] + store = InMemoryFieldCacheStore(clock=lambda: now[0], max_entries=2) + await store.set("a", 1, ttl_seconds=10) + now[0] = 1.0 + await store.set("b", 2, ttl_seconds=10) + now[0] = 2.0 + assert await store.get("a") == 1 + now[0] = 3.0 + await store.set("c", 3, ttl_seconds=10) + + assert await store.get("b") is None + assert await store.get("a") == 1 + now[0] = 20.0 + await store.set("d", 4, ttl_seconds=10) + assert await store.get("a") is None + assert await store.get("c") is None + assert await store.get("d") == 4 + + +@pytest.mark.asyncio +async def test_all_mode_enforces_value_count_and_byte_bounds() -> None: + rule = FieldCacheRule( + name="bounded", + source="response", + path="value", + mode="all", + max_values=2, + max_bytes=32, + inject=FieldCacheInjection(target="request", path="metadata.values", as_list=True), + ) + engine = FieldCacheEngine([rule]) + for value in ("one", "two", "three"): + await engine.extract("response", {"value": value}, _context()) + updated, _ = await engine.inject("request", {"metadata": {}}, _context()) + + assert updated["metadata"]["values"] == ["two", "three"] + with pytest.raises(ValueError, match="exceeds max_bytes"): + await engine.extract("response", {"value": "x" * 64}, _context()) + + +@pytest.mark.asyncio +async def test_last_mode_rejects_oversized_opaque_state() -> None: + rule = FieldCacheRule( + name="provider_response_id", + source="response", + path="id", + max_bytes=16, + inject=FieldCacheInjection( + target="request", + path="previous_response_id", + when_missing_only=True, + ), + metadata={"provider_continuation": True}, + ) + + with pytest.raises(ValueError, match="exceeds max_bytes"): + await FieldCacheEngine([rule]).extract( + "response", + {"id": "resp_" + "x" * 64}, + _context(), + ) + + +@pytest.mark.asyncio +async def test_per_tool_call_mode_bounds_correlated_entries() -> None: + rule = FieldCacheRule( + name="tool_state", + source="response", + path="tool_calls.*", + mode="per_tool_call", + max_values=2, + max_bytes=256, + inject=FieldCacheInjection(target="request", path="metadata.state"), + metadata={"tool_call_id_path": "id"}, + ) + engine = FieldCacheEngine([rule]) + await engine.extract( + "response", + {"tool_calls": [ + {"id": "a", "state": "one"}, + {"id": "b", "state": "two"}, + {"id": "c", "state": "three"}, + ]}, + _context(), + ) + missing, missing_ops = await engine.inject( + "request", + {"metadata": {}}, + _context(metadata={"tool_call_id": "a"}), + ) + retained, retained_ops = await engine.inject( + "request", + {"metadata": {}}, + _context(metadata={"tool_call_id": "c"}), + ) + + assert missing == {"metadata": {}} + assert missing_ops[0].skipped is True + assert missing_ops[0].reason == "tool_call_cache_miss" + assert retained_ops[0].hit is True + assert retained["metadata"]["state"] == {"id": "c", "state": "three"} + + +@pytest.mark.asyncio +async def test_provider_cache_append_is_atomic() -> None: + class YieldingProviderCache: + def __init__(self) -> None: + self.values = {} + + async def retrieve_async(self, key): + await asyncio.sleep(0) + return self.values.get(key) + + async def store_async(self, key, value): + await asyncio.sleep(0) + self.values[key] = value + + async def clear(self): + self.values.clear() + + backend = YieldingProviderCache() + first_store = ProviderCacheFieldStore(backend) + second_store = ProviderCacheFieldStore(backend) + await asyncio.gather( + first_store.append("key", ["a"], max_values=8, max_bytes=128), + second_store.append("key", ["b"], max_values=8, max_bytes=128), + ) + + assert await first_store.get("key") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_legacy_append_fallback_remains_bounded() -> None: + class LegacyStore: + def __init__(self): + self.values = {} + + async def get(self, key): + return self.values.get(key) + + async def set(self, key, value): + self.values[key] = value + + async def append(self, key, values): + self.values.setdefault(key, []).extend(values) + + rule = FieldCacheRule( + name="legacy_all", + source="response", + path="value", + mode="all", + max_values=2, + max_bytes=64, + inject=FieldCacheInjection(target="request", path="metadata.values", as_list=True), + ) + engine = FieldCacheEngine([rule], store=LegacyStore()) + for value in ("one", "two", "three"): + await engine.extract("response", {"value": value}, _context()) + updated, _ = await engine.inject("request", {"metadata": {}}, _context()) + + assert updated["metadata"]["values"] == ["two", "three"] + + +@pytest.mark.asyncio +async def test_store_internal_type_error_is_not_treated_as_legacy_signature() -> None: + class BrokenStore: + async def get(self, key): + return None + + async def set(self, key, value, *, ttl_seconds=None): + raise TypeError("internal serialization failure") + + async def append(self, key, values, **kwargs): + raise TypeError("internal serialization failure") + + with pytest.raises(TypeError, match="internal serialization"): + await FieldCacheEngine([_reasoning_rule()], store=BrokenStore()).extract( + "response", + {"choices": [{"message": {"reasoning_content": "x"}}]}, + _context(), + ) + + +@pytest.mark.asyncio +async def test_last_user_turn_uses_latest_user_message() -> None: + rule = FieldCacheRule( + name="user_signature", + source="request", + path="messages.*.metadata.signature", + mode="last_user_turn", + inject=FieldCacheInjection(target="request", path="metadata.signature"), + allow_missing_session=True, + ) + engine = FieldCacheEngine([rule]) + + operations = await engine.extract( + "request", + { + "messages": [ + {"role": "user", "metadata": {"signature": "first-user"}}, + {"role": "assistant", "metadata": {"signature": "assistant"}}, + {"role": "user", "metadata": {"signature": "last-user"}}, + ] + }, + _context(session_id=None), + ) + updated, _ = await engine.inject("request", {"metadata": {}}, _context(session_id=None)) + + assert operations[0].changed is True + assert updated["metadata"]["signature"] == "last-user" + + +@pytest.mark.asyncio +async def test_last_mode_preserves_list_valued_field() -> None: + rule = FieldCacheRule( + name="list_value", + source="response", + path="metadata.signatures", + inject=FieldCacheInjection(target="request", path="metadata.signatures"), + allow_missing_session=True, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"metadata": {"signatures": ["a", "b"]}}, _context(session_id=None)) + updated, _ = await engine.inject("request", {"metadata": {}}, _context(session_id=None)) + + assert updated["metadata"]["signatures"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_as_list_unwraps_last_mode_value_envelope() -> None: + rule = FieldCacheRule( + name="value", + source="response", + path="value", + inject=FieldCacheInjection(target="request", path="metadata.values", as_list=True), + allow_missing_session=True, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"value": "sig"}, _context(session_id=None)) + updated, _ = await engine.inject("request", {"metadata": {}}, _context(session_id=None)) + + assert updated["metadata"]["values"] == ["sig"] + + +@pytest.mark.asyncio +async def test_last_assistant_turn_skips_without_turn_context() -> None: + rule = FieldCacheRule( + name="assistant_signature", + source="response", + path="choices.*.message.signature", + mode="last_assistant_turn", + inject=FieldCacheInjection(target="request", path="metadata.signature"), + allow_missing_session=True, + ) + engine = FieldCacheEngine([rule]) + + operations = await engine.extract("response", {"choices": [{"message": {"signature": "sig"}}]}, _context(session_id=None)) + + assert operations[0].skipped is True + assert operations[0].reason == "turn_context_not_found" + + +@pytest.mark.asyncio +async def test_turn_mode_uses_metadata_configured_relative_paths() -> None: + rule = FieldCacheRule( + name="assistant_signature", + source="response", + path="unused.global.path", + mode="last_assistant_turn", + inject=FieldCacheInjection(target="request", path="metadata.signature"), + allow_missing_session=True, + metadata={"turn_container_path": "messages", "turn_role_path": "kind", "turn_value_path": "parts.*.signature"}, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract( + "response", + {"messages": [{"kind": "assistant", "parts": [{"signature": "first"}]}, {"kind": "assistant", "parts": [{"signature": "second"}]}]}, + _context(session_id=None), + ) + updated, _ = await engine.inject("request", {"metadata": {}}, _context(session_id=None)) + + assert updated["metadata"]["signature"] == "second" + + +@pytest.mark.asyncio +async def test_per_tool_call_correlates_sibling_id_and_value_for_injection() -> None: + rule = FieldCacheRule( + name="tool_signature", + source="response", + path="tool_calls.*.signature", + mode="per_tool_call", + inject=FieldCacheInjection(target="request", path="metadata.signature"), + allow_missing_session=True, + metadata={ + "tool_container_path": "tool_calls", + "tool_call_id_path": "id", + "tool_value_path": "signature", + }, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"tool_calls": [{"id": "call_a", "signature": "sig-a"}, {"id": "call_b", "signature": "sig-b"}]}, _context(session_id=None)) + updated, operations = await engine.inject("request", {"metadata": {}}, _context(session_id=None, metadata={"tool_call_id": "call_b"})) + + assert operations[0].hit is True + assert updated["metadata"]["signature"] == "sig-b" + + +@pytest.mark.asyncio +async def test_per_tool_call_as_list_injects_matching_values() -> None: + rule = FieldCacheRule( + name="tool_signature", + source="response", + path="tool_calls.*", + mode="per_tool_call", + inject=FieldCacheInjection(target="request", path="metadata.signatures", as_list=True), + allow_missing_session=True, + metadata={"tool_call_id_path": "id", "inject_tool_call_id_path": "tool_ids.*"}, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"tool_calls": [{"id": "a", "signature": "sig-a"}, {"id": "b", "signature": "sig-b"}]}, _context(session_id=None)) + updated, _ = await engine.inject("request", {"metadata": {}, "tool_ids": ["a", "b"]}, _context(session_id=None)) + + assert updated["metadata"]["signatures"] == [{"id": "a", "signature": "sig-a"}, {"id": "b", "signature": "sig-b"}] + + +@pytest.mark.asyncio +async def test_per_tool_call_preserves_list_valued_match() -> None: + rule = FieldCacheRule( + name="tool_signatures", + source="response", + path="tool_calls.*.signatures", + mode="per_tool_call", + inject=FieldCacheInjection(target="request", path="metadata.signatures"), + allow_missing_session=True, + metadata={"tool_container_path": "tool_calls", "tool_call_id_path": "id", "tool_value_path": "signatures"}, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"tool_calls": [{"id": "call", "signatures": ["a", "b"]}]}, _context(session_id=None)) + updated, _ = await engine.inject("request", {"metadata": {}}, _context(session_id=None, metadata={"tool_call_id": "call"})) + + assert updated["metadata"]["signatures"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_engine_supports_legacy_store_without_ttl_keyword() -> None: + class LegacyStore: + def __init__(self) -> None: + self.values = {} + + async def get(self, key): + return self.values.get(key) + + async def set(self, key, value): + self.values[key] = value + + async def append(self, key, values): + self.values.setdefault(key, []).extend(values) + return self.values[key] + + async def clear(self): + self.values.clear() + + engine = FieldCacheEngine([_reasoning_rule()], store=LegacyStore()) + + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "legacy"}}]}, _context()) + updated, _ = await engine.inject("request", {"messages": [{}]}, _context()) + + assert updated["messages"][-1]["reasoning_content"] == "legacy" + + +@pytest.mark.asyncio +async def test_per_tool_call_skips_when_current_tool_id_is_ambiguous() -> None: + rule = FieldCacheRule( + name="tool_signature", + source="response", + path="tool_calls.*", + mode="per_tool_call", + inject=FieldCacheInjection(target="request", path="metadata.signature"), + allow_missing_session=True, + metadata={"tool_call_id_path": "id"}, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"tool_calls": [{"id": "call_a", "signature": "sig-a"}]}, _context(session_id=None)) + updated, operations = await engine.inject("request", {"metadata": {}}, _context(session_id=None)) + + assert operations[0].skipped is True + assert operations[0].reason == "tool_call_id_not_found" + assert updated == {"metadata": {}} + + +@pytest.mark.asyncio +async def test_insert_injection_adds_list_entry() -> None: + rule = FieldCacheRule( + name="prefix_message", + source="response", + path="message", + inject=FieldCacheInjection(target="request", path="messages.0", insert=True), + allow_missing_session=True, + ) + engine = FieldCacheEngine([rule]) + + await engine.extract("response", {"message": {"role": "system", "content": "cached"}}, _context(session_id=None)) + updated, _ = await engine.inject("request", {"messages": [{"role": "user", "content": "hi"}]}, _context(session_id=None)) + + assert updated["messages"] == [{"role": "system", "content": "cached"}, {"role": "user", "content": "hi"}] diff --git a/tests/test_field_cache_paths.py b/tests/test_field_cache_paths.py new file mode 100644 index 000000000..f34d99577 --- /dev/null +++ b/tests/test_field_cache_paths.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import pytest + +from rotator_library.field_cache.paths import FieldCachePathError, extract_path, inject_path, parse_path +from rotator_library.field_cache.types import FieldCacheInjection, FieldCacheRule + + +def test_parse_path_supports_keys_indexes_wildcards_and_tail_index() -> None: + tokens = parse_path("choices.*.message.parts[-1]") + + assert [token.kind for token in tokens] == ["key", "wildcard", "key", "key", "index"] + assert tokens[-1].value == -1 + + +def test_extract_path_handles_nested_dict_list_and_wildcard() -> None: + payload = { + "choices": [ + {"message": {"reasoning_content": "a"}}, + {"message": {"reasoning_content": "b"}}, + ] + } + + assert extract_path(payload, "choices.*.message.reasoning_content") == ["a", "b"] + assert extract_path(payload, "choices.1.message.reasoning_content") == ["b"] + + +def test_extract_path_missing_values_are_noop() -> None: + assert extract_path({"choices": []}, "choices.*.message.reasoning_content") == [] + assert extract_path({}, "missing.path") == [] + + +def test_extract_path_tail_index() -> None: + payload = {"messages": [{"content": "first"}, {"content": "last"}]} + + assert extract_path(payload, "messages[-1].content") == ["last"] + + +def test_inject_path_creates_dict_containers() -> None: + payload = {"messages": [{"role": "assistant"}]} + + changed = inject_path(payload, "messages[-1].reasoning_content", "hidden") + + assert changed is True + assert payload["messages"][-1]["reasoning_content"] == "hidden" + + +def test_inject_path_respects_when_missing_only() -> None: + payload = {"metadata": {"prompt_cache_key": "existing"}} + + changed = inject_path(payload, "metadata.prompt_cache_key", "new", when_missing_only=True) + + assert changed is False + assert payload["metadata"]["prompt_cache_key"] == "existing" + + +def test_inject_path_can_insert_at_final_list_index() -> None: + payload = {"messages": [{"role": "user"}]} + + changed = inject_path(payload, "messages.0", {"role": "system"}, insert=True) + + assert changed is True + assert payload["messages"] == [{"role": "system"}, {"role": "user"}] + + +def test_inject_path_can_insert_into_empty_list() -> None: + payload = {"messages": []} + + changed = inject_path(payload, "messages.0", {"role": "system"}, insert=True) + + assert changed is True + assert payload["messages"] == [{"role": "system"}] + + +def test_inject_path_rejects_wildcards_and_missing_lists() -> None: + with pytest.raises(FieldCachePathError): + inject_path({"choices": []}, "choices.*.message.reasoning_content", "x") + with pytest.raises(FieldCachePathError): + inject_path({}, "messages[-1].reasoning_content", "x") + + +def test_malformed_paths_and_rules_raise_useful_errors() -> None: + with pytest.raises(FieldCachePathError): + parse_path("choices..message") + with pytest.raises(FieldCachePathError): + parse_path("messages[abc]") + with pytest.raises(ValueError): + FieldCacheRule(name="bad/name", source="response", path="x") + with pytest.raises(ValueError, match="cache_key"): + FieldCacheRule(name="valid", cache_key="bad/key", source="response", path="x") + assert FieldCacheRule( + name="reasoning_content", + source="response", + path="choices.*.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="messages[-1].reasoning_content"), + ).scope == ("provider", "model", "credential", "session") + assert FieldCacheRule("positional", "response", "x", "all").mode == "all" diff --git a/tests/test_field_cache_trace.py b/tests/test_field_cache_trace.py new file mode 100644 index 000000000..154706c1e --- /dev/null +++ b/tests/test_field_cache_trace.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.adapters import AdapterContext, get_adapter, run_adapter_chain +from rotator_library.field_cache import FieldCacheContext, FieldCacheEngine, FieldCacheInjection, FieldCacheRule +from rotator_library.transaction_logger import TransactionLogger + + +def _trace_entries(log_dir): + return [json.loads(line) for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +@pytest.mark.asyncio +async def test_adapter_chain_emits_before_after_trace_entries(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + context = AdapterContext( + provider="openai", + model="gpt-test", + protocol="openai_chat", + credential_id="cred_1", + transport="http", + transaction_logger=logger, + adapter_config={"model_override": {"model": "native"}}, + ) + + result = await run_adapter_chain([get_adapter("model_override")], {"model": "public"}, context, stage="request") + + entries = _trace_entries(logger.log_dir) + assert result["model"] == "native" + assert [entry["pass_name"] for entry in entries] == ["before_adapter_chain", "after_adapter", "after_adapter_chain"] + assert entries[1]["metadata"]["adapter"] == "model_override" + assert entries[1]["metadata"]["changed"] is True + assert entries[1]["credential_id"] == "cred_1" + + +@pytest.mark.asyncio +async def test_field_cache_extract_and_inject_emit_before_after_trace_entries(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + rule = FieldCacheRule( + name="reasoning_content", + source="response", + path="choices.*.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="messages[-1].reasoning_content"), + ) + engine = FieldCacheEngine([rule]) + context = FieldCacheContext(provider="openai", model="gpt-test", credential_id="credential_1", session_id="session_1", classifier="global") + + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "hidden"}}]}, context, transaction_logger=logger) + updated, _ = await engine.inject("request", {"messages": [{"role": "user"}]}, context, transaction_logger=logger) + + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert updated["messages"][-1]["reasoning_content"] == "hidden" + assert pass_names == [ + "field_cache_extraction_start", + "before_field_cache_extraction", + "after_field_cache_extraction", + "field_cache_extraction_complete", + "field_cache_injection_start", + "before_field_cache_injection", + "after_field_cache_injection", + "field_cache_injection_complete", + ] + assert entries[2]["metadata"]["rule_name"] == "reasoning_content" + assert entries[2]["metadata"]["matched"] == 1 + assert entries[6]["metadata"]["hit"] is True + assert entries[6]["metadata"]["changed"] is True + + +@pytest.mark.asyncio +async def test_stream_sourced_rule_injection_trace_uses_request_direction(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + rule = FieldCacheRule( + name="provider_session_id", + source="stream_event", + path="metadata.provider_session_id", + inject=FieldCacheInjection(target="request", path="metadata.provider_session_id"), + ) + engine = FieldCacheEngine([rule]) + context = FieldCacheContext(provider="openai", model="gpt-test", credential_id="credential_1", session_id="session_1", classifier="global") + + await engine.extract("stream_event", {"metadata": {"provider_session_id": "sid_1"}}, context, transaction_logger=logger) + await engine.inject("request", {"metadata": {}}, context, transaction_logger=logger) + + entries = _trace_entries(logger.log_dir) + injection_entries = [entry for entry in entries if "injection" in entry["pass_name"]] + assert {entry["direction"] for entry in injection_entries} == {"request"} + + +@pytest.mark.asyncio +async def test_field_cache_errors_emit_transform_log_error(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + rule = FieldCacheRule( + name="bad_injection", + source="response", + path="choices.*.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="messages.*.reasoning_content"), + ) + engine = FieldCacheEngine([rule]) + context = FieldCacheContext(provider="openai", model="gpt-test", credential_id="credential_1", session_id="session_1", classifier="global") + + await engine.extract("response", {"choices": [{"message": {"reasoning_content": "hidden"}}]}, context, transaction_logger=logger) + with pytest.raises(Exception): + await engine.inject("request", {"messages": [{"role": "user"}]}, context, transaction_logger=logger) + + entries = _trace_entries(logger.log_dir) + error_entry = next(entry for entry in entries if entry["pass_name"] == "transform_log_error") + assert error_entry["data"]["failed_pass_name"] == "field_cache_inject" + assert error_entry["metadata"]["rule_name"] == "bad_injection" diff --git a/tests/test_native_protocol_runtime_matrix.py b/tests/test_native_protocol_runtime_matrix.py new file mode 100644 index 000000000..a513e0703 --- /dev/null +++ b/tests/test_native_protocol_runtime_matrix.py @@ -0,0 +1,738 @@ +from __future__ import annotations + +import json +from copy import deepcopy +from types import MethodType, SimpleNamespace +from typing import Any + +import pytest + +from rotator_library.native_provider import NativeProviderContext, NativeProviderExecutor +from rotator_library.client.executor import RequestExecutor +from rotator_library.client.request_builder import RequestContextBuilder +from rotator_library.client.rotating_client import RotatingClient +from rotator_library.core.errors import StructuredAPIResponseError, is_structured_error_payload +from rotator_library.core.types import ErrorAction, RequestContext +from rotator_library.providers.codex_provider import CodexProvider +from rotator_library.providers.antigravity_provider import AntigravityProvider +from rotator_library.protocols import ProtocolContext, get_protocol +from rotator_library.protocols.types import first_text +from rotator_library.routing import parse_route_target + + +PROTOCOLS = ("openai_chat", "anthropic_messages", "responses", "gemini") + +REQUESTS: dict[str, dict[str, Any]] = { + "openai_chat": { + "model": "provider/model-test", + "messages": [ + {"role": "system", "content": "follow the rule"}, + {"role": "user", "content": "hello"}, + ], + }, + "anthropic_messages": { + "model": "provider/model-test", + "system": "follow the rule", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 256, + }, + "responses": { + "model": "provider/model-test", + "instructions": "follow the rule", + "input": "hello", + }, + "gemini": { + "model": "provider/model-test", + "systemInstruction": {"parts": [{"text": "follow the rule"}]}, + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + }, +} + + +def test_success_dictionaries_are_not_misclassified_as_api_errors() -> None: + assert is_structured_error_payload(RESPONSES["openai_chat"]) is False + assert is_structured_error_payload({"error": {"message": "failed"}}) is True + +RESPONSES: dict[str, dict[str, Any]] = { + "openai_chat": { + "id": "chat_1", + "model": "model-test", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + "anthropic_messages": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "model-test", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 2, "output_tokens": 1}, + }, + "responses": { + "id": "resp_1", + "object": "response", + "model": "model-test", + "status": "completed", + "output": [ + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "answer"}], + } + ], + "usage": {"input_tokens": 2, "output_tokens": 1, "total_tokens": 3}, + }, + "gemini": { + "responseId": "gemini_1", + "modelVersion": "model-test", + "candidates": [ + { + "index": 0, + "content": {"role": "model", "parts": [{"text": "answer"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, + }, +} + +OPERATIONS = { + "openai_chat": "chat", + "anthropic_messages": "messages", + "responses": "responses", + "gemini": "generate", +} + + +class RecordingTransport: + """Return one provider response while retaining the native request.""" + + def __init__(self, response: dict[str, Any]) -> None: + self.response = deepcopy(response) + self.payload: dict[str, Any] | None = None + + async def post_json(self, endpoint: str, *, headers: dict[str, str], payload: dict[str, Any]) -> dict[str, Any]: + self.payload = deepcopy(payload) + return deepcopy(self.response) + + +class RecordingHTTPClient: + """Expose the HTTPX-style post seam used by RequestExecutor.""" + + def __init__(self, response: dict[str, Any]) -> None: + self.response = response + self.calls: list[dict[str, Any]] = [] + + async def post(self, endpoint: str, *, headers: dict[str, str], json: dict[str, Any]): + self.calls.append({"endpoint": endpoint, "headers": headers, "json": deepcopy(json)}) + response = deepcopy(self.response) + + class Result: + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, Any]: + return response + + return Result() + + +def _assert_provider_shape(protocol_name: str, payload: dict[str, Any]) -> None: + """Assert that a provider hook receives only its declared wire format.""" + + if protocol_name == "openai_chat": + assert "messages" in payload and "contents" not in payload and "input" not in payload + elif protocol_name == "anthropic_messages": + assert "messages" in payload and "system" in payload and "contents" not in payload + payload.setdefault("max_tokens", 256) + elif protocol_name == "responses": + assert "input" in payload and "messages" not in payload and "contents" not in payload + else: + assert "contents" in payload and "messages" not in payload and "input" not in payload + + +def _output_text(protocol_name: str, payload: dict[str, Any]) -> str | None: + """Parse a client payload again to prove output-language validity.""" + + response = get_protocol(protocol_name).parse_response( + payload, + ProtocolContext(source_protocol=protocol_name, target_protocol=protocol_name), + ) + for message in response.messages: + text = first_text(message.content) + if text: + return text + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_protocol", PROTOCOLS) +@pytest.mark.parametrize("provider_protocol", PROTOCOLS) +@pytest.mark.parametrize("output_protocol", PROTOCOLS) +async def test_runtime_keeps_input_provider_and_output_protocols_independent( + input_protocol: str, + provider_protocol: str, + output_protocol: str, +) -> None: + preparer_calls: list[dict[str, Any]] = [] + + def prepare(payload: dict[str, Any], *, model: str, operation: str) -> dict[str, Any]: + _assert_provider_shape(provider_protocol, payload) + preparer_calls.append(deepcopy(payload)) + return payload + + context = NativeProviderContext( + provider="provider", + model="model-test", + protocol_name=provider_protocol, + endpoint="https://provider.test/generate", + operation=OPERATIONS[provider_protocol], + input_protocol_name=input_protocol, + output_protocol_name=output_protocol, + credential_id="credential-1", + session_id="session-1", + metadata={"public_model": "provider/model-test", "input_provider": "provider"}, + request_preparer=prepare, + ) + transport = RecordingTransport(RESPONSES[provider_protocol]) + + result = await NativeProviderExecutor().execute( + deepcopy(REQUESTS[input_protocol]), + context, + transport, + ) + + assert len(preparer_calls) == 1 + prepared_payload = deepcopy(preparer_calls[0]) + assert prepared_payload.pop("_proxy_model") == "provider/model-test" + assert "_proxy_model" not in transport.payload + assert transport.payload == prepared_payload + assert _output_text(output_protocol, result) == "answer" + + +@pytest.mark.parametrize("output_protocol", PROTOCOLS) +def test_litellm_and_custom_chat_results_use_the_selected_output_protocol(output_protocol: str) -> None: + context = RequestContext( + model="provider/model-test", + provider="provider", + kwargs=deepcopy(REQUESTS["openai_chat"]), + streaming=False, + credentials=["credential-1"], + deadline=9999999999.0, + input_protocol_name="gemini", + output_protocol_name=output_protocol, + ) + + result = RequestExecutor._format_execution_response( + deepcopy(RESPONSES["openai_chat"]), + "openai_chat", + context, + ) + + assert _output_text(output_protocol, result) == "answer" + + +@pytest.mark.asyncio +async def test_structured_errors_raise_before_cross_protocol_success_formatting() -> None: + class ErrorPlugin: + def has_custom_logic(self) -> bool: + return True + + async def acompletion(self, client, **kwargs): + return {"error": {"type": "rate_limit_error", "message": "slow down", "status_code": 429}} + + target = parse_route_target("provider/model-test@custom") + context = RequestContext( + model="provider/model-test", + provider="provider", + kwargs=deepcopy(REQUESTS["openai_chat"]), + streaming=False, + credentials=["credential-1"], + deadline=9999999999.0, + output_protocol_name="anthropic_messages", + routing_targets=(target,), + ) + context.routing_target_index = 0 + executor = RequestExecutor.__new__(RequestExecutor) + executor._http_client = object() + + with pytest.raises(StructuredAPIResponseError) as raised: + await executor._execute_provider_request( + "provider", + context.model, + ErrorPlugin(), + "secret", + "credential-1", + deepcopy(context.kwargs), + context, + ) + + assert raised.value.error_type == "rate_limit" + + +@pytest.mark.asyncio +async def test_structured_error_survives_real_credential_exhaustion_loop() -> None: + class CredentialContext: + def __init__(self, credential: str) -> None: + self.credential = credential + self.stable_id = f"stable-{credential}" + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + class UsageManager: + states = {} + + async def get_availability_stats(self, model, quota_group=None): + return {"available": 2, "total": 2, "blocked": 0, "blocked_by": {}, "rotation_mode": "sequential"} + + async def acquire_credential(self, *, candidates, **kwargs): + return CredentialContext(candidates[0]) + + class ErrorPlugin: + def has_custom_logic(self): + return True + + async def acompletion(self, client, **kwargs): + return {"error": {"status": 429, "message": "busy"}} + + executor = RequestExecutor.__new__(RequestExecutor) + executor._max_retries = 1 + executor._http_client = object() + executor._cooldown = None + manager = UsageManager() + + async def prepare_execution(self, context): + filter_result = SimpleNamespace(priorities={}, tier_names={}) + return manager, filter_result, list(context.credentials), None, {} + + async def prepare_kwargs(self, provider, model, credential, context, **kwargs): + return deepcopy(context.kwargs) + + async def handle_error( + self, + error, + credential_context, + model, + provider, + attempt, + accumulator, + retry_state, + request_headers, + context, + ): + accumulator.normal_errors.append({"error_type": error.error_type}) + return ErrorAction.ROTATE + + executor._prepare_execution = MethodType(prepare_execution, executor) + executor._prepare_request_kwargs = MethodType(prepare_kwargs, executor) + executor._handle_error_with_context = MethodType(handle_error, executor) + executor._get_plugin_instance = MethodType(lambda self, provider: ErrorPlugin(), executor) + context = RequestContext( + model="provider/model-test", + provider="provider", + kwargs=deepcopy(REQUESTS["openai_chat"]), + streaming=False, + credentials=["credential-1", "credential-2"], + deadline=9999999999.0, + ) + + with pytest.raises(StructuredAPIResponseError) as raised: + await executor._execute_non_streaming(context) + + assert raised.value.error_type == "rate_limit" + + +@pytest.mark.asyncio +async def test_native_structured_errors_raise_before_provider_response_parsing() -> None: + context = NativeProviderContext( + provider="provider", + model="model-test", + protocol_name="gemini", + endpoint="https://provider.test/generate", + operation="generate", + input_protocol_name="openai_chat", + output_protocol_name="responses", + ) + transport = RecordingTransport( + {"error": {"code": 429, "status": "RESOURCE_EXHAUSTED", "message": "quota exhausted"}} + ) + + with pytest.raises(StructuredAPIResponseError) as raised: + await NativeProviderExecutor().execute( + deepcopy(REQUESTS["openai_chat"]), + context, + transport, + ) + + assert raised.value.error_type == "quota_exceeded" + + +@pytest.mark.parametrize( + ("error_payload", "expected_type", "expected_status"), + [ + ({"error": "upstream failed"}, "invalid_request", 400), + ({"error": {"status": 429, "message": "busy"}}, "rate_limit", 429), + ({"error": {"code": 403, "status": "PERMISSION_DENIED"}}, "forbidden", 403), + ], +) +@pytest.mark.asyncio +async def test_structured_error_variants_are_normalized_before_success_parsing( + error_payload: dict[str, Any], + expected_type: str, + expected_status: int, +) -> None: + context = NativeProviderContext( + provider="provider", + model="model-test", + protocol_name="gemini", + endpoint="https://provider.test/generate", + operation="generate", + ) + + with pytest.raises(StructuredAPIResponseError) as raised: + await NativeProviderExecutor().execute({}, context, RecordingTransport(error_payload)) + + assert raised.value.error_type == expected_type + assert raised.value.http_status == expected_status + + +@pytest.mark.asyncio +async def test_response_adapters_run_on_selected_client_protocol_payload() -> None: + context = NativeProviderContext( + provider="provider", + model="model-test", + protocol_name="anthropic_messages", + endpoint="https://provider.test/messages", + operation="messages", + input_protocol_name="openai_chat", + output_protocol_name="openai_chat", + adapter_names=("reasoning_content",), + adapter_config={ + "reasoning_content": { + "source_fields": ["reasoning_content"], + "output_field": "analysis", + } + }, + ) + response = deepcopy(RESPONSES["anthropic_messages"]) + response["content"].insert(0, {"type": "thinking", "thinking": "provider reasoning"}) + + result = await NativeProviderExecutor().execute( + deepcopy(REQUESTS["openai_chat"]), + context, + RecordingTransport(response), + ) + + assert result["choices"][0]["message"]["analysis"] == "provider reasoning" + + +@pytest.mark.asyncio +async def test_opaque_thought_signatures_are_cached_but_never_returned_to_clients() -> None: + provider = AntigravityProvider() + rules = provider.get_field_cache_rules("antigravity/gemini-3-pro") + adapter_names = provider.get_adapter_names("antigravity/gemini-3-pro") + adapter_config = provider.get_adapter_config("antigravity/gemini-3-pro") + context = NativeProviderContext( + provider="antigravity", + model="gemini-3-pro", + protocol_name="gemini", + endpoint="https://provider.test/generate", + operation="generate", + input_protocol_name="gemini", + output_protocol_name="gemini", + credential_id="credential-1", + session_id="session-1", + field_cache_rules=rules, + adapter_names=adapter_names, + adapter_config=adapter_config, + metadata={"public_model": "antigravity/gemini-3-pro", "input_provider": "antigravity"}, + ) + signed_response = deepcopy(RESPONSES["gemini"]) + signed_response["candidates"][0]["content"]["parts"].insert( + 0, + {"text": "hidden reasoning", "thought": True, "thoughtSignature": "provider-secret-signature"}, + ) + executor = NativeProviderExecutor() + + first_result = await executor.execute( + deepcopy(REQUESTS["gemini"]), + context, + RecordingTransport(signed_response), + ) + + assert "provider-secret-signature" not in json.dumps(first_result) + + second_transport = RecordingTransport(deepcopy(RESPONSES["gemini"])) + await executor.execute(deepcopy(REQUESTS["gemini"]), context, second_transport) + + assert second_transport.payload["request"]["metadata"]["thoughtSignatures"] == [ + "provider-secret-signature" + ] + + +@pytest.mark.asyncio +async def test_real_callback_nested_edits_preserve_source_native_metadata() -> None: + source_payload = deepcopy(REQUESTS["anthropic_messages"]) + source_payload["system"] = [ + { + "type": "text", + "text": "follow the rule", + "cache_control": {"type": "ephemeral"}, + } + ] + source_payload["tools"] = [ + { + "name": "lookup", + "description": "original description", + "input_schema": {"type": "object", "properties": {}}, + } + ] + anthropic = get_protocol("anthropic_messages") + chat = get_protocol("openai_chat") + unified = anthropic.parse_request(source_payload) + chat_view = chat.build_request( + unified, + ProtocolContext(source_protocol="anthropic_messages", target_protocol="openai_chat"), + ) + context = RequestContext( + model="provider/model-test", + provider="provider", + kwargs=deepcopy(chat_view), + streaming=False, + credentials=["credential-1"], + deadline=9999999999.0, + input_protocol_name="anthropic_messages", + output_protocol_name="anthropic_messages", + protocol_request=source_payload, + unified_request=unified, + input_provider="provider", + ) + async def callback(request, kwargs): + kwargs["messages"][-1]["content"] = "changed by callback" + kwargs["tools"][0]["function"]["description"] = "changed tool description" + + context.pre_request_callback = callback + executor = RequestExecutor.__new__(RequestExecutor) + attempted = await executor._prepare_request_kwargs( + "provider", + context.model, + "credential-1", + context, + native_execution=True, + ) + await executor._run_pre_request_callback(context, attempted) + assert context.kwargs["messages"][-1]["content"] == "hello" + assert context.kwargs["tools"][0]["function"]["description"] == "original description" + + merged = RequestExecutor._canonical_request_for_native(context, attempted) + provider_payload = anthropic.build_request( + merged, + ProtocolContext( + source_protocol="anthropic_messages", + target_protocol="anthropic_messages", + source_provider="provider", + target_provider="provider", + provider_state_compatible=True, + ), + ) + + assert provider_payload["system"][0]["cache_control"] == {"type": "ephemeral"} + assert provider_payload["messages"][-1]["content"][0]["text"] == "changed by callback" + assert provider_payload["tools"][0]["description"] == "changed tool description" + + +def test_internal_attempt_fields_do_not_become_same_protocol_extensions() -> None: + source_payload = deepcopy(REQUESTS["openai_chat"]) + source_payload["vendor_setting"] = {"enabled": True} + chat = get_protocol("openai_chat") + unified = chat.parse_request(source_payload) + context = RequestContext( + model="provider/model-test", + provider="provider", + kwargs=deepcopy(source_payload), + streaming=False, + credentials=["credential-1"], + deadline=9999999999.0, + input_protocol_name="openai_chat", + output_protocol_name="openai_chat", + protocol_request=source_payload, + unified_request=unified, + input_provider="provider", + ) + attempted = deepcopy(source_payload) + attempted["transaction_context"] = {"request_id": "internal"} + attempted["litellm_params"] = {"api_base": "internal"} + + merged = RequestExecutor._canonical_request_for_native(context, attempted) + provider_payload = chat.build_request( + merged, + ProtocolContext( + source_protocol="openai_chat", + target_protocol="openai_chat", + source_provider="provider", + target_provider="provider", + ), + ) + + assert provider_payload["vendor_setting"] == {"enabled": True} + assert "transaction_context" not in provider_payload + assert "litellm_params" not in provider_payload + + +@pytest.mark.asyncio +async def test_proxy_expanded_responses_history_suppresses_cached_provider_continuation() -> None: + rules = CodexProvider().get_field_cache_rules("codex/gpt-5.1-codex") + executor = NativeProviderExecutor() + base_context = NativeProviderContext( + provider="codex", + model="gpt-5.1-codex", + protocol_name="responses", + endpoint="https://provider.test/responses", + operation="responses", + input_protocol_name="responses", + output_protocol_name="responses", + credential_id="credential-1", + session_id="session-1", + field_cache_rules=rules, + metadata={"public_model": "codex/gpt-5.1-codex", "input_provider": "codex"}, + ) + first_transport = RecordingTransport(deepcopy(RESPONSES["responses"])) + await executor.execute(deepcopy(REQUESTS["responses"]), base_context, first_transport) + + expanded_request = deepcopy(REQUESTS["responses"]) + expanded_request["input"] = [ + "earlier input", + RESPONSES["responses"]["output"][0], + "current input", + ] + second_context = NativeProviderContext( + **{ + **base_context.__dict__, + "metadata": {**base_context.metadata, "disable_provider_continuation": True}, + } + ) + second_transport = RecordingTransport({**deepcopy(RESPONSES["responses"]), "id": "resp_2"}) + await executor.execute(expanded_request, second_context, second_transport) + + assert "previous_response_id" not in second_transport.payload + assert len(second_transport.payload["input"]) == 3 + assert "earlier input" in str(second_transport.payload["input"]) + assert "current input" in str(second_transport.payload["input"]) + assert "answer" in str(second_transport.payload["input"]) + + +@pytest.mark.asyncio +async def test_agenerate_runs_real_builder_to_native_executor_handoff() -> None: + validated_provider_requests: list[dict[str, Any]] = [] + + class Resolver: + def resolve_model_id(self, model: str, provider: str) -> str: + return model + + class Session: + session_id = "session-1" + affinity_key = "affinity-1" + tracking_namespace = "namespace" + + class SessionTracker: + def infer_session(self, *args, **kwargs): + return Session() + + async def scope(provider, classifier, request_api_keys, request_providers, private): + return { + "credentials": ["credential-1"], + "usage_manager_key": provider, + "provider_config": {}, + "credential_secrets": {"credential-1": "secret"}, + "classifier": classifier or "global", + } + + class GeminiProvider: + def has_custom_logic(self) -> bool: + return False + + def get_protocol_name(self, model="") -> str: + return "gemini" + + def get_native_operation(self, model="", request=None, stream=False) -> str: + return "generate" + + def get_native_endpoint(self, model="", operation="generate") -> str: + return "https://provider.test/generate" + + def get_native_headers(self, credential_identifier, model="", operation="generate"): + return {"Authorization": f"Bearer {credential_identifier}"} + + def normalize_native_model(self, model="") -> str: + return model.split("/", 1)[-1] + + def get_adapter_names(self, model=""): + return () + + def get_adapter_config(self, model=""): + return {} + + def get_field_cache_rules(self, model=""): + return () + + def validate_request(self, request, model=""): + validated_provider_requests.append(deepcopy(request)) + return True + + builder = RequestContextBuilder( + resolve_scope_for_provider=scope, + model_resolver=Resolver(), + session_tracker=SessionTracker(), + get_global_timeout=lambda: 30, + get_enable_request_logging=lambda: False, + ) + http_client = RecordingHTTPClient(RESPONSES["gemini"]) + real_executor = RequestExecutor.__new__(RequestExecutor) + real_executor._http_client = http_client + + class HandoffExecutor: + async def execute(self, context: RequestContext): + target = parse_route_target("provider/model-test@native") + context.routing_targets = (target,) + context.routing_target_index = 0 + attempt_kwargs = deepcopy(context.kwargs) + attempt_kwargs["messages"][-1]["content"] = "changed by callback" + attempt_kwargs["transaction_context"] = {"request_id": "internal"} + attempt_kwargs["litellm_params"] = {"api_base": "internal"} + return await real_executor._execute_provider_request( + "provider", + context.model, + GeminiProvider(), + "secret", + "credential-1", + attempt_kwargs, + context, + ) + + client = RotatingClient.__new__(RotatingClient) + client._request_builder = builder + client._executor = HandoffExecutor() + + result = await client.agenerate( + deepcopy(REQUESTS["anthropic_messages"]), + input_protocol="anthropic_messages", + output_protocol="responses", + ) + + assert "contents" in http_client.calls[0]["json"] + assert "messages" not in http_client.calls[0]["json"] + assert "transaction_context" not in http_client.calls[0]["json"] + assert "litellm_params" not in http_client.calls[0]["json"] + assert "contents" in validated_provider_requests[0] + assert "messages" not in validated_provider_requests[0] + assert http_client.calls[0]["json"]["systemInstruction"]["parts"][0]["text"] == "follow the rule" + assert http_client.calls[0]["json"]["contents"][0]["parts"][0]["text"] == "changed by callback" + assert result["object"] == "response" + assert _output_text("responses", result) == "answer" diff --git a/tests/test_native_provider_executor.py b/tests/test_native_provider_executor.py new file mode 100644 index 000000000..675a7beb5 --- /dev/null +++ b/tests/test_native_provider_executor.py @@ -0,0 +1,550 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.adapters import PayloadAdapter, register_adapter +from rotator_library.field_cache import FieldCacheInjection, FieldCacheRule +from rotator_library.native_provider import NativeHTTPTransport, NativeProviderContext, NativeProviderExecutor +from rotator_library.protocols import ProtocolError +from rotator_library.core.errors import StructuredAPIResponseError, classify_error +from rotator_library.transaction_logger import TransactionLogger + + +class FakeHTTPResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + +class FakeHTTPErrorResponse(FakeHTTPResponse): + status_code = 429 + text = "quota exhausted" + headers = {"Retry-After": "17"} + + +class FakeHTTPClient: + def __init__(self, response): + self.response = response + self.calls = [] + + async def post(self, endpoint, *, headers, json): + self.calls.append({"endpoint": endpoint, "headers": headers, "json": json}) + return FakeHTTPResponse(self.response) + + +@pytest.mark.asyncio +async def test_native_transport_preserves_non_success_status_and_body() -> None: + class ErrorClient: + async def post(self, endpoint, *, headers, json): + return FakeHTTPErrorResponse( + {"error": {"status": "RESOURCE_EXHAUSTED", "message": "quota exhausted"}} + ) + + with pytest.raises(StructuredAPIResponseError) as raised: + await NativeHTTPTransport(ErrorClient()).post_json( + "https://example.test/generate", + headers={"Authorization": "Bearer secret"}, + payload={"contents": []}, + ) + + assert raised.value.http_status == 429 + assert raised.value.error_type == "quota_exceeded" + assert raised.value.response["error"]["message"] == "quota exhausted" + assert classify_error(raised.value).retry_after == 17 + + +@pytest.mark.asyncio +async def test_native_transport_preserves_body_retry_delay_without_header() -> None: + class DelayResponse(FakeHTTPErrorResponse): + headers = {} + + class DelayClient: + async def post(self, endpoint, *, headers, json): + return DelayResponse( + { + "error": { + "status": "RESOURCE_EXHAUSTED", + "message": "quota exhausted", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.ErrorInfo", + "metadata": {"quotaResetDelay": "39s"}, + } + ], + } + } + ) + + with pytest.raises(StructuredAPIResponseError) as raised: + await NativeHTTPTransport(DelayClient()).post_json( + "https://example.test/generate", + headers={}, + payload={}, + ) + + assert classify_error(raised.value).retry_after == 39 + + +@pytest.mark.asyncio +async def test_native_transport_rejects_redirects_as_non_success() -> None: + class RedirectResponse(FakeHTTPResponse): + status_code = 302 + + class RedirectClient: + async def post(self, endpoint, *, headers, json): + return RedirectResponse({"location": "https://other.example"}) + + with pytest.raises(StructuredAPIResponseError) as raised: + await NativeHTTPTransport(RedirectClient()).post_json( + "https://example.test/generate", + headers={}, + payload={}, + ) + + assert raised.value.status_code == 302 + assert raised.value.error_type == "invalid_request" + + +def _trace_entries(log_dir): + return [json.loads(line) for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +@pytest.mark.asyncio +async def test_native_provider_executor_runs_protocol_adapter_cache_and_trace(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + rule = FieldCacheRule( + name="reasoning", + source="response", + path="choices.0.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="metadata.reasoning_content"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + headers={"authorization": "Bearer test"}, + adapter_names=("model_override",), + adapter_config={"model_override": {"model": "provider/gpt-test"}}, + field_cache_rules=(rule,), + transaction_logger=logger, + ) + response = { + "id": "chat_1", + "model": "provider/gpt-test", + "choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "hidden"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_cost": 0.01}, + } + client = FakeHTTPClient(response) + + result = await NativeProviderExecutor().execute( + {"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}]}, + context, + NativeHTTPTransport(client), + ) + + assert result["id"] == "chat_1" + assert client.calls[0]["json"]["model"] == "provider/gpt-test" + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "hidden" not in trace_text + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "native_protocol_selected" in pass_names + assert "raw_native_client_request" in pass_names + assert "parsed_native_unified_request" in pass_names + assert "built_native_provider_request" in pass_names + assert "after_request_adapter_chain" in pass_names + assert "field_cache_injection_start" in pass_names + assert "after_field_cache_injection" in pass_names + assert "field_cache_injection_complete" in pass_names + assert "native_provider_request" in pass_names + assert "raw_native_provider_response" in pass_names + assert "parsed_native_unified_response" in pass_names + assert "formatted_native_response" in pass_names + assert "after_response_adapter_chain" in pass_names + assert "field_cache_extraction_start" in pass_names + assert "after_field_cache_extraction" in pass_names + assert "field_cache_extraction_complete" in pass_names + usage_entries = [entry for entry in _trace_entries(logger.log_dir) if entry["pass_name"] == "usage_accounting_summary"] + assert usage_entries[-1]["data"]["cost"]["provider_reported_cost"] == 0.01 + assert "final_client_response" in pass_names + + +@pytest.mark.asyncio +async def test_native_provider_default_field_cache_persists_across_requests() -> None: + rule = FieldCacheRule( + name="reasoning", + source="response", + path="choices.0.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="metadata.reasoning_content"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=(rule,), + ) + first_client = FakeHTTPClient( + { + "id": "chat_1", + "model": "gpt-test", + "choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "cached-state"}}], + } + ) + second_client = FakeHTTPClient( + { + "id": "chat_2", + "model": "gpt-test", + "choices": [{"message": {"role": "assistant", "content": "ok"}}], + } + ) + executor = NativeProviderExecutor() + + await executor.execute({"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}]}, context, NativeHTTPTransport(first_client)) + await executor.execute({"model": "gpt-test", "messages": [{"role": "user", "content": "again"}]}, context, NativeHTTPTransport(second_client)) + + assert second_client.calls[0]["json"]["metadata"]["reasoning_content"] == "cached-state" + + +@pytest.mark.asyncio +async def test_native_provider_trace_redacts_configured_injection_paths(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + rule = FieldCacheRule( + name="state", + source="response", + path="choices.0.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="metadata.state"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=(rule,), + transaction_logger=logger, + ) + executor = NativeProviderExecutor() + + await executor.execute( + {"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}]}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "chat_1", "choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "opaque-state"}}]})), + ) + await executor.execute( + {"model": "gpt-test", "messages": [{"role": "user", "content": "again"}]}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "chat_2", "choices": [{"message": {"role": "assistant", "content": "ok"}}]})), + ) + + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "opaque-state" not in trace_text + assert '"state": "[REDACTED]"' in trace_text + + +@pytest.mark.asyncio +async def test_native_provider_executor_logs_transform_errors(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + adapter_names=("missing_adapter",), + transaction_logger=logger, + ) + + with pytest.raises(KeyError): + await NativeProviderExecutor().execute({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(FakeHTTPClient({}))) + + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "transform_log_error" in pass_names + + +@pytest.mark.asyncio +async def test_native_provider_executor_rejects_unsupported_operation_before_transport() -> None: + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + operation="embeddings", + endpoint="https://example.test/chat", + ) + client = FakeHTTPClient({"id": "should_not_call"}) + + with pytest.raises(ProtocolError, match="unsupported operation"): + await NativeProviderExecutor().execute( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(client), + ) + + assert client.calls == [] + + +@pytest.mark.asyncio +async def test_native_runtime_executes_unified_request_source_and_target() -> None: + rule = FieldCacheRule( + name="unified_request_state", + source="unified_request", + path="metadata.client_state", + inject=FieldCacheInjection(target="unified_request", path="metadata.cached_client_state"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=(rule,), + ) + executor = NativeProviderExecutor() + + await executor.execute( + {"model": "gpt-test", "messages": [], "metadata": {"client_state": "state-a"}}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "chat_1", "choices": []})), + ) + second_client = FakeHTTPClient({"id": "chat_2", "choices": []}) + await executor.execute( + {"model": "gpt-test", "messages": [], "metadata": {}}, + context, + NativeHTTPTransport(second_client), + ) + + assert second_client.calls[0]["json"]["metadata"]["cached_client_state"] == "state-a" + + +@pytest.mark.asyncio +async def test_native_runtime_executes_unified_response_source() -> None: + rule = FieldCacheRule( + name="response_object", + source="unified_response", + path="metadata.object", + inject=FieldCacheInjection(target="request", path="metadata.cached_object"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=(rule,), + ) + executor = NativeProviderExecutor() + + await executor.execute({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(FakeHTTPClient({"id": "chat_1", "object": "chat.completion", "choices": []}))) + second_client = FakeHTTPClient({"id": "chat_2", "choices": []}) + await executor.execute({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(second_client)) + + assert second_client.calls[0]["json"]["metadata"]["cached_object"] == "chat.completion" + + +@pytest.mark.asyncio +async def test_native_response_source_extracts_raw_provider_response_before_client_formatting() -> None: + rule = FieldCacheRule( + name="anthropic_signature", + source="response", + path="content.*.signature", + mode="all", + inject=FieldCacheInjection(target="request", path="metadata.signatures", as_list=True), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="claude_code", + model="claude-sonnet-4-5", + protocol_name="anthropic_messages", + input_protocol_name="anthropic_messages", + output_protocol_name="openai_chat", + endpoint="https://example.test/messages", + operation="messages", + field_cache_rules=(rule,), + ) + executor = NativeProviderExecutor() + + await executor.execute( + {"model": "claude-sonnet-4-5", "messages": [], "max_tokens": 1}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "ok", "signature": "sig-1"}]})), + ) + second_client = FakeHTTPClient({"id": "msg_2", "type": "message", "role": "assistant", "content": []}) + await executor.execute({"model": "claude-sonnet-4-5", "messages": [], "max_tokens": 1}, context, NativeHTTPTransport(second_client)) + + assert second_client.calls[0]["json"]["metadata"]["signatures"] == ["sig-1"] + + +@pytest.mark.asyncio +async def test_native_runtime_executes_metadata_target_before_adapters() -> None: + class MetadataEchoAdapter(PayloadAdapter): + name = "test_metadata_echo" + supported_stages = ("request",) + + async def transform_request(self, payload, context): + payload.setdefault("metadata", {})["adapter_seen_state"] = context.metadata.get("cached_state") + return payload + + register_adapter(MetadataEchoAdapter, replace=True) + rule = FieldCacheRule( + name="metadata_state", + source="response", + path="choices.0.message.reasoning_content", + inject=FieldCacheInjection(target="metadata", path="cached_state"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + adapter_names=("test_metadata_echo",), + field_cache_rules=(rule,), + ) + executor = NativeProviderExecutor() + + await executor.execute({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(FakeHTTPClient({"id": "chat_1", "choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "meta-state"}}]}))) + second_client = FakeHTTPClient({"id": "chat_2", "choices": []}) + await executor.execute({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(second_client)) + + assert second_client.calls[0]["json"]["metadata"]["adapter_seen_state"] == "meta-state" + + +@pytest.mark.asyncio +async def test_native_adapter_generic_traces_are_suppressed_for_field_cache_safety(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + rule = FieldCacheRule( + name="hidden_state", + source="response", + path="choices.0.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="metadata.hidden_state"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + adapter_names=("model_override",), + adapter_config={"model_override": {"model": "provider/gpt-test"}}, + field_cache_rules=(rule,), + transaction_logger=logger, + ) + + await NativeProviderExecutor().execute( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "chat_1", "choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "opaque-state"}}]})), + ) + + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert "before_adapter_chain" not in pass_names + assert "after_adapter" not in pass_names + assert "after_request_adapter_chain" in pass_names + assert "opaque-state" not in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + + +@pytest.mark.asyncio +async def test_native_runtime_executes_request_source() -> None: + rule = FieldCacheRule( + name="request_state", + source="request", + path="metadata.outgoing_state", + inject=FieldCacheInjection(target="unified_request", path="metadata.reused_request_state"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=(rule,), + ) + executor = NativeProviderExecutor() + + await executor.execute( + {"model": "gpt-test", "messages": [], "metadata": {"outgoing_state": "request-state"}}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "chat_1", "choices": []})), + ) + second_client = FakeHTTPClient({"id": "chat_2", "choices": []}) + await executor.execute({"model": "gpt-test", "messages": [], "metadata": {}}, context, NativeHTTPTransport(second_client)) + + assert second_client.calls[0]["json"]["metadata"]["reused_request_state"] == "request-state" + + +@pytest.mark.asyncio +async def test_native_metadata_injection_trace_redacts_configured_paths(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + rule = FieldCacheRule( + name="metadata_secret", + source="response", + path="choices.0.message.reasoning_content", + inject=FieldCacheInjection(target="metadata", path="cached_blob"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=(rule,), + transaction_logger=logger, + ) + executor = NativeProviderExecutor() + + await executor.execute( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "chat_1", "choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "metadata-secret"}}]})), + ) + await executor.execute( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(FakeHTTPClient({"id": "chat_2", "choices": []})), + ) + + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "metadata-secret" not in trace_text + metadata_entries = [entry for entry in _trace_entries(logger.log_dir) if entry["pass_name"] == "after_metadata_field_cache_injection"] + assert metadata_entries[-1]["data"]["cached_blob"] == "[REDACTED]" + + +@pytest.mark.asyncio +async def test_native_provider_stream_rejects_unsupported_operation_before_transport() -> None: + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + operation="embeddings", + endpoint="https://example.test/chat", + ) + client = FakeHTTPClient({"id": "should_not_call"}) + + with pytest.raises(ProtocolError, match="unsupported operation"): + async for _ in NativeProviderExecutor().stream( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(client), + ): + pass + + assert client.calls == [] diff --git a/tests/test_native_provider_streaming.py b/tests/test_native_provider_streaming.py new file mode 100644 index 000000000..4cadef0df --- /dev/null +++ b/tests/test_native_provider_streaming.py @@ -0,0 +1,376 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.adapters import PayloadAdapter, register_adapter +from rotator_library.core.errors import StreamedAPIError +from rotator_library.field_cache import FieldCacheInjection, FieldCacheRule +from rotator_library.native_provider import NativeHTTPTransport, NativeProviderContext, NativeProviderExecutor +from rotator_library.providers.antigravity_provider import AntigravityProvider +from rotator_library.providers.claude_code_provider import ClaudeCodeProvider +from rotator_library.providers.codex_provider import CodexProvider +from rotator_library.transaction_logger import TransactionLogger + + +class FakeStreamingClient: + def __init__(self, chunks): + self.chunks = chunks + self.calls = [] + + async def stream_json_lines(self, endpoint, *, headers, json): + self.calls.append({"endpoint": endpoint, "headers": headers, "json": json}) + for chunk in self.chunks: + yield chunk + + +def _trace_entries(log_dir): + return [json.loads(line) for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +@pytest.mark.asyncio +async def test_native_provider_stream_traces_and_yields_formatted_events(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=( + FieldCacheRule(name="stream_reasoning", source="stream_event", path="raw.choices.0.delta.reasoning_content", allow_missing_session=True, scope=("provider", "model")), + FieldCacheRule(name="stream_vendor_state", source="stream_event", path="raw.choices.0.delta.vendor_state", allow_missing_session=True, scope=("provider", "model")), + ), + transaction_logger=logger, + ) + chunks = [ + {"choices": [{"delta": {"content": "hi", "reasoning_content": "hidden", "vendor_state": "opaque-vendor-state"}}]}, + "[DONE]", + ] + client = FakeStreamingClient(chunks) + + events = [event async for event in NativeProviderExecutor().stream({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(client))] + + assert len(events) == 2 + first_payload = json.loads(events[0].removeprefix("data: ")) + assert first_payload["choices"][0]["delta"]["content"] == "hi" + assert first_payload["choices"][0]["delta"]["reasoning_content"] == "hidden" + assert "vendor_state" not in first_payload["choices"][0]["delta"] + assert events[-1] == "data: [DONE]\n\n" + assert client.calls[0]["json"]["stream"] is True + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "native_provider_stream_request" in pass_names + assert pass_names.count("raw_native_provider_stream_chunk") == 2 + assert pass_names.count("parsed_native_stream_event") == 2 + assert "after_field_cache_extraction" in pass_names + assert "after_field_cache_stream_extraction" in pass_names + assert pass_names.count("formatted_client_stream_event") == 2 + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "opaque-vendor-state" not in trace_text + + +@pytest.mark.asyncio +async def test_native_provider_stream_traces_usage_accounting_summary(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + transaction_logger=logger, + ) + chunks = [ + { + "choices": [], + "usage": { + "prompt_tokens": 2, + "completion_tokens": 3, + "cost_details": {"total_cost": 0.04, "currency": "USD", "source": "stream_usage"}, + }, + }, + "[DONE]", + ] + + _ = [event async for event in NativeProviderExecutor().stream({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(FakeStreamingClient(chunks)))] + + entries = _trace_entries(logger.log_dir) + summaries = [entry for entry in entries if entry["pass_name"] == "usage_accounting_summary"] + assert summaries + assert summaries[-1]["data"]["usage"]["input_tokens"] == 2 + assert summaries[-1]["data"]["usage"]["completion_tokens"] == 3 + assert summaries[-1]["data"]["usage"]["provider_reported_cost"] == 0.04 + assert summaries[-1]["data"]["cost"]["provider_reported_cost"] == 0.04 + + +@pytest.mark.asyncio +async def test_native_provider_stream_preserves_earlier_cost_when_later_usage_arrives(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + transaction_logger=logger, + ) + chunks = [ + {"choices": [], "usage": {"cost_details": {"total_cost": 0.07, "source": "early_cost"}}}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + "[DONE]", + ] + + _ = [event async for event in NativeProviderExecutor().stream({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(FakeStreamingClient(chunks)))] + + summaries = [entry for entry in _trace_entries(logger.log_dir) if entry["pass_name"] == "usage_accounting_summary"] + assert summaries[-1]["data"]["usage"]["input_tokens"] == 2 + assert summaries[-1]["data"]["usage"]["completion_tokens"] == 3 + assert summaries[-1]["data"]["usage"]["provider_reported_cost"] == 0.07 + + +@pytest.mark.asyncio +async def test_native_provider_stream_preserves_cost_when_later_raw_usage_arrives(tmp_path) -> None: + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + transaction_logger=logger, + ) + chunks = [ + {"choices": [], "usage": {"cost_details": {"total_cost": 0.08, "source": "early_cost"}}}, + {"choices": [], "prompt_tokens": 2, "completion_tokens": 3}, + "[DONE]", + ] + + _ = [event async for event in NativeProviderExecutor().stream({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(FakeStreamingClient(chunks)))] + + summaries = [entry for entry in _trace_entries(logger.log_dir) if entry["pass_name"] == "usage_accounting_summary"] + assert summaries[-1]["data"]["usage"]["input_tokens"] == 2 + assert summaries[-1]["data"]["usage"]["completion_tokens"] == 3 + assert summaries[-1]["data"]["usage"]["provider_reported_cost"] == 0.08 + + +@pytest.mark.asyncio +async def test_native_provider_stream_logs_errors(tmp_path) -> None: + class BrokenClient: + async def stream_json_lines(self, endpoint, *, headers, json): + raise RuntimeError("broken stream") + yield None + + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext(provider="native", model="gpt-test", protocol_name="openai_chat", endpoint="https://example.test/chat", transaction_logger=logger) + + with pytest.raises(RuntimeError): + [event async for event in NativeProviderExecutor().stream({"model": "gpt-test", "messages": []}, context, NativeHTTPTransport(BrokenClient()))] + + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "transform_log_error" in pass_names + + +@pytest.mark.asyncio +async def test_native_provider_error_event_raises_for_rotation_before_client_formatting() -> None: + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + ) + + with pytest.raises(StreamedAPIError) as raised: + _ = [ + event + async for event in NativeProviderExecutor().stream( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(FakeStreamingClient([{"error": {"type": "rate_limit", "message": "rotate me"}}])), + ) + ] + + assert raised.value.data["error"]["type"] == "rate_limit" + + +@pytest.mark.asyncio +async def test_native_provider_stream_runs_stream_event_adapter_chain(tmp_path) -> None: + class StreamTextAdapter(PayloadAdapter): + name = "test_stream_text_adapter" + supported_stages = ("stream_event",) + + async def transform_stream_event(self, payload, context): + payload.delta.content[0].text = "adapted" + return payload + + register_adapter(StreamTextAdapter, replace=True) + logger = TransactionLogger("native", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + adapter_names=("test_stream_text_adapter",), + transaction_logger=logger, + ) + + events = [ + event + async for event in NativeProviderExecutor().stream( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(FakeStreamingClient([{"choices": [{"delta": {"content": "before"}}]}, "[DONE]"])), + ) + ] + + payload = json.loads(events[0].removeprefix("data: ")) + assert payload["choices"][0]["delta"]["content"] == "adapted" + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "after_stream_event_adapter_chain" in pass_names + + +@pytest.mark.asyncio +async def test_native_cross_protocol_stream_formats_openai_chat_sse() -> None: + context = NativeProviderContext( + provider="claude_code", + model="claude-sonnet-4-5", + protocol_name="anthropic_messages", + input_protocol_name="anthropic_messages", + output_protocol_name="openai_chat", + endpoint="https://example.test/messages", + operation="messages", + ) + chunks = [ + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + "[DONE]", + ] + + events = [event async for event in NativeProviderExecutor().stream({"model": "claude-sonnet-4-5", "messages": [], "max_tokens": 1}, context, NativeHTTPTransport(FakeStreamingClient(chunks)))] + + assert events[0].startswith("data: ") + payload = json.loads(events[0][len("data: ") :].strip()) + assert payload["object"] == "chat.completion.chunk" + assert payload["choices"][0]["delta"]["content"] == "hi" + assert "content_block_delta" not in events[0] + + +@pytest.mark.asyncio +async def test_native_provider_stream_extracts_unified_stream_events_for_later_requests() -> None: + rule = FieldCacheRule( + name="unified_stream_text", + source="unified_stream_event", + path="delta.content.0.text", + inject=FieldCacheInjection(target="request", path="metadata.cached_stream_text"), + allow_missing_session=True, + scope=("provider", "model"), + ) + context = NativeProviderContext( + provider="native", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + field_cache_rules=(rule,), + ) + executor = NativeProviderExecutor() + + _ = [ + event + async for event in executor.stream( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(FakeStreamingClient([{"choices": [{"delta": {"content": "stream-state"}}]}, "[DONE]"])), + ) + ] + second_client = FakeStreamingClient(["[DONE]"]) + _ = [ + event + async for event in executor.stream( + {"model": "gpt-test", "messages": []}, + context, + NativeHTTPTransport(second_client), + ) + ] + + assert second_client.calls[0]["json"]["metadata"]["cached_stream_text"] == "stream-state" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("provider", "protocol", "operation", "request_payload", "chunks", "injected_path", "expected"), + [ + ( + AntigravityProvider(), + "gemini", + "stream_generate", + {"model": "gpt-test", "contents": []}, + [{"candidates": [{"content": {"role": "model", "parts": [{"text": "private", "thought": True, "thoughtSignature": "gem-signature"}]}}]}, "[DONE]"], + ("request", "metadata", "thoughtSignatures"), + ["gem-signature"], + ), + ( + ClaudeCodeProvider(), + "anthropic_messages", + "messages", + {"model": "gpt-test", "messages": [], "max_tokens": 1}, + [{"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "claude-signature"}}, {"type": "message_stop"}], + ("metadata", "thinking_signatures"), + ["claude-signature"], + ), + ( + CodexProvider(), + "responses", + "responses", + {"model": "gpt-test", "input": []}, + [{"type": "response.completed", "response": {"id": "resp-continuation", "status": "completed", "model": "gpt-test", "output": []}}, "[DONE]"], + ("previous_response_id",), + "resp-continuation", + ), + ], +) +async def test_provider_stream_state_is_cached_for_followups_but_not_exposed( + provider, + protocol, + operation, + request_payload, + chunks, + injected_path, + expected, +) -> None: + stream_rules = tuple(rule for rule in provider.get_field_cache_rules("gpt-test") if rule.source == "stream_event") + context = NativeProviderContext( + provider=provider.provider_env_name, + model="gpt-test", + protocol_name=protocol, + input_protocol_name=protocol, + output_protocol_name="openai_chat", + endpoint="https://example.test/stream", + operation=operation, + credential_id="credential-1", + session_id="session-1", + scope_key="scope-1", + field_cache_rules=stream_rules, + ) + executor = NativeProviderExecutor() + + first_output = [ + frame + async for frame in executor.stream( + request_payload, + context, + NativeHTTPTransport(FakeStreamingClient(chunks)), + ) + ] + second_client = FakeStreamingClient(["[DONE]"]) + _ = [ + frame + async for frame in executor.stream( + request_payload, + context, + NativeHTTPTransport(second_client), + ) + ] + + current = second_client.calls[0]["json"] + for key in injected_path: + current = current[key] + assert current == expected + output_text = "".join(first_output) + assert "gem-signature" not in output_text + assert "claude-signature" not in output_text + assert "resp-continuation" not in output_text diff --git a/tests/test_native_streaming_transport_seam.py b/tests/test_native_streaming_transport_seam.py new file mode 100644 index 000000000..5cef50670 --- /dev/null +++ b/tests/test_native_streaming_transport_seam.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import pytest + +from rotator_library.native_provider.http import NativeHTTPTransport +from rotator_library.native_provider.streaming import native_stream_event_from_formatted, provider_supports_native_streaming +from rotator_library.providers.provider_interface import ProviderInterface + + +class DefaultProvider(ProviderInterface): + async def get_models(self, api_key, client): + return [] + + +class StreamingProvider(DefaultProvider): + native_streaming_supported = True + + +def test_provider_native_streaming_support_defaults_false() -> None: + assert provider_supports_native_streaming(DefaultProvider(), model="gpt-test") is False + assert provider_supports_native_streaming(StreamingProvider(), model="gpt-test") is True + + +def test_native_formatted_sse_chunk_uses_common_stream_event_seam() -> None: + event = native_stream_event_from_formatted('data: {"choices":[{"delta":{"content":"hi"}}]}\n\n') + + assert event.visible_output is True + assert event.event_type == "delta" + + +class FakeStreamResponse: + def __init__(self, lines): + self.lines = lines + self.raised = False + + def raise_for_status(self): + self.raised = True + + async def aiter_lines(self): + for line in self.lines: + yield line + + +class FakeStreamContext: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class FakeHttpxClient: + def __init__(self, response): + self.response = response + self.calls = [] + + def stream(self, method, endpoint, *, headers, json): + self.calls.append((method, endpoint, headers, json)) + return FakeStreamContext(self.response) + + +@pytest.mark.asyncio +async def test_native_http_transport_streams_httpx_lines() -> None: + response = FakeStreamResponse(["", ": heartbeat", "event: content_block_delta", 'data: {"delta":', 'data: "hi"}', "", "data: [DONE]", ""]) + client = FakeHttpxClient(response) + + chunks = [chunk async for chunk in NativeHTTPTransport(client).stream_json_lines("https://example.test/stream", headers={"h": "v"}, payload={"p": True})] + + assert client.calls == [("POST", "https://example.test/stream", {"h": "v"}, {"p": True})] + assert response.raised is True + assert chunks == [{"delta": "hi", "type": "content_block_delta"}, "[DONE]"] + + +class FakeByteResponse: + def raise_for_status(self): + pass + + async def aiter_bytes(self): + yield b'event: response.output_text.delta\ndata: {"a":' + yield b' 1}\n\ndata: [DO' + yield b'NE]\n\n' + + +@pytest.mark.asyncio +async def test_native_http_transport_streams_httpx_bytes() -> None: + chunks = [chunk async for chunk in NativeHTTPTransport(FakeHttpxClient(FakeByteResponse())).stream_json_lines("url", headers={}, payload={})] + + assert chunks == [{"a": 1, "type": "response.output_text.delta"}, "[DONE]"] diff --git a/tests/test_native_usage_accounting.py b/tests/test_native_usage_accounting.py new file mode 100644 index 000000000..62116a6d4 --- /dev/null +++ b/tests/test_native_usage_accounting.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.native_provider import NativeHTTPTransport, NativeProviderContext, NativeProviderExecutor +from rotator_library.transaction_logger import TransactionLogger + + +class FakeNativeTransport(NativeHTTPTransport): + def __init__(self): + pass + + async def post_json(self, endpoint, *, headers, payload): + return { + "id": "chatcmpl_1", + "model": "gpt-test", + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 6, "completion_tokens_details": {"reasoning_tokens": 2}}, + } + + +@pytest.mark.asyncio +async def test_native_executor_traces_normalized_usage(tmp_path) -> None: + logger = TransactionLogger("openai", "gpt-test", parent_dir=tmp_path) + context = NativeProviderContext( + provider="openai", + model="gpt-test", + protocol_name="openai_chat", + endpoint="https://example.test/chat", + headers={}, + transaction_logger=logger, + ) + + await NativeProviderExecutor().execute({"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}]}, context, FakeNativeTransport()) + + entries = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + usage_entries = [entry for entry in entries if entry["pass_name"] == "usage_accounting_summary"] + assert usage_entries[-1]["data"]["usage"]["completion_tokens"] == 4 + assert usage_entries[-1]["data"]["usage"]["reasoning_tokens"] == 2 diff --git a/tests/test_nvidia_provider.py b/tests/test_nvidia_provider.py new file mode 100644 index 000000000..6360cf486 --- /dev/null +++ b/tests/test_nvidia_provider.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import pytest + +from rotator_library.providers.nvidia_provider import NvidiaProvider + + +MODEL = "nvidia_nim/google/diffusiongemma-26b-a4b-it" + + +@pytest.mark.parametrize( + ("reasoning_effort", "expected"), + [ + (None, True), + ("low", True), + ("medium", True), + ("high", True), + ("off", False), + ("disable", False), + ("disabled", False), + ("none", False), + ], +) +def test_diffusion_gemma_maps_reasoning_effort_to_enable_thinking( + reasoning_effort: str | None, + expected: bool, +) -> None: + """DiffusionGemma receives NVIDIA's required top-level template toggle.""" + payload = {"chat_template_kwargs": {"preserved": True}} + if reasoning_effort is not None: + payload["reasoning_effort"] = reasoning_effort + + NvidiaProvider().handle_thinking_parameter(payload, MODEL) + + assert payload["chat_template_kwargs"] == { + "preserved": True, + "enable_thinking": expected, + } + assert "reasoning_effort" not in payload + assert "extra_body" not in payload + + +def test_diffusion_gemma_mapping_does_not_affect_other_nvidia_models() -> None: + """The exact model match prevents template flags leaking to other models.""" + payload = {"reasoning_effort": "high"} + + NvidiaProvider().handle_thinking_parameter( + payload, + "nvidia_nim/google/gemma-3-27b-it", + ) + + assert payload == {"reasoning_effort": "high"} + + +@pytest.mark.parametrize( + "initial_template_kwargs", + [ + None, + [], + {"enable_thinking": False, "preserved": True}, + ], +) +def test_diffusion_gemma_builds_or_overwrites_template_toggle( + initial_template_kwargs: object, +) -> None: + """Malformed/missing template containers cannot suppress the model toggle.""" + + payload = {} + if initial_template_kwargs is not None: + payload["chat_template_kwargs"] = initial_template_kwargs + + NvidiaProvider().handle_thinking_parameter(payload, MODEL) + + expected = {"enable_thinking": True} + if isinstance(initial_template_kwargs, dict): + expected["preserved"] = True + assert payload["chat_template_kwargs"] == expected + + +@pytest.mark.parametrize( + ("model", "assert_disabled"), + [ + ( + "nvidia_nim/moonshotai/kimi-k2.5", + lambda payload: payload["chat_template_kwargs"]["thinking"] is False, + ), + ( + "nvidia_nim/mistralai/mistral-medium-3.5-instruct", + lambda payload: "extra_body" not in payload, + ), + ( + "nvidia_nim/deepseek-ai/deepseek-v3.2", + lambda payload: payload["extra_body"]["chat_template_kwargs"]["thinking"] + is False, + ), + ( + "nvidia_nim/deepseek-ai/deepseek-v4-flash", + lambda payload: payload["extra_body"]["chat_template_kwargs"]["thinking"] + is False, + ), + ], +) +def test_disabled_alias_applies_to_existing_reasoning_branches( + model: str, + assert_disabled, +) -> None: + """The shared disabled alias intentionally applies to every supported family.""" + + payload = {"reasoning_effort": "disabled"} + + NvidiaProvider().handle_thinking_parameter(payload, model) + + assert "reasoning_effort" not in payload + assert assert_disabled(payload) diff --git a/tests/test_protocol_anthropic_messages.py b/tests/test_protocol_anthropic_messages.py new file mode 100644 index 000000000..9c3872fc8 --- /dev/null +++ b/tests/test_protocol_anthropic_messages.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import json + +from rotator_library.protocols import ProtocolContext, get_protocol, list_protocols + + +def _compatible_context() -> ProtocolContext: + """Return a provider-state domain that may restore hidden signatures.""" + + return ProtocolContext( + source_protocol="anthropic_messages", + target_protocol="anthropic_messages", + source_provider="anthropic", + target_provider="anthropic", + ) + + +def test_anthropic_build_uses_mutated_unified_thinking_signature() -> None: + adapter = get_protocol("anthropic_messages") + raw = { + "model": "anthropic/claude-test", + "messages": [{"role": "assistant", "content": [{"type": "thinking", "thinking": "before", "signature": "sig_1"}]}], + } + + unified = adapter.parse_request(raw) + unified.messages[0].content[0].reasoning.text = "after" + rebuilt = adapter.build_request(unified, _compatible_context()) + + assert rebuilt["messages"][0]["content"][0]["thinking"] == "after" + assert rebuilt["messages"][0]["content"][0]["signature"] == "sig_1" + + +def test_anthropic_messages_protocol_is_discovered_with_aliases() -> None: + assert "anthropic_messages" in list_protocols() + assert get_protocol("anthropic") is get_protocol("anthropic_messages") + assert get_protocol("messages") is get_protocol("anthropic_messages") + + +def test_anthropic_request_round_trip_preserves_thinking_tools_and_cache_metadata() -> None: + adapter = get_protocol("anthropic_messages") + raw = { + "model": "anthropic/claude-test", + "system": [{"type": "text", "text": "system"}], + "max_tokens": 100, + "stream": True, + "messages": [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "worked", "signature": "sig_1"}, + {"type": "text", "text": "answer"}, + {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {"q": "x"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "result"}]}, + ], + "tools": [{"name": "lookup", "description": "Lookup", "input_schema": {"type": "object"}}], + "vendor_extension": {"kept": True}, + } + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified, _compatible_context()) + + assert unified.model == "anthropic/claude-test" + assert unified.system[0].text == "system" + assert unified.messages[1].reasoning[0].signature == "sig_1" + assert unified.messages[1].tool_calls[0].name == "lookup" + assert unified.messages[2].content[0].tool_result.tool_call_id == "toolu_1" + assert unified.tools[0].input_schema == {"type": "object"} + assert rebuilt["vendor_extension"] == {"kept": True} + assert isinstance(rebuilt["system"], list) + assert rebuilt["system"][0]["text"] == "system" + assert rebuilt["messages"][1]["content"][0]["signature"] == "sig_1" + + +def test_anthropic_system_preserves_block_metadata() -> None: + adapter = get_protocol("anthropic_messages") + raw = { + "model": "anthropic/claude-test", + "system": [{"type": "text", "text": "system", "cache_control": {"type": "ephemeral"}}], + "messages": [{"role": "user", "content": "hello"}], + } + + unified = adapter.parse_request(raw) + unified.system[0].text = "updated" + rebuilt = adapter.build_request(unified) + + assert rebuilt["system"] == [{"type": "text", "text": "updated", "cache_control": {"type": "ephemeral"}}] + + +def test_anthropic_response_extracts_content_and_usage() -> None: + adapter = get_protocol("anthropic_messages") + raw = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "anthropic/claude-test", + "content": [ + {"type": "redacted_thinking", "signature": "sig_2"}, + {"type": "text", "text": "answer"}, + ], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "cache_creation_input_tokens": 2, + "cache_read_input_tokens": 3, + }, + } + + unified = adapter.parse_response(raw) + + assert unified.id == "msg_1" + assert unified.messages[0].content[1].text == "answer" + assert unified.messages[0].reasoning[0].signature == "sig_2" + assert unified.messages[0].reasoning[0].redacted is True + assert unified.usage is not None + assert unified.usage.input_tokens == 10 + assert unified.usage.output_tokens == 5 + assert unified.usage.cache_write_tokens == 2 + assert unified.usage.cache_read_tokens == 3 + + +def test_anthropic_stream_event_parses_text_delta_and_usage() -> None: + adapter = get_protocol("anthropic_messages") + text_event = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}} + usage_event = {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 7}} + + parsed_text = adapter.parse_stream_event(f"data: {json.dumps(text_event)}\n\n") + parsed_usage = adapter.parse_stream_event(usage_event) + + assert parsed_text.type == "content_block_delta" + assert parsed_text.delta is not None + assert parsed_text.delta.content[0].text == "Hi" + assert parsed_usage.usage is not None + assert parsed_usage.usage.output_tokens == 7 diff --git a/tests/test_protocol_client_surfaces.py b/tests/test_protocol_client_surfaces.py new file mode 100644 index 000000000..a93e818f1 --- /dev/null +++ b/tests/test_protocol_client_surfaces.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from fastapi.testclient import TestClient + +from proxy_app import main as proxy_main +from rotator_library.client.gemini import GeminiHandler +from rotator_library.client.rotating_client import RotatingClient +from rotator_library.client.protocol_selection import request_output_protocol + + +class SurfaceClient: + """Minimal route client that records client-protocol handoff details.""" + + def __init__(self) -> None: + self.calls: list[dict] = [] + + def resolve_output_protocol(self, payload, *, input_protocol, request=None, explicit=None): + return explicit or request_output_protocol(request) or input_protocol + + async def agenerate(self, payload, *, input_protocol, request=None, **kwargs): + self.resolve_output_protocol( + payload, + input_protocol=input_protocol, + request=request, + explicit=kwargs.get("output_protocol"), + ) + self.calls.append( + { + "payload": payload, + "input_protocol": input_protocol, + "output_header": request.headers.get("X-Proxy-Output-Protocol"), + } + ) + if payload.get("stream"): + async def stream(): + yield 'data: {"candidates":[{"content":{"role":"model","parts":[{"text":"streamed"}]}}]}\n\n' + return stream() + if request.headers.get("X-Proxy-Output-Protocol") == "anthropic_messages": + return { + "id": "msg_surface", + "type": "message", + "role": "assistant", + "model": payload["model"], + "content": [{"type": "text", "text": "selected"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + return { + "id": "chat_surface", + "object": "chat.completion", + "model": payload["model"], + "choices": [{"message": {"role": "assistant", "content": "chat"}, "finish_reason": "stop"}], + } + + async def gemini_generate(self, payload, *, model, raw_request=None): + self.calls.append( + { + "payload": payload, + "model": model, + "input_protocol": "gemini", + "output_header": raw_request.headers.get("X-Proxy-Output-Protocol"), + } + ) + return { + "responseId": "gemini_surface", + "modelVersion": model, + "candidates": [{"content": {"role": "model", "parts": [{"text": "gemini"}]}, "finishReason": "STOP"}], + } + + async def gemini_stream_generate(self, payload, *, model, raw_request=None): + self.calls.append({"payload": payload, "model": model, "operation": "stream_generate"}) + + async def stream(): + yield 'data: {"candidates":[{"content":{"role":"model","parts":[{"text":"gemini-stream"}]}}]}\n\n' + + return stream() + + def gemini_count_tokens(self, payload, *, model): + self.calls.append({"payload": payload, "model": model, "operation": "count_tokens"}) + return {"totalTokens": 7} + + +def _surface_client() -> tuple[TestClient, SurfaceClient]: + proxy_main.PROXY_API_KEY = None + proxy_main.ENABLE_RAW_LOGGING = False + rotating = SurfaceClient() + proxy_main.app.state.rotating_client = rotating + return TestClient(proxy_main.app), rotating + + +def test_chat_route_forwards_independent_output_protocol_selector() -> None: + client, rotating = _surface_client() + + response = client.post( + "/v1/chat/completions", + headers={"X-Proxy-Output-Protocol": "anthropic_messages"}, + json={"model": "openai/gpt-test", "messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == 200 + assert response.json()["type"] == "message" + assert rotating.calls[0]["input_protocol"] == "openai_chat" + assert rotating.calls[0]["output_header"] == "anthropic_messages" + + +def test_gemini_generate_and_count_routes_preserve_native_client_shape() -> None: + client, rotating = _surface_client() + payload = {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]} + + generated = client.post("/v1beta/models/gemini-2.5-pro:generateContent", json=payload) + counted = client.post("/v1beta/models/gemini-2.5-pro:countTokens", json=payload) + + assert generated.status_code == 200 + assert generated.json()["candidates"][0]["content"]["parts"][0]["text"] == "gemini" + assert counted.json() == {"totalTokens": 7} + assert rotating.calls[0]["input_protocol"] == "gemini" + assert rotating.calls[0]["payload"] == payload + assert rotating.calls[1]["operation"] == "count_tokens" + + +def test_gemini_stream_generate_route_preserves_native_stream_shape() -> None: + client, rotating = _surface_client() + + response = client.post( + "/v1beta/models/gemini-2.5-pro:streamGenerateContent", + json={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + ) + + assert response.status_code == 200 + assert "gemini-stream" in response.text + assert rotating.calls[0]["operation"] == "stream_generate" + + +def test_cross_protocol_stream_selection_reaches_selected_wire_format() -> None: + client, _ = _surface_client() + + response = client.post( + "/v1/chat/completions", + headers={"X-Proxy-Output-Protocol": "gemini"}, + json={ + "model": "openai/gpt-test", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + }, + ) + + assert response.status_code == 200 + assert '"candidates"' in response.text + assert "streamed" in response.text + + +def test_gemini_generate_rejects_stream_flag_until_stream_route_exists() -> None: + client, _ = _surface_client() + + response = client.post( + "/v1beta/models/gemini-2.5-pro:generateContent", + json={"contents": [{"role": "user", "parts": [{"text": "hello"}]}], "stream": True}, + ) + + assert response.status_code == 400 + assert response.json()["error"]["status"] == "INVALID_ARGUMENT" + + +def test_output_protocol_precedence_is_explicit_then_header_then_provider_then_input() -> None: + client = RotatingClient.__new__(RotatingClient) + client._get_provider_instance = lambda provider: SimpleNamespace( + get_default_output_protocol=lambda model: "responses" + ) + request = SimpleNamespace(headers={"X-Proxy-Output-Protocol": "anthropic"}) + payload = {"model": "configured/model"} + + assert client.resolve_output_protocol(payload, input_protocol="gemini", request=request, explicit="chat") == "openai_chat" + assert client.resolve_output_protocol(payload, input_protocol="gemini", request=request) == "anthropic_messages" + assert client.resolve_output_protocol(payload, input_protocol="gemini") == "responses" + assert client.resolve_output_protocol({"model": "bare"}, input_protocol="gemini") == "gemini" + + +def test_output_protocol_provider_default_resolves_through_model_alias(monkeypatch) -> None: + monkeypatch.setenv("MODEL_ROUTE_ALIAS", "configured/model") + client = RotatingClient.__new__(RotatingClient) + client._get_provider_instance = lambda provider: SimpleNamespace( + get_default_output_protocol=lambda model: "anthropic_messages" + ) + + assert client.resolve_output_protocol({"model": "alias"}, input_protocol="gemini") == "anthropic_messages" + + +def test_output_header_is_case_insensitive_and_unknown_values_are_client_errors() -> None: + assert request_output_protocol(SimpleNamespace(headers={"X-Proxy-Output-Protocol": "Anthropic"})) == "anthropic_messages" + try: + request_output_protocol(SimpleNamespace(headers={"X-Proxy-Output-Protocol": "unknown"})) + except ValueError as error: + assert "Unsupported output protocol" in str(error) + else: + raise AssertionError("unknown output protocol must fail") + + +class FailingSurfaceClient(SurfaceClient): + async def agenerate(self, payload, **kwargs): + raise ValueError("local validation failed") + + async def anthropic_messages(self, body, **kwargs): + raise ValueError("local validation failed") + + async def gemini_generate(self, payload, **kwargs): + raise ValueError("local validation failed") + + +@pytest.mark.parametrize( + ("path", "headers", "body", "assertion"), + ( + ( + "/v1/chat/completions", + {"X-Proxy-Output-Protocol": "anthropic_messages"}, + {"model": "openai/gpt-test", "messages": [{"role": "user", "content": "hello"}]}, + lambda payload: payload["type"] == "error" and payload["error"]["type"] == "invalid_request_error", + ), + ( + "/v1/messages", + {"X-Proxy-Output-Protocol": "gemini"}, + {"model": "claude-test", "max_tokens": 8, "messages": [{"role": "user", "content": "hello"}]}, + lambda payload: payload["error"]["status"] == "INVALID_ARGUMENT", + ), + ( + "/v1beta/models/gemini-2.5-pro:generateContent", + {"X-Proxy-Output-Protocol": "openai_chat"}, + {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + lambda payload: payload["error"]["type"] == "invalid_request", + ), + ), +) +def test_proxy_side_errors_use_selected_output_protocol(path, headers, body, assertion) -> None: + proxy_main.PROXY_API_KEY = None + proxy_main.ENABLE_RAW_LOGGING = False + proxy_main.app.state.rotating_client = FailingSurfaceClient() + + response = TestClient(proxy_main.app).post(path, headers=headers, json=body) + + assert response.status_code == 400 + assert assertion(response.json()) + assert "detail" not in response.json() + + +def test_unknown_output_header_is_a_400_in_input_protocol_shape() -> None: + client, _ = _surface_client() + + response = client.post( + "/v1/chat/completions", + headers={"X-Proxy-Output-Protocol": "unknown"}, + json={"model": "openai/gpt-test", "messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == 400 + assert response.json()["error"]["type"] == "invalid_request" + assert "Unsupported output protocol" in response.json()["error"]["message"] + + +@pytest.mark.parametrize( + ("output_protocol", "assertion"), + ( + ( + "anthropic_messages", + lambda payload: payload["type"] == "error" and payload["error"]["type"] == "invalid_request_error", + ), + ( + "gemini", + lambda payload: payload["error"]["status"] == "INVALID_ARGUMENT", + ), + ), +) +def test_responses_malformed_json_uses_selected_output_protocol(output_protocol, assertion) -> None: + client, _ = _surface_client() + + response = client.post( + "/v1/responses", + headers={ + "Content-Type": "application/json", + "X-Proxy-Output-Protocol": output_protocol, + }, + content="{", + ) + + assert response.status_code == 400 + assert assertion(response.json()) + assert "detail" not in response.json() + + +def test_library_accepts_all_generative_cross_protocol_stream_pairs() -> None: + from rotator_library.client.protocol_selection import require_same_protocol_stream + + for source in ("openai_chat", "anthropic_messages", "responses", "gemini"): + for output in ("openai_chat", "anthropic_messages", "responses", "gemini"): + require_same_protocol_stream(source, output) + + +class GeminiRuntimeClient: + def __init__(self) -> None: + self.calls = [] + + async def agenerate(self, payload, **kwargs): + self.calls.append((payload, kwargs)) + return {"candidates": []} + + def token_count(self, *, model, messages=None, text=None): + return 5 if messages is not None else 2 + + +def test_gemini_handler_defaults_bare_models_but_preserves_model_routes(monkeypatch) -> None: + monkeypatch.setenv("MODEL_ROUTE_ALIAS", "openai/gpt-test") + client = GeminiRuntimeClient() + handler = GeminiHandler(client) + + assert handler._routable_model("gemini-2.5-pro") == "gemini/gemini-2.5-pro" + assert handler._routable_model("alias") == "alias" + assert handler._routable_model("configured/model") == "configured/model" + + +@pytest.mark.asyncio +async def test_gemini_handler_forwards_raw_request_to_generic_runtime_selector() -> None: + client = GeminiRuntimeClient() + handler = GeminiHandler(client) + request = SimpleNamespace(headers={"X-Proxy-Output-Protocol": "anthropic_messages"}) + + await handler.generate( + {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + model="gemini-2.5-pro", + raw_request=request, + ) + + assert client.calls[0][1]["input_protocol"] == "gemini" + assert client.calls[0][1]["request"] is request diff --git a/tests/test_protocol_gemini.py b/tests/test_protocol_gemini.py new file mode 100644 index 000000000..4dd696707 --- /dev/null +++ b/tests/test_protocol_gemini.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import json + +from rotator_library.protocols import get_protocol, list_protocols + + +def test_gemini_protocol_is_discovered_with_aliases() -> None: + assert "gemini" in list_protocols() + assert get_protocol("google_gemini") is get_protocol("gemini") + assert get_protocol("generate_content") is get_protocol("gemini") + + +def test_gemini_request_round_trip_preserves_parts_tools_and_settings() -> None: + adapter = get_protocol("gemini") + raw = { + "model": "gemini/gemini-test", + "systemInstruction": {"parts": [{"text": "system"}]}, + "contents": [ + {"role": "user", "parts": [{"text": "hello"}, {"inlineData": {"mimeType": "image/png", "data": "abc"}}]}, + { + "role": "model", + "parts": [ + {"text": "thought", "thought": True, "thoughtSignature": "sig_1"}, + {"functionCall": {"name": "lookup", "args": {"q": "x"}}}, + ], + }, + {"role": "user", "parts": [{"functionResponse": {"name": "lookup", "response": {"value": 1}}}]}, + ], + "tools": [{"functionDeclarations": [{"name": "lookup", "description": "Lookup", "parameters": {"type": "object"}}]}], + "generationConfig": {"temperature": 0.3}, + "safetySettings": [{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"}], + "vendor_extension": {"kept": True}, + } + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified) + + assert unified.model == "gemini/gemini-test" + assert unified.system[0].text == "system" + assert unified.messages[1].role == "assistant" + assert unified.messages[1].reasoning[0].signature == "sig_1" + assert unified.messages[1].tool_calls[0].name == "lookup" + assert unified.messages[2].content[0].tool_result.content == {"value": 1} + assert unified.tools[0].name == "lookup" + assert rebuilt["generationConfig"] == {"temperature": 0.3} + assert rebuilt["safetySettings"][0]["threshold"] == "BLOCK_NONE" + assert rebuilt["vendor_extension"] == {"kept": True} + + +def test_gemini_parses_multiple_function_declarations() -> None: + adapter = get_protocol("gemini") + raw = { + "contents": [], + "tools": [ + { + "functionDeclarations": [ + {"name": "one", "parameters": {"type": "object"}}, + {"name": "two", "parameters": {"type": "object"}}, + ] + } + ], + } + + unified = adapter.parse_request(raw) + unified.tools[1].description = "Second tool" + rebuilt = adapter.build_request(unified) + + assert [tool.name for tool in unified.tools] == ["one", "two"] + assert len(rebuilt["tools"]) == 1 + assert [declaration["name"] for declaration in rebuilt["tools"][0]["functionDeclarations"]] == ["one", "two"] + assert rebuilt["tools"][0]["functionDeclarations"][1]["description"] == "Second tool" + + +def test_gemini_response_extracts_usage_and_thought_signature() -> None: + adapter = get_protocol("gemini") + raw = { + "responseId": "resp_1", + "modelVersion": "gemini-test-001", + "candidates": [ + { + "finishReason": "STOP", + "content": {"role": "model", "parts": [{"text": "answer", "thought": True, "thoughtSignature": "sig_2"}]}, + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 2, + "cachedContentTokenCount": 3, + "totalTokenCount": 17, + }, + } + + unified = adapter.parse_response(raw) + + assert unified.id == "resp_1" + assert unified.stop_reason == "stop" + assert unified.metadata["native_stop_reason"] == "STOP" + assert unified.messages[0].content[0].text == "answer" + assert unified.messages[0].reasoning[0].signature == "sig_2" + assert unified.usage is not None + assert unified.usage.input_tokens == 10 + assert unified.usage.output_tokens == 5 + assert unified.usage.reasoning_tokens == 2 + assert unified.usage.cache_read_tokens == 3 + + +def test_gemini_stream_event_parses_candidate_delta() -> None: + adapter = get_protocol("gemini") + event = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}, "finishReason": None}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + } + + parsed = adapter.parse_stream_event(f"data: {json.dumps(event)}\n\n") + + assert parsed.type == "message_delta" + assert parsed.delta is not None + assert parsed.delta.content[0].text == "Hi" + assert parsed.usage is not None + assert parsed.usage.total_tokens == 2 diff --git a/tests/test_protocol_interoperability.py b/tests/test_protocol_interoperability.py new file mode 100644 index 000000000..fcd0d80ec --- /dev/null +++ b/tests/test_protocol_interoperability.py @@ -0,0 +1,933 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from rotator_library.protocols import ProtocolContext, ProtocolError, UnifiedMessage, get_protocol +from rotator_library.protocols.canonical import ( + instruction_blocks, + message_reasoning, + message_tool_calls, + message_tool_results, + ordered_message_blocks, +) + + +PROTOCOLS = ("openai_chat", "anthropic_messages", "responses", "gemini") + + +REQUEST_FIXTURES: dict[str, dict[str, Any]] = { + "openai_chat": { + "model": "model-test", + "messages": [ + {"role": "system", "content": "system rule"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": "calling", + "reasoning_content": "reasoned", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q":"x"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "name": "lookup", "content": '{"value":1}'}, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Lookup", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + } + ], + "max_completion_tokens": 123, + "temperature": 0.2, + "stop": ["END"], + "tool_choice": "auto", + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "answer", + "strict": True, + "schema": {"type": "object", "properties": {"value": {"type": "integer"}}}, + }, + }, + "vendor_extension": {"source_only": True}, + }, + "anthropic_messages": { + "model": "model-test", + "system": [{"type": "text", "text": "system rule"}], + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "reasoned", "signature": "sig_anthropic"}, + {"type": "text", "text": "calling"}, + {"type": "tool_use", "id": "call_1", "name": "lookup", "input": {"q": "x"}}, + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "call_1", "content": {"value": 1}} + ], + }, + ], + "tools": [ + { + "name": "lookup", + "description": "Lookup", + "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ], + "max_tokens": 123, + "temperature": 0.2, + "stop_sequences": ["END"], + "tool_choice": {"type": "auto"}, + "output_config": { + "format": { + "type": "json_schema", + "name": "answer", + "strict": True, + "schema": {"type": "object", "properties": {"value": {"type": "integer"}}}, + } + }, + "vendor_extension": {"source_only": True}, + }, + "responses": { + "model": "model-test", + "instructions": "system rule", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "look"}, + {"type": "input_image", "image_url": "data:image/png;base64,aW1hZ2U="}, + ], + }, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "reasoned"}], + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": '{"q":"x"}', + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": {"value": 1}, + }, + ], + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "Lookup", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ], + "max_output_tokens": 123, + "temperature": 0.2, + "tool_choice": "auto", + "text": { + "format": { + "type": "json_schema", + "name": "answer", + "strict": True, + "schema": {"type": "object", "properties": {"value": {"type": "integer"}}}, + } + }, + "vendor_extension": {"source_only": True}, + }, + "gemini": { + "model": "model-test", + "systemInstruction": {"parts": [{"text": "system rule"}]}, + "contents": [ + { + "role": "user", + "parts": [ + {"text": "look"}, + {"inlineData": {"mimeType": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "model", + "parts": [ + {"text": "reasoned", "thought": True, "thoughtSignature": "sig_gemini"}, + {"text": "calling"}, + {"functionCall": {"id": "call_1", "name": "lookup", "args": {"q": "x"}}}, + ], + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"id": "call_1", "name": "lookup", "response": {"value": 1}}} + ], + }, + ], + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Lookup", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ] + } + ], + "generationConfig": { + "maxOutputTokens": 123, + "temperature": 0.2, + "stopSequences": ["END"], + "responseMimeType": "application/json", + "responseJsonSchema": {"type": "object", "properties": {"value": {"type": "integer"}}}, + }, + "toolConfig": {"functionCallingConfig": {"mode": "AUTO"}}, + "vendor_extension": {"source_only": True}, + }, +} + + +RESPONSE_FIXTURES: dict[str, dict[str, Any]] = { + "openai_chat": { + "id": "source_response", + "model": "model-test", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": "answer", + "reasoning_content": "reasoned", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q":"x"}'}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}, + "vendor_extension": {"source_only": True}, + }, + "anthropic_messages": { + "id": "source_response", + "type": "message", + "role": "assistant", + "model": "model-test", + "content": [ + {"type": "thinking", "thinking": "reasoned", "signature": "sig_anthropic"}, + {"type": "text", "text": "answer"}, + {"type": "tool_use", "id": "call_1", "name": "lookup", "input": {"q": "x"}}, + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 4, "output_tokens": 3}, + "vendor_extension": {"source_only": True}, + }, + "responses": { + "id": "source_response", + "object": "response", + "model": "model-test", + "status": "completed", + "output": [ + {"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "reasoned"}]}, + {"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "answer"}]}, + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": '{"q":"x"}'}, + ], + "usage": {"input_tokens": 4, "output_tokens": 3, "total_tokens": 7}, + "vendor_extension": {"source_only": True}, + }, + "gemini": { + "responseId": "source_response", + "modelVersion": "model-test", + "candidates": [ + { + "finishReason": "STOP", + "content": { + "role": "model", + "parts": [ + {"text": "reasoned", "thought": True, "thoughtSignature": "sig_gemini"}, + {"text": "answer"}, + {"functionCall": {"id": "call_1", "name": "lookup", "args": {"q": "x"}}}, + ], + }, + } + ], + "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 3, "totalTokenCount": 7}, + "vendor_extension": {"source_only": True}, + }, +} + + +def _context(source: str, target: str) -> ProtocolContext: + return ProtocolContext(source_protocol=source, target_protocol=target) + + +def _request_semantics(request: Any) -> dict[str, Any]: + instruction_text = "".join(block.text or "" for block in instruction_blocks(request)) + visible_text: list[str] = [] + reasoning_text: list[str] = [] + calls: list[tuple[str | None, str | None, Any]] = [] + results: list[tuple[str | None, str | None, Any]] = [] + media: list[tuple[str, str | None]] = [] + sequence: list[tuple[str, str, Any]] = [] + for message in request.messages: + if message.role in {"system", "developer"}: + continue + for block in message.content: + if block.text and not block.reasoning: + visible_text.append(block.text) + if block.type in {"image", "file", "document", "audio", "video"}: + media_type = getattr(block.source, "media_type", None) + media.append((block.type, media_type)) + reasoning_text.extend(block.text or "" for block in message_reasoning(message)) + calls.extend((call.id, call.name, call.arguments) for call in message_tool_calls(message)) + results.extend((result.tool_call_id, result.name, result.content) for result in message_tool_results(message)) + for block in ordered_message_blocks(message): + if block.reasoning: + sequence.append((message.role, "reasoning", block.reasoning.text)) + elif block.tool_call: + sequence.append((message.role, "tool_call", (block.tool_call.name, block.tool_call.arguments))) + elif block.tool_result: + sequence.append((message.role, "tool_result", (block.tool_result.name, block.tool_result.content))) + elif block.text: + sequence.append((message.role, "text", block.text)) + elif block.type in {"image", "file", "document", "audio", "video"}: + sequence.append((message.role, block.type, getattr(block.source, "media_type", None))) + return { + "instructions": instruction_text, + "visible_text": visible_text, + "reasoning_text": reasoning_text, + "calls": calls, + "results": results, + "media": media, + "tools": [(tool.name, tool.input_schema) for tool in request.tools], + "max_output_tokens": request.generation_params.get("max_output_tokens"), + "temperature": request.generation_params.get("temperature"), + "stop_sequences": request.generation_params.get("stop_sequences"), + "tool_choice": request.generation_params.get("tool_choice"), + "structured_output": request.generation_params.get("structured_output"), + "sequence": sequence, + } + + +def _response_semantics(messages: list[UnifiedMessage]) -> dict[str, Any]: + text: list[str] = [] + reasoning: list[str] = [] + calls: list[tuple[str | None, str | None, Any]] = [] + sequence: list[tuple[str, Any]] = [] + for message in messages: + text.extend(block.text or "" for block in message.content if block.text and not block.reasoning) + reasoning.extend(block.text or "" for block in message_reasoning(message)) + calls.extend((call.id, call.name, call.arguments) for call in message_tool_calls(message)) + for block in ordered_message_blocks(message): + if block.reasoning: + sequence.append(("reasoning", block.reasoning.text)) + elif block.tool_call: + sequence.append(("tool_call", (block.tool_call.name, block.tool_call.arguments))) + elif block.text: + sequence.append(("text", block.text)) + return {"text": text, "reasoning": reasoning, "calls": calls, "sequence": sequence} + + +@pytest.mark.parametrize("source_protocol", PROTOCOLS) +@pytest.mark.parametrize("target_protocol", PROTOCOLS) +def test_request_semantics_survive_every_protocol_pair(source_protocol: str, target_protocol: str) -> None: + source = get_protocol(source_protocol) + target = get_protocol(target_protocol) + unified = source.parse_request(REQUEST_FIXTURES[source_protocol], _context(source_protocol, target_protocol)) + before = _request_semantics(unified) + + provider_payload = target.build_request(unified, _context(source_protocol, target_protocol)) + reparsed = target.parse_request(provider_payload, _context(target_protocol, target_protocol)) + after = _request_semantics(reparsed) + + assert after["instructions"] == before["instructions"] + assert after["visible_text"] == before["visible_text"] + assert after["reasoning_text"] == before["reasoning_text"] + assert after["calls"] == before["calls"] + assert after["results"] == before["results"] + assert after["media"] == before["media"] + assert after["tools"] == before["tools"] + assert after["max_output_tokens"] == before["max_output_tokens"] + assert after["temperature"] == before["temperature"] + assert after["tool_choice"] == before["tool_choice"] + assert after["structured_output"]["type"] == before["structured_output"]["type"] + assert after["structured_output"].get("schema") == before["structured_output"].get("schema") + assert after["sequence"] == before["sequence"] + if target_protocol != "responses": + assert after["stop_sequences"] == before["stop_sequences"] + if source_protocol != target_protocol: + assert "vendor_extension" not in provider_payload + + +@pytest.mark.parametrize("source_protocol", PROTOCOLS) +@pytest.mark.parametrize("target_protocol", PROTOCOLS) +def test_response_semantics_survive_every_output_protocol(source_protocol: str, target_protocol: str) -> None: + source = get_protocol(source_protocol) + target = get_protocol(target_protocol) + unified = source.parse_response(RESPONSE_FIXTURES[source_protocol], _context(source_protocol, target_protocol)) + before = _response_semantics(unified.messages) + + client_payload = target.format_response(unified, _context(source_protocol, target_protocol)) + reparsed = target.parse_response(client_payload, _context(target_protocol, target_protocol)) + after = _response_semantics(reparsed.messages) + + assert after["text"] == before["text"] + assert after["reasoning"] == before["reasoning"] + assert after["calls"] == before["calls"] + assert after["sequence"] == before["sequence"] + assert reparsed.stop_reason == unified.stop_reason + if source_protocol != target_protocol: + assert "vendor_extension" not in client_payload + + +def test_required_unknown_content_is_rejected_before_cross_protocol_transport() -> None: + source = get_protocol("openai_chat") + target = get_protocol("anthropic_messages") + request = source.parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": [{"type": "future_required_media", "value": "x"}]}], + } + ) + + with pytest.raises(ProtocolError, match="future_required_media"): + target.build_request(request, _context("openai_chat", "anthropic_messages")) + + +def test_same_protocol_unknown_content_is_preserved() -> None: + adapter = get_protocol("openai_chat") + raw = { + "model": "model-test", + "messages": [{"role": "user", "content": [{"type": "future_required_media", "value": "x"}]}], + } + + request = adapter.parse_request(raw) + assert adapter.build_request(request) == raw + + +def test_optional_unsupported_controls_are_dropped_with_conversion_warnings() -> None: + source = get_protocol("anthropic_messages") + target = get_protocol("openai_chat") + request = source.parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": "hello"}], + "top_k": 12, + } + ) + + payload = target.build_request(request, _context("anthropic_messages", "openai_chat")) + + assert "top_k" not in payload + assert [(warning.field, warning.target_protocol) for warning in request.warnings] == [("top_k", "openai_chat")] + + +def test_responses_stop_sequence_loss_is_explicit() -> None: + source = get_protocol("openai_chat") + target = get_protocol("responses") + request = source.parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": "hello"}], + "stop": ["END"], + } + ) + + payload = target.build_request(request, _context("openai_chat", "responses")) + + assert "stop" not in payload + assert any(warning.field == "stop_sequences" for warning in request.warnings) + + +def test_source_identity_prevents_foreign_passthrough_without_context() -> None: + source = get_protocol("openai_chat") + target = get_protocol("gemini") + request = source.parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": "hello", "chat_only": True}], + "vendor_extension": {"source_only": True}, + } + ) + + payload = target.build_request(request) + + assert "vendor_extension" not in payload + assert "chat_only" not in payload["contents"][0] + + +def test_opaque_reasoning_signatures_do_not_cross_protocol_families() -> None: + anthropic = get_protocol("anthropic_messages") + gemini = get_protocol("gemini") + anthropic_request = anthropic.parse_request(REQUEST_FIXTURES["anthropic_messages"]) + gemini_request = gemini.parse_request(REQUEST_FIXTURES["gemini"]) + + gemini_payload = gemini.build_request(anthropic_request, _context("anthropic_messages", "gemini")) + anthropic_payload = anthropic.build_request(gemini_request, _context("gemini", "anthropic_messages")) + + assert "thoughtSignature" not in str(gemini_payload) + assert "sig_anthropic" not in str(gemini_payload) + assert "signature" not in str(anthropic_payload) + assert "sig_gemini" not in str(anthropic_payload) + assert anthropic_request.messages[1].reasoning[0].signature == "sig_anthropic" + assert gemini_request.messages[1].reasoning[0].signature == "sig_gemini" + + +def test_unknown_system_content_is_validated() -> None: + source = get_protocol("anthropic_messages") + target = get_protocol("openai_chat") + request = source.parse_request( + { + "model": "model-test", + "system": [{"type": "future_required_instruction", "value": "x"}], + "messages": [{"role": "user", "content": "hello"}], + } + ) + + with pytest.raises(ProtocolError, match="future_required_instruction"): + target.build_request(request, _context("anthropic_messages", "openai_chat")) + + +def test_provider_bound_responses_continuation_is_not_silently_dropped() -> None: + source = get_protocol("responses") + target = get_protocol("anthropic_messages") + request = source.parse_request( + { + "model": "model-test", + "previous_response_id": "resp_provider_bound", + "input": "continue", + } + ) + + with pytest.raises(ProtocolError, match="previous_response_id"): + target.build_request(request, _context("responses", "anthropic_messages")) + + +def test_gemini_safety_policy_is_not_silently_dropped() -> None: + source = get_protocol("gemini") + target = get_protocol("openai_chat") + request = source.parse_request( + { + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "safetySettings": [{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_LOW_AND_ABOVE"}], + } + ) + + with pytest.raises(ProtocolError, match="safety settings"): + target.build_request(request, _context("gemini", "openai_chat")) + + +def test_missing_media_identity_is_rejected() -> None: + source = get_protocol("openai_chat") + target = get_protocol("gemini") + request = source.parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {}}]}], + } + ) + + with pytest.raises(ProtocolError, match="without URL, data, or file identity"): + target.build_request(request, _context("openai_chat", "gemini")) + + +def test_tool_argument_and_result_wire_types_are_destination_valid() -> None: + source = get_protocol("openai_chat") + request = source.parse_request(REQUEST_FIXTURES["openai_chat"]) + + anthropic_payload = get_protocol("anthropic_messages").build_request( + request, + _context("openai_chat", "anthropic_messages"), + ) + gemini_payload = get_protocol("gemini").build_request( + request, + _context("openai_chat", "gemini"), + ) + responses_payload = get_protocol("responses").build_request( + request, + _context("openai_chat", "responses"), + ) + + anthropic_call = next(block for message in anthropic_payload["messages"] for block in message["content"] if block["type"] == "tool_use") + anthropic_result = next(block for message in anthropic_payload["messages"] for block in message["content"] if block["type"] == "tool_result") + gemini_call = next(part for content in gemini_payload["contents"] for part in content["parts"] if "functionCall" in part) + gemini_result = next(part for content in gemini_payload["contents"] for part in content["parts"] if "functionResponse" in part) + responses_call = next(item for item in responses_payload["input"] if item["type"] == "function_call") + responses_result = next(item for item in responses_payload["input"] if item["type"] == "function_call_output") + + assert anthropic_call["input"] == {"q": "x"} + assert isinstance(anthropic_result["content"], str) + assert gemini_call["functionCall"]["args"] == {"q": "x"} + assert gemini_result["functionResponse"]["response"] == {"value": 1} + assert responses_call["arguments"] == '{"q":"x"}' + assert responses_result["output"] == '{"value":1}' + + +def test_structured_output_has_one_canonical_request_field_for_every_source() -> None: + for source_name in PROTOCOLS: + request = get_protocol(source_name).parse_request(REQUEST_FIXTURES[source_name]) + assert request.response_format == request.generation_params["structured_output"] + + +def test_unknown_native_stop_reason_never_leaks_as_invalid_chat_value() -> None: + response = get_protocol("gemini").parse_response( + { + "responseId": "resp_unknown", + "candidates": [{"finishReason": "OTHER", "content": {"role": "model", "parts": [{"text": "answer"}]}}], + } + ) + + payload = get_protocol("openai_chat").format_response(response, _context("gemini", "openai_chat")) + + assert payload["choices"][0]["finish_reason"] is None + assert response.metadata["native_stop_reason"] == "OTHER" + + +@pytest.mark.parametrize( + ("source_choice", "expected_mode"), + [ + ({"type": "function", "function": {"name": "lookup"}}, "named"), + ("required", "required"), + ("none", "none"), + ], +) +def test_tool_choice_modes_survive_all_destinations(source_choice: Any, expected_mode: str) -> None: + source = get_protocol("openai_chat") + request = source.parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": "hello"}], + "tools": REQUEST_FIXTURES["openai_chat"]["tools"], + "tool_choice": source_choice, + } + ) + + for target_name in PROTOCOLS: + target = get_protocol(target_name) + payload = target.build_request(request, _context("openai_chat", target_name)) + reparsed = target.parse_request(payload) + assert reparsed.generation_params["tool_choice"]["mode"] == expected_mode + + +def test_anthropic_system_preserves_instruction_boundaries_and_order() -> None: + request = get_protocol("openai_chat").parse_request( + { + "model": "model-test", + "messages": [ + {"role": "system", "content": "system A"}, + {"role": "developer", "content": "developer B"}, + {"role": "user", "content": "hello"}, + ], + } + ) + + payload = get_protocol("anthropic_messages").build_request( + request, + _context("openai_chat", "anthropic_messages"), + ) + + assert payload["system"] == [ + {"type": "text", "text": "system A"}, + {"type": "text", "text": "developer B"}, + ] + + +def test_opaque_state_requires_compatible_provider_domain_even_for_same_protocol() -> None: + adapter = get_protocol("anthropic_messages") + request = adapter.parse_request(REQUEST_FIXTURES["anthropic_messages"]) + incompatible = ProtocolContext( + source_protocol="anthropic_messages", + target_protocol="anthropic_messages", + source_provider="provider-a", + target_provider="provider-b", + ) + compatible = ProtocolContext( + source_protocol="anthropic_messages", + target_protocol="anthropic_messages", + source_provider="provider-a", + target_provider="provider-a", + ) + + incompatible_payload = adapter.build_request(request, incompatible) + compatible_payload = adapter.build_request(request, compatible) + unscoped_payload = adapter.build_request(request) + + assert "sig_anthropic" not in str(incompatible_payload) + assert "sig_anthropic" not in str(unscoped_payload) + assert "sig_anthropic" in str(compatible_payload) + + +def test_foreign_reasoning_extensions_never_enter_target_payload() -> None: + request = get_protocol("gemini").parse_request( + { + "contents": [ + { + "role": "model", + "parts": [ + { + "text": "thinking", + "thought": True, + "thoughtSignature": "sig_private", + "vendor_reasoning_extension": "must-not-leak", + } + ], + } + ] + } + ) + + payload = get_protocol("anthropic_messages").build_request( + request, + _context("gemini", "anthropic_messages"), + ) + + assert "vendor_reasoning_extension" not in str(payload) + assert "sig_private" not in str(payload) + + +def test_repeated_name_only_gemini_calls_get_distinct_cross_protocol_ids_without_same_protocol_leak() -> None: + adapter = get_protocol("gemini") + raw = { + "contents": [ + { + "role": "model", + "parts": [ + {"functionCall": {"name": "lookup", "args": {"q": "a"}}}, + {"functionCall": {"name": "lookup", "args": {"q": "b"}}}, + ], + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"name": "lookup", "response": {"value": "a"}}}, + {"functionResponse": {"name": "lookup", "response": {"value": "b"}}}, + ], + }, + ] + } + request = adapter.parse_request(raw) + + same_protocol = adapter.build_request(request) + chat_payload = get_protocol("openai_chat").build_request( + request, + _context("gemini", "openai_chat"), + ) + + assert "id" not in same_protocol["contents"][0]["parts"][0]["functionCall"] + assert "id" not in same_protocol["contents"][0]["parts"][1]["functionCall"] + assert [call["id"] for call in chat_payload["messages"][0]["tool_calls"]] == ["call_0", "call_1"] + assert [message["tool_call_id"] for message in chat_payload["messages"][1:]] == ["call_0", "call_1"] + + +def test_tool_result_error_state_survives_supported_protocol_conversions() -> None: + source = get_protocol("anthropic_messages") + request = source.parse_request( + { + "model": "model-test", + "messages": [ + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "lookup", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": [{"type": "text", "text": "failed"}], + "is_error": True, + } + ], + }, + ], + } + ) + + same_payload = source.build_request(request) + assert same_payload["messages"][1]["content"][0]["content"] == [{"type": "text", "text": "failed"}] + assert same_payload["messages"][1]["content"][0]["is_error"] is True + + chat_payload = get_protocol("openai_chat").build_request( + request, + _context("anthropic_messages", "openai_chat"), + ) + responses_payload = get_protocol("responses").build_request( + request, + _context("anthropic_messages", "responses"), + ) + gemini_payload = get_protocol("gemini").build_request( + request, + _context("anthropic_messages", "gemini"), + ) + + assert chat_payload["messages"][1]["content"] == '{"error":[{"type":"text","text":"failed"}]}' + assert responses_payload["input"][1]["output"] == '{"error":[{"type":"text","text":"failed"}]}' + assert gemini_payload["contents"][1]["parts"][0]["functionResponse"]["response"] == { + "error": [{"type": "text", "text": "failed"}] + } + + +@pytest.mark.parametrize("protocol_name", PROTOCOLS) +@pytest.mark.parametrize( + "result_content", + [ + {"error": None}, + {"error": "domain-status"}, + {"error": None, "data": "ok"}, + ], +) +def test_successful_tool_result_with_error_field_is_not_corrupted( + protocol_name: str, + result_content: dict[str, Any], +) -> None: + request = get_protocol("openai_chat").parse_request( + { + "model": "model-test", + "messages": [ + {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "call_1", "content": result_content}, + ], + } + ) + target = get_protocol(protocol_name) + payload = target.build_request(request, _context("openai_chat", protocol_name)) + reparsed = target.parse_request(payload) + results = [result for message in reparsed.messages for result in message_tool_results(message)] + + assert results[0].is_error is not True + assert results[0].content == result_content + + +def test_responses_same_protocol_function_arguments_remain_json_text() -> None: + adapter = get_protocol("responses") + response = adapter.parse_response(RESPONSE_FIXTURES["responses"]) + + payload = adapter.format_response(response) + call = next(item for item in payload["output"] if item["type"] == "function_call") + + assert call["arguments"] == '{"q":"x"}' + + +def test_required_response_modalities_map_or_reject() -> None: + source = get_protocol("openai_chat") + request = source.parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": "speak"}], + "modalities": ["text", "audio"], + "audio": {"voice": "alloy", "format": "wav"}, + } + ) + + gemini_payload = get_protocol("gemini").build_request(request, _context("openai_chat", "gemini")) + assert gemini_payload["generationConfig"]["responseModalities"] == ["TEXT", "AUDIO"] + + with pytest.raises(ProtocolError, match="response modalities"): + get_protocol("anthropic_messages").build_request( + request, + _context("openai_chat", "anthropic_messages"), + ) + + +def test_gemini_schema_enforcement_round_trips_as_strict() -> None: + request = get_protocol("openai_chat").parse_request(REQUEST_FIXTURES["openai_chat"]) + payload = get_protocol("gemini").build_request(request, _context("openai_chat", "gemini")) + reparsed = get_protocol("gemini").parse_request(payload) + + assert payload["generationConfig"]["responseJsonSchema"] == request.response_format["schema"] + assert reparsed.response_format["strict"] is True + + +def test_explicit_non_strict_schema_records_gemini_strengthening() -> None: + request = get_protocol("openai_chat").parse_request(REQUEST_FIXTURES["openai_chat"]) + request.response_format["strict"] = False + request.generation_params["structured_output"]["strict"] = False + + get_protocol("gemini").build_request(request, _context("openai_chat", "gemini")) + + assert any(warning.code == "structured_output_strictness_strengthened" for warning in request.warnings) + + +def test_failed_response_is_not_relabelled_as_success() -> None: + response = get_protocol("responses").parse_response( + {"id": "resp_failed", "object": "response", "status": "failed", "output": []} + ) + + for target_name in ("openai_chat", "anthropic_messages", "gemini"): + with pytest.raises(ProtocolError, match="failed provider response"): + get_protocol(target_name).format_response( + response, + _context("responses", target_name), + ) + assert get_protocol("responses").format_response(response)["status"] == "failed" + + +def test_named_tool_choice_requires_a_declared_name() -> None: + request = get_protocol("openai_chat").parse_request( + { + "model": "model-test", + "messages": [{"role": "user", "content": "hello"}], + "tools": [{"type": "function", "function": {"name": "lookup", "parameters": {"type": "object"}}}], + "tool_choice": {"type": "function", "function": {}}, + } + ) + + with pytest.raises(ProtocolError, match="requires a tool name"): + get_protocol("anthropic_messages").build_request( + request, + _context("openai_chat", "anthropic_messages"), + ) + + +def test_provider_bound_responses_controls_reject_cross_protocol_conversion() -> None: + request = get_protocol("responses").parse_request( + {"model": "model-test", "input": "hello", "background": True} + ) + + with pytest.raises(ProtocolError, match="Provider-bound Responses controls"): + get_protocol("openai_chat").build_request( + request, + _context("responses", "openai_chat"), + ) + + +def test_optional_responses_controls_warn_when_target_cannot_represent_them() -> None: + request = get_protocol("responses").parse_request( + {"model": "model-test", "input": "hello", "prompt_cache_key": "cache-key"} + ) + + payload = get_protocol("anthropic_messages").build_request( + request, + _context("responses", "anthropic_messages"), + ) + + assert "prompt_cache_key" not in payload + assert any(warning.field == "prompt_cache_key" for warning in request.warnings) diff --git a/tests/test_protocol_ollama_mcp.py b/tests/test_protocol_ollama_mcp.py new file mode 100644 index 000000000..7c859203b --- /dev/null +++ b/tests/test_protocol_ollama_mcp.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from rotator_library.protocols import OPERATION_EMBEDDINGS, OPERATION_MCP, OPERATION_OLLAMA_CHAT, OPERATION_OLLAMA_GENERATE, ProtocolContext, get_protocol + + +def test_ollama_chat_generate_and_stream_shapes() -> None: + adapter = get_protocol("ollama") + chat = adapter.parse_request({"model": "llama3", "messages": [{"role": "user", "content": "hi"}], "stream": True}) + generate = adapter.parse_request({"model": "llama3", "prompt": "write", "options": {"temperature": 0.1}}) + explicit_stream_false = adapter.parse_request({"model": "llama3", "prompt": "write", "stream": False}) + embeddings = adapter.parse_request({"model": "llama3", "operation": OPERATION_EMBEDDINGS, "prompt": "embed this"}) + contextual_embeddings = adapter.parse_request( + {"model": "llama3", "prompt": "embed via context"}, + ProtocolContext(provider_options={"operation": OPERATION_EMBEDDINGS}), + ) + new_embeddings = adapter.parse_request({"model": "llama3", "input": "embed input"}) + final_chat = adapter.parse_response({"model": "llama3", "message": {"role": "assistant", "content": "hello"}, "done": True}) + final_generate = adapter.parse_response({"model": "llama3", "response": "generated", "done": True}) + embeddings_response = adapter.parse_response({"model": "llama3", "embeddings": [[0.1, 0.2]]}) + chunk = adapter.parse_stream_event('{"model":"llama3","response":"he","done":false,"eval_count":2}') + + assert chat.operation == OPERATION_OLLAMA_CHAT + assert adapter.build_request(chat)["messages"][0]["content"] == "hi" + assert generate.operation == OPERATION_OLLAMA_GENERATE + assert adapter.build_request(generate)["prompt"] == "write" + assert "stream" not in adapter.build_request(generate) + assert adapter.build_request(explicit_stream_false)["stream"] is False + assert embeddings.operation == OPERATION_EMBEDDINGS + assert adapter.build_request(embeddings)["prompt"] == "embed this" + assert contextual_embeddings.operation == OPERATION_EMBEDDINGS + assert adapter.build_request(contextual_embeddings)["prompt"] == "embed via context" + assert adapter.build_request(new_embeddings)["input"] == "embed input" + assert final_chat.messages[0].content[0].text == "hello" + assert final_generate.output == ["generated"] + assert final_generate.messages == [] + assert embeddings_response.data == [[0.1, 0.2]] + assert chunk.delta is not None + assert chunk.delta.content[0].text == "he" + assert chunk.usage is not None + assert chunk.usage.output_tokens == 2 + + +def test_ollama_format_response_uses_mutated_unified_fields() -> None: + adapter = get_protocol("ollama") + chat = adapter.parse_response({"model": "llama3", "message": {"role": "assistant", "content": "before"}, "done": True}) + chat.messages[0].content[0].text = "after" + + generate = adapter.parse_response({"model": "llama3", "response": "before", "eval_count": 2}) + generate.output[0] = "after" + generate.usage.output_tokens = 5 + + embeddings = adapter.parse_response({"model": "llama3", "embedding": [0.1, 0.2]}) + embeddings.data = [0.3, 0.4] + + assert adapter.format_response(chat)["message"]["content"] == "after" + assert adapter.format_response(generate)["response"] == "after" + assert adapter.format_response(generate)["eval_count"] == 5 + assert adapter.format_response(embeddings)["embedding"] == [0.3, 0.4] + + +def test_mcp_jsonrpc_round_trip_and_error_preservation() -> None: + adapter = get_protocol("mcp") + request = {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "lookup"}} + unified = adapter.parse_request(request) + rebuilt = adapter.build_request(unified) + response = adapter.parse_response({"jsonrpc": "2.0", "id": 1, "result": {"content": []}}) + error = adapter.parse_response({"jsonrpc": "2.0", "id": 1, "error": {"code": -1, "message": "failed"}}) + + assert unified.operation == OPERATION_MCP + assert rebuilt == request + assert adapter.format_response(response)["result"] == {"content": []} + assert adapter.format_response(error)["error"]["message"] == "failed" + + +def test_mcp_preserves_notifications_and_falsey_params() -> None: + adapter = get_protocol("mcp") + notification = {"jsonrpc": "2.0", "method": "notifications/initialized"} + false_params = {"jsonrpc": "2.0", "id": 0, "method": "tools/call", "params": False} + + assert adapter.build_request(adapter.parse_request(notification)) == notification + assert adapter.build_request(adapter.parse_request(false_params)) == false_params + + +def test_mcp_preserves_jsonrpc_batches() -> None: + adapter = get_protocol("mcp") + batch = [ + {"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}}, + {"jsonrpc": "2.0", "id": 2, "method": "resources/list", "params": None}, + ] + response_batch = [ + {"jsonrpc": "2.0", "id": 1, "result": {"tools": []}}, + {"jsonrpc": "2.0", "id": 2, "error": {"code": -1, "message": "failed"}}, + ] + + assert adapter.build_request(adapter.parse_request(batch)) == batch + assert adapter.format_response(adapter.parse_response(response_batch)) == response_batch diff --git a/tests/test_protocol_openai_chat.py b/tests/test_protocol_openai_chat.py new file mode 100644 index 000000000..7bbb5dfab --- /dev/null +++ b/tests/test_protocol_openai_chat.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +import json + +from rotator_library.protocols import get_protocol, list_protocols + + +def test_openai_chat_protocol_is_discovered_with_aliases() -> None: + assert "openai_chat" in list_protocols() + assert get_protocol("openai") is get_protocol("openai_chat") + assert get_protocol("chat_completions") is get_protocol("openai_chat") + + +def test_litellm_fallback_protocol_preserves_raw_payload() -> None: + adapter = get_protocol("litellm_fallback") + raw = {"model": "custom/model", "messages": [{"role": "user", "content": "hi"}], "vendor_flag": True} + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified) + + assert unified.extra["messages"] == raw["messages"] + assert rebuilt == raw + + +def test_openai_chat_request_round_trip_preserves_tools_and_reasoning() -> None: + adapter = get_protocol("openai_chat") + raw = { + "model": "openai/gpt-test", + "stream": True, + "temperature": 0.2, + "messages": [ + {"role": "system", "content": "system"}, + {"role": "developer", "content": "dev"}, + {"role": "user", "content": [{"type": "text", "text": "hello"}]}, + { + "role": "assistant", + "content": "thinking done", + "reasoning_content": "internal chain", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{\"q\":\"x\"}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "result"}, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Lookup a value", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + } + ], + "vendor_extension": {"kept": True}, + } + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified) + + assert unified.model == "openai/gpt-test" + assert unified.stream is True + assert unified.generation_params["temperature"] == 0.2 + assert unified.messages[3].reasoning[0].text == "internal chain" + assert unified.messages[3].tool_calls[0].name == "lookup" + assert unified.tools[0].input_schema["properties"]["q"]["type"] == "string" + assert rebuilt["vendor_extension"] == {"kept": True} + assert rebuilt["messages"][3]["reasoning_content"] == "internal chain" + + +def test_openai_chat_build_uses_mutated_unified_fields_and_preserves_block_extras() -> None: + adapter = get_protocol("openai_chat") + raw = { + "model": "openai/gpt-test", + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "before", "cache_control": {"type": "ephemeral"}}], + } + ], + } + + unified = adapter.parse_request(raw) + unified.messages[0].content[0].text = "after" + rebuilt = adapter.build_request(unified) + + assert isinstance(rebuilt["messages"][0]["content"], list) + assert rebuilt["messages"][0]["content"][0]["text"] == "after" + assert rebuilt["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_openai_chat_response_extracts_usage_cost_and_reasoning() -> None: + adapter = get_protocol("openai_chat") + raw = { + "id": "chatcmpl_1", + "object": "chat.completion", + "created": 123, + "model": "openai/gpt-test", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "answer", + "reasoning_content": "reasoned", + }, + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 18, + "prompt_tokens_details": {"cached_tokens": 3, "cache_creation_tokens": 2}, + "completion_tokens_details": {"reasoning_tokens": 3}, + "cost_details": {"total_cost": 0.01, "currency": "USD"}, + }, + } + + unified = adapter.parse_response(raw) + + assert unified.id == "chatcmpl_1" + assert unified.messages[0].content[0].text == "answer" + assert unified.messages[0].reasoning[0].text == "reasoned" + assert unified.usage is not None + assert unified.usage.input_tokens == 10 + assert unified.usage.output_tokens == 5 + assert unified.usage.reasoning_tokens == 3 + assert unified.usage.cache_read_tokens == 3 + assert unified.usage.cache_write_tokens == 2 + assert unified.usage.cost is not None + assert unified.usage.cost.provider_reported_cost == 0.01 + + formatted = adapter.format_response(unified) + assert formatted["usage"]["prompt_tokens"] == 10 + assert formatted["usage"]["completion_tokens"] == 5 + assert formatted["usage"]["total_tokens"] == 18 + assert formatted["usage"]["prompt_tokens_details"] == {"cached_tokens": 3, "cache_creation_tokens": 2} + assert formatted["usage"]["completion_tokens_details"] == {"reasoning_tokens": 3} + assert formatted["usage"]["cost_details"]["total_cost"] == 0.01 + assert "input_tokens" not in formatted["usage"] + assert "output_tokens" not in formatted["usage"] + + +def test_openai_legacy_function_call_is_unified_and_round_trips() -> None: + adapter = get_protocol("openai_chat") + raw = { + "model": "openai/gpt-test", + "messages": [ + { + "role": "assistant", + "content": None, + "function_call": {"name": "lookup", "arguments": "{\"q\":\"x\"}"}, + } + ], + } + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified) + + assert unified.messages[0].tool_calls[0].name == "lookup" + assert unified.messages[0].tool_calls[0].extra["legacy_function_call"] is True + unified.messages[0].tool_calls[0].arguments = {"q": "changed"} + assert rebuilt["messages"][0]["function_call"] == {"name": "lookup", "arguments": "{\"q\":\"x\"}"} + assert "tool_calls" not in rebuilt["messages"][0] + + rebuilt = adapter.build_request(unified) + assert rebuilt["messages"][0]["function_call"] == {"name": "lookup", "arguments": "{\"q\":\"changed\"}"} + + +def test_openai_chat_stream_event_parses_sse_delta_and_done() -> None: + adapter = get_protocol("openai_chat") + event = { + "id": "chunk_1", + "model": "openai/gpt-test", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}, "finish_reason": None}], + } + + parsed = adapter.parse_stream_event(f"data: {json.dumps(event)}\n\n") + done = adapter.parse_stream_event("data: [DONE]\n\n") + + assert parsed.type == "message_delta" + assert parsed.delta is not None + assert parsed.delta.content[0].text == "Hel" + assert done.type == "done" diff --git a/tests/test_protocol_openai_embeddings.py b/tests/test_protocol_openai_embeddings.py new file mode 100644 index 000000000..f398f7a33 --- /dev/null +++ b/tests/test_protocol_openai_embeddings.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from rotator_library.protocols import OPERATION_EMBEDDINGS, get_protocol, list_protocols + + +def test_openai_embeddings_protocol_round_trip_and_usage() -> None: + adapter = get_protocol("openai_embeddings") + raw = {"model": "text-embedding-test", "input": ["one", "two"], "dimensions": 128, "custom": True} + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified) + response = adapter.parse_response( + { + "model": "text-embedding-test", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + } + ) + + assert "openai_embeddings" in list_protocols() + assert get_protocol("embeddings") is adapter + assert adapter.supports_operation(OPERATION_EMBEDDINGS) + assert unified.operation == OPERATION_EMBEDDINGS + assert unified.input == ["one", "two"] + assert rebuilt == raw + assert response.data[0]["embedding"] == [0.1, 0.2] + assert response.usage is not None + assert response.usage.input_tokens == 4 + assert adapter.format_response(response)["data"] == response.data + assert adapter.format_response(response)["usage"]["total_tokens"] == 4 diff --git a/tests/test_protocol_openai_images_audio.py b/tests/test_protocol_openai_images_audio.py new file mode 100644 index 000000000..4fd0cd3ce --- /dev/null +++ b/tests/test_protocol_openai_images_audio.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from rotator_library.protocols import ( + OPERATION_AUDIO_TRANSCRIPTION, + OPERATION_AUDIO_TRANSLATION, + OPERATION_IMAGE_EDIT, + OPERATION_IMAGE_GENERATION, + OPERATION_IMAGE_VARIATION, + OPERATION_SPEECH, + ProtocolContext, + get_protocol, +) + + +def test_openai_images_generation_and_edit_shapes() -> None: + adapter = get_protocol("openai_images") + generation = adapter.parse_request({"model": "gpt-image-test", "prompt": "draw a red cube", "size": "1024x1024"}) + edit = adapter.parse_request({"model": "gpt-image-test", "prompt": "make it blue", "image": "image-ref", "mask": "mask-ref"}) + variation = adapter.parse_request({"model": "gpt-image-test", "image": "image-ref"}) + response = adapter.parse_response({"data": [{"url": "https://example.test/image.png", "revised_prompt": "cube"}]}) + edit_response = adapter.parse_response( + {"data": [{"url": "https://example.test/edit.png"}]}, + ProtocolContext(provider_options={"operation": OPERATION_IMAGE_EDIT}), + ) + + assert generation.operation == OPERATION_IMAGE_GENERATION + assert adapter.build_request(generation)["prompt"] == "draw a red cube" + assert edit.operation == OPERATION_IMAGE_EDIT + assert variation.operation == OPERATION_IMAGE_VARIATION + assert {entry["field"] for entry in edit.files} == {"image", "mask"} + assert adapter.build_request(edit)["image"] == "image-ref" + assert response.data[0]["revised_prompt"] == "cube" + assert edit_response.operation == OPERATION_IMAGE_EDIT + + +def test_openai_audio_transcription_and_speech_shapes() -> None: + adapter = get_protocol("openai_audio") + transcription = adapter.parse_request({"model": "whisper-test", "file": "audio-ref", "language": "en"}) + translation = adapter.parse_request( + {"model": "whisper-test", "file": "audio-ref", "prompt": "domain hint"}, + ProtocolContext(provider_options={"operation": OPERATION_AUDIO_TRANSLATION}), + ) + speech = adapter.parse_request({"model": "tts-test", "input": "hello", "voice": "alloy"}) + text_response = adapter.parse_response({"text": "hello world"}) + translation_response = adapter.parse_response( + {"text": "bonjour"}, + ProtocolContext(provider_options={"operation": "audio_translation"}), + ) + binary_response = adapter.parse_response(b"RIFF") + + assert transcription.operation == OPERATION_AUDIO_TRANSCRIPTION + assert transcription.files[0]["value"] == "audio-ref" + assert translation.operation == OPERATION_AUDIO_TRANSLATION + translation.input = "mutated hint" + assert adapter.build_request(translation)["prompt"] == "mutated hint" + assert speech.operation == OPERATION_SPEECH + assert adapter.build_request(speech)["voice"] == "alloy" + assert text_response.operation == OPERATION_AUDIO_TRANSCRIPTION + assert text_response.output == ["hello world"] + assert adapter.format_response(text_response)["text"] == "hello world" + assert translation_response.operation == "audio_translation" + assert binary_response.content_type == "application/octet-stream" + assert adapter.format_response(binary_response) == b"RIFF" diff --git a/tests/test_protocol_operation_model.py b/tests/test_protocol_operation_model.py new file mode 100644 index 000000000..c2c5d0e93 --- /dev/null +++ b/tests/test_protocol_operation_model.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from rotator_library.protocols import ( + OPERATION_CHAT, + OPERATION_COUNT_TOKENS, + OPERATION_EMBEDDINGS, + OPERATION_MESSAGES, + OPERATION_RESPONSES, + OPERATION_UNKNOWN, + ProtocolAdapter, + ProtocolContext, + UnifiedRequest, + UnifiedResponse, + get_protocol, + normalize_operation, +) + + +def test_operation_names_are_extensible_strings() -> None: + assert normalize_operation(" Chat ") == OPERATION_CHAT + assert normalize_operation(None) == OPERATION_UNKNOWN + assert normalize_operation("custom_gateway_op") == "custom_gateway_op" + + +def test_unified_request_response_carry_non_chat_operation_fields() -> None: + request = UnifiedRequest( + operation=OPERATION_EMBEDDINGS, + model="text-embedding-test", + input=["a", "b"], + modalities=["text"], + files=[{"name": "audio.wav", "content_type": "audio/wav"}], + ) + response = UnifiedResponse( + operation=OPERATION_EMBEDDINGS, + model="text-embedding-test", + data=[{"embedding": [0.1, 0.2]}], + content_type="application/json", + ) + + assert request.to_dict()["operation"] == OPERATION_EMBEDDINGS + assert request.to_dict()["input"] == ["a", "b"] + assert response.to_dict()["data"][0]["embedding"] == [0.1, 0.2] + assert response.to_dict()["content_type"] == "application/json" + + +def test_base_adapter_preserves_operation_fields() -> None: + adapter = ProtocolAdapter() + raw = { + "operation": "custom_op", + "model": "provider/model", + "input": "payload", + "modalities": ["text"], + "files": [{"name": "file.bin"}], + "custom": True, + } + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified) + + assert unified.operation == "custom_op" + assert unified.input == "payload" + assert rebuilt == raw + + +def test_protocols_advertise_supported_operations() -> None: + assert get_protocol("openai_chat").supports_operation(OPERATION_CHAT) + assert not get_protocol("openai_chat").supports_operation(OPERATION_EMBEDDINGS) + assert get_protocol("litellm_fallback").supports_operation(OPERATION_UNKNOWN) + + +def test_core_protocols_stamp_parsed_operation_fields() -> None: + chat = get_protocol("openai_chat").parse_request({"model": "m", "messages": []}) + messages = get_protocol("anthropic_messages").parse_request({"model": "m", "messages": []}) + responses = get_protocol("responses").parse_request({"model": "m", "input": "hi"}) + gemini = get_protocol("gemini").parse_request({"contents": []}) + + assert chat.operation == OPERATION_CHAT + assert messages.operation == OPERATION_MESSAGES + assert responses.operation == OPERATION_RESPONSES + assert gemini.operation == OPERATION_CHAT + + +def test_count_tokens_operation_can_be_context_selected() -> None: + context = ProtocolContext(provider_options={"operation": OPERATION_COUNT_TOKENS}) + anthropic = get_protocol("anthropic_messages").parse_response({"input_tokens": 12}, context) + gemini = get_protocol("gemini").parse_response({"totalTokens": 14}, context) + + assert anthropic.operation == OPERATION_COUNT_TOKENS + assert anthropic.usage is not None + assert anthropic.usage.input_tokens == 12 + assert get_protocol("anthropic_messages").format_response(anthropic) == {"input_tokens": 12} + assert gemini.operation == OPERATION_COUNT_TOKENS + assert gemini.usage is not None + assert gemini.usage.total_tokens == 14 + assert get_protocol("gemini").format_response(gemini) == {"totalTokens": 14} + + anthropic.usage.input_tokens = 13 + gemini.usage.total_tokens = 15 + assert get_protocol("anthropic_messages").format_response(anthropic)["input_tokens"] == 13 + assert get_protocol("gemini").format_response(gemini)["totalTokens"] == 15 + + +def test_context_operation_helpers_reject_unsupported_operations() -> None: + context = ProtocolContext(provider_options={"operation": OPERATION_EMBEDDINGS}) + + anthropic = get_protocol("anthropic_messages").parse_request({"model": "m", "messages": []}, context) + gemini = get_protocol("gemini").parse_request({"contents": []}, context) + + assert anthropic.operation == OPERATION_MESSAGES + assert gemini.operation == OPERATION_CHAT diff --git a/tests/test_protocol_registry.py b/tests/test_protocol_registry.py new file mode 100644 index 000000000..0b449eaf6 --- /dev/null +++ b/tests/test_protocol_registry.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import json +from datetime import datetime +from decimal import Decimal + +import pytest + +from rotator_library.protocols import ( + ContentBlock, + ProtocolAdapter, + ProtocolContext, + ProtocolError, + ToolCall, + UnifiedMessage, + UnifiedRequest, + get_protocol, + get_protocol_class, + list_protocols, + register_protocol, + serialize_value, +) + + +class ExampleProtocol(ProtocolAdapter): + name = "example_test_protocol" + aliases = ("example_test_alias",) + + +def test_protocol_serialization_preserves_nested_values() -> None: + request = UnifiedRequest( + model="provider/model", + messages=[ + UnifiedMessage( + role="assistant", + content=[ContentBlock(type="text", text="hello")], + tool_calls=[ToolCall(id="call_1", name="lookup", arguments={"q": "x"})], + extra={"provider_field": {"kept": True}}, + ) + ], + metadata={"request_id": "req_1"}, + ) + + serialized = request.to_dict() + + assert serialized["messages"][0]["content"][0]["text"] == "hello" + assert serialized["messages"][0]["tool_calls"][0]["arguments"] == {"q": "x"} + assert serialized["messages"][0]["extra"]["provider_field"] == {"kept": True} + json.dumps(serialize_value(request)) + + +def test_base_protocol_preserves_unknown_request_fields() -> None: + adapter = ProtocolAdapter() + raw = {"model": "provider/model", "stream": True, "custom": {"value": 1}} + + unified = adapter.parse_request(raw, ProtocolContext(provider="provider")) + rebuilt = adapter.build_request(unified) + + assert unified.model == "provider/model" + assert unified.stream is True + assert unified.extra == {"custom": {"value": 1}} + assert rebuilt == raw + + +def test_register_protocol_resolves_alias_and_reuses_instance() -> None: + register_protocol(ExampleProtocol, replace=True) + + assert get_protocol_class("example_test_protocol") is ExampleProtocol + assert get_protocol_class("example_test_alias") is ExampleProtocol + assert get_protocol("example_test_protocol") is get_protocol("example_test_alias") + assert "example_test_protocol" in list_protocols() + + +def test_register_protocol_rejects_duplicate_alias() -> None: + class FirstProtocol(ProtocolAdapter): + name = "duplicate_alias_first" + aliases = ("duplicate_alias",) + + class SecondProtocol(ProtocolAdapter): + name = "duplicate_alias_second" + aliases = ("duplicate_alias",) + + register_protocol(FirstProtocol, replace=True) + + with pytest.raises(ValueError, match="alias already registered"): + register_protocol(SecondProtocol) + + +def test_register_protocol_rejects_name_that_conflicts_with_existing_alias() -> None: + class AliasOwnerProtocol(ProtocolAdapter): + name = "alias_owner_protocol" + aliases = ("reserved_protocol_name",) + + class ConflictingNameProtocol(ProtocolAdapter): + name = "reserved_protocol_name" + + register_protocol(AliasOwnerProtocol, replace=True) + + with pytest.raises(ValueError, match="name conflicts with registered alias"): + register_protocol(ConflictingNameProtocol) + + +def test_protocol_serialization_handles_common_non_json_values() -> None: + value = serialize_value( + { + "bytes": b"hello", + "decimal": Decimal("1.25"), + "datetime": datetime(2026, 1, 2, 3, 4, 5), + "object": object(), + } + ) + + assert value["bytes"] == "hello" + assert value["decimal"] == 1.25 + assert value["datetime"] == "2026-01-02T03:04:05" + json.dumps(value) + + +def test_usage_total_does_not_double_count_reasoning_by_default() -> None: + from rotator_library.protocols import Usage + + usage = Usage(input_tokens=10, output_tokens=5, reasoning_tokens=3) + + assert usage.total_tokens == 15 + + +def test_protocol_error_includes_pass_and_payload_preview() -> None: + error = ProtocolError( + "failed", + protocol="example", + pass_name="parse_request", + payload={"secret": "not-redacted-here", "value": 1}, + ) + + assert "example.parse_request" in str(error) + assert "value" in str(error) diff --git a/tests/test_protocol_responses.py b/tests/test_protocol_responses.py new file mode 100644 index 000000000..77476adf3 --- /dev/null +++ b/tests/test_protocol_responses.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import json + +from rotator_library.protocols import get_protocol, list_protocols + + +def test_responses_protocol_is_discovered_with_aliases_and_websocket_support() -> None: + adapter = get_protocol("responses") + + assert "responses" in list_protocols() + assert get_protocol("openai_responses") is adapter + assert adapter.supports_transport("websocket") is False + assert adapter.is_future_transport("websocket") is True + + +def test_responses_request_round_trip_preserves_previous_response_and_tools() -> None: + adapter = get_protocol("responses") + raw = { + "model": "openai/gpt-test", + "instructions": "system", + "previous_response_id": "resp_prev", + "stream": True, + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hello"}]}, + {"type": "function_call_output", "call_id": "call_1", "output": "result"}, + ], + "tools": [{"type": "function", "name": "lookup", "description": "Lookup", "parameters": {"type": "object"}}], + "reasoning": {"effort": "medium"}, + "vendor_extension": {"kept": True}, + } + + unified = adapter.parse_request(raw) + rebuilt = adapter.build_request(unified) + + assert unified.previous_response_id == "resp_prev" + assert unified.system[0].text == "system" + assert unified.messages[0].content[0].text == "hello" + assert unified.messages[1].tool_call_id == "call_1" + assert unified.tools[0].name == "lookup" + assert rebuilt["previous_response_id"] == "resp_prev" + assert rebuilt["reasoning"] == {"effort": "medium"} + assert rebuilt["vendor_extension"] == {"kept": True} + + +def test_responses_response_extracts_output_items_reasoning_calls_and_usage() -> None: + adapter = get_protocol("responses") + raw = { + "id": "resp_1", + "object": "response", + "created_at": 123, + "model": "openai/gpt-test", + "status": "completed", + "output": [ + {"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "answer"}]}, + {"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "reasoned"}]}, + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{\"q\":\"x\"}"}, + ], + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 18, + "input_tokens_details": {"cached_tokens": 3, "cache_creation_tokens": 2}, + "output_tokens_details": {"reasoning_tokens": 3}, + "cost_details": {"total_cost": 0.02, "currency": "USD"}, + }, + } + + unified = adapter.parse_response(raw) + + assert unified.id == "resp_1" + assert len(unified.output) == 3 + assert unified.messages[0].content[0].text == "answer" + assert unified.messages[1].reasoning[0].text == "reasoned" + assert unified.messages[2].tool_calls[0].name == "lookup" + assert unified.usage is not None + assert unified.usage.input_tokens == 10 + assert unified.usage.reasoning_tokens == 3 + assert unified.usage.cache_read_tokens == 3 + assert unified.usage.cache_write_tokens == 2 + assert unified.usage.cost is not None + assert unified.usage.cost.provider_reported_cost == 0.02 + + formatted = adapter.format_response(unified) + assert formatted["usage"]["input_tokens"] == 10 + assert formatted["usage"]["output_tokens"] == 5 + assert formatted["usage"]["total_tokens"] == 18 + assert formatted["usage"]["input_tokens_details"] == {"cached_tokens": 3, "cache_creation_tokens": 2} + assert formatted["usage"]["output_tokens_details"] == {"reasoning_tokens": 3} + assert formatted["usage"]["cost_details"]["total_cost"] == 0.02 + assert "raw" not in formatted["usage"] + assert "extra" not in formatted["usage"] + + +def test_responses_format_preserves_unknown_output_items() -> None: + adapter = get_protocol("responses") + raw = { + "id": "resp_1", + "model": "openai/gpt-test", + "output": [ + {"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "before"}]}, + {"id": "future_1", "type": "future_item", "custom": {"kept": True}}, + ], + } + + unified = adapter.parse_response(raw) + unified.messages[0].content[0].text = "after" + rebuilt = adapter.format_response(unified) + + assert rebuilt["output"][0]["content"][0]["text"] == "after" + assert rebuilt["output"][1] == {"id": "future_1", "type": "future_item", "custom": {"kept": True}} + + +def test_responses_stream_event_parses_text_delta_and_completed_response() -> None: + adapter = get_protocol("responses") + delta = {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "Hi"} + completed = { + "type": "response.completed", + "response": { + "id": "resp_1", + "model": "openai/gpt-test", + "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Hi"}]}], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } + + parsed_delta = adapter.parse_stream_event(f"data: {json.dumps(delta)}\n\n") + parsed_completed = adapter.parse_stream_event(completed) + + assert parsed_delta.type == "message_delta" + assert parsed_delta.delta is not None + assert parsed_delta.delta.content[0].text == "Hi" + assert parsed_completed.type == "response.completed" + assert parsed_completed.usage is not None + assert parsed_completed.usage.total_tokens == 2 diff --git a/tests/test_protocol_streaming_matrix.py b/tests/test_protocol_streaming_matrix.py new file mode 100644 index 000000000..65cff803a --- /dev/null +++ b/tests/test_protocol_streaming_matrix.py @@ -0,0 +1,635 @@ +from __future__ import annotations + +import json +from types import SimpleNamespace + +import httpx +import pytest + +from rotator_library.client.executor import RequestExecutor +from rotator_library.client.filters import CredentialFilter +from rotator_library.client.transforms import ProviderTransforms +from rotator_library.core.types import RequestContext +from rotator_library.native_provider import NativeProviderContext, NativeProviderExecutor +from rotator_library.providers.provider_interface import ProviderInterface +from rotator_library.protocols import ProtocolContext, get_protocol +from rotator_library.protocols.streaming import ProtocolStreamConverter + + +PROTOCOLS = ("openai_chat", "anthropic_messages", "responses", "gemini") + + +def _source_frames(protocol: str) -> list[object]: + if protocol == "openai_chat": + return [ + {"id": "chat_1", "model": "model-a", "choices": [{"delta": {"reasoning_content": "think"}}]}, + {"id": "chat_1", "model": "model-a", "choices": [{"delta": {"content": "hi"}}]}, + {"id": "chat_1", "model": "model-a", "choices": [{"delta": {"tool_calls": [{"index": 0, "id": "call_1", "type": "function", "function": {"name": "weather", "arguments": '{"city":'}}]}}]}, + {"id": "chat_1", "model": "model-a", "choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '"Paris"}'}}]}, "finish_reason": "tool_calls"}], "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}}, + "[DONE]", + ] + if protocol == "anthropic_messages": + return [ + {"type": "message_start", "message": {"id": "msg_1", "role": "assistant", "content": [], "model": "model-a", "usage": {"input_tokens": 2, "output_tokens": 0}}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "think"}}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "hi"}}, + {"type": "content_block_start", "index": 2, "content_block": {"type": "tool_use", "id": "call_1", "name": "weather", "input": {}}}, + {"type": "content_block_delta", "index": 2, "delta": {"type": "input_json_delta", "partial_json": '{"city":"Paris"}'}}, + {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 3}}, + {"type": "message_stop"}, + ] + if protocol == "responses": + return [ + {"type": "response.created", "response": {"id": "resp_1", "status": "in_progress", "model": "model-a", "output": []}}, + {"type": "response.output_item.added", "output_index": 0, "item": {"id": "rs_0", "type": "reasoning", "summary": [{"type": "summary_text", "text": "think"}]}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 1, "content_index": 0, "delta": "hi"}, + {"type": "response.output_item.added", "output_index": 2, "item": {"id": "fc_2", "type": "function_call", "call_id": "call_1", "name": "weather", "arguments": ""}}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_2", "call_id": "call_1", "output_index": 2, "delta": '{"city":"Paris"}'}, + {"type": "response.completed", "response": {"id": "resp_1", "status": "completed", "model": "model-a", "output": [], "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5}}}, + ] + return [ + {"responseId": "gem_1", "modelVersion": "model-a", "candidates": [{"content": {"role": "model", "parts": [{"text": "think", "thought": True}, {"text": "hi"}, {"functionCall": {"id": "call_1", "name": "weather", "args": {"city": "Paris"}}}]}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 3, "totalTokenCount": 5}}, + "[DONE]", + ] + + +@pytest.mark.parametrize("source_protocol", PROTOCOLS) +@pytest.mark.parametrize("output_protocol", PROTOCOLS) +def test_streaming_matrix_preserves_text_reasoning_tools_usage_and_terminal_lifecycle( + source_protocol: str, + output_protocol: str, +) -> None: + context = ProtocolContext( + model="model-a", + source_protocol=source_protocol, + target_protocol=output_protocol, + input_protocol=source_protocol, + provider_protocol=source_protocol, + output_protocol=output_protocol, + transport="sse", + ) + converter = ProtocolStreamConverter( + get_protocol(source_protocol), + get_protocol(output_protocol), + context, + ) + + frames = [formatted for raw in _source_frames(source_protocol) for formatted in converter.convert(raw)] + output = "".join(str(frame) for frame in frames) + + assert "hi" in output + assert "think" in output + assert "weather" in output + assert "Paris" in output + if output_protocol == "openai_chat": + assert "chat.completion.chunk" in output + assert "reasoning_content" in output + assert 'data: [DONE]' in output + elif output_protocol == "anthropic_messages": + assert "event: message_start" in output + assert "thinking_delta" in output + assert "input_json_delta" in output + assert "event: message_stop" in output + elif output_protocol == "responses": + assert "event: response.created" in output + assert "response.reasoning_summary_text.delta" in output + assert "response.function_call_arguments.delta" in output + assert "event: response.completed" in output + else: + assert "modelVersion" in output + assert "functionCall" in output + assert '"thought": true' in output + assert "finishReason" in output + + +@pytest.mark.parametrize("output_protocol", PROTOCOLS) +def test_streaming_errors_are_formatted_in_selected_protocol(output_protocol: str) -> None: + context = ProtocolContext(model="model-a", source_protocol="openai_chat", target_protocol=output_protocol, transport="sse") + converter = ProtocolStreamConverter(get_protocol("openai_chat"), get_protocol(output_protocol), context) + + output = "".join(converter.convert({"error": {"type": "rate_limit", "message": "slow down"}})) + + assert "slow down" in output + if output_protocol == "anthropic_messages": + assert "event: error" in output + elif output_protocol == "responses": + assert "event: response.failed" in output + else: + assert '"error"' in output + + +@pytest.mark.parametrize("output_protocol", ("openai_chat", "anthropic_messages", "gemini")) +def test_responses_failed_stream_is_an_error_in_every_alternate_output(output_protocol: str) -> None: + context = ProtocolContext(model="model-a", source_protocol="responses", target_protocol=output_protocol, transport="sse") + converter = ProtocolStreamConverter(get_protocol("responses"), get_protocol(output_protocol), context) + + output = "".join(converter.convert({ + "type": "response.failed", + "response": { + "id": "resp_failed", + "status": "failed", + "model": "model-a", + "output": [], + "error": {"type": "rate_limit", "message": "try later"}, + }, + })) + + assert "try later" in output + assert "response.completed" not in output + + +def test_gemini_rejects_incomplete_foreign_tool_arguments_at_terminal() -> None: + context = ProtocolContext(model="model-a", source_protocol="openai_chat", target_protocol="gemini", transport="sse") + converter = ProtocolStreamConverter(get_protocol("openai_chat"), get_protocol("gemini"), context) + converter.convert({"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "call_1", "type": "function", "function": {"name": "weather", "arguments": '{"city":'}}]}}]}) + + with pytest.raises(ValueError, match="incomplete streamed tool-call"): + converter.convert("[DONE]") + + +def test_gemini_emits_fragmented_tool_call_exactly_once() -> None: + context = ProtocolContext(model="model-a", source_protocol="responses", target_protocol="gemini", transport="sse") + converter = ProtocolStreamConverter(get_protocol("responses"), get_protocol("gemini"), context) + source = [ + {"type": "response.output_item.added", "output_index": 0, "item": {"id": "fc_0", "type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": ""}}, + {"type": "response.function_call_arguments.delta", "output_index": 0, "item_id": "fc_0", "call_id": "call_1", "delta": '{"city":'}, + {"type": "response.function_call_arguments.delta", "output_index": 0, "item_id": "fc_0", "call_id": "call_1", "delta": '"Paris"}'}, + {"type": "response.completed", "response": {"id": "resp_1", "status": "completed", "model": "model-a", "output": []}}, + ] + + output = "".join(frame for event in source for frame in converter.convert(event)) + + assert output.count("functionCall") == 1 + assert '"city": "Paris"' in output + + +def test_gemini_flushes_genuine_zero_argument_tool_call_once_at_terminal() -> None: + context = ProtocolContext(model="model-a", source_protocol="responses", target_protocol="gemini", transport="sse") + converter = ProtocolStreamConverter(get_protocol("responses"), get_protocol("gemini"), context) + converter.convert({"type": "response.output_item.added", "output_index": 0, "item": {"id": "fc_0", "type": "function_call", "call_id": "call_1", "name": "ping", "arguments": ""}}) + + output = "".join(converter.convert({"type": "response.completed", "response": {"id": "resp_1", "status": "completed", "model": "model-a", "output": []}})) + + assert output.count("functionCall") == 1 + assert '"args": {}' in output + + +def test_stream_formatter_never_exposes_foreign_thought_signatures() -> None: + context = ProtocolContext(model="model-a", source_protocol="gemini", target_protocol="gemini", transport="sse") + converter = ProtocolStreamConverter(get_protocol("gemini"), get_protocol("gemini"), context) + + output = "".join(converter.convert({"candidates": [{"content": {"role": "model", "parts": [{"text": "hidden", "thought": True, "thoughtSignature": "opaque-secret"}]}}]})) + + assert "hidden" in output + assert "opaque-secret" not in output + + +@pytest.mark.parametrize("output_protocol", PROTOCOLS) +def test_bare_done_never_fabricates_success(output_protocol: str) -> None: + context = ProtocolContext(model="model-a", source_protocol="openai_chat", target_protocol=output_protocol, transport="sse") + converter = ProtocolStreamConverter(get_protocol("openai_chat"), get_protocol(output_protocol), context) + + output = "".join(converter.convert("[DONE]")) + + assert '"finish_reason": "stop"' not in output + assert '"stop_reason": "end_turn"' not in output + assert '"finishReason": "STOP"' not in output + if output_protocol == "responses": + assert "response.incomplete" in output + + +def test_terminal_responses_error_reuses_active_stream_lifecycle() -> None: + protocol_context = ProtocolContext(model="model-a", source_protocol="openai_chat", target_protocol="responses", transport="sse") + converter = ProtocolStreamConverter(get_protocol("openai_chat"), get_protocol("responses"), protocol_context) + initial = "".join(converter.convert({"choices": [{"delta": {"content": "partial"}}]})) + response_id = json.loads(initial.split("data: ", 1)[1].split("\n", 1)[0])["response"]["id"] + request_context = RequestContext( + model="provider/model-a", + provider="provider", + kwargs={}, + streaming=True, + credentials=[], + deadline=9999999999.0, + input_protocol_name="openai_chat", + output_protocol_name="responses", + ) + + terminal = "".join(RequestExecutor.__new__(RequestExecutor)._terminal_stream_error_lines( + request_context, + {"error": {"type": "server_error", "message": "failed after partial"}}, + protocol_context=protocol_context, + )) + + assert "response.created" not in terminal + assert response_id in terminal + assert "partial" in terminal + assert "failed after partial" in terminal + + +class _NativeStreamTransport: + def __init__(self, chunks: list[object]) -> None: + self.chunks = chunks + self.requests = [] + + async def stream_json_lines(self, endpoint, *, headers, payload): + self.requests.append(payload) + for chunk in self.chunks: + yield chunk + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_protocol", PROTOCOLS) +@pytest.mark.parametrize("output_protocol", PROTOCOLS) +async def test_native_executor_stream_matrix_uses_provider_reader_and_selected_writer( + provider_protocol: str, + output_protocol: str, +) -> None: + operations = { + "openai_chat": "chat", + "anthropic_messages": "messages", + "responses": "responses", + "gemini": "generate", + } + context = NativeProviderContext( + provider="provider", + model="model-a", + protocol_name=provider_protocol, + input_protocol_name="openai_chat", + output_protocol_name=output_protocol, + endpoint="https://provider.test/stream", + operation=operations[provider_protocol], + ) + transport = _NativeStreamTransport(_source_frames(provider_protocol)) + + frames = [ + frame + async for frame in NativeProviderExecutor().stream( + {"model": "model-a", "messages": [{"role": "user", "content": "hello"}], "stream": True}, + context, + transport, + ) + ] + output = "".join(frames) + + assert "hi" in output + assert "weather" in output + assert transport.requests + + +class _RuntimeProvider(ProviderInterface): + provider_env_name = "stream_runtime" + + async def get_models(self, api_key, client): + return [] + + +class _CredentialContext: + def __init__(self, credential: str = "credential-1") -> None: + self.credential = credential + self.stable_id = credential + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def mark_success(self, **kwargs): + self.success = kwargs + + def mark_failure(self, error): + self.failure = error + + +class _UsageManager: + def __init__(self, credentials: tuple[str, ...] = ("credential-1",)) -> None: + self.initialized = False + self.states = { + credential: SimpleNamespace( + tier="standard", + priority=1, + totals=SimpleNamespace(request_count=0), + usage=SimpleNamespace(), + model_usage={}, + group_usage={}, + get_usage_for_scope=lambda *args, **kwargs: None, + ) + for credential in credentials + } + self.window_manager = SimpleNamespace(get_primary_definition=lambda: None) + + async def initialize(self, credentials, priorities=None, tiers=None): + self.initialized = True + + async def acquire_credential(self, **kwargs): + candidates = kwargs.get("candidates") or ["credential-1"] + return _CredentialContext(candidates[0]) + + async def get_availability_stats(self, model, quota_group=None): + return { + "total": 1, + "available": 1, + "blocked": 0, + "blocked_by": {"cooldowns": 0, "window_limits": 0, "custom_caps": 0, "fair_cycle": 0, "concurrent": 0}, + "rotation_mode": "sequential", + } + + def get_model_quota_group(self, model): + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("output_protocol", PROTOCOLS) +async def test_request_executor_formats_litellm_stream_in_selected_protocol( + monkeypatch, + output_protocol: str, +) -> None: + usage_manager = _UsageManager() + executor = RequestExecutor( + usage_managers={"stream_runtime": usage_manager}, + cooldown_manager=None, + credential_filter=CredentialFilter({"stream_runtime": _RuntimeProvider}), + provider_transforms=ProviderTransforms({"stream_runtime": _RuntimeProvider}, None), + provider_plugins={"stream_runtime": _RuntimeProvider}, + http_client=httpx.AsyncClient(), + max_retries=1, + global_timeout=5, + litellm_logger_fn=lambda payload: None, + ) + + async def upstream(): + yield {"id": "chat_runtime", "model": "model-a", "choices": [{"delta": {"content": "runtime"}}]} + yield {"id": "chat_runtime", "model": "model-a", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + + async def fake_acompletion(**kwargs): + return upstream() + + monkeypatch.setattr("rotator_library.client.executor.litellm.acompletion", fake_acompletion) + context = RequestContext( + model="stream_runtime/model-a", + provider="stream_runtime", + kwargs={"model": "stream_runtime/model-a", "messages": [], "stream": True}, + streaming=True, + credentials=["credential-1"], + credential_secrets={"credential-1": "secret"}, + deadline=9999999999.0, + input_protocol_name="gemini", + output_protocol_name=output_protocol, + ) + + chunks = [chunk async for chunk in executor._execute_streaming(context)] + output = "".join(chunks) + + assert "runtime" in output + markers = { + "openai_chat": "chat.completion.chunk", + "anthropic_messages": "event: message_start", + "responses": "event: response.created", + "gemini": '"candidates"', + } + assert markers[output_protocol] in output + + +class _NativeErrorRuntimeProvider(_RuntimeProvider): + provider_env_name = "native_stream_runtime" + protocol_name = "openai_chat" + native_streaming_supported = True + + def get_native_headers(self, credential_identifier, model="", operation="chat"): + return {"Authorization": f"Bearer {credential_identifier}"} + + def get_native_endpoint(self, model="", operation="chat"): + return "https://provider.test/chat" + + def normalize_native_model(self, model=""): + return model.split("/", 1)[-1] + + +class _RotatingNativeStreamClient: + def __init__(self) -> None: + self.credentials = [] + + async def stream_json_lines(self, endpoint, *, headers, json): + credential = headers["Authorization"].removeprefix("Bearer ") + self.credentials.append(credential) + if credential == "secret-1": + yield {"error": {"type": "rate_limit", "message": "rotate me", "status_code": 429}} + return + yield {"choices": [{"delta": {"content": "after rotation"}}]} + yield {"choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + yield "[DONE]" + + +@pytest.mark.asyncio +async def test_native_stream_error_rotates_credentials_before_selected_output(monkeypatch) -> None: + monkeypatch.setenv("TRANSIENT_RETRY_DELAY", "0") + monkeypatch.setenv("TRANSIENT_RETRY_JITTER", "0") + credentials = ("credential-1", "credential-2") + usage_manager = _UsageManager(credentials) + http_client = _RotatingNativeStreamClient() + executor = RequestExecutor( + usage_managers={"native_stream_runtime": usage_manager}, + cooldown_manager=None, + credential_filter=CredentialFilter({"native_stream_runtime": _NativeErrorRuntimeProvider}), + provider_transforms=ProviderTransforms({"native_stream_runtime": _NativeErrorRuntimeProvider}, None), + provider_plugins={"native_stream_runtime": _NativeErrorRuntimeProvider}, + http_client=http_client, + max_retries=1, + global_timeout=5, + litellm_logger_fn=lambda payload: None, + ) + kwargs = { + "model": "native_stream_runtime/model-a", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + } + context = RequestContext( + model=kwargs["model"], + provider="native_stream_runtime", + kwargs=kwargs, + streaming=True, + credentials=list(credentials), + credential_secrets={"credential-1": "secret-1", "credential-2": "secret-2"}, + deadline=9999999999.0, + input_protocol_name="openai_chat", + output_protocol_name="responses", + protocol_request=dict(kwargs), + unified_request=get_protocol("openai_chat").parse_request(kwargs), + input_provider="native_stream_runtime", + ) + + output = "".join([chunk async for chunk in executor._execute_streaming(context)]) + + assert http_client.credentials == ["secret-1", "secret-2"] + assert "after rotation" in output + assert "event: response.completed" in output + assert "rotate me" not in output + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("output_protocol", "start_marker"), + (("anthropic_messages", "event: message_start"), ("responses", "event: response.created")), +) +async def test_litellm_in_band_error_rotates_without_duplicate_destination_start( + monkeypatch, + output_protocol: str, + start_marker: str, +) -> None: + monkeypatch.setenv("TRANSIENT_RETRY_DELAY", "0") + monkeypatch.setenv("TRANSIENT_RETRY_JITTER", "0") + credentials = ("credential-1", "credential-2") + usage_manager = _UsageManager(credentials) + calls = [] + + async def fake_acompletion(**kwargs): + calls.append(kwargs["api_key"]) + + async def failed_stream(): + yield {"choices": [{"delta": {"role": "assistant"}}]} + yield 'event: error\ndata: {"error":{"type":"rate_limit","message":"retry this credential"}}\n\n' + + async def successful_stream(): + yield {"choices": [{"delta": {"content": "one lifecycle"}}]} + yield {"choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + + return failed_stream() if kwargs["api_key"] == "secret-1" else successful_stream() + + monkeypatch.setattr("rotator_library.client.executor.litellm.acompletion", fake_acompletion) + executor = RequestExecutor( + usage_managers={"stream_runtime": usage_manager}, + cooldown_manager=None, + credential_filter=CredentialFilter({"stream_runtime": _RuntimeProvider}), + provider_transforms=ProviderTransforms({"stream_runtime": _RuntimeProvider}, None), + provider_plugins={"stream_runtime": _RuntimeProvider}, + http_client=httpx.AsyncClient(), + max_retries=1, + global_timeout=5, + litellm_logger_fn=lambda payload: None, + ) + context = RequestContext( + model="stream_runtime/model-a", + provider="stream_runtime", + kwargs={"model": "stream_runtime/model-a", "messages": [], "stream": True}, + streaming=True, + credentials=list(credentials), + credential_secrets={"credential-1": "secret-1", "credential-2": "secret-2"}, + deadline=9999999999.0, + input_protocol_name="openai_chat", + output_protocol_name=output_protocol, + ) + + output = "".join([chunk async for chunk in executor._execute_streaming(context)]) + + assert calls == ["secret-1", "secret-2"] + assert output.count(start_marker) == 1 + assert "one lifecycle" in output + assert "retry this credential" not in output + + +class _CustomErrorRuntimeProvider(_RuntimeProvider): + provider_env_name = "custom_stream_runtime" + credentials_seen: list[str] = [] + + def has_custom_logic(self): + return True + + async def acompletion(self, client, **kwargs): + credential = kwargs["credential_identifier"] + type(self).credentials_seen.append(credential) + + async def stream(): + if credential == "secret-1": + yield {"error": {"type": "rate_limit", "message": "custom retry"}} + return + yield {"choices": [{"delta": {"content": "custom success"}}]} + yield {"choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + + return stream() + + +@pytest.mark.asyncio +async def test_custom_in_band_mapping_error_rotates_before_output(monkeypatch) -> None: + monkeypatch.setenv("TRANSIENT_RETRY_DELAY", "0") + monkeypatch.setenv("TRANSIENT_RETRY_JITTER", "0") + _CustomErrorRuntimeProvider.credentials_seen = [] + credentials = ("credential-1", "credential-2") + usage_manager = _UsageManager(credentials) + executor = RequestExecutor( + usage_managers={"custom_stream_runtime": usage_manager}, + cooldown_manager=None, + credential_filter=CredentialFilter({"custom_stream_runtime": _CustomErrorRuntimeProvider}), + provider_transforms=ProviderTransforms({"custom_stream_runtime": _CustomErrorRuntimeProvider}, None), + provider_plugins={"custom_stream_runtime": _CustomErrorRuntimeProvider}, + http_client=httpx.AsyncClient(), + max_retries=1, + global_timeout=5, + litellm_logger_fn=lambda payload: None, + ) + context = RequestContext( + model="custom_stream_runtime/model-a", + provider="custom_stream_runtime", + kwargs={"model": "custom_stream_runtime/model-a", "messages": [], "stream": True}, + streaming=True, + credentials=list(credentials), + credential_secrets={"credential-1": "secret-1", "credential-2": "secret-2"}, + deadline=9999999999.0, + input_protocol_name="openai_chat", + output_protocol_name="gemini", + ) + + output = "".join([chunk async for chunk in executor._execute_streaming(context)]) + + assert _CustomErrorRuntimeProvider.credentials_seen == ["secret-1", "secret-2"] + assert "custom success" in output + assert "custom retry" not in output + + +@pytest.mark.asyncio +async def test_in_band_error_after_visible_output_closes_same_responses_lifecycle(monkeypatch) -> None: + credentials = ("credential-1", "credential-2") + usage_manager = _UsageManager(credentials) + calls = [] + + async def fake_acompletion(**kwargs): + calls.append(kwargs["api_key"]) + + async def stream(): + yield {"choices": [{"delta": {"content": "already visible"}}]} + yield 'data: {"error":{"type":"rate_limit","message":"failed after visible"}}\n\n' + + return stream() + + monkeypatch.setattr("rotator_library.client.executor.litellm.acompletion", fake_acompletion) + executor = RequestExecutor( + usage_managers={"stream_runtime": usage_manager}, + cooldown_manager=None, + credential_filter=CredentialFilter({"stream_runtime": _RuntimeProvider}), + provider_transforms=ProviderTransforms({"stream_runtime": _RuntimeProvider}, None), + provider_plugins={"stream_runtime": _RuntimeProvider}, + http_client=httpx.AsyncClient(), + max_retries=1, + global_timeout=5, + litellm_logger_fn=lambda payload: None, + ) + context = RequestContext( + model="stream_runtime/model-a", + provider="stream_runtime", + kwargs={"model": "stream_runtime/model-a", "messages": [], "stream": True}, + streaming=True, + credentials=list(credentials), + credential_secrets={"credential-1": "secret-1", "credential-2": "secret-2"}, + deadline=9999999999.0, + input_protocol_name="openai_chat", + output_protocol_name="responses", + ) + + output = "".join([chunk async for chunk in executor._execute_streaming(context)]) + + assert calls == ["secret-1"] + assert output.count("event: response.created") == 1 + assert "already visible" in output + assert "event: response.failed" in output + assert "failed after visible" in output diff --git a/tests/test_provider_field_cache_contracts.py b/tests/test_provider_field_cache_contracts.py new file mode 100644 index 000000000..ff8fa0b21 --- /dev/null +++ b/tests/test_provider_field_cache_contracts.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from rotator_library.native_provider import NativeProviderContext, NativeProviderExecutor +from rotator_library.providers.antigravity_provider import AntigravityProvider +from rotator_library.providers.claude_code_provider import ClaudeCodeProvider +from rotator_library.providers.codex_provider import CodexProvider +from rotator_library.providers.copilot_provider import CopilotProvider + + +OUTPUT_PROTOCOLS = ("openai_chat", "anthropic_messages", "responses", "gemini") +STATEFUL_PROVIDERS = (AntigravityProvider(), ClaudeCodeProvider(), CodexProvider()) + + +class RecordingResponseTransport: + def __init__(self, responses: list[dict[str, Any]]) -> None: + self.responses = list(responses) + self.requests: list[dict[str, Any]] = [] + + async def post_json(self, endpoint, *, headers, payload): + self.requests.append(payload) + return self.responses.pop(0) + + +class RecordingStreamTransport: + def __init__(self, chunks: list[Any]) -> None: + self.chunks = chunks + self.requests: list[dict[str, Any]] = [] + + async def stream_json_lines(self, endpoint, *, headers, payload): + self.requests.append(payload) + for chunk in self.chunks: + yield chunk + + +def _context( + provider, + *, + output_protocol: str = "openai_chat", + credential_id: str = "credential-a", + session_id: str = "session-a", + model: str = "model-a", + stream: bool = False, + disable_provider_continuation: bool = False, +) -> NativeProviderContext: + operation = provider.get_native_operation(model, None, stream=stream) + return NativeProviderContext( + provider=provider.provider_env_name, + model=model, + protocol_name=provider.get_protocol_name(model), + input_protocol_name=provider.get_protocol_name(model), + output_protocol_name=output_protocol, + endpoint="https://provider.test/generate", + operation=operation, + credential_id=credential_id, + session_id=session_id, + scope_key="scope-a", + adapter_names=provider.get_adapter_names(model), + adapter_config=provider.get_adapter_config(model), + field_cache_rules=provider.get_field_cache_rules(model), + request_preparer=provider.prepare_native_request, + request_validator=getattr(provider, "validate_request", None), + metadata={"disable_provider_continuation": disable_provider_continuation}, + ) + + +def _request(provider) -> dict[str, Any]: + if isinstance(provider, AntigravityProvider): + return {"model": "model-a", "contents": [{"role": "user", "parts": [{"text": "hello"}]}]} + if isinstance(provider, ClaudeCodeProvider): + return {"model": "model-a", "max_tokens": 32, "messages": [{"role": "user", "content": "hello"}]} + if isinstance(provider, CodexProvider): + return {"model": "model-a", "input": "hello"} + return {"model": "model-a", "messages": [{"role": "user", "content": "hello"}]} + + +def _response(provider, state: str) -> dict[str, Any]: + if isinstance(provider, AntigravityProvider): + return { + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "private thought", "thought": True, "thoughtSignature": state}, {"text": "answer"}]}, + "finishReason": "STOP", + }], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + } + if isinstance(provider, ClaudeCodeProvider): + return { + "id": "msg_provider", + "type": "message", + "role": "assistant", + "model": "model-a", + "content": [{"type": "thinking", "thinking": "private thought", "signature": state}, {"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + if isinstance(provider, CodexProvider): + return { + "id": state, + "object": "response", + "status": "completed", + "model": "model-a", + "output": [{"id": "msg_0", "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "answer"}]}], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + return {"id": "chat_1", "model": "model-a", "choices": [{"message": {"role": "assistant", "content": "answer"}, "finish_reason": "stop"}]} + + +def _stream_chunks(provider, state: str) -> list[Any]: + if isinstance(provider, AntigravityProvider): + return [ + {"candidates": [{"content": {"role": "model", "parts": [{"text": "private stream thought", "thought": True, "thoughtSignature": state}, {"text": "stream answer"}]}, "finishReason": "STOP"}]}, + "[DONE]", + ] + if isinstance(provider, ClaudeCodeProvider): + return [ + {"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": state}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "stream answer"}}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}, + {"type": "message_stop"}, + ] + return [ + {"type": "response.output_text.delta", "item_id": "msg_0", "output_index": 0, "content_index": 0, "delta": "stream answer"}, + {"type": "response.completed", "response": _response(provider, state)}, + "[DONE]", + ] + + +def _injected_state(provider, payload: dict[str, Any]) -> Any: + if isinstance(provider, AntigravityProvider): + return payload["request"]["metadata"]["thoughtSignatures"] + if isinstance(provider, ClaudeCodeProvider): + return payload["metadata"]["thinking_signatures"] + return payload["previous_response_id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", STATEFUL_PROVIDERS, ids=lambda provider: provider.provider_env_name) +@pytest.mark.parametrize("output_protocol", OUTPUT_PROTOCOLS) +async def test_real_provider_non_stream_state_is_output_protocol_independent(provider, output_protocol: str) -> None: + executor = NativeProviderExecutor() + transport = RecordingResponseTransport([_response(provider, "state-one"), _response(provider, "state-two")]) + context = _context(provider, output_protocol=output_protocol) + + first = await executor.execute(_request(provider), context, transport) + await executor.execute(_request(provider), context, transport) + + assert _injected_state(provider, transport.requests[1]) == (["state-one"] if not isinstance(provider, CodexProvider) else "state-one") + serialized = json.dumps(first) + if not isinstance(provider, CodexProvider): + assert "state-one" not in serialized + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", STATEFUL_PROVIDERS, ids=lambda provider: provider.provider_env_name) +async def test_real_provider_mixed_stream_and_non_stream_share_latest_logical_state(provider) -> None: + executor = NativeProviderExecutor() + first_transport = RecordingResponseTransport([_response(provider, "non-stream-state")]) + await executor.execute(_request(provider), _context(provider), first_transport) + + stream_transport = RecordingStreamTransport(_stream_chunks(provider, "stream-state")) + stream_output = [frame async for frame in executor.stream(_request(provider), _context(provider, stream=True), stream_transport)] + + final_transport = RecordingResponseTransport([_response(provider, "final-state")]) + await executor.execute(_request(provider), _context(provider), final_transport) + + expected = "stream-state" if isinstance(provider, CodexProvider) else ["non-stream-state", "stream-state"] + assert _injected_state(provider, final_transport.requests[0]) == expected + assert "stream-state" not in "".join(stream_output) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", STATEFUL_PROVIDERS, ids=lambda provider: provider.provider_env_name) +@pytest.mark.parametrize( + "isolated_context", + ( + {"credential_id": "credential-b"}, + {"session_id": "session-b"}, + {"model": "model-b"}, + ), +) +async def test_real_provider_state_isolated_by_model_credential_and_session(provider, isolated_context) -> None: + executor = NativeProviderExecutor() + await executor.execute( + _request(provider), + _context(provider), + RecordingResponseTransport([_response(provider, "private-state")]), + ) + isolated_transport = RecordingResponseTransport([_response(provider, "other-state")]) + + await executor.execute( + _request(provider), + _context(provider, **isolated_context), + isolated_transport, + ) + + with pytest.raises((KeyError, TypeError)): + _injected_state(provider, isolated_transport.requests[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", STATEFUL_PROVIDERS, ids=lambda provider: provider.provider_env_name) +@pytest.mark.parametrize("missing_scope", ({"credential_id": None}, {"session_id": None})) +async def test_real_provider_state_skips_when_required_scope_is_missing(provider, missing_scope) -> None: + executor = NativeProviderExecutor() + context = _context(provider, **missing_scope) + transport = RecordingResponseTransport([ + _response(provider, "private-state"), + _response(provider, "second-state"), + ]) + + await executor.execute(_request(provider), context, transport) + await executor.execute(_request(provider), context, transport) + + with pytest.raises((KeyError, TypeError)): + _injected_state(provider, transport.requests[1]) + + +@pytest.mark.asyncio +async def test_codex_local_lineage_disables_both_continuation_sources() -> None: + provider = CodexProvider() + executor = NativeProviderExecutor() + await executor.execute( + _request(provider), + _context(provider), + RecordingResponseTransport([_response(provider, "resp_non_stream")]), + ) + _ = [ + frame + async for frame in executor.stream( + _request(provider), + _context(provider, stream=True), + RecordingStreamTransport(_stream_chunks(provider, "resp_stream")), + ) + ] + suppressed = RecordingResponseTransport([_response(provider, "resp_final")]) + + await executor.execute( + _request(provider), + _context(provider, disable_provider_continuation=True), + suppressed, + ) + + assert "previous_response_id" not in suppressed.requests[0] + + +@pytest.mark.asyncio +async def test_copilot_native_contract_has_no_invented_provider_state() -> None: + provider = CopilotProvider() + executor = NativeProviderExecutor() + transport = RecordingResponseTransport([_response(provider, "unused"), _response(provider, "unused")]) + context = _context(provider) + + await executor.execute(_request(provider), context, transport) + await executor.execute(_request(provider), context, transport) + + assert transport.requests[1] == transport.requests[0] diff --git a/tests/test_provider_protocol_declarations.py b/tests/test_provider_protocol_declarations.py new file mode 100644 index 000000000..a16ee56f3 --- /dev/null +++ b/tests/test_provider_protocol_declarations.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from rotator_library.field_cache import FieldCacheRule +from rotator_library.providers.provider_interface import ProviderInterface + + +class DeclarationProvider(ProviderInterface): + protocol_name = "openai_chat" + adapter_names = ("model_override", "suppress_developer_role") + field_cache_rules = ( + FieldCacheRule(name="reasoning_content", source="response", path="choices.*.message.reasoning_content"), + ) + + async def get_models(self, api_key, client): + return [] + + +class BareProvider(ProviderInterface): + async def get_models(self, api_key, client): + return [] + + +def test_provider_interface_plain_defaults_are_noop() -> None: + provider = BareProvider() + + assert provider.get_protocol_name("model") is None + assert provider.get_adapter_names("model") == () + assert provider.get_adapter_config("model") == {} + assert provider.get_field_cache_rules("model") == () + + +def test_provider_interface_defaults_are_noop_for_protocol_stack() -> None: + provider = DeclarationProvider() + + assert provider.get_protocol_name("model") == "openai_chat" + assert provider.get_adapter_names("model") == ("model_override", "suppress_developer_role") + assert provider.get_adapter_config("model") == {} + assert provider.get_field_cache_rules("model")[0].name == "reasoning_content" + + +def test_provider_interface_methods_can_be_model_specific() -> None: + class ModelSpecificProvider(DeclarationProvider): + def get_protocol_name(self, model: str = ""): + return "responses" if "response" in model else super().get_protocol_name(model) + + def get_adapter_config(self, model: str = ""): + return {"model_override": {"model": f"native/{model}"}} + + provider = ModelSpecificProvider() + + assert provider.get_protocol_name("response-model") == "responses" + assert provider.get_protocol_name("chat-model") == "openai_chat" + assert provider.get_adapter_config("chat-model") == {"model_override": {"model": "native/chat-model"}} diff --git a/tests/test_provider_runtime_config.py b/tests/test_provider_runtime_config.py new file mode 100644 index 000000000..33b32dfdb --- /dev/null +++ b/tests/test_provider_runtime_config.py @@ -0,0 +1,355 @@ +from __future__ import annotations + +import json +import pytest + +from rotator_library.providers.provider_interface import ProviderInterface +from rotator_library.providers import PROVIDER_PLUGINS, _create_dynamic_plugin_class, _register_providers +from rotator_library.client.rotating_client import _add_configured_no_auth_credentials +from rotator_library.client.scopes import NO_AUTH_CREDENTIAL, ScopeManager +from rotator_library.config.experimental import ExperimentalConfigError, load_config_from_mapping +from rotator_library.providers.claude_code_provider import ClaudeCodeProvider +from rotator_library.providers.codex_provider import CodexProvider +from rotator_library.providers.copilot_provider import CopilotProvider + + +class ConfiguredProvider(ProviderInterface): + provider_env_name = "configured" + protocol_name = "litellm_fallback" + adapter_names = ("noop",) + model_quota_groups = {"base": ["base-model"]} + + async def get_models(self, api_key, client): + return [] + + +def _write_config(tmp_path, payload: dict) -> str: + path = tmp_path / "config.json" + path.write_text(json.dumps(payload), encoding="utf-8") + return str(path) + + +def test_provider_json_protocol_adapters_field_cache_and_quota_groups_are_wired(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + { + "providers": { + "configured": { + "protocol_name": "openai_chat", + "default_output_protocol": "anthropic_messages", + "adapter_names": ["model_override"], + "adapter_config": {"model_override": {"model": "upstream-model"}}, + "native_streaming_supported": True, + "field_cache": [ + {"name": "state", "source": "response", "path": "metadata.state", "target_path": "metadata.cached_state"} + ], + "model_quota_groups": {"json_group": ["json-model"]}, + } + } + }, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + provider = ConfiguredProvider() + + assert provider.get_protocol_name("configured/gpt") == "openai_chat" + assert provider.get_default_output_protocol("configured/gpt") == "anthropic_messages" + assert provider.get_adapter_names("configured/gpt") == ("model_override",) + assert provider.get_adapter_config("configured/gpt") == {"model_override": {"model": "upstream-model"}} + assert provider.supports_native_streaming("configured/gpt", "chat") is True + assert [rule.name for rule in provider.get_field_cache_rules("configured/gpt")] == ["state"] + assert provider.get_model_quota_group("json-model") == "json_group" + assert provider.get_model_quota_group("base-model") == "base" + + +def test_provider_runtime_config_can_be_bound_to_startup_snapshot(tmp_path, monkeypatch) -> None: + class SnapshotProvider(ProviderInterface): + provider_env_name = "snapshot_provider" + + async def get_models(self, api_key, client): + return [] + + startup = load_config_from_mapping({ + "providers": {"snapshot_provider": {"protocol_name": "gemini"}} + }) + changed_path = _write_config( + tmp_path, + {"providers": {"snapshot_provider": {"protocol_name": "responses"}}}, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", changed_path) + provider = SnapshotProvider() + provider.bind_runtime_config(startup) + + assert provider.get_protocol_name("model") == "gemini" + + +def test_provider_singleton_rejects_incompatible_snapshot_rebinding() -> None: + class BoundProvider(ProviderInterface): + provider_env_name = "bound_provider" + + async def get_models(self, api_key, client): + return [] + + first = load_config_from_mapping({"providers": {"bound_provider": {"protocol_name": "gemini"}}}) + equivalent = load_config_from_mapping({"providers": {"bound_provider": {"protocol_name": "gemini"}}}) + incompatible = load_config_from_mapping({"providers": {"bound_provider": {"protocol_name": "responses"}}}) + provider = BoundProvider() + + provider.bind_runtime_config(first) + provider.bind_runtime_config(equivalent) + with pytest.raises(RuntimeError, match="already bound"): + provider.bind_runtime_config(incompatible) + + +def test_provider_json_quota_groups_still_allow_env_override(tmp_path, monkeypatch) -> None: + config_path = _write_config(tmp_path, {"providers": {"configured": {"model_quota_groups": {"json_group": ["json-model"]}}}}) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + monkeypatch.setenv("QUOTA_GROUPS_CONFIGURED_JSON_GROUP", "env-model") + + assert ConfiguredProvider().get_model_quota_group("env-model") == "json_group" + assert ConfiguredProvider().get_model_quota_group("json-model") is None + + +def test_priority_provider_overrides_respect_json_adapter_config_and_streaming(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + { + "providers": { + "claude_code": { + "native_streaming_supported": True, + "adapter_config": {"suppress_developer_role": {"mode": "assistant"}}, + }, + "copilot": { + "native_streaming_supported": True, + "adapter_config": {"suppress_developer_role": {"mode": "user"}}, + }, + "codex": {"native_streaming_supported": True}, + } + }, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + + assert ClaudeCodeProvider().supports_native_streaming("claude_code/claude", "messages") is True + assert ClaudeCodeProvider().get_adapter_config("claude_code/claude")["suppress_developer_role"]["mode"] == "assistant" + assert CopilotProvider().supports_native_streaming("copilot/gpt", "chat") is True + assert CopilotProvider().get_adapter_config("copilot/gpt")["suppress_developer_role"]["mode"] == "user" + assert CodexProvider().supports_native_streaming("codex/gpt", "responses") is True + + +@pytest.mark.asyncio +async def test_config_defined_provider_can_use_any_native_protocol(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + { + "providers": { + "configured_native": { + "api_base": "https://native.example/v1beta", + "protocol_name": "gemini", + "default_output_protocol": "anthropic_messages", + "auth_mode": "x-goog-api-key", + "models": ["gemini-custom"], + "native_streaming_supported": True, + "endpoint_paths": { + "generate": "/models/{model}:generateContent", + "stream_generate": "/models/{model}:streamGenerateContent?alt=sse", + }, + } + } + }, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + provider = _create_dynamic_plugin_class("configured_native")() + + assert provider.get_protocol_name("configured_native/gemini-custom") == "gemini" + assert provider.get_default_output_protocol("configured_native/gemini-custom") == "anthropic_messages" + assert provider.get_native_operation("gemini-custom", stream=True) == "stream_generate" + assert provider.get_native_endpoint("configured_native/gemini-custom", "generate") == ( + "https://native.example/v1beta/models/gemini-custom:generateContent" + ) + assert provider.get_native_headers("secret", operation="generate")["x-goog-api-key"] == "secret" + assert provider.supports_native_streaming("configured_native/gemini-custom", "stream_generate") is True + assert await provider.get_models("secret", object()) == ["configured_native/gemini-custom"] + + +def test_config_defined_provider_is_registered_without_api_base_env(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + { + "providers": { + "config_only_native": { + "api_base": "https://native.example/v1", + "protocol_name": "anthropic_messages", + "models": ["claude-custom"], + } + } + }, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + monkeypatch.delenv("CONFIG_ONLY_NATIVE_API_BASE", raising=False) + + try: + _register_providers() + plugin = PROVIDER_PLUGINS["config_only_native"]() + assert plugin.get_protocol_name("config_only_native/claude-custom") == "anthropic_messages" + assert plugin.get_api_base() == "https://native.example/v1" + finally: + PROVIDER_PLUGINS.pop("config_only_native", None) + + +def test_file_provider_rejects_silently_ignored_custom_transport_keys(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + {"providers": {"openai": {"api_base": "https://ignored.example/v1"}}}, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + + with pytest.raises(ValueError, match="implemented in code"): + _register_providers() + + +def test_custom_provider_rejects_non_generative_protocol(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + { + "providers": { + "embedding_only": { + "api_base": "https://embedding.example/v1", + "protocol_name": "embeddings", + } + } + }, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + + with pytest.raises(ValueError, match="supported generative protocol"): + _register_providers() + + +def test_builtin_provider_rejects_non_generative_protocol_override() -> None: + with pytest.raises(ExperimentalConfigError, match="supported generative protocol"): + load_config_from_mapping({ + "providers": { + "claude_code": {"protocol_name": "openai_embeddings"} + } + }) + + +def test_dynamic_provider_transport_config_is_snapshotted_at_construction(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + { + "providers": { + "snapshot_native": { + "api_base": "https://first.example/v1", + "protocol_name": "responses", + } + } + }, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + provider = _create_dynamic_plugin_class("snapshot_native")() + (tmp_path / "config.json").write_text( + json.dumps( + { + "providers": { + "snapshot_native": { + "api_base": "https://second.example/v1", + "protocol_name": "gemini", + } + } + } + ), + encoding="utf-8", + ) + + assert provider.get_api_base() == "https://first.example/v1" + assert provider.get_protocol_name("snapshot_native/model") == "responses" + + +def test_no_auth_custom_provider_gets_internal_selection_slot_without_dummy_secret() -> None: + credentials = {} + config = load_config_from_mapping( + { + "providers": { + "local_native": { + "api_base": "http://localhost:9000/v1", + "protocol_name": "openai_chat", + "auth_mode": "none", + } + } + } + ) + + _add_configured_no_auth_credentials( + credentials, + {"local_native"}, + config=config, + ) + + assert credentials == {"local_native": ["__proxy_no_auth__"]} + + +def test_registered_dynamic_provider_binds_config_before_lazy_instantiation(tmp_path, monkeypatch) -> None: + config_path = _write_config( + tmp_path, + { + "providers": { + "registered_snapshot": { + "api_base": "https://first.example/v1", + "protocol_name": "responses", + } + } + }, + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", config_path) + + try: + _register_providers() + plugin_class = PROVIDER_PLUGINS["registered_snapshot"] + (tmp_path / "config.json").write_text( + json.dumps( + { + "providers": { + "registered_snapshot": { + "api_base": "https://second.example/v1", + "protocol_name": "gemini", + } + } + } + ), + encoding="utf-8", + ) + provider = plugin_class() + + assert provider.get_api_base() == "https://first.example/v1" + assert provider.get_protocol_name("registered_snapshot/model") == "responses" + finally: + PROVIDER_PLUGINS.pop("registered_snapshot", None) + + +@pytest.mark.asyncio +async def test_no_auth_slot_is_preserved_in_classifier_scope(tmp_path) -> None: + created = [] + + async def ensure(provider, classifier, credentials): + created.append((provider, classifier, credentials)) + return f"classifier:{classifier}:{provider}" + + manager = ScopeManager( + all_credentials={"local_native": [NO_AUTH_CREDENTIAL]}, + usage_base_path=tmp_path, + fingerprint_key=b"test", + model_list_cache={}, + ensure_scoped_usage_manager=ensure, + ) + + scope = await manager.resolve_scope_for_provider( + "local_native", + "tenant-a", + None, + None, + True, + ) + + assert scope["credentials"] == [NO_AUTH_CREDENTIAL] + assert scope["credential_secrets"] == {} + assert created == [("local_native", "tenant-a", [NO_AUTH_CREDENTIAL])] diff --git a/tests/test_request_builder_routing.py b/tests/test_request_builder_routing.py new file mode 100644 index 000000000..459c6c941 --- /dev/null +++ b/tests/test_request_builder_routing.py @@ -0,0 +1,331 @@ +from __future__ import annotations + +import pytest + +from rotator_library.client.request_builder import RequestContextBuilder +from rotator_library.session_tracking import SessionTrackingHints + + +class FakeModelResolver: + def resolve_model_id(self, model, provider): + return model + + +class FakeSession: + session_id = "session" + affinity_key = "affinity" + possible_compaction = False + lineage_parent_session_id = None + tracking_namespace = "namespace" + + +class FakeSessionTracker: + def __init__(self): + self.calls = [] + + def infer_session(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return FakeSession() + + +async def _scope(provider, classifier, request_api_keys, request_providers, private): + return { + "credentials": [f"{provider}-cred"], + "usage_manager_key": provider, + "provider_config": {"provider": provider}, + "credential_secrets": {f"{provider}-cred": f"{provider}-secret"}, + "classifier": classifier or "global", + } + + +def _builder(session_tracker=None) -> RequestContextBuilder: + return RequestContextBuilder( + resolve_scope_for_provider=_scope, + model_resolver=FakeModelResolver(), + session_tracker=session_tracker or FakeSessionTracker(), + get_global_timeout=lambda: 30, + get_enable_request_logging=lambda: False, + ) + + +@pytest.mark.asyncio +async def test_request_builder_leaves_no_config_provider_model_unrouted(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + context = await _builder().build_completion_context(None, None, {"model": "openai/gpt-5.1", "messages": []}) + + assert context.routing_targets is None + assert context.provider == "openai" + assert context.credentials == ["openai-cred"] + + +@pytest.mark.asyncio +async def test_request_builder_populates_fallback_group_targets_from_env(monkeypatch) -> None: + monkeypatch.setenv("FALLBACK_GROUPS", "code_chain") + monkeypatch.setenv("FALLBACK_GROUP_CODE_CHAIN", "codex/gpt-5.1-codex,openai/gpt-5.1") + monkeypatch.setenv("MODEL_ROUTE_CODEX", "group:code_chain") + + context = await _builder().build_completion_context(None, None, {"model": "codex", "messages": []}) + + assert context.provider == "codex" + assert context.model == "codex/gpt-5.1-codex" + assert context.routing_group_name == "code_chain" + assert context.routing_group is not None + assert context.routing_group.name == "code_chain" + assert [target.prefixed_model for target in context.routing_targets] == ["codex/gpt-5.1-codex", "openai/gpt-5.1"] + assert context.routing_targets[1].metadata["request_scope"]["credentials"] == ["openai-cred"] + + +@pytest.mark.asyncio +async def test_request_builder_rejects_unprefixed_model_without_route(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + + with pytest.raises(ValueError): + await _builder().build_completion_context(None, None, {"model": "gpt-5.1", "messages": []}) + + +@pytest.mark.asyncio +async def test_request_builder_consumes_internal_session_tracking_hints(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + tracker = FakeSessionTracker() + kwargs = { + "model": "openai/gpt-5.1", + "messages": [], + "_session_tracking_hints": SessionTrackingHints( + global_strong_anchors=["responses_previous_response_id:resp_parent"], + affinity_key="responses_previous_response_id:resp_parent", + ), + } + + context = await _builder(tracker).build_completion_context(None, None, kwargs) + + assert "_session_tracking_hints" not in context.kwargs + hints = tracker.calls[0][1]["hints"] + assert hints.global_strong_anchors == ["responses_previous_response_id:resp_parent"] + assert hints.affinity_key == "responses_previous_response_id:resp_parent" + + +@pytest.mark.asyncio +async def test_request_builder_passes_provider_independent_session_domain(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + tracker = FakeSessionTracker() + context = await _builder(tracker).build_completion_context( + None, + None, + {"model": "openai/gpt-5.1", "messages": []}, + ) + + assert tracker.calls[0][1]["scope_key"] == "public" + assert context.session_isolation_key == "public" + + +@pytest.mark.asyncio +async def test_provider_session_hook_receives_provider_native_shape(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + seen = [] + + class GeminiPlugin: + def get_protocol_name(self, model=""): + return "gemini" + + def normalize_native_model(self, model=""): + return model.split("/", 1)[-1] + + def get_native_operation(self, model="", request=None, stream=False): + assert request is None + return "generate" + + def prepare_native_request(self, request, model="", operation=""): + return request + + def get_session_tracking_hints(self, request, model=""): + seen.append(request) + return None + + plugin = GeminiPlugin() + builder = RequestContextBuilder( + resolve_scope_for_provider=_scope, + model_resolver=FakeModelResolver(), + session_tracker=FakeSessionTracker(), + get_global_timeout=lambda: 30, + get_enable_request_logging=lambda: False, + get_provider_instance=lambda provider: plugin, + ) + + await builder.build_completion_context( + None, + None, + { + "_input_protocol": "anthropic_messages", + "model": "gemini/gemini-3-pro", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert seen + assert "contents" in seen[0] + assert "messages" not in seen[0] + + +@pytest.mark.asyncio +async def test_provider_default_output_applies_only_without_explicit_override(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + + class Provider: + def get_default_output_protocol(self, model=""): + return "anthropic_messages" + + builder = RequestContextBuilder( + resolve_scope_for_provider=_scope, + model_resolver=FakeModelResolver(), + session_tracker=FakeSessionTracker(), + get_global_timeout=lambda: 30, + get_enable_request_logging=lambda: False, + get_provider_instance=lambda provider: Provider(), + ) + implicit = await builder.build_completion_context( + None, + None, + {"model": "openai/gpt-5.1", "messages": [{"role": "user", "content": "hello"}]}, + ) + explicit = await builder.build_completion_context( + None, + None, + { + "_output_protocol": "responses", + "model": "openai/gpt-5.1", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert implicit.output_protocol_name == "anthropic_messages" + assert explicit.output_protocol_name == "responses" + + +@pytest.mark.asyncio +async def test_request_builder_carries_independent_cross_protocol_stream_selection(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + builder = RequestContextBuilder( + resolve_scope_for_provider=_scope, + model_resolver=FakeModelResolver(), + session_tracker=FakeSessionTracker(), + get_global_timeout=lambda: 30, + get_enable_request_logging=lambda: False, + ) + + context = await builder.build_completion_context( + None, + None, + { + "_input_protocol": "openai_chat", + "_output_protocol": "gemini", + "model": "openai/gpt-5.1", + "messages": [], + "stream": True, + }, + ) + + assert context.streaming is True + assert context.input_protocol_name == "openai_chat" + assert context.output_protocol_name == "gemini" + +@pytest.mark.asyncio +async def test_request_builder_classifier_domain_is_stable_across_providers(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + openai_tracker = FakeSessionTracker() + anthropic_tracker = FakeSessionTracker() + openai_context = await _builder(openai_tracker).build_completion_context( + None, + None, + {"model": "openai/gpt-5.1", "messages": [], "classifier": "agent-one"}, + ) + anthropic_context = await _builder(anthropic_tracker).build_completion_context( + None, + None, + {"model": "anthropic/claude", "messages": [], "classifier": "agent-one"}, + ) + + assert openai_context.session_isolation_key == anthropic_context.session_isolation_key + assert openai_context.session_isolation_key.startswith("classifier:") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("input_protocol", "payload", "source_only_key"), + [ + ( + "anthropic_messages", + {"model": "anthropic/claude", "system": "rule", "messages": [{"role": "user", "content": "hello"}]}, + "system", + ), + ( + "responses", + {"model": "openai/gpt-5.1", "instructions": "rule", "input": "hello"}, + "input", + ), + ( + "gemini", + {"model": "gemini/gemini-3", "systemInstruction": {"parts": [{"text": "rule"}]}, "contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + "contents", + ), + ], +) +async def test_request_builder_keeps_raw_and_canonical_protocol_views( + monkeypatch, + input_protocol, + payload, + source_only_key, +) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + context = await _builder().build_completion_context( + None, + None, + {**payload, "_input_protocol": input_protocol, "_output_protocol": "anthropic_messages"}, + ) + + assert context.input_protocol_name == input_protocol + assert context.output_protocol_name == "anthropic_messages" + assert context.protocol_request[source_only_key] == payload[source_only_key] + assert context.unified_request.source_protocol == input_protocol + assert "messages" in context.kwargs + if source_only_key != "system": + assert source_only_key not in context.kwargs + + +@pytest.mark.asyncio +async def test_request_builder_consumes_provider_continuation_control(monkeypatch) -> None: + monkeypatch.delenv("FALLBACK_GROUPS", raising=False) + context = await _builder().build_completion_context( + None, + None, + { + "model": "openai/gpt-5.1", + "messages": [{"role": "user", "content": "hello"}], + "_disable_provider_continuation": True, + }, + ) + + assert context.disable_provider_continuation is True + assert "_disable_provider_continuation" not in context.protocol_request + assert "_disable_provider_continuation" not in context.kwargs + + +def test_only_typed_internal_hints_can_contribute_global_identity() -> None: + forged_dict = { + "global_strong_anchors": ["forged-global"], + "strong_anchors": ["ordinary-internal"], + } + provider_hints = SessionTrackingHints( + global_strong_anchors=["forged-provider-global"], + strong_anchors=["provider-native"], + ) + + rejected = RequestContextBuilder._merge_session_hints(forged_dict, provider_hints) + self_owned = RequestContextBuilder._merge_session_hints( + SessionTrackingHints(global_strong_anchors=["proxy-owned-global"]), + provider_hints, + ) + + assert rejected.global_strong_anchors == [] + assert rejected.strong_anchors == ["ordinary-internal", "provider-native"] + assert self_owned.global_strong_anchors == ["proxy-owned-global"] diff --git a/tests/test_request_executor_fallback_error_summary.py b/tests/test_request_executor_fallback_error_summary.py new file mode 100644 index 000000000..bab3ea55b --- /dev/null +++ b/tests/test_request_executor_fallback_error_summary.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from types import MethodType + +import pytest + +from rotator_library.client.executor import RequestExecutor +from rotator_library.core.types import RequestContext +from rotator_library.routing import parse_route_target + + +def _context() -> RequestContext: + return RequestContext( + model="code", + provider="requested", + kwargs={"model": "code", "messages": []}, + streaming=False, + credentials=["cred-a"], + deadline=9999999999.0, + routing_targets=(parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1")), + routing_group_name="code_chain", + ) + + +@pytest.mark.asyncio +async def test_fallback_summary_includes_all_structured_target_failures_without_credentials() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_execute(self, context): + attempts.append(context.provider) + return { + "error": { + "type": "proxy_all_credentials_exhausted", + "message": f"{context.provider} failed for cred-secret-value", + "details": {"normal_error_summary": "1 rate_limit"}, + } + } + + executor._execute_non_streaming = MethodType(fake_execute, executor) + + response = await executor._execute_non_streaming_with_fallback(_context()) + + fallback_targets = response["error"]["details"]["fallback_targets"] + assert attempts == ["codex", "openai"] + assert [failure["provider"] for failure in fallback_targets] == ["codex", "openai"] + assert [failure["error_type"] for failure in fallback_targets] == ["rate_limit", "rate_limit"] + assert "credential" not in str(fallback_targets).lower() + assert "cred-secret-value" not in str(fallback_targets) diff --git a/tests/test_request_executor_fallback_groups.py b/tests/test_request_executor_fallback_groups.py new file mode 100644 index 000000000..d35e79035 --- /dev/null +++ b/tests/test_request_executor_fallback_groups.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +import json +from types import MethodType + +import pytest + +from rotator_library.client.executor import RequestExecutor +from rotator_library.core.types import RequestContext +from rotator_library.core.errors import StructuredAPIResponseError, structured_api_response_error +from rotator_library.routing import parse_route_target +from rotator_library.routing.types import FallbackGroup +from rotator_library.transaction_logger import TransactionLogger + + +class ClassifiedFailure(Exception): + def __init__(self, error_type: str) -> None: + super().__init__(error_type) + self.error_type = error_type + + +def _context(*, routing_targets=None, logger=None, routing_group=None) -> RequestContext: + return RequestContext( + model="code", + provider="requested", + kwargs={"model": "code", "messages": []}, + streaming=False, + credentials=["cred-a"], + deadline=9999999999.0, + transaction_logger=logger, + routing_targets=routing_targets, + routing_group_name="code_chain" if routing_targets else None, + routing_group=routing_group, + ) + + +def _executor_with_attempts(attempts): + executor = RequestExecutor.__new__(RequestExecutor) + + async def fake_execute(self, context): + attempts.append(context) + if len(attempts) == 1: + raise ClassifiedFailure("rate_limit") + return {"id": "ok", "model": context.model} + + executor._execute_non_streaming = MethodType(fake_execute, executor) + return executor + + +@pytest.mark.asyncio +async def test_non_streaming_fallback_group_tries_next_target_on_retryable_error() -> None: + attempts = [] + targets = (parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1")) + + context = _context(routing_targets=targets) + result = await _executor_with_attempts(attempts)._execute_non_streaming_with_fallback(context) + + assert result == {"id": "ok", "model": "openai/gpt-5.1"} + assert [attempt.provider for attempt in attempts] == ["codex", "openai"] + assert [attempt.kwargs["model"] for attempt in attempts] == ["codex/gpt-5.1-codex", "openai/gpt-5.1"] + assert context.routing_attempt_history[0]["error_type"] == "rate_limit" + assert context.routing_attempt_history[0]["fallback_allowed"] is True + assert context.routing_attempt_history[1]["success"] is True + + +@pytest.mark.asyncio +async def test_structured_numeric_rate_limit_error_reaches_fallback_policy() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_execute(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + raise structured_api_response_error({"error": {"status": 429, "message": "busy"}}) + return {"id": "ok", "model": context.model} + + executor._execute_non_streaming = MethodType(fake_execute, executor) + targets = (parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1")) + context = _context(routing_targets=targets) + + result = await executor._execute_non_streaming_with_fallback(context) + + assert result["id"] == "ok" + assert attempts == ["codex", "openai"] + assert context.routing_attempt_history[0]["error_type"] == "rate_limit" + + +@pytest.mark.asyncio +async def test_non_streaming_fallback_group_falls_back_on_quota_exceeded() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_execute(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + raise ClassifiedFailure("quota_exceeded") + return {"id": "ok", "model": context.model} + + executor._execute_non_streaming = MethodType(fake_execute, executor) + + result = await executor._execute_non_streaming_with_fallback( + _context(routing_targets=(parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1"))) + ) + + assert result == {"id": "ok", "model": "openai/gpt-5.1"} + assert attempts == ["codex", "openai"] + + +@pytest.mark.asyncio +async def test_non_streaming_fallback_group_stops_on_permanent_error() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_execute(self, context): + attempts.append(context) + raise ClassifiedFailure("validation") + + executor._execute_non_streaming = MethodType(fake_execute, executor) + + with pytest.raises(ClassifiedFailure): + await executor._execute_non_streaming_with_fallback( + _context(routing_targets=(parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1"))) + ) + + assert len(attempts) == 1 + + +@pytest.mark.asyncio +async def test_non_streaming_fallback_group_hard_stops_group_override() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_execute(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + raise ClassifiedFailure("authentication") + return {"id": "ok", "model": context.model} + + executor._execute_non_streaming = MethodType(fake_execute, executor) + + targets = (parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1")) + with pytest.raises(ClassifiedFailure): + await executor._execute_non_streaming_with_fallback( + _context( + routing_targets=targets, + routing_group=FallbackGroup( + name="code_chain", + targets=targets, + failover_on=frozenset({"authentication"}), + stop_on=frozenset({"validation"}), + ), + ) + ) + + assert attempts == ["codex"] + + +@pytest.mark.asyncio +async def test_non_streaming_fallback_group_handles_structured_error_response() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_execute(self, context): + attempts.append(context) + if len(attempts) == 1: + return {"error": {"type": "proxy_all_credentials_exhausted", "details": {"normal_error_summary": "2 rate_limit"}}} + return {"id": "ok", "model": context.model} + + executor._execute_non_streaming = MethodType(fake_execute, executor) + + result = await executor._execute_non_streaming_with_fallback( + _context(routing_targets=(parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1"))) + ) + + assert result == {"id": "ok", "model": "openai/gpt-5.1"} + + +@pytest.mark.asyncio +async def test_raised_credential_exhaustion_keeps_fallback_and_final_summary() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + def exhausted(provider: str) -> StructuredAPIResponseError: + response = { + "error": { + "type": "proxy_all_credentials_exhausted", + "message": f"{provider} exhausted", + "details": {"normal_error_summary": "2 rate_limit"}, + } + } + return StructuredAPIResponseError( + f"{provider} exhausted", + error_type="proxy_all_credentials_exhausted", + status_code=503, + response=response, + ) + + async def fake_execute(self, context): + attempts.append(context.provider) + raise exhausted(context.provider) + + executor._execute_non_streaming = MethodType(fake_execute, executor) + targets = ( + parse_route_target("codex/gpt-5.1-codex"), + parse_route_target("openai/gpt-5.1"), + ) + context = _context(routing_targets=targets) + + with pytest.raises(StructuredAPIResponseError) as raised: + await executor._execute_non_streaming_with_fallback(context) + + assert attempts == ["codex", "openai"] + assert [entry["error_type"] for entry in context.routing_attempt_history] == [ + "rate_limit", + "rate_limit", + ] + assert [ + failure["provider"] + for failure in raised.value.response["error"]["details"]["fallback_targets"] + ] == ["codex", "openai"] + + +@pytest.mark.asyncio +async def test_non_streaming_fallback_group_emits_routing_trace(tmp_path) -> None: + attempts = [] + logger = TransactionLogger("routing", "code", parent_dir=tmp_path) + targets = (parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1")) + + await _executor_with_attempts(attempts)._execute_non_streaming_with_fallback(_context(routing_targets=targets, logger=logger)) + + pass_names = [json.loads(line)["pass_name"] for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + assert "routing_decision" in pass_names + assert pass_names.count("routing_target_attempt_started") == 2 + assert "routing_target_attempt_failed" in pass_names + assert "routing_fallback_selected" in pass_names + assert "routing_target_attempt_succeeded" in pass_names + + +@pytest.mark.asyncio +async def test_non_streaming_fallback_group_records_usage_only_on_successful_target() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + recorded = [] + + async def fake_execute(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + raise ClassifiedFailure("rate_limit") + recorded.append(context.provider) + return {"id": "ok", "model": context.model} + + executor._execute_non_streaming = MethodType(fake_execute, executor) + + result = await executor._execute_non_streaming_with_fallback( + _context(routing_targets=(parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1"))) + ) + + assert result == {"id": "ok", "model": "openai/gpt-5.1"} + assert attempts == ["codex", "openai"] + assert recorded == ["openai"] diff --git a/tests/test_request_executor_native_routing.py b/tests/test_request_executor_native_routing.py new file mode 100644 index 000000000..72d409ef6 --- /dev/null +++ b/tests/test_request_executor_native_routing.py @@ -0,0 +1,642 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.client import executor as executor_module +from rotator_library.client.executor import RequestExecutor, RoutingExecutionError +from rotator_library.core.types import RequestContext +from rotator_library.field_cache import FieldCacheInjection, FieldCacheRule +from rotator_library.config.experimental import load_config_from_mapping +from rotator_library.providers.antigravity_provider import AntigravityProvider +from rotator_library.providers.claude_code_provider import ClaudeCodeProvider +from rotator_library.providers.codex_provider import CodexProvider +from rotator_library.providers.copilot_provider import CopilotProvider +from rotator_library.protocols import get_protocol +from rotator_library.routing import parse_route_target + + +class FakeNativeResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + +class FakeHTTPClient: + def __init__(self): + self.calls = [] + + async def post(self, endpoint, *, headers, json): + self.calls.append({"endpoint": endpoint, "headers": headers, "json": json}) + return FakeNativeResponse({"id": "chat_native", "choices": [{"message": {"role": "assistant", "content": "ok"}}]}) + + +class SequencedHTTPClient: + def __init__(self, responses): + self.responses = list(responses) + self.calls = [] + + async def post(self, endpoint, *, headers, json): + self.calls.append({"endpoint": endpoint, "headers": headers, "json": json}) + return FakeNativeResponse(self.responses.pop(0)) + + +class NativePlugin: + def has_custom_logic(self): + return False + + def get_protocol_name(self, model=""): + return "openai_chat" + + def get_native_endpoint(self, model="", operation="chat"): + return "https://native.test/chat" + + def get_native_headers(self, credential_identifier, model="", operation="chat"): + return {"Authorization": f"Bearer {credential_identifier}", "X-Operation": operation, "X-Model": model} + + def get_native_operation(self, model="", request=None, stream=False): + return "messages" if stream else "chat" + + def normalize_native_model(self, model=""): + return model.split("/", 1)[1] if "/" in model else model + + def get_adapter_names(self, model=""): + return () + + def get_adapter_config(self, model=""): + return {} + + def get_field_cache_rules(self, model=""): + return () + + +class NativePluginWithRule(NativePlugin): + def get_field_cache_rules(self, model=""): + return ( + FieldCacheRule( + name="state", + source="response", + path="choices.0.message.reasoning_content", + inject=FieldCacheInjection(target="request", path="metadata.state"), + ), + ) + + +class NativePluginWithVendorRule(NativePlugin): + def get_field_cache_rules(self, model=""): + return ( + FieldCacheRule( + name="vendor_state", + source="response", + path="choices.0.message.vendor_state", + inject=FieldCacheInjection(target="request", path="metadata.vendor_state"), + allow_missing_session=True, + scope=("provider", "model"), + ), + ) + + +class NativePluginWithStreamVendorRule(NativePlugin): + def get_field_cache_rules(self, model=""): + return ( + FieldCacheRule( + name="vendor_state", + source="stream_event", + path="raw.choices.0.message.vendor_state", + allow_missing_session=True, + scope=("provider", "model"), + ), + ) + + +class NativeOptOutPlugin(NativePlugin): + def should_use_native_protocol(self, model="", operation="chat", *, stream=False, execution="auto"): + return False + + +class CustomPlugin: + def __init__(self): + self.calls = [] + + def has_custom_logic(self): + return True + + async def acompletion(self, client, **kwargs): + self.calls.append(kwargs) + return {"id": "custom"} + + def get_protocol_name(self, model=""): + return "gemini" + + def get_native_endpoint(self, model="", operation="chat"): + return "https://native.test/should-not-run" + + def get_native_headers(self, credential_identifier, model="", operation="chat"): + return {"Authorization": f"Bearer {credential_identifier}"} + + +def _context(target=None) -> RequestContext: + kwargs = {"model": "provider/gpt-test", "messages": [{"role": "user", "content": "hi"}]} + return RequestContext( + model="provider/gpt-test", + provider="provider", + kwargs=kwargs, + streaming=False, + credentials=["cred"], + deadline=9999999999.0, + routing_targets=(target,) if target else None, + protocol_request=dict(kwargs), + unified_request=get_protocol("openai_chat").parse_request(kwargs), + input_provider="provider", + session_id="session-a", + ) + + +def _provider_context(provider: str, model: str, kwargs: dict, target=None) -> RequestContext: + return RequestContext( + model=model, + provider=provider, + kwargs=kwargs, + streaming=False, + credentials=["cred"], + deadline=9999999999.0, + routing_targets=(target,) if target else None, + protocol_request=dict(kwargs), + unified_request=get_protocol("openai_chat").parse_request(kwargs), + input_provider=provider, + session_id="session-a", + ) + + +def _executor(http_client=None) -> RequestExecutor: + executor = RequestExecutor.__new__(RequestExecutor) + executor._http_client = http_client or FakeHTTPClient() + executor._apply_litellm_logger = lambda kwargs: None + return executor + + +@pytest.mark.asyncio +async def test_native_declared_provider_uses_native_executor_in_auto_mode() -> None: + http_client = FakeHTTPClient() + target = parse_route_target("provider/gpt-test") + context = _context(target) + context.routing_target_index = 0 + + response = await _executor(http_client)._execute_provider_request( + "provider", "provider/gpt-test", NativePlugin(), "secret", "stable", dict(context.kwargs), context + ) + + assert response["id"] == "chat_native" + assert http_client.calls[0]["endpoint"] == "https://native.test/chat" + assert http_client.calls[0]["headers"]["Authorization"] == "Bearer secret" + assert http_client.calls[0]["headers"]["X-Operation"] == "chat" + assert http_client.calls[0]["headers"]["X-Model"] == "gpt-test" + assert http_client.calls[0]["json"]["model"] == "gpt-test" + + +@pytest.mark.asyncio +async def test_auto_mode_prefers_custom_logic_over_native_declaration() -> None: + http_client = FakeHTTPClient() + plugin = CustomPlugin() + context = _context(parse_route_target("provider/gpt-test")) + context.routing_target_index = 0 + + response = await _executor(http_client)._execute_provider_request( + "provider", "provider/gpt-test", plugin, "secret", "stable", dict(context.kwargs), context + ) + + assert response == {"id": "custom"} + assert plugin.calls[0]["credential_identifier"] == "secret" + assert http_client.calls == [] + + +@pytest.mark.asyncio +async def test_request_executor_reuses_native_field_cache_store() -> None: + http_client = SequencedHTTPClient( + [ + {"id": "chat_1", "choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "cached"}}]}, + {"id": "chat_2", "choices": [{"message": {"role": "assistant", "content": "ok"}}]}, + ] + ) + executor = _executor(http_client) + context = _context(parse_route_target("provider/gpt-test")) + context.routing_target_index = 0 + + await executor._execute_provider_request("provider", "provider/gpt-test", NativePluginWithRule(), "secret", "stable", dict(context.kwargs), context) + await executor._execute_provider_request("provider", "provider/gpt-test", NativePluginWithRule(), "secret", "stable", dict(context.kwargs), context) + + assert http_client.calls[1]["json"]["metadata"]["state"] == "cached" + + +def test_native_context_merges_json_field_cache_rules(monkeypatch, tmp_path) -> None: + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "field_cache": { + "provider": { + "*": [ + { + "name": "state", + "source": "response", + "path": "json.path", + "inject": {"target": "request", "path": "metadata.state"}, + }, + { + "name": "extra", + "source": "response", + "path": "json.extra", + "inject": {"target": "request", "path": "metadata.extra"}, + }, + ] + } + } + } + ), + encoding="utf-8", + ) + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", str(config_path)) + + native_context = _executor()._build_native_provider_context( + "provider", + "provider/gpt-test", + NativePluginWithRule(), + "secret", + "stable", + _context(), + None, + ) + + assert [rule.name for rule in native_context.field_cache_rules] == ["state", "extra"] + assert native_context.field_cache_rules[0].path == "json.path" + + +def test_native_context_uses_bound_field_cache_snapshot(monkeypatch, tmp_path) -> None: + startup = load_config_from_mapping({ + "field_cache": { + "provider": { + "*": [{"name": "startup", "source": "response", "path": "startup.path"}] + } + } + }) + changed_path = tmp_path / "changed.json" + changed_path.write_text(json.dumps({ + "field_cache": { + "provider": { + "*": [{"name": "changed", "source": "response", "path": "changed.path"}] + } + } + }), encoding="utf-8") + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", str(changed_path)) + executor = _executor() + executor._experimental_config = startup + + native_context = executor._build_native_provider_context( + "provider", + "provider/gpt-test", + NativePlugin(), + "secret", + "stable", + _context(), + None, + ) + + assert [rule.name for rule in native_context.field_cache_rules] == ["startup"] + + +def test_native_context_rejects_behavior_weakening_rule_override() -> None: + config = load_config_from_mapping({ + "field_cache": { + "provider": { + "*": [{ + "name": "state", + "source": "response", + "path": "changed.path", + "inject": {"target": "request", "path": "metadata.different"}, + }] + } + } + }) + executor = _executor() + executor._experimental_config = config + + with pytest.raises(RoutingExecutionError, match="cannot weaken"): + executor._build_native_provider_context( + "provider", + "provider/gpt-test", + NativePluginWithRule(), + "secret", + "stable", + _context(), + None, + ) + + +@pytest.mark.asyncio +async def test_claude_code_provider_runs_mock_live_native_request(monkeypatch) -> None: + monkeypatch.setenv("CLAUDE_CODE_API_BASE", "https://claude-code.test") + http_client = SequencedHTTPClient([ + {"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "ok"}], "usage": {"input_tokens": 1, "output_tokens": 1}} + ]) + provider = ClaudeCodeProvider() + target = parse_route_target("claude_code/claude-sonnet-4-5") + context = _provider_context( + "claude_code", + "claude_code/claude-sonnet-4-5", + {"model": "claude_code/claude-sonnet-4-5", "messages": [{"role": "developer", "content": "rules"}, {"role": "user", "content": "hi"}]}, + target, + ) + context.routing_target_index = 0 + + response = await _executor(http_client)._execute_provider_request("claude_code", context.model, provider, "secret", "stable", dict(context.kwargs), context) + + assert response["id"] == "msg_1" + assert response["choices"][0]["message"]["content"] == "ok" + assert http_client.calls[0]["endpoint"] == "https://claude-code.test/v1/messages" + assert http_client.calls[0]["json"]["model"] == "claude-sonnet-4-5" + assert http_client.calls[0]["json"]["max_tokens"] == 4096 + assert http_client.calls[0]["json"]["messages"][0]["role"] == "user" + + +@pytest.mark.asyncio +async def test_codex_provider_runs_mock_live_native_request(monkeypatch) -> None: + monkeypatch.setenv("CODEX_API_BASE", "https://codex.test") + http_client = SequencedHTTPClient([ + {"id": "resp_1", "object": "response", "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}]} + ]) + provider = CodexProvider() + target = parse_route_target("codex/gpt-5.1-codex") + context = _provider_context("codex", "codex/gpt-5.1-codex", {"model": "codex/gpt-5.1-codex", "messages": [{"role": "user", "content": "hi"}]}, target) + context.routing_target_index = 0 + + response = await _executor(http_client)._execute_provider_request("codex", context.model, provider, "secret", "stable", dict(context.kwargs), context) + + assert response["id"] == "resp_1" + assert response["choices"][0]["message"]["content"] == "ok" + assert http_client.calls[0]["endpoint"] == "https://codex.test/v1/responses" + assert http_client.calls[0]["json"]["model"] == "gpt-5.1-codex" + assert http_client.calls[0]["json"]["input"][0]["content"] == [{"type": "input_text", "text": "hi"}] + assert "messages" not in http_client.calls[0]["json"] + + +@pytest.mark.asyncio +async def test_copilot_provider_runs_mock_live_native_request(monkeypatch) -> None: + monkeypatch.setenv("COPILOT_API_BASE", "https://copilot.test") + http_client = SequencedHTTPClient([ + {"id": "chat_1", "choices": [{"message": {"role": "assistant", "content": "ok"}}]} + ]) + provider = CopilotProvider() + target = parse_route_target("copilot/gpt-4.1") + context = _provider_context("copilot", "copilot/gpt-4.1", {"model": "copilot/gpt-4.1", "messages": [{"role": "developer", "content": "rules"}, {"role": "user", "content": "hi"}]}, target) + context.routing_target_index = 0 + + response = await _executor(http_client)._execute_provider_request("copilot", context.model, provider, "secret", "stable", dict(context.kwargs), context) + + assert response["id"] == "chat_1" + assert http_client.calls[0]["endpoint"] == "https://copilot.test/chat/completions" + assert http_client.calls[0]["json"]["model"] == "gpt-4.1" + assert http_client.calls[0]["json"]["messages"][0]["role"] == "system" + + +@pytest.mark.asyncio +async def test_antigravity_provider_runs_mock_live_native_request(monkeypatch) -> None: + monkeypatch.setenv("ANTIGRAVITY_API_BASE", "https://antigravity.test/v1internal") + http_client = SequencedHTTPClient([ + {"candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}], "usageMetadata": {"totalTokenCount": 2}} + ]) + provider = AntigravityProvider() + target = parse_route_target("antigravity/claude-sonnet-4.5") + context = _provider_context("antigravity", "antigravity/claude-sonnet-4.5", {"model": "antigravity/claude-sonnet-4.5", "messages": [{"role": "user", "content": "hi"}]}, target) + context.routing_target_index = 0 + + response = await _executor(http_client)._execute_provider_request("antigravity", context.model, provider, "secret", "stable", dict(context.kwargs), context) + + assert response["choices"][0]["message"]["content"] == "ok" + assert http_client.calls[0]["endpoint"] == "https://antigravity.test/v1internal:generateContent" + assert http_client.calls[0]["json"]["model"] == "claude-sonnet-4-5" + assert http_client.calls[0]["json"]["request"]["contents"][0]["parts"][0]["text"] == "hi" + assert http_client.calls[0]["json"]["requestType"] == "CHAT_COMPLETION" + assert "requestId" in http_client.calls[0]["json"] + + +def test_claude_code_header_modes(monkeypatch) -> None: + provider = ClaudeCodeProvider() + + monkeypatch.setenv("CLAUDE_CODE_AUTH_HEADER", "auto") + assert provider.get_native_headers("sk-ant-test")["x-api-key"] == "sk-ant-test" + assert provider.get_native_headers("oauth-token")["Authorization"] == "Bearer oauth-token" + monkeypatch.setenv("CLAUDE_CODE_AUTH_HEADER", "x-api-key") + assert provider.get_native_headers("any-token")["x-api-key"] == "any-token" + + +def test_antigravity_alias_normalization_preserves_thinking_level() -> None: + provider = AntigravityProvider() + request = {"_proxy_model": "antigravity/gemini-3-pro-low", "model": "gemini-3-pro-preview", "contents": [{"role": "user", "parts": [{"text": "hi"}]}]} + + prepared = provider.prepare_native_request(request, model=provider.normalize_native_model("antigravity/gemini-3-pro-low"), operation="generate") + + assert provider.normalize_native_model("antigravity/gemini-3-pro-low") == "gemini-3-pro-preview" + assert prepared["model"] == "gemini-3-pro-low" + assert prepared["generationConfig"]["thinkingConfig"]["thinkingLevel"] == "low" + assert prepared["metadata"]["thinking_level"] == "low" + assert "_proxy_model" not in prepared + assert prepared["contents"][0]["parts"][0]["text"] == "hi" + + +def test_auto_native_selection_honors_provider_opt_out() -> None: + assert executor_module._should_use_native_protocol( + NativeOptOutPlugin(), + "provider/gpt-test", + None, + {"model": "provider/gpt-test", "messages": []}, + stream=False, + execution="auto", + ) is False + + +def test_auto_native_streaming_selection_honors_provider_opt_out() -> None: + assert executor_module._should_use_native_streaming( + NativeOptOutPlugin(), + "provider/gpt-test", + parse_route_target("provider/gpt-test"), + "auto", + "provider", + ) is False + + +def test_native_operation_selection_never_receives_client_or_litellm_payload() -> None: + class OperationProbe(NativePlugin): + request = object() + + def get_native_operation(self, model="", request=None, stream=False): + self.request = request + return "chat" + + plugin = OperationProbe() + + assert executor_module._should_use_native_protocol( + plugin, + "provider/gpt-test", + parse_route_target("provider/gpt-test"), + {"messages": [], "custom_llm_provider": "openai", "api_base": "https://litellm-only.test"}, + stream=False, + execution="auto", + ) is True + assert plugin.request is None + + +def test_route_error_type_from_response_hard_stop_wins_over_retry_summary() -> None: + response = { + "error": { + "type": "authentication", + "message": "provider said quota secret-token", + "details": {"normal_error_summary": "rate_limit quota capacity", "status_code": 401}, + } + } + + assert executor_module._route_error_type_from_response(response) == "authentication" + + +def test_route_error_type_from_response_uses_structured_status_codes() -> None: + assert executor_module._route_error_type_from_response({"error": {"code": 403}}) == "forbidden" + assert executor_module._route_error_type_from_response({"error": {"status": 401}}) == "authentication" + assert executor_module._route_error_type_from_response({"error": {"details": {"status": 403}}}) == "forbidden" + + +def test_route_error_type_from_response_uses_structured_aliases() -> None: + assert executor_module._route_error_type_from_response({"error": {"type": "invalid_api_key"}}) == "authentication" + assert executor_module._route_error_type_from_response({"error": {"code": "invalid_argument"}}) == "invalid_request" + assert executor_module._route_error_type_from_response({"error": {"code": "resource_exhausted"}}) == "quota_exceeded" + assert executor_module._route_error_type_from_response({"error": {"code": "unavailable"}}) == "server_error" + assert executor_module._route_error_type_from_response({"error": {"code": "deadline_exceeded"}}) == "api_connection" + assert executor_module._route_error_type_from_response({"error": {"details": {"status_code": 503}}}) == "server_error" + + +def test_route_error_type_from_response_reads_abnormal_errors_before_proxy_summary() -> None: + response = { + "error": { + "type": "proxy_all_credentials_exhausted", + "details": { + "normal_error_summary": "rate_limit quota", + "abnormal_errors": [{"error_type": "authentication", "status_code": 401}], + }, + } + } + + assert executor_module._route_error_type_from_response(response) == "authentication" + + +def test_stream_chunk_error_type_detects_terminal_error_frames() -> None: + assert executor_module._stream_chunk_error_type('data: {"error":{"type":"rate_limit"}}\n\n') == "rate_limit" + assert executor_module._stream_chunk_error_type('event: response.failed\ndata: {"error":{"type":"authentication"}}\n\n') == "authentication" + assert executor_module._stream_chunk_error_type('event: error\ndata: {"type":"error","code":429}\n\n') == "rate_limit" + assert executor_module._stream_chunk_error_type('event: error\ndata: {"type":"error","code":"context_length_exceeded"}\n\n') == "context_window_exceeded" + + +def test_target_failure_summary_is_structural_and_sanitized() -> None: + summary = executor_module._target_failure_summary(parse_route_target("openai/gpt"), "rate-limit", status_code=429) + + assert summary["error_type"] == "rate_limit" + assert summary["status_code"] == 429 + assert summary["message"] == "" + + +def test_explicit_native_streaming_fails_when_provider_does_not_support_it() -> None: + target = parse_route_target("provider/gpt-test@native") + + with pytest.raises(RoutingExecutionError) as exc: + executor_module._should_use_native_streaming(NativePlugin(), "provider/gpt-test", target, "native", "provider") + + assert exc.value.error_type == "configuration_error" + + +def test_antigravity_cache_injection_targets_safe_envelope() -> None: + provider = AntigravityProvider() + + rules = provider.get_field_cache_rules("gemini-3-flash") + + assert rules[0].inject.path == "request.metadata.thoughtSignatures" + + +def test_native_context_raises_on_invalid_field_cache_config(monkeypatch, tmp_path) -> None: + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"field_cache": {"provider": {"*": [{"name": "bad", "source": "response"}]}}}), encoding="utf-8") + monkeypatch.setenv("LLM_PROXY_CONFIG_FILE", str(config_path)) + + with pytest.raises(RoutingExecutionError) as exc: + _executor()._build_native_provider_context("provider", "provider/gpt-test", NativePlugin(), "secret", "stable", _context(), None) + + assert exc.value.error_type == "configuration_error" + + +def test_executor_trace_redaction_uses_native_field_cache_response_paths() -> None: + context = _context(parse_route_target("provider/gpt-test")) + payload = {"choices": [{"message": {"role": "assistant", "content": "ok", "vendor_state": "opaque-vendor-state"}}]} + + redacted = executor_module._redact_context_field_cache_paths(payload, context, "response", NativePluginWithVendorRule()) + + assert redacted["choices"][0]["message"]["vendor_state"] == "[REDACTED]" + assert payload["choices"][0]["message"]["vendor_state"] == "opaque-vendor-state" + + +def test_executor_stream_trace_redaction_uses_native_field_cache_paths() -> None: + context = _context(parse_route_target("provider/gpt-test")) + sse_line = 'data: {"choices":[{"delta":{"content":"ok"},"message":{"vendor_state":"opaque-vendor-state"}}]}\n\n' + + redacted = executor_module._redact_stream_sse_for_trace(sse_line, context, NativePluginWithStreamVendorRule()) + + parsed = json.loads(redacted[6:].strip()) + + assert "opaque-vendor-state" not in redacted + assert parsed["choices"][0]["message"]["vendor_state"] == "[REDACTED]" + + +@pytest.mark.asyncio +async def test_litellm_fallback_execution_is_explicit(monkeypatch) -> None: + calls = [] + + async def fake_acompletion(**kwargs): + calls.append(kwargs) + return {"id": "litellm"} + + monkeypatch.setattr(executor_module.litellm, "acompletion", fake_acompletion) + target = parse_route_target("provider/gpt-test@litellm_fallback") + context = _context(target) + context.routing_target_index = 0 + + response = await _executor()._execute_provider_request("provider", "provider/gpt-test", NativePlugin(), "secret", "stable", dict(context.kwargs), context) + + assert response == {"id": "litellm"} + assert calls[0]["api_key"] == "secret" + + +@pytest.mark.asyncio +async def test_custom_execution_mode_requires_custom_plugin() -> None: + target = parse_route_target("provider/gpt-test@custom") + context = _context(target) + context.routing_target_index = 0 + + plugin = CustomPlugin() + + response = await _executor()._execute_provider_request("provider", "provider/gpt-test", plugin, "secret", "stable", dict(context.kwargs), context) + + assert response == {"id": "custom"} + assert plugin.calls[0]["credential_identifier"] == "secret" + + +@pytest.mark.asyncio +async def test_native_execution_mode_fails_when_provider_has_no_native_declaration() -> None: + target = parse_route_target("provider/gpt-test@native") + context = _context(target) + context.routing_target_index = 0 + + with pytest.raises(RoutingExecutionError) as exc: + await _executor()._execute_provider_request("provider", "provider/gpt-test", None, "secret", "stable", dict(context.kwargs), context) + + assert exc.value.error_type == "configuration_error" diff --git a/tests/test_request_executor_stream_metrics.py b/tests/test_request_executor_stream_metrics.py new file mode 100644 index 000000000..c7a04b0a9 --- /dev/null +++ b/tests/test_request_executor_stream_metrics.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +import json +import asyncio + +import pytest + +from rotator_library.client.streaming import StreamingHandler +from rotator_library.core.errors import StreamedAPIError +from rotator_library.transaction_logger import TransactionLogger + + +async def _chunks(): + yield {"id": "chunk_1", "choices": [{"delta": {"content": "hi"}}]} + yield {"id": "chunk_2", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + + +class HangingStream: + def __init__(self) -> None: + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + await asyncio.sleep(1) + return {"choices": [{"delta": {"content": "late"}}]} + + async def aclose(self) -> None: + self.closed = True + + +class DelayedStream: + def __init__(self, delay: float = 0.03) -> None: + self.delay = delay + self.index = 0 + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index == 0: + self.index += 1 + await asyncio.sleep(self.delay) + return {"id": "chunk_1", "choices": [{"delta": {"content": "hi"}}]} + if self.index == 1: + self.index += 1 + return {"id": "chunk_2", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + raise StopAsyncIteration + + async def aclose(self) -> None: + self.closed = True + + +class FirstThenHangStream(HangingStream): + def __init__(self) -> None: + super().__init__() + self.index = 0 + + async def __anext__(self): + if self.index == 0: + self.index += 1 + return {"id": "chunk_1", "choices": [{"delta": {}}]} + return await super().__anext__() + + +class DisconnectedRequest: + async def is_disconnected(self) -> bool: + return True + + +class DelayedDisconnectedRequest: + def __init__(self) -> None: + self.calls = 0 + + async def is_disconnected(self) -> bool: + self.calls += 1 + await asyncio.sleep(0.01) + return self.calls >= 1 + + +def _trace_passes(log_dir): + return [json.loads(line)["pass_name"] for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +@pytest.mark.asyncio +async def test_streaming_handler_emits_lifecycle_metrics_without_changing_output(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(_chunks(), "cred", "openai/gpt-test", transaction_logger=logger)] + + assert chunks[0].startswith("data: ") + assert chunks[-1] == "data: [DONE]\n\n" + pass_names = _trace_passes(logger.log_dir) + assert "stream_started" in pass_names + assert "stream_first_byte" in pass_names + assert "stream_first_visible_output" in pass_names + assert "stream_completed" in pass_names + assert "stream_metrics_final" in pass_names + + +@pytest.mark.asyncio +async def test_stream_trace_metrics_can_be_disabled_without_changing_output(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("STREAM_TRACE_METRICS", "false") + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(_chunks(), "cred", "openai/gpt-test", transaction_logger=logger)] + + assert chunks[0].startswith("data: ") + assert chunks[-1] == "data: [DONE]\n\n" + if (logger.log_dir / "transform_trace.jsonl").exists(): + pass_names = _trace_passes(logger.log_dir) + assert "stream_started" in pass_names + assert "stream_metrics_final" in pass_names + + +@pytest.mark.asyncio +async def test_streaming_handler_closes_upstream_on_client_disconnect(monkeypatch) -> None: + monkeypatch.delenv("STREAM_TTFB_TIMEOUT_SECONDS", raising=False) + stream = HangingStream() + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(stream, "cred", "openai/gpt-test", request=DisconnectedRequest())] + + assert chunks == [] + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_streaming_handler_closes_upstream_when_disconnect_happens_during_wait(monkeypatch) -> None: + monkeypatch.delenv("STREAM_TTFB_TIMEOUT_SECONDS", raising=False) + monkeypatch.delenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", raising=False) + stream = HangingStream() + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(stream, "cred", "openai/gpt-test", request=DelayedDisconnectedRequest())] + + assert chunks == [] + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_streaming_handler_emits_configured_heartbeats(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + monkeypatch.delenv("STREAM_TTFB_TIMEOUT_SECONDS", raising=False) + monkeypatch.delenv("STREAM_STALL_TIMEOUT_SECONDS", raising=False) + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(DelayedStream(), "cred", "openai/gpt-test")] + + assert any(chunk.startswith(": heartbeat") for chunk in chunks) + assert chunks[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_streaming_handler_ttfb_timeout_closes_upstream(monkeypatch) -> None: + monkeypatch.setenv("STREAM_TTFB_TIMEOUT_SECONDS", "0.01") + monkeypatch.delenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", raising=False) + stream = HangingStream() + + with pytest.raises(StreamedAPIError) as exc: + _ = [chunk async for chunk in StreamingHandler().wrap_stream(stream, "cred", "openai/gpt-test")] + + assert stream.closed is True + assert exc.value.data["error"]["details"]["timeout_type"] == "ttfb" + + +@pytest.mark.asyncio +async def test_stream_timeout_closes_upstream_even_when_disconnect_close_disabled(monkeypatch) -> None: + monkeypatch.setenv("STREAM_TTFB_TIMEOUT_SECONDS", "0.01") + monkeypatch.setenv("STREAM_CANCEL_UPSTREAM_ON_DISCONNECT", "false") + stream = HangingStream() + + with pytest.raises(StreamedAPIError): + _ = [chunk async for chunk in StreamingHandler().wrap_stream(stream, "cred", "openai/gpt-test")] + + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_streaming_handler_stall_timeout_after_first_byte(monkeypatch) -> None: + monkeypatch.setenv("STREAM_STALL_TIMEOUT_SECONDS", "0.01") + monkeypatch.delenv("STREAM_TTFB_TIMEOUT_SECONDS", raising=False) + stream = FirstThenHangStream() + chunks = [] + + with pytest.raises(StreamedAPIError) as exc: + async for chunk in StreamingHandler().wrap_stream(stream, "cred", "openai/gpt-test"): + chunks.append(chunk) + + assert chunks and chunks[0].startswith("data: ") + assert stream.closed is True + assert exc.value.data["error"]["details"]["timeout_type"] == "stall" + + +@pytest.mark.asyncio +async def test_streaming_handler_passes_through_formatted_sse_chunks(monkeypatch) -> None: + async def formatted_stream(): + yield 'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n' + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(formatted_stream(), "cred", "openai/gpt-test")] + + assert chunks[0] == 'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n' + assert chunks[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_streaming_handler_does_not_duplicate_direct_done_sentinel(monkeypatch) -> None: + async def done_stream(): + yield "data: [DONE]\n\n" + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(done_stream(), "cred", "openai/gpt-test")] + + assert chunks == ["data: [DONE]\n\n"] + + +@pytest.mark.asyncio +async def test_streaming_handler_splits_mixed_reasoning_and_content_delta() -> None: + """Clients receive the reasoning-to-answer transition as distinct events.""" + + async def mixed_stream(): + yield { + "id": "chatcmpl-diffusion", + "model": "google/diffusiongemma-26b-a4b-it", + "choices": [ + { + "index": 0, + "finish_reason": None, + "delta": { + "reasoning_content": "end of thought", + "content": "beginning of answer", + }, + } + ], + } + yield { + "id": "chatcmpl-diffusion", + "choices": [{"index": 0, "finish_reason": "stop", "delta": {}}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3, + }, + } + + completed_responses = [] + chunks = [ + chunk + async for chunk in StreamingHandler().wrap_stream( + mixed_stream(), + "cred", + "nvidia_nim/google/diffusiongemma-26b-a4b-it", + response_callback=completed_responses.append, + ) + ] + + reasoning_chunk = json.loads(chunks[0][len("data: ") :]) + content_chunk = json.loads(chunks[1][len("data: ") :]) + reasoning_delta = reasoning_chunk["choices"][0]["delta"] + content_delta = content_chunk["choices"][0]["delta"] + + assert reasoning_delta["reasoning_content"] == "end of thought" + assert "content" not in reasoning_delta + assert reasoning_chunk["choices"][0]["finish_reason"] is None + assert "usage" not in reasoning_chunk + assert content_delta["content"] == "beginning of answer" + assert "reasoning_content" not in content_delta + assert completed_responses[0]["choices"][0]["message"]["content"] == "beginning of answer" + assert chunks[-1] == "data: [DONE]\n\n" diff --git a/tests/test_responses_bridge.py b/tests/test_responses_bridge.py new file mode 100644 index 000000000..3d4aa742c --- /dev/null +++ b/tests/test_responses_bridge.py @@ -0,0 +1,230 @@ +from __future__ import annotations + +from rotator_library.protocols.responses import ResponsesProtocol +from rotator_library.responses import ResponsesBridge + + +def test_bridge_converts_responses_request_to_chat_kwargs() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request( + { + "model": "gpt-test", + "instructions": "Follow rules.", + "input": "Hello", + "max_output_tokens": 20, + "metadata": {"trace": "yes"}, + } + ) + + kwargs = bridge.to_chat_kwargs(unified) + + assert kwargs["model"] == "gpt-test" + assert kwargs["messages"] == [ + {"role": "system", "content": "Follow rules."}, + {"role": "user", "content": "Hello"}, + ] + assert kwargs["max_tokens"] == 20 + assert kwargs["metadata"] == {"trace": "yes"} + + +def test_bridge_adds_parent_response_messages_for_previous_response_id() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Continue", "previous_response_id": "resp_parent"}) + parent = {"output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Earlier"}]}]} + + kwargs = bridge.to_chat_kwargs(unified, parent_response=parent) + + assert kwargs["messages"] == [ + {"role": "assistant", "content": "Earlier"}, + {"role": "user", "content": "Continue"}, + ] + assert kwargs["_responses_bridge"]["previous_response_id"] == "resp_parent" + assert kwargs["_session_tracking_hints"].global_strong_anchors == [ + "responses_previous_response_id:resp_parent" + ] + assert kwargs["_session_tracking_hints"].affinity_key == "responses_previous_response_id:resp_parent" + assert kwargs["_session_tracking_hints"].session_scope is None + + +def test_bridge_replays_parent_input_and_output_lineage() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Now"}) + lineage = [ + { + "request": {"model": "gpt-test", "input": "First user"}, + "response": {"output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "First assistant"}]}]}, + }, + { + "request": {"model": "gpt-test", "input": "Second user", "previous_response_id": "resp_1"}, + "response": {"output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Second assistant"}]}]}, + }, + ] + + kwargs = bridge.to_chat_kwargs(unified, parent_responses=lineage) + + assert kwargs["messages"] == [ + {"role": "user", "content": "First user"}, + {"role": "assistant", "content": "First assistant"}, + {"role": "user", "content": "Second user"}, + {"role": "assistant", "content": "Second assistant"}, + {"role": "user", "content": "Now"}, + ] + + +def test_bridge_replays_parent_tool_call_outputs() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Continue"}) + lineage = [ + { + "request": {"model": "gpt-test", "input": "Use tool"}, + "response": {"output": [{"id": "call_1", "type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"}]}, + } + ] + + kwargs = bridge.to_chat_kwargs(unified, parent_responses=lineage) + + assert kwargs["messages"] == [ + {"role": "user", "content": "Use tool"}, + {"role": "assistant", "content": None, "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}]}, + {"role": "user", "content": "Continue"}, + ] + + +def test_bridge_replays_parent_tool_result_inputs_as_tool_messages() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Continue"}) + lineage = [ + { + "request": { + "model": "gpt-test", + "input": [ + {"type": "message", "role": "user", "content": "Use tool"}, + {"type": "function_call_output", "call_id": "call_1", "output": "tool result"}, + ], + }, + "response": {"output": []}, + } + ] + + kwargs = bridge.to_chat_kwargs(unified, parent_responses=lineage) + + assert kwargs["messages"] == [ + {"role": "user", "content": "Use tool"}, + {"role": "tool", "content": "tool result", "tool_call_id": "call_1"}, + {"role": "user", "content": "Continue"}, + ] + + +def test_bridge_preserves_tool_definitions() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request( + { + "model": "gpt-test", + "input": "Use tool", + "tools": [{"type": "function", "name": "lookup", "description": "Lookup", "parameters": {"type": "object"}}], + } + ) + + kwargs = bridge.to_chat_kwargs(unified) + + assert kwargs["tools"] == [ + {"type": "function", "function": {"name": "lookup", "description": "Lookup", "parameters": {"type": "object"}}} + ] + + +def test_bridge_preserves_unsupported_fields_for_trace_metadata() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Hi", "custom_unsupported": 42}) + + kwargs = bridge.to_chat_kwargs(unified) + + assert kwargs["_responses_bridge"]["extra"]["custom_unsupported"] == 42 + + +def test_bridge_converts_chat_response_to_responses_payload() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Hello"}) + chat_response = { + "id": "chat_1", + "model": "gpt-test", + "created": 123, + "choices": [{"message": {"role": "assistant", "content": "Hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + response = bridge.from_chat_response(chat_response, unified) + + assert response["id"] == "chat_1" + assert response["object"] == "response" + assert response["status"] == "completed" + assert response["output"][0]["content"] == [{"type": "output_text", "text": "Hi"}] + assert response["usage"]["input_tokens"] == 1 + assert response["usage"]["output_tokens"] == 2 + assert response["usage"]["total_tokens"] == 3 + + +def test_bridge_preserves_chat_response_top_level_cost_with_usage() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Hello"}) + chat_response = { + "id": "chat_1", + "model": "gpt-test", + "choices": [{"message": {"role": "assistant", "content": "Hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2}, + "request_cost_usd": 0.012, + } + + response = bridge.from_chat_response(chat_response, unified) + + assert response["usage"]["cost_details"]["request_cost_usd"] == 0.012 + assert response["usage"]["cost_details"]["total_cost"] == 0.012 + + +def test_bridge_preserves_chat_response_estimated_cost_with_usage() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Hello"}) + chat_response = { + "id": "chat_1", + "model": "gpt-test", + "choices": [{"message": {"role": "assistant", "content": "Hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2}, + "estimated_cost": 0.013, + } + + response = bridge.from_chat_response(chat_response, unified) + + assert response["usage"]["cost_details"]["estimated_cost"] == 0.013 + + +def test_bridge_converts_chat_tool_calls_to_responses_output_items() -> None: + protocol = ResponsesProtocol() + bridge = ResponsesBridge(protocol) + unified = protocol.parse_request({"model": "gpt-test", "input": "Hello"}) + chat_response = { + "model": "gpt-test", + "choices": [ + { + "message": { + "role": "assistant", + "tool_calls": [{"id": "call_1", "function": {"name": "lookup", "arguments": "{}"}}], + } + } + ], + } + + response = bridge.from_chat_response(chat_response, unified, response_id="resp_tool") + + assert response["id"] == "resp_tool" + assert response["output"] == [ + {"id": "call_1", "type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"} + ] diff --git a/tests/test_responses_routes.py b/tests/test_responses_routes.py new file mode 100644 index 000000000..b5d3e028c --- /dev/null +++ b/tests/test_responses_routes.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +import json + +from fastapi.testclient import TestClient + +from proxy_app import main as proxy_main +from proxy_app.detailed_logger import RawIOLogger +from rotator_library.responses import InMemoryResponsesStore, ResponsesService, ResponsesServiceError + + +class FakeClient: + async def acompletion(self, **kwargs): + if kwargs.get("stream"): + async def chunks(): + yield 'data: {"choices":[{"delta":{"content":"route"}}]}\n\n' + yield 'data: {"choices":[{"delta":{"content":" stream"}}]}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + return { + "id": "chat_route_1", + "model": kwargs["model"], + "choices": [{"message": {"role": "assistant", "content": "route ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + +def _client() -> TestClient: + proxy_main.PROXY_API_KEY = None + proxy_main.ENABLE_RAW_LOGGING = False + proxy_main.app.state.rotating_client = FakeClient() + proxy_main.app.state.responses_service = ResponsesService(store=InMemoryResponsesStore()) + return TestClient(proxy_main.app) + + +def test_raw_logger_never_persists_authentication_headers() -> None: + logger = RawIOLogger.__new__(RawIOLogger) + logger.request_id = "request" + logger.streaming = False + written = {} + logger._write_json = lambda filename, data: written.update( + {"filename": filename, "data": data} + ) + + logger.log_request( + { + "Authorization": "Bearer proxy-secret", + "X-Api-Key": "provider-secret", + "X-Proxy-Session-Domain": "bundle:private.capability", + "Content-Type": "application/json", + }, + {"model": "gpt-test"}, + ) + + assert written["data"]["headers"] == { + "Authorization": "", + "X-Api-Key": "", + "X-Proxy-Session-Domain": "", + "Content-Type": "application/json", + } + + +def test_post_responses_non_stream_success() -> None: + client = _client() + + response = client.post("/v1/responses", json={"model": "gpt-test", "input": "hello"}) + + assert response.status_code == 200 + body = response.json() + assert body["id"] == "chat_route_1" + assert body["object"] == "response" + assert body["output"][0]["content"][0]["text"] == "route ok" + assert response.headers["x-proxy-session-domain"] == "public" + + +def test_post_responses_missing_model_returns_400() -> None: + client = _client() + + response = client.post("/v1/responses", json={"input": "hello"}) + + assert response.status_code == 400 + assert response.json()["error"]["type"] == "invalid_request_error" + + +def test_responses_provider_error_uses_selected_output_protocol() -> None: + class ErrorService(ResponsesService): + async def create_response(self, *args, **kwargs): + raise ResponsesServiceError( + "provider busy", + status_code=429, + error_type="rate_limit", + ) + + client = _client() + proxy_main.app.state.responses_service = ErrorService(store=InMemoryResponsesStore()) + + response = client.post( + "/v1/responses", + headers={"X-Proxy-Output-Protocol": "gemini"}, + json={"model": "openai/gpt-test", "input": "hello"}, + ) + + assert response.status_code == 429 + assert response.json() == { + "error": {"code": 429, "message": "provider busy", "status": "RESOURCE_EXHAUSTED"} + } + + +def test_post_responses_stream_missing_model_returns_400_before_sse() -> None: + client = _client() + + response = client.post("/v1/responses", json={"input": "hello", "stream": True}) + + assert response.status_code == 400 + assert response.json()["error"]["type"] == "invalid_request_error" + + +def test_post_responses_stream_missing_previous_response_returns_404_before_sse() -> None: + client = _client() + + response = client.post("/v1/responses", json={"model": "gpt-test", "input": "hello", "stream": True, "previous_response_id": "missing"}) + + assert response.status_code == 404 + assert response.json()["error"]["type"] == "not_found_error" + + +def test_get_delete_and_input_items_routes() -> None: + client = _client() + created = client.post("/v1/responses", json={"model": "gpt-test", "input": ["hello"]}).json() + + get_response = client.get(f"/v1/responses/{created['id']}") + input_items = client.get(f"/v1/responses/{created['id']}/input_items") + deleted = client.delete(f"/v1/responses/{created['id']}") + missing = client.get(f"/v1/responses/{created['id']}") + + assert get_response.status_code == 200 + assert get_response.json()["id"] == created["id"] + assert input_items.status_code == 200 + assert input_items.json() == {"object": "list", "data": ["hello"]} + assert deleted.status_code == 200 + assert deleted.json() == {"id": created["id"], "object": "response.deleted", "deleted": True} + assert missing.status_code == 404 + assert missing.json()["error"]["type"] == "not_found_error" + + +def test_scoped_response_retrieval_requires_creation_domain_header() -> None: + client = _client() + created_response = client.post( + "/v1/responses", + json={ + "model": "gpt-test", + "input": "private", + "api_keys": {"openai": ["private-key"]}, + "private": True, + }, + ) + response_id = created_response.json()["id"] + domain = created_response.headers["x-proxy-session-domain"] + raw_domain = domain.rsplit(".", 1)[0] + + assert domain.startswith("bundle:") + assert raw_domain.startswith("bundle:") + assert client.get(f"/v1/responses/{response_id}").status_code == 404 + assert ( + client.get( + f"/v1/responses/{response_id}", + headers={"X-Proxy-Session-Domain": raw_domain}, + ).status_code + == 404 + ) + assert ( + client.get( + f"/v1/responses/{response_id}", + headers={"X-Proxy-Session-Domain": domain}, + ).status_code + == 200 + ) + assert ( + client.get( + f"/v1/responses/{response_id}", + headers={"X-Proxy-Session-Domain": "bundle:wrong"}, + ).status_code + == 404 + ) + + +def test_scoped_response_access_capabilities_are_unique_per_response() -> None: + client = _client() + payload = { + "model": "gpt-test", + "input": "private", + "api_keys": {"openai": ["private-key"]}, + "private": True, + } + + first = client.post("/v1/responses", json=payload) + second = client.post("/v1/responses", json=payload) + + assert first.headers["x-proxy-session-domain"] != second.headers["x-proxy-session-domain"] + + +def test_scoped_continuation_requires_parent_access_capability() -> None: + client = _client() + payload = { + "model": "gpt-test", + "input": "private", + "api_keys": {"openai": ["private-key"]}, + "private": True, + } + parent = client.post("/v1/responses", json=payload) + continuation = {**payload, "previous_response_id": parent.json()["id"]} + + denied = client.post("/v1/responses", json=continuation) + allowed = client.post( + "/v1/responses", + json=continuation, + headers={ + "X-Proxy-Session-Domain": parent.headers["x-proxy-session-domain"] + }, + ) + + assert denied.status_code == 404 + assert allowed.status_code == 200 + + +def test_post_responses_stream_returns_sse_events() -> None: + client = _client() + + response = client.post("/v1/responses", json={"model": "gpt-test", "input": "hello", "stream": True}) + + assert response.status_code == 200 + assert response.headers["x-proxy-session-domain"] == "public" + assert "event: response.created" in response.text + assert "event: response.completed" in response.text + + +def test_scoped_stream_uses_same_access_capability_for_stored_response() -> None: + client = _client() + + response = client.post( + "/v1/responses", + json={ + "model": "gpt-test", + "input": "private stream", + "stream": True, + "api_keys": {"openai": ["private-key"]}, + "private": True, + }, + ) + capability = response.headers["x-proxy-session-domain"] + created_payload = next( + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + response_id = created_payload["id"] + + assert response.status_code == 200 + assert capability.startswith("bundle:") + assert ( + client.get( + f"/v1/responses/{response_id}", + headers={"X-Proxy-Session-Domain": capability}, + ).status_code + == 200 + ) diff --git a/tests/test_responses_service.py b/tests/test_responses_service.py new file mode 100644 index 000000000..acc5aea56 --- /dev/null +++ b/tests/test_responses_service.py @@ -0,0 +1,686 @@ +from __future__ import annotations + +import json + +import pytest + +import rotator_library.responses.service as responses_service_module +from rotator_library.responses import InMemoryResponsesStore, ResponsesService, ResponsesServiceError, ResponsesStoreSettings, StoredResponse, create_configured_responses_store +from rotator_library.core.errors import StructuredAPIResponseError +from rotator_library.transaction_logger import TransactionLogger + + +class FakeClient: + def __init__(self) -> None: + self.calls = [] + + async def acompletion(self, **kwargs): + self.calls.append(kwargs) + return { + "id": "chat_response_1", + "model": kwargs["model"], + "choices": [{"message": {"role": "assistant", "content": "Hello back"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + +class FakeCostClient(FakeClient): + async def acompletion(self, **kwargs): + response = await super().acompletion(**kwargs) + response["usage"]["cost_details"] = {"total_cost": 0.033, "source": "responses_provider"} + return response + + +class FakeInternalClient(FakeClient): + def __init__(self) -> None: + super().__init__() + self._request_builder = object() + self._executor = object() + + async def acompletion(self, **kwargs): + callback = kwargs.pop("_request_context_callback", None) + hints = kwargs.pop("_session_tracking_hints", None) + if callback: + callback( + type( + "Context", + (), + { + "session_id": "session-parent", + "session_affinity_key": "affinity-parent", + "usage_manager_key": "scope-parent", + "classifier": "global", + "session_tracker": None, + "provider": "openai", + "model": "gpt-test", + "session_tracking_namespace": "namespace", + }, + )() + ) + self.internal_hints = hints + return await super().acompletion(**kwargs) + + +class FakeNativeProtocolClient: + """Exercise the internal Responses protocol path without a chat bridge.""" + + def __init__(self) -> None: + self.calls = [] + self._request_builder = object() + self._executor = object() + + async def agenerate(self, payload, *, input_protocol, output_protocol, request=None, **kwargs): + callback = kwargs.pop("_request_context_callback", None) + if callback: + callback( + type( + "Context", + (), + { + "session_id": "session-native", + "session_affinity_key": "affinity-native", + "usage_manager_key": "scope-native", + "session_isolation_key": "public", + "classifier": "global", + "session_tracker": None, + "provider": "openai", + "model": payload["model"], + "session_tracking_namespace": "namespace", + }, + )() + ) + self.calls.append( + { + "payload": payload, + "input_protocol": input_protocol, + "output_protocol": output_protocol, + "kwargs": kwargs, + } + ) + index = len(self.calls) + return { + "id": f"resp_native_{index}", + "object": "response", + "model": payload["model"], + "status": "completed", + "output": [ + { + "id": f"msg_{index}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": f"answer {index}"}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + + +class FakeNativeProtocolErrorClient(FakeNativeProtocolClient): + async def agenerate(self, payload, *, input_protocol, output_protocol, request=None, **kwargs): + raise StructuredAPIResponseError( + "provider busy", + error_type="rate_limit", + status_code=429, + response={"error": {"status": 429, "message": "provider busy"}}, + ) + + +class FakeSelectedOutputClient(FakeNativeProtocolClient): + def resolve_output_protocol(self, payload, *, input_protocol, request=None, explicit=None): + return "anthropic_messages" + + +@pytest.mark.asyncio +async def test_native_structured_errors_keep_responses_status_and_type() -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + + with pytest.raises(ResponsesServiceError) as raised: + await service.create_response( + {"model": "openai/gpt-test", "input": "hello"}, + FakeNativeProtocolErrorClient(), + ) + + assert raised.value.status_code == 429 + assert raised.value.error_type == "rate_limit" + + +def test_responses_service_errors_follow_selected_output_protocol() -> None: + error = ResponsesServiceError("provider busy", status_code=429, error_type="rate_limit") + + assert error.to_protocol_payload("gemini") == { + "error": {"code": 429, "message": "provider busy", "status": "RESOURCE_EXHAUSTED"} + } + assert error.to_protocol_payload("anthropic_messages")["error"]["type"] == "rate_limit_error" + + +def _trace_entries(log_dir): + return [json.loads(line) for line in (log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + + +def test_responses_service_owns_request_scope_derivation() -> None: + service = ResponsesService() + + assert service.request_scope_key({"model": "gpt-test"}) == "public" + assert service.request_scope_key( + { + "model": "gpt-test", + "api_keys": {"openai": ["private-key"]}, + "private": True, + } + ).startswith("bundle:") + + +def test_responses_service_recursively_redacts_transport_logging_payload() -> None: + service = ResponsesService() + request = { + "model": "gpt-test", + "api_keys": {"openai": ["top-level-secret"]}, + "metadata": { + "providers": {"private": {"api_key": "nested-secret"}}, + "keep": "visible", + }, + "input": [{"authorization": "Bearer secret", "value": "visible"}], + } + + redacted = service.redact_request_for_logging(request) + + assert redacted == { + "model": "gpt-test", + "metadata": {"keep": "visible"}, + "input": [{"value": "visible"}], + } + assert request["api_keys"]["openai"] == ["top-level-secret"] + assert request["metadata"]["providers"]["private"]["api_key"] == "nested-secret" + + +@pytest.mark.asyncio +async def test_internal_responses_path_uses_native_protocol_and_expands_local_continuation() -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + client = FakeNativeProtocolClient() + + first = await service.create_response( + {"model": "openai/gpt-test", "input": "first"}, + client, + ) + second = await service.create_response( + {"model": "openai/gpt-test", "input": "second", "previous_response_id": first["id"]}, + client, + ) + + assert second["id"] == "resp_native_2" + assert client.calls[0]["input_protocol"] == "responses" + assert client.calls[0]["output_protocol"] == "responses" + assert client.calls[0]["payload"]["input"] == ["first"] + assert client.calls[1]["payload"]["input"] == [ + "first", + first["output"][0], + "second", + ] + assert "previous_response_id" not in client.calls[1]["payload"] + assert client.calls[1]["kwargs"]["_disable_provider_continuation"] is True + + +@pytest.mark.asyncio +async def test_responses_storage_stays_native_before_selected_output_conversion() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + client = FakeSelectedOutputClient() + + result = await service.create_response( + {"model": "openai/gpt-test", "input": "hello"}, + client, + ) + + assert result["type"] == "message" + assert result["content"][0]["text"] == "answer 1" + stored = await store.get("resp_native_1", "public") + assert stored is not None + assert stored.response["object"] == "response" + assert stored.response["output"][0]["type"] == "message" + + +@pytest.mark.asyncio +async def test_scoped_response_requires_unforgeable_access_capability() -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + request = { + "model": "gpt-test", + "input": "private", + "api_keys": {"openai": ["private-key"]}, + "private": True, + } + request_scope = service.prepare_request_scope(request) + response = await service.create_response( + request, + FakeClient(), + request_scope=request_scope, + ) + raw_domain = request_scope.access_token.rsplit(".", 1)[0] + + loaded = await service.get_response_with_access_token( + response["id"], + request_scope.access_token, + ) + + assert loaded["id"] == response["id"] + with pytest.raises(ResponsesServiceError) as raw_error: + await service.get_response_with_access_token(response["id"], raw_domain) + assert raw_error.value.status_code == 404 + with pytest.raises(ResponsesServiceError) as forged_error: + await service.get_response_with_access_token( + response["id"], + request_scope.access_token[:-1] + "x", + ) + assert forged_error.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_scoped_access_capability_survives_durable_store_restart(tmp_path) -> None: + env = { + "RESPONSES_STORE_BACKEND": "provider_cache", + "RESPONSES_STORE_CACHE_NAME": "responses_capability_test", + "RESPONSES_STORE_CACHE_PREFIX": "responses", + "RESPONSES_STORE_CACHE_DIR": str(tmp_path), + "RESPONSES_STORE_CACHE_MEMORY_TTL_SECONDS": "60", + "RESPONSES_STORE_CACHE_DISK_TTL_SECONDS": "60", + } + request = { + "model": "gpt-test", + "input": "private", + "api_keys": {"openai": ["private-key"]}, + "private": True, + } + first = ResponsesService(store=create_configured_responses_store(env=env)) + request_scope = first.prepare_request_scope(request) + response = await first.create_response( + request, + FakeClient(), + request_scope=request_scope, + ) + restarted = ResponsesService(store=create_configured_responses_store(env=env)) + + loaded = await restarted.get_response_with_access_token( + response["id"], + request_scope.access_token, + ) + + assert loaded["id"] == response["id"] + + +@pytest.mark.asyncio +async def test_create_response_stores_non_streaming_response() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + client = FakeClient() + + response = await service.create_response({"model": "gpt-test", "input": "Hello"}, client) + + assert response["id"] == "chat_response_1" + assert response["output"][0]["content"][0]["text"] == "Hello back" + assert (await store.get("chat_response_1")) is not None + assert client.calls[0]["messages"] == [{"role": "user", "content": "Hello"}] + + +@pytest.mark.asyncio +async def test_store_false_does_not_persist_response() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + response = await service.create_response({"model": "gpt-test", "input": "Hello", "store": False}, FakeClient()) + + assert await store.get(response["id"]) is None + + +@pytest.mark.asyncio +async def test_create_response_applies_storage_ttl() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store, store_settings=ResponsesStoreSettings(ttl_seconds=60)) + + response = await service.create_response({"model": "gpt-test", "input": "Hello"}, FakeClient()) + stored = await store.get(response["id"]) + + assert stored is not None + assert stored.expires_at is not None + assert stored.expires_at > stored.created_at + assert stored.metadata["response_id"] == response["id"] + + +@pytest.mark.asyncio +async def test_service_default_store_honors_max_items() -> None: + class SequencedClient(FakeClient): + def __init__(self) -> None: + super().__init__() + self.index = 0 + + async def acompletion(self, **kwargs): + self.index += 1 + response = await super().acompletion(**kwargs) + response["id"] = f"chat_response_{self.index}" + return response + + service = ResponsesService(store_settings=ResponsesStoreSettings(max_items=1)) + client = SequencedClient() + + first = await service.create_response({"model": "gpt-test", "input": "one"}, client) + second = await service.create_response({"model": "gpt-test", "input": "two"}, client) + + assert await service.store.get(first["id"]) is None + assert await service.store.get(second["id"]) is not None + + +@pytest.mark.asyncio +async def test_previous_response_id_loads_parent_context() -> None: + store = InMemoryResponsesStore() + await store.save( + StoredResponse( + id="resp_parent", + model="gpt-test", + status="completed", + scope_key="public", + request={"input": "Earlier"}, + response={ + "id": "resp_parent", + "object": "response", + "model": "gpt-test", + "status": "completed", + "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Earlier"}]}], + }, + input_items=["Earlier"], + output_items=[{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Earlier"}]}], + ) + ) + client = FakeClient() + service = ResponsesService(store=store) + + await service.create_response({"model": "gpt-test", "input": "Continue", "previous_response_id": "resp_parent"}, client) + + assert client.calls[0]["messages"] == [ + {"role": "user", "content": "Earlier"}, + {"role": "assistant", "content": "Earlier"}, + {"role": "user", "content": "Continue"}, + ] + + +@pytest.mark.asyncio +async def test_previous_response_id_loads_full_lineage_oldest_first() -> None: + store = InMemoryResponsesStore() + await store.save( + StoredResponse( + id="resp_grandparent", + model="gpt-test", + status="completed", + scope_key="public", + request={"model": "gpt-test", "input": "First"}, + response={"id": "resp_grandparent", "object": "response", "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "First answer"}]}]}, + ) + ) + await store.save( + StoredResponse( + id="resp_parent", + model="gpt-test", + status="completed", + scope_key="public", + request={"model": "gpt-test", "input": "Second", "previous_response_id": "resp_grandparent"}, + response={"id": "resp_parent", "object": "response", "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Second answer"}]}]}, + ) + ) + client = FakeClient() + + await ResponsesService(store=store).create_response({"model": "gpt-test", "input": "Third", "previous_response_id": "resp_parent"}, client) + + assert client.calls[0]["messages"] == [ + {"role": "user", "content": "First"}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Second"}, + {"role": "assistant", "content": "Second answer"}, + {"role": "user", "content": "Third"}, + ] + + +@pytest.mark.asyncio +async def test_internal_session_hints_do_not_leak_to_direct_clients_or_traces(tmp_path) -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + await store.save( + StoredResponse( + id="resp_parent", + model="gpt-test", + status="completed", + scope_key="public", + response={"id": "resp_parent", "object": "response", "output": []}, + metadata={"session_affinity_key": "affinity-parent"}, + ) + ) + client = FakeClient() + + await service.create_response({"model": "gpt-test", "input": "Continue", "previous_response_id": "resp_parent"}, client, transaction_logger=logger) + + assert "_session_tracking_hints" not in client.calls[0] + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "_session_tracking_hints" not in trace_text + assert "has_session_hints" in trace_text + + +@pytest.mark.asyncio +async def test_internal_client_context_metadata_is_stored_with_response() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + await store.save( + StoredResponse( + id="resp_parent", + model="gpt-test", + status="completed", + scope_key="public", + response={"id": "resp_parent", "object": "response", "output": []}, + metadata={"session_affinity_key": "affinity-parent"}, + ) + ) + client = FakeInternalClient() + + response = await service.create_response({"model": "gpt-test", "input": "Continue", "previous_response_id": "resp_parent"}, client) + stored = await store.get(response["id"]) + + assert client.internal_hints.affinity_key == "responses_previous_response_id:resp_parent" + assert stored is not None + assert stored.session_id == "session-parent" + assert "session_affinity_key" not in stored.metadata + + +@pytest.mark.asyncio +async def test_scoped_responses_preserve_routing_but_never_store_or_trace_secrets(tmp_path) -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + client = FakeClient() + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + raw_request = { + "model": "gpt-test", + "input": "Scoped request", + "api_keys": {"openai": ["super-secret-routing-key"]}, + "providers": { + "openai": { + "api_base": "https://private.example", + "authorization": "provider-secret-header", + } + }, + "private": True, + } + + request_scope = service.prepare_request_scope(raw_request) + response = await service.create_response( + raw_request, + client, + transaction_logger=logger, + request_scope=request_scope, + ) + scope_key = responses_service_module._request_isolation_key(raw_request) + stored = await store.get(response["id"], scope_key) + + assert stored is not None + assert stored.scope_key is not None and stored.scope_key.startswith("bundle:") + assert client.calls[0]["api_keys"] == raw_request["api_keys"] + assert client.calls[0]["providers"] == raw_request["providers"] + persisted_text = json.dumps(stored.to_dict()) + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "super-secret-routing-key" not in persisted_text + assert "provider-secret-header" not in persisted_text + assert "super-secret-routing-key" not in trace_text + assert "provider-secret-header" not in trace_text + + with pytest.raises(ResponsesServiceError) as public_access: + await service.get_response(response["id"]) + assert public_access.value.status_code == 404 + assert (await service.get_response(response["id"], scope_key=stored.scope_key))["id"] == response["id"] + + mismatched = dict(raw_request) + mismatched["api_keys"] = {"openai": ["different-private-key"]} + mismatched["previous_response_id"] = response["id"] + with pytest.raises(ResponsesServiceError) as cross_scope: + await service.create_response(mismatched, FakeClient()) + assert cross_scope.value.status_code == 404 + + continued = dict(raw_request) + continued["previous_response_id"] = response["id"] + await service.create_response( + continued, + FakeClient(), + request_scope=service.prepare_request_scope(continued), + previous_response_access_token=request_scope.access_token, + ) + + +@pytest.mark.asyncio +async def test_responses_service_records_response_id_session_anchor() -> None: + class Tracker: + def __init__(self) -> None: + self.calls = [] + + def record_response(self, *args, **kwargs): + self.calls.append((args, kwargs)) + + tracker = Tracker() + + class Client(FakeInternalClient): + async def acompletion(self, **kwargs): + callback = kwargs.pop("_request_context_callback", None) + kwargs.pop("_session_tracking_hints", None) + if callback: + callback( + type( + "Context", + (), + { + "session_id": "session-parent", + "session_affinity_key": "affinity-parent", + "usage_manager_key": "scope-parent", + "classifier": "global", + "session_tracker": tracker, + "provider": "openai", + "model": "gpt-test", + "session_tracking_namespace": "namespace", + }, + )() + ) + self.calls.append(kwargs) + return { + "id": "resp_parent", + "model": kwargs["model"], + "choices": [{"message": {"role": "assistant", "content": "Hello back"}, "finish_reason": "stop"}], + } + + await ResponsesService(store=InMemoryResponsesStore()).create_response({"model": "gpt-test", "input": "Hello"}, Client()) + + assert tracker.calls[0][0][0] == "session-parent" + assert tracker.calls[0][1]["response"] == {"id": "resp_parent", "object": "response"} + + +@pytest.mark.asyncio +async def test_missing_previous_response_id_raises_not_found() -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + + with pytest.raises(ResponsesServiceError) as exc_info: + await service.create_response({"model": "gpt-test", "input": "Continue", "previous_response_id": "missing"}, FakeClient()) + + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_get_delete_and_list_input_items() -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + response = await service.create_response({"model": "gpt-test", "input": ["Hello"]}, FakeClient()) + + assert (await service.get_response(response["id"]))["id"] == response["id"] + assert await service.list_input_items(response["id"]) == {"object": "list", "data": ["Hello"]} + assert await service.delete_response(response["id"]) == {"id": response["id"], "object": "response.deleted", "deleted": True} + with pytest.raises(ResponsesServiceError): + await service.get_response(response["id"]) + + +@pytest.mark.asyncio +async def test_service_emits_transform_trace_passes(tmp_path) -> None: + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + service = ResponsesService(store=InMemoryResponsesStore()) + + await service.create_response({"model": "gpt-test", "input": "Hello"}, FakeClient(), transaction_logger=logger) + + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert pass_names == [ + "responses_raw_request", + "responses_parsed_request", + "responses_bridge_chat_request", + "responses_bridge_chat_response", + "responses_parsed_response", + "usage_accounting_summary", + "responses_stored_response", + "responses_final_response", + ] + + +@pytest.mark.asyncio +async def test_service_usage_trace_includes_provider_reported_cost(tmp_path) -> None: + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + service = ResponsesService(store=InMemoryResponsesStore()) + + await service.create_response({"model": "gpt-test", "input": "Hello"}, FakeCostClient(), transaction_logger=logger) + + usage_entry = [entry for entry in _trace_entries(logger.log_dir) if entry["pass_name"] == "usage_accounting_summary"][-1] + assert usage_entry["data"]["cost"]["provider_reported_cost"] == 0.033 + assert usage_entry["metadata"]["pricing_source"] == "usage.cost_details" + + +def test_trace_responses_usage_returns_before_conversion_without_logger(monkeypatch) -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + + def fail_extract(*args, **kwargs): + raise AssertionError("usage conversion should be skipped when tracing is disabled") + + monkeypatch.setattr(responses_service_module, "extract_usage_record", fail_extract) + + service._trace_responses_usage(None, {"usage": {"input_tokens": 1}}, "gpt-test", source="test") + + +@pytest.mark.asyncio +async def test_previous_response_trace_payload_skipped_without_logger() -> None: + class Parent: + id = "resp_parent" + scope_key = "public" + response = {"output": []} + output_items = [] + input_items = [] + + def to_dict(self): + raise AssertionError("previous response trace payload should not be built without a logger") + + class Store(InMemoryResponsesStore): + async def get(self, response_id, scope_key="public"): + return Parent() + + service = ResponsesService(store=Store()) + + parent = await service._load_previous_response( + "resp_parent", + None, + expected_scope_key="public", + ) + + assert parent.id == "resp_parent" diff --git a/tests/test_responses_store.py b/tests/test_responses_store.py new file mode 100644 index 000000000..963f6dce2 --- /dev/null +++ b/tests/test_responses_store.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import json +import time + +import pytest + +from rotator_library.responses import InMemoryResponsesStore, ProviderCacheResponsesStore, StoredResponse, create_configured_responses_store, generate_response_id + + +def _stored(response_id: str = "resp_test") -> StoredResponse: + return StoredResponse( + id=response_id, + model="gpt-test", + status="completed", + request={"model": "gpt-test", "input": "hello"}, + response={"id": response_id, "object": "response", "output": []}, + input_items=[{"type": "message", "role": "user", "content": "hello"}], + output_items=[{"type": "message", "role": "assistant", "content": []}], + usage={"input_tokens": 1, "output_tokens": 2, "total_tokens": 3}, + metadata={"previous_response_id": None}, + ) + + +def test_generate_response_id_uses_responses_prefix() -> None: + assert generate_response_id().startswith("resp_") + assert generate_response_id() != generate_response_id() + + +@pytest.mark.asyncio +async def test_in_memory_store_save_get_delete_and_input_items() -> None: + store = InMemoryResponsesStore() + stored = _stored() + + await store.save(stored) + loaded = await store.get(stored.id) + input_items = await store.list_input_items(stored.id) + deleted = await store.delete(stored.id) + + assert loaded is not None + assert loaded.to_dict() == stored.to_dict() + assert input_items == stored.input_items + assert deleted is True + assert await store.get(stored.id) is None + + +@pytest.mark.asyncio +async def test_response_storage_is_composite_by_domain_and_response_id() -> None: + store = InMemoryResponsesStore() + public = _stored("resp_shared") + public.scope_key = "public" + private = _stored("resp_shared") + private.scope_key = "bundle:private" + private.response["domain"] = "private" + + await store.save(public) + await store.save(private) + + assert (await store.get("resp_shared", "public")).response.get("domain") is None + assert (await store.get("resp_shared", "bundle:private")).response["domain"] == "private" + assert await store.delete("resp_shared", "public") is True + assert await store.get("resp_shared", "public") is None + assert await store.get("resp_shared", "bundle:private") is not None + + +@pytest.mark.asyncio +async def test_in_memory_store_returns_copies_and_expires() -> None: + store = InMemoryResponsesStore() + stored = _stored("resp_expiring") + stored.expires_at = time.time() + 100 + await store.save(stored) + + loaded = await store.get(stored.id) + assert loaded is not None + loaded.response["mutated"] = True + assert (await store.get(stored.id)).response.get("mutated") is None + + stored.expires_at = time.time() - 1 + await store.save(stored) + assert await store.get(stored.id) is None + + +@pytest.mark.asyncio +async def test_in_memory_store_prunes_oldest_when_max_items_exceeded() -> None: + store = InMemoryResponsesStore(max_items=2) + first = _stored("resp_first") + first.created_at = 1 + second = _stored("resp_second") + second.created_at = 2 + third = _stored("resp_third") + third.created_at = 3 + + await store.save(first) + await store.save(second) + await store.save(third) + + assert await store.get("resp_first") is None + assert await store.get("resp_second") is not None + assert await store.get("resp_third") is not None + + +@pytest.mark.asyncio +async def test_provider_cache_store_serializes_json_and_does_not_clear_without_key_delete() -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.values = {} + + async def store_async(self, key: str, value: str) -> None: + json.loads(value) + self.values[key] = value + + async def retrieve_async(self, key: str): + return self.values.get(key) + + cache = FakeProviderCache() + store = ProviderCacheResponsesStore(cache) + stored = _stored("resp_provider_cache") + + await store.save(stored) + loaded = await store.get(stored.id) + + assert loaded is not None + assert loaded.id == stored.id + assert await store.list_input_items(stored.id) == stored.input_items + assert await store.delete(stored.id) is False + assert cache.values + + +@pytest.mark.asyncio +async def test_provider_cache_response_ids_cannot_collide_after_encoding() -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.values = {} + + async def store_async(self, key: str, value: str) -> None: + self.values[key] = value + + async def retrieve_async(self, key: str): + return self.values.get(key) + + cache = FakeProviderCache() + store = ProviderCacheResponsesStore(cache) + slash = _stored("resp/a") + underscore = _stored("resp_a") + + await store.save(slash) + await store.save(underscore) + + assert len(cache.values) == 2 + assert (await store.get("resp/a")).id == "resp/a" + assert (await store.get("resp_a")).id == "resp_a" + + +@pytest.mark.asyncio +async def test_provider_cache_store_safely_migrates_matching_legacy_entry() -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.values = {} + + async def store_async(self, key: str, value: str) -> None: + self.values[key] = value + + async def retrieve_async(self, key: str): + return self.values.get(key) + + async def delete_async(self, key: str) -> bool: + return self.values.pop(key, None) is not None + + cache = FakeProviderCache() + store = ProviderCacheResponsesStore(cache) + stored = _stored("resp_legacy") + cache.values["responses:resp_legacy"] = json.dumps(stored.to_dict()) + + loaded = await store.get(stored.id) + + assert loaded is not None + assert loaded.id == stored.id + assert "responses:resp_legacy" not in cache.values + assert len(cache.values) == 1 + + +@pytest.mark.asyncio +async def test_provider_cache_store_rejects_ambiguous_legacy_collision() -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.values = {} + + async def store_async(self, key: str, value: str) -> None: + self.values[key] = value + + async def retrieve_async(self, key: str): + return self.values.get(key) + + cache = FakeProviderCache() + store = ProviderCacheResponsesStore(cache) + collided = _stored("resp_a") + cache.values["responses:resp_a"] = json.dumps(collided.to_dict()) + + assert await store.get("resp/a") is None + assert cache.values["responses:resp_a"] + + +@pytest.mark.asyncio +async def test_provider_cache_store_uses_key_delete_when_available() -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.values = {} + + async def store_async(self, key: str, value: str) -> None: + self.values[key] = value + + async def retrieve_async(self, key: str): + return self.values.get(key) + + async def delete_async(self, key: str) -> bool: + return self.values.pop(key, None) is not None + + store = ProviderCacheResponsesStore(FakeProviderCache()) + stored = _stored("resp_delete") + await store.save(stored) + + assert await store.delete(stored.id) is True + assert await store.get(stored.id) is None + + +@pytest.mark.asyncio +async def test_provider_cache_store_closes_owned_cache_tasks() -> None: + class FakeProviderCache: + def __init__(self) -> None: + self.closed = False + + async def shutdown(self) -> None: + self.closed = True + + cache = FakeProviderCache() + store = ProviderCacheResponsesStore(cache) + + await store.close() + + assert cache.closed is True + + +@pytest.mark.asyncio +async def test_configured_provider_cache_store_persists_between_instances(tmp_path) -> None: + env = { + "RESPONSES_STORE_BACKEND": "provider_cache", + "RESPONSES_STORE_CACHE_NAME": "responses_test", + "RESPONSES_STORE_CACHE_PREFIX": "responses", + "RESPONSES_STORE_CACHE_DIR": str(tmp_path), + "RESPONSES_STORE_CACHE_MEMORY_TTL_SECONDS": "60", + "RESPONSES_STORE_CACHE_DISK_TTL_SECONDS": "60", + } + first_store = create_configured_responses_store(env=env) + stored = _stored("resp_durable") + + await first_store.save(stored) + second_store = create_configured_responses_store(env=env) + loaded = await second_store.get(stored.id) + + assert loaded is not None + assert loaded.to_dict() == stored.to_dict() + + +@pytest.mark.asyncio +async def test_configured_provider_cache_store_deletes_durable_entry(tmp_path) -> None: + env = { + "RESPONSES_STORE_BACKEND": "provider_cache", + "RESPONSES_STORE_CACHE_NAME": "responses_delete_test", + "RESPONSES_STORE_CACHE_PREFIX": "responses", + "RESPONSES_STORE_CACHE_DIR": str(tmp_path), + "RESPONSES_STORE_CACHE_MEMORY_TTL_SECONDS": "60", + "RESPONSES_STORE_CACHE_DISK_TTL_SECONDS": "60", + } + first_store = create_configured_responses_store(env=env) + stored = _stored("resp_durable_delete") + await first_store.save(stored) + + assert await first_store.delete(stored.id) is True + second_store = create_configured_responses_store(env=env) + + assert await second_store.get(stored.id) is None diff --git a/tests/test_responses_streaming.py b/tests/test_responses_streaming.py new file mode 100644 index 000000000..dfa7be08d --- /dev/null +++ b/tests/test_responses_streaming.py @@ -0,0 +1,778 @@ +from __future__ import annotations + +import asyncio +import json + +import pytest + +from rotator_library.responses import InMemoryResponsesStore, ResponsesSSEFormatter, ResponsesService, ResponsesStoreSettings, ResponsesStreamEvent, ResponsesWebSocketFormatter +from rotator_library.transaction_logger import TransactionLogger + + +def _event_names(events: list[str]) -> list[str]: + return [line.removeprefix("event: ") for event in events for line in event.splitlines() if line.startswith("event: ")] + + +class FakeStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield 'data: {"id":"chat_stream","model":"gpt-test","choices":[{"delta":{"role":"assistant"}}]}\n\n' + yield 'data: {"id":"chat_stream","model":"gpt-test","choices":[{"delta":{"content":"Hel"}}]}\n\n' + yield 'data: {"id":"chat_stream","model":"gpt-test","choices":[{"delta":{"content":"lo"}}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"cost_details":{"total_cost":0.044,"source":"stream_provider"}}}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + + +class NativeResponsesStreamingClient: + def __init__(self, output_protocol: str = "responses") -> None: + self.output_protocol = output_protocol + self.calls = [] + + def resolve_output_protocol(self, payload, *, input_protocol, request=None): + return self.output_protocol + + async def agenerate(self, payload, **kwargs): + self.calls.append((payload, kwargs)) + + async def chunks(): + yield 'event: response.created\ndata: {"type":"response.created","response":{"id":"resp_native","object":"response","status":"in_progress","model":"gpt-test","output":[]}}\n\n' + yield 'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","item_id":"msg_0","output_index":0,"content_index":0,"delta":"native"}\n\n' + yield 'event: response.completed\ndata: {"type":"response.completed","response":{"id":"resp_native","object":"response","status":"completed","model":"gpt-test","output":[{"id":"msg_0","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"native"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}\n\n' + + return chunks() + + +class DelayedCloseableStream: + def __init__(self, chunks: list[str], delays: list[float]) -> None: + self.chunks = chunks + self.delays = delays + self.index = 0 + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index >= len(self.chunks): + raise StopAsyncIteration + delay = self.delays[self.index] + chunk = self.chunks[self.index] + self.index += 1 + await asyncio.sleep(delay) + return chunk + + async def aclose(self) -> None: + self.closed = True + + +class DelayedStreamingClient: + def __init__(self, stream: DelayedCloseableStream) -> None: + self.stream = stream + + async def acompletion(self, **kwargs): + return self.stream + + +class SlowAcquireStreamingClient: + def __init__(self, stream: DelayedCloseableStream, delay: float) -> None: + self.stream = stream + self.delay = delay + self.calls = 0 + + async def acompletion(self, **kwargs): + self.calls += 1 + await asyncio.sleep(self.delay) + return self.stream + + +class CancelAwareAcquireClient: + def __init__(self, delay: float = 1.0) -> None: + self.cancelled = False + self.delay = delay + + async def acompletion(self, **kwargs): + try: + await asyncio.sleep(self.delay) + except asyncio.CancelledError: + self.cancelled = True + raise + return DelayedCloseableStream([], []) + + +class DisconnectRequest: + async def is_disconnected(self) -> bool: + return True + + +class DisconnectAfterAcquireRequest: + def __init__(self) -> None: + self.calls = 0 + + async def is_disconnected(self) -> bool: + self.calls += 1 + return self.calls > 1 + + +class FailingStreamingClient: + def __init__(self, message: str = "stream exploded") -> None: + self.message = message + + async def acompletion(self, **kwargs): + message = self.message + + async def chunks(): + yield 'data: {"choices":[{"delta":{"content":"before"}}]}\n\n' + raise RuntimeError(message) + + return chunks() + + +class ErrorChunkStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield 'data: {"choices":[{"delta":{"content":"partial"}}]}\n\n' + yield 'data: {"error":{"message":"provider failed"}}\n\n' + + return chunks() + + +class EventErrorStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield 'data: {"choices":[{"delta":{"content":"partial"}}]}\n\n' + yield 'event: error\ndata: {"message":"event failed"}\n\n' + + return chunks() + + +class CostCommentStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield ': cost {"total_cost":0.031,"currency":"USD","source":"responses_sse"}\n\n' + yield 'data: {"choices":[{"delta":{"content":"priced"}}]}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + + +class CostEventStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield 'event: cost\ndata: {"total_cost":0.017,"currency":"EUR","source":"responses_event_cost"}\n\n' + yield 'data: {"choices":[{"delta":{"content":"priced"}}]}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + + +class RequestCostCommentStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield ': cost {"request_cost_usd":0.023,"source":"reference_sse"}\n\n' + yield 'data: {"choices":[{"delta":{"content":"priced"}}]}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + + +class EstimatedCostCommentStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield ': cost {"estimated_cost":0.024,"source":"reference_estimate"}\n\n' + yield 'data: {"choices":[{"delta":{"content":"priced"}}]}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + + +class TopLevelUsageCostStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield 'data: {"choices":[{"delta":{"content":"priced"}}],"usage":{"prompt_tokens":1,"completion_tokens":1},"total_cost":0.029}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + + +class CostCommentFinalOverrideStreamingClient: + async def acompletion(self, **kwargs): + async def chunks(): + yield ': cost 0.031\n\n' + yield 'data: {"choices":[{"delta":{"content":"priced"}}],"usage":{"prompt_tokens":1,"completion_tokens":1,"cost_details":{"total_cost":0.062,"source":"final_usage"}}}\n\n' + yield "data: [DONE]\n\n" + + return chunks() + + +class FailingStore: + async def save(self, response): + raise RuntimeError("store failed Authorization: Bearer secret-token") + + async def get(self, response_id): + return None + + async def delete(self, response_id): + return False + + async def list_input_items(self, response_id): + return None + + +@pytest.mark.asyncio +async def test_stream_response_emits_responses_sse_events_and_stores_final_response() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, FakeStreamingClient())] + + event_text = "".join(events) + assert _event_names(events) == [ + "response.created", + "response.output_item.added", + "response.output_text.delta", + "response.output_text.delta", + "response.output_item.done", + "response.completed", + ] + assert event_text.endswith("data: [DONE]\n\n") + + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.output_items[0]["content"][0]["text"] == "Hello" + assert stored.usage == {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3, "cost_details": {"total_cost": 0.044, "source": "stream_provider"}} + + +@pytest.mark.asyncio +async def test_native_responses_stream_uses_agenerate_and_stores_terminal_object() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + client = NativeResponsesStreamingClient() + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, client)] + + assert client.calls[0][1]["input_protocol"] == "responses" + assert client.calls[0][1]["output_protocol"] == "responses" + assert "event: response.output_text.delta" in "".join(events) + stored = await store.get("resp_native") + assert stored is not None + assert stored.output_items[0]["content"][0]["text"] == "native" + + +@pytest.mark.asyncio +async def test_native_responses_stream_reformats_after_storing_responses_object() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + client = NativeResponsesStreamingClient("anthropic_messages") + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, client)] + + output = "".join(events) + assert "event: message_start" in output + assert "native" in output + assert "event: message_stop" in output + assert await store.get("resp_native") is not None + + +@pytest.mark.asyncio +async def test_stream_response_store_false_does_not_persist(tmp_path) -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True, "store": False}, FakeStreamingClient(), transaction_logger=logger)] + + response_id = events[0].split('"id": "')[1].split('"')[0] + assert await store.get(response_id) is None + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "responses_store_skipped" in trace_text + assert "responses_stored_stream_response" not in trace_text + + +@pytest.mark.asyncio +async def test_stream_response_errors_emit_failed_event() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, FailingStreamingClient())] + + event_text = "".join(events) + assert "event: response.failed" in event_text + assert "stream exploded" in event_text + assert event_text.endswith("data: [DONE]\n\n") + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.status == "failed" + + +@pytest.mark.asyncio +async def test_stream_response_error_chunks_store_failed_with_partial_output() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, ErrorChunkStreamingClient())] + + event_text = "".join(events) + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert "event: response.failed" in event_text + assert stored is not None + assert stored.status == "failed" + assert stored.output_items[0]["content"][0]["text"] == "partial" + + +@pytest.mark.asyncio +async def test_stream_response_event_error_frames_are_failed() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, EventErrorStreamingClient())] + + event_text = "".join(events) + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert "event: response.failed" in event_text + assert stored is not None + assert stored.response["error"]["message"] == "event failed" + + +@pytest.mark.asyncio +async def test_stream_response_preserves_sse_cost_comment() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, CostCommentStreamingClient())] + + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.usage["cost_details"]["total_cost"] == 0.031 + assert stored.usage["cost_details"]["source"] == "responses_sse" + + +@pytest.mark.asyncio +async def test_stream_response_preserves_sse_cost_event() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, CostEventStreamingClient())] + + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.usage["cost_details"]["total_cost"] == 0.017 + assert stored.usage["cost_details"]["currency"] == "EUR" + + +@pytest.mark.asyncio +async def test_stream_response_preserves_reference_request_cost_comment() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, RequestCostCommentStreamingClient())] + + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.usage["cost_details"]["request_cost_usd"] == 0.023 + + +@pytest.mark.asyncio +async def test_stream_response_preserves_estimated_cost_comment() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, EstimatedCostCommentStreamingClient())] + + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.usage["cost_details"]["estimated_cost"] == 0.024 + + +@pytest.mark.asyncio +async def test_stream_response_preserves_top_level_chunk_cost_with_usage() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, TopLevelUsageCostStreamingClient())] + + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.usage["total_cost"] == 0.029 + + +@pytest.mark.asyncio +async def test_stream_response_final_usage_cost_overrides_sse_cost_comment() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, CostCommentFinalOverrideStreamingClient())] + + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stored is not None + assert stored.usage["cost_details"]["total_cost"] == 0.062 + assert stored.usage["cost_details"]["source"] == "final_usage" + + +@pytest.mark.asyncio +async def test_stream_response_can_skip_failed_storage() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store, store_settings=ResponsesStoreSettings(store_failed=False)) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, FailingStreamingClient())] + + response_id = events[0].split('"id": "')[1].split('"')[0] + assert await store.get(response_id) is None + + +@pytest.mark.asyncio +async def test_stream_events_can_store_in_progress_state() -> None: + store = InMemoryResponsesStore() + service = ResponsesService(store=store, store_settings=ResponsesStoreSettings(store_in_progress=True)) + stream = service.stream_events({"model": "gpt-test", "input": "Hello", "stream": True}, FakeStreamingClient()) + + created = await anext(stream) + stored = await store.get(created.payload["id"]) + await stream.aclose() + + assert stored is not None + assert stored.status == "in_progress" + + +@pytest.mark.asyncio +async def test_scoped_in_progress_state_is_immediately_capability_accessible() -> None: + store = InMemoryResponsesStore() + service = ResponsesService( + store=store, + store_settings=ResponsesStoreSettings(store_in_progress=True), + ) + request = { + "model": "gpt-test", + "input": "private stream", + "stream": True, + "api_keys": {"openai": ["private-key"]}, + "private": True, + } + request_scope = service.prepare_request_scope(request) + stream = service.stream_events( + request, + FakeStreamingClient(), + request_scope=request_scope, + ) + + created = await anext(stream) + stored = await service.get_response_with_access_token( + created.payload["id"], + request_scope.access_token, + ) + await stream.aclose() + + assert stored["status"] == "in_progress" + + +@pytest.mark.asyncio +async def test_stream_response_store_failures_emit_store_specific_trace(tmp_path) -> None: + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + service = ResponsesService(store=FailingStore()) + + with pytest.raises(RuntimeError): + _ = [ + chunk + async for chunk in service.stream_response( + {"model": "gpt-test", "input": "Hello", "stream": True}, + FakeStreamingClient(), + transaction_logger=logger, + ) + ] + + entries = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + errors = [entry for entry in entries if entry["pass_name"] == "transform_log_error"] + assert any(entry["data"]["failed_pass_name"] == "responses_store_stream_response" for entry in errors) + assert "secret-token" not in json.dumps(errors) + + +@pytest.mark.asyncio +async def test_stream_current_state_store_failures_emit_store_specific_trace(tmp_path) -> None: + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + service = ResponsesService(store=FailingStore(), store_settings=ResponsesStoreSettings(store_in_progress=True)) + + with pytest.raises(RuntimeError): + _ = [ + event + async for event in service.stream_events( + {"model": "gpt-test", "input": "Hello", "stream": True}, + FakeStreamingClient(), + transaction_logger=logger, + ) + ] + + entries = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + errors = [entry for entry in entries if entry["pass_name"] == "transform_log_error"] + assert any(entry["data"]["failed_pass_name"] == "responses_store_stream_current_state" for entry in errors) + assert "secret-token" not in json.dumps(errors) + + +@pytest.mark.asyncio +async def test_stream_response_failure_trace_scrubs_header_like_secret_text(tmp_path) -> None: + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + service = ResponsesService(store=InMemoryResponsesStore()) + + _ = [ + chunk + async for chunk in service.stream_response( + {"model": "gpt-test", "input": "Hello", "stream": True}, + FailingStreamingClient("{'Authorization': 'Bearer secret-token'}"), + transaction_logger=logger, + ) + ] + + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "secret-token" not in trace_text + assert "[REDACTED]" in trace_text + + +def test_transport_formatters_expose_sse_and_websocket_seam() -> None: + assert ResponsesSSEFormatter().transport == "sse" + assert ResponsesSSEFormatter().format_stream_event(ResponsesStreamEvent("heartbeat", {"comment": "heartbeat"})) == ": heartbeat\n\n" + websocket = ResponsesWebSocketFormatter() + assert websocket.transport == "websocket" + assert websocket.future_supported is True + assert websocket.format_stream_event(ResponsesStreamEvent("response.created", {"id": "resp"})) == '{"event": "response.created", "data": {"id": "resp"}}' + + +@pytest.mark.asyncio +async def test_stream_response_emits_non_visible_heartbeat(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + stream = DelayedCloseableStream( + ['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', "data: [DONE]\n\n"], + [0.03, 0.0], + ) + service = ResponsesService(store=InMemoryResponsesStore()) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, DelayedStreamingClient(stream))] + + assert ": heartbeat\n\n" in events + assert "event: response.output_text.delta" in "".join(events) + + +@pytest.mark.asyncio +async def test_stream_response_ttfb_timeout_closes_upstream(monkeypatch) -> None: + monkeypatch.setenv("STREAM_TTFB_TIMEOUT_SECONDS", "0.01") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], [0.05]) + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, DelayedStreamingClient(stream))] + + event_text = "".join(events) + assert stream.closed is True + assert "event: response.failed" in event_text + assert "ttfb" in event_text.lower() + assert event_text.endswith("data: [DONE]\n\n") + + +@pytest.mark.asyncio +async def test_stream_response_heartbeat_does_not_reset_ttfb_timeout(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + monkeypatch.setenv("STREAM_TTFB_TIMEOUT_SECONDS", "0.03") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], [0.1]) + service = ResponsesService(store=InMemoryResponsesStore()) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, DelayedStreamingClient(stream))] + + event_text = "".join(events) + assert ": heartbeat\n\n" in events + assert stream.closed is True + assert "event: response.failed" in event_text + assert "ttfb" in event_text.lower() + + +@pytest.mark.asyncio +async def test_stream_response_acquire_wait_honors_ttfb_timeout(monkeypatch) -> None: + monkeypatch.setenv("STREAM_TTFB_TIMEOUT_SECONDS", "0.01") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], [0.0]) + service = ResponsesService(store=InMemoryResponsesStore()) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, SlowAcquireStreamingClient(stream, 0.05))] + + event_text = "".join(events) + assert "event: response.failed" in event_text + assert "ttfb" in event_text.lower() + + +@pytest.mark.asyncio +async def test_stream_response_ttfb_spans_acquire_and_first_chunk(monkeypatch) -> None: + monkeypatch.setenv("STREAM_TTFB_TIMEOUT_SECONDS", "0.05") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], [0.03]) + service = ResponsesService(store=InMemoryResponsesStore()) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, SlowAcquireStreamingClient(stream, 0.03))] + + event_text = "".join(events) + assert stream.closed is True + assert "event: response.failed" in event_text + assert "ttfb" in event_text.lower() + + +@pytest.mark.asyncio +async def test_stream_response_acquire_wait_can_emit_heartbeat(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', "data: [DONE]\n\n"], [0.0, 0.0]) + service = ResponsesService(store=InMemoryResponsesStore()) + + client = SlowAcquireStreamingClient(stream, 0.03) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, client)] + + assert ": heartbeat\n\n" in events + assert client.calls == 1 + assert "event: response.completed" in "".join(events) + + +@pytest.mark.asyncio +async def test_stream_response_heartbeat_does_not_drop_completed_first_chunk(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + stream = DelayedCloseableStream( + ['data: {"choices":[{"delta":{"content":"kept"}}]}\n\n', "data: [DONE]\n\n"], + [0.05, 0.0], + ) + service = ResponsesService(store=InMemoryResponsesStore()) + events = service.stream_events({"model": "gpt-test", "input": "Hello", "stream": True}, DelayedStreamingClient(stream)) + + created = await anext(events) + heartbeat = await anext(events) + await asyncio.sleep(0.06) + added = await anext(events) + delta = await anext(events) + await events.aclose() + + assert created.event_name == "response.created" + assert heartbeat.heartbeat is True + assert added.event_name == "response.output_item.added" + assert delta.event_name == "response.output_text.delta" + assert delta.payload["delta"] == "kept" + + +@pytest.mark.asyncio +async def test_stream_events_aclose_cancels_pending_read_after_heartbeat(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], [1.0]) + service = ResponsesService(store=InMemoryResponsesStore()) + events = service.stream_events({"model": "gpt-test", "input": "Hello", "stream": True}, DelayedStreamingClient(stream)) + + assert (await anext(events)).event_name == "response.created" + assert (await anext(events)).heartbeat is True + await events.aclose() + + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_stream_events_aclose_cancels_pending_acquire_after_heartbeat(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + client = CancelAwareAcquireClient() + service = ResponsesService(store=InMemoryResponsesStore()) + events = service.stream_events({"model": "gpt-test", "input": "Hello", "stream": True}, client) + + assert (await anext(events)).event_name == "response.created" + assert (await anext(events)).heartbeat is True + await events.aclose() + + assert client.cancelled is True + + +@pytest.mark.asyncio +async def test_stream_events_aclose_closes_completed_acquire_stream_after_heartbeat(monkeypatch) -> None: + monkeypatch.setenv("STREAM_HEARTBEAT_INTERVAL_SECONDS", "0.01") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], [1.0]) + client = SlowAcquireStreamingClient(stream, 0.015) + service = ResponsesService(store=InMemoryResponsesStore()) + events = service.stream_events({"model": "gpt-test", "input": "Hello", "stream": True}, client) + + assert (await anext(events)).event_name == "response.created" + assert (await anext(events)).heartbeat is True + await asyncio.sleep(0.02) + await events.aclose() + + assert client.calls == 1 + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_stream_response_stall_timeout_preserves_partial_output(monkeypatch) -> None: + monkeypatch.setenv("STREAM_STALL_TIMEOUT_SECONDS", "0.01") + stream = DelayedCloseableStream( + ['data: {"choices":[{"delta":{"content":"partial"}}]}\n\n', 'data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], + [0.0, 0.05], + ) + store = InMemoryResponsesStore() + service = ResponsesService(store=store) + + events = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, DelayedStreamingClient(stream))] + + event_text = "".join(events) + response_id = events[0].split('"id": "')[1].split('"')[0] + stored = await store.get(response_id) + assert stream.closed is True + assert "event: response.failed" in event_text + assert "stall" in event_text.lower() + assert stored is not None + assert stored.output_items[0]["content"][0]["text"] == "partial" + + +@pytest.mark.asyncio +async def test_stream_events_disconnect_closes_upstream(monkeypatch) -> None: + monkeypatch.setenv("STREAM_CANCEL_UPSTREAM_ON_DISCONNECT", "true") + stream = DelayedCloseableStream(['data: {"choices":[{"delta":{"content":"late"}}]}\n\n'], [0.05]) + service = ResponsesService(store=InMemoryResponsesStore()) + + events = [event async for event in service.stream_events({"model": "gpt-test", "input": "Hello", "stream": True}, DelayedStreamingClient(stream), request=DisconnectAfterAcquireRequest())] + + assert [event.event_name for event in events] == ["response.created"] + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_responses_stream_records_common_stream_metrics(tmp_path) -> None: + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + service = ResponsesService(store=InMemoryResponsesStore()) + + _ = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "Hello", "stream": True}, FakeStreamingClient(), transaction_logger=logger)] + + trace_text = (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8") + assert "stream_started" in trace_text + assert "stream_first_byte" in trace_text + assert "stream_first_visible_output" in trace_text + assert "stream_completed" in trace_text + assert "stream_metrics_final" in trace_text + assert "stream_done_event" in trace_text + assert "responses_stream_event_created" in trace_text + assert "responses_stream_event_output_text_delta" in trace_text + assert "responses_stream_event_completed" in trace_text + assert "responses_sse_formatted_event" in trace_text + usage_entry = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines() if '"usage_accounting_summary"' in line][-1] + assert usage_entry["data"]["cost"]["provider_reported_cost"] == 0.044 + + +@pytest.mark.asyncio +async def test_stream_events_are_transport_neutral_and_sse_wraps_them() -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + + events = [event async for event in service.stream_events({"model": "gpt-test", "input": "Hello", "stream": True}, FakeStreamingClient())] + sse = [ResponsesSSEFormatter().format_stream_event(event) for event in events] + + assert [event.event_name for event in events if not event.terminal] == [ + "response.created", + "response.output_item.added", + "response.output_text.delta", + "response.output_text.delta", + "response.output_item.done", + "response.completed", + ] + assert events[-1].terminal is True + assert "".join(sse).endswith("data: [DONE]\n\n") diff --git a/tests/test_responses_usage_accounting.py b/tests/test_responses_usage_accounting.py new file mode 100644 index 000000000..789868bb0 --- /dev/null +++ b/tests/test_responses_usage_accounting.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.responses import InMemoryResponsesStore, ResponsesService +from rotator_library.transaction_logger import TransactionLogger + + +class FakeUsageClient: + async def acompletion(self, **kwargs): + return { + "id": "chatcmpl_1", + "model": "gpt-test", + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 20, + "completion_tokens": 8, + "completion_tokens_details": {"reasoning_tokens": 3}, + }, + } + + +class FakeStreamingUsageClient: + async def acompletion(self, **kwargs): + async def gen(): + yield 'data: {"id":"chunk_1","choices":[{"delta":{"content":"hi"}}]}\n\n' + yield 'data: {"id":"chunk_2","choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":20,"completion_tokens":8,"total_tokens":28,"completion_tokens_details":{"reasoning_tokens":3}}}\n\n' + + return gen() + + +@pytest.mark.asyncio +async def test_responses_create_traces_normalized_usage(tmp_path) -> None: + logger = TransactionLogger("responses", "gpt-test", parent_dir=tmp_path) + service = ResponsesService(store=InMemoryResponsesStore()) + + await service.create_response({"model": "gpt-test", "input": "hello"}, FakeUsageClient(), transaction_logger=logger) + + entries = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + usage_entries = [entry for entry in entries if entry["pass_name"] == "usage_accounting_summary"] + assert usage_entries + assert usage_entries[-1]["data"]["usage"]["completion_tokens"] == 5 + assert usage_entries[-1]["data"]["usage"]["reasoning_tokens"] == 3 + + +@pytest.mark.asyncio +async def test_responses_stream_preserves_usage_details_in_completed_event() -> None: + service = ResponsesService(store=InMemoryResponsesStore()) + + chunks = [chunk async for chunk in service.stream_response({"model": "gpt-test", "input": "hello", "stream": True}, FakeStreamingUsageClient())] + completed = [chunk for chunk in chunks if "response.completed" in chunk][0] + payload = json.loads(completed.split("data: ", 1)[1]) + + assert payload["usage"]["output_tokens_details"] == {"reasoning_tokens": 3} diff --git a/tests/test_retry_policy.py b/tests/test_retry_policy.py new file mode 100644 index 000000000..bcfaab5a1 --- /dev/null +++ b/tests/test_retry_policy.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import asyncio + +from rotator_library.error_handler import ClassifiedError, PreRequestCallbackError, classify_error +from rotator_library.core.errors import structured_api_response_error +from rotator_library.retry_policy import ( + FailureHistory, + classify_route_error, + decide_provider_cooldown, + is_model_capacity_error, + is_target_failover_eligible, + should_retry_same_credential, + should_rotate_credential, +) + + +class ExplicitRouteError(Exception): + error_type = "unsupported_operation" + + +def _classified(error_type: str, **kwargs) -> ClassifiedError: + return ClassifiedError(error_type, original_exception=Exception(error_type), **kwargs) + + +def test_classifier_output_maps_to_fallback_policy_categories() -> None: + retryable = ["rate_limit", "quota_exceeded", "server_error", "api_connection", "unsupported_operation"] + stopped = [ + "authentication", + "forbidden", + "invalid_request", + "context_window_exceeded", + "credential_reauth_needed", + "pre_request_callback_error", + "cancelled", + ] + + assert all(is_target_failover_eligible(error_type) for error_type in retryable) + assert all(not is_target_failover_eligible(error_type) for error_type in stopped) + assert not is_target_failover_eligible("unknown") + + +def test_classify_route_error_preserves_explicit_and_cancelled_types() -> None: + assert classify_route_error(ExplicitRouteError()) == "unsupported_operation" + assert classify_route_error(asyncio.CancelledError()) == "cancelled" + assert classify_route_error(PreRequestCallbackError("boom")) == "pre_request_callback_error" + + +def test_retry_and_rotation_helpers_delegate_to_existing_semantics() -> None: + small_rate_limit = _classified("rate_limit", retry_after=3) + invalid_request = _classified("invalid_request") + + assert should_retry_same_credential(small_rate_limit, small_cooldown_threshold=10) is True + assert should_rotate_credential(small_rate_limit) is True + assert should_rotate_credential(invalid_request) is False + + +def test_provider_cooldown_uses_large_retry_after_not_small_retry_after() -> None: + small = decide_provider_cooldown( + _classified("rate_limit", retry_after=3), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + ) + large = decide_provider_cooldown( + _classified("rate_limit", retry_after=60), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + ) + + assert small.should_start is False + assert small.reason == "small_retry_after" + assert large.should_start is True + assert large.duration == 60 + assert large.scope == "provider" + + +def test_provider_cooldown_is_conservative_for_quota_by_default() -> None: + disabled = decide_provider_cooldown( + _classified("quota_exceeded", retry_after=3600), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + ) + enabled = decide_provider_cooldown( + _classified("quota_exceeded", retry_after=3600), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + cooldown_on_quota=True, + ) + + assert disabled.should_start is False + assert disabled.reason == "quota_cooldown_disabled" + assert enabled.should_start is True + + +def test_model_capacity_error_uses_model_scoped_cooldown() -> None: + error = Exception("503 MODEL_CAPACITY_EXHAUSTED") + decision = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=30, + model="gpt-test", + original_error=error, + ) + + assert is_model_capacity_error(error) is True + assert decision.should_start is True + assert decision.scope == "model" + assert decision.model == "gpt-test" + assert decision.reason == "model_capacity_cooldown" + + +def test_failure_history_escalates_repeated_transient_backoff(monkeypatch) -> None: + now = 1000.0 + history = FailureHistory(clock=lambda: now) + monkeypatch.setenv("PROVIDER_BACKOFF_THRESHOLD", "3") + monkeypatch.setenv("PROVIDER_BACKOFF_BASE_SECONDS", "10") + monkeypatch.setenv("PROVIDER_BACKOFF_MAX_SECONDS", "40") + history.record(provider="openai", model=None, error_type="server_error", scope="provider", duration=10, reason="test") + history.record(provider="openai", model=None, error_type="server_error", scope="provider", duration=10, reason="test") + + decision = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=10, + provider="openai", + failure_history=history, + ) + + assert decision.duration == 10 + assert decision.backoff_level == 1 + history.record(provider="openai", model=None, error_type="server_error", scope="provider", duration=10, reason="test") + decision = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=10, + provider="openai", + failure_history=history, + ) + assert decision.duration == 20 + assert decision.backoff_level == 2 + + +def test_single_generic_transient_without_retry_after_does_not_cooldown(monkeypatch) -> None: + history = FailureHistory(clock=lambda: 1000.0) + monkeypatch.setenv("PROVIDER_BACKOFF_THRESHOLD", "3") + + decision = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=10, + provider="openai", + failure_history=history, + ) + + assert decision.should_start is False + assert decision.reason == "transient_backoff_threshold_not_met" + + +def test_failure_history_clear_resets_repeated_transient_backoff(monkeypatch) -> None: + history = FailureHistory(clock=lambda: 1000.0) + monkeypatch.setenv("PROVIDER_BACKOFF_THRESHOLD", "2") + history.record(provider="openai", model=None, error_type="server_error", scope="provider", duration=0, reason="first") + + history.clear(provider="openai", model="gpt-test") + decision = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=10, + provider="openai", + failure_history=history, + ) + + assert history.snapshot() == () + assert decision.should_start is False + + +def test_failure_history_backoff_is_provider_scoped(monkeypatch) -> None: + history = FailureHistory(clock=lambda: 1000.0) + monkeypatch.setenv("PROVIDER_BACKOFF_THRESHOLD", "2") + monkeypatch.setenv("PROVIDER_BACKOFF_BASE_SECONDS", "10") + history.record(provider="provider-a", model=None, error_type="server_error", scope="provider", duration=10, reason="test") + + provider_b = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=10, + provider="provider-b", + failure_history=history, + ) + provider_a = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=10, + provider="provider-a", + failure_history=history, + ) + + assert provider_b.backoff_level == 0 + assert provider_a.backoff_level == 1 + + +def test_failure_history_backoff_requires_matching_provider(monkeypatch) -> None: + history = FailureHistory(clock=lambda: 1000.0) + monkeypatch.setenv("PROVIDER_BACKOFF_THRESHOLD", "2") + history.record(provider="provider-a", model=None, error_type="server_error", scope="provider", duration=10, reason="test") + + decision = decide_provider_cooldown( + _classified("server_error"), + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + default_duration=10, + failure_history=history, + ) + + assert decision.backoff_level == 0 + + +def test_shared_classifier_handles_structured_dict_status_codes() -> None: + assert classify_error({"error": {"status": 401}}).error_type == "authentication" + assert classify_error({"error": {"details": {"status_code": 403}}}).error_type == "forbidden" + assert classify_error({"error": {"code": 429}}).error_type == "rate_limit" + + +def test_shared_classifier_handles_structured_dict_type_and_code_text() -> None: + assert classify_error({"error": {"type": "authentication"}}).error_type == "authentication" + assert classify_error({"error": {"code": "permission_denied"}}).error_type == "forbidden" + assert classify_error({"error": {"code": "context_length_exceeded"}}).error_type == "context_window_exceeded" + assert classify_error({"error": {"status_code": 400, "code": "context_length_exceeded"}}).error_type == "context_window_exceeded" + assert classify_error({"error": {"type": "rate_limit"}}).error_type == "rate_limit" + + +def test_structured_context_length_error_is_not_generic_invalid_request() -> None: + error = structured_api_response_error({ + "error": {"status_code": 400, "message": "Context length exceeded"} + }) + + assert error is not None + assert error.error_type == "context_window_exceeded" + + +def test_shared_classifier_preserves_explicit_error_type_attributes() -> None: + class ConfigurationFailure(Exception): + error_type = "configuration_error" + + assert classify_error(ConfigurationFailure("raw secret message")).error_type == "configuration_error" diff --git a/tests/test_routing_attempts.py b/tests/test_routing_attempts.py new file mode 100644 index 000000000..e084b022e --- /dev/null +++ b/tests/test_routing_attempts.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from rotator_library.core.types import RequestContext +from rotator_library.routing import clone_context_for_target, parse_route_target + + +def test_clone_context_for_target_updates_model_provider_without_mutating_original() -> None: + original = RequestContext( + model="requested", + provider="original", + kwargs={"model": "requested", "messages": []}, + streaming=False, + credentials=["cred-a"], + deadline=123.0, + usage_manager_key="original", + input_provider="original", + ) + target = parse_route_target("codex/gpt-5.1-codex@native") + + cloned = clone_context_for_target(original, target, credentials=["cred-b"], target_index=1) + + assert cloned.model == "codex/gpt-5.1-codex" + assert cloned.provider == "codex" + assert cloned.kwargs["model"] == "codex/gpt-5.1-codex" + assert cloned.credentials == ["cred-b"] + assert cloned.usage_manager_key == "codex" + assert cloned.routing_target_index == 1 + assert cloned.input_provider == "original" + assert original.model == "requested" + assert original.kwargs["model"] == "requested" + assert original.credentials == ["cred-a"] + + +def test_clone_context_for_target_preserves_request_metadata() -> None: + original = RequestContext( + model="requested", + provider="original", + kwargs={"model": "requested"}, + streaming=True, + credentials=["cred-a"], + deadline=123.0, + session_id="session-1", + classifier="global", + routing_group_name="chain", + ) + + cloned = clone_context_for_target(original, parse_route_target("openai/gpt-5.1")) + + assert cloned.streaming is True + assert cloned.session_id == "session-1" + assert cloned.classifier == "global" + assert cloned.routing_group_name == "chain" + + +def test_clone_context_for_target_preserves_global_session_domain() -> None: + original = RequestContext( + model="openai/gpt-5", + provider="openai", + kwargs={"model": "openai/gpt-5"}, + streaming=False, + credentials=["cred-a"], + deadline=123.0, + session_tracking_namespace="session-domain:public", + session_affinity_key="provider-native-affinity", + ) + + cloned = clone_context_for_target(original, parse_route_target("anthropic/claude")) + + assert cloned.session_tracking_namespace == "session-domain:public" + assert cloned.session_affinity_key is None + + +def test_clone_context_for_target_preserves_custom_session_namespace() -> None: + original = RequestContext( + model="openai/gpt-5", + provider="openai", + kwargs={"model": "openai/gpt-5"}, + streaming=False, + credentials=["cred-a"], + deadline=123.0, + session_tracking_namespace="custom-namespace", + ) + + cloned = clone_context_for_target(original, parse_route_target("anthropic/claude")) + + assert cloned.session_tracking_namespace == "custom-namespace" diff --git a/tests/test_routing_config.py b/tests/test_routing_config.py new file mode 100644 index 000000000..353f0da60 --- /dev/null +++ b/tests/test_routing_config.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import pytest + +from rotator_library.routing import RoutingConfigError, load_routing_config_from_env, parse_route_target + + +def test_parse_route_target_supports_execution_suffix() -> None: + target = parse_route_target("codex/gpt-5.1-codex@native") + + assert target.provider == "codex" + assert target.model == "gpt-5.1-codex" + assert target.execution == "native" + assert target.prefixed_model == "codex/gpt-5.1-codex" + + +def test_parse_route_target_rejects_missing_provider() -> None: + with pytest.raises(RoutingConfigError): + parse_route_target("gpt-5.1") + + +def test_load_routing_config_from_env_parses_group_and_model_route() -> None: + config = load_routing_config_from_env( + { + "FALLBACK_GROUPS": "sonnet_chain", + "FALLBACK_GROUP_SONNET_CHAIN": "claude_code/claude-sonnet-4-5,copilot/claude-sonnet-4-5@litellm_fallback", + "MODEL_ROUTE_CLAUDE_SONNET": "group:sonnet_chain", + } + ) + + assert tuple(config.fallback_groups) == ("sonnet_chain",) + assert config.fallback_groups["sonnet_chain"].targets[1].execution == "litellm_fallback" + assert config.model_routes["claude_sonnet"] == "group:sonnet_chain" + + +def test_load_routing_config_rejects_empty_group() -> None: + with pytest.raises(RoutingConfigError): + load_routing_config_from_env({"FALLBACK_GROUPS": "empty", "FALLBACK_GROUP_EMPTY": ""}) + + +def test_load_routing_config_rejects_duplicate_group_names() -> None: + with pytest.raises(RoutingConfigError): + load_routing_config_from_env({"FALLBACK_GROUPS": "a,a", "FALLBACK_GROUP_A": "openai/gpt"}) + + +def test_load_routing_config_rejects_unknown_group_route() -> None: + with pytest.raises(RoutingConfigError): + load_routing_config_from_env({"MODEL_ROUTE_CODEX": "group:missing"}) + + +def test_load_routing_config_parses_group_policy_overrides() -> None: + config = load_routing_config_from_env( + { + "FALLBACK_GROUPS": "chain", + "FALLBACK_GROUP_CHAIN": "openai/gpt,copilot/gpt", + "FALLBACK_GROUP_CHAIN_FAILOVER_ON": "network,transient", + "FALLBACK_GROUP_CHAIN_STOP_ON": "validation", + "FALLBACK_GROUP_CHAIN_STREAMING_POLICY": "never", + } + ) + + group = config.fallback_groups["chain"] + assert group.failover_on == frozenset({"api_connection", "server_error"}) + assert group.stop_on == frozenset({"invalid_request"}) + assert group.streaming_policy == "never" + + +def test_load_routing_config_rejects_hard_stop_failover() -> None: + with pytest.raises(RoutingConfigError): + load_routing_config_from_env( + { + "FALLBACK_GROUPS": "chain", + "FALLBACK_GROUP_CHAIN": "openai/gpt,copilot/gpt", + "FALLBACK_GROUP_CHAIN_FAILOVER_ON": "auth", + } + ) + + with pytest.raises(RoutingConfigError): + load_routing_config_from_env( + { + "FALLBACK_GROUPS": "chain", + "FALLBACK_GROUP_CHAIN": "openai/gpt,copilot/gpt", + "FALLBACK_GROUP_CHAIN_FAILOVER_ON": "pre_request_callback", + } + ) + + +def test_load_routing_config_rejects_unknown_streaming_policy() -> None: + with pytest.raises(RoutingConfigError): + load_routing_config_from_env( + { + "FALLBACK_GROUPS": "chain", + "FALLBACK_GROUP_CHAIN": "openai/gpt,copilot/gpt", + "FALLBACK_GROUP_CHAIN_STREAMING_POLICY": "always", + } + ) diff --git a/tests/test_selection_engine.py b/tests/test_selection_engine.py index 80d5768be..91508a47c 100644 --- a/tests/test_selection_engine.py +++ b/tests/test_selection_engine.py @@ -293,6 +293,103 @@ def test_session_sticky_env_config_is_parsed(self): self.assertEqual(config.session_sticky_entry_ttl_seconds, 123) self.assertEqual(config.session_sticky_max_entries, 456) + def test_session_sticky_default_matches_five_minute_cache_idle_window(self): + with patch.dict(os.environ, {}, clear=True): + config = load_provider_usage_config("gemini", {}) + + self.assertEqual(config.session_sticky_entry_ttl_seconds, 300) + + def test_provider_can_override_default_binding_idle_ttl(self): + class ProviderPolicy: + default_session_sticky_entry_ttl_seconds = 90 + + with patch.dict(os.environ, {}, clear=True): + config = load_provider_usage_config("custom", {"custom": ProviderPolicy}) + + self.assertEqual(config.session_sticky_entry_ttl_seconds, 90) + + def test_zero_binding_ttl_disables_sticky_entries(self): + strategy = SequentialStrategy(sticky_entry_ttl_seconds=0) + states = { + "cred_a": self._make_state("cred_a", "a"), + "cred_b": self._make_state("cred_b", "b"), + } + context = SelectionContext( + provider="custom", + model="model", + quota_group=None, + candidates=list(states), + priorities={key: 1 for key in states}, + usage_counts={key: 0 for key in states}, + rotation_mode=RotationMode.SEQUENTIAL, + rotation_tolerance=0.0, + deadline=0.0, + session_id="session", + ) + + strategy.select(context, states) + + self.assertIsNone(strategy.get_current("custom", "model", "session")) + self.assertFalse(strategy._current) + + def test_provider_binding_idle_clocks_are_independent(self): + session_id = "global-session" + deepseek = SequentialStrategy(sticky_entry_ttl_seconds=300) + anthropic = SequentialStrategy(sticky_entry_ttl_seconds=300) + deepseek_states = { + "deepseek_a": self._make_state("deepseek_a", "da"), + "deepseek_b": self._make_state("deepseek_b", "db"), + } + anthropic_states = { + "anthropic_a": self._make_state("anthropic_a", "aa"), + "anthropic_b": self._make_state("anthropic_b", "ab"), + } + + def context(provider: str, model: str, states) -> SelectionContext: + return SelectionContext( + provider=provider, + model=model, + quota_group=None, + candidates=list(states), + priorities={key: 1 for key in states}, + usage_counts={key: 0 for key in states}, + rotation_mode=RotationMode.SEQUENTIAL, + rotation_tolerance=0.0, + deadline=0.0, + session_id=session_id, + session_affinity_key="shared-global-affinity", + ) + + with patch("rotator_library.usage.selection.strategies.sequential.time.time", return_value=1000.0): + deepseek_selected = deepseek.select( + context("deepseek", "deepseek/chat", deepseek_states), + deepseek_states, + ) + anthropic_selected = anthropic.select( + context("anthropic", "anthropic/claude", anthropic_states), + anthropic_states, + ) + self.assertNotEqual(deepseek_selected, anthropic_selected) + + with patch("rotator_library.usage.selection.strategies.sequential.time.time", return_value=1200.0): + self.assertEqual( + anthropic.select( + context("anthropic", "anthropic/claude", anthropic_states), + anthropic_states, + ), + anthropic_selected, + ) + + with patch("rotator_library.usage.selection.strategies.sequential.time.time", return_value=1300.0): + self.assertIsNone( + deepseek.get_current("deepseek", "deepseek/chat", session_id) + ) + self.assertEqual( + anthropic.get_current("anthropic", "anthropic/claude", session_id), + anthropic_selected, + ) + self.assertIsNotNone(deepseek_selected) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_session_tracking.py b/tests/test_session_tracking.py index 22262252a..5c04fe13b 100644 --- a/tests/test_session_tracking.py +++ b/tests/test_session_tracking.py @@ -1,8 +1,12 @@ +import asyncio import os import sys import json import tempfile +import threading +import time import unittest +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from unittest.mock import patch @@ -15,234 +19,2953 @@ _AnchorRecord, ) from rotator_library.client.streaming import StreamingHandler +from rotator_library.client.rotating_client import _resolve_session_persistence_settings +from rotator_library.client.request_builder import RequestContextBuilder class SessionTrackerTests(unittest.TestCase): - def test_weak_first_user_only_does_not_reuse_session(self): - tracker = SessionTracker(ttl_seconds=3600) - request = {"messages": [{"role": "user", "content": "hello world"}]} + @staticmethod + def _long_text(label, repeats=12): + return " ".join([f"{label} carries deterministic substantial history evidence"] * repeats) - session_a = tracker.infer_session(request, provider="gemini", model="pro") - session_b = tracker.infer_session(request, provider="gemini", model="pro") + @staticmethod + def _response(content, response_id=None): + response = {"choices": [{"message": {"role": "assistant", "content": content}}]} + if response_id: + response["id"] = response_id + return response - self.assertIsNotNone(session_a.session_id) - self.assertIsNotNone(session_b.session_id) - self.assertNotEqual(session_a.session_id, session_b.session_id) - self.assertIsNone(session_a.affinity_key) + @staticmethod + def _tool_call_message(call_id, name, arguments="{}"): + return { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + ], + } - def test_single_long_prompt_does_not_become_sticky_by_itself(self): + @staticmethod + def _tool_result(call_id, content="ok"): + return {"role": "tool", "tool_call_id": call_id, "content": content} + + @staticmethod + def _persistent_tracker(path): + """Create a tracker whose every mutation is immediately restart-safe.""" + + return SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + + def test_proxy_session_persistence_environment_settings(self): + """Proxy defaults expose persistence without overriding explicit callers.""" + + with patch.dict(os.environ, {}, clear=True): + self.assertEqual( + _resolve_session_persistence_settings(None, None), + (False, 5.0), + ) + + with patch.dict( + os.environ, + { + "SESSION_PERSISTENCE_ENABLED": "true", + "SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS": "0", + }, + ): + self.assertEqual( + _resolve_session_persistence_settings(None, None), + (True, 0.0), + ) + self.assertEqual( + _resolve_session_persistence_settings(False, 2.5), + (False, 2.5), + ) + + with patch.dict( + os.environ, + {"SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS": "invalid"}, + ): + with self.assertLogs("rotator_library", level="WARNING") as captured: + enabled, interval = _resolve_session_persistence_settings(False, None) + self.assertFalse(enabled) + self.assertEqual(interval, 5.0) + self.assertIn("Invalid SESSION_PERSISTENCE_FLUSH_INTERVAL_SECONDS", captured.output[0]) + + def test_warning_lineage_logs_new_and_continued_session_ids(self): tracker = SessionTracker(ttl_seconds=3600) request = { "messages": [ - { - "role": "user", - "content": "Please perform a careful detailed analysis of this standalone prompt and provide a structured answer with several paragraphs.", - } + {"role": "user", "content": self._long_text("lineage-user")}, + {"role": "assistant", "content": self._long_text("lineage-assistant")}, ] } - session_a = tracker.infer_session(request, provider="gemini", model="pro") - session_b = tracker.infer_session(request, provider="gemini", model="pro") + with self.assertLogs("rotator_library", level="WARNING") as captured: + first = tracker.infer_session(request, provider="gemini", model="pro") + second = tracker.infer_session(request, provider="gemini", model="pro") - self.assertNotEqual(session_a.session_id, session_b.session_id) - self.assertIsNone(session_a.affinity_key) + self.assertEqual(first.session_id, second.session_id) + self.assertIn(f"action=new session_id={first.session_id}", captured.output[0]) + self.assertIn("matched_session_id=-", captured.output[0]) + self.assertIn("candidate_session_id=-", captured.output[0]) + self.assertIn(f"action=continue session_id={first.session_id}", captured.output[1]) + self.assertIn(f"matched_session_id={first.session_id}", captured.output[1]) + self.assertIn("origin=memory", captured.output[1]) - def test_shared_system_and_user_prompt_does_not_become_sticky_by_itself(self): + def test_warning_lineage_separates_rejected_candidate_from_accepted_match(self): tracker = SessionTracker(ttl_seconds=3600) - request = { + opening = "Tell me a story now" + established = { + "messages": [ + {"role": "user", "content": opening}, + {"role": "assistant", "content": self._long_text("candidate-assistant")}, + {"role": "user", "content": "Add another paragraph to the story."}, + ] + } + first = tracker.infer_session(established, provider="deepseek", model="chat") + + with self.assertLogs("rotator_library", level="WARNING") as captured: + independent = tracker.infer_session( + {"messages": [{"role": "user", "content": opening}]}, + provider="deepseek", + model="chat", + ) + + self.assertNotEqual(first.session_id, independent.session_id) + self.assertIn(f"action=new session_id={independent.session_id}", captured.output[0]) + self.assertIn("matched_session_id=-", captured.output[0]) + self.assertIn(f"candidate_session_id={first.session_id}", captured.output[0]) + self.assertIn("confidence=weak score=5", captured.output[0]) + + def test_warning_lineage_logs_compaction_child_and_exact_replay(self): + tracker = SessionTracker(ttl_seconds=3600) + original = { + "messages": [ + {"role": "user", "content": self._long_text("log-parent-user")}, + {"role": "assistant", "content": self._long_text("log-parent-assistant")}, + {"role": "user", "content": self._long_text("log-parent-follow-up")}, + {"role": "assistant", "content": self._long_text("log-parent-later")}, + ] + } + parent = tracker.infer_session(original, provider="gemini", model="pro") + compacted = { "messages": [ { "role": "system", - "content": "You are a shared coding harness with stable instructions used by many independent sessions.", + "content": "Summary of previous conversation: " + + original["messages"][1]["content"], }, + {"role": "user", "content": "Continue from compacted context."}, + ] + } + + with self.assertLogs("rotator_library", level="WARNING") as captured: + child = tracker.infer_session(compacted, provider="gemini", model="pro") + replay = tracker.infer_session(compacted, provider="gemini", model="pro") + + self.assertEqual(child.session_id, replay.session_id) + self.assertIn(f"action=compaction_child session_id={child.session_id}", captured.output[0]) + self.assertIn(f"parent_session_id={parent.session_id}", captured.output[0]) + self.assertIn("marker=True", captured.output[0]) + self.assertIn(f"action=compaction_replay session_id={child.session_id}", captured.output[1]) + self.assertIn(f"matched_session_id={child.session_id}", captured.output[1]) + self.assertIn(f"parent_session_id={parent.session_id}", captured.output[1]) + + def test_warning_lineage_logs_persisted_session_origin(self): + request = { + "messages": [ + {"role": "user", "content": self._long_text("persisted-log-user")}, { - "role": "user", - "content": "Please review this standalone request carefully and provide the exact structured output.", + "role": "assistant", + "content": self._long_text("persisted-log-assistant"), }, ] } + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + original = tracker.infer_session(request, provider="gemini", model="pro") + tracker.flush() + restored = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) - session_a = tracker.infer_session(request, provider="gemini", model="pro") - session_b = tracker.infer_session(request, provider="gemini", model="pro") - - self.assertNotEqual(session_a.session_id, session_b.session_id) - self.assertIsNone(session_a.affinity_key) + with self.assertLogs("rotator_library", level="WARNING") as captured: + continued = restored.infer_session(request, provider="gemini", model="pro") - def test_explicit_ids_are_weak_unless_configured_as_trusted(self): - request = {"conversation_id": "stable-client-id"} + self.assertEqual(original.session_id, continued.session_id) + self.assertIn(f"action=continue session_id={original.session_id}", captured.output[0]) + self.assertIn(f"matched_session_id={original.session_id}", captured.output[0]) + self.assertIn("origin=persisted", captured.output[0]) - conservative = SessionTracker(ttl_seconds=3600) - first = conservative.infer_session(request, provider="gemini", model="pro") - second = conservative.infer_session(request, provider="gemini", model="pro") - self.assertNotEqual(first.session_id, second.session_id) + def test_warning_lineage_logs_untracked_request(self): + tracker = SessionTracker(ttl_seconds=3600) - trusted = SessionTracker( - ttl_seconds=3600, - trusted_explicit_fields=["conversation_id"], - ) - first = trusted.infer_session(request, provider="gemini", model="pro") - second = trusted.infer_session(request, provider="gemini", model="pro") - self.assertEqual(first.session_id, second.session_id) + with self.assertLogs("rotator_library", level="WARNING") as captured: + inferred = tracker.infer_session({}, provider="gemini", model="pro") - def test_trusted_explicit_fields_can_come_from_env(self): - request = {"conversation_id": "stable-client-id"} + self.assertIsNone(inferred.session_id) + self.assertIn("action=untracked session_id=-", captured.output[0]) + self.assertIn("origin=none", captured.output[0]) - with patch.dict(os.environ, {"TRUSTED_SESSION_ID_FIELDS": "conversation_id"}): - tracker = SessionTracker(ttl_seconds=3600) + def test_agentic_tool_loop_persists_compacts_replays_and_continues_child(self): + """Exercise a complete persisted agent loop and compacted-child lifecycle.""" - first = tracker.infer_session(request, provider="gemini", model="pro") - second = tracker.infer_session(request, provider="gemini", model="pro") + routes = [ + ("deepseek", "deepseek/v4"), + ("deepseek", "deepseek/v4-flash"), + ("anthropic", "anthropic/claude-sonnet"), + ("anthropic", "anthropic/claude-opus"), + ("openai", "openai/gpt-5"), + ("deepseek", "deepseek/v4"), + ] + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "agentic_session_stickiness.json" + tracker = self._persistent_tracker(path) + history = [ + { + "role": "user", + "content": self._long_text("agentic-root-task", repeats=16), + } + ] + root_session_id = None + assistant_responses = [] - self.assertEqual(first.session_id, second.session_id) + for round_index in range(6): + provider, model = routes[round_index] + if round_index in {2, 4}: + tracker.flush() + tracker = self._persistent_tracker(path) + self.assertTrue(tracker._sessions[root_session_id].loaded_from_persistence) - def test_multiple_message_anchors_reuse_session_when_tools_are_pruned(self): - tracker = SessionTracker(ttl_seconds=3600) + inferred = tracker.infer_session( + {"messages": list(history)}, + provider=provider, + model=model, + ) + if root_session_id is None: + root_session_id = inferred.session_id + self.assertEqual(inferred.session_id, root_session_id) + self.assertFalse(inferred.possible_compaction) + self.assertEqual(inferred.tracking_namespace, "session-domain:public") + if round_index: + self.assertIn(inferred.confidence, {"probable", "strong"}) - original = { - "messages": [ - {"role": "user", "content": "Please inspect the account quota and explain the cooldown status in detail."}, - { + tool_id = f"call_agentic_{round_index}" + assistant_text = self._long_text( + f"agentic-assistant-analysis-{round_index}", + repeats=14, + ) + assistant_responses.append(assistant_text) + assistant_message = { "role": "assistant", - "content": "I will inspect the quota files and compare the cooldown windows.", + "content": assistant_text, "tool_calls": [ - {"id": "call_123", "type": "function", "function": {"name": "read"}} + { + "id": tool_id, + "type": "function", + "function": { + "name": f"agent_tool_{round_index}", + "arguments": "{}", + }, + } ], + } + tracker.record_response( + inferred.session_id, + provider=provider, + model=model, + tracking_namespace=inferred.tracking_namespace, + response={"choices": [{"message": assistant_message}]}, + ) + history.extend( + [ + assistant_message, + { + "role": "tool", + "tool_call_id": tool_id, + "content": self._long_text( + f"agentic-tool-result-{round_index}", + repeats=12, + ), + }, + { + "role": "user", + "content": self._long_text( + f"agentic-next-instruction-{round_index}", + repeats=10, + ), + }, + ] + ) + + # Six responses were appended, but the sixth response/tool/user tail + # has not yet appeared in a subsequent request. The high-water request + # therefore contains 16 rather than all 19 locally accumulated messages. + parent_high_water = tracker._sessions[root_session_id].history_signatures + self.assertEqual(len(parent_high_water), 16) + tracker.flush() + tracker = self._persistent_tracker(path) + + # A real compacted payload drops every strong tool ID and replaces + # almost all prior messages with one durable summary. + compacted = { + "messages": [ + { + "role": "system", + "content": "Compacted context: " + + assistant_responses[1] + + " " + + assistant_responses[4], + }, + { + "role": "user", + "content": self._long_text("agentic-resume-after-compaction"), + }, + ] + } + child = tracker.infer_session( + compacted, + provider="anthropic", + model="anthropic/claude-opus", + ) + self.assertTrue(child.possible_compaction) + self.assertEqual(child.lineage_parent_session_id, root_session_id) + self.assertNotEqual(child.session_id, root_session_id) + + child_response = self._long_text("agentic-child-response", repeats=14) + tracker.record_response( + child.session_id, + provider="anthropic", + model="anthropic/claude-opus", + tracking_namespace=child.tracking_namespace, + response=self._response(child_response), + ) + tracker.flush() + tracker = self._persistent_tracker(path) + + replay = tracker.infer_session( + compacted, + provider="openai", + model="openai/gpt-5", + ) + self.assertEqual(replay.session_id, child.session_id) + self.assertEqual(replay.lineage_parent_session_id, root_session_id) + self.assertTrue(replay.possible_compaction) + self.assertEqual(replay.confidence, "strong") + + child_follow_up = { + "messages": [ + *compacted["messages"], + {"role": "assistant", "content": child_response}, + { + "role": "user", + "content": self._long_text("agentic-child-next-step"), + }, + ] + } + continued_child = tracker.infer_session( + child_follow_up, + provider="deepseek", + model="deepseek/v4-flash", + ) + self.assertEqual(continued_child.session_id, child.session_id) + self.assertNotEqual(continued_child.session_id, root_session_id) + self.assertFalse(continued_child.possible_compaction) + self.assertEqual(continued_child.lineage_parent_session_id, root_session_id) + + parent_continuation = tracker.infer_session( + {"messages": list(history)}, + provider="openai", + model="openai/gpt-5", + ) + self.assertEqual(parent_continuation.session_id, root_session_id) + self.assertNotEqual(parent_continuation.session_id, child.session_id) + + def test_long_normal_conversation_tracks_every_turn_across_restarts(self): + """Grow a long ordinary chat turn-by-turn without false compaction.""" + + routes = [ + ("deepseek", "deepseek/chat"), + ("openai", "openai/gpt-5"), + ("anthropic", "anthropic/claude-sonnet"), + ("deepseek", "deepseek/reasoner"), + ("openai", "openai/gpt-5-mini"), + ("anthropic", "anthropic/claude-opus"), + ("deepseek", "deepseek/chat"), + ("openai", "openai/gpt-5"), + ] + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "normal_session_stickiness.json" + tracker = self._persistent_tracker(path) + history = [] + session_id = None + first_response = None + request_lengths = [] + + for turn_index in range(8): + provider, model = routes[turn_index] + if turn_index in {3, 6}: + tracker.flush() + tracker = self._persistent_tracker(path) + self.assertTrue(tracker._sessions[session_id].loaded_from_persistence) + + history.append( + { + "role": "user", + "content": self._long_text( + f"normal-conversation-question-{turn_index}", + repeats=14, + ), + } + ) + inferred = tracker.infer_session( + {"messages": list(history)}, + provider=provider, + model=model, + ) + if session_id is None: + session_id = inferred.session_id + self.assertEqual(inferred.session_id, session_id) + self.assertFalse(inferred.possible_compaction) + self.assertEqual(inferred.tracking_namespace, "session-domain:public") + request_lengths.append(len(history)) + self.assertEqual( + len(tracker._sessions[session_id].history_signatures), + len(history), + ) + + response_text = ( + first_response + if turn_index == 4 + else self._long_text( + f"normal-conversation-answer-{turn_index}", + repeats=16, + ) + ) + if first_response is None: + first_response = response_text + tracker.record_response( + inferred.session_id, + provider=provider, + model=model, + tracking_namespace=inferred.tracking_namespace, + response=self._response(response_text), + ) + history.append({"role": "assistant", "content": response_text}) + + self.assertEqual(request_lengths, [1, 3, 5, 7, 9, 11, 13, 15]) + self.assertEqual( + len(tracker._sessions[session_id].history_signatures), + request_lengths[-1], + ) + response_groups = { + record.group + for record in tracker._anchors.values() + if record.session_id == session_id + and record.source == "response" + and record.group + and record.group.startswith("response_event:") + } + self.assertEqual(len(response_groups), 7) + + def test_roleplay_redo_edits_rollback_and_middle_rewrite_stay_on_session(self): + """Model realistic roleplay regeneration and branch editing behavior.""" + + routes = [ + ("mistral", "mistral/large"), + ("anthropic", "anthropic/claude-sonnet"), + ("mistral", "mistral/large"), + ("openai", "openai/gpt-5"), + ("anthropic", "anthropic/claude-opus"), + ("mistral", "mistral/large"), + ] + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "roleplay_session_stickiness.json" + tracker = self._persistent_tracker(path) + history = [ + { + "role": "system", + "content": self._long_text("roleplay-world-rules", repeats=14), + }, + { + "role": "user", + "content": self._long_text("roleplay-opening-scene", repeats=12), }, - {"role": "tool", "tool_call_id": "call_123", "content": "quota payload"}, - {"role": "user", "content": "Now summarize the result and note what changed since yesterday."}, ] - } - pruned_tools = { - "messages": [ - {"role": "user", "content": "Please inspect the account quota and explain the cooldown status in detail."}, + session_id = None + request_snapshots = [] + + for turn_index in range(6): + provider, model = routes[turn_index] + request_snapshot = [dict(message) for message in history] + request_snapshots.append(request_snapshot) + inferred = tracker.infer_session( + {"messages": request_snapshot}, + provider=provider, + model=model, + ) + if session_id is None: + session_id = inferred.session_id + self.assertEqual(inferred.session_id, session_id) + self.assertFalse(inferred.possible_compaction) + self.assertIsNone(inferred.lineage_parent_session_id) + + original_response = self._long_text( + f"roleplay-assistant-scene-{turn_index}", + repeats=15, + ) + tracker.record_response( + session_id, + provider=provider, + model=model, + tracking_namespace=inferred.tracking_namespace, + response=self._response(original_response), + ) + + if turn_index == 2: + # Retrying generation sends the byte-identical request without + # appending the response that is being regenerated. + same_provider_redo = tracker.infer_session( + {"messages": [dict(message) for message in request_snapshot]}, + provider=provider, + model=model, + ) + self.assertEqual(same_provider_redo.session_id, session_id) + self.assertFalse(same_provider_redo.possible_compaction) + + exact_redo = tracker.infer_session( + {"messages": [dict(message) for message in request_snapshot]}, + provider="openai", + model="openai/gpt-5-mini", + ) + self.assertEqual(exact_redo.session_id, session_id) + self.assertFalse(exact_redo.possible_compaction) + chosen_response = self._long_text( + "roleplay-assistant-scene-2-edited-redo", + repeats=15, + ) + tracker.record_response( + session_id, + provider="openai", + model="openai/gpt-5-mini", + tracking_namespace=exact_redo.tracking_namespace, + response=self._response(chosen_response), + ) + else: + chosen_response = original_response + + history.append({"role": "assistant", "content": chosen_response}) + if turn_index < 5: + history.append( + { + "role": "user", + "content": self._long_text( + f"roleplay-user-action-{turn_index + 1}", + repeats=11, + ), + } + ) + + if turn_index == 3: + tracker.flush() + tracker = self._persistent_tracker(path) + self.assertTrue(tracker._sessions[session_id].loaded_from_persistence) + + high_water_before_edits = tracker._sessions[session_id].history_signatures + self.assertEqual(len(high_water_before_edits), len(request_snapshots[-1])) + + # Rewrite an assistant response in the middle while preserving every + # later turn. Seven-plus surviving messages must dominate the edit. + middle_edited_history = [dict(message) for message in history] + middle_assistant_index = 6 + self.assertEqual(middle_edited_history[middle_assistant_index]["role"], "assistant") + middle_edited_history[middle_assistant_index] = { + "role": "assistant", + "content": self._long_text("roleplay-middle-response-rewritten", repeats=15), + } + middle_edit = tracker.infer_session( + {"messages": middle_edited_history}, + provider="anthropic", + model="anthropic/claude-opus", + ) + self.assertEqual(middle_edit.session_id, session_id) + self.assertFalse(middle_edit.possible_compaction) + self.assertIsNone(middle_edit.lineage_parent_session_id) + + # Roll back to an older request snapshot, then resume a different + # branch. The shorter request must not reduce the high-water profile. + # This rollback retains less than half the high-water history. It is + # still ordinary branching because it has neither a summary marker nor + # response-derived probe evidence. + rollback_request = [dict(message) for message in request_snapshots[2]] + high_water_before_rollback = tracker._sessions[session_id].history_signatures + rolled_back = tracker.infer_session( + {"messages": rollback_request}, + provider="deepseek", + model="deepseek/v4-flash", + ) + self.assertEqual(rolled_back.session_id, session_id) + self.assertFalse(rolled_back.possible_compaction) + self.assertIsNone(rolled_back.lineage_parent_session_id) + self.assertEqual( + tracker._sessions[session_id].history_signatures, + high_water_before_rollback, + ) + + branch_response = self._long_text("roleplay-rollback-new-response", repeats=15) + tracker.record_response( + session_id, + provider="deepseek", + model="deepseek/v4-flash", + tracking_namespace=rolled_back.tracking_namespace, + response=self._response(branch_response), + ) + resumed_branch = [ + *rollback_request, + {"role": "assistant", "content": branch_response}, { - "role": "assistant", - "content": "I will inspect the quota files and compare the cooldown windows.", + "role": "user", + "content": self._long_text("roleplay-rollback-new-user-action"), }, - {"role": "user", "content": "Now summarize the result and note what changed since yesterday."}, ] - } + resumed = tracker.infer_session( + {"messages": resumed_branch}, + provider="openai", + model="openai/gpt-5", + ) + self.assertEqual(resumed.session_id, session_id) + self.assertFalse(resumed.possible_compaction) + self.assertIsNone(resumed.lineage_parent_session_id) + self.assertEqual( + tracker._sessions[session_id].history_signatures, + high_water_before_rollback, + ) - session_a = tracker.infer_session(original, provider="gemini", model="pro") - session_b = tracker.infer_session(pruned_tools, provider="gemini", model="pro") + def test_long_changed_tail_continues_persisted_compaction_context(self): + """A long changing user tail must not become part of the summary key.""" + + provider = "gemini" + model = "pro" + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "long_tail_session_stickiness.json" + tracker = self._persistent_tracker(path) + parent_request = { + "messages": [ + {"role": "user", "content": self._long_text("long-tail-parent-user-1")}, + { + "role": "assistant", + "content": self._long_text("long-tail-parent-assistant-1"), + }, + {"role": "user", "content": self._long_text("long-tail-parent-user-2")}, + { + "role": "assistant", + "content": self._long_text("long-tail-parent-assistant-2"), + }, + ] + } + parent = tracker.infer_session(parent_request, provider=provider, model=model) + summary = { + "role": "system", + "content": "Summary of previous conversation: " + + parent_request["messages"][1]["content"], + } + first_compacted = { + "messages": [ + summary, + { + "role": "user", + "content": self._long_text("long-tail-first-instruction", repeats=20), + }, + ] + } + child = tracker.infer_session(first_compacted, provider=provider, model=model) + self.assertTrue(child.possible_compaction) + self.assertEqual(child.lineage_parent_session_id, parent.session_id) + tracker.flush() + tracker = self._persistent_tracker(path) + + changed_tail = { + "messages": [ + summary, + { + "role": "user", + "content": self._long_text("long-tail-second-instruction", repeats=20), + }, + ] + } + continued = tracker.infer_session(changed_tail, provider=provider, model=model) + + self.assertEqual(continued.session_id, child.session_id) + self.assertEqual(continued.lineage_parent_session_id, parent.session_id) + self.assertFalse(continued.possible_compaction) + self.assertEqual(continued.confidence, "strong") + + def test_expired_compaction_child_is_recreated_from_live_parent(self): + """Expired child bindings cannot revive, while a live parent can fork anew.""" + + with patch("rotator_library.session_tracking.time.time", return_value=1000.0): + tracker = SessionTracker(ttl_seconds=10) + parent_request = { + "messages": [ + {"role": "user", "content": self._long_text("ttl-parent-user")}, + { + "role": "assistant", + "content": self._long_text("ttl-parent-assistant"), + }, + {"role": "user", "content": self._long_text("ttl-parent-follow-up")}, + ] + } + parent = tracker.infer_session(parent_request, provider="gemini", model="pro") + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + tracking_namespace=parent.tracking_namespace, + response=self._response(parent_request["messages"][1]["content"]), + ) + compacted = { + "messages": [ + { + "role": "system", + "content": "Compacted context: " + + parent_request["messages"][1]["content"], + }, + {"role": "user", "content": "Continue after compaction."}, + ] + } + child = tracker.infer_session(compacted, provider="gemini", model="pro") + + with patch("rotator_library.session_tracking.time.time", return_value=1005.0): + refreshed_parent = tracker.infer_session( + parent_request, + provider="openai", + model="gpt", + ) + self.assertEqual(refreshed_parent.session_id, parent.session_id) + + with patch("rotator_library.session_tracking.time.time", return_value=1011.0): + after_expiry = tracker.infer_session(compacted, provider="gemini", model="pro") + + self.assertNotEqual(after_expiry.session_id, child.session_id) + self.assertTrue(after_expiry.possible_compaction) + self.assertEqual(after_expiry.lineage_parent_session_id, parent.session_id) + + def test_compaction_context_survives_anchor_caps_and_restart(self): + """Anchor caps retain both durable child bindings deterministically.""" + + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "capped_context_stickiness.json" + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + max_anchors_per_session=16, + ) + parent_request = { + "messages": [ + {"role": "user", "content": self._long_text("cap-parent-user")}, + { + "role": "assistant", + "content": self._long_text("cap-parent-assistant"), + }, + {"role": "user", "content": self._long_text("cap-parent-follow-up")}, + ] + } + tracker.infer_session(parent_request, provider="gemini", model="pro") + compacted = { + "messages": [ + { + "role": "system", + "content": "Summary of previous conversation: " + + parent_request["messages"][1]["content"], + }, + {"role": "user", "content": "Continue capped context."}, + ] + } + child = tracker.infer_session(compacted, provider="gemini", model="pro") + child_response = self._long_text("cap-child-response") + tracker.record_response( + child.session_id, + provider="gemini", + model="pro", + tracking_namespace=child.tracking_namespace, + response=self._response(child_response), + ) + pressure_messages = [ + *compacted["messages"], + {"role": "assistant", "content": child_response}, + ] + for index in range(8): + pressure_messages.append( + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"cap-pressure-{index}"), + } + ) + pressured = tracker.infer_session( + {"messages": pressure_messages}, + provider="gemini", + model="pro", + ) + self.assertEqual(pressured.session_id, child.session_id) + child_sources = { + tracker._anchors[value].source + for value in tracker._sessions[child.session_id].anchors + } + self.assertIn("compaction_context", child_sources) + self.assertIn("compaction_replay", child_sources) + self.assertLessEqual(len(tracker._sessions[child.session_id].anchors), 16) + tracker.flush() + + restored = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + max_anchors_per_session=16, + ) + replay = restored.infer_session(compacted, provider="gemini", model="pro") + self.assertEqual(replay.session_id, child.session_id) + self.assertTrue(replay.possible_compaction) + changed_tail = { + "messages": [ + compacted["messages"][0], + {"role": "user", "content": "Continue capped context differently."}, + ] + } + continued = restored.infer_session( + changed_tail, + provider="gemini", + model="pro", + ) + + self.assertEqual(continued.session_id, child.session_id) + self.assertFalse(continued.possible_compaction) + + def test_compaction_context_cannot_cross_usage_scope(self): + """A validated child binding remains isolated to its usage namespace.""" + + tracker = SessionTracker(ttl_seconds=3600) + parent_request = { + "messages": [ + {"role": "user", "content": self._long_text("context-scope-user")}, + { + "role": "assistant", + "content": self._long_text("context-scope-assistant"), + }, + {"role": "user", "content": self._long_text("context-scope-follow-up")}, + ] + } + parent = tracker.infer_session( + parent_request, + provider="gemini", + model="pro", + scope_key="A", + ) + compacted = { + "messages": [ + { + "role": "system", + "content": "Compacted context: " + + parent_request["messages"][1]["content"], + }, + {"role": "user", "content": "Continue scoped context."}, + ] + } + child_a = tracker.infer_session( + compacted, + provider="gemini", + model="pro", + scope_key="A", + ) + same_payload_b = tracker.infer_session( + compacted, + provider="gemini", + model="pro", + scope_key="B", + ) + + self.assertEqual(child_a.lineage_parent_session_id, parent.session_id) + self.assertNotEqual(same_payload_b.session_id, child_a.session_id) + self.assertFalse(same_payload_b.possible_compaction) + self.assertIsNone(same_payload_b.lineage_parent_session_id) + + def test_trusted_and_provider_identity_override_compaction_bindings(self): + """Authoritative identity must win over an older replay/context anchor.""" + + tracker = SessionTracker( + ttl_seconds=3600, + trusted_explicit_fields=["conversation_id"], + ) + parent_request = { + "messages": [ + {"role": "user", "content": self._long_text("identity-parent-user")}, + { + "role": "assistant", + "content": self._long_text("identity-parent-assistant"), + }, + {"role": "user", "content": self._long_text("identity-parent-follow-up")}, + { + "role": "assistant", + "content": self._long_text("identity-parent-later"), + }, + ] + } + tracker.infer_session(parent_request, provider="gemini", model="pro") + compacted = { + "messages": [ + { + "role": "system", + "content": "Summary of previous conversation: " + + parent_request["messages"][1]["content"], + }, + {"role": "user", "content": "Continue from the compacted state."}, + ] + } + compacted_child = tracker.infer_session(compacted, provider="gemini", model="pro") + + explicit_session = tracker.infer_session( + { + "conversation_id": "trusted-conversation-B", + "messages": [ + {"role": "user", "content": self._long_text("identity-explicit-session")} + ], + }, + provider="gemini", + model="pro", + ) + explicit_replay = tracker.infer_session( + {**compacted, "conversation_id": "trusted-conversation-B"}, + provider="gemini", + model="pro", + ) + self.assertEqual(explicit_replay.session_id, explicit_session.session_id) + self.assertNotEqual(explicit_replay.session_id, compacted_child.session_id) + self.assertFalse(explicit_replay.possible_compaction) + self.assertIsNone(explicit_replay.lineage_parent_session_id) + + unseen_explicit = tracker.infer_session( + {**compacted, "conversation_id": "trusted-conversation-C"}, + provider="gemini", + model="pro", + ) + self.assertNotEqual(unseen_explicit.session_id, compacted_child.session_id) + self.assertNotEqual(unseen_explicit.session_id, explicit_session.session_id) + + provider_hints = SessionTrackingHints(strong_anchors=["provider-native-session-B"]) + provider_session = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": self._long_text("identity-provider-session")} + ] + }, + provider="gemini", + model="pro", + hints=provider_hints, + ) + provider_replay = tracker.infer_session( + compacted, + provider="gemini", + model="pro", + hints=provider_hints, + ) + self.assertEqual(provider_replay.session_id, provider_session.session_id) + self.assertNotEqual(provider_replay.session_id, compacted_child.session_id) + self.assertFalse(provider_replay.possible_compaction) + self.assertIsNone(provider_replay.lineage_parent_session_id) + + opaque_tool_id = "call_92d1e7b45ac843f6" + tool_session = tracker.infer_session( + { + "messages": [ + { + "role": "assistant", + "tool_calls": [{"id": opaque_tool_id, "type": "function"}], + } + ] + }, + provider="gemini", + model="pro", + ) + tool_bound_context = tracker.infer_session( + { + "messages": [ + *compacted["messages"], + { + "role": "tool", + "tool_call_id": opaque_tool_id, + "content": "Opaque tool continuation evidence.", + }, + ] + }, + provider="gemini", + model="pro", + ) + self.assertNotEqual(tool_bound_context.session_id, tool_session.session_id) + self.assertEqual(tool_bound_context.session_id, compacted_child.session_id) + + def test_record_response_normalizes_cross_namespace_fallback_across_restart(self): + """A fallback response stays with its original logical session scope.""" + + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "namespace_session_stickiness.json" + tracker = self._persistent_tracker(path) + session_a = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": self._long_text("namespace-a-request")} + ] + }, + provider="gemini", + model="pro", + scope_key="A", + ) + namespace_b = tracker._namespace("gemini", scope_key="B") + response_text = self._long_text("namespace-mismatched-response") + + with self.assertLogs("rotator_library", level="WARNING") as captured: + tracker.record_response( + session_a.session_id, + provider="gemini", + model="pro", + tracking_namespace=namespace_b, + response=self._response(response_text), + ) + + self.assertIn("normalized response namespace mismatch", captured.output[0]) + self.assertEqual( + tracker._sessions[session_a.session_id].namespace, + session_a.tracking_namespace, + ) + tracker.flush() + tracker = self._persistent_tracker(path) + self.assertEqual( + tracker._sessions[session_a.session_id].namespace, + session_a.tracking_namespace, + ) + + scope_a_continuation = tracker.infer_session( + { + "messages": [ + {"role": "assistant", "content": response_text}, + {"role": "user", "content": "Continue in logical scope A."}, + ] + }, + provider="gemini", + model="pro", + scope_key="A", + ) + self.assertEqual(scope_a_continuation.session_id, session_a.session_id) + + scope_b_request = tracker.infer_session( + { + "messages": [ + {"role": "assistant", "content": response_text}, + {"role": "user", "content": "Continue in scope B."}, + ] + }, + provider="gemini", + model="pro", + scope_key="B", + ) + self.assertNotEqual(scope_b_request.session_id, session_a.session_id) + + def test_raw_tool_ids_do_not_bind_independent_sessions(self): + """Neither short nor entropy-looking request tool IDs are authoritative.""" + + tracker = SessionTracker(ttl_seconds=3600) + generic_owner = tracker.infer_session( + { + "messages": [ + { + "role": "assistant", + "tool_calls": [{"id": "call_1", "type": "function"}], + } + ] + }, + provider="gemini", + model="pro", + ) + unrelated_generic = tracker.infer_session( + { + "messages": [ + { + "role": "tool", + "tool_call_id": "call_1", + "content": "An unrelated tool result from another conversation.", + } + ] + }, + provider="gemini", + model="pro", + ) + self.assertNotEqual(unrelated_generic.session_id, generic_owner.session_id) + self.assertEqual(unrelated_generic.confidence, "weak") + + opaque_id = "call_7f4e2a91c8b64d37" + opaque_owner = tracker.infer_session( + { + "messages": [ + { + "role": "assistant", + "tool_calls": [{"id": opaque_id, "type": "function"}], + } + ] + }, + provider="gemini", + model="pro", + ) + opaque_result = tracker.infer_session( + { + "messages": [ + { + "role": "tool", + "tool_call_id": opaque_id, + "content": "The matching opaque tool result.", + } + ] + }, + provider="gemini", + model="pro", + ) + self.assertNotEqual(opaque_result.session_id, opaque_owner.session_id) + self.assertEqual(opaque_result.confidence, "weak") + + compound_tracker = SessionTracker(ttl_seconds=3600) + compound_owner = compound_tracker.infer_session( + { + "messages": [ + { + "role": "assistant", + "tool_calls": [ + {"id": "call_1", "type": "function"}, + {"id": "call_2", "type": "function"}, + ], + } + ] + }, + provider="gemini", + model="pro", + ) + compound_unrelated = compound_tracker.infer_session( + { + "messages": [ + {"role": "tool", "tool_call_id": "call_1", "content": "first"}, + {"role": "tool", "tool_call_id": "call_2", "content": "second"}, + ] + }, + provider="gemini", + model="pro", + ) + self.assertNotEqual(compound_unrelated.session_id, compound_owner.session_id) + self.assertEqual(compound_unrelated.confidence, "weak") + + def test_closed_tool_event_plus_message_group_continues_across_provider(self): + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("closed-tool-event-user"), + } + call = self._tool_call_message( + "call_read_1", + "read_file", + '{"path":"src/session.py","line":10}', + ) + root = tracker.infer_session( + {"messages": [user_message]}, + provider="deepseek", + model="chat", + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": call}]}, + ) + pending_events = [ + record + for record in tracker._anchors.values() + if record.session_id == root.session_id + and record.source == "tool_event" + ] + self.assertEqual(len(pending_events), 1) + self.assertEqual(pending_events[0].strength, "weak") + + continued = tracker.infer_session( + {"messages": [user_message, call, self._tool_result("call_read_1")]}, + provider="anthropic", + model="claude", + ) + + self.assertEqual(continued.session_id, root.session_id) + self.assertEqual(continued.confidence, "probable") + + def test_two_distinct_closed_tool_events_continue_sparse_agentic_turn(self): + tracker = SessionTracker(ttl_seconds=3600) + first_call = self._tool_call_message("call_1", "read_file", '{"path":"a.py"}') + second_call = self._tool_call_message("call_2", "grep", '{"pattern":"needle"}') + response_message = { + "role": "assistant", + "content": None, + "tool_calls": [*first_call["tool_calls"], *second_call["tool_calls"]], + } + root = tracker.infer_session( + {"messages": [{"role": "user", "content": "continue"}]}, + provider="openai", + model="gpt", + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": response_message}]}, + ) + + continued = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": "continue"}, + response_message, + self._tool_result("call_1", "a"), + self._tool_result("call_2", "b"), + ] + }, + provider="deepseek", + model="chat", + ) + + self.assertEqual(continued.session_id, root.session_id) + self.assertEqual(continued.confidence, "probable") + + def test_one_closed_tool_event_without_other_medium_group_remains_weak(self): + tracker = SessionTracker(ttl_seconds=3600) + call = self._tool_call_message("call_1", "read_file", '{"path":"a.py"}') + root = tracker.infer_session( + {"messages": [{"role": "user", "content": "continue"}]}, + provider="openai", + model="gpt", + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": call}]}, + ) + + ambiguous = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": "continue"}, + call, + self._tool_result("call_1"), + ] + }, + provider="openai", + model="gpt", + ) + + self.assertNotEqual(ambiguous.session_id, root.session_id) + self.assertEqual(ambiguous.confidence, "weak") + + def test_tool_event_arguments_are_canonical_and_result_edits_do_not_fragment(self): + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("canonical-tool-event-user"), + } + emitted = self._tool_call_message( + "call_1", + "read_file", + '{"path":"a.py","line":10}', + ) + reordered = self._tool_call_message( + "call_1", + "read_file", + '{ "line": 10, "path": "a.py" }', + ) + root = tracker.infer_session( + {"messages": [user_message]}, provider="openai", model="gpt" + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": emitted}]}, + ) + first = tracker.infer_session( + { + "messages": [ + user_message, + reordered, + self._tool_result("call_1", "original result"), + ] + }, + provider="anthropic", + model="claude", + ) + edited = tracker.infer_session( + { + "messages": [ + user_message, + reordered, + self._tool_result("call_1", "edited result"), + ] + }, + provider="deepseek", + model="chat", + ) + + self.assertEqual(first.session_id, root.session_id) + self.assertEqual(edited.session_id, root.session_id) + + def test_same_raw_id_with_different_call_structure_does_not_bind(self): + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("tool-event-collision-user"), + } + owner_call = self._tool_call_message( + "call_1", "read_file", '{"path":"owner.py"}' + ) + owner = tracker.infer_session( + {"messages": [user_message]}, provider="openai", model="gpt" + ) + tracker.record_response( + owner.session_id, + tracking_namespace=owner.tracking_namespace, + response={"choices": [{"message": owner_call}]}, + ) + + unrelated = tracker.infer_session( + { + "messages": [ + user_message, + self._tool_call_message( + "call_1", "read_file", '{"path":"other.py"}' + ), + self._tool_result("call_1"), + ] + }, + provider="openai", + model="gpt", + ) + + self.assertNotEqual(unrelated.session_id, owner.session_id) + self.assertEqual(unrelated.confidence, "weak") + + def test_duplicate_tool_results_count_as_one_closed_event(self): + tracker = SessionTracker(ttl_seconds=3600) + call = self._tool_call_message("call_1", "read_file", '{"path":"a.py"}') + root = tracker.infer_session( + {"messages": [{"role": "user", "content": "continue"}]}, + provider="openai", + model="gpt", + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": call}]}, + ) + + duplicated = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": "continue"}, + call, + self._tool_result("call_1", "first"), + self._tool_result("call_1", "duplicate"), + ] + }, + provider="openai", + model="gpt", + ) + + self.assertNotEqual(duplicated.session_id, root.session_id) + self.assertEqual(duplicated.confidence, "weak") + + def test_multiple_unpaired_structural_tool_calls_remain_weak(self): + tracker = SessionTracker(ttl_seconds=3600) + calls = { + "role": "assistant", + "content": None, + "tool_calls": [ + self._tool_call_message( + "call_1", "read_file", '{"path":"a.py"}' + )["tool_calls"][0], + self._tool_call_message( + "call_2", "grep", '{"pattern":"needle"}' + )["tool_calls"][0], + ], + } + + first = tracker.infer_session( + {"messages": [calls]}, provider="openai", model="gpt" + ) + second = tracker.infer_session( + {"messages": [calls]}, provider="openai", model="gpt" + ) + + self.assertNotEqual(second.session_id, first.session_id) + self.assertEqual(second.confidence, "weak") + + def test_previously_closed_event_does_not_upgrade_later_unpaired_call(self): + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("request-local-tool-closure"), + } + call = self._tool_call_message("call_1", "read_file", '{"path":"a.py"}') + root = tracker.infer_session( + {"messages": [user_message]}, provider="openai", model="gpt" + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": call}]}, + ) + closed = tracker.infer_session( + { + "messages": [ + user_message, + call, + self._tool_result("call_1"), + ] + }, + provider="openai", + model="gpt", + ) + self.assertEqual(closed.session_id, root.session_id) + + unpaired = tracker.infer_session( + {"messages": [user_message, call]}, + provider="anthropic", + model="claude", + ) + + self.assertNotEqual(unpaired.session_id, root.session_id) + self.assertEqual(unpaired.confidence, "weak") + + def test_one_result_closes_only_one_of_multiple_same_id_calls(self): + tracker = SessionTracker(ttl_seconds=3600) + calls = { + "role": "assistant", + "content": None, + "tool_calls": [ + self._tool_call_message( + "call_1", "read_file", '{"path":"a.py"}' + )["tool_calls"][0], + self._tool_call_message( + "call_1", "read_file", '{"path":"b.py"}' + )["tool_calls"][0], + ], + } + root = tracker.infer_session( + {"messages": [{"role": "user", "content": "continue"}]}, + provider="openai", + model="gpt", + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": calls}]}, + ) + + ambiguous = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": "continue"}, + calls, + self._tool_result("call_1"), + ] + }, + provider="openai", + model="gpt", + ) + + self.assertNotEqual(ambiguous.session_id, root.session_id) + self.assertEqual(ambiguous.confidence, "weak") + + def test_non_json_and_mapping_tool_arguments_are_canonical(self): + cases = ( + ("path = a.py", "path = a.py"), + ({"path": "a.py", "line": 10}, {"line": 10, "path": "a.py"}), + ) + for emitted_arguments, returned_arguments in cases: + with self.subTest(arguments=emitted_arguments): + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("non-json-tool-arguments"), + } + emitted = self._tool_call_message( + "call_1", "read_file", emitted_arguments + ) + returned = self._tool_call_message( + "call_1", "read_file", returned_arguments + ) + root = tracker.infer_session( + {"messages": [user_message]}, + provider="openai", + model="gpt", + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": emitted}]}, + ) + + continued = tracker.infer_session( + { + "messages": [ + user_message, + returned, + self._tool_result("call_1"), + ] + }, + provider="openai", + model="gpt", + ) + + self.assertEqual(continued.session_id, root.session_id) + + def test_function_name_is_part_of_tool_event_identity(self): + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("tool-function-name-collision"), + } + emitted = self._tool_call_message("call_1", "read_file", '{"path":"a.py"}') + root = tracker.infer_session( + {"messages": [user_message]}, provider="openai", model="gpt" + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": emitted}]}, + ) + + unrelated = tracker.infer_session( + { + "messages": [ + user_message, + self._tool_call_message( + "call_1", "write_file", '{"path":"a.py"}' + ), + self._tool_result("call_1"), + ] + }, + provider="openai", + model="gpt", + ) + + self.assertNotEqual(unrelated.session_id, root.session_id) + self.assertEqual(unrelated.confidence, "weak") + + def test_legacy_function_result_closes_matching_tool_event(self): + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("legacy-function-tool-event"), + } + call = self._tool_call_message("call_1", "lookup", '{"query":"state"}') + root = tracker.infer_session( + {"messages": [user_message]}, provider="openai", model="gpt" + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": call}]}, + ) + + continued = tracker.infer_session( + { + "messages": [ + user_message, + call, + { + "role": "function", + "tool_call_id": "call_1", + "content": "ok", + }, + ] + }, + provider="openai", + model="gpt", + ) + + self.assertEqual(continued.session_id, root.session_id) + + def test_nameless_tool_call_never_creates_structural_event(self): + tracker = SessionTracker(ttl_seconds=3600) + anchors = tracker._build_anchors( + { + "messages": [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "function": {"arguments": '{"path":"a.py"}'}, + } + ], + }, + self._tool_result("call_1"), + ] + }, + "session-domain:public", + None, + ) + + self.assertFalse(any(anchor.source == "tool_event" for anchor in anchors)) + + def test_closed_tool_events_support_continuity_but_not_compaction_probes(self): + tracker = SessionTracker(ttl_seconds=3600) + calls = { + "role": "assistant", + "content": None, + "tool_calls": [ + self._tool_call_message( + "call_1", "read_file", '{"path":"a.py"}' + )["tool_calls"][0], + self._tool_call_message( + "call_2", "grep", '{"pattern":"needle"}' + )["tool_calls"][0], + ], + } + history = [ + calls, + self._tool_result("call_1", "a"), + self._tool_result("call_2", "b"), + ] + root = tracker.infer_session( + {"messages": history}, provider="openai", model="gpt" + ) + + replay = tracker.infer_session( + {"messages": history}, provider="anthropic", model="claude" + ) + + self.assertEqual(replay.session_id, root.session_id) + self.assertEqual(replay.confidence, "probable") + self.assertFalse(replay.possible_compaction) + + def test_closed_tool_event_survives_persistence_and_never_crosses_domain(self): + user_message = { + "role": "user", + "content": self._long_text("persisted-tool-event-user"), + } + call = self._tool_call_message( + "call_persisted", "lookup", '{"query":"status"}' + ) + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "tool_event_session.json" + tracker = self._persistent_tracker(path) + root = tracker.infer_session( + {"messages": [user_message]}, + provider="openai", + model="gpt", + scope_key="domain-a", + ) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response={"choices": [{"message": call}]}, + ) + tracker.flush() + tracker = self._persistent_tracker(path) + request = { + "messages": [ + user_message, + call, + self._tool_result("call_persisted", "changed safely"), + ] + } + + continued = tracker.infer_session( + request, + provider="anthropic", + model="claude", + scope_key="domain-a", + ) + isolated = tracker.infer_session( + request, + provider="anthropic", + model="claude", + scope_key="domain-b", + ) + + self.assertEqual(continued.session_id, root.session_id) + self.assertNotEqual(isolated.session_id, root.session_id) + self.assertFalse(continued.possible_compaction) + + def test_weak_first_user_only_does_not_reuse_session(self): + tracker = SessionTracker(ttl_seconds=3600) + request = {"messages": [{"role": "user", "content": "hello world"}]} + + session_a = tracker.infer_session(request, provider="gemini", model="pro") + session_b = tracker.infer_session(request, provider="gemini", model="pro") + + self.assertIsNotNone(session_a.session_id) + self.assertIsNotNone(session_b.session_id) + self.assertNotEqual(session_a.session_id, session_b.session_id) + self.assertIsNone(session_a.affinity_key) + + def test_single_long_prompt_does_not_become_sticky_by_itself(self): + tracker = SessionTracker(ttl_seconds=3600) + request = { + "messages": [ + { + "role": "user", + "content": "Please perform a careful detailed analysis of this standalone prompt and provide a structured answer with several paragraphs.", + } + ] + } + + session_a = tracker.infer_session(request, provider="gemini", model="pro") + session_b = tracker.infer_session(request, provider="gemini", model="pro") + + self.assertNotEqual(session_a.session_id, session_b.session_id) + self.assertIsNone(session_a.affinity_key) + + def test_shared_system_and_user_prompt_does_not_become_sticky_by_itself(self): + tracker = SessionTracker(ttl_seconds=3600) + request = { + "messages": [ + { + "role": "system", + "content": "You are a shared coding harness with stable instructions used by many independent sessions.", + }, + { + "role": "user", + "content": "Please review this standalone request carefully and provide the exact structured output.", + }, + ] + } + + session_a = tracker.infer_session(request, provider="gemini", model="pro") + session_b = tracker.infer_session(request, provider="gemini", model="pro") + + self.assertNotEqual(session_a.session_id, session_b.session_id) + self.assertIsNone(session_a.affinity_key) + + def test_explicit_ids_are_weak_unless_configured_as_trusted(self): + request = {"conversation_id": "stable-client-id"} + + conservative = SessionTracker(ttl_seconds=3600) + first = conservative.infer_session(request, provider="gemini", model="pro") + second = conservative.infer_session(request, provider="gemini", model="pro") + self.assertNotEqual(first.session_id, second.session_id) + + trusted = SessionTracker( + ttl_seconds=3600, + trusted_explicit_fields=["conversation_id"], + ) + first = trusted.infer_session(request, provider="gemini", model="pro") + second = trusted.infer_session(request, provider="gemini", model="pro") + self.assertEqual(first.session_id, second.session_id) + + def test_trusted_explicit_fields_can_come_from_env(self): + request = {"conversation_id": "stable-client-id"} + + with patch.dict(os.environ, {"TRUSTED_SESSION_ID_FIELDS": "conversation_id"}): + tracker = SessionTracker(ttl_seconds=3600) + + first = tracker.infer_session(request, provider="gemini", model="pro") + second = tracker.infer_session(request, provider="gemini", model="pro") + + self.assertEqual(first.session_id, second.session_id) + + def test_multiple_message_anchors_reuse_session_when_tools_are_pruned(self): + tracker = SessionTracker(ttl_seconds=3600) + + original = { + "messages": [ + {"role": "user", "content": "Please inspect the account quota and explain the cooldown status in detail."}, + { + "role": "assistant", + "content": "I will inspect the quota files and compare the cooldown windows.", + "tool_calls": [ + {"id": "call_123", "type": "function", "function": {"name": "read"}} + ], + }, + {"role": "tool", "tool_call_id": "call_123", "content": "quota payload"}, + {"role": "user", "content": "Now summarize the result and note what changed since yesterday."}, + ] + } + pruned_tools = { + "messages": [ + {"role": "user", "content": "Please inspect the account quota and explain the cooldown status in detail."}, + { + "role": "assistant", + "content": "I will inspect the quota files and compare the cooldown windows.", + }, + {"role": "user", "content": "Now summarize the result and note what changed since yesterday."}, + ] + } + + session_a = tracker.infer_session(original, provider="gemini", model="pro") + session_b = tracker.infer_session(pruned_tools, provider="gemini", model="pro") + + self.assertEqual(session_a.session_id, session_b.session_id) + self.assertIn(session_b.confidence, {"strong", "probable"}) + + def test_response_anchors_bridge_next_request(self): + tracker = SessionTracker(ttl_seconds=3600) + request = { + "messages": [ + {"role": "user", "content": "Investigate why the nightly quota reset failed for the Gemini account."} + ] + } + inferred = tracker.infer_session(request, provider="gemini", model="pro") + tracker.record_response( + inferred.session_id, + provider="gemini", + model="pro", + response={ + "choices": [ + { + "message": { + "role": "assistant", + "content": "The nightly quota reset failed because the reset timestamp was parsed in local time.", + } + } + ] + }, + ) + + next_request = { + "messages": [ + { + "role": "assistant", + "content": "The nightly quota reset failed because the reset timestamp was parsed in local time.", + }, + {"role": "user", "content": "Patch the timestamp parser now."}, + ] + } + continued = tracker.infer_session(next_request, provider="gemini", model="pro") + + self.assertEqual(inferred.session_id, continued.session_id) + + def test_response_id_anchor_bridges_responses_previous_response_id(self): + tracker = SessionTracker(ttl_seconds=3600) + inferred = tracker.infer_session( + {"messages": [{"role": "user", "content": "Start a Responses API conversation."}]}, + provider="openai", + model="gpt-test", + ) + tracker.record_response( + inferred.session_id, + provider="openai", + model="gpt-test", + response={"id": "resp_parent", "object": "response"}, + ) + + continued = tracker.infer_session( + {"messages": [{"role": "user", "content": "Continue."}]}, + provider="anthropic", + model="claude-test", + hints={"global_strong_anchors": ["responses_previous_response_id:resp_parent"]}, + ) + + self.assertEqual(inferred.session_id, continued.session_id) + self.assertEqual(continued.confidence, "strong") + + def test_one_response_user_copy_does_not_create_compaction_lineage(self): + tracker = SessionTracker(ttl_seconds=3600) + request = { + "messages": [ + {"role": "user", "content": "Prepare a detailed durable summary for later compaction."} + ] + } + parent = tracker.infer_session(request, provider="gemini", model="pro") + summary = " ".join( + [ + "The compaction summary says the routing investigation found that response anchors", + "should identify a parent session without continuing sticky routing for the child context.", + ] + * 8 + ) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response={"choices": [{"message": {"role": "assistant", "content": summary}}]}, + ) + + independent = tracker.infer_session( + {"messages": [{"role": "user", "content": summary}]}, + provider="gemini", + model="pro", + ) + repeated = tracker.infer_session( + {"messages": [{"role": "user", "content": summary}]}, + provider="gemini", + model="pro", + ) + + self.assertFalse(independent.possible_compaction) + self.assertIsNone(independent.lineage_parent_session_id) + self.assertFalse(repeated.possible_compaction) + self.assertNotEqual(parent.session_id, independent.session_id) + + def test_unmarked_summary_requires_two_distinct_response_events(self): + tracker = SessionTracker(ttl_seconds=3600) + request_evidence = " ".join( + f"initialrequest{i:02d}durableevidenceword" for i in range(8) + ) + original = { + "messages": [ + {"role": "system", "content": "Stable harness instructions."}, + {"role": "user", "content": request_evidence}, + {"role": "assistant", "content": self._long_text("initial-assistant")}, + {"role": "user", "content": self._long_text("follow-up-user")}, + ] + } + parent = tracker.infer_session(original, provider="gemini", model="pro") + response_a = " ".join(f"alpha{i:02d}longword" for i in range(8)) + response_b = " ".join(f"bravo{i:02d}longword" for i in range(8)) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_a), + ) + continued = tracker.infer_session( + { + "messages": [ + *original["messages"], + {"role": "assistant", "content": response_a}, + {"role": "user", "content": "Continue the investigation with another result."}, + ] + }, + provider="gemini", + model="pro", + ) + self.assertEqual(parent.session_id, continued.session_id) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_b), + ) + + summary = f"{request_evidence} {response_a} {response_b}" + child_request = {"messages": [{"role": "user", "content": summary}]} + child = tracker.infer_session(child_request, provider="gemini", model="pro") + replay = tracker.infer_session(child_request, provider="gemini", model="pro") + + self.assertTrue(child.possible_compaction) + self.assertEqual(parent.session_id, child.lineage_parent_session_id) + self.assertNotEqual(parent.session_id, child.session_id) + self.assertEqual(child.session_id, replay.session_id) + self.assertEqual(parent.session_id, replay.lineage_parent_session_id) + + def test_unmarked_aggregator_quoting_two_responses_is_not_compaction(self): + """Two quoted outputs without request-side history remain independent.""" + + tracker = SessionTracker(ttl_seconds=3600) + parent = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": self._long_text("aggregator-parent-user")}, + { + "role": "assistant", + "content": self._long_text("aggregator-parent-assistant"), + }, + {"role": "user", "content": self._long_text("aggregator-parent-follow-up")}, + ] + }, + provider="gemini", + model="pro", + ) + response_a = " ".join(f"aggregatealpha{i:02d}longword" for i in range(8)) + response_b = " ".join(f"aggregatebravo{i:02d}longword" for i in range(8)) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_a), + ) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_b), + ) + filler_a = " ".join(f"aggregatefillera{i:02d}longword" for i in range(8)) + filler_b = " ".join(f"aggregatefillerb{i:02d}longword" for i in range(8)) + request = { + "messages": [ + { + "role": "user", + "content": f"{response_a} {filler_a} {response_b} {filler_b}", + } + ] + } + namespace = parent.tracking_namespace + probe_indexes = tracker._compaction_probe_indexes(request) + probe_anchors = tracker._build_compaction_probe_anchors( + request, + namespace, + probe_indexes=probe_indexes, + ) + candidate = tracker._best_match(probe_anchors, namespace, time.time()) + self.assertEqual(len(candidate.response_groups), 2) + self.assertFalse(candidate.request_groups) + + independent = tracker.infer_session(request, provider="gemini", model="pro") + + self.assertFalse(independent.possible_compaction) + self.assertNotEqual(independent.session_id, parent.session_id) + self.assertIsNone(independent.lineage_parent_session_id) + + def test_unmarked_compaction_does_not_bind_shared_system_harness(self): + """Only the probe carrying parent evidence may become a context key.""" + + tracker = SessionTracker(ttl_seconds=3600) + request_evidence = " ".join( + f"sharedrequest{i:02d}durableevidenceword" for i in range(8) + ) + parent = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": request_evidence}, + {"role": "assistant", "content": self._long_text("shared-parent-a")}, + {"role": "user", "content": self._long_text("shared-parent-u")}, + ] + }, + provider="gemini", + model="pro", + ) + response_a = " ".join(f"sharedalpha{i:02d}longword" for i in range(8)) + response_b = " ".join(f"sharedbravo{i:02d}longword" for i in range(8)) + for response_text in (response_a, response_b): + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_text), + ) + shared_harness = self._long_text("shared-static-system-harness", repeats=18) + summary = f"{request_evidence} {response_a} {response_b}" + compacted = { + "messages": [ + {"role": "system", "content": shared_harness}, + {"role": "user", "content": summary}, + ] + } + child = tracker.infer_session(compacted, provider="gemini", model="pro") + self.assertTrue(child.possible_compaction) + + independent = tracker.infer_session( + { + "messages": [ + {"role": "system", "content": shared_harness}, + {"role": "user", "content": "Start a distinct unrelated task."}, + ] + }, + provider="gemini", + model="pro", + ) + + self.assertNotEqual(independent.session_id, child.session_id) + self.assertFalse(independent.possible_compaction) + self.assertIsNone(independent.lineage_parent_session_id) + + def test_unmarked_compaction_does_not_bind_shared_user_harness(self): + """Request overlap qualifies lineage but response overlap keys its child.""" + + tracker = SessionTracker(ttl_seconds=3600) + shared_harness = self._long_text("shared-user-harness", repeats=18) + request_evidence = " ".join( + f"shareduserrequest{i:02d}durableword" for i in range(8) + ) + parent = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": shared_harness}, + {"role": "assistant", "content": self._long_text("shared-user-parent-a")}, + {"role": "user", "content": request_evidence}, + {"role": "assistant", "content": self._long_text("shared-user-parent-b")}, + ] + }, + provider="gemini", + model="pro", + ) + response_a = " ".join(f"useralpha{i:02d}durableword" for i in range(8)) + response_b = " ".join(f"userbravo{i:02d}durableword" for i in range(8)) + for response_text in (response_a, response_b): + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_text), + ) + summary = f"{request_evidence} {response_a} {response_b}" + child = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": shared_harness}, + {"role": "user", "content": summary}, + ] + }, + provider="gemini", + model="pro", + ) + self.assertTrue(child.possible_compaction) + + independent = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": shared_harness}, + {"role": "user", "content": "Start another unrelated task."}, + ] + }, + provider="gemini", + model="pro", + ) + + self.assertNotEqual(independent.session_id, child.session_id) + self.assertFalse(independent.possible_compaction) + + def test_marked_compaction_does_not_bind_retained_ordinary_user_probe(self): + """A marked summary may bind itself, never another retained probe.""" + + tracker = SessionTracker(ttl_seconds=3600) + shared_user = self._long_text("marked-retained-user", repeats=18) + parent_request = { + "messages": [ + {"role": "user", "content": shared_user}, + { + "role": "assistant", + "content": self._long_text("marked-retained-assistant"), + }, + {"role": "user", "content": self._long_text("marked-retained-follow-up")}, + { + "role": "assistant", + "content": self._long_text("marked-retained-later"), + }, + ] + } + parent = tracker.infer_session(parent_request, provider="gemini", model="pro") + child = tracker.infer_session( + { + "messages": [ + { + "role": "system", + "content": "Summary of previous conversation: " + + parent_request["messages"][1]["content"], + }, + {"role": "user", "content": shared_user}, + ] + }, + provider="gemini", + model="pro", + ) + self.assertTrue(child.possible_compaction) + self.assertEqual(child.lineage_parent_session_id, parent.session_id) + + standalone = tracker.infer_session( + {"messages": [{"role": "user", "content": shared_user}]}, + provider="gemini", + model="pro", + ) + + self.assertNotEqual(standalone.session_id, child.session_id) + self.assertFalse(standalone.possible_compaction) + + def test_unmarked_compaction_context_continues_changed_tail_after_restart(self): + """The evidence-bearing unmarked summary remains a stable child base.""" + + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "unmarked_context_stickiness.json" + tracker = self._persistent_tracker(path) + request_evidence = " ".join( + f"unmarkedcontext{i:02d}requestevidence" for i in range(8) + ) + parent = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": request_evidence}, + { + "role": "assistant", + "content": self._long_text("unmarked-context-parent"), + }, + {"role": "user", "content": self._long_text("unmarked-context-next")}, + ] + }, + provider="gemini", + model="pro", + ) + response_a = " ".join(f"contextalpha{i:02d}longword" for i in range(8)) + response_b = " ".join(f"contextbravo{i:02d}longword" for i in range(8)) + for response_text in (response_a, response_b): + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_text), + ) + summary = f"{request_evidence} {response_a} {response_b}" + first = {"messages": [{"role": "user", "content": summary}]} + child = tracker.infer_session(first, provider="gemini", model="pro") + self.assertTrue(child.possible_compaction) + tracker.flush() + tracker = self._persistent_tracker(path) + + changed_tail = { + "messages": [ + {"role": "user", "content": summary}, + {"role": "user", "content": "Take a different next action."}, + ] + } + continued = tracker.infer_session( + changed_tail, + provider="gemini", + model="pro", + ) + + self.assertEqual(continued.session_id, child.session_id) + self.assertEqual(continued.lineage_parent_session_id, parent.session_id) + self.assertFalse(continued.possible_compaction) + + def test_duplicate_response_content_counts_as_one_response_event(self): + tracker = SessionTracker(ttl_seconds=3600) + parent_request = { + "messages": [ + {"role": "user", "content": self._long_text("parent-user")}, + {"role": "assistant", "content": self._long_text("parent-assistant")}, + {"role": "user", "content": self._long_text("parent-follow-up")}, + ] + } + parent = tracker.infer_session(parent_request, provider="gemini", model="pro") + repeated_response = self._long_text("same-response", repeats=16) + for _ in range(2): + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(repeated_response), + ) + + inferred = tracker.infer_session( + {"messages": [{"role": "user", "content": repeated_response}]}, + provider="gemini", + model="pro", + ) + + self.assertFalse(inferred.possible_compaction) + + def test_normal_history_with_recorded_assistant_response_is_not_compaction(self): + tracker = SessionTracker(ttl_seconds=3600) + user_text = self._long_text("ordinary-user", repeats=16) + assistant_text = self._long_text("ordinary-assistant", repeats=16) + parent = tracker.infer_session( + {"messages": [{"role": "user", "content": user_text}]}, + provider="deepseek", + model="chat", + ) + tracker.record_response( + parent.session_id, + provider="deepseek", + model="chat", + response=self._response(assistant_text), + ) + + continued = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": assistant_text}, + {"role": "user", "content": "Continue with the next ordinary turn."}, + ] + }, + provider="deepseek", + model="chat", + ) + + self.assertFalse(continued.possible_compaction) + self.assertEqual(parent.session_id, continued.session_id) + + def test_long_classifier_instruction_quoting_one_response_is_independent(self): + tracker = SessionTracker(ttl_seconds=3600) + parent = tracker.infer_session( + { + "messages": [ + {"role": "system", "content": self._long_text("roleplay-system")}, + {"role": "user", "content": self._long_text("roleplay-user")}, + {"role": "assistant", "content": self._long_text("roleplay-history")}, + {"role": "user", "content": self._long_text("roleplay-next")}, + ] + }, + provider="mistral", + model="large", + ) + response_text = " ".join(f"scene{i:02d}longword" for i in range(8)) + tracker.record_response( + parent.session_id, + provider="mistral", + model="large", + response=self._response(response_text), + ) + preamble = " ".join(f"classifier{i:02d}longword" for i in range(8)) + suffix = " ".join(f"location{i:02d}longword" for i in range(16)) + classifier_prompt = f"{preamble} {response_text} {suffix}" + response_chunks = set(tracker._content_chunk_hashes(response_text)) + classifier_chunks = set(tracker._content_chunk_hashes(classifier_prompt)) + self.assertTrue(response_chunks & classifier_chunks) + + classified = tracker.infer_session( + {"messages": [{"role": "user", "content": classifier_prompt}]}, + provider="mistral", + model="large", + ) + + self.assertFalse(classified.possible_compaction) + self.assertIsNone(classified.lineage_parent_session_id) + self.assertNotEqual(parent.session_id, classified.session_id) + + def test_ordinary_context_contraction_without_summary_stays_on_parent(self): + tracker = SessionTracker(ttl_seconds=3600) + original = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"branch-history-{index}"), + } + for index in range(5) + ] + parent = tracker.infer_session( + {"messages": original}, + provider="deepseek", + model="chat", + ) + contracted = {"messages": [original[0], original[1], original[4]]} + + first = tracker.infer_session(contracted, provider="deepseek", model="chat") + replay = tracker.infer_session(contracted, provider="deepseek", model="chat") + + self.assertFalse(first.possible_compaction) + self.assertEqual(parent.session_id, first.session_id) + self.assertEqual(parent.session_id, replay.session_id) + self.assertEqual(len(tracker._sessions[parent.session_id].history_signatures), 5) + + def test_response_provenance_survives_echo_in_request_history(self): + tracker = SessionTracker(ttl_seconds=3600) + user_text = self._long_text("provenance-user") + response_text = self._long_text("provenance-response") + parent = tracker.infer_session( + {"messages": [{"role": "user", "content": user_text}]}, + provider="gemini", + model="pro", + ) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_text), + ) + response_records = { + value: record.group + for value, record in tracker._anchors.items() + if record.session_id == parent.session_id + and record.source == "response" + and record.group + and record.group.startswith("response_event:") + } + + tracker.infer_session( + { + "messages": [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": response_text}, + {"role": "user", "content": "Continue after preserving response provenance."}, + ] + }, + provider="gemini", + model="pro", + ) + + self.assertTrue(response_records) + for value, group in response_records.items(): + self.assertEqual(tracker._anchors[value].source, "response") + self.assertEqual(tracker._anchors[value].group, group) + + def test_system_or_developer_quote_of_one_response_does_not_bridge_session(self): + for role in ("system", "developer"): + with self.subTest(role=role): + tracker = SessionTracker(ttl_seconds=3600) + parent = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": self._long_text(f"{role}-quote-user")}, + { + "role": "assistant", + "content": self._long_text(f"{role}-quote-history"), + }, + {"role": "user", "content": self._long_text(f"{role}-quote-next")}, + ] + }, + provider="gemini", + model="pro", + ) + response_text = self._long_text(f"{role}-quoted-response", repeats=16) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_text), + ) + + quoted = tracker.infer_session( + { + "messages": [ + {"role": role, "content": response_text}, + {"role": "user", "content": "Start an independent request."}, + ] + }, + provider="gemini", + model="pro", + ) + + self.assertFalse(quoted.possible_compaction) + self.assertNotEqual(parent.session_id, quoted.session_id) + + def test_size_only_long_prompt_overlap_does_not_create_compaction_lineage(self): + tracker = SessionTracker(ttl_seconds=3600) + long_text = " ".join( + [ + "This is a normal long prompt without compaction markers that repeats", + "ordinary request context and should not be treated as a summary descendant.", + ] + * 8 + ) + original = { + "messages": [ + {"role": "user", "content": long_text}, + {"role": "assistant", "content": "A normal second message adds continuity evidence."}, + ] + } + repeated = {"messages": [{"role": "user", "content": long_text}]} + + tracker.infer_session(original, provider="gemini", model="pro") + inferred = tracker.infer_session(repeated, provider="gemini", model="pro") + + self.assertFalse(inferred.possible_compaction) + self.assertIsNone(inferred.lineage_parent_session_id) + + def test_size_only_two_message_history_overlap_does_not_create_compaction_lineage(self): + tracker = SessionTracker(ttl_seconds=3600) + first = " ".join(["First ordinary long request message without summary markers."] * 30) + second = " ".join(["Second ordinary long request message without summary markers."] * 30) + request = { + "messages": [ + {"role": "user", "content": first}, + {"role": "assistant", "content": second}, + ] + } + + tracker.infer_session(request, provider="gemini", model="pro") + inferred = tracker.infer_session(request, provider="gemini", model="pro") + + self.assertFalse(inferred.possible_compaction) + + def test_middle_only_summary_replacement_remains_same_session(self): + tracker = SessionTracker(ttl_seconds=3600) + original_messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"middle-history-{index}"), + } + for index in range(8) + ] + parent = tracker.infer_session( + {"messages": original_messages}, + provider="gemini", + model="pro", + ) + middle_summary = { + "role": "developer", + "content": "Summary of previous conversation: " + original_messages[3]["content"], + } + partially_replaced = { + "messages": [ + *original_messages[:3], + middle_summary, + *original_messages[4:], + ] + } + + continued = tracker.infer_session( + partially_replaced, + provider="gemini", + model="pro", + ) + + self.assertFalse(continued.possible_compaction) + self.assertEqual(parent.session_id, continued.session_id) + self.assertEqual(len(tracker._sessions[parent.session_id].history_signatures), 8) + + def test_early_summary_replacing_a_minority_of_history_remains_same_session(self): + tracker = SessionTracker(ttl_seconds=3600) + original_messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"minority-history-{index}"), + } + for index in range(8) + ] + parent = tracker.infer_session( + {"messages": original_messages}, + provider="gemini", + model="pro", + ) + summary = { + "role": "developer", + "content": "Summary of previous conversation: " + original_messages[1]["content"], + } + partially_replaced = { + "messages": [summary, original_messages[0], *original_messages[3:]] + } + + continued = tracker.infer_session( + partially_replaced, + provider="gemini", + model="pro", + ) + + self.assertFalse(continued.possible_compaction) + self.assertEqual(parent.session_id, continued.session_id) + self.assertEqual(len(tracker._sessions[parent.session_id].history_signatures), 8) + + def test_retaining_exactly_half_of_history_is_not_compaction(self): + tracker = SessionTracker(ttl_seconds=3600) + original_messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"half-history-{index}"), + } + for index in range(8) + ] + parent = tracker.infer_session( + {"messages": original_messages}, + provider="gemini", + model="pro", + ) + summary = { + "role": "system", + "content": "Compacted context: " + original_messages[1]["content"], + } + half_retained = {"messages": [summary, *original_messages[4:]]} + + continued = tracker.infer_session( + half_retained, + provider="gemini", + model="pro", + ) + + self.assertFalse(continued.possible_compaction) + self.assertEqual(parent.session_id, continued.session_id) + + def test_marked_summary_replacing_most_history_creates_child(self): + tracker = SessionTracker(ttl_seconds=3600) + original_messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"most-history-{index}"), + } + for index in range(8) + ] + parent = tracker.infer_session( + {"messages": original_messages}, + provider="gemini", + model="pro", + ) + summary = { + "role": "system", + "content": "Compacted context: " + original_messages[1]["content"], + } + mostly_replaced = {"messages": [summary, *original_messages[-2:]]} + + child = tracker.infer_session( + mostly_replaced, + provider="gemini", + model="pro", + ) + + self.assertTrue(child.possible_compaction) + self.assertEqual(parent.session_id, child.lineage_parent_session_id) + self.assertNotEqual(parent.session_id, child.session_id) + + def test_unmarked_two_response_summary_replacing_only_a_minority_is_not_compaction(self): + tracker = SessionTracker(ttl_seconds=3600) + original_messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"unmarked-partial-{index}"), + } + for index in range(8) + ] + parent = tracker.infer_session( + {"messages": original_messages}, + provider="gemini", + model="pro", + ) + response_a = " ".join(f"partialalpha{i:02d}word" for i in range(8)) + response_b = " ".join(f"partialbravo{i:02d}word" for i in range(8)) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_a), + ) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_b), + ) + filler_a = " ".join(f"partialmiddle{i:02d}word" for i in range(8)) + filler_b = " ".join(f"partialending{i:02d}word" for i in range(8)) + summary = { + "role": "user", + "content": f"{response_a} {filler_a} {response_b} {filler_b}", + } + mostly_retained = {"messages": [summary, *original_messages[3:]]} + + continued = tracker.infer_session( + mostly_retained, + provider="gemini", + model="pro", + ) - self.assertEqual(session_a.session_id, session_b.session_id) - self.assertIn(session_b.confidence, {"strong", "probable"}) + self.assertFalse(continued.possible_compaction) + self.assertEqual(parent.session_id, continued.session_id) - def test_response_anchors_bridge_next_request(self): + def test_unmarked_two_response_summary_replacing_most_history_creates_child(self): tracker = SessionTracker(ttl_seconds=3600) - request = { - "messages": [ - {"role": "user", "content": "Investigate why the nightly quota reset failed for the Gemini account."} - ] - } - inferred = tracker.infer_session(request, provider="gemini", model="pro") + request_evidence = " ".join( + f"mostrequest{i:02d}durableevidenceword" for i in range(8) + ) + original_messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": ( + request_evidence + if index == 0 + else self._long_text(f"unmarked-most-{index}") + ), + } + for index in range(8) + ] + parent = tracker.infer_session( + {"messages": original_messages}, + provider="gemini", + model="pro", + ) + response_a = " ".join(f"mostalpha{i:02d}longword" for i in range(8)) + response_b = " ".join(f"mostbravo{i:02d}longword" for i in range(8)) tracker.record_response( - inferred.session_id, + parent.session_id, provider="gemini", model="pro", - response={ - "choices": [ - { - "message": { - "role": "assistant", - "content": "The nightly quota reset failed because the reset timestamp was parsed in local time.", - } - } - ] - }, + response=self._response(response_a), ) - - next_request = { + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_b), + ) + mostly_replaced = { "messages": [ { - "role": "assistant", - "content": "The nightly quota reset failed because the reset timestamp was parsed in local time.", + "role": "user", + "content": f"{request_evidence} {response_a} {response_b}", }, - {"role": "user", "content": "Patch the timestamp parser now."}, + original_messages[-1], ] } - continued = tracker.infer_session(next_request, provider="gemini", model="pro") - self.assertEqual(inferred.session_id, continued.session_id) + child = tracker.infer_session( + mostly_replaced, + provider="gemini", + model="pro", + ) + + self.assertTrue(child.possible_compaction) + self.assertEqual(parent.session_id, child.lineage_parent_session_id) + self.assertNotEqual(parent.session_id, child.session_id) - def test_compaction_probe_detects_user_summary_from_prior_response(self): + def test_compaction_evidence_cannot_cross_usage_scope(self): tracker = SessionTracker(ttl_seconds=3600) - request = { - "messages": [ - {"role": "user", "content": "Prepare a detailed durable summary for later compaction."} - ] - } - parent = tracker.infer_session(request, provider="gemini", model="pro") - summary = " ".join( - [ - "The compaction summary says the routing investigation found that response anchors", - "should identify a parent session without continuing sticky routing for the child context.", - ] - * 8 + parent = tracker.infer_session( + { + "messages": [ + {"role": "user", "content": self._long_text("scope-parent-user")}, + {"role": "assistant", "content": self._long_text("scope-parent-assistant")}, + {"role": "user", "content": self._long_text("scope-parent-follow-up")}, + ] + }, + provider="gemini", + model="pro", + scope_key="scope-a", ) + response_a = " ".join(f"scopealpha{i:02d}longword" for i in range(8)) + response_b = " ".join(f"scopebravo{i:02d}longword" for i in range(8)) tracker.record_response( parent.session_id, + tracking_namespace=parent.tracking_namespace, + response=self._response(response_a), + ) + tracker.record_response( + parent.session_id, + tracking_namespace=parent.tracking_namespace, + response=self._response(response_b), + ) + filler_a = " ".join(f"scopemiddle{i:02d}longword" for i in range(8)) + filler_b = " ".join(f"scopeending{i:02d}longword" for i in range(8)) + summary = f"{response_a} {filler_a} {response_b} {filler_b}" + + inferred = tracker.infer_session( + {"messages": [{"role": "user", "content": summary}]}, provider="gemini", model="pro", - response={"choices": [{"message": {"role": "assistant", "content": summary}}]}, + scope_key="scope-b", ) - child = tracker.infer_session( - {"messages": [{"role": "user", "content": summary}]}, + self.assertFalse(inferred.possible_compaction) + self.assertNotEqual(parent.session_id, inferred.session_id) + + def test_trusted_identity_keeps_live_session_through_compaction(self): + tracker = SessionTracker( + ttl_seconds=3600, + trusted_explicit_fields=["conversation_id"], + ) + original_messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"trusted-history-{index}"), + } + for index in range(6) + ] + parent = tracker.infer_session( + {"conversation_id": "trusted-session", "messages": original_messages}, provider="gemini", model="pro", ) - repeated_child = tracker.infer_session( - {"messages": [{"role": "user", "content": summary}]}, + compacted = tracker.infer_session( + { + "conversation_id": "trusted-session", + "messages": [ + { + "role": "system", + "content": "Summary of previous conversation: " + + original_messages[1]["content"], + }, + {"role": "user", "content": "Continue from the summary."}, + ], + }, provider="gemini", model="pro", ) - self.assertTrue(child.possible_compaction) - self.assertEqual(parent.session_id, child.lineage_parent_session_id) - self.assertNotEqual(parent.session_id, child.session_id) - self.assertEqual(parent.session_id, repeated_child.lineage_parent_session_id) - self.assertNotEqual(child.session_id, repeated_child.session_id) + self.assertTrue(compacted.possible_compaction) + self.assertEqual(parent.session_id, compacted.session_id) + self.assertIsNone(compacted.lineage_parent_session_id) - def test_size_only_long_prompt_overlap_does_not_create_compaction_lineage(self): + def test_raw_tool_identity_does_not_override_validated_compaction(self): tracker = SessionTracker(ttl_seconds=3600) - long_text = " ".join( - [ - "This is a normal long prompt without compaction markers that repeats", - "ordinary request context and should not be treated as a summary descendant.", - ] - * 8 + original_messages = [ + {"role": "user", "content": self._long_text("tool-parent-user")}, + { + "role": "assistant", + "content": self._long_text("tool-parent-assistant"), + "tool_calls": [ + {"id": "call_7f4e2a91c8b64d37", "type": "function"} + ], + }, + { + "role": "tool", + "tool_call_id": "call_7f4e2a91c8b64d37", + "content": self._long_text("tool-parent-result"), + }, + {"role": "user", "content": self._long_text("tool-parent-follow-up")}, + {"role": "assistant", "content": self._long_text("tool-parent-later")}, + {"role": "user", "content": self._long_text("tool-parent-final")}, + ] + parent = tracker.infer_session( + {"messages": original_messages}, + provider="gemini", + model="pro", + ) + compacted = tracker.infer_session( + { + "messages": [ + { + "role": "system", + "content": "Compacted context: " + original_messages[1]["content"], + }, + { + "role": "tool", + "tool_call_id": "call_7f4e2a91c8b64d37", + "content": "The retained tool result remains authoritative.", + }, + ] + }, + provider="gemini", + model="pro", ) + + self.assertTrue(compacted.possible_compaction) + self.assertNotEqual(parent.session_id, compacted.session_id) + self.assertEqual(compacted.lineage_parent_session_id, parent.session_id) + + def test_changed_post_summary_history_continues_validated_compaction_child(self): + tracker = SessionTracker(ttl_seconds=3600) original = { "messages": [ - {"role": "user", "content": long_text}, - {"role": "assistant", "content": "A normal second message adds continuity evidence."}, + {"role": "user", "content": self._long_text("replay-parent-user")}, + {"role": "assistant", "content": self._long_text("replay-parent-assistant")}, + {"role": "user", "content": self._long_text("replay-parent-follow-up")}, + {"role": "assistant", "content": self._long_text("replay-parent-later")}, ] } - repeated = {"messages": [{"role": "user", "content": long_text}]} + parent = tracker.infer_session(original, provider="gemini", model="pro") + summary = { + "role": "system", + "content": "Summary of previous conversation: " + + original["messages"][1]["content"], + } + first_request = { + "messages": [summary, {"role": "user", "content": "Continue branch one."}] + } + changed_request = { + "messages": [summary, {"role": "user", "content": "Continue branch two."}] + } - tracker.infer_session(original, provider="gemini", model="pro") - inferred = tracker.infer_session(repeated, provider="gemini", model="pro") + first_child = tracker.infer_session( + first_request, + provider="gemini", + model="pro", + ) + with self.assertLogs("rotator_library", level="WARNING") as captured: + changed_child = tracker.infer_session( + changed_request, + provider="gemini", + model="pro", + ) + + self.assertEqual(parent.session_id, first_child.lineage_parent_session_id) + self.assertEqual(parent.session_id, changed_child.lineage_parent_session_id) + self.assertEqual(first_child.session_id, changed_child.session_id) + self.assertFalse(changed_child.possible_compaction) + self.assertEqual(changed_child.confidence, "strong") + self.assertIn("action=compaction_continue", captured.output[0]) + + def test_marker_without_parent_evidence_does_not_create_compaction(self): + tracker = SessionTracker(ttl_seconds=3600) + + inferred = tracker.infer_session( + { + "messages": [ + { + "role": "system", + "content": "Summary of previous conversation: no known parent evidence exists.", + }, + {"role": "user", "content": "Continue."}, + ] + }, + provider="gemini", + model="pro", + ) self.assertFalse(inferred.possible_compaction) self.assertIsNone(inferred.lineage_parent_session_id) - def test_size_only_two_message_history_overlap_does_not_create_compaction_lineage(self): + def test_all_compaction_markers_are_recognized(self): tracker = SessionTracker(ttl_seconds=3600) - first = " ".join(["First ordinary long request message without summary markers."] * 30) - second = " ".join(["Second ordinary long request message without summary markers."] * 30) + markers = ( + "summary of previous conversation", + "summary of the previous conversation", + "summarized conversation", + "compressed context", + "compacted context", + "conversation so far", + "previous conversation", + "context reminder", + ) + + for marker in markers: + with self.subTest(marker=marker): + self.assertTrue(tracker._has_compaction_marker(f"Prefix {marker} suffix")) + + def test_compaction_probe_size_boundaries(self): + tracker = SessionTracker(ttl_seconds=3600) + + self.assertFalse(tracker._is_compaction_probe_text("x" * 399)) + self.assertTrue(tracker._is_compaction_probe_text("x" * 400)) + self.assertFalse(tracker._is_compaction_probe_text(" ".join(["x"] * 79))) + self.assertTrue(tracker._is_compaction_probe_text(" ".join(["x"] * 80))) + + def test_oversized_assistant_history_is_never_a_compaction_probe(self): + tracker = SessionTracker(ttl_seconds=3600) + assistant_text = self._long_text("oversized-assistant", repeats=20) + parent = tracker.infer_session( + { + "messages": [ + {"role": "assistant", "content": assistant_text}, + {"role": "user", "content": "Continue this established conversation normally."}, + ] + }, + provider="gemini", + model="pro", + ) + repeated = tracker.infer_session( + { + "messages": [ + {"role": "assistant", "content": assistant_text}, + {"role": "user", "content": "Continue this established conversation normally."}, + ] + }, + provider="gemini", + model="pro", + ) + + self.assertFalse(repeated.possible_compaction) + self.assertEqual(parent.session_id, repeated.session_id) + + def test_oversized_tool_result_is_never_a_compaction_probe(self): + tracker = SessionTracker(ttl_seconds=3600) + tool_text = self._long_text("oversized-tool", repeats=20) request = { "messages": [ - {"role": "user", "content": first}, - {"role": "assistant", "content": second}, + { + "role": "tool", + "tool_call_id": "call_stable", + "content": tool_text, + }, + {"role": "user", "content": "Continue after the tool result."}, ] } + parent = tracker.infer_session(request, provider="gemini", model="pro") + repeated = tracker.infer_session(request, provider="gemini", model="pro") - tracker.infer_session(request, provider="gemini", model="pro") - inferred = tracker.infer_session(request, provider="gemini", model="pro") - - self.assertFalse(inferred.possible_compaction) + self.assertFalse(repeated.possible_compaction) + self.assertEqual(parent.session_id, repeated.session_id) def test_stale_persistence_job_does_not_overwrite_newer_snapshot(self): writes = [] @@ -271,13 +2994,181 @@ def write(self, payload): tracker._write_save_job(new_job) tracker._write_save_job(old_job) - self.assertEqual(len(writes), 1) - self.assertEqual(tracker._last_persisted_generation, 6) + self.assertEqual(len(writes), 1) + self.assertEqual(tracker._last_persisted_generation, 6) + + def test_failed_persistence_write_stays_dirty_and_recovers(self): + class FlakyWriter: + def __init__(self): + self.outcomes = [False, True] + + def write(self, payload): + return self.outcomes.pop(0) + + with tempfile.TemporaryDirectory() as temp_dir: + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=Path(temp_dir) / "session_stickiness.json", + persistence_flush_interval_seconds=0, + ) + tracker._writer = FlakyWriter() + tracker.infer_session( + {"messages": [{"role": "user", "content": self._long_text("failed-write")}]}, + provider="gemini", + model="pro", + ) + + self.assertTrue(tracker._dirty) + self.assertEqual(tracker._last_persisted_generation, 0) + generation = tracker._dirty_generation + tracker.flush() + + self.assertFalse(tracker._dirty) + self.assertEqual(tracker._last_persisted_generation, generation) + + def test_unexpected_writer_exception_does_not_escape_request(self): + class RaisingWriter: + def write(self, payload): + raise RuntimeError("synthetic writer failure") + + with tempfile.TemporaryDirectory() as temp_dir: + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=Path(temp_dir) / "session_stickiness.json", + persistence_flush_interval_seconds=0, + ) + tracker._writer = RaisingWriter() + + inferred = tracker.infer_session( + {"messages": [{"role": "user", "content": self._long_text("writer-exception")}]}, + provider="gemini", + model="pro", + ) + + self.assertIsNotNone(inferred.session_id) + self.assertTrue(tracker._dirty) + + def test_flush_interval_throttles_automatic_write_but_force_flushes(self): + writes = [] + + class FakeWriter: + def write(self, payload): + writes.append(payload) + return True + + with tempfile.TemporaryDirectory() as temp_dir: + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=Path(temp_dir) / "session_stickiness.json", + persistence_flush_interval_seconds=100, + ) + tracker._writer = FakeWriter() + with patch("rotator_library.session_tracking.time.time", return_value=1000.0): + tracker.infer_session( + {"messages": [{"role": "user", "content": self._long_text("throttle-one")}]}, + provider="gemini", + model="pro", + ) + with patch("rotator_library.session_tracking.time.time", return_value=1050.0): + tracker.infer_session( + {"messages": [{"role": "user", "content": self._long_text("throttle-two")}]}, + provider="gemini", + model="pro", + ) + + self.assertEqual(len(writes), 1) + self.assertTrue(tracker._dirty) + tracker.flush() + + self.assertEqual(len(writes), 2) + self.assertFalse(tracker._dirty) + + def test_disk_write_does_not_hold_session_state_lock(self): + entered_write = threading.Event() + release_write = threading.Event() + + class BlockingWriter: + def write(self, payload): + entered_write.set() + if not release_write.wait(timeout=5): + raise TimeoutError("test did not release blocked writer") + return True + + with tempfile.TemporaryDirectory() as temp_dir: + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=Path(temp_dir) / "session_stickiness.json", + persistence_flush_interval_seconds=0, + ) + tracker._writer = BlockingWriter() + worker = threading.Thread( + target=lambda: tracker.infer_session( + {"messages": [{"role": "user", "content": self._long_text("blocked-write")}]}, + provider="gemini", + model="pro", + ) + ) + worker.start() + self.assertTrue(entered_write.wait(timeout=2)) + + acquired = tracker._lock.acquire(timeout=1) + if acquired: + tracker._lock.release() + release_write.set() + worker.join(timeout=5) + + self.assertTrue(acquired) + self.assertFalse(worker.is_alive()) + + def test_concurrent_infer_record_and_flush_preserve_anchor_ownership(self): + tracker = SessionTracker(ttl_seconds=3600) + + def exercise(index): + user_text = self._long_text(f"concurrent-user-{index}") + response_text = self._long_text(f"concurrent-response-{index}") + inferred = tracker.infer_session( + {"messages": [{"role": "user", "content": user_text}]}, + provider="gemini", + model="pro", + ) + tracker.record_response( + inferred.session_id, + provider="gemini", + model="pro", + response=self._response(response_text), + ) + tracker.infer_session( + { + "messages": [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": response_text}, + {"role": "user", "content": f"Continue unique conversation {index}."}, + ] + }, + provider="gemini", + model="pro", + ) + tracker.flush() + + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(exercise, range(24))) + + for value, record in tracker._anchors.items(): + self.assertIn(record.session_id, tracker._sessions) + self.assertIn(value, tracker._sessions[record.session_id].anchors) + for state in tracker._sessions.values(): + for value in state.anchors: + self.assertIn(value, tracker._anchors) + self.assertEqual(tracker._anchors[value].session_id, state.session_id) def test_best_match_tie_breaks_deterministically(self): tracker = SessionTracker(ttl_seconds=3600) now = 1000.0 - namespace = "scope:test:provider:gemini:model:pro" + namespace = "session-domain:test" tracker._anchors["anchor-a"] = _AnchorRecord( session_id="session-a", namespace=namespace, @@ -308,6 +3199,156 @@ def test_best_match_tie_breaks_deterministically(self): self.assertEqual(match.session_id, "session-b") + def test_best_match_prefers_distinct_response_events_on_equal_score(self): + tracker = SessionTracker(ttl_seconds=3600) + namespace = "session-domain:test" + now = 1000.0 + records = { + "single-a": ("session-z-single", "response_event:one"), + "single-b": ("session-z-single", "response_event:one"), + "diverse-a": ("session-a-diverse", "response_event:one"), + "diverse-b": ("session-a-diverse", "response_event:two"), + } + for value, (session_id, group) in records.items(): + tracker._anchors[value] = _AnchorRecord( + session_id=session_id, + namespace=namespace, + strength="medium", + source="response", + group=group, + expires_at=now + 100, + last_seen=now, + ) + anchors = [ + SessionAnchor( + value, + "medium", + source="compaction_probe", + group="compaction_probe:0", + ) + for value in records + ] + + match = tracker._best_match(anchors, namespace, now) + + self.assertEqual(match.session_id, "session-a-diverse") + self.assertEqual(len(match.response_groups), 2) + + def test_ttl_boundary_prunes_session_and_late_response_cannot_resurrect_it(self): + tracker = SessionTracker(ttl_seconds=10) + with patch("rotator_library.session_tracking.time.time", return_value=1000.0): + inferred = tracker.infer_session( + {"messages": [{"role": "user", "content": self._long_text("ttl-parent")}]}, + provider="gemini", + model="pro", + ) + with patch("rotator_library.session_tracking.time.time", return_value=1010.0): + tracker.record_response( + inferred.session_id, + provider="gemini", + model="pro", + response=self._response(self._long_text("late-response")), + ) + + self.assertNotIn(inferred.session_id, tracker._sessions) + self.assertFalse( + any(record.session_id == inferred.session_id for record in tracker._anchors.values()) + ) + + def test_shared_content_anchor_keeps_first_live_owner(self): + tracker = SessionTracker(ttl_seconds=3600) + shared = self._long_text("shared-content", repeats=16) + first = tracker.infer_session( + {"messages": [{"role": "user", "content": shared}]}, + provider="gemini", + model="pro", + ) + first_owned = { + value + for value, record in tracker._anchors.items() + if record.session_id == first.session_id and ":chunk:" in value + } + second = tracker.infer_session( + {"messages": [{"role": "user", "content": shared}]}, + provider="gemini", + model="pro", + ) + + self.assertNotEqual(first.session_id, second.session_id) + self.assertTrue(first_owned) + for value in first_owned: + self.assertEqual(tracker._anchors[value].session_id, first.session_id) + self.assertNotIn(value, tracker._sessions[second.session_id].anchors) + + def test_session_anchor_trim_preserves_bidirectional_ownership(self): + tracker = SessionTracker( + ttl_seconds=3600, + max_anchors_per_session=16, + ) + messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": self._long_text(f"trim-history-{index}", repeats=16), + } + for index in range(12) + ] + inferred = tracker.infer_session( + {"messages": messages}, + provider="gemini", + model="pro", + ) + state = tracker._sessions[inferred.session_id] + + self.assertLessEqual(len(state.anchors), 16) + for value in state.anchors: + self.assertEqual(tracker._anchors[value].session_id, state.session_id) + for value, record in tracker._anchors.items(): + if record.session_id == state.session_id: + self.assertIn(value, state.anchors) + + def test_global_anchor_trim_preserves_bidirectional_ownership(self): + tracker = SessionTracker( + ttl_seconds=3600, + max_anchor_records=100, + ) + for index in range(24): + unique_text = " ".join( + f"global{index:02d}word{word:02d}" for word in range(24) + ) + tracker.infer_session( + {"messages": [{"role": "user", "content": unique_text}]}, + provider="gemini", + model="pro", + ) + + self.assertLessEqual(len(tracker._anchors), 100) + for value, record in tracker._anchors.items(): + self.assertIn(value, tracker._sessions[record.session_id].anchors) + for state in tracker._sessions.values(): + for value in state.anchors: + self.assertEqual(tracker._anchors[value].session_id, state.session_id) + + def test_affinity_is_deterministic_across_tracker_instances(self): + request = { + "messages": [ + {"role": "user", "content": self._long_text("affinity-user")}, + {"role": "assistant", "content": self._long_text("affinity-assistant")}, + ] + } + first = SessionTracker(ttl_seconds=3600).infer_session( + request, + provider="gemini", + model="pro", + ) + second = SessionTracker(ttl_seconds=3600).infer_session( + request, + provider="gemini", + model="pro", + ) + + self.assertIsNotNone(first.affinity_key) + self.assertEqual(first.affinity_key, second.affinity_key) + def test_record_response_uses_stored_namespace_when_tracking_namespace_omitted(self): tracker = SessionTracker(ttl_seconds=3600) hints = SessionTrackingHints(session_scope="quota-group-pro") @@ -355,12 +3396,12 @@ def test_record_response_uses_stored_namespace_when_tracking_namespace_omitted(s self.assertEqual(inferred.session_id, continued.session_id) - def test_provider_model_scope_is_isolated(self): + def test_provider_and_model_changes_keep_global_logical_session(self): tracker = SessionTracker(ttl_seconds=3600) request = { "messages": [ {"role": "user", "content": "Compare the provider model cache behavior using the same long prompt."}, - {"role": "assistant", "content": "The same text should not cross provider or model scopes."}, + {"role": "assistant", "content": "The same logical session should cross providers and models."}, ] } @@ -368,8 +3409,82 @@ def test_provider_model_scope_is_isolated(self): openai = tracker.infer_session(request, provider="openai", model="pro") gemini_flash = tracker.infer_session(request, provider="gemini", model="flash") - self.assertNotEqual(gemini.session_id, openai.session_id) - self.assertNotEqual(gemini.session_id, gemini_flash.session_id) + self.assertEqual(gemini.session_id, openai.session_id) + self.assertEqual(gemini.session_id, gemini_flash.session_id) + self.assertEqual(gemini.tracking_namespace, "session-domain:public") + + def test_cross_provider_first_turn_reroll_remains_ambiguous(self): + """Removing the sole response leaves only one conversational group.""" + + tracker = SessionTracker(ttl_seconds=3600) + request = { + "messages": [ + {"role": "system", "content": self._long_text("shared-agent-harness")}, + {"role": "developer", "content": self._long_text("shared-tool-rules")}, + {"role": "user", "content": "Write one opening scene for this story."}, + ] + } + deepseek = tracker.infer_session(request, provider="deepseek", model="chat") + response = self._long_text("first-turn-response") + tracker.record_response( + deepseek.session_id, + provider="deepseek", + model="chat", + tracking_namespace=deepseek.tracking_namespace, + response=self._response(response), + ) + + reroll = tracker.infer_session(request, provider="openai", model="gpt") + + self.assertNotEqual(reroll.session_id, deepseek.session_id) + self.assertEqual(reroll.confidence, "weak") + + def test_cross_provider_complete_first_turn_and_follow_up_continue(self): + tracker = SessionTracker(ttl_seconds=3600) + initial = { + "messages": [ + {"role": "system", "content": self._long_text("shared-agent-harness")}, + {"role": "user", "content": self._long_text("global-first-user")}, + ] + } + deepseek = tracker.infer_session(initial, provider="deepseek", model="chat") + response = self._long_text("global-first-response") + tracker.record_response( + deepseek.session_id, + provider="deepseek", + model="chat", + tracking_namespace=deepseek.tracking_namespace, + response=self._response(response), + ) + continued_request = { + "messages": [ + *initial["messages"], + {"role": "assistant", "content": response}, + {"role": "user", "content": self._long_text("global-follow-up")}, + ] + } + + anthropic = tracker.infer_session( + continued_request, + provider="anthropic", + model="claude", + ) + + self.assertEqual(anthropic.session_id, deepseek.session_id) + self.assertEqual(anthropic.confidence, "probable") + + def test_trusted_client_identity_crosses_provider_inside_domain(self): + tracker = SessionTracker( + ttl_seconds=3600, + trusted_explicit_fields=["conversation_id"], + ) + request = {"conversation_id": "client-global-session"} + + deepseek = tracker.infer_session(request, provider="deepseek", model="chat") + openai = tracker.infer_session(request, provider="openai", model="gpt") + + self.assertEqual(openai.session_id, deepseek.session_id) + self.assertEqual(openai.confidence, "strong") def test_allowed_usage_scope_is_isolated(self): tracker = SessionTracker(ttl_seconds=3600) @@ -389,6 +3504,114 @@ def test_allowed_usage_scope_is_isolated(self): self.assertNotEqual(public.session_id, scoped.session_id) + def test_request_isolation_keys_are_stable_secret_free_and_strict(self): + self.assertEqual( + RequestContextBuilder._session_isolation_key(None, None, None, False), + "public", + ) + classifier_a = RequestContextBuilder._session_isolation_key( + "agent-a", None, None, False + ) + classifier_a_again = RequestContextBuilder._session_isolation_key( + "agent-a", None, None, False + ) + classifier_b = RequestContextBuilder._session_isolation_key( + "agent-b", None, None, False + ) + classifier_a_with_rotated_credentials = RequestContextBuilder._session_isolation_key( + "agent-a", + {"openai": ["rotated-secret"]}, + {"openai": {"api_base": "https://rotated.example"}}, + True, + ) + self.assertEqual(classifier_a, classifier_a_again) + self.assertEqual(classifier_a, classifier_a_with_rotated_credentials) + self.assertNotEqual(classifier_a, classifier_b) + + bundle_a = RequestContextBuilder._session_isolation_key( + None, + {"deepseek": ["secret-b", "secret-a"], "openai": "secret-c"}, + {"openai": {"api_base": "https://one.example"}}, + True, + ) + bundle_a_reordered = RequestContextBuilder._session_isolation_key( + None, + {"openai": "secret-c", "deepseek": ["secret-a", "secret-b"]}, + {"openai": {"api_base": "https://one.example"}}, + True, + ) + bundle_a_mixed_case = RequestContextBuilder._session_isolation_key( + None, + {"OpenAI": "secret-c", "DeepSeek": ["secret-a", "secret-b"]}, + {"OpenAI": {"api_base": "https://one.example"}}, + True, + ) + bundle_b = RequestContextBuilder._session_isolation_key( + None, + {"deepseek": ["different-secret"], "openai": "secret-c"}, + {"openai": {"api_base": "https://one.example"}}, + True, + ) + self.assertEqual(bundle_a, bundle_a_reordered) + self.assertEqual(bundle_a, bundle_a_mixed_case) + self.assertNotEqual(bundle_a, bundle_b) + self.assertNotIn("secret", bundle_a) + + def test_global_evidence_never_crosses_isolation_domains(self): + tracker = SessionTracker(ttl_seconds=3600) + request = { + "messages": [ + {"role": "user", "content": self._long_text("strict-domain-user")}, + {"role": "assistant", "content": self._long_text("strict-domain-assistant")}, + ] + } + + domain_a = tracker.infer_session( + request, + provider="deepseek", + model="chat", + scope_key="classifier:a", + ) + domain_b = tracker.infer_session( + request, + provider="openai", + model="gpt", + scope_key="classifier:b", + ) + + self.assertNotEqual(domain_a.session_id, domain_b.session_id) + self.assertNotEqual(domain_a.tracking_namespace, domain_b.tracking_namespace) + + def test_raw_scope_text_cannot_impersonate_derived_bundle_domain(self): + tracker = SessionTracker(ttl_seconds=3600) + derived_domain = "bundle:" + ("a" * 64) + request = { + "messages": [ + {"role": "user", "content": self._long_text("scope-marker-collision-user")}, + { + "role": "assistant", + "content": self._long_text("scope-marker-collision-assistant"), + }, + ] + } + + trusted = tracker.infer_session( + request, + provider="openai", + model="gpt", + scope_key=derived_domain, + _trusted_isolation_key=True, + ) + raw = tracker.infer_session( + request, + provider="openai", + model="gpt", + scope_key=derived_domain, + ) + + self.assertNotEqual(raw.session_id, trusted.session_id) + self.assertNotEqual(raw.tracking_namespace, trusted.tracking_namespace) + def test_provider_strong_hint_reuses_session(self): tracker = SessionTracker(ttl_seconds=3600) hints = SessionTrackingHints(strong_anchors=["native-session-abc"]) @@ -399,96 +3622,508 @@ def test_provider_strong_hint_reuses_session(self): self.assertEqual(first.session_id, second.session_id) self.assertIsNotNone(first.affinity_key) - def test_provider_session_scope_can_override_model_inside_allowed_scope(self): + def test_provider_native_hint_does_not_cross_provider(self): + tracker = SessionTracker(ttl_seconds=3600) + hints = SessionTrackingHints(strong_anchors=["native-session-abc"]) + + gemini = tracker.infer_session({}, provider="gemini", model="pro", hints=hints) + openai = tracker.infer_session({}, provider="openai", model="gpt", hints=hints) + + self.assertNotEqual(gemini.session_id, openai.session_id) + self.assertNotEqual(gemini.affinity_key, openai.affinity_key) + + def test_provider_session_scope_cannot_cross_caller_isolation_domain(self): tracker = SessionTracker(ttl_seconds=3600) hints = SessionTrackingHints( strong_anchors=["native-session-abc"], session_scope="quota-group-pro", ) - pro = tracker.infer_session( - {}, provider="gemini", model="pro", scope_key="gemini", hints=hints - ) - flash = tracker.infer_session( - {}, provider="gemini", model="flash", scope_key="gemini", hints=hints - ) - other_scope = tracker.infer_session( - {}, - provider="gemini", - model="flash", - scope_key="classifier:user-a:gemini", - hints=hints, - ) + pro = tracker.infer_session( + {}, provider="gemini", model="pro", scope_key="gemini", hints=hints + ) + flash = tracker.infer_session( + {}, provider="gemini", model="flash", scope_key="gemini", hints=hints + ) + other_scope = tracker.infer_session( + {}, + provider="gemini", + model="flash", + scope_key="classifier:user-a:gemini", + hints=hints, + ) + + self.assertEqual(pro.session_id, flash.session_id) + self.assertNotEqual(pro.session_id, other_scope.session_id) + + def test_provider_session_scope_partitions_only_native_evidence(self): + tracker = SessionTracker(ttl_seconds=3600) + model_a_hints = SessionTrackingHints( + strong_anchors=["native-session-abc"], + session_scope="model-a", + ) + model_b_hints = SessionTrackingHints( + strong_anchors=["native-session-abc"], + session_scope="model-b", + ) + + model_a = tracker.infer_session( + {}, provider="custom", model="a", hints=model_a_hints + ) + model_b = tracker.infer_session( + {}, provider="custom", model="b", hints=model_b_hints + ) + + self.assertNotEqual(model_a.session_id, model_b.session_id) + self.assertEqual(model_a.tracking_namespace, model_b.tracking_namespace) + self.assertNotEqual(model_a.affinity_key, model_b.affinity_key) + + def test_provider_native_scope_partition_survives_persistence_restart(self): + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "provider_native_scope.json" + tracker = self._persistent_tracker(path) + hints_a = SessionTrackingHints( + strong_anchors=["provider-native-secret-a"], + session_scope="family-a", + ) + hints_b = SessionTrackingHints( + strong_anchors=["provider-native-secret-a"], + session_scope="family-b", + ) + session_a = tracker.infer_session( + {}, provider="custom", model="a", hints=hints_a + ) + tracker.flush() + restored = self._persistent_tracker(path) + + continued_a = restored.infer_session( + {}, provider="custom", model="a2", hints=hints_a + ) + independent_b = restored.infer_session( + {}, provider="custom", model="b", hints=hints_b + ) + + self.assertEqual(continued_a.session_id, session_a.session_id) + self.assertNotEqual(independent_b.session_id, session_a.session_id) + + def test_compaction_lineage_creates_new_session_without_strong_anchor(self): + tracker = SessionTracker(ttl_seconds=3600) + original = { + "messages": [ + {"role": "user", "content": "Please analyze the routing state and remember this detailed anchor text."}, + {"role": "assistant", "content": "I analyzed the routing state and found a useful continuity anchor."}, + ] + } + compacted = { + "messages": [ + { + "role": "system", + "content": "Summary of previous conversation: I analyzed the routing state and found a useful continuity anchor.", + }, + {"role": "user", "content": "Continue from the compressed context."}, + ] + } + + parent = tracker.infer_session(original, provider="gemini", model="pro") + child = tracker.infer_session(compacted, provider="gemini", model="pro") + + self.assertNotEqual(parent.session_id, child.session_id) + self.assertTrue(child.possible_compaction) + self.assertEqual(parent.session_id, child.lineage_parent_session_id) + + def test_compaction_detection_is_conservative_to_early_system_messages(self): + tracker = SessionTracker(ttl_seconds=3600) + user_summary = { + "messages": [ + {"role": "user", "content": "Summary of previous conversation: continue this task."} + ] + } + + inferred = tracker.infer_session(user_summary, provider="gemini", model="pro") + + self.assertFalse(inferred.possible_compaction) + + def test_persistence_round_trips_current_schema_with_anchor_metadata(self): + request = { + "messages": [ + {"role": "user", "content": "Persist this first detailed user anchor for the session tracker."}, + {"role": "assistant", "content": "Persist this second detailed assistant anchor as well."}, + ] + } + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + first = tracker.infer_session(request, provider="gemini", model="pro") + tracker.flush() + persisted = json.loads(path.read_text(encoding="utf-8")) + original_state = tracker._sessions[first.session_id] + original_records = { + value: (record.strength, record.source, record.group) + for value, record in tracker._anchors.items() + if record.session_id == first.session_id + } + + restored = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + second = restored.infer_session(request, provider="gemini", model="pro") + restored_state = restored._sessions[first.session_id] + + self.assertEqual(persisted["schema_version"], 3) + self.assertNotIn("anchors", persisted["sessions"][first.session_id]) + self.assertEqual(first.session_id, second.session_id) + self.assertEqual(first.affinity_key, second.affinity_key) + self.assertEqual(original_state.history_signatures, restored_state.history_signatures) + for value, metadata in original_records.items(): + restored_record = restored._anchors[value] + self.assertEqual( + (restored_record.strength, restored_record.source, restored_record.group), + metadata, + ) + + def test_schema_three_persistence_hashes_all_external_identifiers(self): + secrets = { + "explicit": "client-conversation-secret", + "provider": "provider-native-secret", + "affinity": "provider-affinity-secret", + "tool": "tool-call-secret", + "response": "response-id-secret", + } + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + tracker = SessionTracker( + ttl_seconds=3600, + trusted_explicit_fields=["conversation_id"], + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + inferred = tracker.infer_session( + { + "conversation_id": secrets["explicit"], + "messages": [ + { + "role": "assistant", + "tool_calls": [{"id": secrets["tool"], "type": "function"}], + } + ], + }, + provider="custom", + model="model-a", + hints=SessionTrackingHints( + strong_anchors=[secrets["provider"]], + affinity_key=secrets["affinity"], + ), + ) + tracker.record_response( + inferred.session_id, + tracking_namespace=inferred.tracking_namespace, + response={"id": secrets["response"], "object": "response"}, + ) + tracker.flush() + persisted_text = path.read_text(encoding="utf-8") + + for secret in secrets.values(): + self.assertNotIn(secret, persisted_text) + + def test_schema_three_loader_bounds_sessions_and_rejects_unknown_sources(self): + now = time.time() + namespace = "session-domain:test" + sessions = { + f"session-{index}": { + "namespace": namespace, + "expires_at": now + 3600, + "last_seen": now + index, + "history_signatures": [], + } + for index in range(3) + } + evil_anchor = f"{namespace}:evil:{'a' * 64}" + payload = { + "schema_version": 3, + "sessions": sessions, + "anchors": { + evil_anchor: { + "session_id": "session-2", + "namespace": namespace, + "strength": "strong", + "source": "compaction_probe", + "group": "forged", + "expires_at": now + 3600, + "last_seen": now, + } + }, + } + + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + path.write_text(json.dumps(payload), encoding="utf-8") + with patch.object(SessionTracker, "_MAX_PERSISTED_SESSIONS", 2): + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + ) + + self.assertEqual(set(tracker._sessions), {"session-1", "session-2"}) + self.assertNotIn(evil_anchor, tracker._anchors) - self.assertEqual(pro.session_id, flash.session_id) - self.assertNotEqual(pro.session_id, other_scope.session_id) + def test_schema_three_loader_rejects_oversized_state_before_json_parsing(self): + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + path.write_text("{" + ("x" * 128), encoding="utf-8") + with patch.object(SessionTracker, "_MAX_PERSISTED_FILE_BYTES", 64): + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + ) - def test_compaction_lineage_creates_new_session_without_strong_anchor(self): - tracker = SessionTracker(ttl_seconds=3600) + self.assertFalse(tracker._sessions) + self.assertFalse(tracker._anchors) + + def test_persistence_restart_preserves_compaction_and_replay_binding(self): + request_evidence = " ".join( + f"persistrequest{i:02d}durableevidenceword" for i in range(8) + ) original = { "messages": [ - {"role": "user", "content": "Please analyze the routing state and remember this detailed anchor text."}, - {"role": "assistant", "content": "I analyzed the routing state and found a useful continuity anchor."}, + {"role": "user", "content": request_evidence}, + {"role": "assistant", "content": self._long_text("persist-parent-assistant")}, + {"role": "user", "content": self._long_text("persist-parent-follow-up")}, ] } - compacted = { + response_a = " ".join(f"persistalpha{i:02d}word" for i in range(8)) + response_b = " ".join(f"persistbravo{i:02d}word" for i in range(8)) + child_request = { "messages": [ { - "role": "system", - "content": "Summary of previous conversation: I analyzed the routing state and found a useful continuity anchor.", - }, - {"role": "user", "content": "Continue from the compressed context."}, + "role": "user", + "content": f"{request_evidence} {response_a} {response_b}", + } ] } - parent = tracker.infer_session(original, provider="gemini", model="pro") - child = tracker.infer_session(compacted, provider="gemini", model="pro") + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + parent = tracker.infer_session(original, provider="gemini", model="pro") + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_a), + ) + tracker.record_response( + parent.session_id, + provider="gemini", + model="pro", + response=self._response(response_b), + ) + tracker.flush() + + restored_parent = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + child = restored_parent.infer_session( + child_request, + provider="gemini", + model="pro", + ) + restored_parent.flush() + + restored_child = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + persistence_flush_interval_seconds=0, + ) + replay = restored_child.infer_session( + child_request, + provider="gemini", + model="pro", + ) - self.assertNotEqual(parent.session_id, child.session_id) self.assertTrue(child.possible_compaction) self.assertEqual(parent.session_id, child.lineage_parent_session_id) + self.assertEqual(child.session_id, replay.session_id) + self.assertEqual(parent.session_id, replay.lineage_parent_session_id) - def test_compaction_detection_is_conservative_to_early_system_messages(self): - tracker = SessionTracker(ttl_seconds=3600) - user_summary = { - "messages": [ - {"role": "user", "content": "Summary of previous conversation: continue this task."} - ] - } - - inferred = tracker.infer_session(user_summary, provider="gemini", model="pro") + def test_malformed_persistence_payloads_are_ignored_without_startup_failure(self): + payloads = [ + "{not-json", + json.dumps([]), + json.dumps({"schema_version": 1, "sessions": {}, "anchors": {}}), + json.dumps({"schema_version": 2, "sessions": {}, "anchors": {}}), + json.dumps({"schema_version": 3, "sessions": [], "anchors": {}}), + json.dumps( + { + "schema_version": 3, + "sessions": { + "bad": { + "namespace": "session-domain:x", + "expires_at": "not-a-number", + } + }, + "anchors": {}, + } + ), + ] - self.assertFalse(inferred.possible_compaction) + for payload in payloads: + with self.subTest(payload=payload[:30]), tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + path.write_text(payload, encoding="utf-8") + tracker = SessionTracker( + ttl_seconds=3600, + persist_to_disk=True, + persistence_path=path, + ) + self.assertEqual(tracker._sessions, {}) + self.assertEqual(tracker._anchors, {}) - def test_persistence_round_trips_current_schema_with_anchor_metadata(self): - request = { - "messages": [ - {"role": "user", "content": "Persist this first detailed user anchor for the session tracker."}, - {"role": "assistant", "content": "Persist this second detailed assistant anchor as well."}, - ] + def test_persistence_loader_rebuilds_only_valid_anchor_ownership(self): + namespace = "session-domain:test" + expires_at = time.time() + 3600 + good_value = f"{namespace}:message:user:{'a' * 64}" + payload = { + "schema_version": 3, + "sessions": { + "good-session": { + "namespace": namespace, + "expires_at": expires_at, + "last_seen": expires_at - 10, + "affinity_key": "stable-affinity", + "history_signatures": ["b" * 64, "invalid"], + }, + "expired-session": { + "namespace": namespace, + "expires_at": time.time() - 1, + "last_seen": time.time() - 10, + "history_signatures": ["c" * 64], + }, + }, + "anchors": { + good_value: { + "session_id": "good-session", + "namespace": namespace, + "strength": "medium", + "source": "message", + "group": "message:0:user", + "expires_at": expires_at + 100, + "last_seen": expires_at - 10, + }, + f"{namespace}:orphan": { + "session_id": "missing-session", + "namespace": namespace, + "strength": "medium", + "source": "message", + "group": "message:0:user", + "expires_at": expires_at, + "last_seen": expires_at - 10, + }, + f"{namespace}:wrong-namespace": { + "session_id": "good-session", + "namespace": "session-domain:other", + "strength": "medium", + "source": "message", + "group": "message:0:user", + "expires_at": expires_at, + "last_seen": expires_at - 10, + }, + f"{namespace}:bad-strength": { + "session_id": "good-session", + "namespace": namespace, + "strength": "certain", + "source": "message", + "group": "message:0:user", + "expires_at": expires_at, + "last_seen": expires_at - 10, + }, + }, } + with tempfile.TemporaryDirectory() as temp_dir: path = Path(temp_dir) / "session_stickiness.json" + path.write_text(json.dumps(payload), encoding="utf-8") tracker = SessionTracker( ttl_seconds=3600, persist_to_disk=True, persistence_path=path, - persistence_flush_interval_seconds=0, ) - first = tracker.infer_session(request, provider="gemini", model="pro") - tracker.flush() - persisted = json.loads(path.read_text(encoding="utf-8")) - restored = SessionTracker( + state = tracker._sessions["good-session"] + self.assertNotIn("expired-session", tracker._sessions) + self.assertEqual(state.history_signatures, ("b" * 64,)) + self.assertEqual(state.anchors, {good_value}) + self.assertEqual(set(tracker._anchors), {good_value}) + self.assertEqual(tracker._anchors[good_value].expires_at, state.expires_at) + + def test_persistence_loader_enforces_caps_without_orphaning_session_sets(self): + namespace = "session-domain:test" + expires_at = time.time() + 3600 + sessions = {} + anchors = {} + for session_index in range(12): + session_id = f"session-{session_index}" + sessions[session_id] = { + "namespace": namespace, + "expires_at": expires_at, + "last_seen": expires_at - session_index, + "history_signatures": [f"{session_index:064x}"], + } + for anchor_index in range(24): + value = f"{namespace}:loaded:{session_index}:{anchor_index}" + anchors[value] = { + "session_id": session_id, + "namespace": namespace, + "strength": "medium", + "source": "message", + "group": f"message:{anchor_index}:user", + "expires_at": expires_at, + "last_seen": expires_at - anchor_index, + } + payload = { + "schema_version": 3, + "sessions": sessions, + "anchors": anchors, + } + + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "session_stickiness.json" + path.write_text(json.dumps(payload), encoding="utf-8") + tracker = SessionTracker( ttl_seconds=3600, persist_to_disk=True, persistence_path=path, - persistence_flush_interval_seconds=0, + max_anchors_per_session=16, + max_anchor_records=100, ) - second = restored.infer_session(request, provider="gemini", model="pro") - self.assertEqual(persisted["schema_version"], 1) - self.assertEqual(first.session_id, second.session_id) + self.assertLessEqual(len(tracker._anchors), 100) + for state in tracker._sessions.values(): + self.assertLessEqual(len(state.anchors), 16) + for value in state.anchors: + self.assertIn(value, tracker._anchors) + self.assertEqual(tracker._anchors[value].session_id, state.session_id) + for value, record in tracker._anchors.items(): + self.assertIn(value, tracker._sessions[record.session_id].anchors) def test_unversioned_persistence_is_ignored(self): with tempfile.TemporaryDirectory() as temp_dir: @@ -531,6 +4166,449 @@ def test_streaming_chunk_collector_preserves_response_anchors(self): self.assertEqual(assistant_parts, ["hello "]) self.assertEqual(tool_call_ids, ["call_1"]) + def test_streaming_collector_reassembles_structural_tool_event(self): + handler = StreamingHandler() + assistant_parts = [] + tool_call_ids = [] + tool_call_events = {} + handler._collect_session_response_anchors( + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1",' + '"function":{"name":"read_file","arguments":"{\\"path\\":"}}]}}]}' + "\n\n", + assistant_parts, + tool_call_ids, + tool_call_events, + ) + handler._collect_session_response_anchors( + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,' + '"function":{"arguments":"\\"a.py\\"}"}}]}}]}\n\n', + assistant_parts, + tool_call_ids, + tool_call_events, + ) + + assembled = handler._assembled_tool_calls(tool_call_ids, tool_call_events) + + self.assertEqual( + assembled, + [ + { + "id": "call_1", + "function": { + "name": "read_file", + "arguments": '{"path":"a.py"}', + }, + } + ], + ) + + def test_streaming_tool_events_separate_choices_with_shared_indexes(self): + handler = StreamingHandler() + tool_call_ids = [] + tool_call_events = {} + handler._collect_session_response_anchors( + "data: " + + json.dumps( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_a", + "function": { + "name": "read_file", + "arguments": '{"path":"a.py"}', + }, + } + ] + } + }, + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_b", + "function": { + "name": "read_file", + "arguments": '{"path":"b.py"}', + }, + } + ] + } + }, + ] + } + ) + + "\n\n", + [], + tool_call_ids, + tool_call_events, + ) + + assembled = handler._assembled_tool_calls(tool_call_ids, tool_call_events) + + self.assertEqual([call["id"] for call in assembled], ["call_a", "call_b"]) + self.assertEqual( + [call["function"]["arguments"] for call in assembled], + ['{"path":"a.py"}', '{"path":"b.py"}'], + ) + + def test_streaming_tool_events_use_choice_index_across_separate_frames(self): + handler = StreamingHandler() + tool_call_ids = [] + tool_call_events = {} + for choice_index, call_id, path in ( + (0, "call_a", "a.py"), + (1, "call_b", "b.py"), + ): + handler._collect_session_response_anchors( + "data: " + + json.dumps( + { + "choices": [ + { + "index": choice_index, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": call_id, + "function": { + "name": "read_file", + "arguments": json.dumps({"path": path}), + }, + } + ] + }, + } + ] + } + ) + + "\n\n", + [], + tool_call_ids, + tool_call_events, + ) + + assembled = handler._assembled_tool_calls(tool_call_ids, tool_call_events) + + self.assertEqual([call["id"] for call in assembled], ["call_a", "call_b"]) + self.assertEqual( + [call["function"]["arguments"] for call in assembled], + ['{"path": "a.py"}', '{"path": "b.py"}'], + ) + + def test_streaming_cumulative_argument_snapshots_replace_instead_of_append(self): + handler = StreamingHandler() + tool_call_ids = [] + tool_call_events = {} + for arguments in ('{"path":', '{"path":"a.py"}'): + handler._collect_session_response_anchors( + "data: " + + json.dumps( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "function": { + "name": "read_file", + "arguments": arguments, + }, + } + ] + } + } + ] + } + ) + + "\n\n", + [], + tool_call_ids, + tool_call_events, + ) + + assembled = handler._assembled_tool_calls(tool_call_ids, tool_call_events) + + self.assertEqual( + assembled[0]["function"]["arguments"], + '{"path":"a.py"}', + ) + + def test_completed_streamed_tool_call_bridges_closed_event_across_provider(self): + handler = StreamingHandler() + tracker = SessionTracker(ttl_seconds=3600) + user_message = { + "role": "user", + "content": self._long_text("streamed-tool-event-user"), + } + root = tracker.infer_session( + {"messages": [user_message]}, provider="openai", model="gpt" + ) + responses = [] + + def record(response): + responses.append(response) + tracker.record_response( + root.session_id, + tracking_namespace=root.tracking_namespace, + response=response, + ) + + async def stream(): + yield "data: " + json.dumps( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_streamed", + "function": { + "name": "read_file", + "arguments": '{"path":', + }, + } + ] + } + } + ] + } + ) + "\n\n" + yield "data: " + json.dumps( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "function": {"arguments": '"a.py"}'}, + } + ] + } + } + ] + } + ) + "\n\n" + yield "data: " + json.dumps( + { + "choices": [{"delta": {}, "finish_reason": "tool_calls"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + } + ) + "\n\n" + + async def consume(): + return [ + chunk + async for chunk in handler.wrap_stream( + stream(), + "credential", + "model", + response_callback=record, + ) + ] + + asyncio.run(consume()) + assistant = responses[0]["choices"][0]["message"] + continued = tracker.infer_session( + { + "messages": [ + user_message, + assistant, + self._tool_result("call_streamed"), + ] + }, + provider="anthropic", + model="claude", + ) + + self.assertEqual( + assistant["tool_calls"][0]["function"]["arguments"], + '{"path":"a.py"}', + ) + self.assertEqual(continued.session_id, root.session_id) + self.assertEqual(continued.confidence, "probable") + + def test_streaming_chunk_collector_ignores_non_evidence_payloads(self): + handler = StreamingHandler() + assistant_parts = [] + tool_call_ids = [] + payloads = ( + "event: message\n\n", + "data: [DONE]\n\n", + "data: not-json\n\n", + "data: []\n\n", + "data: null\n\n", + "data: 42\n\n", + 'data: {"choices":[null,"invalid",{"delta":"invalid"}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":42}}]}\n\n', + ) + + for payload in payloads: + handler._collect_session_response_anchors( + payload, + assistant_parts, + tool_call_ids, + ) + + self.assertEqual(assistant_parts, []) + self.assertEqual(tool_call_ids, []) + + def test_streaming_chunk_collector_accepts_event_frames_and_data_without_space(self): + handler = StreamingHandler() + assistant_parts = [] + tool_call_ids = [] + handler._collect_session_response_anchors( + 'event: message\ndata:{"choices":[{"delta":{"content":"hello ",' + '"tool_calls":[{"id":"call_1"}]}}]}\n\n', + assistant_parts, + tool_call_ids, + ) + handler._collect_session_response_anchors( + 'data: {"choices":[{"delta":{"content":"world",' + '"tool_calls":[{"id":"call_1"}]}}]}\n\n', + assistant_parts, + tool_call_ids, + ) + + self.assertEqual(assistant_parts, ["hello ", "world"]) + self.assertEqual(tool_call_ids, ["call_1"]) + + def test_stream_eof_without_completion_signal_does_not_record_response_identity(self): + handler = StreamingHandler() + responses = [] + + async def partial_stream(): + yield { + "choices": [ + { + "delta": {"content": "partial assistant evidence"}, + "finish_reason": None, + } + ] + } + + async def consume(): + return [ + item + async for item in handler.wrap_stream( + partial_stream(), + "credential", + "model", + response_callback=responses.append, + ) + ] + + output = asyncio.run(consume()) + + self.assertEqual(responses, []) + self.assertEqual(output[-1], "data: [DONE]\n\n") + + def test_raw_stream_early_finish_reason_without_usage_does_not_record_identity(self): + handler = StreamingHandler() + responses = [] + + async def truncated_stream(): + yield ( + 'data: {"choices":[{"delta":{"content":"truncated evidence"},' + '"finish_reason":"stop"}]}\n\n' + ) + + async def consume(): + return [ + item + async for item in handler.wrap_stream( + truncated_stream(), + "credential", + "model", + response_callback=responses.append, + ) + ] + + output = asyncio.run(consume()) + + self.assertEqual(responses, []) + self.assertEqual(output[-1], "data: [DONE]\n\n") + + def test_raw_sse_completion_requires_done_or_usage_backing(self): + handler = StreamingHandler() + + self.assertFalse( + handler._sse_has_completion_signal( + 'data: {"choices":[{"finish_reason":"stop"}]}\n\n' + ) + ) + self.assertTrue( + handler._sse_has_completion_signal( + 'data: {"choices":[{"finish_reason":"stop"}],"usage":{}}\n\n' + ) + ) + self.assertTrue( + handler._sse_has_completion_signal( + 'data: {"choices":[],"usage":{"completion_tokens":3}}\n\n' + ) + ) + self.assertTrue(handler._sse_has_completion_signal("data: [DONE]\n\n")) + + def test_stream_finish_reason_records_completed_response_identity(self): + handler = StreamingHandler() + responses = [] + + async def complete_stream(): + yield { + "choices": [ + { + "delta": {"content": "complete assistant evidence"}, + "finish_reason": None, + } + ] + } + yield { + "choices": [ + { + "delta": {}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 4, + "completion_tokens": 3, + "total_tokens": 7, + }, + } + + async def consume(): + return [ + item + async for item in handler.wrap_stream( + complete_stream(), + "credential", + "model", + response_callback=responses.append, + ) + ] + + output = asyncio.run(consume()) + + self.assertEqual(len(responses), 1) + self.assertEqual( + responses[0]["choices"][0]["message"]["content"], + "complete assistant evidence", + ) + self.assertEqual(output[-1], "data: [DONE]\n\n") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_startup_display.py b/tests/test_startup_display.py new file mode 100644 index 000000000..a8538f88e --- /dev/null +++ b/tests/test_startup_display.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from proxy_app.startup_display import mask_secret_for_display + + +def test_proxy_api_key_display_is_redacted() -> None: + value = "sk-proxy-secret-value" + + display = mask_secret_for_display(value) + + assert value not in display + assert display == "Set (sk-p...alue, 21 chars)" + + +def test_short_proxy_api_key_display_is_fully_redacted() -> None: + value = "secret" + + display = mask_secret_for_display(value) + + assert value not in display + assert display == "Set (redacted, 6 chars)" diff --git a/tests/test_stream_events.py b/tests/test_stream_events.py new file mode 100644 index 000000000..365e9cfbe --- /dev/null +++ b/tests/test_stream_events.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from rotator_library.streaming import stream_event_from_sse_chunk + + +def test_stream_event_from_sse_chunk_detects_visible_chat_delta() -> None: + event = stream_event_from_sse_chunk('data: {"choices":[{"delta":{"content":"hi"}}]}\n\n') + + assert event.event_type == "delta" + assert event.visible_output is True + + +def test_stream_event_from_sse_chunk_treats_error_and_done_as_not_visible() -> None: + error_event = stream_event_from_sse_chunk('data: {"error":{"type":"rate_limit"}}\n\n') + done_event = stream_event_from_sse_chunk("data: [DONE]\n\n") + + assert error_event.event_type == "error" + assert error_event.visible_output is False + assert done_event.event_type == "completed" + assert done_event.visible_output is False + + +def test_stream_event_from_sse_chunk_malformed_fails_closed() -> None: + event = stream_event_from_sse_chunk("data: not-json\n\n") + + assert event.event_type == "metadata" + assert event.visible_output is False diff --git a/tests/test_stream_metrics.py b/tests/test_stream_metrics.py new file mode 100644 index 000000000..5f082f0c1 --- /dev/null +++ b/tests/test_stream_metrics.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from rotator_library.streaming import StreamEvent, StreamMonitor + + +class FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +def test_stream_metrics_records_ttfb_ttft_and_counts() -> None: + clock = FakeClock() + monitor = StreamMonitor(clock=clock) + clock.advance(0.5) + monitor.record_event(StreamEvent("parsed_chunk")) + clock.advance(0.7) + monitor.record_event(StreamEvent("delta", visible_output=True)) + clock.advance(0.3) + monitor.complete() + + assert monitor.metrics.ttfb_seconds == 0.5 + assert monitor.metrics.ttft_seconds == 1.2 + assert monitor.metrics.duration_seconds == 1.5 + assert monitor.metrics.chunk_count == 2 + assert monitor.metrics.visible_chunk_count == 1 + + +def test_stream_monitor_records_cancel_and_stall() -> None: + clock = FakeClock() + monitor = StreamMonitor(clock=clock) + monitor.record_event(StreamEvent("parsed_chunk")) + clock.advance(5) + + assert monitor.is_stalled(2) is True + monitor.cancel() + assert monitor.metrics.cancelled is True diff --git a/tests/test_stream_policy.py b/tests/test_stream_policy.py new file mode 100644 index 000000000..1edc8aac1 --- /dev/null +++ b/tests/test_stream_policy.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from rotator_library.client.stream_retry_policy import can_retry_stream_after_error as compat_retry_policy +from rotator_library.client.streaming import StreamingHandler +from rotator_library.core.errors import StreamedAPIError +from rotator_library.streaming.policy import can_retry_stream_after_error, is_visible_stream_output + + +def test_reasoning_only_retry_policy_is_preserved() -> None: + reasoning_chunk = 'data: {"choices":[{"delta":{"reasoning_content":"thinking"}}]}\n\n' + text_chunk = 'data: {"choices":[{"delta":{"content":"visible"}}]}\n\n' + + assert can_retry_stream_after_error(None, False) is True + assert can_retry_stream_after_error(reasoning_chunk, True) is True + assert can_retry_stream_after_error(reasoning_chunk, False) is False + assert can_retry_stream_after_error(text_chunk, True) is False + assert compat_retry_policy(reasoning_chunk, True) is True + + +def test_heartbeat_comments_do_not_block_stream_retry() -> None: + assert can_retry_stream_after_error(': heartbeat\n\n', False) is True + assert can_retry_stream_after_error('event: message_start\ndata: {"type":"message_start","message":{"role":"assistant","content":[]}}\n\n', False) is True + assert can_retry_stream_after_error('event: response.created\ndata: {"type":"response.created","response":{"status":"in_progress","output":[]}}\n\n', False) is True + assert can_retry_stream_after_error('data: {"choices":[{"delta":{"role":"assistant"}}]}\n\n', False) is True + + +def test_cost_events_do_not_count_as_visible_output() -> None: + cost_event = 'event: cost\ndata: {"total_cost":0.01}\n\n' + scalar_cost_event = "event: cost\ndata: 0.01\n\n" + + assert is_visible_stream_output(cost_event) is False + assert is_visible_stream_output(cost_event, protocol="responses") is False + assert can_retry_stream_after_error(cost_event, False) is True + assert is_visible_stream_output(scalar_cost_event) is False + assert can_retry_stream_after_error(scalar_cost_event, False) is True + + +def test_visible_output_detection_for_chat_chunks() -> None: + assert is_visible_stream_output('data: {"choices":[{"delta":{"content":"hello"}}]}\n\n') is True + assert is_visible_stream_output('data: {"choices":[{"delta":{"tool_calls":[{"id":"call_1"}]}}]}\n\n') is True + assert is_visible_stream_output('data: {"error":{"type":"rate_limit"}}\n\n') is False + assert is_visible_stream_output('event: error\ndata: {"type":"error","error":{"message":"x"}}\n\n') is False + assert is_visible_stream_output('data: {"type":"error","error":{"message":"x"}}\n\n') is False + assert is_visible_stream_output(': heartbeat\n\n') is False + assert is_visible_stream_output("data: [DONE]\n\n") is False + assert is_visible_stream_output("not-sse") is True + + +def test_visible_output_detection_for_responses_events() -> None: + assert is_visible_stream_output('data: {"event_type":"response.output_text.delta","delta":"hi"}\n\n', protocol="responses") is True + assert is_visible_stream_output('data: {"event_type":"response.output_text.delta","delta":"hi"}\n\n') is True + assert is_visible_stream_output('data: {"event_type":"response.function_call_arguments.delta","delta":"{\\"x\\":"}\n\n') is True + assert is_visible_stream_output('data: {"event_type":"response.failed","error":{"message":"x"}}\n\n', protocol="responses") is False + assert is_visible_stream_output('event: response.failed\ndata: {"error":{"message":"x"}}\n\n', protocol="responses") is False + assert is_visible_stream_output('event: response.output_item.added\ndata: {"type":"response.output_item.added","item":{"id":"fc_0","type":"function_call","call_id":"call_1","name":"lookup","arguments":""}}\n\n') is True + + +def test_visible_output_detection_for_anthropic_and_gemini_events() -> None: + assert is_visible_stream_output('event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"type":"thinking_delta","thinking":"visible thought"}}\n\n') is True + assert is_visible_stream_output('event: content_block_start\ndata: {"type":"content_block_start","content_block":{"type":"tool_use","id":"call_1","name":"lookup","input":{}}}\n\n') is True + assert is_visible_stream_output('data: {"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"lookup","args":{}}}]}}]}\n\n') is True + + +def test_stream_handler_rejects_string_and_event_level_error_payloads() -> None: + handler = StreamingHandler() + + for chunk in ( + 'data: {"error":"provider failed"}\n\n', + 'event: error\ndata: {"type":"rate_limit","message":"slow down"}\n\n', + {"type": "response.failed", "response": {"error": {"type": "server_error", "message": "unavailable"}}}, + ): + try: + handler._process_chunk(chunk, None, False) + except StreamedAPIError as error: + assert error.data["error"]["message"] + else: + raise AssertionError("in-band stream error must raise") diff --git a/tests/test_stream_transport.py b/tests/test_stream_transport.py new file mode 100644 index 000000000..4b52e1e3f --- /dev/null +++ b/tests/test_stream_transport.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from rotator_library.streaming import JSONLineStreamFormatter, SSEStreamFormatter, StreamEvent, WebSocketStreamFormatter + + +def test_sse_formatter_outputs_named_event() -> None: + formatted = SSEStreamFormatter().format_event(StreamEvent("delta", data={"text": "hi"}, visible_output=True)) + + assert formatted.startswith("event: delta\n") + assert "data:" in formatted + + +def test_sse_formatter_outputs_comment_heartbeat() -> None: + formatted = SSEStreamFormatter().format_heartbeat("keepalive") + + assert formatted == ": keepalive\n\n" + + +def test_websocket_formatter_exposes_future_transport_shape() -> None: + formatted = WebSocketStreamFormatter().format_event(StreamEvent("delta", data={"text": "hi"})) + + assert formatted["type"] == "delta" + assert formatted["payload"]["transport"] == "sse" + assert WebSocketStreamFormatter().format_heartbeat()["type"] == "heartbeat" + + +def test_jsonl_formatter_outputs_one_line_json() -> None: + formatted = JSONLineStreamFormatter().format_event(StreamEvent("metadata", data={"ok": True})) + + assert formatted.endswith("\n") + assert "metadata" in formatted + assert "heartbeat" in JSONLineStreamFormatter().format_heartbeat() diff --git a/tests/test_streaming_error_handler.py b/tests/test_streaming_error_handler.py new file mode 100644 index 000000000..f85b513d1 --- /dev/null +++ b/tests/test_streaming_error_handler.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from litellm import RateLimitError + +from rotator_library.streaming.errors import decide_streaming_error_action + + +def _rate_limit(retry_after: int | None = None) -> Exception: + error = RateLimitError("rate limited", llm_provider="openai", model="gpt-test") + if retry_after is not None: + error.retry_after = retry_after + return error + + +def test_streaming_error_decision_retries_same_key_for_small_retry_after() -> None: + decision = decide_streaming_error_action( + _rate_limit(3), + provider="openai", + last_streamed_chunk=None, + attempt=0, + max_retries=2, + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + provider_cooldown_default_seconds=30, + ) + + assert decision.action == "retry_same" + assert decision.start_provider_cooldown is False + + +def test_streaming_error_decision_starts_cooldown_before_visible_output() -> None: + decision = decide_streaming_error_action( + _rate_limit(60), + provider="openai", + last_streamed_chunk=None, + attempt=1, + max_retries=2, + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + provider_cooldown_default_seconds=30, + ) + + assert decision.action == "rotate" + assert decision.start_provider_cooldown is True + assert decision.provider_cooldown_duration == 60 + assert decision.provider_cooldown_scope == "provider" + + +def test_streaming_error_decision_blocks_after_visible_output_and_skips_cooldown() -> None: + decision = decide_streaming_error_action( + _rate_limit(60), + provider="openai", + last_streamed_chunk='data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', + attempt=0, + max_retries=2, + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + provider_cooldown_default_seconds=30, + ) + + assert decision.action == "fallback_blocked_after_output" + assert decision.start_provider_cooldown is False + + +def test_streaming_error_decision_blocks_after_prior_visible_output() -> None: + decision = decide_streaming_error_action( + _rate_limit(60), + provider="openai", + last_streamed_chunk='data: {"usage":{"total_tokens":1}}\n\n', + emitted_output=True, + attempt=0, + max_retries=2, + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + provider_cooldown_default_seconds=30, + ) + + assert decision.action == "fallback_blocked_after_output" + assert decision.start_provider_cooldown is False + + +def test_streaming_error_decision_allows_reasoning_only_retry_when_enabled() -> None: + decision = decide_streaming_error_action( + _rate_limit(3), + provider="openai", + last_streamed_chunk='data: {"choices":[{"delta":{"reasoning_content":"thinking"}}]}\n\n', + attempt=0, + max_retries=2, + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + provider_cooldown_default_seconds=30, + allow_reasoning_only_retry=True, + ) + + assert decision.action == "retry_same" + + +def test_streaming_error_decision_reports_model_cooldown_scope() -> None: + decision = decide_streaming_error_action( + Exception("MODEL_CAPACITY_EXHAUSTED"), + provider="openai", + model="gpt-5", + last_streamed_chunk=None, + attempt=1, + max_retries=2, + small_cooldown_threshold=10, + provider_cooldown_min_seconds=10, + provider_cooldown_default_seconds=30, + ) + + assert decision.start_provider_cooldown is True + assert decision.provider_cooldown_scope == "model" + assert decision.provider_cooldown_model == "gpt-5" diff --git a/tests/test_streaming_fallback_policy.py b/tests/test_streaming_fallback_policy.py new file mode 100644 index 000000000..81576d205 --- /dev/null +++ b/tests/test_streaming_fallback_policy.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import json +from types import MethodType + +import pytest + +from rotator_library.client import executor as executor_module +from rotator_library.client.executor import RequestExecutor +from rotator_library.core.types import RequestContext +from rotator_library.routing import parse_route_target +from rotator_library.routing.types import FallbackGroup +from rotator_library.transaction_logger import TransactionLogger + + +class StreamFailure(Exception): + def __init__(self, error_type: str) -> None: + super().__init__(error_type) + self.error_type = error_type + + +def _context(*, logger=None) -> RequestContext: + targets = (parse_route_target("codex/gpt-5.1-codex"), parse_route_target("openai/gpt-5.1")) + return RequestContext( + model="code", + provider="requested", + kwargs={"model": "code", "messages": [], "stream": True}, + streaming=True, + credentials=["cred-a"], + deadline=9999999999.0, + transaction_logger=logger, + routing_targets=targets, + routing_group_name="code_chain", + routing_group=FallbackGroup(name="code_chain", targets=targets, failover_on=frozenset({"authentication", "rate_limit"}), stop_on=frozenset({"validation"})), + ) + + +def _context_never_streaming_fallback(*, logger=None) -> RequestContext: + context = _context(logger=logger) + context.routing_group = FallbackGroup(name="code_chain", targets=context.routing_targets, failover_on=frozenset({"rate_limit"}), streaming_policy="never") + return context + + +@pytest.mark.asyncio +async def test_streaming_fallback_hard_stops_auth_even_with_group_override() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_stream(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + raise StreamFailure("authentication") + yield "data: [DONE]\n\n" + + executor._execute_streaming = MethodType(fake_stream, executor) + + with pytest.raises(StreamFailure): + [chunk async for chunk in executor._execute_streaming_with_fallback(_context())] + + assert attempts == ["codex"] + + +@pytest.mark.asyncio +async def test_streaming_fallback_tries_next_target_before_output() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_stream(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + raise StreamFailure("rate_limit") + yield 'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n' + yield "data: [DONE]\n\n" + + executor._execute_streaming = MethodType(fake_stream, executor) + + context = _context() + chunks = [chunk async for chunk in executor._execute_streaming_with_fallback(context)] + + assert attempts == ["codex", "openai"] + assert chunks[-1] == "data: [DONE]\n\n" + assert context.routing_attempt_history[0]["error_type"] == "rate_limit" + assert context.routing_attempt_history[0]["fallback_allowed"] is True + assert context.routing_attempt_history[1]["success"] is True + + +@pytest.mark.asyncio +async def test_streaming_fallback_blocks_after_visible_output() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + + async def fake_stream(self, context): + yield 'data: {"choices":[{"delta":{"content":"partial"}}]}\n\n' + raise StreamFailure("rate_limit") + + executor._execute_streaming = MethodType(fake_stream, executor) + + with pytest.raises(StreamFailure): + [chunk async for chunk in executor._execute_streaming_with_fallback(_context())] + + +@pytest.mark.asyncio +async def test_streaming_fallback_blocks_after_zero_argument_responses_tool_call() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_stream(self, context): + attempts.append(context.provider) + yield 'event: response.output_item.added\ndata: {"type":"response.output_item.added","item":{"id":"fc_0","type":"function_call","call_id":"call_1","name":"lookup","arguments":""}}\n\n' + raise StreamFailure("rate_limit") + + executor._execute_streaming = MethodType(fake_stream, executor) + + with pytest.raises(StreamFailure): + _ = [chunk async for chunk in executor._execute_streaming_with_fallback(_context())] + + assert attempts == ["codex"] + + +@pytest.mark.asyncio +async def test_streaming_fallback_trace_records_blocked_after_output(tmp_path) -> None: + executor = RequestExecutor.__new__(RequestExecutor) + logger = TransactionLogger("routing", "code", parent_dir=tmp_path) + + async def fake_stream(self, context): + yield 'data: {"choices":[{"delta":{"content":"partial"}}]}\n\n' + raise StreamFailure("rate_limit") + + executor._execute_streaming = MethodType(fake_stream, executor) + + with pytest.raises(StreamFailure): + [chunk async for chunk in executor._execute_streaming_with_fallback(_context(logger=logger))] + + pass_names = [json.loads(line)["pass_name"] for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + assert "routing_stream_target_attempt_started" in pass_names + assert "routing_stream_target_attempt_failed" in pass_names + assert "routing_stream_fallback_blocked_after_output" in pass_names + + +@pytest.mark.asyncio +async def test_streaming_fallback_treats_error_chunk_as_not_visible_output() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_stream(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + yield 'data: {"error":{"type":"rate_limit"}}\n\n' + yield "data: [DONE]\n\n" + return + yield "data: [DONE]\n\n" + + executor._execute_streaming = MethodType(fake_stream, executor) + + chunks = [chunk async for chunk in executor._execute_streaming_with_fallback(_context())] + + assert attempts == ["codex", "openai"] + assert chunks == ["data: [DONE]\n\n"] + + +@pytest.mark.asyncio +async def test_streaming_fallback_classifies_nested_responses_failure() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_stream(self, context): + attempts.append(context.provider) + if len(attempts) == 1: + yield 'event: response.failed\ndata: {"type":"response.failed","response":{"status":"failed","error":{"type":"rate_limit","message":"slow down"}}}\n\n' + return + yield "data: [DONE]\n\n" + + executor._execute_streaming = MethodType(fake_stream, executor) + + chunks = [chunk async for chunk in executor._execute_streaming_with_fallback(_context())] + + assert attempts == ["codex", "openai"] + assert chunks == ["data: [DONE]\n\n"] + + +@pytest.mark.asyncio +async def test_streaming_fallback_respects_never_policy_before_output() -> None: + executor = RequestExecutor.__new__(RequestExecutor) + attempts = [] + + async def fake_stream(self, context): + attempts.append(context.provider) + raise StreamFailure("rate_limit") + yield "" + + executor._execute_streaming = MethodType(fake_stream, executor) + + with pytest.raises(StreamFailure): + [chunk async for chunk in executor._execute_streaming_with_fallback(_context_never_streaming_fallback())] + + assert attempts == ["codex"] + + +@pytest.mark.asyncio +async def test_streaming_fallback_exhaustion_trace_uses_sanitized_summaries(tmp_path) -> None: + executor = RequestExecutor.__new__(RequestExecutor) + logger = TransactionLogger("routing", "code", parent_dir=tmp_path) + + async def fake_stream(self, context): + raise StreamFailure("rate_limit") + yield "" + + executor._execute_streaming = MethodType(fake_stream, executor) + + with pytest.raises(StreamFailure): + [chunk async for chunk in executor._execute_streaming_with_fallback(_context_never_streaming_fallback(logger=logger))] + + entries = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + exhausted = [entry for entry in entries if entry["pass_name"] == "routing_fallback_exhausted"][-1] + assert exhausted["metadata"]["fallback_targets"][0]["message"] == "" + assert exhausted["metadata"]["streaming_policy"] == "never" + + +def test_stream_timeout_details_merge_into_aggregate_error() -> None: + error_data = {"error": {"type": "proxy_error", "details": {"attempts": 1}}} + stream_error = {"error": {"type": "api_connection", "details": {"timeout_type": "ttfb", "timeout_seconds": 0.1}}} + + executor_module._merge_stream_error_details(error_data, stream_error) + + assert error_data["error"]["type"] == "api_connection" + assert error_data["error"]["details"]["attempts"] == 1 + assert error_data["error"]["details"]["timeout_type"] == "ttfb" diff --git a/tests/test_streaming_usage_accounting.py b/tests/test_streaming_usage_accounting.py new file mode 100644 index 000000000..8077ee634 --- /dev/null +++ b/tests/test_streaming_usage_accounting.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.client.streaming import StreamingHandler +from rotator_library.transaction_logger import TransactionLogger + + +class FakeCredentialContext: + def __init__(self) -> None: + self.success_kwargs = None + + def mark_success(self, **kwargs) -> None: + self.success_kwargs = kwargs + + +async def _usage_chunks(): + yield {"id": "chunk_1", "choices": [{"delta": {"content": "hi"}}]} + yield { + "id": "chunk_2", + "choices": [{"delta": {}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 30, + "prompt_tokens_details": {"cached_tokens": 40, "cache_creation_tokens": 5}, + "completion_tokens_details": {"reasoning_tokens": 10}, + }, + } + + +async def _zero_usage_chunks(): + yield {"id": "chunk_1", "choices": [{"delta": {"content": "hi"}}]} + yield {"id": "chunk_2", "choices": [{"delta": {}, "finish_reason": "stop"}]} + + +async def _cost_comment_chunks(): + yield ': cost {"total_cost":0.042,"currency":"USD","source":"provider_sse"}\n\n' + yield {"id": "chunk_1", "choices": [{"delta": {"content": "hi"}}]} + yield {"id": "chunk_2", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1}} + + +async def _cost_event_chunks(): + yield 'event: cost\ndata: {"total_cost":0.021,"currency":"EUR","source":"event_cost"}\n\n' + yield {"id": "chunk_1", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1}} + + +async def _scalar_cost_event_chunks(): + yield "event: cost\ndata: 0.033\n\n" + yield {"id": "chunk_1", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1}} + + +async def _request_cost_comment_chunks(): + yield ': cost {"request_cost_usd":0.044,"source":"reference_sse"}\n\n' + yield {"id": "chunk_1", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1}} + + +async def _estimated_cost_comment_chunks(): + yield ': cost {"estimated_cost":0.045,"source":"reference_estimate"}\n\n' + yield {"id": "chunk_1", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1}} + + +async def _top_level_cost_usage_sse_chunks(): + yield 'data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1},"total_cost":0.055}\n\n' + + +async def _top_level_cost_usage_dict_chunks(): + yield {"id": "chunk_1", "choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1}, "total_cost": 0.066} + + +async def _cost_comment_overridden_by_final_usage_chunks(): + yield ': cost 0.042\n\n' + yield { + "id": "chunk_2", + "choices": [{"delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "cost_details": {"total_cost": 0.084, "source": "final_usage"}}, + } + + +@pytest.mark.asyncio +async def test_streaming_usage_uses_normalized_accounting_and_trace(tmp_path, monkeypatch) -> None: + monkeypatch.setattr( + "rotator_library.usage.costs.litellm.get_model_info", + lambda model: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + cred_context = FakeCredentialContext() + logger = TransactionLogger("openai", "gpt-test", parent_dir=tmp_path) + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(_usage_chunks(), "cred", "gpt-test", cred_context=cred_context, transaction_logger=logger)] + + assert chunks[-1] == "data: [DONE]\n\n" + assert cred_context.success_kwargs["prompt_tokens"] == 55 + assert cred_context.success_kwargs["prompt_tokens_cache_read"] == 40 + assert cred_context.success_kwargs["prompt_tokens_cache_write"] == 5 + assert cred_context.success_kwargs["completion_tokens"] == 20 + assert cred_context.success_kwargs["thinking_tokens"] == 10 + assert cred_context.success_kwargs["approx_cost"] > 0 + entries = [json.loads(line) for line in (logger.log_dir / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines()] + assert any(entry["pass_name"] == "usage_accounting_summary" for entry in entries) + + +@pytest.mark.asyncio +async def test_streaming_usage_skip_cost_returns_zero() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_usage_chunks(), "cred", "gpt-test", cred_context=cred_context, skip_cost_calculation=True)] + + assert cred_context.success_kwargs["approx_cost"] == 0.0 + + +@pytest.mark.asyncio +async def test_streaming_without_usage_still_marks_success_with_zero_usage() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_zero_usage_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["prompt_tokens"] == 0 + assert cred_context.success_kwargs["completion_tokens"] == 0 + assert cred_context.success_kwargs["thinking_tokens"] == 0 + assert cred_context.success_kwargs["prompt_tokens_cache_read"] == 0 + assert cred_context.success_kwargs["prompt_tokens_cache_write"] == 0 + + +@pytest.mark.asyncio +async def test_streaming_completed_calls_success_callback() -> None: + called = [] + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_zero_usage_chunks(), "cred", "gpt-test", success_callback=lambda: called.append(True))] + + assert called == [True] + + +@pytest.mark.asyncio +async def test_streaming_usage_uses_configured_env_pricing(monkeypatch) -> None: + monkeypatch.setenv("MODEL_PRICE_OPENAI_GPT_TEST_INPUT", "2.0") + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_usage_chunks(), "cred", "openai/gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 110.0 + + +@pytest.mark.asyncio +async def test_streaming_cost_comment_updates_approx_cost() -> None: + cred_context = FakeCredentialContext() + + chunks = [chunk async for chunk in StreamingHandler().wrap_stream(_cost_comment_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert chunks[0].startswith(": cost") + assert cred_context.success_kwargs["approx_cost"] == 0.042 + + +@pytest.mark.asyncio +async def test_streaming_cost_event_updates_approx_cost() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_cost_event_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 0.021 + + +@pytest.mark.asyncio +async def test_streaming_scalar_cost_event_updates_approx_cost() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_scalar_cost_event_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 0.033 + + +@pytest.mark.asyncio +async def test_streaming_reference_request_cost_comment_updates_approx_cost() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_request_cost_comment_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 0.044 + + +@pytest.mark.asyncio +async def test_streaming_estimated_cost_comment_updates_approx_cost() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_estimated_cost_comment_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 0.045 + + +@pytest.mark.asyncio +async def test_streaming_sse_usage_preserves_top_level_cost() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_top_level_cost_usage_sse_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 0.055 + + +@pytest.mark.asyncio +async def test_streaming_dict_usage_preserves_top_level_cost() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_top_level_cost_usage_dict_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 0.066 + + +@pytest.mark.asyncio +async def test_streaming_final_usage_cost_overrides_comment_cost() -> None: + cred_context = FakeCredentialContext() + + _ = [chunk async for chunk in StreamingHandler().wrap_stream(_cost_comment_overridden_by_final_usage_chunks(), "cred", "gpt-test", cred_context=cred_context)] + + assert cred_context.success_kwargs["approx_cost"] == 0.084 diff --git a/tests/test_transaction_logger_transform_trace.py b/tests/test_transaction_logger_transform_trace.py new file mode 100644 index 000000000..250a0df60 --- /dev/null +++ b/tests/test_transaction_logger_transform_trace.py @@ -0,0 +1,298 @@ +from __future__ import annotations + +import json + +import pytest + +from rotator_library.client.executor import RequestExecutor +from rotator_library.client.transforms import ProviderTransforms +from rotator_library.transaction_logger import ProviderLogger, TransactionLogger +from rotator_library.transform_trace import REDACTED + + +def _trace_entries(log_dir): + trace_file = log_dir / "transform_trace.jsonl" + return [json.loads(line) for line in trace_file.read_text(encoding="utf-8").splitlines()] + + +def test_log_request_writes_legacy_file_and_raw_trace(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + + logger.log_request({"model": "gpt-test", "api_key": "secret", "messages": [{"role": "user", "content": "hi"}]}) + + assert (logger.log_dir / "request.json").exists() + entries = _trace_entries(logger.log_dir) + assert entries[0]["pass_name"] == "raw_client_request" + assert entries[0]["request_id"] == logger.request_id + assert entries[0]["direction"] == "request" + assert entries[0]["data"]["api_key"] == REDACTED + assert (logger.log_dir / "transforms" / "0001_raw_client_request.json").exists() + + +def test_log_transformed_request_records_trace_even_when_legacy_file_skips(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + logger.set_trace_context(session_id="session_1", scope_key="scope_1", classifier="class_a") + request = {"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}]} + + logger.log_transformed_request(request, dict(request), credential_id="cred_1") + + entries = _trace_entries(logger.log_dir) + assert entries[0]["pass_name"] == "prepared_provider_request" + assert entries[0]["changed_from_previous"] is False + assert entries[0]["credential_id"] == "cred_1" + assert entries[0]["session_id"] == "session_1" + assert entries[0]["scope_key"] == "scope_1" + assert entries[0]["classifier"] == "class_a" + assert not (logger.log_dir / "request_transformed.json").exists() + + +def test_log_response_and_stream_chunk_write_trace_entries(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + logger.log_request({"model": "gpt-test", "stream": True}) + + logger.log_stream_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + logger.log_response({"model": "gpt-test", "choices": [], "usage": {"total_tokens": 2}}) + + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert "parsed_stream_chunk" in pass_names + assert "final_client_response" in pass_names + stream_entry = next(entry for entry in entries if entry["pass_name"] == "parsed_stream_chunk") + assert stream_entry["direction"] == "stream" + assert not (logger.log_dir / "transforms" / "0002_parsed_stream_chunk.json").exists() + + +def test_provider_logger_writes_provider_trace_entries(tmp_path) -> None: + logger = TransactionLogger("gemini_cli", "gemini_cli/gemini-test", parent_dir=tmp_path) + provider_logger = ProviderLogger(logger.get_context()) + + provider_logger.log_request({"credential_identifier": "secret", "body": {"text": "hi"}}) + provider_logger.log_response_chunk("data: chunk") + provider_logger.log_final_response({"ok": True}) + provider_logger.log_error("provider failed") + + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert pass_names == [ + "provider_request_payload", + "provider_raw_stream_chunk", + "provider_final_response", + "provider_error", + ] + assert entries[0]["component"] == "provider" + assert entries[0]["request_id"] == logger.request_id + assert entries[0]["data"]["credential_identifier"] == REDACTED + assert (logger.log_dir / "provider" / "request_payload.json").exists() + + +@pytest.mark.asyncio +async def test_stream_wrapper_records_raw_parsed_and_assembled_trace(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + executor = RequestExecutor.__new__(RequestExecutor) + + async def stream(): + yield 'data: {"id":"chunk_1","choices":[{"delta":{"content":"Hi"},"finish_reason":null}]}\n\n' + yield 'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n\n' + yield "data: [DONE]\n\n" + + chunks = [chunk async for chunk in executor._transaction_logging_stream_wrapper(stream(), logger, {})] + + assert chunks[-1] == "data: [DONE]\n\n" + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert pass_names.count("raw_stream_chunk") == 3 + assert pass_names.count("parsed_stream_chunk") == 2 + assert pass_names.count("stream_done_event") == 1 + assert "assembled_stream_response" in pass_names + assert "final_client_response" in pass_names + + +@pytest.mark.asyncio +async def test_stream_wrapper_records_error_events(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + executor = RequestExecutor.__new__(RequestExecutor) + + async def stream(): + yield 'data: {"error":{"type":"rate_limit","message":"slow down"}}\n\n' + yield "data: [DONE]\n\n" + + chunks = [chunk async for chunk in executor._transaction_logging_stream_wrapper(stream(), logger, {})] + + assert chunks[-1] == "data: [DONE]\n\n" + pass_names = [entry["pass_name"] for entry in _trace_entries(logger.log_dir)] + assert "stream_error_event" in pass_names + assert "stream_done_event" in pass_names + + +def test_executor_terminal_stream_errors_are_traced(tmp_path) -> None: + class Context: + pass + + context = Context() + context.transaction_logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + context.streaming = True + context.provider = "openai" + context.model = "openai/gpt-test" + context.session_id = None + context.usage_manager_key = "openai" + context.classifier = None + + executor = RequestExecutor.__new__(RequestExecutor) + + lines = executor._terminal_stream_error_lines(context, {"error": {"type": "proxy_error"}}) + + assert lines[-1] == "data: [DONE]\n\n" + pass_names = [entry["pass_name"] for entry in _trace_entries(context.transaction_logger.log_dir)] + assert pass_names == ["stream_error_event", "stream_done_event"] + + +def test_transaction_logger_disabled_writes_no_trace(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", enabled=False, parent_dir=tmp_path) + + logger.log_request({"model": "gpt-test"}) + logger.log_response({"model": "gpt-test"}) + + assert logger.log_dir is None + assert not list(tmp_path.iterdir()) + + +def test_provider_error_trace_scrubs_header_like_secret_text(tmp_path) -> None: + logger = TransactionLogger("gemini_cli", "gemini_cli/gemini-test", parent_dir=tmp_path) + provider_logger = ProviderLogger(logger.get_context()) + + provider_logger.log_error("upstream failed Authorization: Bearer secret-token") + + entries = _trace_entries(logger.log_dir) + assert entries[0]["pass_name"] == "provider_error" + assert "secret-token" not in entries[0]["data"]["message"] + assert "[REDACTED]" in entries[0]["data"]["message"] + + +def test_log_transform_error_uses_standard_shape_and_scrubs_text(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + + logger.log_transform_error( + "after_field_cache_injection", + RuntimeError("bad Authorization: Bearer secret"), + payload={"cookie": "sid=secret"}, + ) + + entries = _trace_entries(logger.log_dir) + assert entries[0]["pass_name"] == "transform_log_error" + assert entries[0]["data"]["failed_pass_name"] == "after_field_cache_injection" + assert "secret" not in json.dumps(entries[0]["data"]) + + +def test_trace_redacts_camel_case_secret_keys(tmp_path) -> None: + logger = TransactionLogger("openai", "openai/gpt-test", parent_dir=tmp_path) + + logger.log_transform_pass( + "camel_secret_payload", + {"apiKey": "a", "accessToken": "b", "refreshToken": "c", "clientSecret": "d", "idToken": "e"}, + direction="request", + stage="client", + ) + + data = _trace_entries(logger.log_dir)[0]["data"] + assert set(data.values()) == {REDACTED} + + +@pytest.mark.asyncio +async def test_provider_transforms_trace_each_live_boundary(tmp_path) -> None: + class HookPlugin: + async def transform_request(self, kwargs, model, credential): + kwargs["hooked"] = credential + return ["hooked request"] + + def get_model_options(self, model): + return {"reasoning_effort": "low", "temperature": 0.2} + + class Config: + def convert_for_litellm(self, provider_override=None, **kwargs): + converted = dict(kwargs) + converted["converted_for_litellm"] = True + return converted + + logger = TransactionLogger("dedaluslabs", "dedaluslabs/test", parent_dir=tmp_path) + transforms = ProviderTransforms( + {"dedaluslabs": HookPlugin()}, + provider_config=Config(), + ) + + result = await transforms.apply( + "dedaluslabs", + "dedaluslabs/test", + "secret-credential", + {"model": "dedaluslabs/test", "tool_choice": "auto"}, + transaction_logger=logger, + credential_id="stable_cred", + transport="http", + trace_metadata={"scope_key": "scope"}, + ) + + assert "tool_choice" not in result + assert result["hooked"] == "secret-credential" + assert result["reasoning_effort"] == "low" + assert result["converted_for_litellm"] is True + entries = _trace_entries(logger.log_dir) + pass_names = [entry["pass_name"] for entry in entries] + assert pass_names == [ + "pre_provider_transform_request", + "after_builtin_provider_transform", + "after_provider_hook_transform", + "after_provider_model_options", + "before_litellm_conversion", + "after_litellm_conversion", + ] + assert all(entry["credential_id"] == "stable_cred" for entry in entries) + assert entries[1]["metadata"]["transform_provider"] == "dedaluslabs" + assert entries[-1]["changed_from_previous"] is True + + +@pytest.mark.asyncio +async def test_provider_builtin_transform_errors_are_traced(tmp_path) -> None: + def broken_transform(kwargs, model, provider): + raise RuntimeError("bad apiKey: secret") + + logger = TransactionLogger("broken", "broken/test", parent_dir=tmp_path) + transforms = ProviderTransforms({}) + transforms._transforms["broken"] = [broken_transform] + + with pytest.raises(RuntimeError): + await transforms.apply( + "broken", + "broken/test", + "cred", + {"model": "broken/test", "apiKey": "secret"}, + transaction_logger=logger, + credential_id="cred_1", + ) + + entries = _trace_entries(logger.log_dir) + error_entry = [entry for entry in entries if entry["pass_name"] == "transform_log_error"][-1] + assert error_entry["data"]["failed_pass_name"] == "builtin_provider_transform" + assert "secret" not in json.dumps(error_entry["data"]) + + +@pytest.mark.asyncio +async def test_provider_transforms_do_not_deepcopy_for_trace_when_disabled() -> None: + class NoCopyValue: + def __deepcopy__(self, memo): + raise AssertionError("trace comparison should not copy when tracing is disabled") + + class HookPlugin: + async def transform_request(self, kwargs, model, credential): + kwargs["hooked"] = True + return ["hooked request"] + + transforms = ProviderTransforms({"dedaluslabs": HookPlugin()}) + + result = await transforms.apply( + "dedaluslabs", + "dedaluslabs/test", + "secret-credential", + {"model": "dedaluslabs/test", "tool_choice": "auto", "opaque": NoCopyValue()}, + ) + + assert "tool_choice" not in result + assert result["hooked"] is True diff --git a/tests/test_transform_trace.py b/tests/test_transform_trace.py new file mode 100644 index 000000000..682ae6b1f --- /dev/null +++ b/tests/test_transform_trace.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from datetime import datetime +from decimal import Decimal + +from rotator_library.transform_trace import ( + REDACTED, + TransformTraceWriter, + sanitize_filename, + sanitize_for_trace, + scrub_sensitive_text, +) + + +@dataclass +class ExamplePayload: + when: datetime + amount: Decimal + + +def test_sanitize_for_trace_redacts_sensitive_keys_recursively() -> None: + payload = { + "api_key": "secret-key", + "headers": {"Authorization": "Bearer secret", "Cookie": "sid=secret", "normal": "token in normal text"}, + "items": [{"refresh_token": "refresh", "text": "token should remain in value"}], + } + + sanitized = sanitize_for_trace(payload) + + assert sanitized["api_key"] == REDACTED + assert sanitized["headers"]["Authorization"] == REDACTED + assert sanitized["headers"]["Cookie"] == REDACTED + assert sanitized["headers"]["normal"] == "token in normal text" + assert sanitized["items"][0]["refresh_token"] == REDACTED + assert sanitize_for_trace({"reasoning_content": "provider-state"})["reasoning_content"] == REDACTED + assert sanitize_for_trace({"thoughtSignature": "provider-state"})["thoughtSignature"] == REDACTED + assert sanitize_for_trace({"metadata": {"state": "provider-state"}})["metadata"]["state"] == REDACTED + assert sanitize_for_trace({"metadata": {"prompt_cache_key": "cache-state"}})["metadata"]["prompt_cache_key"] == REDACTED + assert sanitized["items"][0]["text"] == "token should remain in value" + + +def test_scrub_sensitive_text_targets_header_like_fragments_only() -> None: + text = "normal token text remains\nAuthorization: Bearer abc123\nset-cookie: sid=secret\n{'Authorization': 'Bearer quoted'}" + + scrubbed = scrub_sensitive_text(text) + + assert "normal token text remains" in scrubbed + assert "Authorization: [REDACTED]" in scrubbed + assert "set-cookie: [REDACTED]" in scrubbed + assert "quoted" not in scrubbed + + +def test_sanitize_for_trace_extracts_sdk_like_objects_before_repr() -> None: + class SdkError: + def __init__(self) -> None: + self.status_code = 401 + self.headers = {"Authorization": "Bearer secret"} + + def __repr__(self) -> str: + return "SdkError(Authorization: Bearer leaked)" + + sanitized = sanitize_for_trace(SdkError(), scrub_strings=True) + + assert sanitized["status_code"] == 401 + assert sanitized["headers"]["Authorization"] == REDACTED + assert "leaked" not in json.dumps(sanitized) + + +def test_sanitize_for_trace_serializes_common_non_json_values() -> None: + payload = ExamplePayload(when=datetime(2026, 1, 2, 3, 4, 5), amount=Decimal("1.5")) + + sanitized = sanitize_for_trace({"payload": payload, "bytes": b"hello"}) + + assert sanitized["payload"]["when"] == "2026-01-02T03:04:05" + assert sanitized["payload"]["amount"] == 1.5 + assert sanitized["bytes"] == "hello" + json.dumps(sanitized) + + +def test_sanitize_filename_is_stable_and_filesystem_safe() -> None: + assert sanitize_filename("Raw Client Request") == "raw_client_request" + assert sanitize_filename("provider/request:payload") == "provider_request_payload" + + +def test_transform_trace_writer_records_jsonl_and_snapshots(tmp_path) -> None: + writer = TransformTraceWriter(tmp_path, component="client", provider="openai", model="gpt-test", request_id="req_1") + + first = writer.record("raw_client_request", {"model": "gpt-test"}, direction="request", stage="client") + second = writer.record("parsed_stream_chunk", {"delta": "hi"}, direction="stream", stage="client", snapshot=True) + + assert first is not None + assert second is not None + assert first.sequence == 1 + assert second.sequence == 2 + + lines = (tmp_path / "transform_trace.jsonl").read_text(encoding="utf-8").splitlines() + assert len(lines) == 2 + assert json.loads(lines[0])["pass_name"] == "raw_client_request" + assert json.loads(lines[0])["request_id"] == "req_1" + assert json.loads(lines[1])["direction"] == "stream" + assert (tmp_path / "transforms" / "0001_raw_client_request.json").exists() + assert not (tmp_path / "transforms" / "0002_parsed_stream_chunk.json").exists() + + +def test_transform_trace_writer_disabled_writes_nothing(tmp_path) -> None: + writer = TransformTraceWriter(tmp_path, component="client", enabled=False) + + assert writer.record("raw_client_request", {}, direction="request", stage="client") is None + assert not (tmp_path / "transform_trace.jsonl").exists() + + +def test_transform_trace_writer_snapshot_namespace_prevents_collisions(tmp_path) -> None: + first = TransformTraceWriter(tmp_path, component="provider", snapshot_namespace="provider_a") + second = TransformTraceWriter(tmp_path, component="provider", snapshot_namespace="provider_b") + + first.record("provider_request_payload", {"a": 1}, direction="request", stage="provider") + second.record("provider_request_payload", {"b": 2}, direction="request", stage="provider") + + assert (tmp_path / "transforms" / "0001_provider_a_provider_request_payload.json").exists() + assert (tmp_path / "transforms" / "0001_provider_b_provider_request_payload.json").exists() diff --git a/tests/test_usage_accounting.py b/tests/test_usage_accounting.py new file mode 100644 index 000000000..3c8e86272 --- /dev/null +++ b/tests/test_usage_accounting.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from rotator_library.protocols.types import CostDetails, Usage +from rotator_library.usage.accounting import UsageRecord, extract_usage_record + + +def test_openai_dict_usage_extracts_cache_and_reasoning_without_double_counting() -> None: + record = extract_usage_record( + { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 30, + "total_tokens": 130, + "prompt_tokens_details": {"cached_tokens": 40, "cache_creation_tokens": 5}, + "completion_tokens_details": {"reasoning_tokens": 10}, + } + }, + provider="openai", + model="gpt-test", + ) + + assert record.input_tokens == 55 + assert record.cache_read_tokens == 40 + assert record.cache_write_tokens == 5 + assert record.completion_tokens == 20 + assert record.reasoning_tokens == 10 + assert record.output_tokens == 30 + assert record.raw_total_tokens == 130 + assert record.total_tokens == 130 + + +def test_openai_usage_extracts_provider_reported_cost_details() -> None: + record = extract_usage_record( + { + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "cost_details": {"total_cost": 0.0123, "currency": "EUR", "source": "provider_usage"}, + } + } + ) + + assert record.provider_reported_cost == 0.0123 + assert record.cost_currency == "EUR" + assert record.cost_source == "provider_usage" + + +def test_top_level_cost_is_preserved_when_usage_exists() -> None: + record = extract_usage_record( + { + "usage": {"prompt_tokens": 10, "completion_tokens": 5}, + "total_cost": 0.03, + "currency": "EUR", + } + ) + + assert record.provider_reported_cost == 0.03 + assert record.cost_currency == "EUR" + + +def test_reference_request_cost_usd_is_preserved() -> None: + record = extract_usage_record({"usage": {"prompt_tokens": 1, "completion_tokens": 1}, "request_cost_usd": 0.019}) + + assert record.provider_reported_cost == 0.019 + + +def test_top_level_estimated_cost_is_preserved() -> None: + record = extract_usage_record({"usage": {"prompt_tokens": 1, "completion_tokens": 1}, "estimated_cost": 0.027}) + + assert record.provider_reported_cost == 0.027 + + +def test_structured_cost_breakdown_without_total_is_summed() -> None: + record = extract_usage_record( + { + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "cost_details": {"cached_input_cost": 0.01, "upstream_inference_cost": 0.02, "web_search_cost": 0.003}, + } + } + ) + + assert record.provider_reported_cost == 0.033 + assert record.cost_source == "provider_reported_breakdown" + + +def test_reference_extended_cost_breakdown_aliases_are_summed() -> None: + record = extract_usage_record( + { + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "cost_details": { + "upstream_inference_input_cost": 0.01, + "upstream_inference_output_cost": 0.02, + "image_input_cost": 0.003, + "audio_input_cost": 0.004, + "data_storage_cost": 0.005, + "estimated_cost": 0.006, + "cost_in_usd_ticks": 70_000_000, + }, + } + } + ) + + assert record.provider_reported_cost == 0.055 + + +def test_upstream_inference_total_wins_over_split_fields() -> None: + record = extract_usage_record( + { + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "cost_details": { + "upstream_inference_cost": 0.04, + "upstream_inference_prompt_cost": 0.01, + "upstream_inference_completions_cost": 0.03, + }, + } + } + ) + + assert record.provider_reported_cost == 0.04 + + +def test_openai_object_usage_extracts_attributes() -> None: + response = SimpleNamespace( + usage=SimpleNamespace( + prompt_tokens=12, + completion_tokens=7, + prompt_tokens_details=SimpleNamespace(cached_tokens=2, cache_creation_tokens=1), + completion_tokens_details=SimpleNamespace(reasoning_tokens=3), + ) + ) + + record = extract_usage_record(response) + + assert record.input_tokens == 9 + assert record.cache_read_tokens == 2 + assert record.cache_write_tokens == 1 + assert record.completion_tokens == 4 + assert record.reasoning_tokens == 3 + + +def test_anthropic_usage_extracts_cache_buckets() -> None: + record = extract_usage_record( + { + "input_tokens": 100, + "output_tokens": 20, + "cache_creation_input_tokens": 8, + "cache_read_input_tokens": 12, + } + ) + + assert record.input_tokens == 80 + assert record.cache_read_tokens == 12 + assert record.cache_write_tokens == 8 + assert record.completion_tokens == 20 + assert record.metadata["shape"] == "anthropic" + + +def test_gemini_usage_metadata_extracts_thought_and_cached_tokens() -> None: + record = extract_usage_record( + { + "usageMetadata": { + "promptTokenCount": 80, + "candidatesTokenCount": 25, + "thoughtsTokenCount": 5, + "cachedContentTokenCount": 30, + "totalTokenCount": 110, + } + } + ) + + assert record.input_tokens == 50 + assert record.cache_read_tokens == 30 + assert record.completion_tokens == 20 + assert record.reasoning_tokens == 5 + assert record.metadata["shape"] == "gemini" + + +def test_gemini_usage_extracts_cost_metadata() -> None: + record = extract_usage_record( + { + "usageMetadata": { + "promptTokenCount": 1, + "totalTokenCount": 1, + "costMetadata": {"cost": "0.004", "currency": "USD", "source": "gemini_usage"}, + } + } + ) + + assert record.provider_reported_cost == 0.004 + assert record.cost_source == "gemini_usage" + + +def test_protocol_usage_cost_details_are_preserved() -> None: + record = extract_usage_record( + Usage(input_tokens=2, output_tokens=3, cost=CostDetails(provider_reported_cost=0.02, currency="GBP", source="protocol_usage")) + ) + + assert record.input_tokens == 2 + assert record.completion_tokens == 3 + assert record.provider_reported_cost == 0.02 + assert record.cost_currency == "GBP" + assert record.cost_source == "protocol_usage" + + +def test_responses_usage_extracts_nested_details() -> None: + record = extract_usage_record( + { + "input_tokens": 42, + "output_tokens": 13, + "input_tokens_details": {"cached_tokens": 10}, + "output_tokens_details": {"reasoning_tokens": 4}, + }, + source="responses", + ) + + assert record.input_tokens == 32 + assert record.cache_read_tokens == 10 + assert record.completion_tokens == 9 + assert record.reasoning_tokens == 4 + + +def test_unknown_usage_shape_returns_empty_record() -> None: + record = extract_usage_record({"not_usage": True}, provider="x", model="y") + + assert record == UsageRecord(provider="x", model="y", source="response", metadata={"shape": "openai_like"}) diff --git a/tests/test_usage_costs.py b/tests/test_usage_costs.py new file mode 100644 index 000000000..0ffe6ec5d --- /dev/null +++ b/tests/test_usage_costs.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from rotator_library.usage.accounting import UsageRecord +from rotator_library.usage.costs import CostCalculator, ModelPricing + + +class PricingProvider: + def get_model_pricing(self, model: str): + return ModelPricing( + input_cost_per_token=0.001, + cache_read_cost_per_token=0.0001, + cache_write_cost_per_token=0.0005, + output_cost_per_token=0.002, + reasoning_cost_per_token=0.003, + source="provider_test", + ) + + +class SkipCostProvider: + skip_cost_calculation = True + + +def test_explicit_pricing_calculates_all_usage_buckets() -> None: + usage = UsageRecord(input_tokens=10, cache_read_tokens=20, cache_write_tokens=5, completion_tokens=7, reasoning_tokens=3) + + cost = CostCalculator(provider_plugin=PricingProvider(), use_litellm_fallback=False).calculate(usage, model="test") + + assert cost.input_cost == 0.01 + assert cost.cache_read_cost == 0.002 + assert cost.cache_write_cost == 0.0025 + assert cost.output_cost == 0.014 + assert cost.reasoning_cost == 0.009000000000000001 + assert cost.pricing_source == "provider_test" + + +def test_skip_cost_provider_returns_zero_skipped_breakdown() -> None: + cost = CostCalculator(provider_plugin=SkipCostProvider()).calculate(UsageRecord(input_tokens=100), model="test") + + assert cost.total_cost == 0.0 + assert cost.pricing_source == "skipped" + + +def test_provider_reported_cost_wins_over_advisory_pricing() -> None: + usage = UsageRecord(input_tokens=10, completion_tokens=10, provider_reported_cost=0.123, cost_currency="EUR", cost_source="provider_actual") + + cost = CostCalculator(provider_plugin=PricingProvider(), use_litellm_fallback=False).calculate(usage, model="test") + + assert cost.total_cost == 0.123 + assert cost.provider_reported_cost == 0.123 + assert cost.currency == "EUR" + assert cost.pricing_source == "provider_actual" + assert cost.input_cost == 0.0 + + +def test_skip_cost_still_wins_over_provider_reported_cost() -> None: + usage = UsageRecord(provider_reported_cost=0.123) + + cost = CostCalculator(provider_plugin=SkipCostProvider()).calculate(usage, model="test") + + assert cost.total_cost == 0.0 + assert cost.pricing_source == "skipped" + + +def test_missing_pricing_returns_unavailable_zero() -> None: + cost = CostCalculator(use_litellm_fallback=False).calculate(UsageRecord(input_tokens=100), model="unknown-model") + + assert cost.total_cost == 0.0 + assert cost.pricing_source == "unavailable" + + +def test_litellm_model_info_fallback(monkeypatch) -> None: + monkeypatch.setattr( + "rotator_library.usage.costs.litellm.get_model_info", + lambda model: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + + cost = CostCalculator().calculate(UsageRecord(input_tokens=2, completion_tokens=3), model="gpt-test") + + assert cost.total_cost == 0.008 + assert cost.pricing_source == "litellm_model_info" diff --git a/tests/test_usage_quota_snapshots.py b/tests/test_usage_quota_snapshots.py new file mode 100644 index 000000000..f110ac4ef --- /dev/null +++ b/tests/test_usage_quota_snapshots.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from rotator_library.usage.quota import build_quota_snapshots +from rotator_library.usage.types import CredentialState, WindowStats + + +def _state() -> CredentialState: + state = CredentialState(stable_id="credential-secret", provider="openai", accessor="credential-secret") + model_stats = state.get_model_stats("gpt-test") + model_stats.windows["daily"] = WindowStats(name="daily", request_count=3, limit=10, reset_at=123.0) + group_stats = state.get_group_stats("chat") + group_stats.windows["daily"] = WindowStats(name="daily", request_count=5, limit=20, reset_at=456.0) + return state + + +def test_build_quota_snapshots_for_model_window() -> None: + snapshots = build_quota_snapshots(provider="openai", states={"credential-secret": _state()}, model="gpt-test") + + assert len(snapshots) == 1 + snapshot = snapshots[0] + assert snapshot.provider == "openai" + assert snapshot.model == "gpt-test" + assert snapshot.quota_group is None + assert snapshot.used == 3 + assert snapshot.remaining == 7 + assert snapshot.credential_id != "credential-secret" + assert snapshot.metadata == {"scope": "request_token_window"} + assert "cost_used" not in snapshot.to_dict() + + +def test_build_quota_snapshots_for_group_window_without_credentials() -> None: + snapshots = build_quota_snapshots( + provider="openai", + states={"credential-secret": _state()}, + model="gpt-test", + quota_group="chat", + include_credentials=False, + ) + + group_snapshot = [snapshot for snapshot in snapshots if snapshot.source == "group"][0] + assert group_snapshot.quota_group == "chat" + assert group_snapshot.used == 5 + assert group_snapshot.remaining == 15 + assert group_snapshot.credential_id is None + + +def test_build_quota_snapshots_missing_windows_returns_empty() -> None: + assert build_quota_snapshots(provider="openai", states={"credential-secret": _state()}, model="missing") == []