diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 00000000..ea687980 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,61 @@ +# LLM Proxy + +This context describes the language used to discuss interactive model requests +and their performance through routatic-proxy. + +## Language + +**Time to First Token (TTFT)**: +For a streaming request, the elapsed time from when the proxy starts reading +the request until it writes the first non-empty text, thinking, or tool content +to the client. +_Avoid_: First-byte time, total response time + +**Repeated Conversation**: +A model request that includes a stable prefix from earlier turns plus a new +turn. +_Avoid_: Warm request, duplicate request + +**Cache Directive**: +A client-supplied instruction that marks a prompt boundary as eligible for +reuse by the Provider Prompt Cache. +_Avoid_: Automatic cache rule, proxy cache marker + +**Cache Usage**: +Provider-reported token counts describing prompt data read from or written to +the Provider Prompt Cache. +_Avoid_: Token Count Cache hit, inferred cache use + +**Provider Prompt Cache**: +An upstream provider feature that reuses an unchanged prompt prefix across +model requests. +_Avoid_: Token Count Cache, response cache + +**Token Count Cache**: +A process-local cache that reuses tokenizer results without storing raw prompt +text. +_Avoid_: Provider Prompt Cache, response cache + +**Completion Record**: +Best-effort analytics data describing one finished model request. Losing this +record must not change the model response seen by the client. +_Avoid_: Response, durable event + +**Catalog Snapshot**: +The last successfully loaded set of providers, models, and routing scenarios +used for model selection. +_Avoid_: Live catalog, catalog request + +**Known Model**: +A model present in the active configuration or catalog snapshot. +_Avoid_: Requested model, arbitrary model name + +**Model Identity**: +The canonical `provider/model` name used to distinguish a model across +providers. +_Avoid_: Short model name, display name + +**Unknown Content Block**: +A client content block whose type is not yet modeled by the proxy but whose raw +data and position remain intact. +_Avoid_: Invalid block, ignored block diff --git a/docs/adr/0001-user-responses-before-analytics.md b/docs/adr/0001-user-responses-before-analytics.md new file mode 100644 index 00000000..af4afc5b --- /dev/null +++ b/docs/adr/0001-user-responses-before-analytics.md @@ -0,0 +1,7 @@ +# User responses take priority over analytics completeness + +Completion records are best-effort analytics. The proxy never delays a model +response when the storage queue is full, so it drops the newest analytics +record and reports the drop. During normal shutdown it drains accepted records +until the existing shutdown deadline, then reports any remaining loss. This +trades complete analytics for predictable user-facing latency. diff --git a/docs/adr/0002-client-owned-cache-directives.md b/docs/adr/0002-client-owned-cache-directives.md new file mode 100644 index 00000000..f158555d --- /dev/null +++ b/docs/adr/0002-client-owned-cache-directives.md @@ -0,0 +1,6 @@ +# Provider Prompt Cache directives are owned by the client + +The proxy preserves cache directives supplied at the request, tool, system, and +message-block levels, but it does not add, move, or infer cache boundaries. +Provider adapters may omit a directive only when their wire format cannot use +it. This favors predictable request behavior over automatic cache tuning. diff --git a/docs/adr/0003-ordered-normalized-content.md b/docs/adr/0003-ordered-normalized-content.md new file mode 100644 index 00000000..34b55d24 --- /dev/null +++ b/docs/adr/0003-ordered-normalized-content.md @@ -0,0 +1,8 @@ +# Normalized content keeps client block order + +The canonical request model uses one ordered list of content blocks for system +and message content. Text, images, thinking, tool calls, tool results, and cache +directives stay at their original positions. This requires a larger core +refactor, but it prevents separate convenience fields from changing request +order or losing cache boundaries. Unknown block types remain as ordered raw +JSON for compatible provider formats and are never silently dropped. diff --git a/docs/adr/0004-last-valid-catalog-remains-usable.md b/docs/adr/0004-last-valid-catalog-remains-usable.md new file mode 100644 index 00000000..35190795 --- /dev/null +++ b/docs/adr/0004-last-valid-catalog-remains-usable.md @@ -0,0 +1,7 @@ +# The last valid catalog remains usable + +After the proxy loads a valid catalog, refresh failures never make that catalog +unusable solely because of age. The proxy reports snapshot age and refresh +errors, but keeps routing with the last valid data. Legacy routing is used only +when no valid catalog has ever loaded, favoring request availability over +catalog freshness. diff --git a/docs/cache-and-latency-implementation-plan.md b/docs/cache-and-latency-implementation-plan.md new file mode 100644 index 00000000..8f3cd68a --- /dev/null +++ b/docs/cache-and-latency-implementation-plan.md @@ -0,0 +1,757 @@ +# Cache Hit Rate and Request Latency Implementation Plan + +Status: Proposed + +Created: 2026-08-25 + +Branch: `perf/cache-hit-and-request-latency-plan` + +Base: `main` at `d74b160` + +## Objective + +Improve repeated-request efficiency and reduce proxy-added latency without +changing routing decisions, provider-visible request semantics, or response +correctness. + +The primary success measure is lower p95 time to first token for repeated +streaming conversations. Proxy-added latency and Provider Prompt Cache read +tokens are supporting measures. Request correctness is a hard requirement. +Non-streaming requests use total response time and are never included in TTFT +results. + +Use a controlled repeated-conversation benchmark as the release gate. The +proxy has no reliable conversation identifier, so production metrics must not +guess that separate requests belong to one conversation. Production reports +show TTFT and Token Count Cache reuse as separate measurements. + +The work covers five areas: + +1. Preserve Provider Prompt Cache directives and cache-usage telemetry. +2. Cache repeated token-count results. +3. Remove SQLite telemetry writes from the response critical path. +4. Correct and extend performance measurements. +5. Make catalog refresh and cost-based routing cheaper. + +General LLM response caching and in-flight request coalescing are explicitly +out of scope. Responses can be non-deterministic, streamed, tool-bearing, and +authorization-sensitive. This program does not store or share model responses. + +## Delivery Principles + +- Measure before and after every optimization. +- Change one bottleneck per performance commit. +- Keep caches bounded by an explicit capacity suited to the data they retain. +- Publish immutable data to request handlers. +- Treat analytics persistence as best-effort telemetry, not part of response + correctness. +- Preserve stale-but-valid catalog data when a refresh fails. +- Keep provider-specific cache behavior behind explicit capability decisions. +- Include benchmark evidence in performance commit messages. + +## Success Criteria + +The implementation is complete when: + +- p95 time to first token improves for repeated streaming conversations. +- Anthropic request, tool, system-block, and message-block cache directives + survive normalization and reach supported upstream wire formats unchanged. +- Unsupported wire formats still omit cache directives intentionally. +- Cache read and cache creation token counts are visible in in-memory metrics, + request history, and SQLite history when the upstream supplies them. +- Repeated conversation history produces Token Count Cache hits. +- The Token Count Cache has bounded memory, exposes hits, misses, and evictions, + and remains race-free. +- A slow or locked SQLite database cannot delay a successful HTTP response. +- Storage queue overflow and shutdown-drain failures are observable. +- p95 and p99 calculations operate on sorted samples and are covered by tests. +- Catalog cache hits require no mutex acquisition. +- Catalog refresh does not make concurrent requests wait for disk or SQLite. +- Cost selection uses the provider index and a one-pass minimum rather than + scanning the entire catalog once per provider and sorting all candidates. +- Each optimized path shows a statistically significant improvement under its + targeted benchmark. +- Cold and unrelated paths do not regress by more than 5%. +- `go test ./... -race`, `make lint`, and `make lint-strict` pass. + +## Measurement Foundation + +This phase is delivered first because the current latency percentile methods +index samples in arrival order rather than sorted order. + +### Benchmarks + +Add focused benchmarks using Go 1.25 `b.Loop()`: + +- `BenchmarkCountMessages` + - cold cache + - warm cache + - 10, 50, and 200 message histories + - growing conversation where one message is appended per turn + - repeated system prompt + - parallel callers +- `BenchmarkMetricsRecordSuccess` + - empty buffer + - full buffer + - parallel writers + - concurrent snapshot readers +- `BenchmarkSelectorSelectCheapest` + - small, medium, and large catalogs + - one and several eligible providers + - restrictive and permissive constraints +- `BenchmarkModelRouterCatalogHit` + - fresh snapshot + - current snapshot while a background refresh is active +- `BenchmarkAsyncStorageWriter` + - enqueue only + - batch sizes + - queue saturation + - graceful drain +- request-path integration benchmark using a fake provider with a fixed response + and no external network. +- streaming repeated-conversation integration benchmark that appends one new + turn per request and reports p95 TTFT + +Run benchmarks serially and retain reports outside the repository: + +```bash +GOCACHE=/tmp/routatic-perf-go-cache \ + go test -run '^$' -bench=. -benchmem -count=10 ./internal/token ./internal/metrics ./internal/router ./internal/handlers \ + | tee /tmp/routatic-perf-before.txt + +GOCACHE=/tmp/routatic-perf-go-cache \ + go test -run '^$' -bench=. -benchmem -count=10 ./internal/token ./internal/metrics ./internal/router ./internal/handlers \ + | tee /tmp/routatic-perf-after.txt + +benchstat /tmp/routatic-perf-before.txt /tmp/routatic-perf-after.txt +``` + +Do not claim an improvement from a single run or from statistically +insignificant output. + +### Request-stage timing + +Record separate durations for: + +- body read and JSON parsing +- message extraction +- token counting +- request-fact analysis and routing +- provider request transformation +- upstream time to first non-empty model content +- upstream total duration +- response transformation +- storage enqueue +- total proxy duration + +Measure TTFT from the moment the proxy begins reading the request until it +writes the first non-empty text, thinking, or tool content to the client. +Response headers, empty SSE events, and metadata-only events do not complete the +TTFT measurement. Record TTFT only for streaming requests. Record total response +time for non-streaming requests. + +Use monotonic `time.Time` values already carried by Go timestamps. Avoid logging +every stage per request at info level; aggregate measurements in +`internal/metrics`. + +## Workstream 1: Preserve Provider Prompt Cache Directives + +### Current problem + +Anthropic requests may carry cache directives at the top request level and on +tools, system blocks, and message content blocks. The current request types +only model system-block directives. `core.NormalizeRequest` then flattens the +system field and message content, so the provider registry path cannot preserve +the complete client request. + +### Design + +Use one ordered normalized content-block model for system and message content: + +```go +type NormalizedCacheControl struct { + Type string +} + +type NormalizedContentBlock struct { + Type NormalizedContentType + Text string + Image *NormalizedImage + ToolCall *NormalizedToolCall + ToolResult *NormalizedToolResult + Thinking string + CacheControl *NormalizedCacheControl + Raw json.RawMessage +} +``` + +`NormalizedRequest` stores ordered system blocks. `NormalizedMessage` stores +ordered message blocks. Replace the current separate text, image, thinking, +tool-call, and tool-result fields rather than retaining two representations +that can diverge. + +Add read-only helpers such as `SystemText()` and `MessageText()` for routing and +token counting. Provider transformers iterate the ordered blocks directly so +they preserve the client's content order. + +Extend the provider interface with a pure request-compatibility check. Each +provider adapter validates the normalized request against the selected wire +format before any network call. Return a typed compatibility error containing +the unsupported block or feature. The router does not contain wire-format +rules. + +Normalization rules: + +- A top-level cache directive remains attached to the normalized request. +- A string system prompt becomes one cacheless normalized block. +- An array remains an ordered set of blocks. +- `cache_control.type` is copied without provider interpretation. +- Tool cache directives remain attached to their tool definitions. +- Message cache directives remain attached to the ordered content blocks where + the client placed them. +- An unknown content-block type remains in order as raw JSON. + +Denormalization rules: + +- Emit a JSON string for one cacheless text block to preserve the common wire + shape. +- Emit an ordered block array when any cache directive is present or multiple + blocks must be retained. +- Anthropic-format providers receive the block array unchanged. +- OpenAI Chat transformations retain the existing DeepSeek support and existing + stripping behavior for unsupported models. +- Responses and Gemini transformations omit the directive until their provider + contracts explicitly support an equivalent. +- Anthropic-format providers may round-trip an unknown raw content block. +- Other formats mark the block as unsupported before making an upstream call. + The fallback handler then tries only models whose provider format can + preserve the block. +- If no compatible fallback exists, return a clear client error naming the + unsupported block type. +- Compatibility failures do not count as provider failures and do not affect + circuit breakers. +- When a selected provider format cannot use a cache directive, remove the + directive and continue the request. Increment a bounded-cardinality metric + by provider and wire format, and emit a sampled debug log. Never fail a + request only because its cache directive is unsupported. + +The provider or wire-format boundary owns the support decision. Core +normalization only preserves information. + +Cache boundaries remain client-owned. The proxy must not add, move, or infer a +cache directive. It only preserves directives already present in the incoming +request and forwards them when the selected provider format supports them. + +### Cache-usage telemetry + +Extend `history.RequestRecord` with: + +- `CacheReadTokens` +- `CacheCreationTokens` +- `CacheUsageReported` + +Add nullable cache-token columns and a `cache_usage_reported` column through an +idempotent SQLite migration. When `cache_usage_reported` is false, persist the +token columns as null. When it is true, zero is a real provider-reported value. +Populate the fields from both streaming and non-streaming responses. + +Extend the Responses usage type only after verifying the actual upstream +payload shape with provider fixtures. Do not infer a cached-token JSON field. + +Expose raw cache counters per provider and model. Avoid a universal "hit rate" +formula until the provider-specific token accounting denominator is defined. +Do not include records with unreported cache usage in provider cache totals. + +### Tests + +- Normalize string system prompts. +- Normalize multiple system blocks while retaining order. +- Round-trip mixed text, image, thinking, tool-call, and tool-result blocks + without reordering them. +- Round-trip unknown content blocks as raw JSON through Anthropic formats. +- Round-trip top-level, tool, system-block, and message-block cache directives. +- Verify the Zen Anthropic request body retains cache directives. +- Verify the DeepSeek Chat request retains supported directives. +- Verify unsupported Chat, Responses, and Gemini bodies omit them. +- Verify unsupported directives increment the omission metric without changing + response behavior. +- Verify an unsupported content block skips incompatible models, uses a + compatible fallback, and does not affect circuit breakers. +- Verify no compatible fallback returns a clear client error containing the + block type. +- Verify provider compatibility validation performs no network or mutable + provider-state work. +- Verify stream and non-stream usage populate history and storage. +- Verify reported zero remains distinct from unreported cache usage. +- Verify migrations work on both new and existing databases. +- Add golden provider-body fixtures where practical. + +### Live provider verification + +Before releasing Provider Prompt Cache support for a provider format: + +- run an opt-in test outside CI with user-supplied credentials +- send the same stable prompt prefix twice +- change only the latest turn +- confirm the second response reports provider cache reuse +- record provider, model, wire format, TTFT, cache-read tokens, and + cache-creation tokens +- never persist the prompt text in the test report + +Unit and golden-body tests prove request fidelity but do not prove that an +upstream provider actually reuses the prefix. + +If a provider accepts the request but does not expose observable cache-read +data, mark that cache path as unverified or experimental. Do not advertise a +cache-hit improvement without provider-reported evidence. + +### Acceptance + +- No cache directive disappears before provider capability handling. +- Existing provider stripping tests continue to pass. +- Requests without cache directives keep their existing common wire shape. +- Every advertised provider cache path has a successful live verification + result. Unverified paths are labeled experimental and make no cache-hit claim. + +## Workstream 2: Cache Repeated Token Counts + +### Current problem + +Every request tokenizes the system prompt and every message again. In an +interactive session, most previous message text is identical to the prior turn. + +### Design + +Add a bounded cache owned by `token.Counter`. + +Each cache entry represents one system block or one message. Do not cache a +whole conversation as one entry because adding a new turn would invalidate the +entire key and prevent reuse of the stable history. + +Initial implementation: + +- one concurrency-safe LRU protected by one lock +- fixed-size SHA-256 keys built from the tokenizer name and input text +- no retained copy of the raw prompt text +- maximum of 8,192 entries by default +- skip entries below a measured minimum string length +- no unbounded `sync.Map` +- no `sync.Pool` unless a profile identifies allocation churn it can solve + +The count remains deterministic and encoding-specific. Including the tokenizer +name in the fingerprint prevents a future model-specific tokenizer from +reusing an incompatible count. A cryptographic hash collision is treated as +negligible; the cache must not keep raw text only to check collisions. + +Add explicit configuration: + +```json +{ + "performance": { + "token_cache_enabled": true, + "token_cache_max_entries": 8192 + } +} +``` + +Defaults must be safe for desktop use. A disabled cache must preserve current +behavior exactly. The cache is enabled by default only after its benchmark and +race-test gates pass. Keep the disable switch for troubleshooting. + +When cache settings change during config reload, build a new empty cache and +publish it atomically. Do not resize or mutate the active cache in place. Old +entries are discarded. + +Do not add miss coalescing initially. Add it only if a parallel benchmark shows +that concurrent identical misses are common enough to outweigh coordination +cost. Do not shard the LRU initially. Add shards only if the parallel benchmark +shows lock contention. + +### Metrics + +Record: + +- hits +- misses +- evictions +- skipped-small-input counts +- current entries +- tokenization duration + +### Tests + +- deterministic hit after first count +- distinct strings and encodings produce distinct fingerprints +- entry limit evicts +- disabled cache bypasses storage +- small-entry policy works +- parallel race test +- large input does not cause unbounded retention or retain raw prompt text +- cached and uncached `CountMessages` return identical totals + +### Acceptance + +- Warm repeated-history benchmarks improve significantly. +- Cold-cache performance has no meaningful regression. +- Memory remains within configured bounds under adversarial unique input. + +## Workstream 3: Move SQLite Writes Off the Response Path + +### Current problem + +Successful requests execute separate request and latency inserts synchronously. +SQLite is configured with one open connection, so concurrent request +completions serialize before non-streaming response bodies are written. + +### Design + +Use `requests.duration_ms` as the single source for latency reports. Stop +writing new rows to `latency_samples`, and update latency queries to read from +`requests`. This removes the duplicate write before asynchronous batching is +introduced. Leave the old table and its existing rows untouched in this branch +for rollback safety. A later release may remove the table after the new reports +have been proven stable. + +Replace the two-method handler-facing storage interface with one completion +operation: + +```go +type CompletionRecorder interface { + RecordCompletion(history.RequestRecord) + Shutdown(context.Context) error +} +``` + +`RecordCompletion` enqueues without waiting for SQLite. A dedicated writer: + +- owns one bounded channel +- batches by maximum count or short flush interval +- writes request rows in one transaction +- preserves record order within each batch +- keeps the existing single SQLite writer connection +- emits counters for enqueued, persisted, dropped, failed, queue depth, batch + size, and drain duration +- samples repeated error logs + +Do not add application-level write retries. SQLite already applies its +configured busy timeout for temporary lock contention. If a batch still fails, +count and drop it, emit a sampled error log, and continue with the next batch. + +Queue capacity, batch size, and flush interval are internal constants selected +by benchmarks. Do not add user-facing settings for them in this work. Expose +queue depth, drops, batch size, failures, and drain time so later production +evidence can justify configuration if needed. + +Because this data is analytics telemetry, queue saturation should not block a +user response. When the queue is full, drop the newest analytics record and +increment an explicit counter. Older already-accepted records remain ordered. +Expose the drop count through metrics and show a dashboard warning when it is +non-zero. + +Keep the in-memory history update synchronous because it is O(1) and supplies +the live dashboard immediately. + +Enable async storage by default after its test and benchmark gates pass. Keep a +temporary setting that switches back to the current synchronous writer for one +stable release. Remove the setting and synchronous path after that release if +no rollback is needed. + +### Shutdown lifecycle + +Change server shutdown ordering: + +1. Stop accepting new HTTP requests and wait for active handlers. +2. Close the completion recorder to new entries. +3. Drain accepted records within the caller's existing shutdown deadline. +4. Stop retention work. +5. Close SQLite. + +Use the same lifecycle for signal-based and programmatic shutdown. The current +paths must not close SQLite while request handlers can still enqueue work. If +the deadline expires before the queue drains, log and count the records that +were not persisted. + +### Tests + +- a blocking storage backend cannot delay an HTTP response +- latency reports use `requests.duration_ms` +- one completion produces one SQLite write +- existing `latency_samples` data and schema remain untouched +- batches flush by size and by interval +- queue saturation follows the documented policy +- persistence errors do not stop later batches +- failed batches are counted and are not retried +- shutdown drains accepted records +- shutdown deadline returns a clear error +- enqueue after shutdown is safe and observable +- race tests cover enqueue, flush, and shutdown + +### Acceptance + +- Handler storage-enqueue time remains bounded and independent of SQLite delay. +- No successful request fails because telemetry persistence fails. +- Accepted records drain on normal shutdown. + +## Workstream 4: Correct and Extend Performance Metrics + +### Correctness fixes + +- Replace latency slice shifting with fixed ring buffers holding the latest + 1,000 samples for each global metric. +- Calculate percentiles from a sorted copy. +- Sort once when calculating multiple percentiles. +- Copy samples while holding the lock, then release the lock before sorting. +- Use the same ring-buffer implementation with 200 samples for each known + `provider/model`. +- Report the sample count with every percentile. +- Report p50, p90, and p95 for per-model data. Do not report per-model p99 from + only 200 samples. +- Add table tests for empty, one-element, ordered, reverse-ordered, and repeated + samples. + +### New measurements + +Add counters and bounded timing samples for the request stages listed in +Measurement Foundation. + +Keep detailed stage timings in bounded memory only. Do not add SQLite columns +for stage breakdowns. Persistent request records continue to store total +duration, input and output tokens, and Provider Prompt Cache usage. + +Create per-model metric state only for models known by the active config or +catalog. Group arbitrary or unknown requested model names under `other` so +request input cannot grow metric maps without a bound. + +Key per-model metrics by canonical `provider/model` identity. Do not combine +the same model name across providers because their latency and cache behavior +may differ. + +For streaming requests: + +- record time to first non-empty model content separately from total stream + duration +- use `sync.Once` or equivalent so first-content timing is recorded exactly once +- distinguish headers and metadata events from real model content + +Do not add DNS, TLS, or connection-level `httptrace` instrumentation in this +program. Add it later only if request-stage timing shows connection setup is a +meaningful bottleneck. Do not alter HTTP pool sizes or enable a protocol based +only on intuition. + +### Exposure + +- Keep `/health` compact. +- Add detailed performance data to the existing metrics/dashboard API. +- Include Token Count Cache and storage queue state. +- Include raw provider cache-token counters. +- Document whether every duration includes or excludes upstream time. + +### Tests and benchmarks + +- percentile correctness independent of insertion order +- ring-buffer eviction order +- concurrent record/snapshot race tests +- unknown model names remain grouped under `other` +- identical model names on different providers remain separate +- first-SSE timing recorded once +- `RecordSuccess` full-buffer benchmark +- metrics snapshot benchmark with several models + +### Acceptance + +- Reported percentiles match a reference implementation. +- Metrics collection does not become a top allocation or lock-contention source. +- Proxy overhead and upstream latency can be distinguished. + +## Workstream 5: Speed Up Catalog and Cost-Based Routing + +### Current problem + +All catalog hits acquire one mutex. When the 30-second entry expires, the +request holding that mutex performs SQLite or file loading while concurrent +requests wait. Cost selection then scans the full model map for every eligible +provider and sorts every candidate to select one. + +### Catalog snapshot design + +Publish an immutable snapshot through `atomic.Pointer`: + +```go +type catalogSnapshot struct { + Catalog *catalog.IndexedCatalog + LoadedAt time.Time + Err error +} +``` + +Request behavior: + +- Every request returns the current snapshot with one atomic load. +- Requests never start refresh work, wait for refresh work, or check files or + SQLite for freshness. +- A successful refresh atomically publishes a new immutable snapshot. +- A failed refresh records the error and retains the last valid snapshot with + no maximum age. +- Startup performs one bounded synchronous load when a catalog source exists. +- One background loop refreshes every 30 seconds. +- Catalog update events signal the same loop to refresh early. + +Do not mutate an `IndexedCatalog` after publishing it. +Expose snapshot age and the latest refresh error. Use legacy routing only when +the process has never loaded a valid catalog. + +### Selector design + +Use `IndexedCatalog.ListProviderModels(provider)` or an equivalent precomputed +provider-keyed resolved-model index. Remove the nested full-catalog scan. + +Replace candidate collection and sorting with a one-pass best-candidate +comparison using the existing deterministic ordering: + +1. lower effective cost +2. larger context window +3. lexicographically smaller model ID + +Build enabled-provider state once per immutable configuration/catalog +generation rather than once per request. Register an `AtomicConfig.OnReload` +callback to publish a new selector state. + +Do not cache the final selected model. Token count, tools, images, reasoning +needs, provider state, and configuration can differ per request. Reuse immutable +indexes, then run the one-pass comparison for each request. + +Pass already-computed `RequestFacts` and constraints through routing rather than +re-running message scans and lowercasing. + +### Tests + +- fresh catalog hit performs no refresh +- requests do not trigger refresh or block during refresh +- background timer and catalog-update events trigger one refresh at a time +- callers continue using the current snapshot during refresh +- failed refresh preserves the last valid catalog +- an arbitrarily old valid catalog remains usable and reports its age +- first startup failure falls back to legacy config +- indexed selector matches current selector results +- one-pass tie breaking exactly matches current sort order +- config reload rebuilds enabled-provider state +- routing facts are computed once without changing scenario results +- race tests cover refresh, config reload, and selection + +### Acceptance + +- Catalog-hit benchmark has no request-path mutex contention. +- Selection scales with models belonging to eligible providers, not the full + catalog multiplied by provider count. +- Routing output is unchanged for the existing fixture suite. + +## Recommended Commit Sequence + +Keep the work reviewable and reversible: + +1. `test(perf): add request-path performance baselines` +2. `fix(metrics): correct latency percentiles and ring buffers` +3. `feat(metrics): record request stages and cache usage` +4. `refactor(core): preserve ordered normalized content blocks` +5. `feat(core): preserve all client cache directives` +6. `feat(storage): persist provider cache token usage` +7. `perf(token): cache repeated token counts` +8. `refactor(storage): remove duplicate latency writes` +9. `refactor(storage): add completion recorder` +10. `perf(storage): batch telemetry writes asynchronously` +11. `perf(router): publish immutable catalog snapshots` +12. `perf(router): use indexed one-pass model selection` +13. `perf(router): reuse analyzed request facts` +14. `docs(perf): record benchmark and rollout results` + +Run the affected unit and benchmark suites after every commit. Run the complete +verification suite before merging. + +## Pull Request Delivery + +Ship the work as one pull request with five ordered phases: + +1. **Metrics and benchmarks** + - benchmark foundations + - percentile correctness + - bounded metric rings + - TTFT and request-stage measurements +2. **Ordered content and Provider Prompt Cache fidelity** + - ordered normalized content blocks + - request, tool, system, and message cache directives + - cache-usage storage and metrics +3. **Token Count Cache** + - SHA-256 fingerprints + - 8,192-entry LRU + - cache metrics and configuration +4. **Async SQLite storage** + - read latency from `requests.duration_ms` + - stop writing duplicate latency samples + - bounded completion queue, batching, and shutdown drain +5. **Catalog and routing speed** + - background catalog refresh + - atomic immutable snapshots + - indexed one-pass selection + - one request-fact analysis pass + +Keep the commits small and in the recommended order. Each phase must pass its +focused tests and benchmarks before work moves to the next phase. The pull +request is ready to merge only when all five phases and the full verification +suite pass. + +Performance gates are specific to each phase: + +- the targeted benchmark must show a statistically significant improvement +- cold and unrelated paths must not regress by more than 5% +- correctness-only fixes may ship without claiming a speed improvement +- the complete pull request must improve repeated-conversation p95 TTFT + +## Rollout Strategy + +1. Build and verify metric correctness and stage timing first within the pull + request. +2. Observe a representative workload before enabling new optimizations by + default. +3. Enable Provider Prompt Cache preservation because it is a fidelity fix, + guarded by provider capability tests. +4. Enable the bounded Token Count Cache by default with 8,192 entries after + benchmark and race-test gates pass. Keep the disable switch. +5. Enable async persistence with queue-depth and dropped-record visibility. +6. Enable the routing snapshot and indexed selector after result-equivalence + tests pass. +7. Compare proxy overhead, TTFT, cache tokens, queue behavior, CPU, allocations, + and memory before and after. +8. Run the opt-in live Provider Prompt Cache checks before advertising provider + support. + +Keep the Token Count Cache disable switch for troubleshooting. Keep the async +storage rollback switch for one stable release, then remove it together with +the synchronous writer if no rollback is needed. + +## Risks and Mitigations + +| Risk | Mitigation | +| --- | --- | +| Cache metadata changes provider request shape | Golden provider-body tests and explicit capability handling | +| Token Count Cache retains prompt text | Store only fixed SHA-256 fingerprints, enforce the entry limit, and never persist entries | +| Token Count Cache lock becomes contended | Parallel benchmark first; shard only with evidence | +| Async storage drops analytics | Bounded queue, dropped counter, dashboard warning, graceful drain | +| Shutdown loses accepted records | One lifecycle owner and deadline-aware drain tests | +| Stale catalog persists after refresh failure | Expose snapshot age and refresh error; retain correctness-preserving legacy fallback | +| Selector optimization changes tie breaking | Differential tests against the existing implementation | +| Metrics create their own hot path | Bounded storage, ring buffers, and allocation benchmarks | +| Benchmark noise produces false wins | Ten serial runs and `benchstat`; retain hardware and Go version context | + +## Final Verification + +```bash +GOCACHE=/tmp/routatic-perf-go-cache go test ./... -count=1 +GOCACHE=/tmp/routatic-perf-go-cache go test ./... -count=1 -race +make lint +make lint-strict +git diff --check +``` + +Attach the final `benchstat` comparison and the request-path latency breakdown +to the pull request. Clearly distinguish local benchmark results from deployed +production observations. diff --git a/internal/config/config.go b/internal/config/config.go index 252875d1..a3d6c9ed 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -30,10 +30,17 @@ type Config struct { Logging LoggingConfig `json:"logging"` Debug DebugConfig `json:"debug"` Catalog CatalogConfig `json:"catalog"` + Performance PerformanceConfig `json:"performance,omitempty"` Storage *StorageConfig `json:"storage,omitempty"` UpdateChannel string `json:"update_channel,omitempty"` } +// PerformanceConfig controls bounded in-process latency optimizations. +type PerformanceConfig struct { + TokenCountCacheEnabled *bool `json:"token_count_cache_enabled,omitempty"` + TokenCountCacheCapacity int `json:"token_count_cache_capacity,omitempty"` +} + // CostRoutingConfig controls cost-aware model selection. type CostRoutingConfig struct { Enabled bool `json:"enabled"` diff --git a/internal/core/errors.go b/internal/core/errors.go index 70662a05..5796d20d 100644 --- a/internal/core/errors.go +++ b/internal/core/errors.go @@ -1,6 +1,9 @@ package core -import "errors" +import ( + "errors" + "fmt" +) // Sentinel errors for common provider and routing failures. var ( @@ -22,6 +25,21 @@ type NormalizedError struct { ModelID string } +// CompatibilityError means the selected provider/model cannot represent the +// request. It is safe to skip during fallback without recording a provider +// failure. +type CompatibilityError struct { + Provider string + ModelID string + Reason string +} + +func (e *CompatibilityError) Error() string { + return fmt.Sprintf("model %s/%s is incompatible with request: %s", e.Provider, e.ModelID, e.Reason) +} + +func (e *CompatibilityError) IsCompatibility() bool { return true } + // Error implements the error interface. func (e *NormalizedError) Error() string { return e.Message diff --git a/internal/core/normalize.go b/internal/core/normalize.go index e505e6c1..1f8d403d 100644 --- a/internal/core/normalize.go +++ b/internal/core/normalize.go @@ -17,13 +17,15 @@ type thinkingConfig struct { // This is a lossless extraction: all data from the Anthropic format survives. func NormalizeRequest(anthropicReq *types.MessageRequest) *NormalizedRequest { nr := &NormalizedRequest{ - Model: anthropicReq.Model, - MaxTokens: anthropicReq.MaxTokens, - Stream: anthropicReq.Stream != nil && *anthropicReq.Stream, + Model: anthropicReq.Model, + MaxTokens: anthropicReq.MaxTokens, + Stream: anthropicReq.Stream != nil && *anthropicReq.Stream, + CacheControl: anthropicReq.CacheControl, } // Extract system prompt (string or array of content blocks). nr.SystemPrompt = anthropicReq.SystemText() + nr.SystemBlocks = normalizeSystemBlocks(anthropicReq.System) // Set temperature if provided. if anthropicReq.Temperature != nil { @@ -47,35 +49,7 @@ func NormalizeRequest(anthropicReq *types.MessageRequest) *NormalizedRequest { blocks := msg.ContentBlocks() for _, block := range blocks { - switch block.Type { - case "text": - nm.Content += block.Text - case "tool_use": - nm.ToolCalls = append(nm.ToolCalls, NormalizedToolCall{ - ID: block.ID, - Name: block.Name, - Arguments: string(block.Input), - }) - case "tool_result": - nm.ToolResults = append(nm.ToolResults, NormalizedToolResult{ - ToolCallID: block.ToolUseID, - Content: block.TextContent(), - }) - case "thinking": - nm.Thinking += block.Thinking - case "image": - // Preserve image data so the downstream transformer can convert - // to image_url (or append a [Image] placeholder if the model - // does not support vision). Previously this was collapsed to - // the literal text "[Image]" which destroyed the image bytes - // before the transformer could inspect them. - if block.Source != nil && block.Source.Data != "" { - nm.Images = append(nm.Images, NormalizedImage{ - MediaType: block.Source.MediaType, - Data: block.Source.Data, - }) - } - } + nm.Blocks = append(nm.Blocks, normalizeContentBlock(block)) } nr.Messages = append(nr.Messages, nm) @@ -84,9 +58,10 @@ func NormalizeRequest(anthropicReq *types.MessageRequest) *NormalizedRequest { // Convert tools. for _, tool := range anthropicReq.Tools { nt := NormalizedToolDef{ - Name: tool.Name, - Description: tool.Description, - InputSchema: tool.InputSchema, + Name: tool.Name, + Description: tool.Description, + InputSchema: tool.InputSchema, + CacheControl: tool.CacheControl, } nr.Tools = append(nr.Tools, nt) } @@ -94,6 +69,56 @@ func NormalizeRequest(anthropicReq *types.MessageRequest) *NormalizedRequest { return nr } +func normalizeSystemBlocks(raw json.RawMessage) []NormalizedContentBlock { + if len(raw) == 0 { + return nil + } + var text string + if json.Unmarshal(raw, &text) == nil { + return []NormalizedContentBlock{{Type: "text", Text: text}} + } + var blocks []types.ContentBlock + if json.Unmarshal(raw, &blocks) != nil { + return nil + } + out := make([]NormalizedContentBlock, 0, len(blocks)) + for _, block := range blocks { + out = append(out, normalizeContentBlock(block)) + } + return out +} + +func normalizeContentBlock(block types.ContentBlock) NormalizedContentBlock { + content := block.Content + if len(content) == 0 && len(block.Output) > 0 { + content = block.Output + } + + return NormalizedContentBlock{ + Type: block.Type, + Text: block.Text, + ID: block.ID, + ToolUseID: block.ToolUseID, + Name: block.Name, + Input: append(json.RawMessage(nil), block.Input...), + Content: append(json.RawMessage(nil), content...), + IsError: block.IsError, + Thinking: block.Thinking, + Signature: block.Signature, + Image: func() *NormalizedImage { + if block.Source == nil { + return nil + } + return &NormalizedImage{ + MediaType: block.Source.MediaType, + Data: block.Source.Data, + } + }(), + CacheControl: block.CacheControl, + Raw: append(json.RawMessage(nil), block.Raw...), + } +} + // DenormalizeResponse converts a NormalizedResponse to an Anthropic MessageResponse. func DenormalizeResponse(nr *NormalizedResponse) *types.MessageResponse { resp := &types.MessageResponse{ @@ -113,30 +138,14 @@ func DenormalizeResponse(nr *NormalizedResponse) *types.MessageResponse { switch msg.Role { case "assistant": resp.Role = "assistant" - - // Add thinking block if present. - if msg.Thinking != "" { - resp.Content = append(resp.Content, types.ContentBlock{ - Type: "thinking", - Thinking: msg.Thinking, - }) - } - - // Add text block if present. - if msg.Content != "" { - resp.Content = append(resp.Content, types.ContentBlock{ - Type: "text", - Text: msg.Content, - }) - } - - // Add tool_use blocks. - for _, tc := range msg.ToolCalls { + for _, block := range msg.Blocks { resp.Content = append(resp.Content, types.ContentBlock{ - Type: "tool_use", - ID: tc.ID, - Name: tc.Name, - Input: []byte(tc.Arguments), + Type: block.Type, Text: block.Text, ID: block.ID, + ToolUseID: block.ToolUseID, Name: block.Name, + Input: block.Input, Content: block.Content, + IsError: block.IsError, Thinking: block.Thinking, + Signature: block.Signature, CacheControl: block.CacheControl, + Raw: block.Raw, }) } } diff --git a/internal/core/normalize_test.go b/internal/core/normalize_test.go new file mode 100644 index 00000000..c19c7b31 --- /dev/null +++ b/internal/core/normalize_test.go @@ -0,0 +1,63 @@ +package core + +import ( + "encoding/json" + "testing" + + "github.com/routatic/proxy/pkg/types" +) + +func TestNormalizeRequestPreservesOrderedBlocksAndCacheDirectives(t *testing.T) { + cache := &types.CacheControl{Type: "ephemeral"} + req := &types.MessageRequest{ + Model: "test", + Messages: []types.Message{{ + Role: "user", + Content: json.RawMessage(`[ + {"type":"text","text":"before","cache_control":{"type":"ephemeral"}}, + {"type":"custom_provider_block","payload":{"value":42}}, + {"type":"text","text":"after"} + ]`), + }}, + Tools: []types.Tool{{ + Name: "lookup", InputSchema: json.RawMessage(`{"type":"object"}`), CacheControl: cache, + }}, + } + + normalized := NormalizeRequest(req) + if len(normalized.Messages) != 1 || len(normalized.Messages[0].Blocks) != 3 { + t.Fatalf("ordered blocks were not retained: %+v", normalized.Messages) + } + if normalized.Messages[0].Blocks[0].CacheControl == nil || + normalized.Messages[0].Blocks[0].CacheControl.Type != "ephemeral" { + t.Fatalf("text cache directive was lost: %+v", normalized.Messages[0].Blocks[0]) + } + if got := string(normalized.Messages[0].Blocks[1].Raw); got == "" || + normalized.Messages[0].Blocks[1].Type != "custom_provider_block" { + t.Fatalf("unknown block was not preserved: type=%q raw=%q", + normalized.Messages[0].Blocks[1].Type, got) + } + if normalized.Tools[0].CacheControl == nil { + t.Fatal("tool cache directive was lost") + } +} + +func TestNormalizeRequestPreservesLegacyToolResultOutput(t *testing.T) { + req := &types.MessageRequest{ + Model: "test", + Messages: []types.Message{{ + Role: "tool", + Content: json.RawMessage(`[ + {"type":"tool_result","tool_use_id":"call_1","output":"legacy result"} + ]`), + }}, + } + + normalized := NormalizeRequest(req) + if got, want := normalized.Messages[0].ToolResultsList()[0].Content, "legacy result"; got != want { + t.Fatalf("legacy tool result content = %q, want %q", got, want) + } + if got, want := string(normalized.Messages[0].Blocks[0].Content), `"legacy result"`; got != want { + t.Fatalf("normalized tool result content = %s, want %s", got, want) + } +} diff --git a/internal/core/normalized.go b/internal/core/normalized.go index 1608004d..dc4e9f08 100644 --- a/internal/core/normalized.go +++ b/internal/core/normalized.go @@ -1,5 +1,29 @@ package core +import ( + "encoding/json" + + "github.com/routatic/proxy/pkg/types" +) + +// NormalizedContentBlock is one ordered content block. Raw preserves unknown +// provider-specific JSON so adapters can decide whether and how to forward it. +type NormalizedContentBlock struct { + Type string + Text string + ID string + ToolUseID string + Name string + Input json.RawMessage + Content json.RawMessage + IsError *bool + Thinking string + Signature string + Image *NormalizedImage + CacheControl *types.CacheControl + Raw json.RawMessage +} + // NormalizedToolResult represents a single tool result in the normalized format. type NormalizedToolResult struct { ToolCallID string @@ -16,13 +40,84 @@ type NormalizedImage struct { // All wire formats (Anthropic, OpenAI, Responses, Gemini) map to and from // this representation. type NormalizedMessage struct { - Role string // "user", "assistant", "system", "tool" - Content string // Concatenated text content - Images []NormalizedImage // Image attachments (user messages only) - ToolCalls []NormalizedToolCall // Present on assistant messages - ToolResults []NormalizedToolResult // Present on user messages with tool results - ToolCallID string // Deprecated: use ToolResults instead. Kept for backward compat. - Thinking string // Reasoning/thinking content (assistant only) + Role string // "user", "assistant", "system", "tool" + Blocks []NormalizedContentBlock // Ordered content, including unknown blocks. +} + +func (m NormalizedMessage) TextContent() string { + var text string + for _, block := range m.Blocks { + if block.Type == "text" { + text += block.Text + } + } + return text +} + +func (m NormalizedMessage) ThinkingContent() string { + var thinking string + for _, block := range m.Blocks { + if block.Type == "thinking" { + thinking += block.Thinking + } + } + return thinking +} + +func (m NormalizedMessage) ToolCallsList() []NormalizedToolCall { + var calls []NormalizedToolCall + for _, block := range m.Blocks { + if block.Type == "tool_use" { + calls = append(calls, NormalizedToolCall{ + ID: block.ID, Name: block.Name, Arguments: string(block.Input), + }) + } + } + return calls +} + +func (m NormalizedMessage) ToolResultsList() []NormalizedToolResult { + var results []NormalizedToolResult + for _, block := range m.Blocks { + if block.Type == "tool_result" { + content := toolResultText(block.Content) + results = append(results, NormalizedToolResult{ + ToolCallID: block.ToolUseID, Content: content, + }) + } + } + return results +} + +func toolResultText(raw json.RawMessage) string { + var text string + if json.Unmarshal(raw, &text) == nil { + return text + } + var blocks []struct { + Type string `json:"type"` + Text string `json:"text"` + } + if json.Unmarshal(raw, &blocks) == nil { + for _, block := range blocks { + if block.Type == "text" { + text += block.Text + } + } + if text != "" { + return text + } + } + return string(raw) +} + +func (m NormalizedMessage) HasToolCallID() bool { + for _, block := range m.Blocks { + if block.Type == "tool_result" && block.ToolUseID != "" { + return true + } + } + return false } // NormalizedToolCall represents a tool invocation in the internal format. @@ -36,6 +131,8 @@ type NormalizedToolCall struct { type NormalizedRequest struct { Model string SystemPrompt string + SystemBlocks []NormalizedContentBlock + CacheControl *types.CacheControl Messages []NormalizedMessage MaxTokens int Temperature *float64 @@ -48,9 +145,10 @@ type NormalizedRequest struct { // NormalizedToolDef is a tool definition in the internal format. type NormalizedToolDef struct { - Name string - Description string - InputSchema []byte // JSON bytes of the schema + Name string + Description string + InputSchema []byte // JSON bytes of the schema + CacheControl *types.CacheControl } // NormalizedResponse is the canonical internal response format. diff --git a/internal/core/provider.go b/internal/core/provider.go index 9e2e66da..3fd1a504 100644 --- a/internal/core/provider.go +++ b/internal/core/provider.go @@ -116,3 +116,10 @@ type Provider interface { // stream before it is treated as stuck and aborted. StreamIdleTimeout(model config.ModelConfig) time.Duration } + +// RequestValidator is optionally implemented by providers that can reject +// normalized content before an upstream call. Compatibility failures are +// client errors and must not affect circuit-breaker state. +type RequestValidator interface { + ValidateRequest(req *NormalizedRequest, model config.ModelConfig) error +} diff --git a/internal/core/request_compat.go b/internal/core/request_compat.go new file mode 100644 index 00000000..837bdf1c --- /dev/null +++ b/internal/core/request_compat.go @@ -0,0 +1,58 @@ +package core + +import ( + "fmt" + + "github.com/routatic/proxy/internal/config" +) + +// ValidateRequestCompatibility applies provider capability rules to ordered +// normalized blocks. Provider implementations call this after resolving their +// wire format and model capabilities. +func ValidateRequestCompatibility(req *NormalizedRequest, model config.ModelConfig, caps ProviderCapabilities, wire WireFormat) error { + if req == nil { + return &CompatibilityError{Provider: model.Provider, ModelID: model.ModelID, Reason: "request is nil"} + } + for _, block := range req.SystemBlocks { + if err := validateBlock(block, model, caps, wire); err != nil { + return err + } + } + for _, msg := range req.Messages { + for _, block := range msg.Blocks { + if err := validateBlock(block, model, caps, wire); err != nil { + return err + } + } + } + return nil +} + +func validateBlock(block NormalizedContentBlock, model config.ModelConfig, caps ProviderCapabilities, wire WireFormat) error { + switch block.Type { + case "", "text", "tool_use", "tool_result", "thinking", "image": + default: + if wire != WireFormatAnthropic { + return &CompatibilityError{ + Provider: model.Provider, + ModelID: model.ModelID, + Reason: fmt.Sprintf("content block type %q is only supported by the Anthropic wire format", block.Type), + } + } + } + switch block.Type { + case "tool_use", "tool_result": + if !caps.SupportsTools { + return &CompatibilityError{Provider: model.Provider, ModelID: model.ModelID, Reason: "tools are not supported"} + } + case "thinking": + if !caps.SupportsThinking { + return &CompatibilityError{Provider: model.Provider, ModelID: model.ModelID, Reason: "thinking blocks are not supported"} + } + case "image": + if !caps.SupportsImageInput { + return &CompatibilityError{Provider: model.Provider, ModelID: model.ModelID, Reason: "image input is not supported"} + } + } + return nil +} diff --git a/internal/core/validate.go b/internal/core/validate.go index 41764a90..c60d620b 100644 --- a/internal/core/validate.go +++ b/internal/core/validate.go @@ -21,13 +21,13 @@ func ValidateRequest(req *NormalizedRequest) error { } // Tool-result messages must have a ToolCallID. - if msg.Role == "tool" && msg.ToolCallID == "" { + if msg.Role == "tool" && !msg.HasToolCallID() { return fmt.Errorf("messages[%d]: tool-result message missing tool_call_id", i) } // Assistant messages with tool calls must have non-empty tool calls. - if msg.Role == "assistant" && len(msg.ToolCalls) > 0 { - for j, tc := range msg.ToolCalls { + if msg.Role == "assistant" && len(msg.ToolCallsList()) > 0 { + for j, tc := range msg.ToolCallsList() { if tc.ID == "" { return fmt.Errorf("messages[%d].tool_calls[%d]: missing id", i, j) } diff --git a/internal/gui/assets/app.js b/internal/gui/assets/app.js index ba056970..3abc48a6 100644 --- a/internal/gui/assets/app.js +++ b/internal/gui/assets/app.js @@ -6,6 +6,7 @@ const TRANSLATIONS = { 'status.running': 'Running', 'status.stopped': 'Stopped', 'status.connected': 'Connected', + 'warning.storageDrops': 'Warning: {count} completion record(s) were dropped before storage.', 'tab.overview': 'Overview', 'tab.history': 'History', 'tab.performance': 'Performance', @@ -125,6 +126,7 @@ const TRANSLATIONS = { 'status.running': '运行中', 'status.stopped': '已停止', 'status.connected': '已连接', + 'warning.storageDrops': '警告：有 {count} 条完成记录在写入存储前被丢弃。', 'tab.overview': '概览', 'tab.history': '历史请求', 'tab.fallback': '降级策略', @@ -431,6 +433,15 @@ async function refreshMetrics() { if (!r.ok) return; const d = await r.json(); + const storageWarning = document.getElementById('storage-warning'); + const storageDropped = Number(d.storage_dropped || 0); + if (storageWarning) { + storageWarning.textContent = storageDropped > 0 + ? t('warning.storageDrops').replace('{count}', fmt(storageDropped)) + : ''; + storageWarning.classList.toggle('hidden', storageDropped === 0); + } + // status badge const running = d.proxy_running; const connected = d.connected_to_existing; @@ -1919,4 +1930,3 @@ const AnalyticsModule = { setTimeout(() => { AnalyticsModule.init(); }, 250); - diff --git a/internal/gui/assets/index.html b/internal/gui/assets/index.html index 8c7b93e5..77215a79 100644 --- a/internal/gui/assets/index.html +++ b/internal/gui/assets/index.html @@ -43,6 +43,7 @@
+
Total Requests
diff --git a/internal/gui/perf.go b/internal/gui/perf.go index 74217967..6fa62af4 100644 --- a/internal/gui/perf.go +++ b/internal/gui/perf.go @@ -15,18 +15,20 @@ type modelPerf struct { AvgMs int64 `json:"avg_ms"` P50Ms int64 `json:"p50_ms"` P90Ms int64 `json:"p90_ms"` + P95Ms int64 `json:"p95_ms"` P99Ms int64 `json:"p99_ms"` MinMs int64 `json:"min_ms"` MaxMs int64 `json:"max_ms"` } -func modelPerfFromFields(model string, count int64, avg, p50, p90, p99, min, max time.Duration) modelPerf { +func modelPerfFromFields(model string, count int64, avg, p50, p90, p95, p99, min, max time.Duration) modelPerf { return modelPerf{ Model: model, Count: count, AvgMs: avg.Milliseconds(), P50Ms: p50.Milliseconds(), P90Ms: p90.Milliseconds(), + P95Ms: p95.Milliseconds(), P99Ms: p99.Milliseconds(), MinMs: min.Milliseconds(), MaxMs: max.Milliseconds(), @@ -45,7 +47,7 @@ func (s *Server) handlePerformance(w http.ResponseWriter, r *http.Request) { modelStats, err := latency.GetStats(since) if err == nil { for _, stat := range modelStats { - result[stat.Model] = modelPerfFromFields(stat.Model, stat.Count, stat.Avg, stat.P50, stat.P90, stat.P99, stat.Min, stat.Max) + result[stat.Model] = modelPerfFromFields(stat.Model, stat.Count, stat.Avg, stat.P50, stat.P90, stat.P95, stat.P99, stat.Min, stat.Max) } } @@ -78,7 +80,7 @@ func (s *Server) handlePerformance(w http.ResponseWriter, r *http.Request) { modelStats := s.met.GetModelLatencyStats() for _, stat := range modelStats { - result[stat.Model] = modelPerfFromFields(stat.Model, stat.Count, stat.Avg, stat.P50, stat.P90, stat.P99, stat.Min, stat.Max) + result[stat.Model] = modelPerfFromFields(stat.Model, stat.Count, stat.Avg, stat.P50, stat.P90, stat.P95, stat.P99, stat.Min, stat.Max) } for model, count := range snap.ModelCounts { diff --git a/internal/gui/server.go b/internal/gui/server.go index 8f79ceb1..ae110bf5 100644 --- a/internal/gui/server.go +++ b/internal/gui/server.go @@ -248,6 +248,7 @@ type metricsResponse struct { RequestsStreamed int64 `json:"requests_streamed"` RequestsSuccess int64 `json:"requests_success"` RequestsFailed int64 `json:"requests_failed"` + StorageDropped int64 `json:"storage_dropped"` ModelCounts map[string]int64 `json:"model_counts"` } @@ -264,6 +265,7 @@ func (s *Server) handleMetrics(w http.ResponseWriter, _ *http.Request) { RequestsStreamed: snap.RequestsStreamed, RequestsSuccess: snap.RequestsSuccess, RequestsFailed: snap.RequestsFailed, + StorageDropped: snap.StorageDropped, ModelCounts: snap.ModelCounts, } writeJSON(w, resp) diff --git a/internal/handlers/health.go b/internal/handlers/health.go index 27ae9303..4f8fd317 100644 --- a/internal/handlers/health.go +++ b/internal/handlers/health.go @@ -34,6 +34,12 @@ func NewHealthHandler(tokenCounter *token.Counter, fallbackHandler *router.Fallb func (h *HealthHandler) HandleHealth(w http.ResponseWriter, r *http.Request) { // Get metrics snapshot snapshot := h.metrics.GetSnapshot() + p95, p99 := snapshot.Percentiles() + ttftP95, _ := snapshot.TTFTPercentiles() + warnings := make([]string, 0, 1) + if snapshot.StorageDropped > 0 { + warnings = append(warnings, "storage completion records have been dropped") + } // Get circuit breaker states cbStates := map[string]string{} @@ -56,11 +62,15 @@ func (h *HealthHandler) HandleHealth(w http.ResponseWriter, r *http.Request) { "upstream_calls": snapshot.UpstreamCalls, "rate_limited": snapshot.RateLimited, "deduplicated": snapshot.Deduplicated, - "p95_latency_ms": snapshot.CalculateP95().Milliseconds(), - "p99_latency_ms": snapshot.CalculateP99().Milliseconds(), + "storage_dropped": snapshot.StorageDropped, + "p95_latency_ms": p95.Milliseconds(), + "p99_latency_ms": p99.Milliseconds(), + "ttft_p95_ms": ttftP95.Milliseconds(), + "ttft_samples": len(snapshot.TTFT), }, "circuit_breakers": cbStates, "models": snapshot.ModelCounts, + "warnings": warnings, } w.Header().Set("Content-Type", "application/json") diff --git a/internal/handlers/messages.go b/internal/handlers/messages.go index 4b471993..88ac9edb 100644 --- a/internal/handlers/messages.go +++ b/internal/handlers/messages.go @@ -60,6 +60,7 @@ type responseWriter struct { wroteHeader bool ssePayloadWritten bool contentWritten bool + firstContentAt time.Time usage struct { inputTokens int outputTokens int @@ -148,14 +149,26 @@ func (w *responseWriter) extractUsageFromSSE(b []byte) { func (w *responseWriter) detectContentInSSE(b []byte) { data := string(b) - if strings.Contains(data, `"content_block_start"`) || - strings.Contains(data, `"content_block_delta"`) || - strings.Contains(data, `"text_delta"`) || - strings.Contains(data, `"content":"`) || - strings.Contains(data, `"tool_use"`) || - strings.Contains(data, `"thinking_delta"`) { + if nonEmptyJSONField(data, "text") || + nonEmptyJSONField(data, "thinking") || + nonEmptyJSONField(data, "content") || + nonEmptyJSONField(data, "partial_json") || + (nonEmptyJSONField(data, "name") && strings.Contains(data, `"tool_use"`)) { w.contentWritten = true + if w.firstContentAt.IsZero() { + w.firstContentAt = time.Now() + } + } +} + +func nonEmptyJSONField(data, field string) bool { + prefix := `"` + field + `":"` + idx := strings.Index(data, prefix) + if idx < 0 { + return false } + start := idx + len(prefix) + return start < len(data) && data[start] != '"' } func (w *responseWriter) hasContent() bool { @@ -164,6 +177,12 @@ func (w *responseWriter) hasContent() bool { return w.contentWritten } +func (w *responseWriter) firstContentTime() time.Time { + w.mu.Lock() + defer w.mu.Unlock() + return w.firstContentAt +} + func (w *responseWriter) getOutputTokens() int { w.mu.Lock() defer w.mu.Unlock() @@ -358,6 +377,7 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) return } + requestStart := time.Now() // Generate or get request ID for correlation. // Cap externally-provided IDs at 256 bytes to prevent header abuse. @@ -382,6 +402,7 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) // Read the raw request body with a size limit to prevent memory exhaustion. const maxBodySize = 104857600 // 100 MB r.Body = http.MaxBytesReader(w, r.Body, maxBodySize) + parseStart := time.Now() var rawBody json.RawMessage if err := json.NewDecoder(r.Body).Decode(&rawBody); err != nil { var maxErr *http.MaxBytesError @@ -419,6 +440,7 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) h.sendError(w, http.StatusBadRequest, err.Error(), nil) return } + h.metrics.RecordStage(metrics.StageRequestParse, time.Since(parseStart)) // Record metrics isStreaming := anthropicReq.Stream != nil && *anthropicReq.Stream @@ -455,13 +477,16 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) } // Count tokens. + tokenStart := time.Now() tokenCount, err := h.tokenCounter.CountMessages(systemText, tokenMessages) if err != nil { h.logger.Warn("failed to count tokens", "error", err) tokenCount = 0 } + h.metrics.RecordStage(metrics.StageTokenCount, time.Since(tokenStart)) // Route to appropriate model and build fallback chain. + routeStart := time.Now() facts := router.AnalyzeRequestFacts(routerMessages) needsTools := len(anthropicReq.Tools) > 0 modelChain, routeResult, err := h.buildModelChain(anthropicReq.Model, routerMessages, tokenCount, isStreaming, anthropicReq.MaxTokens, facts.NeedsVision, needsTools) @@ -475,6 +500,7 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) h.sendError(w, status, message, err) return } + h.metrics.RecordStage(metrics.StageRouting, time.Since(routeStart)) h.logger.Info("routing request", "scenario", routeResult.Scenario, @@ -484,8 +510,15 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) "reason", routeResult.Reason, ) + normalizeStart := time.Now() normalizedReq := core.NormalizeRequest(&anthropicReq) normalizedReq.Stream = isStreaming + h.metrics.RecordStage(metrics.StageNormalization, time.Since(normalizeStart)) + modelChain, err = h.filterCompatibleModels(modelChain, normalizedReq) + if err != nil { + h.sendError(w, http.StatusBadRequest, err.Error(), err) + return + } if h.captureLogger != nil && len(modelChain) > 0 { provider := modelChain[0].Provider @@ -494,12 +527,49 @@ func (h *MessagesHandler) HandleMessages(w http.ResponseWriter, r *http.Request) } if isStreaming { - h.handleStreaming(w, r, &anthropicReq, normalizedReq, modelChain, rawBody, routeResult.Scenario, requestID) + h.handleStreaming(w, r, &anthropicReq, normalizedReq, modelChain, rawBody, routeResult.Scenario, requestID, requestStart) } else { h.handleNonStreaming(w, r, &anthropicReq, normalizedReq, modelChain, rawBody, routeResult.Scenario, requestID) } } +func (h *MessagesHandler) filterCompatibleModels(chain []config.ModelConfig, req *core.NormalizedRequest) ([]config.ModelConfig, error) { + if h.providerRegistry == nil { + return chain, nil + } + compatible := make([]config.ModelConfig, 0, len(chain)) + var reasons []string + for _, model := range chain { + prov, ok := h.providerRegistry.Get(client.Provider(model)) + if !ok { + compatible = append(compatible, model) + continue + } + validator, ok := prov.(core.RequestValidator) + if !ok { + compatible = append(compatible, model) + continue + } + if err := validator.ValidateRequest(req, model); err != nil { + var compatErr *core.CompatibilityError + if errors.As(err, &compatErr) { + reasons = append(reasons, compatErr.Error()) + h.logger.Info("model skipped by request compatibility", "model", model.ModelID, "reason", compatErr.Reason) + continue + } + return nil, err + } + compatible = append(compatible, model) + } + if len(compatible) == 0 { + if len(reasons) > 0 { + return nil, fmt.Errorf("no compatible model found: %s", strings.Join(reasons, "; ")) + } + return nil, fmt.Errorf("no compatible model found") + } + return compatible, nil +} + // buildModelChain resolves the request to a model chain (primary + fallbacks), // honoring model_overrides (with a deduplicated scenario safety-net) and // respecting the streaming-scenario-routing toggle. @@ -606,8 +676,13 @@ func (h *MessagesHandler) handleStreaming( rawBody json.RawMessage, scenario router.Scenario, requestID string, + requestStarts ...time.Time, ) { clientCtx := r.Context() + requestStart := time.Now() + if len(requestStarts) > 0 { + requestStart = requestStarts[0] + } rw := &responseWriter{ResponseWriter: w} @@ -653,7 +728,12 @@ func (h *MessagesHandler) handleStreaming( recordStreamSuccess := func(model config.ModelConfig) { cancelAttempt() latency := time.Since(streamStart) - h.metrics.RecordSuccess(model.ModelID, latency) + modelKey := metrics.ModelKey(model.Provider, model.ModelID) + h.metrics.RecordSuccess(modelKey, latency) + h.metrics.RecordStage(metrics.StageUpstream, latency) + if firstContentAt := rw.firstContentTime(); !firstContentAt.IsZero() { + h.metrics.RecordTTFT(firstContentAt.Sub(requestStart)) + } h.logger.Info("streaming completed", "model", model.ModelID, "latency", latency, @@ -679,12 +759,9 @@ func (h *MessagesHandler) handleStreaming( h.history.Add(rec) } if h.storage != nil { - if err := h.storage.InsertRequest(rec); err != nil { - h.logger.Warn("failed to insert request into storage", "error", err) - } - if err := h.storage.InsertLatency(model.ModelID, latency); err != nil { - h.logger.Warn("failed to insert latency sample into storage", "error", err) - } + storageStart := time.Now() + h.storage.RecordCompletion(rec) + h.metrics.RecordStage(metrics.StageStorageEnqueue, time.Since(storageStart)) } } @@ -703,7 +780,7 @@ func (h *MessagesHandler) handleStreaming( "model", model.ModelID, "idle_timeout", idleTimeout) if rw.ssePayloadWritten { h.sendStreamError(rw, "stream idle after SSE payload started") - h.metrics.RecordFailureForModel(model.ModelID) + h.metrics.RecordFailureForModel(metrics.ModelKey(model.Provider, model.ModelID)) return false // abort } return true // continue to next model @@ -713,7 +790,7 @@ func (h *MessagesHandler) handleStreaming( "model", model.ModelID) if rw.ssePayloadWritten { h.sendStreamError(rw, "empty stream after SSE payload started") - h.metrics.RecordFailureForModel(model.ModelID) + h.metrics.RecordFailureForModel(metrics.ModelKey(model.Provider, model.ModelID)) return false // abort } return true // continue to next model @@ -721,7 +798,7 @@ func (h *MessagesHandler) handleStreaming( h.logger.Warn(action+" streaming failed", "model", model.ModelID, "error", err) if rw.ssePayloadWritten { h.sendStreamError(rw, "all upstream models failed after SSE payload started") - h.metrics.RecordFailureForModel(model.ModelID) + h.metrics.RecordFailureForModel(metrics.ModelKey(model.Provider, model.ModelID)) return false // abort — cannot fallback after SSE payload started } return true // continue to next model @@ -758,7 +835,9 @@ func (h *MessagesHandler) handleStreaming( if wireFormat == core.WireFormatAnthropic { atomic.StoreInt32(&heartbeatPaused, 1) } + transformStart := time.Now() errProxy := h.streamProxy.ProxyStream(rw, streamReader, wireFormat, model.ModelID, attemptCtx, idleTimeout, cancelAttempt) + h.metrics.RecordStage(metrics.StageResponseTransform, time.Since(transformStart)) if wireFormat == core.WireFormatAnthropic { atomic.StoreInt32(&heartbeatPaused, 0) } @@ -890,7 +969,9 @@ func (h *MessagesHandler) handleStreaming( // Bind body read to attemptCtx so streaming_timeout_ms aborts mid-stream. streamReader := transformer.NewCtxReadCloser(attemptCtx, streamBody) + transformStart := time.Now() if err := h.streamHandler.ProxyStream(rw, streamReader, model.ModelID, attemptCtx, idleTimeout, cancelAttempt); err != nil { + h.metrics.RecordStage(metrics.StageResponseTransform, time.Since(transformStart)) if err == transformer.ErrClientDisconnected { if clientCtx.Err() != nil { h.logger.Debug("client disconnected during stream") @@ -903,6 +984,7 @@ func (h *MessagesHandler) handleStreaming( } continue } + h.metrics.RecordStage(metrics.StageResponseTransform, time.Since(transformStart)) recordStreamSuccess(model) return @@ -1176,6 +1258,7 @@ func (h *MessagesHandler) handleNonStreaming( ) { ctx := r.Context() startTime := time.Now() + upstreamStart := time.Now() result, responseBody, err := h.fallbackHandler.ExecuteWithFallback( ctx, @@ -1231,13 +1314,16 @@ func (h *MessagesHandler) handleNonStreaming( h.logger.Info("request context canceled during non-streaming fallback", "error", err) return } - h.metrics.RecordFailureForModel(result.ModelID) + if result != nil { + h.metrics.RecordFailureForModel(modelMetricKey(modelChain, result.ModelID)) + } h.sendError(w, http.StatusBadGateway, "all models failed", err) return } + h.metrics.RecordStage(metrics.StageUpstream, time.Since(upstreamStart)) latency := time.Since(startTime) - h.metrics.RecordSuccess(result.ModelID, latency) + h.metrics.RecordSuccess(modelMetricKey(modelChain, result.ModelID), latency) h.logger.Info("request completed", "model", result.ModelID, @@ -1277,12 +1363,9 @@ func (h *MessagesHandler) handleNonStreaming( h.history.Add(rec) } if h.storage != nil { - if err := h.storage.InsertRequest(rec); err != nil { - h.logger.Warn("failed to insert request into storage", "error", err) - } - if err := h.storage.InsertLatency(result.ModelID, latency); err != nil { - h.logger.Warn("failed to insert latency sample into storage", "error", err) - } + storageStart := time.Now() + h.storage.RecordCompletion(rec) + h.metrics.RecordStage(metrics.StageStorageEnqueue, time.Since(storageStart)) } w.Header().Set("Content-Type", "application/json") @@ -1290,6 +1373,15 @@ func (h *MessagesHandler) handleNonStreaming( _, _ = w.Write(responseBody) } +func modelMetricKey(chain []config.ModelConfig, modelID string) string { + for _, model := range chain { + if model.ModelID == modelID { + return metrics.ModelKey(model.Provider, modelID) + } + } + return modelID +} + // executeAnthropicRequest executes a request to the Anthropic endpoint (for MiniMax models). func (h *MessagesHandler) executeAnthropicRequest( ctx context.Context, diff --git a/internal/handlers/storage_adapter.go b/internal/handlers/storage_adapter.go index aa723528..f3c64b26 100644 --- a/internal/handlers/storage_adapter.go +++ b/internal/handlers/storage_adapter.go @@ -1,33 +1,116 @@ package handlers import ( - "time" + "context" + "log/slog" + "sync" "github.com/routatic/proxy/internal/history" + "github.com/routatic/proxy/internal/metrics" "github.com/routatic/proxy/internal/storage" ) type StorageWriter interface { - InsertRequest(rec history.RequestRecord) error - InsertLatency(model string, latency time.Duration) error + RecordCompletion(rec history.RequestRecord) + Shutdown(ctx context.Context) error } type StorageAdapter struct { - requests *storage.Requests - latency *storage.Latency + requests *storage.Requests + metrics *metrics.Metrics + queue chan history.RequestRecord + stop chan struct{} + stopOnce sync.Once + closeOnce sync.Once + queueMu sync.RWMutex + closed bool + wg sync.WaitGroup } -func NewStorageAdapter(db *storage.Database) *StorageAdapter { - return &StorageAdapter{ +func NewStorageAdapter(db *storage.Database, metricSink *metrics.Metrics) *StorageAdapter { + s := &StorageAdapter{ requests: storage.NewRequests(db), - latency: storage.NewLatency(db), + metrics: metricSink, + queue: make(chan history.RequestRecord, 1024), + stop: make(chan struct{}), } + s.wg.Add(1) + go s.run() + return s } -func (s *StorageAdapter) InsertRequest(rec history.RequestRecord) error { - return s.requests.Insert(rec) +func (s *StorageAdapter) run() { + defer s.wg.Done() + for { + select { + case rec, ok := <-s.queue: + if !ok { + return + } + if err := s.requests.Insert(rec); err != nil { + slog.Warn("failed to persist request completion", "request_id", rec.ID, "error", err) + } + case <-s.stop: + s.drainBuffered() + return + } + } +} + +// RecordCompletion enqueues a completed request without blocking the request +// path. If the bounded queue is full, the newest record is dropped. +func (s *StorageAdapter) RecordCompletion(rec history.RequestRecord) { + s.queueMu.RLock() + defer s.queueMu.RUnlock() + if s.closed { + return + } + select { + case s.queue <- rec: + default: + if s.metrics != nil { + s.metrics.RecordStorageDrop() + } + slog.Warn("storage completion queue full; dropping newest record", "request_id", rec.ID) + } } -func (s *StorageAdapter) InsertLatency(model string, latency time.Duration) error { - return s.latency.Insert(model, latency) +func (s *StorageAdapter) drainBuffered() { + for { + select { + case rec, ok := <-s.queue: + if !ok { + return + } + if err := s.requests.Insert(rec); err != nil { + slog.Warn("failed to persist request completion", "request_id", rec.ID, "error", err) + } + default: + return + } + } +} + +// Shutdown stops accepting completions and drains accepted records until ctx +// expires. +func (s *StorageAdapter) Shutdown(ctx context.Context) error { + s.closeOnce.Do(func() { + s.queueMu.Lock() + s.closed = true + close(s.queue) + s.queueMu.Unlock() + }) + done := make(chan struct{}) + go func() { + s.wg.Wait() + close(done) + }() + select { + case <-done: + return nil + case <-ctx.Done(): + s.stopOnce.Do(func() { close(s.stop) }) + s.wg.Wait() + return ctx.Err() + } } diff --git a/internal/handlers/ttft_test.go b/internal/handlers/ttft_test.go new file mode 100644 index 00000000..8f6c992c --- /dev/null +++ b/internal/handlers/ttft_test.go @@ -0,0 +1,31 @@ +package handlers + +import "testing" + +func TestNonEmptyJSONField(t *testing.T) { + if nonEmptyJSONField(`{"type":"text_delta","text":""}`, "text") { + t.Fatal("empty text was treated as content") + } + if !nonEmptyJSONField(`{"type":"text_delta","text":"hello"}`, "text") { + t.Fatal("non-empty text was not detected") + } +} + +func TestResponseWriterDetectsFirstNonEmptyContent(t *testing.T) { + rw := &responseWriter{} + rw.detectContentInSSE([]byte(`event: content_block_start +data: {"type":"content_block_start","content_block":{"type":"text","text":""}} + +`)) + if rw.hasContent() { + t.Fatal("empty content block start was treated as content") + } + + rw.detectContentInSSE([]byte(`event: content_block_delta +data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hello"}} + +`)) + if !rw.hasContent() { + t.Fatal("non-empty text delta was not detected") + } +} diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index 77ec1ef0..e2d73278 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -4,11 +4,20 @@ package metrics import ( "math" "sort" + "strings" "sync" "sync/atomic" "time" ) +// ModelKey returns the canonical provider/model label used by metrics. +func ModelKey(provider, model string) string { + if strings.TrimSpace(provider) == "" { + return model + } + return provider + "/" + model +} + // Metrics holds in-memory metrics for the proxy. type Metrics struct { // Counters (atomic) @@ -19,12 +28,18 @@ type Metrics struct { upstreamCalls atomic.Int64 rateLimited atomic.Int64 deduplicated atomic.Int64 + storageDropped atomic.Int64 // Latency tracking mu sync.RWMutex - latencies []time.Duration + latencies durationRing maxLatencySamples int + // Request-stage timing + stageMu sync.RWMutex + stageLatencies map[string]durationRing + ttft durationRing + // By model modelCounts map[string]*atomic.Int64 modelMu sync.RWMutex @@ -36,20 +51,73 @@ type Metrics struct { modelFailedMu sync.RWMutex // Per-model latency tracking - modelLatencies map[string][]time.Duration + modelLatencies map[string]*durationRing modelLatMu sync.RWMutex maxPerModelSamples int } +const ( + defaultMaxLatencySamples = 1000 + defaultMaxPerModelSamples = 200 + defaultMaxStageSamples = 1000 + StageRequestParse = "request_parse" + StageTokenCount = "token_count" + StageRouting = "routing" + StageNormalization = "normalization" + StageUpstream = "upstream" + StageResponseTransform = "response_transform" + StageStorageEnqueue = "storage_enqueue" +) + +// durationRing stores a bounded rolling window of durations. +type durationRing struct { + values []time.Duration + next int + count int +} + +func newDurationRing(capacity int) durationRing { + return durationRing{values: make([]time.Duration, capacity)} +} + +func (r *durationRing) Add(value time.Duration) { + if len(r.values) == 0 { + return + } + r.values[r.next] = value + r.next = (r.next + 1) % len(r.values) + if r.count < len(r.values) { + r.count++ + } +} + +func (r *durationRing) Snapshot() []time.Duration { + if r.count == 0 { + return nil + } + out := make([]time.Duration, r.count) + start := 0 + if r.count == len(r.values) { + start = r.next + } + for i := range out { + out[i] = r.values[(start+i)%len(r.values)] + } + return out +} + // New creates a new metrics instance. func New() *Metrics { return &Metrics{ - maxLatencySamples: 1000, - maxPerModelSamples: 100, + maxLatencySamples: defaultMaxLatencySamples, + maxPerModelSamples: defaultMaxPerModelSamples, + latencies: newDurationRing(defaultMaxLatencySamples), + ttft: newDurationRing(defaultMaxLatencySamples), + stageLatencies: make(map[string]durationRing), modelCounts: make(map[string]*atomic.Int64), modelSuccess: make(map[string]int64), modelFailed: make(map[string]int64), - modelLatencies: make(map[string][]time.Duration), + modelLatencies: make(map[string]*durationRing), } } @@ -98,27 +166,51 @@ func (m *Metrics) RecordDeduplicated() { m.deduplicated.Add(1) } +// RecordStorageDrop records a completion record dropped from the bounded +// asynchronous storage queue. +func (m *Metrics) RecordStorageDrop() { + m.storageDropped.Add(1) +} + func (m *Metrics) recordLatency(latency time.Duration) { m.mu.Lock() defer m.mu.Unlock() - - // Keep last N samples for p95/p99 - if len(m.latencies) >= m.maxLatencySamples { - // Shift and add new - m.latencies = m.latencies[1:] - } - m.latencies = append(m.latencies, latency) + m.latencies.Add(latency) } func (m *Metrics) recordModelLatency(model string, latency time.Duration) { m.modelLatMu.Lock() defer m.modelLatMu.Unlock() - samples := m.modelLatencies[model] - if len(samples) >= m.maxPerModelSamples { - samples = samples[1:] + ring, ok := m.modelLatencies[model] + if !ok { + r := newDurationRing(m.maxPerModelSamples) + ring = &r + m.modelLatencies[model] = ring + } + ring.Add(latency) +} + +// RecordStage records a bounded timing sample for a known request stage. +func (m *Metrics) RecordStage(stage string, duration time.Duration) { + m.stageMu.Lock() + defer m.stageMu.Unlock() + + ring, ok := m.stageLatencies[stage] + if !ok { + ring = newDurationRing(defaultMaxStageSamples) + m.stageLatencies[stage] = ring } - m.modelLatencies[model] = append(samples, latency) + ring.Add(duration) + m.stageLatencies[stage] = ring +} + +// RecordTTFT records the time to first non-empty model content for a streaming +// request. +func (m *Metrics) RecordTTFT(duration time.Duration) { + m.stageMu.Lock() + defer m.stageMu.Unlock() + m.ttft.Add(duration) } func (m *Metrics) recordModel(model string) { @@ -134,10 +226,18 @@ func (m *Metrics) recordModel(model string) { // GetSnapshot returns a snapshot of current metrics. func (m *Metrics) GetSnapshot() Snapshot { m.mu.RLock() - latencies := make([]time.Duration, len(m.latencies)) - copy(latencies, m.latencies) + latencies := m.latencies.Snapshot() m.mu.RUnlock() + m.stageMu.RLock() + ttft := m.ttft.Snapshot() + stageLatencies := make(map[string][]time.Duration, len(m.stageLatencies)) + for stage, ring := range m.stageLatencies { + samples := ring.Snapshot() + stageLatencies[stage] = samples + } + m.stageMu.RUnlock() + modelCounts := make(map[string]int64) m.modelMu.RLock() for k, v := range m.modelCounts { @@ -168,10 +268,13 @@ func (m *Metrics) GetSnapshot() Snapshot { UpstreamCalls: m.upstreamCalls.Load(), RateLimited: m.rateLimited.Load(), Deduplicated: m.deduplicated.Load(), + StorageDropped: m.storageDropped.Load(), Latencies: latencies, ModelCounts: modelCounts, ModelSuccess: modelSuccess, ModelFailed: modelFailed, + TTFT: ttft, + StageLatencies: stageLatencies, } } @@ -184,10 +287,13 @@ type Snapshot struct { UpstreamCalls int64 RateLimited int64 Deduplicated int64 + StorageDropped int64 Latencies []time.Duration ModelCounts map[string]int64 ModelSuccess map[string]int64 // Per-model success counts ModelFailed map[string]int64 // Per-model failure counts + TTFT []time.Duration + StageLatencies map[string][]time.Duration } // ModelLatencyStats holds latency statistics for a single model. @@ -197,6 +303,7 @@ type ModelLatencyStats struct { Avg time.Duration P50 time.Duration P90 time.Duration + P95 time.Duration P99 time.Duration Min time.Duration Max time.Duration @@ -205,13 +312,18 @@ type ModelLatencyStats struct { // GetModelLatencyStats returns latency statistics for all models. func (m *Metrics) GetModelLatencyStats() []ModelLatencyStats { m.modelLatMu.RLock() - defer m.modelLatMu.RUnlock() - - var stats []ModelLatencyStats - for model, samples := range m.modelLatencies { + samplesByModel := make(map[string][]time.Duration, len(m.modelLatencies)) + for model, ring := range m.modelLatencies { + samples := ring.Snapshot() if len(samples) == 0 { continue } + samplesByModel[model] = samples + } + m.modelLatMu.RUnlock() + + stats := make([]ModelLatencyStats, 0, len(samplesByModel)) + for model, samples := range samplesByModel { stats = append(stats, calculateModelStats(model, samples)) } return stats @@ -236,6 +348,7 @@ func calculateModelStats(model string, samples []time.Duration) ModelLatencyStat p50Idx := int(math.Ceil(float64(count)*0.50)) - 1 p90Idx := int(math.Ceil(float64(count)*0.90)) - 1 + p95Idx := int(math.Ceil(float64(count)*0.95)) - 1 p99Idx := int(math.Ceil(float64(count)*0.99)) - 1 if p50Idx < 0 { p50Idx = 0 @@ -243,6 +356,9 @@ func calculateModelStats(model string, samples []time.Duration) ModelLatencyStat if p90Idx < 0 { p90Idx = 0 } + if p95Idx < 0 { + p95Idx = 0 + } if p99Idx < 0 { p99Idx = 0 } @@ -252,6 +368,9 @@ func calculateModelStats(model string, samples []time.Duration) ModelLatencyStat if p90Idx >= count { p90Idx = count - 1 } + if p95Idx >= count { + p95Idx = count - 1 + } if p99Idx >= count { p99Idx = count - 1 } @@ -262,6 +381,7 @@ func calculateModelStats(model string, samples []time.Duration) ModelLatencyStat Avg: avg, P50: sorted[p50Idx], P90: sorted[p90Idx], + P95: sorted[p95Idx], P99: sorted[p99Idx], Min: sorted[0], Max: sorted[count-1], @@ -270,24 +390,43 @@ func calculateModelStats(model string, samples []time.Duration) ModelLatencyStat // CalculateP95 calculates the p95 latency from the snapshot. func (s Snapshot) CalculateP95() time.Duration { - if len(s.Latencies) == 0 { - return 0 - } - index := int(float64(len(s.Latencies)) * 0.95) - if index >= len(s.Latencies) { - index = len(s.Latencies) - 1 - } - return s.Latencies[index] + p95, _ := s.Percentiles() + return p95 } // CalculateP99 calculates the p99 latency from the snapshot. func (s Snapshot) CalculateP99() time.Duration { - if len(s.Latencies) == 0 { + _, p99 := s.Percentiles() + return p99 +} + +// Percentiles returns p95 and p99 from the snapshot's latency samples. +// It sorts one copy so callers can safely provide samples in any order. +func (s Snapshot) Percentiles() (time.Duration, time.Duration) { + return percentiles(s.Latencies) +} + +// TTFTPercentiles returns p95 and p99 for streaming time-to-first-token samples. +func (s Snapshot) TTFTPercentiles() (time.Duration, time.Duration) { + return percentiles(s.TTFT) +} + +func percentiles(samples []time.Duration) (time.Duration, time.Duration) { + sortedSamples := append([]time.Duration(nil), samples...) + sort.Slice(sortedSamples, func(i, j int) bool { return sortedSamples[i] < sortedSamples[j] }) + return percentile(sortedSamples, 0.95), percentile(sortedSamples, 0.99) +} + +func percentile(sortedSamples []time.Duration, fraction float64) time.Duration { + if len(sortedSamples) == 0 { return 0 } - index := int(float64(len(s.Latencies)) * 0.99) - if index >= len(s.Latencies) { - index = len(s.Latencies) - 1 + index := int(math.Ceil(float64(len(sortedSamples))*fraction)) - 1 + if index < 0 { + index = 0 + } + if index >= len(sortedSamples) { + index = len(sortedSamples) - 1 } - return s.Latencies[index] + return sortedSamples[index] } diff --git a/internal/metrics/metrics_test.go b/internal/metrics/metrics_test.go new file mode 100644 index 00000000..ebaefadb --- /dev/null +++ b/internal/metrics/metrics_test.go @@ -0,0 +1,107 @@ +package metrics + +import ( + "testing" + "time" +) + +func TestSnapshotPercentilesUseSortedSamples(t *testing.T) { + t.Parallel() + + snapshot := Snapshot{Latencies: []time.Duration{ + 90 * time.Millisecond, + 10 * time.Millisecond, + 50 * time.Millisecond, + 20 * time.Millisecond, + 80 * time.Millisecond, + }} + + if got, want := snapshot.CalculateP95(), 90*time.Millisecond; got != want { + t.Fatalf("CalculateP95() = %s, want %s", got, want) + } + if got, want := snapshot.CalculateP99(), 90*time.Millisecond; got != want { + t.Fatalf("CalculateP99() = %s, want %s", got, want) + } +} + +func TestMetricsRetainsLatestGlobalSamples(t *testing.T) { + t.Parallel() + + ring := newDurationRing(3) + ring.Add(time.Second) + ring.Add(2 * time.Second) + ring.Add(3 * time.Second) + ring.Add(4 * time.Second) + + got := ring.Snapshot() + want := []time.Duration{2 * time.Second, 3 * time.Second, 4 * time.Second} + if !equalDurations(got, want) { + t.Fatalf("latencies = %v, want %v", got, want) + } +} + +func TestMetricsModelLatencyUsesBoundedSamples(t *testing.T) { + t.Parallel() + + m := New() + m.recordModelLatency("provider/model", time.Second) + m.recordModelLatency("provider/model", 2*time.Second) + m.recordModelLatency("provider/model", 3*time.Second) + + stats := m.GetModelLatencyStats() + if len(stats) != 1 { + t.Fatalf("got %d model stats, want 1", len(stats)) + } + if got, want := stats[0].Count, int64(3); got != want { + t.Fatalf("Count = %d, want %d", got, want) + } +} + +func TestMetricsRecordsStageAndTTFT(t *testing.T) { + t.Parallel() + + m := New() + m.RecordStage(StageTokenCount, 2*time.Millisecond) + m.RecordTTFT(12 * time.Millisecond) + + snapshot := m.GetSnapshot() + if got, want := percentile(snapshot.TTFT, 0.95), 12*time.Millisecond; got != want { + t.Fatalf("TTFT p95 = %s, want %s", got, want) + } + if got := snapshot.StageLatencies["token_count"]; len(got) != 1 || got[0] != 2*time.Millisecond { + t.Fatalf("token_count stage samples = %v", got) + } +} + +func TestMetricsRecordsStorageDrops(t *testing.T) { + t.Parallel() + + m := New() + m.RecordStorageDrop() + m.RecordStorageDrop() + + if got, want := m.GetSnapshot().StorageDropped, int64(2); got != want { + t.Fatalf("StorageDropped = %d, want %d", got, want) + } +} + +func TestModelKey(t *testing.T) { + if got := ModelKey("opencode-go", "model-a"); got != "opencode-go/model-a" { + t.Fatalf("ModelKey() = %q", got) + } + if got := ModelKey("", "model-a"); got != "model-a" { + t.Fatalf("ModelKey() without provider = %q", got) + } +} + +func equalDurations(a, b []time.Duration) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/internal/provider/aws_bedrock.go b/internal/provider/aws_bedrock.go index db41c853..7e0c18ef 100644 --- a/internal/provider/aws_bedrock.go +++ b/internal/provider/aws_bedrock.go @@ -31,6 +31,12 @@ func NewAWSBedrockProvider(atomic *config.AtomicConfig) *AWSBedrockProvider { // Name returns the provider identifier. func (p *AWSBedrockProvider) Name() string { return "aws-bedrock" } +// ValidateRequest checks ordered content against this provider's capabilities +// before any upstream request is attempted. +func (p *AWSBedrockProvider) ValidateRequest(req *core.NormalizedRequest, model config.ModelConfig) error { + return validateRequest(p, req, model) +} + // Capabilities returns provider-level capabilities. func (p *AWSBedrockProvider) Capabilities() core.ProviderCapabilities { return core.ProviderCapabilities{ @@ -230,7 +236,7 @@ func (p *AWSBedrockProvider) streamResponses(ctx context.Context, req *core.Norm func (p *AWSBedrockProvider) buildResponsesRequest(req *core.NormalizedRequest, model config.ModelConfig) *types.ResponsesRequest { var inputs []types.ResponsesInput for _, msg := range req.Messages { - contentBytes, _ := json.Marshal(msg.Content) + contentBytes, _ := json.Marshal(msg.TextContent()) inputs = append(inputs, types.ResponsesInput{ Role: msg.Role, Content: contentBytes, diff --git a/internal/provider/aws_bedrock_test.go b/internal/provider/aws_bedrock_test.go index baeab355..41010128 100644 --- a/internal/provider/aws_bedrock_test.go +++ b/internal/provider/aws_bedrock_test.go @@ -155,7 +155,7 @@ func TestAWSBedrockProvider_Execute(t *testing.T) { req := &core.NormalizedRequest{ Model: "moonshotai.kimi-k2.5", - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, } model := config.ModelConfig{ModelID: "moonshotai.kimi-k2.5"} @@ -204,7 +204,7 @@ func TestAWSBedrockProvider_Stream(t *testing.T) { req := &core.NormalizedRequest{ Model: "moonshotai.kimi-k2.5", - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, Stream: true, } model := config.ModelConfig{ModelID: "moonshotai.kimi-k2.5"} @@ -243,7 +243,7 @@ func TestAWSBedrockProvider_ExecuteAnthropic_UsesUserAgent(t *testing.T) { req := &core.NormalizedRequest{ Model: "anthropic.claude-sonnet-4", - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, } model := config.ModelConfig{ModelID: "anthropic.claude-sonnet-4"} @@ -274,7 +274,7 @@ func TestAWSBedrockProvider_StreamAnthropic_UsesUserAgent(t *testing.T) { req := &core.NormalizedRequest{ Model: "anthropic.claude-sonnet-4", - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, Stream: true, } model := config.ModelConfig{ModelID: "anthropic.claude-sonnet-4"} @@ -319,7 +319,7 @@ func TestAWSBedrockProvider_Execute_NoProjectID(t *testing.T) { req := &core.NormalizedRequest{ Model: "test-model", - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, } model := config.ModelConfig{ModelID: "test-model"} diff --git a/internal/provider/opencode_go.go b/internal/provider/opencode_go.go index 576d3f96..da6d8f41 100644 --- a/internal/provider/opencode_go.go +++ b/internal/provider/opencode_go.go @@ -29,6 +29,12 @@ func NewOpenCodeGoProvider(atomic *config.AtomicConfig) *OpenCodeGoProvider { // Name returns the provider identifier. func (p *OpenCodeGoProvider) Name() string { return "opencode-go" } +// ValidateRequest checks ordered content against this provider's capabilities +// before any upstream request is attempted. +func (p *OpenCodeGoProvider) ValidateRequest(req *core.NormalizedRequest, model config.ModelConfig) error { + return validateRequest(p, req, model) +} + // Capabilities returns provider-level capabilities. func (p *OpenCodeGoProvider) Capabilities() core.ProviderCapabilities { return core.ProviderCapabilities{ diff --git a/internal/provider/opencode_go_wireformat_test.go b/internal/provider/opencode_go_wireformat_test.go index 22e653d0..299ffbec 100644 --- a/internal/provider/opencode_go_wireformat_test.go +++ b/internal/provider/opencode_go_wireformat_test.go @@ -91,7 +91,7 @@ func TestOpenCodeGoProvider_WireFormatOverride_MatchesEndpoint(t *testing.T) { model := config.ModelConfig{ModelID: "deepseek-v4-pro", WireFormat: "responses"} req := &core.NormalizedRequest{ Model: model.ModelID, - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, } if got := p.WireFormat(model); got != core.WireFormatOpenAIResponses { @@ -119,7 +119,7 @@ func TestOpenCodeGoProvider_Responses_MissingBaseURL(t *testing.T) { model := config.ModelConfig{ModelID: "deepseek-v4-pro", WireFormat: "responses"} req := &core.NormalizedRequest{ Model: model.ModelID, - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, } // An unset responses_base_url must name the missing config key rather than @@ -166,7 +166,7 @@ func TestOpenCodeGoProvider_ExecuteResponses_Override(t *testing.T) { model := config.ModelConfig{ModelID: "muse-spark-1.2-contributor", WireFormat: "responses"} req := &core.NormalizedRequest{ Model: model.ModelID, - Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, + Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, } result, err := p.Execute(context.Background(), req, model) diff --git a/internal/provider/opencode_useragent_test.go b/internal/provider/opencode_useragent_test.go index 2c712a2b..154bc0d0 100644 --- a/internal/provider/opencode_useragent_test.go +++ b/internal/provider/opencode_useragent_test.go @@ -59,7 +59,7 @@ func TestOpenCodeZenProvider_Execute_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeZenProvider(atomic) - req := &core.NormalizedRequest{Model: "deepseek-v4-flash-free", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}} + req := &core.NormalizedRequest{Model: "deepseek-v4-flash-free", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}} model := config.ModelConfig{ModelID: "deepseek-v4-flash-free"} if got := p.WireFormat(model); got != core.WireFormatOpenAIChat { t.Fatalf("WireFormat(%q) = %v, want OpenAIChat", model.ModelID, got) @@ -77,7 +77,7 @@ func TestOpenCodeZenProvider_Stream_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeZenProvider(atomic) - req := &core.NormalizedRequest{Model: "deepseek-v4-flash-free", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, Stream: true} + req := &core.NormalizedRequest{Model: "deepseek-v4-flash-free", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, Stream: true} model := config.ModelConfig{ModelID: "deepseek-v4-flash-free"} body, err := p.Stream(context.Background(), req, model) @@ -109,7 +109,7 @@ func TestOpenCodeZenProvider_ExecuteAnthropic_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeZenProvider(atomic) - req := &core.NormalizedRequest{Model: "claude-sonnet-4.5", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}} + req := &core.NormalizedRequest{Model: "claude-sonnet-4.5", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}} model := config.ModelConfig{ModelID: "claude-sonnet-4.5"} if got := p.WireFormat(model); got != core.WireFormatAnthropic { t.Fatalf("WireFormat(%q) = %v, want Anthropic", model.ModelID, got) @@ -131,7 +131,7 @@ func TestOpenCodeZenProvider_StreamAnthropic_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeZenProvider(atomic) - req := &core.NormalizedRequest{Model: "claude-sonnet-4.5", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, Stream: true} + req := &core.NormalizedRequest{Model: "claude-sonnet-4.5", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, Stream: true} model := config.ModelConfig{ModelID: "claude-sonnet-4.5"} body, err := p.Stream(context.Background(), req, model) @@ -161,7 +161,7 @@ func TestOpenCodeZenProvider_ExecuteResponses_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeZenProvider(atomic) - req := &core.NormalizedRequest{Model: "gpt-5.4", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}} + req := &core.NormalizedRequest{Model: "gpt-5.4", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}} model := config.ModelConfig{ModelID: "gpt-5.4"} if got := p.WireFormat(model); got != core.WireFormatOpenAIResponses { t.Fatalf("WireFormat(%q) = %v, want OpenAIResponses", model.ModelID, got) @@ -179,7 +179,7 @@ func TestOpenCodeGoProvider_Execute_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeGoProvider(atomic) - req := &core.NormalizedRequest{Model: "deepseek-v4-pro", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}} + req := &core.NormalizedRequest{Model: "deepseek-v4-pro", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}} model := config.ModelConfig{ModelID: "deepseek-v4-pro"} if got := p.WireFormat(model); got != core.WireFormatOpenAIChat { t.Fatalf("WireFormat(%q) = %v, want OpenAIChat", model.ModelID, got) @@ -197,7 +197,7 @@ func TestOpenCodeGoProvider_Stream_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeGoProvider(atomic) - req := &core.NormalizedRequest{Model: "deepseek-v4-pro", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, Stream: true} + req := &core.NormalizedRequest{Model: "deepseek-v4-pro", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, Stream: true} model := config.ModelConfig{ModelID: "deepseek-v4-pro"} body, err := p.Stream(context.Background(), req, model) @@ -226,7 +226,7 @@ func TestOpenCodeGoProvider_ExecuteAnthropic_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeGoProvider(atomic) - req := &core.NormalizedRequest{Model: "qwen3.5-plus", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}} + req := &core.NormalizedRequest{Model: "qwen3.5-plus", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}} model := config.ModelConfig{ModelID: "qwen3.5-plus"} if got := p.WireFormat(model); got != core.WireFormatAnthropic { t.Fatalf("WireFormat(%q) = %v, want Anthropic", model.ModelID, got) @@ -248,7 +248,7 @@ func TestOpenCodeGoProvider_StreamAnthropic_OpencodeUserAgent(t *testing.T) { atomic := config.NewAtomicConfig(cfg, "") p := NewOpenCodeGoProvider(atomic) - req := &core.NormalizedRequest{Model: "qwen3.5-plus", Messages: []core.NormalizedMessage{{Role: "user", Content: "Hi"}}, Stream: true} + req := &core.NormalizedRequest{Model: "qwen3.5-plus", Messages: []core.NormalizedMessage{{Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hi"}}}}, Stream: true} model := config.ModelConfig{ModelID: "qwen3.5-plus"} body, err := p.Stream(context.Background(), req, model) diff --git a/internal/provider/opencode_zen.go b/internal/provider/opencode_zen.go index 357ee427..7475e1e4 100644 --- a/internal/provider/opencode_zen.go +++ b/internal/provider/opencode_zen.go @@ -38,6 +38,12 @@ func NewOpenCodeZenProvider(atomic *config.AtomicConfig) *OpenCodeZenProvider { // Name returns the provider identifier. func (p *OpenCodeZenProvider) Name() string { return "opencode-zen" } +// ValidateRequest checks ordered content against this provider's capabilities +// before any upstream request is attempted. +func (p *OpenCodeZenProvider) ValidateRequest(req *core.NormalizedRequest, model config.ModelConfig) error { + return validateRequest(p, req, model) +} + // Capabilities returns provider-level capabilities. func (p *OpenCodeZenProvider) Capabilities() core.ProviderCapabilities { return core.ProviderCapabilities{ diff --git a/internal/provider/provider.go b/internal/provider/provider.go index 08295f0a..307e3e74 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -8,6 +8,7 @@ import ( "time" "github.com/routatic/proxy/internal/config" + "github.com/routatic/proxy/internal/core" ) const routaticUserAgent = "routatic-proxy" @@ -20,6 +21,14 @@ type baseProvider struct { keyCounter atomic.Uint64 } +func validateRequest(p core.Provider, req *core.NormalizedRequest, model config.ModelConfig) error { + caps, ok := p.ModelCapabilities(model.ModelID) + if !ok { + return &core.CompatibilityError{Provider: p.Name(), ModelID: model.ModelID, Reason: "model is not known by provider"} + } + return core.ValidateRequestCompatibility(req, model, caps, p.WireFormat(model)) +} + // newBaseProvider creates a baseProvider with a shared HTTP transport tuned // for high-concurrency upstream calls. func newBaseProvider(atomic *config.AtomicConfig) baseProvider { diff --git a/internal/router/model_router.go b/internal/router/model_router.go index 26cffd33..7c817b71 100644 --- a/internal/router/model_router.go +++ b/internal/router/model_router.go @@ -8,6 +8,7 @@ import ( "sort" "strings" "sync" + "sync/atomic" "time" "github.com/routatic/proxy/internal/catalog" @@ -18,25 +19,94 @@ import ( var ErrUnknownProvider = errors.New("unknown provider") type ModelRouter struct { - atomic *config.AtomicConfig - db *storage.Database - catalogPath string - catMu sync.Mutex - cat *catalog.IndexedCatalog - catErr error - catCache time.Time + atomic *config.AtomicConfig + db *storage.Database + catalogPath string + catMu sync.Mutex + cat atomic.Pointer[catalog.IndexedCatalog] + catErr error + catCache atomic.Int64 + refreshing atomic.Bool + refreshStop chan struct{} + refreshDone chan struct{} + updateSignal chan struct{} + refreshWG sync.WaitGroup } func NewModelRouter(atomic *config.AtomicConfig) *ModelRouter { - return &ModelRouter{atomic: atomic} + return newModelRouter(atomic, nil, "") } func NewModelRouterWithDB(atomic *config.AtomicConfig, db *storage.Database) *ModelRouter { - return &ModelRouter{atomic: atomic, db: db} + return newModelRouter(atomic, db, "") } func NewModelRouterWithCatalog(atomic *config.AtomicConfig, catalogPath string) *ModelRouter { - return &ModelRouter{atomic: atomic, catalogPath: catalogPath} + return newModelRouter(atomic, nil, catalogPath) +} + +func newModelRouter(atomic *config.AtomicConfig, db *storage.Database, catalogPath string) *ModelRouter { + return &ModelRouter{ + atomic: atomic, db: db, catalogPath: catalogPath, + updateSignal: make(chan struct{}, 1), + } +} + +// StartCatalogRefresh keeps a valid catalog snapshot warm in the background. +// Requests continue using the last valid snapshot while a refresh is running. +func (r *ModelRouter) StartCatalogRefresh(ctx context.Context) { + r.catMu.Lock() + if r.refreshStop != nil { + r.catMu.Unlock() + return + } + r.refreshStop = make(chan struct{}) + r.refreshDone = make(chan struct{}) + stop := r.refreshStop + done := r.refreshDone + updates := r.updateSignal + r.catMu.Unlock() + + go func() { + defer close(done) + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + r.refreshCatalogAsync(context.Background()) + for { + select { + case <-ctx.Done(): + return + case <-stop: + return + case <-ticker.C: + r.refreshCatalogAsync(context.Background()) + case <-updates: + r.refreshCatalogAsync(context.Background()) + } + } + }() +} + +// SignalCatalogUpdate asks the background refresher to publish a new snapshot. +func (r *ModelRouter) SignalCatalogUpdate() { + select { + case r.updateSignal <- struct{}{}: + default: + } +} + +// StopCatalogRefresh stops the background refresh loop. +func (r *ModelRouter) StopCatalogRefresh() { + r.catMu.Lock() + stop, done := r.refreshStop, r.refreshDone + r.refreshStop, r.refreshDone = nil, nil + r.catMu.Unlock() + if stop == nil { + return + } + close(stop) + <-done + r.refreshWG.Wait() } func (r *ModelRouter) catalog(ctx context.Context) (*catalog.IndexedCatalog, error) { @@ -45,25 +115,56 @@ func (r *ModelRouter) catalog(ctx context.Context) (*catalog.IndexedCatalog, err return nil, nil } + current := r.cat.Load() + lastRefresh := time.Unix(0, r.catCache.Load()) + if current != nil && time.Since(lastRefresh) < 30*time.Second { + return current, nil + } + if current != nil { + r.refreshCatalogAsync(ctx) + return current, nil + } + r.catMu.Lock() defer r.catMu.Unlock() - - if r.cat != nil && time.Since(r.catCache) < 30*time.Second { - return r.cat, nil + if current = r.cat.Load(); current != nil { + return current, nil } + return r.refreshCatalogLocked(ctx) +} +func (r *ModelRouter) refreshCatalogAsync(ctx context.Context) { + if !r.refreshing.CompareAndSwap(false, true) { + return + } + ctx = context.WithoutCancel(ctx) + r.refreshWG.Add(1) + go func() { + defer r.refreshWG.Done() + defer r.refreshing.Store(false) + r.catMu.Lock() + defer r.catMu.Unlock() + _, _ = r.refreshCatalogLocked(ctx) + }() +} + +func (r *ModelRouter) refreshCatalogLocked(ctx context.Context) (*catalog.IndexedCatalog, error) { ctx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() + var loaded *catalog.IndexedCatalog + var err error if r.db != nil { - r.cat, r.catErr = catalog.LoadFromSQLite(ctx, r.db) + loaded, err = catalog.LoadFromSQLite(ctx, r.db) } else if r.catalogPath != "" { - r.cat, r.catErr = catalog.Load(r.catalogPath) + loaded, err = catalog.Load(r.catalogPath) } - if r.catErr == nil { - r.catCache = time.Now() + r.catErr = err + if err == nil { + r.cat.Store(loaded) + r.catCache.Store(time.Now().UnixNano()) } - return r.cat, r.catErr + return loaded, err } // isRespectRequestedModel returns true when the client-specified model should be diff --git a/internal/router/policy.go b/internal/router/policy.go index 2003384b..66615073 100644 --- a/internal/router/policy.go +++ b/internal/router/policy.go @@ -158,7 +158,7 @@ func (p *ScenarioPolicy) Evaluate(ctx *EvaluationContext) ([]config.ModelConfig, messages = append(messages, MessageContent{Role: "system", Content: systemText}) } for _, msg := range ctx.Request.Messages { - messages = append(messages, MessageContent{Role: msg.Role, Content: msg.Content}) + messages = append(messages, MessageContent{Role: msg.Role, Content: msg.TextContent()}) } isStreaming := ctx.Request.Stream diff --git a/internal/router/selector.go b/internal/router/selector.go index ac8b2099..9fc3ac1d 100644 --- a/internal/router/selector.go +++ b/internal/router/selector.go @@ -3,7 +3,6 @@ package router import ( "errors" "fmt" - "sort" "github.com/routatic/proxy/internal/catalog" "github.com/routatic/proxy/internal/config" @@ -57,31 +56,17 @@ func (s *Selector) SelectCheapest(scenario string, constraints ScenarioConstrain return catalog.ResolvedModel{}, fmt.Errorf("unknown scenario %q", scenario) } - candidates := s.resolveCandidates(scen, constraints) - if len(candidates) == 0 { + best := s.resolveCandidates(scen, constraints) + if best == nil { return catalog.ResolvedModel{}, fmt.Errorf("%w: scenario %q", ErrNoCandidateModel, scenario) } - - sort.Slice(candidates, func(i, j int) bool { - a, b := candidates[i], candidates[j] - costA := a.CostInputPerM + a.CostOutputPerM + s.effectivePenalty(a.Provider) - costB := b.CostInputPerM + b.CostOutputPerM + s.effectivePenalty(b.Provider) - if costA != costB { - return costA < costB - } - if a.ContextWindow != b.ContextWindow { - return a.ContextWindow > b.ContextWindow - } - return a.ModelID < b.ModelID - }) - - return candidates[0], nil + return *best, nil } // resolveCandidates enumerates all enabled provider/model pairs for a scenario // and returns the resolved models that match the scenario requirements and // constraints. -func (s *Selector) resolveCandidates(scen catalog.Scenario, constraints ScenarioConstraints) []catalog.ResolvedModel { +func (s *Selector) resolveCandidates(scen catalog.Scenario, constraints ScenarioConstraints) *catalog.ResolvedModel { providers := s.providerSet(scen) minContext := max(scen.MinContextWindow, constraints.Context) @@ -90,40 +75,52 @@ func (s *Selector) resolveCandidates(scen catalog.Scenario, constraints Scenario maxContext = s.cfg.CostRouting.MaxContextWindow } - var candidates []catalog.ResolvedModel + var best *catalog.ResolvedModel for providerName := range providers { - provider, ok := s.catalog.Providers[providerName] - if !ok { - continue - } - for modelKey, model := range s.catalog.Models { - if !modelSupportsProvider(modelKey, providerName) { + for _, candidate := range s.catalog.ListProviderModels(providerName) { + if maxContext > 0 && candidate.ContextWindow > maxContext { continue } - if maxContext > 0 && model.ContextWindow() > maxContext { + if !resolvedModelMatches(candidate, scen, constraints, minContext) { continue } - if !modelMatches(model, scen, constraints, minContext) { - continue + if best == nil || s.betterResolvedModel(candidate, *best) { + copy := candidate + best = © } - candidates = append(candidates, catalog.ResolvedModel{ - Provider: provider.Name, - ModelID: catalog.ModelNameFromKey(modelKey), - CanonicalName: modelKey, - DisplayName: model.DisplayName(), - BaseURL: provider.BaseURL, - APIKey: provider.APIKey, - AnthropicToolsDisabled: provider.AnthropicToolsDisabled, - ContextWindow: model.ContextWindow(), - CostInputPerM: model.CostInputPerM(), - CostOutputPerM: model.CostOutputPerM(), - Tools: model.SupportsTools(), - Vision: model.SupportsVision(), - Reasoning: model.Reasoning, - }) } } - return candidates + return best +} + +func resolvedModelMatches(model catalog.ResolvedModel, scen catalog.Scenario, constraints ScenarioConstraints, minContext int64) bool { + if model.ContextWindow < minContext { + return false + } + if scen.RequiresTools != nil && *scen.RequiresTools && !model.Tools { + return false + } + if scen.RequiresVision != nil && *scen.RequiresVision && !model.Vision { + return false + } + if scen.RequiresReasoning != nil && *scen.RequiresReasoning && !model.Reasoning { + return false + } + return (!constraints.Tools || model.Tools) && + (!constraints.Vision || model.Vision) && + (!constraints.Reasoning || model.Reasoning) +} + +func (s *Selector) betterResolvedModel(a, b catalog.ResolvedModel) bool { + costA := a.CostInputPerM + a.CostOutputPerM + s.effectivePenalty(a.Provider) + costB := b.CostInputPerM + b.CostOutputPerM + s.effectivePenalty(b.Provider) + if costA != costB { + return costA < costB + } + if a.ContextWindow != b.ContextWindow { + return a.ContextWindow > b.ContextWindow + } + return a.ModelID < b.ModelID } // providerSet returns the enabled providers that should be considered for a @@ -191,35 +188,6 @@ func (s *Selector) globalPreferProviders() []string { return s.cfg.CostRouting.PreferProviders } -func modelSupportsProvider(modelKey string, provider string) bool { - return catalog.ProviderFromModelKey(modelKey) == provider -} - -func modelMatches(model catalog.Model, scen catalog.Scenario, constraints ScenarioConstraints, minContext int64) bool { - if model.ContextWindow() < minContext { - return false - } - if scen.RequiresTools != nil && *scen.RequiresTools && !model.SupportsTools() { - return false - } - if scen.RequiresVision != nil && *scen.RequiresVision && !model.SupportsVision() { - return false - } - if scen.RequiresReasoning != nil && *scen.RequiresReasoning && !model.Reasoning { - return false - } - if constraints.Tools && !model.SupportsTools() { - return false - } - if constraints.Vision && !model.SupportsVision() { - return false - } - if constraints.Reasoning && !model.Reasoning { - return false - } - return true -} - // enabledProviders returns the providers that have an effective API key in the // active config. A non-empty global API key enables all known providers. func enabledProviders(cfg *config.Config) map[string]bool { diff --git a/internal/server/server.go b/internal/server/server.go index eb30694e..9a9bbd3a 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -29,16 +29,18 @@ import ( // Server represents the proxy server. type Server struct { - atomic *config.AtomicConfig - httpSrv *http.Server - mux http.Handler - mu sync.Mutex - logger *slog.Logger - levelVar *slog.LevelVar - History *history.History // exported so the ui command can read it - metrics *metrics.Metrics // stored for Metrics() getter - storage *storage.Database - retention *storage.Retention + atomic *config.AtomicConfig + httpSrv *http.Server + mux http.Handler + mu sync.Mutex + logger *slog.Logger + levelVar *slog.LevelVar + History *history.History // exported so the ui command can read it + metrics *metrics.Metrics // stored for Metrics() getter + storage *storage.Database + retention *storage.Retention + storageWriter handlers.StorageWriter + modelRouter *router.ModelRouter } // NewServer creates a new proxy server. @@ -57,6 +59,7 @@ func NewServer(atomic *config.AtomicConfig, captureLogger *debug.CaptureLogger) if err != nil { return nil, fmt.Errorf("failed to create token counter: %w", err) } + configureTokenCache(tokenCounter, cfg) // Create metrics metrics := metrics.New() @@ -111,7 +114,7 @@ func NewServer(atomic *config.AtomicConfig, captureLogger *debug.CaptureLogger) // Create handlers. var storageWriter handlers.StorageWriter if db != nil { - storageWriter = handlers.NewStorageAdapter(db) + storageWriter = handlers.NewStorageAdapter(db, metrics) } messagesHandler := handlers.NewMessagesHandler( @@ -163,26 +166,40 @@ func NewServer(atomic *config.AtomicConfig, captureLogger *debug.CaptureLogger) } srv := &Server{ - atomic: atomic, - httpSrv: httpSrv, - mux: mux, - logger: logger, - levelVar: levelVar, - History: hist, - metrics: metrics, - storage: db, - retention: retention, + atomic: atomic, + httpSrv: httpSrv, + mux: mux, + logger: logger, + levelVar: levelVar, + History: hist, + metrics: metrics, + storage: db, + retention: retention, + storageWriter: storageWriter, + modelRouter: modelRouter, } // Register callback to update log level on config reload atomic.OnReload(func(newCfg *config.Config) { levelVar.Set(parseLogLevel(newCfg.Logging.Level)) + configureTokenCache(tokenCounter, newCfg) logger.Info("log level updated", "level", newCfg.Logging.Level) }) return srv, nil } +func configureTokenCache(counter *token.Counter, cfg *config.Config) { + if counter == nil || cfg == nil { + return + } + enabled := true + if cfg.Performance.TokenCountCacheEnabled != nil { + enabled = *cfg.Performance.TokenCountCacheEnabled + } + counter.ConfigureCache(enabled, cfg.Performance.TokenCountCacheCapacity) +} + // Metrics returns the in-process metrics collector. func (s *Server) Metrics() *metrics.Metrics { return s.metrics @@ -216,21 +233,19 @@ func (s *Server) Start() error { // Graceful shutdown. ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() + s.modelRouter.StartCatalogRefresh(context.Background()) go func() { <-ctx.Done() s.logger.Info("shutting down server...") + s.modelRouter.StopCatalogRefresh() + if s.retention != nil { s.retention.Stop() } - if s.storage != nil { - _ = s.storage.Close() - } - shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() s.mu.Lock() srvToShutdown := s.httpSrv @@ -241,6 +256,19 @@ func (s *Server) Start() error { s.logger.Error("server shutdown failed", "error", err) } } + cancel() + + if s.storageWriter != nil { + drainCtx, drainCancel := context.WithTimeout(context.Background(), 10*time.Second) + if err := s.storageWriter.Shutdown(drainCtx); err != nil { + s.logger.Warn("storage completion drain did not finish", "error", err) + } + drainCancel() + } + + if s.storage != nil { + _ = s.storage.Close() + } }() s.mu.Lock() @@ -258,12 +286,23 @@ func (s *Server) Start() error { // Shutdown gracefully shuts down the proxy server. func (s *Server) Shutdown(ctx context.Context) error { s.logger.Info("programmatic shutdown requested") + s.modelRouter.StopCatalogRefresh() s.mu.Lock() srvToShutdown := s.httpSrv s.mu.Unlock() if srvToShutdown != nil { - return srvToShutdown.Shutdown(ctx) + if err := srvToShutdown.Shutdown(ctx); err != nil { + return err + } + } + if s.storageWriter != nil { + if err := s.storageWriter.Shutdown(ctx); err != nil { + return err + } + } + if s.storage != nil { + return s.storage.Close() } return nil } diff --git a/internal/storage/latency.go b/internal/storage/latency.go index 64bac17e..1ea01992 100644 --- a/internal/storage/latency.go +++ b/internal/storage/latency.go @@ -2,6 +2,7 @@ package storage import ( "context" + "database/sql" "math" "sort" "time" @@ -37,6 +38,7 @@ type ModelLatencyStats struct { Avg time.Duration P50 time.Duration P90 time.Duration + P95 time.Duration P99 time.Duration Min time.Duration Max time.Duration @@ -48,9 +50,9 @@ func (l *Latency) GetStats(since time.Time) ([]ModelLatencyStats, error) { defer cancel() query := ` - SELECT model, latency_ms - FROM latency_samples - WHERE recorded_at >= ? + SELECT model, duration_ms + FROM requests + WHERE start_time >= ? AND success = 1 ORDER BY model ` @@ -63,11 +65,13 @@ func (l *Latency) GetStats(since time.Time) ([]ModelLatencyStats, error) { samplesByModel := make(map[string][]int64) for rows.Next() { var model string - var latencyMs int64 + var latencyMs sql.NullInt64 if err := rows.Scan(&model, &latencyMs); err != nil { return nil, err } - samplesByModel[model] = append(samplesByModel[model], latencyMs) + if latencyMs.Valid { + samplesByModel[model] = append(samplesByModel[model], latencyMs.Int64) + } } if err := rows.Err(); err != nil { @@ -153,6 +157,7 @@ func calculateStats(model string, samples []int64) ModelLatencyStats { p50Idx := int(float64(count)*0.50) - 1 p90Idx := int(math.Ceil(float64(count)*0.90)) - 1 + p95Idx := int(math.Ceil(float64(count)*0.95)) - 1 p99Idx := int(math.Ceil(float64(count)*0.99)) - 1 if p50Idx < 0 { p50Idx = 0 @@ -160,6 +165,9 @@ func calculateStats(model string, samples []int64) ModelLatencyStats { if p90Idx < 0 { p90Idx = 0 } + if p95Idx < 0 { + p95Idx = 0 + } if p99Idx < 0 { p99Idx = 0 } @@ -169,6 +177,9 @@ func calculateStats(model string, samples []int64) ModelLatencyStats { if p90Idx >= count { p90Idx = count - 1 } + if p95Idx >= count { + p95Idx = count - 1 + } if p99Idx >= count { p99Idx = count - 1 } @@ -179,6 +190,7 @@ func calculateStats(model string, samples []int64) ModelLatencyStats { Avg: time.Duration(avg) * time.Millisecond, P50: time.Duration(sorted[p50Idx]) * time.Millisecond, P90: time.Duration(sorted[p90Idx]) * time.Millisecond, + P95: time.Duration(sorted[p95Idx]) * time.Millisecond, P99: time.Duration(sorted[p99Idx]) * time.Millisecond, Min: time.Duration(sorted[0]) * time.Millisecond, Max: time.Duration(sorted[count-1]) * time.Millisecond, diff --git a/internal/token/counter.go b/internal/token/counter.go index 4c499007..2b978fde 100644 --- a/internal/token/counter.go +++ b/internal/token/counter.go @@ -2,9 +2,12 @@ package token import ( + "container/list" + "crypto/sha256" "fmt" "os" "path/filepath" + "sync" "github.com/pkoukk/tiktoken-go" ) @@ -12,6 +15,67 @@ import ( // Counter handles token counting for text and message arrays. type Counter struct { tiktoken *tiktoken.Tiktoken + encoding string + cacheMu sync.Mutex + cache *tokenCache +} + +const defaultCacheCapacity = 8192 + +type tokenCache struct { + enabled bool + capacity int + ll *list.List + items map[[32]byte]*list.Element +} + +type tokenCacheEntry struct { + key [32]byte + count int +} + +func newTokenCache(enabled bool, capacity int) *tokenCache { + if capacity <= 0 { + capacity = defaultCacheCapacity + } + return &tokenCache{ + enabled: enabled, + capacity: capacity, + ll: list.New(), + items: make(map[[32]byte]*list.Element, capacity), + } +} + +func (c *tokenCache) get(key [32]byte) (int, bool) { + if !c.enabled { + return 0, false + } + elem, ok := c.items[key] + if !ok { + return 0, false + } + c.ll.MoveToFront(elem) + return elem.Value.(tokenCacheEntry).count, true +} + +func (c *tokenCache) put(key [32]byte, count int) { + if !c.enabled { + return + } + if elem, ok := c.items[key]; ok { + elem.Value = tokenCacheEntry{key: key, count: count} + c.ll.MoveToFront(elem) + return + } + elem := c.ll.PushFront(tokenCacheEntry{key: key, count: count}) + c.items[key] = elem + if c.ll.Len() > c.capacity { + oldest := c.ll.Back() + if oldest != nil { + c.ll.Remove(oldest) + delete(c.items, oldest.Value.(tokenCacheEntry).key) + } + } } // defaultCacheDir returns a user-writable cache directory for tiktoken files. @@ -41,13 +105,57 @@ func NewCounter() (*Counter, error) { if err != nil { return nil, fmt.Errorf("failed to get encoding: %w", err) } - return &Counter{tiktoken: enc}, nil + return &Counter{ + tiktoken: enc, + encoding: "cl100k_base", + cache: newTokenCache(true, defaultCacheCapacity), + }, nil +} + +// ConfigureCache atomically replaces the token count cache. Replacing the +// cache drops old entries so a config reload cannot retain stale state. +func (c *Counter) ConfigureCache(enabled bool, capacity int) { + c.cacheMu.Lock() + c.cache = newTokenCache(enabled, capacity) + c.cacheMu.Unlock() +} + +// CacheStats returns the current cache size, capacity, and hit/miss counters. +// It is intended for diagnostics and tests. +func (c *Counter) CacheStats() (size, capacity int, enabled bool) { + c.cacheMu.Lock() + defer c.cacheMu.Unlock() + if c.cache == nil { + return 0, 0, false + } + return c.cache.ll.Len(), c.cache.capacity, c.cache.enabled } // CountTokens counts tokens in a string. func (c *Counter) CountTokens(text string) (int, error) { + encoding := c.encoding + if encoding == "" { + encoding = "unknown" + } + key := sha256.Sum256(append([]byte(encoding+"\x00"), []byte(text)...)) + c.cacheMu.Lock() + if c.cache != nil { + if count, ok := c.cache.get(key); ok { + c.cacheMu.Unlock() + return count, nil + } + } + c.cacheMu.Unlock() + tokens := c.tiktoken.Encode(text, nil, nil) - return len(tokens), nil + count := len(tokens) + + c.cacheMu.Lock() + if c.cache != nil { + c.cache.put(key, count) + } + c.cacheMu.Unlock() + return count, nil } // MessageContent represents a single message in a conversation. diff --git a/internal/token/counter_test.go b/internal/token/counter_test.go index 609c2ca7..cc312ec7 100644 --- a/internal/token/counter_test.go +++ b/internal/token/counter_test.go @@ -87,3 +87,44 @@ func mustHome() string { } return h } + +func TestCountTokensUsesBoundedCache(t *testing.T) { + counter, err := NewCounter() + if err != nil { + t.Fatalf("NewCounter: %v", err) + } + counter.ConfigureCache(true, 2) + + first, err := counter.CountTokens("repeated prompt") + if err != nil { + t.Fatalf("first CountTokens: %v", err) + } + second, err := counter.CountTokens("repeated prompt") + if err != nil { + t.Fatalf("second CountTokens: %v", err) + } + if first != second { + t.Fatalf("cached count changed: first=%d second=%d", first, second) + } + + size, capacity, enabled := counter.CacheStats() + if !enabled || capacity != 2 || size != 1 { + t.Fatalf("unexpected cache stats: size=%d capacity=%d enabled=%v", size, capacity, enabled) + } +} + +func TestConfigureCacheReplacesCache(t *testing.T) { + counter, err := NewCounter() + if err != nil { + t.Fatalf("NewCounter: %v", err) + } + if _, err := counter.CountTokens("old prompt"); err != nil { + t.Fatalf("CountTokens: %v", err) + } + + counter.ConfigureCache(false, 10) + size, capacity, enabled := counter.CacheStats() + if enabled || capacity != 10 || size != 0 { + t.Fatalf("cache was not replaced: size=%d capacity=%d enabled=%v", size, capacity, enabled) + } +} diff --git a/internal/transformer/cache_control_test.go b/internal/transformer/cache_control_test.go new file mode 100644 index 00000000..8cfc1ac3 --- /dev/null +++ b/internal/transformer/cache_control_test.go @@ -0,0 +1,50 @@ +package transformer + +import ( + "encoding/json" + "testing" + + "github.com/routatic/proxy/internal/config" + "github.com/routatic/proxy/pkg/types" +) + +func TestTransformRequestPreservesCacheControlOnNonTextBlocks(t *testing.T) { + req := &types.MessageRequest{ + Model: "deepseek-v4-pro", + MaxTokens: 100, + Messages: []types.Message{ + { + Role: "user", + Content: json.RawMessage(`[ + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"abc"},"cache_control":{"type":"ephemeral"}}, + {"type":"tool_result","tool_use_id":"toolu_1","content":"done","cache_control":{"type":"ephemeral"}} + ]`), + }, + { + Role: "assistant", + Content: json.RawMessage(`[{"type":"thinking","thinking":"reason","cache_control":{"type":"ephemeral"}}]`), + }, + }, + } + + out, err := NewRequestTransformer().TransformRequest(req, config.ModelConfig{ + ModelID: "deepseek-v4-pro", + Vision: true, + }) + if err != nil { + t.Fatalf("TransformRequest: %v", err) + } + + if len(out.Messages) < 3 { + t.Fatalf("got %d messages, want image, tool, and assistant messages", len(out.Messages)) + } + if out.Messages[0].CacheControl == nil { + t.Fatal("image cache directive was not preserved on the user message") + } + if out.Messages[1].CacheControl == nil { + t.Fatal("tool-result cache directive was not preserved") + } + if out.Messages[2].CacheControl == nil { + t.Fatal("thinking cache directive was not preserved on the assistant message") + } +} diff --git a/internal/transformer/normalized_bridge.go b/internal/transformer/normalized_bridge.go index 51f0337a..efaab093 100644 --- a/internal/transformer/normalized_bridge.go +++ b/internal/transformer/normalized_bridge.go @@ -58,13 +58,11 @@ func NormalizedToResponses(req *core.NormalizedRequest, model config.ModelConfig // Convert messages. for _, msg := range req.Messages { input := types.ResponsesInput{Role: msg.Role} - content := msg.Content + content := msg.TextContent() // For assistant messages with tool calls, serialize as text. - if len(msg.ToolCalls) > 0 { - for _, tc := range msg.ToolCalls { - content += "[Tool: " + tc.Name + "(" + tc.Arguments + ")]" - } + for _, tc := range msg.ToolCallsList() { + content += "[Tool: " + tc.Name + "(" + tc.Arguments + ")]" } if content != "" { @@ -110,7 +108,7 @@ func NormalizedToGemini(req *core.NormalizedRequest, model config.ModelConfig) * // Convert messages. for _, msg := range req.Messages { gc := types.GeminiContent{Role: msg.Role} - gc.Parts = append(gc.Parts, types.GeminiPart{Text: msg.Content}) + gc.Parts = append(gc.Parts, types.GeminiPart{Text: msg.TextContent()}) contents = append(contents, gc) } @@ -150,20 +148,23 @@ func OpenAIResponseToNormalized(openaiResp *types.ChatCompletionResponse, modelI // Extract text content. if msg.Content != nil { - nm.Content = msg.ContentText() + nm.Blocks = append(nm.Blocks, core.NormalizedContentBlock{ + Type: "text", Text: msg.ContentText(), + }) } // Extract reasoning content (pointer field). if msg.ReasoningContent != nil { - nm.Thinking = *msg.ReasoningContent + nm.Blocks = append(nm.Blocks, core.NormalizedContentBlock{ + Type: "thinking", Thinking: *msg.ReasoningContent, + }) } // Extract tool calls. for _, tc := range msg.ToolCalls { - nm.ToolCalls = append(nm.ToolCalls, core.NormalizedToolCall{ - ID: tc.ID, - Name: tc.Function.Name, - Arguments: tc.Function.Arguments, + nm.Blocks = append(nm.Blocks, core.NormalizedContentBlock{ + Type: "tool_use", ID: tc.ID, Name: tc.Function.Name, + Input: []byte(tc.Function.Arguments), }) } @@ -208,20 +209,19 @@ func ResponsesToNormalized(responsesResp *types.ResponsesResponse, modelID strin nm := core.NormalizedMessage{Role: output.Role} for _, c := range output.Content { if c.Type == "output_text" { - nm.Content += c.Text + nm.Blocks = append(nm.Blocks, core.NormalizedContentBlock{ + Type: "text", Text: c.Text, + }) } } nr.Messages = append(nr.Messages, nm) case "function_call": nm := core.NormalizedMessage{ Role: "assistant", - ToolCalls: []core.NormalizedToolCall{ - { - ID: output.CallID, - Name: output.Name, - Arguments: output.Arguments, - }, - }, + Blocks: []core.NormalizedContentBlock{{ + Type: "tool_use", ID: output.CallID, Name: output.Name, + Input: []byte(output.Arguments), + }}, } nr.Messages = append(nr.Messages, nm) } @@ -249,7 +249,9 @@ func GeminiToNormalized(geminiResp *types.GeminiResponse, modelID string) *core. for _, part := range candidate.Content.Parts { if part.Text != "" { - nm.Content += part.Text + nm.Blocks = append(nm.Blocks, core.NormalizedContentBlock{ + Type: "text", Text: part.Text, + }) } } @@ -282,12 +284,17 @@ func GeminiToNormalized(geminiResp *types.GeminiResponse, modelID string) *core. // pipeline. func normalizedToMessageRequest(req *core.NormalizedRequest) *types.MessageRequest { anthropicReq := &types.MessageRequest{ - Model: req.Model, - MaxTokens: req.MaxTokens, + Model: req.Model, + MaxTokens: req.MaxTokens, + CacheControl: req.CacheControl, } // Set system prompt. - if req.SystemPrompt != "" { + if len(req.SystemBlocks) > 0 { + if b, err := json.Marshal(normalizedBlocksToAnthropic(req.SystemBlocks)); err == nil { + anthropicReq.System = b + } + } else if req.SystemPrompt != "" { if b, err := json.Marshal(req.SystemPrompt); err == nil { anthropicReq.System = json.RawMessage(b) } @@ -319,51 +326,7 @@ func normalizedToMessageRequest(req *core.NormalizedRequest) *types.MessageReque for _, nm := range req.Messages { msg := types.Message{Role: nm.Role} - var blocks []types.ContentBlock - if nm.Content != "" { - blocks = append(blocks, types.ContentBlock{Type: "text", Text: nm.Content}) - } - // Reconstruct image blocks from the normalized representation so the - // downstream transformer can decide whether to convert them to - // image_url (vision-capable model) or to a [Image] text placeholder. - for _, img := range nm.Images { - blocks = append(blocks, types.ContentBlock{ - Type: "image", - Source: &types.ImageSource{ - Type: "base64", - MediaType: img.MediaType, - Data: img.Data, - }, - }) - } - if nm.Thinking != "" { - blocks = append(blocks, types.ContentBlock{Type: "thinking", Thinking: nm.Thinking}) - } - for _, tc := range nm.ToolCalls { - blocks = append(blocks, types.ContentBlock{ - Type: "tool_use", - ID: tc.ID, - Name: tc.Name, - Input: []byte(tc.Arguments), - }) - } - if len(nm.ToolResults) > 0 { - for _, tr := range nm.ToolResults { - content, _ := json.Marshal(tr.Content) - blocks = append(blocks, types.ContentBlock{ - Type: "tool_result", - ToolUseID: tr.ToolCallID, - Content: content, - }) - } - } else if nm.ToolCallID != "" { - content, _ := json.Marshal(nm.Content) - blocks = append(blocks, types.ContentBlock{ - Type: "tool_result", - ToolUseID: nm.ToolCallID, - Content: content, - }) - } + blocks := normalizedMessageBlocks(nm) if len(blocks) > 0 { b, _ := json.Marshal(blocks) @@ -378,15 +341,39 @@ func normalizedToMessageRequest(req *core.NormalizedRequest) *types.MessageReque // Convert tools. for _, nt := range req.Tools { anthropicReq.Tools = append(anthropicReq.Tools, types.Tool{ - Name: nt.Name, - Description: nt.Description, - InputSchema: nt.InputSchema, + Name: nt.Name, + Description: nt.Description, + InputSchema: nt.InputSchema, + CacheControl: nt.CacheControl, }) } return anthropicReq } +func normalizedMessageBlocks(nm core.NormalizedMessage) []types.ContentBlock { + return normalizedBlocksToAnthropic(nm.Blocks) +} + +func normalizedBlocksToAnthropic(blocks []core.NormalizedContentBlock) []types.ContentBlock { + out := make([]types.ContentBlock, 0, len(blocks)) + for _, block := range blocks { + converted := types.ContentBlock{ + Type: block.Type, Text: block.Text, ID: block.ID, ToolUseID: block.ToolUseID, + Name: block.Name, Input: block.Input, Content: block.Content, + IsError: block.IsError, Thinking: block.Thinking, Signature: block.Signature, + CacheControl: block.CacheControl, Raw: block.Raw, + } + if block.Image != nil { + converted.Source = &types.ImageSource{ + Type: "base64", MediaType: block.Image.MediaType, Data: block.Image.Data, + } + } + out = append(out, converted) + } + return out +} + func rawJSONString(s string) json.RawMessage { b, err := json.Marshal(s) if err != nil { @@ -400,11 +387,11 @@ func rawJSONString(s string) json.RawMessage { func joinMessageText(messages []core.NormalizedMessage) string { var text string for _, m := range messages { - if m.Content != "" { + if content := m.TextContent(); content != "" { if text != "" { text += "\n" } - text += m.Role + ": " + m.Content + text += m.Role + ": " + content } } return text diff --git a/internal/transformer/normalized_bridge_test.go b/internal/transformer/normalized_bridge_test.go index 5384eb56..34ce86b1 100644 --- a/internal/transformer/normalized_bridge_test.go +++ b/internal/transformer/normalized_bridge_test.go @@ -14,7 +14,7 @@ func TestNormalizedToAnthropic_SystemPromptWithNewline(t *testing.T) { SystemPrompt: "Line one\nLine two\nLine three", MaxTokens: 100, Messages: []core.NormalizedMessage{ - {Role: "user", Content: "Hello"}, + {Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hello"}}}, }, } @@ -38,7 +38,7 @@ func TestNormalizedToAnthropic_MessageContentWithNewline(t *testing.T) { Model: "minimax-m3", MaxTokens: 100, Messages: []core.NormalizedMessage{ - {Role: "user", Content: "Hello\nWorld"}, + {Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hello\nWorld"}}}, }, } @@ -61,7 +61,7 @@ func TestNormalizedToResponses_SystemPromptWithNewline(t *testing.T) { SystemPrompt: "Line one\nLine two", MaxTokens: 100, Messages: []core.NormalizedMessage{ - {Role: "user", Content: "Hello\nWorld"}, + {Role: "user", Blocks: []core.NormalizedContentBlock{{Type: "text", Text: "Hello\nWorld"}}}, }, } @@ -88,7 +88,7 @@ func TestNormalizedToResponses_SystemPromptWithNewline(t *testing.T) { if err := json.Unmarshal(responsesReq.Input[1].Content, &messageContent); err != nil { t.Fatalf("message content was not valid JSON: %v", err) } - if messageContent != req.Messages[0].Content { - t.Fatalf("message content mismatch: got %q, want %q", messageContent, req.Messages[0].Content) + if messageContent != req.Messages[0].TextContent() { + t.Fatalf("message content mismatch: got %q, want %q", messageContent, req.Messages[0].TextContent()) } } diff --git a/internal/transformer/request.go b/internal/transformer/request.go index 8a0cb7b8..bdcc8df0 100644 --- a/internal/transformer/request.go +++ b/internal/transformer/request.go @@ -70,6 +70,16 @@ func constrainTemperature(modelID string, temp float64) float64 { func stripCacheControl(messages []types.ChatMessage) { for i := range messages { messages[i].CacheControl = nil + var parts []types.ChatContentPart + if err := json.Unmarshal(messages[i].Content, &parts); err != nil { + continue + } + for j := range parts { + parts[j].CacheControl = nil + } + if content, err := json.Marshal(parts); err == nil { + messages[i].Content = content + } } } @@ -140,6 +150,11 @@ func (t *RequestTransformer) TransformRequest( // Transform tools if present if len(anthropicReq.Tools) > 0 { openaiReq.Tools = t.transformTools(anthropicReq.Tools) + if !isDeepSeekModel(model.ModelID) { + for i := range openaiReq.Tools { + openaiReq.Tools[i].CacheControl = nil + } + } } return openaiReq, nil @@ -345,6 +360,9 @@ func (t *RequestTransformer) transformMessages(anthropicReq *types.MessageReques } } } + if systemMsg.CacheControl == nil { + systemMsg.CacheControl = anthropicReq.CacheControl + } result = append(result, systemMsg) } @@ -431,8 +449,12 @@ func (t *RequestTransformer) transformUserMessage(blocks []types.ContentBlock, v var textParts []string var imageParts []types.ChatContentPart hasImage := false + var messageCacheControl *types.CacheControl for _, block := range blocks { + if messageCacheControl == nil { + messageCacheControl = block.CacheControl + } switch block.Type { case "text": textParts = append(textParts, block.Text) @@ -440,9 +462,10 @@ func (t *RequestTransformer) transformUserMessage(blocks []types.ContentBlock, v // In OpenAI, tool results are separate messages with role "tool" toolContent := block.TextContent() result = append(result, types.ChatMessage{ - Role: "tool", - Content: contentText(toolContent), - ToolCallID: block.GetToolID(), + Role: "tool", + Content: contentText(toolContent), + ToolCallID: block.GetToolID(), + CacheControl: block.CacheControl, }) case "image": if block.Source != nil { @@ -452,6 +475,7 @@ func (t *RequestTransformer) transformUserMessage(blocks []types.ContentBlock, v ImageURL: &types.ImageURL{ URL: fmt.Sprintf("data:%s;base64,%s", block.Source.MediaType, block.Source.Data), }, + CacheControl: block.CacheControl, }) } else { hasImage = true @@ -471,8 +495,9 @@ func (t *RequestTransformer) transformUserMessage(blocks []types.ContentBlock, v var parts []types.ChatContentPart if len(textParts) > 0 { parts = append(parts, types.ChatContentPart{ - Type: "text", - Text: strings.Join(textParts, ""), + Type: "text", + Text: strings.Join(textParts, ""), + CacheControl: messageCacheControl, }) } parts = append(parts, imageParts...) @@ -480,7 +505,9 @@ func (t *RequestTransformer) transformUserMessage(blocks []types.ContentBlock, v if err != nil { return nil, fmt.Errorf("failed to marshal multimodal content: %w", err) } - result = append(result, types.ChatMessage{Role: "user", Content: contentJSON}) + result = append(result, types.ChatMessage{ + Role: "user", Content: contentJSON, CacheControl: messageCacheControl, + }) } else { // Text-only message (possibly with image placeholder for non-vision models) text := strings.Join(textParts, "") @@ -492,8 +519,9 @@ func (t *RequestTransformer) transformUserMessage(blocks []types.ContentBlock, v } } result = append(result, types.ChatMessage{ - Role: "user", - Content: contentText(text), + Role: "user", + Content: contentText(text), + CacheControl: messageCacheControl, }) } } @@ -506,8 +534,12 @@ func (t *RequestTransformer) transformAssistantMessage(blocks []types.ContentBlo var textParts []string var thinkingParts []string var toolCalls []types.ToolCall + var messageCacheControl *types.CacheControl for _, block := range blocks { + if messageCacheControl == nil { + messageCacheControl = block.CacheControl + } switch block.Type { case "text": textParts = append(textParts, block.Text) @@ -581,6 +613,7 @@ func (t *RequestTransformer) transformAssistantMessage(blocks []types.ContentBlo Content: contentText(content), ReasoningContent: reasoningContentPtr, ToolCalls: toolCalls, + CacheControl: messageCacheControl, } return []types.ChatMessage{msg}, nil @@ -635,7 +668,8 @@ func (t *RequestTransformer) transformTools(tools []types.Tool) []types.ToolDef } result = append(result, types.ToolDef{ - Type: "function", + Type: "function", + CacheControl: tool.CacheControl, Function: types.FunctionDef{ Name: tool.Name, Description: tool.Description, diff --git a/pkg/types/anthropic.go b/pkg/types/anthropic.go index 8295e950..eaae6d78 100644 --- a/pkg/types/anthropic.go +++ b/pkg/types/anthropic.go @@ -11,16 +11,17 @@ import ( // MessageRequest represents a request to the Anthropic Messages API. type MessageRequest struct { - Model string `json:"model"` - MaxTokens int `json:"max_tokens"` - System json.RawMessage `json:"system,omitempty"` - Messages []Message `json:"messages"` - Stream *bool `json:"stream,omitempty"` - Tools []Tool `json:"tools,omitempty"` - Temperature *float64 `json:"temperature,omitempty"` - TopP *float64 `json:"top_p,omitempty"` - Metadata *Metadata `json:"metadata,omitempty"` - Thinking json.RawMessage `json:"thinking,omitempty"` + Model string `json:"model"` + MaxTokens int `json:"max_tokens"` + System json.RawMessage `json:"system,omitempty"` + Messages []Message `json:"messages"` + Stream *bool `json:"stream,omitempty"` + Tools []Tool `json:"tools,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + Metadata *Metadata `json:"metadata,omitempty"` + Thinking json.RawMessage `json:"thinking,omitempty"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } // SystemText extracts the system prompt text from the raw system field. @@ -108,18 +109,33 @@ func (m *Message) ContentBlocks() []ContentBlock { // - "thinking": Thinking, Signature are populated // - "image": Source is populated type ContentBlock struct { - Type string `json:"type"` - Text string `json:"text"` - ID string `json:"id,omitempty"` // For tool_use (the tool call ID) - ToolUseID string `json:"tool_use_id,omitempty"` // For tool_result (references the tool_use ID) - Name string `json:"name,omitempty"` - Input json.RawMessage `json:"input,omitempty"` - Output json.RawMessage `json:"output,omitempty"` // Deprecated: use Content - Content json.RawMessage `json:"content,omitempty"` // For tool_result inner content - IsError *bool `json:"is_error,omitempty"` // For tool_result - Thinking string `json:"thinking"` // For thinking blocks - Signature string `json:"signature,omitempty"` // For thinking blocks - Source *ImageSource `json:"source,omitempty"` // For image blocks + Type string `json:"type"` + Text string `json:"text"` + ID string `json:"id,omitempty"` // For tool_use (the tool call ID) + ToolUseID string `json:"tool_use_id,omitempty"` // For tool_result (references the tool_use ID) + Name string `json:"name,omitempty"` + Input json.RawMessage `json:"input,omitempty"` + Output json.RawMessage `json:"output,omitempty"` // Deprecated: use Content + Content json.RawMessage `json:"content,omitempty"` // For tool_result inner content + IsError *bool `json:"is_error,omitempty"` // For tool_result + Thinking string `json:"thinking"` // For thinking blocks + Signature string `json:"signature,omitempty"` // For thinking blocks + Source *ImageSource `json:"source,omitempty"` // For image blocks + CacheControl *CacheControl `json:"cache_control,omitempty"` + Raw json.RawMessage `json:"-"` +} + +// UnmarshalJSON preserves the original block bytes so unknown fields survive +// normalization and provider transformation. +func (b *ContentBlock) UnmarshalJSON(data []byte) error { + type alias ContentBlock + var decoded alias + if err := json.Unmarshal(data, &decoded); err != nil { + return err + } + *b = ContentBlock(decoded) + b.Raw = append([]byte(nil), data...) + return nil } // GetToolID returns the appropriate tool ID for this block type. @@ -168,19 +184,20 @@ func (b ContentBlock) MarshalJSON() ([]byte, error) { switch b.Type { case "text": type TextBlock struct { - Type string `json:"type"` - Text string `json:"text"` + Type string `json:"type"` + Text string `json:"text"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } return json.Marshal(TextBlock{ - Type: b.Type, - Text: b.Text, + Type: b.Type, Text: b.Text, CacheControl: b.CacheControl, }) case "tool_use": type ToolUseBlock struct { - Type string `json:"type"` - ID string `json:"id"` - Name string `json:"name"` - Input json.RawMessage `json:"input"` + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + Input json.RawMessage `json:"input"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } input := b.Input if len(input) == 0 { @@ -190,42 +207,51 @@ func (b ContentBlock) MarshalJSON() ([]byte, error) { Type: b.Type, ID: b.ID, Name: b.Name, - Input: input, + Input: input, CacheControl: b.CacheControl, }) case "tool_result": type ToolResultBlock struct { - Type string `json:"type"` - ToolUseID string `json:"tool_use_id"` - Content json.RawMessage `json:"content,omitempty"` - IsError *bool `json:"is_error,omitempty"` + Type string `json:"type"` + ToolUseID string `json:"tool_use_id"` + Content json.RawMessage `json:"content,omitempty"` + IsError *bool `json:"is_error,omitempty"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } return json.Marshal(ToolResultBlock{ - Type: b.Type, - ToolUseID: b.ToolUseID, - Content: b.Content, - IsError: b.IsError, + Type: b.Type, + ToolUseID: b.ToolUseID, + Content: b.Content, + IsError: b.IsError, + CacheControl: b.CacheControl, }) case "thinking": type ThinkingBlock struct { - Type string `json:"type"` - Thinking string `json:"thinking"` - Signature string `json:"signature,omitempty"` + Type string `json:"type"` + Thinking string `json:"thinking"` + Signature string `json:"signature,omitempty"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } return json.Marshal(ThinkingBlock{ - Type: b.Type, - Thinking: b.Thinking, - Signature: b.Signature, + Type: b.Type, + Thinking: b.Thinking, + Signature: b.Signature, + CacheControl: b.CacheControl, }) case "image": type ImageBlock struct { - Type string `json:"type"` - Source *ImageSource `json:"source"` + Type string `json:"type"` + Source *ImageSource `json:"source"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } return json.Marshal(ImageBlock{ - Type: b.Type, - Source: b.Source, + Type: b.Type, + Source: b.Source, + CacheControl: b.CacheControl, }) default: + if len(b.Raw) > 0 { + return append([]byte(nil), b.Raw...), nil + } type Alias ContentBlock return json.Marshal(Alias(b)) } @@ -241,9 +267,10 @@ type ImageSource struct { // Tool represents a tool definition for function calling. type Tool struct { - Name string `json:"name"` - Description string `json:"description,omitempty"` - InputSchema json.RawMessage `json:"input_schema"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + InputSchema json.RawMessage `json:"input_schema"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } // ToolResult represents the result of a tool execution. diff --git a/pkg/types/openai.go b/pkg/types/openai.go index 0dbc0c73..9f8d7032 100644 --- a/pkg/types/openai.go +++ b/pkg/types/openai.go @@ -8,9 +8,10 @@ import "encoding/json" // ChatContentPart represents a single part in a multimodal message content array. type ChatContentPart struct { - Type string `json:"type"` - Text string `json:"text,omitempty"` - ImageURL *ImageURL `json:"image_url,omitempty"` + Type string `json:"type"` + Text string `json:"text,omitempty"` + ImageURL *ImageURL `json:"image_url,omitempty"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } // ImageURL represents the URL source for an image in a multimodal message. @@ -103,8 +104,9 @@ type FunctionCall struct { // ToolDef represents a tool definition for function calling. type ToolDef struct { - Type string `json:"type"` - Function FunctionDef `json:"function"` + Type string `json:"type"` + Function FunctionDef `json:"function"` + CacheControl *CacheControl `json:"cache_control,omitempty"` } // FunctionDef represents the function definition schema.