From c1af290e737865ba73a759da7d32bc52bc67d4ea Mon Sep 17 00:00:00 2001 From: FFengIll Date: Wed, 5 Aug 2026 13:34:09 +0800 Subject: [PATCH 1/8] docs(protocol-stage): port protocol-stage / recording / harness-matrix design notes Front-load the design and plan documents for the Protocol Stage feature port (from codex/protocol-stage-hardening) as the first commit of the branch, so every subsequent code commit is traceable to a written plan: - .design/protocol-stage-chain.md: composable protocol stage chain rationale - .design/protocol-stage-tool-loop.md: stage tool-loop semantics (max rounds, side-effect commit, failover replay safety) - .design/protocol-recording-redesign.md: recording boundary redesign - .design/harness-matrix.md: cross-stage test matrix layout - docs/guardrails.md, README.md, cli/harness/README.md: user-facing notes No code changes; documentation only. --- .design/harness-matrix.md | 81 ++- .design/protocol-recording-redesign.md | 434 +++++++++++++ .design/protocol-stage-chain.md | 823 +++++++++++++++++++++++++ .design/protocol-stage-tool-loop.md | 463 ++++++++++++++ README.md | 5 + cli/harness/README.md | 33 +- docs/guardrails.md | 22 + 7 files changed, 1844 insertions(+), 17 deletions(-) create mode 100644 .design/protocol-recording-redesign.md create mode 100644 .design/protocol-stage-chain.md create mode 100644 .design/protocol-stage-tool-loop.md diff --git a/.design/harness-matrix.md b/.design/harness-matrix.md index 4c011ef9e..701f51c2f 100644 --- a/.design/harness-matrix.md +++ b/.design/harness-matrix.md @@ -213,28 +213,29 @@ go test -tags e2e ./internal/protocoltest/... -run TestContentShapes executor (`ExecuteAll*`) so the CLI can run it directly — including idempotence and the rule-flag suite, which would otherwise be go-test-only. -| `--mode` | single (A→B) | transitive (A→B→C) | idempotent (`g(f(A))==A`) | flags (per-rule) | content_shapes (§10.1) | cache_controls (§10.2) | -|----------|:---:|:---:|:---:|:---:|:---:|:---:| -| `default` *(no flag)* | ✅ | — | ✅ | — | — | — | -| `all` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | -| `single` | ✅ | — | — | — | — | — | -| `transitive` | — | ✅ | — | — | — | — | -| `idempotent` | — | — | ✅ | — | — | — | -| `flags` | — | — | — | ✅ | — | — | -| `content_shapes` | — | — | — | — | ✅ | — | -| `cache_controls` | — | — | — | — | — | ✅ | +| `--mode` | single (A→B) | transitive (A→B→C) | idempotent (`g(f(A))==A`) | flags (per-rule) | content_shapes (§10.1) | cache_controls (§10.2) | dormant Bridges | +|----------|:---:|:---:|:---:|:---:|:---:|:---:|:---:| +| `default` *(no flag)* | ✅ | — | ✅ | — | — | — | ✅ | +| `all` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| `single` | ✅ | — | — | — | — | — | — | +| `transitive` | — | ✅ | — | — | — | — | — | +| `idempotent` | — | — | ✅ | — | — | — | — | +| `flags` | — | — | — | ✅ | — | — | — | +| `content_shapes` | — | — | — | — | ✅ | — | — | +| `cache_controls` | — | — | — | — | — | ✅ | — | +| `bridges` | — | — | — | — | — | — | ✅ | This mode → section mapping is declared in one place: the `matrixSections` registry in `cli/harness/matrix.go`. Each entry names the section, lists the `--mode` values that include it, marks whether it is http-only (`flags`, -`content_shapes`, and `cache_controls` drive raw requests directly), and points -at its `ExecuteAll*` executor. Adding a section = one registry entry + extending -the `--mode` enum (see §8 for why this replaced a hand-maintained if-chain). +`content_shapes`, and `cache_controls` drive raw requests directly, while +`bridges` runs in-process), and points at its `ExecuteAll*` executor. Adding a +section = one registry entry + extending the `--mode` enum (see §8 for why this +replaced a hand-maintained if-chain). ```bash -# Default: single-hop + idempotent round-trips. Two-hop and flags are OFF by -# default (two-hop is the slowest and overlaps single-hop; flags are an -# orthogonal axis). +# Default: production single-hop + idempotent round-trips + the cheap dormant +# Bridge matrix. Two-hop and flags are OFF by default. go run ./cli/harness matrix # Everything @@ -247,6 +248,26 @@ go run ./cli/harness matrix --mode=idempotent go run ./cli/harness matrix --mode=flags # per-rule flag behavior go run ./cli/harness matrix --mode=content_shapes # request content-shape regression go run ./cli/harness matrix --mode=cache_controls # single-hop + ABA cache/no-cache +go run ./cli/harness matrix --mode=bridges # in-process Stage/Bridge topology +go run ./cli/harness matrix --mode=bridges \ + --source=anthropic_v1 --target=anthropic_beta + +# Real HTTP/server path with production Stage selection enabled +go run ./cli/harness matrix --mode=single --stage \ + --source=openai_chat --target=anthropic_beta +go run ./cli/harness matrix --mode=single --stage \ + --source=anthropic_beta --target=anthropic_beta +go run ./cli/harness matrix --mode=single --stage \ + --source=anthropic_beta --target=openai_chat +go run ./cli/harness matrix --mode=single --stage \ + --source=anthropic_v1 --target=anthropic_v1 +go run ./cli/harness matrix --mode=single --stage \ + --source=anthropic_v1 --target=openai_chat + +# Enable an active allow-only Guardrails runtime while exercising the +# production Beta Stage routes. Scenario outputs remain unchanged. +go run ./cli/harness matrix --mode=single --stage --guardrails \ + --source=anthropic_beta # Filter by scenario / source / target go run ./cli/harness matrix --scenario text --source anthropic_v1 @@ -272,6 +293,34 @@ Other CLI flags (`--batch`, `--record-dir`, `--mcp`, `-v`) are documented in `--record-dir` additionally force the sections above to run sequentially (§3.1). +The `bridges` section is deliberately separate from single-hop. Single-hop +traverses the production gateway over HTTP. It validates legacy by default; +`--stage` enables the server's twelve explicitly registered production routes: +three provider targets for each Beta, V1, Chat, and Responses source. V1 +remains a separate concrete protocol with its own HTTP response and SSE +adapters; the V1 → Beta matrix compatibility label resolves to V1 identity at +runtime and is not a production V1/Beta Bridge. The separate dormant +`--mode=bridges` matrix does contain the real V1→Beta Bridge. `--record-dir` +enables `recording_v2` and persists RequestRecord envelopes for supported +single-service Stage cases. +`--guardrails` injects an active allow-only test runtime. Combined +with `--stage`, it verifies that Beta → Beta and Beta → Chat remain on the +production Stage path while Guardrails evaluates complete and stream +lifecycles; dedicated real HTTP tests cover blocking mutations. It is a harness +fixture, not a production Guardrails configuration shortcut. Bridges runs the +dormant `stage.BuildTopology`/`stage.Adapt` path in-process and labels every +direct result `bridges//...`; concrete multi-level results use +`bridges/chain///...`. It must not be cited as production-path +proof. The matrix covers exact Anthropic v1/beta/OpenAI Chat identity, +Anthropic V1→Beta, Anthropic v1/beta → OpenAI Chat, OpenAI Chat → Anthropic +Beta, and the concrete V1 → Beta-native Stage → Chat plus Chat → Beta-native +Stage → Chat chains. Every route runs text, tool use, tool result, stream, and +non-stream (54 cells total). +Because it has no client transport, `--mode=bridges` only accepts +`--client=http`; it reuses scenario/source/target/streaming/batch filters but +does not claim support for `--mcp`, `--stage`, `--guardrails`, or +`--record-dir`. + ### Client drivers (`--client`) By default the matrix sends hand-crafted JSON over Go's `net/http`. That diff --git a/.design/protocol-recording-redesign.md b/.design/protocol-recording-redesign.md new file mode 100644 index 000000000..fcb815e23 --- /dev/null +++ b/.design/protocol-recording-redesign.md @@ -0,0 +1,434 @@ +# Request Recording Redesign + +> Status: R1–R5 and the additive R8 protocol-wide rollout are implemented. +> RequestRecord remains opt-in and requires both `--stage` and an enabled +> scenario `recording_v2` flag. R6 is proven for both the lifecycle foundation +> and the Beta-native MCP Tool Loop. MCP production routing and the persisted +> 26-case real-path matrix are complete behind `--stage`. +> +> Scope: the protocol request/response content retained for one incoming +> request. `UsageRecord`, request logging, and stage tracing are separate. + +## Decision + +The new unit is `RequestRecord`. It records the original request received at +the client-facing endpoint, the provider calls that actually happened, and the +final response returned outward after Stage and Bridge processing. + +Recording attaches only at the two stable client/provider boundaries: + +```text +HTTP Adapter + → Request Scope: capture input_request + → Stages / Bridges + → Provider Observer + → Provider Endpoint + +Provider Endpoint response + → Provider Observer: capture provider_response + → Stages / Bridges + → Client Output Adapter: final rewrites + capture final_response + → HTTP/SSE writer +``` + +- **Request Scope** records the original client-protocol request before any + Stage or Bridge and owns the completed `RequestRecord` across failover + attempts. +- **Provider Observer** records the provider-native request passed into the + terminal endpoint and the untouched provider-native response returned from + it. +- **Client Output Adapter** records the final response after outward Stages, + Bridges, response transforms, public-model rewriting, and other + client-facing adjustments. + +Recording does not snapshot every Stage. Stage insertion, removal, or +reordering must not change the persisted `RequestRecord` shape. + +## What “Original Input Request” Means + +`input_request` is the first protocol-native request created from the inbound +HTTP request and handed to the client-facing Endpoint. It is captured before +client preparation, Guardrail, Tool Loop, Bridge conversion, consistency, or +provider transforms. + +It is the original endpoint request, not a Gin object and not a byte-for-byte +HTTP body/header capture. The HTTP Adapter remains responsible for parsing the +wire request; Recording receives an immutable protocol payload from it. + +## What “Provider Request and Response” Means + +The provider request is the final protocol-native value handed to the +`Provider Endpoint`, after all inward Stages, Bridges, consistency transforms, +and provider transforms. + +The provider response is the protocol-native value returned by the +`Provider Endpoint`, before any outward Bridge or response-processing Stage. + +These are the closest stable values owned by Tingly-Box around a provider +call. They are not an HTTP transport dump of SDK headers and bytes. Exact SDK +wire capture would be a separate transport feature and is not part of this +design. + +## Data Model + +One incoming request produces one `RequestRecord`: + +```text +RequestRecord +├── request identity / scenario / outcome / duration +├── input_request +├── provider_exchanges[] +│ ├── sequence / attempt +│ ├── provider / model / protocol +│ ├── provider_request +│ ├── provider_response +│ └── outcome / error / duration +└── final_response? // present for a successful request +``` + +```go +type RequestRecord struct { + RequestID string + InputRequest Payload + ProviderExchanges []ProviderExchange + FinalResponse *Payload + Outcome Outcome +} + +type ProviderExchange struct { + Sequence int + Attempt int + Provider string + Model string + Protocol protocol.APIType + Request Payload + Response *Payload + Error string +} + +type Payload struct { + Protocol protocol.APIType + ContentType string + Body json.RawMessage +} +``` + +`ProviderExchange` is flat and ordered. Failover creates entries with different +attempt numbers. A Tool Loop creates several ordered entries under the same +attempt. A separate nested attempt model is unnecessary for message recording. + +## Three Stable Payload Boundaries + +The core record contains only these payload boundaries: + +1. `input_request`: the original client-protocol request. +2. `provider_exchanges[n]`: each actual provider request and raw provider + response. +3. `final_response`: the client-protocol response after all outward processing. + +A successful request always records `final_response`, including identity paths +where it equals the provider response. Readers never need fallback or equality +rules. A request that fails before producing a response may leave it empty. + +No intermediate Stage request or response is stored. + +## Ownership and Lifecycle + +```text +request scope: BeginRequestRecord(input request) + └── provider attempt + └── Provider Observer: begin ProviderExchange + └── Provider Endpoint + └── Provider Observer: finish ProviderExchange + └── optional failover / additional Tool Loop exchanges + └── Client Output Adapter: capture final outward response +request execution scope: FinishRequestRecord exactly once +``` + +Rules: + +1. The outer request-execution scope is the only owner of + `FinishRequestRecord`. +2. A provider error finishes only its `ProviderExchange`; it does not finish + the overall record while failover may continue. +3. Exchanges are appended in actual provider-call order. +4. The recorder never calls the provider and never controls failover. +5. Recording objects contain no Gin context and never write HTTP/SSE. +6. Recording cannot change request execution. + +## Placement in Protocol Stage Topology + +The Provider Observer wraps the terminal before topology construction. The +request scope is created before the failover loop, while the Client Output +Adapter records the value immediately before HTTP/SSE serialization: + +```text +recorder := BeginRequestRecord(input) +for each provider attempt: + provider := ObserveProvider(terminal, recorder, attempt) + endpoint := BuildTopology(provider, stages, bridges) + ServeClientOutput(endpoint, recorder) +FinishRequestRecord(recorder) +``` + +Capturing only the value returned by `BuildTopology` is too early. Existing +client adapters still apply protocol response transforms, public-model +rewriting, and stream-event adjustments before writing; `final_response` must +reflect those operations. + +This placement has stable semantics for all topologies: + +| Topology | Client boundary records | Provider Observer records | +| --- | --- | --- | +| Identity | original request and final response | provider request/response | +| Cross-protocol Bridge | source-protocol input and converted output | target-protocol request/response | +| Guardrail | request before inbound policy; response after outbound policy | request after inbound policy; raw provider response | +| Tool Loop | original request and only the final outward response | every provider round | +| Failover | original request and only the response ultimately returned | every provider call by attempt | + +Recording every Stage boundary is deliberately rejected for the core record: + +- payload count would grow with topology depth; +- internal Stage order would become a storage contract; +- the same content would usually be duplicated; +- Guardrail and Tool Loop implementation details would leak into the request + artifact. + +An ordered stage trace may separately record stage name, protocol, duration, +and outcome. It is diagnostics metadata, not request/response content. + +## Complete and Streaming + +Complete calls snapshot the input request in the request scope and the +request/response pair at each Provider Observer invocation. The Client Output +Adapter snapshots the final response after its last response rewrite. + +For streaming, each observer wraps the stream returned by the next endpoint +and assembles events in that observer's native protocol while the normal caller +pulls them. The observer does not consume the stream independently: + +- Provider Observer assembles the complete raw provider response for every + provider exchange. +- Client Output Adapter assembles the final outward events after client-facing + rewrites into one complete response. +- Raw stream chunks are not stored in the first implementation. + +Recording reuses protocol-owned typed values, Wire DTOs, and stream assemblers. +It never performs protocol conversion itself and does not maintain a second +recording-specific codec registry. + +## Separation from Usage + +`RequestRecord` and `UsageRecord` are independent: + +- neither creates, updates, or finalizes the other; +- either feature works when the other is disabled; +- both may use the same `request_id` for correlation; +- Recording does not parse or normalize token counts. + +A provider response may contain a native `usage` field inside its captured +body. It remains ordinary response content and does not make the records share +ownership. + +## Legacy Field Mapping + +The target semantics map as follows: + +| Legacy concept | New field | +| --- | --- | +| Client/pre-transform snapshot (`original_request`) | `input_request` | +| Final provider-bound request (`transformed_request` was the closest implementation) | `provider_exchanges[n].provider_request` | +| Raw provider response (`provider_response`, previously not reliably populated) | `provider_exchanges[n].provider_response` | +| Final outward result (`final_response`) | `final_response` | +| Transform step names | separate stage trace | + +## Implementation Containment + +The new implementation starts in one isolated package: + +```text +internal/record/ +├── record.go // RequestRecord, ProviderExchange, Payload +├── recorder.go // request-scoped lifecycle +├── provider_endpoint.go // terminal Endpoint observer +└── stream.go // observes EventStream through protocol assemblers +``` + +This package has no Gin, HTTP writer, routing, Guardrail, MCP, or legacy +recorder dependency. It may wrap `protocol/stage.Endpoint`; the Stage core does +not import Recording. + +Protocol handling stays in the existing protocol packages: + +- complete request/response values remain the typed values already carried by + `stage.Call`, `stage.Response`, and `wire` DTOs; +- Provider observation delegates through the existing `stage.Endpoint` and + `stage.EventStream` contracts; +- stream reconstruction reuses the existing Anthropic V1/Beta, OpenAI Chat, + and OpenAI Responses assemblers in `internal/protocol/assembler`; +- client output capture reuses the same Wire DTO/event values already produced + by the HTTP/SSE adapters. + +If Recording exposes a real common gap, the protocol layer is enhanced once. +The expected addition is a small protocol-owned assembler interface/factory +over the existing implementations, for example: + +```go +type StreamAssembler interface { + Add(value any) error + Finish() (any, error) +} + +func NewStreamAssembler(api protocol.APIType) (StreamAssembler, error) +``` + +The concrete type switches for SDK and Wire events belong to the protocol +assembler package, not `internal/record`. No source→target protocol-pair logic +is added for Recording. + +Production integration is limited to three seams per client protocol: + +1. The protocol prologue creates one recorder and captures `input_request` + before the failover loop. +2. The Stage target builder wraps the selected Provider Endpoint. +3. The client output adapter captures the post-rewrite complete response or + outgoing stream events. + +The first Beta identity canary therefore touches only +`anthropic_message.go` and `protocol_stage_anthropic_beta.go` outside the new +package. It does not modify Stage/Bridge contracts, Guardrail, MCP, or the +legacy recorder. + +Failover integration is a later, central change in `failover_dispatch.go`: it +provides the current attempt number while keeping the same request-scoped +recorder. Tool Loop requires no recording hook; repeated calls through the +already-wrapped Provider Endpoint naturally append exchanges. + +Other source protocols are integrated one at a time through their existing +prologue and client output adapter, with an independent test and commit for +each. No all-protocol handler rewrite is required. + +## Additive Migration Plan + +| Checkpoint | Status | Change | Production effect | +| --- | --- | --- | --- | +| R1 — Foundation | Complete | `RequestRecord`, ordered exchanges, lifecycle, in-memory tests | None | +| R2 — Protocol capture support | Complete | Common interface over existing Beta, V1, Chat, and Responses assemblers | None | +| R3 — Boundary harness | Complete | Verify input, provider, and output snapshots across all Stage routes | No persisted output | +| R4 — Single-route canary | Complete | Beta identity, single service, no MCP | Opt-in only | +| R5 — Failover | Complete | Ordered attempt exchanges across homogeneous and cross-protocol failover, one final record | Opt-in only | +| R6 — Tool Loop | Complete | Beta complete/stream Tool Loop records multiple exchanges in one attempt through the real Provider Observer; the persisted real-path matrix verifies one original input, two provider exchanges, and one final response | Opt-in behind `--stage`; no default cutover | +| R7 — Persistence/UI | Partial | Native reader and request inspection surface; R4 already writes an additive `request_record` envelope through the existing sink | Opt-in only | +| R8 — Protocol-wide rollout | Complete | All twelve production Stage routes, complete and stream, Stage-compatible service sets and no MCP | Opt-in only; no default cutover | +| R9 — Cleanup | Not started | Remove Gin recorder, transform recorder, stream hooks, and MCP recorder interface | After parity proof | + +### Current Activation Boundary + +The new recording path is selected only when all of these are true: + +- the server starts with `--stage`; +- the request scenario has `recording_v2` enabled; +- the route is one of the twelve explicitly registered production Stage routes; +- every active service resolves to a registered Stage provider protocol + (current exclusion: Google-style providers); + +Every other recording combination keeps the complete legacy lifecycle. Without +`recording_v2`, no new recorder or sink work is performed. Restarting without +`--stage` is the rollback. The harness compatibility label V1 → Beta resolves +to the V1 provider protocol at runtime; the production recording path still +uses V1 identity and does not register the dormant V1→Beta Bridge. + +The completed `RequestRecord` is persisted through the existing asynchronous +obs sink as an additive `request_record` envelope. Legacy readers may ignore +that field; non-Stage or unsupported feature combinations continue to use the +legacy behavior. + +The original HTTP body is retained only when both Protocol Stage and the +scenario recording flag are enabled. This keeps the default path free of a +second full-body copy while preserving unknown client fields for Chat, +Responses, Anthropic V1, and Anthropic Beta recordings. + +The existing `recording_v2` modes project the RequestRecord at persistence: + +| Mode | Persisted RequestRecord boundaries | +| --- | --- | +| `request` | original input and ordered provider requests | +| `request_response` | request boundaries plus final client response | +| `staged_request_response` | all boundaries, including raw provider responses | + +Provider stream recording requires a protocol terminal event before a clean +close is classified as success. Anthropic uses `message_stop`, Chat uses a +non-empty `finish_reason`, and Responses uses a terminal response/error event. +A close or EOF before that boundary is recorded as failure without persisting a +partial provider response. `response.failed` preserves its response payload +but marks both the provider exchange and request as failed. + +Provider-bound normalization also completes before the observer snapshots the +request. In particular, the OpenAI Chat cleanup that removes gateway-only +`x_thinking` and empty tools runs as the final provider transform, so the +recorded provider request is the value actually handed to transport. + +R8 verification ran the text matrix through the real server path with +`--stage --record-dir`: all 26 labeled complete/stream cases passed. The +persisted output contained 26 successful `RequestRecord` envelopes, each with +one input, one provider exchange, and one final response. Those 26 cases cover +the twelve production Stage routes twice, plus the two V1 → Beta compatibility +cases that normalize to V1 identity. + +R5 keeps the recorder at request scope while the failover orchestrator exposes +the current attempt number only when recording is active. Each provider +terminal wrapper appends one exchange with that attempt number. The request is +finished once, after the failover gate has committed the winning response or +flushed the terminal error. Complete and stream tests cover all four source +protocols with cross-protocol failure → success, plus exhausted two-provider +failure. Rules containing a provider protocol outside the registered Stage +surface do not enter the new recording path. + +R6 requires no new recording hook. Both Tool Loop implementations call the +already observed Provider Endpoint once per model round, so each round appends +an ordered exchange under the same attempt number. The authoritative Beta +complete/stream tests prove that an internal tool round followed by a final +model round produces two successful provider-native Beta exchanges and one +final outward response; the original input is captured before tool injection. +Production handler selection is wired behind `--stage`. The persisted +real-HTTP MCP matrix now covers all 13 registered source/target labels in both +complete and streaming modes. Each case flushes and reads the gzip JSONL +artifact, correlates the new record by request ID, and verifies: + +- the source-protocol input remains pre-injection; +- exactly two successful provider-native exchanges occur in attempt 1; +- the first exchange contains the injected owned tool and its call; +- the second exchange contains the local tool result and final provider answer; +- the final response is recorded in the original client protocol. + +The executable command is: + +```bash +go run ./cli/harness matrix --mode=single --stage --mcp \ + --scenario=mcp_owned_tool --record-dir=/tmp/tingly-mcp-records +``` + +## Required Verification + +- identity complete/stream records the input, one provider pair, and the final + response; +- every route preserves the original client-protocol request in + `input_request` before any Stage or Bridge mutation; +- each cross-protocol route records target-protocol provider payloads and the + source-protocol output response; +- same-protocol response mutation is reflected in `final_response`; +- retryable failure followed by success retains both provider exchanges and + finishes one `RequestRecord`; +- Tool Loop retains every provider exchange in order but only one final output; +- stream assembly produces the same logical payload as complete recording; +- Recording works with Usage disabled and Usage works with Recording disabled; +- the harness covers all registered Stage route labels in complete and stream + modes, including persisted MCP multi-round boundaries. + +## Explicit Non-Goals + +- Recording does not drive routing, retries, health, usage, or affinity. +- Recording does not snapshot every Stage payload. +- Recording does not become a canonical protocol AST. +- Recording does not own HTTP response writing or SSE framing. +- R1–R3 do not replace or modify the existing recorder. diff --git a/.design/protocol-stage-chain.md b/.design/protocol-stage-chain.md new file mode 100644 index 000000000..fc3dea905 --- /dev/null +++ b/.design/protocol-stage-chain.md @@ -0,0 +1,823 @@ +# Protocol Stage Chain + +> Status: Phases 1–5 are active additively behind `--stage`. The Beta working +> boundary is connected to production handlers with exact two-boundary +> selection, V1 request promotion, Beta-native Guardrail and Tool Loop stages, +> MCP/runtime + servertool adapters, recording, and side-effect-aware failover. +> `tingly-box start --stage` selects the supported Chat/Beta/V1 production +> routes plus OpenAI Responses → Responses/Anthropic Beta/OpenAI Chat, +> Anthropic Beta/V1 → OpenAI Responses, and OpenAI Chat identity/Beta/Responses. +> All four supported ingress protocols can include the Beta-native Guardrail +> Stage; unsupported routes remain on legacy. Request recording is +> available on all twelve routes, including failover across Stage-compatible +> services, with `recording_v2`. +> +> Scope: LLM request/response data plane for non-streaming and streaming calls. + +## Current Status + +The migration remains additive. Starting without `--stage` selects legacy for +every request. Starting with `--stage` makes a per-provider-attempt decision; +an unsupported route or feature combination selects the complete legacy +lifecycle before the provider is called. + +| Phase | Status | Current boundary | +| --- | --- | --- | +| 1 — Endpoint/Stage foundation | Complete | Contracts, ordering, stream ownership, and per-call state | +| 2 — Bridges and production routes | Complete for planned protocol surface | Twelve opt-in routes listed below | +| 3 — Guardrails canary | Complete for all four supported ingress protocols | Request, complete response, and stream events; Beta, Chat, and Responses targets | +| 3b — Request recording rollout | Complete for Stage routes, failover, and MCP multi-round calls | Original input, ordered provider exchanges, and final complete/stream response through the existing sink; persisted MCP matrix covers all 26 route/mode cases | +| 4 — Tool Loop canary | Active behind `--stage` | Exact source→Beta→provider topology, complete/stream Tool Loop, V1 request promotion, tool-result-correlated mixed continuation, recording, and dispatch-aware failover | +| 5 — Opt-in handler integration | Active | Existing handlers may select Stage only from the immutable `--stage` startup choice; default traffic remains legacy | +| 6 — Legacy removal | Not started | No legacy feature path has been removed | + +Production route selection with `--stage`: + +| Client protocol | Provider protocol | Plain request | Guardrails enabled | MCP enabled | Protocol recording | +| --- | --- | --- | --- | --- | --- | +| `anthropic_beta` | `anthropic_beta` | Stage | Stage with `guardrail_anthropic_beta` | Beta Tool Loop | Stage `RequestRecord`, including failover, when every active service is Stage-compatible; otherwise Legacy | +| `anthropic_beta` | `openai_chat` | Stage through Beta→Chat Bridge | Stage with the same Beta Guardrail | Beta Tool Loop→Chat | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `anthropic_beta` | `openai_responses` | Stage through Beta→Responses Bridge | Stage with the same Beta Guardrail | Beta Tool Loop→Responses | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `anthropic_v1` | `anthropic_v1` | Stage | V1→Beta Guardrail; provider request uses Beta | V1→Beta Tool Loop; provider request uses Beta | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `anthropic_v1` | `openai_chat` | Stage through V1→Chat Bridge | V1→Beta Guardrail→Chat | V1→Beta Tool Loop→Chat | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `anthropic_v1` | `openai_responses` | Stage through V1→Responses Bridge | V1→Beta Guardrail→Responses | V1→Beta Tool Loop→Responses | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `openai_chat` | `openai_chat` | Stage | Chat→Beta Guardrail→Chat | Chat→Beta Tool Loop→Chat | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `openai_chat` | `anthropic_beta` | Stage through Chat→Beta Bridge | Chat→Beta Guardrail | Chat→Beta Tool Loop | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `openai_chat` | `openai_responses` | Stage through Chat→Responses Bridge | Chat→Beta Guardrail→Responses | Chat→Beta Tool Loop→Responses | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `openai_responses` | `openai_responses` | Stage | Responses→Beta Guardrail→Responses | Responses→Beta Tool Loop→Responses | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `openai_responses` | `anthropic_beta` | Stage through Responses→Beta Bridge | Responses→Beta Guardrail | Responses→Beta Tool Loop | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| `openai_responses` | `openai_chat` | Stage through Responses→Chat Bridge | Responses→Beta Guardrail→Chat | Responses→Beta Tool Loop→Chat | Stage `RequestRecord` for a compatible service set; otherwise Legacy | +| Any other pair | Any | Legacy | Legacy | Legacy | Legacy | + +The planned protocol-pair rollout is complete: the three provider-facing +protocols (Beta, Chat, and Responses) are available from each supported source, +while Anthropic V1 remains deliberately distinct and native only to V1 clients. +The production candidate now uses Anthropic Beta as the Tool Loop working +protocol because Beta has the complete MCP/server-tool surface. V1 requests +enter it through the lossless V1→Beta promotion Bridge; Chat and Responses use +their existing Bridges. The earlier Chat-native Tool Loop remains a useful +lifecycle proof; production MCP selection now uses the Beta-native topology +behind the explicit `--stage` canary. + +The Chat and Responses Guardrail canaries read the existing Guardrails runtime +and scenario/global feature flag through a Stage-specific gate. They do not add +the `openai` scenario to the legacy Guardrails support list: doing so would +change converted legacy behavior without `--stage`. The Stage-specific gate +keeps the new OpenAI ingress behavior opt-in while preserving legacy rollback. + +## Decision + +Tingly-Box will evolve protocol-bound features into an ordered chain of +in-process **Protocol Stages**. A stage exposes the same complete and streaming +operations as the endpoint it wraps. Bidirectional protocol bridges adapt the +client protocol to the stage protocol and the stage protocol to the selected +provider protocol. + +The foundations were deliberately built旁路 before handler integration. The +current canaries connect only the route, Guardrail, and Tool Loop combinations +listed in Current Status, behind `--stage`; rules containing an unsupported +provider protocol and every unsupported combination retain whole-request +legacy ownership. + +## Why This Change + +Guardrails, MCP, and server-tool behavior currently attach to a mixture of: + +- concrete SDK request/response types; +- request transform chains; +- `protocol.HandleContext` stream hooks; +- protocol-pair dispatch branches; +- an MCP-owned multi-round stream loop; +- Gin response and SSE writers. + +This spreads lifecycle ownership across features. In particular, a stream may +be consumed, filtered, recorded, guarded, and written by different layers. A +feature that should conceptually be “one processing level” must understand +several protocols and several execution modes. + +Normal endpoint composition makes the order explicit: requests travel inward; +responses and stream events travel outward through the same wrappers. + +## Terminology + +| Term | Meaning | +| --- | --- | +| Protocol | A concrete API shape such as `openai_chat`, `openai_responses`, `anthropic_v1`, or `anthropic_beta` | +| Endpoint | A complete non-stream and stream implementation of one protocol | +| Protocol Stage | A named full-duplex endpoint wrapper implementing one native protocol | +| Bridge | A bidirectional adapter between two protocols: request inward, response/events/errors outward | +| Terminal Endpoint | The innermost provider-facing endpoint | +| HTTP Adapter | The only component allowed to parse ingress HTTP and commit response bytes/SSE | +| Wire DTO | A typed, serializable client-protocol response/event contract used only at the outward protocol boundary | + +Do not call Protocol Stages “tiers”. Tier already means provider failover +priority in Tingly-Box. + +## Target Shape + +```text +client + | client protocol + v +HTTP Adapter + | +Ingress Bridge client protocol <-> stage protocol + | +Guardrails Stage request pre-check / final response post-check + | +Tool Loop Stage tool catalog, interception, execution, continuation + | +Provider Bridge stage protocol <-> selected provider protocol + | +Provider Endpoint +``` + +The arrows for requests point downward. Complete responses, stream events, and +errors return upward through the same wrappers. + +### Wire DTO Boundary + +Wire DTOs are the typed hand-off between a Bridge's outward conversion and the +client protocol's HTTP/SSE adapter. They are not a shared internal protocol or +a canonical model for Guardrails, Tool Loop, routing, or provider calls. + +- request conversion and provider execution keep using the concrete protocol + SDK types plus `Call` and `ProtocolState`; +- complete/stream Bridge return paths may emit the source client protocol's + `wire.*` value directly; +- the outer HTTP/SSE adapter serializes that value and remains the sole owner of + headers, public model rewriting, response transforms, and framing; +- Bridges must construct typed wire values directly. A + `map → JSON → SDK/wire` round-trip is not a protocol boundary and can silently + discard extensions such as refusal or detailed usage. + +Wire types are named after the protocol they serialize, never after a specific +conversion route. This keeps `Responses → Chat` and any future provider path +able to reuse the same Chat output contract without coupling the DTO to its +origin. + +## Core Contracts + +The first foundation defines four concepts: + +```go +type Endpoint interface { + Protocol() protocol.APIType + Complete(context.Context, Call) (*Response, error) + Stream(context.Context, Call) (EventStream, error) +} + +type Stage interface { + Name() string + Protocol() protocol.APIType + Wrap(next Endpoint) Endpoint +} + +type EventStream interface { + Next(context.Context) (Event, error) + Close() error + Result() StreamResult +} + +func Compose(terminal Endpoint, stages ...Stage) (Endpoint, error) +``` + +`Compose(terminal, guardrails, tools)` means: + +```text +guardrails(tools(terminal)) +``` + +The stage list is written in request order, outermost to innermost. Composition +fails before execution when a stage protocol differs from the endpoint it +wraps. This keeps an accidental implicit conversion out of feature code. + +## Contract Invariants + +1. **One native protocol per endpoint** — `Protocol()` is concrete and never an + alias such as “OpenAI-compatible”. +2. **Both execution modes are required** — a stage implements `Complete` and + `Stream`; pass-through implementations are valid. +3. **No transport ownership** — a stage cannot depend on Gin, write headers, or + frame SSE. +4. **One stream consumer chain** — each wrapper pulls from the next stream and + returns an event outward. Only the HTTP adapter drives the outermost stream. +5. **Explicit close** — the caller closes any successfully returned stream. + Wrappers propagate close to the upstream stream. +6. **Per-call state** — mutable conversion, Guardrail, tool-call, and usage state + is created for one call/attempt, never stored globally in a shared Stage. +7. **Structured outcome** — usage, response model, trace, and committed side + effects travel in `Response` or `StreamResult`, not in Gin context fields. +8. **No hidden bridge** — protocol changes happen only in a named Bridge. +9. **No retry ownership** — stages report failures and commitment state; + routing/failover decides whether another provider attempt is allowed. + +## Bidirectional Bridge Contract + +A bridge must convert the whole protocol surface, not merely the request: + +- parsed request and request-derived state; +- complete response; +- stream events, including one-to-many event expansion; +- terminal usage and finish reason; +- response model and identifiers; +- typed errors and retry hints. + +Bridge instances are immutable configuration and must be concurrency-safe. Any +mutable assembler or correlation state belongs to a per-call bridge session +created while converting the request: + +```go +type Bridge interface { + Source() protocol.APIType + Target() protocol.APIType + Capabilities() Capabilities + Open(context.Context, Call, Operation) (BridgeSession, error) +} + +type BridgeSession interface { + TargetCall() Call + ConvertComplete(context.Context, *Response) (*Response, error) + ConvertStream(context.Context, EventStream) (EventStream, error) + ConvertError(context.Context, error) error +} +``` + +`Open` converts the request inward and creates exactly one session for that +call. `OperationComplete` and `OperationStream` are explicit because request +conversion may set different stream/usage fields. The session converts complete +responses, streams, and target errors outward. After successful stream +conversion, the converted stream owns and closes the target stream. + +`Call.State` is a bounded `ProtocolState` carrier for request-derived facts that +remain necessary after changing protocol. The initial `OpenAIChat` field holds +the `OpenAIConfig` produced by Anthropic-to-Chat conversion for later provider +transforms. This state is per-call and is intentionally not an extensible +property bag. + +An immutable `BridgeRegistry` resolves exact protocol pairs and capabilities. +The topology builder works from the provider terminal outward and inserts a +registered bridge whenever adjacent stages speak different protocols. This +allows every stage to implement one native protocol without requiring all +stages to choose the same protocol. + +The generic Bridge foundation is additive. Concrete bridges must reuse existing +request/nonstream/stream converters rather than create a second conversion +implementation. + +### Canonical Stage Protocol + +`anthropic_beta` is the leading initial candidate because it is already the +normalization target for non-Anthropic requests sent to Anthropic providers and +can represent rich content and tools. It is not hard-coded as a universal +promise. + +Before enabling a chain, a capability check must prove both bridge legs for the +requested features: + +```text +source -> stage protocol -> provider target +``` + +Capabilities include complete response, streaming, tools, tool results, usage, +finish reason, and error fidelity. Known OpenAI Responses tool-use defects mean +some combinations must remain on the legacy path during migration. + +## Feature Ownership + +### Guardrails Stage + +The Guardrails Stage owns: + +- inbound user-content evaluation and permitted request mutation; +- outbound evaluation of the final client-visible response; +- stream accumulation required by a configured policy; +- credential masking state and cleanup; +- Guardrail-specific trace facts that do not expose protected content. + +It does not inject tools, execute tools, select providers, or write errors to +HTTP. + +### Tool Loop Stage + +The focused Phase 4 contract is +[`protocol-stage-tool-loop.md`](./protocol-stage-tool-loop.md). This section is +the architectural summary; the focused document owns lifecycle, V1 request +promotion, activation, and acceptance details. + +The Tool Loop Stage owns: + +- server-visible tool catalog injection; +- complete and streaming tool-call assembly; +- classification of client/external versus server/internal tools; +- policy checks immediately before execution; +- invocation through a protocol-neutral `ToolExecutor`; +- appending tool results and continuing the next model round; +- max-round enforcement and usage accumulation. + +Two native implementations now exist. The earlier `openai_chat` Stage proves +the generic lifecycle contracts. The production candidate is +`anthropic_beta`: it keeps Anthropic MCP/tool/thinking structures native and +avoids a neutral DTO or an OpenAI-shaped intermediate. Existing Bridges make +that Beta Stage available to V1, Chat, and Responses without teaching the Tool +Loop those protocols. A request tool and server tool with the same name fail +explicitly because name-only ownership would be ambiguous. + +The Beta stream implementation buffers one provider round before exposure. +Anthropic may emit visible text or thinking before a later internal +`tool_use`, so no prefix proves that the round is safe to release. Pure +server-owned rounds are consumed internally; external rounds are replayed; +mixed rounds execute only owned calls, hide their blocks, renumber remaining +indexes, and retain a single-consume continuation segment correlated to the +external tool-result IDs expected from the next client turn. Explicit session +identity, TTL cleanup, and a capacity bound prevent cross-client splicing and +unbounded abandoned state. This correctness requirement is explicit; future +TTFT optimization must not assume tool blocks arrive first. + +MCP remains a tool catalog/runtime source. It now lists virtual tools directly +as Beta tool definitions. `servertool.Executor` is reused through the existing +execution boundary, and the provider-scoped continuation store is reused with +typed Beta messages. None of these dependencies needs to understand every +client/provider protocol; only the Tool Loop Stage understands Beta. + +Splitting MCP and servertool into two protocol stages immediately would create +a false boundary because both would compete to own the same model tool-call +loop. They can become separately composable later only if their request and +response lifecycles become independently meaningful. + +## Intended Stage Order + +Default order: + +```text +Guardrails(ToolLoop(Provider)) +``` + +Consequences: + +- the inbound Guardrail sees the original user content before tool injection; +- the Tool Loop consumes internal model tool calls and produces a final answer; +- the outbound Guardrail sees what will actually be returned to the user; +- tool authorization runs inside the Tool Loop before any executor call. + +Current stream and non-stream paths do not express this order identically. +Migration must record the difference and intentionally converge on the order +above rather than silently claiming byte-for-byte parity. + +## Streaming Lifecycle + +1. The HTTP Adapter parses the request without committing SSE headers. +2. The composed endpoint returns an `EventStream` or a pre-stream error. +3. The HTTP Adapter installs the existing failover commit gate when applicable. +4. It pulls events from the outer stream. +5. The first real client-visible event commits the response. +6. It frames and writes that event in the client protocol. +7. On completion or error it reads `StreamResult`, records usage/trace, and + closes the stream. + +Converters and stages must support cancellation and backpressure by performing +work only when `Next` is called and honoring the passed context. Full-response +buffering must not be used merely to simplify transport handling. A Stage may +buffer one bounded provider round when protocol semantics make an earlier +ownership decision impossible, as documented for the Beta Tool Loop above. + +## Failover and Irreversible Side Effects + +Two commitment boundaries matter: + +1. **Output committed** — the first client-visible chunk has left the process. +2. **Side effects committed** — a server tool has successfully performed work + that cannot be safely replayed. + +After either boundary, the outer orchestrator must not discard the attempt and +restart the full chain on another provider. + +The first tool-stage implementation will conservatively mark successful tool +execution as committed. Future work may allow retry for tools that explicitly +declare idempotency and use a stable `(request ID, tool call ID)` deduplication +key. + +## Observability + +Every execution should eventually emit an ordered stage trace such as: + +```text +openai_chat -> anthropic_beta -> guardrails -> tool_loop + -> openai_responses -> provider +``` + +Each entry records concrete protocol, duration, outcome, and safe counters. It +must not record prompts, credentials, masked content, or raw tool arguments by +default. Diagnostics must traverse the production chain; a direct provider +probe remains useful only as an explicit comparison path. + +## Incremental Migration + +### Phase 1 — Foundation, no traffic + +- Add `internal/protocol/stage` contracts and composition validation. +- Add unit tests for complete and stream ordering, protocol mismatch, and close. +- Do not import the package from existing server code. + +### Phase 2 — Bridge and concrete-chain harness + +- Adapt existing request/nonstream/stream converters behind a bidirectional + bridge session. +- Split protocol conversion from Consistency/Vendor provider finalization. +- Add a real composed A -> B -> C request path to the harness. +- Add capability checks and keep unsupported combinations on legacy. + +Generic Bridge sessions, capability checks, an immutable exact-pair registry, +identity bridges, and mixed-protocol in-memory topology tests are implemented. +Concrete Anthropic v1/beta → OpenAI Chat and OpenAI Chat → Anthropic Beta +bridges are implemented for complete and stream. The dormant matrix now runs a +real Chat → Beta-native Stage → Chat topology in both modes. Production-path +server and harness validation is tracked separately; the dormant matrix alone +must not be treated as evidence of production dispatch wiring. + +### Phase 3 — Guardrails canary + +- Implement Guardrails for one native stage protocol. +- Compare request, complete-response, and stream-event mutations with legacy + behavior in independent and real-path tests. +- Enable only when both `--stage` and the existing Guardrails scenario gate are + active; retain whole-attempt legacy fallback for unsupported feature mixes. + +The first canary is authoritative rather than shadowed: a response policy can +only inspect the result of the one provider call that will be returned to the +client. Running a second legacy response path would either call the provider +twice or compare a synthetic lifecycle, neither of which is a safe dry run. +Anthropic Beta Guardrails therefore own request and response processing only +after the Stage topology is selected. If selection declines, the untouched +attempt enters the complete legacy Guardrail lifecycle. + +### Phase 4 — Tool Loop canary + +Implementation and production-entry requirements for this phase are specified +in [`protocol-stage-tool-loop.md`](./protocol-stage-tool-loop.md). + +- Move generic complete/stream MCP loops behind one Stage. +- Inject `ToolCatalog`, `ToolPolicy`, and `ToolExecutor` dependencies. +- Validate using deterministic mocks and read-only tools before a production + canary. Never dual-execute tools for shadow comparison. + +The in-process portion is complete. The Chat-native Stage established the +lifecycle contracts; the Beta-native Stage now owns complete and streaming +continuation using the existing Anthropic MCP adapter. The existing runtime +injects Beta tools directly, the existing servertool executor is reused, and +provider-scoped mixed continuation is typed, result-correlated, bounded, and +single-consume. V1 requests +enter Beta through JSON marshal/unmarshal of the V1 subset. Deterministic tests +cover internal, external, and mixed ownership, max-round/side-effect +boundaries, usage aggregation, stream close, composed V1→Beta execution, and +RequestRecord multi-exchange behavior. Production selection now validates both +exact boundaries before constructing the per-attempt topology, and the real +HTTP harness covers all 13 source/target labels in complete and streaming MCP +owned-tool modes (26/26 executed, no skips). + +### Phase 5 — Handler integration behind `--stage` + +- Compose a fresh chain for each provider attempt from a pristine request. +- Preserve existing routing, load balancing, and first-chunk gate behavior. +- Keep default traffic on legacy. Protocol matrix, official SDK, Duo, and + failover validation harden the explicit `--stage` path; they do not activate + it implicitly. + +The first opt-in integration selects Stage per provider attempt after routing +has resolved the concrete target protocol but before legacy Base conversion. +`--stage` is immutable for the server process. The Stage path reuses client +preparation, target consistency, rule, and vendor transforms as native protocol +stages; the provider endpoint and HTTP adapter retain their existing ownership. +Unsupported protocol pairs and incomplete MCP topologies remain on legacy. +Once a Stage attempt has started, it is never replayed through legacy. + +There is no automatic rollout condition in this phase. The command path passes +`--stage` into `server.WithProtocolStage`, which is copied once into the model +handler and its immutable selector. No scenario flag, MCP flag, Guardrail flag, +or compatible protocol pair may enable Stage when that startup choice is false. + +The native and bridged routes are explicitly enumerated in Current Status; +registration is always per exact source/target pair. +`anthropic_v1` remains a separate protocol with its own request, response, +stream, terminal, and identity registration. MCP and Guardrails promote only +its request into their shared Beta working boundary. When both are enabled, the +same opt-in chain places the Beta Guardrail outside the Tool Loop. All twelve +registered routes may use Stage recording with failover when every active +service resolves to a registered provider protocol. The request scope emits one +record; each attempt appends one ordered exchange. +Guardrails are native on Beta-source routes and on V1 requests promoted at the +Beta boundary. Without `--stage`, Guardrails-only V1 still selects the entire +legacy pipeline. + +### Phase 6 — Legacy removal + +- Remove protocol-specific feature hooks and MCP transforms in separate changes. +- Remove per-protocol MCP experiment toggles after all supported traffic uses + the new chain. +- Keep rollback available until the cleanup change itself is proven stable. + +## Rollout and Rollback + +- `legacy` remains the default at the beginning. +- `tingly-box start --stage` activates Stage selection; restart without the flag + is the rollback artifact. +- Pure conversion and dry-run Guardrail behavior may run in shadow. +- Tool execution may only run once and therefore uses explicit canaries. +- Fallback is allowed only before output or side effects are committed. +- No persisted schema change is required for the foundation or first canaries. +- Removing legacy code is never combined with initially enabling a migrated + Stage. + +## Security Requirements + +- Stage metadata cannot become an unbounded feature-specific property bag. +- Sensitive request/response data must not enter stage traces. +- Tool execution requires the same callable/permission checks as the current + server-tool pipeline. +- Cancellation and close paths must release Guardrail buffers and provider + streams. +- Unsupported protocol capabilities fail closed or stay on legacy during the + migration; they are never silently omitted. +- Side-effect commitment must be propagated even when a later model round + fails. + +## Feature Stage Contract + +Guardrail and Tool Loop integrations must behave as typed multi-level services, +not renamed handler hooks: + +- each feature is a full-duplex `Stage` that wraps one `Endpoint` and owns both + complete and stream lifecycles in one concrete protocol; +- feature code declares its concrete protocol and never invokes protocol + conversion directly; `BuildTopology` inserts capability-complete Bridges; +- one Bridge session converts the request inward and the response or events + outward for a single call; +- adjacent feature Stages using the same protocol compose without another + conversion; +- a feature Stage can be tested without Gin, server routing, or a provider + transport and can later be replaced by a remote Endpoint implementation; +- missing capabilities fail topology construction or select the entire legacy + pipeline; one request never mixes Stage and legacy ownership; +- tools are dependencies of a `ToolLoopStage`: MCP, server tools, and builtins + implement `ToolExecutor` rather than becoming protocols themselves. + +The generic Guardrail foundation remains an observe-only, fail-open Stage for +new evaluators. The first authoritative adapter implements the existing +Anthropic Beta request mutation, complete-response blocking/restoration, and +stream `tool_use` buffering/rewrite while preserving endpoint errors, stream +ownership, usage, model, and monotonic side-effect facts. It contains no Gin, +provider, or Bridge dependency. + +## UX-First Review + +- **Vocabulary**: “Protocol Stage” avoids collision with routing Tier. +- **Smart defaults**: existing feature enablement builds one default order; no + mode picker is introduced. +- **Concrete values**: diagnostics show `anthropic_beta`, not an alias. +- **Real path**: harness and diagnostics exercise the same composed endpoint + used by requests. +- **Reversibility**: legacy remains available through canary rollout and cleanup + is a later decision. +- **Scoped effects**: shadow mode excludes tool execution and fallback stops at + commitment boundaries. + +## Alternatives Rejected for the Initial Migration + +- Expanding raw protocol hooks as the target architecture. +- Creating a user-configurable stage-order editor. +- Running every stage as a separate HTTP process before in-process semantics are + complete. +- Making servertool a separate protocol stage while the Tool Loop still owns + tool-call assembly and continuation. +- Replacing all protocol dispatch paths in one change. + +## Implementation Checkpoint — 2026-07-13 + +The following foundations are implemented under `internal/protocol/stage`: + +- complete and streaming Endpoint contracts; +- ordered same-protocol Stage composition; +- per-call bidirectional Bridge sessions; +- core and semantic capability checks; +- exact-pair immutable Bridge registry plus identity fallback; +- topology construction that inserts Bridges between differently typed Stages; +- monotonic propagation of usage/model fallback and committed side effects; +- complete and streaming in-memory multi-hop harnesses. + +Bidirectional Anthropic Beta/OpenAI Chat Bridges and the V1→Beta request +promotion Bridge are implemented. Their response directions expose +transport-neutral complete and stream conversion entrypoints while existing +`Handle*` functions remain the production wrappers. The dormant 54-cell Bridge +matrix includes concrete V1 → Beta-native Stage → Chat and Chat → Beta-native +Stage → Chat topologies and verifies text, tool-use, and tool-result semantics +in complete and streaming modes. + +Runtime integration is opt-in through `--stage`. For each OpenAI Chat provider +attempt whose concrete target is Anthropic Beta, the server builds a fresh Chat +preparation → Bridge → Beta provider-finalization → provider endpoint topology. +For an Anthropic Beta request, it builds Beta preparation followed by either +the Beta identity path or the Beta → OpenAI Chat Bridge, then the concrete +provider-finalization and endpoint. Streaming and complete responses return +through the same endpoint chain and the outer Beta HTTP adapter. Anthropic V1 +uses separate V1 preparation followed by either V1 identity or the V1 → OpenAI +Chat Bridge, then the concrete provider-finalization and endpoint. +OpenAI Responses identity uses Responses preparation and finalization around a +transport-free Responses provider endpoint; its outer adapter preserves raw +provider JSON fields, Responses SSE event names, usage-detail compatibility, +public model rewriting, and pre-stream failover semantics. +Responses → Anthropic Beta uses a per-call Bridge session. Requests cross into +Beta before provider finalization; complete messages and stream events cross +back into Responses wire DTOs before the same Responses HTTP adapter runs. +Responses → OpenAI Chat follows the same boundary: the Bridge converts requests +to Chat, carries explicit Chat state into provider finalization, and converts +complete responses or Chat chunks back to Responses wire shapes. +The reverse Beta → Responses Bridge restores Responses complete results and +stream events to Beta before outer stages run. Consequently the existing +`guardrail_anthropic_beta` Stage now governs Beta-, Chat-, and +Responses-backed providers without acquiring Responses-specific logic. +Chat → Responses converts the Chat request before provider finalization and +restores Responses complete/stream results to Chat wire DTOs. Its outer Chat +adapter therefore remains unchanged. Chat identity uses an explicit identity +registration; the outer adapter accepts both provider SDK values and +Bridge-produced wire DTOs, then applies the same public-model rewrite. +V1 → Responses uses the same Responses provider boundary but a distinct V1 +Bridge registration and typed V1 response recovery. Under `--stage`, V1 +Guardrails and MCP promote requests through the Beta working boundary; +unsupported feature combinations continue to select the entire legacy +lifecycle before Stage starts. Protocol recording is available on the twelve +registered routes under its separate opt-in gate. +Capability-missing pairs, feature-owned legacy lifecycles, +and the explicit response-roundtrip diagnostic remain on legacy. Debug routing +exposes the concrete `X-Tingly-Protocol-Pipeline: stage|legacy` decision. + +The first feature canary composes `guardrail_anthropic_beta` at the Beta working +boundary. `BuildTopology` inserts source→Beta above it and Beta→provider below +it as required, so provider responses return to Beta before response policy +runs. The same one-protocol Guardrail therefore covers native Beta, promoted +V1, Chat, and Responses ingress across Beta-, Chat-, and Responses-backed +providers in complete and stream modes. Real HTTP tests verify response +blocking and streamed tool-use rewriting; `harness matrix --stage --guardrails` supplies an +allow-only runtime for full semantic compatibility matrices. + +Verification recorded for the Phase 3 checkpoint: + +- `go test ./internal/protocoltest -count=1` passes the complete real HTTP + protocol test package; +- Guardrail Stage unit tests and the real Beta Guardrail HTTP canaries pass + under `-race`; +- `go vet` passes for the Guardrail Stage, protocoltest, and harness packages; +- `harness matrix --mode=single --stage --guardrails + --source=anthropic_beta` reports 72 cases: 66 passed, 6 expected + streaming-only skips, and 0 failures. + +Verification recorded for the Responses Guardrail checkpoint: + +- real HTTP complete and stream tests cover Responses→Beta Guardrail→provider + across Responses, Anthropic Beta, and OpenAI Chat targets; +- the no-`--stage` negative test remains on legacy; +- the MCP composition test proves the Guardrail remains outside the Beta Tool + Loop and blocks only the final external tool call after one owned execution; +- `harness matrix --mode=single --stage --guardrails + --source=openai_responses --scenario=text --client=http` passes 6/6 cases. + +Verification recorded for the Chat Guardrail checkpoint: + +- real HTTP complete and stream tests cover Chat→Beta Guardrail→provider across + OpenAI Chat, Anthropic Beta, and OpenAI Responses targets; +- the no-`--stage` negative test remains on legacy; +- the shared MCP composition test covers Chat ingress and preserves + `Guardrail(ToolLoop(Provider))` ordering; +- `harness matrix --mode=single --stage --guardrails --source=openai_chat + --scenario=text --client=http` passes 6/6 cases. + +Verification recorded for the native Responses checkpoint: + +- selector and wire-shape unit tests pass; +- real HTTP selection tests cover complete, stream, unsupported-route fallback, + and MCP fallback; +- raw HTTP and official OpenAI Go SDK text matrices each pass 2/2; +- the full Responses identity matrix reports 24 cases: 19 passed, 5 expected + capability/scenario skips, and 0 failures. + +Verification recorded for the Responses → Anthropic Beta checkpoint: + +- Bridge complete/stream tests cover request conversion, response recovery, + usage, side effects, event ordering, and stream close ownership; +- real HTTP selection tests cover complete and stream routing; +- raw HTTP and official OpenAI Go SDK text matrices each pass 2/2; +- the full route matrix reports 24 cases: 19 passed, 5 expected Responses + source/scenario skips, and 0 failures. + +Verification recorded for the Responses → OpenAI Chat checkpoint: + +- Bridge complete/stream tests cover Chat state, request conversion, response + recovery, usage, side effects, event ordering, and close ownership; +- real HTTP selection tests cover complete and stream routing; +- raw HTTP and official OpenAI Go SDK text matrices each pass 2/2; +- the full route matrix reports 24 cases: 19 passed, 5 expected Responses + source/scenario skips, and 0 failures. + +Verification recorded for the Anthropic Beta → Responses checkpoint: + +- Bridge complete/stream tests cover request conversion, typed Beta recovery, + usage, side effects, normalized events, and close ownership; +- real HTTP selection and authoritative Guardrail tests cover complete and + stream routing through a Responses provider; +- raw HTTP, official Go SDK, and allow-only Guardrails text matrices each pass + 2/2; +- plain and Guardrails full route matrices each report 24 cases: 22 passed, 2 + expected streaming-only skips, and 0 failures. + +Verification recorded for the OpenAI Chat → Responses checkpoint: + +- Bridge complete/stream tests cover request conversion, Chat wire recovery, + usage, side effects, model rewriting, and idempotent stream close ownership; +- real HTTP selection tests cover complete and stream routing through a + Responses provider; +- raw HTTP and official OpenAI Go SDK text matrices each pass 2/2; +- the combined raw/Go SDK full route matrix reports 40 cases: 33 passed, 7 + expected client/scenario capability skips, and 0 failures. + +Verification recorded for the Anthropic V1 → Responses checkpoint: + +- the distinct V1 Bridge test covers request conversion, typed V1 response + recovery, normalized usage, public model rewriting, and side effects; +- real HTTP selection tests cover complete and stream routing through a + Responses provider; +- the combined raw/Go SDK full route matrix reports 40 cases: 33 passed, 7 + expected client/scenario capability skips, and 0 failures. + +Verification recorded for the OpenAI Chat identity checkpoint: + +- the exact-pair selector and real HTTP tests cover complete and stream Stage + routing plus public model rewriting; +- the complete route/harness section reports 36 cases: 34 passed, 2 expected + streaming-only skips, and 0 failures; +- the official OpenAI Go SDK text matrix passes complete and stream 2/2. + +Verification recorded for the typed Wire DTO boundary checkpoint: + +- Chat and Responses complete builders preserve tool calls, refusal, cache and + reasoning usage details in typed output contracts; +- Chat → Responses, Responses → Chat, and Responses → Anthropic Beta real-path + text harnesses each pass complete and stream 2/2; +- their full raw/Go SDK matrices report respectively 33/40, 30/40, and 19/24 + passed, with only the documented capability/scenario skips and no failures; +- `go test ./internal/protocoltest -count=1` passes the full real HTTP protocol + suite. + +Verification recorded for the in-process Beta Tool Loop checkpoint: + +- Beta-native complete and stream tests cover owned, external, and mixed tool + calls, continuation, max rounds, usage, side effects, and stream ownership; +- MCP runtime tests cover direct Beta tool injection, Advisor filtering, the + existing servertool executor boundary, and provider-scoped continuation; +- V1→Beta tests preserve the full V1 wire subset and source-visible model for + complete and stream responses; +- composed topology tests run V1 → Beta Tool Loop → Beta provider in complete + and stream modes without production handler registration; +- RequestRecord tests retain two provider exchanges in one attempt and one + final response for both complete and stream; +- `harness matrix --mode=bridges --source=anthropic_v1 + --target=anthropic_beta` passes 6/6, and the V1 → Beta Stage → Chat filtered + matrix passes alongside direct V1→Chat (12/12); +- the full `internal/protocoltest` suite, targeted race suites, and `go vet` + pass. Production `--stage + MCP` selection is intentionally not claimed. + +Commit checkpoints, oldest to newest: + +| Commit | Checkpoint | +| --- | --- | +| `173509393` | Native Anthropic Beta Stage route | +| `0fc34defc` | Anthropic Beta → OpenAI Chat Stage route | +| `13babf39e` | Native Anthropic V1 Stage route | +| `580a28964` | Anthropic V1 → OpenAI Chat Stage route | +| `7eadb37ca` | Observe-only Guardrail Stage foundation | +| `430303114` | Authoritative Anthropic Beta Guardrail canary | +| `43ef808fd` | Native OpenAI Responses Stage route | +| `86c83e37d` | OpenAI Responses → Anthropic Beta Stage route | +| `58cc33247` | OpenAI Responses → OpenAI Chat Stage route | +| `e3bb6ba72` | Anthropic Beta → OpenAI Responses Stage route | +| `9dcdf3f7e` | OpenAI Chat → OpenAI Responses Stage route | +| `690012613` | Anthropic V1 → OpenAI Responses Stage route | +| `e52c9ab36` | Native OpenAI Chat Stage route | +| current checkpoint | Typed Wire DTO boundary for complete Bridge responses | +| `7d98b523f` | RequestRecord lifecycle foundation | +| `5471ccbb1` | Shared stream assembly for recording | +| `1eceb7474` | Provider boundary observer | +| `acc51481f` | RequestRecord persistence envelope | +| `b4b57a084` | Anthropic Beta identity recording canary | +| `ac3eef070` | Anthropic Beta recording across all targets | +| `5831bc81d` | Anthropic V1 recording across all targets | +| `5c058199e` | OpenAI Chat recording across all targets | +| `6827a4dbf` | OpenAI Responses recording across all targets | +| `b52d75af7` | Request-scoped recording across provider failover attempts | +| `71bbca421` | Beta-native Tool Loop complete lifecycle | +| `9abb86590` | Beta mixed ownership continuation | +| `cb255719d` | Beta-native Tool Loop streaming | +| `90e40dc2c` | MCP/runtime and servertool Beta adapters | +| `3aebef338` | Lossless V1 request projection to Beta | +| `7cee576ee` | V1→Beta Stage Bridge | +| `ef0f76a62` | Composed V1 ingress through Beta Tool Loop | +| `e29b71171` | Beta Tool Loop RequestRecord proof | +| `6ee32a2ae` | V1→Beta dormant harness matrix | diff --git a/.design/protocol-stage-tool-loop.md b/.design/protocol-stage-tool-loop.md new file mode 100644 index 000000000..454c95b0f --- /dev/null +++ b/.design/protocol-stage-tool-loop.md @@ -0,0 +1,463 @@ +# Protocol Stage Tool Loop + +> Status: Phase 4 canary active behind `--stage`. The in-process Beta +> implementation, production handler selection, and real HTTP harness path are +> complete; default rollout remains deferred. +> +> Canonical scope: MCP and server-tool model loops expressed as one +> `anthropic_beta` Protocol Stage. The broader chain and rollout remain defined +> in [`protocol-stage-chain.md`](./protocol-stage-chain.md). + +## Decision + +The production Tool Loop candidate uses `anthropic_beta` as its concrete +working protocol. + +This is not a new canonical AST and not a claim that Beta replaces every +protocol boundary. It is a deliberate protocol choice for one Stage: + +- Anthropic Beta has the most complete MCP, tool, thinking, and structured + content surface already used by Tingly-Box; +- MCP tools can remain Anthropic-native instead of passing through an OpenAI + DTO or a new neutral representation; +- existing Bridges convert other client/provider protocols at the Stage + boundary; +- MCP remains a tool source/runtime and `servertool` remains an executor; +- one Tool Loop owns injection, interception, execution, and continuation. + +The resulting shape has the same separation properties as multiple services or +HTTP middleware while remaining in-process: + +```text +client protocol + -> ingress Bridge + -> Guardrail Stage (optional) + -> Tool Loop Stage [anthropic_beta] + -> provider Bridge + -> Provider Endpoint +``` + +Requests move inward. Complete responses, stream events, errors, usage, and +side-effect facts return outward through the same levels. + +## Why Beta Instead of a Neutral Tool DTO + +Tool invocation is not just a function name plus JSON arguments. The native +protocol also defines: + +- request tool definitions and tool-choice behavior; +- `tool_use` and `tool_result` content blocks; +- thinking/signature preservation; +- structured and nested tool-result content; +- stream event ordering and block indexes; +- MCP, Advisor, and server-tool extensions; +- stop reasons and usage fields. + +A neutral DTO would have to reproduce this protocol surface and then maintain +another pair of conversions. That reduces type visibility while increasing +semantic conversion work. Beta already provides the required vocabulary, so +the Stage should use it directly. + +The earlier OpenAI Chat Tool Loop remains useful as a lifecycle proof. It is +not the production MCP normalization layer. + +## Goals + +- Give MCP/server tools one complete non-stream and stream lifecycle. +- Keep the Tool Loop independent of Gin, HTTP/SSE writers, routing, and + provider selection. +- Make ownership of every tool call explicit and deterministic. +- Reuse the existing MCP runtime, Anthropic adapter, servertool executor, and + protocol assemblers. +- Support other protocols through explicit, capability-complete Bridges. +- Preserve request identity, usage, model, cancellation, stream ownership, and + irreversible side-effect facts across multiple provider rounds. +- Keep activation additive behind the existing `--stage` process choice. + +## Non-Goals + +- No new user-facing stage mode, order editor, or MCP-specific startup flag. +- No replacement of all protocols with Anthropic Beta. +- No protocol-neutral tool/content AST. +- No separate MCP Stage and servertool Stage while both participate in the same + model continuation loop. +- No default cutover or removal of the legacy MCP path in Phase 4. +- No UI, API, Swagger, database, or persistence-format change. +- No remote service boundary yet; the Stage contract merely keeps that option + open. + +## Component Boundaries + +| Component | Owns | Does not own | +| --- | --- | --- | +| `tool_loop_anthropic_beta` | Round lifecycle, tool classification, execution sequencing, continuation, usage aggregation | Routing, HTTP, persistence | +| Beta Tool Provider | Injecting the exact enabled tools and returning their owned names | Executing calls or selecting providers | +| MCP runtime | Tool discovery, visibility, Advisor/server-tool configuration | Model round lifecycle | +| Servertool executor | Policy and actual tool execution | Protocol conversion or stream handling | +| Continuation store | Bounded provider/session-scoped mixed continuation segments, correlated by external tool-result IDs | General conversation history | +| Bridge | Request conversion inward and response/event conversion outward for one call | Tool ownership and execution | +| Provider Endpoint | One provider invocation in its concrete protocol | Tool continuation | +| Provider Observer | One provider-native request/response exchange per invocation | Stage snapshots | + +The Tool Loop is one Stage because injection, interception, execution, and +continuation are one indivisible model lifecycle. Splitting MCP and servertool +into peer Stages would give two levels competing to consume the same +`tool_use`. They may become separate dependencies or remote services, but not +independent protocol levels until their lifecycles are independently +meaningful. + +## Native Contracts and Names + +The concrete implementation names are part of the design vocabulary: + +| Name | Meaning | +| --- | --- | +| `tool_loop_anthropic_beta` | Stage name used in topology and diagnostics | +| `AnthropicBetaToolProvider` | Prepares a Beta request and returns exact owned tool names | +| `AnthropicBetaStageExecutor` | Existing server-tool execution boundary used by the Stage | +| `AnthropicBetaContinuationStore` | Typed Beta mixed-continuation interface | +| `ProtocolStageBetaToolProvider` | Server adapter from MCP runtime to the Stage provider | +| `ProviderBetaContinuationStore` | Provider/session-scoped bounded store implementation | +| `WithServertoolProviders` | Server option for registering additional in-process tool providers across startup and config reload | + +Ownership is an exact per-request snapshot returned by the tool provider. It is +not inferred from an MCP name prefix. A client tool and a server tool with the +same name are ambiguous and fail before the provider call. + +## Complete Lifecycle + +For each call, the Stage: + +1. Validates and deep-clones the Beta request so caller-owned input is not + mutated. +2. Applies one pending mixed continuation segment, if present. +3. Asks the Beta Tool Provider to inject enabled tools and return exact owned + names. +4. Rejects empty names, duplicate ownership, claimed-but-not-injected tools, + and collisions with client-declared tools. +5. Invokes the next Endpoint once for the current model round. +6. Extracts tool calls from the provider-native Beta message. +7. Classifies the response: + +| Provider result | Behavior | +| --- | --- | +| No owned tool calls | Return the response outward | +| Owned tool calls only | Execute them, append Beta tool results, continue another provider round | +| External/client tool calls only | Return unchanged for the client to execute | +| Mixed owned + external, store available | Execute owned calls, store the owned continuation, hide owned calls, return external calls | +| Mixed owned + external, no store | Return unchanged and execute nothing | + +8. Aggregates usage and model facts over every provider round. +9. Stops at the configured maximum round count. + +Mixed behavior is conservative. Without a continuation store, executing only +the internal half would create a conversation that cannot be resumed safely, +so the Stage leaves the entire response outward. + +## Mixed Continuation + +When one response contains both server-owned and client-owned calls, the Stage +cannot send internal tool results back to the provider until the client returns +its external results. + +The first request therefore: + +1. executes only owned calls; +2. builds a typed Beta assistant + internal-result continuation segment; +3. stores the segment under session, provider, and protocol identity; +4. removes owned calls from the outward response; +5. returns only client-owned calls. + +The following request: + +1. consumes the segment exactly once; +2. merges the client tool results with the stored internal results; +3. resumes the provider conversation. + +The store is bounded, provider-scoped, and single-consume. A session may have +more than one pending segment; the current request consumes only the segment +whose expected external tool IDs are present in the trailing client +`tool_result` turn. An unrelated request leaves every segment untouched. + +Cross-request continuation requires an explicit user or +`X-Tingly-Session-ID` identity. Client IP fallback is deliberately rejected: +an IP address is not a conversation identity and may represent many users +behind one NAT. Expired entries are swept during reads and writes, and a global +capacity bound evicts the oldest expiry so abandoned continuations cannot grow +without limit. The Stage itself does not know the storage key or routing +identity. + +## Streaming Lifecycle + +### Round buffering is required + +Anthropic streams may emit text or thinking blocks before a later `tool_use`. +No visible prefix proves that the round contains no internal tool call. If the +prefix were emitted immediately, a later internal call could not be hidden +without producing an invalid or partially leaked client stream. + +The Beta Tool Loop therefore buffers exactly one provider round: + +1. pull events only when the outer caller invokes `Next`; +2. assemble the round while retaining its native events; +3. close the inner provider stream exactly once; +4. classify the completed Beta message; +5. either continue internally or replay a client-visible round. + +This is a protocol-correctness decision, not an implementation shortcut. +Future TTFT optimization must preserve the invariant and cannot assume tool +blocks arrive before text/thinking. + +### Stream classification + +| Round type | Outward events | +| --- | --- | +| No owned calls | Replay the complete native round | +| Owned calls only | Replay nothing; execute, append results, start next round | +| External calls only | Replay the complete native round | +| Mixed with store | Remove owned block start/delta/stop events and renumber remaining indexes | +| Mixed without store | Replay unchanged and execute nothing | + +Usage, model, and `SideEffectsCommitted` are monotonic across all hidden and +visible rounds. Cancellation and provider errors propagate through the outer +stream. Every acquired inner stream is closed once. + +## Protocol Topology + +### Client ingress to the Beta Tool Loop + +| Client protocol | Ingress | +| --- | --- | +| `anthropic_beta` | Identity | +| `anthropic_v1` | V1→Beta subset Bridge | +| `openai_chat` | Chat→Beta Bridge | +| `openai_responses` | Responses→Beta Bridge | + +### Beta Tool Loop to provider + +| Provider protocol | Provider boundary | +| --- | --- | +| `anthropic_beta` | Identity | +| `openai_chat` | Beta→Chat Bridge | +| `openai_responses` | Beta→Responses Bridge | +| `anthropic_v1` | Not currently supported below the Beta Stage | + +Exact-pair registration remains mandatory. The topology builder does not search +for arbitrary transitive conversion paths at runtime. + +## Anthropic V1 Compatibility Contract + +V1 and Beta are distinct protocols even though V1 create-message requests are +a structural Beta subset. + +### Request direction + +V1→Beta request conversion uses JSON marshal/unmarshal as the contract. It is +lossless for the currently pinned V1 create-message request surface and avoids +manual field copying. + +### Response direction + +The compatibility guarantee stops at request promotion. Beta→V1 complete +responses and stream events keep the existing permissive JSON projection in +this phase; they are not checked against a maintained V1 response subset. +Beta-only output may therefore lack equivalent V1 typed semantics. + +This is an explicit scope decision rather than a production blocker for the +V1-with-MCP canary. Strict response/event subset validation is deferred until +a concrete compatibility requirement or production evidence justifies it. + +This decision does not make the Bridge unconditional bidirectional +compatibility. In particular, a provider whose concrete protocol is +`anthropic_v1` still cannot sit below the Beta Tool Loop without a separately +designed Beta→V1 request Bridge. + +SDK regeneration must continue to verify V1→Beta request compatibility. A +future strict response contract would require its own maintained response and +event compatibility suite. + +## Guardrail Ordering + +The intended order is: + +```text +Guardrails(ToolLoop(Provider)) +``` + +Consequences: + +- inbound Guardrails inspect original user content before tool injection; +- ToolLoop consumes internal rounds and produces the actual final response; +- outbound Guardrails evaluate only the client-visible final result; +- tool authorization occurs immediately before executor invocation. + +This order is fixed by product semantics. Phase 4 does not introduce a user +setting for stage order. + +Under `--stage`, V1 Guardrails independently promote the request to the Beta +working protocol. With MCP enabled, both Beta-native stages compose at that +same boundary. Chat and Responses ingress reach the same ordering through their +source→Beta Bridges and return through the provider-side Bridge. V1 +promotion applies only to requests; outward V1 responses +continue to use the existing permissive projection and do not establish a new +strict Beta-to-V1 compatibility contract. Without `--stage`, V1 Guardrails +retain the complete legacy lifecycle. + +## Recording and Usage + +`RequestRecord` and `UsageRecord` remain separate concerns. + +Recording attaches only at stable boundaries: + +```text +original client request + -> Stages / Bridges + -> Provider Observer: exchange 1..N + -> Stages / Bridges + -> final client response +``` + +The Tool Loop adds no recording hook. Each provider round naturally crosses +the already-observed Provider Endpoint and becomes one ordered +`ProviderExchange` under the same attempt. One incoming request still has one +original input and one final outward response. Intermediate Stage responses are +not persisted. + +Usage aggregation remains a protocol-neutral Stage result. Disabling recording +must not disable usage, and disabling usage must not disable recording. + +## Errors, Side Effects, and Failover + +Two commitment facts are independent: + +- output committed: a client-visible event has been written; +- side effects committed: execution crossed into the server-tool runtime. + +After runtime dispatch begins, later tool, provider, conversion, or stream +failures carry `SideEffectsCommitted=true`. A runtime may perform an external +action and then lose its response, so treating only successful returns as the +boundary could replay that action during failover. Validation, disabled-tool, +and policy failures before dispatch do not commit side effects. The outer +failover orchestrator must not restart the whole attempt on another provider +after either commitment boundary. + +## Production Activation + +Activation reuses the existing `--stage` server startup choice. No new mode or +flag is introduced. + +Selection is per provider attempt: + +| Condition | Pipeline | +| --- | --- | +| `--stage` disabled | Complete legacy lifecycle | +| No MCP/tool-loop feature | Existing plain Stage selection | +| MCP enabled and exact Beta Tool Loop topology is complete | Beta Tool Loop Stage | +| Missing Bridge/capability/dependency | Complete legacy lifecycle before provider invocation | + +Once a Stage attempt begins, it never falls back into legacy mid-attempt. +Rollback is a restart without `--stage`. + +Compatibility never activates this path by itself: MCP enablement, Guardrails, +and a complete Beta topology are necessary feature conditions, but the +process-level `--stage` choice remains the first and mandatory gate. + +The production wiring checkpoint implements: + +1. replace the current unconditional MCP→legacy selection with an exact + topology eligibility check; +2. construct the Beta tool provider, existing executor, and provider-bound + continuation store for each attempt; +3. compose Guardrail, Tool Loop, provider Bridge, Provider Observer, and + terminal endpoint in the documented order; +4. promote V1 MCP requests to Beta while preserving the current permissive + Beta→V1 response/event projection; +5. add a debug-level entry naming `tool_loop_anthropic_beta` and every concrete + protocol boundary; +6. verify through the real HTTP path with `harness matrix --stage --mcp` before + calling the canary active. + +All six items are implemented. The owned-tool production-path matrix is: + +```bash +go run ./cli/harness matrix --mode=single --stage --mcp \ + --scenario=mcp_owned_tool +``` + +It covers all 13 registered source/target labels in complete and streaming +modes. A selection that produces no executable cases is an error; an all-skip +run can no longer report false success. + +## Diagnostics and UX-First Review + +Diagnostics must answer: + +1. Did this request use Stage or legacy? +2. Which concrete protocols and named Stages did it traverse? +3. How many provider rounds and tool executions occurred? +4. Where did failure or commitment happen? + +The design follows the product UX principles: + +- reuse `--stage`; do not add a mode picker; +- display concrete protocol names such as `anthropic_beta`, not aliases; +- keep MCP enablement and Stage rollout as orthogonal axes; +- make smart fallback automatic when topology is unsupported; +- make harness diagnostics traverse the real path; dormant Bridge tests are + labeled as such and cannot be presented as production evidence; +- keep rollback and re-entry explicit through process restart; +- scope side effects to the current attempt and stop failover after commitment. + +## Current Implementation Checkpoint + +Implemented and committed: + +- Beta complete and streaming Tool Loop; +- exact owned/external/mixed tool classification; +- provider-scoped mixed continuation correlated to the current external + tool-result turn, with explicit-session, TTL, and capacity bounds; +- direct MCP runtime→Beta tool definitions; +- reuse of the existing Anthropic adapter and servertool executor; +- lossless V1 request promotion and V1→Beta Bridge; +- complete/stream RequestRecord multi-exchange proof; +- composed V1→Beta Tool Loop tests; +- dormant 54-cell Bridge matrix including V1→Beta and V1→Beta Stage→Chat; +- exact two-boundary production selection for all four ingress protocols; +- production Beta Tool Loop assembly with MCP runtime, servertool executor, + provider-scoped continuation, Guardrail ordering, and provider observation; +- failover suppression after server-tool runtime dispatch, including runtime + errors whose external side effects may already have happened; +- real HTTP complete/stream V1→Beta MCP canary through `harness matrix`; +- persisted real-HTTP RequestRecord validation across all 26 owned-tool route + cases, proving original input, two provider exchanges, and final output; +- real HTTP owned-tool fixture, driven by raw HTTP plus the official Anthropic + and OpenAI Go SDKs, proving provider round 1 → local execution → provider + round 2 → final response across all four ingress protocols; +- normalized converted Anthropic events expose their wire payload to protocol + assemblers, so cross-protocol streams remain interceptable by the Beta loop; +- additional `servertool.ToolProvider` instances can enter through the formal + server option and survive config-driven pipeline rebuilds. + +Pending by design: + +- optional strict Beta→V1 response/event validation, deferred until a concrete + compatibility need appears; +- default rollout and legacy removal. + +## Acceptance Criteria for Production Wiring + +- Exact supported source/target pairs select the Beta Tool Loop only with + `--stage` and the existing MCP feature enabled. +- Unsupported pairs choose legacy before any provider or tool side effect. +- Complete and stream paths produce the same final semantics as the existing + MCP lifecycle for owned, external, and mixed calls. +- Provider rounds are recorded as ordered exchanges in one request/attempt. +- Usage and source-visible model remain correct across hidden rounds. +- Tool name collisions fail before the provider call. +- Later failures after tool success prevent failover replay. +- V1 MCP requests are promoted to Beta; outward responses/events retain the + current permissive V1 projection behavior. +- `harness matrix --mode=single --stage --mcp + --scenario=mcp_owned_tool` executes all 26 complete/stream route cases with + no skips and exposes the concrete Stage path in debug logs. +- Starting without `--stage` leaves current behavior unchanged. diff --git a/README.md b/README.md index 6feca5257..a416dff52 100644 --- a/README.md +++ b/README.md @@ -351,6 +351,11 @@ Run the Go server (hot-reload via `go run`): task start # or directly: go run ./cli/tingly-box --verbose start --debug --port 12580 --browser=false + +# Opt into the Protocol Stage pipeline (currently Chat→Beta, Beta→Beta, +# Beta→Chat, V1→V1, and V1→Chat). Beta Guardrails are Stage-native; +# MCP, recording, V1 Guardrails, and unsupported routes remain on legacy. +go run ./cli/tingly-box start --stage ``` Open http://localhost:12580 in your browser (serves the last built frontend bundle). diff --git a/cli/harness/README.md b/cli/harness/README.md index 8ef827734..c09e476a0 100644 --- a/cli/harness/README.md +++ b/cli/harness/README.md @@ -69,6 +69,13 @@ go build -o harness ./cli/harness # Tier A — exhaustive protocol-transform matrix ./harness matrix ./harness matrix --scenario text --source anthropic_v1 --target openai_chat +./harness matrix --mode=bridges # dormant Stage/Bridge topology, in-process +./harness matrix --mode=single --stage --source=openai_chat --target=anthropic_beta # real server Stage path +./harness matrix --mode=single --stage --source=openai_chat --target=openai_responses # Chat→Responses Stage path +./harness matrix --mode=single --stage --source=openai_chat --target=openai_chat # native Chat Stage path +./harness matrix --mode=single --stage --source=anthropic_v1 --target=openai_responses # V1→Responses Stage path +./harness matrix --mode=single --stage --source=anthropic_beta --target=anthropic_beta # native Beta Stage path +./harness matrix --mode=single --stage --mcp --source=anthropic_v1 --target=anthropic_v1 # real V1→Beta Tool Loop path # Tier A through real client stacks (--client; see .design/harness-matrix.md # "Client drivers"): official Go SDKs in-process, or real Python/Node SDKs @@ -120,10 +127,34 @@ functions. - Known-broken cells are centralized in `protocoltest.skipSourceScenarios` (e.g. `openai_responses|tool_use`). - `--json` for CI; `-v` / `-vv` to raise log verbosity; `--record-dir` to dump - request/response pairs; `--batch N` for stability runs. + request/response pairs; `--batch N` for stability runs. In the ephemeral + harness environment, `--record-dir` enables `recording_v2` for the Anthropic + and OpenAI scenarios and flushes every sink before shutdown, so the returned + files are complete test artifacts. **Use it for:** catching transform regressions exhaustively and instantly. +### Dormant Stage/Bridge section + +`./harness matrix --mode=bridges` validates the additive protocol Stage path +without routing production gateway traffic through it. Results are visibly +prefixed with `bridges/`; concrete multi-level results use `bridges/chain/`. +The 42-cell section covers Anthropic v1/beta/OpenAI Chat identity, +Anthropic v1/beta → OpenAI Chat, OpenAI Chat → Anthropic Beta, and a real +OpenAI Chat → Anthropic Beta-native Stage → OpenAI Chat topology for text, +tool-use, and tool-result requests in both execution modes. It reuses the +normal matrix filters and batch option. +The standalone mode rejects external client drivers, MCP, and HTTP recording, +because none of those surfaces are traversed by the in-process topology. + +This is converter/topology evidence, not a substitute for the production HTTP +single-hop section. Use `matrix --mode=single --stage` to enable the real server +selector; the current production Stage routes are OpenAI Chat → OpenAI Chat/Anthropic Beta/OpenAI Responses, +Anthropic Beta → Anthropic Beta/OpenAI Chat/OpenAI Responses, Anthropic V1 → +Anthropic V1/OpenAI Chat/OpenAI Responses, and OpenAI Responses → OpenAI +Responses/Anthropic Beta/OpenAI Chat. V1 remains a distinct protocol with +separate Bridges. This planned protocol-pair surface is now complete. + --- ## Tier B — `replay` diff --git a/docs/guardrails.md b/docs/guardrails.md index 4c1f16fd5..c556a387f 100644 --- a/docs/guardrails.md +++ b/docs/guardrails.md @@ -2,6 +2,28 @@ Guardrails adds rule-based safety checks around model output, tool calls, tool results, and protected credentials. +## Protocol Stage canary + +Guardrails behavior remains enabled through the existing scenario setting; no +new user-facing Guardrails mode was added. When the server is also started with +`--stage`, Anthropic Beta requests routed to either an Anthropic Beta provider +or an OpenAI Chat provider use the Beta-native `guardrail_anthropic_beta` +Stage. The Stage owns request masking/filtering, complete-response evaluation, +and streaming tool-use evaluation as one full-duplex lifecycle. + +All other combinations keep their existing behavior: + +- without `--stage`, every Guardrails request uses the legacy pipeline; +- Anthropic V1 Guardrails still use legacy even when `--stage` is active; +- MCP-enabled or protocol-recorded Beta requests select the complete legacy + lifecycle rather than mixing Stage and legacy ownership; +- unsupported protocol pairs stay legacy. + +Rollback is a server restart without `--stage`. No Guardrails config or stored +policy migration is required. For development verification, +`harness matrix --stage --guardrails` uses an allow-only test runtime; that +harness flag does not enable or configure production policies. + ## What Guardrails manages Guardrails is organized into three user-facing areas: From 25383fcc8e312b4bfb67caf10c2641946acde0f4 Mon Sep 17 00:00:00 2001 From: FFengIll Date: Tue, 4 Aug 2026 19:42:50 +0800 Subject: [PATCH 2/8] feat(protocol,stage): port composable protocol stage chain + recording from hardening Port the self-contained stage core from codex/protocol-stage-hardening (off #1491): - internal/protocol/stage: bridge/compose/registry/topology + anthropic/openai/responses bridges, guardrail, toolloop - internal/protocol/{assembler,nonstream,stream,transform,wire,request}: protocol-layer additions and adjustments the stage depends on - internal/record: new request recording lifecycle (recorder, provider_endpoint, boundary matrix) - internal/obs: record/sink/slim additions + request_record test - internal/guardrails/mutate: RewriteAnthropicToolUseEventDecision used by stage guardrail All dependencies (protocol/*, guardrails, record, obs) land in this commit so the unit builds clean: go vet ./internal/protocol/... ./internal/guardrails/... ./internal/record/... ./internal/obs/ Wiring into the server (protocol_stage_*.go glue, forwarding path fix, server skeleton 3-way merge) and cross-stage test matrix follow in subsequent commits. Batch 1+2 of the protocol-stage-hardening port. --- .../guardrails/mutate/anthropic_stream.go | 35 +- .../mutate/anthropic_stream_test.go | 28 + internal/obs/noop_exporter.go | 15 + internal/obs/record.go | 11 +- internal/obs/request_record_test.go | 110 +++ internal/obs/sink.go | 38 + internal/obs/slim.go | 46 +- .../assembler/openai_responses_assembler.go | 13 +- .../protocol/assembler/stream_assembler.go | 219 ++++++ .../assembler/stream_assembler_test.go | 260 +++++++ internal/protocol/json_snapshot.go | 50 ++ internal/protocol/nonstream/anthropic.go | 12 +- internal/protocol/nonstream/nonstream_test.go | 23 + .../protocol/nonstream/openai_to_anthropic.go | 112 ++- .../openai_to_anthropic_semantics_test.go | 73 ++ internal/protocol/nonstream/openai_to_chat.go | 108 +-- .../protocol/nonstream/openai_to_chat_test.go | 26 + .../protocol/nonstream/openai_to_responses.go | 193 +++-- .../protocol/nonstream/openai_usage_test.go | 88 +++ .../protocol/request/anthropic_v1_to_beta.go | 33 +- .../request/anthropic_v1_to_beta_test.go | 69 ++ .../protocol/stage/anthropicbridge/bridge.go | 214 ++++++ .../stage/anthropicbridge/bridge_test.go | 615 +++++++++++++++ .../stage/anthropicbridge/responses.go | 147 ++++ .../stage/anthropicbridge/responses_stream.go | 110 +++ .../stage/anthropicbridge/responses_test.go | 195 +++++ .../protocol/stage/anthropicbridge/stream.go | 131 ++++ .../protocol/stage/anthropicbridge/v1_beta.go | 219 ++++++ .../stage/anthropicbridge/v1_beta_test.go | 193 +++++ internal/protocol/stage/bridge.go | 246 ++++++ internal/protocol/stage/bridge_test.go | 699 ++++++++++++++++++ internal/protocol/stage/capabilities.go | 76 ++ internal/protocol/stage/compose.go | 88 +++ internal/protocol/stage/compose_test.go | 431 +++++++++++ internal/protocol/stage/doc.go | 8 + internal/protocol/stage/endpoint.go | 91 +++ .../stage/guardrail/anthropic_beta.go | 328 ++++++++ .../stage/guardrail/anthropic_beta_test.go | 311 ++++++++ .../protocol/stage/guardrail/guardrail.go | 222 ++++++ .../stage/guardrail/guardrail_test.go | 268 +++++++ internal/protocol/stage/identity.go | 53 ++ .../protocol/stage/openaibridge/bridge.go | 153 ++++ .../stage/openaibridge/bridge_test.go | 529 +++++++++++++ .../protocol/stage/openaibridge/responses.go | 119 +++ .../stage/openaibridge/responses_stream.go | 113 +++ .../stage/openaibridge/responses_test.go | 150 ++++ .../protocol/stage/openaibridge/stream.go | 162 ++++ internal/protocol/stage/registry.go | 98 +++ .../protocol/stage/responsesbridge/bridge.go | 138 ++++ .../stage/responsesbridge/bridge_test.go | 203 +++++ .../protocol/stage/responsesbridge/chat.go | 120 +++ .../stage/responsesbridge/chat_stream.go | 111 +++ .../stage/responsesbridge/chat_test.go | 186 +++++ .../protocol/stage/responsesbridge/stream.go | 153 ++++ .../protocol/stage/toolloop/openai_chat.go | 324 ++++++++ .../stage/toolloop/openai_chat_stream.go | 290 ++++++++ .../stage/toolloop/openai_chat_test.go | 518 +++++++++++++ internal/protocol/stage/toolloop/runtime.go | 123 +++ .../protocol/stage/toolloop/runtime_test.go | 34 + internal/protocol/stage/topology.go | 90 +++ internal/protocol/stage/topology_test.go | 231 ++++++ .../anthropic_beta_to_openai_responses.go | 2 +- ...opic_beta_to_openai_responses_converter.go | 121 +-- ...ic_beta_to_openai_responses_golden_test.go | 12 + .../stream/anthropic_to_openai_converter.go | 25 +- .../anthropic_to_openai_converter_test.go | 99 +++ .../openai_chat_to_responses_converter.go | 68 +- .../openai_chat_to_responses_golden_test.go | 24 +- .../stream/openai_chat_to_responses_test.go | 45 +- ...ai_responses_to_anthropic_assembly_test.go | 4 +- ...openai_responses_to_anthropic_converter.go | 23 +- ...enai_responses_to_anthropic_golden_test.go | 2 + .../openai_responses_to_chat_converter.go | 10 + .../protocol/stream/openai_to_anthropic.go | 6 + .../stream/openai_to_anthropic_beta.go | 6 + .../stream/openai_to_anthropic_converter.go | 45 +- .../protocol/transform/provider_cleanup.go | 28 + .../transform/provider_cleanup_test.go | 29 + internal/protocol/wire/openai_chat.go | 106 ++- internal/protocol/wire/openai_responses.go | 162 +++- internal/record/boundary_matrix_test.go | 224 ++++++ internal/record/provider_endpoint.go | 171 +++++ internal/record/provider_endpoint_test.go | 327 ++++++++ internal/record/record.go | 114 +++ internal/record/recorder.go | 249 +++++++ internal/record/recorder_test.go | 208 ++++++ 86 files changed, 11832 insertions(+), 330 deletions(-) create mode 100644 internal/obs/noop_exporter.go create mode 100644 internal/obs/request_record_test.go create mode 100644 internal/protocol/assembler/stream_assembler.go create mode 100644 internal/protocol/assembler/stream_assembler_test.go create mode 100644 internal/protocol/json_snapshot.go create mode 100644 internal/protocol/nonstream/openai_to_anthropic_semantics_test.go create mode 100644 internal/protocol/stage/anthropicbridge/bridge.go create mode 100644 internal/protocol/stage/anthropicbridge/bridge_test.go create mode 100644 internal/protocol/stage/anthropicbridge/responses.go create mode 100644 internal/protocol/stage/anthropicbridge/responses_stream.go create mode 100644 internal/protocol/stage/anthropicbridge/responses_test.go create mode 100644 internal/protocol/stage/anthropicbridge/stream.go create mode 100644 internal/protocol/stage/anthropicbridge/v1_beta.go create mode 100644 internal/protocol/stage/anthropicbridge/v1_beta_test.go create mode 100644 internal/protocol/stage/bridge.go create mode 100644 internal/protocol/stage/bridge_test.go create mode 100644 internal/protocol/stage/capabilities.go create mode 100644 internal/protocol/stage/compose.go create mode 100644 internal/protocol/stage/compose_test.go create mode 100644 internal/protocol/stage/doc.go create mode 100644 internal/protocol/stage/endpoint.go create mode 100644 internal/protocol/stage/guardrail/anthropic_beta.go create mode 100644 internal/protocol/stage/guardrail/anthropic_beta_test.go create mode 100644 internal/protocol/stage/guardrail/guardrail.go create mode 100644 internal/protocol/stage/guardrail/guardrail_test.go create mode 100644 internal/protocol/stage/identity.go create mode 100644 internal/protocol/stage/openaibridge/bridge.go create mode 100644 internal/protocol/stage/openaibridge/bridge_test.go create mode 100644 internal/protocol/stage/openaibridge/responses.go create mode 100644 internal/protocol/stage/openaibridge/responses_stream.go create mode 100644 internal/protocol/stage/openaibridge/responses_test.go create mode 100644 internal/protocol/stage/openaibridge/stream.go create mode 100644 internal/protocol/stage/registry.go create mode 100644 internal/protocol/stage/responsesbridge/bridge.go create mode 100644 internal/protocol/stage/responsesbridge/bridge_test.go create mode 100644 internal/protocol/stage/responsesbridge/chat.go create mode 100644 internal/protocol/stage/responsesbridge/chat_stream.go create mode 100644 internal/protocol/stage/responsesbridge/chat_test.go create mode 100644 internal/protocol/stage/responsesbridge/stream.go create mode 100644 internal/protocol/stage/toolloop/openai_chat.go create mode 100644 internal/protocol/stage/toolloop/openai_chat_stream.go create mode 100644 internal/protocol/stage/toolloop/openai_chat_test.go create mode 100644 internal/protocol/stage/toolloop/runtime.go create mode 100644 internal/protocol/stage/toolloop/runtime_test.go create mode 100644 internal/protocol/stage/topology.go create mode 100644 internal/protocol/stage/topology_test.go create mode 100644 internal/protocol/stream/anthropic_to_openai_converter_test.go create mode 100644 internal/protocol/transform/provider_cleanup.go create mode 100644 internal/protocol/transform/provider_cleanup_test.go create mode 100644 internal/record/boundary_matrix_test.go create mode 100644 internal/record/provider_endpoint.go create mode 100644 internal/record/provider_endpoint_test.go create mode 100644 internal/record/record.go create mode 100644 internal/record/recorder.go create mode 100644 internal/record/recorder_test.go diff --git a/internal/guardrails/mutate/anthropic_stream.go b/internal/guardrails/mutate/anthropic_stream.go index 9e618c27f..5b4b7c2ed 100644 --- a/internal/guardrails/mutate/anthropic_stream.go +++ b/internal/guardrails/mutate/anthropic_stream.go @@ -2,6 +2,7 @@ package mutate import ( "encoding/json" + "fmt" "strings" "github.com/anthropics/anthropic-sdk-go" @@ -50,6 +51,17 @@ func RewriteAnthropicToolUseEvent( streamState *protocol.GuardrailsStreamState, event interface{}, ) (bool, []AnthropicBufferedEvent, error) { + kind, rewritten, err := RewriteAnthropicToolUseEventDecision(credentialMask, streamState, event) + return kind != AnthropicToolUseDecisionNone, rewritten, err +} + +// RewriteAnthropicToolUseEventDecision is the decision-preserving form used by +// callers that must distinguish a policy block from buffered passthrough. +func RewriteAnthropicToolUseEventDecision( + credentialMask *guardrailscore.CredentialMaskState, + streamState *protocol.GuardrailsStreamState, + event interface{}, +) (AnthropicToolUseDecisionKind, []AnthropicBufferedEvent, error) { var ( eventType string index int @@ -60,7 +72,7 @@ func RewriteAnthropicToolUseEvent( switch evt := event.(type) { case *anthropic.MessageStreamEventUnion: if evt == nil { - return false, nil, nil + return AnthropicToolUseDecisionNone, nil, nil } eventType = evt.Type index = int(evt.Index) @@ -68,23 +80,26 @@ func RewriteAnthropicToolUseEvent( rawJSON = strings.Clone(evt.RawJSON()) case *anthropic.BetaRawMessageStreamEventUnion: if evt == nil { - return false, nil, nil + return AnthropicToolUseDecisionNone, nil, nil } eventType = evt.Type index = int(evt.Index) block = evt.ContentBlock rawJSON = strings.Clone(evt.RawJSON()) default: - return false, nil, nil + return AnthropicToolUseDecisionNone, nil, nil } if !ShouldRewriteAnthropicEvent(streamState, eventType, block) { - return false, nil, nil + return AnthropicToolUseDecisionNone, nil, nil } var eventMap map[string]interface{} if err := json.Unmarshal([]byte(rawJSON), &eventMap); err != nil { - return false, nil, err + return AnthropicToolUseDecisionNone, nil, err + } + if eventMap == nil { + return AnthropicToolUseDecisionNone, nil, fmt.Errorf("decode Anthropic stream event: payload is not an object") } if eventType != "" { eventMap["type"] = eventType @@ -93,15 +108,15 @@ func RewriteAnthropicToolUseEvent( decision := HandleAnthropicToolUseBuffer(credentialMask, streamState, eventType, index, block, eventMap) switch decision.Kind { case AnthropicToolUseDecisionBuffer: - return true, nil, nil + return decision.Kind, nil, nil case AnthropicToolUseDecisionBlock: if decision.BlockMessage == "" { - return true, nil, nil + return decision.Kind, nil, nil } if streamState != nil { streamState.RewroteBlockedToolUse = true } - return true, []AnthropicBufferedEvent{ + return decision.Kind, []AnthropicBufferedEvent{ { EventType: anthropicEventTypeContentBlockStart, Payload: map[string]interface{}{ @@ -133,9 +148,9 @@ func RewriteAnthropicToolUseEvent( }, }, nil case AnthropicToolUseDecisionPassthrough: - return true, decision.Passthrough, nil + return decision.Kind, decision.Passthrough, nil default: - return false, nil, nil + return AnthropicToolUseDecisionNone, nil, nil } } diff --git a/internal/guardrails/mutate/anthropic_stream_test.go b/internal/guardrails/mutate/anthropic_stream_test.go index 4f62e5791..cbf56bc69 100644 --- a/internal/guardrails/mutate/anthropic_stream_test.go +++ b/internal/guardrails/mutate/anthropic_stream_test.go @@ -1,11 +1,39 @@ package mutate import ( + "encoding/json" "testing" + "github.com/anthropics/anthropic-sdk-go" "github.com/tingly-dev/tingly-box/internal/protocol" ) +func TestRewriteAnthropicToolUseEventDecisionDistinguishesAllowedPassthrough(t *testing.T) { + state := &protocol.GuardrailsStreamState{ + PendingBlockMessages: make(map[string]string), PendingBlockedIndex: make(map[int]string), + AnthropicToolEvents: make(map[int][]protocol.GuardrailsBufferedEvent), AnthropicToolIDs: make(map[int]string), + } + events := []string{ + `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"tool-1","name":"lookup","input":{}}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":\"safe\"}"}}`, + `{"type":"content_block_stop","index":0}`, + } + want := []AnthropicToolUseDecisionKind{AnthropicToolUseDecisionBuffer, AnthropicToolUseDecisionBuffer, AnthropicToolUseDecisionPassthrough} + for i, raw := range events { + var event anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal([]byte(raw), &event); err != nil { + t.Fatalf("decode event %d: %v", i, err) + } + kind, _, err := RewriteAnthropicToolUseEventDecision(nil, state, &event) + if err != nil { + t.Fatalf("event %d: %v", i, err) + } + if kind != want[i] { + t.Fatalf("event %d decision = %q, want %q", i, kind, want[i]) + } + } +} + func TestHandleAnthropicToolUseBuffer_RewritesBlockedMessageDeltaStopReason(t *testing.T) { streamState := &protocol.GuardrailsStreamState{ RewroteBlockedToolUse: true, diff --git a/internal/obs/noop_exporter.go b/internal/obs/noop_exporter.go new file mode 100644 index 000000000..2ed3fd48a --- /dev/null +++ b/internal/obs/noop_exporter.go @@ -0,0 +1,15 @@ +package obs + +import "context" + +// noopExporter is a sink exporter used by experiments/tests that need a valid +// Sink without writing records anywhere. +type noopExporter struct{} + +// NewNoopExporter returns a RecordExporter that discards all records. +func NewNoopExporter() RecordExporter { + return noopExporter{} +} + +func (noopExporter) Export(context.Context, []*Record) error { return nil } +func (noopExporter) Shutdown(context.Context) error { return nil } diff --git a/internal/obs/record.go b/internal/obs/record.go index 9097a4d18..b8c6b9010 100644 --- a/internal/obs/record.go +++ b/internal/obs/record.go @@ -1,6 +1,10 @@ package obs -import "time" +import ( + "time" + + requestrecord "github.com/tingly-dev/tingly-box/internal/record" +) // Record is the canonical data model for one LLM request/response cycle. // Construct it on the hot path and pass to Sink.Emit; the only cost is a @@ -23,6 +27,11 @@ type Record struct { ProviderResponse *RecordResponse FinalResponse *RecordResponse + // RequestRecord is the additive Protocol Stage recording envelope. Legacy + // request/response fields remain unchanged while the new recorder is + // canaried behind the Stage and scenario recording switches. + RequestRecord *requestrecord.RequestRecord + Duration time.Duration Err string Steps []string diff --git a/internal/obs/request_record_test.go b/internal/obs/request_record_test.go new file mode 100644 index 000000000..ea14e7601 --- /dev/null +++ b/internal/obs/request_record_test.go @@ -0,0 +1,110 @@ +package obs + +import ( + "context" + "testing" + "time" + + "github.com/tingly-dev/tingly-box/internal/protocol" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" +) + +func TestSinkEmitRequestRecordUsesExistingPipeline(t *testing.T) { + exporter := &recordingExporter{} + sink := NewSink("", RecordModeStagedRequestResponse, WithExporters(exporter)) + if sink == nil { + t.Fatal("NewSink returned nil") + } + t.Cleanup(sink.Close) + + started := time.Now().UTC() + requestRecord := &requestrecord.RequestRecord{ + Timestamp: started, + RequestID: "request-id", + SessionID: "session-id", + Scenario: "claude_code", + Outcome: requestrecord.OutcomeSucceeded, + Duration: time.Second, + InputRequest: requestrecord.Payload{ + Protocol: protocol.TypeAnthropicBeta, + }, + ProviderExchanges: []requestrecord.ProviderExchange{{ + Provider: "provider", + Model: "provider-model", + Protocol: protocol.TypeAnthropicBeta, + }}, + } + + sink.EmitRequestRecord(requestRecord) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := sink.ForceFlush(ctx); err != nil { + t.Fatalf("ForceFlush: %v", err) + } + if len(exporter.batches) != 1 || len(exporter.batches[0]) != 1 { + t.Fatalf("exported batches = %#v", exporter.batches) + } + + got := exporter.batches[0][0] + if got.RequestRecord == requestRecord { + t.Fatal("request record was not detached before asynchronous export") + } + if got.Provider != "provider" || got.Model != "provider-model" { + t.Fatalf("provider/model = %q/%q", got.Provider, got.Model) + } + if full := FullRecord(got); full.RequestRecord != got.RequestRecord { + t.Fatal("full exporter shape dropped request_record") + } +} + +func TestSinkEmitRequestRecordHonorsRecordingMode(t *testing.T) { + for _, testCase := range []struct { + name string + mode RecordMode + wantProviderResponse bool + wantFinalResponse bool + }{ + {name: "request", mode: RecordModeRequestOnly}, + {name: "request response", mode: RecordModeRequestResponse, wantFinalResponse: true}, + {name: "staged", mode: RecordModeStagedRequestResponse, wantProviderResponse: true, wantFinalResponse: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + exporter := &recordingExporter{} + sink := NewSink("", testCase.mode, WithExporters(exporter)) + t.Cleanup(sink.Close) + requestRecord := &requestrecord.RequestRecord{ + Timestamp: time.Now().UTC(), + InputRequest: requestrecord.Payload{ + Protocol: protocol.TypeOpenAIChat, + }, + ProviderExchanges: []requestrecord.ProviderExchange{{ + Protocol: protocol.TypeOpenAIChat, + Request: requestrecord.Payload{Protocol: protocol.TypeOpenAIChat}, + Response: &requestrecord.Payload{Protocol: protocol.TypeOpenAIChat}, + }}, + FinalResponse: &requestrecord.Payload{Protocol: protocol.TypeOpenAIChat}, + } + + sink.EmitRequestRecord(requestRecord) + requestRecord.InputRequest.Body = []byte(`{"mutated":true}`) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := sink.ForceFlush(ctx); err != nil { + t.Fatalf("ForceFlush: %v", err) + } + got := exporter.batches[0][0].RequestRecord + if string(got.InputRequest.Body) == string(requestRecord.InputRequest.Body) { + t.Fatal("exported record retained caller-owned payload storage") + } + if (got.ProviderExchanges[0].Response != nil) != testCase.wantProviderResponse { + t.Fatalf("provider response present = %v, want %v", got.ProviderExchanges[0].Response != nil, testCase.wantProviderResponse) + } + if (got.FinalResponse != nil) != testCase.wantFinalResponse { + t.Fatalf("final response present = %v, want %v", got.FinalResponse != nil, testCase.wantFinalResponse) + } + if requestRecord.ProviderExchanges[0].Response == nil || requestRecord.FinalResponse == nil { + t.Fatal("mode projection mutated the completed RequestRecord") + } + }) + } +} diff --git a/internal/obs/sink.go b/internal/obs/sink.go index 5af6a41b5..89c3896f9 100644 --- a/internal/obs/sink.go +++ b/internal/obs/sink.go @@ -7,6 +7,7 @@ import ( "github.com/google/uuid" "github.com/sirupsen/logrus" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" ) // RecordMode defines which fields are captured by the Sink. @@ -145,6 +146,43 @@ func (s *Sink) Emit(r *Record) { s.processor.Emit(r) } +// EmitRequestRecord writes the new request-boundary recording model through +// the existing asynchronous exporter pipeline. The legacy Record fields stay +// available during the additive Stage canary; readers can distinguish the new +// shape by the request_record field. +func (s *Sink) EmitRequestRecord(requestRecord *requestrecord.RequestRecord) { + if s == nil || requestRecord == nil { + return + } + + persisted := requestRecord + switch s.mode { + case RecordModeRequestOnly: + persisted = requestRecord.Project(false, false) + case RecordModeRequestResponse: + persisted = requestRecord.Project(false, true) + case RecordModeStagedRequestResponse: + persisted = requestRecord.Project(true, true) + default: + return + } + record := &Record{ + Timestamp: persisted.Timestamp, + RequestID: persisted.RequestID, + SessionID: persisted.SessionID, + Scenario: persisted.Scenario, + Duration: persisted.Duration, + Err: persisted.Error, + RequestRecord: persisted, + } + if exchanges := persisted.ProviderExchanges; len(exchanges) > 0 { + last := exchanges[len(exchanges)-1] + record.Provider = last.Provider + record.Model = last.Model + } + s.Emit(record) +} + // RecordWithScenario builds a single-stage Record (original request + final // response) and emits it. Used by client-side roundtrippers that don't go // through the transform pipeline. Server-side code should construct a *Record diff --git a/internal/obs/slim.go b/internal/obs/slim.go index b19a24e71..671c48f9e 100644 --- a/internal/obs/slim.go +++ b/internal/obs/slim.go @@ -3,6 +3,8 @@ package obs import ( "encoding/json" "time" + + requestrecord "github.com/tingly-dev/tingly-box/internal/record" ) const inlineThreshold = 256 // bytes; values smaller than this stay inline @@ -21,10 +23,11 @@ type SlimRecord struct { Scenario string `json:"scenario,omitempty"` Model string `json:"model,omitempty"` - OriginalRequest *SlimHTTPData `json:"original_request,omitempty"` - TransformedRequest *SlimHTTPData `json:"transformed_request,omitempty"` - ProviderResponse *SlimHTTPData `json:"provider_response,omitempty"` - FinalResponse *SlimHTTPData `json:"final_response,omitempty"` + OriginalRequest *SlimHTTPData `json:"original_request,omitempty"` + TransformedRequest *SlimHTTPData `json:"transformed_request,omitempty"` + ProviderResponse *SlimHTTPData `json:"provider_response,omitempty"` + FinalResponse *SlimHTTPData `json:"final_response,omitempty"` + RequestRecord *requestrecord.RequestRecord `json:"request_record,omitempty"` DurationMs int64 `json:"duration_ms"` Error string `json:"error,omitempty"` @@ -35,12 +38,12 @@ type SlimRecord struct { // SlimHTTPData mirrors RecordRequest / RecordResponse with a body that may // contain {"$ref":"sha256:"} markers instead of large inline values. type SlimHTTPData struct { - Method string `json:"method,omitempty"` - URL string `json:"url,omitempty"` - Headers map[string]string `json:"headers,omitempty"` - StatusCode int `json:"status_code,omitempty"` - Body interface{} `json:"body,omitempty"` - IsStreaming bool `json:"is_streaming,omitempty"` + Method string `json:"method,omitempty"` + URL string `json:"url,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + StatusCode int `json:"status_code,omitempty"` + Body interface{} `json:"body,omitempty"` + IsStreaming bool `json:"is_streaming,omitempty"` } // SlimifyRecord converts a Record to a SlimRecord by replacing large JSON @@ -64,17 +67,18 @@ func FullRecord(r *Record) *SlimRecord { func recordToSlim(r *Record, knownBlobs map[string]struct{}, threshold int) (*SlimRecord, map[string][]byte) { newBlobs := make(map[string][]byte) slim := &SlimRecord{ - V: 3, - Timestamp: r.Timestamp.UTC().Format(time.RFC3339), - RequestID: r.RequestID, - SessionID: r.SessionID, - SessionSrc: r.SessionSrc, - Provider: r.Provider, - Scenario: r.Scenario, - Model: r.Model, - DurationMs: r.Duration.Milliseconds(), - Error: r.Err, - Steps: r.Steps, + V: 3, + Timestamp: r.Timestamp.UTC().Format(time.RFC3339), + RequestID: r.RequestID, + SessionID: r.SessionID, + SessionSrc: r.SessionSrc, + Provider: r.Provider, + Scenario: r.Scenario, + Model: r.Model, + DurationMs: r.Duration.Milliseconds(), + Error: r.Err, + Steps: r.Steps, + RequestRecord: r.RequestRecord, } if r.OriginalRequest != nil { diff --git a/internal/protocol/assembler/openai_responses_assembler.go b/internal/protocol/assembler/openai_responses_assembler.go index 4b198d658..0f4e834ff 100644 --- a/internal/protocol/assembler/openai_responses_assembler.go +++ b/internal/protocol/assembler/openai_responses_assembler.go @@ -205,6 +205,9 @@ func (a *ResponsesAssembler) Accumulate(event responses.ResponseStreamEventUnion case "response.failed": a.status = "failed" a.finished = true + if responseHasPayload(&event.Response) { + a.response = &event.Response + } return true case "response.incomplete": @@ -374,7 +377,15 @@ func responseHasPayload(resp *responses.Response) bool { if resp == nil { return false } - return resp.ID != "" || len(resp.Output) > 0 || resp.Usage.InputTokens != 0 || resp.Usage.OutputTokens != 0 || resp.Usage.TotalTokens != 0 + return resp.ID != "" || + resp.Status != "" || + resp.Model != "" || + resp.Error.Code != "" || + resp.Error.Message != "" || + len(resp.Output) > 0 || + resp.Usage.InputTokens != 0 || + resp.Usage.OutputTokens != 0 || + resp.Usage.TotalTokens != 0 } func (a *ResponsesAssembler) ensureResponseHasAccumulatedOutput(resp *responses.Response) { diff --git a/internal/protocol/assembler/stream_assembler.go b/internal/protocol/assembler/stream_assembler.go new file mode 100644 index 000000000..d2b9dc041 --- /dev/null +++ b/internal/protocol/assembler/stream_assembler.go @@ -0,0 +1,219 @@ +package assembler + +import ( + "encoding/json" + "errors" + "fmt" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/responses" + "github.com/tingly-dev/tingly-box/internal/protocol" +) + +// StreamAssembler is the protocol-owned common surface for reconstructing one +// complete response from native stream events. It accepts SDK events, Wire +// DTOs, or json.RawMessage values; protocol-specific handling remains here +// rather than in observers such as Recording. +type StreamAssembler interface { + Add(value any) error + Finish() (any, error) + Terminal() bool + TerminalError() error +} + +// NewStreamAssembler adapts the existing protocol assemblers to one common +// interface. Protocol conversion is intentionally out of scope: the caller +// must select the protocol already spoken at its observation boundary. +func NewStreamAssembler(api protocol.APIType) (StreamAssembler, error) { + return newStreamAssembler(api, 1) +} + +// NewStreamAssemblerForRequest configures protocol-specific expectations from +// the provider-bound request. Today this is used for OpenAI Chat's n choices. +func NewStreamAssemblerForRequest(api protocol.APIType, request any) (StreamAssembler, error) { + return newStreamAssembler(api, OpenAIChatChoiceCount(request)) +} + +func newStreamAssembler(api protocol.APIType, expectedChatChoices int) (StreamAssembler, error) { + switch api { + case protocol.TypeAnthropicV1: + return &anthropicV1StreamAssembler{inner: NewAnthropicSDKAssembler()}, nil + case protocol.TypeAnthropicBeta: + return &anthropicBetaStreamAssembler{inner: NewAnthropicBetaSDKAssembler()}, nil + case protocol.TypeOpenAIChat: + return &openAIChatStreamAssembler{ + inner: NewOpenAIStreamAssembler(), + expectedChoices: expectedChatChoices, + }, nil + case protocol.TypeOpenAIResponses: + return &openAIResponsesStreamAssembler{inner: NewResponsesAssembler()}, nil + default: + return nil, fmt.Errorf("stream assembler: unsupported protocol %q", api) + } +} + +// OpenAIChatChoiceCount returns the request's expected number of choices, +// defaulting invalid, omitted, or non-Chat values to one. +func OpenAIChatChoiceCount(request any) int { + chatRequest, ok := request.(*openai.ChatCompletionNewParams) + if !ok || chatRequest == nil { + return 1 + } + choices := chatRequest.N.Or(1) + if choices < 1 { + return 1 + } + maxInt := int(^uint(0) >> 1) + if uint64(choices) > uint64(maxInt) { + return maxInt + } + return int(choices) +} + +type anthropicV1StreamAssembler struct { + inner *AnthropicSDKAssembler + started bool + terminal bool +} + +func (a *anthropicV1StreamAssembler) Add(value any) error { + raw, err := streamEventJSON(value) + if err != nil { + return err + } + var event anthropic.MessageStreamEventUnion + if err := json.Unmarshal(raw, &event); err != nil { + return fmt.Errorf("decode Anthropic V1 stream event: %w", err) + } + if err := a.inner.Accumulate(event); err != nil { + return err + } + if event.Type == "message_start" { + a.started = true + } + if event.Type == "message_stop" && a.started { + a.terminal = true + } + return nil +} + +func (a *anthropicV1StreamAssembler) Terminal() bool { return a.terminal } +func (*anthropicV1StreamAssembler) TerminalError() error { return nil } + +func (a *anthropicV1StreamAssembler) Finish() (any, error) { + return a.inner.Finish(), nil +} + +type anthropicBetaStreamAssembler struct { + inner *AnthropicBetaSDKAssembler + started bool + terminal bool +} + +func (a *anthropicBetaStreamAssembler) Add(value any) error { + raw, err := streamEventJSON(value) + if err != nil { + return err + } + var event anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(raw, &event); err != nil { + return fmt.Errorf("decode Anthropic Beta stream event: %w", err) + } + if err := a.inner.Accumulate(event); err != nil { + return err + } + if event.Type == "message_start" { + a.started = true + } + if event.Type == "message_stop" && a.started { + a.terminal = true + } + return nil +} + +func (a *anthropicBetaStreamAssembler) Terminal() bool { return a.terminal } +func (*anthropicBetaStreamAssembler) TerminalError() error { return nil } + +func (a *anthropicBetaStreamAssembler) Finish() (any, error) { + return a.inner.Finish(), nil +} + +type openAIChatStreamAssembler struct { + inner *OpenAIChatStreamAssembler + expectedChoices int + finishedChoices map[int64]struct{} +} + +func (a *openAIChatStreamAssembler) Add(value any) error { + raw, err := streamEventJSON(value) + if err != nil { + return err + } + var event openai.ChatCompletionChunk + if err := json.Unmarshal(raw, &event); err != nil { + return fmt.Errorf("decode OpenAI Chat stream event: %w", err) + } + if !a.inner.AddChunk(event) { + return fmt.Errorf("accumulate OpenAI Chat stream chunk %q", event.ID) + } + for _, choice := range event.Choices { + if choice.FinishReason != "" && choice.Index >= 0 && choice.Index < int64(a.expectedChoices) { + if a.finishedChoices == nil { + a.finishedChoices = make(map[int64]struct{}) + } + a.finishedChoices[choice.Index] = struct{}{} + } + } + return nil +} + +func (a *openAIChatStreamAssembler) Terminal() bool { + return len(a.finishedChoices) >= a.expectedChoices +} +func (*openAIChatStreamAssembler) TerminalError() error { return nil } + +func (a *openAIChatStreamAssembler) Finish() (any, error) { + return a.inner.Finish(), nil +} + +type openAIResponsesStreamAssembler struct { + inner *ResponsesAssembler + terminalErr error +} + +func (a *openAIResponsesStreamAssembler) Add(value any) error { + raw, err := streamEventJSON(value) + if err != nil { + return err + } + var event responses.ResponseStreamEventUnion + if err := json.Unmarshal(raw, &event); err != nil { + return fmt.Errorf("decode OpenAI Responses stream event: %w", err) + } + a.inner.Accumulate(event) + switch event.Type { + case "response.failed", "error": + a.terminalErr = fmt.Errorf("OpenAI Responses stream ended with %s", event.Type) + } + return nil +} + +func (a *openAIResponsesStreamAssembler) Finish() (any, error) { + return a.inner.Finish(), nil +} + +func (a *openAIResponsesStreamAssembler) Terminal() bool { return a.inner.IsFinished() } +func (a *openAIResponsesStreamAssembler) TerminalError() error { return a.terminalErr } + +func streamEventJSON(value any) ([]byte, error) { + if value == nil { + return nil, errors.New("stream event is nil") + } + + raw, err := protocol.SnapshotJSON(value) + if err != nil { + return nil, fmt.Errorf("snapshot stream event %T: %w", value, err) + } + return raw, nil +} diff --git a/internal/protocol/assembler/stream_assembler_test.go b/internal/protocol/assembler/stream_assembler_test.go new file mode 100644 index 000000000..6f33845fc --- /dev/null +++ b/internal/protocol/assembler/stream_assembler_test.go @@ -0,0 +1,260 @@ +package assembler + +import ( + "encoding/json" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/packages/param" + "github.com/openai/openai-go/v3/responses" + "github.com/stretchr/testify/require" + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestNewStreamAssemblerRejectsUnsupportedProtocol(t *testing.T) { + assembler, err := NewStreamAssembler("") + require.Nil(t, assembler) + require.ErrorContains(t, err, "unsupported protocol") + + assembler, err = NewStreamAssembler(protocol.TypeGoogle) + require.Nil(t, assembler) + require.ErrorContains(t, err, "google") +} + +func TestCommonStreamAssemblerAnthropicV1(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeAnthropicV1) + require.NoError(t, err) + + require.NoError(t, assembler.Add(anthropic.MessageStreamEventUnion{ + Type: "message_start", + Message: anthropic.Message{ + ID: "msg-v1", + Type: "message", + Role: "assistant", + Model: "claude-v1", + }, + })) + + result, err := assembler.Finish() + require.NoError(t, err) + message, ok := result.(*anthropic.Message) + require.True(t, ok) + require.Equal(t, "msg-v1", string(message.ID)) +} + +func TestCommonStreamAssemblerAnthropicBetaFromJSON(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeAnthropicBeta) + require.NoError(t, err) + + event := json.RawMessage(`{ + "type":"message_start", + "message":{ + "id":"msg-beta", + "type":"message", + "role":"assistant", + "content":[], + "model":"claude-beta", + "stop_reason":null, + "stop_sequence":null, + "usage":{"input_tokens":1,"output_tokens":0} + } + }`) + require.NoError(t, assembler.Add(event)) + + result, err := assembler.Finish() + require.NoError(t, err) + message, ok := result.(*anthropic.BetaMessage) + require.True(t, ok) + require.Equal(t, "msg-beta", string(message.ID)) +} + +func TestCommonStreamAssemblerAnthropicBetaFromConvertedEvent(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeAnthropicBeta) + require.NoError(t, err) + + event := protocolstream.AnthropicEvent{ + Type: "message_start", + Data: map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": "msg-converted-beta", "type": "message", "role": "assistant", + "content": []any{}, "model": "claude-beta", "stop_reason": nil, "stop_sequence": nil, + "usage": map[string]any{"input_tokens": 1, "output_tokens": 0}, + }, + }, + } + require.NoError(t, assembler.Add(event)) + + result, err := assembler.Finish() + require.NoError(t, err) + message, ok := result.(*anthropic.BetaMessage) + require.True(t, ok) + require.Equal(t, "msg-converted-beta", string(message.ID)) +} + +func TestCommonStreamAssemblerOpenAIChatAcceptsWireDTO(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeOpenAIChat) + require.NoError(t, err) + + require.NoError(t, assembler.Add(wire.ChatStreamChunk{ + ID: "chat-1", + Object: "chat.completion.chunk", + Created: 1, + Model: "public-model", + Choices: []wire.ChatStreamChoice{{ + Index: 0, + Delta: wire.ChatStreamDelta{Role: "assistant", Content: "hello "}, + }}, + })) + stop := "stop" + require.NoError(t, assembler.Add(wire.ChatStreamChunk{ + ID: "chat-1", + Object: "chat.completion.chunk", + Model: "public-model", + Choices: []wire.ChatStreamChoice{{ + Index: 0, + Delta: wire.ChatStreamDelta{Content: "world"}, + FinishReason: &stop, + }}, + })) + + result, err := assembler.Finish() + require.NoError(t, err) + completion, ok := result.(*openai.ChatCompletion) + require.True(t, ok) + require.Equal(t, "chat-1", completion.ID) + require.Equal(t, "public-model", completion.Model) + require.Equal(t, "hello world", completion.Choices[0].Message.Content) + require.True(t, assembler.Terminal()) +} + +func TestCommonStreamAssemblerOpenAIChatRejectsMismatchedIDs(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeOpenAIChat) + require.NoError(t, err) + require.NoError(t, assembler.Add(wire.ChatStreamChunk{ID: "chat-a", Object: "chat.completion.chunk"})) + require.Error(t, assembler.Add(wire.ChatStreamChunk{ID: "chat-b", Object: "chat.completion.chunk"})) +} + +func TestCommonStreamAssemblerOpenAIChatWaitsForEveryChoice(t *testing.T) { + request := &openai.ChatCompletionNewParams{N: param.NewOpt(int64(2))} + assembler, err := NewStreamAssemblerForRequest(protocol.TypeOpenAIChat, request) + require.NoError(t, err) + stop := "stop" + require.NoError(t, assembler.Add(wire.ChatStreamChunk{ + ID: "chat-multi", Object: "chat.completion.chunk", + Choices: []wire.ChatStreamChoice{{ + Index: 0, FinishReason: &stop, + }}, + })) + require.False(t, assembler.Terminal()) + require.NoError(t, assembler.Add(wire.ChatStreamChunk{ + ID: "chat-multi", Object: "chat.completion.chunk", + Choices: []wire.ChatStreamChoice{{Index: 2, FinishReason: &stop}, {Index: 3, FinishReason: &stop}}, + })) + require.False(t, assembler.Terminal(), "out-of-range choice indexes must not complete the stream") + require.NoError(t, assembler.Add(wire.ChatStreamChunk{ + ID: "chat-multi", Object: "chat.completion.chunk", + Choices: []wire.ChatStreamChoice{{ + Index: 1, FinishReason: &stop, + }}, + })) + require.True(t, assembler.Terminal()) +} + +func TestCommonStreamAssemblerOpenAIChatDoesNotPreallocateUntrustedN(t *testing.T) { + request := &openai.ChatCompletionNewParams{N: param.NewOpt(int64(^uint64(0) >> 1))} + assembler, err := NewStreamAssemblerForRequest(protocol.TypeOpenAIChat, request) + require.NoError(t, err) + stop := "stop" + require.NoError(t, assembler.Add(wire.ChatStreamChunk{ + ID: "chat-large-n", Object: "chat.completion.chunk", + Choices: []wire.ChatStreamChoice{{Index: 0, FinishReason: &stop}}, + })) + require.False(t, assembler.Terminal()) +} + +func TestCommonStreamAssemblerAnthropicRequiresMessageStart(t *testing.T) { + for _, api := range []protocol.APIType{protocol.TypeAnthropicV1, protocol.TypeAnthropicBeta} { + t.Run(string(api), func(t *testing.T) { + assembler, err := NewStreamAssembler(api) + require.NoError(t, err) + require.NoError(t, assembler.Add(json.RawMessage(`{"type":"message_stop"}`))) + require.False(t, assembler.Terminal()) + }) + } +} + +func TestCommonStreamAssemblerOpenAIResponsesAcceptsWireDTO(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeOpenAIResponses) + require.NoError(t, err) + + require.NoError(t, assembler.Add(wire.ResponsesCreatedEvent{ + Type: "response.created", + Response: wire.ResponsesWireResponse{ + ID: "resp-1", + Object: "response", + Status: "in_progress", + Model: "public-model", + }, + })) + require.NoError(t, assembler.Add(wire.ResponsesOutputTextDeltaEvent{ + Type: "response.output_text.delta", + OutputIndex: 0, + ContentIndex: 0, + Delta: "hello responses", + })) + require.NoError(t, assembler.Add(wire.ResponsesCompletedEvent{ + Type: "response.completed", + Response: wire.ResponsesWireResponse{ + ID: "resp-1", + Object: "response", + Status: "completed", + Model: "public-model", + Output: []wire.ResponsesOutputItemWire{}, + }, + })) + + result, err := assembler.Finish() + require.NoError(t, err) + response, ok := result.(*responses.Response) + require.True(t, ok) + require.Equal(t, "resp-1", response.ID) + require.Equal(t, responses.ResponsesModel("public-model"), response.Model) + require.Equal(t, "hello responses", response.OutputText()) +} + +func TestCommonStreamAssemblerRejectsInvalidEvent(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeOpenAIChat) + require.NoError(t, err) + require.ErrorContains(t, assembler.Add([]byte(`not-json`)), "not valid JSON") + require.ErrorContains(t, assembler.Add(nil), "nil") +} + +func TestCommonStreamAssemblerRejectsTypedNilEvent(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeOpenAIChat) + require.NoError(t, err) + var event *openai.ChatCompletionChunk + require.NotPanics(t, func() { + require.ErrorContains(t, assembler.Add(event), "nil") + }) +} + +func TestCommonStreamAssemblerPreservesResponsesFailedPayload(t *testing.T) { + assembler, err := NewStreamAssembler(protocol.TypeOpenAIResponses) + require.NoError(t, err) + require.NoError(t, assembler.Add(json.RawMessage( + `{"type":"response.failed","sequence_number":1,"response":{"id":"resp-failed","object":"response","status":"failed","model":"provider-model","output":[],"error":{"code":"server_error","message":"provider failed"}}}`, + ))) + require.True(t, assembler.Terminal()) + require.ErrorContains(t, assembler.TerminalError(), "response.failed") + + result, err := assembler.Finish() + require.NoError(t, err) + response, ok := result.(*responses.Response) + require.True(t, ok) + require.Equal(t, "resp-failed", response.ID) + require.Equal(t, responses.ResponseStatusFailed, response.Status) +} diff --git a/internal/protocol/json_snapshot.go b/internal/protocol/json_snapshot.go new file mode 100644 index 000000000..b1a104581 --- /dev/null +++ b/internal/protocol/json_snapshot.go @@ -0,0 +1,50 @@ +package protocol + +import ( + "encoding/json" + "fmt" + "reflect" +) + +// SnapshotJSON returns an owned JSON snapshot of one protocol value. SDK +// values may expose RawJSON to preserve fields unknown to their typed DTOs; +// typed nils are rejected before calling that method. +func SnapshotJSON(value any) ([]byte, error) { + if isNilJSONValue(value) { + return nil, fmt.Errorf("protocol JSON value %T is nil", value) + } + + var raw []byte + switch typed := value.(type) { + case json.RawMessage: + raw = typed + case []byte: + raw = typed + case interface{ RawJSON() string }: + raw = []byte(typed.RawJSON()) + } + if len(raw) == 0 { + var err error + raw, err = json.Marshal(value) + if err != nil { + return nil, err + } + } + if !json.Valid(raw) { + return nil, fmt.Errorf("protocol JSON value %T is not valid JSON", value) + } + return append([]byte(nil), raw...), nil +} + +func isNilJSONValue(value any) bool { + if value == nil { + return true + } + v := reflect.ValueOf(value) + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return v.IsNil() + default: + return false + } +} diff --git a/internal/protocol/nonstream/anthropic.go b/internal/protocol/nonstream/anthropic.go index 94671342b..103c378a3 100644 --- a/internal/protocol/nonstream/anthropic.go +++ b/internal/protocol/nonstream/anthropic.go @@ -13,9 +13,9 @@ import ( "github.com/tingly-dev/tingly-box/internal/protocol/wire" ) -// HandleAnthropicBetaToOpenAIResponse converts an Anthropic BetaMessage to the -// OpenAI Chat Completions wire format. -func HandleAnthropicBetaToOpenAIResponse(bm *anthropic.BetaMessage, responseModel string) wire.ChatCompletionWire { +// ConvertAnthropicBetaToOpenAIChat converts an Anthropic Beta message to the +// transport-neutral OpenAI Chat Completions wire value. +func ConvertAnthropicBetaToOpenAIChat(bm *anthropic.BetaMessage, responseModel string) wire.ChatCompletionWire { var toolCalls []wire.ChatCompletionToolCallWire var textContent string var thinking string @@ -88,6 +88,12 @@ func HandleAnthropicBetaToOpenAIResponse(bm *anthropic.BetaMessage, responseMode } } +// HandleAnthropicBetaToOpenAIResponse preserves the legacy conversion entry +// point while delegating to the transport-neutral value converter. +func HandleAnthropicBetaToOpenAIResponse(bm *anthropic.BetaMessage, responseModel string) wire.ChatCompletionWire { + return ConvertAnthropicBetaToOpenAIChat(bm, responseModel) +} + // HandleAnthropicV1 handles Anthropic v1 non-streaming response. // Returns (UsageStat, error) func HandleAnthropicV1(hc *protocol.HandleContext, m *anthropic.Message) (*protocol.TokenUsage, error) { diff --git a/internal/protocol/nonstream/nonstream_test.go b/internal/protocol/nonstream/nonstream_test.go index eafa0c1c6..3883e96f5 100644 --- a/internal/protocol/nonstream/nonstream_test.go +++ b/internal/protocol/nonstream/nonstream_test.go @@ -111,6 +111,29 @@ func TestHandleAnthropicToOpenAI(t *testing.T) { } } +func TestConvertAnthropicBetaToOpenAIChatMatchesLegacyEntry(t *testing.T) { + message := &anthropic.BetaMessage{ + ID: "msg_transport_neutral", + Role: "assistant", + Content: []anthropic.BetaContentBlockUnion{ + {Type: "text", Text: "parallel path"}, + }, + StopReason: "end_turn", + Usage: anthropic.BetaUsage{ + InputTokens: 4, + OutputTokens: 2, + }, + } + + pure := ConvertAnthropicBetaToOpenAIChat(message, "client-visible-model") + legacy := HandleAnthropicBetaToOpenAIResponse(message, "client-visible-model") + // Created is intentionally generated at conversion time; normalize it when + // comparing two sequential calls to the same pure implementation. + pure.Created = 0 + legacy.Created = 0 + assert.Equal(t, pure, legacy) +} + func TestHandleOpenAIToAnthropic(t *testing.T) { tests := []struct { name string diff --git a/internal/protocol/nonstream/openai_to_anthropic.go b/internal/protocol/nonstream/openai_to_anthropic.go index 4ae05ded7..2c97014ea 100644 --- a/internal/protocol/nonstream/openai_to_anthropic.go +++ b/internal/protocol/nonstream/openai_to_anthropic.go @@ -3,7 +3,6 @@ package nonstream import ( "encoding/json" "fmt" - "time" "github.com/anthropics/anthropic-sdk-go" "github.com/google/uuid" @@ -29,8 +28,37 @@ func anthropicUsageWire(u *protocol.TokenUsage) wire.AnthropicUsageWire { } func HandleOpenAIChatToAnthropic(chat *openai.ChatCompletion, model string) *anthropic.BetaMessage { - wire := wire.AnthropicMsgWire{ - ID: fmt.Sprintf("msg_%d", time.Now().Unix()), + value, err := marshalOpenAIChatToAnthropic(chat, model) + if err != nil { + return &anthropic.BetaMessage{} + } + var msg anthropic.BetaMessage + if err := json.Unmarshal(value, &msg); err != nil { + return &anthropic.BetaMessage{} + } + return &msg +} + +// ConvertOpenAIChatToAnthropicV1 converts an OpenAI Chat completion to a typed +// Anthropic v1 response without writing it to an HTTP transport. +func ConvertOpenAIChatToAnthropicV1(chat *openai.ChatCompletion, model string) (*anthropic.Message, error) { + value, err := marshalOpenAIChatToAnthropic(chat, model) + if err != nil { + return nil, err + } + var msg anthropic.Message + if err := json.Unmarshal(value, &msg); err != nil { + return nil, fmt.Errorf("decode Anthropic v1 response: %w", err) + } + return &msg, nil +} + +func marshalOpenAIChatToAnthropic(chat *openai.ChatCompletion, model string) ([]byte, error) { + if chat == nil { + return nil, fmt.Errorf("convert OpenAI Chat response to Anthropic: response is nil") + } + result := wire.AnthropicMsgWire{ + ID: "msg_" + uuid.NewString(), Type: "message", Role: "assistant", Content: []interface{}{}, @@ -43,7 +71,7 @@ func HandleOpenAIChatToAnthropic(chat *openai.ChatCompletion, model string) *ant // Preserve server_tool_use from ExtraFields if present if chat.JSON.ExtraFields != nil { if serverToolUse, exists := chat.JSON.ExtraFields["server_tool_use"]; exists && serverToolUse.Valid() { - wire.ServerToolUse = json.RawMessage(serverToolUse.Raw()) + result.ServerToolUse = json.RawMessage(serverToolUse.Raw()) } } @@ -68,22 +96,38 @@ func HandleOpenAIChatToAnthropic(chat *openai.ChatCompletion, model string) *ant contentBlocks = append(contentBlocks, anthropic.NewToolUseBlock(toolCall.ID, input, toolCall.Function.Name)) } if choice.FinishReason == "tool_calls" { - wire.StopReason = "tool_use" + result.StopReason = "tool_use" + } else if choice.FinishReason == "length" { + result.StopReason = "max_tokens" } break } - wire.Content = contentBlocks + result.Content = contentBlocks - jsonBytes, _ := json.Marshal(wire) - var msg anthropic.BetaMessage - json.Unmarshal(jsonBytes, &msg) - return &msg + jsonBytes, err := json.Marshal(result) + if err != nil { + return nil, fmt.Errorf("encode Anthropic response: %w", err) + } + return jsonBytes, nil } // HandleOpenAIChatToAnthropicBeta converts OpenAI response to Anthropic beta format func HandleOpenAIChatToAnthropicBeta(chat *openai.ChatCompletion, model string) anthropic.BetaMessage { - wire := wire.AnthropicMsgWire{ - ID: fmt.Sprintf("msg_%d", time.Now().Unix()), + msg, err := ConvertOpenAIChatToAnthropicBeta(chat, model) + if err != nil { + return anthropic.BetaMessage{} + } + return *msg +} + +// ConvertOpenAIChatToAnthropicBeta converts an OpenAI Chat completion to a +// typed Anthropic beta response without writing it to an HTTP transport. +func ConvertOpenAIChatToAnthropicBeta(chat *openai.ChatCompletion, model string) (*anthropic.BetaMessage, error) { + if chat == nil { + return nil, fmt.Errorf("convert OpenAI Chat response to Anthropic beta: response is nil") + } + result := wire.AnthropicMsgWire{ + ID: "msg_" + uuid.NewString(), Type: "message", Role: "assistant", Content: []interface{}{}, @@ -95,7 +139,7 @@ func HandleOpenAIChatToAnthropicBeta(chat *openai.ChatCompletion, model string) if chat.JSON.ExtraFields != nil { if serverToolUse, exists := chat.JSON.ExtraFields["server_tool_use"]; exists && serverToolUse.Valid() { - wire.ServerToolUse = json.RawMessage(serverToolUse.Raw()) + result.ServerToolUse = json.RawMessage(serverToolUse.Raw()) } } @@ -120,16 +164,23 @@ func HandleOpenAIChatToAnthropicBeta(chat *openai.ChatCompletion, model string) contentBlocks = append(contentBlocks, anthropic.NewBetaToolUseBlock(toolCall.ID, input, toolCall.Function.Name)) } if choice.FinishReason == "tool_calls" { - wire.StopReason = string(anthropic.BetaStopReasonToolUse) + result.StopReason = string(anthropic.BetaStopReasonToolUse) + } else if choice.FinishReason == "length" { + result.StopReason = string(anthropic.BetaStopReasonMaxTokens) } break } - wire.Content = contentBlocks + result.Content = contentBlocks - jsonBytes, _ := json.Marshal(wire) + jsonBytes, err := json.Marshal(result) + if err != nil { + return nil, fmt.Errorf("encode Anthropic beta response: %w", err) + } var msg anthropic.BetaMessage - json.Unmarshal(jsonBytes, &msg) - return msg + if err := json.Unmarshal(jsonBytes, &msg); err != nil { + return nil, fmt.Errorf("decode Anthropic beta response: %w", err) + } + return &msg, nil } // HandleResponsesToAnthropicBeta converts OpenAI Responses API response to Anthropic beta format @@ -159,7 +210,7 @@ func HandleResponsesToAnthropicBeta(rs *responses.Response, model string) anthro if err := json.Unmarshal([]byte(argsStr), &arguments); err != nil { arguments = make(map[string]interface{}) } - contentBlocks = append(contentBlocks, anthropic.NewBetaToolUseBlock(output.ID, arguments, output.Name)) + contentBlocks = append(contentBlocks, anthropic.NewBetaToolUseBlock(responsesToolCallID(output), arguments, output.Name)) wire.StopReason = string(anthropic.BetaStopReasonToolUse) } } @@ -171,6 +222,13 @@ func HandleResponsesToAnthropicBeta(rs *responses.Response, model string) anthro } } } + if rs.Status == "incomplete" { + if rs.IncompleteDetails.Reason == "content_filter" { + wire.StopReason = string(anthropic.BetaStopReasonRefusal) + } else { + wire.StopReason = string(anthropic.BetaStopReasonMaxTokens) + } + } for _, output := range rs.Output { for _, content := range output.Content { @@ -216,7 +274,7 @@ func HandleResponsesToAnthropicV1(rs *responses.Response, model string) anthropi if err := json.Unmarshal([]byte(argsStr), &arguments); err != nil { arguments = make(map[string]interface{}) } - contentBlocks = append(contentBlocks, anthropic.NewToolUseBlock(output.ID, arguments, output.Name)) + contentBlocks = append(contentBlocks, anthropic.NewToolUseBlock(responsesToolCallID(output), arguments, output.Name)) wire.StopReason = "tool_use" } } @@ -228,6 +286,13 @@ func HandleResponsesToAnthropicV1(rs *responses.Response, model string) anthropi } } } + if rs.Status == "incomplete" { + if rs.IncompleteDetails.Reason == "content_filter" { + wire.StopReason = "refusal" + } else { + wire.StopReason = "max_tokens" + } + } for _, output := range rs.Output { for _, content := range output.Content { @@ -246,6 +311,13 @@ func HandleResponsesToAnthropicV1(rs *responses.Response, model string) anthropi return msg } +func responsesToolCallID(output responses.ResponseOutputItemUnion) string { + if output.CallID != "" { + return output.CallID + } + return output.ID +} + // resolveResponsesArguments extracts the arguments string from a Responses API output item. func resolveResponsesArguments(rs *responses.Response, output responses.ResponseOutputItemUnion) string { if output.Arguments.OfString != "" { diff --git a/internal/protocol/nonstream/openai_to_anthropic_semantics_test.go b/internal/protocol/nonstream/openai_to_anthropic_semantics_test.go new file mode 100644 index 000000000..8fdb6472c --- /dev/null +++ b/internal/protocol/nonstream/openai_to_anthropic_semantics_test.go @@ -0,0 +1,73 @@ +package nonstream + +import ( + "encoding/json" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/responses" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResponsesToAnthropicUsesCallID(t *testing.T) { + var response responses.Response + require.NoError(t, json.Unmarshal([]byte(`{ + "id":"resp_tool","object":"response","status":"completed","model":"provider-model", + "output":[{"id":"fc_item","call_id":"call_weather","type":"function_call","name":"get_weather","arguments":"{\"city\":\"Paris\"}","status":"completed"}], + "usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2,"input_tokens_details":{"cached_tokens":0},"output_tokens_details":{"reasoning_tokens":0}} + }`), &response)) + + beta := HandleResponsesToAnthropicBeta(&response, "public-model") + require.Len(t, beta.Content, 1) + assert.Equal(t, "tool_use", beta.Content[0].Type) + assert.Equal(t, "call_weather", beta.Content[0].ID) + + v1 := HandleResponsesToAnthropicV1(&response, "public-model") + require.Len(t, v1.Content, 1) + assert.Equal(t, "tool_use", v1.Content[0].Type) + assert.Equal(t, "call_weather", v1.Content[0].ID) +} + +func TestCompleteConversionsPreserveTruncation(t *testing.T) { + chat := &openai.ChatCompletion{ + ID: "chatcmpl-length", + Choices: []openai.ChatCompletionChoice{{ + FinishReason: "length", + Message: openai.ChatCompletionMessage{Role: "assistant", Content: "partial"}, + }}, + } + v1, err := ConvertOpenAIChatToAnthropicV1(chat, "public-model") + require.NoError(t, err) + assert.Equal(t, "max_tokens", string(v1.StopReason)) + beta, err := ConvertOpenAIChatToAnthropicBeta(chat, "public-model") + require.NoError(t, err) + assert.Equal(t, "max_tokens", string(beta.StopReason)) + + var response responses.Response + require.NoError(t, json.Unmarshal([]byte(`{ + "id":"resp_incomplete","object":"response","status":"incomplete","model":"provider-model", + "incomplete_details":{"reason":"max_output_tokens"}, + "output":[{"id":"msg_1","type":"message","role":"assistant","status":"incomplete","content":[{"type":"output_text","text":"partial","annotations":[]}]}], + "usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2,"input_tokens_details":{"cached_tokens":0},"output_tokens_details":{"reasoning_tokens":0}} + }`), &response)) + responsesV1 := HandleResponsesToAnthropicV1(&response, "public-model") + assert.Equal(t, "max_tokens", string(responsesV1.StopReason)) + responsesBeta := HandleResponsesToAnthropicBeta(&response, "public-model") + assert.Equal(t, "max_tokens", string(responsesBeta.StopReason)) +} + +func TestChatToAnthropicSyntheticIDsAreUnique(t *testing.T) { + chat := &openai.ChatCompletion{ + ID: "chatcmpl-source", + Choices: []openai.ChatCompletionChoice{{ + FinishReason: "stop", + Message: openai.ChatCompletionMessage{Role: "assistant", Content: "hello"}, + }}, + } + first, err := ConvertOpenAIChatToAnthropicV1(chat, "public-model") + require.NoError(t, err) + second, err := ConvertOpenAIChatToAnthropicV1(chat, "public-model") + require.NoError(t, err) + assert.NotEqual(t, first.ID, second.ID) +} diff --git a/internal/protocol/nonstream/openai_to_chat.go b/internal/protocol/nonstream/openai_to_chat.go index 1b41285f7..72d5fef7c 100644 --- a/internal/protocol/nonstream/openai_to_chat.go +++ b/internal/protocol/nonstream/openai_to_chat.go @@ -8,51 +8,82 @@ import ( "github.com/tingly-dev/tingly-box/internal/protocol" usageconv "github.com/tingly-dev/tingly-box/internal/protocol/usage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" ) type responsesToChatNonStreamState struct { content strings.Builder refusal strings.Builder - toolCalls []map[string]any + toolCalls []wire.ChatCompletionToolCallWire } // HandleResponsesToOpenAIChat writes a Responses API response as OpenAI Chat format. // Corresponds to stream.HandleResponsesToOpenAIChatStream. func HandleResponsesToOpenAIChat(hc *protocol.HandleContext, rs *responses.Response) (map[string]any, *protocol.TokenUsage, error) { + chatResp := BuildOpenAIChatPayloadFromResponses(rs, hc.ResponseModel) + hc.GinContext.JSON(http.StatusOK, chatResp) + return chatResp, usageconv.FromOpenAIResponses(rs.Usage), nil +} + +// BuildOpenAIChatPayloadFromResponses converts a complete Responses result to +// the minimal Chat Completions wire shape without writing an HTTP response. +func BuildOpenAIChatPayloadFromResponses(rs *responses.Response, responseModel string) map[string]any { + return ConvertResponsesToOpenAIChat(rs, responseModel).ToMap() +} + +// ConvertResponsesToOpenAIChat builds the typed Chat Completions wire +// contract without coupling protocol conversion to HTTP or generic maps. +func ConvertResponsesToOpenAIChat(rs *responses.Response, responseModel string) wire.ChatCompletionWire { state := buildResponsesToChatNonStreamState(rs) - message := state.message() - - choices := []map[string]any{ - { - "index": 0, - "message": message, - "finish_reason": mapResponsesFinishReason(rs, len(state.toolCalls) > 0), - }, + message := wire.ChatCompletionMessageWire{ + Role: "assistant", + Content: state.content.String(), + Refusal: state.refusal.String(), + ToolCalls: state.toolCalls, } // The canonical type owns the Chat usage shape; only total_tokens differs, // preferring the value upstream actually reported when it sent one. normalizedUsage := usageconv.FromOpenAIResponses(rs.Usage) - usage := normalizedUsage.ToOpenAIChatUsageMap() - if reported := int(rs.Usage.TotalTokens); reported != 0 { - usage["total_tokens"] = reported + totalInputTokens := normalizedUsage.PromptTotalTokens() + totalTokens := int(rs.Usage.TotalTokens) + if totalTokens == 0 { + totalTokens = totalInputTokens + normalizedUsage.OutputTokens + } + usage := wire.ChatCompletionUsageWire{ + PromptTokens: int64(totalInputTokens), + CompletionTokens: int64(normalizedUsage.OutputTokens), + TotalTokens: int64(totalTokens), + } + if normalizedUsage.CacheReadTokens > 0 || normalizedUsage.CacheWriteTokens > 0 { + usage.PromptTokensDetails = &wire.ChatCompletionPromptDetailsWire{ + CachedTokens: int64(normalizedUsage.CacheReadTokens), + CacheWriteTokens: int64(normalizedUsage.CacheWriteTokens), + } + } + if normalizedUsage.ReasoningTokens > 0 { + usage.CompletionTokensDetails = &wire.ChatCompletionOutputDetailsWire{ + ReasoningTokens: int64(normalizedUsage.ReasoningTokens), + } } - chatResp := map[string]any{ - "id": rs.ID, - "object": "chat.completion", - "created": int64(rs.CreatedAt), - "model": hc.ResponseModel, - "choices": choices, - "usage": usage, + return wire.ChatCompletionWire{ + ID: rs.ID, + Object: "chat.completion", + Created: int64(rs.CreatedAt), + Model: responseModel, + Choices: []wire.ChatCompletionChoiceWire{{ + Index: 0, + Message: message, + FinishReason: mapResponsesFinishReason(rs, len(state.toolCalls) > 0), + }}, + Usage: usage, } - hc.GinContext.JSON(http.StatusOK, chatResp) - return chatResp, usageconv.FromOpenAIResponses(rs.Usage), nil } func buildResponsesToChatNonStreamState(rs *responses.Response) *responsesToChatNonStreamState { state := &responsesToChatNonStreamState{ - toolCalls: make([]map[string]any, 0), + toolCalls: make([]wire.ChatCompletionToolCallWire, 0), } if rs == nil { return state @@ -70,12 +101,12 @@ func buildResponsesToChatNonStreamState(rs *responses.Response) *responsesToChat } } case "function_call", "custom_tool_call", "mcp_call": - state.toolCalls = append(state.toolCalls, map[string]any{ - "id": firstNonEmpty(output.CallID, output.ID), - "type": "function", - "function": map[string]any{ - "name": output.Name, - "arguments": output.Arguments.OfString, + state.toolCalls = append(state.toolCalls, wire.ChatCompletionToolCallWire{ + ID: firstNonEmpty(output.CallID, output.ID), + Type: "function", + Function: wire.ChatCompletionFunctionWire{ + Name: output.Name, + Arguments: output.Arguments.OfString, }, }) } @@ -84,27 +115,6 @@ func buildResponsesToChatNonStreamState(rs *responses.Response) *responsesToChat return state } -func (s *responsesToChatNonStreamState) message() map[string]any { - message := map[string]any{ - "role": "assistant", - } - if s == nil { - return message - } - - if content := s.content.String(); content != "" { - message["content"] = content - } - if refusal := s.refusal.String(); refusal != "" { - message["refusal"] = refusal - } - if len(s.toolCalls) > 0 { - message["tool_calls"] = s.toolCalls - } - - return message -} - func mapResponsesFinishReason(rs *responses.Response, hasToolCalls bool) string { if hasToolCalls { return "tool_calls" diff --git a/internal/protocol/nonstream/openai_to_chat_test.go b/internal/protocol/nonstream/openai_to_chat_test.go index 91e9280bb..77bfa2564 100644 --- a/internal/protocol/nonstream/openai_to_chat_test.go +++ b/internal/protocol/nonstream/openai_to_chat_test.go @@ -117,3 +117,29 @@ func TestOpenAIResponsesToChatIncompleteReasons(t *testing.T) { }) } } + +func TestConvertResponsesToOpenAIChatPreservesTypedExtensions(t *testing.T) { + raw := []byte(`{ + "id":"resp_extensions","created_at":1710000000,"model":"gpt-4.1", + "object":"response","status":"completed", + "output":[{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"refusal","refusal":"cannot comply"}]}], + "usage":{"input_tokens":12,"output_tokens":7,"total_tokens":19,"input_tokens_details":{"cached_tokens":4},"output_tokens_details":{"reasoning_tokens":3}} + }`) + var resp responses.Response + require.NoError(t, json.Unmarshal(raw, &resp)) + + converted := ConvertResponsesToOpenAIChat(&resp, "public-model") + require.Len(t, converted.Choices, 1) + assert.Equal(t, "cannot comply", converted.Choices[0].Message.Refusal) + require.NotNil(t, converted.Usage.PromptTokensDetails) + assert.EqualValues(t, 4, converted.Usage.PromptTokensDetails.CachedTokens) + require.NotNil(t, converted.Usage.CompletionTokensDetails) + assert.EqualValues(t, 3, converted.Usage.CompletionTokensDetails.ReasoningTokens) + + payload := converted.ToMap() + choices := payload["choices"].([]map[string]any) + message := choices[0]["message"].(map[string]any) + assert.Equal(t, "cannot comply", message["refusal"]) + usage := payload["usage"].(map[string]any) + assert.EqualValues(t, 3, usage["completion_tokens_details"].(map[string]any)["reasoning_tokens"]) +} diff --git a/internal/protocol/nonstream/openai_to_responses.go b/internal/protocol/nonstream/openai_to_responses.go index 513050cae..e284e41f6 100644 --- a/internal/protocol/nonstream/openai_to_responses.go +++ b/internal/protocol/nonstream/openai_to_responses.go @@ -9,6 +9,7 @@ import ( "github.com/openai/openai-go/v3" "github.com/tingly-dev/tingly-box/internal/protocol" "github.com/tingly-dev/tingly-box/internal/protocol/usage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" ) // HandleOpenAIChatToResponses writes an OpenAI Chat response as Responses API format. @@ -29,6 +30,12 @@ func HandleAnthropicBetaToResponses(hc *protocol.HandleContext, resp *anthropic. // BuildResponsesPayloadFromChat converts a Chat completion response to Responses API format. func BuildResponsesPayloadFromChat(resp *openai.ChatCompletion, responseModel, actualModel string) map[string]any { + return ConvertChatToResponsesWire(resp, responseModel, actualModel).ToMap() +} + +// ConvertChatToResponsesWire builds the typed Responses API output contract +// without routing protocol conversion through generic maps or JSON decoding. +func ConvertChatToResponsesWire(resp *openai.ChatCompletion, responseModel, actualModel string) wire.ResponsesWireResponse { model := responseModel if model == "" { model = actualModel @@ -41,135 +48,177 @@ func BuildResponsesPayloadFromChat(resp *openai.ChatCompletion, responseModel, a finishReason = string(resp.Choices[0].FinishReason) } - status, incompleteDetails := chatFinishReasonToResponsesStatus(finishReason) + status, incompleteReason := chatFinishReasonToResponsesStatus(finishReason) itemStatus := status - output := []map[string]any{} + output := []wire.ResponsesOutputItemWire{} if messageContent != "" { - output = append(output, map[string]any{ + output = append(output, wire.ResponsesOutputItemWire{ // The real Responses API always assigns output items an id; // strict clients (AI SDK zod) require it. - "id": "msg_" + resp.ID, - "type": "message", - "role": "assistant", - "status": itemStatus, - "content": []map[string]any{ + ID: "msg_" + resp.ID, + Type: "message", + Role: "assistant", + Status: itemStatus, + Content: []wire.ResponsesContentPartWire{ // The real Responses API always includes annotations on // output_text; strict clients (AI SDK zod) require it. - {"type": "output_text", "text": messageContent, "annotations": []any{}}, + {Type: "output_text", Text: messageContent, Annotations: []any{}}, }, }) } + if len(resp.Choices) > 0 { + for _, toolCall := range resp.Choices[0].Message.ToolCalls { + arguments := toolCall.Function.Arguments + output = append(output, wire.ResponsesOutputItemWire{ + ID: toolCall.ID, + CallID: toolCall.ID, + Type: "function_call", + Name: toolCall.Function.Name, + Arguments: &arguments, + Status: itemStatus, + }) + } + } // The canonical type owns the Responses usage shape, including the cache // read/write and reasoning detail a client would otherwise read as zeros. - usageMap := usage.FromOpenAIChatCompletion(resp.Usage).ToOpenAIResponsesUsageMap() - - result := map[string]any{ - "id": resp.ID, - "object": "response", - "created_at": time.Now().Unix(), - "model": model, - "status": status, - "output": output, - "usage": usageMap, + usageWire := tokenUsageToResponsesWire(usage.FromOpenAIChatCompletion(resp.Usage)) + + result := wire.ResponsesWireResponse{ + ID: resp.ID, + Object: "response", + CreatedAt: time.Now().Unix(), + Model: model, + Status: status, + Output: output, + Usage: usageWire, } - if incompleteDetails != nil { - result["incomplete_details"] = incompleteDetails + if incompleteReason != "" { + result.IncompleteDetails = &wire.ResponsesIncompleteDetailsWire{ + Reason: incompleteReason, + } } return result } -func chatFinishReasonToResponsesStatus(finishReason string) (string, map[string]any) { +func chatFinishReasonToResponsesStatus(finishReason string) (status, incompleteReason string) { switch finishReason { case "length": - return "incomplete", map[string]any{"reason": "max_output_tokens"} + return "incomplete", "max_output_tokens" case "content_filter": - return "incomplete", map[string]any{"reason": "content_filter"} + return "incomplete", "content_filter" default: - return "completed", nil + return "completed", "" } } // BuildResponsesPayloadFromAnthropicBeta converts an Anthropic Beta message response to Responses API format. func BuildResponsesPayloadFromAnthropicBeta(resp *anthropic.BetaMessage, responseModel, actualModel string) map[string]any { + return ConvertAnthropicBetaToResponsesWire(resp, responseModel, actualModel).ToMap() +} + +// ConvertAnthropicBetaToResponsesWire builds the typed Responses API output +// contract without coupling the Bridge to transport or SDK JSON internals. +func ConvertAnthropicBetaToResponsesWire(resp *anthropic.BetaMessage, responseModel, actualModel string) wire.ResponsesWireResponse { model := responseModel if model == "" { model = actualModel } - status, incompleteDetails := anthropicStopReasonToResponsesStatus(string(resp.StopReason)) + status, incompleteReason := anthropicStopReasonToResponsesStatus(string(resp.StopReason)) - output := []map[string]any{} - outputIndex := 0 + output := []wire.ResponsesOutputItemWire{} - var textParts []map[string]any + var textParts []wire.ResponsesContentPartWire for _, block := range resp.Content { - switch block.Type { - case "text": + if block.Type == "text" { if block.Text == "" { continue } - textParts = append(textParts, map[string]any{ - "type": "output_text", - "text": block.Text, - "annotations": []any{}, - }) - case "tool_use": - argsJSON := "{}" - if block.Input != nil { - if raw, err := json.Marshal(block.Input); err == nil { - argsJSON = string(raw) - } - } - output = append(output, map[string]any{ - "type": "function_call", - "id": block.ID, - "name": block.Name, - "arguments": argsJSON, - "output_index": outputIndex, + textParts = append(textParts, wire.ResponsesContentPartWire{ + Type: "output_text", + Text: block.Text, + Annotations: []any{}, }) - outputIndex++ } } if len(textParts) > 0 { - msgItem := map[string]any{ - "id": "msg_" + resp.ID, - "type": "message", - "role": "assistant", - "status": status, - "content": textParts, + msgItem := wire.ResponsesOutputItemWire{ + ID: "msg_" + resp.ID, + Type: "message", + Role: "assistant", + Status: status, + Content: textParts, } - output = append([]map[string]any{msgItem}, output...) + output = append(output, msgItem) + } + + for _, block := range resp.Content { + if block.Type != "tool_use" { + continue + } + argsJSON := "{}" + if block.Input != nil { + if raw, err := json.Marshal(block.Input); err == nil { + argsJSON = string(raw) + } + } + output = append(output, wire.ResponsesOutputItemWire{ + Type: "function_call", + ID: block.ID, + CallID: block.ID, + Name: block.Name, + Arguments: &argsJSON, + Status: status, + }) } // Responses-API input_tokens is the TOTAL prompt cost, so normalizing the // Anthropic usage folds cache_creation into it and reports cache_read // separately — Anthropic cache_creation lands in the cache_write_tokens slot, // both being the premium-rate write portion of the prompt total. - usageMap := usage.FromAnthropicBetaMessage(resp.Usage).ToOpenAIResponsesUsageMap() - - result := map[string]any{ - "id": resp.ID, - "object": "response", - "created_at": time.Now().Unix(), - "model": model, - "status": status, - "output": output, - "usage": usageMap, + usageWire := tokenUsageToResponsesWire(usage.FromAnthropicBetaMessage(resp.Usage)) + + result := wire.ResponsesWireResponse{ + ID: resp.ID, + Object: "response", + CreatedAt: time.Now().Unix(), + Model: model, + Status: status, + Output: output, + Usage: usageWire, } - if incompleteDetails != nil { - result["incomplete_details"] = incompleteDetails + if incompleteReason != "" { + result.IncompleteDetails = &wire.ResponsesIncompleteDetailsWire{ + Reason: incompleteReason, + } } return result } -func anthropicStopReasonToResponsesStatus(stopReason string) (string, map[string]any) { +func anthropicStopReasonToResponsesStatus(stopReason string) (status, incompleteReason string) { switch stopReason { case "max_tokens": - return "incomplete", map[string]any{"reason": "max_output_tokens"} + return "incomplete", "max_output_tokens" default: - return "completed", nil + return "completed", "" + } +} + +func tokenUsageToResponsesWire(normalized *protocol.TokenUsage) *wire.ResponsesUsageWire { + inputTokens := normalized.PromptTotalTokens() + return &wire.ResponsesUsageWire{ + InputTokens: int64(inputTokens), + OutputTokens: int64(normalized.OutputTokens), + TotalTokens: int64(inputTokens + normalized.OutputTokens), + InputTokensDetails: wire.ResponsesInputTokensDetailsWire{ + CachedTokens: int64(normalized.CacheReadTokens), + CacheWriteTokens: int64(normalized.CacheWriteTokens), + }, + OutputTokensDetails: wire.ResponsesOutputTokensDetailsWire{ + ReasoningTokens: int64(normalized.ReasoningTokens), + }, } } diff --git a/internal/protocol/nonstream/openai_usage_test.go b/internal/protocol/nonstream/openai_usage_test.go index 7ac666d7f..583dd4f28 100644 --- a/internal/protocol/nonstream/openai_usage_test.go +++ b/internal/protocol/nonstream/openai_usage_test.go @@ -1,10 +1,13 @@ package nonstream import ( + "encoding/json" + "strings" "testing" "github.com/anthropics/anthropic-sdk-go" "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/responses" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -32,6 +35,10 @@ func TestBuildResponsesPayloadFromChat_UsageDetails(t *testing.T) { } payload := BuildResponsesPayloadFromChat(resp, "gpt-x", "gpt-x") + typed := ConvertChatToResponsesWire(resp, "gpt-x", "gpt-x") + require.NotNil(t, typed.Usage) + assert.EqualValues(t, 30, typed.Usage.InputTokensDetails.CachedTokens) + assert.EqualValues(t, 12, typed.Usage.OutputTokensDetails.ReasoningTokens) usage, _ := payload["usage"].(map[string]any) require.NotNil(t, usage) @@ -47,6 +54,83 @@ func TestBuildResponsesPayloadFromChat_UsageDetails(t *testing.T) { assert.EqualValues(t, 12, outDetails["reasoning_tokens"]) } +func TestConvertAnthropicBetaToResponsesWireToolCall(t *testing.T) { + resp := &anthropic.BetaMessage{ + ID: "msg_tool", + Role: "assistant", + Type: "message", + Content: []anthropic.BetaContentBlockUnion{ + {Type: "text", Text: "checking"}, + { + Type: "tool_use", ID: "call_1", Name: "lookup", + Input: json.RawMessage(`{"query":"typed wire"}`), + }, + }, + StopReason: "tool_use", + } + + converted := ConvertAnthropicBetaToResponsesWire(resp, "public-model", "provider-model") + require.Len(t, converted.Output, 2) + assert.Equal(t, "message", converted.Output[0].Type) + item := converted.Output[1] + assert.Equal(t, "function_call", item.Type) + assert.Equal(t, "completed", item.Status) + assert.Equal(t, "call_1", item.CallID) + require.NotNil(t, item.Arguments) + assert.Contains(t, *item.Arguments, "typed wire") + + encoded, err := json.Marshal(converted) + require.NoError(t, err) + assert.True(t, strings.Contains(string(encoded), `"arguments":"{\"query\":\"typed wire\"}"`), string(encoded)) + assert.NotContains(t, string(encoded), `"output_index"`) +} + +func TestConvertChatToResponsesWirePreservesToolCalls(t *testing.T) { + resp := &openai.ChatCompletion{ + ID: "chatcmpl_tool", + Choices: []openai.ChatCompletionChoice{{ + FinishReason: "tool_calls", + Message: openai.ChatCompletionMessage{ToolCalls: []openai.ChatCompletionMessageToolCallUnion{{ + ID: "call_chat_1", + Type: "function", + Function: openai.ChatCompletionMessageFunctionToolCallFunction{ + Name: "lookup", Arguments: `{"query":"typed wire"}`, + }, + }}}, + }}, + } + + converted := ConvertChatToResponsesWire(resp, "public-model", "provider-model") + require.Len(t, converted.Output, 1) + item := converted.Output[0] + assert.Equal(t, "function_call", item.Type) + assert.Equal(t, "call_chat_1", item.ID) + assert.Equal(t, "call_chat_1", item.CallID) + assert.Equal(t, "lookup", item.Name) + require.NotNil(t, item.Arguments) + assert.JSONEq(t, `{"query":"typed wire"}`, *item.Arguments) +} + +func TestResponsesToAnthropicUsesCallIDForToolResultRoundTrip(t *testing.T) { + resp := &responses.Response{ + ID: "resp_tool", + Output: []responses.ResponseOutputItemUnion{{ + ID: "fc_item_1", Type: "function_call", CallID: "call_provider_1", Name: "lookup", + Arguments: responses.ResponseOutputItemUnionArguments{OfString: `{"query":"round trip"}`}, + }}, + } + + beta := HandleResponsesToAnthropicBeta(resp, "public-model") + require.Len(t, beta.Content, 1) + assert.Equal(t, "tool_use", beta.Content[0].Type) + assert.Equal(t, "call_provider_1", beta.Content[0].ID) + + v1 := HandleResponsesToAnthropicV1(resp, "public-model") + require.Len(t, v1.Content, 1) + assert.Equal(t, "tool_use", v1.Content[0].Type) + assert.Equal(t, "call_provider_1", v1.Content[0].ID) +} + // TestBuildResponsesPayloadFromAnthropicBeta_UsageDetails verifies that the // Responses-API input_tokens is the TOTAL prompt cost (uncached + cache-read + // cache-creation), matching the streaming converter, and that cache-read is @@ -69,6 +153,10 @@ func TestBuildResponsesPayloadFromAnthropicBeta_UsageDetails(t *testing.T) { } payload := BuildResponsesPayloadFromAnthropicBeta(resp, "claude-x", "claude-x") + typed := ConvertAnthropicBetaToResponsesWire(resp, "claude-x", "claude-x") + require.NotNil(t, typed.Usage) + assert.EqualValues(t, 66, typed.Usage.InputTokens) + assert.EqualValues(t, 11, typed.Usage.InputTokensDetails.CachedTokens) usage, _ := payload["usage"].(map[string]any) require.NotNil(t, usage) diff --git a/internal/protocol/request/anthropic_v1_to_beta.go b/internal/protocol/request/anthropic_v1_to_beta.go index c59479a7f..8d282d5bd 100644 --- a/internal/protocol/request/anthropic_v1_to_beta.go +++ b/internal/protocol/request/anthropic_v1_to_beta.go @@ -2,6 +2,7 @@ package request import ( "encoding/json" + "fmt" "github.com/anthropics/anthropic-sdk-go" "github.com/sirupsen/logrus" @@ -17,24 +18,32 @@ import ( // tool_result content and image data; the round-trip has no such gaps and // needs no updates when the SDK adds new block types. // -// Returns nil (rather than erroring) if req is nil or the round-trip fails, -// so callers doing best-effort context extraction (smart-routing) degrade -// gracefully instead of panicking. +// This compatibility wrapper keeps the historical nil-on-failure behavior for +// context-extraction callers. Protocol boundaries that need an actionable +// error should use ConvertAnthropicV1ToBetaRequestWithError. func ConvertAnthropicV1ToBetaRequest(req *anthropic.MessageNewParams) *anthropic.BetaMessageNewParams { - if req == nil { + converted, err := ConvertAnthropicV1ToBetaRequestWithError(req) + if err != nil { + logrus.WithError(err).Warn("ConvertAnthropicV1ToBetaRequest: wire conversion failed") return nil } + return converted +} - b, err := json.Marshal(req) - if err != nil { - logrus.WithError(err).Warn("ConvertAnthropicV1ToBetaRequest: marshal v1 request failed") - return nil +// ConvertAnthropicV1ToBetaRequestWithError performs the same wire conversion +// and reports malformed or non-JSON parameter values to the caller. +func ConvertAnthropicV1ToBetaRequestWithError(req *anthropic.MessageNewParams) (*anthropic.BetaMessageNewParams, error) { + if req == nil { + return nil, nil } + data, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("marshal Anthropic v1 request: %w", err) + } var beta anthropic.BetaMessageNewParams - if err := json.Unmarshal(b, &beta); err != nil { - logrus.WithError(err).Warn("ConvertAnthropicV1ToBetaRequest: unmarshal into beta shape failed") - return nil + if err := json.Unmarshal(data, &beta); err != nil { + return nil, fmt.Errorf("unmarshal Anthropic v1 request as Beta: %w", err) } - return &beta + return &beta, nil } diff --git a/internal/protocol/request/anthropic_v1_to_beta_test.go b/internal/protocol/request/anthropic_v1_to_beta_test.go index 4515fad73..7728f8474 100644 --- a/internal/protocol/request/anthropic_v1_to_beta_test.go +++ b/internal/protocol/request/anthropic_v1_to_beta_test.go @@ -1,6 +1,8 @@ package request import ( + "encoding/json" + "reflect" "testing" "github.com/anthropics/anthropic-sdk-go" @@ -130,3 +132,70 @@ func TestConvertAnthropicV1ToBetaRequest_ModelAndThinking(t *testing.T) { func TestConvertAnthropicV1ToBetaRequest_Nil(t *testing.T) { assert.Nil(t, ConvertAnthropicV1ToBetaRequest(nil)) } + +func TestConvertAnthropicV1ToBetaRequestPreservesWireSubset(t *testing.T) { + raw := []byte(`{ + "model":"claude-sonnet", + "max_tokens":1024, + "metadata":{"user_id":"user-1"}, + "system":[{"type":"text","text":"system","cache_control":{"type":"ephemeral"}}], + "messages":[ + {"role":"user","content":[ + {"type":"text","text":"look"}, + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}} + ]}, + {"role":"assistant","content":[{"type":"tool_use","id":"tool-1","name":"lookup","input":{"city":"Paris"}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"tool-1","content":[{"type":"text","text":"sunny"}]}]} + ], + "tools":[{"name":"lookup","description":"weather lookup","input_schema":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}], + "tool_choice":{"type":"tool","name":"lookup","disable_parallel_tool_use":true}, + "stop_sequences":["done"], + "temperature":0.2, + "top_k":20, + "top_p":0.8, + "thinking":{"type":"enabled","budget_tokens":256} + }`) + + var v1 anthropic.MessageNewParams + if err := json.Unmarshal(raw, &v1); err != nil { + t.Fatalf("decode v1 request: %v", err) + } + beta, err := ConvertAnthropicV1ToBetaRequestWithError(&v1) + if err != nil { + t.Fatalf("ConvertAnthropicV1ToBetaRequestWithError() error = %v", err) + } + + want := decodeJSONObject(t, mustMarshalJSON(t, &v1)) + got := decodeJSONObject(t, mustMarshalJSON(t, beta)) + if !reflect.DeepEqual(got, want) { + t.Fatalf("Beta wire request differs from v1 subset\n got: %#v\nwant: %#v", got, want) + } +} + +func TestConvertAnthropicV1ToBetaRequestNil(t *testing.T) { + if got := ConvertAnthropicV1ToBetaRequest(nil); got != nil { + t.Fatalf("ConvertAnthropicV1ToBetaRequest(nil) = %#v, want nil", got) + } + got, err := ConvertAnthropicV1ToBetaRequestWithError(nil) + if err != nil || got != nil { + t.Fatalf("ConvertAnthropicV1ToBetaRequestWithError(nil) = %#v, %v", got, err) + } +} + +func mustMarshalJSON(t *testing.T, value any) []byte { + t.Helper() + data, err := json.Marshal(value) + if err != nil { + t.Fatalf("json.Marshal(%T): %v", value, err) + } + return data +} + +func decodeJSONObject(t *testing.T, data []byte) map[string]any { + t.Helper() + var value map[string]any + if err := json.Unmarshal(data, &value); err != nil { + t.Fatalf("decode JSON object: %v", err) + } + return value +} diff --git a/internal/protocol/stage/anthropicbridge/bridge.go b/internal/protocol/stage/anthropicbridge/bridge.go new file mode 100644 index 000000000..aa5f8b2d0 --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/bridge.go @@ -0,0 +1,214 @@ +// Package anthropicbridge adapts Anthropic Messages calls to OpenAI Chat while +// keeping the outward response in the exact Anthropic source protocol. +package anthropicbridge + +import ( + "context" + "fmt" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/nonstream" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" +) + +// ChatOptions configures the existing Anthropic-to-OpenAI-Chat request +// conversion. Options are immutable after the Bridge is constructed. +type ChatOptions struct { + Compatible bool + DisableStreamUsage bool + // ResponseModel overrides the source-visible Anthropic response model while + // leaving the provider-bound request model unchanged. + ResponseModel string +} + +// NewV1ToOpenAIChat returns an immutable Anthropic v1 -> OpenAI Chat Bridge. +func NewV1ToOpenAIChat(options ChatOptions) stage.Bridge { + return &chatBridge{source: protocol.TypeAnthropicV1, options: options} +} + +// NewBetaToOpenAIChat returns an immutable Anthropic beta -> OpenAI Chat Bridge. +func NewBetaToOpenAIChat(options ChatOptions) stage.Bridge { + return &chatBridge{source: protocol.TypeAnthropicBeta, options: options} +} + +type chatBridge struct { + source protocol.APIType + options ChatOptions +} + +func (b *chatBridge) Source() protocol.APIType { return b.source } + +func (b *chatBridge) Target() protocol.APIType { return protocol.TypeOpenAIChat } + +func (b *chatBridge) Capabilities() stage.Capabilities { + return stage.AllBridgeCapabilities +} + +func (b *chatBridge) Open(_ context.Context, call stage.Call, operation stage.Operation) (stage.BridgeSession, error) { + isStreaming, err := operationStreaming(operation) + if err != nil { + return nil, fmt.Errorf("open Anthropic to OpenAI Chat bridge: %w", err) + } + + var ( + chatRequest *openai.ChatCompletionNewParams + config *protocol.OpenAIConfig + sourceModel string + ) + switch b.source { + case protocol.TypeAnthropicV1: + anthropicRequest, err := v1Request(call.Request) + if err != nil { + return nil, err + } + chatRequest, config = request.ConvertAnthropicToOpenAIRequest( + anthropicRequest, + b.options.Compatible, + isStreaming, + b.options.DisableStreamUsage, + ) + sourceModel = string(anthropicRequest.Model) + case protocol.TypeAnthropicBeta: + anthropicRequest, err := betaRequest(call.Request) + if err != nil { + return nil, err + } + chatRequest, config = request.ConvertAnthropicBetaToOpenAIRequest( + anthropicRequest, + b.options.Compatible, + isStreaming, + b.options.DisableStreamUsage, + ) + sourceModel = string(anthropicRequest.Model) + default: + return nil, fmt.Errorf("open Anthropic to OpenAI Chat bridge: unsupported source protocol %q", b.source) + } + if chatRequest == nil { + return nil, fmt.Errorf("open Anthropic to OpenAI Chat bridge %q: request conversion returned nil", b.source) + } + if b.options.ResponseModel != "" { + sourceModel = b.options.ResponseModel + } + + targetCall := call + targetCall.Request = chatRequest + targetCall.State.OpenAIChat = config + return &chatSession{ + source: b.source, + operation: operation, + targetCall: targetCall, + targetRequest: chatRequest, + sourceModel: sourceModel, + }, nil +} + +func operationStreaming(operation stage.Operation) (bool, error) { + switch operation { + case stage.OperationComplete: + return false, nil + case stage.OperationStream: + return true, nil + default: + return false, fmt.Errorf("unsupported operation %s", operation) + } +} + +func v1Request(value any) (*anthropic.MessageNewParams, error) { + switch request := value.(type) { + case *anthropic.MessageNewParams: + if request == nil { + return nil, fmt.Errorf("open Anthropic v1 to OpenAI Chat bridge: request is nil") + } + return request, nil + case anthropic.MessageNewParams: + return &request, nil + default: + return nil, fmt.Errorf("open Anthropic v1 to OpenAI Chat bridge: request has type %T, want anthropic.MessageNewParams", value) + } +} + +func betaRequest(value any) (*anthropic.BetaMessageNewParams, error) { + switch request := value.(type) { + case *anthropic.BetaMessageNewParams: + if request == nil { + return nil, fmt.Errorf("open Anthropic beta to OpenAI Chat bridge: request is nil") + } + return request, nil + case anthropic.BetaMessageNewParams: + return &request, nil + default: + return nil, fmt.Errorf("open Anthropic beta to OpenAI Chat bridge: request has type %T, want anthropic.BetaMessageNewParams", value) + } +} + +type chatSession struct { + source protocol.APIType + operation stage.Operation + targetCall stage.Call + targetRequest *openai.ChatCompletionNewParams + sourceModel string +} + +func (s *chatSession) TargetCall() stage.Call { return s.targetCall } + +func (s *chatSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + if s.operation != stage.OperationComplete { + return nil, fmt.Errorf("convert Anthropic complete response: session was opened for %s", s.operation) + } + chat, err := chatCompletion(response) + if err != nil { + return nil, err + } + + var value any + switch s.source { + case protocol.TypeAnthropicV1: + value, err = nonstream.ConvertOpenAIChatToAnthropicV1(chat, s.sourceModel) + case protocol.TypeAnthropicBeta: + value, err = nonstream.ConvertOpenAIChatToAnthropicBeta(chat, s.sourceModel) + default: + err = fmt.Errorf("convert OpenAI Chat response: unsupported Anthropic source protocol %q", s.source) + } + if err != nil { + return nil, err + } + normalizedUsage := protocolusage.FromOpenAIChatCompletion(chat.Usage) + if !normalizedUsage.HasUsage() { + normalizedUsage = nil + } + return &stage.Response{ + Value: value, + Usage: normalizedUsage, + Model: s.sourceModel, + }, nil +} + +func chatCompletion(response *stage.Response) (*openai.ChatCompletion, error) { + if response == nil { + return nil, fmt.Errorf("convert OpenAI Chat response to Anthropic: response is nil") + } + switch value := response.Value.(type) { + case *openai.ChatCompletion: + if value == nil { + return nil, fmt.Errorf("convert OpenAI Chat response to Anthropic: value is nil") + } + return value, nil + case openai.ChatCompletion: + return &value, nil + default: + return nil, fmt.Errorf("convert OpenAI Chat response to Anthropic: value has type %T, want openai.ChatCompletion", response.Value) + } +} + +func (s *chatSession) ConvertStream(_ context.Context, target stage.EventStream) (stage.EventStream, error) { + if s.operation != stage.OperationStream { + return nil, fmt.Errorf("convert OpenAI Chat stream: session was opened for %s", s.operation) + } + return newAnthropicStream(target, s.source, s.sourceModel, s.targetRequest) +} + +func (s *chatSession) ConvertError(_ context.Context, err error) error { return err } diff --git a/internal/protocol/stage/anthropicbridge/bridge_test.go b/internal/protocol/stage/anthropicbridge/bridge_test.go new file mode 100644 index 000000000..1144a1c78 --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/bridge_test.go @@ -0,0 +1,615 @@ +package anthropicbridge + +import ( + "context" + "encoding/json" + "errors" + "io" + "reflect" + "strings" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" +) + +func TestAnthropicIdentityBridges(t *testing.T) { + t.Parallel() + + v1Request := &anthropic.MessageNewParams{Model: "claude-v1", MaxTokens: 32} + v1Response := &anthropic.Message{ID: "msg-v1", Model: "claude-v1"} + v1Event := anthropic.MessageStreamEventUnion{Type: "message_stop"} + betaRequest := &anthropic.BetaMessageNewParams{Model: "claude-beta", MaxTokens: 32} + betaResponse := &anthropic.BetaMessage{ID: "msg-beta", Model: "claude-beta"} + betaEvent := anthropic.BetaRawMessageStreamEventUnion{Type: "message_stop"} + + tests := []struct { + name string + api protocol.APIType + request any + response any + event any + }{ + {name: "v1", api: protocol.TypeAnthropicV1, request: v1Request, response: v1Response, event: v1Event}, + {name: "beta", api: protocol.TypeAnthropicBeta, request: betaRequest, response: betaResponse, event: betaEvent}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + usage := protocol.NewTokenUsage(3, 2) + targetStream := &memoryStream{ + events: []stage.Event{{Value: tt.event}}, + result: stage.StreamResult{Usage: usage, Model: "identity-model", SideEffectsCommitted: true}, + } + terminal := &memoryEndpoint{ + api: tt.api, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + if call.Request != tt.request { + t.Fatalf("identity request = %T %v, want same value", call.Request, call.Request) + } + return &stage.Response{Value: tt.response, Usage: usage, Model: "identity-model", SideEffectsCommitted: true}, nil + }, + stream: func(_ context.Context, call stage.Call) (stage.EventStream, error) { + if call.Request != tt.request { + t.Fatalf("identity stream request = %T, want same value", call.Request) + } + return targetStream, nil + }, + } + adapted := mustAdapt(t, terminal, stage.NewIdentityBridge(tt.api)) + state := stage.ProtocolState{OpenAIChat: &protocol.OpenAIConfig{HasThinking: true}} + call := stage.Call{Request: tt.request, Metadata: stage.CallMetadata{RequestID: "identity", Attempt: 2}, State: state} + + response, err := adapted.Complete(context.Background(), call) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if response.Value != tt.response || terminal.lastCall.State.OpenAIChat != state.OpenAIChat { + t.Fatalf("identity complete changed value/state: response=%T state=%p", response.Value, terminal.lastCall.State.OpenAIChat) + } + + stream, err := adapted.Stream(context.Background(), call) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + event, err := stream.Next(context.Background()) + if err != nil || !reflect.DeepEqual(event.Value, tt.event) { + t.Fatalf("Next() = (%T %+v, %v), want identity event", event.Value, event.Value, err) + } + if _, err := stream.Next(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("second Next() error = %v, want io.EOF", err) + } + if got := stream.Result(); got.Usage != usage || got.Model != "identity-model" || !got.SideEffectsCommitted { + t.Fatalf("Result() = %+v", got) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if targetStream.closeCount != 1 || terminal.lastCall.State.OpenAIChat != state.OpenAIChat { + t.Fatalf("identity close/state = %d/%p", targetStream.closeCount, terminal.lastCall.State.OpenAIChat) + } + }) + } +} + +func TestAnthropicToOpenAIChatComplete(t *testing.T) { + t.Parallel() + + completion := decodeChatCompletion(t, `{ + "id":"chat-1", + "model":"provider-model", + "choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":"","tool_calls":[{"id":"tool-1","type":"function","function":{"name":"lookup","arguments":"{\"q\":\"x\"}"}}]}}], + "usage":{"prompt_tokens":10,"completion_tokens":4,"total_tokens":14,"prompt_tokens_details":{"cached_tokens":2},"completion_tokens_details":{"reasoning_tokens":1}} + }`) + + tests := []struct { + name string + bridge stage.Bridge + request any + wantType any + sourceModel string + }{ + { + name: "v1", + bridge: NewV1ToOpenAIChat(ChatOptions{}), + request: &anthropic.MessageNewParams{Model: "client-v1", MaxTokens: 64, Messages: []anthropic.MessageParam{anthropic.NewUserMessage(anthropic.NewTextBlock("hello"))}}, + wantType: (*anthropic.Message)(nil), + sourceModel: "client-v1", + }, + { + name: "beta", + bridge: NewBetaToOpenAIChat(ChatOptions{}), + request: &anthropic.BetaMessageNewParams{Model: "client-beta", MaxTokens: 64, Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hello"))}}, + wantType: (*anthropic.BetaMessage)(nil), + sourceModel: "client-beta", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIChat, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + chatRequest := requireChatRequest(t, call) + if chatRequest.StreamOptions.IncludeUsage.Valid() { + t.Fatal("complete request unexpectedly enabled stream usage") + } + if call.Metadata.RequestID != "complete-request" || call.Metadata.Attempt != 3 { + t.Fatalf("metadata = %+v", call.Metadata) + } + return &stage.Response{ + Value: completion, + Usage: protocol.NewTokenUsage(999, 999), + Model: "provider-model", + SideEffectsCommitted: true, + }, nil + }, + } + adapted := mustAdapt(t, terminal, tt.bridge) + response, err := adapted.Complete(context.Background(), stage.Call{ + Request: tt.request, + Metadata: stage.CallMetadata{RequestID: "complete-request", Attempt: 3}, + State: stage.ProtocolState{OpenAIChat: &protocol.OpenAIConfig{ReasoningEffort: "xhigh"}}, + }) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if reflect.TypeOf(response.Value) != reflect.TypeOf(tt.wantType) { + t.Fatalf("response type = %T, want %T", response.Value, tt.wantType) + } + if response.Model != tt.sourceModel || !response.SideEffectsCommitted { + t.Fatalf("response facts = %+v", response) + } + if response.Usage == nil || response.Usage.InputTokens != 8 || response.Usage.OutputTokens != 4 || response.Usage.CacheReadTokens != 2 || response.Usage.ReasoningTokens != 1 { + t.Fatalf("normalized usage = %+v", response.Usage) + } + var wire map[string]any + value, err := json.Marshal(response.Value) + if err != nil { + t.Fatalf("marshal response: %v", err) + } + if err := json.Unmarshal(value, &wire); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + if wire["model"] != tt.sourceModel || wire["stop_reason"] != "tool_use" { + t.Fatalf("response wire model/stop = %v/%v", wire["model"], wire["stop_reason"]) + } + content, ok := wire["content"].([]any) + if !ok || len(content) != 1 || content[0].(map[string]any)["type"] != "tool_use" { + t.Fatalf("response content = %#v", wire["content"]) + } + if terminal.lastCall.State.OpenAIChat == nil || terminal.lastCall.State.OpenAIChat.ReasoningEffort != "medium" { + t.Fatalf("target OpenAIConfig = %+v", terminal.lastCall.State.OpenAIChat) + } + }) + } +} + +func TestAnthropicToOpenAIChatCompletePreservesTargetUsageWhenWireUsageIsMissing(t *testing.T) { + t.Parallel() + + completion := decodeChatCompletion(t, `{ + "id":"chat-no-usage", + "model":"provider-model", + "choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}] + }`) + targetUsage := protocol.NewTokenUsage(13, 5) + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIChat, + complete: func(context.Context, stage.Call) (*stage.Response, error) { + return &stage.Response{Value: completion, Usage: targetUsage}, nil + }, + } + adapted := mustAdapt(t, terminal, NewV1ToOpenAIChat(ChatOptions{})) + response, err := adapted.Complete(context.Background(), stage.Call{Request: &anthropic.MessageNewParams{Model: "client-model"}}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if response.Usage != targetUsage { + t.Fatalf("response usage = %+v, want target usage %+v", response.Usage, targetUsage) + } +} + +func TestAnthropicToOpenAIChatSeparatesProviderAndResponseModels(t *testing.T) { + t.Parallel() + + request := &anthropic.BetaMessageNewParams{ + Model: "provider-model", + MaxTokens: 64, + Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hello"))}, + } + completion := decodeChatCompletion(t, `{ + "id":"chat-models", + "model":"provider-model", + "choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}], + "usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4} + }`) + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIChat, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + chatRequest := requireChatRequest(t, call) + if chatRequest.Model != "provider-model" { + t.Fatalf("provider request model = %q", chatRequest.Model) + } + return &stage.Response{Value: completion}, nil + }, + } + bridge := NewBetaToOpenAIChat(ChatOptions{Compatible: true, ResponseModel: "client-alias"}) + response, err := mustAdapt(t, terminal, bridge).Complete(context.Background(), stage.Call{Request: request}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + message, ok := response.Value.(*anthropic.BetaMessage) + if !ok || string(message.Model) != "client-alias" || response.Model != "client-alias" { + t.Fatalf("source-visible models = message:%v response:%q", message, response.Model) + } +} + +func TestAnthropicIdentityStageThenOpenAIChatTopology(t *testing.T) { + t.Parallel() + + completion := decodeChatCompletion(t, `{ + "id":"chat-topology", + "model":"provider-model", + "choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"topology ok"}}], + "usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5} + }`) + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIChat, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + requireChatRequest(t, call) + return &stage.Response{Value: completion, Model: "provider-model"}, nil + }, + } + registry, err := stage.NewBridgeRegistry(NewV1ToOpenAIChat(ChatOptions{})) + if err != nil { + t.Fatalf("NewBridgeRegistry() error = %v", err) + } + identity := &anthropicPassthroughStage{api: protocol.TypeAnthropicV1} + topology, err := stage.BuildTopology(stage.TopologyConfig{ + Terminal: terminal, + Stages: []stage.Stage{identity}, + ClientProtocol: protocol.TypeAnthropicV1, + Registry: registry, + RequiredCapabilities: stage.CapabilityUsage | stage.CapabilityFinishReason, + }) + if err != nil { + t.Fatalf("BuildTopology() error = %v", err) + } + request := &anthropic.MessageNewParams{ + Model: "client-topology", + MaxTokens: 32, + Messages: []anthropic.MessageParam{anthropic.NewUserMessage(anthropic.NewTextBlock("hello"))}, + } + response, err := topology.Complete(context.Background(), stage.Call{Request: request}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if _, ok := identity.request.(*anthropic.MessageNewParams); !ok { + t.Fatalf("identity stage request = %T", identity.request) + } + if _, ok := identity.response.(*anthropic.Message); !ok { + t.Fatalf("identity stage response = %T", identity.response) + } + if _, ok := response.Value.(*anthropic.Message); !ok || response.Model != "client-topology" { + t.Fatalf("topology response = %T %+v", response.Value, response) + } +} + +func TestAnthropicToOpenAIChatStream(t *testing.T) { + t.Parallel() + + chunks := []stage.Event{ + {Value: decodeChatChunk(t, `{"id":"chunk-1","model":"provider-model","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":""}]}`)}, + {Value: decodeChatChunk(t, `{"id":"chunk-1","model":"provider-model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`)}, + {Value: decodeChatChunk(t, `{"id":"chunk-1","model":"provider-model","choices":[],"usage":{"prompt_tokens":7,"completion_tokens":2,"total_tokens":9,"prompt_tokens_details":{"cached_tokens":1}}}`)}, + } + + tests := []struct { + name string + bridge stage.Bridge + request any + sourceModel string + }{ + {name: "v1", bridge: NewV1ToOpenAIChat(ChatOptions{}), request: &anthropic.MessageNewParams{Model: "stream-v1", MaxTokens: 64, Messages: []anthropic.MessageParam{anthropic.NewUserMessage(anthropic.NewTextBlock("hello"))}}, sourceModel: "stream-v1"}, + {name: "beta", bridge: NewBetaToOpenAIChat(ChatOptions{}), request: &anthropic.BetaMessageNewParams{Model: "stream-beta", MaxTokens: 64, Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hello"))}}, sourceModel: "stream-beta"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + target := &memoryStream{ + events: append([]stage.Event(nil), chunks...), + result: stage.StreamResult{Usage: protocol.NewTokenUsage(999, 999), Model: "provider-model", SideEffectsCommitted: true}, + } + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIChat, + stream: func(_ context.Context, call stage.Call) (stage.EventStream, error) { + chatRequest := requireChatRequest(t, call) + if !chatRequest.StreamOptions.IncludeUsage.Valid() || !chatRequest.StreamOptions.IncludeUsage.Value { + t.Fatalf("stream request include_usage = %+v", chatRequest.StreamOptions.IncludeUsage) + } + if call.State.OpenAIChat == nil { + t.Fatal("stream target call lost OpenAIConfig") + } + return target, nil + }, + } + adapted := mustAdapt(t, terminal, tt.bridge) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: tt.request}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + var eventTypes []string + var first protocolstream.AnthropicEvent + for { + event, err := stream.Next(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("Next() error = %v", err) + } + value, ok := event.Value.(protocolstream.AnthropicEvent) + if !ok { + t.Fatalf("event type = %T", event.Value) + } + if len(eventTypes) == 0 { + first = value + } + eventTypes = append(eventTypes, value.Type) + } + wantTypes := []string{"message_start", "content_block_start", "content_block_delta", "content_block_stop", "message_delta", "message_stop"} + if !reflect.DeepEqual(eventTypes, wantTypes) { + t.Fatalf("event types = %v, want %v", eventTypes, wantTypes) + } + firstJSON, err := json.Marshal(first.Data) + if err != nil || !strings.Contains(string(firstJSON), `"model":"`+tt.sourceModel+`"`) { + t.Fatalf("message_start = %s, err=%v", firstJSON, err) + } + result := stream.Result() + if result.Usage == nil || result.Usage.InputTokens != 6 || result.Usage.CacheReadTokens != 1 || result.Usage.OutputTokens != 2 { + t.Fatalf("stream usage = %+v", result.Usage) + } + if result.Model != tt.sourceModel || !result.SideEffectsCommitted { + t.Fatalf("stream result = %+v", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("second Close() error = %v", err) + } + if target.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", target.closeCount) + } + }) + } +} + +func TestAnthropicStreamTreatsZeroUsageAsMissing(t *testing.T) { + t.Parallel() + + stream := &anthropicStream{ + converter: staticUsageConverter{usage: protocol.ZeroTokenUsage()}, + model: "client-model", + } + if result := stream.Result(); result.Usage != nil { + t.Fatalf("Result().Usage = %+v, want nil", result.Usage) + } +} + +func TestAnthropicToOpenAIChatFailuresAndCancellation(t *testing.T) { + t.Parallel() + + t.Run("wrong source request", func(t *testing.T) { + terminal := &memoryEndpoint{api: protocol.TypeOpenAIChat} + adapted := mustAdapt(t, terminal, NewV1ToOpenAIChat(ChatOptions{})) + _, err := adapted.Complete(context.Background(), stage.Call{Request: "wrong"}) + if err == nil || !strings.Contains(err.Error(), "request has type string") { + t.Fatalf("Complete() error = %v", err) + } + }) + + t.Run("wrong complete response", func(t *testing.T) { + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIChat, + complete: func(context.Context, stage.Call) (*stage.Response, error) { + return &stage.Response{Value: "wrong"}, nil + }, + } + adapted := mustAdapt(t, terminal, NewV1ToOpenAIChat(ChatOptions{})) + _, err := adapted.Complete(context.Background(), stage.Call{Request: &anthropic.MessageNewParams{Model: "m"}}) + if err == nil || !strings.Contains(err.Error(), "value has type string") { + t.Fatalf("Complete() error = %v", err) + } + }) + + t.Run("wrong stream event", func(t *testing.T) { + target := &memoryStream{events: []stage.Event{{Value: "wrong"}}} + terminal := &memoryEndpoint{api: protocol.TypeOpenAIChat, stream: func(context.Context, stage.Call) (stage.EventStream, error) { return target, nil }} + adapted := mustAdapt(t, terminal, NewV1ToOpenAIChat(ChatOptions{})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &anthropic.MessageNewParams{Model: "m"}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + _, err = stream.Next(context.Background()) + if err == nil || !strings.Contains(err.Error(), "event has type string") { + t.Fatalf("Next() error = %v", err) + } + _ = stream.Close() + if target.closeCount != 1 { + t.Fatalf("target close count = %d", target.closeCount) + } + }) + + t.Run("target iterator error", func(t *testing.T) { + upstreamErr := errors.New("upstream stream failed") + target := &memoryStream{nextErr: upstreamErr} + terminal := &memoryEndpoint{api: protocol.TypeOpenAIChat, stream: func(context.Context, stage.Call) (stage.EventStream, error) { return target, nil }} + adapted := mustAdapt(t, terminal, NewV1ToOpenAIChat(ChatOptions{})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &anthropic.MessageNewParams{Model: "m"}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + _, err = stream.Next(context.Background()) + if !errors.Is(err, upstreamErr) { + t.Fatalf("Next() error = %v", err) + } + _ = stream.Close() + }) + + t.Run("canceled before pull", func(t *testing.T) { + target := &memoryStream{events: []stage.Event{{Value: openai.ChatCompletionChunk{}}}} + terminal := &memoryEndpoint{api: protocol.TypeOpenAIChat, stream: func(context.Context, stage.Call) (stage.EventStream, error) { return target, nil }} + adapted := mustAdapt(t, terminal, NewV1ToOpenAIChat(ChatOptions{})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &anthropic.MessageNewParams{Model: "m"}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = stream.Next(ctx) + if !errors.Is(err, context.Canceled) || target.nextCount != 0 { + t.Fatalf("Next() = %v, target pulls=%d", err, target.nextCount) + } + _ = stream.Close() + }) +} + +func requireChatRequest(t *testing.T, call stage.Call) *openai.ChatCompletionNewParams { + t.Helper() + request, ok := call.Request.(*openai.ChatCompletionNewParams) + if !ok || request == nil { + t.Fatalf("target request = %T, want *openai.ChatCompletionNewParams", call.Request) + } + if call.State.OpenAIChat == nil { + t.Fatal("target call has nil OpenAIConfig") + } + return request +} + +func decodeChatCompletion(t *testing.T, value string) *openai.ChatCompletion { + t.Helper() + var completion openai.ChatCompletion + if err := json.Unmarshal([]byte(value), &completion); err != nil { + t.Fatalf("decode Chat completion: %v", err) + } + return &completion +} + +func decodeChatChunk(t *testing.T, value string) openai.ChatCompletionChunk { + t.Helper() + var chunk openai.ChatCompletionChunk + if err := json.Unmarshal([]byte(value), &chunk); err != nil { + t.Fatalf("decode Chat chunk: %v", err) + } + return chunk +} + +func mustAdapt(t *testing.T, terminal stage.Endpoint, bridge stage.Bridge) stage.Endpoint { + t.Helper() + adapted, err := stage.Adapt(terminal, bridge) + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + return adapted +} + +type memoryEndpoint struct { + api protocol.APIType + complete func(context.Context, stage.Call) (*stage.Response, error) + stream func(context.Context, stage.Call) (stage.EventStream, error) + lastCall stage.Call +} + +type staticUsageConverter struct { + usage *protocol.TokenUsage +} + +func (staticUsageConverter) Next() (any, bool, error) { return nil, true, nil } + +func (c staticUsageConverter) Usage() *protocol.TokenUsage { return c.usage } + +func (e *memoryEndpoint) Protocol() protocol.APIType { return e.api } + +func (e *memoryEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + e.lastCall = call + if e.complete == nil { + return nil, errors.New("unexpected complete call") + } + return e.complete(ctx, call) +} + +func (e *memoryEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + e.lastCall = call + if e.stream == nil { + return nil, errors.New("unexpected stream call") + } + return e.stream(ctx, call) +} + +type memoryStream struct { + events []stage.Event + result stage.StreamResult + nextErr error + nextCount int + closeCount int +} + +func (s *memoryStream) Next(ctx context.Context) (stage.Event, error) { + s.nextCount++ + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + if len(s.events) > 0 { + event := s.events[0] + s.events = s.events[1:] + return event, nil + } + if s.nextErr != nil { + return stage.Event{}, s.nextErr + } + return stage.Event{}, io.EOF +} + +func (s *memoryStream) Close() error { + s.closeCount++ + return nil +} + +func (s *memoryStream) Result() stage.StreamResult { return s.result } + +type anthropicPassthroughStage struct { + api protocol.APIType + request any + response any +} + +func (s *anthropicPassthroughStage) Name() string { return "anthropic_identity" } +func (s *anthropicPassthroughStage) Protocol() protocol.APIType { return s.api } +func (s *anthropicPassthroughStage) Wrap(next stage.Endpoint) stage.Endpoint { + return &anthropicPassthroughEndpoint{stage: s, next: next} +} + +type anthropicPassthroughEndpoint struct { + stage *anthropicPassthroughStage + next stage.Endpoint +} + +func (e *anthropicPassthroughEndpoint) Protocol() protocol.APIType { return e.stage.api } +func (e *anthropicPassthroughEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + e.stage.request = call.Request + response, err := e.next.Complete(ctx, call) + if response != nil { + e.stage.response = response.Value + } + return response, err +} +func (e *anthropicPassthroughEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + e.stage.request = call.Request + return e.next.Stream(ctx, call) +} diff --git a/internal/protocol/stage/anthropicbridge/responses.go b/internal/protocol/stage/anthropicbridge/responses.go new file mode 100644 index 000000000..b248bd56f --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/responses.go @@ -0,0 +1,147 @@ +package anthropicbridge + +import ( + "context" + "fmt" + + "github.com/openai/openai-go/v3/responses" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/nonstream" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" +) + +// ResponsesOptions configures Anthropic to OpenAI Responses conversion. +type ResponsesOptions struct { + ResponseModel string +} + +// NewBetaToOpenAIResponses returns an Anthropic Beta to Responses Bridge. +func NewBetaToOpenAIResponses(options ResponsesOptions) stage.Bridge { + return &responsesBridge{source: protocol.TypeAnthropicBeta, options: options} +} + +// NewV1ToOpenAIResponses returns an Anthropic v1 to Responses Bridge. +// Production registration is deliberately separate from construction. +func NewV1ToOpenAIResponses(options ResponsesOptions) stage.Bridge { + return &responsesBridge{source: protocol.TypeAnthropicV1, options: options} +} + +type responsesBridge struct { + source protocol.APIType + options ResponsesOptions +} + +func (b *responsesBridge) Source() protocol.APIType { return b.source } +func (*responsesBridge) Target() protocol.APIType { return protocol.TypeOpenAIResponses } +func (*responsesBridge) Capabilities() stage.Capabilities { + return stage.AllBridgeCapabilities +} + +func (b *responsesBridge) Open(_ context.Context, call stage.Call, operation stage.Operation) (stage.BridgeSession, error) { + switch operation { + case stage.OperationComplete, stage.OperationStream: + default: + return nil, fmt.Errorf("open Anthropic to OpenAI Responses bridge: unsupported operation %s", operation) + } + var ( + targetRequest *responses.ResponseNewParams + sourceModel string + ) + switch b.source { + case protocol.TypeAnthropicBeta: + sourceRequest, err := betaRequest(call.Request) + if err != nil { + return nil, err + } + targetRequest = request.ConvertAnthropicBetaToResponsesRequest(sourceRequest) + sourceModel = string(sourceRequest.Model) + case protocol.TypeAnthropicV1: + sourceRequest, err := v1Request(call.Request) + if err != nil { + return nil, err + } + targetRequest = request.ConvertAnthropicV1ToResponsesRequest(sourceRequest) + sourceModel = string(sourceRequest.Model) + default: + return nil, fmt.Errorf("open Anthropic to OpenAI Responses bridge: unsupported source protocol %q", b.source) + } + if targetRequest == nil { + return nil, fmt.Errorf("open Anthropic to OpenAI Responses bridge %q: request conversion returned nil", b.source) + } + if b.options.ResponseModel != "" { + sourceModel = b.options.ResponseModel + } + targetCall := call + targetCall.Request = targetRequest + targetCall.State.OpenAIChat = nil + return &responsesSession{ + source: b.source, + operation: operation, + targetCall: targetCall, + sourceModel: sourceModel, + }, nil +} + +type responsesSession struct { + source protocol.APIType + operation stage.Operation + targetCall stage.Call + sourceModel string +} + +func (s *responsesSession) TargetCall() stage.Call { return s.targetCall } + +func (s *responsesSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + if s.operation != stage.OperationComplete { + return nil, fmt.Errorf("convert OpenAI Responses complete response to Anthropic: session was opened for %s", s.operation) + } + value, err := responsesValue(response) + if err != nil { + return nil, err + } + var converted any + switch s.source { + case protocol.TypeAnthropicBeta: + message := nonstream.HandleResponsesToAnthropicBeta(value, s.sourceModel) + converted = &message + case protocol.TypeAnthropicV1: + message := nonstream.HandleResponsesToAnthropicV1(value, s.sourceModel) + converted = &message + default: + return nil, fmt.Errorf("convert OpenAI Responses complete response: unsupported Anthropic source %q", s.source) + } + usage := protocolusage.FromOpenAIResponses(value.Usage) + if !usage.HasUsage() { + usage = nil + } + return &stage.Response{Value: converted, Usage: usage, Model: s.sourceModel}, nil +} + +func responsesValue(response *stage.Response) (*responses.Response, error) { + if response == nil { + return nil, fmt.Errorf("convert OpenAI Responses response to Anthropic: response is nil") + } + switch value := response.Value.(type) { + case *responses.Response: + if value == nil { + return nil, fmt.Errorf("convert OpenAI Responses response to Anthropic: value is nil") + } + return value, nil + case responses.Response: + return &value, nil + default: + return nil, fmt.Errorf("convert OpenAI Responses response to Anthropic: value has type %T, want responses.Response", response.Value) + } +} + +func (s *responsesSession) ConvertStream(ctx context.Context, target stage.EventStream) (stage.EventStream, error) { + if s.operation != stage.OperationStream { + return nil, fmt.Errorf("convert OpenAI Responses stream to Anthropic: session was opened for %s", s.operation) + } + return newAnthropicResponsesStream(ctx, target, s.sourceModel) +} + +func (*responsesSession) ConvertError(_ context.Context, err error) error { return err } diff --git a/internal/protocol/stage/anthropicbridge/responses_stream.go b/internal/protocol/stage/anthropicbridge/responses_stream.go new file mode 100644 index 000000000..a611f453a --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/responses_stream.go @@ -0,0 +1,110 @@ +package anthropicbridge + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + + "github.com/openai/openai-go/v3/responses" + + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" +) + +func newAnthropicResponsesStream(ctx context.Context, target stage.EventStream, sourceModel string) (stage.EventStream, error) { + if target == nil { + return nil, fmt.Errorf("convert OpenAI Responses stream to Anthropic: target stream is nil") + } + iterator := &responsesStreamIterator{target: target, ctx: ctx} + return &anthropicResponsesStream{ + iterator: iterator, + converter: protocolstream.NewOpenAIResponsesToAnthropicConverter(ctx, iterator, sourceModel), + model: sourceModel, + }, nil +} + +type responsesStreamIterator struct { + target stage.EventStream + ctx context.Context + current responses.ResponseStreamEventUnion + err error + + closeOnce sync.Once + closeErr error +} + +func (s *responsesStreamIterator) setContext(ctx context.Context) { s.ctx = ctx } +func (s *responsesStreamIterator) Next() bool { + if s.err != nil { + return false + } + ctx := s.ctx + if ctx == nil { + ctx = context.Background() + } + event, err := s.target.Next(ctx) + if err != nil { + if !errors.Is(err, io.EOF) { + s.err = err + } + return false + } + switch value := event.Value.(type) { + case responses.ResponseStreamEventUnion: + s.current = value + case *responses.ResponseStreamEventUnion: + if value == nil { + s.err = fmt.Errorf("convert OpenAI Responses stream to Anthropic: event is nil") + return false + } + s.current = *value + default: + s.err = fmt.Errorf("convert OpenAI Responses stream to Anthropic: event has type %T", event.Value) + return false + } + return true +} +func (s *responsesStreamIterator) Current() responses.ResponseStreamEventUnion { return s.current } +func (s *responsesStreamIterator) Err() error { return s.err } +func (s *responsesStreamIterator) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.target.Close() }) + return s.closeErr +} + +type anthropicResponsesStream struct { + iterator *responsesStreamIterator + converter protocolstream.StreamConverter + model string +} + +func (s *anthropicResponsesStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + s.iterator.setContext(ctx) + value, done, err := s.converter.Next() + if err != nil { + return stage.Event{}, err + } + if done { + if err := s.iterator.Err(); err != nil { + return stage.Event{}, err + } + return stage.Event{}, io.EOF + } + event, ok := protocolstream.AsAnthropicEvent(value) + if !ok { + return stage.Event{}, fmt.Errorf("convert OpenAI Responses stream to Anthropic: converter emitted %T", value) + } + return stage.Event{Value: event}, nil +} +func (s *anthropicResponsesStream) Close() error { return s.iterator.Close() } +func (s *anthropicResponsesStream) Result() stage.StreamResult { + usage := s.converter.Usage() + if usage != nil && !usage.HasUsage() { + usage = nil + } + return stage.StreamResult{Usage: usage, Model: s.model} +} diff --git a/internal/protocol/stage/anthropicbridge/responses_test.go b/internal/protocol/stage/anthropicbridge/responses_test.go new file mode 100644 index 000000000..01107ad92 --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/responses_test.go @@ -0,0 +1,195 @@ +package anthropicbridge + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3/responses" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" +) + +func TestAnthropicBetaToOpenAIResponsesComplete(t *testing.T) { + t.Parallel() + + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIResponses, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + request, ok := call.Request.(*responses.ResponseNewParams) + if !ok || request == nil { + t.Fatalf("request type = %T", call.Request) + } + if request.Model != "provider-model" || call.State.OpenAIChat != nil { + t.Fatalf("target call = %#v state=%+v", request, call.State) + } + return &stage.Response{ + Value: decodeResponsesResponse(t, `{ + "id":"resp_1","object":"response","model":"provider-model","status":"completed", + "output":[{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello from responses","annotations":[]}]}], + "usage":{"input_tokens":9,"output_tokens":4,"total_tokens":13,"input_tokens_details":{"cached_tokens":2},"output_tokens_details":{"reasoning_tokens":0}} + }`), + SideEffectsCommitted: true, + }, nil + }, + } + adapted := mustAdapt(t, terminal, NewBetaToOpenAIResponses(ResponsesOptions{ResponseModel: "public-model"})) + result, err := adapted.Complete(context.Background(), stage.Call{Request: &anthropic.BetaMessageNewParams{ + Model: "provider-model", + MaxTokens: 321, + Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hello"))}, + }}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + message, ok := result.Value.(*anthropic.BetaMessage) + if !ok || message == nil { + t.Fatalf("response type = %T", result.Value) + } + if message.Model != "public-model" || len(message.Content) != 1 || !strings.Contains(message.Content[0].Text, "hello from responses") { + t.Fatalf("message = %#v", message) + } + if result.Usage == nil || result.Usage.InputTokens != 7 || result.Usage.CacheReadTokens != 2 || result.Usage.OutputTokens != 4 { + t.Fatalf("usage = %#v", result.Usage) + } + if !result.SideEffectsCommitted { + t.Fatal("side effects were not preserved") + } +} + +func TestAnthropicV1ToOpenAIResponsesComplete(t *testing.T) { + t.Parallel() + + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIResponses, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + request, ok := call.Request.(*responses.ResponseNewParams) + if !ok || request == nil { + t.Fatalf("request type = %T", call.Request) + } + if request.Model != "provider-v1-model" || call.Metadata.RequestID != "v1-responses" { + t.Fatalf("target call = %#v metadata=%+v", request, call.Metadata) + } + return &stage.Response{ + Value: decodeResponsesResponse(t, `{ + "id":"resp_v1","object":"response","model":"provider-v1-model","status":"completed", + "output":[{"id":"msg_v1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello v1 from responses","annotations":[]}]}], + "usage":{"input_tokens":7,"output_tokens":3,"total_tokens":10,"input_tokens_details":{"cached_tokens":1},"output_tokens_details":{"reasoning_tokens":0}} + }`), + SideEffectsCommitted: true, + }, nil + }, + } + adapted := mustAdapt(t, terminal, NewV1ToOpenAIResponses(ResponsesOptions{ResponseModel: "public-v1-model"})) + result, err := adapted.Complete(context.Background(), stage.Call{ + Request: &anthropic.MessageNewParams{ + Model: "provider-v1-model", + MaxTokens: 128, + Messages: []anthropic.MessageParam{anthropic.NewUserMessage(anthropic.NewTextBlock("hello"))}, + }, + Metadata: stage.CallMetadata{RequestID: "v1-responses"}, + }) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + message, ok := result.Value.(*anthropic.Message) + if !ok || message == nil { + t.Fatalf("response type = %T", result.Value) + } + if message.Model != "public-v1-model" || len(message.Content) != 1 || !strings.Contains(message.Content[0].Text, "hello v1 from responses") { + t.Fatalf("message = %#v", message) + } + if result.Usage == nil || result.Usage.InputTokens != 6 || result.Usage.CacheReadTokens != 1 || result.Usage.OutputTokens != 3 { + t.Fatalf("usage = %#v", result.Usage) + } + if !result.SideEffectsCommitted { + t.Fatal("side effects were not preserved") + } +} + +func TestAnthropicBetaToOpenAIResponsesStream(t *testing.T) { + t.Parallel() + + target := &memoryStream{ + events: responsesStageEvents(t, + `{"type":"response.created","sequence_number":0,"response":{"id":"resp_stream","object":"response","model":"provider-model","status":"in_progress","output":[]}}`, + `{"type":"response.output_text.delta","sequence_number":1,"item_id":"msg_1","output_index":0,"content_index":0,"delta":"stream responses"}`, + `{"type":"response.output_text.done","sequence_number":2,"item_id":"msg_1","output_index":0,"content_index":0,"text":"stream responses"}`, + `{"type":"response.completed","sequence_number":3,"response":{"id":"resp_stream","object":"response","model":"provider-model","status":"completed","output":[],"usage":{"input_tokens":6,"output_tokens":2,"total_tokens":8,"input_tokens_details":{"cached_tokens":0},"output_tokens_details":{"reasoning_tokens":0}}}}`, + ), + result: stage.StreamResult{SideEffectsCommitted: true}, + } + terminal := &memoryEndpoint{api: protocol.TypeOpenAIResponses, stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return target, nil + }} + adapted := mustAdapt(t, terminal, NewBetaToOpenAIResponses(ResponsesOptions{ResponseModel: "public-stream-model"})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &anthropic.BetaMessageNewParams{ + Model: "provider-model", + MaxTokens: 128, + }}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + var eventTypes []string + var sawText bool + for { + event, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + t.Fatalf("Next() error = %v", nextErr) + } + anthropicEvent, ok := event.Value.(protocolstream.AnthropicEvent) + if !ok { + t.Fatalf("event type = %T", event.Value) + } + eventTypes = append(eventTypes, anthropicEvent.Type) + encoded, _ := json.Marshal(anthropicEvent.Data) + if anthropicEvent.Type == "content_block_delta" && strings.Contains(string(encoded), "stream responses") { + sawText = true + } + } + if !sawText || len(eventTypes) == 0 || eventTypes[0] != "message_start" || eventTypes[len(eventTypes)-1] != "message_stop" { + t.Fatalf("events = %v, saw text = %v", eventTypes, sawText) + } + result := stream.Result() + if result.Model != "public-stream-model" || result.Usage == nil || result.Usage.InputTokens != 6 || result.Usage.OutputTokens != 2 || !result.SideEffectsCommitted { + t.Fatalf("Result() = %+v", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if target.closeCount != 1 { + t.Fatalf("target close count = %d", target.closeCount) + } +} + +func decodeResponsesResponse(t *testing.T, raw string) *responses.Response { + t.Helper() + var response responses.Response + if err := json.Unmarshal([]byte(raw), &response); err != nil { + t.Fatalf("decode Responses response: %v", err) + } + return &response +} + +func responsesStageEvents(t *testing.T, values ...string) []stage.Event { + t.Helper() + events := make([]stage.Event, 0, len(values)) + for _, raw := range values { + var event responses.ResponseStreamEventUnion + if err := json.Unmarshal([]byte(raw), &event); err != nil { + t.Fatalf("decode Responses event: %v", err) + } + events = append(events, stage.Event{Value: event}) + } + return events +} diff --git a/internal/protocol/stage/anthropicbridge/stream.go b/internal/protocol/stage/anthropicbridge/stream.go new file mode 100644 index 000000000..7b72f2e6b --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/stream.go @@ -0,0 +1,131 @@ +package anthropicbridge + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + + "github.com/openai/openai-go/v3" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" +) + +func newAnthropicStream( + target stage.EventStream, + source protocol.APIType, + sourceModel string, + targetRequest *openai.ChatCompletionNewParams, +) (stage.EventStream, error) { + if target == nil { + return nil, fmt.Errorf("convert OpenAI Chat stream to Anthropic: target stream is nil") + } + iterator := &chatStreamIterator{target: target} + var converter protocolstream.StreamConverter + switch source { + case protocol.TypeAnthropicV1: + converter = protocolstream.NewOpenAIChatToAnthropicV1Converter(iterator, sourceModel, targetRequest) + case protocol.TypeAnthropicBeta: + converter = protocolstream.NewOpenAIChatToAnthropicBetaConverter(iterator, sourceModel, targetRequest) + default: + return nil, fmt.Errorf("convert OpenAI Chat stream: unsupported Anthropic source protocol %q", source) + } + return &anthropicStream{ + iterator: iterator, + converter: converter, + model: sourceModel, + }, nil +} + +type chatStreamIterator struct { + target stage.EventStream + ctx context.Context + current openai.ChatCompletionChunk + err error + + closeOnce sync.Once + closeErr error +} + +func (s *chatStreamIterator) setContext(ctx context.Context) { s.ctx = ctx } + +func (s *chatStreamIterator) Next() bool { + if s.err != nil { + return false + } + ctx := s.ctx + if ctx == nil { + ctx = context.Background() + } + event, err := s.target.Next(ctx) + if err != nil { + if !errors.Is(err, io.EOF) { + s.err = err + } + return false + } + + switch value := event.Value.(type) { + case openai.ChatCompletionChunk: + s.current = value + case *openai.ChatCompletionChunk: + if value == nil { + s.err = fmt.Errorf("convert OpenAI Chat stream to Anthropic: chunk is nil") + return false + } + s.current = *value + default: + s.err = fmt.Errorf("convert OpenAI Chat stream to Anthropic: event has type %T, want openai.ChatCompletionChunk", event.Value) + return false + } + return true +} + +func (s *chatStreamIterator) Current() openai.ChatCompletionChunk { return s.current } + +func (s *chatStreamIterator) Err() error { return s.err } + +func (s *chatStreamIterator) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.target.Close() }) + return s.closeErr +} + +type anthropicStream struct { + iterator *chatStreamIterator + converter protocolstream.StreamConverter + model string +} + +func (s *anthropicStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + s.iterator.setContext(ctx) + value, done, err := s.converter.Next() + if err != nil { + return stage.Event{}, err + } + if done { + if err := s.iterator.Err(); err != nil { + return stage.Event{}, err + } + return stage.Event{}, io.EOF + } + event, ok := protocolstream.AsAnthropicEvent(value) + if !ok { + return stage.Event{}, fmt.Errorf("convert OpenAI Chat stream to Anthropic: converter emitted %T", value) + } + return stage.Event{Value: event}, nil +} + +func (s *anthropicStream) Close() error { return s.iterator.Close() } + +func (s *anthropicStream) Result() stage.StreamResult { + usage := s.converter.Usage() + if usage != nil && !usage.HasUsage() { + usage = nil + } + return stage.StreamResult{Usage: usage, Model: s.model} +} diff --git a/internal/protocol/stage/anthropicbridge/v1_beta.go b/internal/protocol/stage/anthropicbridge/v1_beta.go new file mode 100644 index 000000000..126d98c25 --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/v1_beta.go @@ -0,0 +1,219 @@ +package anthropicbridge + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + + "github.com/anthropics/anthropic-sdk-go" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" +) + +// NewV1ToBeta returns an Anthropic v1 -> Beta Bridge. V1 request promotion is +// the guaranteed compatibility direction because the V1 request is a Beta wire +// subset. Reverse response/event projection intentionally remains permissive; +// Beta-only output may not have equivalent V1 typed semantics. +func NewV1ToBeta() stage.Bridge { return v1ToBetaBridge{} } + +type v1ToBetaBridge struct{} + +func (v1ToBetaBridge) Source() protocol.APIType { return protocol.TypeAnthropicV1 } + +func (v1ToBetaBridge) Target() protocol.APIType { return protocol.TypeAnthropicBeta } + +func (v1ToBetaBridge) Capabilities() stage.Capabilities { return stage.AllBridgeCapabilities } + +func (v1ToBetaBridge) Open(_ context.Context, call stage.Call, operation stage.Operation) (stage.BridgeSession, error) { + if _, err := operationStreaming(operation); err != nil { + return nil, fmt.Errorf("open Anthropic v1 to Beta bridge: %w", err) + } + v1, err := v1BetaRequest(call.Request) + if err != nil { + return nil, err + } + beta, err := request.ConvertAnthropicV1ToBetaRequestWithError(v1) + if err != nil { + return nil, fmt.Errorf("open Anthropic v1 to Beta bridge: %w", err) + } + targetCall := call + targetCall.Request = beta + return &v1ToBetaSession{ + operation: operation, + targetCall: targetCall, + sourceModel: string(v1.Model), + }, nil +} + +func v1BetaRequest(value any) (*anthropic.MessageNewParams, error) { + switch typed := value.(type) { + case *anthropic.MessageNewParams: + if typed == nil { + return nil, errors.New("open Anthropic v1 to Beta bridge: request is nil") + } + return typed, nil + case anthropic.MessageNewParams: + return &typed, nil + default: + return nil, fmt.Errorf("open Anthropic v1 to Beta bridge: request has type %T, want anthropic.MessageNewParams", value) + } +} + +type v1ToBetaSession struct { + operation stage.Operation + targetCall stage.Call + sourceModel string +} + +func (s *v1ToBetaSession) TargetCall() stage.Call { return s.targetCall } + +// TODO: Add strict Beta-output subset validation only if a lossless V1 response +// contract becomes necessary. This phase intentionally leaves responses and +// stream events unconstrained and preserves the existing JSON projection. +func (s *v1ToBetaSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + if s.operation != stage.OperationComplete { + return nil, fmt.Errorf("convert Anthropic Beta response to v1: session was opened for %s", s.operation) + } + beta, err := betaMessage(response) + if err != nil { + return nil, err + } + v1, err := convertJSON[anthropic.Message](beta, "Anthropic Beta response to v1") + if err != nil { + return nil, err + } + model := response.Model + if s.sourceModel != "" { + v1.Model = s.sourceModel + model = s.sourceModel + } + return &stage.Response{ + Value: v1, + Usage: response.Usage, + Model: model, + SideEffectsCommitted: response.SideEffectsCommitted, + }, nil +} + +func betaMessage(response *stage.Response) (*anthropic.BetaMessage, error) { + if response == nil { + return nil, errors.New("convert Anthropic Beta response to v1: response is nil") + } + switch typed := response.Value.(type) { + case *anthropic.BetaMessage: + if typed == nil { + return nil, errors.New("convert Anthropic Beta response to v1: value is nil") + } + return typed, nil + case anthropic.BetaMessage: + return &typed, nil + default: + return nil, fmt.Errorf("convert Anthropic Beta response to v1: value has type %T, want anthropic.BetaMessage", response.Value) + } +} + +func (s *v1ToBetaSession) ConvertStream(_ context.Context, target stage.EventStream) (stage.EventStream, error) { + if s.operation != stage.OperationStream { + return nil, fmt.Errorf("convert Anthropic Beta stream to v1: session was opened for %s", s.operation) + } + if target == nil { + return nil, errors.New("convert Anthropic Beta stream to v1: target stream is nil") + } + return &v1BetaStream{target: target, sourceModel: s.sourceModel}, nil +} + +func (s *v1ToBetaSession) ConvertError(_ context.Context, err error) error { return err } + +type v1BetaStream struct { + target stage.EventStream + sourceModel string + + closeOnce sync.Once + closeErr error +} + +func (s *v1BetaStream) Next(ctx context.Context) (stage.Event, error) { + event, err := s.target.Next(ctx) + if err != nil { + return stage.Event{}, err + } + beta, err := betaStreamEvent(event.Value) + if err != nil { + return stage.Event{}, err + } + v1, err := convertJSON[anthropic.MessageStreamEventUnion](beta, "Anthropic Beta stream event to v1") + if err != nil { + return stage.Event{}, err + } + if v1.Type == "message_start" && s.sourceModel != "" { + v1.Message.Model = s.sourceModel + } + return stage.Event{Value: *v1}, nil +} + +func (s *v1BetaStream) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.target.Close() }) + return s.closeErr +} + +func (s *v1BetaStream) Result() stage.StreamResult { + result := s.target.Result() + if s.sourceModel != "" { + result.Model = s.sourceModel + } + return result +} + +func betaStreamEvent(value any) (*anthropic.BetaRawMessageStreamEventUnion, error) { + switch typed := value.(type) { + case anthropic.BetaRawMessageStreamEventUnion: + return &typed, nil + case *anthropic.BetaRawMessageStreamEventUnion: + if typed == nil { + return nil, errors.New("convert Anthropic Beta stream event to v1: event is nil") + } + return typed, nil + case json.RawMessage: + return decodeBetaStreamEvent(typed) + case []byte: + return decodeBetaStreamEvent(typed) + case interface{ RawJSON() string }: + raw := typed.RawJSON() + if raw != "" { + return decodeBetaStreamEvent([]byte(raw)) + } + case protocolstream.AnthropicEvent: + converted, err := convertJSON[anthropic.BetaRawMessageStreamEventUnion](typed.Data, "Anthropic Beta transport event") + if err != nil { + return nil, err + } + return converted, nil + } + return nil, fmt.Errorf("convert Anthropic Beta stream event to v1: event has type %T, want Anthropic Beta event or JSON", value) +} + +func decodeBetaStreamEvent(raw []byte) (*anthropic.BetaRawMessageStreamEventUnion, error) { + var event anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(raw, &event); err != nil { + return nil, fmt.Errorf("convert Anthropic Beta stream event to v1: decode %T: %w", raw, err) + } + return &event, nil +} + +func convertJSON[T any](value any, label string) (*T, error) { + data, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("convert %s: marshal: %w", label, err) + } + var converted T + if err := json.Unmarshal(data, &converted); err != nil { + return nil, fmt.Errorf("convert %s: unmarshal: %w", label, err) + } + return &converted, nil +} + +var _ stage.EventStream = (*v1BetaStream)(nil) diff --git a/internal/protocol/stage/anthropicbridge/v1_beta_test.go b/internal/protocol/stage/anthropicbridge/v1_beta_test.go new file mode 100644 index 000000000..0de4c5470 --- /dev/null +++ b/internal/protocol/stage/anthropicbridge/v1_beta_test.go @@ -0,0 +1,193 @@ +package anthropicbridge + +import ( + "context" + "encoding/json" + "errors" + "io" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" +) + +func TestV1ToBetaCompletePreservesWireAndFacts(t *testing.T) { + t.Parallel() + + usage := protocol.NewTokenUsage(9, 4) + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + beta, ok := call.Request.(*anthropic.BetaMessageNewParams) + if !ok || beta == nil { + t.Fatalf("request type = %T", call.Request) + } + if call.Metadata.RequestID != "v1-beta" || beta.Model != "client-model" || len(beta.Tools) != 1 { + t.Fatalf("target call = %#v metadata=%+v", beta, call.Metadata) + } + return &stage.Response{ + Value: decodeV1BetaMessage(t, `{ + "id":"msg_1","type":"message","role":"assistant","model":"claude-provider", + "content":[{"type":"tool_use","id":"tool-1","name":"lookup","input":{"city":"Paris"}}], + "stop_reason":"tool_use","stop_sequence":null, + "usage":{"input_tokens":9,"output_tokens":4} + }`), + Usage: usage, Model: "claude-provider", SideEffectsCommitted: true, + }, nil + }, + } + adapted := mustAdapt(t, terminal, NewV1ToBeta()) + v1 := decodeV1BetaRequest(t, `{ + "model":"client-model","max_tokens":128, + "messages":[{"role":"user","content":[{"type":"text","text":"weather"}]}], + "tools":[{"name":"lookup","input_schema":{"type":"object","properties":{"city":{"type":"string"}}}}] + }`) + response, err := adapted.Complete(context.Background(), stage.Call{ + Request: v1, Metadata: stage.CallMetadata{RequestID: "v1-beta", Attempt: 2}, + }) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + message, ok := response.Value.(*anthropic.Message) + if !ok || message == nil { + t.Fatalf("response type = %T", response.Value) + } + if message.StopReason != "tool_use" || len(message.Content) != 1 || message.Content[0].Type != "tool_use" { + t.Fatalf("response = %#v", message) + } + if message.Model != "client-model" || response.Usage != usage || response.Model != "client-model" || !response.SideEffectsCommitted { + t.Fatalf("response facts = %+v", response) + } +} + +func TestV1ToBetaStreamConvertsLifecycleAndOwnsTarget(t *testing.T) { + t.Parallel() + + wires := []string{ + `{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-provider","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":5,"output_tokens":0}}}`, + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}`, + `{"type":"content_block_stop","index":0}`, + `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}`, + `{"type":"message_stop"}`, + } + events := make([]stage.Event, 0, len(wires)) + for i, wire := range wires { + value := any(decodeV1BetaEvent(t, wire)) + if i%2 == 0 { + value = json.RawMessage(wire) + } + events = append(events, stage.Event{Value: value}) + } + usage := protocol.NewTokenUsage(5, 2) + target := &memoryStream{ + events: events, + result: stage.StreamResult{Usage: usage, Model: "claude-provider", SideEffectsCommitted: true}, + } + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(_ context.Context, call stage.Call) (stage.EventStream, error) { + if _, ok := call.Request.(*anthropic.BetaMessageNewParams); !ok { + t.Fatalf("request type = %T", call.Request) + } + return target, nil + }, + } + adapted := mustAdapt(t, terminal, NewV1ToBeta()) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &anthropic.MessageNewParams{Model: "claude-provider", MaxTokens: 32}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + for i, wire := range wires { + event, err := stream.Next(context.Background()) + if err != nil { + t.Fatalf("Next(%d) error = %v", i, err) + } + v1, ok := event.Value.(anthropic.MessageStreamEventUnion) + if !ok { + t.Fatalf("event %d type = %T", i, event.Value) + } + beta := decodeV1BetaEvent(t, wire) + if v1.Type != beta.Type || v1.Index != beta.Index { + t.Fatalf("event %d identity = (%q,%d), want (%q,%d)", i, v1.Type, v1.Index, beta.Type, beta.Index) + } + switch v1.Type { + case "message_start": + if v1.Message.ID != "msg_1" || v1.Message.Model != "claude-provider" || v1.Message.Usage.InputTokens != 5 { + t.Fatalf("message_start = %#v", v1.Message) + } + case "content_block_start": + if v1.ContentBlock.Type != "text" { + t.Fatalf("content_block_start = %#v", v1.ContentBlock) + } + case "content_block_delta": + if v1.Delta.Type != "text_delta" || v1.Delta.Text != "hello" { + t.Fatalf("content_block_delta = %#v", v1.Delta) + } + case "message_delta": + if v1.Delta.StopReason != "end_turn" || v1.Usage.OutputTokens != 2 { + t.Fatalf("message_delta = delta=%#v usage=%#v", v1.Delta, v1.Usage) + } + } + } + if _, err := stream.Next(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("final Next() error = %v", err) + } + if got := stream.Result(); got.Usage != usage || got.Model != "claude-provider" || !got.SideEffectsCommitted { + t.Fatalf("Result() = %+v", got) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("second Close() error = %v", err) + } + if target.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", target.closeCount) + } +} + +func TestV1ToBetaRejectsWrongTypes(t *testing.T) { + t.Parallel() + + bridge := NewV1ToBeta() + if _, err := bridge.Open(context.Background(), stage.Call{Request: &anthropic.BetaMessageNewParams{}}, stage.OperationComplete); err == nil { + t.Fatal("Open() accepted a Beta request on the v1 side") + } + + session, err := bridge.Open(context.Background(), stage.Call{Request: &anthropic.MessageNewParams{}}, stage.OperationComplete) + if err != nil { + t.Fatalf("Open() error = %v", err) + } + if _, err := session.ConvertComplete(context.Background(), &stage.Response{Value: &anthropic.Message{}}); err == nil { + t.Fatal("ConvertComplete() accepted a v1 response on the Beta side") + } +} + +func decodeV1BetaRequest(t *testing.T, raw string) *anthropic.MessageNewParams { + t.Helper() + var value anthropic.MessageNewParams + if err := json.Unmarshal([]byte(raw), &value); err != nil { + t.Fatalf("decode v1 request: %v", err) + } + return &value +} + +func decodeV1BetaMessage(t *testing.T, raw string) *anthropic.BetaMessage { + t.Helper() + var value anthropic.BetaMessage + if err := json.Unmarshal([]byte(raw), &value); err != nil { + t.Fatalf("decode Beta message: %v", err) + } + return &value +} + +func decodeV1BetaEvent(t *testing.T, raw string) anthropic.BetaRawMessageStreamEventUnion { + t.Helper() + var value anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal([]byte(raw), &value); err != nil { + t.Fatalf("decode Beta event: %v", err) + } + return value +} diff --git a/internal/protocol/stage/bridge.go b/internal/protocol/stage/bridge.go new file mode 100644 index 000000000..3fd5b8b16 --- /dev/null +++ b/internal/protocol/stage/bridge.go @@ -0,0 +1,246 @@ +package stage + +import ( + "context" + "errors" + "fmt" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +// Bridge describes an immutable, concurrency-safe bidirectional protocol +// adapter. Open converts one source Call to the target protocol and creates all +// mutable response/stream correlation state for that call. +type Bridge interface { + Source() protocol.APIType + Target() protocol.APIType + Capabilities() Capabilities + Open(ctx context.Context, call Call, operation Operation) (BridgeSession, error) +} + +// Operation identifies which endpoint operation the Bridge session is being +// opened for. Request conversion may legitimately differ between complete and +// streaming calls (for example, stream flags and usage options). +type Operation uint8 + +const ( + OperationComplete Operation = iota + 1 + OperationStream +) + +func (o Operation) String() string { + switch o { + case OperationComplete: + return "complete" + case OperationStream: + return "stream" + default: + return fmt.Sprintf("unknown(%d)", o) + } +} + +// BridgeSession is the per-call reverse path created while converting a request +// inward. A session is used for exactly one Complete or Stream invocation. +// +// ConvertStream must return a source-protocol stream that owns the target stream: +// it converts runtime Next errors and closes the target stream from Close. +type BridgeSession interface { + TargetCall() Call + ConvertComplete(ctx context.Context, response *Response) (*Response, error) + ConvertStream(ctx context.Context, stream EventStream) (EventStream, error) + ConvertError(ctx context.Context, err error) error +} + +// Adapt exposes next in bridge.Source() while calling next in bridge.Target(). +// It validates the structural and core-capability boundary without executing a +// request. +func Adapt(next Endpoint, bridge Bridge) (Endpoint, error) { + if isNil(next) { + return nil, fmt.Errorf("adapt protocol bridge: target endpoint is nil") + } + if isNil(bridge) { + return nil, fmt.Errorf("adapt protocol bridge: bridge is nil") + } + + source := bridge.Source() + if source == "" { + return nil, fmt.Errorf("adapt protocol bridge: bridge has empty source protocol") + } + target := bridge.Target() + if target == "" { + return nil, fmt.Errorf("adapt protocol bridge %q -> ?: bridge has empty target protocol", source) + } + if next.Protocol() == "" { + return nil, fmt.Errorf("adapt protocol bridge %q -> %q: target endpoint has empty protocol", source, target) + } + if next.Protocol() != target { + return nil, fmt.Errorf( + "adapt protocol bridge %q -> %q: cannot call endpoint speaking %q", + source, + target, + next.Protocol(), + ) + } + if missing := bridge.Capabilities().Missing(CoreBridgeCapabilities); missing != 0 { + return nil, fmt.Errorf( + "adapt protocol bridge %q -> %q: missing core capabilities: %s", + source, + target, + missing, + ) + } + + return &bridgeEndpoint{ + next: next, + bridge: bridge, + source: source, + target: target, + }, nil +} + +type bridgeEndpoint struct { + next Endpoint + bridge Bridge + source protocol.APIType + target protocol.APIType +} + +func (e *bridgeEndpoint) Protocol() protocol.APIType { + return e.source +} + +func (e *bridgeEndpoint) Complete(ctx context.Context, call Call) (*Response, error) { + session, targetCall, err := e.open(ctx, call, OperationComplete) + if err != nil { + return nil, err + } + + response, err := e.next.Complete(ctx, targetCall) + if err != nil { + return nil, e.convertError(ctx, session, err) + } + if response == nil { + return nil, fmt.Errorf("protocol bridge %q -> %q: target endpoint returned a nil response", e.source, e.target) + } + + converted, err := session.ConvertComplete(ctx, response) + if err != nil { + return nil, err + } + if converted == nil { + return nil, fmt.Errorf("protocol bridge %q -> %q: session returned a nil converted response", e.source, e.target) + } + + result := *converted + mergeResponseFacts(&result, response) + return &result, nil +} + +func (e *bridgeEndpoint) Stream(ctx context.Context, call Call) (EventStream, error) { + session, targetCall, err := e.open(ctx, call, OperationStream) + if err != nil { + return nil, err + } + + targetStream, err := e.next.Stream(ctx, targetCall) + if err != nil { + return nil, e.convertError(ctx, session, err) + } + if isNil(targetStream) { + return nil, fmt.Errorf("protocol bridge %q -> %q: target endpoint returned a nil stream", e.source, e.target) + } + + converted, err := session.ConvertStream(ctx, targetStream) + if err != nil { + return nil, closeAfterConversionFailure(targetStream, err, e.source, e.target) + } + if isNil(converted) { + err := fmt.Errorf("protocol bridge %q -> %q: session returned a nil converted stream", e.source, e.target) + return nil, closeAfterConversionFailure(targetStream, err, e.source, e.target) + } + + return &factPreservingStream{ + converted: converted, + target: targetStream, + }, nil +} + +func (e *bridgeEndpoint) open(ctx context.Context, call Call, operation Operation) (BridgeSession, Call, error) { + session, err := e.bridge.Open(ctx, call, operation) + if err != nil { + return nil, Call{}, err + } + if isNil(session) { + return nil, Call{}, fmt.Errorf("protocol bridge %q -> %q: Open returned a nil session", e.source, e.target) + } + + targetCall := session.TargetCall() + // Protocol conversion must not erase attempt identity. Any future metadata + // transformation needs an explicit field and policy rather than a hidden + // bridge-local mutation. + targetCall.Metadata = call.Metadata + return session, targetCall, nil +} + +func (e *bridgeEndpoint) convertError(ctx context.Context, session BridgeSession, targetErr error) error { + converted := session.ConvertError(ctx, targetErr) + if converted != nil { + return converted + } + return fmt.Errorf( + "protocol bridge %q -> %q swallowed target error: %w", + e.source, + e.target, + targetErr, + ) +} + +func mergeResponseFacts(converted, target *Response) { + if converted.Usage == nil { + converted.Usage = target.Usage + } + if converted.Model == "" { + converted.Model = target.Model + } + converted.SideEffectsCommitted = converted.SideEffectsCommitted || target.SideEffectsCommitted +} + +func mergeStreamFacts(converted, target StreamResult) StreamResult { + if converted.Usage == nil { + converted.Usage = target.Usage + } + if converted.Model == "" { + converted.Model = target.Model + } + converted.SideEffectsCommitted = converted.SideEffectsCommitted || target.SideEffectsCommitted + return converted +} + +func closeAfterConversionFailure(stream EventStream, conversionErr error, source, target protocol.APIType) error { + if closeErr := stream.Close(); closeErr != nil { + return errors.Join( + conversionErr, + fmt.Errorf("close target stream after bridge %q -> %q conversion failure: %w", source, target, closeErr), + ) + } + return conversionErr +} + +// factPreservingStream keeps protocol-neutral facts monotonic while delegating +// event conversion and target-stream ownership to the BridgeSession's stream. +type factPreservingStream struct { + converted EventStream + target EventStream +} + +func (s *factPreservingStream) Next(ctx context.Context) (Event, error) { + return s.converted.Next(ctx) +} + +func (s *factPreservingStream) Close() error { + return s.converted.Close() +} + +func (s *factPreservingStream) Result() StreamResult { + return mergeStreamFacts(s.converted.Result(), s.target.Result()) +} diff --git a/internal/protocol/stage/bridge_test.go b/internal/protocol/stage/bridge_test.go new file mode 100644 index 000000000..e08888a61 --- /dev/null +++ b/internal/protocol/stage/bridge_test.go @@ -0,0 +1,699 @@ +package stage + +import ( + "context" + "errors" + "fmt" + "io" + "reflect" + "strings" + "testing" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +func TestBuildTopologyCompleteAndStreamFlow(t *testing.T) { + t.Parallel() + + var calls []string + usage := protocol.NewTokenUsage(17, 9) + terminalStream := &recordingEventStream{ + calls: &calls, + events: []Event{{Value: "terminal event"}}, + result: StreamResult{ + Usage: usage, + Model: "provider-model", + SideEffectsCommitted: true, + }, + } + terminal := &recordingEndpoint{ + protocol: protocol.TypeOpenAIResponses, + calls: &calls, + response: &Response{ + Value: "terminal response", + Usage: usage, + Model: "provider-model", + SideEffectsCommitted: true, + }, + stream: terminalStream, + } + providerBridge := &testingBridge{ + name: "provider_bridge", + source: protocol.TypeAnthropicBeta, + target: protocol.TypeOpenAIResponses, + caps: AllBridgeCapabilities, + calls: &calls, + dropFacts: true, + } + ingressBridge := &testingBridge{ + name: "ingress_bridge", + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: AllBridgeCapabilities, + calls: &calls, + dropFacts: true, + } + + registry, err := NewBridgeRegistry(providerBridge, ingressBridge) + if err != nil { + t.Fatalf("NewBridgeRegistry() error = %v", err) + } + chain, err := BuildTopology(TopologyConfig{ + Terminal: terminal, + Stages: []Stage{ + &recordingStage{name: "guardrails", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + &recordingStage{name: "tool_loop", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + }, + ClientProtocol: protocol.TypeOpenAIChat, + Registry: registry, + RequiredCapabilities: CapabilityUsage | + CapabilityToolUse | + CapabilityToolResult, + }) + if err != nil { + t.Fatalf("BuildTopology() error = %v", err) + } + if chain.Protocol() != protocol.TypeOpenAIChat { + t.Fatalf("chain.Protocol() = %q, want %q", chain.Protocol(), protocol.TypeOpenAIChat) + } + if len(calls) != 0 { + t.Fatalf("BuildTopology() executed chain, calls = %v", calls) + } + + call := Call{ + Request: "client request", + Metadata: CallMetadata{ + RequestID: "req-chain", + Attempt: 3, + }, + } + response, err := chain.Complete(context.Background(), call) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if response.Value != "ingress_bridge(provider_bridge(terminal response))" { + t.Fatalf("response.Value = %v", response.Value) + } + assertResponseFacts(t, response, usage, "provider-model", true) + if terminal.lastCall.Metadata != call.Metadata { + t.Fatalf("terminal metadata = %+v, want %+v", terminal.lastCall.Metadata, call.Metadata) + } + + wantCompleteCalls := []string{ + "ingress_bridge:request", + "guardrails:request", + "tool_loop:request", + "provider_bridge:request", + "terminal:request", + "terminal:response", + "provider_bridge:response", + "tool_loop:response", + "guardrails:response", + "ingress_bridge:response", + } + if !reflect.DeepEqual(calls, wantCompleteCalls) { + t.Fatalf("complete calls = %v, want %v", calls, wantCompleteCalls) + } + + calls = nil + stream, err := chain.Stream(context.Background(), call) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + event, err := stream.Next(context.Background()) + if err != nil { + t.Fatalf("Next() error = %v", err) + } + if event.Value != "ingress_bridge(provider_bridge(terminal event))" { + t.Fatalf("event.Value = %v", event.Value) + } + _, err = stream.Next(context.Background()) + if !errors.Is(err, io.EOF) { + t.Fatalf("second Next() error = %v, want io.EOF", err) + } + result := stream.Result() + if result.Usage != usage || result.Model != "provider-model" || !result.SideEffectsCommitted { + t.Fatalf("Result() = %+v, want preserved terminal facts", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if terminalStream.closeCount != 1 { + t.Fatalf("terminal close count = %d, want 1", terminalStream.closeCount) + } + + wantStreamCalls := []string{ + "ingress_bridge:request", + "guardrails:stream_request", + "tool_loop:stream_request", + "provider_bridge:request", + "terminal:stream_request", + "terminal:event", + "provider_bridge:event", + "tool_loop:event", + "guardrails:event", + "ingress_bridge:event", + "terminal:eof", + "provider_bridge:eof", + "tool_loop:eof", + "guardrails:eof", + "ingress_bridge:eof", + "ingress_bridge:close", + "guardrails:close", + "tool_loop:close", + "provider_bridge:close", + "terminal:close", + } + if !reflect.DeepEqual(calls, wantStreamCalls) { + t.Fatalf("stream calls = %v, want %v", calls, wantStreamCalls) + } + + if providerBridge.openCount != 2 || ingressBridge.openCount != 2 { + t.Fatalf("bridge opens: provider=%d ingress=%d, want 2 each", providerBridge.openCount, ingressBridge.openCount) + } + wantOperations := []Operation{OperationComplete, OperationStream} + if !reflect.DeepEqual(providerBridge.operations, wantOperations) || !reflect.DeepEqual(ingressBridge.operations, wantOperations) { + t.Fatalf("bridge operations: provider=%v ingress=%v, want %v", providerBridge.operations, ingressBridge.operations, wantOperations) + } + if providerBridge.sessions[0] == providerBridge.sessions[1] || ingressBridge.sessions[0] == ingressBridge.sessions[1] { + t.Fatal("Bridge.Open() reused a session across calls") + } +} + +func TestAdaptRejectsInvalidBoundary(t *testing.T) { + t.Parallel() + + validEndpoint := &recordingEndpoint{protocol: protocol.TypeAnthropicBeta} + validBridge := func() *testingBridge { + return &testingBridge{ + name: "bridge", + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: AllBridgeCapabilities, + } + } + + var typedNilEndpoint *recordingEndpoint + var typedNilBridge *testingBridge + tests := []struct { + name string + next Endpoint + bridge Bridge + want string + }{ + {name: "nil endpoint", bridge: validBridge(), want: "target endpoint is nil"}, + {name: "typed nil endpoint", next: typedNilEndpoint, bridge: validBridge(), want: "target endpoint is nil"}, + {name: "nil bridge", next: validEndpoint, want: "bridge is nil"}, + {name: "typed nil bridge", next: validEndpoint, bridge: typedNilBridge, want: "bridge is nil"}, + { + name: "empty source", + next: validEndpoint, + bridge: &testingBridge{target: protocol.TypeAnthropicBeta, caps: AllBridgeCapabilities}, + want: "empty source protocol", + }, + { + name: "empty target", + next: validEndpoint, + bridge: &testingBridge{source: protocol.TypeOpenAIChat, caps: AllBridgeCapabilities}, + want: "empty target protocol", + }, + { + name: "empty endpoint protocol", + next: &recordingEndpoint{}, + bridge: validBridge(), + want: "target endpoint has empty protocol", + }, + { + name: "target mismatch", + next: &recordingEndpoint{protocol: protocol.TypeOpenAIResponses}, + bridge: &testingBridge{ + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: AllBridgeCapabilities, + }, + want: `cannot call endpoint speaking "openai_responses"`, + }, + { + name: "missing core capability", + next: validEndpoint, + bridge: &testingBridge{ + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: CapabilityComplete | CapabilityError, + }, + want: "missing core capabilities: stream", + }, + {name: "valid baseline", next: validEndpoint, bridge: validBridge()}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := Adapt(tt.next, tt.bridge) + if tt.want == "" { + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + if got == nil { + t.Fatal("Adapt() returned nil endpoint") + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("Adapt() error = %v, want containing %q", err, tt.want) + } + }) + } +} + +func TestAdaptRuntimeFailures(t *testing.T) { + t.Parallel() + + upstreamErr := errors.New("upstream failed") + conversionErr := errors.New("stream conversion failed") + + t.Run("request conversion error", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta} + bridge := validTestingBridge() + bridge.openErr = errors.New("request conversion failed") + adapted := mustAdapt(t, endpoint, bridge) + + _, err := adapted.Complete(context.Background(), Call{}) + if !errors.Is(err, bridge.openErr) { + t.Fatalf("Complete() error = %v", err) + } + if endpoint.completeCalls != 0 { + t.Fatal("target endpoint executed after request conversion error") + } + }) + + t.Run("nil session", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta} + bridge := validTestingBridge() + bridge.nilSession = true + adapted := mustAdapt(t, endpoint, bridge) + + _, err := adapted.Complete(context.Background(), Call{}) + if err == nil || !strings.Contains(err.Error(), "Open returned a nil session") { + t.Fatalf("Complete() error = %v", err) + } + }) + + t.Run("target error converted", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta, completeErr: upstreamErr} + bridge := validTestingBridge() + adapted := mustAdapt(t, endpoint, bridge) + + _, err := adapted.Complete(context.Background(), Call{}) + if !errors.Is(err, upstreamErr) || !strings.Contains(err.Error(), "bridge: ") { + t.Fatalf("Complete() error = %v", err) + } + }) + + t.Run("target error cannot be swallowed", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta, completeErr: upstreamErr} + bridge := validTestingBridge() + bridge.swallowError = true + adapted := mustAdapt(t, endpoint, bridge) + + _, err := adapted.Complete(context.Background(), Call{}) + if !errors.Is(err, upstreamErr) || !strings.Contains(err.Error(), "swallowed target error") { + t.Fatalf("Complete() error = %v", err) + } + }) + + t.Run("nil target response", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta} + adapted := mustAdapt(t, endpoint, validTestingBridge()) + + _, err := adapted.Complete(context.Background(), Call{}) + if err == nil || !strings.Contains(err.Error(), "target endpoint returned a nil response") { + t.Fatalf("Complete() error = %v", err) + } + }) + + t.Run("nil converted response", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta, response: &Response{Value: "ok"}} + bridge := validTestingBridge() + bridge.nilResponse = true + adapted := mustAdapt(t, endpoint, bridge) + + _, err := adapted.Complete(context.Background(), Call{}) + if err == nil || !strings.Contains(err.Error(), "nil converted response") { + t.Fatalf("Complete() error = %v", err) + } + }) + + t.Run("target stream open error converted", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta, streamErr: upstreamErr} + adapted := mustAdapt(t, endpoint, validTestingBridge()) + + _, err := adapted.Stream(context.Background(), Call{}) + if !errors.Is(err, upstreamErr) || !strings.Contains(err.Error(), "bridge: ") { + t.Fatalf("Stream() error = %v", err) + } + }) + + t.Run("nil target stream", func(t *testing.T) { + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta} + adapted := mustAdapt(t, endpoint, validTestingBridge()) + + _, err := adapted.Stream(context.Background(), Call{}) + if err == nil || !strings.Contains(err.Error(), "target endpoint returned a nil stream") { + t.Fatalf("Stream() error = %v", err) + } + }) + + t.Run("stream conversion error closes target", func(t *testing.T) { + targetStream := &recordingEventStream{} + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta, stream: targetStream} + bridge := validTestingBridge() + bridge.streamConversionErr = conversionErr + adapted := mustAdapt(t, endpoint, bridge) + + _, err := adapted.Stream(context.Background(), Call{}) + if !errors.Is(err, conversionErr) { + t.Fatalf("Stream() error = %v", err) + } + if targetStream.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", targetStream.closeCount) + } + }) + + t.Run("nil converted stream closes target", func(t *testing.T) { + targetStream := &recordingEventStream{} + endpoint := &failureEndpoint{protocol: protocol.TypeAnthropicBeta, stream: targetStream} + bridge := validTestingBridge() + bridge.nilStream = true + adapted := mustAdapt(t, endpoint, bridge) + + _, err := adapted.Stream(context.Background(), Call{}) + if err == nil || !strings.Contains(err.Error(), "nil converted stream") { + t.Fatalf("Stream() error = %v", err) + } + if targetStream.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", targetStream.closeCount) + } + }) +} + +func TestIdentityBridgePreservesValues(t *testing.T) { + t.Parallel() + + usage := protocol.NewTokenUsage(5, 2) + targetStream := &recordingEventStream{ + events: []Event{{Value: "event"}}, + result: StreamResult{Usage: usage, Model: "model", SideEffectsCommitted: true}, + } + terminal := &recordingEndpoint{ + protocol: protocol.TypeAnthropicBeta, + response: &Response{Value: "response", Usage: usage, Model: "model", SideEffectsCommitted: true}, + stream: targetStream, + } + adapted := mustAdapt(t, terminal, NewIdentityBridge(protocol.TypeAnthropicBeta)) + + config := &protocol.OpenAIConfig{HasThinking: true} + call := Call{Request: "request", State: ProtocolState{OpenAIChat: config}} + response, err := adapted.Complete(context.Background(), call) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if response.Value != "response" { + t.Fatalf("response.Value = %v", response.Value) + } + assertResponseFacts(t, response, usage, "model", true) + + if terminal.lastCall.State.OpenAIChat != config { + t.Fatal("identity complete call did not preserve protocol state") + } + + stream, err := adapted.Stream(context.Background(), call) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + event, err := stream.Next(context.Background()) + if err != nil || event.Value != "event" { + t.Fatalf("Next() = (%+v, %v)", event, err) + } + if got := stream.Result(); got.Usage != usage || got.Model != "model" || !got.SideEffectsCommitted { + t.Fatalf("Result() = %+v", got) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if targetStream.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", targetStream.closeCount) + } + if terminal.lastCall.State.OpenAIChat != config { + t.Fatal("identity stream call did not preserve protocol state") + } +} + +func TestOperationString(t *testing.T) { + t.Parallel() + + if OperationComplete.String() != "complete" || OperationStream.String() != "stream" { + t.Fatalf("operation strings = %q, %q", OperationComplete, OperationStream) + } + if got := Operation(99).String(); got != "unknown(99)" { + t.Fatalf("unknown operation string = %q", got) + } +} + +func TestCapabilities(t *testing.T) { + t.Parallel() + + got := CapabilityComplete | CapabilityUsage | CapabilityToolUse + if !got.Supports(CapabilityComplete | CapabilityUsage) { + t.Fatal("Capabilities.Supports() = false for contained set") + } + if got.Supports(CapabilityStream) { + t.Fatal("Capabilities.Supports() = true for missing capability") + } + if missing := got.Missing(CapabilityComplete | CapabilityStream | CapabilityToolResult); missing != CapabilityStream|CapabilityToolResult { + t.Fatalf("Missing() = %v", missing) + } + if got.String() != "complete,usage,tool_use" { + t.Fatalf("String() = %q", got) + } + if Capabilities(0).String() != "none" { + t.Fatalf("zero String() = %q", Capabilities(0)) + } + unknown := Capabilities(1 << 20) + if unknown.String() != "unknown(0x100000)" { + t.Fatalf("unknown String() = %q", unknown) + } +} + +func TestBuildTopologyRejectsMissingSemanticCapability(t *testing.T) { + t.Parallel() + + terminal := &recordingEndpoint{protocol: protocol.TypeOpenAIResponses} + provider := &testingBridge{ + name: "provider", + source: protocol.TypeAnthropicBeta, + target: protocol.TypeOpenAIResponses, + caps: CoreBridgeCapabilities | CapabilityUsage, + } + ingress := &testingBridge{ + name: "ingress", + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: AllBridgeCapabilities, + } + + registry, err := NewBridgeRegistry(provider, ingress) + if err != nil { + t.Fatalf("NewBridgeRegistry() error = %v", err) + } + _, err = BuildTopology(TopologyConfig{ + Terminal: terminal, + Stages: []Stage{ + &recordingStage{name: "guardrails", protocol: protocol.TypeAnthropicBeta}, + }, + ClientProtocol: protocol.TypeOpenAIChat, + Registry: registry, + RequiredCapabilities: CapabilityToolUse, + }) + if err == nil || !strings.Contains(err.Error(), "bridge below stage") || !strings.Contains(err.Error(), "tool_use") { + t.Fatalf("BuildTopology() error = %v", err) + } +} + +func assertResponseFacts(t *testing.T, response *Response, usage *protocol.TokenUsage, model string, committed bool) { + t.Helper() + if response.Usage != usage || response.Model != model || response.SideEffectsCommitted != committed { + t.Fatalf("response facts = %+v, want usage=%p model=%q committed=%v", response, usage, model, committed) + } +} + +func mustAdapt(t *testing.T, endpoint Endpoint, bridge Bridge) Endpoint { + t.Helper() + adapted, err := Adapt(endpoint, bridge) + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + return adapted +} + +func validTestingBridge() *testingBridge { + return &testingBridge{ + name: "bridge", + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: AllBridgeCapabilities, + } +} + +type testingBridge struct { + name string + source protocol.APIType + target protocol.APIType + caps Capabilities + calls *[]string + dropFacts bool + openErr error + nilSession bool + swallowError bool + nilResponse bool + streamConversionErr error + nilStream bool + openCount int + sessions []*testingBridgeSession + operations []Operation +} + +func (b *testingBridge) Source() protocol.APIType { + return b.source +} + +func (b *testingBridge) Target() protocol.APIType { + return b.target +} + +func (b *testingBridge) Capabilities() Capabilities { + return b.caps +} + +func (b *testingBridge) Open(_ context.Context, call Call, operation Operation) (BridgeSession, error) { + if b.openErr != nil { + return nil, b.openErr + } + b.append(b.name + ":request") + b.operations = append(b.operations, operation) + b.openCount++ + if b.nilSession { + return nil, nil + } + + session := &testingBridgeSession{ + bridge: b, + call: Call{ + Request: fmt.Sprintf("%s(%v)", b.name, call.Request), + // Metadata is intentionally omitted. Adapt must restore it. + }, + } + b.sessions = append(b.sessions, session) + return session, nil +} + +func (b *testingBridge) append(value string) { + if b.calls != nil { + *b.calls = append(*b.calls, value) + } +} + +type testingBridgeSession struct { + bridge *testingBridge + call Call +} + +func (s *testingBridgeSession) TargetCall() Call { + return s.call +} + +func (s *testingBridgeSession) ConvertComplete(_ context.Context, response *Response) (*Response, error) { + s.bridge.append(s.bridge.name + ":response") + if s.bridge.nilResponse { + return nil, nil + } + converted := &Response{Value: fmt.Sprintf("%s(%v)", s.bridge.name, response.Value)} + if !s.bridge.dropFacts { + converted.Usage = response.Usage + converted.Model = response.Model + converted.SideEffectsCommitted = response.SideEffectsCommitted + } + return converted, nil +} + +func (s *testingBridgeSession) ConvertStream(_ context.Context, stream EventStream) (EventStream, error) { + if s.bridge.streamConversionErr != nil { + return nil, s.bridge.streamConversionErr + } + if s.bridge.nilStream { + return nil, nil + } + return &testingBridgeStream{bridge: s.bridge, target: stream}, nil +} + +func (s *testingBridgeSession) ConvertError(_ context.Context, err error) error { + if s.bridge.swallowError { + return nil + } + return fmt.Errorf("%s: %w", s.bridge.name, err) +} + +type testingBridgeStream struct { + bridge *testingBridge + target EventStream +} + +func (s *testingBridgeStream) Next(ctx context.Context) (Event, error) { + event, err := s.target.Next(ctx) + switch { + case err == nil: + s.bridge.append(s.bridge.name + ":event") + event.Value = fmt.Sprintf("%s(%v)", s.bridge.name, event.Value) + case errors.Is(err, io.EOF): + s.bridge.append(s.bridge.name + ":eof") + default: + err = fmt.Errorf("%s: %w", s.bridge.name, err) + } + return event, err +} + +func (s *testingBridgeStream) Close() error { + s.bridge.append(s.bridge.name + ":close") + return s.target.Close() +} + +func (s *testingBridgeStream) Result() StreamResult { + if s.bridge.dropFacts { + return StreamResult{} + } + return s.target.Result() +} + +type failureEndpoint struct { + protocol protocol.APIType + response *Response + completeErr error + stream EventStream + streamErr error + completeCalls int + streamCalls int +} + +func (e *failureEndpoint) Protocol() protocol.APIType { + return e.protocol +} + +func (e *failureEndpoint) Complete(_ context.Context, _ Call) (*Response, error) { + e.completeCalls++ + return e.response, e.completeErr +} + +func (e *failureEndpoint) Stream(_ context.Context, _ Call) (EventStream, error) { + e.streamCalls++ + return e.stream, e.streamErr +} diff --git a/internal/protocol/stage/capabilities.go b/internal/protocol/stage/capabilities.go new file mode 100644 index 000000000..6d8fe4eb5 --- /dev/null +++ b/internal/protocol/stage/capabilities.go @@ -0,0 +1,76 @@ +package stage + +import ( + "fmt" + "strings" +) + +// Capabilities is a deterministic bit set describing which semantic surfaces a +// Bridge preserves. Complete, stream, and error support are mandatory for every +// Bridge adapted as an Endpoint; a chain may require additional capabilities. +type Capabilities uint64 + +const ( + CapabilityComplete Capabilities = 1 << iota + CapabilityStream + CapabilityError + CapabilityUsage + CapabilityFinishReason + CapabilityToolUse + CapabilityToolResult +) + +// CoreBridgeCapabilities are required for every Bridge used by Adapt. +const CoreBridgeCapabilities = CapabilityComplete | CapabilityStream | CapabilityError + +// AllBridgeCapabilities contains every capability currently understood by the +// stage package. Identity bridges support this complete set. +const AllBridgeCapabilities = CoreBridgeCapabilities | + CapabilityUsage | + CapabilityFinishReason | + CapabilityToolUse | + CapabilityToolResult + +var orderedCapabilities = []struct { + capability Capabilities + name string +}{ + {CapabilityComplete, "complete"}, + {CapabilityStream, "stream"}, + {CapabilityError, "error"}, + {CapabilityUsage, "usage"}, + {CapabilityFinishReason, "finish_reason"}, + {CapabilityToolUse, "tool_use"}, + {CapabilityToolResult, "tool_result"}, +} + +// Supports reports whether c contains every required capability. +func (c Capabilities) Supports(required Capabilities) bool { + return c&required == required +} + +// Missing returns the required capabilities not present in c. +func (c Capabilities) Missing(required Capabilities) Capabilities { + return required &^ c +} + +// String returns stable, comma-separated concrete capability names. +func (c Capabilities) String() string { + if c == 0 { + return "none" + } + + remaining := c + names := make([]string, 0, len(orderedCapabilities)+1) + for _, item := range orderedCapabilities { + if c&item.capability == 0 { + continue + } + names = append(names, item.name) + remaining &^= item.capability + } + if remaining != 0 { + names = append(names, fmt.Sprintf("unknown(%#x)", uint64(remaining))) + } + return strings.Join(names, ",") +} diff --git a/internal/protocol/stage/compose.go b/internal/protocol/stage/compose.go new file mode 100644 index 000000000..697d91b1f --- /dev/null +++ b/internal/protocol/stage/compose.go @@ -0,0 +1,88 @@ +package stage + +import ( + "fmt" + "reflect" + "strings" +) + +// Compose wraps terminal with stages written in request order, from outermost +// to innermost. For example: +// +// Compose(provider, guardrails, tools) +// +// produces guardrails(tools(provider)). Responses and stream events naturally +// return through tools and then guardrails. +// +// Compose performs structural validation only. It never invokes Complete or +// Stream and it never inserts an implicit protocol conversion. +func Compose(terminal Endpoint, stages ...Stage) (Endpoint, error) { + if isNil(terminal) { + return nil, fmt.Errorf("compose protocol stages: terminal endpoint is nil") + } + + currentProtocol := terminal.Protocol() + if currentProtocol == "" { + return nil, fmt.Errorf("compose protocol stages: terminal endpoint has empty protocol") + } + + current := terminal + for i := len(stages) - 1; i >= 0; i-- { + stage := stages[i] + if isNil(stage) { + return nil, fmt.Errorf("compose protocol stages: stage at index %d is nil", i) + } + + name := strings.TrimSpace(stage.Name()) + if name == "" { + return nil, fmt.Errorf("compose protocol stages: stage at index %d has empty name", i) + } + + stageProtocol := stage.Protocol() + if stageProtocol == "" { + return nil, fmt.Errorf("compose protocol stages: stage %q has empty protocol", name) + } + if stageProtocol != currentProtocol { + return nil, fmt.Errorf( + "compose protocol stages: stage %q speaks %q and cannot wrap endpoint speaking %q", + name, + stageProtocol, + currentProtocol, + ) + } + + wrapped := stage.Wrap(current) + if isNil(wrapped) { + return nil, fmt.Errorf("compose protocol stages: stage %q returned a nil endpoint", name) + } + if wrapped.Protocol() != stageProtocol { + return nil, fmt.Errorf( + "compose protocol stages: stage %q returned endpoint speaking %q, want %q", + name, + wrapped.Protocol(), + stageProtocol, + ) + } + + current = wrapped + currentProtocol = stageProtocol + } + + return current, nil +} + +// isNil recognizes typed nil pointers stored in an interface so Compose can +// report a validation error instead of panicking while calling their methods. +func isNil(value any) bool { + if value == nil { + return true + } + + v := reflect.ValueOf(value) + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return v.IsNil() + default: + return false + } +} diff --git a/internal/protocol/stage/compose_test.go b/internal/protocol/stage/compose_test.go new file mode 100644 index 000000000..cdf032597 --- /dev/null +++ b/internal/protocol/stage/compose_test.go @@ -0,0 +1,431 @@ +package stage + +import ( + "context" + "errors" + "io" + "reflect" + "strings" + "testing" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +func TestComposeCompleteOrder(t *testing.T) { + t.Parallel() + + var calls []string + terminal := &recordingEndpoint{ + protocol: protocol.TypeAnthropicBeta, + calls: &calls, + response: &Response{ + Value: "terminal response", + Usage: protocol.NewTokenUsage(7, 3), + Model: "provider-model", + }, + } + + composed, err := Compose( + terminal, + &recordingStage{name: "guardrails", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + &recordingStage{name: "tool_loop", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + ) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + if len(calls) != 0 { + t.Fatalf("Compose() executed endpoint, calls = %v", calls) + } + + call := Call{ + Request: "native request", + Metadata: CallMetadata{ + RequestID: "req-1", + Attempt: 2, + }, + } + response, err := composed.Complete(context.Background(), call) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + + wantCalls := []string{ + "guardrails:request", + "tool_loop:request", + "terminal:request", + "terminal:response", + "tool_loop:response", + "guardrails:response", + } + if !reflect.DeepEqual(calls, wantCalls) { + t.Fatalf("calls = %v, want %v", calls, wantCalls) + } + if response != terminal.response { + t.Fatal("Complete() did not preserve terminal response") + } + if terminal.lastCall.Metadata != call.Metadata { + t.Fatalf("metadata = %+v, want %+v", terminal.lastCall.Metadata, call.Metadata) + } +} + +func TestComposeStreamOrderAndClose(t *testing.T) { + t.Parallel() + + var calls []string + terminalStream := &recordingEventStream{ + calls: &calls, + events: []Event{ + {Value: "event-1"}, + }, + result: StreamResult{ + Usage: protocol.NewTokenUsage(11, 5), + Model: "provider-model", + SideEffectsCommitted: true, + }, + } + terminal := &recordingEndpoint{ + protocol: protocol.TypeAnthropicBeta, + calls: &calls, + stream: terminalStream, + } + + composed, err := Compose( + terminal, + &recordingStage{name: "guardrails", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + &recordingStage{name: "tool_loop", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + ) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + + stream, err := composed.Stream(context.Background(), Call{Request: "native request"}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + event, err := stream.Next(context.Background()) + if err != nil { + t.Fatalf("Next() error = %v", err) + } + if event.Value != "event-1" { + t.Fatalf("event.Value = %v, want event-1", event.Value) + } + if got := stream.Result(); !reflect.DeepEqual(got, terminalStream.result) { + t.Fatalf("Result() = %+v, want %+v", got, terminalStream.result) + } + + _, err = stream.Next(context.Background()) + if !errors.Is(err, io.EOF) { + t.Fatalf("second Next() error = %v, want io.EOF", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if terminalStream.closeCount != 1 { + t.Fatalf("terminal close count = %d, want 1", terminalStream.closeCount) + } + + wantCalls := []string{ + "guardrails:stream_request", + "tool_loop:stream_request", + "terminal:stream_request", + "terminal:event", + "tool_loop:event", + "guardrails:event", + "terminal:eof", + "tool_loop:eof", + "guardrails:eof", + "guardrails:close", + "tool_loop:close", + "terminal:close", + } + if !reflect.DeepEqual(calls, wantCalls) { + t.Fatalf("calls = %v, want %v", calls, wantCalls) + } +} + +func TestComposeRejectsInvalidChain(t *testing.T) { + t.Parallel() + + validEndpoint := &recordingEndpoint{protocol: protocol.TypeAnthropicBeta} + validStage := func() *recordingStage { + return &recordingStage{name: "guardrails", protocol: protocol.TypeAnthropicBeta} + } + + var typedNilEndpoint *recordingEndpoint + var typedNilStage *recordingStage + + tests := []struct { + name string + terminal Endpoint + stages []Stage + want string + }{ + { + name: "nil terminal", + want: "terminal endpoint is nil", + }, + { + name: "typed nil terminal", + terminal: typedNilEndpoint, + want: "terminal endpoint is nil", + }, + { + name: "empty terminal protocol", + terminal: &recordingEndpoint{}, + want: "terminal endpoint has empty protocol", + }, + { + name: "nil stage", + terminal: validEndpoint, + stages: []Stage{nil}, + want: "stage at index 0 is nil", + }, + { + name: "typed nil stage", + terminal: validEndpoint, + stages: []Stage{typedNilStage}, + want: "stage at index 0 is nil", + }, + { + name: "empty stage name", + terminal: validEndpoint, + stages: []Stage{&recordingStage{protocol: protocol.TypeAnthropicBeta}}, + want: "stage at index 0 has empty name", + }, + { + name: "empty stage protocol", + terminal: validEndpoint, + stages: []Stage{&recordingStage{name: "guardrails"}}, + want: `stage "guardrails" has empty protocol`, + }, + { + name: "protocol mismatch", + terminal: validEndpoint, + stages: []Stage{&recordingStage{ + name: "guardrails", + protocol: protocol.TypeOpenAIResponses, + }}, + want: `stage "guardrails" speaks "openai_responses" and cannot wrap endpoint speaking "anthropic_beta"`, + }, + { + name: "nil wrapped endpoint", + terminal: validEndpoint, + stages: []Stage{&recordingStage{ + name: "guardrails", + protocol: protocol.TypeAnthropicBeta, + returnNil: true, + }}, + want: `stage "guardrails" returned a nil endpoint`, + }, + { + name: "wrapped endpoint changed protocol", + terminal: validEndpoint, + stages: []Stage{&recordingStage{ + name: "guardrails", + protocol: protocol.TypeAnthropicBeta, + wrappedProtocol: protocol.TypeOpenAIChat, + }}, + want: `stage "guardrails" returned endpoint speaking "openai_chat", want "anthropic_beta"`, + }, + { + name: "valid baseline", + terminal: validEndpoint, + stages: []Stage{validStage()}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := Compose(tt.terminal, tt.stages...) + if tt.want == "" { + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + if got == nil { + t.Fatal("Compose() returned nil endpoint") + } + return + } + + if err == nil { + t.Fatalf("Compose() error = nil, want containing %q", tt.want) + } + if !strings.Contains(err.Error(), tt.want) { + t.Fatalf("Compose() error = %q, want containing %q", err, tt.want) + } + }) + } +} + +type recordingEndpoint struct { + protocol protocol.APIType + calls *[]string + response *Response + stream EventStream + lastCall Call +} + +func (e *recordingEndpoint) Protocol() protocol.APIType { + return e.protocol +} + +func (e *recordingEndpoint) Complete(_ context.Context, call Call) (*Response, error) { + e.lastCall = call + e.append("terminal:request") + e.append("terminal:response") + return e.response, nil +} + +func (e *recordingEndpoint) Stream(_ context.Context, call Call) (EventStream, error) { + e.lastCall = call + e.append("terminal:stream_request") + return e.stream, nil +} + +func (e *recordingEndpoint) append(value string) { + if e.calls != nil { + *e.calls = append(*e.calls, value) + } +} + +type recordingStage struct { + name string + protocol protocol.APIType + calls *[]string + returnNil bool + wrappedProtocol protocol.APIType +} + +func (s *recordingStage) Name() string { + return s.name +} + +func (s *recordingStage) Protocol() protocol.APIType { + return s.protocol +} + +func (s *recordingStage) Wrap(next Endpoint) Endpoint { + if s.returnNil { + return nil + } + wrappedProtocol := s.protocol + if s.wrappedProtocol != "" { + wrappedProtocol = s.wrappedProtocol + } + return &recordingStageEndpoint{ + name: s.name, + protocol: wrappedProtocol, + calls: s.calls, + next: next, + } +} + +type recordingStageEndpoint struct { + name string + protocol protocol.APIType + calls *[]string + next Endpoint +} + +func (e *recordingStageEndpoint) Protocol() protocol.APIType { + return e.protocol +} + +func (e *recordingStageEndpoint) Complete(ctx context.Context, call Call) (*Response, error) { + e.append(e.name + ":request") + response, err := e.next.Complete(ctx, call) + if err != nil { + return nil, err + } + e.append(e.name + ":response") + return response, nil +} + +func (e *recordingStageEndpoint) Stream(ctx context.Context, call Call) (EventStream, error) { + e.append(e.name + ":stream_request") + stream, err := e.next.Stream(ctx, call) + if err != nil { + return nil, err + } + return &recordingStageStream{name: e.name, calls: e.calls, next: stream}, nil +} + +func (e *recordingStageEndpoint) append(value string) { + if e.calls != nil { + *e.calls = append(*e.calls, value) + } +} + +type recordingStageStream struct { + name string + calls *[]string + next EventStream +} + +func (s *recordingStageStream) Next(ctx context.Context) (Event, error) { + event, err := s.next.Next(ctx) + switch { + case err == nil: + s.append(s.name + ":event") + case errors.Is(err, io.EOF): + s.append(s.name + ":eof") + } + return event, err +} + +func (s *recordingStageStream) Close() error { + s.append(s.name + ":close") + return s.next.Close() +} + +func (s *recordingStageStream) Result() StreamResult { + return s.next.Result() +} + +func (s *recordingStageStream) append(value string) { + if s.calls != nil { + *s.calls = append(*s.calls, value) + } +} + +type recordingEventStream struct { + calls *[]string + events []Event + result StreamResult + next int + closeCount int +} + +func (s *recordingEventStream) Next(ctx context.Context) (Event, error) { + if err := ctx.Err(); err != nil { + return Event{}, err + } + if s.next >= len(s.events) { + s.append("terminal:eof") + return Event{}, io.EOF + } + + event := s.events[s.next] + s.next++ + s.append("terminal:event") + return event, nil +} + +func (s *recordingEventStream) Close() error { + s.closeCount++ + s.append("terminal:close") + return nil +} + +func (s *recordingEventStream) Result() StreamResult { + return s.result +} + +func (s *recordingEventStream) append(value string) { + if s.calls != nil { + *s.calls = append(*s.calls, value) + } +} diff --git a/internal/protocol/stage/doc.go b/internal/protocol/stage/doc.go new file mode 100644 index 000000000..29860c238 --- /dev/null +++ b/internal/protocol/stage/doc.go @@ -0,0 +1,8 @@ +// Package stage defines transport-independent protocol endpoints and ordered +// endpoint wrappers. +// +// The package is the additive foundation for the Protocol Stage Chain design in +// .design/protocol-stage-chain.md. It intentionally has no integration with the +// current server dispatch path yet: adding these contracts must not change +// production traffic until bridges and feature stages are validated separately. +package stage diff --git a/internal/protocol/stage/endpoint.go b/internal/protocol/stage/endpoint.go new file mode 100644 index 000000000..9404815d6 --- /dev/null +++ b/internal/protocol/stage/endpoint.go @@ -0,0 +1,91 @@ +package stage + +import ( + "context" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +// Call is one invocation of an Endpoint in that endpoint's native protocol. +// Request remains protocol-native; a Bridge is responsible for changing its +// concrete type before it crosses a protocol boundary. +type Call struct { + Request any + Metadata CallMetadata + State ProtocolState +} + +// ProtocolState carries typed, request-derived facts that a later endpoint or +// stage still needs after the request changes protocol. It deliberately avoids +// an open-ended property bag: every carried value needs an explicit contract. +// +// Values are per-call and may be mutated by inner request transforms. A Bridge +// must never retain them on the shared Bridge instance. +type ProtocolState struct { + // OpenAIChat is populated when a request is converted to OpenAI Chat. It + // preserves reasoning/thinking facts used by provider-specific transforms. + OpenAIChat *protocol.OpenAIConfig +} + +// CallMetadata carries the small set of attempt identity fields that every +// stage may need. It is immutable by convention: a stage should copy Call +// before changing metadata for an inner invocation. +type CallMetadata struct { + RequestID string + // Attempt is zero for the first provider attempt and increments for retries. + Attempt int +} + +// Response is the complete result returned by an Endpoint. Value is expressed +// in the endpoint's native protocol. The remaining fields are protocol-neutral +// facts that outer stages and the eventual failover adapter must preserve. +type Response struct { + Value any + Usage *protocol.TokenUsage + Model string + SideEffectsCommitted bool +} + +// Event is one native-protocol streaming event. +type Event struct { + Value any +} + +// StreamResult is the latest protocol-neutral summary of an EventStream. It is +// valid before completion, but usage and model data may only become final after +// Next returns io.EOF. +type StreamResult struct { + Usage *protocol.TokenUsage + Model string + SideEffectsCommitted bool +} + +// EventStream is a pull-based stream in one concrete protocol. +// +// Next returns io.EOF on normal completion and must honor ctx cancellation. +// The caller must call Close exactly once for every successfully returned +// EventStream. +// Result returns the latest terminal summary and must not advance the stream. +type EventStream interface { + Next(ctx context.Context) (Event, error) + Close() error + Result() StreamResult +} + +// Endpoint is a complete non-streaming and streaming implementation of one +// concrete protocol. It does not own HTTP parsing, response headers, or SSE +// framing. +type Endpoint interface { + Protocol() protocol.APIType + Complete(ctx context.Context, call Call) (*Response, error) + Stream(ctx context.Context, call Call) (EventStream, error) +} + +// Stage is a named full-duplex wrapper implemented in one concrete protocol. +// Wrap must not execute next. Compose validates the protocol reported by both +// the Stage and the wrapped Endpoint. +type Stage interface { + Name() string + Protocol() protocol.APIType + Wrap(next Endpoint) Endpoint +} diff --git a/internal/protocol/stage/guardrail/anthropic_beta.go b/internal/protocol/stage/guardrail/anthropic_beta.go new file mode 100644 index 000000000..1963cee2d --- /dev/null +++ b/internal/protocol/stage/guardrail/anthropic_beta.go @@ -0,0 +1,328 @@ +package guardrail + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "sync" + + "github.com/anthropics/anthropic-sdk-go" + + "github.com/tingly-dev/tingly-box/internal/constant" + guardrailsruntime "github.com/tingly-dev/tingly-box/internal/guardrails" + guardrailsadapter "github.com/tingly-dev/tingly-box/internal/guardrails/adapter" + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + guardrailsmutate "github.com/tingly-dev/tingly-box/internal/guardrails/mutate" + guardrailspipeline "github.com/tingly-dev/tingly-box/internal/guardrails/pipeline" + protocol "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" +) + +// AnthropicBetaConfig supplies the existing Guardrails policy runtime and the +// request metadata envelope. The Stage owns a fresh credential/stream state for +// every Complete or Stream call. +type AnthropicBetaConfig struct { + Name string + Runtime *guardrailsruntime.Guardrails + BaseInput guardrailscore.Input + Observe Observer +} + +// NewAnthropicBeta constructs an authoritative Anthropic Beta Guardrail Stage. +// It preserves the legacy fail-open behavior for request and complete-response +// evaluation errors while keeping stream rewrite errors visible to the caller. +func NewAnthropicBeta(config AnthropicBetaConfig) (protocolstage.Stage, error) { + name := config.Name + if name == "" { + name = "guardrail_anthropic_beta" + } + if config.Runtime == nil || config.Runtime.PolicyEngine() == nil { + return nil, fmt.Errorf("construct Anthropic Beta Guardrail Stage %q: runtime is unavailable", name) + } + return &anthropicBetaStage{ + name: name, + runtime: config.Runtime, + baseInput: config.BaseInput, + observe: config.Observe, + }, nil +} + +type anthropicBetaStage struct { + name string + runtime *guardrailsruntime.Guardrails + baseInput guardrailscore.Input + observe Observer +} + +func (s *anthropicBetaStage) Name() string { return s.name } +func (*anthropicBetaStage) Protocol() protocol.APIType { return protocol.TypeAnthropicBeta } +func (s *anthropicBetaStage) Wrap(next protocolstage.Endpoint) protocolstage.Endpoint { + return &anthropicBetaEndpoint{stage: s, next: next} +} + +type anthropicBetaEndpoint struct { + stage *anthropicBetaStage + next protocolstage.Endpoint +} + +func (*anthropicBetaEndpoint) Protocol() protocol.APIType { return protocol.TypeAnthropicBeta } + +func (e *anthropicBetaEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + session, prepared, err := e.prepare(ctx, call) + if err != nil { + return nil, err + } + response, err := e.next.Complete(ctx, prepared) + if err != nil { + return response, err + } + if response == nil { + return nil, fmt.Errorf("Anthropic Beta Guardrail Stage received a nil response") + } + message, ok := response.Value.(*anthropic.BetaMessage) + if !ok || message == nil { + return nil, fmt.Errorf("Anthropic Beta Guardrail Stage received response %T", response.Value) + } + + input := session.responseInput() + mutation, evaluationErr := guardrailspipeline.ProcessAnthropicV1BetaNonStreamResponse(ctx, e.stage.runtime, input, message) + if evaluationErr != nil { + guardrailsmutate.RestoreAnthropicV1BetaResponseCredentials(session.mask, message) + e.report(PhaseResponse, Decision{}, evaluationErr) + return response, nil + } + if !mutation.Changed { + guardrailsmutate.RestoreAnthropicV1BetaResponseCredentials(session.mask, message) + } + e.report(PhaseResponse, decisionFromResult(mutation.Evaluation.Result), nil) + return response, nil +} + +func (e *anthropicBetaEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + session, prepared, err := e.prepare(ctx, call) + if err != nil { + return nil, err + } + stream, err := e.next.Stream(ctx, prepared) + if err != nil { + return nil, err + } + if stream == nil { + return nil, fmt.Errorf("Anthropic Beta Guardrail Stage received a nil stream") + } + baseInput := session.responseInput() + streamState := newGuardrailsStreamState() + onEvent, onError := guardrailspipeline.NewGuardrailsHooks(ctx, e.stage.runtime, baseInput, streamState) + return &anthropicBetaGuardrailStream{ + parent: e, + stream: stream, + mask: session.mask, + streamState: streamState, + onEvent: onEvent, + onError: onError, + }, nil +} + +type anthropicBetaSession struct { + request *anthropic.BetaMessageNewParams + base guardrailscore.Input + mask *guardrailscore.CredentialMaskState +} + +func (e *anthropicBetaEndpoint) prepare(ctx context.Context, call protocolstage.Call) (*anthropicBetaSession, protocolstage.Call, error) { + request, ok := call.Request.(*anthropic.BetaMessageNewParams) + if !ok || request == nil { + return nil, protocolstage.Call{}, fmt.Errorf("Anthropic Beta Guardrail Stage received request %T", call.Request) + } + mask := guardrailscore.NewCredentialMaskState() + base := e.stage.baseInput + base.Direction = guardrailscore.DirectionRequest + base.State.CredentialMask = mask + base.Payload.Protocol = "anthropic_beta" + base.Payload.Request = request + base.SetContextValue(constant.CtxKeyCredentialMaskState, mask) + + evaluationErr := guardrailspipeline.ProcessAnthropicBetaRequest(ctx, e.stage.runtime, base) + if evaluationErr != nil { + e.report(PhaseRequest, Decision{}, evaluationErr) + } else { + e.report(PhaseRequest, Decision{Verdict: VerdictAllow}, nil) + } + prepared := call + prepared.Request = request + return &anthropicBetaSession{request: request, base: base, mask: mask}, prepared, nil +} + +func (s *anthropicBetaSession) responseInput() guardrailscore.Input { + input := s.base + input.Direction = guardrailscore.DirectionResponse + input.State.CredentialMask = s.mask + input.Content = guardrailscore.Content{ + Messages: guardrailsadapter.AdaptMessagesFromAnthropicV1Beta(s.request.System, s.request.Messages), + } + input.Payload.Request = s.request + return input +} + +func (e *anthropicBetaEndpoint) report(phase Phase, decision Decision, err error) { + if e.stage.observe == nil { + return + } + e.stage.observe(Observation{ + Stage: e.stage.name, + Protocol: protocol.TypeAnthropicBeta, + Phase: phase, + Decision: decision, + Err: err, + }) +} + +func decisionFromResult(result guardrailscore.Result) Decision { + decision := Decision{Verdict: VerdictAllow} + if result.Verdict == guardrailscore.VerdictBlock { + decision.Verdict = VerdictBlock + } + if len(result.Reasons) > 0 { + decision.Reason = result.Reasons[0].Reason + } + return decision +} + +func newGuardrailsStreamState() *protocol.GuardrailsStreamState { + return &protocol.GuardrailsStreamState{ + PendingBlockMessages: make(map[string]string), + PendingBlockedIndex: make(map[int]string), + AnthropicToolEvents: make(map[int][]protocol.GuardrailsBufferedEvent), + AnthropicToolIDs: make(map[int]string), + } +} + +type anthropicBetaGuardrailStream struct { + parent *anthropicBetaEndpoint + stream protocolstage.EventStream + mask *guardrailscore.CredentialMaskState + streamState *protocol.GuardrailsStreamState + onEvent func(event interface{}) error + onError func(error) + pending []protocolstage.Event + + closeOnce sync.Once + closeErr error +} + +func (s *anthropicBetaGuardrailStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + for { + if len(s.pending) > 0 { + event := s.pending[0] + s.pending = s.pending[1:] + return event, nil + } + + event, err := s.stream.Next(ctx) + if err != nil { + if !errors.Is(err, io.EOF) && s.onError != nil { + s.onError(err) + } + return event, err + } + betaEvent, normalizeErr := normalizeAnthropicBetaEvent(event.Value) + if normalizeErr != nil { + return protocolstage.Event{}, normalizeErr + } + if s.onEvent != nil { + if hookErr := s.onEvent(betaEvent); hookErr != nil { + if s.onError != nil { + s.onError(hookErr) + } + return protocolstage.Event{}, hookErr + } + } + + decision, rewritten, rewriteErr := guardrailsmutate.RewriteAnthropicToolUseEventDecision(s.mask, s.streamState, betaEvent) + if rewriteErr != nil { + if s.onError != nil { + s.onError(rewriteErr) + } + return protocolstage.Event{}, rewriteErr + } + if decision == guardrailsmutate.AnthropicToolUseDecisionNone { + s.parent.report(PhaseEvent, Decision{Verdict: VerdictAllow}, nil) + return event, nil + } + if len(rewritten) == 0 { + continue + } + for _, value := range rewritten { + s.pending = append(s.pending, protocolstage.Event{Value: protocolstream.AnthropicEvent{ + Type: value.EventType, + Data: value.Payload, + }}) + } + verdict := VerdictAllow + reason := "" + if decision == guardrailsmutate.AnthropicToolUseDecisionBlock { + verdict = VerdictBlock + reason = "stream tool use rewritten" + } + s.parent.report(PhaseEvent, Decision{Verdict: verdict, Reason: reason}, nil) + } +} + +func (s *anthropicBetaGuardrailStream) Close() error { + s.closeOnce.Do(func() { + s.closeErr = s.stream.Close() + }) + return s.closeErr +} + +func (s *anthropicBetaGuardrailStream) Result() protocolstage.StreamResult { return s.stream.Result() } + +func normalizeAnthropicBetaEvent(value any) (*anthropic.BetaRawMessageStreamEventUnion, error) { + switch event := value.(type) { + case anthropic.BetaRawMessageStreamEventUnion: + copy := event + return ©, nil + case *anthropic.BetaRawMessageStreamEventUnion: + if event == nil { + return nil, fmt.Errorf("Anthropic Beta Guardrail Stage received a nil stream event") + } + return event, nil + case protocolstream.AnthropicEvent: + payload, err := json.Marshal(event.Data) + if err != nil { + return nil, fmt.Errorf("marshal converted Anthropic Beta Guardrail event: %w", err) + } + var object map[string]json.RawMessage + if err := json.Unmarshal(payload, &object); err != nil { + return nil, fmt.Errorf("decode converted Anthropic Beta Guardrail event: %w", err) + } + if object == nil { + return nil, fmt.Errorf("decode converted Anthropic Beta Guardrail event: payload is not an object") + } + eventType, err := json.Marshal(event.Type) + if err != nil { + return nil, fmt.Errorf("marshal converted Anthropic Beta Guardrail event type: %w", err) + } + object["type"] = eventType + payload, err = json.Marshal(object) + if err != nil { + return nil, fmt.Errorf("encode converted Anthropic Beta Guardrail event: %w", err) + } + var decoded anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(payload, &decoded); err != nil { + return nil, fmt.Errorf("unmarshal converted Anthropic Beta Guardrail event: %w", err) + } + return &decoded, nil + default: + return nil, fmt.Errorf("Anthropic Beta Guardrail Stage received stream event %T", value) + } +} + +var _ protocolstage.Stage = (*anthropicBetaStage)(nil) +var _ protocolstage.EventStream = (*anthropicBetaGuardrailStream)(nil) diff --git a/internal/protocol/stage/guardrail/anthropic_beta_test.go b/internal/protocol/stage/guardrail/anthropic_beta_test.go new file mode 100644 index 000000000..ff5c44ccf --- /dev/null +++ b/internal/protocol/stage/guardrail/anthropic_beta_test.go @@ -0,0 +1,311 @@ +package guardrail + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + + guardrailsruntime "github.com/tingly-dev/tingly-box/internal/guardrails" + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + protocol "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" +) + +func TestAnthropicBetaGuardrailStageMasksRequestAndRestoresCompleteResponse(t *testing.T) { + t.Parallel() + + const secret = "sk-live-secret-value" + credential := guardrailscore.ProtectedCredential{ + ID: "credential-1", + Name: "test key", + Type: guardrailscore.ProtectedCredentialTypeAPIKey, + Secret: secret, + AliasToken: "TINGLY_CRED_API_KEY_TEST", + Enabled: true, + } + runtime := testGuardrailsRuntime(policyRunnerFunc(func(context.Context, guardrailscore.Input) (guardrailscore.Result, error) { + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + })) + runtime.SetCredentialCache(guardrailsruntime.BuildCredentialCache([]guardrailscore.ProtectedCredential{credential}, []string{"anthropic"})) + + request := &anthropic.BetaMessageNewParams{ + Model: "client-model", + MaxTokens: 64, + Messages: []anthropic.BetaMessageParam{ + anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("use " + secret)), + }, + } + var providerText string + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, protocolstage.Call) (*protocolstage.Response, error) { + providerText = request.Messages[0].Content[0].OfText.Text + return &protocolstage.Response{Value: &anthropic.BetaMessage{ + Content: []anthropic.BetaContentBlockUnion{{Type: "text", Text: providerText}}, + }}, nil + }, + } + endpoint := composeAnthropicBetaGuardrail(t, terminal, runtime) + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if strings.Contains(providerText, secret) || !strings.Contains(providerText, credential.AliasToken) { + t.Fatalf("provider request text = %q", providerText) + } + message := response.Value.(*anthropic.BetaMessage) + if got := message.Content[0].Text; got != "use "+secret { + t.Fatalf("client response text = %q, want restored secret", got) + } +} + +func TestAnthropicBetaGuardrailStageRestoresCredentialsWhenResponseEvaluationFailsOpen(t *testing.T) { + t.Parallel() + + const secret = "sk-fail-open-secret" + credential := guardrailscore.ProtectedCredential{ + ID: "credential-fail-open", Name: "test key", Type: guardrailscore.ProtectedCredentialTypeAPIKey, + Secret: secret, AliasToken: "TINGLY_CRED_API_KEY_FAIL_OPEN", Enabled: true, + } + evaluationErr := errors.New("policy backend unavailable") + runtime := testGuardrailsRuntime(policyRunnerFunc(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse { + return guardrailscore.Result{}, evaluationErr + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + })) + runtime.SetCredentialCache(guardrailsruntime.BuildCredentialCache([]guardrailscore.ProtectedCredential{credential}, []string{"anthropic"})) + + request := &anthropic.BetaMessageNewParams{ + Model: "client-model", MaxTokens: 64, + Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("use " + secret))}, + } + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, protocolstage.Call) (*protocolstage.Response, error) { + return &protocolstage.Response{Value: &anthropic.BetaMessage{ + Content: []anthropic.BetaContentBlockUnion{{Type: "text", Text: "use " + credential.AliasToken}}, + }}, nil + }, + } + response, err := composeAnthropicBetaGuardrail(t, terminal, runtime).Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatalf("Complete() error = %v, want fail-open response", err) + } + if got := response.Value.(*anthropic.BetaMessage).Content[0].Text; got != "use "+secret { + t.Fatalf("client response text = %q, want restored secret", got) + } +} + +func TestAnthropicBetaGuardrailStageBlocksCompleteResponse(t *testing.T) { + t.Parallel() + + runtime := testGuardrailsRuntime(policyRunnerFunc(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse && strings.Contains(input.Content.Text, "danger") { + return blockedResult("dangerous output"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + })) + request := &anthropic.BetaMessageNewParams{Model: "client-model", MaxTokens: 64} + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, protocolstage.Call) (*protocolstage.Response, error) { + return &protocolstage.Response{Value: &anthropic.BetaMessage{ + Content: []anthropic.BetaContentBlockUnion{{Type: "text", Text: "danger"}}, + StopReason: anthropic.BetaStopReasonToolUse, + }}, nil + }, + } + endpoint := composeAnthropicBetaGuardrail(t, terminal, runtime) + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + message := response.Value.(*anthropic.BetaMessage) + if len(message.Content) != 1 || message.Content[0].Type != "text" || !strings.Contains(message.Content[0].Text, "Blocked by guardrails") { + t.Fatalf("blocked response content = %+v", message.Content) + } + if message.StopReason != anthropic.BetaStopReasonEndTurn { + t.Fatalf("stop reason = %q, want end_turn", message.StopReason) + } +} + +func TestAnthropicBetaGuardrailStageRejectsNilCompleteResponse(t *testing.T) { + t.Parallel() + + runtime := testGuardrailsRuntime(policyRunnerFunc(func(context.Context, guardrailscore.Input) (guardrailscore.Result, error) { + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + })) + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, protocolstage.Call) (*protocolstage.Response, error) { return nil, nil }, + } + request := &anthropic.BetaMessageNewParams{Model: "client-model", MaxTokens: 64} + if _, err := composeAnthropicBetaGuardrail(t, terminal, runtime).Complete(context.Background(), protocolstage.Call{Request: request}); err == nil { + t.Fatal("Complete() error = nil, want nil-response contract error") + } +} + +func TestAnthropicBetaGuardrailPendingEventsHonorCancellation(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + stream := &anthropicBetaGuardrailStream{ + stream: &fakeStream{}, + pending: []protocolstage.Event{{Value: "pending"}}, + } + if _, err := stream.Next(ctx); !errors.Is(err, context.Canceled) { + t.Fatalf("Next() error = %v, want context.Canceled", err) + } +} + +func TestNormalizeAnthropicBetaEventRejectsNullPayload(t *testing.T) { + t.Parallel() + + if _, err := normalizeAnthropicBetaEvent(protocolstream.AnthropicEvent{Type: "message_stop", Data: nil}); err == nil { + t.Fatal("normalizeAnthropicBetaEvent() error = nil, want object payload error") + } +} + +func TestAnthropicBetaGuardrailStageRewritesBlockedToolUseStream(t *testing.T) { + t.Parallel() + + runtime := testGuardrailsRuntime(policyRunnerFunc(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse && input.Content.Command != nil { + return blockedResult("command denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + })) + request := &anthropic.BetaMessageNewParams{Model: "client-model", MaxTokens: 64} + target := &fakeStream{ + events: []protocolstage.Event{ + {Value: decodeBetaEvent(t, `{"type":"message_start","message":{"id":"msg-1","type":"message","role":"assistant","model":"provider","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":0}}}`)}, + {Value: decodeBetaEvent(t, `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"tool-1","name":"shell","input":{}}}`)}, + {Value: decodeBetaEvent(t, `{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"command\":\"rm -rf /\"}"}}`)}, + {Value: decodeBetaEvent(t, `{"type":"content_block_stop","index":0}`)}, + {Value: decodeBetaEvent(t, `{"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":4}}`)}, + {Value: decodeBetaEvent(t, `{"type":"message_stop"}`)}, + }, + result: protocolstage.StreamResult{Usage: protocol.NewTokenUsage(3, 4), Model: "provider"}, + } + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, protocolstage.Call) (protocolstage.EventStream, error) { + return target, nil + }, + } + endpoint := composeAnthropicBetaGuardrail(t, terminal, runtime) + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + defer stream.Close() + + var payloads []map[string]any + for { + event, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + t.Fatalf("Next() error = %v", nextErr) + } + payloads = append(payloads, eventMap(t, event.Value)) + } + if len(payloads) != 6 { + t.Fatalf("output event count = %d, want 6: %#v", len(payloads), payloads) + } + block, _ := payloads[1]["content_block"].(map[string]any) + if block["type"] != "text" { + t.Fatalf("rewritten block = %#v", block) + } + delta, _ := payloads[2]["delta"].(map[string]any) + if text, _ := delta["text"].(string); !strings.Contains(text, "Blocked by guardrails") { + t.Fatalf("rewritten delta = %#v", delta) + } + messageDelta, _ := payloads[4]["delta"].(map[string]any) + if messageDelta["stop_reason"] != "end_turn" { + t.Fatalf("message delta = %#v", messageDelta) + } + if got := stream.Result(); got.Model != "provider" || got.Usage == nil { + t.Fatalf("Result() = %+v", got) + } +} + +type policyRunnerFunc func(context.Context, guardrailscore.Input) (guardrailscore.Result, error) + +func (f policyRunnerFunc) Evaluate(ctx context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + return f(ctx, input) +} + +func testGuardrailsRuntime(runner guardrailsruntime.PolicyRunner) *guardrailsruntime.Guardrails { + return &guardrailsruntime.Guardrails{Policy: runner, HasActivePolicies: true} +} + +func blockedResult(reason string) guardrailscore.Result { + return guardrailscore.Result{ + Verdict: guardrailscore.VerdictBlock, + Reasons: []guardrailscore.PolicyResult{{ + PolicyID: "test-policy", + Verdict: guardrailscore.VerdictBlock, + Reason: reason, + }}, + } +} + +func composeAnthropicBetaGuardrail(t *testing.T, terminal protocolstage.Endpoint, runtime *guardrailsruntime.Guardrails) protocolstage.Endpoint { + t.Helper() + guardrail, err := NewAnthropicBeta(AnthropicBetaConfig{ + Runtime: runtime, + BaseInput: guardrailscore.Input{ + Scenario: "anthropic", + Model: "provider-model", + }, + }) + if err != nil { + t.Fatalf("NewAnthropicBeta() error = %v", err) + } + endpoint, err := protocolstage.Compose(terminal, guardrail) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + return endpoint +} + +func decodeBetaEvent(t *testing.T, raw string) anthropic.BetaRawMessageStreamEventUnion { + t.Helper() + var event anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal([]byte(raw), &event); err != nil { + t.Fatalf("decode Beta event: %v", err) + } + return event +} + +func eventMap(t *testing.T, value any) map[string]any { + t.Helper() + var raw []byte + switch event := value.(type) { + case anthropic.BetaRawMessageStreamEventUnion: + raw = []byte(event.RawJSON()) + case protocolstream.AnthropicEvent: + var err error + raw, err = json.Marshal(event.Data) + if err != nil { + t.Fatalf("marshal Anthropic event: %v", err) + } + default: + t.Fatalf("event type = %T", value) + } + var payload map[string]any + if err := json.Unmarshal(raw, &payload); err != nil { + t.Fatalf("decode event payload: %v", err) + } + return payload +} diff --git a/internal/protocol/stage/guardrail/guardrail.go b/internal/protocol/stage/guardrail/guardrail.go new file mode 100644 index 000000000..00eeec868 --- /dev/null +++ b/internal/protocol/stage/guardrail/guardrail.go @@ -0,0 +1,222 @@ +// Package guardrail implements a protocol-native, full-duplex Guardrail Stage. +// +// The foundation is deliberately observe-only: evaluators inspect requests, +// complete responses, and stream events without changing live traffic. Policy +// enforcement and mutation are separate adapters built on this lifecycle. +package guardrail + +import ( + "context" + "fmt" + "strings" + "sync" + + protocol "github.com/tingly-dev/tingly-box/ai" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" +) + +// Phase identifies the concrete lifecycle surface inspected by an Evaluator. +type Phase string + +const ( + PhaseRequest Phase = "request" + PhaseResponse Phase = "response" + PhaseEvent Phase = "event" +) + +// Verdict is an evaluator's protocol-neutral policy decision. +type Verdict string + +const ( + VerdictAllow Verdict = "allow" + VerdictBlock Verdict = "block" +) + +// Decision is an immutable observation result. Reason is diagnostic context; +// it must not contain full request or response bodies. +type Decision struct { + Verdict Verdict + Reason string +} + +// Observation reports one dry-run evaluation without affecting live traffic. +// Err is evaluator failure, not downstream provider failure. +type Observation struct { + Stage string + Protocol protocol.APIType + Phase Phase + Decision Decision + Err error +} + +// Evaluator inspects one concrete protocol. Implementations must be +// concurrency-safe and must not mutate Call, Response, Event, or their native +// protocol values. Mutable correlation belongs in the Session returned by +// Open, never on the shared Evaluator. +type Evaluator interface { + Protocol() protocol.APIType + Open(ctx context.Context, call protocolstage.Call) (Session, error) +} + +// Session owns all mutable state for one Complete or Stream invocation. +type Session interface { + EvaluateRequest(ctx context.Context, call protocolstage.Call) (Decision, error) + EvaluateResponse(ctx context.Context, response *protocolstage.Response) (Decision, error) + EvaluateEvent(ctx context.Context, event protocolstage.Event) (Decision, error) +} + +// Observer receives dry-run facts. It must not retain native protocol values; +// Observation intentionally contains only bounded diagnostic metadata. +type Observer func(Observation) + +// Config constructs one observe-only Guardrail Stage. +type Config struct { + Name string + Evaluator Evaluator + Observe Observer +} + +// New constructs a protocol-native Guardrail Stage. Evaluation failures are +// fail-open and reported to Observe; downstream failures retain their normal +// endpoint semantics. +func New(config Config) (protocolstage.Stage, error) { + name := strings.TrimSpace(config.Name) + if name == "" { + name = "guardrail" + } + if config.Evaluator == nil { + return nil, fmt.Errorf("construct Guardrail Stage %q: evaluator is nil", name) + } + api := config.Evaluator.Protocol() + if api == "" { + return nil, fmt.Errorf("construct Guardrail Stage %q: evaluator protocol is empty", name) + } + return &guardrailStage{name: name, api: api, evaluator: config.Evaluator, observe: config.Observe}, nil +} + +type guardrailStage struct { + name string + api protocol.APIType + evaluator Evaluator + observe Observer +} + +func (s *guardrailStage) Name() string { return s.name } +func (s *guardrailStage) Protocol() protocol.APIType { return s.api } +func (s *guardrailStage) Wrap(next protocolstage.Endpoint) protocolstage.Endpoint { + return &guardrailEndpoint{stage: s, next: next} +} + +type guardrailEndpoint struct { + stage *guardrailStage + next protocolstage.Endpoint +} + +func (e *guardrailEndpoint) Protocol() protocol.APIType { return e.stage.api } + +func (e *guardrailEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + session := e.open(ctx, call) + e.evaluateRequest(ctx, session, call) + response, err := e.next.Complete(ctx, call) + if err != nil { + return response, err + } + if response == nil { + return nil, fmt.Errorf("Guardrail Stage %q received a nil response", e.stage.name) + } + e.evaluateResponse(ctx, session, response) + return response, nil +} + +func (e *guardrailEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + session := e.open(ctx, call) + e.evaluateRequest(ctx, session, call) + stream, err := e.next.Stream(ctx, call) + if err != nil { + return nil, err + } + if stream == nil { + return nil, fmt.Errorf("Guardrail Stage %q received a nil stream", e.stage.name) + } + return &observedStream{parent: e, session: session, stream: stream}, nil +} + +func (e *guardrailEndpoint) open(ctx context.Context, call protocolstage.Call) Session { + session, err := e.stage.evaluator.Open(ctx, call) + if err != nil { + e.report(PhaseRequest, Decision{}, err) + return nil + } + if session == nil { + e.report(PhaseRequest, Decision{}, fmt.Errorf("Guardrail Stage evaluator returned a nil session")) + } + return session +} + +func (e *guardrailEndpoint) evaluateRequest(ctx context.Context, session Session, call protocolstage.Call) { + if session == nil { + return + } + decision, err := session.EvaluateRequest(ctx, call) + e.report(PhaseRequest, decision, err) +} + +func (e *guardrailEndpoint) evaluateResponse(ctx context.Context, session Session, response *protocolstage.Response) { + if session == nil { + return + } + decision, err := session.EvaluateResponse(ctx, response) + e.report(PhaseResponse, decision, err) +} + +func (e *guardrailEndpoint) evaluateEvent(ctx context.Context, session Session, event protocolstage.Event) { + if session == nil { + return + } + decision, err := session.EvaluateEvent(ctx, event) + e.report(PhaseEvent, decision, err) +} + +func (e *guardrailEndpoint) report(phase Phase, decision Decision, err error) { + if e.stage.observe == nil { + return + } + e.stage.observe(Observation{ + Stage: e.stage.name, + Protocol: e.stage.api, + Phase: phase, + Decision: decision, + Err: err, + }) +} + +type observedStream struct { + parent *guardrailEndpoint + session Session + stream protocolstage.EventStream + + closeOnce sync.Once + closeErr error +} + +func (s *observedStream) Next(ctx context.Context) (protocolstage.Event, error) { + event, err := s.stream.Next(ctx) + if err != nil { + return event, err + } + s.parent.evaluateEvent(ctx, s.session, event) + return event, nil +} + +func (s *observedStream) Close() error { + s.closeOnce.Do(func() { + s.closeErr = s.stream.Close() + }) + return s.closeErr +} + +func (s *observedStream) Result() protocolstage.StreamResult { return s.stream.Result() } + +var _ protocolstage.Stage = (*guardrailStage)(nil) +var _ protocolstage.Endpoint = (*guardrailEndpoint)(nil) +var _ protocolstage.EventStream = (*observedStream)(nil) diff --git a/internal/protocol/stage/guardrail/guardrail_test.go b/internal/protocol/stage/guardrail/guardrail_test.go new file mode 100644 index 000000000..4166ae709 --- /dev/null +++ b/internal/protocol/stage/guardrail/guardrail_test.go @@ -0,0 +1,268 @@ +package guardrail + +import ( + "context" + "errors" + "io" + "reflect" + "testing" + + protocol "github.com/tingly-dev/tingly-box/ai" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" +) + +func TestGuardrailStageObservesCompleteLifecycleWithoutMutation(t *testing.T) { + t.Parallel() + + request := &struct{ Text string }{Text: "request"} + responseValue := &struct{ Text string }{Text: "response"} + evaluator := &fakeEvaluator{api: protocol.TypeAnthropicBeta} + var observations []Observation + guardrail := mustGuardrail(t, Config{ + Name: "guardrail_beta", + Evaluator: evaluator, + Observe: func(observation Observation) { observations = append(observations, observation) }, + }) + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(_ context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + if call.Request != request { + t.Fatalf("terminal request = %p, want %p", call.Request, request) + } + return &protocolstage.Response{Value: responseValue, Model: "provider", SideEffectsCommitted: true}, nil + }, + } + endpoint, err := protocolstage.Compose(terminal, guardrail) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + if response.Value != responseValue || response.Model != "provider" || !response.SideEffectsCommitted { + t.Fatalf("response changed = %+v", response) + } + if got := evaluator.sessions; got != 1 { + t.Fatalf("sessions = %d, want 1", got) + } + wantPhases := []Phase{PhaseRequest, PhaseResponse} + if got := observationPhases(observations); !reflect.DeepEqual(got, wantPhases) { + t.Fatalf("phases = %v, want %v", got, wantPhases) + } +} + +func TestGuardrailStageObservesStreamAndPreservesOwnership(t *testing.T) { + t.Parallel() + + evaluator := &fakeEvaluator{api: protocol.TypeAnthropicBeta} + var observations []Observation + guardrail := mustGuardrail(t, Config{ + Evaluator: evaluator, + Observe: func(observation Observation) { observations = append(observations, observation) }, + }) + target := &fakeStream{ + events: []protocolstage.Event{{Value: "first"}, {Value: "second"}}, + result: protocolstage.StreamResult{ + Usage: protocol.NewTokenUsage(3, 2), + Model: "stream-model", + SideEffectsCommitted: true, + }, + } + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, protocolstage.Call) (protocolstage.EventStream, error) { + return target, nil + }, + } + endpoint, err := protocolstage.Compose(terminal, guardrail) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: "request"}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + for index, want := range []string{"first", "second"} { + event, nextErr := stream.Next(context.Background()) + if nextErr != nil || event.Value != want { + t.Fatalf("Next(%d) = (%v, %v), want %q", index, event.Value, nextErr, want) + } + } + if _, err := stream.Next(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("final Next() error = %v, want io.EOF", err) + } + if got := stream.Result(); got.Model != "stream-model" || got.Usage == nil || !got.SideEffectsCommitted { + t.Fatalf("Result() = %+v", got) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("second Close() error = %v", err) + } + if target.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", target.closeCount) + } + wantPhases := []Phase{PhaseRequest, PhaseEvent, PhaseEvent} + if got := observationPhases(observations); !reflect.DeepEqual(got, wantPhases) { + t.Fatalf("phases = %v, want %v", got, wantPhases) + } +} + +func TestGuardrailStageFailsOpenOnEvaluatorErrors(t *testing.T) { + t.Parallel() + + evaluationErr := errors.New("evaluation unavailable") + evaluator := &fakeEvaluator{api: protocol.TypeAnthropicBeta, openErr: evaluationErr} + var observations []Observation + guardrail := mustGuardrail(t, Config{ + Evaluator: evaluator, + Observe: func(observation Observation) { observations = append(observations, observation) }, + }) + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, protocolstage.Call) (*protocolstage.Response, error) { + return &protocolstage.Response{Value: "ok"}, nil + }, + } + endpoint, err := protocolstage.Compose(terminal, guardrail) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: "request"}) + if err != nil || response.Value != "ok" { + t.Fatalf("Complete() = (%+v, %v)", response, err) + } + if len(observations) != 1 || !errors.Is(observations[0].Err, evaluationErr) { + t.Fatalf("observations = %+v", observations) + } +} + +func TestGuardrailStageRejectsNilDownstreamStream(t *testing.T) { + t.Parallel() + + guardrail := mustGuardrail(t, Config{Evaluator: &fakeEvaluator{api: protocol.TypeAnthropicBeta}}) + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, protocolstage.Call) (protocolstage.EventStream, error) { + return nil, nil + }, + } + endpoint, err := protocolstage.Compose(terminal, guardrail) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + if _, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: "request"}); err == nil { + t.Fatal("Stream() error = nil, want nil-stream contract error") + } +} + +func TestGuardrailStageRejectsNilDownstreamResponse(t *testing.T) { + t.Parallel() + + guardrail := mustGuardrail(t, Config{Evaluator: &fakeEvaluator{api: protocol.TypeAnthropicBeta}}) + terminal := &fakeEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, protocolstage.Call) (*protocolstage.Response, error) { + return nil, nil + }, + } + endpoint, err := protocolstage.Compose(terminal, guardrail) + if err != nil { + t.Fatalf("Compose() error = %v", err) + } + if _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: "request"}); err == nil { + t.Fatal("Complete() error = nil, want nil-response contract error") + } +} + +func TestGuardrailStageRejectsInvalidConstruction(t *testing.T) { + t.Parallel() + + if _, err := New(Config{}); err == nil { + t.Fatal("New() error = nil, want evaluator error") + } + if _, err := New(Config{Evaluator: &fakeEvaluator{}}); err == nil { + t.Fatal("New() error = nil, want protocol error") + } +} + +type fakeEvaluator struct { + api protocol.APIType + openErr error + sessions int +} + +func (e *fakeEvaluator) Protocol() protocol.APIType { return e.api } +func (e *fakeEvaluator) Open(context.Context, protocolstage.Call) (Session, error) { + if e.openErr != nil { + return nil, e.openErr + } + e.sessions++ + return fakeSession{}, nil +} + +type fakeSession struct{} + +func (fakeSession) EvaluateRequest(context.Context, protocolstage.Call) (Decision, error) { + return Decision{Verdict: VerdictAllow, Reason: "request observed"}, nil +} +func (fakeSession) EvaluateResponse(context.Context, *protocolstage.Response) (Decision, error) { + return Decision{Verdict: VerdictAllow, Reason: "response observed"}, nil +} +func (fakeSession) EvaluateEvent(context.Context, protocolstage.Event) (Decision, error) { + return Decision{Verdict: VerdictAllow, Reason: "event observed"}, nil +} + +type fakeEndpoint struct { + api protocol.APIType + complete func(context.Context, protocolstage.Call) (*protocolstage.Response, error) + stream func(context.Context, protocolstage.Call) (protocolstage.EventStream, error) +} + +func (e *fakeEndpoint) Protocol() protocol.APIType { return e.api } +func (e *fakeEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + return e.complete(ctx, call) +} +func (e *fakeEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + return e.stream(ctx, call) +} + +type fakeStream struct { + events []protocolstage.Event + index int + result protocolstage.StreamResult + closeCount int +} + +func (s *fakeStream) Next(context.Context) (protocolstage.Event, error) { + if s.index >= len(s.events) { + return protocolstage.Event{}, io.EOF + } + event := s.events[s.index] + s.index++ + return event, nil +} +func (s *fakeStream) Close() error { + s.closeCount++ + return nil +} +func (s *fakeStream) Result() protocolstage.StreamResult { return s.result } + +func mustGuardrail(t *testing.T, config Config) protocolstage.Stage { + t.Helper() + guardrail, err := New(config) + if err != nil { + t.Fatalf("New() error = %v", err) + } + return guardrail +} + +func observationPhases(observations []Observation) []Phase { + phases := make([]Phase, 0, len(observations)) + for _, observation := range observations { + phases = append(phases, observation.Phase) + } + return phases +} diff --git a/internal/protocol/stage/identity.go b/internal/protocol/stage/identity.go new file mode 100644 index 000000000..4f5c7a983 --- /dev/null +++ b/internal/protocol/stage/identity.go @@ -0,0 +1,53 @@ +package stage + +import ( + "context" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +// NewIdentityBridge returns a stateless same-protocol Bridge. It is useful when +// a topology requires an explicit boundary but no wire conversion is needed. +func NewIdentityBridge(api protocol.APIType) Bridge { + return identityBridge{api: api} +} + +type identityBridge struct { + api protocol.APIType +} + +func (b identityBridge) Source() protocol.APIType { + return b.api +} + +func (b identityBridge) Target() protocol.APIType { + return b.api +} + +func (b identityBridge) Capabilities() Capabilities { + return AllBridgeCapabilities +} + +func (b identityBridge) Open(_ context.Context, call Call, _ Operation) (BridgeSession, error) { + return &identityBridgeSession{call: call}, nil +} + +type identityBridgeSession struct { + call Call +} + +func (s *identityBridgeSession) TargetCall() Call { + return s.call +} + +func (s *identityBridgeSession) ConvertComplete(_ context.Context, response *Response) (*Response, error) { + return response, nil +} + +func (s *identityBridgeSession) ConvertStream(_ context.Context, stream EventStream) (EventStream, error) { + return stream, nil +} + +func (s *identityBridgeSession) ConvertError(_ context.Context, err error) error { + return err +} diff --git a/internal/protocol/stage/openaibridge/bridge.go b/internal/protocol/stage/openaibridge/bridge.go new file mode 100644 index 000000000..4305e394e --- /dev/null +++ b/internal/protocol/stage/openaibridge/bridge.go @@ -0,0 +1,153 @@ +// Package openaibridge adapts OpenAI calls to a target provider protocol while +// keeping the outward response in the exact OpenAI source protocol. +package openaibridge + +import ( + "context" + "fmt" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/nonstream" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" +) + +const defaultAnthropicMaxTokens int64 = 4096 + +// AnthropicOptions configures the existing OpenAI Chat to Anthropic Beta +// request and response conversion. Options are immutable after construction. +type AnthropicOptions struct { + DefaultMaxTokens int64 + DisableStreamUsage bool + // ResponseModel overrides the source-visible OpenAI response model while + // leaving the provider-bound request model unchanged. Production routing + // uses this when a public model alias resolves to a different provider model. + ResponseModel string +} + +// NewChatToAnthropicBeta returns an immutable OpenAI Chat -> Anthropic Beta +// Bridge. It is dormant until explicitly registered in a Stage topology. +func NewChatToAnthropicBeta(options AnthropicOptions) stage.Bridge { + if options.DefaultMaxTokens <= 0 { + options.DefaultMaxTokens = defaultAnthropicMaxTokens + } + return &anthropicBetaBridge{options: options} +} + +type anthropicBetaBridge struct { + options AnthropicOptions +} + +func (*anthropicBetaBridge) Source() protocol.APIType { return protocol.TypeOpenAIChat } + +func (*anthropicBetaBridge) Target() protocol.APIType { return protocol.TypeAnthropicBeta } + +func (*anthropicBetaBridge) Capabilities() stage.Capabilities { + return stage.AllBridgeCapabilities +} + +func (b *anthropicBetaBridge) Open(_ context.Context, call stage.Call, operation stage.Operation) (stage.BridgeSession, error) { + switch operation { + case stage.OperationComplete, stage.OperationStream: + default: + return nil, fmt.Errorf("open OpenAI Chat to Anthropic Beta bridge: unsupported operation %s", operation) + } + + chatRequest, err := chatRequest(call.Request) + if err != nil { + return nil, err + } + targetRequest := request.ConvertOpenAIToAnthropicRequest(chatRequest, b.options.DefaultMaxTokens) + if targetRequest == nil { + return nil, fmt.Errorf("open OpenAI Chat to Anthropic Beta bridge: request conversion returned nil") + } + + targetCall := call + targetCall.Request = targetRequest + // OpenAIChat state describes a Chat target request. It must not leak across + // the protocol boundary into an Anthropic-native endpoint. + targetCall.State.OpenAIChat = nil + sourceModel := string(chatRequest.Model) + if b.options.ResponseModel != "" { + sourceModel = b.options.ResponseModel + } + return &anthropicBetaSession{ + operation: operation, + targetCall: targetCall, + sourceModel: sourceModel, + disableStreamUsage: b.options.DisableStreamUsage, + }, nil +} + +func chatRequest(value any) (*openai.ChatCompletionNewParams, error) { + switch request := value.(type) { + case *openai.ChatCompletionNewParams: + if request == nil { + return nil, fmt.Errorf("open OpenAI Chat to Anthropic Beta bridge: request is nil") + } + return request, nil + case openai.ChatCompletionNewParams: + return &request, nil + default: + return nil, fmt.Errorf("open OpenAI Chat to Anthropic Beta bridge: request has type %T, want openai.ChatCompletionNewParams", value) + } +} + +type anthropicBetaSession struct { + operation stage.Operation + targetCall stage.Call + sourceModel string + disableStreamUsage bool +} + +func (s *anthropicBetaSession) TargetCall() stage.Call { return s.targetCall } + +func (s *anthropicBetaSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + if s.operation != stage.OperationComplete { + return nil, fmt.Errorf("convert Anthropic Beta complete response to OpenAI Chat: session was opened for %s", s.operation) + } + message, err := betaMessage(response) + if err != nil { + return nil, err + } + + value := nonstream.ConvertAnthropicBetaToOpenAIChat(message, s.sourceModel) + normalizedUsage := protocolusage.FromAnthropicBetaMessage(message.Usage) + if !normalizedUsage.HasUsage() { + normalizedUsage = nil + } + return &stage.Response{ + Value: value, + Usage: normalizedUsage, + Model: s.sourceModel, + }, nil +} + +func betaMessage(response *stage.Response) (*anthropic.BetaMessage, error) { + if response == nil { + return nil, fmt.Errorf("convert Anthropic Beta response to OpenAI Chat: response is nil") + } + switch value := response.Value.(type) { + case *anthropic.BetaMessage: + if value == nil { + return nil, fmt.Errorf("convert Anthropic Beta response to OpenAI Chat: value is nil") + } + return value, nil + case anthropic.BetaMessage: + return &value, nil + default: + return nil, fmt.Errorf("convert Anthropic Beta response to OpenAI Chat: value has type %T, want anthropic.BetaMessage", response.Value) + } +} + +func (s *anthropicBetaSession) ConvertStream(_ context.Context, target stage.EventStream) (stage.EventStream, error) { + if s.operation != stage.OperationStream { + return nil, fmt.Errorf("convert Anthropic Beta stream to OpenAI Chat: session was opened for %s", s.operation) + } + return newChatStream(target, s.sourceModel, s.disableStreamUsage) +} + +func (*anthropicBetaSession) ConvertError(_ context.Context, err error) error { return err } diff --git a/internal/protocol/stage/openaibridge/bridge_test.go b/internal/protocol/stage/openaibridge/bridge_test.go new file mode 100644 index 000000000..fcde6857e --- /dev/null +++ b/internal/protocol/stage/openaibridge/bridge_test.go @@ -0,0 +1,529 @@ +package openaibridge + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestOpenAIChatToAnthropicBetaComplete(t *testing.T) { + t.Parallel() + + message := decodeBetaMessage(t, map[string]any{ + "id": "msg_reverse", "type": "message", "role": "assistant", + "model": "provider-model", "stop_reason": "tool_use", + "content": []any{map[string]any{ + "type": "tool_use", "id": "tool_reverse", "name": "lookup", + "input": map[string]any{"q": "parallel"}, + }}, + "usage": map[string]any{ + "input_tokens": 8, "output_tokens": 4, + "cache_read_input_tokens": 2, "cache_creation_input_tokens": 1, + }, + }) + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + request := requireBetaRequest(t, call.Request) + if string(request.Model) != "client-model" || request.MaxTokens != 123 { + t.Fatalf("target request model/max = %q/%d", request.Model, request.MaxTokens) + } + if call.Metadata.RequestID != "reverse-complete" || call.Metadata.Attempt != 2 { + t.Fatalf("metadata = %+v", call.Metadata) + } + if call.State.OpenAIChat != nil { + t.Fatalf("target OpenAIChat state leaked: %+v", call.State.OpenAIChat) + } + return &stage.Response{ + Value: message, + Usage: protocol.NewTokenUsage(999, 999), + Model: "provider-model", + SideEffectsCommitted: true, + }, nil + }, + } + adapted := mustAdapt(t, terminal, NewChatToAnthropicBeta(AnthropicOptions{})) + request := &openai.ChatCompletionNewParams{ + Model: openai.ChatModel("client-model"), + Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("use lookup")}, + MaxTokens: openai.Opt(int64(123)), + } + response, err := adapted.Complete(context.Background(), stage.Call{ + Request: request, + Metadata: stage.CallMetadata{RequestID: "reverse-complete", Attempt: 2}, + State: stage.ProtocolState{OpenAIChat: &protocol.OpenAIConfig{HasThinking: true}}, + }) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + chat, ok := response.Value.(wire.ChatCompletionWire) + if !ok { + t.Fatalf("response type = %T, want wire.ChatCompletionWire", response.Value) + } + if chat.Model != "client-model" || response.Model != "client-model" || !response.SideEffectsCommitted { + t.Fatalf("response model/facts = %q %+v", chat.Model, response) + } + if len(chat.Choices) != 1 || len(chat.Choices[0].Message.ToolCalls) != 1 || chat.Choices[0].FinishReason != "tool_calls" { + t.Fatalf("response choices = %+v", chat.Choices) + } + tool := chat.Choices[0].Message.ToolCalls[0] + if tool.ID != "tool_reverse" || tool.Function.Name != "lookup" || !strings.Contains(tool.Function.Arguments, "parallel") { + t.Fatalf("tool call = %+v", tool) + } + if response.Usage == nil || response.Usage.InputTokens != 9 || response.Usage.CacheReadTokens != 2 || response.Usage.OutputTokens != 4 { + t.Fatalf("normalized usage = %+v", response.Usage) + } +} + +func TestOpenAIChatToAnthropicBetaUsesDefaultMaxTokens(t *testing.T) { + t.Parallel() + + bridge := NewChatToAnthropicBeta(AnthropicOptions{DefaultMaxTokens: 777}) + session, err := bridge.Open(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{ + Model: openai.ChatModel("model"), + Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("hello")}, + }}, stage.OperationComplete) + if err != nil { + t.Fatalf("Open() error = %v", err) + } + request := requireBetaRequest(t, session.TargetCall().Request) + if request.MaxTokens != 777 { + t.Fatalf("MaxTokens = %d, want 777", request.MaxTokens) + } +} + +func TestOpenAIChatToAnthropicBetaSeparatesProviderAndResponseModels(t *testing.T) { + t.Parallel() + + bridge := NewChatToAnthropicBeta(AnthropicOptions{ResponseModel: "public-model"}) + session, err := bridge.Open(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{ + Model: openai.ChatModel("provider-model"), + }}, stage.OperationComplete) + if err != nil { + t.Fatalf("Open() error = %v", err) + } + request := requireBetaRequest(t, session.TargetCall().Request) + if request.Model != "provider-model" { + t.Fatalf("target model = %q, want provider-model", request.Model) + } + response, err := session.ConvertComplete(context.Background(), &stage.Response{Value: decodeBetaMessage(t, map[string]any{ + "id": "msg_model", "type": "message", "role": "assistant", "model": "provider-model", + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + "stop_reason": "end_turn", "usage": map[string]any{"input_tokens": 1, "output_tokens": 1}, + })}) + if err != nil { + t.Fatalf("ConvertComplete() error = %v", err) + } + chat := response.Value.(wire.ChatCompletionWire) + if chat.Model != "public-model" || response.Model != "public-model" { + t.Fatalf("source-visible models = %q/%q", chat.Model, response.Model) + } +} + +func TestOpenAIChatToAnthropicBetaStream(t *testing.T) { + t.Parallel() + + targetStream := &memoryStream{ + events: betaStageEvents(t, + map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": "msg_stream", "type": "message", "role": "assistant", + "model": "provider-model", "content": []any{}, + "usage": map[string]any{"input_tokens": 5, "output_tokens": 0}, + }, + }, + map[string]any{"type": "content_block_start", "index": 0, "content_block": map[string]any{"type": "tool_use", "id": "tool_stream", "name": "lookup", "input": map[string]any{}}}, + map[string]any{"type": "content_block_delta", "index": 0, "delta": map[string]any{"type": "input_json_delta", "partial_json": `{"q":"parallel"}`}}, + map[string]any{"type": "content_block_stop", "index": 0}, + map[string]any{"type": "message_delta", "delta": map[string]any{"stop_reason": "tool_use"}, "usage": map[string]any{"output_tokens": 3}}, + map[string]any{"type": "message_stop"}, + ), + result: stage.StreamResult{ + Usage: protocol.NewTokenUsage(999, 999), + Model: "provider-model", + SideEffectsCommitted: true, + }, + } + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(_ context.Context, call stage.Call) (stage.EventStream, error) { + requireBetaRequest(t, call.Request) + if call.Metadata.RequestID != "reverse-stream" { + t.Fatalf("metadata = %+v", call.Metadata) + } + return targetStream, nil + }, + } + adapted := mustAdapt(t, terminal, NewChatToAnthropicBeta(AnthropicOptions{})) + stream, err := adapted.Stream(context.Background(), stage.Call{ + Request: &openai.ChatCompletionNewParams{ + Model: openai.ChatModel("client-stream-model"), + Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("lookup")}, + }, + Metadata: stage.CallMetadata{RequestID: "reverse-stream"}, + }) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + var chunks []wire.ChatStreamChunk + for { + event, err := stream.Next(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("Next() error = %v", err) + } + chunk, ok := event.Value.(wire.ChatStreamChunk) + if !ok { + t.Fatalf("event type = %T", event.Value) + } + chunks = append(chunks, chunk) + } + if len(chunks) != 4 { + t.Fatalf("chunk count = %d, want 4", len(chunks)) + } + if len(chunks[1].Choices[0].Delta.ToolCalls) != 1 || chunks[1].Choices[0].Delta.ToolCalls[0].Function.Name != "lookup" { + t.Fatalf("tool start chunk = %+v", chunks[1]) + } + if len(chunks[2].Choices[0].Delta.ToolCalls) != 1 || chunks[2].Choices[0].Delta.ToolCalls[0].Function.Arguments == nil || !strings.Contains(*chunks[2].Choices[0].Delta.ToolCalls[0].Function.Arguments, "parallel") { + t.Fatalf("tool delta chunk = %+v", chunks[2]) + } + if chunks[3].Choices[0].FinishReason == nil || *chunks[3].Choices[0].FinishReason != "tool_calls" || chunks[3].Usage == nil { + t.Fatalf("final chunk = %+v", chunks[3]) + } + result := stream.Result() + if result.Model != "client-stream-model" || result.Usage == nil || result.Usage.InputTokens != 5 || result.Usage.OutputTokens != 3 || !result.SideEffectsCommitted { + t.Fatalf("Result() = %+v", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("second Close() error = %v", err) + } + if targetStream.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", targetStream.closeCount) + } +} + +func TestOpenAIChatToAnthropicBetaStreamAcceptsNormalizedStageEvents(t *testing.T) { + t.Parallel() + + targetStream := &memoryStream{events: []stage.Event{ + {Value: protocolstream.AnthropicEvent{Type: "message_start", Data: map[string]any{ + "message": map[string]any{ + "id": "msg_normalized", "type": "message", "role": "assistant", + "model": "provider-model", "content": []any{}, + "usage": map[string]any{"input_tokens": 3, "output_tokens": 0}, + }, + }}}, + {Value: protocolstream.AnthropicEvent{Type: "content_block_start", Data: map[string]any{ + "type": "content_block_start", "index": 0, + "content_block": map[string]any{"type": "text", "text": ""}, + }}}, + {Value: protocolstream.AnthropicEvent{Type: "content_block_delta", Data: map[string]any{ + "type": "content_block_delta", "index": 0, + "delta": map[string]any{"type": "text_delta", "text": "parallel path"}, + }}}, + {Value: protocolstream.AnthropicEvent{Type: "content_block_stop", Data: map[string]any{ + "type": "content_block_stop", "index": 0, + }}}, + {Value: protocolstream.AnthropicEvent{Type: "message_delta", Data: map[string]any{ + "type": "message_delta", "delta": map[string]any{"stop_reason": "end_turn"}, + "usage": map[string]any{"output_tokens": 2}, + }}}, + {Value: protocolstream.AnthropicEvent{Type: "message_stop", Data: map[string]any{"type": "message_stop"}}}, + }} + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return targetStream, nil + }, + } + adapted := mustAdapt(t, terminal, NewChatToAnthropicBeta(AnthropicOptions{})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{Model: openai.ChatModel("client-model")}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + var content strings.Builder + for { + event, err := stream.Next(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("Next() error = %v", err) + } + chunk := event.Value.(wire.ChatStreamChunk) + content.WriteString(chunk.Choices[0].Delta.Content) + } + if content.String() != "parallel path" { + t.Fatalf("content = %q", content.String()) + } + if result := stream.Result(); result.Usage == nil || result.Usage.InputTokens != 3 || result.Usage.OutputTokens != 2 { + t.Fatalf("Result() = %+v", result) + } +} + +func TestOpenAIChatToAnthropicBetaErrors(t *testing.T) { + t.Parallel() + + bridge := NewChatToAnthropicBeta(AnthropicOptions{}) + if _, err := bridge.Open(context.Background(), stage.Call{Request: "wrong"}, stage.OperationComplete); err == nil || !strings.Contains(err.Error(), "request has type string") { + t.Fatalf("wrong request error = %v", err) + } + if _, err := bridge.Open(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{}}, stage.Operation(99)); err == nil || !strings.Contains(err.Error(), "unsupported operation") { + t.Fatalf("wrong operation error = %v", err) + } + + want := errors.New("provider unavailable") + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, stage.Call) (*stage.Response, error) { + return nil, want + }, + } + adapted := mustAdapt(t, terminal, bridge) + _, err := adapted.Complete(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{Model: openai.ChatModel("model")}}) + if !errors.Is(err, want) { + t.Fatalf("Complete() error = %v, want provider error", err) + } + + wrongResponse := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, stage.Call) (*stage.Response, error) { + return &stage.Response{Value: "wrong"}, nil + }, + } + adapted = mustAdapt(t, wrongResponse, bridge) + _, err = adapted.Complete(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{Model: openai.ChatModel("model")}}) + if err == nil || !strings.Contains(err.Error(), "value has type string") { + t.Fatalf("wrong response error = %v", err) + } +} + +func TestOpenAIChatToAnthropicBetaStreamDisablesWireUsage(t *testing.T) { + t.Parallel() + + targetStream := &memoryStream{events: betaStageEvents(t, + map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": "msg_no_usage", "type": "message", "role": "assistant", + "model": "provider-model", "content": []any{}, + "usage": map[string]any{"input_tokens": 7, "output_tokens": 0}, + }, + }, + map[string]any{"type": "message_delta", "delta": map[string]any{"stop_reason": "end_turn"}, "usage": map[string]any{"output_tokens": 2}}, + map[string]any{"type": "message_stop"}, + )} + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return targetStream, nil + }, + } + adapted := mustAdapt(t, terminal, NewChatToAnthropicBeta(AnthropicOptions{DisableStreamUsage: true})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{Model: openai.ChatModel("model")}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + var last wire.ChatStreamChunk + for { + event, err := stream.Next(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("Next() error = %v", err) + } + last = event.Value.(wire.ChatStreamChunk) + } + if last.Usage != nil { + t.Fatalf("final wire usage = %+v, want nil", last.Usage) + } + if result := stream.Result(); result.Usage == nil || result.Usage.InputTokens != 7 || result.Usage.OutputTokens != 2 { + t.Fatalf("internal stream result = %+v", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } +} + +func TestOpenAIChatToAnthropicBetaStreamPropagatesTargetError(t *testing.T) { + t.Parallel() + + want := errors.New("target stream failed") + targetStream := &memoryStream{err: want} + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return targetStream, nil + }, + } + adapted := mustAdapt(t, terminal, NewChatToAnthropicBeta(AnthropicOptions{})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{Model: openai.ChatModel("model")}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + if _, err := stream.Next(context.Background()); !errors.Is(err, want) { + t.Fatalf("Next() error = %v, want target error", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if targetStream.closeCount != 1 { + t.Fatalf("target close count = %d, want 1", targetStream.closeCount) + } +} + +func TestOpenAIChatToAnthropicBetaStreamRejectsWrongEventAndHonorsCancellation(t *testing.T) { + t.Parallel() + + targetStream := &memoryStream{events: []stage.Event{{Value: "wrong"}}} + terminal := &memoryEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return targetStream, nil + }, + } + adapted := mustAdapt(t, terminal, NewChatToAnthropicBeta(AnthropicOptions{})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{Model: openai.ChatModel("model")}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + canceled, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := stream.Next(canceled); !errors.Is(err, context.Canceled) { + t.Fatalf("canceled Next() error = %v", err) + } + if targetStream.index != 0 { + t.Fatalf("canceled Next() consumed %d target events", targetStream.index) + } + if _, err := stream.Next(context.Background()); err == nil || !strings.Contains(err.Error(), "event has type string") { + t.Fatalf("wrong event error = %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } +} + +func mustAdapt(t *testing.T, terminal stage.Endpoint, bridge stage.Bridge) stage.Endpoint { + t.Helper() + adapted, err := stage.Adapt(terminal, bridge) + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + return adapted +} + +func requireBetaRequest(t *testing.T, value any) *anthropic.BetaMessageNewParams { + t.Helper() + switch request := value.(type) { + case *anthropic.BetaMessageNewParams: + if request == nil { + t.Fatal("target request is nil") + } + return request + case anthropic.BetaMessageNewParams: + return &request + default: + t.Fatalf("target request type = %T", value) + return nil + } +} + +func decodeBetaMessage(t *testing.T, body map[string]any) *anthropic.BetaMessage { + t.Helper() + raw, err := json.Marshal(body) + if err != nil { + t.Fatalf("marshal fixture: %v", err) + } + var message anthropic.BetaMessage + if err := json.Unmarshal(raw, &message); err != nil { + t.Fatalf("decode fixture: %v", err) + } + return &message +} + +func betaStageEvents(t *testing.T, bodies ...map[string]any) []stage.Event { + t.Helper() + events := make([]stage.Event, 0, len(bodies)) + for _, body := range bodies { + raw, err := json.Marshal(body) + if err != nil { + t.Fatalf("marshal event fixture: %v", err) + } + var event anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(raw, &event); err != nil { + t.Fatalf("decode event fixture: %v", err) + } + events = append(events, stage.Event{Value: event}) + } + return events +} + +type memoryEndpoint struct { + api protocol.APIType + complete func(context.Context, stage.Call) (*stage.Response, error) + stream func(context.Context, stage.Call) (stage.EventStream, error) +} + +func (e *memoryEndpoint) Protocol() protocol.APIType { return e.api } + +func (e *memoryEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + if e.complete == nil { + return nil, errors.New("complete not configured") + } + return e.complete(ctx, call) +} + +func (e *memoryEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + if e.stream == nil { + return nil, errors.New("stream not configured") + } + return e.stream(ctx, call) +} + +type memoryStream struct { + events []stage.Event + index int + err error + result stage.StreamResult + closeCount int +} + +func (s *memoryStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + if s.index >= len(s.events) { + if s.err != nil { + return stage.Event{}, s.err + } + return stage.Event{}, io.EOF + } + event := s.events[s.index] + s.index++ + return event, nil +} + +func (s *memoryStream) Close() error { + s.closeCount++ + return nil +} + +func (s *memoryStream) Result() stage.StreamResult { return s.result } diff --git a/internal/protocol/stage/openaibridge/responses.go b/internal/protocol/stage/openaibridge/responses.go new file mode 100644 index 000000000..f6ea00bc6 --- /dev/null +++ b/internal/protocol/stage/openaibridge/responses.go @@ -0,0 +1,119 @@ +package openaibridge + +import ( + "context" + "fmt" + + "github.com/openai/openai-go/v3/responses" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/nonstream" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" +) + +// ResponsesOptions configures Chat to OpenAI Responses conversion. +type ResponsesOptions struct { + DefaultMaxTokens int64 + DisableStreamUsage bool + ResponseModel string +} + +// NewChatToOpenAIResponses returns an OpenAI Chat to Responses Bridge. +func NewChatToOpenAIResponses(options ResponsesOptions) stage.Bridge { + if options.DefaultMaxTokens <= 0 { + options.DefaultMaxTokens = defaultAnthropicMaxTokens + } + return &responsesTargetBridge{options: options} +} + +type responsesTargetBridge struct { + options ResponsesOptions +} + +func (*responsesTargetBridge) Source() protocol.APIType { return protocol.TypeOpenAIChat } +func (*responsesTargetBridge) Target() protocol.APIType { return protocol.TypeOpenAIResponses } +func (*responsesTargetBridge) Capabilities() stage.Capabilities { + return stage.AllBridgeCapabilities +} + +func (b *responsesTargetBridge) Open(_ context.Context, call stage.Call, operation stage.Operation) (stage.BridgeSession, error) { + switch operation { + case stage.OperationComplete, stage.OperationStream: + default: + return nil, fmt.Errorf("open OpenAI Chat to Responses bridge: unsupported operation %s", operation) + } + sourceRequest, err := chatRequest(call.Request) + if err != nil { + return nil, err + } + targetRequest := request.ConvertChatToOpenAIResponses(sourceRequest, b.options.DefaultMaxTokens) + if targetRequest == nil { + return nil, fmt.Errorf("open OpenAI Chat to Responses bridge: request conversion returned nil") + } + targetCall := call + targetCall.Request = targetRequest + targetCall.State.OpenAIChat = nil + sourceModel := string(sourceRequest.Model) + if b.options.ResponseModel != "" { + sourceModel = b.options.ResponseModel + } + return &responsesTargetSession{ + operation: operation, + targetCall: targetCall, + sourceModel: sourceModel, + disableStreamUsage: b.options.DisableStreamUsage, + }, nil +} + +type responsesTargetSession struct { + operation stage.Operation + targetCall stage.Call + sourceModel string + disableStreamUsage bool +} + +func (s *responsesTargetSession) TargetCall() stage.Call { return s.targetCall } + +func (s *responsesTargetSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + if s.operation != stage.OperationComplete { + return nil, fmt.Errorf("convert OpenAI Responses complete response to Chat: session was opened for %s", s.operation) + } + value, err := openAIResponsesValue(response) + if err != nil { + return nil, err + } + converted := nonstream.ConvertResponsesToOpenAIChat(value, s.sourceModel) + usage := protocolusage.FromOpenAIResponses(value.Usage) + if !usage.HasUsage() { + usage = nil + } + return &stage.Response{Value: converted, Usage: usage, Model: s.sourceModel}, nil +} + +func openAIResponsesValue(response *stage.Response) (*responses.Response, error) { + if response == nil { + return nil, fmt.Errorf("convert OpenAI Responses response to Chat: response is nil") + } + switch value := response.Value.(type) { + case *responses.Response: + if value == nil { + return nil, fmt.Errorf("convert OpenAI Responses response to Chat: value is nil") + } + return value, nil + case responses.Response: + return &value, nil + default: + return nil, fmt.Errorf("convert OpenAI Responses response to Chat: value has type %T, want responses.Response", response.Value) + } +} + +func (s *responsesTargetSession) ConvertStream(ctx context.Context, target stage.EventStream) (stage.EventStream, error) { + if s.operation != stage.OperationStream { + return nil, fmt.Errorf("convert OpenAI Responses stream to Chat: session was opened for %s", s.operation) + } + return newResponsesChatStream(ctx, target, s.sourceModel, s.disableStreamUsage) +} + +func (*responsesTargetSession) ConvertError(_ context.Context, err error) error { return err } diff --git a/internal/protocol/stage/openaibridge/responses_stream.go b/internal/protocol/stage/openaibridge/responses_stream.go new file mode 100644 index 000000000..b07f78b16 --- /dev/null +++ b/internal/protocol/stage/openaibridge/responses_stream.go @@ -0,0 +1,113 @@ +package openaibridge + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + + "github.com/openai/openai-go/v3/responses" + + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func newResponsesChatStream(ctx context.Context, target stage.EventStream, sourceModel string, disableUsage bool) (stage.EventStream, error) { + if target == nil { + return nil, fmt.Errorf("convert OpenAI Responses stream to Chat: target stream is nil") + } + iterator := &openAIResponsesStreamIterator{target: target, ctx: ctx} + return &responsesChatStream{ + iterator: iterator, + converter: protocolstream.NewOpenAIResponsesToChatConverter(iterator, sourceModel, disableUsage), + model: sourceModel, + }, nil +} + +type openAIResponsesStreamIterator struct { + target stage.EventStream + ctx context.Context + current responses.ResponseStreamEventUnion + err error + + closeOnce sync.Once + closeErr error +} + +func (s *openAIResponsesStreamIterator) setContext(ctx context.Context) { s.ctx = ctx } +func (s *openAIResponsesStreamIterator) Next() bool { + if s.err != nil { + return false + } + ctx := s.ctx + if ctx == nil { + ctx = context.Background() + } + event, err := s.target.Next(ctx) + if err != nil { + if !errors.Is(err, io.EOF) { + s.err = err + } + return false + } + switch value := event.Value.(type) { + case responses.ResponseStreamEventUnion: + s.current = value + case *responses.ResponseStreamEventUnion: + if value == nil { + s.err = fmt.Errorf("convert OpenAI Responses stream to Chat: event is nil") + return false + } + s.current = *value + default: + s.err = fmt.Errorf("convert OpenAI Responses stream to Chat: event has type %T", event.Value) + return false + } + return true +} +func (s *openAIResponsesStreamIterator) Current() responses.ResponseStreamEventUnion { + return s.current +} +func (s *openAIResponsesStreamIterator) Err() error { return s.err } +func (s *openAIResponsesStreamIterator) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.target.Close() }) + return s.closeErr +} + +type responsesChatStream struct { + iterator *openAIResponsesStreamIterator + converter protocolstream.StreamConverter + model string +} + +func (s *responsesChatStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + s.iterator.setContext(ctx) + value, done, err := s.converter.Next() + if err != nil { + return stage.Event{}, err + } + if done { + if err := s.iterator.Err(); err != nil { + return stage.Event{}, err + } + return stage.Event{}, io.EOF + } + chunk, ok := value.(wire.ChatStreamChunk) + if !ok { + return stage.Event{}, fmt.Errorf("convert OpenAI Responses stream to Chat: converter emitted %T", value) + } + return stage.Event{Value: chunk}, nil +} +func (s *responsesChatStream) Close() error { return s.iterator.Close() } +func (s *responsesChatStream) Result() stage.StreamResult { + usage := s.converter.Usage() + if usage != nil && !usage.HasUsage() { + usage = nil + } + return stage.StreamResult{Usage: usage, Model: s.model} +} diff --git a/internal/protocol/stage/openaibridge/responses_test.go b/internal/protocol/stage/openaibridge/responses_test.go new file mode 100644 index 000000000..b0c41297f --- /dev/null +++ b/internal/protocol/stage/openaibridge/responses_test.go @@ -0,0 +1,150 @@ +package openaibridge + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/responses" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestOpenAIChatToOpenAIResponsesComplete(t *testing.T) { + t.Parallel() + + terminal := &memoryEndpoint{ + api: protocol.TypeOpenAIResponses, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + request, ok := call.Request.(*responses.ResponseNewParams) + if !ok || request == nil { + t.Fatalf("request type = %T", call.Request) + } + if request.Model != "provider-model" || call.Metadata.RequestID != "chat-responses-complete" { + t.Fatalf("target call = %#v metadata=%+v", request, call.Metadata) + } + return &stage.Response{ + Value: decodeOpenAIResponsesResponse(t, `{ + "id":"resp_1","object":"response","created_at":123,"model":"provider-model","status":"completed", + "output":[{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello from responses","annotations":[]}]}], + "usage":{"input_tokens":9,"output_tokens":4,"total_tokens":13,"input_tokens_details":{"cached_tokens":2},"output_tokens_details":{"reasoning_tokens":0}} + }`), + SideEffectsCommitted: true, + }, nil + }, + } + adapted := mustAdapt(t, terminal, NewChatToOpenAIResponses(ResponsesOptions{ResponseModel: "public-model"})) + result, err := adapted.Complete(context.Background(), stage.Call{ + Request: &openai.ChatCompletionNewParams{ + Model: openai.ChatModel("provider-model"), + Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("hello")}, + }, + Metadata: stage.CallMetadata{RequestID: "chat-responses-complete"}, + }) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + chat, ok := result.Value.(wire.ChatCompletionWire) + if !ok { + t.Fatalf("response type = %T", result.Value) + } + if chat.Model != "public-model" || result.Model != "public-model" || !result.SideEffectsCommitted { + t.Fatalf("response = %+v result=%+v", chat, result) + } + if len(chat.Choices) != 1 || !strings.Contains(chat.Choices[0].Message.Content, "hello from responses") { + t.Fatalf("choices = %+v", chat.Choices) + } + if result.Usage == nil || result.Usage.InputTokens != 7 || result.Usage.CacheReadTokens != 2 || result.Usage.OutputTokens != 4 { + t.Fatalf("usage = %+v", result.Usage) + } +} + +func TestOpenAIChatToOpenAIResponsesStream(t *testing.T) { + t.Parallel() + + target := &memoryStream{ + events: openAIResponsesStageEvents(t, + `{"type":"response.created","sequence_number":0,"response":{"id":"resp_stream","object":"response","model":"provider-model","status":"in_progress","output":[]}}`, + `{"type":"response.output_text.delta","sequence_number":1,"item_id":"msg_1","output_index":0,"content_index":0,"delta":"stream responses"}`, + `{"type":"response.output_text.done","sequence_number":2,"item_id":"msg_1","output_index":0,"content_index":0,"text":"stream responses"}`, + `{"type":"response.completed","sequence_number":3,"response":{"id":"resp_stream","object":"response","model":"provider-model","status":"completed","output":[],"usage":{"input_tokens":6,"output_tokens":2,"total_tokens":8,"input_tokens_details":{"cached_tokens":0},"output_tokens_details":{"reasoning_tokens":0}}}}`, + ), + result: stage.StreamResult{SideEffectsCommitted: true}, + } + terminal := &memoryEndpoint{api: protocol.TypeOpenAIResponses, stream: func(_ context.Context, call stage.Call) (stage.EventStream, error) { + if _, ok := call.Request.(*responses.ResponseNewParams); !ok { + t.Fatalf("request type = %T", call.Request) + } + return target, nil + }} + adapted := mustAdapt(t, terminal, NewChatToOpenAIResponses(ResponsesOptions{ResponseModel: "public-stream-model"})) + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &openai.ChatCompletionNewParams{Model: openai.ChatModel("provider-model")}}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + var sawText, sawFinal bool + for { + event, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + t.Fatalf("Next() error = %v", nextErr) + } + chunk, ok := event.Value.(wire.ChatStreamChunk) + if !ok { + t.Fatalf("event type = %T", event.Value) + } + if len(chunk.Choices) > 0 && strings.Contains(chunk.Choices[0].Delta.Content, "stream responses") { + sawText = true + } + if len(chunk.Choices) > 0 && chunk.Choices[0].FinishReason != nil { + sawFinal = true + } + } + if !sawText || !sawFinal { + t.Fatalf("saw text/final = %v/%v", sawText, sawFinal) + } + result := stream.Result() + if result.Model != "public-stream-model" || result.Usage == nil || result.Usage.InputTokens != 6 || result.Usage.OutputTokens != 2 || !result.SideEffectsCommitted { + t.Fatalf("Result() = %+v", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if err := stream.Close(); err != nil { + t.Fatalf("second Close() error = %v", err) + } + if target.closeCount != 1 { + t.Fatalf("target close count = %d", target.closeCount) + } +} + +func decodeOpenAIResponsesResponse(t *testing.T, raw string) *responses.Response { + t.Helper() + var response responses.Response + if err := json.Unmarshal([]byte(raw), &response); err != nil { + t.Fatalf("decode Responses response: %v", err) + } + return &response +} + +func openAIResponsesStageEvents(t *testing.T, values ...string) []stage.Event { + t.Helper() + events := make([]stage.Event, 0, len(values)) + for _, raw := range values { + var event responses.ResponseStreamEventUnion + if err := json.Unmarshal([]byte(raw), &event); err != nil { + t.Fatalf("decode Responses event: %v", err) + } + events = append(events, stage.Event{Value: event}) + } + return events +} diff --git a/internal/protocol/stage/openaibridge/stream.go b/internal/protocol/stage/openaibridge/stream.go new file mode 100644 index 000000000..9ed8fedbc --- /dev/null +++ b/internal/protocol/stage/openaibridge/stream.go @@ -0,0 +1,162 @@ +package openaibridge + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "sync" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func newChatStream(target stage.EventStream, sourceModel string, disableStreamUsage bool) (stage.EventStream, error) { + if target == nil { + return nil, fmt.Errorf("convert Anthropic Beta stream to OpenAI Chat: target stream is nil") + } + iterator := &anthropicBetaStreamIterator{target: target} + return &chatStream{ + iterator: iterator, + converter: protocolstream.NewAnthropicBetaToOpenAIChatConverter(iterator, sourceModel, disableStreamUsage), + model: sourceModel, + }, nil +} + +type anthropicBetaStreamIterator struct { + target stage.EventStream + ctx context.Context + current anthropic.BetaRawMessageStreamEventUnion + err error + + closeOnce sync.Once + closeErr error +} + +func (s *anthropicBetaStreamIterator) setContext(ctx context.Context) { s.ctx = ctx } + +func (s *anthropicBetaStreamIterator) Next() bool { + if s.err != nil { + return false + } + ctx := s.ctx + if ctx == nil { + ctx = context.Background() + } + event, err := s.target.Next(ctx) + if err != nil { + if !errors.Is(err, io.EOF) { + s.err = err + } + return false + } + + switch value := event.Value.(type) { + case anthropic.BetaRawMessageStreamEventUnion: + s.current = value + case *anthropic.BetaRawMessageStreamEventUnion: + if value == nil { + s.err = fmt.Errorf("convert Anthropic Beta stream to OpenAI Chat: event is nil") + return false + } + s.current = *value + case protocolstream.AnthropicEvent: + if err := s.setNormalizedEvent(value); err != nil { + s.err = err + return false + } + default: + s.err = fmt.Errorf("convert Anthropic Beta stream to OpenAI Chat: event has type %T, want anthropic.BetaRawMessageStreamEventUnion or stream.AnthropicEvent", event.Value) + return false + } + return true +} + +func (s *anthropicBetaStreamIterator) setNormalizedEvent(event protocolstream.AnthropicEvent) error { + raw, err := json.Marshal(event.Data) + if err != nil { + return fmt.Errorf("convert normalized Anthropic event %q: marshal data: %w", event.Type, err) + } + var fields map[string]json.RawMessage + if err := json.Unmarshal(raw, &fields); err != nil || fields == nil { + return fmt.Errorf("convert normalized Anthropic event %q: data must be a JSON object", event.Type) + } + if _, ok := fields["type"]; !ok { + typeJSON, err := json.Marshal(event.Type) + if err != nil { + return fmt.Errorf("convert normalized Anthropic event type: %w", err) + } + fields["type"] = typeJSON + raw, err = json.Marshal(fields) + if err != nil { + return fmt.Errorf("convert normalized Anthropic event %q: marshal envelope: %w", event.Type, err) + } + } + var decoded anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(raw, &decoded); err != nil { + return fmt.Errorf("convert normalized Anthropic event %q: decode Beta event: %w", event.Type, err) + } + if decoded.Type != event.Type { + return fmt.Errorf("convert normalized Anthropic event: envelope type %q does not match data type %q", event.Type, decoded.Type) + } + s.current = decoded + return nil +} + +func (s *anthropicBetaStreamIterator) Current() anthropic.BetaRawMessageStreamEventUnion { + return s.current +} + +func (s *anthropicBetaStreamIterator) Err() error { return s.err } + +func (s *anthropicBetaStreamIterator) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.target.Close() }) + return s.closeErr +} + +type chatStream struct { + iterator *anthropicBetaStreamIterator + converter protocolstream.StreamConverter + model string +} + +func (s *chatStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + s.iterator.setContext(ctx) + value, done, err := s.converter.Next() + if err != nil { + return stage.Event{}, err + } + if done { + if err := s.iterator.Err(); err != nil { + return stage.Event{}, err + } + return stage.Event{}, io.EOF + } + switch chunk := value.(type) { + case wire.ChatStreamChunk: + return stage.Event{Value: chunk}, nil + case *wire.ChatStreamChunk: + if chunk == nil { + return stage.Event{}, fmt.Errorf("convert Anthropic Beta stream to OpenAI Chat: converter emitted a nil chunk") + } + return stage.Event{Value: *chunk}, nil + default: + return stage.Event{}, fmt.Errorf("convert Anthropic Beta stream to OpenAI Chat: converter emitted %T", value) + } +} + +func (s *chatStream) Close() error { return s.iterator.Close() } + +func (s *chatStream) Result() stage.StreamResult { + usage := s.converter.Usage() + if usage != nil && !usage.HasUsage() { + usage = nil + } + return stage.StreamResult{Usage: usage, Model: s.model} +} diff --git a/internal/protocol/stage/registry.go b/internal/protocol/stage/registry.go new file mode 100644 index 000000000..ecd25cb31 --- /dev/null +++ b/internal/protocol/stage/registry.go @@ -0,0 +1,98 @@ +package stage + +import ( + "fmt" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +type bridgeKey struct { + source protocol.APIType + target protocol.APIType +} + +// BridgeRegistry is an immutable exact-pair registry. Build a new registry for +// configuration reloads rather than mutating a registry used by active calls. +type BridgeRegistry struct { + bridges map[bridgeKey]Bridge +} + +// NewBridgeRegistry validates and registers exact source/target pairs. Duplicate +// pairs are rejected so bridge selection never depends on registration order. +func NewBridgeRegistry(bridges ...Bridge) (*BridgeRegistry, error) { + registry := &BridgeRegistry{bridges: make(map[bridgeKey]Bridge, len(bridges))} + for i, bridge := range bridges { + if isNil(bridge) { + return nil, fmt.Errorf("create bridge registry: bridge at index %d is nil", i) + } + source := bridge.Source() + target := bridge.Target() + if source == "" || target == "" { + return nil, fmt.Errorf( + "create bridge registry: bridge at index %d has empty protocol (%q -> %q)", + i, + source, + target, + ) + } + if missing := bridge.Capabilities().Missing(CoreBridgeCapabilities); missing != 0 { + return nil, fmt.Errorf( + "create bridge registry: bridge %q -> %q missing core capabilities: %s", + source, + target, + missing, + ) + } + + key := bridgeKey{source: source, target: target} + if _, exists := registry.bridges[key]; exists { + return nil, fmt.Errorf("create bridge registry: duplicate bridge %q -> %q", source, target) + } + registry.bridges[key] = bridge + } + return registry, nil +} + +// Resolve returns the exact source/target Bridge with the required semantic +// capabilities. Same-protocol pairs use an identity bridge unless an explicit +// same-protocol bridge was registered. +func (r *BridgeRegistry) Resolve(source, target protocol.APIType, required Capabilities) (Bridge, error) { + return r.resolve(source, target, required, true) +} + +// ResolveRegistered returns only a Bridge explicitly registered for the exact +// source/target pair. Unlike Resolve, it does not synthesize an identity Bridge +// for same-protocol pairs. Runtime selectors use this method so enabling a +// production route always requires an intentional registry entry. +func (r *BridgeRegistry) ResolveRegistered(source, target protocol.APIType, required Capabilities) (Bridge, error) { + return r.resolve(source, target, required, false) +} + +func (r *BridgeRegistry) resolve(source, target protocol.APIType, required Capabilities, allowIdentity bool) (Bridge, error) { + if r == nil { + return nil, fmt.Errorf("resolve protocol bridge %q -> %q: registry is nil", source, target) + } + if source == "" || target == "" { + return nil, fmt.Errorf("resolve protocol bridge: protocols must be concrete (%q -> %q)", source, target) + } + + bridge, exists := r.bridges[bridgeKey{source: source, target: target}] + if !exists && allowIdentity && source == target { + bridge = NewIdentityBridge(source) + exists = true + } + if !exists { + return nil, fmt.Errorf("resolve protocol bridge %q -> %q: not registered", source, target) + } + + required |= CoreBridgeCapabilities + if missing := bridge.Capabilities().Missing(required); missing != 0 { + return nil, fmt.Errorf( + "resolve protocol bridge %q -> %q: missing capabilities: %s", + source, + target, + missing, + ) + } + return bridge, nil +} diff --git a/internal/protocol/stage/responsesbridge/bridge.go b/internal/protocol/stage/responsesbridge/bridge.go new file mode 100644 index 000000000..679fc5736 --- /dev/null +++ b/internal/protocol/stage/responsesbridge/bridge.go @@ -0,0 +1,138 @@ +// Package responsesbridge adapts OpenAI Responses calls to provider protocols +// while keeping the outward response in the Responses source protocol. +package responsesbridge + +import ( + "context" + "fmt" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3/responses" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/nonstream" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" +) + +const defaultAnthropicMaxTokens int64 = 4096 + +// AnthropicOptions configures Responses to Anthropic Beta conversion. +type AnthropicOptions struct { + DefaultMaxTokens int64 + // ResponseModel is the source-visible Responses model alias. + ResponseModel string +} + +// NewToAnthropicBeta returns an immutable OpenAI Responses to Anthropic Beta +// Bridge. Mutable response and stream correlation state is created per call. +func NewToAnthropicBeta(options AnthropicOptions) stage.Bridge { + if options.DefaultMaxTokens <= 0 { + options.DefaultMaxTokens = defaultAnthropicMaxTokens + } + return &anthropicBetaBridge{options: options} +} + +type anthropicBetaBridge struct { + options AnthropicOptions +} + +func (*anthropicBetaBridge) Source() protocol.APIType { return protocol.TypeOpenAIResponses } +func (*anthropicBetaBridge) Target() protocol.APIType { return protocol.TypeAnthropicBeta } +func (*anthropicBetaBridge) Capabilities() stage.Capabilities { + return stage.AllBridgeCapabilities +} + +func (b *anthropicBetaBridge) Open(_ context.Context, call stage.Call, operation stage.Operation) (stage.BridgeSession, error) { + switch operation { + case stage.OperationComplete, stage.OperationStream: + default: + return nil, fmt.Errorf("open OpenAI Responses to Anthropic Beta bridge: unsupported operation %s", operation) + } + responsesRequest, err := responsesRequest(call.Request) + if err != nil { + return nil, err + } + targetRequest := request.ConvertOpenAIResponsesToAnthropicBetaRequest(*responsesRequest, b.options.DefaultMaxTokens) + if targetRequest == nil { + return nil, fmt.Errorf("open OpenAI Responses to Anthropic Beta bridge: request conversion returned nil") + } + + targetCall := call + targetCall.Request = targetRequest + targetCall.State.OpenAIChat = nil + sourceModel := string(responsesRequest.Model) + if b.options.ResponseModel != "" { + sourceModel = b.options.ResponseModel + } + return &anthropicBetaSession{ + operation: operation, + targetCall: targetCall, + sourceModel: sourceModel, + }, nil +} + +func responsesRequest(value any) (*responses.ResponseNewParams, error) { + switch request := value.(type) { + case *responses.ResponseNewParams: + if request == nil { + return nil, fmt.Errorf("open OpenAI Responses to Anthropic Beta bridge: request is nil") + } + return request, nil + case responses.ResponseNewParams: + return &request, nil + default: + return nil, fmt.Errorf("open OpenAI Responses to Anthropic Beta bridge: request has type %T, want responses.ResponseNewParams", value) + } +} + +type anthropicBetaSession struct { + operation stage.Operation + targetCall stage.Call + sourceModel string +} + +func (s *anthropicBetaSession) TargetCall() stage.Call { return s.targetCall } + +func (s *anthropicBetaSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + if s.operation != stage.OperationComplete { + return nil, fmt.Errorf("convert Anthropic Beta complete response to OpenAI Responses: session was opened for %s", s.operation) + } + message, err := betaMessage(response) + if err != nil { + return nil, err + } + converted := nonstream.ConvertAnthropicBetaToResponsesWire(message, s.sourceModel, string(message.Model)) + usage := protocolusage.FromAnthropicBetaMessage(message.Usage) + if !usage.HasUsage() { + usage = nil + } + return &stage.Response{Value: converted, Usage: usage, Model: s.sourceModel}, nil +} + +func betaMessage(response *stage.Response) (*anthropic.BetaMessage, error) { + if response == nil { + return nil, fmt.Errorf("convert Anthropic Beta response to OpenAI Responses: response is nil") + } + switch value := response.Value.(type) { + case *anthropic.BetaMessage: + if value == nil { + return nil, fmt.Errorf("convert Anthropic Beta response to OpenAI Responses: value is nil") + } + return value, nil + case anthropic.BetaMessage: + return &value, nil + default: + return nil, fmt.Errorf("convert Anthropic Beta response to OpenAI Responses: value has type %T, want anthropic.BetaMessage", response.Value) + } +} + +func (s *anthropicBetaSession) ConvertStream(_ context.Context, target stage.EventStream) (stage.EventStream, error) { + if s.operation != stage.OperationStream { + return nil, fmt.Errorf("convert Anthropic Beta stream to OpenAI Responses: session was opened for %s", s.operation) + } + return newResponsesStream(target, s.sourceModel) +} + +func (*anthropicBetaSession) ConvertError(_ context.Context, err error) error { return err } diff --git a/internal/protocol/stage/responsesbridge/bridge_test.go b/internal/protocol/stage/responsesbridge/bridge_test.go new file mode 100644 index 000000000..b61a5a7a0 --- /dev/null +++ b/internal/protocol/stage/responsesbridge/bridge_test.go @@ -0,0 +1,203 @@ +package responsesbridge + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3/packages/param" + "github.com/openai/openai-go/v3/responses" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestResponsesToAnthropicBetaComplete(t *testing.T) { + t.Parallel() + + terminal := &memoryEndpoint{ + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + request, ok := call.Request.(*anthropic.BetaMessageNewParams) + if !ok || request == nil { + t.Fatalf("request type = %T", call.Request) + } + if request.Model != "provider-model" || request.MaxTokens != 321 { + t.Fatalf("provider request model/max = %q/%d", request.Model, request.MaxTokens) + } + return &stage.Response{ + Value: decodeBetaMessage(t, map[string]any{ + "id": "msg_1", "type": "message", "role": "assistant", + "model": "provider-model", "stop_reason": "end_turn", + "content": []any{map[string]any{"type": "text", "text": "hello from beta"}}, + "usage": map[string]any{"input_tokens": 7, "output_tokens": 3, "cache_read_input_tokens": 2}, + }), + SideEffectsCommitted: true, + }, nil + }, + } + adapted, err := stage.Adapt(terminal, NewToAnthropicBeta(AnthropicOptions{ResponseModel: "public-model"})) + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + result, err := adapted.Complete(context.Background(), stage.Call{Request: &responses.ResponseNewParams{ + Model: "provider-model", + Input: responses.ResponseNewParamsInputUnion{OfString: param.NewOpt("hello")}, + MaxOutputTokens: param.NewOpt(int64(321)), + }}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + response, ok := result.Value.(wire.ResponsesWireResponse) + if !ok { + t.Fatalf("response type = %T", result.Value) + } + encoded, err := json.Marshal(response) + if err != nil { + t.Fatalf("marshal response: %v", err) + } + if response.Model != "public-model" || !strings.Contains(string(encoded), "hello from beta") { + t.Fatalf("response = %#v json=%s", response, encoded) + } + if result.Usage == nil || result.Usage.InputTokens != 7 || result.Usage.CacheReadTokens != 2 || result.Usage.OutputTokens != 3 { + t.Fatalf("usage = %#v", result.Usage) + } + if !result.SideEffectsCommitted { + t.Fatal("side effects were not preserved") + } +} + +func TestResponsesToAnthropicBetaStream(t *testing.T) { + t.Parallel() + + target := &memoryStream{events: betaEvents(t, + map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": "msg_stream", "type": "message", "role": "assistant", + "model": "provider-model", "content": []any{}, + "usage": map[string]any{"input_tokens": 5, "output_tokens": 0}, + }, + }, + map[string]any{"type": "content_block_start", "index": 0, "content_block": map[string]any{"type": "text", "text": ""}}, + map[string]any{"type": "content_block_delta", "index": 0, "delta": map[string]any{"type": "text_delta", "text": "stream text"}}, + map[string]any{"type": "content_block_stop", "index": 0}, + map[string]any{"type": "message_delta", "delta": map[string]any{"stop_reason": "end_turn"}, "usage": map[string]any{"output_tokens": 2}}, + map[string]any{"type": "message_stop"}, + )} + terminal := &memoryEndpoint{stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return target, nil + }} + adapted, err := stage.Adapt(terminal, NewToAnthropicBeta(AnthropicOptions{ResponseModel: "public-stream-model"})) + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &responses.ResponseNewParams{ + Model: "provider-model", + Input: responses.ResponseNewParamsInputUnion{OfString: param.NewOpt("hello")}, + }}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + var types []string + var sawText bool + for { + event, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + t.Fatalf("Next() error = %v", nextErr) + } + responseEvent, ok := event.Value.(wire.ResponsesEvent) + if !ok { + t.Fatalf("event type = %T", event.Value) + } + types = append(types, responseEvent.EventType()) + if responseEvent.EventType() == "response.output_text.delta" { + encoded, _ := json.Marshal(responseEvent) + sawText = strings.Contains(string(encoded), "stream text") + } + } + if !sawText || len(types) == 0 || types[0] != "response.created" || types[len(types)-1] != "response.completed" { + t.Fatalf("event types = %v, saw text = %v", types, sawText) + } + result := stream.Result() + if result.Model != "public-stream-model" || result.Usage == nil || result.Usage.InputTokens != 5 || result.Usage.OutputTokens != 2 { + t.Fatalf("Result() = %+v", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if target.closeCount != 1 { + t.Fatalf("target close count = %d", target.closeCount) + } +} + +type memoryEndpoint struct { + complete func(context.Context, stage.Call) (*stage.Response, error) + stream func(context.Context, stage.Call) (stage.EventStream, error) +} + +func (*memoryEndpoint) Protocol() protocol.APIType { return protocol.TypeAnthropicBeta } +func (e *memoryEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + return e.complete(ctx, call) +} +func (e *memoryEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + return e.stream(ctx, call) +} + +type memoryStream struct { + events []stage.Event + index int + closeCount int +} + +func (s *memoryStream) Next(context.Context) (stage.Event, error) { + if s.index >= len(s.events) { + return stage.Event{}, io.EOF + } + event := s.events[s.index] + s.index++ + return event, nil +} +func (s *memoryStream) Close() error { + s.closeCount++ + return nil +} +func (*memoryStream) Result() stage.StreamResult { return stage.StreamResult{} } + +func decodeBetaMessage(t *testing.T, value map[string]any) *anthropic.BetaMessage { + t.Helper() + raw, err := json.Marshal(value) + if err != nil { + t.Fatalf("marshal Beta message: %v", err) + } + var message anthropic.BetaMessage + if err := json.Unmarshal(raw, &message); err != nil { + t.Fatalf("decode Beta message: %v", err) + } + return &message +} + +func betaEvents(t *testing.T, values ...map[string]any) []stage.Event { + t.Helper() + events := make([]stage.Event, 0, len(values)) + for _, value := range values { + raw, err := json.Marshal(value) + if err != nil { + t.Fatalf("marshal Beta event: %v", err) + } + var event anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(raw, &event); err != nil { + t.Fatalf("decode Beta event: %v", err) + } + events = append(events, stage.Event{Value: event}) + } + return events +} diff --git a/internal/protocol/stage/responsesbridge/chat.go b/internal/protocol/stage/responsesbridge/chat.go new file mode 100644 index 000000000..6e64b84d5 --- /dev/null +++ b/internal/protocol/stage/responsesbridge/chat.go @@ -0,0 +1,120 @@ +package responsesbridge + +import ( + "context" + "fmt" + + "github.com/openai/openai-go/v3" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/nonstream" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" +) + +// ChatOptions configures Responses to OpenAI Chat conversion. +type ChatOptions struct { + DefaultMaxTokens int64 + ResponseModel string +} + +// NewToOpenAIChat returns an immutable OpenAI Responses to Chat Completions +// Bridge with per-call reverse conversion state. +func NewToOpenAIChat(options ChatOptions) stage.Bridge { + if options.DefaultMaxTokens <= 0 { + options.DefaultMaxTokens = defaultAnthropicMaxTokens + } + return &openAIChatBridge{options: options} +} + +type openAIChatBridge struct { + options ChatOptions +} + +func (*openAIChatBridge) Source() protocol.APIType { return protocol.TypeOpenAIResponses } +func (*openAIChatBridge) Target() protocol.APIType { return protocol.TypeOpenAIChat } +func (*openAIChatBridge) Capabilities() stage.Capabilities { + return stage.AllBridgeCapabilities +} + +func (b *openAIChatBridge) Open(_ context.Context, call stage.Call, operation stage.Operation) (stage.BridgeSession, error) { + switch operation { + case stage.OperationComplete, stage.OperationStream: + default: + return nil, fmt.Errorf("open OpenAI Responses to Chat bridge: unsupported operation %s", operation) + } + responsesRequest, err := responsesRequest(call.Request) + if err != nil { + return nil, err + } + targetRequest := request.ConvertOpenAIResponsesToChat(responsesRequest, b.options.DefaultMaxTokens) + if targetRequest == nil { + return nil, fmt.Errorf("open OpenAI Responses to Chat bridge: request conversion returned nil") + } + targetCall := call + targetCall.Request = targetRequest + targetCall.State.OpenAIChat = &protocol.OpenAIConfig{ + HasThinking: false, + ReasoningEffort: "none", + } + sourceModel := string(responsesRequest.Model) + if b.options.ResponseModel != "" { + sourceModel = b.options.ResponseModel + } + return &openAIChatSession{ + operation: operation, + targetCall: targetCall, + sourceModel: sourceModel, + }, nil +} + +type openAIChatSession struct { + operation stage.Operation + targetCall stage.Call + sourceModel string +} + +func (s *openAIChatSession) TargetCall() stage.Call { return s.targetCall } + +func (s *openAIChatSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + if s.operation != stage.OperationComplete { + return nil, fmt.Errorf("convert Chat complete response to OpenAI Responses: session was opened for %s", s.operation) + } + completion, err := chatCompletion(response) + if err != nil { + return nil, err + } + converted := nonstream.ConvertChatToResponsesWire(completion, s.sourceModel, string(completion.Model)) + usage := protocolusage.FromOpenAIChatCompletion(completion.Usage) + if !usage.HasUsage() { + usage = nil + } + return &stage.Response{Value: converted, Usage: usage, Model: s.sourceModel}, nil +} + +func chatCompletion(response *stage.Response) (*openai.ChatCompletion, error) { + if response == nil { + return nil, fmt.Errorf("convert Chat response to OpenAI Responses: response is nil") + } + switch value := response.Value.(type) { + case *openai.ChatCompletion: + if value == nil { + return nil, fmt.Errorf("convert Chat response to OpenAI Responses: value is nil") + } + return value, nil + case openai.ChatCompletion: + return &value, nil + default: + return nil, fmt.Errorf("convert Chat response to OpenAI Responses: value has type %T, want openai.ChatCompletion", response.Value) + } +} + +func (s *openAIChatSession) ConvertStream(_ context.Context, target stage.EventStream) (stage.EventStream, error) { + if s.operation != stage.OperationStream { + return nil, fmt.Errorf("convert Chat stream to OpenAI Responses: session was opened for %s", s.operation) + } + return newChatResponsesStream(target, s.sourceModel) +} + +func (*openAIChatSession) ConvertError(_ context.Context, err error) error { return err } diff --git a/internal/protocol/stage/responsesbridge/chat_stream.go b/internal/protocol/stage/responsesbridge/chat_stream.go new file mode 100644 index 000000000..2f8acb1a8 --- /dev/null +++ b/internal/protocol/stage/responsesbridge/chat_stream.go @@ -0,0 +1,111 @@ +package responsesbridge + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + + "github.com/openai/openai-go/v3" + + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func newChatResponsesStream(target stage.EventStream, sourceModel string) (stage.EventStream, error) { + if target == nil { + return nil, fmt.Errorf("convert Chat stream to OpenAI Responses: target stream is nil") + } + iterator := &openAIChatStreamIterator{target: target} + return &chatResponsesStream{ + iterator: iterator, + converter: protocolstream.NewChatToResponsesConverter(iterator, sourceModel), + model: sourceModel, + }, nil +} + +type openAIChatStreamIterator struct { + target stage.EventStream + ctx context.Context + current openai.ChatCompletionChunk + err error + + closeOnce sync.Once + closeErr error +} + +func (s *openAIChatStreamIterator) setContext(ctx context.Context) { s.ctx = ctx } +func (s *openAIChatStreamIterator) Next() bool { + if s.err != nil { + return false + } + ctx := s.ctx + if ctx == nil { + ctx = context.Background() + } + event, err := s.target.Next(ctx) + if err != nil { + if !errors.Is(err, io.EOF) { + s.err = err + } + return false + } + switch value := event.Value.(type) { + case openai.ChatCompletionChunk: + s.current = value + case *openai.ChatCompletionChunk: + if value == nil { + s.err = fmt.Errorf("convert Chat stream to OpenAI Responses: event is nil") + return false + } + s.current = *value + default: + s.err = fmt.Errorf("convert Chat stream to OpenAI Responses: event has type %T", event.Value) + return false + } + return true +} +func (s *openAIChatStreamIterator) Current() openai.ChatCompletionChunk { return s.current } +func (s *openAIChatStreamIterator) Err() error { return s.err } +func (s *openAIChatStreamIterator) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.target.Close() }) + return s.closeErr +} + +type chatResponsesStream struct { + iterator *openAIChatStreamIterator + converter protocolstream.StreamConverter + model string +} + +func (s *chatResponsesStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + s.iterator.setContext(ctx) + value, done, err := s.converter.Next() + if err != nil { + return stage.Event{}, err + } + if done { + if err := s.iterator.Err(); err != nil { + return stage.Event{}, err + } + return stage.Event{}, io.EOF + } + event, ok := value.(wire.ResponsesEvent) + if !ok { + return stage.Event{}, fmt.Errorf("convert Chat stream to OpenAI Responses: converter emitted %T", value) + } + return stage.Event{Value: event}, nil +} +func (s *chatResponsesStream) Close() error { return s.iterator.Close() } +func (s *chatResponsesStream) Result() stage.StreamResult { + usage := s.converter.Usage() + if usage != nil && !usage.HasUsage() { + usage = nil + } + return stage.StreamResult{Usage: usage, Model: s.model} +} diff --git a/internal/protocol/stage/responsesbridge/chat_test.go b/internal/protocol/stage/responsesbridge/chat_test.go new file mode 100644 index 000000000..1589d41d9 --- /dev/null +++ b/internal/protocol/stage/responsesbridge/chat_test.go @@ -0,0 +1,186 @@ +package responsesbridge + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/packages/param" + "github.com/openai/openai-go/v3/responses" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestResponsesToOpenAIChatComplete(t *testing.T) { + t.Parallel() + + terminal := &chatMemoryEndpoint{complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + request, ok := call.Request.(*openai.ChatCompletionNewParams) + if !ok || request == nil { + t.Fatalf("request type = %T", call.Request) + } + if request.Model != "provider-model" || !request.MaxTokens.Valid() || request.MaxTokens.Value != 222 { + t.Fatalf("provider request = %#v", request) + } + if call.State.OpenAIChat == nil { + t.Fatal("OpenAI Chat state was not populated") + } + return &stage.Response{ + Value: decodeChatCompletion(t, map[string]any{ + "id": "chatcmpl_1", "object": "chat.completion", "model": "provider-model", + "choices": []any{map[string]any{ + "index": 0, "finish_reason": "stop", + "message": map[string]any{"role": "assistant", "content": "hello from chat"}, + }}, + "usage": map[string]any{"prompt_tokens": 8, "completion_tokens": 4, "total_tokens": 12}, + }), + SideEffectsCommitted: true, + }, nil + }} + adapted, err := stage.Adapt(terminal, NewToOpenAIChat(ChatOptions{ResponseModel: "public-model"})) + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + result, err := adapted.Complete(context.Background(), stage.Call{Request: &responses.ResponseNewParams{ + Model: "provider-model", + Input: responses.ResponseNewParamsInputUnion{OfString: param.NewOpt("hello")}, + MaxOutputTokens: param.NewOpt(int64(222)), + }}) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + response, ok := result.Value.(wire.ResponsesWireResponse) + if !ok { + t.Fatalf("response type = %T", result.Value) + } + encoded, err := json.Marshal(response) + if err != nil { + t.Fatalf("marshal response: %v", err) + } + if response.Model != "public-model" || !strings.Contains(string(encoded), "hello from chat") { + t.Fatalf("response = %#v json=%s", response, encoded) + } + if result.Usage == nil || result.Usage.InputTokens != 8 || result.Usage.OutputTokens != 4 { + t.Fatalf("usage = %#v", result.Usage) + } + if !result.SideEffectsCommitted { + t.Fatal("side effects were not preserved") + } +} + +func TestResponsesToOpenAIChatStream(t *testing.T) { + t.Parallel() + + target := &memoryStream{events: chatEvents(t, + map[string]any{ + "id": "chatcmpl_stream", "object": "chat.completion.chunk", "model": "provider-model", + "choices": []any{map[string]any{"index": 0, "delta": map[string]any{"role": "assistant"}, "finish_reason": nil}}, + }, + map[string]any{ + "id": "chatcmpl_stream", "object": "chat.completion.chunk", "model": "provider-model", + "choices": []any{map[string]any{"index": 0, "delta": map[string]any{"content": "stream chat"}, "finish_reason": nil}}, + }, + map[string]any{ + "id": "chatcmpl_stream", "object": "chat.completion.chunk", "model": "provider-model", + "choices": []any{map[string]any{"index": 0, "delta": map[string]any{}, "finish_reason": "stop"}}, + "usage": map[string]any{"prompt_tokens": 6, "completion_tokens": 2, "total_tokens": 8}, + }, + )} + terminal := &chatMemoryEndpoint{stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return target, nil + }} + adapted, err := stage.Adapt(terminal, NewToOpenAIChat(ChatOptions{ResponseModel: "public-stream-model"})) + if err != nil { + t.Fatalf("Adapt() error = %v", err) + } + stream, err := adapted.Stream(context.Background(), stage.Call{Request: &responses.ResponseNewParams{ + Model: "provider-model", + Input: responses.ResponseNewParamsInputUnion{OfString: param.NewOpt("hello")}, + }}) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + + var types []string + var sawText bool + for { + event, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + t.Fatalf("Next() error = %v", nextErr) + } + responseEvent, ok := event.Value.(wire.ResponsesEvent) + if !ok { + t.Fatalf("event type = %T", event.Value) + } + types = append(types, responseEvent.EventType()) + if responseEvent.EventType() == "response.output_text.delta" { + encoded, _ := json.Marshal(responseEvent) + sawText = strings.Contains(string(encoded), "stream chat") + } + } + if !sawText || len(types) == 0 || types[0] != "response.created" || types[len(types)-1] != "response.completed" { + t.Fatalf("event types = %v, saw text = %v", types, sawText) + } + result := stream.Result() + if result.Model != "public-stream-model" || result.Usage == nil || result.Usage.InputTokens != 6 || result.Usage.OutputTokens != 2 { + t.Fatalf("Result() = %+v", result) + } + if err := stream.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if target.closeCount != 1 { + t.Fatalf("target close count = %d", target.closeCount) + } +} + +type chatMemoryEndpoint struct { + complete func(context.Context, stage.Call) (*stage.Response, error) + stream func(context.Context, stage.Call) (stage.EventStream, error) +} + +func (*chatMemoryEndpoint) Protocol() protocol.APIType { return protocol.TypeOpenAIChat } +func (e *chatMemoryEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + return e.complete(ctx, call) +} +func (e *chatMemoryEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + return e.stream(ctx, call) +} + +func decodeChatCompletion(t *testing.T, value map[string]any) *openai.ChatCompletion { + t.Helper() + raw, err := json.Marshal(value) + if err != nil { + t.Fatalf("marshal Chat response: %v", err) + } + var response openai.ChatCompletion + if err := json.Unmarshal(raw, &response); err != nil { + t.Fatalf("decode Chat response: %v", err) + } + return &response +} + +func chatEvents(t *testing.T, values ...map[string]any) []stage.Event { + t.Helper() + events := make([]stage.Event, 0, len(values)) + for _, value := range values { + raw, err := json.Marshal(value) + if err != nil { + t.Fatalf("marshal Chat event: %v", err) + } + var event openai.ChatCompletionChunk + if err := json.Unmarshal(raw, &event); err != nil { + t.Fatalf("decode Chat event: %v", err) + } + events = append(events, stage.Event{Value: event}) + } + return events +} diff --git a/internal/protocol/stage/responsesbridge/stream.go b/internal/protocol/stage/responsesbridge/stream.go new file mode 100644 index 000000000..6752be165 --- /dev/null +++ b/internal/protocol/stage/responsesbridge/stream.go @@ -0,0 +1,153 @@ +package responsesbridge + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "sync" + + "github.com/anthropics/anthropic-sdk-go" + + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func newResponsesStream(target stage.EventStream, sourceModel string) (stage.EventStream, error) { + if target == nil { + return nil, fmt.Errorf("convert Anthropic Beta stream to OpenAI Responses: target stream is nil") + } + iterator := &anthropicBetaStreamIterator{target: target} + return &responsesStream{ + iterator: iterator, + converter: protocolstream.NewAnthropicBetaToOpenAIResponsesConverter(iterator, sourceModel), + model: sourceModel, + }, nil +} + +type anthropicBetaStreamIterator struct { + target stage.EventStream + ctx context.Context + current anthropic.BetaRawMessageStreamEventUnion + err error + + closeOnce sync.Once + closeErr error +} + +func (s *anthropicBetaStreamIterator) setContext(ctx context.Context) { s.ctx = ctx } + +func (s *anthropicBetaStreamIterator) Next() bool { + if s.err != nil { + return false + } + ctx := s.ctx + if ctx == nil { + ctx = context.Background() + } + event, err := s.target.Next(ctx) + if err != nil { + if !errors.Is(err, io.EOF) { + s.err = err + } + return false + } + switch value := event.Value.(type) { + case anthropic.BetaRawMessageStreamEventUnion: + s.current = value + case *anthropic.BetaRawMessageStreamEventUnion: + if value == nil { + s.err = fmt.Errorf("convert Anthropic Beta stream to OpenAI Responses: event is nil") + return false + } + s.current = *value + case protocolstream.AnthropicEvent: + if err := s.setNormalizedEvent(value); err != nil { + s.err = err + return false + } + default: + s.err = fmt.Errorf("convert Anthropic Beta stream to OpenAI Responses: event has type %T", event.Value) + return false + } + return true +} + +func (s *anthropicBetaStreamIterator) setNormalizedEvent(event protocolstream.AnthropicEvent) error { + raw, err := json.Marshal(event.Data) + if err != nil { + return fmt.Errorf("convert normalized Anthropic event %q: marshal data: %w", event.Type, err) + } + var fields map[string]json.RawMessage + if err := json.Unmarshal(raw, &fields); err != nil || fields == nil { + return fmt.Errorf("convert normalized Anthropic event %q: data must be a JSON object", event.Type) + } + if _, ok := fields["type"]; !ok { + fields["type"], err = json.Marshal(event.Type) + if err != nil { + return fmt.Errorf("convert normalized Anthropic event type: %w", err) + } + raw, err = json.Marshal(fields) + if err != nil { + return fmt.Errorf("convert normalized Anthropic event %q: marshal envelope: %w", event.Type, err) + } + } + var decoded anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(raw, &decoded); err != nil { + return fmt.Errorf("convert normalized Anthropic event %q: decode Beta event: %w", event.Type, err) + } + if decoded.Type != event.Type { + return fmt.Errorf("convert normalized Anthropic event: envelope type %q does not match data type %q", event.Type, decoded.Type) + } + s.current = decoded + return nil +} + +func (s *anthropicBetaStreamIterator) Current() anthropic.BetaRawMessageStreamEventUnion { + return s.current +} +func (s *anthropicBetaStreamIterator) Err() error { return s.err } +func (s *anthropicBetaStreamIterator) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.target.Close() }) + return s.closeErr +} + +type responsesStream struct { + iterator *anthropicBetaStreamIterator + converter protocolstream.StreamConverter + model string +} + +func (s *responsesStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + s.iterator.setContext(ctx) + value, done, err := s.converter.Next() + if err != nil { + return stage.Event{}, err + } + if done { + if err := s.iterator.Err(); err != nil { + return stage.Event{}, err + } + return stage.Event{}, io.EOF + } + event, ok := value.(wire.ResponsesEvent) + if !ok { + return stage.Event{}, fmt.Errorf("convert Anthropic Beta stream to OpenAI Responses: converter emitted %T", value) + } + return stage.Event{Value: event}, nil +} + +func (s *responsesStream) Close() error { return s.iterator.Close() } + +func (s *responsesStream) Result() stage.StreamResult { + usage := s.converter.Usage() + if usage != nil && !usage.HasUsage() { + usage = nil + } + return stage.StreamResult{Usage: usage, Model: s.model} +} diff --git a/internal/protocol/stage/toolloop/openai_chat.go b/internal/protocol/stage/toolloop/openai_chat.go new file mode 100644 index 000000000..1d92f74f4 --- /dev/null +++ b/internal/protocol/stage/toolloop/openai_chat.go @@ -0,0 +1,324 @@ +package toolloop + +import ( + "context" + "errors" + "fmt" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/packages/param" + "github.com/openai/openai-go/v3/shared" + + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +const defaultMaxRounds = 8 + +// OpenAIChatConfig constructs a Chat-native ToolLoop Stage. Catalog, policy, +// and executor are protocol-neutral dependencies; only this adapter understands +// OpenAI Chat request and response types. +type OpenAIChatConfig struct { + Name string + Catalog ToolCatalog + Policy ToolPolicy + Executor ToolExecutor + MaxRounds int +} + +// NewOpenAIChat returns a full-duplex Stage without attaching it to production +// routing. BuildTopology may later insert Bridges around this Chat-native level. +func NewOpenAIChat(config OpenAIChatConfig) (protocolstage.Stage, error) { + if err := validateDependencies(config.Catalog, config.Executor); err != nil { + return nil, fmt.Errorf("construct OpenAI Chat ToolLoop Stage: %w", err) + } + name := config.Name + if name == "" { + name = "tool_loop_openai_chat" + } + policy := config.Policy + if policy == nil { + policy = AllowAllPolicy{} + } + maxRounds := config.MaxRounds + if maxRounds <= 0 { + maxRounds = defaultMaxRounds + } + return &openAIChatStage{ + name: name, + catalog: config.Catalog, + policy: policy, + executor: config.Executor, + maxRounds: maxRounds, + }, nil +} + +type openAIChatStage struct { + name string + catalog ToolCatalog + policy ToolPolicy + executor ToolExecutor + maxRounds int +} + +func (s *openAIChatStage) Name() string { return s.name } +func (*openAIChatStage) Protocol() protocol.APIType { return protocol.TypeOpenAIChat } +func (s *openAIChatStage) Wrap(next protocolstage.Endpoint) protocolstage.Endpoint { + return &openAIChatEndpoint{stage: s, next: next} +} + +type openAIChatEndpoint struct { + stage *openAIChatStage + next protocolstage.Endpoint +} + +func (*openAIChatEndpoint) Protocol() protocol.APIType { return protocol.TypeOpenAIChat } + +func (e *openAIChatEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + prepared, owned, err := e.prepare(ctx, call) + if err != nil { + return nil, err + } + + runCtx := ctx + current := prepared + var totalUsage *protocol.TokenUsage + sideEffectsCommitted := false + for round := 1; round <= e.stage.maxRounds; round++ { + response, callErr := e.next.Complete(runCtx, current) + if callErr != nil { + return nil, WrapError(callErr, sideEffectsCommitted) + } + if response == nil { + return nil, WrapError(errors.New("OpenAI Chat ToolLoop received a nil response"), sideEffectsCommitted) + } + totalUsage = mergeTokenUsage(totalUsage, response.Usage) + sideEffectsCommitted = sideEffectsCommitted || response.SideEffectsCommitted + + roundResponse, parseErr := parseChatRound(response.Value) + if parseErr != nil { + return nil, WrapError(parseErr, sideEffectsCommitted) + } + if !allCallsOwned(roundResponse.calls, owned) { + response.Usage = totalUsage + response.SideEffectsCommitted = sideEffectsCommitted + return response, nil + } + if round == e.stage.maxRounds { + return nil, WrapError(ErrMaxRounds, sideEffectsCommitted) + } + + results, nextCtx, committed, executeErr := e.executeCalls(runCtx, roundResponse.calls) + sideEffectsCommitted = sideEffectsCommitted || committed + if executeErr != nil { + return nil, WrapError(executeErr, sideEffectsCommitted) + } + runCtx = nextCtx + nextRequest, appendErr := appendChatToolResults(current.Request, roundResponse.assistant, results) + if appendErr != nil { + return nil, WrapError(appendErr, sideEffectsCommitted) + } + current.Request = nextRequest + } + return nil, WrapError(ErrMaxRounds, sideEffectsCommitted) +} + +func (e *openAIChatEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + prepared, owned, err := e.prepare(ctx, call) + if err != nil { + return nil, err + } + return newOpenAIChatToolLoopStream(ctx, e, prepared, owned) +} + +func (e *openAIChatEndpoint) prepare(ctx context.Context, call protocolstage.Call) (protocolstage.Call, map[string]struct{}, error) { + request, ok := call.Request.(*openai.ChatCompletionNewParams) + if !ok || request == nil { + return protocolstage.Call{}, nil, fmt.Errorf("OpenAI Chat ToolLoop received request %T", call.Request) + } + definitions, err := e.stage.catalog.ListTools(ctx) + if err != nil { + return protocolstage.Call{}, nil, fmt.Errorf("list ToolLoop catalog: %w", err) + } + if err := validateDefinitions(definitions); err != nil { + return protocolstage.Call{}, nil, err + } + + cloned := *request + cloned.Messages = append([]openai.ChatCompletionMessageParamUnion(nil), request.Messages...) + cloned.Tools = append([]openai.ChatCompletionToolUnionParam(nil), request.Tools...) + existing := make(map[string]struct{}, len(cloned.Tools)) + for _, tool := range cloned.Tools { + if function := tool.GetFunction(); function != nil && function.Name != "" { + existing[function.Name] = struct{}{} + } + } + owned := make(map[string]struct{}, len(definitions)) + for _, definition := range definitions { + if _, exists := existing[definition.Name]; exists { + return protocolstage.Call{}, nil, fmt.Errorf("%w: %q", ErrToolNameCollision, definition.Name) + } + owned[definition.Name] = struct{}{} + function := shared.FunctionDefinitionParam{ + Name: definition.Name, + Parameters: cloneParameters(definition.Parameters), + } + if definition.Description != "" { + function.Description = param.NewOpt(definition.Description) + } + cloned.Tools = append(cloned.Tools, openai.ChatCompletionFunctionTool(function)) + } + prepared := call + prepared.Request = &cloned + return prepared, owned, nil +} + +func (e *openAIChatEndpoint) executeCalls(ctx context.Context, calls []ToolCall) ([]ToolResult, context.Context, bool, error) { + for _, call := range calls { + if err := e.stage.policy.Authorize(ctx, call); err != nil { + return nil, ctx, false, fmt.Errorf("authorize tool %q: %w", call.Name, err) + } + } + + results := make([]ToolResult, 0, len(calls)) + runCtx := ctx + committed := false + for _, call := range calls { + nextCtx, result, err := e.stage.executor.Execute(runCtx, call) + if nextCtx != nil { + runCtx = nextCtx + } + if result.ToolCallID == "" { + result.ToolCallID = call.ID + } + if err != nil { + result.IsError = true + if result.Content == "" { + result.Content = err.Error() + } + } else { + committed = true + } + results = append(results, result) + } + return results, runCtx, committed, nil +} + +type chatRound struct { + assistant openai.ChatCompletionMessageParamUnion + calls []ToolCall +} + +func parseChatRound(value any) (chatRound, error) { + switch response := value.(type) { + case openai.ChatCompletion: + return parseSDKChatRound(&response) + case *openai.ChatCompletion: + return parseSDKChatRound(response) + case wire.ChatCompletionWire: + return parseWireChatRound(&response) + case *wire.ChatCompletionWire: + return parseWireChatRound(response) + default: + return chatRound{}, fmt.Errorf("OpenAI Chat ToolLoop received response %T", value) + } +} + +func parseSDKChatRound(response *openai.ChatCompletion) (chatRound, error) { + if response == nil || len(response.Choices) == 0 { + return chatRound{assistant: openai.AssistantMessage("")}, nil + } + message := response.Choices[0].Message + calls := make([]ToolCall, 0, len(message.ToolCalls)) + for _, call := range message.ToolCalls { + if call.Type != "function" { + continue + } + calls = append(calls, ToolCall{ID: call.ID, Name: call.Function.Name, Arguments: call.Function.Arguments}) + } + return chatRound{assistant: message.ToParam(), calls: calls}, nil +} + +func parseWireChatRound(response *wire.ChatCompletionWire) (chatRound, error) { + if response == nil || len(response.Choices) == 0 { + return chatRound{assistant: openai.AssistantMessage("")}, nil + } + message := response.Choices[0].Message + assistant := openai.AssistantMessage(message.Content) + calls := make([]ToolCall, 0, len(message.ToolCalls)) + for _, call := range message.ToolCalls { + if call.Type != "" && call.Type != "function" { + continue + } + calls = append(calls, ToolCall{ID: call.ID, Name: call.Function.Name, Arguments: call.Function.Arguments}) + assistant.OfAssistant.ToolCalls = append(assistant.OfAssistant.ToolCalls, openai.ChatCompletionMessageToolCallUnionParam{ + OfFunction: &openai.ChatCompletionMessageFunctionToolCallParam{ + ID: call.ID, + Function: openai.ChatCompletionMessageFunctionToolCallFunctionParam{ + Name: call.Function.Name, + Arguments: call.Function.Arguments, + }, + }, + }) + } + return chatRound{assistant: assistant, calls: calls}, nil +} + +func allCallsOwned(calls []ToolCall, owned map[string]struct{}) bool { + if len(calls) == 0 { + return false + } + for _, call := range calls { + if _, ok := owned[call.Name]; !ok { + return false + } + } + return true +} + +func appendChatToolResults(request any, assistant openai.ChatCompletionMessageParamUnion, results []ToolResult) (*openai.ChatCompletionNewParams, error) { + params, ok := request.(*openai.ChatCompletionNewParams) + if !ok || params == nil { + return nil, fmt.Errorf("append ToolLoop results to request %T", request) + } + cloned := *params + cloned.Messages = append([]openai.ChatCompletionMessageParamUnion(nil), params.Messages...) + cloned.Messages = append(cloned.Messages, assistant) + for _, result := range results { + content := result.Content + if result.IsError && content == "" { + content = "tool execution failed" + } + cloned.Messages = append(cloned.Messages, openai.ToolMessage(content, result.ToolCallID)) + } + return &cloned, nil +} + +func mergeTokenUsage(total, current *protocol.TokenUsage) *protocol.TokenUsage { + if current == nil { + return total + } + if total == nil { + copy := *current + return © + } + total.InputTokens += current.InputTokens + total.OutputTokens += current.OutputTokens + total.CacheReadTokens += current.CacheReadTokens + total.CacheWriteTokens += current.CacheWriteTokens + total.ReasoningTokens += current.ReasoningTokens + total.SystemTokens += current.SystemTokens + return total +} + +func cloneParameters(parameters map[string]any) shared.FunctionParameters { + if parameters == nil { + return nil + } + cloned := make(shared.FunctionParameters, len(parameters)) + for key, value := range parameters { + cloned[key] = value + } + return cloned +} diff --git a/internal/protocol/stage/toolloop/openai_chat_stream.go b/internal/protocol/stage/toolloop/openai_chat_stream.go new file mode 100644 index 000000000..ccf6aabc2 --- /dev/null +++ b/internal/protocol/stage/toolloop/openai_chat_stream.go @@ -0,0 +1,290 @@ +package toolloop + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/assembler" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +// openAIChatToolLoopStream only buffers a round while deciding whether it is +// an internal tool round. Once visible content appears, events are passed +// through incrementally and that public round is never intercepted. +type openAIChatToolLoopStream struct { + endpoint *openAIChatEndpoint + call protocolstage.Call + owned map[string]struct{} + runCtx context.Context + + round int + current protocolstage.EventStream + assembler assembler.StreamAssembler + buffered []protocolstage.Event + pending []protocolstage.Event + public bool + + usage *protocol.TokenUsage + model string + sideEffects bool + done bool + closed bool + eofPending bool +} + +func newOpenAIChatToolLoopStream( + ctx context.Context, + endpoint *openAIChatEndpoint, + call protocolstage.Call, + owned map[string]struct{}, +) (protocolstage.EventStream, error) { + stream := &openAIChatToolLoopStream{ + endpoint: endpoint, + call: call, + owned: owned, + runCtx: ctx, + } + if err := stream.startRound(ctx); err != nil { + return nil, err + } + return stream, nil +} + +func (s *openAIChatToolLoopStream) Next(ctx context.Context) (protocolstage.Event, error) { + for { + if len(s.pending) > 0 { + event := s.pending[0] + s.pending = s.pending[1:] + if len(s.pending) == 0 && s.eofPending { + s.done = true + } + return event, nil + } + if s.done || s.closed { + return protocolstage.Event{}, io.EOF + } + + event, err := s.current.Next(ctx) + if err == nil { + if assembleErr := s.assembler.Add(event.Value); assembleErr != nil { + return protocolstage.Event{}, s.fail(assembleErr) + } + if s.public { + return event, nil + } + + visible, classifyErr := chatStreamEventVisible(event.Value) + if classifyErr != nil { + return protocolstage.Event{}, s.fail(classifyErr) + } + s.buffered = append(s.buffered, event) + if visible { + s.public = true + s.pending = s.buffered + s.buffered = nil + } + continue + } + + s.absorbCurrentResult() + closeErr := s.closeCurrent() + if !errors.Is(err, io.EOF) { + if closeErr != nil { + err = errors.Join(err, closeErr) + } + return protocolstage.Event{}, s.fail(err) + } + if closeErr != nil { + return protocolstage.Event{}, s.fail(closeErr) + } + if s.public { + s.done = true + return protocolstage.Event{}, io.EOF + } + + complete, finishErr := s.assembler.Finish() + if finishErr != nil { + return protocolstage.Event{}, s.fail(finishErr) + } + roundResponse, parseErr := parseChatRound(complete) + if parseErr != nil { + return protocolstage.Event{}, s.fail(parseErr) + } + if !allCallsOwned(roundResponse.calls, s.owned) { + s.pending = s.buffered + s.buffered = nil + s.eofPending = true + if len(s.pending) == 0 { + s.done = true + return protocolstage.Event{}, io.EOF + } + continue + } + if s.round >= s.endpoint.stage.maxRounds { + return protocolstage.Event{}, s.fail(ErrMaxRounds) + } + + results, nextCtx, committed, executeErr := s.endpoint.executeCalls(s.runCtx, roundResponse.calls) + s.sideEffects = s.sideEffects || committed + if executeErr != nil { + return protocolstage.Event{}, s.fail(executeErr) + } + s.runCtx = nextCtx + nextRequest, appendErr := appendChatToolResults(s.call.Request, roundResponse.assistant, results) + if appendErr != nil { + return protocolstage.Event{}, s.fail(appendErr) + } + s.call.Request = nextRequest + s.buffered = nil + s.public = false + if startErr := s.startRound(s.runCtx); startErr != nil { + return protocolstage.Event{}, s.fail(startErr) + } + } +} + +func (s *openAIChatToolLoopStream) Close() error { + if s.closed { + return nil + } + s.closed = true + s.done = true + s.absorbCurrentResult() + return s.closeCurrent() +} + +func (s *openAIChatToolLoopStream) Result() protocolstage.StreamResult { + usage := cloneTokenUsage(s.usage) + model := s.model + sideEffects := s.sideEffects + if s.current != nil { + current := s.current.Result() + usage = mergeTokenUsage(usage, current.Usage) + if current.Model != "" { + model = current.Model + } + sideEffects = sideEffects || current.SideEffectsCommitted + } + return protocolstage.StreamResult{ + Usage: usage, + Model: model, + SideEffectsCommitted: sideEffects, + } +} + +func (s *openAIChatToolLoopStream) startRound(ctx context.Context) error { + stream, err := s.endpoint.next.Stream(ctx, s.call) + if err != nil { + return err + } + if stream == nil { + return errors.New("OpenAI Chat ToolLoop received a nil stream") + } + streamAssembler, err := assembler.NewStreamAssembler(protocol.TypeOpenAIChat) + if err != nil { + _ = stream.Close() + return err + } + s.current = stream + s.assembler = streamAssembler + s.round++ + return nil +} + +func (s *openAIChatToolLoopStream) absorbCurrentResult() { + if s.current == nil { + return + } + result := s.current.Result() + s.usage = mergeTokenUsage(s.usage, result.Usage) + if result.Model != "" { + s.model = result.Model + } + s.sideEffects = s.sideEffects || result.SideEffectsCommitted +} + +func (s *openAIChatToolLoopStream) closeCurrent() error { + if s.current == nil { + return nil + } + current := s.current + s.current = nil + return current.Close() +} + +func (s *openAIChatToolLoopStream) fail(err error) error { + s.done = true + if s.current != nil { + s.absorbCurrentResult() + if closeErr := s.closeCurrent(); closeErr != nil { + err = errors.Join(err, closeErr) + } + } + return WrapError(err, s.sideEffects) +} + +func cloneTokenUsage(usage *protocol.TokenUsage) *protocol.TokenUsage { + if usage == nil { + return nil + } + cloned := *usage + return &cloned +} + +// chatStreamEventVisible reports whether an event commits the round to public +// streaming. Content, refusal, or reasoning wins over tool-call data in the +// same event so a mixed visible round can never be consumed internally. +func chatStreamEventVisible(value any) (bool, error) { + raw, err := chatStreamEventJSON(value) + if err != nil { + return false, err + } + var chunk struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + Refusal string `json:"refusal"` + ReasoningContent string `json:"reasoning_content"` + } `json:"delta"` + } `json:"choices"` + } + if err := json.Unmarshal(raw, &chunk); err != nil { + return false, fmt.Errorf("classify OpenAI Chat stream event %T: %w", value, err) + } + for _, choice := range chunk.Choices { + if choice.Delta.Content != "" || choice.Delta.Refusal != "" || choice.Delta.ReasoningContent != "" { + return true, nil + } + } + return false, nil +} + +func chatStreamEventJSON(value any) ([]byte, error) { + if value == nil { + return nil, errors.New("OpenAI Chat ToolLoop received a nil stream event") + } + switch event := value.(type) { + case json.RawMessage: + return event, nil + case []byte: + return event, nil + case wire.ChatStreamChunk: + return json.Marshal(event) + case *wire.ChatStreamChunk: + return json.Marshal(event) + case interface{ RawJSON() string }: + if raw := event.RawJSON(); raw != "" { + return []byte(raw), nil + } + } + raw, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("marshal OpenAI Chat stream event %T: %w", value, err) + } + return raw, nil +} diff --git a/internal/protocol/stage/toolloop/openai_chat_test.go b/internal/protocol/stage/toolloop/openai_chat_test.go new file mode 100644 index 000000000..740d171bd --- /dev/null +++ b/internal/protocol/stage/toolloop/openai_chat_test.go @@ -0,0 +1,518 @@ +package toolloop + +import ( + "context" + "encoding/json" + "errors" + "io" + "reflect" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/shared" + + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" + "github.com/tingly-dev/tingly-box/internal/record" +) + +func TestMergeTokenUsagePreservesCurrentCacheDetails(t *testing.T) { + total := &protocol.TokenUsage{ + InputTokens: 3, OutputTokens: 2, CacheReadTokens: 1, CacheWriteTokens: 2, ReasoningTokens: 1, SystemTokens: 1, + } + current := &protocol.TokenUsage{ + InputTokens: 5, OutputTokens: 4, CacheReadTokens: 3, CacheWriteTokens: 4, ReasoningTokens: 2, SystemTokens: 2, + } + + got := mergeTokenUsage(total, current) + if got.InputTokens != 8 || got.OutputTokens != 6 || got.CacheReadTokens != 4 || got.CacheWriteTokens != 6 || got.ReasoningTokens != 3 || got.SystemTokens != 3 { + t.Fatalf("merged usage = %#v", got) + } +} + +func TestOpenAIChatCompleteRunsOwnedToolAndContinues(t *testing.T) { + catalog := staticCatalog{{Name: "lookup", Description: "Look up a value", Parameters: map[string]any{"type": "object"}}} + executor := &fakeExecutor{results: map[string]ToolResult{"lookup": {Content: "Paris"}}} + terminal := &scriptedChatEndpoint{completeResponses: []*protocolstage.Response{ + {Value: sdkToolCallCompletion(t, "call-1", "lookup", `{"city":"France"}`), Usage: protocol.NewTokenUsage(3, 2), Model: "provider"}, + {Value: textCompletion("The capital is Paris."), Usage: protocol.NewTokenUsage(5, 4), Model: "provider"}, + }} + stage, err := NewOpenAIChat(OpenAIChatConfig{Catalog: catalog, Executor: executor}) + if err != nil { + t.Fatal(err) + } + endpoint, err := protocolstage.Compose(terminal, stage) + if err != nil { + t.Fatal(err) + } + request := &openai.ChatCompletionNewParams{Model: "client", Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("hello")}} + + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatal(err) + } + if got := response.Value.(wire.ChatCompletionWire).Choices[0].Message.Content; got != "The capital is Paris." { + t.Fatalf("final content = %q", got) + } + if response.Usage == nil || response.Usage.InputTokens != 8 || response.Usage.OutputTokens != 6 { + t.Fatalf("aggregate usage = %#v", response.Usage) + } + if !response.SideEffectsCommitted { + t.Fatal("successful tool execution did not commit side effects") + } + if len(terminal.completeCalls) != 2 { + t.Fatalf("provider calls = %d, want 2", len(terminal.completeCalls)) + } + firstRequest := terminal.completeCalls[0].Request.(*openai.ChatCompletionNewParams) + if len(firstRequest.Tools) != 1 || firstRequest.Tools[0].GetFunction().Name != "lookup" { + t.Fatalf("injected tools = %#v", firstRequest.Tools) + } + secondRequest := terminal.completeCalls[1].Request.(*openai.ChatCompletionNewParams) + if len(secondRequest.Messages) != 3 { + t.Fatalf("continuation messages = %d, want user + assistant + tool", len(secondRequest.Messages)) + } + if len(executor.calls) != 1 || executor.calls[0].ID != "call-1" { + t.Fatalf("executed calls = %#v", executor.calls) + } + if len(request.Tools) != 0 || len(request.Messages) != 1 { + t.Fatal("ToolLoop mutated the caller's original request") + } +} + +func TestOpenAIChatCompleteLeavesExternalAndMixedCallsOutward(t *testing.T) { + for _, tt := range []struct { + name string + response wire.ChatCompletionWire + }{ + {name: "external", response: toolCallCompletion("call-ext", "client_tool", `{}`)}, + {name: "mixed", response: multiToolCallCompletion( + wire.ChatCompletionToolCallWire{ID: "call-owned", Type: "function", Function: wire.ChatCompletionFunctionWire{Name: "lookup", Arguments: `{}`}}, + wire.ChatCompletionToolCallWire{ID: "call-ext", Type: "function", Function: wire.ChatCompletionFunctionWire{Name: "client_tool", Arguments: `{}`}}, + )}, + } { + t.Run(tt.name, func(t *testing.T) { + executor := &fakeExecutor{} + terminal := &scriptedChatEndpoint{completeResponses: []*protocolstage.Response{{Value: tt.response}}} + stage, _ := NewOpenAIChat(OpenAIChatConfig{Catalog: staticCatalog{{Name: "lookup"}}, Executor: executor}) + endpoint, _ := protocolstage.Compose(terminal, stage) + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if err != nil { + t.Fatal(err) + } + if response == nil || len(terminal.completeCalls) != 1 || len(executor.calls) != 0 { + t.Fatalf("external/mixed call was consumed: response=%#v calls=%d executed=%d", response, len(terminal.completeCalls), len(executor.calls)) + } + }) + } +} + +func TestOpenAIChatCompletePreservesSideEffectBoundaryAfterLaterFailure(t *testing.T) { + providerErr := errors.New("second round failed") + terminal := &scriptedChatEndpoint{ + completeResponses: []*protocolstage.Response{{Value: toolCallCompletion("call-1", "lookup", `{}`)}}, + completeErrors: []error{nil, providerErr}, + } + stage, _ := NewOpenAIChat(OpenAIChatConfig{ + Catalog: staticCatalog{{Name: "lookup"}}, + Executor: &fakeExecutor{results: map[string]ToolResult{"lookup": {Content: "ok"}}}, + }) + endpoint, _ := protocolstage.Compose(terminal, stage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if !errors.Is(err, providerErr) || !HasCommittedSideEffects(err) { + t.Fatalf("later error = %v, committed=%v", err, HasCommittedSideEffects(err)) + } +} + +func TestOpenAIChatCompleteRecordsToolRoundsAsExchangesInOneAttempt(t *testing.T) { + request := &openai.ChatCompletionNewParams{Model: "public", Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("hello")}} + recorder, err := record.New(record.Config{ + Enabled: true, + RequestID: "req-complete-tool-loop", + InputProtocol: protocol.TypeOpenAIChat, + Input: request, + }) + if err != nil { + t.Fatal(err) + } + terminal := &scriptedChatEndpoint{completeResponses: []*protocolstage.Response{ + {Value: toolCallCompletion("call-1", "lookup", `{}`)}, + {Value: textCompletion("done")}, + }} + observed := record.ObserveProvider(terminal, recorder, record.ExchangeMetadata{Attempt: 3}) + toolStage, _ := NewOpenAIChat(OpenAIChatConfig{ + Catalog: staticCatalog{{Name: "lookup"}}, + Executor: &fakeExecutor{results: map[string]ToolResult{"lookup": {Content: "ok"}}}, + }) + endpoint, _ := protocolstage.Compose(observed, toolStage) + + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatal(err) + } + if err := recorder.SetFinalResponse(protocol.TypeOpenAIChat, response.Value); err != nil { + t.Fatal(err) + } + completed, first := recorder.Finish(nil) + if !first { + t.Fatal("recorder was already finished") + } + assertToolLoopRecord(t, completed) +} + +func TestOpenAIChatRejectsAmbiguousToolNameOwnership(t *testing.T) { + request := &openai.ChatCompletionNewParams{Tools: []openai.ChatCompletionToolUnionParam{ + openai.ChatCompletionFunctionTool(shared.FunctionDefinitionParam{Name: "lookup"}), + }} + stage, _ := NewOpenAIChat(OpenAIChatConfig{ + Catalog: staticCatalog{{Name: "lookup"}}, + Executor: &fakeExecutor{}, + }) + endpoint, _ := protocolstage.Compose(&scriptedChatEndpoint{}, stage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if !errors.Is(err, ErrToolNameCollision) { + t.Fatalf("tool name collision error = %v", err) + } +} + +func TestOpenAIChatStreamHidesOwnedToolRoundAndContinues(t *testing.T) { + toolEvents := toolCallStreamEvents("call-1", "lookup", `{"city":"France"}`) + textEvents := textStreamEvents("The capital is Paris.") + first := &memoryEventStream{events: toolEvents, result: protocolstage.StreamResult{Usage: protocol.NewTokenUsage(3, 2), Model: "provider"}} + second := &memoryEventStream{events: textEvents, result: protocolstage.StreamResult{Usage: protocol.NewTokenUsage(5, 4), Model: "provider"}} + terminal := &scriptedChatEndpoint{streams: []*memoryEventStream{first, second}} + executor := &fakeExecutor{results: map[string]ToolResult{"lookup": {Content: "Paris"}}} + stage, _ := NewOpenAIChat(OpenAIChatConfig{Catalog: staticCatalog{{Name: "lookup"}}, Executor: executor}) + endpoint, _ := protocolstage.Compose(terminal, stage) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if err != nil { + t.Fatal(err) + } + got := collectEvents(t, stream) + if !reflect.DeepEqual(got, textEvents) { + t.Fatalf("outward events = %#v, want only final text events %#v", got, textEvents) + } + result := stream.Result() + if result.Usage == nil || result.Usage.InputTokens != 8 || result.Usage.OutputTokens != 6 { + t.Fatalf("aggregate stream usage = %#v", result.Usage) + } + if !result.SideEffectsCommitted || result.Model != "provider" { + t.Fatalf("stream result = %#v", result) + } + if len(terminal.streamCalls) != 2 || len(executor.calls) != 1 { + t.Fatalf("provider calls=%d executor calls=%d", len(terminal.streamCalls), len(executor.calls)) + } + continuation := terminal.streamCalls[1].Request.(*openai.ChatCompletionNewParams) + if len(continuation.Messages) != 2 { + t.Fatalf("continuation messages = %d, want assistant + tool", len(continuation.Messages)) + } + if first.closeCalls != 1 || second.closeCalls != 1 { + t.Fatalf("inner close calls = %d, %d", first.closeCalls, second.closeCalls) + } + if err := stream.Close(); err != nil { + t.Fatal(err) + } + if first.closeCalls != 1 || second.closeCalls != 1 { + t.Fatal("outer Close closed an already completed inner stream twice") + } +} + +func TestOpenAIChatStreamReplaysExternalToolRound(t *testing.T) { + events := toolCallStreamEvents("call-ext", "client_tool", `{}`) + inner := &memoryEventStream{events: events} + terminal := &scriptedChatEndpoint{streams: []*memoryEventStream{inner}} + executor := &fakeExecutor{} + stage, _ := NewOpenAIChat(OpenAIChatConfig{Catalog: staticCatalog{{Name: "lookup"}}, Executor: executor}) + endpoint, _ := protocolstage.Compose(terminal, stage) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if err != nil { + t.Fatal(err) + } + got := collectEvents(t, stream) + if !reflect.DeepEqual(got, events) { + t.Fatalf("replayed events = %#v, want %#v", got, events) + } + if len(executor.calls) != 0 || len(terminal.streamCalls) != 1 { + t.Fatalf("external tool was consumed: executor=%d provider=%d", len(executor.calls), len(terminal.streamCalls)) + } + _ = stream.Close() +} + +func TestOpenAIChatStreamKeepsVisibleContentPullBased(t *testing.T) { + events := textStreamEvents("hello") + inner := &memoryEventStream{events: events} + terminal := &scriptedChatEndpoint{streams: []*memoryEventStream{inner}} + stage, _ := NewOpenAIChat(OpenAIChatConfig{Catalog: staticCatalog{{Name: "lookup"}}, Executor: &fakeExecutor{}}) + endpoint, _ := protocolstage.Compose(terminal, stage) + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if err != nil { + t.Fatal(err) + } + + first, err := stream.Next(context.Background()) + if err != nil || !reflect.DeepEqual(first, events[0]) { + t.Fatalf("first event = %#v, %v", first, err) + } + if inner.nextCalls != 2 { + t.Fatalf("provider pulls after first outward event = %d, want role + first visible event only", inner.nextCalls) + } + second, err := stream.Next(context.Background()) + if err != nil || !reflect.DeepEqual(second, events[1]) { + t.Fatalf("second event = %#v, %v", second, err) + } + if inner.nextCalls != 2 { + t.Fatalf("buffered visible event caused another provider pull: %d", inner.nextCalls) + } + _ = stream.Close() +} + +func TestOpenAIChatStreamPreservesSideEffectBoundaryAfterLaterFailure(t *testing.T) { + providerErr := errors.New("second stream failed") + terminal := &scriptedChatEndpoint{ + streams: []*memoryEventStream{{events: toolCallStreamEvents("call-1", "lookup", `{}`)}}, + streamErrors: []error{nil, providerErr}, + } + stage, _ := NewOpenAIChat(OpenAIChatConfig{ + Catalog: staticCatalog{{Name: "lookup"}}, + Executor: &fakeExecutor{results: map[string]ToolResult{"lookup": {Content: "ok"}}}, + }) + endpoint, _ := protocolstage.Compose(terminal, stage) + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if err != nil { + t.Fatal(err) + } + + _, err = stream.Next(context.Background()) + if !errors.Is(err, providerErr) || !HasCommittedSideEffects(err) { + t.Fatalf("later stream error = %v, committed=%v", err, HasCommittedSideEffects(err)) + } + if !stream.Result().SideEffectsCommitted { + t.Fatal("stream result lost committed side effects") + } + _ = stream.Close() +} + +func TestOpenAIChatStreamEnforcesMaxRoundsBeforeToolExecution(t *testing.T) { + executor := &fakeExecutor{} + terminal := &scriptedChatEndpoint{streams: []*memoryEventStream{{events: toolCallStreamEvents("call-1", "lookup", `{}`)}}} + stage, _ := NewOpenAIChat(OpenAIChatConfig{ + Catalog: staticCatalog{{Name: "lookup"}}, + Executor: executor, + MaxRounds: 1, + }) + endpoint, _ := protocolstage.Compose(terminal, stage) + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if err != nil { + t.Fatal(err) + } + + _, err = stream.Next(context.Background()) + if !errors.Is(err, ErrMaxRounds) || HasCommittedSideEffects(err) { + t.Fatalf("max-round error = %v, committed=%v", err, HasCommittedSideEffects(err)) + } + if len(executor.calls) != 0 { + t.Fatalf("executed %d tools after reaching max rounds", len(executor.calls)) + } + _ = stream.Close() +} + +func TestOpenAIChatStreamRecordsToolRoundsAsExchangesInOneAttempt(t *testing.T) { + request := &openai.ChatCompletionNewParams{Model: "public", Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("hello")}} + recorder, err := record.New(record.Config{ + Enabled: true, + RequestID: "req-tool-loop", + InputProtocol: protocol.TypeOpenAIChat, + Input: request, + }) + if err != nil { + t.Fatal(err) + } + terminal := &scriptedChatEndpoint{streams: []*memoryEventStream{ + {events: toolCallStreamEvents("call-1", "lookup", `{}`)}, + {events: textStreamEvents("done")}, + }} + observed := record.ObserveProvider(terminal, recorder, record.ExchangeMetadata{ + Attempt: 3, + Provider: "provider-a", + Model: "provider-model", + }) + toolStage, _ := NewOpenAIChat(OpenAIChatConfig{ + Catalog: staticCatalog{{Name: "lookup"}}, + Executor: &fakeExecutor{results: map[string]ToolResult{"lookup": {Content: "ok"}}}, + }) + endpoint, _ := protocolstage.Compose(observed, toolStage) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatal(err) + } + _ = collectEvents(t, stream) + if err := stream.Close(); err != nil { + t.Fatal(err) + } + if err := recorder.SetFinalResponse(protocol.TypeOpenAIChat, textCompletion("done")); err != nil { + t.Fatal(err) + } + completed, first := recorder.Finish(nil) + if !first { + t.Fatal("recorder was already finished") + } + assertToolLoopRecord(t, completed) +} + +func assertToolLoopRecord(t *testing.T, completed *record.RequestRecord) { + t.Helper() + if len(completed.ProviderExchanges) != 2 { + t.Fatalf("provider exchanges = %d, want 2", len(completed.ProviderExchanges)) + } + for index, exchange := range completed.ProviderExchanges { + if exchange.Sequence != index+1 || exchange.Attempt != 3 || exchange.Outcome != record.OutcomeSucceeded || exchange.Response == nil { + t.Fatalf("exchange %d = %#v", index, exchange) + } + } + if completed.FinalResponse == nil || completed.FinalResponse.Protocol != protocol.TypeOpenAIChat { + t.Fatalf("final response = %#v", completed.FinalResponse) + } +} + +type staticCatalog []ToolDefinition + +func (c staticCatalog) ListTools(context.Context) ([]ToolDefinition, error) { + return append([]ToolDefinition(nil), c...), nil +} + +type fakeExecutor struct { + results map[string]ToolResult + calls []ToolCall +} + +func (e *fakeExecutor) Execute(ctx context.Context, call ToolCall) (context.Context, ToolResult, error) { + e.calls = append(e.calls, call) + return ctx, e.results[call.Name], nil +} + +type scriptedChatEndpoint struct { + completeResponses []*protocolstage.Response + completeErrors []error + completeCalls []protocolstage.Call + streams []*memoryEventStream + streamErrors []error + streamCalls []protocolstage.Call +} + +func (*scriptedChatEndpoint) Protocol() protocol.APIType { return protocol.TypeOpenAIChat } +func (e *scriptedChatEndpoint) Complete(_ context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + index := len(e.completeCalls) + e.completeCalls = append(e.completeCalls, call) + if index < len(e.completeErrors) && e.completeErrors[index] != nil { + return nil, e.completeErrors[index] + } + if index >= len(e.completeResponses) { + return nil, errors.New("unexpected provider call") + } + return e.completeResponses[index], nil +} +func (e *scriptedChatEndpoint) Stream(_ context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + index := len(e.streamCalls) + e.streamCalls = append(e.streamCalls, call) + if index < len(e.streamErrors) && e.streamErrors[index] != nil { + return nil, e.streamErrors[index] + } + if index >= len(e.streams) { + return nil, errors.New("unexpected provider stream") + } + return e.streams[index], nil +} + +type memoryEventStream struct { + events []protocolstage.Event + result protocolstage.StreamResult + nextCalls int + closeCalls int +} + +func (s *memoryEventStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + s.nextCalls++ + if s.nextCalls > len(s.events) { + return protocolstage.Event{}, io.EOF + } + return s.events[s.nextCalls-1], nil +} + +func (s *memoryEventStream) Close() error { + s.closeCalls++ + return nil +} + +func (s *memoryEventStream) Result() protocolstage.StreamResult { return s.result } + +func collectEvents(t *testing.T, stream protocolstage.EventStream) []protocolstage.Event { + t.Helper() + var events []protocolstage.Event + for { + event, err := stream.Next(context.Background()) + if errors.Is(err, io.EOF) { + return events + } + if err != nil { + t.Fatal(err) + } + events = append(events, event) + } +} + +func toolCallStreamEvents(id, name, arguments string) []protocolstage.Event { + finish := "tool_calls" + return []protocolstage.Event{ + {Value: wire.ChatStreamChunk{ID: "chat-1", Model: "provider", Choices: []wire.ChatStreamChoice{{Delta: wire.ChatStreamDelta{Role: "assistant"}}}}}, + {Value: wire.ChatStreamChunk{ID: "chat-1", Model: "provider", Choices: []wire.ChatStreamChoice{{Delta: wire.ChatStreamDelta{ToolCalls: []wire.ChatStreamToolCall{{Index: 0, ID: id, Type: "function", Function: wire.ChatStreamToolFunction{Name: name, Arguments: &arguments}}}}}}}}, + {Value: wire.ChatStreamChunk{ID: "chat-1", Model: "provider", Choices: []wire.ChatStreamChoice{{FinishReason: &finish}}}}, + } +} + +func textStreamEvents(content string) []protocolstage.Event { + finish := "stop" + return []protocolstage.Event{ + {Value: wire.ChatStreamChunk{ID: "chat-2", Model: "provider", Choices: []wire.ChatStreamChoice{{Delta: wire.ChatStreamDelta{Role: "assistant"}}}}}, + {Value: wire.ChatStreamChunk{ID: "chat-2", Model: "provider", Choices: []wire.ChatStreamChoice{{Delta: wire.ChatStreamDelta{Content: content}}}}}, + {Value: wire.ChatStreamChunk{ID: "chat-2", Model: "provider", Choices: []wire.ChatStreamChoice{{FinishReason: &finish}}}}, + } +} + +func toolCallCompletion(id, name, arguments string) wire.ChatCompletionWire { + return multiToolCallCompletion(wire.ChatCompletionToolCallWire{ + ID: id, Type: "function", Function: wire.ChatCompletionFunctionWire{Name: name, Arguments: arguments}, + }) +} + +func multiToolCallCompletion(calls ...wire.ChatCompletionToolCallWire) wire.ChatCompletionWire { + return wire.ChatCompletionWire{Choices: []wire.ChatCompletionChoiceWire{{ + Message: wire.ChatCompletionMessageWire{Role: "assistant", ToolCalls: calls}, FinishReason: "tool_calls", + }}} +} + +func textCompletion(content string) wire.ChatCompletionWire { + return wire.ChatCompletionWire{Choices: []wire.ChatCompletionChoiceWire{{ + Message: wire.ChatCompletionMessageWire{Role: "assistant", Content: content}, FinishReason: "stop", + }}} +} + +func sdkToolCallCompletion(t *testing.T, id, name, arguments string) *openai.ChatCompletion { + t.Helper() + raw, err := json.Marshal(toolCallCompletion(id, name, arguments)) + if err != nil { + t.Fatal(err) + } + var completion openai.ChatCompletion + if err := json.Unmarshal(raw, &completion); err != nil { + t.Fatal(err) + } + return &completion +} diff --git a/internal/protocol/stage/toolloop/runtime.go b/internal/protocol/stage/toolloop/runtime.go new file mode 100644 index 000000000..a78e35690 --- /dev/null +++ b/internal/protocol/stage/toolloop/runtime.go @@ -0,0 +1,123 @@ +package toolloop + +import ( + "context" + "errors" + "fmt" +) + +// ToolDefinition is the protocol-neutral description injected into the +// ToolLoop Stage's native request protocol. +type ToolDefinition struct { + Name string + Description string + Parameters map[string]any +} + +// ToolCall is one model-requested invocation after protocol-native assembly. +// Arguments contains the model's JSON text unchanged; the execution backend is +// responsible for schema validation. +type ToolCall struct { + ID string + Name string + Arguments string +} + +// ToolResult is the protocol-neutral value appended to the next model round. +type ToolResult struct { + ToolCallID string + Content string + IsError bool +} + +// ToolCatalog lists the server-visible tools for one request. The returned +// definitions also form the ownership set: calls whose names are absent remain +// client/external tool calls and are returned outward unchanged. +type ToolCatalog interface { + ListTools(ctx context.Context) ([]ToolDefinition, error) +} + +// ToolPolicy authorizes a server-owned call immediately before execution. +type ToolPolicy interface { + Authorize(ctx context.Context, call ToolCall) error +} + +// ToolExecutor invokes one server-owned tool. It may return an updated context +// for request-scoped state such as advisor depth or credentials. +type ToolExecutor interface { + Execute(ctx context.Context, call ToolCall) (context.Context, ToolResult, error) +} + +// AllowAllPolicy is the explicit default used when no additional authorization +// layer is configured. +type AllowAllPolicy struct{} + +func (AllowAllPolicy) Authorize(context.Context, ToolCall) error { return nil } + +var ( + ErrMaxRounds = errors.New("tool loop reached the maximum number of rounds") + ErrToolNameCollision = errors.New("server tool name collides with request tool") +) + +// ExecutionError preserves the irreversible-side-effect boundary when a later +// provider round or policy check fails. Failover code can inspect it without +// importing a concrete ToolLoop implementation. +type ExecutionError struct { + Err error + SideEffectsCommitted bool +} + +func (e *ExecutionError) Error() string { + if e == nil || e.Err == nil { + return "tool loop failed" + } + return e.Err.Error() +} + +func (e *ExecutionError) Unwrap() error { + if e == nil { + return nil + } + return e.Err +} + +// WrapError annotates err only when a successful tool execution has already +// committed side effects. Before that boundary the original error is retained +// so existing retry classification remains unchanged. +func WrapError(err error, sideEffectsCommitted bool) error { + if err == nil || !sideEffectsCommitted { + return err + } + return &ExecutionError{Err: err, SideEffectsCommitted: true} +} + +// HasCommittedSideEffects reports whether retrying the whole provider attempt +// could replay an already successful tool action. +func HasCommittedSideEffects(err error) bool { + var executionErr *ExecutionError + return errors.As(err, &executionErr) && executionErr.SideEffectsCommitted +} + +func validateDependencies(catalog ToolCatalog, executor ToolExecutor) error { + if catalog == nil { + return errors.New("tool loop catalog is nil") + } + if executor == nil { + return errors.New("tool loop executor is nil") + } + return nil +} + +func validateDefinitions(definitions []ToolDefinition) error { + seen := make(map[string]struct{}, len(definitions)) + for i, definition := range definitions { + if definition.Name == "" { + return fmt.Errorf("tool definition at index %d has an empty name", i) + } + if _, exists := seen[definition.Name]; exists { + return fmt.Errorf("tool definition %q is duplicated", definition.Name) + } + seen[definition.Name] = struct{}{} + } + return nil +} diff --git a/internal/protocol/stage/toolloop/runtime_test.go b/internal/protocol/stage/toolloop/runtime_test.go new file mode 100644 index 000000000..f081a81a4 --- /dev/null +++ b/internal/protocol/stage/toolloop/runtime_test.go @@ -0,0 +1,34 @@ +package toolloop + +import ( + "errors" + "testing" +) + +func TestWrapErrorPreservesCommittedSideEffects(t *testing.T) { + providerErr := errors.New("provider failed") + + if got := WrapError(providerErr, false); !errors.Is(got, providerErr) || HasCommittedSideEffects(got) { + t.Fatalf("uncommitted error = %#v", got) + } + + got := WrapError(providerErr, true) + if !errors.Is(got, providerErr) { + t.Fatalf("wrapped error does not preserve cause: %v", got) + } + if !HasCommittedSideEffects(got) { + t.Fatal("wrapped error lost committed side-effect state") + } +} + +func TestValidateDefinitionsRejectsEmptyAndDuplicateNames(t *testing.T) { + if err := validateDefinitions([]ToolDefinition{{}}); err == nil { + t.Fatal("empty tool name was accepted") + } + if err := validateDefinitions([]ToolDefinition{{Name: "lookup"}, {Name: "lookup"}}); err == nil { + t.Fatal("duplicate tool name was accepted") + } + if err := validateDefinitions([]ToolDefinition{{Name: "lookup"}, {Name: "calculate"}}); err != nil { + t.Fatalf("valid definitions rejected: %v", err) + } +} diff --git a/internal/protocol/stage/topology.go b/internal/protocol/stage/topology.go new file mode 100644 index 000000000..abcb33a93 --- /dev/null +++ b/internal/protocol/stage/topology.go @@ -0,0 +1,90 @@ +package stage + +import ( + "fmt" + "strings" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +// TopologyConfig describes an arbitrary protocol-stage path. Stages are listed +// in client request order, outermost to innermost. Adjacent stages may speak +// different protocols; BuildTopology resolves an explicit Bridge for each +// mismatch while constructing from Terminal outward. +type TopologyConfig struct { + Terminal Endpoint + Stages []Stage + ClientProtocol protocol.APIType + Registry *BridgeRegistry + RequiredCapabilities Capabilities +} + +// BuildTopology constructs a client-facing Endpoint without executing it. For +// a client A, outer stage B, inner stage C, and provider D, the result is: +// +// bridge A->B ( +// stage B ( +// bridge B->C ( +// stage C ( +// bridge C->D (terminal D))))) +func BuildTopology(config TopologyConfig) (Endpoint, error) { + if isNil(config.Terminal) { + return nil, fmt.Errorf("build protocol stage topology: terminal endpoint is nil") + } + if config.Terminal.Protocol() == "" { + return nil, fmt.Errorf("build protocol stage topology: terminal endpoint has empty protocol") + } + if config.ClientProtocol == "" { + return nil, fmt.Errorf("build protocol stage topology: client protocol is empty") + } + if config.Registry == nil { + return nil, fmt.Errorf("build protocol stage topology: bridge registry is nil") + } + + required := config.RequiredCapabilities | CoreBridgeCapabilities + current := config.Terminal + for i := len(config.Stages) - 1; i >= 0; i-- { + stage := config.Stages[i] + if isNil(stage) { + return nil, fmt.Errorf("build protocol stage topology: stage at index %d is nil", i) + } + name := strings.TrimSpace(stage.Name()) + if name == "" { + return nil, fmt.Errorf("build protocol stage topology: stage at index %d has empty name", i) + } + stageProtocol := stage.Protocol() + if stageProtocol == "" { + return nil, fmt.Errorf("build protocol stage topology: stage %q has empty protocol", name) + } + + if stageProtocol != current.Protocol() { + bridge, err := config.Registry.Resolve(stageProtocol, current.Protocol(), required) + if err != nil { + return nil, fmt.Errorf("build protocol stage topology: bridge below stage %q: %w", name, err) + } + current, err = Adapt(current, bridge) + if err != nil { + return nil, fmt.Errorf("build protocol stage topology: adapt below stage %q: %w", name, err) + } + } + + var err error + current, err = Compose(current, stage) + if err != nil { + return nil, fmt.Errorf("build protocol stage topology: compose stage %q: %w", name, err) + } + } + + if config.ClientProtocol == current.Protocol() { + return current, nil + } + ingress, err := config.Registry.Resolve(config.ClientProtocol, current.Protocol(), required) + if err != nil { + return nil, fmt.Errorf("build protocol stage topology: ingress bridge: %w", err) + } + current, err = Adapt(current, ingress) + if err != nil { + return nil, fmt.Errorf("build protocol stage topology: adapt ingress bridge: %w", err) + } + return current, nil +} diff --git a/internal/protocol/stage/topology_test.go b/internal/protocol/stage/topology_test.go new file mode 100644 index 000000000..639223637 --- /dev/null +++ b/internal/protocol/stage/topology_test.go @@ -0,0 +1,231 @@ +package stage + +import ( + "context" + "reflect" + "strings" + "testing" + + protocol "github.com/tingly-dev/tingly-box/ai" +) + +func TestBuildTopologyMixedProtocols(t *testing.T) { + t.Parallel() + + var calls []string + usage := protocol.NewTokenUsage(23, 12) + terminal := &recordingEndpoint{ + protocol: protocol.TypeAnthropicV1, + calls: &calls, + response: &Response{ + Value: "terminal response", + Usage: usage, + Model: "provider-model", + SideEffectsCommitted: true, + }, + } + ingress := &testingBridge{ + name: "bridge_chat_beta", + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: AllBridgeCapabilities, + calls: &calls, + } + between := &testingBridge{ + name: "bridge_beta_responses", + source: protocol.TypeAnthropicBeta, + target: protocol.TypeOpenAIResponses, + caps: AllBridgeCapabilities, + calls: &calls, + } + provider := &testingBridge{ + name: "bridge_responses_v1", + source: protocol.TypeOpenAIResponses, + target: protocol.TypeAnthropicV1, + caps: AllBridgeCapabilities, + calls: &calls, + } + registry, err := NewBridgeRegistry(ingress, between, provider) + if err != nil { + t.Fatalf("NewBridgeRegistry() error = %v", err) + } + + topology, err := BuildTopology(TopologyConfig{ + Terminal: terminal, + Stages: []Stage{ + &recordingStage{name: "guardrails", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + &recordingStage{name: "tool_loop", protocol: protocol.TypeOpenAIResponses, calls: &calls}, + }, + ClientProtocol: protocol.TypeOpenAIChat, + Registry: registry, + RequiredCapabilities: CapabilityUsage | CapabilityToolUse, + }) + if err != nil { + t.Fatalf("BuildTopology() error = %v", err) + } + if topology.Protocol() != protocol.TypeOpenAIChat { + t.Fatalf("topology.Protocol() = %q", topology.Protocol()) + } + if len(calls) != 0 { + t.Fatalf("BuildTopology() executed chain, calls = %v", calls) + } + + call := Call{ + Request: "client request", + Metadata: CallMetadata{ + RequestID: "req-topology", + Attempt: 4, + }, + } + response, err := topology.Complete(context.Background(), call) + if err != nil { + t.Fatalf("Complete() error = %v", err) + } + wantValue := "bridge_chat_beta(bridge_beta_responses(bridge_responses_v1(terminal response)))" + if response.Value != wantValue { + t.Fatalf("response.Value = %v, want %v", response.Value, wantValue) + } + assertResponseFacts(t, response, usage, "provider-model", true) + if terminal.lastCall.Metadata != call.Metadata { + t.Fatalf("terminal metadata = %+v, want %+v", terminal.lastCall.Metadata, call.Metadata) + } + + wantCalls := []string{ + "bridge_chat_beta:request", + "guardrails:request", + "bridge_beta_responses:request", + "tool_loop:request", + "bridge_responses_v1:request", + "terminal:request", + "terminal:response", + "bridge_responses_v1:response", + "tool_loop:response", + "bridge_beta_responses:response", + "guardrails:response", + "bridge_chat_beta:response", + } + if !reflect.DeepEqual(calls, wantCalls) { + t.Fatalf("calls = %v, want %v", calls, wantCalls) + } +} + +func TestBridgeRegistryResolution(t *testing.T) { + t.Parallel() + + bridge := &testingBridge{ + name: "chat_beta", + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: CoreBridgeCapabilities | CapabilityUsage, + } + registry, err := NewBridgeRegistry(bridge) + if err != nil { + t.Fatalf("NewBridgeRegistry() error = %v", err) + } + + resolved, err := registry.Resolve(protocol.TypeOpenAIChat, protocol.TypeAnthropicBeta, CapabilityUsage) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if resolved != bridge { + t.Fatal("Resolve() did not return the registered bridge") + } + + identity, err := registry.Resolve(protocol.TypeAnthropicBeta, protocol.TypeAnthropicBeta, AllBridgeCapabilities) + if err != nil { + t.Fatalf("identity Resolve() error = %v", err) + } + if identity.Source() != protocol.TypeAnthropicBeta || identity.Target() != protocol.TypeAnthropicBeta { + t.Fatalf("identity bridge = %q -> %q", identity.Source(), identity.Target()) + } + if _, err := registry.ResolveRegistered(protocol.TypeAnthropicBeta, protocol.TypeAnthropicBeta, AllBridgeCapabilities); err == nil { + t.Fatal("ResolveRegistered() accepted an implicit identity bridge") + } + + _, err = registry.Resolve(protocol.TypeOpenAIChat, protocol.TypeAnthropicBeta, CapabilityToolUse) + if err == nil || !strings.Contains(err.Error(), "missing capabilities: tool_use") { + t.Fatalf("capability Resolve() error = %v", err) + } + _, err = registry.Resolve(protocol.TypeOpenAIResponses, protocol.TypeAnthropicBeta, 0) + if err == nil || !strings.Contains(err.Error(), "not registered") { + t.Fatalf("missing Resolve() error = %v", err) + } +} + +func TestNewBridgeRegistryRejectsInvalidEntries(t *testing.T) { + t.Parallel() + + valid := &testingBridge{ + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: AllBridgeCapabilities, + } + tests := []struct { + name string + bridges []Bridge + want string + }{ + {name: "nil", bridges: []Bridge{nil}, want: "index 0 is nil"}, + { + name: "empty protocol", + bridges: []Bridge{&testingBridge{caps: AllBridgeCapabilities}}, + want: "has empty protocol", + }, + { + name: "missing core", + bridges: []Bridge{&testingBridge{ + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + caps: CapabilityComplete, + }}, + want: "missing core capabilities: stream,error", + }, + {name: "duplicate", bridges: []Bridge{valid, valid}, want: "duplicate bridge"}, + {name: "valid", bridges: []Bridge{valid}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + registry, err := NewBridgeRegistry(tt.bridges...) + if tt.want == "" { + if err != nil || registry == nil { + t.Fatalf("NewBridgeRegistry() = (%v, %v)", registry, err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("NewBridgeRegistry() error = %v, want containing %q", err, tt.want) + } + }) + } +} + +func TestBuildTopologyRejectsMissingBridgeBeforeExecution(t *testing.T) { + t.Parallel() + + var calls []string + registry, err := NewBridgeRegistry() + if err != nil { + t.Fatalf("NewBridgeRegistry() error = %v", err) + } + terminal := &recordingEndpoint{ + protocol: protocol.TypeAnthropicV1, + calls: &calls, + } + + _, err = BuildTopology(TopologyConfig{ + Terminal: terminal, + Stages: []Stage{ + &recordingStage{name: "guardrails", protocol: protocol.TypeAnthropicBeta, calls: &calls}, + }, + ClientProtocol: protocol.TypeOpenAIChat, + Registry: registry, + }) + if err == nil || !strings.Contains(err.Error(), `"anthropic_beta" -> "anthropic_v1": not registered`) { + t.Fatalf("BuildTopology() error = %v", err) + } + if len(calls) != 0 { + t.Fatalf("BuildTopology() executed chain, calls = %v", calls) + } +} diff --git a/internal/protocol/stream/anthropic_beta_to_openai_responses.go b/internal/protocol/stream/anthropic_beta_to_openai_responses.go index ff4449e25..9630ca09f 100644 --- a/internal/protocol/stream/anthropic_beta_to_openai_responses.go +++ b/internal/protocol/stream/anthropic_beta_to_openai_responses.go @@ -28,7 +28,7 @@ func HandleAnthropicBetaToOpenAIResponsesStream( } }() - conv := newAnthropicBetaToResponsesConverter(stream, responseModel) + conv := NewAnthropicBetaToOpenAIResponsesConverter(stream, responseModel) usage, err := RunConverter(hc, conv, responsesSSEWriter(c)) diff --git a/internal/protocol/stream/anthropic_beta_to_openai_responses_converter.go b/internal/protocol/stream/anthropic_beta_to_openai_responses_converter.go index 7d81a3536..4da61e6d6 100644 --- a/internal/protocol/stream/anthropic_beta_to_openai_responses_converter.go +++ b/internal/protocol/stream/anthropic_beta_to_openai_responses_converter.go @@ -2,11 +2,11 @@ package stream import ( "fmt" + "sort" "strings" "time" "github.com/anthropics/anthropic-sdk-go" - anthropicstream "github.com/anthropics/anthropic-sdk-go/packages/ssestream" "github.com/tingly-dev/tingly-box/internal/protocol" usagepkg "github.com/tingly-dev/tingly-box/internal/protocol/usage" "github.com/tingly-dev/tingly-box/internal/protocol/wire" @@ -15,22 +15,24 @@ import ( // anthropicBetaToResponsesConverter converts an Anthropic Beta stream into // a sequence of Responses API wire events. It implements StreamConverter. type anthropicBetaToResponsesConverter struct { - stream *anthropicstream.Stream[anthropic.BetaRawMessageStreamEventUnion] + stream AnthropicBetaStream responseModel string acc *usagepkg.AnthropicAccumulator // state (formerly responsesConverterState) - responseID string - itemID string - outputIndex int - accumulatedText string - finished bool - pendingToolCalls map[int]*pendingResponseToolCall - hasSentCreated bool - sequenceNumber int - createdAt int64 - currentBlockType string - stopReason string + responseID string + itemID string + outputIndex int + textOutputIndex int + accumulatedText string + finished bool + pendingToolCalls map[int]*pendingResponseToolCall + hasSentCreated bool + sequenceNumber int + createdAt int64 + currentBlockType string + currentBlockOutputIndex int + stopReason string // internal event queue pending []wire.ResponsesEvent @@ -38,15 +40,16 @@ type anthropicBetaToResponsesConverter struct { // pendingResponseToolCall tracks a tool call being assembled from Anthropic stream chunks type pendingResponseToolCall struct { - itemID string - name string - arguments strings.Builder + itemID string + name string + outputIndex int + arguments strings.Builder } // newAnthropicBetaToResponsesConverter creates a converter that reads from an // Anthropic Beta stream and yields Responses API wire events. func newAnthropicBetaToResponsesConverter( - stream *anthropicstream.Stream[anthropic.BetaRawMessageStreamEventUnion], + stream AnthropicBetaStream, responseModel string, ) *anthropicBetaToResponsesConverter { ts := time.Now().Unix() @@ -56,11 +59,22 @@ func newAnthropicBetaToResponsesConverter( acc: usagepkg.NewAnthropicAccumulator(), responseID: fmt.Sprintf("resp_%d", ts), itemID: fmt.Sprintf("item_%d", ts), + textOutputIndex: -1, pendingToolCalls: make(map[int]*pendingResponseToolCall), createdAt: ts, } } +// NewAnthropicBetaToOpenAIResponsesConverter creates a transport-neutral +// converter. HTTP/SSE framing and stream close ownership remain with the +// caller, so the same converter can serve legacy handlers and Stage Bridges. +func NewAnthropicBetaToOpenAIResponsesConverter( + stream AnthropicBetaStream, + responseModel string, +) StreamConverter { + return newAnthropicBetaToResponsesConverter(stream, responseModel) +} + func (c *anthropicBetaToResponsesConverter) Next() (interface{}, bool, error) { if len(c.pending) > 0 { evt := c.pending[0] @@ -147,12 +161,16 @@ func (c *anthropicBetaToResponsesConverter) emitContentBlockStart(event *anthrop index := event.Index blockType := event.ContentBlock.Type c.currentBlockType = blockType + c.currentBlockOutputIndex = c.outputIndex if blockType == "text" { + if c.textOutputIndex == -1 { + c.textOutputIndex = c.currentBlockOutputIndex + } c.pending = append(c.pending, wire.ResponsesOutputItemAddedEvent{ Type: "response.output_item.added", SequenceNumber: int64(c.nextSeq()), - OutputIndex: c.outputIndex, + OutputIndex: c.currentBlockOutputIndex, Item: wire.ResponsesOutputItemWire{ ID: c.itemID, Type: "message", @@ -165,20 +183,21 @@ func (c *anthropicBetaToResponsesConverter) emitContentBlockStart(event *anthrop Type: "response.content_part.added", SequenceNumber: int64(c.nextSeq()), ItemID: c.itemID, - OutputIndex: c.outputIndex, + OutputIndex: c.currentBlockOutputIndex, ContentIndex: 0, Part: wire.ResponsesContentPartWire{Type: "output_text", Text: ""}, }) + c.outputIndex++ } else if blockType == "tool_use" { toolID := event.ContentBlock.ID toolName := event.ContentBlock.Name - c.pendingToolCalls[int(index)] = &pendingResponseToolCall{itemID: toolID, name: toolName} + c.pendingToolCalls[int(index)] = &pendingResponseToolCall{itemID: toolID, name: toolName, outputIndex: c.currentBlockOutputIndex} arguments := "" c.pending = append(c.pending, wire.ResponsesOutputItemAddedEvent{ Type: "response.output_item.added", SequenceNumber: int64(c.nextSeq()), - OutputIndex: c.outputIndex, + OutputIndex: c.currentBlockOutputIndex, Item: wire.ResponsesOutputItemWire{ Type: "function_call", ID: toolID, @@ -203,7 +222,7 @@ func (c *anthropicBetaToResponsesConverter) emitContentBlockDelta(event *anthrop Type: "response.output_text.delta", Delta: text, ItemID: c.itemID, - OutputIndex: c.outputIndex, + OutputIndex: c.currentBlockOutputIndex, ContentIndex: 0, SequenceNumber: int64(c.nextSeq()), }) @@ -215,7 +234,7 @@ func (c *anthropicBetaToResponsesConverter) emitContentBlockDelta(event *anthrop Type: "response.function_call_arguments.delta", Delta: argsDelta, ItemID: pending.itemID, - OutputIndex: c.outputIndex, + OutputIndex: pending.outputIndex, SequenceNumber: int64(c.nextSeq()), }) } @@ -231,7 +250,7 @@ func (c *anthropicBetaToResponsesConverter) emitContentBlockStop(event *anthropi wire.ResponsesOutputTextDoneEvent{ Type: "response.output_text.done", ItemID: c.itemID, - OutputIndex: c.outputIndex, + OutputIndex: c.currentBlockOutputIndex, ContentIndex: 0, Text: c.accumulatedText, SequenceNumber: int64(c.nextSeq()), @@ -240,14 +259,14 @@ func (c *anthropicBetaToResponsesConverter) emitContentBlockStop(event *anthropi Type: "response.content_part.done", SequenceNumber: int64(c.nextSeq()), ItemID: c.itemID, - OutputIndex: c.outputIndex, + OutputIndex: c.currentBlockOutputIndex, ContentIndex: 0, Part: wire.ResponsesContentPartWire{Type: "output_text", Text: c.accumulatedText}, }, wire.ResponsesOutputItemDoneEvent{ Type: "response.output_item.done", SequenceNumber: int64(c.nextSeq()), - OutputIndex: c.outputIndex, + OutputIndex: c.currentBlockOutputIndex, Item: wire.ResponsesOutputItemWire{ ID: c.itemID, Type: "message", @@ -266,14 +285,14 @@ func (c *anthropicBetaToResponsesConverter) emitContentBlockStop(event *anthropi wire.ResponsesFunctionCallArgumentsDoneEvent{ Type: "response.function_call_arguments.done", ItemID: pending.itemID, - OutputIndex: c.outputIndex, + OutputIndex: pending.outputIndex, Arguments: argumentsStr, SequenceNumber: int64(c.nextSeq()), }, wire.ResponsesOutputItemDoneEvent{ Type: "response.output_item.done", SequenceNumber: int64(c.nextSeq()), - OutputIndex: c.outputIndex, + OutputIndex: pending.outputIndex, Item: wire.ResponsesOutputItemWire{ Type: "function_call", ID: pending.itemID, @@ -305,29 +324,44 @@ func (c *anthropicBetaToResponsesConverter) emitCompletionEvents() { itemStatus = "incomplete" } - var output []wire.ResponsesOutputItemWire + type indexedOutput struct { + index int + item wire.ResponsesOutputItemWire + } + indexed := make([]indexedOutput, 0, len(c.pendingToolCalls)+1) if c.accumulatedText != "" { - output = append(output, wire.ResponsesOutputItemWire{ - ID: c.itemID, - Type: "message", - Status: itemStatus, - Role: "assistant", - Content: []wire.ResponsesContentPartWire{ - {Type: "output_text", Text: c.accumulatedText}, + indexed = append(indexed, indexedOutput{ + index: c.textOutputIndex, + item: wire.ResponsesOutputItemWire{ + ID: c.itemID, + Type: "message", + Status: itemStatus, + Role: "assistant", + Content: []wire.ResponsesContentPartWire{ + {Type: "output_text", Text: c.accumulatedText}, + }, }, }) } for _, pending := range c.pendingToolCalls { argumentsStr := pending.arguments.String() - output = append(output, wire.ResponsesOutputItemWire{ - Type: "function_call", - ID: pending.itemID, - CallID: pending.itemID, - Name: pending.name, - Arguments: &argumentsStr, - Status: "completed", + indexed = append(indexed, indexedOutput{ + index: pending.outputIndex, + item: wire.ResponsesOutputItemWire{ + Type: "function_call", + ID: pending.itemID, + CallID: pending.itemID, + Name: pending.name, + Arguments: &argumentsStr, + Status: itemStatus, + }, }) } + sort.Slice(indexed, func(i, j int) bool { return indexed[i].index < indexed[j].index }) + output := make([]wire.ResponsesOutputItemWire, 0, len(indexed)) + for _, entry := range indexed { + output = append(output, entry.item) + } u := c.acc.Result() resp := wire.ResponsesWireResponse{ @@ -337,6 +371,7 @@ func (c *anthropicBetaToResponsesConverter) emitCompletionEvents() { CompletedAt: c.createdAt, Output: output, Usage: responsesUsageWire(u), + Model: c.responseModel, } if isIncomplete { diff --git a/internal/protocol/stream/anthropic_beta_to_openai_responses_golden_test.go b/internal/protocol/stream/anthropic_beta_to_openai_responses_golden_test.go index 978231e53..b83340499 100644 --- a/internal/protocol/stream/anthropic_beta_to_openai_responses_golden_test.go +++ b/internal/protocol/stream/anthropic_beta_to_openai_responses_golden_test.go @@ -167,12 +167,24 @@ func TestAnthropicBetaToResponsesConverter_GoldenSequence(t *testing.T) { textDone := got[6].(wire.ResponsesOutputTextDoneEvent) assert.Equal(t, "Hello, World!", textDone.Text) + assert.Equal(t, 0, textDone.OutputIndex) + assert.Equal(t, 0, got[2].(wire.ResponsesOutputItemAddedEvent).OutputIndex) + assert.Equal(t, 0, got[3].(wire.ResponsesContentPartAddedEvent).OutputIndex) argsDone := got[12].(wire.ResponsesFunctionCallArgumentsDoneEvent) assert.Equal(t, `{"city":"Paris"}`, argsDone.Arguments) + assert.Equal(t, 1, got[9].(wire.ResponsesOutputItemAddedEvent).OutputIndex) + assert.Equal(t, 1, got[10].(wire.ResponsesFunctionCallArgumentsDeltaEvent).OutputIndex) + assert.Equal(t, 1, got[11].(wire.ResponsesFunctionCallArgumentsDeltaEvent).OutputIndex) + assert.Equal(t, 1, argsDone.OutputIndex) + assert.Equal(t, 1, got[13].(wire.ResponsesOutputItemDoneEvent).OutputIndex) completed := got[14].(wire.ResponsesCompletedEvent) assert.Equal(t, "completed", completed.Response.Status) + assert.Equal(t, "claude-3-5-sonnet-20241022", completed.Response.Model) + require.Len(t, completed.Response.Output, 2) + assert.Equal(t, "message", completed.Response.Output[0].Type) + assert.Equal(t, "function_call", completed.Response.Output[1].Type) // 4. Usage reflects upstream message_start (input) + message_delta (output). usage := conv.Usage() diff --git a/internal/protocol/stream/anthropic_to_openai_converter.go b/internal/protocol/stream/anthropic_to_openai_converter.go index 534f3e4d8..b89ec383b 100644 --- a/internal/protocol/stream/anthropic_to_openai_converter.go +++ b/internal/protocol/stream/anthropic_to_openai_converter.go @@ -6,13 +6,21 @@ import ( "time" "github.com/anthropics/anthropic-sdk-go" - anthropicstream "github.com/anthropics/anthropic-sdk-go/packages/ssestream" "github.com/tingly-dev/tingly-box/internal/protocol" usagepkg "github.com/tingly-dev/tingly-box/internal/protocol/usage" "github.com/tingly-dev/tingly-box/internal/protocol/wire" ) +// AnthropicBetaStream is the transport-neutral iterator surface consumed by +// the Anthropic Beta to OpenAI Chat stream converter. The Anthropic SDK stream +// and dormant Stage adapters both implement this contract. +type AnthropicBetaStream interface { + Next() bool + Current() anthropic.BetaRawMessageStreamEventUnion + Err() error +} + // anthropicToOpenAIConverter is a stateful iterator that reads Anthropic Beta stream // events and emits OpenAI Chat Completion wire chunks. // @@ -21,7 +29,7 @@ import ( // values ("role":"", "finish_reason":"", zero usage on every chunk), which // strict clients reject. type anthropicToOpenAIConverter struct { - stream *anthropicstream.Stream[anthropic.BetaRawMessageStreamEventUnion] + stream AnthropicBetaStream responseModel string disableStreamUsage bool hooks *AnthropicToOpenAIMCPHooks @@ -46,7 +54,7 @@ type anthropicToOpenAIConverter struct { } func newAnthropicToOpenAIConverter( - stream *anthropicstream.Stream[anthropic.BetaRawMessageStreamEventUnion], + stream AnthropicBetaStream, responseModel string, disableStreamUsage bool, hooks *AnthropicToOpenAIMCPHooks, @@ -62,6 +70,17 @@ func newAnthropicToOpenAIConverter( } } +// NewAnthropicBetaToOpenAIChatConverter creates a hook-free, transport-neutral +// converter. HTTP/SSE framing, MCP hooks, and stream close ownership remain +// with the caller. +func NewAnthropicBetaToOpenAIChatConverter( + stream AnthropicBetaStream, + responseModel string, + disableStreamUsage bool, +) StreamConverter { + return newAnthropicToOpenAIConverter(stream, responseModel, disableStreamUsage, nil) +} + func (c *anthropicToOpenAIConverter) Next() (interface{}, bool, error) { // Drain buffered events first if len(c.pending) > 0 { diff --git a/internal/protocol/stream/anthropic_to_openai_converter_test.go b/internal/protocol/stream/anthropic_to_openai_converter_test.go new file mode 100644 index 000000000..6c7c2605c --- /dev/null +++ b/internal/protocol/stream/anthropic_to_openai_converter_test.go @@ -0,0 +1,99 @@ +package stream + +import ( + "encoding/json" + "errors" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestNewAnthropicBetaToOpenAIChatConverter(t *testing.T) { + stream := &anthropicBetaSliceStream{events: anthropicBetaEvents(t, + map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": "msg_parallel", "type": "message", "role": "assistant", + "content": []any{}, "model": "provider-model", + "usage": map[string]any{"input_tokens": 4, "output_tokens": 0}, + }, + }, + map[string]any{"type": "content_block_start", "index": 0, "content_block": map[string]any{"type": "text", "text": ""}}, + map[string]any{"type": "content_block_delta", "index": 0, "delta": map[string]any{"type": "text_delta", "text": "parallel path"}}, + map[string]any{"type": "content_block_stop", "index": 0}, + map[string]any{"type": "message_delta", "delta": map[string]any{"stop_reason": "end_turn"}, "usage": map[string]any{"output_tokens": 2}}, + map[string]any{"type": "message_stop"}, + )} + converter := NewAnthropicBetaToOpenAIChatConverter(stream, "client-visible-model", false) + + var chunks []wire.ChatStreamChunk + for { + event, done, err := converter.Next() + require.NoError(t, err) + if done { + break + } + chunk, ok := event.(wire.ChatStreamChunk) + require.Truef(t, ok, "event type = %T", event) + chunks = append(chunks, chunk) + } + + require.Len(t, chunks, 3) + assert.Equal(t, "assistant", chunks[0].Choices[0].Delta.Role) + assert.Equal(t, "parallel path", chunks[1].Choices[0].Delta.Content) + require.NotNil(t, chunks[2].Choices[0].FinishReason) + assert.Equal(t, "stop", *chunks[2].Choices[0].FinishReason) + require.NotNil(t, chunks[2].Usage) + assert.EqualValues(t, 4, chunks[2].Usage.PromptTokens) + assert.EqualValues(t, 2, chunks[2].Usage.CompletionTokens) + require.NotNil(t, converter.Usage()) + assert.Equal(t, 4, converter.Usage().InputTokens) + assert.Equal(t, 2, converter.Usage().OutputTokens) +} + +func TestNewAnthropicBetaToOpenAIChatConverterPropagatesIteratorError(t *testing.T) { + want := errors.New("upstream failed before an event") + converter := NewAnthropicBetaToOpenAIChatConverter(&anthropicBetaSliceStream{err: want}, "model", false) + + event, done, err := converter.Next() + assert.Nil(t, event) + assert.False(t, done) + assert.ErrorIs(t, err, want) +} + +type anthropicBetaSliceStream struct { + events []anthropic.BetaRawMessageStreamEventUnion + index int + err error +} + +func (s *anthropicBetaSliceStream) Next() bool { + if s.index >= len(s.events) { + return false + } + s.index++ + return true +} + +func (s *anthropicBetaSliceStream) Current() anthropic.BetaRawMessageStreamEventUnion { + return s.events[s.index-1] +} + +func (s *anthropicBetaSliceStream) Err() error { return s.err } + +func anthropicBetaEvents(t *testing.T, bodies ...map[string]any) []anthropic.BetaRawMessageStreamEventUnion { + t.Helper() + events := make([]anthropic.BetaRawMessageStreamEventUnion, 0, len(bodies)) + for _, body := range bodies { + raw, err := json.Marshal(body) + require.NoError(t, err) + var event anthropic.BetaRawMessageStreamEventUnion + require.NoError(t, json.Unmarshal(raw, &event)) + events = append(events, event) + } + return events +} diff --git a/internal/protocol/stream/openai_chat_to_responses_converter.go b/internal/protocol/stream/openai_chat_to_responses_converter.go index 3f62d0aa1..bbb8fcdf8 100644 --- a/internal/protocol/stream/openai_chat_to_responses_converter.go +++ b/internal/protocol/stream/openai_chat_to_responses_converter.go @@ -7,7 +7,6 @@ import ( "time" "github.com/openai/openai-go/v3" - openaistream "github.com/openai/openai-go/v3/packages/ssestream" "github.com/tingly-dev/tingly-box/internal/protocol" protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" @@ -17,7 +16,7 @@ import ( // chatToResponsesConverter converts an OpenAI Chat Completions stream into // a sequence of Responses API events. It implements StreamConverter. type chatToResponsesConverter struct { - stream *openaistream.Stream[openai.ChatCompletionChunk] + stream OpenAIChatStream responseModel string // internal state @@ -26,6 +25,7 @@ type chatToResponsesConverter struct { sequenceNumber int64 outputIndex int textItemID string + textOutputIndex int hasTextItem bool pendingToolCalls map[int]*pendingToolCallResponse accumulatedText strings.Builder @@ -51,13 +51,14 @@ type pendingToolCallResponse struct { // NewChatToResponsesConverter creates a converter that reads from an OpenAI // Chat Completions stream and yields Responses API wire events. -func NewChatToResponsesConverter(stream *openaistream.Stream[openai.ChatCompletionChunk], responseModel string) *chatToResponsesConverter { +func NewChatToResponsesConverter(stream OpenAIChatStream, responseModel string) *chatToResponsesConverter { return &chatToResponsesConverter{ stream: stream, responseModel: responseModel, responseID: fmt.Sprintf("resp_%d", time.Now().Unix()), createdAt: time.Now().Unix(), textItemID: fmt.Sprintf("msg_%d", time.Now().UnixNano()), + textOutputIndex: -1, usage: protocol.ZeroTokenUsage(), pendingToolCalls: make(map[int]*pendingToolCallResponse), } @@ -127,6 +128,8 @@ func (c *chatToResponsesConverter) processChunk(chunk *openai.ChatCompletionChun // Handle content delta if choice.Delta.Content != "" { if !c.hasTextItem { + c.textOutputIndex = c.outputIndex + c.outputIndex++ c.emitTextItemAdded() c.hasTextItem = true } @@ -135,7 +138,7 @@ func (c *chatToResponsesConverter) processChunk(chunk *openai.ChatCompletionChun Type: "response.output_text.delta", SequenceNumber: c.nextSeq(), ItemID: c.textItemID, - OutputIndex: 0, + OutputIndex: c.textOutputIndex, ContentIndex: 0, Delta: choice.Delta.Content, Logprobs: []interface{}{}, @@ -152,10 +155,6 @@ func (c *chatToResponsesConverter) processChunk(chunk *openai.ChatCompletionChun itemID = truncateToolCallID(toolCall.ID) } - // Reserve OutputIndex 0 for the text message item; tool calls start at 1. - if c.outputIndex == 0 { - c.outputIndex = 1 - } toolOutputIndex := c.outputIndex c.outputIndex++ @@ -235,15 +234,27 @@ func (c *chatToResponsesConverter) emitCompletionEvents() { Type: "response.output_text.done", SequenceNumber: c.nextSeq(), ItemID: c.textItemID, - OutputIndex: 0, + OutputIndex: c.textOutputIndex, ContentIndex: 0, Text: text, Logprobs: []interface{}{}, }) + c.pending = append(c.pending, wire.ResponsesContentPartDoneEvent{ + Type: "response.content_part.done", + SequenceNumber: c.nextSeq(), + OutputIndex: c.textOutputIndex, + ItemID: c.textItemID, + ContentIndex: 0, + Part: wire.ResponsesContentPartWire{ + Type: "output_text", + Text: text, + Annotations: []interface{}{}, + }, + }) c.pending = append(c.pending, wire.ResponsesOutputItemDoneEvent{ Type: "response.output_item.done", SequenceNumber: c.nextSeq(), - OutputIndex: 0, + OutputIndex: c.textOutputIndex, Item: newResponsesMessageItem(c.textItemID, "completed", text), }) } @@ -283,9 +294,16 @@ func (c *chatToResponsesConverter) emitCompletionEvents() { itemStatus = "incomplete" } - var output []wire.ResponsesOutputItemWire + type indexedOutput struct { + index int + item wire.ResponsesOutputItemWire + } + indexed := make([]indexedOutput, 0, len(c.pendingToolCalls)+1) if c.accumulatedText.Len() > 0 { - output = append(output, newResponsesMessageItem(c.textItemID, itemStatus, c.accumulatedText.String())) + indexed = append(indexed, indexedOutput{ + index: c.textOutputIndex, + item: newResponsesMessageItem(c.textItemID, itemStatus, c.accumulatedText.String()), + }) } for _, idx := range sortedIndexes { ptc := c.pendingToolCalls[idx] @@ -293,7 +311,15 @@ func (c *chatToResponsesConverter) emitCompletionEvents() { if callID == "" { callID = ptc.itemID } - output = append(output, newResponsesFunctionCallItem(ptc.itemID, callID, ptc.name, ptc.arguments.String(), itemStatus)) + indexed = append(indexed, indexedOutput{ + index: ptc.outputIdx, + item: newResponsesFunctionCallItem(ptc.itemID, callID, ptc.name, ptc.arguments.String(), itemStatus), + }) + } + sort.Slice(indexed, func(i, j int) bool { return indexed[i].index < indexed[j].index }) + output := make([]wire.ResponsesOutputItemWire, 0, len(indexed)) + for _, entry := range indexed { + output = append(output, entry.item) } if isIncomplete { @@ -328,15 +354,23 @@ func chatFinishReasonToIncomplete(finishReason string) (bool, string) { } func (c *chatToResponsesConverter) emitTextItemAdded() { - if c.outputIndex == 0 { - c.outputIndex = 1 - } c.pending = append(c.pending, wire.ResponsesOutputItemAddedEvent{ Type: "response.output_item.added", SequenceNumber: c.nextSeq(), - OutputIndex: 0, + OutputIndex: c.textOutputIndex, Item: newResponsesMessageItem(c.textItemID, "in_progress", ""), }) + c.pending = append(c.pending, wire.ResponsesContentPartAddedEvent{ + Type: "response.content_part.added", + SequenceNumber: c.nextSeq(), + ItemID: c.textItemID, + OutputIndex: c.textOutputIndex, + ContentIndex: 0, + Part: wire.ResponsesContentPartWire{ + Type: "output_text", + Annotations: []interface{}{}, + }, + }) } func (c *chatToResponsesConverter) nextSeq() int64 { diff --git a/internal/protocol/stream/openai_chat_to_responses_golden_test.go b/internal/protocol/stream/openai_chat_to_responses_golden_test.go index d0ca53601..c1b322dd1 100644 --- a/internal/protocol/stream/openai_chat_to_responses_golden_test.go +++ b/internal/protocol/stream/openai_chat_to_responses_golden_test.go @@ -110,16 +110,18 @@ func TestChatToResponsesConverter_GoldenSequence(t *testing.T) { // 1. Exact ordered event sequence — the heart of the oracle. want := []string{ "response.created", - "response.output_item.added", // text item (index 0) + "response.output_item.added", // text item (index 0) + "response.content_part.added", "response.output_text.delta", // "Hello" "response.output_text.delta", // ", World!" "response.output_item.added", // function call (index 1) "response.function_call_arguments.delta", // {"city": "response.function_call_arguments.delta", // "Paris"} "response.output_text.done", // text done before tool done - "response.output_item.done", // text message item - "response.function_call_arguments.done", // full arguments - "response.output_item.done", // function call item + "response.content_part.done", + "response.output_item.done", // text message item + "response.function_call_arguments.done", // full arguments + "response.output_item.done", // function call item "response.completed", } gotTypes := make([]string, len(got)) @@ -134,17 +136,17 @@ func TestChatToResponsesConverter_GoldenSequence(t *testing.T) { } // 3. Spot-check key payloads. - assert.Equal(t, "Hello", got[2].(wire.ResponsesOutputTextDeltaEvent).Delta) - assert.Equal(t, ", World!", got[3].(wire.ResponsesOutputTextDeltaEvent).Delta) + assert.Equal(t, "Hello", got[3].(wire.ResponsesOutputTextDeltaEvent).Delta) + assert.Equal(t, ", World!", got[4].(wire.ResponsesOutputTextDeltaEvent).Delta) - textDone := got[7].(wire.ResponsesOutputTextDoneEvent) + textDone := got[8].(wire.ResponsesOutputTextDoneEvent) assert.Equal(t, "Hello, World!", textDone.Text) - argsDone := got[9].(wire.ResponsesFunctionCallArgumentsDoneEvent) + argsDone := got[11].(wire.ResponsesFunctionCallArgumentsDoneEvent) assert.Equal(t, "get_weather", argsDone.Name) assert.Equal(t, `{"city":"Paris"}`, argsDone.Arguments) - completed := got[11].(wire.ResponsesCompletedEvent) + completed := got[13].(wire.ResponsesCompletedEvent) assert.Equal(t, "completed", completed.Response.Status) require.Len(t, completed.Response.Output, 2, "final output carries text + tool-call items") @@ -163,12 +165,16 @@ func seqOf(t *testing.T, e wire.ResponsesEvent) int64 { return v.SequenceNumber case wire.ResponsesOutputItemAddedEvent: return v.SequenceNumber + case wire.ResponsesContentPartAddedEvent: + return v.SequenceNumber case wire.ResponsesOutputTextDeltaEvent: return v.SequenceNumber case wire.ResponsesFunctionCallArgumentsDeltaEvent: return v.SequenceNumber case wire.ResponsesOutputTextDoneEvent: return v.SequenceNumber + case wire.ResponsesContentPartDoneEvent: + return v.SequenceNumber case wire.ResponsesOutputItemDoneEvent: return v.SequenceNumber case wire.ResponsesFunctionCallArgumentsDoneEvent: diff --git a/internal/protocol/stream/openai_chat_to_responses_test.go b/internal/protocol/stream/openai_chat_to_responses_test.go index 835878020..98824a6e0 100644 --- a/internal/protocol/stream/openai_chat_to_responses_test.go +++ b/internal/protocol/stream/openai_chat_to_responses_test.go @@ -201,13 +201,14 @@ func TestChatToResponsesConverter_TextDelta(t *testing.T) { }, }) - // Should emit: response.created, output_item.added, output_text.delta - require.Len(t, conv.pending, 3) + // Should emit the complete Responses text lifecycle before the delta. + require.Len(t, conv.pending, 4) assert.Equal(t, "response.created", conv.pending[0].(wire.ResponsesCreatedEvent).Type) assert.Equal(t, "response.output_item.added", conv.pending[1].(wire.ResponsesOutputItemAddedEvent).Type) - assert.Equal(t, "response.output_text.delta", conv.pending[2].(wire.ResponsesOutputTextDeltaEvent).Type) + assert.Equal(t, "response.content_part.added", conv.pending[2].(wire.ResponsesContentPartAddedEvent).Type) + assert.Equal(t, "response.output_text.delta", conv.pending[3].(wire.ResponsesOutputTextDeltaEvent).Type) - delta := conv.pending[2].(wire.ResponsesOutputTextDeltaEvent) + delta := conv.pending[3].(wire.ResponsesOutputTextDeltaEvent) assert.Equal(t, "Hello, World!", delta.Delta) } @@ -241,6 +242,42 @@ func TestChatToResponsesConverter_ToolCall(t *testing.T) { assert.Equal(t, `{"loc`, argsDelta.Delta) } +func TestChatToResponsesConverter_ToolOnlyUsesFirstOutputIndex(t *testing.T) { + conv := NewChatToResponsesConverter(nil, "gpt-4o-mini") + conv.hasSentCreated = true + + conv.processChunk(&openai.ChatCompletionChunk{Choices: []openai.ChatCompletionChunkChoice{{ + Delta: openai.ChatCompletionChunkChoiceDelta{ToolCalls: []openai.ChatCompletionChunkChoiceDeltaToolCall{{ + Index: 0, + ID: "call_123", + Function: openai.ChatCompletionChunkChoiceDeltaToolCallFunction{ + Name: "get_weather", Arguments: `{"city":"Paris"}`, + }, + }}}, + }}}) + conv.processChunk(&openai.ChatCompletionChunk{Choices: []openai.ChatCompletionChunkChoice{{ + FinishReason: "tool_calls", + }}}) + + added := conv.pending[0].(wire.ResponsesOutputItemAddedEvent) + assert.Equal(t, 0, added.OutputIndex) + + var done wire.ResponsesOutputItemDoneEvent + var completed wire.ResponsesCompletedEvent + for _, event := range conv.pending { + switch event := event.(type) { + case wire.ResponsesOutputItemDoneEvent: + done = event + case wire.ResponsesCompletedEvent: + completed = event + } + } + assert.Equal(t, 0, done.OutputIndex) + require.Len(t, completed.Response.Output, 1) + assert.Equal(t, "function_call", completed.Response.Output[0].Type) + assert.Equal(t, "call_123", completed.Response.Output[0].CallID) +} + // TestChatToResponsesConverter_CompletedEvent tests usage propagation // into the response.completed event. func TestChatToResponsesConverter_CompletedEvent(t *testing.T) { diff --git a/internal/protocol/stream/openai_responses_to_anthropic_assembly_test.go b/internal/protocol/stream/openai_responses_to_anthropic_assembly_test.go index da0acdbcb..a4f7ffa5a 100644 --- a/internal/protocol/stream/openai_responses_to_anthropic_assembly_test.go +++ b/internal/protocol/stream/openai_responses_to_anthropic_assembly_test.go @@ -105,7 +105,9 @@ func TestHandleResponsesToAnthropicV1Assembly_Golden(t *testing.T) { toolBlock := content[1].(map[string]any) assert.Equal(t, "tool_use", toolBlock["type"]) - assert.Equal(t, "fc_1", toolBlock["id"]) + // Anthropic tool_result correlation uses Responses call_id, while the + // response item id remains an internal stream assembly identifier. + assert.Equal(t, "call_1", toolBlock["id"]) assert.Equal(t, "get_weather", toolBlock["name"]) assert.Equal(t, map[string]any{"city": "Paris"}, toolBlock["input"]) diff --git a/internal/protocol/stream/openai_responses_to_anthropic_converter.go b/internal/protocol/stream/openai_responses_to_anthropic_converter.go index 9c7a2f371..363a119f7 100644 --- a/internal/protocol/stream/openai_responses_to_anthropic_converter.go +++ b/internal/protocol/stream/openai_responses_to_anthropic_converter.go @@ -57,6 +57,17 @@ func newResponsesToAnthropicConverter(ctx context.Context, stream ResponsesStrea } } +// NewOpenAIResponsesToAnthropicConverter creates a transport-neutral +// Responses-to-Anthropic converter. The caller owns stream closure and wire +// framing; emitted values can be normalized with AsAnthropicEvent. +func NewOpenAIResponsesToAnthropicConverter( + ctx context.Context, + stream ResponsesStreamIter, + responseModel string, +) StreamConverter { + return newResponsesToAnthropicConverter(ctx, stream, responseModel) +} + func (r *responsesToAnthropicConverter) Next() (interface{}, bool, error) { if !r.messageStartSent { r.emitMessageStart() @@ -323,7 +334,11 @@ func (r *responsesToAnthropicConverter) processEvent(currentEvent responses.Resp r.emitContentBlockDelta(r.state.textBlockIndex, anthropicTextDelta(textDelta.Delta)) case "function_call", "custom_tool_call", "mcp_call": itemID := itemAdded.Item.ID - truncatedID := truncateToolCallID(itemID) + callID := itemAdded.Item.CallID + if callID == "" { + callID = itemID + } + truncatedID := truncateToolCallID(callID) blockIndex := r.state.nextBlockIndex r.state.nextBlockIndex++ @@ -537,7 +552,11 @@ func (r *responsesToAnthropicConverter) finalize(resp *responses.Response, stopR continue } - truncatedID := truncateToolCallID(itemID) + callID := outputItem.CallID + if callID == "" { + callID = itemID + } + truncatedID := truncateToolCallID(callID) blockIndex := r.state.nextBlockIndex r.state.nextBlockIndex++ diff --git a/internal/protocol/stream/openai_responses_to_anthropic_golden_test.go b/internal/protocol/stream/openai_responses_to_anthropic_golden_test.go index fd35e1481..a9a9c13a2 100644 --- a/internal/protocol/stream/openai_responses_to_anthropic_golden_test.go +++ b/internal/protocol/stream/openai_responses_to_anthropic_golden_test.go @@ -120,7 +120,9 @@ func TestResponsesToAnthropicConverter_GoldenSequence(t *testing.T) { toolBlockStart := eventDataAsMap(t, got[5].data) toolBlock := toolBlockStart["content_block"].(map[string]interface{}) assert.Equal(t, "tool_use", toolBlock["type"]) + assert.Equal(t, "call_1", toolBlock["id"], "Anthropic tool_result must reference the Responses call_id") assert.Equal(t, "get_weather", toolBlock["name"]) + assert.Equal(t, "call_1", toolBlock["id"]) // args delta argsDelta := eventDataAsMap(t, got[6].data)["delta"].(map[string]interface{}) diff --git a/internal/protocol/stream/openai_responses_to_chat_converter.go b/internal/protocol/stream/openai_responses_to_chat_converter.go index 605a6d0b3..3e265529e 100644 --- a/internal/protocol/stream/openai_responses_to_chat_converter.go +++ b/internal/protocol/stream/openai_responses_to_chat_converter.go @@ -57,6 +57,16 @@ func newResponsesToChatConverter(stream ResponsesStreamIter, responseModel strin } } +// NewOpenAIResponsesToChatConverter creates a transport-neutral Responses to +// Chat stream converter. The caller owns stream closure and SSE framing. +func NewOpenAIResponsesToChatConverter( + stream ResponsesStreamIter, + responseModel string, + disableUsage bool, +) StreamConverter { + return newResponsesToChatConverter(stream, responseModel, disableUsage) +} + func (c *responsesToChatConverter) Next() (interface{}, bool, error) { // Drain buffered events first if len(c.pending) > 0 { diff --git a/internal/protocol/stream/openai_to_anthropic.go b/internal/protocol/stream/openai_to_anthropic.go index e645ec2e8..11fa560f2 100644 --- a/internal/protocol/stream/openai_to_anthropic.go +++ b/internal/protocol/stream/openai_to_anthropic.go @@ -99,6 +99,12 @@ type OpenAIToAnthropicMCPHooks struct { var ErrMCPStreamContinue = errors.New("mcp stream should continue") +// NewOpenAIChatToAnthropicV1Converter creates the transport-free V1 stream +// state machine. The caller owns driving and closing the supplied stream. +func NewOpenAIChatToAnthropicV1Converter(stream OpenAIChatStream, responseModel string, req *openai.ChatCompletionNewParams) StreamConverter { + return newOpenAIToAnthropicConverter(stream, responseModel, req, nil, mapOpenAIFinishReasonToAnthropic) +} + // HandleOpenAIToAnthropicStreamResponse processes OpenAI streaming events and converts them to Anthropic format. // Returns UsageStat containing token usage information for tracking. func HandleOpenAIToAnthropicStreamResponse(hc *protocol.HandleContext, req *openai.ChatCompletionNewParams, stream *openaistream.Stream[openai.ChatCompletionChunk], responseModel string) (*protocol.TokenUsage, error) { diff --git a/internal/protocol/stream/openai_to_anthropic_beta.go b/internal/protocol/stream/openai_to_anthropic_beta.go index 9b7387bb2..b798eba70 100644 --- a/internal/protocol/stream/openai_to_anthropic_beta.go +++ b/internal/protocol/stream/openai_to_anthropic_beta.go @@ -13,6 +13,12 @@ import ( "github.com/tingly-dev/tingly-box/internal/protocol" ) +// NewOpenAIChatToAnthropicBetaConverter creates the transport-free beta stream +// state machine. The caller owns driving and closing the supplied stream. +func NewOpenAIChatToAnthropicBetaConverter(stream OpenAIChatStream, responseModel string, req *openai.ChatCompletionNewParams) StreamConverter { + return newOpenAIToAnthropicConverter(stream, responseModel, req, nil, mapOpenAIFinishReasonToAnthropicBeta) +} + // HandleOpenAIToAnthropicBetaStream processes OpenAI streaming events and converts them to Anthropic beta format. // Returns UsageStat containing token usage information for tracking. func HandleOpenAIToAnthropicBetaStream(hc *protocol.HandleContext, req *openai.ChatCompletionNewParams, stream *openaistream.Stream[openai.ChatCompletionChunk], responseModel string) (*protocol.TokenUsage, error) { diff --git a/internal/protocol/stream/openai_to_anthropic_converter.go b/internal/protocol/stream/openai_to_anthropic_converter.go index f173b93ff..46fbab23c 100644 --- a/internal/protocol/stream/openai_to_anthropic_converter.go +++ b/internal/protocol/stream/openai_to_anthropic_converter.go @@ -1,13 +1,13 @@ package stream import ( + "encoding/json" "fmt" "sort" "time" "github.com/gin-gonic/gin" "github.com/openai/openai-go/v3" - openaistream "github.com/openai/openai-go/v3/packages/ssestream" "github.com/sirupsen/logrus" "github.com/tingly-dev/tingly-box/internal/protocol" @@ -23,10 +23,49 @@ type anthropicStreamEvent struct { data any } +// AnthropicEvent is the transport-neutral view of an Anthropic stream event. +// HTTP writers may keep using the internal representation; protocol stages use +// this exported value to carry event name and data without taking over SSE +// framing. +type AnthropicEvent struct { + Type string + Data any +} + +// RawJSON returns the protocol payload rather than the transport wrapper. +// Protocol-owned assemblers use this to reconstruct converted streams without +// mistaking the Type/Data carrier itself for an Anthropic wire event. +func (e AnthropicEvent) RawJSON() string { + raw, err := json.Marshal(e.Data) + if err != nil { + return "" + } + return string(raw) +} + +// AsAnthropicEvent exposes an event emitted by an Anthropic stream converter. +func AsAnthropicEvent(event any) (AnthropicEvent, bool) { + value, ok := event.(anthropicStreamEvent) + if !ok { + return AnthropicEvent{}, false + } + return AnthropicEvent{Type: value.eventType, Data: value.data}, true +} + +// OpenAIChatStream is the minimum iterator surface required by the Chat to +// Anthropic state machine. The OpenAI SDK stream and stage stream adapters both +// implement it. +type OpenAIChatStream interface { + Next() bool + Current() openai.ChatCompletionChunk + Err() error + Close() error +} + // openAIToAnthropicConverter is a stateful iterator that reads OpenAI Chat Completion // chunks and emits Anthropic SSE events (map-based). type openAIToAnthropicConverter struct { - stream *openaistream.Stream[openai.ChatCompletionChunk] + stream OpenAIChatStream responseModel string req *openai.ChatCompletionNewParams hooks *OpenAIToAnthropicMCPHooks @@ -47,7 +86,7 @@ type openAIToAnthropicConverter struct { } func newOpenAIToAnthropicConverter( - stream *openaistream.Stream[openai.ChatCompletionChunk], + stream OpenAIChatStream, responseModel string, req *openai.ChatCompletionNewParams, hooks *OpenAIToAnthropicMCPHooks, diff --git a/internal/protocol/transform/provider_cleanup.go b/internal/protocol/transform/provider_cleanup.go new file mode 100644 index 000000000..e1a36f820 --- /dev/null +++ b/internal/protocol/transform/provider_cleanup.go @@ -0,0 +1,28 @@ +package transform + +import ( + "fmt" + + "github.com/openai/openai-go/v3" + "github.com/tingly-dev/tingly-box/internal/protocol/request" +) + +// OpenAIChatProviderCleanupTransform removes gateway-only fields after vendor +// transforms and before the request reaches provider-bound observers. +type OpenAIChatProviderCleanupTransform struct{} + +func NewOpenAIChatProviderCleanupTransform() *OpenAIChatProviderCleanupTransform { + return &OpenAIChatProviderCleanupTransform{} +} + +func (*OpenAIChatProviderCleanupTransform) Name() string { return "openai_chat_provider_cleanup" } + +func (*OpenAIChatProviderCleanupTransform) Apply(ctx *TransformContext) error { + requestValue, ok := ctx.Request.(*openai.ChatCompletionNewParams) + if !ok || requestValue == nil { + return fmt.Errorf("OpenAI Chat provider cleanup received %T", ctx.Request) + } + request.CleanupOpenaiFields(requestValue) + ctx.Request = requestValue + return nil +} diff --git a/internal/protocol/transform/provider_cleanup_test.go b/internal/protocol/transform/provider_cleanup_test.go new file mode 100644 index 000000000..5490ec1bc --- /dev/null +++ b/internal/protocol/transform/provider_cleanup_test.go @@ -0,0 +1,29 @@ +package transform + +import ( + "encoding/json" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/stretchr/testify/require" +) + +func TestOpenAIChatProviderCleanupRemovesGatewayFields(t *testing.T) { + var request openai.ChatCompletionNewParams + require.NoError(t, json.Unmarshal([]byte(`{ + "model":"provider-model", + "messages":[ + {"role":"assistant","content":"prior","x_thinking":"private"} + ], + "tools":[] + }`), &request)) + ctx := NewTransformContext(&request) + + require.NoError(t, NewOpenAIChatProviderCleanupTransform().Apply(ctx)) + cleaned, ok := ctx.Request.(*openai.ChatCompletionNewParams) + require.True(t, ok) + body, err := json.Marshal(cleaned) + require.NoError(t, err) + require.NotContains(t, string(body), "x_thinking") + require.NotContains(t, string(body), `"tools"`) +} diff --git a/internal/protocol/wire/openai_chat.go b/internal/protocol/wire/openai_chat.go index 55e0d3cee..5c873a9ad 100644 --- a/internal/protocol/wire/openai_chat.go +++ b/internal/protocol/wire/openai_chat.go @@ -1,7 +1,5 @@ package wire -import "encoding/json" - // Chat Completions stream DTOs preserve the minimal outbound JSON shape emitted by this proxy. // Keep these fields checked against openai-go Chat Completions stream types when updating the SDK. type ChatStreamChunk struct { @@ -73,35 +71,95 @@ type ChatStreamError struct { // ChatCompletionWire is the OpenAI Chat Completions response wire format. type ChatCompletionWire struct { - ID string `json:"id"` - Object string `json:"object"` - Created int64 `json:"created"` - Model string `json:"model"` + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` Choices []ChatCompletionChoiceWire `json:"choices"` - Usage ChatCompletionUsageWire `json:"usage"` + Usage ChatCompletionUsageWire `json:"usage"` } -// ToMap serializes to a generic map for callers that apply runtime transforms. +// ToMap converts the typed wire contract to the legacy map surface used by +// runtime response transforms. Keep this explicit so Stage Bridges never need +// a JSON marshal/unmarshal round-trip to obtain a wire DTO. func (r ChatCompletionWire) ToMap() map[string]interface{} { - raw, _ := json.Marshal(r) - var m map[string]interface{} - _ = json.Unmarshal(raw, &m) - return m + choices := make([]map[string]any, 0, len(r.Choices)) + for _, choice := range r.Choices { + message := map[string]any{"role": choice.Message.Role} + if choice.Message.Content != "" { + message["content"] = choice.Message.Content + } + if choice.Message.Refusal != "" { + message["refusal"] = choice.Message.Refusal + } + if choice.Message.ReasoningContent != "" { + message["reasoning_content"] = choice.Message.ReasoningContent + } + if len(choice.Message.ToolCalls) > 0 { + toolCalls := make([]map[string]any, 0, len(choice.Message.ToolCalls)) + for _, toolCall := range choice.Message.ToolCalls { + toolCalls = append(toolCalls, map[string]any{ + "id": toolCall.ID, + "type": toolCall.Type, + "function": map[string]any{ + "name": toolCall.Function.Name, + "arguments": toolCall.Function.Arguments, + }, + }) + } + message["tool_calls"] = toolCalls + } + choices = append(choices, map[string]any{ + "index": choice.Index, + "message": message, + "finish_reason": choice.FinishReason, + }) + } + + usage := map[string]any{ + "prompt_tokens": r.Usage.PromptTokens, + "completion_tokens": r.Usage.CompletionTokens, + "total_tokens": r.Usage.TotalTokens, + } + if r.Usage.PromptTokensDetails != nil { + details := map[string]any{ + "cached_tokens": r.Usage.PromptTokensDetails.CachedTokens, + } + if r.Usage.PromptTokensDetails.CacheWriteTokens > 0 { + details["cache_write_tokens"] = r.Usage.PromptTokensDetails.CacheWriteTokens + } + usage["prompt_tokens_details"] = details + } + if r.Usage.CompletionTokensDetails != nil { + usage["completion_tokens_details"] = map[string]any{ + "reasoning_tokens": r.Usage.CompletionTokensDetails.ReasoningTokens, + } + } + + return map[string]any{ + "id": r.ID, + "object": r.Object, + "created": r.Created, + "model": r.Model, + "choices": choices, + "usage": usage, + } } // ChatCompletionChoiceWire is a single choice in the OpenAI Chat Completions response. type ChatCompletionChoiceWire struct { - Index int `json:"index"` + Index int `json:"index"` Message ChatCompletionMessageWire `json:"message"` - FinishReason string `json:"finish_reason"` + FinishReason string `json:"finish_reason"` } // ChatCompletionMessageWire is the message inside a choice. type ChatCompletionMessageWire struct { - Role string `json:"role"` - Content string `json:"content,omitempty"` + Role string `json:"role"` + Content string `json:"content,omitempty"` + Refusal string `json:"refusal,omitempty"` ToolCalls []ChatCompletionToolCallWire `json:"tool_calls,omitempty"` - ReasoningContent string `json:"reasoning_content,omitempty"` + ReasoningContent string `json:"reasoning_content,omitempty"` } // ChatCompletionToolCallWire is a single tool call inside a message. @@ -121,10 +179,11 @@ type ChatCompletionFunctionWire struct { // prompt_tokens = TOTAL (uncached + cached + written); cached_tokens and // cache_write_tokens are reported, disjoint subsets of it. type ChatCompletionUsageWire struct { - PromptTokens int64 `json:"prompt_tokens"` - CompletionTokens int64 `json:"completion_tokens"` - TotalTokens int64 `json:"total_tokens"` - PromptTokensDetails *ChatCompletionPromptDetailsWire `json:"prompt_tokens_details,omitempty"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + TotalTokens int64 `json:"total_tokens"` + PromptTokensDetails *ChatCompletionPromptDetailsWire `json:"prompt_tokens_details,omitempty"` + CompletionTokensDetails *ChatCompletionOutputDetailsWire `json:"completion_tokens_details,omitempty"` } // ChatCompletionPromptDetailsWire breaks down prompt token categories. @@ -133,3 +192,8 @@ type ChatCompletionPromptDetailsWire struct { CachedTokens int64 `json:"cached_tokens"` CacheWriteTokens int64 `json:"cache_write_tokens,omitempty"` } + +// ChatCompletionOutputDetailsWire breaks down completion token categories. +type ChatCompletionOutputDetailsWire struct { + ReasoningTokens int64 `json:"reasoning_tokens"` +} diff --git a/internal/protocol/wire/openai_responses.go b/internal/protocol/wire/openai_responses.go index 18ef0874a..2ffba8e87 100644 --- a/internal/protocol/wire/openai_responses.go +++ b/internal/protocol/wire/openai_responses.go @@ -26,8 +26,8 @@ func (e ResponsesFunctionCallArgumentsDeltaEvent) EventType() string { return e. func (e ResponsesFunctionCallArgumentsDoneEvent) EventType() string { return e.Type } type ResponsesStreamErrorEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` Error ResponsesStreamErrorBody `json:"error"` } @@ -37,26 +37,26 @@ type ResponsesStreamErrorBody struct { } type ResponsesCreatedEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` Response ResponsesWireResponse `json:"response"` } type ResponsesInProgressEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` Response ResponsesWireResponse `json:"response"` } type ResponsesCompletedEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` Response ResponsesWireResponse `json:"response"` } type ResponsesIncompleteEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` Response ResponsesWireResponse `json:"response"` } @@ -72,14 +72,102 @@ type ResponsesWireResponse struct { IncompleteDetails *ResponsesIncompleteDetailsWire `json:"incomplete_details,omitempty"` } +// ToMap converts the typed response contract to the legacy map surface used by +// existing non-stream handlers. Stage Bridges should pass the typed value. +func (r ResponsesWireResponse) ToMap() map[string]any { + output := make([]map[string]any, 0, len(r.Output)) + for _, item := range r.Output { + value := map[string]any{ + "id": item.ID, + "type": item.Type, + } + if item.Status != "" { + value["status"] = item.Status + } + if item.Role != "" { + value["role"] = item.Role + } + if len(item.Content) > 0 { + content := make([]map[string]any, 0, len(item.Content)) + for _, part := range item.Content { + partValue := map[string]any{"type": part.Type} + if part.Type == "output_text" || part.Text != "" { + partValue["text"] = part.Text + } + if part.Type == "output_text" { + annotations := part.Annotations + if annotations == nil { + annotations = []any{} + } + partValue["annotations"] = annotations + } else if part.Annotations != nil { + partValue["annotations"] = part.Annotations + } + content = append(content, partValue) + } + value["content"] = content + } + if item.CallID != "" { + value["call_id"] = item.CallID + } + if item.Name != "" { + value["name"] = item.Name + } + if item.Arguments != nil { + value["arguments"] = *item.Arguments + } + output = append(output, value) + } + + result := map[string]any{ + "id": r.ID, + "object": r.Object, + "created_at": r.CreatedAt, + "status": r.Status, + "output": output, + } + if r.Model != "" { + result["model"] = r.Model + } + if r.CompletedAt != 0 { + result["completed_at"] = r.CompletedAt + } + if r.IncompleteDetails != nil { + result["incomplete_details"] = map[string]any{"reason": r.IncompleteDetails.Reason} + } + if r.Usage != nil { + usage := map[string]any{ + "input_tokens": r.Usage.InputTokens, + "output_tokens": r.Usage.OutputTokens, + "total_tokens": r.Usage.TotalTokens, + } + if r.Usage.InputTokensDetails.CachedTokens > 0 || r.Usage.InputTokensDetails.CacheWriteTokens > 0 { + details := map[string]any{ + "cached_tokens": r.Usage.InputTokensDetails.CachedTokens, + } + if r.Usage.InputTokensDetails.CacheWriteTokens > 0 { + details["cache_write_tokens"] = r.Usage.InputTokensDetails.CacheWriteTokens + } + usage["input_tokens_details"] = details + } + if r.Usage.OutputTokensDetails.ReasoningTokens > 0 { + usage["output_tokens_details"] = map[string]any{ + "reasoning_tokens": r.Usage.OutputTokensDetails.ReasoningTokens, + } + } + result["usage"] = usage + } + return result +} + type ResponsesIncompleteDetailsWire struct { Reason string `json:"reason"` } type ResponsesUsageWire struct { - InputTokens int64 `json:"input_tokens"` - OutputTokens int64 `json:"output_tokens"` - TotalTokens int64 `json:"total_tokens"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + TotalTokens int64 `json:"total_tokens"` InputTokensDetails ResponsesInputTokensDetailsWire `json:"input_tokens_details,omitempty"` OutputTokensDetails ResponsesOutputTokensDetailsWire `json:"output_tokens_details,omitempty"` } @@ -103,28 +191,28 @@ type ResponsesOutputTokensDetailsWire struct { } type ResponsesOutputItemAddedEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` - OutputIndex int `json:"output_index"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` + OutputIndex int `json:"output_index"` Item ResponsesOutputItemWire `json:"item"` } type ResponsesOutputItemDoneEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` - OutputIndex int `json:"output_index"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` + OutputIndex int `json:"output_index"` Item ResponsesOutputItemWire `json:"item"` } type ResponsesOutputItemWire struct { - ID string `json:"id"` - Type string `json:"type"` - Role string `json:"role,omitempty"` - Status string `json:"status"` - Content []ResponsesContentPartWire `json:"content,omitempty"` - CallID string `json:"call_id,omitempty"` - Name string `json:"name,omitempty"` - Arguments *string `json:"arguments,omitempty"` + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role,omitempty"` + Status string `json:"status,omitempty"` + Content []ResponsesContentPartWire `json:"content,omitempty"` + CallID string `json:"call_id,omitempty"` + Name string `json:"name,omitempty"` + Arguments *string `json:"arguments,omitempty"` } type ResponsesContentPartWire struct { @@ -157,20 +245,20 @@ func (p ResponsesContentPartWire) MarshalJSON() ([]byte, error) { } type ResponsesContentPartAddedEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` - ItemID string `json:"item_id"` - OutputIndex int `json:"output_index"` - ContentIndex int `json:"content_index"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` + ItemID string `json:"item_id"` + OutputIndex int `json:"output_index"` + ContentIndex int `json:"content_index"` Part ResponsesContentPartWire `json:"part"` } type ResponsesContentPartDoneEvent struct { - Type string `json:"type"` - SequenceNumber int64 `json:"sequence_number"` - ItemID string `json:"item_id"` - OutputIndex int `json:"output_index"` - ContentIndex int `json:"content_index"` + Type string `json:"type"` + SequenceNumber int64 `json:"sequence_number"` + ItemID string `json:"item_id"` + OutputIndex int `json:"output_index"` + ContentIndex int `json:"content_index"` Part ResponsesContentPartWire `json:"part"` } diff --git a/internal/record/boundary_matrix_test.go b/internal/record/boundary_matrix_test.go new file mode 100644 index 000000000..85bf54f5e --- /dev/null +++ b/internal/record/boundary_matrix_test.go @@ -0,0 +1,224 @@ +package record + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestProviderObserverBoundaryMatrix(t *testing.T) { + routes := []struct { + name string + source protocol.APIType + target protocol.APIType + }{ + {name: "v1_to_v1", source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicV1}, + {name: "v1_to_chat", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat}, + {name: "v1_to_responses", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIResponses}, + {name: "beta_to_beta", source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta}, + {name: "beta_to_chat", source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIChat}, + {name: "beta_to_responses", source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIResponses}, + {name: "chat_to_chat", source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIChat}, + {name: "chat_to_beta", source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta}, + {name: "chat_to_responses", source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIResponses}, + {name: "responses_to_responses", source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIResponses}, + {name: "responses_to_chat", source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIChat}, + {name: "responses_to_beta", source: protocol.TypeOpenAIResponses, target: protocol.TypeAnthropicBeta}, + } + + for _, route := range routes { + route := route + for _, streaming := range []bool{false, true} { + mode := "complete" + if streaming { + mode = "stream" + } + t.Run(route.name+"/"+mode, func(t *testing.T) { + t.Parallel() + testProviderObserverRoute(t, route.source, route.target, streaming) + }) + } + } +} + +func testProviderObserverRoute(t *testing.T, source, target protocol.APIType, streaming bool) { + t.Helper() + recorder, err := New(Config{ + Enabled: true, + RequestID: "matrix-" + string(source) + "-" + string(target), + InputProtocol: source, + Input: map[string]any{"boundary": "input", "protocol": source}, + }) + require.NoError(t, err) + + terminal := &matrixProviderEndpoint{api: target} + observed := ObserveProvider(terminal, recorder, ExchangeMetadata{ + Attempt: 1, + Provider: "matrix-provider", + Model: "matrix-model", + }) + + var bridges []stage.Bridge + if source != target { + bridges = append(bridges, matrixBridge{source: source, target: target}) + } + registry, err := stage.NewBridgeRegistry(bridges...) + require.NoError(t, err) + topology, err := stage.BuildTopology(stage.TopologyConfig{ + Terminal: observed, + ClientProtocol: source, + Registry: registry, + RequiredCapabilities: stage.AllBridgeCapabilities, + }) + require.NoError(t, err) + + call := stage.Call{ + Request: map[string]any{"boundary": "input", "protocol": source}, + Metadata: stage.CallMetadata{ + RequestID: "matrix-request", + Attempt: 1, + }, + } + if streaming { + stream, streamErr := topology.Stream(context.Background(), call) + require.NoError(t, streamErr) + for { + _, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + require.NoError(t, nextErr) + } + require.NoError(t, stream.Close()) + require.NoError(t, recorder.SetFinalResponse(source, map[string]any{ + "boundary": "final", + "protocol": source, + "stream": true, + })) + } else { + response, completeErr := topology.Complete(context.Background(), call) + require.NoError(t, completeErr) + require.NoError(t, recorder.SetFinalResponse(source, response.Value)) + } + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Equal(t, source, completed.InputRequest.Protocol) + require.Equal(t, source, completed.FinalResponse.Protocol) + require.Len(t, completed.ProviderExchanges, 1) + exchange := completed.ProviderExchanges[0] + require.Equal(t, target, exchange.Protocol) + require.Equal(t, target, exchange.Request.Protocol) + require.NotNil(t, exchange.Response) + require.Equal(t, target, exchange.Response.Protocol) +} + +type matrixProviderEndpoint struct { + api protocol.APIType +} + +func (e *matrixProviderEndpoint) Protocol() protocol.APIType { return e.api } + +func (e *matrixProviderEndpoint) Complete(_ context.Context, call stage.Call) (*stage.Response, error) { + return &stage.Response{Value: map[string]any{ + "boundary": "provider_response", + "protocol": e.api, + "request": call.Request, + }}, nil +} + +func (e *matrixProviderEndpoint) Stream(context.Context, stage.Call) (stage.EventStream, error) { + return &recordingTestStream{events: matrixStreamEvents(e.api)}, nil +} + +type matrixBridge struct { + source protocol.APIType + target protocol.APIType +} + +func (b matrixBridge) Source() protocol.APIType { return b.source } + +func (b matrixBridge) Target() protocol.APIType { return b.target } + +func (matrixBridge) Capabilities() stage.Capabilities { return stage.AllBridgeCapabilities } + +func (b matrixBridge) Open(_ context.Context, call stage.Call, _ stage.Operation) (stage.BridgeSession, error) { + targetCall := call + targetCall.Request = map[string]any{ + "boundary": "provider_request", + "protocol": b.target, + "source_protocol": b.source, + } + return &matrixBridgeSession{source: b.source, targetCall: targetCall}, nil +} + +type matrixBridgeSession struct { + source protocol.APIType + targetCall stage.Call +} + +func (s *matrixBridgeSession) TargetCall() stage.Call { return s.targetCall } + +func (s *matrixBridgeSession) ConvertComplete(_ context.Context, response *stage.Response) (*stage.Response, error) { + result := *response + result.Value = map[string]any{ + "boundary": "final", + "protocol": s.source, + "provider_response": response.Value, + } + return &result, nil +} + +func (*matrixBridgeSession) ConvertStream(_ context.Context, stream stage.EventStream) (stage.EventStream, error) { + return stream, nil +} + +func (*matrixBridgeSession) ConvertError(_ context.Context, err error) error { return err } + +func matrixStreamEvents(api protocol.APIType) []stage.Event { + switch api { + case protocol.TypeAnthropicV1: + return []stage.Event{ + {Value: json.RawMessage(`{"type":"message_start","message":{"id":"msg-v1","type":"message","role":"assistant","content":[],"model":"provider-model","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}`)}, + {Value: json.RawMessage(`{"type":"message_stop"}`)}, + } + case protocol.TypeAnthropicBeta: + return []stage.Event{ + {Value: json.RawMessage(`{"type":"message_start","message":{"id":"msg-beta","type":"message","role":"assistant","content":[],"model":"provider-model","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}`)}, + {Value: json.RawMessage(`{"type":"message_stop"}`)}, + } + case protocol.TypeOpenAIChat: + stop := "stop" + return []stage.Event{{Value: wire.ChatStreamChunk{ + ID: "chat-stream", + Object: "chat.completion.chunk", + Model: "provider-model", + Choices: []wire.ChatStreamChoice{{ + Index: 0, + Delta: wire.ChatStreamDelta{Role: "assistant", Content: "hello"}, + FinishReason: &stop, + }}, + }}} + case protocol.TypeOpenAIResponses: + return []stage.Event{{Value: wire.ResponsesCompletedEvent{ + Type: "response.completed", + Response: wire.ResponsesWireResponse{ + ID: "resp-stream", + Object: "response", + Status: "completed", + Model: "provider-model", + Output: []wire.ResponsesOutputItemWire{}, + }, + }}} + default: + panic(fmt.Sprintf("unsupported matrix protocol %q", api)) + } +} diff --git a/internal/record/provider_endpoint.go b/internal/record/provider_endpoint.go new file mode 100644 index 000000000..698f45969 --- /dev/null +++ b/internal/record/provider_endpoint.go @@ -0,0 +1,171 @@ +package record + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/assembler" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" +) + +var errProviderStreamClosedBeforeTerminal = errors.New("provider stream closed before a terminal event") + +// ObserveProvider wraps the terminal provider Endpoint when recording is +// enabled. A nil Recorder returns next unchanged, keeping the default path free +// of wrapper allocation and stream assembly. +func ObserveProvider(next stage.Endpoint, recorder *Recorder, meta ExchangeMetadata) stage.Endpoint { + if recorder == nil { + return next + } + meta.Protocol = next.Protocol() + return &providerEndpoint{ + next: next, + recorder: recorder, + meta: meta, + } +} + +type providerEndpoint struct { + next stage.Endpoint + recorder *Recorder + meta ExchangeMetadata +} + +func (e *providerEndpoint) Protocol() protocol.APIType { + return e.next.Protocol() +} + +func (e *providerEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + exchange, _ := e.recorder.BeginExchange(e.meta, call.Request) + response, callErr := e.next.Complete(ctx, call) + if exchange != nil { + var value any + if response != nil { + value = response.Value + } + finishExchange(exchange, value, callErr) + } + return response, callErr +} + +func (e *providerEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + exchange, _ := e.recorder.BeginExchange(e.meta, call.Request) + stream, callErr := e.next.Stream(ctx, call) + if callErr != nil { + finishExchange(exchange, nil, callErr) + return nil, callErr + } + if exchange == nil || stream == nil { + return stream, nil + } + + streamAssembler, _ := assembler.NewStreamAssemblerForRequest(e.meta.Protocol, call.Request) + return &providerStream{ + next: stream, + exchange: exchange, + assembler: streamAssembler, + callCtx: ctx, + }, nil +} + +type providerStream struct { + next stage.EventStream + exchange *Exchange + assembler assembler.StreamAssembler + callCtx context.Context + assemblyErr error + + finishOnce sync.Once +} + +func (s *providerStream) Next(ctx context.Context) (stage.Event, error) { + event, err := s.next.Next(ctx) + if err == nil { + if s.assembler != nil { + if assembleErr := s.assembler.Add(event.Value); assembleErr != nil { + // Recording is observational. Stop assembling this response without + // changing the provider stream seen by the caller. + s.assemblyErr = fmt.Errorf("assemble provider stream event: %w", assembleErr) + s.assembler = nil + } + } + return event, nil + } + + if errors.Is(err, io.EOF) { + if s.assemblyErr != nil { + s.finish(s.assemblyErr) + } else if s.assembler != nil && !s.assembler.Terminal() { + s.finish(errProviderStreamClosedBeforeTerminal) + } else { + s.finish(s.terminalError()) + } + } else { + s.finish(err) + } + return event, err +} + +func (s *providerStream) Close() error { + closeErr := s.next.Close() + if closeErr != nil { + s.finish(closeErr) + } else if s.assemblyErr != nil { + s.finish(s.assemblyErr) + } else if s.assembler != nil && !s.assembler.Terminal() { + if s.callCtx != nil && s.callCtx.Err() != nil { + s.finish(s.callCtx.Err()) + } else { + s.finish(errProviderStreamClosedBeforeTerminal) + } + } else { + // A successful outer Stage/Bridge may stop after the provider's terminal + // event without pulling one additional EOF from this inner stream. The + // request driver still closes the chain normally, so preserve the + // assembled provider response as a successful exchange. Cancellation and + // other early termination already reach Next and win finishOnce first. + s.finish(s.terminalError()) + } + return closeErr +} + +func (s *providerStream) Result() stage.StreamResult { + return s.next.Result() +} + +func (s *providerStream) terminalError() error { + if s.assembler == nil { + return nil + } + return s.assembler.TerminalError() +} + +func (s *providerStream) finish(streamErr error) { + s.finishOnce.Do(func() { + var response any + if s.assembler != nil && s.assembler.Terminal() { + var err error + response, err = s.assembler.Finish() + if err != nil { + streamErr = errors.Join(streamErr, fmt.Errorf("assemble provider stream response: %w", err)) + response = nil + } + } + finishExchange(s.exchange, response, streamErr) + }) +} + +func finishExchange(exchange *Exchange, response any, callErr error) { + if exchange == nil { + return + } + if err := exchange.Finish(response, callErr); err != nil && response != nil { + // A capture/serialization failure must not leave the exchange pending or + // affect request execution. Preserve its outcome without the response. + _ = exchange.Finish(nil, callErr) + } +} diff --git a/internal/record/provider_endpoint_test.go b/internal/record/provider_endpoint_test.go new file mode 100644 index 000000000..7800b5323 --- /dev/null +++ b/internal/record/provider_endpoint_test.go @@ -0,0 +1,327 @@ +package record + +import ( + "context" + "encoding/json" + "errors" + "io" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" +) + +func TestObserveProviderDisabledReturnsOriginalEndpoint(t *testing.T) { + endpoint := &recordingTestEndpoint{api: protocol.TypeOpenAIChat} + recorder, err := New(Config{Enabled: false}) + require.NoError(t, err) + require.Same(t, endpoint, ObserveProvider(endpoint, recorder, ExchangeMetadata{})) +} + +func TestObserveProviderCompleteCapturesTerminalBoundaries(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeOpenAIChat) + request := map[string]any{"model": "provider-model", "messages": []any{}} + providerResponse := map[string]any{"id": "provider-response", "model": "provider-model"} + endpoint := &recordingTestEndpoint{ + api: protocol.TypeOpenAIChat, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + require.Equal(t, request, call.Request) + return &stage.Response{Value: providerResponse, Model: "provider-model"}, nil + }, + } + wrapper := ObserveProvider(endpoint, recorder, ExchangeMetadata{ + Attempt: 2, + Provider: "provider", + Model: "provider-model", + }) + + response, err := wrapper.Complete(context.Background(), stage.Call{Request: request}) + require.NoError(t, err) + require.Equal(t, providerResponse, response.Value) + require.NoError(t, recorder.SetFinalResponse(protocol.TypeOpenAIChat, map[string]any{"id": "final"})) + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + exchange := completed.ProviderExchanges[0] + require.Equal(t, protocol.TypeOpenAIChat, exchange.Protocol) + require.Equal(t, 2, exchange.Attempt) + require.Equal(t, OutcomeSucceeded, exchange.Outcome) + require.JSONEq(t, `{"model":"provider-model","messages":[]}`, string(exchange.Request.Body)) + require.JSONEq(t, `{"id":"provider-response","model":"provider-model"}`, string(exchange.Response.Body)) +} + +func TestObserveProviderCaptureFailureDoesNotAffectProviderCall(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeOpenAIChat) + providerResponse := map[string]any{"id": "provider-response"} + providerCalled := false + endpoint := &recordingTestEndpoint{ + api: protocol.TypeOpenAIChat, + complete: func(_ context.Context, call stage.Call) (*stage.Response, error) { + providerCalled = true + require.NotNil(t, call.Request) + return &stage.Response{Value: providerResponse}, nil + }, + } + wrapper := ObserveProvider(endpoint, recorder, ExchangeMetadata{Attempt: 1}) + + response, err := wrapper.Complete(context.Background(), stage.Call{ + Request: map[string]any{"cannot_marshal": func() {}}, + }) + + require.NoError(t, err) + require.True(t, providerCalled) + require.Equal(t, providerResponse, response.Value) + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Empty(t, completed.ProviderExchanges) +} + +func TestObserveProviderPreservesProviderError(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeAnthropicBeta) + providerErr := errors.New("provider failed") + endpoint := &recordingTestEndpoint{ + api: protocol.TypeAnthropicBeta, + complete: func(context.Context, stage.Call) (*stage.Response, error) { + return nil, providerErr + }, + } + wrapper := ObserveProvider(endpoint, recorder, ExchangeMetadata{Attempt: 1}) + + response, err := wrapper.Complete(context.Background(), stage.Call{Request: map[string]any{"model": "m"}}) + require.Nil(t, response) + require.ErrorIs(t, err, providerErr) + + completed, first := recorder.Finish(providerErr) + require.True(t, first) + require.Equal(t, OutcomeFailed, completed.Outcome) + require.Len(t, completed.ProviderExchanges, 1) + require.Equal(t, OutcomeFailed, completed.ProviderExchanges[0].Outcome) + require.Equal(t, providerErr.Error(), completed.ProviderExchanges[0].Error) +} + +func TestObserveProviderStreamAssemblesRawProviderResponse(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeAnthropicBeta) + providerStream := &recordingTestStream{events: []stage.Event{ + {Value: json.RawMessage(`{"type":"message_start","message":{"id":"msg-stream","type":"message","role":"assistant","content":[],"model":"provider-model","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}`)}, + {Value: json.RawMessage(`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`)}, + {Value: json.RawMessage(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}`)}, + {Value: json.RawMessage(`{"type":"content_block_stop","index":0}`)}, + {Value: json.RawMessage(`{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":1}}`)}, + {Value: json.RawMessage(`{"type":"message_stop"}`)}, + }} + endpoint := &recordingTestEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return providerStream, nil + }, + } + wrapper := ObserveProvider(endpoint, recorder, ExchangeMetadata{Attempt: 1}) + stream, err := wrapper.Stream(context.Background(), stage.Call{Request: map[string]any{"model": "provider-model"}}) + require.NoError(t, err) + + for { + _, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + require.NoError(t, nextErr) + } + require.NoError(t, stream.Close()) + require.Equal(t, 1, providerStream.closeCount) + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + exchange := completed.ProviderExchanges[0] + require.Equal(t, OutcomeSucceeded, exchange.Outcome) + require.NotNil(t, exchange.Response) + require.Contains(t, string(exchange.Response.Body), "msg-stream") + require.Contains(t, string(exchange.Response.Body), "hello") +} + +func TestObserveProviderStreamCleanCloseAfterTerminalEventSucceeds(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeAnthropicBeta) + providerStream := &recordingTestStream{events: []stage.Event{ + {Value: json.RawMessage(`{"type":"message_start","message":{"id":"msg-stream","type":"message","role":"assistant","content":[],"model":"provider-model","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}`)}, + {Value: json.RawMessage(`{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":1}}`)}, + {Value: json.RawMessage(`{"type":"message_stop"}`)}, + }} + endpoint := &recordingTestEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return providerStream, nil + }, + } + stream, err := ObserveProvider(endpoint, recorder, ExchangeMetadata{Attempt: 2}).Stream( + context.Background(), + stage.Call{Request: map[string]any{"model": "provider-model"}}, + ) + require.NoError(t, err) + for range providerStream.events { + _, err = stream.Next(context.Background()) + require.NoError(t, err) + } + require.NoError(t, stream.Close()) + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + require.Equal(t, 2, completed.ProviderExchanges[0].Attempt) + require.Equal(t, OutcomeSucceeded, completed.ProviderExchanges[0].Outcome) + require.NotNil(t, completed.ProviderExchanges[0].Response) +} + +func TestObserveProviderStreamEarlyCloseFailsWithoutPartialResponse(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeAnthropicBeta) + providerStream := &recordingTestStream{events: []stage.Event{ + {Value: json.RawMessage(`{"type":"message_start","message":{"id":"partial","type":"message","role":"assistant","content":[],"model":"provider-model","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}}`)}, + }} + endpoint := &recordingTestEndpoint{ + api: protocol.TypeAnthropicBeta, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return providerStream, nil + }, + } + stream, err := ObserveProvider(endpoint, recorder, ExchangeMetadata{Attempt: 1}).Stream( + context.Background(), + stage.Call{Request: map[string]any{"model": "provider-model"}}, + ) + require.NoError(t, err) + _, err = stream.Next(context.Background()) + require.NoError(t, err) + require.NoError(t, stream.Close()) + + completed, first := recorder.Finish(errors.New("client stopped")) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + require.Equal(t, OutcomeFailed, completed.ProviderExchanges[0].Outcome) + require.ErrorContains(t, errors.New(completed.ProviderExchanges[0].Error), "before a terminal event") + require.Nil(t, completed.ProviderExchanges[0].Response) +} + +func TestObserveProviderResponsesFailedPreservesResponseAndFailure(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeOpenAIResponses) + providerStream := &recordingTestStream{events: []stage.Event{ + {Value: json.RawMessage(`{"type":"response.failed","sequence_number":1,"response":{"id":"resp-failed","object":"response","status":"failed","model":"provider-model","output":[],"error":{"code":"server_error","message":"provider failed"}}}`)}, + }} + endpoint := &recordingTestEndpoint{ + api: protocol.TypeOpenAIResponses, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return providerStream, nil + }, + } + stream, err := ObserveProvider(endpoint, recorder, ExchangeMetadata{Attempt: 1}).Stream( + context.Background(), + stage.Call{Request: map[string]any{"model": "provider-model"}}, + ) + require.NoError(t, err) + _, err = stream.Next(context.Background()) + require.NoError(t, err) + _, err = stream.Next(context.Background()) + require.ErrorIs(t, err, io.EOF) + require.NoError(t, stream.Close()) + + completed, first := recorder.Finish(errors.New("provider failed")) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + exchange := completed.ProviderExchanges[0] + require.Equal(t, OutcomeFailed, exchange.Outcome) + require.ErrorContains(t, errors.New(exchange.Error), "response.failed") + require.NotNil(t, exchange.Response) + require.Contains(t, string(exchange.Response.Body), `"id":"resp-failed"`) + require.Contains(t, string(exchange.Response.Body), `"status":"failed"`) +} + +func TestObserveProviderStreamAssemblyErrorDoesNotBecomeSuccess(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeOpenAIChat) + stop := "stop" + providerStream := &recordingTestStream{events: []stage.Event{ + {Value: map[string]any{ + "id": "chat-a", "object": "chat.completion.chunk", + "choices": []map[string]any{{"index": 0, "delta": map[string]any{"content": "partial"}}}, + }}, + {Value: map[string]any{ + "id": "chat-b", "object": "chat.completion.chunk", + "choices": []map[string]any{{"index": 0, "delta": map[string]any{}, "finish_reason": stop}}, + }}, + }} + endpoint := &recordingTestEndpoint{ + api: protocol.TypeOpenAIChat, + stream: func(context.Context, stage.Call) (stage.EventStream, error) { + return providerStream, nil + }, + } + stream, err := ObserveProvider(endpoint, recorder, ExchangeMetadata{Attempt: 1}).Stream( + context.Background(), + stage.Call{Request: map[string]any{"model": "provider-model"}}, + ) + require.NoError(t, err) + for { + _, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + require.NoError(t, nextErr) + } + require.NoError(t, stream.Close()) + + completed, first := recorder.Finish(errors.New("recording assembly failed")) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + exchange := completed.ProviderExchanges[0] + require.Equal(t, OutcomeFailed, exchange.Outcome) + require.Contains(t, exchange.Error, "accumulate OpenAI Chat") + require.Nil(t, exchange.Response) +} + +func newRecordingTestRecorder(t *testing.T, api protocol.APIType) *Recorder { + t.Helper() + recorder, err := New(Config{ + Enabled: true, + RequestID: "request-id", + InputProtocol: api, + Input: map[string]any{"model": "client-model"}, + }) + require.NoError(t, err) + return recorder +} + +type recordingTestEndpoint struct { + api protocol.APIType + complete func(context.Context, stage.Call) (*stage.Response, error) + stream func(context.Context, stage.Call) (stage.EventStream, error) +} + +func (e *recordingTestEndpoint) Protocol() protocol.APIType { return e.api } + +func (e *recordingTestEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + return e.complete(ctx, call) +} + +func (e *recordingTestEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + return e.stream(ctx, call) +} + +type recordingTestStream struct { + events []stage.Event + index int + closeCount int +} + +func (s *recordingTestStream) Next(context.Context) (stage.Event, error) { + if s.index >= len(s.events) { + return stage.Event{}, io.EOF + } + event := s.events[s.index] + s.index++ + return event, nil +} + +func (s *recordingTestStream) Close() error { + s.closeCount++ + return nil +} + +func (*recordingTestStream) Result() stage.StreamResult { return stage.StreamResult{} } diff --git a/internal/record/record.go b/internal/record/record.go new file mode 100644 index 000000000..efd124437 --- /dev/null +++ b/internal/record/record.go @@ -0,0 +1,114 @@ +// Package record defines request-scoped protocol recording primitives. +// +// It is intentionally independent of HTTP, Gin, routing, usage accounting, +// and persistence. Callers create a Recorder only when recording is explicitly +// enabled; the disabled path returns nil and performs no payload capture. +package record + +import ( + "encoding/json" + "time" + + "github.com/tingly-dev/tingly-box/internal/protocol" +) + +// Outcome is the terminal result of a request or provider exchange. +type Outcome string + +const ( + OutcomePending Outcome = "pending" + OutcomeSucceeded Outcome = "succeeded" + OutcomeFailed Outcome = "failed" + OutcomeCancelled Outcome = "cancelled" +) + +// Payload is one complete value at a stable protocol boundary. +type Payload struct { + Protocol protocol.APIType `json:"protocol"` + ContentType string `json:"content_type"` + Body json.RawMessage `json:"body"` +} + +// RequestRecord is the completed record for one incoming client request. +type RequestRecord struct { + Timestamp time.Time `json:"timestamp"` + RequestID string `json:"request_id"` + SessionID string `json:"session_id,omitempty"` + Scenario string `json:"scenario,omitempty"` + InputRequest Payload `json:"input_request"` + ProviderExchanges []ProviderExchange `json:"provider_exchanges,omitempty"` + FinalResponse *Payload `json:"final_response,omitempty"` + Outcome Outcome `json:"outcome"` + Error string `json:"error,omitempty"` + Duration time.Duration `json:"duration"` +} + +// ProviderExchange records one actual invocation of a provider endpoint. +// Sequence is one-based and reflects invocation order within RequestRecord. +type ProviderExchange struct { + Sequence int `json:"sequence"` + Attempt int `json:"attempt"` + Provider string `json:"provider,omitempty"` + Model string `json:"model,omitempty"` + Protocol protocol.APIType `json:"protocol"` + Request Payload `json:"provider_request"` + Response *Payload `json:"provider_response,omitempty"` + Outcome Outcome `json:"outcome"` + Error string `json:"error,omitempty"` + StartedAt time.Time `json:"started_at"` + Duration time.Duration `json:"duration"` +} + +// ExchangeMetadata identifies one provider endpoint invocation. +type ExchangeMetadata struct { + Attempt int + Provider string + Model string + Protocol protocol.APIType +} + +func clonePayload(src Payload) Payload { + dst := src + dst.Body = append(json.RawMessage(nil), src.Body...) + return dst +} + +func clonePayloadPointer(src *Payload) *Payload { + if src == nil { + return nil + } + dst := clonePayload(*src) + return &dst +} + +func cloneRequestRecord(src RequestRecord) RequestRecord { + dst := src + dst.InputRequest = clonePayload(src.InputRequest) + dst.FinalResponse = clonePayloadPointer(src.FinalResponse) + dst.ProviderExchanges = make([]ProviderExchange, len(src.ProviderExchanges)) + for i := range src.ProviderExchanges { + dst.ProviderExchanges[i] = src.ProviderExchanges[i] + dst.ProviderExchanges[i].Request = clonePayload(src.ProviderExchanges[i].Request) + dst.ProviderExchanges[i].Response = clonePayloadPointer(src.ProviderExchanges[i].Response) + } + return dst +} + +// Project returns an immutable copy suitable for a recording mode that may +// omit response payloads. Provider requests and the original input remain part +// of every RequestRecord; response inclusion is controlled independently. +func (r *RequestRecord) Project(includeProviderResponses, includeFinalResponse bool) *RequestRecord { + if r == nil { + return nil + } + projected := cloneRequestRecord(*r) + if !includeProviderResponses { + for index := range projected.ProviderExchanges { + projected.ProviderExchanges[index].Response = nil + } + } + if !includeFinalResponse { + projected.FinalResponse = nil + } + return &projected +} diff --git a/internal/record/recorder.go b/internal/record/recorder.go new file mode 100644 index 000000000..b5d6f6abe --- /dev/null +++ b/internal/record/recorder.go @@ -0,0 +1,249 @@ +package record + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "time" + + "github.com/tingly-dev/tingly-box/internal/protocol" +) + +const jsonContentType = "application/json" + +var ErrFinished = errors.New("request recorder is already finished") + +// Config creates one request-scoped Recorder. Enabled must be explicitly true; +// the zero value keeps recording disabled. +type Config struct { + Enabled bool + RequestID string + SessionID string + Scenario string + InputProtocol protocol.APIType + Input any +} + +// Recorder incrementally builds one RequestRecord. +type Recorder struct { + mu sync.Mutex + startedAt time.Time + record RequestRecord + finished *RequestRecord +} + +// Exchange is the one-shot completion handle returned by BeginExchange. +type Exchange struct { + recorder *Recorder + index int + done bool +} + +// New creates a Recorder only when cfg.Enabled is true. The disabled path +// returns before validation or JSON serialization. +func New(cfg Config) (*Recorder, error) { + if !cfg.Enabled { + return nil, nil + } + + input, err := capturePayload(cfg.InputProtocol, cfg.Input) + if err != nil { + return nil, fmt.Errorf("capture input request: %w", err) + } + + now := time.Now().UTC() + return &Recorder{ + startedAt: now, + record: RequestRecord{ + Timestamp: now, + RequestID: cfg.RequestID, + SessionID: cfg.SessionID, + Scenario: cfg.Scenario, + InputRequest: input, + Outcome: OutcomePending, + }, + }, nil +} + +// Enabled reports whether a non-nil recorder is active. +func (r *Recorder) Enabled() bool { + return r != nil +} + +// BeginExchange captures the provider-bound request and appends an ordered +// provider exchange. It is a no-op for a nil (disabled) Recorder. +func (r *Recorder) BeginExchange(meta ExchangeMetadata, request any) (*Exchange, error) { + if r == nil { + return nil, nil + } + if r.isFinished() { + return nil, ErrFinished + } + + payload, err := capturePayload(meta.Protocol, request) + if err != nil { + return nil, fmt.Errorf("capture provider request: %w", err) + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.finished != nil { + return nil, ErrFinished + } + + now := time.Now().UTC() + r.record.ProviderExchanges = append(r.record.ProviderExchanges, ProviderExchange{ + Sequence: len(r.record.ProviderExchanges) + 1, + Attempt: meta.Attempt, + Provider: meta.Provider, + Model: meta.Model, + Protocol: meta.Protocol, + Request: payload, + Outcome: OutcomePending, + StartedAt: now, + }) + return &Exchange{recorder: r, index: len(r.record.ProviderExchanges) - 1}, nil +} + +// Finish completes one provider exchange. Repeated calls are idempotent. +func (e *Exchange) Finish(response any, callErr error) error { + if e == nil || e.recorder == nil { + return nil + } + r := e.recorder + r.mu.Lock() + if e.done { + r.mu.Unlock() + return nil + } + if r.finished != nil { + r.mu.Unlock() + return ErrFinished + } + if e.index < 0 || e.index >= len(r.record.ProviderExchanges) { + r.mu.Unlock() + return fmt.Errorf("provider exchange index %d is invalid", e.index) + } + api := r.record.ProviderExchanges[e.index].Protocol + r.mu.Unlock() + + var responsePayload *Payload + if response != nil { + payload, err := capturePayload(api, response) + if err != nil { + return fmt.Errorf("capture provider response: %w", err) + } + responsePayload = &payload + } + + r.mu.Lock() + defer r.mu.Unlock() + if e.done { + return nil + } + if r.finished != nil { + return ErrFinished + } + if e.index < 0 || e.index >= len(r.record.ProviderExchanges) { + return fmt.Errorf("provider exchange index %d is invalid", e.index) + } + + exchange := &r.record.ProviderExchanges[e.index] + exchange.Response = responsePayload + exchange.Outcome = outcomeForError(callErr) + if callErr != nil { + exchange.Error = callErr.Error() + } + exchange.Duration = time.Since(exchange.StartedAt) + e.done = true + return nil +} + +// SetFinalResponse captures the client-visible response after every outward +// transformation. It is a no-op for a nil (disabled) Recorder. +func (r *Recorder) SetFinalResponse(api protocol.APIType, response any) error { + if r == nil { + return nil + } + if r.isFinished() { + return ErrFinished + } + payload, err := capturePayload(api, response) + if err != nil { + return fmt.Errorf("capture final response: %w", err) + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.finished != nil { + return ErrFinished + } + r.record.FinalResponse = &payload + return nil +} + +func (r *Recorder) isFinished() bool { + if r == nil { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + return r.finished != nil +} + +// Finish completes the request exactly once. The returned boolean is true only +// for the caller that performed the transition; later callers receive an +// immutable copy and false. +func (r *Recorder) Finish(requestErr error) (*RequestRecord, bool) { + if r == nil { + return nil, false + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.finished != nil { + copy := cloneRequestRecord(*r.finished) + return ©, false + } + + r.record.Outcome = outcomeForError(requestErr) + if requestErr != nil { + r.record.Error = requestErr.Error() + } + r.record.Duration = time.Since(r.startedAt) + completed := cloneRequestRecord(r.record) + r.finished = &completed + copy := cloneRequestRecord(completed) + return ©, true +} + +func capturePayload(api protocol.APIType, value any) (Payload, error) { + if api == "" { + return Payload{}, errors.New("protocol is empty") + } + body, err := payloadJSON(value) + if err != nil { + return Payload{}, err + } + return Payload{ + Protocol: api, + ContentType: jsonContentType, + Body: append(json.RawMessage(nil), body...), + }, nil +} + +func payloadJSON(value any) ([]byte, error) { + return protocol.SnapshotJSON(value) +} + +func outcomeForError(err error) Outcome { + if err == nil { + return OutcomeSucceeded + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return OutcomeCancelled + } + return OutcomeFailed +} diff --git a/internal/record/recorder_test.go b/internal/record/recorder_test.go new file mode 100644 index 000000000..9adebcf18 --- /dev/null +++ b/internal/record/recorder_test.go @@ -0,0 +1,208 @@ +package record + +import ( + "context" + "encoding/json" + "errors" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tingly-dev/tingly-box/internal/protocol" +) + +type countingMarshaler struct { + calls *int +} + +type rawJSONValue struct { + raw string +} + +func (v rawJSONValue) RawJSON() string { return v.raw } + +func (rawJSONValue) MarshalJSON() ([]byte, error) { + return []byte(`{"expanded_zero_value":""}`), nil +} + +func (m countingMarshaler) MarshalJSON() ([]byte, error) { + (*m.calls)++ + return []byte(`{"unexpected":true}`), nil +} + +func TestDisabledRecorderDoesNoCaptureWork(t *testing.T) { + calls := 0 + recorder, err := New(Config{ + Enabled: false, + Input: countingMarshaler{calls: &calls}, + }) + require.NoError(t, err) + require.Nil(t, recorder) + require.Zero(t, calls) + + exchange, err := recorder.BeginExchange(ExchangeMetadata{}, countingMarshaler{calls: &calls}) + require.NoError(t, err) + require.Nil(t, exchange) + require.NoError(t, recorder.SetFinalResponse("", countingMarshaler{calls: &calls})) + require.NoError(t, exchange.Finish(countingMarshaler{calls: &calls}, nil)) + require.Zero(t, calls) + + completed, first := recorder.Finish(nil) + require.Nil(t, completed) + require.False(t, first) +} + +func TestRecorderCapturesOrderedProviderExchanges(t *testing.T) { + recorder, err := New(Config{ + Enabled: true, + RequestID: "req-1", + SessionID: "session-1", + Scenario: "chat", + InputProtocol: protocol.TypeOpenAIChat, + Input: map[string]any{"model": "public", "messages": []any{}}, + }) + require.NoError(t, err) + require.True(t, recorder.Enabled()) + + first, err := recorder.BeginExchange(ExchangeMetadata{ + Attempt: 1, + Provider: "primary", + Model: "provider-model-a", + Protocol: protocol.TypeAnthropicBeta, + }, map[string]any{"model": "provider-model-a"}) + require.NoError(t, err) + require.NoError(t, first.Finish(nil, errors.New("provider unavailable"))) + + second, err := recorder.BeginExchange(ExchangeMetadata{ + Attempt: 2, + Provider: "fallback", + Model: "provider-model-b", + Protocol: protocol.TypeOpenAIChat, + }, map[string]any{"model": "provider-model-b"}) + require.NoError(t, err) + require.NoError(t, second.Finish(map[string]any{"id": "provider-response"}, nil)) + require.NoError(t, recorder.SetFinalResponse( + protocol.TypeOpenAIChat, + map[string]any{"id": "client-response", "model": "public"}, + )) + + completed, firstFinish := recorder.Finish(nil) + require.True(t, firstFinish) + require.Equal(t, "req-1", completed.RequestID) + require.Equal(t, OutcomeSucceeded, completed.Outcome) + require.Len(t, completed.ProviderExchanges, 2) + require.Equal(t, 1, completed.ProviderExchanges[0].Sequence) + require.Equal(t, 1, completed.ProviderExchanges[0].Attempt) + require.Equal(t, OutcomeFailed, completed.ProviderExchanges[0].Outcome) + require.Equal(t, "provider unavailable", completed.ProviderExchanges[0].Error) + require.Nil(t, completed.ProviderExchanges[0].Response) + require.Equal(t, 2, completed.ProviderExchanges[1].Sequence) + require.Equal(t, 2, completed.ProviderExchanges[1].Attempt) + require.Equal(t, OutcomeSucceeded, completed.ProviderExchanges[1].Outcome) + require.Equal(t, protocol.TypeOpenAIChat, completed.FinalResponse.Protocol) + require.JSONEq(t, `{"id":"client-response","model":"public"}`, string(completed.FinalResponse.Body)) + + // Finish is idempotent and returns defensive copies. + completed.InputRequest.Body[0] = 'x' + again, secondFinish := recorder.Finish(nil) + require.False(t, secondFinish) + require.True(t, json.Valid(again.InputRequest.Body)) + require.JSONEq(t, `{"messages":[],"model":"public"}`, string(again.InputRequest.Body)) +} + +func TestRecorderKeepsToolLoopRoundsInOneAttempt(t *testing.T) { + recorder, err := New(Config{ + Enabled: true, + InputProtocol: protocol.TypeAnthropicBeta, + Input: map[string]any{"model": "claude"}, + }) + require.NoError(t, err) + + for round := 1; round <= 3; round++ { + exchange, beginErr := recorder.BeginExchange(ExchangeMetadata{ + Attempt: 1, + Provider: "provider", + Model: "claude", + Protocol: protocol.TypeAnthropicBeta, + }, map[string]any{"round": round}) + require.NoError(t, beginErr) + require.NoError(t, exchange.Finish(map[string]any{"round": round}, nil)) + } + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 3) + for i, exchange := range completed.ProviderExchanges { + require.Equal(t, i+1, exchange.Sequence) + require.Equal(t, 1, exchange.Attempt) + } +} + +func TestRecorderPrefersProtocolRawJSON(t *testing.T) { + recorder, err := New(Config{ + Enabled: true, + InputProtocol: protocol.TypeAnthropicBeta, + Input: map[string]any{"model": "client"}, + }) + require.NoError(t, err) + exchange, err := recorder.BeginExchange(ExchangeMetadata{ + Protocol: protocol.TypeAnthropicBeta, + }, map[string]any{"model": "provider"}) + require.NoError(t, err) + require.NoError(t, exchange.Finish(rawJSONValue{raw: `{"model":"provider","unknown":"preserved"}`}, nil)) + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.JSONEq(t, + `{"model":"provider","unknown":"preserved"}`, + string(completed.ProviderExchanges[0].Response.Body), + ) +} + +func TestRecorderMapsCancellationAndRejectsMutationAfterFinish(t *testing.T) { + recorder, err := New(Config{ + Enabled: true, + InputProtocol: protocol.TypeOpenAIResponses, + Input: map[string]any{"model": "gpt"}, + }) + require.NoError(t, err) + + completed, first := recorder.Finish(context.Canceled) + require.True(t, first) + require.Equal(t, OutcomeCancelled, completed.Outcome) + require.ErrorIs(t, recorder.SetFinalResponse(protocol.TypeOpenAIResponses, map[string]any{}), ErrFinished) + _, err = recorder.BeginExchange(ExchangeMetadata{Protocol: protocol.TypeOpenAIResponses}, map[string]any{}) + require.ErrorIs(t, err, ErrFinished) +} + +func TestEnabledRecorderValidatesProtocolAndJSON(t *testing.T) { + _, err := New(Config{Enabled: true, Input: map[string]any{}}) + require.ErrorContains(t, err, "protocol is empty") + + _, err = New(Config{ + Enabled: true, + InputProtocol: protocol.TypeOpenAIChat, + Input: func() {}, + }) + require.ErrorContains(t, err, "unsupported type") +} + +func TestRecorderRejectsTypedNilWithoutPanicking(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeOpenAIChat) + var response *rawJSONValue + require.NotPanics(t, func() { + require.ErrorContains(t, recorder.SetFinalResponse(protocol.TypeOpenAIChat, response), "nil") + }) +} + +func TestRecorderDoesNotSerializeAfterFinish(t *testing.T) { + recorder := newRecordingTestRecorder(t, protocol.TypeOpenAIChat) + _, first := recorder.Finish(nil) + require.True(t, first) + + calls := 0 + value := countingMarshaler{calls: &calls} + _, err := recorder.BeginExchange(ExchangeMetadata{Protocol: protocol.TypeOpenAIChat}, value) + require.ErrorIs(t, err, ErrFinished) + require.ErrorIs(t, recorder.SetFinalResponse(protocol.TypeOpenAIChat, value), ErrFinished) + require.Zero(t, calls) +} From 055fce085b62d35466a688643cf42356433f8c3a Mon Sep 17 00:00:00 2001 From: FFengIll Date: Tue, 4 Aug 2026 19:47:33 +0800 Subject: [PATCH 3/8] feat(protocolserver,mcpserver): port protocol-stage glue and rewire to new package layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port the server-side stage integration from codex/protocol-stage-hardening, adapting to main's post-#1493..#1495 package extraction (gateway core now lives in protocolserver, MCP engine in mcpserver): - protocol_stage_*.go glue (pipeline/selector/recording/tool_loop/anthropic/openai endpoints): moved into package protocolserver (was package server on hardening), since ProtocolHandler now lives there. Rewired imports: server/{forwarding,recording} -> protocolserver/*; server/transform folded into protocolserver/transform; mcpmodule alias -> internal/mcpserver. - 16 files hardening modified at old internal/server paths (anthropic_message, openai_chat, openai_responses, protocol_{dispatch,handler}, failover_dispatch, guardrails_runtime_ai, mcp_hooks, module/mcp/{continuation_store,format_adapter,generic_loop_processor, generic_stream_interceptor,tool_executor}, servertool/executor): applied via clean 3-way merge (base=#1491) against their new protocolserver/mcpserver homes — 0 conflicts. - anthropic_beta_stage*.go + tool_executor/continuation_store tests: moved to package mcpserver (their referenced symbols Tool/NewAnthropicBetaAdapter live there now); package mcp -> mcpserver. - internal/mcp/runtime: ListServerToolsForAnthropicBetaInjection added by hardening. - Restored the dropped `internal/protocol` import in protocol_handler.go (lost in 3-way merge). Builds clean: go vet ./internal/{server,protocolserver,mcpserver,mcp,protocol,record,obs,guardrails}/... Server-skeleton wiring (server.go/lifecycle/options 3-way merge) follows next. Batch 3 of the protocol-stage-hardening port. --- internal/mcp/runtime/runtime.go | 31 + .../runtime/runtime_beta_injection_test.go | 36 + internal/mcpserver/anthropic_beta_stage.go | 334 +++++++ .../anthropic_beta_stage_record_test.go | 164 ++++ .../mcpserver/anthropic_beta_stage_stream.go | 368 ++++++++ .../mcpserver/anthropic_beta_stage_test.go | 751 ++++++++++++++++ internal/mcpserver/continuation_store.go | 248 +++++- internal/mcpserver/continuation_store_test.go | 88 ++ internal/mcpserver/format_adapter.go | 3 + internal/mcpserver/generic_loop_processor.go | 4 +- .../mcpserver/generic_stream_interceptor.go | 4 +- internal/mcpserver/tool_executor.go | 9 +- internal/mcpserver/tool_executor_test.go | 36 + internal/protocolserver/anthropic_message.go | 97 ++- internal/protocolserver/failover_dispatch.go | 101 ++- .../protocolserver/failover_dispatch_test.go | 83 ++ .../protocolserver/guardrails_runtime_ai.go | 12 +- internal/protocolserver/mcp_hooks.go | 2 +- internal/protocolserver/openai_chat.go | 36 +- internal/protocolserver/openai_responses.go | 37 +- internal/protocolserver/protocol_dispatch.go | 11 +- internal/protocolserver/protocol_handler.go | 35 +- .../protocol_stage_anthropic_beta.go | 521 +++++++++++ ...col_stage_anthropic_beta_recording_test.go | 161 ++++ .../protocol_stage_anthropic_v1.go | 608 +++++++++++++ .../protocol_stage_openai_chat_endpoint.go | 125 +++ .../protocol_stage_openai_responses.go | 627 +++++++++++++ .../protocol_stage_openai_responses_test.go | 116 +++ .../protocolserver/protocol_stage_pipeline.go | 824 ++++++++++++++++++ .../protocol_stage_recording.go | 262 ++++++ .../protocol_stage_recording_test.go | 44 + .../protocolserver/protocol_stage_selector.go | 110 +++ .../protocol_stage_selector_test.go | 73 ++ .../protocol_stage_stream_test.go | 358 ++++++++ .../protocol_stage_tool_loop.go | 50 ++ .../protocolserver/protocol_transform_test.go | 26 + .../protocolserver/servertool/executor.go | 17 + .../servertool/executor_test.go | 41 + .../transform/protocol_stage_beta_tools.go | 72 ++ .../protocol_stage_beta_tools_test.go | 49 ++ 40 files changed, 6511 insertions(+), 63 deletions(-) create mode 100644 internal/mcp/runtime/runtime_beta_injection_test.go create mode 100644 internal/mcpserver/anthropic_beta_stage.go create mode 100644 internal/mcpserver/anthropic_beta_stage_record_test.go create mode 100644 internal/mcpserver/anthropic_beta_stage_stream.go create mode 100644 internal/mcpserver/anthropic_beta_stage_test.go create mode 100644 internal/mcpserver/continuation_store_test.go create mode 100644 internal/mcpserver/tool_executor_test.go create mode 100644 internal/protocolserver/protocol_stage_anthropic_beta.go create mode 100644 internal/protocolserver/protocol_stage_anthropic_beta_recording_test.go create mode 100644 internal/protocolserver/protocol_stage_anthropic_v1.go create mode 100644 internal/protocolserver/protocol_stage_openai_chat_endpoint.go create mode 100644 internal/protocolserver/protocol_stage_openai_responses.go create mode 100644 internal/protocolserver/protocol_stage_openai_responses_test.go create mode 100644 internal/protocolserver/protocol_stage_pipeline.go create mode 100644 internal/protocolserver/protocol_stage_recording.go create mode 100644 internal/protocolserver/protocol_stage_recording_test.go create mode 100644 internal/protocolserver/protocol_stage_selector.go create mode 100644 internal/protocolserver/protocol_stage_selector_test.go create mode 100644 internal/protocolserver/protocol_stage_stream_test.go create mode 100644 internal/protocolserver/protocol_stage_tool_loop.go create mode 100644 internal/protocolserver/servertool/executor_test.go create mode 100644 internal/protocolserver/transform/protocol_stage_beta_tools.go create mode 100644 internal/protocolserver/transform/protocol_stage_beta_tools_test.go diff --git a/internal/mcp/runtime/runtime.go b/internal/mcp/runtime/runtime.go index 07e182bb3..cd4ee6348 100644 --- a/internal/mcp/runtime/runtime.go +++ b/internal/mcp/runtime/runtime.go @@ -8,6 +8,7 @@ import ( "sync" "time" + "github.com/anthropics/anthropic-sdk-go" "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/packages/param" "github.com/openai/openai-go/v3/shared" @@ -160,6 +161,36 @@ func (r *Runtime) ListServerToolsForInjection(ctx context.Context) []openai.Chat return out } +// ListServerToolsForAnthropicBetaInjection returns the same enabled +// server-visible virtual tools directly in the Beta working protocol. The +// ToolLoop Stage can therefore stay Beta-native instead of routing tool +// definitions through an OpenAI DTO first. +func (r *Runtime) ListServerToolsForAnthropicBetaInjection(ctx context.Context) []anthropic.BetaToolUnionParam { + if r == nil || r.virtualRegistry == nil { + return nil + } + virtualTools := r.virtualRegistry.ListVirtualTools() + out := make([]anthropic.BetaToolUnionParam, 0, len(virtualTools)) + for _, vt := range virtualTools { + if !r.isVirtualServerToolInjectable(vt) { + continue + } + schema := anthropic.BetaToolInputSchemaParam{Properties: map[string]any{}} + if schemaBytes, err := json.Marshal(vt.InputSchema); err == nil && string(schemaBytes) != "null" { + var converted anthropic.BetaToolInputSchemaParam + if json.Unmarshal(schemaBytes, &converted) == nil { + schema = converted + } + } + tool := anthropic.BetaToolUnionParamOfTool(schema, NormalizeToolName("builtin", vt.Name)) + if vt.Description != "" { + tool.OfTool.Description = anthropic.Opt(vt.Description) + } + out = append(out, tool) + } + return out +} + func (r *Runtime) isVirtualServerToolInjectable(vt coretool.VirtualTool) bool { if strings.TrimSpace(vt.Name) == "" || !IsServerVisibleVirtualTool(vt) { return false diff --git a/internal/mcp/runtime/runtime_beta_injection_test.go b/internal/mcp/runtime/runtime_beta_injection_test.go new file mode 100644 index 000000000..6eeaa806b --- /dev/null +++ b/internal/mcp/runtime/runtime_beta_injection_test.go @@ -0,0 +1,36 @@ +package runtime + +import ( + "context" + "testing" + + coretool "github.com/tingly-dev/tingly-box/internal/tool" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func TestListServerToolsForAnthropicBetaInjectionUsesBetaTypesDirectly(t *testing.T) { + runtime := NewRuntime(func() *typ.MCPRuntimeConfig { return &typ.MCPRuntimeConfig{} }) + t.Cleanup(runtime.Close) + runtime.VirtualRegistry().Register(coretool.VirtualTool{ + Name: "lookup", + Description: "Look up a value", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{"query": map[string]any{"type": "string"}}, + "required": []string{"query"}, + }, + Visibility: typ.ToolVisibilityServer, + }) + + tools := runtime.ListServerToolsForAnthropicBetaInjection(context.Background()) + if len(tools) != 1 || tools[0].OfTool == nil { + t.Fatalf("Beta tools = %#v", tools) + } + tool := tools[0].OfTool + if tool.Name != NormalizeToolName("builtin", "lookup") || tool.Description.Value != "Look up a value" { + t.Fatalf("Beta tool = %#v", tool) + } + if len(tool.InputSchema.Required) != 1 || tool.InputSchema.Required[0] != "query" { + t.Fatalf("Beta input schema = %#v", tool.InputSchema) + } +} diff --git a/internal/mcpserver/anthropic_beta_stage.go b/internal/mcpserver/anthropic_beta_stage.go new file mode 100644 index 000000000..1bbb52391 --- /dev/null +++ b/internal/mcpserver/anthropic_beta_stage.go @@ -0,0 +1,334 @@ +package mcpserver + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/anthropics/anthropic-sdk-go" + + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + stagetoolloop "github.com/tingly-dev/tingly-box/internal/protocol/stage/toolloop" +) + +// AnthropicBetaToolProvider prepares one Beta request with the exact tools the +// Stage owns. Returning ownership explicitly avoids treating a client-declared +// tool as internal merely because its name resembles an MCP tool name. +type AnthropicBetaToolProvider interface { + PrepareRequest(ctx context.Context, request *anthropic.BetaMessageNewParams) ([]string, error) +} + +// AnthropicBetaStageExecutor is the existing server-tool execution boundary +// needed by the Beta-native Stage. The call remains a Beta tool_use until this +// boundary; no canonical tool DTO is introduced. +type AnthropicBetaStageExecutor interface { + ExecuteToolWithContext(ctx context.Context, tool Tool, messages []map[string]any) (context.Context, ToolExecutionResult, error) +} + +// AnthropicBetaContinuationStore owns the Beta-native continuation segment +// needed when one model response mixes internal and client-owned tool calls. +// A production implementation may bind one instance to a provider and derive +// the session key from ctx; the Stage never knows that storage key. +type AnthropicBetaContinuationStore interface { + Pop(ctx context.Context, request *anthropic.BetaMessageNewParams) ([]anthropic.BetaMessageParam, bool) + Put(ctx context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) +} + +type AnthropicBetaStageConfig struct { + Name string + Tools AnthropicBetaToolProvider + Executor AnthropicBetaStageExecutor + Continuations AnthropicBetaContinuationStore + MaxRounds int +} + +func NewAnthropicBetaStage(config AnthropicBetaStageConfig) (protocolstage.Stage, error) { + if config.Tools == nil { + return nil, errors.New("construct Anthropic Beta ToolLoop Stage: tool provider is nil") + } + if config.Executor == nil { + return nil, errors.New("construct Anthropic Beta ToolLoop Stage: executor is nil") + } + name := config.Name + if name == "" { + name = "tool_loop_anthropic_beta" + } + maxRounds := config.MaxRounds + if maxRounds <= 0 { + maxRounds = defaultMaxRounds + } + return &anthropicBetaToolLoopStage{ + name: name, + tools: config.Tools, + executor: config.Executor, + continuations: config.Continuations, + maxRounds: maxRounds, + adapter: NewAnthropicBetaAdapter(), + }, nil +} + +type anthropicBetaToolLoopStage struct { + name string + tools AnthropicBetaToolProvider + executor AnthropicBetaStageExecutor + continuations AnthropicBetaContinuationStore + maxRounds int + adapter *AnthropicBetaAdapter +} + +func (s *anthropicBetaToolLoopStage) Name() string { return s.name } +func (*anthropicBetaToolLoopStage) Protocol() protocol.APIType { return protocol.TypeAnthropicBeta } +func (s *anthropicBetaToolLoopStage) Wrap(next protocolstage.Endpoint) protocolstage.Endpoint { + return &anthropicBetaToolLoopEndpoint{stage: s, next: next} +} + +type anthropicBetaToolLoopEndpoint struct { + stage *anthropicBetaToolLoopStage + next protocolstage.Endpoint +} + +func (*anthropicBetaToolLoopEndpoint) Protocol() protocol.APIType { + return protocol.TypeAnthropicBeta +} + +func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + prepared, owned, err := e.prepare(ctx, call) + if err != nil { + return nil, err + } + + runCtx := ctx + current := prepared + var totalUsage *protocol.TokenUsage + sideEffectsCommitted := false + for round := 1; round <= e.stage.maxRounds; round++ { + response, callErr := e.next.Complete(runCtx, current) + if callErr != nil { + return nil, stagetoolloop.WrapError(callErr, sideEffectsCommitted) + } + if response == nil { + return nil, stagetoolloop.WrapError(errors.New("Anthropic Beta ToolLoop received a nil response"), sideEffectsCommitted) + } + totalUsage = mergeBetaStageUsage(totalUsage, response.Usage) + sideEffectsCommitted = sideEffectsCommitted || response.SideEffectsCommitted + + message, parseErr := betaStageMessage(response.Value) + if parseErr != nil { + return nil, stagetoolloop.WrapError(parseErr, sideEffectsCommitted) + } + tools, extractErr := e.stage.adapter.ExtractTools(message) + if extractErr != nil { + return nil, stagetoolloop.WrapError(extractErr, sideEffectsCommitted) + } + managed, external, externalIDs := splitBetaStageTools(tools, owned) + if len(managed) == 0 { + response.Usage = totalUsage + response.SideEffectsCommitted = sideEffectsCommitted + return response, nil + } + if len(external) > 0 { + if e.stage.continuations == nil { + response.Usage = totalUsage + response.SideEffectsCommitted = sideEffectsCommitted + return response, nil + } + results, nextCtx, committed := e.executeTools(runCtx, current.Request, managed) + sideEffectsCommitted = sideEffectsCommitted || committed + runCtx = nextCtx + normalized, normalizeErr := validateAndNormalizeMixedStash(externalIDs, results) + if normalizeErr != nil { + return nil, stagetoolloop.WrapError(normalizeErr, sideEffectsCommitted) + } + segmentValue, segmentErr := e.stage.adapter.BuildContinuationSegment(message, normalized) + if segmentErr != nil { + return nil, stagetoolloop.WrapError(segmentErr, sideEffectsCommitted) + } + segment, ok := segmentValue.([]anthropic.BetaMessageParam) + if !ok || len(segment) == 0 { + return nil, stagetoolloop.WrapError(errors.New("Anthropic Beta ToolLoop built an empty mixed continuation"), sideEffectsCommitted) + } + e.stage.continuations.Put(runCtx, segment, externalIDs) + filtered, filterErr := e.stage.adapter.FilterVirtualTools(message, external) + if filterErr != nil { + return nil, stagetoolloop.WrapError(filterErr, sideEffectsCommitted) + } + response.Value = filtered + response.Usage = totalUsage + response.SideEffectsCommitted = sideEffectsCommitted + return response, nil + } + if round == e.stage.maxRounds { + return nil, stagetoolloop.WrapError(stagetoolloop.ErrMaxRounds, sideEffectsCommitted) + } + + results, nextCtx, committed := e.executeTools(runCtx, current.Request, managed) + sideEffectsCommitted = sideEffectsCommitted || committed + runCtx = nextCtx + resultValues := make([]any, len(results)) + for i := range results { + resultValues[i] = results[i] + } + nextRequest, appendErr := e.stage.adapter.AppendToolResults(current.Request, message, resultValues) + if appendErr != nil { + return nil, stagetoolloop.WrapError(appendErr, sideEffectsCommitted) + } + current.Request = nextRequest + } + return nil, stagetoolloop.WrapError(stagetoolloop.ErrMaxRounds, sideEffectsCommitted) +} + +func (e *anthropicBetaToolLoopEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + prepared, owned, err := e.prepare(ctx, call) + if err != nil { + return nil, err + } + return newAnthropicBetaToolLoopStream(ctx, e, prepared, owned) +} + +func (e *anthropicBetaToolLoopEndpoint) prepare(ctx context.Context, call protocolstage.Call) (protocolstage.Call, map[string]struct{}, error) { + request, ok := call.Request.(*anthropic.BetaMessageNewParams) + if !ok || request == nil { + return protocolstage.Call{}, nil, fmt.Errorf("Anthropic Beta ToolLoop received request %T", call.Request) + } + cloned, err := cloneBetaStageRequest(request) + if err != nil { + return protocolstage.Call{}, nil, err + } + existing := betaStageToolNames(cloned.Tools) + ownedNames, err := e.stage.tools.PrepareRequest(ctx, cloned) + if err != nil { + return protocolstage.Call{}, nil, fmt.Errorf("prepare Anthropic Beta ToolLoop tools: %w", err) + } + owned := make(map[string]struct{}, len(ownedNames)) + for _, name := range ownedNames { + if name == "" { + return protocolstage.Call{}, nil, errors.New("Anthropic Beta ToolLoop provider returned an empty owned tool name") + } + if _, duplicated := owned[name]; duplicated { + return protocolstage.Call{}, nil, fmt.Errorf("Anthropic Beta ToolLoop provider returned duplicate owned tool %q", name) + } + if _, collision := existing[name]; collision { + return protocolstage.Call{}, nil, fmt.Errorf("%w: %q", stagetoolloop.ErrToolNameCollision, name) + } + owned[name] = struct{}{} + } + preparedNames := betaStageToolNames(cloned.Tools) + for name := range owned { + if _, injected := preparedNames[name]; !injected { + return protocolstage.Call{}, nil, fmt.Errorf("Anthropic Beta ToolLoop provider claimed tool %q without injecting it", name) + } + } + prepared := call + if e.stage.continuations != nil { + if segment, ok := e.stage.continuations.Pop(ctx, cloned); ok { + continued, applyErr := e.stage.adapter.ApplyContinuation(cloned, segment) + if applyErr != nil { + return protocolstage.Call{}, nil, fmt.Errorf("apply Anthropic Beta ToolLoop continuation: %w", applyErr) + } + var continuedOK bool + cloned, continuedOK = continued.(*anthropic.BetaMessageNewParams) + if !continuedOK || cloned == nil { + return protocolstage.Call{}, nil, fmt.Errorf("apply Anthropic Beta ToolLoop continuation returned %T", continued) + } + } + } + prepared.Request = cloned + return prepared, owned, nil +} + +func (e *anthropicBetaToolLoopEndpoint) executeTools( + ctx context.Context, + request any, + tools []Tool, +) ([]ToolExecutionResult, context.Context, bool) { + messages := extractMessagesForToolCall(request) + results := make([]ToolExecutionResult, 0, len(tools)) + runCtx := ctx + committed := false + for _, tool := range tools { + nextCtx, result, err := e.stage.executor.ExecuteToolWithContext(runCtx, tool, messages) + if nextCtx != nil { + runCtx = nextCtx + } + if result.ToolUseID == "" { + result.ToolUseID = tool.ID() + } + if err != nil { + result.IsError = true + } + if err == nil || result.Dispatched { + committed = true + } + results = append(results, result) + } + return results, runCtx, committed +} + +func betaStageMessage(value any) (*anthropic.BetaMessage, error) { + switch message := value.(type) { + case *anthropic.BetaMessage: + if message == nil { + return nil, errors.New("Anthropic Beta ToolLoop received a nil message") + } + return message, nil + case anthropic.BetaMessage: + return &message, nil + default: + return nil, fmt.Errorf("Anthropic Beta ToolLoop received response %T", value) + } +} + +func splitBetaStageTools(tools []Tool, owned map[string]struct{}) (managed, external []Tool, externalIDs []string) { + managed = make([]Tool, 0, len(tools)) + external = make([]Tool, 0, len(tools)) + for _, tool := range tools { + if _, ok := owned[tool.Name()]; ok { + managed = append(managed, tool) + continue + } + external = append(external, tool) + externalIDs = append(externalIDs, tool.ID()) + } + return managed, external, externalIDs +} + +func betaStageToolNames(tools []anthropic.BetaToolUnionParam) map[string]struct{} { + names := make(map[string]struct{}, len(tools)) + for _, tool := range tools { + if tool.OfTool != nil && tool.OfTool.Name != "" { + names[tool.OfTool.Name] = struct{}{} + } + } + return names +} + +func cloneBetaStageRequest(request *anthropic.BetaMessageNewParams) (*anthropic.BetaMessageNewParams, error) { + raw, err := json.Marshal(request) + if err != nil { + return nil, fmt.Errorf("clone Anthropic Beta ToolLoop request: %w", err) + } + var cloned anthropic.BetaMessageNewParams + if err := json.Unmarshal(raw, &cloned); err != nil { + return nil, fmt.Errorf("clone Anthropic Beta ToolLoop request: %w", err) + } + return &cloned, nil +} + +func mergeBetaStageUsage(total, current *protocol.TokenUsage) *protocol.TokenUsage { + if current == nil { + return total + } + if total == nil { + copy := *current + return © + } + total.InputTokens += current.InputTokens + total.OutputTokens += current.OutputTokens + total.CacheReadTokens += current.CacheReadTokens + total.CacheWriteTokens += current.CacheWriteTokens + total.ReasoningTokens += current.ReasoningTokens + total.SystemTokens += current.SystemTokens + return total +} diff --git a/internal/mcpserver/anthropic_beta_stage_record_test.go b/internal/mcpserver/anthropic_beta_stage_record_test.go new file mode 100644 index 000000000..a25bee449 --- /dev/null +++ b/internal/mcpserver/anthropic_beta_stage_record_test.go @@ -0,0 +1,164 @@ +package mcpserver + +import ( + "context" + "errors" + "io" + "strings" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/assembler" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/record" + coretool "github.com/tingly-dev/tingly-box/internal/tool" +) + +func TestAnthropicBetaToolLoopRecordingComplete(t *testing.T) { + request := &anthropic.BetaMessageNewParams{ + Model: "client", + MaxTokens: 64, + Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hello"))}, + } + recorder := newBetaStageRecorder(t, "beta-complete", request) + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-record", Name: "lookup"})}, + {Value: betaStageTextMessage(t, "recorded final")}, + }} + observed := record.ObserveProvider(terminal, recorder, record.ExchangeMetadata{ + Attempt: 4, Provider: "provider-a", Model: "provider-model", + }) + endpoint := composeRecordedBetaToolLoop(t, observed) + + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatal(err) + } + if err := recorder.SetFinalResponse(protocol.TypeAnthropicBeta, response.Value); err != nil { + t.Fatal(err) + } + completed, first := recorder.Finish(nil) + if !first { + t.Fatal("recorder was already finished") + } + assertBetaStageToolLoopRecord(t, completed, "toolu-record", "recorded final") +} + +func TestAnthropicBetaToolLoopRecordingStream(t *testing.T) { + request := &anthropic.BetaMessageNewParams{Model: "client", MaxTokens: 64} + recorder := newBetaStageRecorder(t, "beta-stream", request) + terminal := &betaStageScriptedEndpoint{streams: []*betaStageMemoryStream{ + {events: betaStageToolStreamEvents(betaStageToolCallSpec{ID: "toolu-record-stream", Name: "lookup"})}, + {events: betaStageTextStreamEvents("recorded stream final")}, + }} + observed := record.ObserveProvider(terminal, recorder, record.ExchangeMetadata{ + Attempt: 4, Provider: "provider-a", Model: "provider-model", + }) + endpoint := composeRecordedBetaToolLoop(t, observed) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatal(err) + } + finalAssembler, err := assembler.NewStreamAssembler(protocol.TypeAnthropicBeta) + if err != nil { + t.Fatal(err) + } + for { + event, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + t.Fatal(nextErr) + } + if err := finalAssembler.Add(event.Value); err != nil { + t.Fatal(err) + } + } + if err := stream.Close(); err != nil { + t.Fatal(err) + } + finalResponse, err := finalAssembler.Finish() + if err != nil { + t.Fatal(err) + } + if err := recorder.SetFinalResponse(protocol.TypeAnthropicBeta, finalResponse); err != nil { + t.Fatal(err) + } + completed, first := recorder.Finish(nil) + if !first { + t.Fatal("recorder was already finished") + } + assertBetaStageToolLoopRecord(t, completed, "toolu-record-stream", "recorded stream final") +} + +func newBetaStageRecorder(t *testing.T, requestID string, input *anthropic.BetaMessageNewParams) *record.Recorder { + t.Helper() + recorder, err := record.New(record.Config{ + Enabled: true, + RequestID: requestID, + InputProtocol: protocol.TypeAnthropicBeta, + Input: input, + }) + if err != nil { + t.Fatal(err) + } + return recorder +} + +func composeRecordedBetaToolLoop(t *testing.T, terminal protocolstage.Endpoint) protocolstage.Endpoint { + t.Helper() + toolStage, err := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }}, + }) + if err != nil { + t.Fatal(err) + } + endpoint, err := protocolstage.Compose(terminal, toolStage) + if err != nil { + t.Fatal(err) + } + return endpoint +} + +func assertBetaStageToolLoopRecord(t *testing.T, completed *record.RequestRecord, toolID, finalText string) { + t.Helper() + if completed == nil || completed.Outcome != record.OutcomeSucceeded { + t.Fatalf("completed record = %#v", completed) + } + if len(completed.ProviderExchanges) != 2 { + t.Fatalf("provider exchanges = %d, want 2", len(completed.ProviderExchanges)) + } + for i, exchange := range completed.ProviderExchanges { + if exchange.Sequence != i+1 || exchange.Attempt != 4 || exchange.Provider != "provider-a" || exchange.Protocol != protocol.TypeAnthropicBeta { + t.Fatalf("exchange %d metadata = %#v", i, exchange) + } + if exchange.Outcome != record.OutcomeSucceeded || exchange.Response == nil { + t.Fatalf("exchange %d result = %#v", i, exchange) + } + } + firstResponse := string(completed.ProviderExchanges[0].Response.Body) + if !strings.Contains(firstResponse, toolID) || !strings.Contains(firstResponse, "lookup") { + t.Fatalf("first provider response = %s", firstResponse) + } + secondRequest := string(completed.ProviderExchanges[1].Request.Body) + if !strings.Contains(secondRequest, toolID) || !strings.Contains(secondRequest, "tool_result") { + t.Fatalf("second provider request = %s", secondRequest) + } + secondResponse := string(completed.ProviderExchanges[1].Response.Body) + if !strings.Contains(secondResponse, finalText) { + t.Fatalf("second provider response = %s", secondResponse) + } + if completed.FinalResponse == nil || completed.FinalResponse.Protocol != protocol.TypeAnthropicBeta || !strings.Contains(string(completed.FinalResponse.Body), finalText) { + t.Fatalf("final response = %#v", completed.FinalResponse) + } + if strings.Contains(string(completed.InputRequest.Body), "lookup") { + t.Fatalf("input request was captured after Stage mutation: %s", completed.InputRequest.Body) + } +} diff --git a/internal/mcpserver/anthropic_beta_stage_stream.go b/internal/mcpserver/anthropic_beta_stage_stream.go new file mode 100644 index 000000000..203c24b82 --- /dev/null +++ b/internal/mcpserver/anthropic_beta_stage_stream.go @@ -0,0 +1,368 @@ +package mcpserver + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "sort" + + "github.com/anthropics/anthropic-sdk-go" + + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/assembler" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + stagetoolloop "github.com/tingly-dev/tingly-box/internal/protocol/stage/toolloop" +) + +// anthropicBetaToolLoopStream buffers one provider round before exposing it. +// Beta permits text/thinking blocks before a later tool_use, so no earlier +// event can prove that the round is safe to expose. Buffering is the only way +// to hide internal tools while preserving one valid Anthropic message stream. +type anthropicBetaToolLoopStream struct { + endpoint *anthropicBetaToolLoopEndpoint + call protocolstage.Call + owned map[string]struct{} + runCtx context.Context + + round int + current protocolstage.EventStream + assembler assembler.StreamAssembler + buffered []protocolstage.Event + pending []protocolstage.Event + + usage *protocol.TokenUsage + model string + sideEffects bool + done bool + closed bool +} + +func newAnthropicBetaToolLoopStream( + ctx context.Context, + endpoint *anthropicBetaToolLoopEndpoint, + call protocolstage.Call, + owned map[string]struct{}, +) (protocolstage.EventStream, error) { + stream := &anthropicBetaToolLoopStream{ + endpoint: endpoint, + call: call, + owned: owned, + runCtx: ctx, + } + if err := stream.startRound(ctx); err != nil { + return nil, err + } + return stream, nil +} + +func (s *anthropicBetaToolLoopStream) Next(ctx context.Context) (protocolstage.Event, error) { + for { + if len(s.pending) > 0 { + event := s.pending[0] + s.pending = s.pending[1:] + return event, nil + } + if s.done || s.closed { + return protocolstage.Event{}, io.EOF + } + + event, err := s.current.Next(ctx) + if err == nil { + if assembleErr := s.assembler.Add(event.Value); assembleErr != nil { + return protocolstage.Event{}, s.fail(assembleErr) + } + s.buffered = append(s.buffered, event) + continue + } + + s.absorbCurrentResult() + closeErr := s.closeCurrent() + if !errors.Is(err, io.EOF) { + if closeErr != nil { + err = errors.Join(err, closeErr) + } + return protocolstage.Event{}, s.fail(err) + } + if closeErr != nil { + return protocolstage.Event{}, s.fail(closeErr) + } + + complete, finishErr := s.assembler.Finish() + if finishErr != nil { + return protocolstage.Event{}, s.fail(finishErr) + } + message, parseErr := betaStageMessage(complete) + if parseErr != nil { + return protocolstage.Event{}, s.fail(parseErr) + } + tools, extractErr := s.endpoint.stage.adapter.ExtractTools(message) + if extractErr != nil { + return protocolstage.Event{}, s.fail(extractErr) + } + managed, external, externalIDs := splitBetaStageTools(tools, s.owned) + if len(managed) == 0 { + s.pending = s.buffered + s.buffered = nil + s.done = true + continue + } + if len(external) > 0 { + if s.endpoint.stage.continuations == nil { + s.pending = s.buffered + s.buffered = nil + s.done = true + continue + } + results, nextCtx, committed := s.endpoint.executeTools(s.runCtx, s.call.Request, managed) + s.sideEffects = s.sideEffects || committed + s.runCtx = nextCtx + normalized, normalizeErr := validateAndNormalizeMixedStash(externalIDs, results) + if normalizeErr != nil { + return protocolstage.Event{}, s.fail(normalizeErr) + } + segmentValue, segmentErr := s.endpoint.stage.adapter.BuildContinuationSegment(message, normalized) + if segmentErr != nil { + return protocolstage.Event{}, s.fail(segmentErr) + } + segment, ok := segmentValue.([]anthropic.BetaMessageParam) + if !ok || len(segment) == 0 { + return protocolstage.Event{}, s.fail(errors.New("Anthropic Beta ToolLoop built an empty mixed continuation")) + } + s.endpoint.stage.continuations.Put(s.runCtx, segment, externalIDs) + filtered, filterErr := filterBetaStageStreamEvents(s.buffered, s.owned) + if filterErr != nil { + return protocolstage.Event{}, s.fail(filterErr) + } + s.pending = filtered + s.buffered = nil + s.done = true + continue + } + if s.round >= s.endpoint.stage.maxRounds { + return protocolstage.Event{}, s.fail(stagetoolloop.ErrMaxRounds) + } + + results, nextCtx, committed := s.endpoint.executeTools(s.runCtx, s.call.Request, managed) + s.sideEffects = s.sideEffects || committed + s.runCtx = nextCtx + resultValues := make([]any, len(results)) + for i := range results { + resultValues[i] = results[i] + } + nextRequest, appendErr := s.endpoint.stage.adapter.AppendToolResults(s.call.Request, message, resultValues) + if appendErr != nil { + return protocolstage.Event{}, s.fail(appendErr) + } + s.call.Request = nextRequest + s.buffered = nil + if startErr := s.startRound(s.runCtx); startErr != nil { + return protocolstage.Event{}, s.fail(startErr) + } + } +} + +func (s *anthropicBetaToolLoopStream) Close() error { + if s.closed { + return nil + } + s.closed = true + s.done = true + s.absorbCurrentResult() + return s.closeCurrent() +} + +func (s *anthropicBetaToolLoopStream) Result() protocolstage.StreamResult { + usage := cloneBetaStageUsage(s.usage) + model := s.model + sideEffects := s.sideEffects + if s.current != nil { + current := s.current.Result() + usage = mergeBetaStageUsage(usage, current.Usage) + if current.Model != "" { + model = current.Model + } + sideEffects = sideEffects || current.SideEffectsCommitted + } + return protocolstage.StreamResult{Usage: usage, Model: model, SideEffectsCommitted: sideEffects} +} + +func (s *anthropicBetaToolLoopStream) startRound(ctx context.Context) error { + stream, err := s.endpoint.next.Stream(ctx, s.call) + if err != nil { + return err + } + if stream == nil { + return errors.New("Anthropic Beta ToolLoop received a nil stream") + } + streamAssembler, err := assembler.NewStreamAssembler(protocol.TypeAnthropicBeta) + if err != nil { + _ = stream.Close() + return err + } + s.current = stream + s.assembler = streamAssembler + s.round++ + return nil +} + +func (s *anthropicBetaToolLoopStream) absorbCurrentResult() { + if s.current == nil { + return + } + result := s.current.Result() + s.usage = mergeBetaStageUsage(s.usage, result.Usage) + if result.Model != "" { + s.model = result.Model + } + s.sideEffects = s.sideEffects || result.SideEffectsCommitted +} + +func (s *anthropicBetaToolLoopStream) closeCurrent() error { + if s.current == nil { + return nil + } + current := s.current + s.current = nil + return current.Close() +} + +func (s *anthropicBetaToolLoopStream) fail(err error) error { + s.done = true + if s.current != nil { + s.absorbCurrentResult() + if closeErr := s.closeCurrent(); closeErr != nil { + err = errors.Join(err, closeErr) + } + } + return stagetoolloop.WrapError(err, s.sideEffects) +} + +func cloneBetaStageUsage(usage *protocol.TokenUsage) *protocol.TokenUsage { + if usage == nil { + return nil + } + cloned := *usage + return &cloned +} + +func filterBetaStageStreamEvents(events []protocolstage.Event, owned map[string]struct{}) ([]protocolstage.Event, error) { + suppressed := make(map[int]struct{}) + for _, event := range events { + value, err := betaStageStreamEvent(event.Value) + if err != nil { + return nil, err + } + tool, ok := NewAnthropicBetaAdapter().ExtractToolFromEvent(value) + if !ok { + continue + } + if _, internal := owned[tool.Name()]; !internal { + continue + } + if index, ok := extractContentBlockIndex(value); ok { + suppressed[index] = struct{}{} + } + } + if len(suppressed) == 0 { + return append([]protocolstage.Event(nil), events...), nil + } + indices := make([]int, 0, len(suppressed)) + for index := range suppressed { + indices = append(indices, index) + } + sort.Ints(indices) + + filtered := make([]protocolstage.Event, 0, len(events)) + for _, event := range events { + value, err := betaStageStreamEvent(event.Value) + if err != nil { + return nil, err + } + index, indexed := extractContentBlockIndex(value) + if indexed { + if _, remove := suppressed[index]; remove { + continue + } + offset := 0 + for _, suppressedIndex := range indices { + if suppressedIndex < index { + offset++ + } + } + if offset > 0 { + value, err = rewriteBetaStageEventIndex(value, index-offset) + if err != nil { + return nil, err + } + } + } + filtered = append(filtered, protocolstage.Event{Value: value}) + } + return filtered, nil +} + +func betaStageStreamEvent(value any) (*anthropic.BetaRawMessageStreamEventUnion, error) { + switch event := value.(type) { + case *anthropic.BetaRawMessageStreamEventUnion: + if event == nil { + return nil, errors.New("Anthropic Beta ToolLoop received a nil stream event") + } + return event, nil + case anthropic.BetaRawMessageStreamEventUnion: + copy := event + return ©, nil + } + raw, err := betaStageStreamEventJSON(value) + if err != nil { + return nil, err + } + var event anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(raw, &event); err != nil { + return nil, fmt.Errorf("decode Anthropic Beta ToolLoop stream event %T: %w", value, err) + } + return &event, nil +} + +func rewriteBetaStageEventIndex(event *anthropic.BetaRawMessageStreamEventUnion, index int) (*anthropic.BetaRawMessageStreamEventUnion, error) { + raw, err := betaStageStreamEventJSON(event) + if err != nil { + return nil, err + } + var value map[string]any + if err := json.Unmarshal(raw, &value); err != nil { + return nil, err + } + value["index"] = index + rewritten, err := json.Marshal(value) + if err != nil { + return nil, err + } + var result anthropic.BetaRawMessageStreamEventUnion + if err := json.Unmarshal(rewritten, &result); err != nil { + return nil, err + } + return &result, nil +} + +func betaStageStreamEventJSON(value any) ([]byte, error) { + if value == nil { + return nil, errors.New("Anthropic Beta ToolLoop received a nil stream event") + } + switch event := value.(type) { + case json.RawMessage: + return event, nil + case []byte: + return event, nil + case interface{ RawJSON() string }: + if raw := event.RawJSON(); raw != "" { + return []byte(raw), nil + } + } + raw, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("marshal Anthropic Beta ToolLoop stream event %T: %w", value, err) + } + return raw, nil +} diff --git a/internal/mcpserver/anthropic_beta_stage_test.go b/internal/mcpserver/anthropic_beta_stage_test.go new file mode 100644 index 000000000..facf20698 --- /dev/null +++ b/internal/mcpserver/anthropic_beta_stage_test.go @@ -0,0 +1,751 @@ +package mcpserver + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/anthropicbridge" + stagetoolloop "github.com/tingly-dev/tingly-box/internal/protocol/stage/toolloop" + coretool "github.com/tingly-dev/tingly-box/internal/tool" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func TestMergeBetaStageUsagePreservesCurrentCacheDetails(t *testing.T) { + total := &protocol.TokenUsage{InputTokens: 3, OutputTokens: 2, CacheReadTokens: 1, CacheWriteTokens: 2, ReasoningTokens: 1} + current := &protocol.TokenUsage{InputTokens: 5, OutputTokens: 4, CacheReadTokens: 3, CacheWriteTokens: 4, ReasoningTokens: 2} + + got := mergeBetaStageUsage(total, current) + if got.InputTokens != 8 || got.OutputTokens != 6 || got.CacheReadTokens != 4 || got.CacheWriteTokens != 6 || got.ReasoningTokens != 3 { + t.Fatalf("merged usage = %#v", got) + } +} + +func TestAnthropicV1TopologyRunsBetaToolLoop(t *testing.T) { + t.Run("complete", func(t *testing.T) { + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-v1", Name: "lookup"}), Usage: protocol.NewTokenUsage(3, 2)}, + {Value: betaStageTextMessage(t, "done through Beta"), Usage: protocol.NewTokenUsage(5, 4), Model: "provider"}, + }} + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }} + endpoint := buildV1BetaToolLoopTopology(t, terminal, executor) + + response, err := endpoint.Complete(context.Background(), protocolstage.Call{ + Request: &anthropic.MessageNewParams{ + Model: "client", + MaxTokens: 64, + Messages: []anthropic.MessageParam{anthropic.NewUserMessage(anthropic.NewTextBlock("hello"))}, + }, + Metadata: protocolstage.CallMetadata{RequestID: "v1-beta-complete", Attempt: 2}, + }) + if err != nil { + t.Fatal(err) + } + message, ok := response.Value.(*anthropic.Message) + if !ok || message == nil || len(message.Content) != 1 || message.Content[0].Text != "done through Beta" { + t.Fatalf("v1 response = %#v (%T)", response.Value, response.Value) + } + if response.Usage == nil || response.Usage.InputTokens != 8 || response.Usage.OutputTokens != 6 || !response.SideEffectsCommitted { + t.Fatalf("response facts = %#v", response) + } + if len(terminal.calls) != 2 || len(executor.calls) != 1 { + t.Fatalf("provider calls=%d executor calls=%d", len(terminal.calls), len(executor.calls)) + } + if terminal.calls[0].Metadata.RequestID != "v1-beta-complete" { + t.Fatalf("provider metadata = %+v", terminal.calls[0].Metadata) + } + continued := terminal.calls[1].Request.(*anthropic.BetaMessageNewParams) + if len(continued.Messages) != 3 { + t.Fatalf("Beta continuation messages = %d", len(continued.Messages)) + } + }) + + t.Run("stream", func(t *testing.T) { + first := &betaStageMemoryStream{ + events: betaStageToolStreamEvents(betaStageToolCallSpec{ID: "toolu-v1-stream", Name: "lookup"}), + result: protocolstage.StreamResult{Usage: protocol.NewTokenUsage(3, 2), Model: "provider"}, + } + second := &betaStageMemoryStream{ + events: betaStageTextStreamEvents("done through Beta stream"), + result: protocolstage.StreamResult{Usage: protocol.NewTokenUsage(5, 4), Model: "provider"}, + } + terminal := &betaStageScriptedEndpoint{streams: []*betaStageMemoryStream{first, second}} + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }} + endpoint := buildV1BetaToolLoopTopology(t, terminal, executor) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &anthropic.MessageNewParams{ + Model: "client", MaxTokens: 64, + }}) + if err != nil { + t.Fatal(err) + } + var bodies []string + for { + event, nextErr := stream.Next(context.Background()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + t.Fatal(nextErr) + } + if _, ok := event.Value.(anthropic.MessageStreamEventUnion); !ok { + t.Fatalf("v1 stream event type = %T", event.Value) + } + raw, marshalErr := json.Marshal(event.Value) + if marshalErr != nil { + t.Fatal(marshalErr) + } + bodies = append(bodies, string(raw)) + } + body := strings.Join(bodies, "\n") + if strings.Contains(body, "toolu-v1-stream") || strings.Contains(body, "lookup") { + t.Fatalf("internal Beta round leaked to v1 stream: %s", body) + } + if !strings.Contains(body, "done through Beta stream") { + t.Fatalf("final Beta round missing from v1 stream: %s", body) + } + result := stream.Result() + if result.Usage == nil || result.Usage.InputTokens != 8 || result.Usage.OutputTokens != 6 || !result.SideEffectsCommitted { + t.Fatalf("stream result = %#v", result) + } + if err := stream.Close(); err != nil { + t.Fatal(err) + } + if first.closeCalls != 1 || second.closeCalls != 1 || len(executor.calls) != 1 { + t.Fatalf("close/execution = %d/%d/%d", first.closeCalls, second.closeCalls, len(executor.calls)) + } + }) +} + +func buildV1BetaToolLoopTopology(t *testing.T, terminal protocolstage.Endpoint, executor AnthropicBetaStageExecutor) protocolstage.Endpoint { + t.Helper() + toolStage, err := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + }) + if err != nil { + t.Fatal(err) + } + registry, err := protocolstage.NewBridgeRegistry(anthropicbridge.NewV1ToBeta()) + if err != nil { + t.Fatal(err) + } + endpoint, err := protocolstage.BuildTopology(protocolstage.TopologyConfig{ + Terminal: terminal, + Stages: []protocolstage.Stage{toolStage}, + ClientProtocol: protocol.TypeAnthropicV1, + Registry: registry, + }) + if err != nil { + t.Fatal(err) + } + return endpoint +} + +func TestAnthropicBetaStageCompleteRunsOwnedToolAndContinues(t *testing.T) { + tools := staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}} + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("Paris").Contents}, + }} + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-1", Name: "lookup", Input: map[string]any{"city": "France"}}), Usage: protocol.NewTokenUsage(3, 2), Model: "provider"}, + {Value: betaStageTextMessage(t, "The capital is Paris."), Usage: protocol.NewTokenUsage(5, 4), Model: "provider"}, + }} + toolStage, err := NewAnthropicBetaStage(AnthropicBetaStageConfig{Tools: tools, Executor: executor}) + if err != nil { + t.Fatal(err) + } + endpoint, err := protocolstage.Compose(terminal, toolStage) + if err != nil { + t.Fatal(err) + } + request := &anthropic.BetaMessageNewParams{ + Model: "client", + MaxTokens: 100, + Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hello"))}, + } + + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if err != nil { + t.Fatal(err) + } + message := response.Value.(*anthropic.BetaMessage) + if len(message.Content) != 1 || message.Content[0].Type != "text" || message.Content[0].Text != "The capital is Paris." { + t.Fatalf("final response = %#v", message.Content) + } + if response.Usage == nil || response.Usage.InputTokens != 8 || response.Usage.OutputTokens != 6 { + t.Fatalf("aggregate usage = %#v", response.Usage) + } + if !response.SideEffectsCommitted { + t.Fatal("successful tool execution did not commit side effects") + } + if len(terminal.calls) != 2 || len(executor.calls) != 1 { + t.Fatalf("provider calls=%d executor calls=%d", len(terminal.calls), len(executor.calls)) + } + firstRequest := terminal.calls[0].Request.(*anthropic.BetaMessageNewParams) + if _, ok := betaStageToolNames(firstRequest.Tools)["lookup"]; !ok { + t.Fatalf("injected tools = %#v", firstRequest.Tools) + } + continuation := terminal.calls[1].Request.(*anthropic.BetaMessageNewParams) + if len(continuation.Messages) != 3 { + t.Fatalf("continuation messages = %d, want user + assistant + tool-result user", len(continuation.Messages)) + } + if len(request.Tools) != 0 || len(request.Messages) != 1 { + t.Fatal("Anthropic Beta ToolLoop mutated the caller request") + } +} + +func TestAnthropicBetaStageCompleteLeavesExternalAndMixedToolsOutward(t *testing.T) { + for _, tt := range []struct { + name string + calls []betaStageToolCallSpec + }{ + {name: "external", calls: []betaStageToolCallSpec{{ID: "toolu-ext", Name: "client_tool"}}}, + {name: "mixed", calls: []betaStageToolCallSpec{{ID: "toolu-owned", Name: "lookup"}, {ID: "toolu-ext", Name: "client_tool"}}}, + } { + t.Run(tt.name, func(t *testing.T) { + executor := &fakeBetaStageExecutor{} + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{{Value: betaStageToolMessage(t, tt.calls...)}}} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + if response == nil || len(terminal.calls) != 1 || len(executor.calls) != 0 { + t.Fatalf("external/mixed tools were consumed: response=%#v provider=%d executor=%d", response, len(terminal.calls), len(executor.calls)) + } + }) + } +} + +func TestAnthropicBetaStageCompleteStoresAndAppliesMixedContinuation(t *testing.T) { + continuations := &memoryBetaStageContinuations{} + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("internal result").Contents}, + }} + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + {Value: betaStageToolMessage(t, + betaStageToolCallSpec{ID: "toolu-owned", Name: "lookup"}, + betaStageToolCallSpec{ID: "toolu-external", Name: "client_tool"}, + )}, + {Value: betaStageTextMessage(t, "combined")}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + Continuations: continuations, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + first, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + filtered := first.Value.(*anthropic.BetaMessage) + filteredTools, err := NewAnthropicBetaAdapter().ExtractTools(filtered) + if err != nil { + t.Fatal(err) + } + if len(filteredTools) != 1 || filteredTools[0].Name() != "client_tool" { + t.Fatalf("filtered mixed tools = %#v", filteredTools) + } + if !first.SideEffectsCommitted || continuations.puts != 1 || len(executor.calls) != 1 { + t.Fatalf("first result committed=%v puts=%d executions=%d", first.SideEffectsCommitted, continuations.puts, len(executor.calls)) + } + + externalResult := anthropic.NewBetaUserMessage(anthropic.NewBetaToolResultBlock("toolu-external", "external result", false)) + second, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{ + Messages: []anthropic.BetaMessageParam{externalResult}, + }}) + if err != nil { + t.Fatal(err) + } + if second.Value.(*anthropic.BetaMessage).Content[0].Text != "combined" { + t.Fatalf("second response = %#v", second.Value) + } + if continuations.pops != 2 || len(terminal.calls) != 2 { + t.Fatalf("continuation pops=%d provider calls=%d", continuations.pops, len(terminal.calls)) + } + continued := terminal.calls[1].Request.(*anthropic.BetaMessageNewParams) + if len(continued.Messages) != 2 { + t.Fatalf("continued messages = %d, want assistant + merged tool-result user", len(continued.Messages)) + } + if got := betaStageToolResultIDs(continued.Messages[1]); len(got) != 2 || got[0] != "toolu-owned" || got[1] != "toolu-external" { + t.Fatalf("merged tool result IDs = %#v", got) + } +} + +func TestAnthropicBetaStageRejectsAmbiguousOwnership(t *testing.T) { + request := &anthropic.BetaMessageNewParams{Tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: &fakeBetaStageExecutor{}, + }) + endpoint, _ := protocolstage.Compose(&betaStageScriptedEndpoint{}, toolStage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: request}) + if !errors.Is(err, stagetoolloop.ErrToolNameCollision) { + t.Fatalf("tool name collision error = %v", err) + } +} + +func TestAnthropicBetaStagePreservesSideEffectBoundaryAfterLaterFailure(t *testing.T) { + providerErr := errors.New("second round failed") + terminal := &betaStageScriptedEndpoint{ + responses: []*protocolstage.Response{{Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-1", Name: "lookup"})}}, + errors: []error{nil, providerErr}, + } + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }}, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if !errors.Is(err, providerErr) || !stagetoolloop.HasCommittedSideEffects(err) { + t.Fatalf("later error = %v, committed=%v", err, stagetoolloop.HasCommittedSideEffects(err)) + } +} + +func TestAnthropicBetaStageTreatsPostDispatchToolErrorAsCommitted(t *testing.T) { + providerErr := errors.New("second round failed") + toolErr := errors.New("tool response was lost after dispatch") + terminal := &betaStageScriptedEndpoint{ + responses: []*protocolstage.Response{{Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-1", Name: "lookup"})}}, + errors: []error{nil, providerErr}, + } + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: &fakeBetaStageExecutor{ + results: map[string]ToolExecutionResult{"lookup": {Dispatched: true}}, + errors: map[string]error{"lookup": toolErr}, + }, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if !errors.Is(err, providerErr) || !stagetoolloop.HasCommittedSideEffects(err) { + t.Fatalf("later error = %v, committed=%v", err, stagetoolloop.HasCommittedSideEffects(err)) + } +} + +func TestAnthropicBetaStageStreamHidesOwnedRoundAndContinues(t *testing.T) { + toolEvents := betaStageToolStreamEvents(betaStageToolCallSpec{ID: "toolu-owned", Name: "lookup", Input: map[string]any{"q": "x"}}) + textEvents := betaStageTextStreamEvents("done") + first := &betaStageMemoryStream{events: toolEvents, result: protocolstage.StreamResult{Usage: protocol.NewTokenUsage(3, 2), Model: "provider"}} + second := &betaStageMemoryStream{events: textEvents, result: protocolstage.StreamResult{Usage: protocol.NewTokenUsage(5, 4), Model: "provider"}} + terminal := &betaStageScriptedEndpoint{streams: []*betaStageMemoryStream{first, second}} + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + got := collectBetaStageEvents(t, stream) + if betaStageEventBodies(t, got) != betaStageEventBodies(t, textEvents) { + t.Fatalf("outward events = %s, want final round %s", betaStageEventBodies(t, got), betaStageEventBodies(t, textEvents)) + } + result := stream.Result() + if result.Usage == nil || result.Usage.InputTokens != 8 || result.Usage.OutputTokens != 6 { + t.Fatalf("aggregate stream usage = %#v", result.Usage) + } + if !result.SideEffectsCommitted || result.Model != "provider" { + t.Fatalf("stream result = %#v", result) + } + if len(terminal.streamCalls) != 2 || len(executor.calls) != 1 { + t.Fatalf("provider streams=%d executions=%d", len(terminal.streamCalls), len(executor.calls)) + } + continued := terminal.streamCalls[1].Request.(*anthropic.BetaMessageNewParams) + if len(continued.Messages) != 2 { + t.Fatalf("continuation messages = %d, want assistant + tool-result user", len(continued.Messages)) + } + if first.closeCalls != 1 || second.closeCalls != 1 { + t.Fatalf("inner close calls = %d, %d", first.closeCalls, second.closeCalls) + } + _ = stream.Close() +} + +func TestAnthropicBetaStageStreamFiltersMixedOwnedBlocks(t *testing.T) { + continuations := &memoryBetaStageContinuations{} + events := betaStageToolStreamEvents( + betaStageToolCallSpec{ID: "toolu-owned", Name: "lookup"}, + betaStageToolCallSpec{ID: "toolu-external", Name: "client_tool"}, + ) + terminal := &betaStageScriptedEndpoint{streams: []*betaStageMemoryStream{{events: events}}} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{"lookup": {Contents: coretool.TextToolResult("ok").Contents}}}, + Continuations: continuations, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + got := collectBetaStageEvents(t, stream) + body := betaStageEventBodies(t, got) + if strings.Contains(body, "lookup") || strings.Contains(body, "toolu-owned") { + t.Fatalf("owned tool leaked into outward stream: %s", body) + } + if !strings.Contains(body, "client_tool") || !strings.Contains(body, "toolu-external") { + t.Fatalf("external tool missing from outward stream: %s", body) + } + for _, event := range got { + raw, _ := betaStageStreamEventJSON(event.Value) + var value map[string]any + _ = json.Unmarshal(raw, &value) + if index, ok := value["index"].(float64); ok && int(index) != 0 { + t.Fatalf("filtered external block index = %v, want 0", index) + } + } + if continuations.puts != 1 || !stream.Result().SideEffectsCommitted { + t.Fatalf("mixed continuation puts=%d result=%#v", continuations.puts, stream.Result()) + } + _ = stream.Close() +} + +func TestAnthropicBetaStageStreamPreservesSideEffectBoundaryAfterLaterFailure(t *testing.T) { + providerErr := errors.New("second stream failed") + terminal := &betaStageScriptedEndpoint{ + streams: []*betaStageMemoryStream{{events: betaStageToolStreamEvents(betaStageToolCallSpec{ID: "toolu-1", Name: "lookup"})}}, + streamErrors: []error{nil, providerErr}, + } + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{"lookup": {Contents: coretool.TextToolResult("ok").Contents}}}, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + + _, err = stream.Next(context.Background()) + if !errors.Is(err, providerErr) || !stagetoolloop.HasCommittedSideEffects(err) { + t.Fatalf("later stream error = %v, committed=%v", err, stagetoolloop.HasCommittedSideEffects(err)) + } + if !stream.Result().SideEffectsCommitted { + t.Fatal("stream result lost committed side effects") + } + _ = stream.Close() +} + +func TestProviderBetaContinuationStoreIsProviderScopedAndSingleConsume(t *testing.T) { + ctx := typ.WithSessionID(context.Background(), typ.SessionID{Source: typ.SessionSourceHeader, Value: "session-a"}) + first := NewProviderBetaContinuationStore("provider-a") + second := NewProviderBetaContinuationStore("provider-b") + segment := []anthropic.BetaMessageParam{{ + Role: anthropic.BetaMessageParamRoleAssistant, + Content: []anthropic.BetaContentBlockParamUnion{anthropic.NewBetaTextBlock("stored")}, + }} + request := &anthropic.BetaMessageNewParams{Messages: []anthropic.BetaMessageParam{ + anthropic.NewBetaUserMessage(anthropic.NewBetaToolResultBlock("toolu-external", "ok", false)), + }} + first.Put(ctx, segment, []string{"toolu-external"}) + + if _, ok := second.Pop(ctx, request); ok { + t.Fatal("continuation leaked across providers") + } + got, ok := first.Pop(ctx, request) + if !ok || len(got) != 1 { + t.Fatalf("stored continuation = %#v, ok=%v", got, ok) + } + if _, ok := first.Pop(ctx, request); ok { + t.Fatal("continuation was consumed more than once") + } +} + +type staticBetaStageTools struct { + tools []anthropic.BetaToolUnionParam +} + +func (s staticBetaStageTools) PrepareRequest(_ context.Context, request *anthropic.BetaMessageNewParams) ([]string, error) { + request.Tools = append(request.Tools, s.tools...) + names := make([]string, 0, len(s.tools)) + for _, tool := range s.tools { + if tool.OfTool != nil { + names = append(names, tool.OfTool.Name) + } + } + return names, nil +} + +type fakeBetaStageExecutor struct { + results map[string]ToolExecutionResult + errors map[string]error + calls []Tool +} + +type memoryBetaStageContinuations struct { + segment []anthropic.BetaMessageParam + expectedIDs []string + puts int + pops int +} + +func (s *memoryBetaStageContinuations) Pop(_ context.Context, request *anthropic.BetaMessageNewParams) ([]anthropic.BetaMessageParam, bool) { + s.pops++ + if len(s.segment) == 0 || !containsAll(continuationResultIDs(request), stringSet(s.expectedIDs)) { + return nil, false + } + segment := s.segment + s.segment = nil + return segment, true +} + +func (s *memoryBetaStageContinuations) Put(_ context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) { + s.puts++ + s.segment = append([]anthropic.BetaMessageParam(nil), segment...) + s.expectedIDs = append([]string(nil), externalIDs...) +} + +func (e *fakeBetaStageExecutor) ExecuteToolWithContext(ctx context.Context, tool Tool, _ []map[string]any) (context.Context, ToolExecutionResult, error) { + e.calls = append(e.calls, tool) + return ctx, e.results[tool.Name()], e.errors[tool.Name()] +} + +type betaStageScriptedEndpoint struct { + responses []*protocolstage.Response + errors []error + calls []protocolstage.Call + streams []*betaStageMemoryStream + streamErrors []error + streamCalls []protocolstage.Call +} + +func (*betaStageScriptedEndpoint) Protocol() protocol.APIType { return protocol.TypeAnthropicBeta } + +func (e *betaStageScriptedEndpoint) Complete(_ context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + index := len(e.calls) + e.calls = append(e.calls, call) + if index < len(e.errors) && e.errors[index] != nil { + return nil, e.errors[index] + } + if index >= len(e.responses) { + return nil, errors.New("unexpected provider call") + } + return e.responses[index], nil +} + +func (e *betaStageScriptedEndpoint) Stream(_ context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + index := len(e.streamCalls) + e.streamCalls = append(e.streamCalls, call) + if index < len(e.streamErrors) && e.streamErrors[index] != nil { + return nil, e.streamErrors[index] + } + if index >= len(e.streams) { + return nil, errors.New("unexpected provider stream") + } + return e.streams[index], nil +} + +type betaStageMemoryStream struct { + events []protocolstage.Event + result protocolstage.StreamResult + next int + closeCalls int +} + +func (s *betaStageMemoryStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + if s.next >= len(s.events) { + return protocolstage.Event{}, io.EOF + } + event := s.events[s.next] + s.next++ + return event, nil +} + +func (s *betaStageMemoryStream) Close() error { + s.closeCalls++ + return nil +} + +func (s *betaStageMemoryStream) Result() protocolstage.StreamResult { return s.result } + +type betaStageToolCallSpec struct { + ID string + Name string + Input map[string]any +} + +func betaStageToolMessage(t *testing.T, calls ...betaStageToolCallSpec) *anthropic.BetaMessage { + t.Helper() + content := make([]map[string]any, 0, len(calls)) + for _, call := range calls { + input := call.Input + if input == nil { + input = map[string]any{} + } + content = append(content, map[string]any{ + "type": "tool_use", "id": call.ID, "name": call.Name, "input": input, + }) + } + return decodeBetaStageMessage(t, map[string]any{ + "id": "msg-tool", "type": "message", "role": "assistant", "content": content, + "model": "provider", "stop_reason": "tool_use", "usage": map[string]any{"input_tokens": 3, "output_tokens": 2}, + }) +} + +func betaStageTextMessage(t *testing.T, text string) *anthropic.BetaMessage { + t.Helper() + return decodeBetaStageMessage(t, map[string]any{ + "id": "msg-text", "type": "message", "role": "assistant", + "content": []map[string]any{{"type": "text", "text": text}}, + "model": "provider", "stop_reason": "end_turn", "usage": map[string]any{"input_tokens": 5, "output_tokens": 4}, + }) +} + +func decodeBetaStageMessage(t *testing.T, value any) *anthropic.BetaMessage { + t.Helper() + raw, err := json.Marshal(value) + if err != nil { + t.Fatal(err) + } + var message anthropic.BetaMessage + if err := json.Unmarshal(raw, &message); err != nil { + t.Fatal(err) + } + return &message +} + +func betaStageToolDefinition(name string) anthropic.BetaToolUnionParam { + return anthropic.BetaToolUnionParamOfTool(anthropic.BetaToolInputSchemaParam{ + Properties: map[string]any{}, + }, name) +} + +func betaStageToolResultIDs(message anthropic.BetaMessageParam) []string { + var ids []string + for _, block := range message.Content { + if block.OfToolResult != nil { + ids = append(ids, block.OfToolResult.ToolUseID) + } + } + return ids +} + +func collectBetaStageEvents(t *testing.T, stream protocolstage.EventStream) []protocolstage.Event { + t.Helper() + var events []protocolstage.Event + for { + event, err := stream.Next(context.Background()) + if errors.Is(err, io.EOF) { + return events + } + if err != nil { + t.Fatal(err) + } + events = append(events, event) + } +} + +func betaStageEventBodies(t *testing.T, events []protocolstage.Event) string { + t.Helper() + var bodies []string + for _, event := range events { + raw, err := betaStageStreamEventJSON(event.Value) + if err != nil { + t.Fatal(err) + } + bodies = append(bodies, string(raw)) + } + return strings.Join(bodies, "\n") +} + +func betaStageToolStreamEvents(calls ...betaStageToolCallSpec) []protocolstage.Event { + events := []protocolstage.Event{betaStageRawEvent(map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": "msg-tool-stream", "type": "message", "role": "assistant", "content": []any{}, + "model": "provider", "stop_reason": nil, "stop_sequence": nil, + "usage": map[string]any{"input_tokens": 3, "output_tokens": 0}, + }, + })} + for index, call := range calls { + input := call.Input + if input == nil { + input = map[string]any{} + } + inputJSON, _ := json.Marshal(input) + events = append(events, + betaStageRawEvent(map[string]any{ + "type": "content_block_start", "index": index, + "content_block": map[string]any{"type": "tool_use", "id": call.ID, "name": call.Name, "input": map[string]any{}}, + }), + betaStageRawEvent(map[string]any{ + "type": "content_block_delta", "index": index, + "delta": map[string]any{"type": "input_json_delta", "partial_json": string(inputJSON)}, + }), + betaStageRawEvent(map[string]any{"type": "content_block_stop", "index": index}), + ) + } + events = append(events, + betaStageRawEvent(map[string]any{ + "type": "message_delta", "delta": map[string]any{"stop_reason": "tool_use", "stop_sequence": nil}, + "usage": map[string]any{"output_tokens": 2}, + }), + betaStageRawEvent(map[string]any{"type": "message_stop"}), + ) + return events +} + +func betaStageTextStreamEvents(text string) []protocolstage.Event { + return []protocolstage.Event{ + betaStageRawEvent(map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": "msg-text-stream", "type": "message", "role": "assistant", "content": []any{}, + "model": "provider", "stop_reason": nil, "stop_sequence": nil, + "usage": map[string]any{"input_tokens": 5, "output_tokens": 0}, + }, + }), + betaStageRawEvent(map[string]any{ + "type": "content_block_start", "index": 0, + "content_block": map[string]any{"type": "text", "text": ""}, + }), + betaStageRawEvent(map[string]any{ + "type": "content_block_delta", "index": 0, + "delta": map[string]any{"type": "text_delta", "text": text}, + }), + betaStageRawEvent(map[string]any{"type": "content_block_stop", "index": 0}), + betaStageRawEvent(map[string]any{ + "type": "message_delta", "delta": map[string]any{"stop_reason": "end_turn", "stop_sequence": nil}, + "usage": map[string]any{"output_tokens": 4}, + }), + betaStageRawEvent(map[string]any{"type": "message_stop"}), + } +} + +func betaStageRawEvent(value any) protocolstage.Event { + raw, _ := json.Marshal(value) + return protocolstage.Event{Value: json.RawMessage(raw)} +} diff --git a/internal/mcpserver/continuation_store.go b/internal/mcpserver/continuation_store.go index a7960f421..dc74985d5 100644 --- a/internal/mcpserver/continuation_store.go +++ b/internal/mcpserver/continuation_store.go @@ -1,75 +1,243 @@ package mcpserver import ( + "context" + "encoding/json" "fmt" "sync" "time" + "github.com/anthropics/anthropic-sdk-go" "github.com/openai/openai-go/v3" "github.com/tingly-dev/tingly-box/internal/typ" ) -const continuationTTL = 10 * time.Minute +const ( + continuationTTL = 10 * time.Minute + maxContinuationItems = 256 +) type continuationItem struct { - Segment any - ExpiresAt time.Time + Segment any + ExpectedIDs map[string]struct{} + ExpiresAt time.Time } type continuationStore struct { mu sync.Mutex - items map[string]continuationItem + items map[string][]continuationItem } func newContinuationStore() *continuationStore { - return &continuationStore{ - items: make(map[string]continuationItem), - } + return &continuationStore{items: make(map[string][]continuationItem)} } func continuationKey(sessionID typ.SessionID, providerUUID string, adapterID string) string { + // An IP address is not a conversation identity. Persisting continuation + // state under it can splice requests from different clients behind one NAT. + if sessionID.IsEmpty() || sessionID.IsIPFallback() || providerUUID == "" || adapterID == "" { + return "" + } return fmt.Sprintf("%s:%s|%s|%s", sessionID.Source, sessionID.Value, providerUUID, adapterID) } -func (s *continuationStore) put(key string, segment any) { +func (s *continuationStore) put(key string, segment any, expectedIDs []string) { if s == nil || key == "" || segment == nil { return } + expected := stringSet(expectedIDs) + if len(expected) == 0 { + return + } + + now := time.Now() s.mu.Lock() defer s.mu.Unlock() - s.items[key] = continuationItem{ - Segment: segment, - ExpiresAt: time.Now().Add(continuationTTL), + s.sweepLocked(now) + for s.sizeLocked() >= maxContinuationItems { + s.evictOldestLocked() } + s.items[key] = append(s.items[key], continuationItem{ + Segment: segment, + ExpectedIDs: expected, + ExpiresAt: now.Add(continuationTTL), + }) } -func (s *continuationStore) pop(key string) (any, bool) { +func (s *continuationStore) pop(key string, request any) (any, bool) { if s == nil || key == "" { return nil, false } + resultIDs := continuationResultIDs(request) + if len(resultIDs) == 0 { + return nil, false + } + s.mu.Lock() defer s.mu.Unlock() - item, ok := s.items[key] - if !ok { - return nil, false + s.sweepLocked(time.Now()) + items := s.items[key] + for index, item := range items { + if !containsAll(resultIDs, item.ExpectedIDs) { + continue + } + items = append(items[:index], items[index+1:]...) + if len(items) == 0 { + delete(s.items, key) + } else { + s.items[key] = items + } + return item.Segment, true } - delete(s.items, key) - if time.Now().After(item.ExpiresAt) { - return nil, false + return nil, false +} + +func (s *continuationStore) sweepLocked(now time.Time) { + for key, items := range s.items { + kept := items[:0] + for _, item := range items { + if now.Before(item.ExpiresAt) { + kept = append(kept, item) + } + } + if len(kept) == 0 { + delete(s.items, key) + } else { + s.items[key] = kept + } + } +} + +func (s *continuationStore) sizeLocked() int { + total := 0 + for _, items := range s.items { + total += len(items) + } + return total +} + +func (s *continuationStore) evictOldestLocked() { + var oldestKey string + oldestIndex := -1 + var oldestExpiry time.Time + for key, items := range s.items { + for index, item := range items { + if oldestIndex < 0 || item.ExpiresAt.Before(oldestExpiry) { + oldestKey, oldestIndex, oldestExpiry = key, index, item.ExpiresAt + } + } + } + if oldestIndex < 0 { + return + } + items := s.items[oldestKey] + items = append(items[:oldestIndex], items[oldestIndex+1:]...) + if len(items) == 0 { + delete(s.items, oldestKey) + } else { + s.items[oldestKey] = items + } +} + +func stringSet(values []string) map[string]struct{} { + result := make(map[string]struct{}, len(values)) + for _, value := range values { + if value != "" { + result[value] = struct{}{} + } + } + return result +} + +func containsAll(actual, expected map[string]struct{}) bool { + for id := range expected { + if _, ok := actual[id]; !ok { + return false + } + } + return true +} + +// continuationResultIDs extracts only results in the current trailing client +// turn. Looking through the whole history could match a stale result from an +// earlier turn and consume an unrelated continuation. +func continuationResultIDs(request any) map[string]struct{} { + raw, err := json.Marshal(request) + if err != nil { + return nil + } + var root map[string]any + if err := json.Unmarshal(raw, &root); err != nil { + return nil + } + if messages, ok := root["messages"].([]any); ok { + return trailingMessageResultIDs(messages) + } + if input, ok := root["input"].([]any); ok { + return trailingResponsesResultIDs(input) + } + return nil +} + +func trailingMessageResultIDs(messages []any) map[string]struct{} { + result := make(map[string]struct{}) + for index := len(messages) - 1; index >= 0; index-- { + message, ok := messages[index].(map[string]any) + if !ok { + break + } + role, _ := message["role"].(string) + if role == "tool" { + if id, _ := message["tool_call_id"].(string); id != "" { + result[id] = struct{}{} + } + continue + } + if index != len(messages)-1 { + break + } + content, ok := message["content"].([]any) + if !ok { + break + } + for _, value := range content { + block, _ := value.(map[string]any) + if block["type"] != "tool_result" { + continue + } + if id, _ := block["tool_use_id"].(string); id != "" { + result[id] = struct{}{} + } + } + break } - return item.Segment, true + return result +} + +func trailingResponsesResultIDs(input []any) map[string]struct{} { + result := make(map[string]struct{}) + for index := len(input) - 1; index >= 0; index-- { + item, ok := input[index].(map[string]any) + if !ok || item["type"] != "function_call_output" { + break + } + if id, _ := item["call_id"].(string); id != "" { + result[id] = struct{}{} + } + } + return result } var mixedContinuationStore = newContinuationStore() -func StoreOpenAIContinuationSegment(sessionID typ.SessionID, providerUUID string, segment []openai.ChatCompletionMessageParamUnion) { +func StoreOpenAIContinuationSegment(sessionID typ.SessionID, providerUUID string, segment []openai.ChatCompletionMessageParamUnion, externalIDs []string) { key := continuationKey(sessionID, providerUUID, "openai-chat") - mixedContinuationStore.put(key, segment) + mixedContinuationStore.put(key, segment, externalIDs) } -func PopOpenAIContinuationSegment(sessionID typ.SessionID, providerUUID string) ([]openai.ChatCompletionMessageParamUnion, bool) { +func PopOpenAIContinuationSegment(sessionID typ.SessionID, providerUUID string, request *openai.ChatCompletionNewParams) ([]openai.ChatCompletionMessageParamUnion, bool) { key := continuationKey(sessionID, providerUUID, "openai-chat") - seg, ok := mixedContinuationStore.pop(key) + seg, ok := mixedContinuationStore.pop(key, request) if !ok { return nil, false } @@ -79,3 +247,37 @@ func PopOpenAIContinuationSegment(sessionID typ.SessionID, providerUUID string) } return messages, true } + +// ProviderBetaContinuationStore adapts the shared bounded, single-consume +// mixed continuation store to the Beta-native ToolLoop Stage. +type ProviderBetaContinuationStore struct { + providerUUID string +} + +func NewProviderBetaContinuationStore(providerUUID string) *ProviderBetaContinuationStore { + return &ProviderBetaContinuationStore{providerUUID: providerUUID} +} + +func (s *ProviderBetaContinuationStore) Pop(ctx context.Context, request *anthropic.BetaMessageNewParams) ([]anthropic.BetaMessageParam, bool) { + if s == nil { + return nil, false + } + key := continuationKey(typ.GetSessionID(ctx), s.providerUUID, "anthropic-beta") + segment, ok := mixedContinuationStore.pop(key, request) + if !ok { + return nil, false + } + messages, ok := segment.([]anthropic.BetaMessageParam) + if !ok || len(messages) == 0 { + return nil, false + } + return messages, true +} + +func (s *ProviderBetaContinuationStore) Put(ctx context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) { + if s == nil || len(segment) == 0 { + return + } + key := continuationKey(typ.GetSessionID(ctx), s.providerUUID, "anthropic-beta") + mixedContinuationStore.put(key, append([]anthropic.BetaMessageParam(nil), segment...), externalIDs) +} diff --git a/internal/mcpserver/continuation_store_test.go b/internal/mcpserver/continuation_store_test.go new file mode 100644 index 000000000..73d5e01af --- /dev/null +++ b/internal/mcpserver/continuation_store_test.go @@ -0,0 +1,88 @@ +package mcpserver + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func TestContinuationStoreOnlyConsumesMatchingCurrentToolResults(t *testing.T) { + store := newContinuationStore() + key := continuationKey(typ.SessionID{Source: typ.SessionSourceHeader, Value: "session"}, "provider", "anthropic-beta") + store.put(key, "first", []string{"toolu-first"}) + store.put(key, "second", []string{"toolu-second"}) + + unrelated := &anthropic.BetaMessageNewParams{Messages: []anthropic.BetaMessageParam{ + anthropic.NewBetaUserMessage(anthropic.NewBetaToolResultBlock("toolu-other", "no", false)), + }} + if _, ok := store.pop(key, unrelated); ok { + t.Fatal("unrelated tool result consumed a continuation") + } + + second := &anthropic.BetaMessageNewParams{Messages: []anthropic.BetaMessageParam{ + anthropic.NewBetaUserMessage(anthropic.NewBetaToolResultBlock("toolu-second", "yes", false)), + }} + if got, ok := store.pop(key, second); !ok || got != "second" { + t.Fatalf("matched continuation = %#v, ok=%v", got, ok) + } + if len(store.items[key]) != 1 || store.items[key][0].Segment != "first" { + t.Fatalf("unmatched continuation was removed: %#v", store.items[key]) + } +} + +func TestContinuationStoreRejectsIPFallbackAndSweepsExpiredItems(t *testing.T) { + if key := continuationKey(typ.SessionID{Source: typ.SessionSourceIP, Value: "127.0.0.1"}, "provider", "anthropic-beta"); key != "" { + t.Fatalf("IP fallback produced continuation key %q", key) + } + + store := newContinuationStore() + store.items["expired"] = []continuationItem{{ + Segment: "old", + ExpectedIDs: stringSet([]string{"toolu-old"}), + ExpiresAt: time.Now().Add(-time.Second), + }} + store.put("active", "new", []string{"toolu-new"}) + if _, ok := store.items["expired"]; ok { + t.Fatal("expired continuation was not swept on write") + } +} + +func TestContinuationStoreIsBoundedAndRecognizesOpenAIToolResults(t *testing.T) { + store := newContinuationStore() + for index := 0; index < maxContinuationItems+10; index++ { + store.put("key", index, []string{fmt.Sprintf("call-%d", index)}) + } + if got := store.sizeLocked(); got != maxContinuationItems { + t.Fatalf("continuation count = %d, want %d", got, maxContinuationItems) + } + + request := &openai.ChatCompletionNewParams{Messages: []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("before"), + openai.ToolMessage("result", "call-42"), + }} + if ids := continuationResultIDs(request); len(ids) != 1 { + t.Fatalf("OpenAI trailing result IDs = %#v", ids) + } else if _, ok := ids["call-42"]; !ok { + t.Fatalf("OpenAI trailing result IDs = %#v", ids) + } +} + +func TestProviderBetaContinuationStoreRequiresExplicitSession(t *testing.T) { + store := NewProviderBetaContinuationStore("provider") + segment := []anthropic.BetaMessageParam{{ + Role: anthropic.BetaMessageParamRoleAssistant, + Content: []anthropic.BetaContentBlockParamUnion{anthropic.NewBetaTextBlock("stored")}, + }} + request := &anthropic.BetaMessageNewParams{Messages: []anthropic.BetaMessageParam{ + anthropic.NewBetaUserMessage(anthropic.NewBetaToolResultBlock("toolu-external", "ok", false)), + }} + store.Put(context.Background(), segment, []string{"toolu-external"}) + if _, ok := store.Pop(context.Background(), request); ok { + t.Fatal("empty session persisted a continuation") + } +} diff --git a/internal/mcpserver/format_adapter.go b/internal/mcpserver/format_adapter.go index 668c5898d..a35a9d458 100644 --- a/internal/mcpserver/format_adapter.go +++ b/internal/mcpserver/format_adapter.go @@ -46,6 +46,9 @@ type ToolExecutionResult struct { ToolUseID string Contents []coretool.ToolContent IsError bool + // Dispatched is true once execution crossed into the tool runtime. An + // error after this point may still have produced irreversible effects. + Dispatched bool } // TextContent returns the concatenated text of all text content items. diff --git a/internal/mcpserver/generic_loop_processor.go b/internal/mcpserver/generic_loop_processor.go index 8e6f91ba6..c72672a20 100644 --- a/internal/mcpserver/generic_loop_processor.go +++ b/internal/mcpserver/generic_loop_processor.go @@ -227,7 +227,7 @@ func (p *GenericLoopProcessor) handleMixed(response, req any) (any, error) { return p.adapter.FilterVirtualTools(response, external) } key := continuationKey(typ.GetSessionID(p.ctx), p.provider.UUID, p.adapterID()) - mixedContinuationStore.put(key, segment) + mixedContinuationStore.put(key, segment, externalIDs) // Filter response to only include external tools filteredResponse, err := p.adapter.FilterVirtualTools(response, external) @@ -312,7 +312,7 @@ func (p *GenericLoopProcessor) adapterID() string { func (p *GenericLoopProcessor) applyStoredContinuation(req any) any { sessionID := typ.GetSessionID(p.ctx) key := continuationKey(sessionID, p.provider.UUID, p.adapterID()) - segment, ok := mixedContinuationStore.pop(key) + segment, ok := mixedContinuationStore.pop(key, req) if !ok { return req } diff --git a/internal/mcpserver/generic_stream_interceptor.go b/internal/mcpserver/generic_stream_interceptor.go index 2827e95e8..816ec11e3 100644 --- a/internal/mcpserver/generic_stream_interceptor.go +++ b/internal/mcpserver/generic_stream_interceptor.go @@ -648,7 +648,7 @@ func (i *GenericStreamInterceptor) handleMixed(response, req any) error { return i.adapter.SendFinalMessage(i.c) } key := continuationKey(typ.GetSessionID(i.c.Request.Context()), i.provider.UUID, i.adapterID()) - mixedContinuationStore.put(key, segment) + mixedContinuationStore.put(key, segment, externalIDs) // Send final message (external tools already streamed to client) return i.adapter.SendFinalMessage(i.c) @@ -893,7 +893,7 @@ func (i *GenericStreamInterceptor) adapterID() string { func (i *GenericStreamInterceptor) applyStoredContinuation() { sessionID := typ.GetSessionID(i.c.Request.Context()) key := continuationKey(sessionID, i.provider.UUID, i.adapterID()) - segment, ok := mixedContinuationStore.pop(key) + segment, ok := mixedContinuationStore.pop(key, i.currentReq) if !ok { return } diff --git a/internal/mcpserver/tool_executor.go b/internal/mcpserver/tool_executor.go index 03d944ca3..a68e9bce7 100644 --- a/internal/mcpserver/tool_executor.go +++ b/internal/mcpserver/tool_executor.go @@ -3,6 +3,7 @@ package mcpserver import ( "context" + "github.com/tingly-dev/tingly-box/internal/protocolserver/servertool" coretool "github.com/tingly-dev/tingly-box/internal/tool" ) @@ -14,7 +15,6 @@ type ServerToolExecutor struct { // ToolExecutorServer defines the server methods needed for tool execution type ToolExecutorServer interface { CallMCPToolWithHooks(ctx context.Context, toolName, arguments string, messages []map[string]any) (context.Context, coretool.ToolResult, error) - CallMCPTool(ctx context.Context, toolName, arguments string, messages []map[string]any) (string, error) } func NewServerToolExecutor(s ToolExecutorServer) *ServerToolExecutor { @@ -48,9 +48,10 @@ func (e *ServerToolExecutor) ExecuteToolWithContext( } result := ToolExecutionResult{ - ToolUseID: tool.ID(), - Contents: toolResult.Contents, - IsError: err != nil || toolResult.IsError, + ToolUseID: tool.ID(), + Contents: toolResult.Contents, + IsError: err != nil || toolResult.IsError, + Dispatched: err == nil || servertool.WasDispatched(err), } return nextCtx, result, err } diff --git a/internal/mcpserver/tool_executor_test.go b/internal/mcpserver/tool_executor_test.go new file mode 100644 index 000000000..e6eb34223 --- /dev/null +++ b/internal/mcpserver/tool_executor_test.go @@ -0,0 +1,36 @@ +package mcpserver + +import ( + "context" + "errors" + "testing" + + "github.com/tingly-dev/tingly-box/internal/protocolserver/servertool" + coretool "github.com/tingly-dev/tingly-box/internal/tool" +) + +type dispatchedErrorServer struct { + err error +} + +func (s dispatchedErrorServer) CallMCPToolWithHooks(context.Context, string, string, []map[string]any) (context.Context, coretool.ToolResult, error) { + return context.Background(), coretool.ToolResult{}, &servertool.DispatchError{Err: s.err} +} + +func TestServerToolExecutorPreservesDispatchBoundary(t *testing.T) { + runtimeErr := errors.New("runtime failed") + executor := NewServerToolExecutor(dispatchedErrorServer{err: runtimeErr}) + _, result, err := executor.ExecuteToolWithContext(context.Background(), testTool{id: "toolu-1", name: "lookup"}, nil) + if !errors.Is(err, runtimeErr) || !result.Dispatched || !result.IsError { + t.Fatalf("result=%#v err=%v", result, err) + } +} + +type testTool struct { + id string + name string +} + +func (t testTool) ID() string { return t.id } +func (t testTool) Name() string { return t.name } +func (testTool) Arguments() string { return "{}" } diff --git a/internal/protocolserver/anthropic_message.go b/internal/protocolserver/anthropic_message.go index da952de5e..c5c0affb1 100644 --- a/internal/protocolserver/anthropic_message.go +++ b/internal/protocolserver/anthropic_message.go @@ -13,6 +13,7 @@ import ( "github.com/tingly-dev/tingly-box/internal/protocol" "github.com/tingly-dev/tingly-box/internal/protocolserver/recording" "github.com/tingly-dev/tingly-box/internal/typ" + pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" ) // HandleAnthropicMessages handles Anthropic v1 messages API requests @@ -56,6 +57,7 @@ func (ph *ProtocolHandler) HandleAnthropicMessages(c *gin.Context) { }) return } + ph.rememberProtocolStageOriginalInput(c, scenarioType, bodyBytes) // Determine provider & requestModel var ( @@ -179,12 +181,26 @@ func (ph *ProtocolHandler) AnthropicMessagesV1(c *gin.Context, req *protocol.Ant // pristine request as received (post-vision-proxy, pre-pre-chain); the // winning attempt's provider/model is re-bound per attempt via SetActiveService. var recorder *recording.ProtocolRecorder + var stageRecording *protocolStageRequestRecording if scenarioConfig.IsRecordingEnable() { bs, err := req.MarshalJSON() if err != nil { bs = []byte("{}") } recorder = ph.EnsureProtocolRecorder(c, string(scenarioType), provider, requestModel, ph.getScenarioRecordMode(scenarioType), bs) + if ph.protocolStageRecordingSupportsRule(rule) { + stageRecording = ph.newProtocolStageRequestRecording( + scenarioType, + protocol.TypeAnthropicV1, + protocolStageOriginalInput(c, req.MessageNewParams), + sessionID, + pkgobs.RequestIDFromContext(c.Request.Context()), + ) + } + } + if stageRecording != nil { + enableProtocolStageAttemptTracking(c, stageRecording) + defer stageRecording.finishFromHTTP(c) } // Snapshot a pristine template only when failover is possible; the single @@ -212,7 +228,7 @@ func (ph *ProtocolHandler) AnthropicMessagesV1(c *gin.Context, req *protocol.Ant } areq = cloned } - ph.runAnthropicV1Attempt(c, areq, responseModel, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig, recorder) + ph.runAnthropicV1Attempt(c, areq, responseModel, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig, recorder, stageRecording) }) } @@ -221,7 +237,7 @@ func (ph *ProtocolHandler) AnthropicMessagesV1(c *gin.Context, req *protocol.Ant // pre-transform chain and guardrails, resolve the target API for this provider's // style, transform, and dispatch. Setup failures route through failAttemptSetup // so the orchestrator can advance to the next candidate. -func (ph *ProtocolHandler) runAnthropicV1Attempt(c *gin.Context, req *protocol.AnthropicMessagesRequest, responseModel string, provider *typ.Provider, requestModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig, recorder *recording.ProtocolRecorder) { +func (ph *ProtocolHandler) runAnthropicV1Attempt(c *gin.Context, req *protocol.AnthropicMessagesRequest, responseModel string, provider *typ.Provider, requestModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig, recorder *recording.ProtocolRecorder, stageRecording *protocolStageRequestRecording) { // Resolve dual endpoint: when the provider has an Anthropic-compatible // dual URL configured, route there natively to avoid a transform. provider = provider.ResolveStyle(protocol.APIStyleAnthropic) @@ -242,11 +258,6 @@ func (ph *ProtocolHandler) runAnthropicV1Attempt(c *gin.Context, req *protocol.A return } - scenario := GetTrackingContextScenario(c) - if ph.guardrailsEnabledForScenario(scenario) { - ApplyGuardrailsToAnthropicV1Request(c, ph.currentGuardrailsRuntime(), req.MessageNewParams, requestModel, provider) - } - // Determine target API type for protocol transformation detection target := protocol.TypeAnthropicV1 switch provider.APIStyle { @@ -266,6 +277,30 @@ func (ph *ProtocolHandler) runAnthropicV1Attempt(c *gin.Context, req *protocol.A // Resolve flags with scenario injection and auto-apply for CleanHeader. // (This also applies the custom User-Agent to the request context.) ruleFlags := ResolveRuleFlagsWithScenario(c, rule, scenarioType, scenarioConfig, protocol.TypeAnthropicV1, target, provider) + if ph.tryProtocolStageAnthropicV1( + c, + req, + responseModel, + target, + provider, + requestModel, + rule, + isStreaming, + scenarioConfig, + ruleFlags, + recorder, + stageRecording, + ) { + return + } + + // The Stage path owns Guardrail request and response processing as one + // full-duplex unit. Apply the legacy request mutation only after Stage + // selection declines the entire attempt, avoiding duplicate policy work. + scenario := GetTrackingContextScenario(c) + if ph.guardrailsEnabledForScenario(scenario) { + ApplyGuardrailsToAnthropicV1Request(c, ph.currentGuardrailsRuntime(), req.MessageNewParams, requestModel, provider) + } reqCtx, err := ph.TransformAnthropicV1(c, req, target, provider, isStreaming, recorder, scenarioType, RulePreBaseTransforms(ruleFlags), RulePreVendorTransforms(ruleFlags)) if err != nil { @@ -306,12 +341,26 @@ func (ph *ProtocolHandler) AnthropicMessagesV1Beta(c *gin.Context, req *protocol // Get or create the recorder for dual-stage recording (pristine request body). var recorder *recording.ProtocolRecorder + var stageRecording *protocolStageRequestRecording if scenarioConfig.IsRecordingEnable() { bs, err := req.MarshalJSON() if err != nil { bs = []byte("{}") } recorder = ph.EnsureProtocolRecorder(c, string(scenarioType), provider, requestModel, ph.getScenarioRecordMode(scenarioType), bs) + if ph.protocolStageRecordingSupportsRule(rule) { + stageRecording = ph.newProtocolStageRequestRecording( + scenarioType, + protocol.TypeAnthropicBeta, + protocolStageOriginalInput(c, req.BetaMessageNewParams), + sessionID, + pkgobs.RequestIDFromContext(c.Request.Context()), + ) + } + } + if stageRecording != nil { + enableProtocolStageAttemptTracking(c, stageRecording) + defer stageRecording.finishFromHTTP(c) } // Snapshot a pristine template only when failover is possible. @@ -338,13 +387,13 @@ func (ph *ProtocolHandler) AnthropicMessagesV1Beta(c *gin.Context, req *protocol } areq = cloned } - ph.runAnthropicBetaAttempt(c, areq, responseModel, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig, recorder) + ph.runAnthropicBetaAttempt(c, areq, responseModel, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig, recorder, stageRecording) }) } // runAnthropicBetaAttempt executes the provider-dependent half of an Anthropic // beta request for one failover attempt. See runAnthropicV1Attempt. -func (ph *ProtocolHandler) runAnthropicBetaAttempt(c *gin.Context, req *protocol.AnthropicBetaMessagesRequest, responseModel string, provider *typ.Provider, requestModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig, recorder *recording.ProtocolRecorder) { +func (ph *ProtocolHandler) runAnthropicBetaAttempt(c *gin.Context, req *protocol.AnthropicBetaMessagesRequest, responseModel string, provider *typ.Provider, requestModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig, recorder *recording.ProtocolRecorder, stageRecording *protocolStageRequestRecording) { // Resolve dual endpoint: when the provider has an Anthropic-compatible // dual URL configured, route there natively to avoid a transform. provider = provider.ResolveStyle(protocol.APIStyleAnthropic) @@ -365,12 +414,6 @@ func (ph *ProtocolHandler) runAnthropicBetaAttempt(c *gin.Context, req *protocol return } - // request guardrails - scenario := GetTrackingContextScenario(c) - if ph.guardrailsEnabledForScenario(scenario) { - ApplyGuardrailsToAnthropicV1BetaRequest(c, ph.currentGuardrailsRuntime(), req.BetaMessageNewParams, requestModel, provider) - } - // Determine target API type for protocol transformation detection target := protocol.TypeAnthropicBeta switch provider.APIStyle { @@ -390,6 +433,30 @@ func (ph *ProtocolHandler) runAnthropicBetaAttempt(c *gin.Context, req *protocol // Resolve flags with scenario injection and auto-apply for CleanHeader. // (This also applies the custom User-Agent to the request context.) ruleFlags := ResolveRuleFlagsWithScenario(c, rule, scenarioType, scenarioConfig, protocol.TypeAnthropicBeta, target, provider) + if ph.tryProtocolStageAnthropicBeta( + c, + req, + responseModel, + target, + provider, + requestModel, + rule, + isStreaming, + scenarioConfig, + ruleFlags, + recorder, + stageRecording, + ) { + return + } + + // The Stage path owns Guardrail request and response processing as one + // full-duplex unit. Apply the legacy request mutation only after Stage + // selection declines the entire attempt, avoiding duplicate policy work. + scenario := GetTrackingContextScenario(c) + if ph.guardrailsEnabledForScenario(scenario) { + ApplyGuardrailsToAnthropicV1BetaRequest(c, ph.currentGuardrailsRuntime(), req.BetaMessageNewParams, requestModel, provider) + } reqCtx, err := ph.TransformAnthropicBeta(c, req, target, provider, isStreaming, recorder, scenarioType, RulePreBaseTransforms(ruleFlags), RulePreVendorTransforms(ruleFlags)) if err != nil { diff --git a/internal/protocolserver/failover_dispatch.go b/internal/protocolserver/failover_dispatch.go index d177222eb..798ce0f14 100644 --- a/internal/protocolserver/failover_dispatch.go +++ b/internal/protocolserver/failover_dispatch.go @@ -21,6 +21,8 @@ package protocolserver import ( "bytes" + "context" + "errors" "net/http" "github.com/gin-gonic/gin" @@ -61,6 +63,7 @@ func (ph *ProtocolHandler) handlePreStreamFailure(c *gin.Context, err error, rec // rejected in the prologue, before the gate is installed, so they remain // non-retryable and reach the client unchanged. func (ph *ProtocolHandler) FailAttemptSetup(c *gin.Context, err error) { + observeProtocolStageSetupFailure(c, err) c.JSON(http.StatusInternalServerError, ErrorResponse{ Error: ErrorDetail{ Message: err.Error(), @@ -108,11 +111,35 @@ func isRetryableStatus(status int) bool { // (valid only while uncommitted). type firstChunkGate struct { gin.ResponseWriter - real gin.ResponseWriter - buf bytes.Buffer - hdr http.Header - status int - committed bool + real gin.ResponseWriter + buf bytes.Buffer + hdr http.Header + status int + committed bool + sideEffectsCommitted bool + attemptErr error +} + +func (g *firstChunkGate) MarkAttemptFailed(err error) { + if g != nil && err != nil { + g.attemptErr = err + } +} + +func (g *firstChunkGate) AttemptError() error { + if g == nil { + return nil + } + return g.attemptErr +} + +func markProtocolStageAttemptFailed(c *gin.Context, err error) { + if c == nil || err == nil { + return + } + if gate, ok := c.Writer.(*firstChunkGate); ok { + gate.MarkAttemptFailed(err) + } } func newFirstChunkGate(w gin.ResponseWriter) *firstChunkGate { @@ -202,6 +229,24 @@ func (g *firstChunkGate) Committed() bool { return g.committed } +// SideEffectsCommitted reports whether retrying the attempt could replay an +// already successful server-owned operation, even though no bytes reached the +// client yet. +func (g *firstChunkGate) SideEffectsCommitted() bool { + return g.sideEffectsCommitted +} + +// MarkSideEffectsCommittedIfGate records an irreversible in-process action on +// the active failover gate. It intentionally does not flush buffered output. +func MarkSideEffectsCommittedIfGate(w gin.ResponseWriter) bool { + gate, ok := w.(*firstChunkGate) + if !ok { + return false + } + gate.sideEffectsCommitted = true + return true +} + // CommitFirstChunk is the producer's "first real chunk arrived" signal. // It flushes captured headers + status + buffered body to the real // writer and switches to pass-through. Idempotent. @@ -268,6 +313,7 @@ func (g *firstChunkGate) Discard() { } g.buf.Reset() g.status = 0 + g.attemptErr = nil for k := range g.hdr { delete(g.hdr, k) } @@ -354,6 +400,7 @@ func (ph *ProtocolHandler) DispatchWithPriorityFailover( ) { activeServices := rule.GetActiveServices() if len(activeServices) <= 1 { + setProtocolStageAttempt(c, 1) attempt(initialProvider, initialModel) return } @@ -393,11 +440,34 @@ func (ph *ProtocolHandler) DispatchWithPriorityFailover( rec.SetActiveService(provider, model) } + setProtocolStageAttempt(c, i+1) attempt(provider, model) // A committed gate means the stream's first real chunk reached // the wire — bytes have left the process, retry is impossible. if gate.Committed() { + if attemptErr := gate.AttemptError(); attemptErr != nil { + if isProtocolStageClientCancellation(c, attemptErr) { + fields := failoverLogFields(c, rule, provider, model, serviceID) + fields["stage"] = "failover_committed_cancelled" + fields["attempt"] = i + 1 + fields["active_services"] = len(activeServices) + logrus.WithContext(c.Request.Context()).WithFields(fields).Debug( + "[failover] client cancelled after response commit", + ) + return + } + loadbalance.RecordServiceFailure(rule.UUID, serviceID) + fields := failoverLogFields(c, rule, provider, model, serviceID) + fields["stage"] = "failover_committed_failure" + fields["attempt"] = i + 1 + fields["active_services"] = len(activeServices) + fields["error"] = attemptErr.Error() + logrus.WithContext(c.Request.Context()).WithFields(fields).Warn( + "[failover] stream failed after response commit; retry is impossible", + ) + return + } loadbalance.RecordServiceSuccess(rule.UUID, serviceID) fields := failoverLogFields(c, rule, provider, model, serviceID) fields["stage"] = "failover_success" @@ -408,6 +478,16 @@ func (ph *ProtocolHandler) DispatchWithPriorityFailover( logrus.WithContext(c.Request.Context()).WithFields(fields).Infof("[failover] succeeded on attempt %d with %s/%s", i+1, provider.UUID, model) return } + if gate.SideEffectsCommitted() { + fields := failoverLogFields(c, rule, provider, model, serviceID) + fields["stage"] = "failover_side_effect_boundary" + fields["attempt"] = i + 1 + fields["active_services"] = len(activeServices) + fields["status"] = gate.Status() + logrus.WithContext(c.Request.Context()).WithFields(fields). + Warn("[failover] retry stopped because tool side effects were committed") + return + } status := gate.Status() if !isRetryableStatus(status) { fields := failoverLogFields(c, rule, provider, model, serviceID) @@ -489,6 +569,17 @@ func (ph *ProtocolHandler) DispatchWithPriorityFailover( } } +func isProtocolStageClientCancellation(c *gin.Context, err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.Canceled) { + return true + } + return c != nil && c.Request != nil && c.Request.Context().Err() != nil && + errors.Is(err, c.Request.Context().Err()) +} + func failoverLogFields(c *gin.Context, rule *typ.Rule, provider *typ.Provider, model, serviceID string) logrus.Fields { fields := logrus.Fields{ "service": serviceID, diff --git a/internal/protocolserver/failover_dispatch_test.go b/internal/protocolserver/failover_dispatch_test.go index 0c1953b57..b9b4dce2a 100644 --- a/internal/protocolserver/failover_dispatch_test.go +++ b/internal/protocolserver/failover_dispatch_test.go @@ -1,15 +1,39 @@ package protocolserver import ( + "context" + "errors" "net/http" "net/http/httptest" "testing" + "time" "github.com/gin-gonic/gin" "github.com/tingly-dev/tingly-box/internal/loadbalance" "github.com/tingly-dev/tingly-box/internal/typ" ) +func TestProtocolStageClientCancellationClassification(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + requestContext, cancel := context.WithCancel(context.Background()) + c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(requestContext) + cancel() + + if !isProtocolStageClientCancellation(c, context.Canceled) { + t.Fatal("context cancellation was not classified as client cancellation") + } + deadlineContext, deadlineCancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) + defer deadlineCancel() + c.Request = c.Request.WithContext(deadlineContext) + if !isProtocolStageClientCancellation(c, context.DeadlineExceeded) { + t.Fatal("request context terminal error was not classified as client cancellation") + } + if isProtocolStageClientCancellation(c, errors.New("provider failed")) { + t.Fatal("provider failure was classified as client cancellation") + } +} + func init() { gin.SetMode(gin.TestMode) } @@ -41,6 +65,22 @@ func TestFirstChunkGate_BufferCaptureBeforeCommit(t *testing.T) { } } +func TestProtocolStageAttemptTrackingRequiresActiveRecording(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + setProtocolStageAttempt(c, 2) + if _, exists := c.Get(protocolStageAttemptKey); exists { + t.Fatal("attempt context was created while recording was disabled") + } + + enableProtocolStageAttemptTracking(c) + setProtocolStageAttempt(c, 2) + if got := currentProtocolStageAttempt(c); got != 2 { + t.Fatalf("current attempt = %d, want 2", got) + } +} + func TestFirstChunkGate_CommitFirstChunkFlushesThenPassesThrough(t *testing.T) { rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) @@ -139,6 +179,7 @@ func TestFirstChunkGate_DiscardResetsThenRetry(t *testing.T) { g.Header().Set("X-Try-1", "yes") g.WriteHeader(429) _, _ = g.WriteString("rate limited") + g.MarkAttemptFailed(errors.New("first attempt failed")) g.Discard() if g.Status() != 0 { @@ -150,6 +191,9 @@ func TestFirstChunkGate_DiscardResetsThenRetry(t *testing.T) { if got := g.Header().Get("X-Try-1"); got != "" { t.Fatalf("after Discard header still present: %q", got) } + if g.AttemptError() != nil { + t.Fatalf("after Discard attempt error = %v, want nil", g.AttemptError()) + } // Next attempt succeeds and commits a fresh response. _, _ = g.WriteString(`{"ok":true}`) @@ -166,6 +210,22 @@ func TestFirstChunkGate_DiscardResetsThenRetry(t *testing.T) { } } +func TestFirstChunkGatePreservesCommittedAttemptFailure(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + g := newFirstChunkGate(c.Writer) + c.Writer = g + _, _ = g.WriteString("event: response.failed\n\n") + g.CommitFirstChunk() + + attemptErr := errors.New("response.failed") + markProtocolStageAttemptFailed(c, attemptErr) + + if !errors.Is(g.AttemptError(), attemptErr) { + t.Fatalf("attempt error = %v, want %v", g.AttemptError(), attemptErr) + } +} + func TestFirstChunkGate_DiscardNoopAfterCommit(t *testing.T) { rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) @@ -223,6 +283,29 @@ func TestFirstChunkGate_FlushNoopUntilCommitted(t *testing.T) { } } +func TestFirstChunkGate_SideEffectsStopRetryWithoutFlushing(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + gate := newFirstChunkGate(c.Writer) + + gate.WriteHeader(http.StatusInternalServerError) + _, _ = gate.Write([]byte("later provider round failed")) + if !MarkSideEffectsCommittedIfGate(gate) { + t.Fatal("MarkSideEffectsCommittedIfGate() = false") + } + if !gate.SideEffectsCommitted() { + t.Fatal("side effects were not recorded") + } + if gate.Committed() || recorder.Body.Len() != 0 { + t.Fatalf("marking side effects flushed output: committed=%v body=%q", gate.Committed(), recorder.Body.String()) + } + + gate.CommitIfBuffered() + if recorder.Code != http.StatusInternalServerError || recorder.Body.String() != "later provider round failed" { + t.Fatalf("committed response = %d %q", recorder.Code, recorder.Body.String()) + } +} + func TestIsRetryableStatus(t *testing.T) { cases := []struct { code int diff --git a/internal/protocolserver/guardrails_runtime_ai.go b/internal/protocolserver/guardrails_runtime_ai.go index c4c2492ed..3576a8052 100644 --- a/internal/protocolserver/guardrails_runtime_ai.go +++ b/internal/protocolserver/guardrails_runtime_ai.go @@ -48,10 +48,18 @@ func GuardrailsSupportsScenario(scenario string) bool { // GuardrailsEnabledForScenario centralizes feature-flag checks so protocol // handlers do not repeat scenario/global guardrails gating logic. func GuardrailsEnabledForScenario(cfg *config.Config, runtime *guardrails.Guardrails, scenario string) bool { - if runtime == nil || runtime.PolicyEngine() == nil || cfg == nil || !runtime.IsActive() { + if !GuardrailsSupportsScenario(scenario) { return false } - if !GuardrailsSupportsScenario(scenario) { + return GuardrailsConfiguredForScenario(cfg, runtime, scenario) +} + +// GuardrailsConfiguredForScenario reports whether the runtime and feature flag +// are active without applying the legacy scenario-support list. Protocol Stage +// uses this narrower primitive to opt in new ingress behavior behind --stage +// without declaring the same scenario supported by legacy endpoints. +func GuardrailsConfiguredForScenario(cfg *config.Config, runtime *guardrails.Guardrails, scenario string) bool { + if runtime == nil || runtime.PolicyEngine() == nil || cfg == nil || !runtime.IsActive() { return false } return cfg.GetScenarioFlag(typ.RuleScenario(scenario), config.ExtensionGuardrails) || diff --git a/internal/protocolserver/mcp_hooks.go b/internal/protocolserver/mcp_hooks.go index 3c44c5fa0..33b2cb1e6 100644 --- a/internal/protocolserver/mcp_hooks.go +++ b/internal/protocolserver/mcp_hooks.go @@ -83,7 +83,7 @@ func (ph *ProtocolHandler) buildOpenAIToAnthropicMCPHooks( return stream.ErrMCPStreamContinue } - mcp.StoreOpenAIContinuationSegment(typ.GetSessionID(ctx), providerUUID, segment) + mcp.StoreOpenAIContinuationSegment(typ.GetSessionID(ctx), providerUUID, segment, externalIDs) return nil }, } diff --git a/internal/protocolserver/openai_chat.go b/internal/protocolserver/openai_chat.go index 0483456f2..9b4cc6bbe 100644 --- a/internal/protocolserver/openai_chat.go +++ b/internal/protocolserver/openai_chat.go @@ -12,6 +12,7 @@ import ( "github.com/tingly-dev/tingly-box/internal/protocol" "github.com/tingly-dev/tingly-box/internal/protocol/transform" "github.com/tingly-dev/tingly-box/internal/typ" + pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" ) // HandleOpenAIChatCompletions handles OpenAI v1 chat completion requests @@ -30,6 +31,7 @@ func (ph *ProtocolHandler) HandleOpenAIChatCompletions(c *gin.Context) { }) return } + ph.rememberProtocolStageOriginalInput(c, typ.RuleScenario(scenario), bodyBytes) // Parse OpenAI-style request var req = &protocol.OpenAIChatCompletionRequest{} @@ -150,6 +152,21 @@ func (ph *ProtocolHandler) OpenAIChatCompletion(c *gin.Context, req *protocol.Op // attempt by the failover loop (UpdateTrackingForFailover). SetTrackingContext(c, rule, provider, actualModel, responseModel, isStreaming) + var stageRecording *protocolStageRequestRecording + if scenarioConfig.IsRecordingEnable() && ph.protocolStageRecordingSupportsRule(rule) { + stageRecording = ph.newProtocolStageRequestRecording( + scenarioType, + protocol.TypeOpenAIChat, + protocolStageOriginalInput(c, req.ChatCompletionNewParams), + sessionID, + pkgobs.RequestIDFromContext(c.Request.Context()), + ) + } + if stageRecording != nil { + enableProtocolStageAttemptTracking(c, stageRecording) + defer stageRecording.finishFromHTTP(c) + } + // Snapshot a pristine template only when failover is possible. multi := len(rule.GetActiveServices()) > 1 var template []byte @@ -176,14 +193,14 @@ func (ph *ProtocolHandler) OpenAIChatCompletion(c *gin.Context, req *protocol.Op } areq = cloned } - ph.runOpenAIChatAttempt(c, areq, responseModel, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig) + ph.runOpenAIChatAttempt(c, areq, responseModel, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig, stageRecording) }) } // runOpenAIChatAttempt executes the provider-dependent half of an OpenAI chat // request for one failover attempt. Setup failures route through // failAttemptSetup so the orchestrator can advance to the next candidate. -func (ph *ProtocolHandler) runOpenAIChatAttempt(c *gin.Context, req *protocol.OpenAIChatCompletionRequest, responseModel string, provider *typ.Provider, actualModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig) { +func (ph *ProtocolHandler) runOpenAIChatAttempt(c *gin.Context, req *protocol.OpenAIChatCompletionRequest, responseModel string, provider *typ.Provider, actualModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig, stageRecording *protocolStageRequestRecording) { // Resolve dual endpoint: when the provider has an OpenAI-compatible // dual URL configured, route there natively to avoid a transform. provider = provider.ResolveStyle(protocol.APIStyleOpenAI) @@ -228,6 +245,21 @@ func (ph *ProtocolHandler) runOpenAIChatAttempt(c *gin.Context, req *protocol.Op // (resolveRuleFlagsWithScenario also applies the custom User-Agent to the // request context, so no separate call is needed here.) ruleFlags := ResolveRuleFlagsWithScenario(c, rule, scenarioType, scenarioConfig, protocol.TypeOpenAIChat, target, provider) + if ph.tryProtocolStageOpenAIChat( + c, + req, + responseModel, + target, + provider, + actualModel, + rule, + isStreaming, + scenarioConfig, + ruleFlags, + stageRecording, + ) { + return + } // === Transform via pipeline === reqCtx, err := ph.TransformOpenAIChat(c, req, target, provider, isStreaming, nil, scenarioType, RulePreBaseTransforms(ruleFlags), RulePreVendorTransforms(ruleFlags)) diff --git a/internal/protocolserver/openai_responses.go b/internal/protocolserver/openai_responses.go index 18078752c..53b8e6477 100644 --- a/internal/protocolserver/openai_responses.go +++ b/internal/protocolserver/openai_responses.go @@ -13,6 +13,7 @@ import ( "github.com/tingly-dev/tingly-box/internal/loadbalance" "github.com/tingly-dev/tingly-box/internal/protocol" "github.com/tingly-dev/tingly-box/internal/typ" + pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" ) // HandleResponsesCreate handles POST /v1/responses @@ -30,6 +31,7 @@ func (ph *ProtocolHandler) HandleResponsesCreate(c *gin.Context) { }) return } + ph.rememberProtocolStageOriginalInput(c, typ.RuleScenario(scenario), bodyBytes) // Parse request (minimal parsing for validation) var req = &protocol.ResponseCreateRequest{} @@ -165,6 +167,21 @@ func (ph *ProtocolHandler) ResponsesCreate(c *gin.Context, scenarioType typ.Rule scenarioConfig := ph.deps.Config.GetScenarioConfig(scenarioType) actualModel := string(req.Model) + var stageRecording *protocolStageRequestRecording + if scenarioConfig.IsRecordingEnable() && ph.protocolStageRecordingSupportsRule(rule) { + stageRecording = ph.newProtocolStageRequestRecording( + scenarioType, + protocol.TypeOpenAIResponses, + protocolStageOriginalInput(c, req.ResponseNewParams), + typ.GetSessionID(c.Request.Context()), + pkgobs.RequestIDFromContext(c.Request.Context()), + ) + } + if stageRecording != nil { + enableProtocolStageAttemptTracking(c, stageRecording) + defer stageRecording.finishFromHTTP(c) + } + // Snapshot a pristine template only when failover is possible. The template // is the typed ResponseNewParams (post-vision-proxy — cloned per attempt so // PreprocessInputData and vision proxy are not re-run). @@ -182,14 +199,14 @@ func (ph *ProtocolHandler) ResponsesCreate(c *gin.Context, scenarioType typ.Rule } areq.ResponseNewParams = clonedParams } - ph.runOpenAIResponsesAttempt(c, areq, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig) + ph.runOpenAIResponsesAttempt(c, areq, responseModel, p, retryModel, rule, isStreaming, scenarioType, scenarioConfig, stageRecording) }) } // runOpenAIResponsesAttempt executes the provider-dependent half of an OpenAI // Responses request for one failover attempt. Setup failures route through // failAttemptSetup so the orchestrator can advance to the next candidate. -func (ph *ProtocolHandler) runOpenAIResponsesAttempt(c *gin.Context, req *protocol.ResponseCreateRequest, provider *typ.Provider, actualModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig) { +func (ph *ProtocolHandler) runOpenAIResponsesAttempt(c *gin.Context, req *protocol.ResponseCreateRequest, responseModel string, provider *typ.Provider, actualModel string, rule *typ.Rule, isStreaming bool, scenarioType typ.RuleScenario, scenarioConfig *typ.ScenarioConfig, stageRecording *protocolStageRequestRecording) { // Resolve dual endpoint: when the provider has an OpenAI-compatible // dual URL configured, route there natively to avoid a transform. provider = provider.ResolveStyle(protocol.APIStyleOpenAI) @@ -224,6 +241,22 @@ func (ph *ProtocolHandler) runOpenAIResponsesAttempt(c *gin.Context, req *protoc // Resolve flags with scenario injection, consistent with the chat/v1/beta // handlers (this also applies the custom User-Agent to the request context). ruleFlags := ResolveRuleFlagsWithScenario(c, rule, scenarioType, scenarioConfig, protocol.TypeOpenAIResponses, target, provider) + if ph.tryProtocolStageOpenAIResponses( + c, + req, + responseModel, + target, + provider, + actualModel, + rule, + isStreaming, + scenarioConfig, + ruleFlags, + maxAllowed, + stageRecording, + ) { + return + } reqCtx, err := ph.TransformOpenAIResponses(c, req, target, provider, isStreaming, nil, scenarioType, maxAllowed, RulePreBaseTransforms(ruleFlags), RulePreVendorTransforms(ruleFlags)) if err != nil { ph.FailAttemptSetup(c, fmt.Errorf("Transform failed: %w", err)) diff --git a/internal/protocolserver/protocol_dispatch.go b/internal/protocolserver/protocol_dispatch.go index 113a6c73d..47a2933c1 100644 --- a/internal/protocolserver/protocol_dispatch.go +++ b/internal/protocolserver/protocol_dispatch.go @@ -75,6 +75,7 @@ func (ph *ProtocolHandler) DispatchChainResult( // rule + applied flags are all known, before any response byte is written. if c.GetHeader("X-Tingly-Debug-Routing") == "1" { setProbeUpstreamHeaders(c, reqCtx, rule, provider) + c.Header(protocolPipelineHeader, "legacy") } switch reqCtx.TargetAPI { @@ -179,9 +180,13 @@ func (ph *ProtocolHandler) dispatchOpenAIResponses( // X-Tingly-* response headers, consumed by the probe's captureRoutingRoundTripper. // Gated by the caller on X-Tingly-Debug-Routing so production traffic is untouched. func setProbeUpstreamHeaders(c *gin.Context, reqCtx *transform.TransformContext, rule *typ.Rule, provider *typ.Provider) { - c.Header("X-Tingly-Upstream-API", string(reqCtx.TargetAPI)) + setProbeUpstreamHeadersForTarget(c, reqCtx.TargetAPI, rule, provider) +} + +func setProbeUpstreamHeadersForTarget(c *gin.Context, target protocol.APIType, rule *typ.Rule, provider *typ.Provider) { + c.Header("X-Tingly-Upstream-API", string(target)) if provider != nil { - c.Header("X-Tingly-Upstream-URL", upstreamURLFor(provider, reqCtx.TargetAPI)) + c.Header("X-Tingly-Upstream-URL", upstreamURLFor(provider, target)) } // Synthetic rules (provider probes) carry no meaningful rule identity. if rule != nil && rule.UUID != ProbeSyntheticRuleUUID { @@ -559,7 +564,7 @@ func (ph *ProtocolHandler) dispatchOpenAIChat( actualModel, responseModel := reqCtx.RequestModel, reqCtx.ResponseModel req := reqCtx.Request.(*openai.ChatCompletionNewParams) - if seg, ok := mcp.PopOpenAIContinuationSegment(typ.GetSessionID(c.Request.Context()), provider.UUID); ok { + if seg, ok := mcp.PopOpenAIContinuationSegment(typ.GetSessionID(c.Request.Context()), provider.UUID, req); ok { req.Messages = append(append([]openai.ChatCompletionMessageParamUnion{}, seg...), req.Messages...) } // AlignToolMessagesForOpenAI is already performed by ConsistencyTransform diff --git a/internal/protocolserver/protocol_handler.go b/internal/protocolserver/protocol_handler.go index b11321be0..4a96ac73f 100644 --- a/internal/protocolserver/protocol_handler.go +++ b/internal/protocolserver/protocol_handler.go @@ -23,6 +23,7 @@ import ( "github.com/tingly-dev/tingly-box/internal/loadbalance" mcpruntime "github.com/tingly-dev/tingly-box/internal/mcp/runtime" "github.com/tingly-dev/tingly-box/internal/obs" + "github.com/tingly-dev/tingly-box/internal/protocol" "github.com/tingly-dev/tingly-box/internal/protocolserver/recording" "github.com/tingly-dev/tingly-box/internal/protocolserver/servertool" "github.com/tingly-dev/tingly-box/internal/routing" @@ -41,6 +42,11 @@ import ( type ProtocolHandlerDeps struct { Config *config.Config + // ProtocolStageEnabled is an immutable process-start choice. When true, + // registered and capability-complete protocol paths may use Protocol Stage; + // unsupported paths remain on the legacy pipeline. + ProtocolStageEnabled bool + // TokenTracker records usage to the OTel meter pipeline (may be nil if // OTel setup failed at startup — callers must nil-check). TokenTracker *tracker.TokenTracker @@ -109,7 +115,8 @@ type ProtocolHandlerDeps struct { // files (openai_*.go, anthropic_*.go, protocol_*.go, etc.) will be moved // here in later steps and become methods on *ProtocolHandler. type ProtocolHandler struct { - deps ProtocolHandlerDeps + deps ProtocolHandlerDeps + protocolStageSelector *ProtocolStageSelector // mcpTC caches the stateless MCP chain transforms (see // protocol_transform.go); they depend only on construction-time deps. @@ -118,7 +125,19 @@ type ProtocolHandler struct { // NewHandler constructs the AI Model API handler from its dependencies. func NewHandler(deps ProtocolHandlerDeps) *ProtocolHandler { - return &ProtocolHandler{deps: deps} + handler := &ProtocolHandler{ + deps: deps, + protocolStageSelector: NewProtocolStageSelector(deps.ProtocolStageEnabled), + } + if deps.ProtocolStageEnabled { + logrus.WithFields(logrus.Fields{ + "protocol_pipeline": "stage", + "stage_routes": "anthropic_v1->anthropic_v1,anthropic_v1->anthropic_beta,anthropic_v1->openai_chat,anthropic_v1->openai_responses,anthropic_beta->anthropic_beta,anthropic_beta->openai_chat,anthropic_beta->openai_responses,openai_chat->openai_chat,openai_chat->anthropic_beta,openai_chat->openai_responses,openai_responses->openai_responses,openai_responses->anthropic_beta,openai_responses->openai_chat", + "tool_loop": "tool_loop_anthropic_beta", + "other_routes": "legacy", + }).Info("Protocol Stage mode enabled") + } + return handler } // The methods below are thin wrappers wiring the Deps callbacks to the @@ -135,6 +154,18 @@ func (ph *ProtocolHandler) guardrailsEnabledForScenario(scenario string) bool { return GuardrailsEnabledForScenario(ph.deps.Config, ph.currentGuardrailsRuntime(), scenario) } +// guardrailsEnabledForProtocolStage keeps new ingress support behind --stage +// without broadening legacy scenario behavior. OpenAI Chat and Responses use +// this opt-in gate; legacy OpenAI paths retain their existing behavior. +func (ph *ProtocolHandler) guardrailsEnabledForProtocolStage(scenario string, source protocol.APIType) bool { + switch source { + case protocol.TypeOpenAIChat, protocol.TypeOpenAIResponses: + return GuardrailsConfiguredForScenario(ph.deps.Config, ph.currentGuardrailsRuntime(), scenario) + default: + return ph.guardrailsEnabledForScenario(scenario) + } +} + func (ph *ProtocolHandler) mcpEnabled() bool { return MCPEnabled(ph.deps.Config) } diff --git a/internal/protocolserver/protocol_stage_anthropic_beta.go b/internal/protocolserver/protocol_stage_anthropic_beta.go new file mode 100644 index 000000000..7b6eefc1c --- /dev/null +++ b/internal/protocolserver/protocol_stage_anthropic_beta.go @@ -0,0 +1,521 @@ +package protocolserver + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/anthropicbridge" + protocolguardrail "github.com/tingly-dev/tingly-box/internal/protocol/stage/guardrail" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/transform" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" + "github.com/tingly-dev/tingly-box/internal/protocolserver/recording" + servertransform "github.com/tingly-dev/tingly-box/internal/protocolserver/transform" + "github.com/tingly-dev/tingly-box/internal/typ" + pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" +) + +// tryProtocolStageAnthropicBeta selects an explicitly registered Beta-source +// route for one provider attempt. Anthropic V1 is intentionally not included: +// V1 and Beta remain distinct request, response, and streaming protocols. +func (ph *ProtocolHandler) tryProtocolStageAnthropicBeta( + c *gin.Context, + req *protocol.AnthropicBetaMessagesRequest, + responseModel string, + target protocol.APIType, + provider *typ.Provider, + actualModel string, + rule *typ.Rule, + isStreaming bool, + scenarioConfig *typ.ScenarioConfig, + ruleFlags typ.RuleFlags, + recorder *recording.ProtocolRecorder, + stageRecording *protocolStageRequestRecording, +) bool { + mcpEnabled := ph.mcpEnabled() + if mcpEnabled { + if !ph.shouldUseProtocolStageBetaChain(c, protocol.TypeAnthropicBeta, target, protocolstage.AllBridgeCapabilities) { + return false + } + if ph.deps.MCPRuntime == nil { + logProtocolStageFallback(c, protocol.TypeAnthropicBeta, target, "MCP runtime is unavailable") + return false + } + } else if !ph.shouldUseProtocolStage(c, protocol.TypeAnthropicBeta, target, protocolstage.AllBridgeCapabilities) { + return false + } + if recorder != nil || stageRecording != nil { + if stageRecording == nil { + logProtocolStageFallback(c, protocol.TypeAnthropicBeta, target, "new recording path requires a fully Stage-compatible service set") + return false + } + } + + var requestErr error + defer func() { + markProtocolStageAttemptFailed(c, requestErr) + if stageRecording != nil { + stageRecording.observeAttempt(requestErr) + } + }() + + if c.GetHeader("X-Tingly-Debug-Routing") == "1" { + setProbeUpstreamHeadersForTarget(c, target, rule, provider) + c.Header(protocolPipelineHeader, "stage") + } + + scenarioFlags, clientTransforms := protocolStageAnthropicBetaClientTransforms(scenarioConfig, ruleFlags) + providerTransforms := protocolStageProviderTransforms(target, + []transform.Transform{transform.NewConsistencyTransform(target)}, + RulePreVendorTransforms(ruleFlags), + []transform.Transform{vendorTransformShared}, + ) + terminal, registry, err := ph.protocolStageAnthropicBetaTarget( + target, + provider, + actualModel, + responseModel, + scenarioFlags, + ) + if err != nil { + requestErr = err + ph.FailAttemptSetup(c, err) + return true + } + terminal = requestrecord.ObserveProvider(terminal, stageRecordingRecorder(stageRecording), requestrecord.ExchangeMetadata{ + Attempt: currentProtocolStageAttempt(c), + Provider: provider.Name, + Model: actualModel, + }) + stages := []protocolstage.Stage{ + newProtocolTransformStage( + "client_prepare", + protocol.TypeAnthropicBeta, + provider, + scenarioFlags, + isStreaming, + clientTransforms, + protocolStageTransformOptions(ph, c)..., + ), + } + if ph.guardrailsEnabledForScenario(GetTrackingContextScenario(c)) { + guardrailStage, guardrailErr := protocolguardrail.NewAnthropicBeta(protocolguardrail.AnthropicBetaConfig{ + Runtime: ph.currentGuardrailsRuntime(), + BaseInput: BuildGuardrailsBaseInput( + c, + actualModel, + provider, + guardrailscore.DirectionRequest, + nil, + ), + Observe: protocolStageGuardrailObserver(c), + }) + if guardrailErr != nil { + requestErr = guardrailErr + ph.FailAttemptSetup(c, guardrailErr) + return true + } + stages = append(stages, guardrailStage) + } + if mcpEnabled { + toolLoop, toolLoopErr := ph.newProtocolStageBetaToolLoop(c, provider, HasNativeAdvisorBeta(req)) + if toolLoopErr != nil { + requestErr = toolLoopErr + ph.FailAttemptSetup(c, toolLoopErr) + return true + } + stages = append(stages, toolLoop) + } + stages = append(stages, + newProtocolTransformStage( + "provider_finalize", + target, + provider, + scenarioFlags, + isStreaming, + providerTransforms, + protocolStageTransformOptions(ph, c)..., + ), + ) + endpoint, err := protocolstage.BuildTopology(protocolstage.TopologyConfig{ + Terminal: terminal, + Stages: stages, + ClientProtocol: protocol.TypeAnthropicBeta, + Registry: registry, + RequiredCapabilities: protocolstage.AllBridgeCapabilities, + }) + if err != nil { + requestErr = fmt.Errorf("build Anthropic Beta Protocol Stage topology: %w", err) + ph.FailAttemptSetup(c, requestErr) + return true + } + + logProtocolStageEntry(c, protocol.TypeAnthropicBeta, target, stages, isStreaming) + + call := protocolstage.Call{ + Request: req.BetaMessageNewParams, + Metadata: protocolstage.CallMetadata{ + RequestID: pkgobs.RequestIDFromContext(c.Request.Context()), + }, + } + legacyRecorder := recorder + if stageRecording != nil { + // The canary emits one new-format envelope. The legacy recorder remains + // the rollback path when Stage is disabled or this route falls back. + legacyRecorder = nil + } + if isStreaming { + requestErr = ph.serveProtocolStageAnthropicBetaStream(c, endpoint, call, responseModel, legacyRecorder, stageRecordingRecorder(stageRecording)) + return true + } + requestErr = ph.serveProtocolStageAnthropicBetaComplete(c, endpoint, call, responseModel, provider, actualModel, rule, legacyRecorder, stageRecordingRecorder(stageRecording)) + return true +} + +func protocolStageGuardrailObserver(c *gin.Context) protocolguardrail.Observer { + return func(observation protocolguardrail.Observation) { + entry := logrus.WithContext(c.Request.Context()).WithFields(logrus.Fields{ + "protocol_pipeline": "stage", + "stage": observation.Stage, + "stage_protocol": observation.Protocol, + "guardrail_phase": observation.Phase, + "guardrail_verdict": observation.Decision.Verdict, + }) + if observation.Err != nil { + entry.WithError(observation.Err).Warn("Protocol Stage Guardrail evaluation failed open") + return + } + if observation.Decision.Verdict == protocolguardrail.VerdictBlock { + entry.Debug("Protocol Stage Guardrail changed the response") + } + } +} + +func (ph *ProtocolHandler) protocolStageAnthropicBetaTarget( + target protocol.APIType, + provider *typ.Provider, + actualModel string, + responseModel string, + scenarioFlags *typ.ScenarioFlags, +) (protocolstage.Endpoint, *protocolstage.BridgeRegistry, error) { + disableStreamUsage := scenarioFlags != nil && scenarioFlags.SkipUsage + betaToChat := anthropicbridge.NewBetaToOpenAIChat(anthropicbridge.ChatOptions{ + Compatible: true, + DisableStreamUsage: disableStreamUsage, + ResponseModel: responseModel, + }) + betaToResponses := anthropicbridge.NewBetaToOpenAIResponses(anthropicbridge.ResponsesOptions{ + ResponseModel: responseModel, + }) + registry, err := protocolstage.NewBridgeRegistry( + protocolstage.NewIdentityBridge(protocol.TypeAnthropicBeta), + betaToChat, + betaToResponses, + ) + if err != nil { + return nil, nil, fmt.Errorf("build Anthropic Beta Protocol Stage registry: %w", err) + } + + switch target { + case protocol.TypeAnthropicBeta: + return &anthropicBetaProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeOpenAIChat: + return &openAIChatProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeOpenAIResponses: + return &openAIResponsesProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + default: + return nil, nil, fmt.Errorf("Anthropic Beta Protocol Stage target %q is not implemented", target) + } +} + +func protocolStageAnthropicBetaClientTransforms( + scenarioConfig *typ.ScenarioConfig, + ruleFlags typ.RuleFlags, +) (*typ.ScenarioFlags, []transform.Transform) { + var scenarioFlags *typ.ScenarioFlags + var transforms []transform.Transform + if scenarioConfig != nil { + flags := scenarioConfig.GetDefaultFlags() + scenarioFlags = &flags + if flags.SmartCompact { + transforms = append(transforms, servertransform.NewThinkingCompactTransform(2)) + } + } + transforms = append(transforms, RulePreBaseTransforms(ruleFlags)...) + return scenarioFlags, transforms +} + +func (ph *ProtocolHandler) serveProtocolStageAnthropicBetaComplete( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + provider *typ.Provider, + actualModel string, + rule *typ.Rule, + recorder *recording.ProtocolRecorder, + requestRecorder *requestrecord.Recorder, +) error { + response, err := endpoint.Complete(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.failRequest(c, recorder, err, "Anthropic Beta Protocol Stage provider request failed") + return err + } + preserveProtocolStageSideEffectBoundary(c, nil, response.SideEffectsCommitted) + message, ok := response.Value.(*anthropic.BetaMessage) + if !ok || message == nil { + responseErr := fmt.Errorf("Anthropic Beta Protocol Stage response has type %T", response.Value) + ph.failRequest(c, recorder, responseErr, "Protocol Stage response conversion failed") + return responseErr + } + body, err := protocolStageAnthropicBetaMessageJSON(message, responseModel) + if err != nil { + ph.failRequest(c, recorder, err, "Protocol Stage response conversion failed") + return err + } + captureProtocolStageFinalResponse(c.Request.Context(), requestRecorder, protocol.TypeAnthropicBeta, json.RawMessage(body)) + if response.Usage != nil { + ph.trackUsageWithTokenUsage(c, response.Usage, nil) + } + ph.updateAffinityMessageID(c, rule, string(message.ID)) + message.Model = anthropic.Model(responseModel) + if recorder != nil { + recorder.SetAssembledResponse(message) + recorder.RecordResponse(provider, actualModel) + } + c.Data(http.StatusOK, "application/json; charset=utf-8", body) + return nil +} + +func protocolStageAnthropicBetaMessageJSON(message *anthropic.BetaMessage, responseModel string) ([]byte, error) { + if message == nil { + return nil, fmt.Errorf("Anthropic Beta Protocol Stage response is nil") + } + structured, err := json.Marshal(message) + if err != nil { + return nil, fmt.Errorf("marshal Anthropic Beta response: %w", err) + } + var object map[string]json.RawMessage + if raw := []byte(message.RawJSON()); len(raw) > 0 { + if err := json.Unmarshal(raw, &object); err != nil { + return nil, fmt.Errorf("decode Anthropic Beta raw response: %w", err) + } + } else { + object = make(map[string]json.RawMessage) + } + // The SDK retains the provider wire payload in RawJSON. Merge current + // structured fields over that payload so Stage mutations (Guardrails, + // transforms) reach the client while unknown provider fields survive. + var current map[string]json.RawMessage + if err := json.Unmarshal(structured, ¤t); err != nil { + return nil, fmt.Errorf("decode structured Anthropic Beta response: %w", err) + } + for key, value := range current { + object[key] = value + } + model, err := json.Marshal(responseModel) + if err != nil { + return nil, fmt.Errorf("marshal Anthropic Beta response model: %w", err) + } + object["model"] = model + body, err := json.Marshal(object) + if err != nil { + return nil, fmt.Errorf("encode Anthropic Beta response: %w", err) + } + return body, nil +} + +func (ph *ProtocolHandler) serveProtocolStageAnthropicBetaStream( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + recorder *recording.ProtocolRecorder, + requestRecorder *requestrecord.Recorder, +) error { + flusher, ok := c.Writer.(http.Flusher) + if !ok { + err := fmt.Errorf("Anthropic Beta Protocol Stage streaming is unsupported by this connection") + ph.FailAttemptSetup(c, err) + return err + } + stream, err := endpoint.Stream(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.failRequest(c, recorder, err, "Anthropic Beta Protocol Stage provider stream failed") + return err + } + defer func() { + if closeErr := stream.Close(); closeErr != nil { + logrus.WithContext(c.Request.Context()).Warnf("close Anthropic Beta Protocol Stage stream: %v", closeErr) + } + }() + finalCapture := newProtocolStageFinalStreamCapture(c.Request.Context(), requestRecorder, protocol.TypeAnthropicBeta) + + wrote := false + sawMessageStart := false + sawMessageStop := false + for { + event, nextErr := stream.Next(c.Request.Context()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nextErr, result.SideEffectsCommitted) + if errors.Is(nextErr, context.Canceled) || protocol.IsContextCanceled(nextErr) { + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + return nextErr + } + if result.Usage != nil && result.Usage.HasUsage() { + ph.trackUsageWithTokenUsage(c, result.Usage, nextErr) + } else { + ph.trackUsageFromContext(c, 0, 0, nextErr) + } + if !wrote { + SendErrorResponse(c, nextErr, "Anthropic Beta Protocol Stage provider stream failed") + } else { + protocolstream.MarshalAndSendErrorEvent(c, "Protocol Stage stream terminated", "stream_error", "stream_failed") + flusher.Flush() + } + if recorder != nil { + recorder.RecordError(nextErr) + } + return nextErr + } + + eventType, payload, eventErr := protocolStageAnthropicBetaEventJSON(event.Value, responseModel) + if eventErr != nil { + streamErr := fmt.Errorf("Anthropic Beta Protocol Stage stream emitted %T", event.Value) + preserveProtocolStageSideEffectBoundary(c, eventErr, stream.Result().SideEffectsCommitted) + if !wrote { + ph.FailAttemptSetup(c, errors.Join(streamErr, eventErr)) + } else { + protocolstream.MarshalAndSendErrorEvent(c, "Protocol Stage stream emitted an invalid event", "stream_error", "stream_failed") + flusher.Flush() + } + if recorder != nil { + recorder.RecordError(streamErr) + } + return errors.Join(streamErr, eventErr) + } + finalCapture.add(c.Request.Context(), json.RawMessage(payload)) + if !wrote { + setProtocolStageAnthropicSSEHeaders(c) + wrote = true + } + switch eventType { + case "message_start": + sawMessageStart = true + case "message_stop": + sawMessageStop = true + case "content_block_delta": + protocol.MarkFirstToken(c) + } + c.SSEvent(eventType, string(payload)) + CommitFirstChunkIfGate(c.Writer) + flusher.Flush() + } + + if !wrote { + setProtocolStageAnthropicSSEHeaders(c) + } + if sawMessageStart && !sawMessageStop { + preserveProtocolStageSideEffectBoundary(c, nil, stream.Result().SideEffectsCommitted) + protocolstream.MarshalAndSendErrorEvent(c, "upstream stream ended before completion", "stream_error", "incomplete_stream") + flusher.Flush() + return errors.New("Anthropic Beta Protocol Stage stream ended before message_stop") + } + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nil, result.SideEffectsCommitted) + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + finalCapture.finish(c.Request.Context()) + return nil +} + +func protocolStageAnthropicBetaEventJSON(value any, responseModel string) (string, []byte, error) { + var eventType string + var raw []byte + switch event := value.(type) { + case anthropic.BetaRawMessageStreamEventUnion: + eventType = event.Type + raw = []byte(event.RawJSON()) + if len(raw) == 0 { + var err error + raw, err = json.Marshal(event) + if err != nil { + return "", nil, fmt.Errorf("marshal Anthropic Beta stream event: %w", err) + } + } + case protocolstream.AnthropicEvent: + eventType = event.Type + var err error + raw, err = json.Marshal(event.Data) + if err != nil { + return "", nil, fmt.Errorf("marshal converted Anthropic Beta stream event: %w", err) + } + default: + return "", nil, fmt.Errorf("unsupported Anthropic Beta stream event %T", value) + } + if eventType == "" { + return "", nil, fmt.Errorf("Anthropic Beta stream event has empty type") + } + if eventType != "message_start" { + return eventType, raw, nil + } + var object map[string]json.RawMessage + if err := json.Unmarshal(raw, &object); err != nil { + return "", nil, fmt.Errorf("decode Anthropic Beta message_start: %w", err) + } + var message map[string]json.RawMessage + if err := json.Unmarshal(object["message"], &message); err != nil { + return "", nil, fmt.Errorf("decode Anthropic Beta message_start.message: %w", err) + } + model, err := json.Marshal(responseModel) + if err != nil { + return "", nil, fmt.Errorf("marshal Anthropic Beta stream response model: %w", err) + } + message["model"] = model + object["message"], err = json.Marshal(message) + if err != nil { + return "", nil, fmt.Errorf("encode Anthropic Beta message_start.message: %w", err) + } + payload, err := json.Marshal(object) + if err != nil { + return "", nil, fmt.Errorf("encode Anthropic Beta message_start: %w", err) + } + return eventType, payload, nil +} + +func setProtocolStageAnthropicSSEHeaders(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("Access-Control-Allow-Origin", "*") + c.Header("Access-Control-Allow-Headers", "Cache-Control") +} diff --git a/internal/protocolserver/protocol_stage_anthropic_beta_recording_test.go b/internal/protocolserver/protocol_stage_anthropic_beta_recording_test.go new file mode 100644 index 000000000..0a5b066dc --- /dev/null +++ b/internal/protocolserver/protocol_stage_anthropic_beta_recording_test.go @@ -0,0 +1,161 @@ +package protocolserver + +import ( + "context" + "encoding/json" + "io" + "net/http/httptest" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func TestProtocolStageAnthropicBetaCompleteRecordsProviderAndFinalBoundaries(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("POST", "/v1/messages?beta=true", nil) + recorder := newBetaBoundaryRecorder(t) + providerMessage := betaMessageFromJSON(t, `{ + "id":"provider-message","type":"message","role":"assistant","content":[], + "model":"provider-model","stop_reason":"end_turn","stop_sequence":null, + "usage":{"input_tokens":1,"output_tokens":1} + }`) + terminal := &betaBoundaryEndpoint{ + completeResponse: &protocolstage.Response{Value: providerMessage}, + } + endpoint := requestrecord.ObserveProvider(terminal, recorder, requestrecord.ExchangeMetadata{ + Attempt: 1, + Provider: "provider", + Model: "provider-model", + }) + + err := (&ProtocolHandler{}).serveProtocolStageAnthropicBetaComplete( + c, + endpoint, + protocolstage.Call{Request: map[string]any{"model": "provider-model"}}, + "public-model", + &typ.Provider{Name: "provider"}, + "provider-model", + &typ.Rule{}, + nil, + recorder, + ) + require.NoError(t, err) + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + require.Contains(t, string(completed.ProviderExchanges[0].Response.Body), `"model":"provider-model"`) + require.NotNil(t, completed.FinalResponse) + require.Contains(t, string(completed.FinalResponse.Body), `"model":"public-model"`) +} + +func TestProtocolStageAnthropicBetaStreamRecordsProviderAndFinalBoundaries(t *testing.T) { + gin.SetMode(gin.TestMode) + response := httptest.NewRecorder() + c, _ := gin.CreateTestContext(response) + gate := newFirstChunkGate(c.Writer) + c.Writer = gate + c.Request = httptest.NewRequest("POST", "/v1/messages?beta=true", nil) + recorder := newBetaBoundaryRecorder(t) + events := []protocolstage.Event{ + {Value: betaStreamEventFromJSON(t, `{ + "type":"message_start","message":{"id":"provider-stream","type":"message","role":"assistant", + "content":[],"model":"provider-model","stop_reason":null,"stop_sequence":null, + "usage":{"input_tokens":1,"output_tokens":0}} + }`)}, + {Value: betaStreamEventFromJSON(t, `{"type":"message_stop"}`)}, + } + terminal := &betaBoundaryEndpoint{streamEvents: events} + endpoint := requestrecord.ObserveProvider(terminal, recorder, requestrecord.ExchangeMetadata{ + Attempt: 1, + Provider: "provider", + Model: "provider-model", + }) + + err := (&ProtocolHandler{}).serveProtocolStageAnthropicBetaStream( + c, + endpoint, + protocolstage.Call{Request: map[string]any{"model": "provider-model"}}, + "public-model", + nil, + recorder, + ) + require.NoError(t, err) + require.True(t, gate.Committed(), "the first valid Stage event must commit the failover gate") + require.NotEmpty(t, response.Body.String(), "committed Stage events must reach the real writer") + + completed, first := recorder.Finish(nil) + require.True(t, first) + require.Len(t, completed.ProviderExchanges, 1) + require.Contains(t, string(completed.ProviderExchanges[0].Response.Body), `"model":"provider-model"`) + require.NotNil(t, completed.FinalResponse) + require.Contains(t, string(completed.FinalResponse.Body), `"model":"public-model"`) +} + +func newBetaBoundaryRecorder(t *testing.T) *requestrecord.Recorder { + t.Helper() + recorder, err := requestrecord.New(requestrecord.Config{ + Enabled: true, + RequestID: "request-id", + InputProtocol: protocol.TypeAnthropicBeta, + Input: map[string]any{"model": "client-model"}, + }) + require.NoError(t, err) + return recorder +} + +func betaMessageFromJSON(t *testing.T, raw string) *anthropic.BetaMessage { + t.Helper() + var message anthropic.BetaMessage + require.NoError(t, json.Unmarshal([]byte(raw), &message)) + return &message +} + +func betaStreamEventFromJSON(t *testing.T, raw string) anthropic.BetaRawMessageStreamEventUnion { + t.Helper() + var event anthropic.BetaRawMessageStreamEventUnion + require.NoError(t, json.Unmarshal([]byte(raw), &event)) + return event +} + +type betaBoundaryEndpoint struct { + completeResponse *protocolstage.Response + streamEvents []protocolstage.Event +} + +func (*betaBoundaryEndpoint) Protocol() protocol.APIType { return protocol.TypeAnthropicBeta } + +func (e *betaBoundaryEndpoint) Complete(context.Context, protocolstage.Call) (*protocolstage.Response, error) { + return e.completeResponse, nil +} + +func (e *betaBoundaryEndpoint) Stream(context.Context, protocolstage.Call) (protocolstage.EventStream, error) { + return &betaBoundaryStream{events: e.streamEvents}, nil +} + +type betaBoundaryStream struct { + events []protocolstage.Event + index int +} + +func (s *betaBoundaryStream) Next(context.Context) (protocolstage.Event, error) { + if s.index >= len(s.events) { + return protocolstage.Event{}, io.EOF + } + event := s.events[s.index] + s.index++ + return event, nil +} + +func (*betaBoundaryStream) Close() error { return nil } + +func (*betaBoundaryStream) Result() protocolstage.StreamResult { + return protocolstage.StreamResult{} +} diff --git a/internal/protocolserver/protocol_stage_anthropic_v1.go b/internal/protocolserver/protocol_stage_anthropic_v1.go new file mode 100644 index 000000000..1f62dc5fc --- /dev/null +++ b/internal/protocolserver/protocol_stage_anthropic_v1.go @@ -0,0 +1,608 @@ +package protocolserver + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "sync" + + "github.com/anthropics/anthropic-sdk-go" + anthropicstream "github.com/anthropics/anthropic-sdk-go/packages/ssestream" + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/anthropicbridge" + protocolguardrail "github.com/tingly-dev/tingly-box/internal/protocol/stage/guardrail" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/transform" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" + "github.com/tingly-dev/tingly-box/internal/protocolserver/forwarding" + "github.com/tingly-dev/tingly-box/internal/protocolserver/recording" + servertransform "github.com/tingly-dev/tingly-box/internal/protocolserver/transform" + "github.com/tingly-dev/tingly-box/internal/typ" + pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" +) + +// tryProtocolStageAnthropicV1 selects only explicitly registered V1-source +// routes. MCP and Guardrail requests promote to the Beta working protocol; +// plain requests retain their existing concrete provider protocol. +func (ph *ProtocolHandler) tryProtocolStageAnthropicV1( + c *gin.Context, + req *protocol.AnthropicMessagesRequest, + responseModel string, + target protocol.APIType, + provider *typ.Provider, + actualModel string, + rule *typ.Rule, + isStreaming bool, + scenarioConfig *typ.ScenarioConfig, + ruleFlags typ.RuleFlags, + recorder *recording.ProtocolRecorder, + stageRecording *protocolStageRequestRecording, +) bool { + mcpEnabled := ph.mcpEnabled() + guardrailsEnabled := ph.guardrailsEnabledForScenario(GetTrackingContextScenario(c)) + usesBetaStages := mcpEnabled || guardrailsEnabled + stageTarget := target + if usesBetaStages && target == protocol.TypeAnthropicV1 { + stageTarget = protocol.TypeAnthropicBeta + } + if usesBetaStages { + if !ph.shouldUseProtocolStageBetaChain(c, protocol.TypeAnthropicV1, stageTarget, protocolstage.AllBridgeCapabilities) { + return false + } + if mcpEnabled && ph.deps.MCPRuntime == nil { + logProtocolStageFallback(c, protocol.TypeAnthropicV1, stageTarget, "MCP runtime is unavailable") + return false + } + } else if !ph.shouldUseProtocolStage(c, protocol.TypeAnthropicV1, stageTarget, protocolstage.AllBridgeCapabilities) { + return false + } + if recorder != nil || stageRecording != nil { + if stageRecording == nil { + logProtocolStageFallback(c, protocol.TypeAnthropicV1, target, "new recording path requires a fully Stage-compatible service set") + return false + } + } + + var requestErr error + defer func() { + markProtocolStageAttemptFailed(c, requestErr) + if stageRecording != nil { + stageRecording.observeAttempt(requestErr) + } + }() + + if c.GetHeader("X-Tingly-Debug-Routing") == "1" { + setProbeUpstreamHeadersForTarget(c, stageTarget, rule, provider) + c.Header(protocolPipelineHeader, "stage") + } + + scenarioFlags, clientTransforms := protocolStageAnthropicV1ClientTransforms(scenarioConfig, ruleFlags) + providerTransforms := protocolStageProviderTransforms(stageTarget, + []transform.Transform{transform.NewConsistencyTransform(stageTarget)}, + RulePreVendorTransforms(ruleFlags), + []transform.Transform{vendorTransformShared}, + ) + terminal, registry, err := ph.protocolStageAnthropicV1Target(stageTarget, provider, actualModel, responseModel, scenarioFlags) + if err != nil { + requestErr = err + ph.FailAttemptSetup(c, err) + return true + } + terminal = requestrecord.ObserveProvider(terminal, stageRecordingRecorder(stageRecording), requestrecord.ExchangeMetadata{ + Attempt: currentProtocolStageAttempt(c), + Provider: provider.Name, + Model: actualModel, + }) + stages := []protocolstage.Stage{ + newProtocolTransformStage( + "client_prepare", + protocol.TypeAnthropicV1, + provider, + scenarioFlags, + isStreaming, + clientTransforms, + protocolStageTransformOptions(ph, c)..., + ), + } + if guardrailsEnabled { + // Promote only the V1 request into the Beta working protocol and + // reuse the Beta-native Guardrail. Outward V1 response compatibility + // remains the existing permissive projection rather than a new strict + // Beta-to-V1 contract. + guardrailStage, guardrailErr := protocolguardrail.NewAnthropicBeta(protocolguardrail.AnthropicBetaConfig{ + Runtime: ph.currentGuardrailsRuntime(), + BaseInput: BuildGuardrailsBaseInput( + c, + actualModel, + provider, + guardrailscore.DirectionRequest, + nil, + ), + Observe: protocolStageGuardrailObserver(c), + }) + if guardrailErr != nil { + requestErr = guardrailErr + ph.FailAttemptSetup(c, guardrailErr) + return true + } + stages = append(stages, guardrailStage) + } + if mcpEnabled { + toolLoop, toolLoopErr := ph.newProtocolStageBetaToolLoop(c, provider, false) + if toolLoopErr != nil { + requestErr = toolLoopErr + ph.FailAttemptSetup(c, toolLoopErr) + return true + } + stages = append(stages, toolLoop) + } + stages = append(stages, + newProtocolTransformStage( + "provider_finalize", + stageTarget, + provider, + scenarioFlags, + isStreaming, + providerTransforms, + protocolStageTransformOptions(ph, c)..., + ), + ) + endpoint, err := protocolstage.BuildTopology(protocolstage.TopologyConfig{ + Terminal: terminal, + Stages: stages, + ClientProtocol: protocol.TypeAnthropicV1, + Registry: registry, + RequiredCapabilities: protocolstage.AllBridgeCapabilities, + }) + if err != nil { + requestErr = fmt.Errorf("build Anthropic V1 Protocol Stage topology: %w", err) + ph.FailAttemptSetup(c, requestErr) + return true + } + + logProtocolStageEntry(c, protocol.TypeAnthropicV1, stageTarget, stages, isStreaming) + + call := protocolstage.Call{ + Request: req.MessageNewParams, + Metadata: protocolstage.CallMetadata{ + RequestID: pkgobs.RequestIDFromContext(c.Request.Context()), + }, + } + legacyRecorder := recorder + if stageRecording != nil { + legacyRecorder = nil + } + if isStreaming { + requestErr = ph.serveProtocolStageAnthropicV1Stream(c, endpoint, call, responseModel, legacyRecorder, stageRecordingRecorder(stageRecording)) + return true + } + requestErr = ph.serveProtocolStageAnthropicV1Complete(c, endpoint, call, responseModel, provider, actualModel, rule, legacyRecorder, stageRecordingRecorder(stageRecording)) + return true +} + +func (ph *ProtocolHandler) protocolStageAnthropicV1Target( + target protocol.APIType, + provider *typ.Provider, + actualModel string, + responseModel string, + scenarioFlags *typ.ScenarioFlags, +) (protocolstage.Endpoint, *protocolstage.BridgeRegistry, error) { + disableStreamUsage := scenarioFlags != nil && scenarioFlags.SkipUsage + registry, err := protocolstage.NewBridgeRegistry( + protocolstage.NewIdentityBridge(protocol.TypeAnthropicV1), + protocolstage.NewIdentityBridge(protocol.TypeAnthropicBeta), + anthropicbridge.NewV1ToBeta(), + anthropicbridge.NewV1ToOpenAIChat(anthropicbridge.ChatOptions{ + Compatible: true, + DisableStreamUsage: disableStreamUsage, + ResponseModel: responseModel, + }), + anthropicbridge.NewV1ToOpenAIResponses(anthropicbridge.ResponsesOptions{ + ResponseModel: responseModel, + }), + anthropicbridge.NewBetaToOpenAIChat(anthropicbridge.ChatOptions{ + Compatible: true, + DisableStreamUsage: disableStreamUsage, + ResponseModel: responseModel, + }), + anthropicbridge.NewBetaToOpenAIResponses(anthropicbridge.ResponsesOptions{ + ResponseModel: responseModel, + }), + ) + if err != nil { + return nil, nil, fmt.Errorf("build Anthropic V1 Protocol Stage registry: %w", err) + } + switch target { + case protocol.TypeAnthropicV1: + return &anthropicV1ProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeAnthropicBeta: + return &anthropicBetaProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeOpenAIChat: + return &openAIChatProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeOpenAIResponses: + return &openAIResponsesProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + default: + return nil, nil, fmt.Errorf("Anthropic V1 Protocol Stage target %q is not implemented", target) + } +} + +func protocolStageAnthropicV1ClientTransforms( + scenarioConfig *typ.ScenarioConfig, + ruleFlags typ.RuleFlags, +) (*typ.ScenarioFlags, []transform.Transform) { + var scenarioFlags *typ.ScenarioFlags + var transforms []transform.Transform + if scenarioConfig != nil { + flags := scenarioConfig.GetDefaultFlags() + scenarioFlags = &flags + if flags.SmartCompact { + transforms = append(transforms, servertransform.NewThinkingCompactTransform(2)) + } + } + transforms = append(transforms, RulePreBaseTransforms(ruleFlags)...) + return scenarioFlags, transforms +} + +type anthropicV1ProviderEndpoint struct { + ph *ProtocolHandler + provider *typ.Provider + model string +} + +func (*anthropicV1ProviderEndpoint) Protocol() protocol.APIType { return protocol.TypeAnthropicV1 } + +func (e *anthropicV1ProviderEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + request, err := protocolStageAnthropicV1Request(call.Request) + if err != nil { + return nil, err + } + wrapper := e.ph.deps.ClientPool.GetAnthropicClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + message, cancel, err := forwarding.ForwardAnthropicV1(fc, wrapper, request) + if cancel != nil { + defer cancel() + } + if err != nil { + return nil, err + } + return &protocolstage.Response{ + Value: message, + Usage: protocolusage.FromAnthropicMessage(message.Usage), + Model: e.model, + }, nil +} + +func (e *anthropicV1ProviderEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + request, err := protocolStageAnthropicV1Request(call.Request) + if err != nil { + return nil, err + } + wrapper := e.ph.deps.ClientPool.GetAnthropicClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + stream, cancel, err := forwarding.ForwardAnthropicV1Stream(fc, wrapper, request) + if err != nil { + if cancel != nil { + cancel() + } + return nil, err + } + return &anthropicV1ProviderStream{ + stream: stream, + cancel: cancel, + model: e.model, + usage: protocolusage.NewAnthropicAccumulator(), + }, nil +} + +func protocolStageAnthropicV1Request(value any) (*anthropic.MessageNewParams, error) { + request, ok := value.(*anthropic.MessageNewParams) + if !ok || request == nil { + return nil, &protocolStageSetupError{err: fmt.Errorf("Anthropic V1 provider endpoint received %T", value)} + } + return request, nil +} + +type anthropicV1ProviderStream struct { + stream *anthropicstream.Stream[anthropic.MessageStreamEventUnion] + cancel context.CancelFunc + model string + usage *protocolusage.AnthropicAccumulator + + closeOnce sync.Once + closeErr error +} + +func (s *anthropicV1ProviderStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + if s.stream == nil { + return protocolstage.Event{}, fmt.Errorf("Anthropic V1 provider stream is nil") + } + if !s.stream.Next() { + if err := s.stream.Err(); err != nil { + return protocolstage.Event{}, err + } + return protocolstage.Event{}, io.EOF + } + event := s.stream.Current() + if s.usage != nil { + s.usage.Consume(&event) + } + return protocolstage.Event{Value: event}, nil +} + +func (s *anthropicV1ProviderStream) Close() error { + s.closeOnce.Do(func() { + if s.stream != nil { + s.closeErr = s.stream.Close() + } + if s.cancel != nil { + s.cancel() + } + }) + return s.closeErr +} + +func (s *anthropicV1ProviderStream) Result() protocolstage.StreamResult { + var usage *protocol.TokenUsage + if s.usage != nil && s.usage.HasUsage() { + usage = s.usage.Result() + } + return protocolstage.StreamResult{Usage: usage, Model: s.model} +} + +func (ph *ProtocolHandler) serveProtocolStageAnthropicV1Complete( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + provider *typ.Provider, + actualModel string, + rule *typ.Rule, + recorder *recording.ProtocolRecorder, + requestRecorder *requestrecord.Recorder, +) error { + response, err := endpoint.Complete(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.failRequest(c, recorder, err, "Anthropic V1 Protocol Stage provider request failed") + return err + } + preserveProtocolStageSideEffectBoundary(c, nil, response.SideEffectsCommitted) + message, ok := response.Value.(*anthropic.Message) + if !ok || message == nil { + responseErr := fmt.Errorf("Anthropic V1 Protocol Stage response has type %T", response.Value) + ph.failRequest(c, recorder, responseErr, "Protocol Stage response conversion failed") + return responseErr + } + body, err := protocolStageAnthropicV1MessageJSON(message, responseModel) + if err != nil { + ph.failRequest(c, recorder, err, "Protocol Stage response conversion failed") + return err + } + captureProtocolStageFinalResponse(c.Request.Context(), requestRecorder, protocol.TypeAnthropicV1, json.RawMessage(body)) + if response.Usage != nil { + ph.trackUsageWithTokenUsage(c, response.Usage, nil) + } + ph.updateAffinityMessageID(c, rule, string(message.ID)) + message.Model = anthropic.Model(responseModel) + if recorder != nil { + recorder.SetAssembledResponse(message) + recorder.RecordResponse(provider, actualModel) + } + c.Data(http.StatusOK, "application/json; charset=utf-8", body) + return nil +} + +func protocolStageAnthropicV1MessageJSON(message *anthropic.Message, responseModel string) ([]byte, error) { + if message == nil { + return nil, fmt.Errorf("Anthropic V1 Protocol Stage response is nil") + } + raw := []byte(message.RawJSON()) + if len(raw) == 0 { + var err error + raw, err = json.Marshal(message) + if err != nil { + return nil, fmt.Errorf("marshal Anthropic V1 response: %w", err) + } + } + var object map[string]json.RawMessage + if err := json.Unmarshal(raw, &object); err != nil { + return nil, fmt.Errorf("decode Anthropic V1 response: %w", err) + } + model, err := json.Marshal(responseModel) + if err != nil { + return nil, fmt.Errorf("marshal Anthropic V1 response model: %w", err) + } + object["model"] = model + body, err := json.Marshal(object) + if err != nil { + return nil, fmt.Errorf("encode Anthropic V1 response: %w", err) + } + return body, nil +} + +func (ph *ProtocolHandler) serveProtocolStageAnthropicV1Stream( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + recorder *recording.ProtocolRecorder, + requestRecorder *requestrecord.Recorder, +) error { + flusher, ok := c.Writer.(http.Flusher) + if !ok { + err := fmt.Errorf("Anthropic V1 Protocol Stage streaming is unsupported by this connection") + ph.FailAttemptSetup(c, err) + return err + } + stream, err := endpoint.Stream(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.failRequest(c, recorder, err, "Anthropic V1 Protocol Stage provider stream failed") + return err + } + defer func() { + if closeErr := stream.Close(); closeErr != nil { + logrus.WithContext(c.Request.Context()).Warnf("close Anthropic V1 Protocol Stage stream: %v", closeErr) + } + }() + finalCapture := newProtocolStageFinalStreamCapture(c.Request.Context(), requestRecorder, protocol.TypeAnthropicV1) + + wrote := false + sawMessageStart := false + sawMessageStop := false + for { + event, nextErr := stream.Next(c.Request.Context()) + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nextErr, result.SideEffectsCommitted) + if errors.Is(nextErr, context.Canceled) || protocol.IsContextCanceled(nextErr) { + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + return nextErr + } + if result.Usage != nil && result.Usage.HasUsage() { + ph.trackUsageWithTokenUsage(c, result.Usage, nextErr) + } else { + ph.trackUsageFromContext(c, 0, 0, nextErr) + } + if !wrote { + SendErrorResponse(c, nextErr, "Anthropic V1 Protocol Stage provider stream failed") + } else { + protocolstream.MarshalAndSendErrorEvent(c, "Protocol Stage stream terminated", "stream_error", "stream_failed") + flusher.Flush() + } + if recorder != nil { + recorder.RecordError(nextErr) + } + return nextErr + } + + eventType, payload, eventErr := protocolStageAnthropicV1EventJSON(event.Value, responseModel) + if eventErr != nil { + streamErr := fmt.Errorf("Anthropic V1 Protocol Stage stream emitted %T", event.Value) + preserveProtocolStageSideEffectBoundary(c, eventErr, stream.Result().SideEffectsCommitted) + if !wrote { + ph.FailAttemptSetup(c, errors.Join(streamErr, eventErr)) + } else { + protocolstream.MarshalAndSendErrorEvent(c, "Protocol Stage stream emitted an invalid event", "stream_error", "stream_failed") + flusher.Flush() + } + if recorder != nil { + recorder.RecordError(streamErr) + } + return errors.Join(streamErr, eventErr) + } + finalCapture.add(c.Request.Context(), json.RawMessage(payload)) + if !wrote { + setProtocolStageAnthropicSSEHeaders(c) + wrote = true + } + switch eventType { + case "message_start": + sawMessageStart = true + case "message_stop": + sawMessageStop = true + case "content_block_delta": + protocol.MarkFirstToken(c) + } + c.SSEvent(eventType, string(payload)) + CommitFirstChunkIfGate(c.Writer) + flusher.Flush() + } + + if !wrote { + setProtocolStageAnthropicSSEHeaders(c) + } + if sawMessageStart && !sawMessageStop { + preserveProtocolStageSideEffectBoundary(c, nil, stream.Result().SideEffectsCommitted) + protocolstream.MarshalAndSendErrorEvent(c, "upstream stream ended before completion", "stream_error", "incomplete_stream") + flusher.Flush() + return errors.New("Anthropic V1 Protocol Stage stream ended before message_stop") + } + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nil, result.SideEffectsCommitted) + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + finalCapture.finish(c.Request.Context()) + return nil +} + +func protocolStageAnthropicV1EventJSON(value any, responseModel string) (string, []byte, error) { + var eventType string + var raw []byte + switch event := value.(type) { + case anthropic.MessageStreamEventUnion: + eventType = event.Type + raw = []byte(event.RawJSON()) + if len(raw) == 0 { + var err error + raw, err = json.Marshal(event) + if err != nil { + return "", nil, fmt.Errorf("marshal Anthropic V1 stream event: %w", err) + } + } + case protocolstream.AnthropicEvent: + eventType = event.Type + var err error + raw, err = json.Marshal(event.Data) + if err != nil { + return "", nil, fmt.Errorf("marshal converted Anthropic V1 stream event: %w", err) + } + default: + return "", nil, fmt.Errorf("unsupported Anthropic V1 stream event %T", value) + } + if eventType == "" { + return "", nil, fmt.Errorf("Anthropic V1 stream event has empty type") + } + if eventType != "message_start" { + return eventType, raw, nil + } + var object map[string]json.RawMessage + if err := json.Unmarshal(raw, &object); err != nil { + return "", nil, fmt.Errorf("decode Anthropic V1 message_start: %w", err) + } + var message map[string]json.RawMessage + if err := json.Unmarshal(object["message"], &message); err != nil { + return "", nil, fmt.Errorf("decode Anthropic V1 message_start.message: %w", err) + } + model, err := json.Marshal(responseModel) + if err != nil { + return "", nil, fmt.Errorf("marshal Anthropic V1 stream response model: %w", err) + } + message["model"] = model + object["message"], err = json.Marshal(message) + if err != nil { + return "", nil, fmt.Errorf("encode Anthropic V1 message_start.message: %w", err) + } + payload, err := json.Marshal(object) + if err != nil { + return "", nil, fmt.Errorf("encode Anthropic V1 message_start: %w", err) + } + return eventType, payload, nil +} diff --git a/internal/protocolserver/protocol_stage_openai_chat_endpoint.go b/internal/protocolserver/protocol_stage_openai_chat_endpoint.go new file mode 100644 index 000000000..aada9b3a8 --- /dev/null +++ b/internal/protocolserver/protocol_stage_openai_chat_endpoint.go @@ -0,0 +1,125 @@ +package protocolserver + +import ( + "context" + "fmt" + "io" + "sync" + + "github.com/openai/openai-go/v3" + openaistream "github.com/openai/openai-go/v3/packages/ssestream" + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/request" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" + "github.com/tingly-dev/tingly-box/internal/protocolserver/forwarding" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +// openAIChatProviderEndpoint is the transport-free Chat provider terminal used +// by protocol paths whose selected provider exposes Chat Completions. +type openAIChatProviderEndpoint struct { + ph *ProtocolHandler + provider *typ.Provider + model string +} + +func (*openAIChatProviderEndpoint) Protocol() protocol.APIType { return protocol.TypeOpenAIChat } + +func (e *openAIChatProviderEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + chatRequest, err := protocolStageOpenAIChatRequest(call.Request) + if err != nil { + return nil, err + } + request.CleanupOpenaiFields(chatRequest) + wrapper := e.ph.deps.ClientPool.GetOpenAIClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + completion, cancel, err := forwarding.ForwardOpenAIChat(fc, wrapper, chatRequest) + if cancel != nil { + defer cancel() + } + if err != nil { + return nil, err + } + return &protocolstage.Response{ + Value: completion, + Usage: protocolusage.FromOpenAIChatCompletion(completion.Usage), + Model: e.model, + }, nil +} + +func (e *openAIChatProviderEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + chatRequest, err := protocolStageOpenAIChatRequest(call.Request) + if err != nil { + return nil, err + } + request.CleanupOpenaiFields(chatRequest) + wrapper := e.ph.deps.ClientPool.GetOpenAIClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + stream, cancel, err := forwarding.ForwardOpenAIChatStream(fc, wrapper, chatRequest) + if err != nil { + if cancel != nil { + cancel() + } + return nil, err + } + return &openAIChatProviderStream{ + stream: stream, + cancel: cancel, + model: e.model, + }, nil +} + +func protocolStageOpenAIChatRequest(value any) (*openai.ChatCompletionNewParams, error) { + request, ok := value.(*openai.ChatCompletionNewParams) + if !ok || request == nil { + return nil, &protocolStageSetupError{err: fmt.Errorf("OpenAI Chat provider endpoint received %T", value)} + } + return request, nil +} + +type openAIChatProviderStream struct { + stream *openaistream.Stream[openai.ChatCompletionChunk] + cancel context.CancelFunc + model string + usage *protocol.TokenUsage + + closeOnce sync.Once + closeErr error +} + +func (s *openAIChatProviderStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + if s.stream == nil { + return protocolstage.Event{}, fmt.Errorf("OpenAI Chat provider stream is nil") + } + if !s.stream.Next() { + if err := s.stream.Err(); err != nil { + return protocolstage.Event{}, err + } + return protocolstage.Event{}, io.EOF + } + chunk := s.stream.Current() + if usage := protocolusage.FromOpenAIChatCompletion(chunk.Usage); usage.HasUsage() { + s.usage = usage + } + return protocolstage.Event{Value: chunk}, nil +} + +func (s *openAIChatProviderStream) Close() error { + s.closeOnce.Do(func() { + if s.stream != nil { + s.closeErr = s.stream.Close() + } + if s.cancel != nil { + s.cancel() + } + }) + return s.closeErr +} + +func (s *openAIChatProviderStream) Result() protocolstage.StreamResult { + return protocolstage.StreamResult{Usage: s.usage, Model: s.model} +} diff --git a/internal/protocolserver/protocol_stage_openai_responses.go b/internal/protocolserver/protocol_stage_openai_responses.go new file mode 100644 index 000000000..13d87b96c --- /dev/null +++ b/internal/protocolserver/protocol_stage_openai_responses.go @@ -0,0 +1,627 @@ +package protocolserver + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "sync" + + "github.com/gin-gonic/gin" + "github.com/openai/openai-go/v3/responses" + "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/anthropicbridge" + protocolguardrail "github.com/tingly-dev/tingly-box/internal/protocol/stage/guardrail" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/responsesbridge" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/transform" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" + "github.com/tingly-dev/tingly-box/internal/protocolserver/forwarding" + "github.com/tingly-dev/tingly-box/internal/typ" + pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" +) + +// tryProtocolStageOpenAIResponses selects an explicitly registered +// Responses-source route and owns its complete request lifecycle. +func (ph *ProtocolHandler) tryProtocolStageOpenAIResponses( + c *gin.Context, + req *protocol.ResponseCreateRequest, + responseModel string, + target protocol.APIType, + provider *typ.Provider, + actualModel string, + rule *typ.Rule, + isStreaming bool, + scenarioConfig *typ.ScenarioConfig, + ruleFlags typ.RuleFlags, + maxAllowed int, + stageRecording *protocolStageRequestRecording, +) bool { + mcpEnabled := ph.mcpEnabled() + guardrailsEnabled := ph.guardrailsEnabledForProtocolStage(GetTrackingContextScenario(c), protocol.TypeOpenAIResponses) + usesBetaStages := mcpEnabled || guardrailsEnabled + if usesBetaStages { + if !ph.shouldUseProtocolStageBetaChain(c, protocol.TypeOpenAIResponses, target, protocolstage.AllBridgeCapabilities) { + return false + } + if mcpEnabled && ph.deps.MCPRuntime == nil { + logProtocolStageFallback(c, protocol.TypeOpenAIResponses, target, "MCP runtime is unavailable") + return false + } + } else if !ph.shouldUseProtocolStage(c, protocol.TypeOpenAIResponses, target, protocolstage.AllBridgeCapabilities) { + return false + } + if scenarioConfig.IsRecordingEnable() && stageRecording == nil { + logProtocolStageFallback(c, protocol.TypeOpenAIResponses, target, "new recording path requires a fully Stage-compatible service set") + return false + } + + var requestErr error + defer func() { + markProtocolStageAttemptFailed(c, requestErr) + if stageRecording != nil { + stageRecording.observeAttempt(requestErr) + } + }() + + if c.GetHeader("X-Tingly-Debug-Routing") == "1" { + setProbeUpstreamHeadersForTarget(c, target, rule, provider) + c.Header(protocolPipelineHeader, "stage") + } + + scenarioFlags := scenarioFlagsOrNil(scenarioConfig) + options := append(protocolStageTransformOptions(ph, c), transform.WithMaxTokens(int64(maxAllowed))) + stages := []protocolstage.Stage{ + newProtocolTransformStage( + "client_prepare", + protocol.TypeOpenAIResponses, + provider, + scenarioFlags, + isStreaming, + RulePreBaseTransforms(ruleFlags), + options..., + ), + } + if guardrailsEnabled { + guardrailStage, guardrailErr := protocolguardrail.NewAnthropicBeta(protocolguardrail.AnthropicBetaConfig{ + Runtime: ph.currentGuardrailsRuntime(), + BaseInput: BuildGuardrailsBaseInput( + c, + actualModel, + provider, + guardrailscore.DirectionRequest, + nil, + ), + Observe: protocolStageGuardrailObserver(c), + }) + if guardrailErr != nil { + requestErr = guardrailErr + ph.FailAttemptSetup(c, guardrailErr) + return true + } + stages = append(stages, guardrailStage) + } + if mcpEnabled { + toolLoop, toolLoopErr := ph.newProtocolStageBetaToolLoop(c, provider, false) + if toolLoopErr != nil { + requestErr = toolLoopErr + ph.FailAttemptSetup(c, toolLoopErr) + return true + } + stages = append(stages, toolLoop) + } + stages = append(stages, + newProtocolTransformStage( + "provider_finalize", + target, + provider, + scenarioFlags, + isStreaming, + protocolStageProviderTransforms(target, + []transform.Transform{transform.NewConsistencyTransform(target)}, + RulePreVendorTransforms(ruleFlags), + []transform.Transform{vendorTransformShared}, + ), + options..., + ), + ) + terminal, registry, err := ph.protocolStageOpenAIResponsesTarget(target, provider, actualModel, responseModel, maxAllowed) + if err != nil { + requestErr = err + ph.FailAttemptSetup(c, err) + return true + } + terminal = requestrecord.ObserveProvider(terminal, stageRecordingRecorder(stageRecording), requestrecord.ExchangeMetadata{ + Attempt: currentProtocolStageAttempt(c), + Provider: provider.Name, + Model: actualModel, + }) + endpoint, err := protocolstage.BuildTopology(protocolstage.TopologyConfig{ + Terminal: terminal, + Stages: stages, + ClientProtocol: protocol.TypeOpenAIResponses, + Registry: registry, + RequiredCapabilities: protocolstage.AllBridgeCapabilities, + }) + if err != nil { + requestErr = fmt.Errorf("build OpenAI Responses Protocol Stage topology: %w", err) + ph.FailAttemptSetup(c, requestErr) + return true + } + + logProtocolStageEntry(c, protocol.TypeOpenAIResponses, target, stages, isStreaming) + call := protocolstage.Call{ + Request: req.ResponseNewParams, + Metadata: protocolstage.CallMetadata{ + RequestID: pkgobs.RequestIDFromContext(c.Request.Context()), + }, + } + if isStreaming { + requestErr = ph.serveProtocolStageOpenAIResponsesStream(c, endpoint, call, responseModel, stageRecordingRecorder(stageRecording)) + return true + } + requestErr = ph.serveProtocolStageOpenAIResponsesComplete(c, endpoint, call, responseModel, stageRecordingRecorder(stageRecording)) + return true +} + +func (ph *ProtocolHandler) protocolStageOpenAIResponsesTarget( + target protocol.APIType, + provider *typ.Provider, + actualModel string, + responseModel string, + maxAllowed int, +) (protocolstage.Endpoint, *protocolstage.BridgeRegistry, error) { + responsesToBeta := responsesbridge.NewToAnthropicBeta(responsesbridge.AnthropicOptions{ + DefaultMaxTokens: int64(maxAllowed), + ResponseModel: responseModel, + }) + responsesToChat := responsesbridge.NewToOpenAIChat(responsesbridge.ChatOptions{ + DefaultMaxTokens: int64(maxAllowed), + ResponseModel: responseModel, + }) + registry, err := protocolstage.NewBridgeRegistry( + protocolstage.NewIdentityBridge(protocol.TypeOpenAIResponses), + protocolstage.NewIdentityBridge(protocol.TypeAnthropicBeta), + responsesToBeta, + responsesToChat, + anthropicbridge.NewBetaToOpenAIChat(anthropicbridge.ChatOptions{ + Compatible: true, + ResponseModel: responseModel, + }), + anthropicbridge.NewBetaToOpenAIResponses(anthropicbridge.ResponsesOptions{ + ResponseModel: responseModel, + }), + ) + if err != nil { + return nil, nil, fmt.Errorf("build OpenAI Responses Protocol Stage registry: %w", err) + } + switch target { + case protocol.TypeOpenAIResponses: + return &openAIResponsesProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeAnthropicBeta: + return &anthropicBetaProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeOpenAIChat: + return &openAIChatProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + default: + return nil, nil, fmt.Errorf("OpenAI Responses Protocol Stage target %q is not implemented", target) + } +} + +// openAIResponsesProviderEndpoint is the transport-free Responses provider +// terminal. HTTP parsing, headers, SSE framing, and public model rewriting stay +// at the outer server adapter. +type openAIResponsesProviderEndpoint struct { + ph *ProtocolHandler + provider *typ.Provider + model string +} + +func (*openAIResponsesProviderEndpoint) Protocol() protocol.APIType { + return protocol.TypeOpenAIResponses +} + +func (e *openAIResponsesProviderEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + request, err := protocolStageOpenAIResponsesRequest(call.Request) + if err != nil { + return nil, err + } + wrapper := e.ph.deps.ClientPool.GetOpenAIClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + response, cancel, err := forwarding.ForwardOpenAIResponses(fc, wrapper, *request) + if cancel != nil { + defer cancel() + } + if err != nil { + return nil, err + } + return &protocolstage.Response{ + Value: response, + Usage: protocolusage.FromOpenAIResponses(response.Usage), + Model: e.model, + }, nil +} + +func (e *openAIResponsesProviderEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + request, err := protocolStageOpenAIResponsesRequest(call.Request) + if err != nil { + return nil, err + } + wrapper := e.ph.deps.ClientPool.GetOpenAIClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + stream, cancel, err := forwarding.ForwardOpenAIResponsesStream(fc, wrapper, *request) + if err != nil { + if cancel != nil { + cancel() + } + return nil, err + } + primed, err := protocolstream.PrimeResponsesStream(stream) + if err != nil { + if stream != nil { + _ = stream.Close() + } + if cancel != nil { + cancel() + } + return nil, err + } + return &openAIResponsesProviderStream{ + stream: primed, + cancel: cancel, + model: e.model, + }, nil +} + +func protocolStageOpenAIResponsesRequest(value any) (*responses.ResponseNewParams, error) { + request, ok := value.(*responses.ResponseNewParams) + if !ok || request == nil { + return nil, &protocolStageSetupError{err: fmt.Errorf("OpenAI Responses provider endpoint received %T", value)} + } + return request, nil +} + +type openAIResponsesProviderStream struct { + stream protocolstream.ResponsesStreamIter + cancel context.CancelFunc + model string + usage *protocol.TokenUsage + + closeOnce sync.Once + closeErr error +} + +func (s *openAIResponsesProviderStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + if s.stream == nil { + return protocolstage.Event{}, fmt.Errorf("OpenAI Responses provider stream is nil") + } + if !s.stream.Next() { + if err := s.stream.Err(); err != nil { + return protocolstage.Event{}, err + } + return protocolstage.Event{}, io.EOF + } + event := s.stream.Current() + if usage := protocolStageOpenAIResponsesUsage([]byte(event.RawJSON())); usage != nil { + s.usage = usage + } + return protocolstage.Event{Value: event}, nil +} + +func (s *openAIResponsesProviderStream) Close() error { + s.closeOnce.Do(func() { + if s.stream != nil { + s.closeErr = s.stream.Close() + } + if s.cancel != nil { + s.cancel() + } + }) + return s.closeErr +} + +func (s *openAIResponsesProviderStream) Result() protocolstage.StreamResult { + return protocolstage.StreamResult{Usage: s.usage, Model: s.model} +} + +func (ph *ProtocolHandler) serveProtocolStageOpenAIResponsesComplete( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + requestRecorder *requestrecord.Recorder, +) error { + result, err := endpoint.Complete(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.failRequest(c, nil, err, "OpenAI Responses Protocol Stage provider request failed") + return err + } + preserveProtocolStageSideEffectBoundary(c, nil, result.SideEffectsCommitted) + body, err := protocolStageOpenAIResponsesValueJSON(result.Value, responseModel) + if err != nil { + ph.FailAttemptSetup(c, err) + return err + } + captureProtocolStageFinalResponse(c.Request.Context(), requestRecorder, protocol.TypeOpenAIResponses, json.RawMessage(body)) + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + c.Data(http.StatusOK, "application/json; charset=utf-8", body) + return nil +} + +func protocolStageOpenAIResponsesValueJSON(value any, responseModel string) ([]byte, error) { + switch response := value.(type) { + case *responses.Response: + return protocolStageOpenAIResponsesJSON(response, responseModel) + case responses.Response: + return protocolStageOpenAIResponsesJSON(&response, responseModel) + case wire.ResponsesWireResponse: + response.Model = responseModel + body, err := json.Marshal(response) + if err != nil { + return nil, fmt.Errorf("marshal Protocol Stage Responses wire response: %w", err) + } + return body, nil + case *wire.ResponsesWireResponse: + if response == nil { + return nil, fmt.Errorf("OpenAI Responses Protocol Stage wire response is nil") + } + converted := *response + converted.Model = responseModel + body, err := json.Marshal(converted) + if err != nil { + return nil, fmt.Errorf("marshal Protocol Stage Responses wire response: %w", err) + } + return body, nil + default: + return nil, fmt.Errorf("OpenAI Responses Protocol Stage response has type %T", value) + } +} + +func protocolStageOpenAIResponsesJSON(response *responses.Response, responseModel string) ([]byte, error) { + if response == nil { + return nil, fmt.Errorf("OpenAI Responses Protocol Stage response is nil") + } + body := []byte(response.RawJSON()) + if len(body) == 0 { + var err error + body, err = json.Marshal(response) + if err != nil { + return nil, fmt.Errorf("marshal OpenAI Responses response: %w", err) + } + } + modified, err := sjson.SetBytes(body, "model", responseModel) + if err != nil { + return nil, fmt.Errorf("rewrite OpenAI Responses response model: %w", err) + } + return modified, nil +} + +func (ph *ProtocolHandler) serveProtocolStageOpenAIResponsesStream( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + requestRecorder *requestrecord.Recorder, +) error { + flusher, ok := c.Writer.(http.Flusher) + if !ok { + err := fmt.Errorf("OpenAI Responses Protocol Stage streaming is unsupported by this connection") + ph.FailAttemptSetup(c, err) + return err + } + stream, err := endpoint.Stream(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.handlePreStreamFailure(c, err, nil) + return err + } + defer func() { + if closeErr := stream.Close(); closeErr != nil { + logrus.WithContext(c.Request.Context()).Warnf("close OpenAI Responses Protocol Stage stream: %v", closeErr) + } + }() + finalCapture := newProtocolStageFinalStreamCapture(c.Request.Context(), requestRecorder, protocol.TypeOpenAIResponses) + + wrote := false + sawTerminalEvent := false + var terminalErr error + for { + event, nextErr := stream.Next(c.Request.Context()) + if errors.Is(nextErr, io.EOF) { + if sawTerminalEvent { + break + } + nextErr = fmt.Errorf("OpenAI Responses Protocol Stage stream ended without a terminal event") + } + if nextErr != nil { + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nextErr, result.SideEffectsCommitted) + if errors.Is(nextErr, context.Canceled) || protocol.IsContextCanceled(nextErr) { + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + return nextErr + } + if result.Usage != nil && result.Usage.HasUsage() { + ph.trackUsageWithTokenUsage(c, result.Usage, nextErr) + } else { + ph.trackUsageFromContext(c, 0, 0, nextErr) + } + if !wrote { + protocolstream.SendStreamingError(c, nextErr) + } else { + protocolstream.OpenAIResponsesEvent(c, "error", map[string]any{ + "error": map[string]any{ + "message": nextErr.Error(), + "type": "stream_error", + "code": "stream_failed", + }, + }) + flusher.Flush() + } + return nextErr + } + + eventType, payload, eventErr := protocolStageOpenAIResponsesEventJSON(event.Value, responseModel) + if eventErr != nil { + streamErr := fmt.Errorf("OpenAI Responses Protocol Stage stream emitted %T: %w", event.Value, eventErr) + preserveProtocolStageSideEffectBoundary(c, streamErr, stream.Result().SideEffectsCommitted) + if !wrote { + ph.FailAttemptSetup(c, streamErr) + } else { + protocolstream.OpenAIResponsesEvent(c, "error", map[string]any{ + "error": map[string]any{ + "message": "Protocol Stage stream emitted an invalid event", + "type": "stream_error", + "code": "stream_failed", + }, + }) + flusher.Flush() + } + return streamErr + } + finalCapture.add(c.Request.Context(), json.RawMessage(payload)) + if !wrote { + setProtocolStageOpenAIResponsesSSEHeaders(c) + wrote = true + } + protocolstream.OpenAIResponsesEvent(c, eventType, payload) + if protocolStageOpenAIResponsesTerminalEvent(eventType) { + sawTerminalEvent = true + if protocolStageOpenAIResponsesFailureEvent(eventType) { + terminalErr = fmt.Errorf("OpenAI Responses stream ended with %s", eventType) + } + } + CommitFirstChunkIfGate(c.Writer) + flusher.Flush() + } + + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nil, result.SideEffectsCommitted) + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, terminalErr) + } + finalCapture.finish(c.Request.Context()) + return terminalErr +} + +func protocolStageOpenAIResponsesFailureEvent(eventType string) bool { + return eventType == "response.failed" || eventType == "error" +} + +func protocolStageOpenAIResponsesTerminalEvent(eventType string) bool { + switch eventType { + case "response.completed", "response.incomplete", "response.failed", "error": + return true + default: + return false + } +} + +func setProtocolStageOpenAIResponsesSSEHeaders(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("Access-Control-Allow-Origin", "*") + c.Header("Access-Control-Allow-Headers", "Cache-Control") +} + +func protocolStageOpenAIResponsesEventJSON(value any, responseModel string) (string, []byte, error) { + var eventType string + var raw []byte + switch event := value.(type) { + case responses.ResponseStreamEventUnion: + eventType = event.Type + raw = []byte(strings.Clone(event.RawJSON())) + if len(raw) == 0 { + var err error + raw, err = json.Marshal(event) + if err != nil { + return "", nil, fmt.Errorf("marshal OpenAI Responses stream event: %w", err) + } + } + case wire.ResponsesEvent: + eventType = event.EventType() + var err error + raw, err = json.Marshal(event) + if err != nil { + return "", nil, fmt.Errorf("marshal converted OpenAI Responses stream event: %w", err) + } + default: + return "", nil, fmt.Errorf("unsupported OpenAI Responses stream event %T", value) + } + if eventType == "" { + return "", nil, fmt.Errorf("OpenAI Responses stream event has empty type") + } + + response := gjson.GetBytes(raw, "response") + if !response.Exists() { + return eventType, raw, nil + } + usage := response.Get("usage") + if inputDetails := usage.Get("input_tokens_details"); inputDetails.Exists() && !inputDetails.Get("cached_tokens").Exists() { + modified, err := sjson.SetBytes(raw, "response.usage.input_tokens_details.cached_tokens", 0) + if err != nil { + return "", nil, fmt.Errorf("backfill OpenAI Responses cached tokens: %w", err) + } + raw = modified + } + if usage.Exists() && !usage.Get("output_tokens_details.reasoning_tokens").Exists() { + modified, err := sjson.SetBytes(raw, "response.usage.output_tokens_details.reasoning_tokens", 0) + if err != nil { + return "", nil, fmt.Errorf("backfill OpenAI Responses reasoning tokens: %w", err) + } + raw = modified + } + if model := response.Get("model"); model.Exists() && model.String() != "" { + modified, err := sjson.SetBytes(raw, "response.model", responseModel) + if err != nil { + return "", nil, fmt.Errorf("rewrite OpenAI Responses stream model: %w", err) + } + raw = modified + } + return eventType, raw, nil +} + +func protocolStageOpenAIResponsesUsage(raw []byte) *protocol.TokenUsage { + response := gjson.GetBytes(raw, "response") + if !response.Exists() { + return nil + } + usage := response.Get("usage") + input := usage.Get("input_tokens").Int() + output := usage.Get("output_tokens").Int() + cacheRead := usage.Get("input_tokens_details.cached_tokens").Int() + cacheWrite := usage.Get("input_tokens_details.cache_write_tokens").Int() + reasoning := usage.Get("output_tokens_details.reasoning_tokens").Int() + if input == 0 && output == 0 && cacheRead == 0 && cacheWrite == 0 && reasoning == 0 { + return nil + } + return protocol.NewTokenUsageFull(int(input-cacheRead), int(output), int(cacheRead), int(cacheWrite), int(reasoning)) +} diff --git a/internal/protocolserver/protocol_stage_openai_responses_test.go b/internal/protocolserver/protocol_stage_openai_responses_test.go new file mode 100644 index 000000000..66bbade97 --- /dev/null +++ b/internal/protocolserver/protocol_stage_openai_responses_test.go @@ -0,0 +1,116 @@ +package protocolserver + +import ( + "encoding/json" + "testing" + + "github.com/openai/openai-go/v3/responses" + "github.com/tidwall/gjson" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +func TestProtocolStageOpenAIResponsesJSONPreservesRawWireFields(t *testing.T) { + t.Parallel() + + raw := `{"id":"resp_1","object":"response","model":"upstream-model","output":[],"usage":{"input_tokens":3,"output_tokens":2},"provider_extension":{"kept":true}}` + var response responses.Response + if err := json.Unmarshal([]byte(raw), &response); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + + body, err := protocolStageOpenAIResponsesJSON(&response, "public-model") + if err != nil { + t.Fatalf("protocolStageOpenAIResponsesJSON: %v", err) + } + if got := gjson.GetBytes(body, "model").String(); got != "public-model" { + t.Fatalf("model = %q, want public-model", got) + } + if !gjson.GetBytes(body, "provider_extension.kept").Bool() { + t.Fatalf("provider extension was not preserved: %s", body) + } +} + +func TestProtocolStageOpenAIResponsesValueJSONAcceptsTypedWireResponse(t *testing.T) { + t.Parallel() + + body, err := protocolStageOpenAIResponsesValueJSON(wire.ResponsesWireResponse{ + ID: "resp_typed", + Object: "response", + Model: "provider-model", + Status: "completed", + Output: []wire.ResponsesOutputItemWire{{ + ID: "msg_1", Type: "message", Role: "assistant", Status: "completed", + Content: []wire.ResponsesContentPartWire{{Type: "output_text", Text: "typed response"}}, + }}, + Usage: &wire.ResponsesUsageWire{InputTokens: 3, OutputTokens: 2, TotalTokens: 5}, + }, "public-model") + if err != nil { + t.Fatalf("protocolStageOpenAIResponsesValueJSON: %v", err) + } + if got := gjson.GetBytes(body, "model").String(); got != "public-model" { + t.Fatalf("model = %q, want public-model", got) + } + if got := gjson.GetBytes(body, "output.0.content.0.text").String(); got != "typed response" { + t.Fatalf("text = %q", got) + } + if !gjson.GetBytes(body, "usage.input_tokens_details.cached_tokens").Exists() || + !gjson.GetBytes(body, "usage.output_tokens_details.reasoning_tokens").Exists() { + t.Fatalf("strict usage details missing: %s", body) + } +} + +func TestProtocolStageOpenAIResponsesEventJSONPreservesWireShape(t *testing.T) { + t.Parallel() + + raw := `{"type":"response.completed","sequence_number":4,"response":{"id":"resp_1","object":"response","model":"upstream-model","output":[],"usage":{"input_tokens":9,"output_tokens":4,"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":1},"output_tokens_details":{"reasoning_tokens":1}},"provider_extension":"kept"}}` + var event responses.ResponseStreamEventUnion + if err := json.Unmarshal([]byte(raw), &event); err != nil { + t.Fatalf("unmarshal event: %v", err) + } + + eventType, body, err := protocolStageOpenAIResponsesEventJSON(event, "public-model") + if err != nil { + t.Fatalf("protocolStageOpenAIResponsesEventJSON: %v", err) + } + if eventType != "response.completed" { + t.Fatalf("event type = %q", eventType) + } + if got := gjson.GetBytes(body, "response.model").String(); got != "public-model" { + t.Fatalf("model = %q, want public-model", got) + } + if got := gjson.GetBytes(body, "response.provider_extension").String(); got != "kept" { + t.Fatalf("provider extension = %q", got) + } + if !gjson.GetBytes(body, "response.usage.input_tokens_details.cached_tokens").Exists() { + t.Fatalf("cached_tokens was not backfilled: %s", body) + } + if !gjson.GetBytes(body, "response.usage.output_tokens_details.reasoning_tokens").Exists() { + t.Fatalf("reasoning_tokens was not backfilled: %s", body) + } + + usage := protocolStageOpenAIResponsesUsage(body) + if usage == nil || usage.InputTokens != 7 || usage.OutputTokens != 4 || usage.CacheReadTokens != 2 || usage.CacheWriteTokens != 1 || usage.ReasoningTokens != 1 { + t.Fatalf("usage = %#v", usage) + } +} + +func TestProtocolStageOpenAIResponsesUsagePreservesCacheDetails(t *testing.T) { + t.Parallel() + + usage := protocolStageOpenAIResponsesUsage([]byte(`{ + "response": { + "usage": { + "input_tokens": 10, + "output_tokens": 4, + "input_tokens_details": {"cached_tokens": 3, "cache_write_tokens": 2}, + "output_tokens_details": {"reasoning_tokens": 1} + } + } + }`)) + if usage == nil { + t.Fatal("usage is nil") + } + if usage.InputTokens != 7 || usage.OutputTokens != 4 || usage.CacheReadTokens != 3 || usage.CacheWriteTokens != 2 || usage.ReasoningTokens != 1 { + t.Fatalf("usage = %#v", usage) + } +} diff --git a/internal/protocolserver/protocol_stage_pipeline.go b/internal/protocolserver/protocol_stage_pipeline.go new file mode 100644 index 000000000..186f42cd8 --- /dev/null +++ b/internal/protocolserver/protocol_stage_pipeline.go @@ -0,0 +1,824 @@ +package protocolserver + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "sync" + + "github.com/anthropics/anthropic-sdk-go" + anthropicstream "github.com/anthropics/anthropic-sdk-go/packages/ssestream" + "github.com/gin-gonic/gin" + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/responses" + "github.com/sirupsen/logrus" + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/assembler" + "github.com/tingly-dev/tingly-box/internal/protocol/ops" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/anthropicbridge" + protocolguardrail "github.com/tingly-dev/tingly-box/internal/protocol/stage/guardrail" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/openaibridge" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/transform" + protocolusage "github.com/tingly-dev/tingly-box/internal/protocol/usage" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" + "github.com/tingly-dev/tingly-box/internal/protocolserver/forwarding" + "github.com/tingly-dev/tingly-box/internal/protocolserver/recording" + "github.com/tingly-dev/tingly-box/internal/typ" + pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" +) + +const protocolPipelineHeader = "X-Tingly-Protocol-Pipeline" + +// tryProtocolStageOpenAIChat selects and executes the Stage path for one +// provider attempt. Returning false means the caller must continue through the +// legacy transform/dispatch path. Returning true means this method has fully +// handled the attempt, including any response or error. +func (ph *ProtocolHandler) tryProtocolStageOpenAIChat( + c *gin.Context, + req *protocol.OpenAIChatCompletionRequest, + responseModel string, + target protocol.APIType, + provider *typ.Provider, + actualModel string, + rule *typ.Rule, + isStreaming bool, + scenarioConfig *typ.ScenarioConfig, + ruleFlags typ.RuleFlags, + stageRecording *protocolStageRequestRecording, +) bool { + mcpEnabled := ph.mcpEnabled() + guardrailsEnabled := ph.guardrailsEnabledForProtocolStage(GetTrackingContextScenario(c), protocol.TypeOpenAIChat) + usesBetaStages := mcpEnabled || guardrailsEnabled + if usesBetaStages { + if !ph.shouldUseProtocolStageBetaChain(c, protocol.TypeOpenAIChat, target, protocolstage.AllBridgeCapabilities) { + return false + } + if mcpEnabled && ph.deps.MCPRuntime == nil { + logProtocolStageFallback(c, protocol.TypeOpenAIChat, target, "MCP runtime is unavailable") + return false + } + } else if !ph.shouldUseProtocolStage(c, protocol.TypeOpenAIChat, target, protocolstage.AllBridgeCapabilities) { + return false + } + // The response-roundtrip header is an explicit legacy diagnostic. Preserve + // that exact experiment instead of changing its semantics under --stage. + if target == protocol.TypeAnthropicBeta && ShouldRoundtripResponse(c, "anthropic") { + logProtocolStageFallback(c, protocol.TypeOpenAIChat, target, "response roundtrip diagnostic requires the legacy pipeline") + return false + } + if !protocolStageOpenAIChatChoiceCompatible(target, req.ChatCompletionNewParams) { + logProtocolStageFallback(c, protocol.TypeOpenAIChat, target, "cross-protocol Chat n > 1 is not representable by the target protocol") + return false + } + if scenarioConfig.IsRecordingEnable() && stageRecording == nil { + logProtocolStageFallback(c, protocol.TypeOpenAIChat, target, "new recording path requires a fully Stage-compatible service set") + return false + } + + var requestErr error + defer func() { + markProtocolStageAttemptFailed(c, requestErr) + if stageRecording != nil { + stageRecording.observeAttempt(requestErr) + } + }() + + if c.GetHeader("X-Tingly-Debug-Routing") == "1" { + setProbeUpstreamHeadersForTarget(c, target, rule, provider) + c.Header(protocolPipelineHeader, "stage") + } + + preBase := RulePreBaseTransforms(ruleFlags) + preVendor := RulePreVendorTransforms(ruleFlags) + disableStreamUsage := ruleFlags.SkipUsage || ruleFlags.CursorCompat + terminal, registry, err := ph.protocolStageOpenAIChatTarget( + target, + provider, + actualModel, + responseModel, + disableStreamUsage, + ) + if err != nil { + requestErr = err + ph.FailAttemptSetup(c, err) + return true + } + terminal = requestrecord.ObserveProvider(terminal, stageRecordingRecorder(stageRecording), requestrecord.ExchangeMetadata{ + Attempt: currentProtocolStageAttempt(c), + Provider: provider.Name, + Model: actualModel, + }) + + scenarioFlags := scenarioFlagsOrNil(scenarioConfig) + stages := []protocolstage.Stage{ + newProtocolTransformStage( + "client_prepare", + protocol.TypeOpenAIChat, + provider, + scenarioFlags, + isStreaming, + preBase, + protocolStageTransformOptions(ph, c)..., + ), + } + if guardrailsEnabled { + guardrailStage, guardrailErr := protocolguardrail.NewAnthropicBeta(protocolguardrail.AnthropicBetaConfig{ + Runtime: ph.currentGuardrailsRuntime(), + BaseInput: BuildGuardrailsBaseInput( + c, + actualModel, + provider, + guardrailscore.DirectionRequest, + nil, + ), + Observe: protocolStageGuardrailObserver(c), + }) + if guardrailErr != nil { + requestErr = guardrailErr + ph.FailAttemptSetup(c, guardrailErr) + return true + } + stages = append(stages, guardrailStage) + } + if mcpEnabled { + toolLoop, toolLoopErr := ph.newProtocolStageBetaToolLoop(c, provider, false) + if toolLoopErr != nil { + requestErr = toolLoopErr + ph.FailAttemptSetup(c, toolLoopErr) + return true + } + stages = append(stages, toolLoop) + } + stages = append(stages, + newProtocolTransformStage( + "provider_finalize", + target, + provider, + scenarioFlags, + isStreaming, + protocolStageProviderTransforms(target, + []transform.Transform{transform.NewConsistencyTransform(target)}, + preVendor, + []transform.Transform{vendorTransformShared}, + ), + protocolStageTransformOptions(ph, c)..., + ), + ) + endpoint, err := protocolstage.BuildTopology(protocolstage.TopologyConfig{ + Terminal: terminal, + Stages: stages, + ClientProtocol: protocol.TypeOpenAIChat, + Registry: registry, + RequiredCapabilities: protocolstage.AllBridgeCapabilities, + }) + if err != nil { + requestErr = fmt.Errorf("build Protocol Stage topology: %w", err) + ph.FailAttemptSetup(c, requestErr) + return true + } + + logProtocolStageEntry(c, protocol.TypeOpenAIChat, target, stages, isStreaming) + + call := protocolstage.Call{ + Request: req.ChatCompletionNewParams, + Metadata: protocolstage.CallMetadata{ + RequestID: pkgobs.RequestIDFromContext(c.Request.Context()), + }, + } + if isStreaming { + requestErr = ph.serveProtocolStageOpenAIChatStream(c, endpoint, call, responseModel, disableStreamUsage, nil, stageRecordingRecorder(stageRecording)) + return true + } + requestErr = ph.serveProtocolStageOpenAIChatComplete(c, endpoint, call, responseModel, provider, actualModel, disableStreamUsage, nil, stageRecordingRecorder(stageRecording)) + return true +} + +func protocolStageOpenAIChatChoiceCompatible(target protocol.APIType, request any) bool { + return target == protocol.TypeOpenAIChat || assembler.OpenAIChatChoiceCount(request) <= 1 +} + +func (ph *ProtocolHandler) protocolStageOpenAIChatTarget( + target protocol.APIType, + provider *typ.Provider, + actualModel string, + responseModel string, + disableStreamUsage bool, +) (protocolstage.Endpoint, *protocolstage.BridgeRegistry, error) { + registry, err := protocolstage.NewBridgeRegistry( + protocolstage.NewIdentityBridge(protocol.TypeOpenAIChat), + protocolstage.NewIdentityBridge(protocol.TypeAnthropicBeta), + openaibridge.NewChatToAnthropicBeta(openaibridge.AnthropicOptions{ + DefaultMaxTokens: 4096, + DisableStreamUsage: disableStreamUsage, + ResponseModel: responseModel, + }), + openaibridge.NewChatToOpenAIResponses(openaibridge.ResponsesOptions{ + DefaultMaxTokens: 4096, + DisableStreamUsage: disableStreamUsage, + ResponseModel: responseModel, + }), + anthropicbridge.NewBetaToOpenAIChat(anthropicbridge.ChatOptions{ + Compatible: true, + DisableStreamUsage: disableStreamUsage, + ResponseModel: responseModel, + }), + anthropicbridge.NewBetaToOpenAIResponses(anthropicbridge.ResponsesOptions{ + ResponseModel: responseModel, + }), + ) + if err != nil { + return nil, nil, fmt.Errorf("build OpenAI Chat Protocol Stage registry: %w", err) + } + switch target { + case protocol.TypeOpenAIChat: + return &openAIChatProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeAnthropicBeta: + return &anthropicBetaProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + case protocol.TypeOpenAIResponses: + return &openAIResponsesProviderEndpoint{ph: ph, provider: provider, model: actualModel}, registry, nil + default: + return nil, nil, fmt.Errorf("OpenAI Chat Protocol Stage target %q is not implemented", target) + } +} + +func (ph *ProtocolHandler) shouldUseProtocolStage( + c *gin.Context, + source, target protocol.APIType, + required protocolstage.Capabilities, +) bool { + selector := ph.protocolStageSelector + if selector == nil { + return false + } + useStage, selectionErr := selector.ShouldUseStage(source, target, required) + if !useStage && selector.Enabled() && selectionErr != nil { + logrus.WithContext(c.Request.Context()).WithFields(logrus.Fields{ + "protocol_pipeline": "legacy", + "source_protocol": source, + "target_protocol": target, + }).Debugf("Protocol Stage route unavailable: %v", selectionErr) + } + return useStage +} + +func logProtocolStageFallback(c *gin.Context, source, target protocol.APIType, reason string) { + logrus.WithContext(c.Request.Context()).WithFields(logrus.Fields{ + "protocol_pipeline": "legacy", + "source_protocol": source, + "target_protocol": target, + "reason": reason, + }).Debug("Protocol Stage request stayed on legacy") +} + +func logProtocolStageEntry( + c *gin.Context, + source, target protocol.APIType, + stages []protocolstage.Stage, + streaming bool, +) { + operation := "complete" + if streaming { + operation = "stream" + } + chain := make([]string, 0, len(stages)) + for _, current := range stages { + chain = append(chain, fmt.Sprintf("%s[%s]", current.Name(), current.Protocol())) + } + logrus.WithContext(c.Request.Context()).WithFields(logrus.Fields{ + "protocol_pipeline": "stage", + "source_protocol": source, + "target_protocol": target, + "operation": operation, + "stage_chain": strings.Join(chain, " -> "), + }).Debug("Entering Protocol Stage pipeline") +} + +func appendProtocolStageTransforms(groups ...[]transform.Transform) []transform.Transform { + var count int + for _, group := range groups { + count += len(group) + } + result := make([]transform.Transform, 0, count) + for _, group := range groups { + result = append(result, group...) + } + return result +} + +func protocolStageProviderTransforms(target protocol.APIType, groups ...[]transform.Transform) []transform.Transform { + result := appendProtocolStageTransforms(groups...) + if target == protocol.TypeOpenAIChat { + result = append(result, transform.NewOpenAIChatProviderCleanupTransform()) + } + return result +} + +func protocolStageTransformOptions(ph *ProtocolHandler, c *gin.Context) []transform.TransformOption { + options := []transform.TransformOption{ + transform.WithDevice(ph.deps.Config.ClaudeCodeDeviceID), + } + if c.GetHeader("X-Tingly-Advisor-Depth") != "" { + options = append(options, transform.WithIsAdvisorRequest(true)) + } + return options +} + +type protocolStageSetupError struct{ err error } + +func (e *protocolStageSetupError) Error() string { return e.err.Error() } +func (e *protocolStageSetupError) Unwrap() error { return e.err } + +// protocolTransformStage reuses the existing non-transport transforms inside a +// native Protocol Stage boundary. It is constructed per provider attempt, while +// all mutable transform state remains per call. +type protocolTransformStage struct { + name string + api protocol.APIType + provider *typ.Provider + scenarioFlags *typ.ScenarioFlags + streaming bool + transforms []transform.Transform + options []transform.TransformOption +} + +func newProtocolTransformStage( + name string, + api protocol.APIType, + provider *typ.Provider, + scenarioFlags *typ.ScenarioFlags, + streaming bool, + transforms []transform.Transform, + options ...transform.TransformOption, +) protocolstage.Stage { + return &protocolTransformStage{ + name: name, + api: api, + provider: provider, + scenarioFlags: scenarioFlags, + streaming: streaming, + transforms: append([]transform.Transform(nil), transforms...), + options: append([]transform.TransformOption(nil), options...), + } +} + +func (s *protocolTransformStage) Name() string { return s.name } +func (s *protocolTransformStage) Protocol() protocol.APIType { return s.api } +func (s *protocolTransformStage) Wrap(next protocolstage.Endpoint) protocolstage.Endpoint { + return &protocolTransformEndpoint{stage: s, next: next} +} + +type protocolTransformEndpoint struct { + stage *protocolTransformStage + next protocolstage.Endpoint +} + +func (e *protocolTransformEndpoint) Protocol() protocol.APIType { return e.stage.api } + +func (e *protocolTransformEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + prepared, release, err := e.stage.prepare(ctx, call) + if err != nil { + return nil, err + } + defer release() + return e.next.Complete(ctx, prepared) +} + +func (e *protocolTransformEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + prepared, release, err := e.stage.prepare(ctx, call) + if err != nil { + return nil, err + } + defer release() + return e.next.Stream(ctx, prepared) +} + +func (s *protocolTransformStage) prepare(ctx context.Context, call protocolstage.Call) (protocolstage.Call, func(), error) { + if len(s.transforms) == 0 { + return call, func() {}, nil + } + opts := []transform.TransformOption{ + transform.WithContext(ctx), + transform.WithProvider(s.provider), + transform.WithScenarioFlags(s.scenarioFlags), + transform.WithStreaming(s.streaming), + } + opts = append(opts, s.options...) + var transformCtx *transform.TransformContext + switch request := call.Request.(type) { + case *openai.ChatCompletionNewParams: + transformCtx = transform.NewTransformContext(request, opts...) + transformCtx.Config.OpenAIConfig = call.State.OpenAIChat + case *responses.ResponseNewParams: + transformCtx = transform.NewTransformContext(request, opts...) + case *anthropic.MessageNewParams: + transformCtx = transform.NewTransformContext(request, opts...) + case *anthropic.BetaMessageNewParams: + transformCtx = transform.NewTransformContext(request, opts...) + default: + return protocolstage.Call{}, func() {}, &protocolStageSetupError{err: fmt.Errorf( + "Protocol Stage %q received request %T for %q", + s.name, + call.Request, + s.api, + )} + } + transformCtx.SourceAPI = s.api + transformCtx.TargetAPI = s.api + finalCtx, err := transform.NewTransformChain(s.transforms).Execute(transformCtx) + if err != nil { + transformCtx.Release() + return protocolstage.Call{}, func() {}, &protocolStageSetupError{err: fmt.Errorf("Protocol Stage %q transform: %w", s.name, err)} + } + prepared := call + prepared.Request = finalCtx.Request + return prepared, finalCtx.Release, nil +} + +// anthropicBetaProviderEndpoint is the transport-free provider terminal used +// by the first production Stage route. +type anthropicBetaProviderEndpoint struct { + ph *ProtocolHandler + provider *typ.Provider + model string +} + +func (*anthropicBetaProviderEndpoint) Protocol() protocol.APIType { return protocol.TypeAnthropicBeta } + +func (e *anthropicBetaProviderEndpoint) Complete(ctx context.Context, call protocolstage.Call) (*protocolstage.Response, error) { + request, err := protocolStageBetaRequest(call.Request) + if err != nil { + return nil, err + } + wrapper := e.ph.deps.ClientPool.GetAnthropicClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + message, cancel, err := forwarding.ForwardAnthropicV1Beta(fc, wrapper, request) + if cancel != nil { + defer cancel() + } + if err != nil { + return nil, err + } + return &protocolstage.Response{ + Value: message, + Usage: protocolusage.FromAnthropicBetaMessage(message.Usage), + Model: e.model, + }, nil +} + +func (e *anthropicBetaProviderEndpoint) Stream(ctx context.Context, call protocolstage.Call) (protocolstage.EventStream, error) { + request, err := protocolStageBetaRequest(call.Request) + if err != nil { + return nil, err + } + wrapper := e.ph.deps.ClientPool.GetAnthropicClient(ctx, e.provider, e.model) + fc := forwarding.NewForwardContext(ctx, e.provider) + stream, cancel, err := forwarding.ForwardAnthropicV1BetaStream(fc, wrapper, request) + if err != nil { + if cancel != nil { + cancel() + } + return nil, err + } + return &anthropicBetaProviderStream{ + stream: stream, + cancel: cancel, + model: e.model, + usage: protocolusage.NewAnthropicAccumulator(), + }, nil +} + +func protocolStageBetaRequest(value any) (*anthropic.BetaMessageNewParams, error) { + request, ok := value.(*anthropic.BetaMessageNewParams) + if !ok || request == nil { + return nil, &protocolStageSetupError{err: fmt.Errorf("Anthropic Beta provider endpoint received %T", value)} + } + return request, nil +} + +type anthropicBetaProviderStream struct { + stream *anthropicstream.Stream[anthropic.BetaRawMessageStreamEventUnion] + cancel context.CancelFunc + model string + usage *protocolusage.AnthropicAccumulator + + closeOnce sync.Once + closeErr error +} + +func (s *anthropicBetaProviderStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + if s.stream == nil { + return protocolstage.Event{}, fmt.Errorf("Anthropic Beta provider stream is nil") + } + if !s.stream.Next() { + if err := s.stream.Err(); err != nil { + return protocolstage.Event{}, err + } + return protocolstage.Event{}, io.EOF + } + event := s.stream.Current() + if s.usage != nil { + s.usage.ConsumeBeta(&event) + } + return protocolstage.Event{Value: event}, nil +} + +func (s *anthropicBetaProviderStream) Close() error { + s.closeOnce.Do(func() { + if s.stream != nil { + s.closeErr = s.stream.Close() + } + if s.cancel != nil { + s.cancel() + } + }) + return s.closeErr +} + +func (s *anthropicBetaProviderStream) Result() protocolstage.StreamResult { + var usage *protocol.TokenUsage + if s.usage != nil && s.usage.HasUsage() { + usage = s.usage.Result() + } + return protocolstage.StreamResult{Usage: usage, Model: s.model} +} + +func (ph *ProtocolHandler) serveProtocolStageOpenAIChatComplete( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + provider *typ.Provider, + actualModel string, + disableUsage bool, + recorder *recording.ProtocolRecorder, + requestRecorder *requestrecord.Recorder, +) error { + response, err := endpoint.Complete(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.failRequest(c, recorder, err, "Protocol Stage provider request failed") + return err + } + preserveProtocolStageSideEffectBoundary(c, nil, response.SideEffectsCommitted) + value, err := protocolStageChatResponseMap(response.Value) + if err != nil { + ph.failRequest(c, recorder, err, "Protocol Stage response conversion failed") + return err + } + value = ops.ApplyResponseTransforms(value, provider.APIBase, actualModel) + value["model"] = responseModel + if disableUsage { + delete(value, "usage") + } + captureProtocolStageFinalResponse(c.Request.Context(), requestRecorder, protocol.TypeOpenAIChat, value) + if response.Usage != nil && response.Usage.HasUsage() { + ph.trackUsageWithTokenUsage(c, response.Usage, nil) + } + if recorder != nil { + recorder.SetAssembledResponse(value) + recorder.RecordResponse(provider, actualModel) + } + c.JSON(http.StatusOK, value) + return nil +} + +func protocolStageChatResponseMap(value any) (map[string]any, error) { + switch response := value.(type) { + case wire.ChatCompletionWire: + return response.ToMap(), nil + case *wire.ChatCompletionWire: + if response == nil { + return nil, fmt.Errorf("Protocol Stage OpenAI Chat response is nil") + } + return response.ToMap(), nil + case openai.ChatCompletion: + return protocolStageChatSDKMap(response) + case *openai.ChatCompletion: + if response == nil { + return nil, fmt.Errorf("Protocol Stage OpenAI Chat response is nil") + } + return protocolStageChatSDKMap(response) + default: + return nil, fmt.Errorf("Protocol Stage OpenAI Chat response has type %T", value) + } +} + +func protocolStageChatSDKMap(value any) (map[string]any, error) { + raw, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("marshal Protocol Stage OpenAI Chat response: %w", err) + } + var result wire.ChatCompletionWire + if err := json.Unmarshal(raw, &result); err != nil { + return nil, fmt.Errorf("decode Protocol Stage OpenAI Chat response: %w", err) + } + return result.ToMap(), nil +} + +func (ph *ProtocolHandler) serveProtocolStageOpenAIChatStream( + c *gin.Context, + endpoint protocolstage.Endpoint, + call protocolstage.Call, + responseModel string, + disableUsage bool, + recorder *recording.ProtocolRecorder, + requestRecorder *requestrecord.Recorder, +) error { + flusher, ok := c.Writer.(http.Flusher) + if !ok { + err := fmt.Errorf("Protocol Stage streaming is unsupported by this connection") + ph.FailAttemptSetup(c, err) + return err + } + stream, err := endpoint.Stream(c.Request.Context(), call) + if err != nil { + preserveProtocolStageSideEffectBoundary(c, err, false) + var setupErr *protocolStageSetupError + if errors.As(err, &setupErr) { + ph.FailAttemptSetup(c, setupErr) + return err + } + ph.failRequest(c, recorder, err, "Protocol Stage provider stream failed") + return err + } + defer func() { + if closeErr := stream.Close(); closeErr != nil { + logrus.WithContext(c.Request.Context()).Warnf("close Protocol Stage stream: %v", closeErr) + } + }() + finalCapture := newProtocolStageFinalStreamCapture(c.Request.Context(), requestRecorder, protocol.TypeOpenAIChat) + + wrote := false + expectedChoices := assembler.OpenAIChatChoiceCount(call.Request) + finishedChoices := make(map[int]struct{}) + for { + event, nextErr := stream.Next(c.Request.Context()) + if errors.Is(nextErr, io.EOF) { + if len(finishedChoices) >= expectedChoices { + break + } + nextErr = errors.New("OpenAI Chat Protocol Stage stream ended without a terminal finish_reason") + } + if nextErr != nil { + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nextErr, result.SideEffectsCommitted) + if errors.Is(nextErr, context.Canceled) || protocol.IsContextCanceled(nextErr) { + if result.Usage != nil { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + return nextErr + } + if result.Usage != nil && result.Usage.HasUsage() { + ph.trackUsageWithTokenUsage(c, result.Usage, nextErr) + } else { + ph.trackUsageFromContext(c, 0, 0, nextErr) + } + if !wrote { + SendErrorResponse(c, nextErr, "Protocol Stage provider stream failed") + } else { + protocolstream.OpenAISSE(c, ErrorResponse{Error: ErrorDetail{ + Message: "Protocol Stage stream terminated", + Type: "protocol_error", + }}) + flusher.Flush() + } + if recorder != nil { + recorder.RecordError(nextErr) + } + return nextErr + } + + chunk, chunkErr := protocolStageChatStreamChunk(event.Value) + if chunkErr != nil { + streamErr := chunkErr + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, streamErr, result.SideEffectsCommitted) + if result.Usage != nil && result.Usage.HasUsage() { + ph.trackUsageWithTokenUsage(c, result.Usage, streamErr) + } else { + ph.trackUsageFromContext(c, 0, 0, streamErr) + } + if !wrote { + ph.FailAttemptSetup(c, streamErr) + } else { + protocolstream.OpenAISSE(c, ErrorResponse{Error: ErrorDetail{ + Message: "Protocol Stage stream emitted an invalid event", + Type: "protocol_error", + }}) + flusher.Flush() + } + if recorder != nil { + recorder.RecordError(streamErr) + } + return streamErr + } + chunk.Model = responseModel + protocolStageTrackFinishedChatChoices(chunk, expectedChoices, finishedChoices) + if disableUsage { + chunk.Usage = nil + if len(chunk.Choices) == 0 { + continue + } + } + finalCapture.add(c.Request.Context(), chunk) + if !wrote { + setProtocolStageSSEHeaders(c) + wrote = true + } + if protocolStageChatChunkHasContent(chunk) { + protocol.MarkFirstToken(c) + } + protocolstream.OpenAISSE(c, chunk) + CommitFirstChunkIfGate(c.Writer) + flusher.Flush() + } + + if !wrote { + setProtocolStageSSEHeaders(c) + } + protocolstream.OpenAISSEDone(c) + CommitFirstChunkIfGate(c.Writer) + flusher.Flush() + result := stream.Result() + preserveProtocolStageSideEffectBoundary(c, nil, result.SideEffectsCommitted) + if result.Usage != nil && result.Usage.HasUsage() { + ph.trackUsageWithTokenUsage(c, result.Usage, nil) + } + finalCapture.finish(c.Request.Context()) + return nil +} + +func protocolStageTrackFinishedChatChoices(chunk wire.ChatStreamChunk, expected int, finished map[int]struct{}) { + for _, choice := range chunk.Choices { + if choice.FinishReason != nil && *choice.FinishReason != "" && choice.Index >= 0 && choice.Index < expected { + finished[choice.Index] = struct{}{} + } + } +} + +func protocolStageChatStreamChunk(value any) (wire.ChatStreamChunk, error) { + switch chunk := value.(type) { + case wire.ChatStreamChunk: + return chunk, nil + case *wire.ChatStreamChunk: + if chunk == nil { + return wire.ChatStreamChunk{}, fmt.Errorf("Protocol Stage OpenAI Chat stream chunk is nil") + } + return *chunk, nil + case openai.ChatCompletionChunk: + return protocolStageChatSDKChunk(chunk) + case *openai.ChatCompletionChunk: + if chunk == nil { + return wire.ChatStreamChunk{}, fmt.Errorf("Protocol Stage OpenAI Chat stream chunk is nil") + } + return protocolStageChatSDKChunk(chunk) + default: + return wire.ChatStreamChunk{}, fmt.Errorf("Protocol Stage stream emitted %T, want Chat stream chunk", value) + } +} + +func protocolStageChatSDKChunk(value any) (wire.ChatStreamChunk, error) { + raw, err := json.Marshal(value) + if err != nil { + return wire.ChatStreamChunk{}, fmt.Errorf("marshal Protocol Stage OpenAI Chat stream chunk: %w", err) + } + var result wire.ChatStreamChunk + if err := json.Unmarshal(raw, &result); err != nil { + return wire.ChatStreamChunk{}, fmt.Errorf("decode Protocol Stage OpenAI Chat stream chunk: %w", err) + } + return result, nil +} + +func setProtocolStageSSEHeaders(c *gin.Context) { + c.Header("Access-Control-Allow-Headers", "Origin, Content-Type, Authorization, Cache-Control") + c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") + c.Header("Access-Control-Allow-Origin", "*") + c.Header("Content-Type", "text/event-stream; charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") +} + +func protocolStageChatChunkHasContent(chunk wire.ChatStreamChunk) bool { + for _, choice := range chunk.Choices { + if choice.Delta.Content != "" || choice.Delta.ReasoningContent != "" || len(choice.Delta.ToolCalls) > 0 { + return true + } + } + return false +} diff --git a/internal/protocolserver/protocol_stage_recording.go b/internal/protocolserver/protocol_stage_recording.go new file mode 100644 index 000000000..c5a330f5d --- /dev/null +++ b/internal/protocolserver/protocol_stage_recording.go @@ -0,0 +1,262 @@ +package protocolserver + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/sirupsen/logrus" + internalobs "github.com/tingly-dev/tingly-box/internal/obs" + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/assembler" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +const protocolStageOriginalInputKey = "protocol_stage_original_input" + +// protocolStageRequestRecording owns the new request-boundary recorder and its +// existing obs sink. It is intentionally separate from the legacy Gin-based +// ProtocolRecorder while the Stage path is canaried additively. +type protocolStageRequestRecording struct { + recorder *requestrecord.Recorder + sink *internalobs.Sink + + mu sync.Mutex + used bool + lastAttemptErr error +} + +func (ph *ProtocolHandler) rememberProtocolStageOriginalInput(c *gin.Context, scenario typ.RuleScenario, body []byte) { + if ph == nil || !ph.deps.ProtocolStageEnabled || ph.deps.Config == nil || c == nil || len(body) == 0 { + return + } + if !ph.deps.Config.GetScenarioConfig(scenario).IsRecordingEnable() { + return + } + c.Set(protocolStageOriginalInputKey, json.RawMessage(append([]byte(nil), body...))) +} + +func protocolStageOriginalInput(c *gin.Context, fallback any) any { + if c == nil { + return fallback + } + if input, ok := c.Get(protocolStageOriginalInputKey); ok { + return input + } + return fallback +} + +func (ph *ProtocolHandler) newProtocolStageRequestRecording( + scenario typ.RuleScenario, + inputProtocol protocol.APIType, + input any, + sessionID typ.SessionID, + requestID string, +) *protocolStageRequestRecording { + if ph == nil || !ph.deps.ProtocolStageEnabled || ph.deps.GetOrCreateScenarioSink == nil { + return nil + } + sink := ph.deps.GetOrCreateScenarioSink(scenario) + if sink == nil { + return nil + } + sessionShort, _ := internalobs.SessionShort(sessionID) + if requestID == "" { + requestID = uuid.NewString() + } + recorder, err := requestrecord.New(requestrecord.Config{ + Enabled: true, + RequestID: requestID, + SessionID: sessionShort, + Scenario: string(scenario), + InputProtocol: inputProtocol, + Input: input, + }) + if err != nil { + logrus.Debugf("obs: failed to build Protocol Stage RequestRecord: %v", err) + return nil + } + return &protocolStageRequestRecording{recorder: recorder, sink: sink} +} + +func (r *protocolStageRequestRecording) finish(requestErr error) { + if r == nil || r.recorder == nil || r.sink == nil { + return + } + completed, first := r.recorder.Finish(requestErr) + if !first || completed == nil { + return + } + r.sink.EmitRequestRecord(completed) +} + +func (r *protocolStageRequestRecording) observeAttempt(requestErr error) { + if r == nil { + return + } + r.mu.Lock() + r.used = true + r.lastAttemptErr = requestErr + r.mu.Unlock() +} + +func (r *protocolStageRequestRecording) finishFromHTTP(c *gin.Context) { + if r == nil { + return + } + r.mu.Lock() + used := r.used + requestErr := r.lastAttemptErr + r.mu.Unlock() + if !used { + return + } + if c != nil && c.Writer.Status() >= http.StatusBadRequest && requestErr == nil { + requestErr = fmt.Errorf("request completed with HTTP status %d", c.Writer.Status()) + } + r.finish(requestErr) +} + +func (ph *ProtocolHandler) protocolStageRecordingSupportsRule(rule *typ.Rule) bool { + if ph == nil || rule == nil { + return false + } + services := rule.GetActiveServices() + if len(services) == 0 { + return false + } + for _, service := range services { + provider, err := ph.deps.Config.GetProviderByUUID(service.Provider) + if err != nil || provider == nil || provider.APIStyle == protocol.APIStyleGoogle { + return false + } + } + return true +} + +const protocolStageAttemptKey = "protocol_stage_attempt" +const protocolStageRecordingActiveKey = "protocol_stage_recording_active" +const protocolStageRecordingKey = "protocol_stage_recording" + +func enableProtocolStageAttemptTracking(c *gin.Context, recordings ...*protocolStageRequestRecording) { + if c != nil { + c.Set(protocolStageRecordingActiveKey, true) + if len(recordings) > 0 && recordings[0] != nil { + c.Set(protocolStageRecordingKey, recordings[0]) + } + } +} + +func observeProtocolStageSetupFailure(c *gin.Context, requestErr error) { + if c == nil || requestErr == nil { + return + } + value, exists := c.Get(protocolStageRecordingKey) + if !exists { + return + } + recording, ok := value.(*protocolStageRequestRecording) + if !ok || recording == nil { + return + } + recording.mu.Lock() + if recording.used { + recording.lastAttemptErr = requestErr + } + recording.mu.Unlock() +} + +func setProtocolStageAttempt(c *gin.Context, attempt int) { + if c == nil { + return + } + if active, ok := c.Get(protocolStageRecordingActiveKey); !ok || active != true { + return + } + c.Set(protocolStageAttemptKey, attempt) +} + +func currentProtocolStageAttempt(c *gin.Context) int { + if c != nil { + if attempt, ok := c.Get(protocolStageAttemptKey); ok { + if value, valid := attempt.(int); valid && value > 0 { + return value + } + } + } + return 1 +} + +func stageRecordingRecorder(recording *protocolStageRequestRecording) *requestrecord.Recorder { + if recording == nil { + return nil + } + return recording.recorder +} + +func captureProtocolStageFinalResponse( + ctx context.Context, + recorder *requestrecord.Recorder, + api protocol.APIType, + response any, +) { + if recorder == nil { + return + } + if err := recorder.SetFinalResponse(api, response); err != nil { + logrus.WithContext(ctx).WithError(err).Debug("Protocol Stage RequestRecord final response capture failed") + } +} + +type protocolStageFinalStreamCapture struct { + recorder *requestrecord.Recorder + protocol protocol.APIType + assembler assembler.StreamAssembler +} + +func newProtocolStageFinalStreamCapture( + ctx context.Context, + recorder *requestrecord.Recorder, + api protocol.APIType, +) *protocolStageFinalStreamCapture { + if recorder == nil { + return nil + } + streamAssembler, err := assembler.NewStreamAssembler(api) + if err != nil { + logrus.WithContext(ctx).WithError(err).Debug("Protocol Stage RequestRecord final stream assembler unavailable") + return nil + } + return &protocolStageFinalStreamCapture{ + recorder: recorder, + protocol: api, + assembler: streamAssembler, + } +} + +func (c *protocolStageFinalStreamCapture) add(ctx context.Context, event any) { + if c == nil || c.assembler == nil { + return + } + if err := c.assembler.Add(event); err != nil { + logrus.WithContext(ctx).WithError(err).Debug("Protocol Stage RequestRecord final stream event capture failed") + c.assembler = nil + } +} + +func (c *protocolStageFinalStreamCapture) finish(ctx context.Context) { + if c == nil || c.assembler == nil || c.recorder == nil { + return + } + response, err := c.assembler.Finish() + if err != nil { + logrus.WithContext(ctx).WithError(err).Debug("Protocol Stage RequestRecord final stream assembly failed") + return + } + captureProtocolStageFinalResponse(ctx, c.recorder, c.protocol, response) +} diff --git a/internal/protocolserver/protocol_stage_recording_test.go b/internal/protocolserver/protocol_stage_recording_test.go new file mode 100644 index 000000000..446c4cc94 --- /dev/null +++ b/internal/protocolserver/protocol_stage_recording_test.go @@ -0,0 +1,44 @@ +package protocolserver + +import ( + "errors" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func TestProtocolStageOriginalInputDisabledDoesNotCaptureBody(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + body := []byte(`{"large":"payload"}`) + + (&ProtocolHandler{}).rememberProtocolStageOriginalInput(c, typ.ScenarioOpenAI, body) + + if _, exists := c.Get(protocolStageOriginalInputKey); exists { + t.Fatal("disabled Protocol Stage recording retained the request body") + } +} + +func TestProtocolStageRecordingTracksSetupFailureAfterStageUse(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + recording := &protocolStageRequestRecording{} + enableProtocolStageAttemptTracking(c, recording) + + observeProtocolStageSetupFailure(c, errors.New("ignored before Stage use")) + recording.mu.Lock() + if recording.lastAttemptErr != nil { + t.Fatalf("unused recording captured setup error: %v", recording.lastAttemptErr) + } + recording.mu.Unlock() + + recording.observeAttempt(errors.New("first provider failed")) + observeProtocolStageSetupFailure(c, errors.New("fallback setup failed")) + recording.mu.Lock() + defer recording.mu.Unlock() + if recording.lastAttemptErr == nil || recording.lastAttemptErr.Error() != "fallback setup failed" { + t.Fatalf("last attempt error = %v", recording.lastAttemptErr) + } +} diff --git a/internal/protocolserver/protocol_stage_selector.go b/internal/protocolserver/protocol_stage_selector.go new file mode 100644 index 000000000..cb15557fa --- /dev/null +++ b/internal/protocolserver/protocol_stage_selector.go @@ -0,0 +1,110 @@ +package protocolserver + +import ( + "fmt" + + "github.com/gin-gonic/gin" + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/anthropicbridge" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/openaibridge" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/responsesbridge" +) + +// ProtocolStageSelector is the immutable process-level selector enabled by +// --stage. It resolves only exact, capability-complete protocol pairs; callers +// keep all other paths on the legacy pipeline. +type ProtocolStageSelector struct { + enabled bool + registry *stage.BridgeRegistry +} + +// NewProtocolStageSelector constructs the production selector. The registry is +// code-defined rather than mutable configuration so active requests cannot see +// a partially updated protocol graph. +func NewProtocolStageSelector(enabled bool) *ProtocolStageSelector { + registry, err := stage.NewBridgeRegistry( + stage.NewIdentityBridge(protocol.TypeAnthropicV1), + anthropicbridge.NewV1ToBeta(), + anthropicbridge.NewV1ToOpenAIChat(anthropicbridge.ChatOptions{}), + anthropicbridge.NewV1ToOpenAIResponses(anthropicbridge.ResponsesOptions{}), + stage.NewIdentityBridge(protocol.TypeAnthropicBeta), + anthropicbridge.NewBetaToOpenAIChat(anthropicbridge.ChatOptions{}), + anthropicbridge.NewBetaToOpenAIResponses(anthropicbridge.ResponsesOptions{}), + stage.NewIdentityBridge(protocol.TypeOpenAIChat), + openaibridge.NewChatToAnthropicBeta(openaibridge.AnthropicOptions{}), + openaibridge.NewChatToOpenAIResponses(openaibridge.ResponsesOptions{}), + stage.NewIdentityBridge(protocol.TypeOpenAIResponses), + responsesbridge.NewToAnthropicBeta(responsesbridge.AnthropicOptions{}), + responsesbridge.NewToOpenAIChat(responsesbridge.ChatOptions{}), + ) + if err != nil { + panic(fmt.Sprintf("construct Protocol Stage selector: %v", err)) + } + return &ProtocolStageSelector{enabled: enabled, registry: registry} +} + +// ShouldUseBetaStageChain returns true only when both explicit boundaries +// around the Beta working protocol are capability-complete. Checking the two +// exact hops prevents a direct source->target Bridge from accidentally claiming +// that a source->Beta->target topology is available. +func (s *ProtocolStageSelector) ShouldUseBetaStageChain( + source, target protocol.APIType, + required stage.Capabilities, +) (bool, error) { + if !s.Enabled() { + return false, nil + } + if s.registry == nil { + return false, fmt.Errorf("Protocol Stage registry is nil") + } + if _, err := s.registry.ResolveRegistered(source, protocol.TypeAnthropicBeta, required); err != nil { + return false, fmt.Errorf("Beta Stage ingress: %w", err) + } + if _, err := s.registry.ResolveRegistered(protocol.TypeAnthropicBeta, target, required); err != nil { + return false, fmt.Errorf("Beta Stage provider boundary: %w", err) + } + return true, nil +} + +func (ph *ProtocolHandler) shouldUseProtocolStageBetaChain( + c *gin.Context, + source, target protocol.APIType, + required stage.Capabilities, +) bool { + selector := ph.protocolStageSelector + if selector == nil { + return false + } + useStage, selectionErr := selector.ShouldUseBetaStageChain(source, target, required) + if !useStage && selector.Enabled() && selectionErr != nil { + logProtocolStageFallback(c, source, target, selectionErr.Error()) + } + return useStage +} + +// Enabled reports the immutable server-start choice. +func (s *ProtocolStageSelector) Enabled() bool { + return s != nil && s.enabled +} + +// ShouldUseStage returns true only when --stage is enabled and an explicitly +// registered exact Bridge satisfies every required capability. A missing route +// is returned as an error so diagnostics can explain why that request stayed on +// legacy. Implicit identity Bridges are intentionally excluded: production +// routes must be opted in one pair at a time. +func (s *ProtocolStageSelector) ShouldUseStage( + source, target protocol.APIType, + required stage.Capabilities, +) (bool, error) { + if !s.Enabled() { + return false, nil + } + if s.registry == nil { + return false, fmt.Errorf("Protocol Stage registry is nil") + } + if _, err := s.registry.ResolveRegistered(source, target, required); err != nil { + return false, err + } + return true, nil +} diff --git a/internal/protocolserver/protocol_stage_selector_test.go b/internal/protocolserver/protocol_stage_selector_test.go new file mode 100644 index 000000000..1d4c96e4d --- /dev/null +++ b/internal/protocolserver/protocol_stage_selector_test.go @@ -0,0 +1,73 @@ +package protocolserver + +import ( + "testing" + + protocol "github.com/tingly-dev/tingly-box/ai" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" +) + +func TestProtocolStageSelector(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + enabled bool + source protocol.APIType + target protocol.APIType + want bool + wantErr bool + }{ + {name: "disabled supported pair", source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta}, + {name: "enabled supported pair", enabled: true, source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta, want: true}, + {name: "enabled chat to responses", enabled: true, source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIResponses, want: true}, + {name: "enabled registered chat identity", enabled: true, source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIChat, want: true}, + {name: "enabled registered beta identity", enabled: true, source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta, want: true}, + {name: "enabled registered v1 identity", enabled: true, source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicV1, want: true}, + {name: "enabled beta to chat", enabled: true, source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIChat, want: true}, + {name: "enabled beta to responses", enabled: true, source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIResponses, want: true}, + {name: "enabled v1 to chat", enabled: true, source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat, want: true}, + {name: "enabled v1 to responses", enabled: true, source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIResponses, want: true}, + {name: "enabled registered responses identity", enabled: true, source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIResponses, want: true}, + {name: "enabled responses to beta", enabled: true, source: protocol.TypeOpenAIResponses, target: protocol.TypeAnthropicBeta, want: true}, + {name: "enabled responses to chat", enabled: true, source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIChat, want: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := NewProtocolStageSelector(tt.enabled).ShouldUseStage(tt.source, tt.target, stage.AllBridgeCapabilities) + if got != tt.want || (err != nil) != tt.wantErr { + t.Fatalf("ShouldUseStage() = %v, %v; want %v, error=%v", got, err, tt.want, tt.wantErr) + } + }) + } +} + +func TestProtocolStageSelectorBetaStageChain(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + enabled bool + source protocol.APIType + target protocol.APIType + want bool + wantErr bool + }{ + {name: "disabled", source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta}, + {name: "beta identity", enabled: true, source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta, want: true}, + {name: "v1 promoted to beta", enabled: true, source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicBeta, want: true}, + {name: "chat through beta to chat", enabled: true, source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIChat, want: true}, + {name: "responses through beta to responses", enabled: true, source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIResponses, want: true}, + {name: "beta cannot target v1", enabled: true, source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicV1, wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := NewProtocolStageSelector(tt.enabled).ShouldUseBetaStageChain(tt.source, tt.target, stage.AllBridgeCapabilities) + if got != tt.want || (err != nil) != tt.wantErr { + t.Fatalf("ShouldUseBetaStageChain() = %v, %v; want %v, error=%v", got, err, tt.want, tt.wantErr) + } + }) + } +} diff --git a/internal/protocolserver/protocol_stage_stream_test.go b/internal/protocolserver/protocol_stage_stream_test.go new file mode 100644 index 000000000..3a7325feb --- /dev/null +++ b/internal/protocolserver/protocol_stage_stream_test.go @@ -0,0 +1,358 @@ +package protocolserver + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/packages/param" + "github.com/openai/openai-go/v3/responses" + "github.com/tingly-dev/tingly-box/internal/protocol" + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" +) + +func TestProtocolStageStreamCommitsFirstClientEvent(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + api protocol.APIType + event protocolstage.Event + serve func(*ProtocolHandler, *gin.Context, protocolstage.Endpoint) + want string + }{ + { + name: "OpenAI Chat", + api: protocol.TypeOpenAIChat, + event: protocolstage.Event{Value: wire.ChatStreamChunk{ + ID: "chatcmpl_test", + Object: "chat.completion.chunk", + Model: "test-model", + Choices: []wire.ChatStreamChoice{{ + Index: 0, + Delta: wire.ChatStreamDelta{Content: "hello"}, + }}, + }}, + serve: func(ph *ProtocolHandler, c *gin.Context, endpoint protocolstage.Endpoint) { + ph.serveProtocolStageOpenAIChatStream(c, endpoint, protocolstage.Call{}, "test-model", false, nil, nil) + }, + want: "hello", + }, + { + name: "Anthropic Beta", + api: protocol.TypeAnthropicBeta, + event: protocolstage.Event{Value: protocolstream.AnthropicEvent{Type: "content_block_delta", Data: map[string]any{"type": "content_block_delta", "delta": map[string]any{"type": "text_delta", "text": "hello"}}}}, + serve: func(ph *ProtocolHandler, c *gin.Context, endpoint protocolstage.Endpoint) { + ph.serveProtocolStageAnthropicBetaStream(c, endpoint, protocolstage.Call{}, "test-model", nil, nil) + }, + want: "content_block_delta", + }, + { + name: "Anthropic V1", + api: protocol.TypeAnthropicV1, + event: protocolstage.Event{Value: protocolstream.AnthropicEvent{Type: "content_block_delta", Data: map[string]any{"type": "content_block_delta", "delta": map[string]any{"type": "text_delta", "text": "hello"}}}}, + serve: func(ph *ProtocolHandler, c *gin.Context, endpoint protocolstage.Endpoint) { + ph.serveProtocolStageAnthropicV1Stream(c, endpoint, protocolstage.Call{}, "test-model", nil, nil) + }, + want: "content_block_delta", + }, + { + name: "OpenAI Responses", + api: protocol.TypeOpenAIResponses, + event: protocolstage.Event{Value: wire.ResponsesOutputTextDeltaEvent{ + Type: "response.output_text.delta", Delta: "hello", + }}, + serve: func(ph *ProtocolHandler, c *gin.Context, endpoint protocolstage.Endpoint) { + ph.serveProtocolStageOpenAIResponsesStream(c, endpoint, protocolstage.Call{}, "test-model", nil) + }, + want: "hello", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c, recorder, gate := newProtocolStageGateContext() + stream := &protocolStageTestStream{events: []protocolstage.Event{tt.event}} + endpoint := protocolStageTestEndpoint{api: tt.api, stream: stream} + + tt.serve(&ProtocolHandler{}, c, endpoint) + + if !gate.Committed() { + t.Fatal("first client-visible stream event did not commit the failover gate") + } + if !strings.Contains(recorder.Body.String(), tt.want) { + t.Fatalf("response body %q does not contain %q", recorder.Body.String(), tt.want) + } + if stream.closeCalls != 1 { + t.Fatalf("stream Close called %d times, want 1", stream.closeCalls) + } + }) + } +} + +func TestProtocolStageOpenAIChatStreamCancellationDoesNotCommitError(t *testing.T) { + gin.SetMode(gin.TestMode) + + c, recorder, gate := newProtocolStageGateContext() + stream := &protocolStageTestStream{terminalErr: context.Canceled} + endpoint := protocolStageTestEndpoint{api: protocol.TypeOpenAIChat, stream: stream} + + err := (&ProtocolHandler{}).serveProtocolStageOpenAIChatStream(c, endpoint, protocolstage.Call{}, "test-model", false, nil, nil) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("stream error = %v, want context.Canceled", err) + } + if gate.Committed() { + t.Fatal("client cancellation committed an error response") + } + if recorder.Body.Len() != 0 { + t.Fatalf("client cancellation wrote response body %q", recorder.Body.String()) + } + if stream.closeCalls != 1 { + t.Fatalf("stream Close called %d times, want 1", stream.closeCalls) + } +} + +func TestProtocolStageOpenAIChatEmptyStreamReturnsRetryableError(t *testing.T) { + gin.SetMode(gin.TestMode) + + c, recorder, gate := newProtocolStageGateContext() + stream := &protocolStageTestStream{} + endpoint := protocolStageTestEndpoint{api: protocol.TypeOpenAIChat, stream: stream} + + err := (&ProtocolHandler{}).serveProtocolStageOpenAIChatStream(c, endpoint, protocolstage.Call{}, "test-model", false, nil, nil) + + if err == nil || !strings.Contains(err.Error(), "without a terminal finish_reason") { + t.Fatalf("stream error = %v, want missing terminal finish_reason", err) + } + if gate.Committed() { + t.Fatal("empty stream error committed the failover gate") + } + if recorder.Body.Len() != 0 { + t.Fatalf("empty stream wrote response body %q", recorder.Body.String()) + } +} + +func TestProtocolStageOpenAIChatTruncatedStreamEmitsError(t *testing.T) { + gin.SetMode(gin.TestMode) + + c, recorder, gate := newProtocolStageGateContext() + stream := &protocolStageTestStream{events: []protocolstage.Event{{ + Value: wire.ChatStreamChunk{ + ID: "chat-truncated", Object: "chat.completion.chunk", Model: "provider-model", + Choices: []wire.ChatStreamChoice{{ + Index: 0, Delta: wire.ChatStreamDelta{Content: "partial"}, + }}, + }, + }}} + endpoint := protocolStageTestEndpoint{api: protocol.TypeOpenAIChat, stream: stream} + + err := (&ProtocolHandler{}).serveProtocolStageOpenAIChatStream( + c, endpoint, protocolstage.Call{}, "public-model", false, nil, nil, + ) + if err == nil || !strings.Contains(err.Error(), "without a terminal finish_reason") { + t.Fatalf("stream error = %v, want missing terminal finish_reason", err) + } + if !gate.Committed() { + t.Fatal("partial Chat event did not commit the failover gate") + } + body := recorder.Body.String() + if !strings.Contains(body, "partial") || !strings.Contains(body, "protocol_error") || strings.Contains(body, "[DONE]") { + t.Fatalf("truncated Chat stream did not expose a terminal error: %q", body) + } +} + +func TestProtocolStageOpenAIChatMultiChoiceRequiresEveryFinishReason(t *testing.T) { + gin.SetMode(gin.TestMode) + + c, recorder, _ := newProtocolStageGateContext() + stop := "stop" + stream := &protocolStageTestStream{events: []protocolstage.Event{{ + Value: wire.ChatStreamChunk{ + ID: "chat-multi", Object: "chat.completion.chunk", Model: "provider-model", + Choices: []wire.ChatStreamChoice{{ + Index: 0, Delta: wire.ChatStreamDelta{Content: "first"}, FinishReason: &stop, + }, { + Index: 1, Delta: wire.ChatStreamDelta{Content: "partial"}, + }}, + }, + }}} + endpoint := protocolStageTestEndpoint{api: protocol.TypeOpenAIChat, stream: stream} + call := protocolstage.Call{Request: &openai.ChatCompletionNewParams{N: param.NewOpt(int64(2))}} + + err := (&ProtocolHandler{}).serveProtocolStageOpenAIChatStream( + c, endpoint, call, "public-model", false, nil, nil, + ) + if err == nil || !strings.Contains(err.Error(), "without a terminal finish_reason") { + t.Fatalf("stream error = %v, want missing terminal finish_reason", err) + } + body := recorder.Body.String() + if !strings.Contains(body, "protocol_error") || strings.Contains(body, "[DONE]") { + t.Fatalf("partial multi-choice stream was marked complete: %q", body) + } +} + +func TestProtocolStageOpenAIChatCrossProtocolRejectsMultipleChoices(t *testing.T) { + request := &openai.ChatCompletionNewParams{N: param.NewOpt(int64(2))} + if protocolStageOpenAIChatChoiceCompatible(protocol.TypeAnthropicBeta, request) { + t.Fatal("Anthropic target accepted an unrepresentable multi-choice Chat request") + } + if protocolStageOpenAIChatChoiceCompatible(protocol.TypeOpenAIResponses, request) { + t.Fatal("Responses target accepted an unrepresentable multi-choice Chat request") + } + if !protocolStageOpenAIChatChoiceCompatible(protocol.TypeOpenAIChat, request) { + t.Fatal("Chat identity target rejected its native multi-choice request") + } +} + +func TestProtocolStageOpenAIResponsesTruncatedStreamEmitsError(t *testing.T) { + gin.SetMode(gin.TestMode) + + c, recorder, gate := newProtocolStageGateContext() + stream := &protocolStageTestStream{events: []protocolstage.Event{{ + Value: wire.ResponsesCreatedEvent{ + Type: "response.created", + Response: wire.ResponsesWireResponse{ + ID: "resp_test", Object: "response", Status: "in_progress", + }, + }, + }}} + endpoint := protocolStageTestEndpoint{api: protocol.TypeOpenAIResponses, stream: stream} + + (&ProtocolHandler{}).serveProtocolStageOpenAIResponsesStream(c, endpoint, protocolstage.Call{}, "test-model", nil) + + if !gate.Committed() { + t.Fatal("client-visible Responses event did not commit the failover gate") + } + body := recorder.Body.String() + if !strings.Contains(body, "response.created") || !strings.Contains(body, "stream_failed") { + t.Fatalf("truncated stream body does not include partial event and terminal error: %q", body) + } + if stream.closeCalls != 1 { + t.Fatalf("stream Close called %d times, want 1", stream.closeCalls) + } +} + +func TestProtocolStageOpenAIResponsesFailedStreamRecordsFailurePayload(t *testing.T) { + gin.SetMode(gin.TestMode) + + c, responseWriter, gate := newProtocolStageGateContext() + recorder, err := requestrecord.New(requestrecord.Config{ + Enabled: true, + InputProtocol: protocol.TypeOpenAIResponses, + Input: map[string]any{"model": "client-model"}, + }) + if err != nil { + t.Fatalf("new recorder: %v", err) + } + var failedEvent responses.ResponseStreamEventUnion + if err := json.Unmarshal([]byte(`{"type":"response.failed","sequence_number":1,"response":{"id":"resp-failed","object":"response","status":"failed","model":"provider-model","output":[],"error":{"code":"server_error","message":"provider failed"}}}`), &failedEvent); err != nil { + t.Fatalf("decode failed event: %v", err) + } + stream := &protocolStageTestStream{events: []protocolstage.Event{{Value: failedEvent}}} + terminal := protocolStageTestEndpoint{api: protocol.TypeOpenAIResponses, stream: stream} + endpoint := requestrecord.ObserveProvider(terminal, recorder, requestrecord.ExchangeMetadata{ + Attempt: 1, Provider: "provider", Model: "provider-model", + }) + + serveErr := (&ProtocolHandler{}).serveProtocolStageOpenAIResponsesStream( + c, + endpoint, + protocolstage.Call{Request: map[string]any{"model": "provider-model"}}, + "public-model", + recorder, + ) + if serveErr == nil || !strings.Contains(serveErr.Error(), "response.failed") { + t.Fatalf("stream error = %v, want response.failed", serveErr) + } + if !gate.Committed() { + t.Fatal("response.failed event did not reach the client") + } + body := responseWriter.Body.String() + if !strings.Contains(body, "response.failed") || strings.Contains(body, "stream_failed") { + t.Fatalf("failed stream must preserve its terminal event without a synthetic error: %q", body) + } + + completed, first := recorder.Finish(serveErr) + if !first { + t.Fatal("recorder did not finish") + } + if completed.Outcome != requestrecord.OutcomeFailed { + t.Fatalf("request outcome = %q, want failed", completed.Outcome) + } + if len(completed.ProviderExchanges) != 1 { + t.Fatalf("provider exchanges = %d, want 1", len(completed.ProviderExchanges)) + } + exchange := completed.ProviderExchanges[0] + if exchange.Outcome != requestrecord.OutcomeFailed || exchange.Response == nil { + t.Fatalf("provider exchange = %#v, want failed with response", exchange) + } + if completed.FinalResponse == nil || !strings.Contains(string(completed.FinalResponse.Body), `"status":"failed"`) { + t.Fatalf("final response = %#v, want failed response payload", completed.FinalResponse) + } +} + +func newProtocolStageGateContext() (*gin.Context, *httptest.ResponseRecorder, *firstChunkGate) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest("POST", "/", nil) + gate := newFirstChunkGate(c.Writer) + c.Writer = gate + return c, recorder, gate +} + +type protocolStageTestEndpoint struct { + api protocol.APIType + stream protocolstage.EventStream +} + +func (e protocolStageTestEndpoint) Protocol() protocol.APIType { return e.api } + +func (protocolStageTestEndpoint) Complete(context.Context, protocolstage.Call) (*protocolstage.Response, error) { + return nil, errors.New("unexpected Complete call") +} + +func (e protocolStageTestEndpoint) Stream(context.Context, protocolstage.Call) (protocolstage.EventStream, error) { + return e.stream, nil +} + +type protocolStageTestStream struct { + events []protocolstage.Event + terminalErr error + closeCalls int +} + +func (s *protocolStageTestStream) Next(ctx context.Context) (protocolstage.Event, error) { + if err := ctx.Err(); err != nil { + return protocolstage.Event{}, err + } + if len(s.events) > 0 { + event := s.events[0] + s.events = s.events[1:] + return event, nil + } + if s.terminalErr != nil { + err := s.terminalErr + s.terminalErr = nil + return protocolstage.Event{}, err + } + return protocolstage.Event{}, io.EOF +} + +func (*protocolStageTestStream) Result() protocolstage.StreamResult { + return protocolstage.StreamResult{} +} + +func (s *protocolStageTestStream) Close() error { + s.closeCalls++ + return nil +} diff --git a/internal/protocolserver/protocol_stage_tool_loop.go b/internal/protocolserver/protocol_stage_tool_loop.go new file mode 100644 index 000000000..d26133873 --- /dev/null +++ b/internal/protocolserver/protocol_stage_tool_loop.go @@ -0,0 +1,50 @@ +package protocolserver + +import ( + "fmt" + + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + + protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" + stagetoolloop "github.com/tingly-dev/tingly-box/internal/protocol/stage/toolloop" + mcpmodule "github.com/tingly-dev/tingly-box/internal/mcpserver" + servertransform "github.com/tingly-dev/tingly-box/internal/protocolserver/transform" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func (ph *ProtocolHandler) newProtocolStageBetaToolLoop( + c *gin.Context, + provider *typ.Provider, + hasNativeAdvisor bool, +) (protocolstage.Stage, error) { + if ph == nil || ph.deps.MCPRuntime == nil { + return nil, fmt.Errorf("construct Beta Tool Loop: MCP runtime is nil") + } + if provider == nil || provider.UUID == "" { + return nil, fmt.Errorf("construct Beta Tool Loop: provider identity is empty") + } + return mcpmodule.NewAnthropicBetaStage(mcpmodule.AnthropicBetaStageConfig{ + Tools: servertransform.NewProtocolStageBetaToolProvider( + ph.deps.MCPRuntime, + c.GetHeader("X-Tingly-Advisor-Depth") != "", + hasNativeAdvisor, + ), + Executor: mcpmodule.NewServerToolExecutor(ph), + Continuations: mcpmodule.NewProviderBetaContinuationStore(provider.UUID), + }) +} + +// preserveProtocolStageSideEffectBoundary prevents provider failover from +// replaying a successful server tool when a later round or outward conversion +// fails before any client-visible output has committed. +func preserveProtocolStageSideEffectBoundary(c *gin.Context, err error, committed bool) { + if c == nil || (!committed && !stagetoolloop.HasCommittedSideEffects(err)) { + return + } + MarkSideEffectsCommittedIfGate(c.Writer) + logrus.WithContext(c.Request.Context()).WithFields(logrus.Fields{ + "protocol_pipeline": "stage", + "side_effects_committed": true, + }).Debug("Protocol Stage disabled failover after tool side effects") +} diff --git a/internal/protocolserver/protocol_transform_test.go b/internal/protocolserver/protocol_transform_test.go index 0fd1c3bf1..0de7847f7 100644 --- a/internal/protocolserver/protocol_transform_test.go +++ b/internal/protocolserver/protocol_transform_test.go @@ -77,3 +77,29 @@ func TestBuildTransformChain_PreBaseBeforeBase(t *testing.T) { assert.Equal(t, 0, cursor, "pre-Base rule transform must be first in the chain") assert.Less(t, cursor, base, "pre-Base rule transform must run before base_convert") } + +func TestProtocolStageProviderTransformsPutChatCleanupLast(t *testing.T) { + transforms := protocolStageProviderTransforms( + protocol.TypeOpenAIChat, + []transform.Transform{ + transform.NewConsistencyTransform(protocol.TypeOpenAIChat), + transform.NewVendorTransform(), + }, + ) + names := make([]string, 0, len(transforms)) + for _, current := range transforms { + names = append(names, current.Name()) + } + assert.Equal(t, []string{ + "consistency_normalize", + "vendor_adjust", + "openai_chat_provider_cleanup", + }, names) + + nonChat := protocolStageProviderTransforms( + protocol.TypeAnthropicBeta, + []transform.Transform{transform.NewVendorTransform()}, + ) + require.Len(t, nonChat, 1) + assert.Equal(t, "vendor_adjust", nonChat[0].Name()) +} diff --git a/internal/protocolserver/servertool/executor.go b/internal/protocolserver/servertool/executor.go index 210c07904..1a256814b 100644 --- a/internal/protocolserver/servertool/executor.go +++ b/internal/protocolserver/servertool/executor.go @@ -3,6 +3,7 @@ package servertool import ( "context" "encoding/json" + "errors" "fmt" coretool "github.com/tingly-dev/tingly-box/internal/tool" @@ -28,6 +29,21 @@ type Executor interface { Execute(ctx context.Context, call ToolCall) (context.Context, coretool.ToolResult, error) } +// DispatchError marks an error returned after the runtime invocation began. +// Callers use this boundary to prevent failover from replaying a tool whose +// remote side effects may have happened even though the call returned an error. +type DispatchError struct { + Err error +} + +func (e *DispatchError) Error() string { return e.Err.Error() } +func (e *DispatchError) Unwrap() error { return e.Err } + +func WasDispatched(err error) bool { + var dispatched *DispatchError + return errors.As(err, &dispatched) +} + // DefaultExecutor is the standard implementation of Executor. type DefaultExecutor struct { runtime RuntimeCaller @@ -73,6 +89,7 @@ func (e *DefaultExecutor) Execute(ctx context.Context, call ToolCall) (context.C result, err := e.runtime.CallTool(ctx, resolvedName, call.Arguments) if err != nil { result = normalizeError(err) + err = &DispatchError{Err: err} } // 6. Restore depth so it doesn't accumulate across sequential tool calls. diff --git a/internal/protocolserver/servertool/executor_test.go b/internal/protocolserver/servertool/executor_test.go new file mode 100644 index 000000000..80eedccb1 --- /dev/null +++ b/internal/protocolserver/servertool/executor_test.go @@ -0,0 +1,41 @@ +package servertool + +import ( + "context" + "errors" + "testing" + + coretool "github.com/tingly-dev/tingly-box/internal/tool" +) + +type errorRuntime struct { + err error + calls int +} + +func (r *errorRuntime) CallTool(context.Context, string, string) (coretool.ToolResult, error) { + r.calls++ + return coretool.ToolResult{}, r.err +} + +func (*errorRuntime) ListCallableServerToolNames(context.Context) map[string]struct{} { + return map[string]struct{}{"tingly_box_mcp__test__run": {}} +} + +func (*errorRuntime) GetAdvisorMaxUses() int { return 0 } + +func TestDefaultExecutorMarksOnlyPostDispatchErrors(t *testing.T) { + runtimeErr := errors.New("runtime failed") + runtime := &errorRuntime{err: runtimeErr} + executor := NewDefaultExecutor(runtime, nil) + + _, _, err := executor.Execute(context.Background(), ToolCall{NormalizedName: "tingly_box_mcp__test__run"}) + if !errors.Is(err, runtimeErr) || !WasDispatched(err) || runtime.calls != 1 { + t.Fatalf("runtime error = %v, dispatched=%v, calls=%d", err, WasDispatched(err), runtime.calls) + } + + _, _, err = executor.Execute(context.Background(), ToolCall{NormalizedName: "not-an-mcp-tool"}) + if err == nil || WasDispatched(err) || runtime.calls != 1 { + t.Fatalf("validation error = %v, dispatched=%v, calls=%d", err, WasDispatched(err), runtime.calls) + } +} diff --git a/internal/protocolserver/transform/protocol_stage_beta_tools.go b/internal/protocolserver/transform/protocol_stage_beta_tools.go new file mode 100644 index 000000000..9e89f2c1a --- /dev/null +++ b/internal/protocolserver/transform/protocol_stage_beta_tools.go @@ -0,0 +1,72 @@ +package transform + +import ( + "context" + + "github.com/anthropics/anthropic-sdk-go" + + "github.com/tingly-dev/tingly-box/internal/mcp/runtime" +) + +// ProtocolStageBetaToolProvider prepares the Beta working request with the +// exact MCP server tools owned by the Protocol Stage ToolLoop. It reuses the +// existing merge and Advisor-system behavior without invoking a generic LLM +// tool representation. +type ProtocolStageBetaToolProvider struct { + runtime *runtime.Runtime + isAdvisorRequest bool + hasNativeAdvisor bool +} + +func NewProtocolStageBetaToolProvider( + rt *runtime.Runtime, + isAdvisorRequest bool, + hasNativeAdvisor bool, +) *ProtocolStageBetaToolProvider { + return &ProtocolStageBetaToolProvider{ + runtime: rt, + isAdvisorRequest: isAdvisorRequest, + hasNativeAdvisor: hasNativeAdvisor, + } +} + +func (p *ProtocolStageBetaToolProvider) PrepareRequest( + ctx context.Context, + request *anthropic.BetaMessageNewParams, +) ([]string, error) { + if p == nil || p.runtime == nil || request == nil || p.isAdvisorRequest { + return nil, nil + } + tools := p.runtime.ListServerToolsForAnthropicBetaInjection(ctx) + if len(tools) == 0 { + return nil, nil + } + if p.hasNativeAdvisor { + filtered := make([]anthropic.BetaToolUnionParam, 0, len(tools)) + for _, tool := range tools { + if tool.OfTool != nil && tool.OfTool.Name == advisorInjectedToolName { + continue + } + filtered = append(filtered, tool) + } + tools = filtered + } + if len(tools) == 0 { + return nil, nil + } + request.Tools = mergeUniqueAnthropicBetaTools(request.Tools, tools) + names := extractAnthropicBetaToolNames(tools) + if containsString(names, advisorInjectedToolName) { + request.System = appendAdvisorBehaviorToAnthropicBetaSystem(request.System) + } + return names, nil +} + +func containsString(values []string, target string) bool { + for _, value := range values { + if value == target { + return true + } + } + return false +} diff --git a/internal/protocolserver/transform/protocol_stage_beta_tools_test.go b/internal/protocolserver/transform/protocol_stage_beta_tools_test.go new file mode 100644 index 000000000..92919a95b --- /dev/null +++ b/internal/protocolserver/transform/protocol_stage_beta_tools_test.go @@ -0,0 +1,49 @@ +package transform + +import ( + "context" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/mark3labs/mcp-go/mcp" + + "github.com/tingly-dev/tingly-box/internal/mcp/runtime" + coretool "github.com/tingly-dev/tingly-box/internal/tool" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func TestProtocolStageBetaToolProviderPreservesAdvisorPreparation(t *testing.T) { + cfg := &typ.MCPRuntimeConfig{Sources: []typ.MCPSourceConfig{{ + ID: "advisor", Name: "advisor", Transport: "advisor", Enabled: typ.BoolPtr(true), + Visibility: typ.ToolVisibilityServer, Tools: []string{"advisor"}, + }}} + runtime := runtime.NewRuntime(func() *typ.MCPRuntimeConfig { return cfg }) + t.Cleanup(runtime.Close) + runtime.VirtualRegistry().Register(coretool.VirtualTool{ + Name: "advisor", Description: "server-side advisor", + InputSchema: mcp.ToolInputSchema{Type: "object"}, Visibility: typ.ToolVisibilityServer, + }) + + provider := NewProtocolStageBetaToolProvider(runtime, false, false) + request := &anthropic.BetaMessageNewParams{} + owned, err := provider.PrepareRequest(context.Background(), request) + if err != nil { + t.Fatal(err) + } + if len(owned) != 1 || owned[0] != advisorInjectedToolName || !hasAnthropicBetaToolName(request.Tools, advisorInjectedToolName) { + t.Fatalf("owned=%#v tools=%#v", owned, request.Tools) + } + if len(request.System) == 0 { + t.Fatal("Advisor behavior was not added to the Beta system prompt") + } + + native := NewProtocolStageBetaToolProvider(runtime, false, true) + nativeRequest := &anthropic.BetaMessageNewParams{} + nativeOwned, err := native.PrepareRequest(context.Background(), nativeRequest) + if err != nil { + t.Fatal(err) + } + if len(nativeOwned) != 0 || len(nativeRequest.Tools) != 0 || len(nativeRequest.System) != 0 { + t.Fatalf("native advisor was duplicated: owned=%#v request=%#v", nativeOwned, nativeRequest) + } +} From 7bd82d133075f6bd54eddb278a2d40ad3effeb1a Mon Sep 17 00:00:00 2001 From: FFengIll Date: Tue, 4 Aug 2026 19:50:44 +0800 Subject: [PATCH 4/8] feat(server): wire Protocol Stage opt-in into the server skeleton Complete the server-side wiring for the ported Protocol Stage pipeline (hardening's server.go/server_lifecycle.go/server_options.go changes, 3-way merged against main's post-#1493..#1495 skeleton): - server.go: pass ProtocolStageEnabled into protocolserver.ProtocolHandlerDeps; add protocolStageEnabled and servertoolProviders fields to the Server struct. - server_options.go: add WithProtocolStage and WithServertoolProviders options (servertool import repointed to internal/protocolserver/servertool). - server_lifecycle.go: add ForceFlushRecordings (auto-merged clean; used by the protocol harness so a run leaves a complete recording artifact before exit). The two real conflicts (server.go handler construction, server_options.go import) resolved in favor of main's protocolserver layout, retaining hardening's additive ProtocolStageEnabled/servertoolProviders surface. go build ./... is clean end-to-end (server + cli). Batch 4 of the protocol-stage-hardening port. --- internal/server/server.go | 9 ++++++++- internal/server/server_lifecycle.go | 26 ++++++++++++++++++++++++++ internal/server/server_options.go | 18 ++++++++++++++++++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/internal/server/server.go b/internal/server/server.go index 508c8701c..21a13dbdb 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -114,13 +114,19 @@ type Server struct { mcpRuntime *mcpruntime.Runtime // servertool pipeline — owns virtual tool providers and hook list - servertoolPipeline *servertool.Pipeline + servertoolPipeline *servertool.Pipeline + servertoolProviders []servertool.ToolProvider // guardrails runtime state (owned by protocolserver; constructed in // NewServer before anything reads it) guardrailsState *protocolserver.GuardrailsState guardrailsConfigMu sync.Mutex + // protocolStageEnabled is an immutable process-start choice gating the + // additive Protocol Stage request pipeline. Unsupported protocol paths + // remain on the legacy pipeline regardless of this flag. + protocolStageEnabled bool + // recording sinks recordSink *obs.Sink @@ -501,6 +507,7 @@ func NewServer(cfg *config.Config, opts ...ServerOption) *Server { GetOrCreateScenarioSink: server.GetOrCreateScenarioSink, GuardrailsState: server.guardrailsState, GetScenarioRecordMode: server.GetScenarioRecordMode, + ProtocolStageEnabled: server.protocolStageEnabled, }) // Setup middleware diff --git a/internal/server/server_lifecycle.go b/internal/server/server_lifecycle.go index 98feb2c28..dd38f4735 100644 --- a/internal/server/server_lifecycle.go +++ b/internal/server/server_lifecycle.go @@ -3,6 +3,7 @@ package server import ( "cmp" "context" + "errors" "fmt" "github.com/tingly-dev/tingly-box/internal/protocolserver" "log" @@ -23,6 +24,31 @@ import ( "github.com/tingly-dev/tingly-box/vmodel/virtualserver" ) +// ForceFlushRecordings waits until every scenario recording sink has exported +// its queued records. It is used by diagnostics and the protocol harness so a +// successful run leaves a complete artifact before the test server exits. +func (s *Server) ForceFlushRecordings(ctx context.Context) error { + if s == nil { + return nil + } + s.scenarioRecordSinksMu.RLock() + sinks := make([]*obs.Sink, 0, len(s.scenarioRecordSinks)) + for _, sink := range s.scenarioRecordSinks { + if sink != nil { + sinks = append(sinks, sink) + } + } + s.scenarioRecordSinksMu.RUnlock() + + var flushErrs []error + for _, sink := range sinks { + if err := sink.ForceFlush(ctx); err != nil { + flushErrs = append(flushErrs, err) + } + } + return errors.Join(flushErrs...) +} + // Start starts the HTTP server func (s *Server) Start(port int) error { // Start token refresher background goroutine diff --git a/internal/server/server_options.go b/internal/server/server_options.go index c889c7028..cfaf7613a 100644 --- a/internal/server/server_options.go +++ b/internal/server/server_options.go @@ -9,6 +9,7 @@ import ( "github.com/tingly-dev/tingly-box/internal/guardrails" "github.com/tingly-dev/tingly-box/internal/obs" "github.com/tingly-dev/tingly-box/internal/protocolserver/recording" + "github.com/tingly-dev/tingly-box/internal/protocolserver/servertool" "github.com/tingly-dev/tingly-box/internal/typ" pkgobs "github.com/tingly-dev/tingly-box/pkg/obs" ) @@ -121,6 +122,23 @@ func WithHTTPTimeouts(t HTTPTimeouts) ServerOption { } } +// WithProtocolStage enables the additive Protocol Stage request pipeline. +// Unsupported protocol paths remain on the legacy pipeline. +func WithProtocolStage(enabled bool) ServerOption { + return func(s *Server) { + s.protocolStageEnabled = enabled + } +} + +// WithServertoolProviders registers in-process server-owned tools with the MCP +// runtime and Tool Loop. Providers remain registered when configuration hot +// reload rebuilds the servertool pipeline. +func WithServertoolProviders(providers ...servertool.ToolProvider) ServerOption { + return func(s *Server) { + s.servertoolProviders = append(s.servertoolProviders, providers...) + } +} + // WithMultiLogger sets the multi-mode logger for the server func WithMultiLogger(logger *pkgobs.MultiLogger) ServerOption { return func(s *Server) { From 84b58926be55785632872998aaaa8491373c24e2 Mon Sep 17 00:00:00 2001 From: FFengIll Date: Tue, 4 Aug 2026 20:59:19 +0800 Subject: [PATCH 5/8] fix(server): wire servertool providers into the pipeline and construct MCP runtime after builtin registration WithServertoolProviders stored providers on the Server but nothing ever read them, so in-process server-owned tools never reached the MCP runtime's virtual registry. Re-register them in registerAdviserFromConfig so they survive pipeline rebuilds and config hot reload. Also restore the integration branch's construction order: RegisterBuiltinTools must run before NewRuntime. The port created the runtime first, so on a first-run config (no persisted MCP runtime yet) NewRuntime saw a nil config and returned nil, disabling the whole stage tool loop and leaving requests to fall back to legacy. Registering builtins first seeds the MCP runtime config and yields a usable runtime on first boot. --- internal/server/server.go | 60 ++++++++++++++++++++++----------------- 1 file changed, 34 insertions(+), 26 deletions(-) diff --git a/internal/server/server.go b/internal/server/server.go index 21a13dbdb..dde8c2979 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -114,8 +114,8 @@ type Server struct { mcpRuntime *mcpruntime.Runtime // servertool pipeline — owns virtual tool providers and hook list - servertoolPipeline *servertool.Pipeline - servertoolProviders []servertool.ToolProvider + servertoolPipeline *servertool.Pipeline + servertoolProviders []servertool.ToolProvider // guardrails runtime state (owned by protocolserver; constructed in // NewServer before anything reads it) @@ -397,12 +397,14 @@ func NewServer(cfg *config.Config, opts ...ServerOption) *Server { // Set template manager in config for model fetching fallback server.config.SetTemplateManager(templateManager) - server.mcpRuntime = mcpruntime.NewRuntime(cfg.GetMCPRuntimeConfig) - server.mcpRuntime.SetClientPool(server.clientPool) // Auto-register built-in tools (e.g., webtools) if not already present if err := mcpruntime.RegisterBuiltinTools(cfg.GetMCPRuntimeConfig, cfg.SetToolConfig); err != nil { logrus.WithError(err).Warn("mcp: failed to register builtin tools") } + // Construct the runtime after registration so a first-run config has the + // same usable MCP dependency graph as subsequent restarts. + server.mcpRuntime = mcpruntime.NewRuntime(cfg.GetMCPRuntimeConfig) + server.mcpRuntime.SetClientPool(server.clientPool) // Register adviser as virtual tool if configured server.registerAdviserFromConfig() @@ -571,34 +573,40 @@ func (s *Server) Cancel() context.CancelFunc { } // registerAdviserFromConfig reads the MCP config and registers the adviser -// virtual tool if an enabled advisor source is found. +// virtual tool if an enabled advisor source is found, plus any in-process +// server tools supplied through ServerOptions. func (s *Server) registerAdviserFromConfig() { - mcpCfg := s.mcpRuntime.GetConfig() - if mcpCfg == nil { - s.servertoolPipeline = servertool.NewPipeline() - return - } - for _, source := range mcpCfg.Sources { - if source.Advisor == nil || source.Enabled == nil || !*source.Enabled { - continue + pipeline := servertool.NewPipeline() + if s.mcpRuntime != nil { + if mcpCfg := s.mcpRuntime.GetConfig(); mcpCfg != nil { + for _, source := range mcpCfg.Sources { + if source.Advisor == nil || source.Enabled == nil || !*source.Enabled { + continue + } + advisorCfg := *source.Advisor + if advisorCfg.ProviderResolver == nil { + advisorCfg.ProviderResolver = s.config.GetProviderByUUID + } + pipeline.Register(advisortool.NewProvider(advisorCfg, s.clientPool, s.mcpRuntime.SessionStore())) + logrus.Info("mcp: registered adviser via servertool pipeline") + break + } } - advisorCfg := *source.Advisor - if advisorCfg.ProviderResolver == nil { - advisorCfg.ProviderResolver = s.config.GetProviderByUUID + // Re-register the in-process server tools supplied by ServerOptions so + // they survive every pipeline rebuild, including config hot reload. The + // virtual registry is name-keyed, so re-registration is idempotent. + for _, provider := range s.servertoolProviders { + if provider != nil { + pipeline.Register(provider) + } } - pipeline := servertool.NewPipeline() - pipeline.Register(advisortool.NewProvider(advisorCfg, s.clientPool, s.mcpRuntime.SessionStore())) - pipeline.RegisterInto(s.mcpRuntime.VirtualRegistry()) - s.servertoolPipeline = pipeline - - logrus.Info("mcp: registered adviser via servertool pipeline") - return + if registry := s.mcpRuntime.VirtualRegistry(); registry != nil { + pipeline.RegisterInto(registry) + } } - - // No advisor configured — empty pipeline. - s.servertoolPipeline = servertool.NewPipeline() + s.servertoolPipeline = pipeline } // setupConfigWatcher initializes the configuration hot-reload watcher From 624ba0f20fc298de44cf1f22ba5716892180dfb5 Mon Sep 17 00:00:00 2001 From: FFengIll Date: Tue, 4 Aug 2026 20:59:15 +0800 Subject: [PATCH 6/8] fix(mcpserver): carry owned-tool results across external rounds and gate stash on session Problem 1: a pure-managed round followed by a pure-external round returned the external round with no continuation, so the server-tool results from the earlier round never reached the model's final answer on the client's next turn. Track internal rounds accumulated during the loop and stash them ahead of the external assistant message; generalize the merge so a trailing assistant segment receives the client's results as a new user turn instead of folding them into a managed-results message. Problem 2: on IP-fallback / no-session requests the continuation store silently no-oped, so a mixed round executed the server-owned tool, filtered its block, and dropped the result with no error. Put now reports failure and the store exposes CanStash; the stage checks it before executing owned tools and returns ErrContinuationUnavailable instead of committing side effects it cannot carry. An external-only round after internal rounds degrades to delivering the round without a stash (guardrail blocking still works). --- internal/mcpserver/anthropic_beta_adapter.go | 45 +++-- internal/mcpserver/anthropic_beta_stage.go | 91 ++++++++- .../mcpserver/anthropic_beta_stage_stream.go | 34 +++- .../mcpserver/anthropic_beta_stage_test.go | 181 +++++++++++++++++- internal/mcpserver/continuation_store.go | 25 ++- internal/protocol/stage/toolloop/runtime.go | 5 +- 6 files changed, 357 insertions(+), 24 deletions(-) diff --git a/internal/mcpserver/anthropic_beta_adapter.go b/internal/mcpserver/anthropic_beta_adapter.go index f1cd17830..ed6c9528f 100644 --- a/internal/mcpserver/anthropic_beta_adapter.go +++ b/internal/mcpserver/anthropic_beta_adapter.go @@ -117,20 +117,7 @@ func (a *AnthropicBetaAdapter) AppendToolResults(req, resp any, results []any) ( newReq := *reqParams newMessages := append([]anthropic.BetaMessageParam{}, reqParams.Messages...) newMessages = append(newMessages, betaMessageToParamPreservingThinking(msg)) - - // Convert results to tool result blocks - resultBlocks := make([]anthropic.BetaContentBlockParamUnion, len(results)) - for i, r := range results { - tr := r.(ToolExecutionResult) - resultBlocks[i] = anthropic.BetaContentBlockParamUnion{ - OfToolResult: &anthropic.BetaToolResultBlockParam{ - ToolUseID: tr.ToolUseID, - Content: toolContentsToAnthropicBeta(tr.Contents), - IsError: anthropic.Bool(tr.IsError), - }, - } - } - newMessages = append(newMessages, anthropic.NewBetaUserMessage(resultBlocks...)) + newMessages = append(newMessages, anthropic.NewBetaUserMessage(betaStageToolResultBlocks(results)...)) newReq.Messages = newMessages return &newReq, nil @@ -209,8 +196,22 @@ func mergeAnthropicBetaContinuation(segment []anthropic.BetaMessageParam, messag } merged := append([]anthropic.BetaMessageParam{}, segment...) - lastIdx := len(merged) - 1 - merged[lastIdx].Content = append(append([]anthropic.BetaContentBlockParamUnion{}, merged[lastIdx].Content...), messages[toolResultIdx].Content...) + externalResults := append([]anthropic.BetaContentBlockParamUnion{}, messages[toolResultIdx].Content...) + last := len(merged) - 1 + if merged[last].Role == anthropic.BetaMessageParamRoleUser && betaMessageHasToolResult(merged[last]) { + // The segment ends with the managed-results user turn; fold the + // client's external results into that same results message. + merged[last].Content = append(merged[last].Content, externalResults...) + } else { + // The segment ends with the current round's assistant message (an + // external-only round after internal rounds). The client's external + // results form a new trailing user turn so every tool_result still + // follows its assistant tool_use in the merged request. + merged = append(merged, anthropic.BetaMessageParam{ + Role: anthropic.BetaMessageParamRoleUser, + Content: externalResults, + }) + } if assistantIdx == -1 || toolResultIdx < assistantIdx { result := append([]anthropic.BetaMessageParam{}, merged...) result = append(result, messages[:toolResultIdx]...) @@ -224,6 +225,18 @@ func mergeAnthropicBetaContinuation(segment []anthropic.BetaMessageParam, messag return result } +// betaMessageHasToolResult reports whether a user message carries any tool +// result blocks. The merge folds incoming results into a trailing results user +// message; anything else receives a fresh results turn instead. +func betaMessageHasToolResult(message anthropic.BetaMessageParam) bool { + for _, block := range message.Content { + if block.OfToolResult != nil { + return true + } + } + return false +} + func betaMessageToParamPreservingThinking(msg *anthropic.BetaMessage) anthropic.BetaMessageParam { if msg == nil { return anthropic.BetaMessageParam{} diff --git a/internal/mcpserver/anthropic_beta_stage.go b/internal/mcpserver/anthropic_beta_stage.go index 1bbb52391..05a35d98d 100644 --- a/internal/mcpserver/anthropic_beta_stage.go +++ b/internal/mcpserver/anthropic_beta_stage.go @@ -33,7 +33,14 @@ type AnthropicBetaStageExecutor interface { // the session key from ctx; the Stage never knows that storage key. type AnthropicBetaContinuationStore interface { Pop(ctx context.Context, request *anthropic.BetaMessageNewParams) ([]anthropic.BetaMessageParam, bool) - Put(ctx context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) + // Put stores a continuation segment that the next request may pop once its + // external tool results are present. It reports failure so the Stage never + // silently drops the internal tool context. + Put(ctx context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) error + // CanStash reports whether a continuation can be persisted for ctx. The + // Stage checks it before committing server-tool side effects for a round + // that must be stashed. + CanStash(ctx context.Context) bool } type AnthropicBetaStageConfig struct { @@ -102,6 +109,11 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto runCtx := ctx current := prepared var totalUsage *protocol.TokenUsage + // internalMessages accumulates the assistant/tool-result turns of internal + // (server-owned) rounds the client never saw. When the model later asks for + // a client-owned tool, these messages are carried into the continuation so + // the managed results still reach the final answer. + var internalMessages []anthropic.BetaMessageParam sideEffectsCommitted := false for round := 1; round <= e.stage.maxRounds; round++ { response, callErr := e.next.Complete(runCtx, current) @@ -124,6 +136,25 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto } managed, external, externalIDs := splitBetaStageTools(tools, owned) if len(managed) == 0 { + // A final answer (no tool calls) or a pure-external round with no + // internal context needs no continuation. + if len(external) == 0 || len(internalMessages) == 0 { + response.Usage = totalUsage + response.SideEffectsCommitted = sideEffectsCommitted + return response, nil + } + // An external-only round after internal rounds: carry the internal + // context to the client's next turn when the continuation can be + // persisted. Without an explicit session the store cannot bind the + // segment safely (IP addresses are not a conversation identity), so + // deliver the external round as-is; guardrail blocking and terminal + // responses still work, and no new side effects are committed here. + if e.canPersistContinuation(runCtx) { + segment := append(internalMessages, betaMessageToParamPreservingThinking(message)) + if stashErr := e.stashContinuation(runCtx, segment, externalIDs); stashErr != nil { + return nil, stagetoolloop.WrapError(stashErr, sideEffectsCommitted) + } + } response.Usage = totalUsage response.SideEffectsCommitted = sideEffectsCommitted return response, nil @@ -134,6 +165,9 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto response.SideEffectsCommitted = sideEffectsCommitted return response, nil } + if !e.stage.continuations.CanStash(runCtx) { + return nil, stagetoolloop.WrapError(stagetoolloop.ErrContinuationUnavailable, sideEffectsCommitted) + } results, nextCtx, committed := e.executeTools(runCtx, current.Request, managed) sideEffectsCommitted = sideEffectsCommitted || committed runCtx = nextCtx @@ -149,7 +183,12 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto if !ok || len(segment) == 0 { return nil, stagetoolloop.WrapError(errors.New("Anthropic Beta ToolLoop built an empty mixed continuation"), sideEffectsCommitted) } - e.stage.continuations.Put(runCtx, segment, externalIDs) + // Prepend earlier internal rounds so results from previous managed + // rounds survive the client's next turn alongside this round's. + segment = append(internalMessages, segment...) + if stashErr := e.stashContinuation(runCtx, segment, externalIDs); stashErr != nil { + return nil, stagetoolloop.WrapError(stashErr, sideEffectsCommitted) + } filtered, filterErr := e.stage.adapter.FilterVirtualTools(message, external) if filterErr != nil { return nil, stagetoolloop.WrapError(filterErr, sideEffectsCommitted) @@ -170,6 +209,7 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto for i := range results { resultValues[i] = results[i] } + internalMessages = appendBetaStageInternalRound(internalMessages, message, resultValues) nextRequest, appendErr := e.stage.adapter.AppendToolResults(current.Request, message, resultValues) if appendErr != nil { return nil, stagetoolloop.WrapError(appendErr, sideEffectsCommitted) @@ -266,6 +306,53 @@ func (e *anthropicBetaToolLoopEndpoint) executeTools( return results, runCtx, committed } +// canPersistContinuation reports whether a continuation segment can be stored +// for this context. A nil store (test-only fallback) or a store that cannot +// bind to the current session cannot persist one. +func (e *anthropicBetaToolLoopEndpoint) canPersistContinuation(ctx context.Context) bool { + return e.stage.continuations != nil && e.stage.continuations.CanStash(ctx) +} + +// stashContinuation stores a continuation segment and surfaces persistence +// failure as a Stage error so the internal tool context is never silently lost. +func (e *anthropicBetaToolLoopEndpoint) stashContinuation(ctx context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) error { + if e.stage.continuations == nil { + return nil + } + if err := e.stage.continuations.Put(ctx, segment, externalIDs); err != nil { + return fmt.Errorf("store Anthropic Beta ToolLoop continuation: %w", err) + } + return nil +} + +// appendBetaStageInternalRound records one server-owned round in the messages +// the client has not seen, mirroring AppendToolResults so the in-loop request +// and any later continuation segment stay in sync. +func appendBetaStageInternalRound(messages []anthropic.BetaMessageParam, message *anthropic.BetaMessage, results []any) []anthropic.BetaMessageParam { + messages = append(messages, betaMessageToParamPreservingThinking(message)) + return append(messages, anthropic.NewBetaUserMessage(betaStageToolResultBlocks(results)...)) +} + +// betaStageToolResultBlocks converts execution results into Beta tool_result +// content blocks shared by the adapter's append and continuation builders. +func betaStageToolResultBlocks(results []any) []anthropic.BetaContentBlockParamUnion { + blocks := make([]anthropic.BetaContentBlockParamUnion, 0, len(results)) + for _, r := range results { + tr, ok := r.(ToolExecutionResult) + if !ok { + continue + } + blocks = append(blocks, anthropic.BetaContentBlockParamUnion{ + OfToolResult: &anthropic.BetaToolResultBlockParam{ + ToolUseID: tr.ToolUseID, + Content: toolContentsToAnthropicBeta(tr.Contents), + IsError: anthropic.Bool(tr.IsError), + }, + }) + } + return blocks +} + func betaStageMessage(value any) (*anthropic.BetaMessage, error) { switch message := value.(type) { case *anthropic.BetaMessage: diff --git a/internal/mcpserver/anthropic_beta_stage_stream.go b/internal/mcpserver/anthropic_beta_stage_stream.go index 203c24b82..26b771d82 100644 --- a/internal/mcpserver/anthropic_beta_stage_stream.go +++ b/internal/mcpserver/anthropic_beta_stage_stream.go @@ -31,6 +31,9 @@ type anthropicBetaToolLoopStream struct { assembler assembler.StreamAssembler buffered []protocolstage.Event pending []protocolstage.Event + // internalMessages accumulates the assistant/tool-result turns of internal + // (server-owned) rounds the client never saw, mirroring the Complete path. + internalMessages []anthropic.BetaMessageParam usage *protocol.TokenUsage model string @@ -103,6 +106,26 @@ func (s *anthropicBetaToolLoopStream) Next(ctx context.Context) (protocolstage.E } managed, external, externalIDs := splitBetaStageTools(tools, s.owned) if len(managed) == 0 { + // A final answer (no tool calls) or a pure-external round with no + // internal context needs no continuation. + if len(external) == 0 || len(s.internalMessages) == 0 { + s.pending = s.buffered + s.buffered = nil + s.done = true + continue + } + // An external-only round after internal rounds: carry the internal + // context to the client's next turn when the continuation can be + // persisted. Without an explicit session the store cannot bind the + // segment safely (IP addresses are not a conversation identity), so + // deliver the external round as-is; guardrail blocking and terminal + // responses still work, and no new side effects are committed here. + if s.endpoint.canPersistContinuation(s.runCtx) { + segment := append(s.internalMessages, betaMessageToParamPreservingThinking(message)) + if stashErr := s.endpoint.stashContinuation(s.runCtx, segment, externalIDs); stashErr != nil { + return protocolstage.Event{}, s.fail(stashErr) + } + } s.pending = s.buffered s.buffered = nil s.done = true @@ -115,6 +138,9 @@ func (s *anthropicBetaToolLoopStream) Next(ctx context.Context) (protocolstage.E s.done = true continue } + if !s.endpoint.stage.continuations.CanStash(s.runCtx) { + return protocolstage.Event{}, s.fail(stagetoolloop.ErrContinuationUnavailable) + } results, nextCtx, committed := s.endpoint.executeTools(s.runCtx, s.call.Request, managed) s.sideEffects = s.sideEffects || committed s.runCtx = nextCtx @@ -130,7 +156,12 @@ func (s *anthropicBetaToolLoopStream) Next(ctx context.Context) (protocolstage.E if !ok || len(segment) == 0 { return protocolstage.Event{}, s.fail(errors.New("Anthropic Beta ToolLoop built an empty mixed continuation")) } - s.endpoint.stage.continuations.Put(s.runCtx, segment, externalIDs) + // Prepend earlier internal rounds so results from previous managed + // rounds survive the client's next turn alongside this round's. + segment = append(s.internalMessages, segment...) + if stashErr := s.endpoint.stashContinuation(s.runCtx, segment, externalIDs); stashErr != nil { + return protocolstage.Event{}, s.fail(stashErr) + } filtered, filterErr := filterBetaStageStreamEvents(s.buffered, s.owned) if filterErr != nil { return protocolstage.Event{}, s.fail(filterErr) @@ -151,6 +182,7 @@ func (s *anthropicBetaToolLoopStream) Next(ctx context.Context) (protocolstage.E for i := range results { resultValues[i] = results[i] } + s.internalMessages = appendBetaStageInternalRound(s.internalMessages, message, resultValues) nextRequest, appendErr := s.endpoint.stage.adapter.AppendToolResults(s.call.Request, message, resultValues) if appendErr != nil { return protocolstage.Event{}, s.fail(appendErr) diff --git a/internal/mcpserver/anthropic_beta_stage_test.go b/internal/mcpserver/anthropic_beta_stage_test.go index facf20698..7a4902f56 100644 --- a/internal/mcpserver/anthropic_beta_stage_test.go +++ b/internal/mcpserver/anthropic_beta_stage_test.go @@ -291,6 +291,179 @@ func TestAnthropicBetaStageCompleteStoresAndAppliesMixedContinuation(t *testing. } } +func TestAnthropicBetaStageCompleteCarriesManagedRoundAcrossExternalRound(t *testing.T) { + continuations := &memoryBetaStageContinuations{} + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("internal result").Contents}, + }} + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + // round 1: pure server-owned tool + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-owned", Name: "lookup"})}, + // round 2: pure client-owned tool + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-external", Name: "client_tool"})}, + // round 3: final answer after the client's follow-up + {Value: betaStageTextMessage(t, "combined")}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + Continuations: continuations, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + first, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + filtered := first.Value.(*anthropic.BetaMessage) + filteredTools, err := NewAnthropicBetaAdapter().ExtractTools(filtered) + if err != nil { + t.Fatal(err) + } + if len(filteredTools) != 1 || filteredTools[0].Name() != "client_tool" { + t.Fatalf("outward tools = %#v", filteredTools) + } + if continuations.puts != 1 || len(executor.calls) != 1 { + t.Fatalf("continuation puts=%d executions=%d", continuations.puts, len(executor.calls)) + } + + // The client replies with the external tool result only; the managed round + // must be carried by the continuation. + externalResult := anthropic.NewBetaUserMessage(anthropic.NewBetaToolResultBlock("toolu-external", "external result", false)) + second, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{ + Messages: []anthropic.BetaMessageParam{externalResult}, + }}) + if err != nil { + t.Fatal(err) + } + if second.Value.(*anthropic.BetaMessage).Content[0].Text != "combined" { + t.Fatalf("second response = %#v", second.Value) + } + continued := terminal.calls[2].Request.(*anthropic.BetaMessageNewParams) + if len(continued.Messages) != 4 { + t.Fatalf("continued messages = %d, want assistant(owned)+user(owned result)+assistant(external)+user(external result)", len(continued.Messages)) + } + if got := betaStageToolResultIDs(continued.Messages[1]); len(got) != 1 || got[0] != "toolu-owned" { + t.Fatalf("managed result IDs = %#v", got) + } + if got := betaStageToolResultIDs(continued.Messages[3]); len(got) != 1 || got[0] != "toolu-external" { + t.Fatalf("external result IDs = %#v", got) + } +} + +func TestAnthropicBetaStageCompleteFailsWhenContinuationCannotBeStashed(t *testing.T) { + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("internal result").Contents}, + }} + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + {Value: betaStageToolMessage(t, + betaStageToolCallSpec{ID: "toolu-owned", Name: "lookup"}, + betaStageToolCallSpec{ID: "toolu-external", Name: "client_tool"}, + )}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + Continuations: &memoryBetaStageContinuations{failStash: true}, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if !errors.Is(err, stagetoolloop.ErrContinuationUnavailable) { + t.Fatalf("error = %v, want ErrContinuationUnavailable", err) + } + // No server tool may run when the mixed round cannot be persisted. + if len(executor.calls) != 0 { + t.Fatalf("executor calls = %d, want 0", len(executor.calls)) + } +} + +func TestAnthropicBetaStageStreamCarriesManagedRoundAcrossExternalRound(t *testing.T) { + continuations := &memoryBetaStageContinuations{} + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }} + terminal := &betaStageScriptedEndpoint{streams: []*betaStageMemoryStream{ + {events: betaStageToolStreamEvents(betaStageToolCallSpec{ID: "toolu-owned", Name: "lookup"})}, + {events: betaStageToolStreamEvents(betaStageToolCallSpec{ID: "toolu-external", Name: "client_tool"})}, + {events: betaStageTextStreamEvents("combined")}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + Continuations: continuations, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + first, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + firstEvents := collectBetaStageEvents(t, first) + if body := betaStageEventBodies(t, firstEvents); !strings.Contains(body, "client_tool") || strings.Contains(body, "lookup") { + t.Fatalf("outward first stream = %s", body) + } + if continuations.puts != 1 || len(executor.calls) != 1 { + t.Fatalf("continuation puts=%d executions=%d", continuations.puts, len(executor.calls)) + } + + // Client replies with the external tool result; the managed round must be + // carried by the continuation into the follow-up provider request. + externalResult := anthropic.NewBetaUserMessage(anthropic.NewBetaToolResultBlock("toolu-external", "external result", false)) + second, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{ + Messages: []anthropic.BetaMessageParam{externalResult}, + }}) + if err != nil { + t.Fatal(err) + } + secondEvents := collectBetaStageEvents(t, second) + if body := betaStageEventBodies(t, secondEvents); !strings.Contains(body, "combined") { + t.Fatalf("outward second stream = %s", body) + } + continued := terminal.streamCalls[2].Request.(*anthropic.BetaMessageNewParams) + if len(continued.Messages) != 4 { + t.Fatalf("continued messages = %d, want assistant(owned)+user(owned result)+assistant(external)+user(external result)", len(continued.Messages)) + } + if got := betaStageToolResultIDs(continued.Messages[1]); len(got) != 1 || got[0] != "toolu-owned" { + t.Fatalf("managed result IDs = %#v", got) + } + if got := betaStageToolResultIDs(continued.Messages[3]); len(got) != 1 || got[0] != "toolu-external" { + t.Fatalf("external result IDs = %#v", got) + } +} + +func TestAnthropicBetaStageStreamFailsWhenContinuationCannotBeStashed(t *testing.T) { + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }} + terminal := &betaStageScriptedEndpoint{streams: []*betaStageMemoryStream{ + {events: betaStageToolStreamEvents( + betaStageToolCallSpec{ID: "toolu-owned", Name: "lookup"}, + betaStageToolCallSpec{ID: "toolu-external", Name: "client_tool"}, + )}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + Continuations: &memoryBetaStageContinuations{failStash: true}, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + stream, err := endpoint.Stream(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + _, err = stream.Next(context.Background()) + if !errors.Is(err, stagetoolloop.ErrContinuationUnavailable) { + t.Fatalf("stream error = %v, want ErrContinuationUnavailable", err) + } + // No server tool may run when the mixed round cannot be persisted. + if len(executor.calls) != 0 { + t.Fatalf("executor calls = %d, want 0", len(executor.calls)) + } + _ = stream.Close() +} + func TestAnthropicBetaStageRejectsAmbiguousOwnership(t *testing.T) { request := &anthropic.BetaMessageNewParams{Tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}} toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ @@ -507,6 +680,11 @@ type memoryBetaStageContinuations struct { expectedIDs []string puts int pops int + failStash bool +} + +func (s *memoryBetaStageContinuations) CanStash(context.Context) bool { + return !s.failStash } func (s *memoryBetaStageContinuations) Pop(_ context.Context, request *anthropic.BetaMessageNewParams) ([]anthropic.BetaMessageParam, bool) { @@ -519,10 +697,11 @@ func (s *memoryBetaStageContinuations) Pop(_ context.Context, request *anthropic return segment, true } -func (s *memoryBetaStageContinuations) Put(_ context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) { +func (s *memoryBetaStageContinuations) Put(_ context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) error { s.puts++ s.segment = append([]anthropic.BetaMessageParam(nil), segment...) s.expectedIDs = append([]string(nil), externalIDs...) + return nil } func (e *fakeBetaStageExecutor) ExecuteToolWithContext(ctx context.Context, tool Tool, _ []map[string]any) (context.Context, ToolExecutionResult, error) { diff --git a/internal/mcpserver/continuation_store.go b/internal/mcpserver/continuation_store.go index dc74985d5..cb40d451b 100644 --- a/internal/mcpserver/continuation_store.go +++ b/internal/mcpserver/continuation_store.go @@ -3,6 +3,7 @@ package mcpserver import ( "context" "encoding/json" + "errors" "fmt" "sync" "time" @@ -274,10 +275,30 @@ func (s *ProviderBetaContinuationStore) Pop(ctx context.Context, request *anthro return messages, true } -func (s *ProviderBetaContinuationStore) Put(ctx context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) { +// CanStash reports whether a continuation can be persisted for this context. +// IP addresses are not a conversation identity, so requests without an explicit +// session cannot store continuation state. The Stage checks this before it +// executes server-owned tools so it never commits side effects it cannot carry +// to the client's next turn. +func (s *ProviderBetaContinuationStore) CanStash(ctx context.Context) bool { + if s == nil || s.providerUUID == "" { + return false + } + sessionID := typ.GetSessionID(ctx) + return !sessionID.IsEmpty() && !sessionID.IsIPFallback() +} + +func (s *ProviderBetaContinuationStore) Put(ctx context.Context, segment []anthropic.BetaMessageParam, externalIDs []string) error { if s == nil || len(segment) == 0 { - return + return errors.New("store Anthropic Beta continuation: nil store or empty segment") } key := continuationKey(typ.GetSessionID(ctx), s.providerUUID, "anthropic-beta") + if key == "" { + return errors.New("store Anthropic Beta continuation: requires an explicit session") + } + if len(stringSet(externalIDs)) == 0 { + return errors.New("store Anthropic Beta continuation: no external tool IDs to correlate") + } mixedContinuationStore.put(key, append([]anthropic.BetaMessageParam(nil), segment...), externalIDs) + return nil } diff --git a/internal/protocol/stage/toolloop/runtime.go b/internal/protocol/stage/toolloop/runtime.go index a78e35690..76511fdcc 100644 --- a/internal/protocol/stage/toolloop/runtime.go +++ b/internal/protocol/stage/toolloop/runtime.go @@ -55,8 +55,9 @@ type AllowAllPolicy struct{} func (AllowAllPolicy) Authorize(context.Context, ToolCall) error { return nil } var ( - ErrMaxRounds = errors.New("tool loop reached the maximum number of rounds") - ErrToolNameCollision = errors.New("server tool name collides with request tool") + ErrMaxRounds = errors.New("tool loop reached the maximum number of rounds") + ErrToolNameCollision = errors.New("server tool name collides with request tool") + ErrContinuationUnavailable = errors.New("tool loop continuation requires an explicit session") ) // ExecutionError preserves the irreversible-side-effect boundary when a later From bbe33d474445a6fb4ea6d3e6d1a9a9635c3ec84b Mon Sep 17 00:00:00 2001 From: FFengIll Date: Tue, 4 Aug 2026 21:10:35 +0800 Subject: [PATCH 7/8] fix(toolloop): count max rounds as tool executions and commit side effects after dispatch Problem 3: a pure-managed tool call on the last allowed round failed with ErrMaxRounds, discarding the tool calls, while a mixed round on the same round succeeded. maxRounds now bounds server-tool executions: the last budgeted round executes its tools and the model gets one extra provider round to produce the final answer. A further server-tool request fails closed. Applies to the Beta complete/stream stages and the OpenAI Chat stage; unify both stage defaults behind DefaultMaxRounds=3 (matching the legacy MCP loop) and set the production Beta loop explicitly. Problem 4: the Chat tool loop marked side effects committed only on a successful executor return. An executor that dispatches an irreversible action and then loses its response left the attempt uncommitted, so failover could replay the action on another provider. ToolResult gains a Dispatched signal (matching the Beta stage's existing behavior); the Chat stage commits when err == nil || result.Dispatched, keeping pre-dispatch validation and policy failures uncommitted. --- internal/mcpserver/anthropic_beta_stage.go | 15 +++-- .../mcpserver/anthropic_beta_stage_stream.go | 8 ++- .../mcpserver/anthropic_beta_stage_test.go | 56 +++++++++++++++++++ .../protocol/stage/toolloop/openai_chat.go | 15 +++-- .../stage/toolloop/openai_chat_stream.go | 4 +- .../stage/toolloop/openai_chat_test.go | 44 ++++++++++++--- internal/protocol/stage/toolloop/runtime.go | 12 ++++ .../protocol_stage_tool_loop.go | 5 +- 8 files changed, 136 insertions(+), 23 deletions(-) diff --git a/internal/mcpserver/anthropic_beta_stage.go b/internal/mcpserver/anthropic_beta_stage.go index 05a35d98d..c0226083d 100644 --- a/internal/mcpserver/anthropic_beta_stage.go +++ b/internal/mcpserver/anthropic_beta_stage.go @@ -64,7 +64,7 @@ func NewAnthropicBetaStage(config AnthropicBetaStageConfig) (protocolstage.Stage } maxRounds := config.MaxRounds if maxRounds <= 0 { - maxRounds = defaultMaxRounds + maxRounds = stagetoolloop.DefaultMaxRounds } return &anthropicBetaToolLoopStage{ name: name, @@ -115,7 +115,9 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto // the managed results still reach the final answer. var internalMessages []anthropic.BetaMessageParam sideEffectsCommitted := false - for round := 1; round <= e.stage.maxRounds; round++ { + // maxRounds bounds server-tool executions; one extra provider round is + // allowed so the model can produce the final answer after the last one. + for round := 1; round <= e.stage.maxRounds+1; round++ { response, callErr := e.next.Complete(runCtx, current) if callErr != nil { return nil, stagetoolloop.WrapError(callErr, sideEffectsCommitted) @@ -135,6 +137,12 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto return nil, stagetoolloop.WrapError(extractErr, sideEffectsCommitted) } managed, external, externalIDs := splitBetaStageTools(tools, owned) + if len(managed) > 0 && round > e.stage.maxRounds { + // The tool-execution budget is exhausted and the model still + // requests server-owned tools; fail closed instead of executing + // beyond the budget or leaking internal tool definitions. + return nil, stagetoolloop.WrapError(stagetoolloop.ErrMaxRounds, sideEffectsCommitted) + } if len(managed) == 0 { // A final answer (no tool calls) or a pure-external round with no // internal context needs no continuation. @@ -198,9 +206,6 @@ func (e *anthropicBetaToolLoopEndpoint) Complete(ctx context.Context, call proto response.SideEffectsCommitted = sideEffectsCommitted return response, nil } - if round == e.stage.maxRounds { - return nil, stagetoolloop.WrapError(stagetoolloop.ErrMaxRounds, sideEffectsCommitted) - } results, nextCtx, committed := e.executeTools(runCtx, current.Request, managed) sideEffectsCommitted = sideEffectsCommitted || committed diff --git a/internal/mcpserver/anthropic_beta_stage_stream.go b/internal/mcpserver/anthropic_beta_stage_stream.go index 26b771d82..97d9c9156 100644 --- a/internal/mcpserver/anthropic_beta_stage_stream.go +++ b/internal/mcpserver/anthropic_beta_stage_stream.go @@ -105,6 +105,11 @@ func (s *anthropicBetaToolLoopStream) Next(ctx context.Context) (protocolstage.E return protocolstage.Event{}, s.fail(extractErr) } managed, external, externalIDs := splitBetaStageTools(tools, s.owned) + // maxRounds bounds server-tool executions; the round after the last + // allowed execution must be a final answer or a client-owned round. + if len(managed) > 0 && s.round > s.endpoint.stage.maxRounds { + return protocolstage.Event{}, s.fail(stagetoolloop.ErrMaxRounds) + } if len(managed) == 0 { // A final answer (no tool calls) or a pure-external round with no // internal context needs no continuation. @@ -171,9 +176,6 @@ func (s *anthropicBetaToolLoopStream) Next(ctx context.Context) (protocolstage.E s.done = true continue } - if s.round >= s.endpoint.stage.maxRounds { - return protocolstage.Event{}, s.fail(stagetoolloop.ErrMaxRounds) - } results, nextCtx, committed := s.endpoint.executeTools(s.runCtx, s.call.Request, managed) s.sideEffects = s.sideEffects || committed diff --git a/internal/mcpserver/anthropic_beta_stage_test.go b/internal/mcpserver/anthropic_beta_stage_test.go index 7a4902f56..06911e07d 100644 --- a/internal/mcpserver/anthropic_beta_stage_test.go +++ b/internal/mcpserver/anthropic_beta_stage_test.go @@ -464,6 +464,62 @@ func TestAnthropicBetaStageStreamFailsWhenContinuationCannotBeStashed(t *testing _ = stream.Close() } +func TestAnthropicBetaStageCompleteAllowsToolExecutionOnLastBudgetRound(t *testing.T) { + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }} + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-1", Name: "lookup"})}, + {Value: betaStageTextMessage(t, "final")}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + MaxRounds: 1, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + // maxRounds bounds tool executions: the single allowed execution runs, then + // the model gets one more round to produce the final answer. + response, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if err != nil { + t.Fatal(err) + } + if response.Value.(*anthropic.BetaMessage).Content[0].Text != "final" { + t.Fatalf("final response = %#v", response.Value) + } + if len(executor.calls) != 1 { + t.Fatalf("executed %d tools, want 1", len(executor.calls)) + } + if len(terminal.calls) != 2 { + t.Fatalf("provider calls = %d, want 2 (execution round + final answer round)", len(terminal.calls)) + } +} + +func TestAnthropicBetaStageCompleteFailsWhenToolExecutionExceedsBudget(t *testing.T) { + executor := &fakeBetaStageExecutor{results: map[string]ToolExecutionResult{ + "lookup": {Contents: coretool.TextToolResult("ok").Contents}, + }} + terminal := &betaStageScriptedEndpoint{responses: []*protocolstage.Response{ + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-1", Name: "lookup"})}, + {Value: betaStageToolMessage(t, betaStageToolCallSpec{ID: "toolu-2", Name: "lookup"})}, + }} + toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ + Tools: staticBetaStageTools{tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}}, + Executor: executor, + MaxRounds: 1, + }) + endpoint, _ := protocolstage.Compose(terminal, toolStage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &anthropic.BetaMessageNewParams{}}) + if !errors.Is(err, stagetoolloop.ErrMaxRounds) || !stagetoolloop.HasCommittedSideEffects(err) { + t.Fatalf("error = %v, committed=%v", err, stagetoolloop.HasCommittedSideEffects(err)) + } + if len(executor.calls) != 1 { + t.Fatalf("executed %d tools, want 1 (the budget allows exactly one execution)", len(executor.calls)) + } +} + func TestAnthropicBetaStageRejectsAmbiguousOwnership(t *testing.T) { request := &anthropic.BetaMessageNewParams{Tools: []anthropic.BetaToolUnionParam{betaStageToolDefinition("lookup")}} toolStage, _ := NewAnthropicBetaStage(AnthropicBetaStageConfig{ diff --git a/internal/protocol/stage/toolloop/openai_chat.go b/internal/protocol/stage/toolloop/openai_chat.go index 1d92f74f4..ea273c41d 100644 --- a/internal/protocol/stage/toolloop/openai_chat.go +++ b/internal/protocol/stage/toolloop/openai_chat.go @@ -14,8 +14,6 @@ import ( "github.com/tingly-dev/tingly-box/internal/protocol/wire" ) -const defaultMaxRounds = 8 - // OpenAIChatConfig constructs a Chat-native ToolLoop Stage. Catalog, policy, // and executor are protocol-neutral dependencies; only this adapter understands // OpenAI Chat request and response types. @@ -43,7 +41,7 @@ func NewOpenAIChat(config OpenAIChatConfig) (protocolstage.Stage, error) { } maxRounds := config.MaxRounds if maxRounds <= 0 { - maxRounds = defaultMaxRounds + maxRounds = DefaultMaxRounds } return &openAIChatStage{ name: name, @@ -85,7 +83,9 @@ func (e *openAIChatEndpoint) Complete(ctx context.Context, call protocolstage.Ca current := prepared var totalUsage *protocol.TokenUsage sideEffectsCommitted := false - for round := 1; round <= e.stage.maxRounds; round++ { + // maxRounds bounds server-tool executions; one extra provider round is + // allowed so the model can produce the final answer after the last one. + for round := 1; round <= e.stage.maxRounds+1; round++ { response, callErr := e.next.Complete(runCtx, current) if callErr != nil { return nil, WrapError(callErr, sideEffectsCommitted) @@ -105,7 +105,7 @@ func (e *openAIChatEndpoint) Complete(ctx context.Context, call protocolstage.Ca response.SideEffectsCommitted = sideEffectsCommitted return response, nil } - if round == e.stage.maxRounds { + if round == e.stage.maxRounds+1 { return nil, WrapError(ErrMaxRounds, sideEffectsCommitted) } @@ -197,7 +197,10 @@ func (e *openAIChatEndpoint) executeCalls(ctx context.Context, calls []ToolCall) if result.Content == "" { result.Content = err.Error() } - } else { + } + if err == nil || result.Dispatched { + // An executor may dispatch an irreversible action and then lose its + // response; only failures before dispatch leave the attempt uncommitted. committed = true } results = append(results, result) diff --git a/internal/protocol/stage/toolloop/openai_chat_stream.go b/internal/protocol/stage/toolloop/openai_chat_stream.go index ccf6aabc2..689b7bdf6 100644 --- a/internal/protocol/stage/toolloop/openai_chat_stream.go +++ b/internal/protocol/stage/toolloop/openai_chat_stream.go @@ -125,7 +125,9 @@ func (s *openAIChatToolLoopStream) Next(ctx context.Context) (protocolstage.Even } continue } - if s.round >= s.endpoint.stage.maxRounds { + // maxRounds bounds server-tool executions; the round after the last + // allowed execution must be a final answer or a client-owned round. + if s.round > s.endpoint.stage.maxRounds { return protocolstage.Event{}, s.fail(ErrMaxRounds) } diff --git a/internal/protocol/stage/toolloop/openai_chat_test.go b/internal/protocol/stage/toolloop/openai_chat_test.go index 740d171bd..891b691ef 100644 --- a/internal/protocol/stage/toolloop/openai_chat_test.go +++ b/internal/protocol/stage/toolloop/openai_chat_test.go @@ -296,9 +296,14 @@ func TestOpenAIChatStreamPreservesSideEffectBoundaryAfterLaterFailure(t *testing _ = stream.Close() } -func TestOpenAIChatStreamEnforcesMaxRoundsBeforeToolExecution(t *testing.T) { - executor := &fakeExecutor{} - terminal := &scriptedChatEndpoint{streams: []*memoryEventStream{{events: toolCallStreamEvents("call-1", "lookup", `{}`)}}} +func TestOpenAIChatStreamAllowsToolExecutionThenRequiresFinalAnswer(t *testing.T) { + executor := &fakeExecutor{results: map[string]ToolResult{ + "lookup": {Content: "ok"}, + }} + terminal := &scriptedChatEndpoint{streams: []*memoryEventStream{ + {events: toolCallStreamEvents("call-1", "lookup", `{}`)}, + {events: toolCallStreamEvents("call-2", "lookup", `{}`)}, + }} stage, _ := NewOpenAIChat(OpenAIChatConfig{ Catalog: staticCatalog{{Name: "lookup"}}, Executor: executor, @@ -310,16 +315,40 @@ func TestOpenAIChatStreamEnforcesMaxRoundsBeforeToolExecution(t *testing.T) { t.Fatal(err) } + // maxRounds bounds tool executions: the single allowed execution happens, + // then a further all-owned round fails closed with side effects committed. _, err = stream.Next(context.Background()) - if !errors.Is(err, ErrMaxRounds) || HasCommittedSideEffects(err) { + if !errors.Is(err, ErrMaxRounds) || !HasCommittedSideEffects(err) { t.Fatalf("max-round error = %v, committed=%v", err, HasCommittedSideEffects(err)) } - if len(executor.calls) != 0 { - t.Fatalf("executed %d tools after reaching max rounds", len(executor.calls)) + if len(executor.calls) != 1 { + t.Fatalf("executed %d tools, want 1 (the budget allows one execution)", len(executor.calls)) } _ = stream.Close() } +func TestOpenAIChatCompleteTreatsPostDispatchToolErrorAsCommitted(t *testing.T) { + providerErr := errors.New("second round failed") + toolErr := errors.New("tool response was lost after dispatch") + terminal := &scriptedChatEndpoint{ + completeResponses: []*protocolstage.Response{{Value: toolCallCompletion("call-1", "lookup", `{}`)}}, + completeErrors: []error{nil, providerErr}, + } + stage, _ := NewOpenAIChat(OpenAIChatConfig{ + Catalog: staticCatalog{{Name: "lookup"}}, + Executor: &fakeExecutor{ + results: map[string]ToolResult{"lookup": {Dispatched: true}}, + errors: map[string]error{"lookup": toolErr}, + }, + }) + endpoint, _ := protocolstage.Compose(terminal, stage) + + _, err := endpoint.Complete(context.Background(), protocolstage.Call{Request: &openai.ChatCompletionNewParams{}}) + if !errors.Is(err, providerErr) || !HasCommittedSideEffects(err) { + t.Fatalf("later error = %v, committed=%v", err, HasCommittedSideEffects(err)) + } +} + func TestOpenAIChatStreamRecordsToolRoundsAsExchangesInOneAttempt(t *testing.T) { request := &openai.ChatCompletionNewParams{Model: "public", Messages: []openai.ChatCompletionMessageParamUnion{openai.UserMessage("hello")}} recorder, err := record.New(record.Config{ @@ -387,12 +416,13 @@ func (c staticCatalog) ListTools(context.Context) ([]ToolDefinition, error) { type fakeExecutor struct { results map[string]ToolResult + errors map[string]error calls []ToolCall } func (e *fakeExecutor) Execute(ctx context.Context, call ToolCall) (context.Context, ToolResult, error) { e.calls = append(e.calls, call) - return ctx, e.results[call.Name], nil + return ctx, e.results[call.Name], e.errors[call.Name] } type scriptedChatEndpoint struct { diff --git a/internal/protocol/stage/toolloop/runtime.go b/internal/protocol/stage/toolloop/runtime.go index 76511fdcc..5dc1aff9a 100644 --- a/internal/protocol/stage/toolloop/runtime.go +++ b/internal/protocol/stage/toolloop/runtime.go @@ -28,6 +28,12 @@ type ToolResult struct { ToolCallID string Content string IsError bool + // Dispatched reports that the executor crossed into the runtime before + // failing. An irreversible action may already have happened, so a later + // error must still mark the attempt's side effects as committed to stop + // failover from replaying it. Leave it false for validation, disabled-tool, + // and policy failures that happen before dispatch. + Dispatched bool } // ToolCatalog lists the server-visible tools for one request. The returned @@ -54,6 +60,12 @@ type AllowAllPolicy struct{} func (AllowAllPolicy) Authorize(context.Context, ToolCall) error { return nil } +// DefaultMaxRounds bounds server-tool executions per attempt for every native +// ToolLoop Stage. One final provider round is allowed after the last execution +// so the model can produce an answer; a further server-tool request fails +// closed instead of executing beyond the budget. +const DefaultMaxRounds = 3 + var ( ErrMaxRounds = errors.New("tool loop reached the maximum number of rounds") ErrToolNameCollision = errors.New("server tool name collides with request tool") diff --git a/internal/protocolserver/protocol_stage_tool_loop.go b/internal/protocolserver/protocol_stage_tool_loop.go index d26133873..878d1e01d 100644 --- a/internal/protocolserver/protocol_stage_tool_loop.go +++ b/internal/protocolserver/protocol_stage_tool_loop.go @@ -6,9 +6,9 @@ import ( "github.com/gin-gonic/gin" "github.com/sirupsen/logrus" + mcpmodule "github.com/tingly-dev/tingly-box/internal/mcpserver" protocolstage "github.com/tingly-dev/tingly-box/internal/protocol/stage" stagetoolloop "github.com/tingly-dev/tingly-box/internal/protocol/stage/toolloop" - mcpmodule "github.com/tingly-dev/tingly-box/internal/mcpserver" servertransform "github.com/tingly-dev/tingly-box/internal/protocolserver/transform" "github.com/tingly-dev/tingly-box/internal/typ" ) @@ -32,6 +32,9 @@ func (ph *ProtocolHandler) newProtocolStageBetaToolLoop( ), Executor: mcpmodule.NewServerToolExecutor(ph), Continuations: mcpmodule.NewProviderBetaContinuationStore(provider.UUID), + // Explicit so production does not silently depend on a package default: + // DefaultMaxRounds bounds server-tool executions per attempt. + MaxRounds: stagetoolloop.DefaultMaxRounds, }) } From 3527eeaae152192a5cadc36cb3f594c49380042d Mon Sep 17 00:00:00 2001 From: FFengIll Date: Tue, 4 Aug 2026 19:54:32 +0800 Subject: [PATCH 8/8] test(protocoltest): port cross-stage test matrix and CLI harness wiring Port the integration-test surface and user-facing wiring for the Protocol Stage feature from codex/protocol-stage-hardening: - internal/protocoltest: bridge/mcp/guardrail/recording matrix helpers + tests (bridge_matrix, mcp_matrix, mcp_recording_*, protocol_stage_server_test, protocol_stage_tool_loop_test, record_artifacts, guardrails); testenv gains rootServer + ProtocolStage/servertool options + ForceFlushRecordings; failover/matrix/virtual_client tests updated. Repointed stale imports (server/servertool -> protocolserver/servertool; server/module/mcp -> mcpserver). - internal/command: expose EnableProtocolStage + WithProtocolStage plumbing (options/server_options{,_test}.go, command/server.go). - cli/harness + cli/tingly-box: matrix and main_test updates. - gui/wails3/run.go: pass WithProtocolStage(opts.EnableProtocolStage) in all launch modes. Design notes for this feature live in a dedicated front-loaded commit (docs(protocol-stage)) at the head of the branch. Full module builds clean: go build ./... and go vet ./internal/... ./cli/... produces the tingly-box binary. (experiments/ is untracked working-tree scratch, not part of this port.) Batch 5 (final) of the protocol-stage-hardening port. --- cli/harness/main_test.go | 73 ++ cli/harness/matrix.go | 126 +- cli/tingly-box/main_test.go | 23 + gui/wails3/run.go | 3 + internal/command/options/server_options.go | 44 +- .../command/options/server_options_test.go | 29 + internal/command/server.go | 25 +- internal/protocoltest/bridge_matrix.go | 1144 +++++++++++++++++ internal/protocoltest/bridge_matrix_test.go | 99 ++ internal/protocoltest/failover_test.go | 72 ++ internal/protocoltest/guardrails.go | 24 + internal/protocoltest/matrix.go | 102 +- internal/protocoltest/mcp_matrix.go | 231 ++++ .../protocoltest/mcp_recording_matrix_test.go | 46 + .../protocoltest/mcp_recording_validation.go | 139 ++ .../protocol_stage_server_test.go | 980 ++++++++++++++ .../protocol_stage_tool_loop_test.go | 408 ++++++ internal/protocoltest/record_artifacts.go | 69 + internal/protocoltest/testenv.go | 82 +- internal/protocoltest/virtual_client_test.go | 11 + 20 files changed, 3658 insertions(+), 72 deletions(-) create mode 100644 internal/command/options/server_options_test.go create mode 100644 internal/protocoltest/bridge_matrix.go create mode 100644 internal/protocoltest/bridge_matrix_test.go create mode 100644 internal/protocoltest/guardrails.go create mode 100644 internal/protocoltest/mcp_matrix.go create mode 100644 internal/protocoltest/mcp_recording_matrix_test.go create mode 100644 internal/protocoltest/mcp_recording_validation.go create mode 100644 internal/protocoltest/protocol_stage_server_test.go create mode 100644 internal/protocoltest/protocol_stage_tool_loop_test.go create mode 100644 internal/protocoltest/record_artifacts.go diff --git a/cli/harness/main_test.go b/cli/harness/main_test.go index 29e8b40ce..8b38bf936 100644 --- a/cli/harness/main_test.go +++ b/cli/harness/main_test.go @@ -116,6 +116,79 @@ func TestMatrixCmdNonStreamingFlag(t *testing.T) { } } +func TestMatrixCmdStageGuardrailsFlags(t *testing.T) { + cli, parser := newTestParser(t) + if _, err := parser.Parse([]string{"matrix", "--mode=single", "--stage", "--guardrails"}); err != nil { + t.Fatalf("Stage Guardrails flags should parse: %v", err) + } + if !cli.Matrix.StageEnabled || !cli.Matrix.Guardrails { + t.Fatalf("parsed matrix = %+v", cli.Matrix) + } +} + +func TestMatrixCmdBridgeMode(t *testing.T) { + cli, parser := newTestParser(t) + if _, err := parser.Parse([]string{"matrix", "--mode=bridges", "--scenario=tool_result", "--streaming"}); err != nil { + t.Fatalf("bridge mode should parse: %v", err) + } + if cli.Matrix.Mode != "bridges" || len(cli.Matrix.Scenarios) != 1 || cli.Matrix.Scenarios[0] != "tool_result" || !cli.Matrix.Streaming { + t.Fatalf("parsed bridge matrix = %+v", cli.Matrix) + } +} + +func TestMatrixCmdBridgeModeRejectsExternalClient(t *testing.T) { + cmd := &MatrixCmd{Mode: "bridges", Client: "gosdk"} + err := cmd.Run() + if err == nil || !strings.Contains(err.Error(), "only supports --client=http") { + t.Fatalf("Run() error = %v", err) + } +} + +func TestMatrixCmdBridgeModeRejectsUnsupportedFeatures(t *testing.T) { + tests := []struct { + name string + cmd MatrixCmd + want string + }{ + {name: "mcp", cmd: MatrixCmd{Mode: "bridges", Client: "http", MCPEnabled: true}, want: "does not support --mcp"}, + {name: "stage", cmd: MatrixCmd{Mode: "bridges", Client: "http", StageEnabled: true}, want: "does not support --stage"}, + {name: "guardrails", cmd: MatrixCmd{Mode: "bridges", Client: "http", Guardrails: true}, want: "does not support --guardrails"}, + {name: "recording", cmd: MatrixCmd{Mode: "bridges", Client: "http", RecordDir: t.TempDir()}, want: "does not support --record-dir"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.cmd.Run() + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("Run() error = %v, want %q", err, tt.want) + } + }) + } +} + +func TestMatrixCmdRejectsFiltersWithNoExecutableCases(t *testing.T) { + cmd := &MatrixCmd{ + Mode: "single", + Client: "http", + Scenarios: []string{"does_not_exist"}, + } + err := cmd.Run() + if err == nil || !strings.Contains(err.Error(), "no executable test cases") { + t.Fatalf("Run() error = %v", err) + } +} + +func TestMatrixCmdOwnedToolScenarioRequiresStageAndMCP(t *testing.T) { + cmd := &MatrixCmd{ + Mode: "single", + Client: "http", + Scenarios: []string{"mcp_owned_tool"}, + } + err := cmd.Run() + if err == nil || !strings.Contains(err.Error(), "requires both --stage and --mcp") { + t.Fatalf("Run() error = %v", err) + } +} + func TestProviderTestCmdReturnsNotImplemented(t *testing.T) { var p ProviderTestCmd err := p.Run() diff --git a/cli/harness/matrix.go b/cli/harness/matrix.go index ada3b29bb..643f8c7b3 100644 --- a/cli/harness/matrix.go +++ b/cli/harness/matrix.go @@ -21,18 +21,20 @@ import ( // // --source and --target filter pairs by their source/target component. type MatrixCmd struct { - Scenarios []string `kong:"name='scenario',sep=',',help='Filter by scenario name (can repeat or comma-separate)'"` - Sources []string `kong:"name='source',sep=',',help='Filter by source protocol (can repeat or comma-separate)'"` - Targets []string `kong:"name='target',sep=',',help='Filter by target protocol (can repeat or comma-separate)'"` - Streaming bool `kong:"name='streaming',help='Run only streaming tests'"` - NonStream bool `kong:"name='non-streaming',help='Run only non-streaming tests'"` - Mode string `kong:"name='mode',default='default',enum='default,all,single,transitive,idempotent,flags,content_shapes,cache_controls',help='Section selection: default (single + idempotent round-trip; two-hop OFF), all (every section), single (A→B only), transitive (A→B→C only), idempotent (round-trip g(f(A))==A only), flags (per-rule flag behavior only), content_shapes (request content-shape regression only), cache_controls (single-hop + ABA cache/no-cache requests)'"` - Client string `kong:"name='client',default='http',enum='http,gosdk,python,node,aisdk',help='Client driver: http (raw JSON over net/http, default), gosdk (official anthropic-sdk-go / openai-go), python (real Python SDKs via subprocess driver), node (real Node SDKs via subprocess driver), aisdk (AI SDK by Vercel via subprocess driver)'"` - JsonOutput bool `kong:"name='json',help='Output results as JSON'"` - Verbose int `kong:"name='verbose',short='v',type='counter',help='Verbose output (repeat for more detail)'"` - RecordDir string `kong:"name='record-dir',env='HARNESS_RECORD_DIR',help='Directory for recording requests/responses (default: disabled)'"` - BatchCount int `kong:"name='batch',default='1',help='Number of times to run each test (for stability/performance testing)'"` - MCPEnabled bool `kong:"name='mcp',help='Enable MCP feature flag in test env'"` + Scenarios []string `kong:"name='scenario',sep=',',help='Filter by scenario name (can repeat or comma-separate)'"` + Sources []string `kong:"name='source',sep=',',help='Filter by source protocol (can repeat or comma-separate)'"` + Targets []string `kong:"name='target',sep=',',help='Filter by target protocol (can repeat or comma-separate)'"` + Streaming bool `kong:"name='streaming',help='Run only streaming tests'"` + NonStream bool `kong:"name='non-streaming',help='Run only non-streaming tests'"` + Mode string `kong:"name='mode',default='default',enum='default,all,single,transitive,idempotent,flags,content_shapes,cache_controls,bridges',help='Section selection: default (single + idempotent + dormant Bridges; two-hop OFF), all (every section), single (production A→B only), transitive (production A→B→C only), idempotent (production round-trip only), flags (per-rule flags only), content_shapes (request content-shape regression only), cache_controls (single-hop + ABA cache/no-cache requests), bridges (dormant Stage/Bridge topology only)'"` + Client string `kong:"name='client',default='http',enum='http,gosdk,python,node,aisdk',help='Client driver: http (raw JSON over net/http, default), gosdk (official anthropic-sdk-go / openai-go), python (real Python SDKs via subprocess driver), node (real Node SDKs via subprocess driver), aisdk (AI SDK by Vercel via subprocess driver)'"` + JsonOutput bool `kong:"name='json',help='Output results as JSON'"` + Verbose int `kong:"name='verbose',short='v',type='counter',help='Verbose output (repeat for more detail)'"` + RecordDir string `kong:"name='record-dir',env='HARNESS_RECORD_DIR',help='Directory for recording requests/responses (default: disabled)'"` + BatchCount int `kong:"name='batch',default='1',help='Number of times to run each test (for stability/performance testing)'"` + MCPEnabled bool `kong:"name='mcp',help='Enable MCP feature flag in test env'"` + StageEnabled bool `kong:"name='stage',help='Enable production Protocol Stage selection in the test server'"` + Guardrails bool `kong:"name='guardrails',help='Enable an active allow-only Guardrails runtime in the test server'"` } // matrixSection describes one runnable section of the validation matrix and @@ -40,10 +42,11 @@ type MatrixCmd struct { // here (plus its ExecuteAll* executor in internal/protocoltest) and extending // the --mode enum on MatrixCmd. type matrixSection struct { - name string - modes []string // --mode values that include this section - httpOnly bool // drives raw requests directly; requires --client=http - exec func(*protocoltest.Matrix) []protocoltest.TestResult + name string + modes []string // --mode values that include this section + httpOnly bool // requires --client=http + exec func(*protocoltest.Matrix) []protocoltest.TestResult + bridgeExec func(*protocoltest.BridgeMatrix) []protocoltest.TestResult } // matrixSections is the section registry: the single source of truth for what @@ -55,16 +58,17 @@ var matrixSections = []matrixSection{ {name: "flags", modes: []string{"all", "flags"}, httpOnly: true, exec: (*protocoltest.Matrix).ExecuteAllFlags}, {name: "content_shapes", modes: []string{"all", "content_shapes"}, httpOnly: true, exec: (*protocoltest.Matrix).ExecuteAllContentShapes}, {name: "cache_controls", modes: []string{"all", "cache_controls"}, httpOnly: true, exec: (*protocoltest.Matrix).ExecuteAllCacheControls}, + {name: "bridges", modes: []string{"default", "all", "bridges"}, httpOnly: true, bridgeExec: (*protocoltest.BridgeMatrix).ExecuteAll}, } // Help returns extended help text shown by `harness matrix --help`. func (*MatrixCmd) Help() string { return `Examples: - # Default: single-hop (A→B) + idempotent round-trips (g(f(A))==A). + # Default: production single-hop + idempotent round-trips + dormant Bridges. # Two-hop (A→B→C) transitive chains are OFF by default. harness matrix - # Run absolutely everything: single + two-hop + idempotent + # Run every section: single + two-hop + idempotent + flags + dormant Bridges harness matrix --mode=all # Run only two-hop (A→B→C) transitive chain tests @@ -82,9 +86,33 @@ func (*MatrixCmd) Help() string { # Run single-hop + ABA prompt-cache request tests harness matrix --mode=cache_controls + # Run only the dormant Stage/Bridge topology (no production dispatch claim) + harness matrix --mode=bridges + harness matrix --mode=bridges --source=anthropic_v1 --target=anthropic_beta + # Run only single-hop (A→B) tests harness matrix --mode=single + # Exercise production Stage selection (Chat/Beta/V1 routes plus Responses routes) + harness matrix --mode=single --stage --source=openai_chat --target=anthropic_beta + harness matrix --mode=single --stage --source=anthropic_beta --target=openai_chat + harness matrix --mode=single --stage --source=openai_responses --target=openai_responses + harness matrix --mode=single --stage --source=openai_responses --target=anthropic_beta + harness matrix --mode=single --stage --source=openai_responses --target=openai_chat + harness matrix --mode=single --stage --source=anthropic_beta --target=openai_responses + + # Exercise every production Stage route with a server-owned MCP tool loop + harness matrix --mode=single --stage --mcp --scenario=mcp_owned_tool + + # Persist and automatically validate RequestRecord boundaries for every MCP round + harness matrix --mode=single --stage --mcp --scenario=mcp_owned_tool --record-dir=/tmp/tingly-mcp-records + + # Exercise the opt-in Beta identity RequestRecord canary and retain artifacts + harness matrix --mode=single --stage --source=anthropic_beta --target=anthropic_beta --record-dir=/tmp/tingly-records + + # Exercise Beta Guardrail as a Stage without changing scenario semantics + harness matrix --mode=single --stage --guardrails --source=anthropic_beta + # Drive requests through real client stacks instead of raw HTTP harness matrix --mode=single --client=gosdk # official Go SDKs, in-process harness matrix --mode=single --client=python # real Python SDKs (subprocess driver) @@ -128,6 +156,9 @@ func (m *MatrixCmd) Run() error { if m.Streaming && m.NonStream { return fmt.Errorf("cannot specify both --streaming and --non-streaming") } + if m.Client != "http" && m.Mode == "bridges" { + return fmt.Errorf("--mode=bridges only supports --client=http (the Bridge matrix runs in-process and has no client transport)") + } if m.Client != "http" { for _, sec := range matrixSections { if sec.httpOnly && m.Mode == sec.name { @@ -135,6 +166,21 @@ func (m *MatrixCmd) Run() error { } } } + if m.Mode == "bridges" && m.MCPEnabled { + return fmt.Errorf("--mode=bridges does not support --mcp (the Bridge matrix validates protocol topology only)") + } + if m.Mode == "bridges" && m.StageEnabled { + return fmt.Errorf("--mode=bridges does not support --stage (use --mode=single to exercise the production Stage path)") + } + if m.Mode == "bridges" && m.Guardrails { + return fmt.Errorf("--mode=bridges does not support --guardrails (use --mode=single to exercise the production Guardrail path)") + } + if m.Mode == "bridges" && m.RecordDir != "" { + return fmt.Errorf("--mode=bridges does not support --record-dir (the Bridge matrix runs in-process without HTTP recording)") + } + if slices.Contains(m.Scenarios, protocoltest.MCPStageOwnedToolScenarioName) && (!m.MCPEnabled || !m.StageEnabled) { + return fmt.Errorf("--scenario=%s requires both --stage and --mcp", protocoltest.MCPStageOwnedToolScenarioName) + } client, err := resolveClient(m.Client) if err != nil { @@ -143,6 +189,9 @@ func (m *MatrixCmd) Run() error { // Build matrix with filters matrix := protocoltest.DefaultMatrix() + if m.MCPEnabled && m.StageEnabled { + matrix = matrix.WithMCPStageCoverage() + } if client != nil { matrix = matrix.WithClient(client) } @@ -171,6 +220,31 @@ func (m *MatrixCmd) Run() error { if m.MCPEnabled { matrix = matrix.WithMCPEnabled() } + if m.StageEnabled { + matrix = matrix.WithProtocolStage() + } + if m.Guardrails { + matrix = matrix.WithGuardrails() + } + bridgeMatrix := protocoltest.DefaultBridgeMatrix() + if len(m.Scenarios) > 0 { + bridgeMatrix = bridgeMatrix.OnlyScenarios(m.Scenarios...) + } + if len(m.Sources) > 0 { + bridgeMatrix = bridgeMatrix.OnlySources(m.Sources...) + } + if len(m.Targets) > 0 { + bridgeMatrix = bridgeMatrix.OnlyTargets(m.Targets...) + } + if m.Streaming { + bridgeMatrix = bridgeMatrix.OnlyStreaming(true) + } + if m.NonStream { + bridgeMatrix = bridgeMatrix.OnlyStreaming(false) + } + if m.BatchCount > 1 { + bridgeMatrix = bridgeMatrix.WithBatchCount(m.BatchCount) + } // Collect results for the sections the selected --mode includes (see // matrixSections for the mode → section mapping). @@ -185,9 +259,22 @@ func (m *MatrixCmd) Run() error { logrus.Warnf("skipping %s section: only supported with --client=http", sec.name) continue } + if sec.bridgeExec != nil { + combined = append(combined, sec.bridgeExec(bridgeMatrix)...) + continue + } combined = append(combined, sec.exec(matrix)...) } results := filterResults(combined, m) + executed := 0 + for _, result := range results { + if !result.Skipped { + executed++ + } + } + if executed == 0 { + return fmt.Errorf("no executable test cases matched the selected matrix filters") + } // Output results if m.JsonOutput { @@ -196,6 +283,9 @@ func (m *MatrixCmd) Run() error { } } else { printTable(results, verbose) + if m.RecordDir != "" && m.StageEnabled && m.MCPEnabled && slices.Contains(m.Scenarios, protocoltest.MCPStageOwnedToolScenarioName) { + fmt.Printf("\n📼 RequestRecord artifacts verified: %d case(s) in %s\n", executed, m.RecordDir) + } } // Determine exit code diff --git a/cli/tingly-box/main_test.go b/cli/tingly-box/main_test.go index 4caa1fa09..71d2c6337 100644 --- a/cli/tingly-box/main_test.go +++ b/cli/tingly-box/main_test.go @@ -180,6 +180,29 @@ func TestStartCmdDebugFlagSet(t *testing.T) { } } +func TestProtocolStageFlagParsesOnServerCommands(t *testing.T) { + tests := []struct { + name string + args []string + get func(*CLI) bool + }{ + {name: "start", args: []string{"start", "--stage"}, get: func(cli *CLI) bool { return cli.Start.EnableProtocolStage }}, + {name: "restart", args: []string{"restart", "--stage"}, get: func(cli *CLI) bool { return cli.Restart.EnableProtocolStage }}, + {name: "open", args: []string{"open", "--stage"}, get: func(cli *CLI) bool { return cli.Open.EnableProtocolStage }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cli, parser := newTestParser(t) + if _, err := parser.Parse(tt.args); err != nil { + t.Fatalf("%v should parse: %v", tt.args, err) + } + if !tt.get(cli) { + t.Fatalf("%v did not enable Protocol Stage", tt.args) + } + }) + } +} + // TestTUICommandAndQuickstartAlias ensures both `tui` (the canonical name) and // `quickstart` (the hidden legacy alias) parse without arguments. func TestTUICommandAndQuickstartAlias(t *testing.T) { diff --git a/gui/wails3/run.go b/gui/wails3/run.go index 9692e8bd9..2e72a4e49 100644 --- a/gui/wails3/run.go +++ b/gui/wails3/run.go @@ -258,6 +258,7 @@ func (l *appLauncher) StartGUI(appManager *command.AppManager, opts options.Star appManager.AppConfig(), server.WithUI(opts.EnableUI), server.WithDebug(opts.EnableDebug), + server.WithProtocolStage(opts.EnableProtocolStage), server.WithOpenBrowser(opts.EnableOpenBrowser), server.WithHost(opts.Host), server.WithRecordMode(recordMode), @@ -305,6 +306,7 @@ func (l *appLauncher) StartTray(appManager *command.AppManager, opts options.Sta appManager.AppConfig(), server.WithUI(opts.EnableUI), server.WithDebug(opts.EnableDebug), + server.WithProtocolStage(opts.EnableProtocolStage), server.WithOpenBrowser(opts.EnableOpenBrowser), server.WithHost(opts.Host), server.WithRecordMode(recordMode), @@ -353,6 +355,7 @@ func (l *appLauncher) StartSlim(appManager *command.AppManager, opts options.Sta appManager.AppConfig(), server.WithUI(opts.EnableUI), server.WithDebug(opts.EnableDebug), + server.WithProtocolStage(opts.EnableProtocolStage), server.WithOpenBrowser(opts.EnableOpenBrowser), server.WithHost(opts.Host), server.WithRecordMode(recordMode), diff --git a/internal/command/options/server_options.go b/internal/command/options/server_options.go index ba88b0e68..57aa3335e 100644 --- a/internal/command/options/server_options.go +++ b/internal/command/options/server_options.go @@ -12,6 +12,7 @@ type StartFlags struct { Host string EnableUI bool EnableDebug bool + EnableProtocolStage bool EnableOpenBrowser bool EnableStyleTransform bool Daemon bool @@ -23,16 +24,17 @@ type StartFlags struct { // StartServerOptions contains resolved options for starting the server type StartServerOptions struct { - Host string - Port int - EnableUI bool - EnableDebug bool - EnableOpenBrowser bool - Daemon bool - LogFile string - PromptRestart bool - RecordMode string - RecordDir string + Host string + Port int + EnableUI bool + EnableDebug bool + EnableProtocolStage bool + EnableOpenBrowser bool + Daemon bool + LogFile string + PromptRestart bool + RecordMode string + RecordDir string } // AddStartFlags adds all start-related flags to a command @@ -42,6 +44,7 @@ func AddStartFlags(cmd *cobra.Command, flags *StartFlags) { cmd.Flags().StringVar(&flags.Host, "host", "localhost", "Server host") cmd.Flags().BoolVarP(&flags.EnableUI, "ui", "u", true, "Enable web UI (default: true)") cmd.Flags().BoolVar(&flags.EnableDebug, "debug", false, "Enable debug mode including gin, low level logging and so on (default: false)") + cmd.Flags().BoolVar(&flags.EnableProtocolStage, "stage", false, "Enable the Protocol Stage request pipeline (default: false)") cmd.Flags().BoolVar(&flags.EnableOpenBrowser, "browser", true, "Auto-open browser when server starts (default: true)") cmd.Flags().BoolVar(&flags.EnableStyleTransform, "adapter", true, "Enable API style transformation (default: true)") cmd.Flags().BoolVar(&flags.Daemon, "daemon", false, "Run as daemon in background (default: false)") @@ -74,15 +77,16 @@ func ResolveStartOptions(cmd *cobra.Command, flags StartFlags, appConfig *config } return StartServerOptions{ - Host: flags.Host, - Port: resolvedPort, - EnableUI: flags.EnableUI, - EnableDebug: resolvedDebug, - EnableOpenBrowser: flags.EnableOpenBrowser, - Daemon: flags.Daemon, - LogFile: flags.LogFile, - PromptRestart: flags.PromptRestart, - RecordMode: flags.RecordMode, - RecordDir: resolvedRecordDir, + Host: flags.Host, + Port: resolvedPort, + EnableUI: flags.EnableUI, + EnableDebug: resolvedDebug, + EnableProtocolStage: flags.EnableProtocolStage, + EnableOpenBrowser: flags.EnableOpenBrowser, + Daemon: flags.Daemon, + LogFile: flags.LogFile, + PromptRestart: flags.PromptRestart, + RecordMode: flags.RecordMode, + RecordDir: resolvedRecordDir, } } diff --git a/internal/command/options/server_options_test.go b/internal/command/options/server_options_test.go new file mode 100644 index 000000000..ed208f227 --- /dev/null +++ b/internal/command/options/server_options_test.go @@ -0,0 +1,29 @@ +package options + +import ( + "testing" + + "github.com/spf13/cobra" + "github.com/tingly-dev/tingly-box/internal/config" +) + +func TestResolveStartOptionsProtocolStage(t *testing.T) { + t.Parallel() + + appConfig, err := config.NewAppConfig(config.WithConfigDir(t.TempDir())) + if err != nil { + t.Fatalf("NewAppConfig() error = %v", err) + } + t.Cleanup(func() { _ = appConfig.GetGlobalConfig().CloseStores() }) + + var flags StartFlags + cmd := &cobra.Command{Use: "start"} + AddStartFlags(cmd, &flags) + if err := cmd.ParseFlags([]string{"--stage"}); err != nil { + t.Fatalf("ParseFlags() error = %v", err) + } + resolved := ResolveStartOptions(cmd, flags, appConfig) + if !resolved.EnableProtocolStage { + t.Fatal("EnableProtocolStage = false, want true") + } +} diff --git a/internal/command/server.go b/internal/command/server.go index 9b6cd07d9..ca8bc500d 100644 --- a/internal/command/server.go +++ b/internal/command/server.go @@ -34,6 +34,7 @@ type StartCmdKong struct { Host string `kong:"flag,name='host',help='Server host'"` EnableUI bool `kong:"flag,name='ui',short='u',default='true',help='Enable web UI'"` EnableDebug bool `kong:"flag,name='debug',help='Enable debug mode'"` + EnableProtocolStage bool `kong:"flag,name='stage',help='Enable the Protocol Stage request pipeline'"` EnableOpenBrowser bool `kong:"flag,name='browser',default='true',help='Auto-open browser'"` EnableStyleTransform bool `kong:"flag,name='adapter',default='true',help='Enable API style transform'"` Daemon bool `kong:"flag,name='daemon',help='Run as daemon'"` @@ -54,6 +55,7 @@ func (s *StartCmdKong) Run(appManager *AppManager, source LaunchSource) error { Host: s.Host, EnableUI: s.EnableUI, EnableDebug: s.EnableDebug, + EnableProtocolStage: s.EnableProtocolStage, EnableOpenBrowser: s.EnableOpenBrowser, EnableStyleTransform: s.EnableStyleTransform, Daemon: s.Daemon, @@ -128,6 +130,7 @@ func (r *RestartCmdKong) Run(appManager *AppManager, source LaunchSource) error Host: r.Host, EnableUI: r.EnableUI, EnableDebug: r.EnableDebug, + EnableProtocolStage: r.EnableProtocolStage, EnableOpenBrowser: r.EnableOpenBrowser, EnableStyleTransform: r.EnableStyleTransform, Daemon: r.Daemon, @@ -302,16 +305,17 @@ func resolveStartCmdKongOptions(start *StartCmdKong, appConfig *config.AppConfig } return options.StartServerOptions{ - Host: start.Host, - Port: resolvedPort, - EnableUI: start.EnableUI, - EnableDebug: resolvedDebug, - EnableOpenBrowser: start.EnableOpenBrowser, - Daemon: start.Daemon, - LogFile: start.LogFile, - PromptRestart: start.PromptRestart, - RecordMode: start.RecordMode, - RecordDir: resolvedRecordDir, + Host: start.Host, + Port: resolvedPort, + EnableUI: start.EnableUI, + EnableDebug: resolvedDebug, + EnableProtocolStage: start.EnableProtocolStage, + EnableOpenBrowser: start.EnableOpenBrowser, + Daemon: start.Daemon, + LogFile: start.LogFile, + PromptRestart: start.PromptRestart, + RecordMode: start.RecordMode, + RecordDir: resolvedRecordDir, } } @@ -572,6 +576,7 @@ func startServerWithHook(appManager *AppManager, opts options.StartServerOptions serverManager := NewServerManager( appConfig, server.WithDebug(opts.EnableDebug), + server.WithProtocolStage(opts.EnableProtocolStage), server.WithUI(opts.EnableUI), server.WithOpenBrowser(opts.EnableOpenBrowser), server.WithHost(opts.Host), diff --git a/internal/protocoltest/bridge_matrix.go b/internal/protocoltest/bridge_matrix.go new file mode 100644 index 000000000..2b6138206 --- /dev/null +++ b/internal/protocoltest/bridge_matrix.go @@ -0,0 +1,1144 @@ +package protocoltest + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "time" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/openai/openai-go/v3" + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocol/sse" + "github.com/tingly-dev/tingly-box/internal/protocol/stage" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/anthropicbridge" + "github.com/tingly-dev/tingly-box/internal/protocol/stage/openaibridge" + protocolstream "github.com/tingly-dev/tingly-box/internal/protocol/stream" + "github.com/tingly-dev/tingly-box/internal/protocol/wire" +) + +const bridgeMatrixModel = "bridge-client-model" + +// BridgeMatrix validates the dormant protocol Stage/Bridge topology without +// claiming to traverse the production gateway dispatch path. It intentionally +// uses the same TestResult shape as Matrix so CLI filtering and output remain +// consistent while result names make the execution surface explicit. +type BridgeMatrix struct { + Pairs []ProtocolPair + Chains []BridgeChain + Scenarios []Scenario + Streaming []bool + BatchCount int +} + +// BridgeChain describes one concrete multi-level topology. Source and Target +// are the client and terminal protocols; Stage is the native middle protocol. +type BridgeChain struct { + Name string + Source protocol.APIType + Stage protocol.APIType + Target protocol.APIType +} + +// DefaultBridgePairs lists only concrete Stage boundaries implemented today. +func DefaultBridgePairs() []ProtocolPair { + return []ProtocolPair{ + {Source: protocol.TypeAnthropicV1, Target: protocol.TypeAnthropicV1}, + {Source: protocol.TypeAnthropicV1, Target: protocol.TypeAnthropicBeta}, + {Source: protocol.TypeAnthropicBeta, Target: protocol.TypeAnthropicBeta}, + {Source: protocol.TypeOpenAIChat, Target: protocol.TypeOpenAIChat}, + {Source: protocol.TypeAnthropicV1, Target: protocol.TypeOpenAIChat}, + {Source: protocol.TypeAnthropicBeta, Target: protocol.TypeOpenAIChat}, + {Source: protocol.TypeOpenAIChat, Target: protocol.TypeAnthropicBeta}, + } +} + +// DefaultBridgeChains contains real concrete Bridges on both sides of an +// Anthropic Beta-native Stage. It remains in-process and carries no production +// traffic. +func DefaultBridgeChains() []BridgeChain { + return []BridgeChain{ + { + Name: "v1_beta_stage_chat", + Source: protocol.TypeAnthropicV1, + Stage: protocol.TypeAnthropicBeta, + Target: protocol.TypeOpenAIChat, + }, + { + Name: "chat_beta_stage_chat", + Source: protocol.TypeOpenAIChat, + Stage: protocol.TypeAnthropicBeta, + Target: protocol.TypeOpenAIChat, + }, + } +} + +// DefaultBridgeMatrix covers request/response semantics that the first +// Anthropic Bridges declare: text, tool use, tool results, complete, and stream. +func DefaultBridgeMatrix() *BridgeMatrix { + return &BridgeMatrix{ + Pairs: DefaultBridgePairs(), + Chains: DefaultBridgeChains(), + Scenarios: []Scenario{ + TextScenario(), + ToolUseScenario(), + ToolResultScenario(), + }, + Streaming: []bool{false, true}, + } +} + +func (m *BridgeMatrix) clone() *BridgeMatrix { + copy := *m + return © +} + +func (m *BridgeMatrix) OnlyScenarios(names ...string) *BridgeMatrix { + wanted := make(map[string]bool, len(names)) + for _, name := range names { + wanted[name] = true + } + filtered := make([]Scenario, 0, len(names)) + for _, scenario := range m.Scenarios { + if wanted[scenario.Name] { + filtered = append(filtered, scenario) + } + } + out := m.clone() + out.Scenarios = filtered + return out +} + +func (m *BridgeMatrix) OnlySources(sources ...string) *BridgeMatrix { + wanted := make(map[protocol.APIType]bool, len(sources)) + for _, source := range sources { + wanted[protocol.APIType(source)] = true + } + filtered := make([]ProtocolPair, 0, len(m.Pairs)) + for _, pair := range m.Pairs { + if wanted[pair.Source] { + filtered = append(filtered, pair) + } + } + out := m.clone() + out.Pairs = filtered + filteredChains := make([]BridgeChain, 0, len(m.Chains)) + for _, chain := range m.Chains { + if wanted[chain.Source] { + filteredChains = append(filteredChains, chain) + } + } + out.Chains = filteredChains + return out +} + +func (m *BridgeMatrix) OnlyTargets(targets ...string) *BridgeMatrix { + wanted := make(map[protocol.APIType]bool, len(targets)) + for _, target := range targets { + wanted[protocol.APIType(target)] = true + } + filtered := make([]ProtocolPair, 0, len(m.Pairs)) + for _, pair := range m.Pairs { + if wanted[pair.Target] { + filtered = append(filtered, pair) + } + } + out := m.clone() + out.Pairs = filtered + filteredChains := make([]BridgeChain, 0, len(m.Chains)) + for _, chain := range m.Chains { + if wanted[chain.Target] { + filteredChains = append(filteredChains, chain) + } + } + out.Chains = filteredChains + return out +} + +func (m *BridgeMatrix) OnlyStreaming(streaming bool) *BridgeMatrix { + out := m.clone() + out.Streaming = []bool{streaming} + return out +} + +func (m *BridgeMatrix) WithBatchCount(count int) *BridgeMatrix { + out := m.clone() + out.BatchCount = count + return out +} + +// ExecuteAll runs the in-process Bridge matrix and returns CLI-shaped results. +func (m *BridgeMatrix) ExecuteAll() []TestResult { + routes := make([]bridgeMatrixRoute, 0, len(m.Pairs)+len(m.Chains)) + for _, pair := range m.Pairs { + routes = append(routes, bridgeMatrixRoute{Pair: pair}) + } + for _, chain := range m.Chains { + routes = append(routes, bridgeMatrixRoute{ + Name: chain.Name, + Pair: ProtocolPair{Source: chain.Source, Target: chain.Target}, + Stage: chain.Stage, + }) + } + results := make([]TestResult, 0, len(routes)*len(m.Scenarios)*len(m.Streaming)) + for _, scenario := range m.Scenarios { + for _, route := range routes { + for _, streaming := range m.Streaming { + if m.BatchCount > 1 { + results = append(results, m.executeBatch(scenario, route, streaming)) + continue + } + results = append(results, m.executeOne(scenario, route, streaming)) + } + } + } + return results +} + +type bridgeMatrixRoute struct { + Name string + Pair ProtocolPair + Stage protocol.APIType +} + +func (r bridgeMatrixRoute) isChain() bool { return r.Name != "" } + +func (r bridgeMatrixRoute) scenarioName(scenario string) string { + if r.isChain() { + return "bridges/chain/" + r.Name + "/" + scenario + } + return "bridges/" + scenario +} + +func (r bridgeMatrixRoute) resultName(scenario string, streaming bool) string { + if r.isChain() { + return fmt.Sprintf("bridges/chain/%s/%s/%s/%s/%s", r.Name, scenario, r.Pair.Source, r.Pair.Target, streamMode(streaming)) + } + return fmt.Sprintf("bridges/%s/%s/%s/%s", scenario, r.Pair.Source, r.Pair.Target, streamMode(streaming)) +} + +func (m *BridgeMatrix) executeOne(scenario Scenario, route bridgeMatrixRoute, streaming bool) TestResult { + start := time.Now() + pair := route.Pair + result := TestResult{ + Name: route.resultName(scenario.Name, streaming), + Scenario: route.scenarioName(scenario.Name), + Source: pair.Source, + Target: pair.Target, + Streaming: streaming, + } + + request, err := bridgeMatrixRequest(pair.Source, scenario.Name) + if err != nil { + return bridgeMatrixFailure(result, start, "fixture/request", err, "") + } + terminal, err := newBridgeMatrixTerminal(route, scenario.Name, streaming) + if err != nil { + return bridgeMatrixFailure(result, start, "fixture/terminal", err, "") + } + endpoint, probe, err := bridgeMatrixEndpoint(terminal, route) + if err != nil { + return bridgeMatrixFailure(result, start, "topology", err, "") + } + + call := stage.Call{ + Request: request, + Metadata: stage.CallMetadata{ + RequestID: fmt.Sprintf("bridge-matrix-%s-%s-%s", scenario.Name, pair.Source, streamMode(streaming)), + }, + } + var semantic *RoundTripResult + var failures []AssertionError + if streaming { + stream, streamErr := endpoint.Stream(context.Background(), call) + if streamErr != nil { + return bridgeMatrixFailure(result, start, "stream/open", streamErr, "") + } + semantic, failures = consumeBridgeMatrixStream(stream, route, scenario.Name) + if terminal.lastStream == nil || terminal.lastStream.closeCount != 1 { + failures = append(failures, AssertionError{ + Assertion: "stream/ownership", + Error: fmt.Sprintf("target close count = %d, want 1", terminal.streamCloseCount()), + }) + } + } else { + response, completeErr := endpoint.Complete(context.Background(), call) + if completeErr != nil { + return bridgeMatrixFailure(result, start, "complete", completeErr, "") + } + semantic, failures = parseBridgeMatrixResponse(response, route, scenario.Name) + } + + failures = append(failures, validateBridgeMatrixTargetCall(terminal.lastCall, request, route, scenario.Name, streaming)...) + failures = append(failures, validateBridgeMatrixProbe(probe, streaming)...) + if semantic != nil { + failures = append(failures, runBridgeMatrixAssertions(semantic, scenario.Name)...) + result.Response = semantic + } + result.Passed = len(failures) == 0 + result.Errors = failures + result.Duration = time.Since(start) + return result +} + +func (m *BridgeMatrix) executeBatch(scenario Scenario, route bridgeMatrixRoute, streaming bool) TestResult { + count := m.BatchCount + runs := make([]TestResult, 0, count) + passed := 0 + var total, min, max time.Duration + uniqueErrors := make(map[string]AssertionError) + for i := 0; i < count; i++ { + run := m.executeOne(scenario, route, streaming) + runs = append(runs, run) + total += run.Duration + if i == 0 || run.Duration < min { + min = run.Duration + } + if run.Duration > max { + max = run.Duration + } + if run.Passed { + passed++ + } + for _, failure := range run.Errors { + uniqueErrors[failure.Assertion+"\x00"+failure.Error] = failure + } + } + errors := make([]AssertionError, 0, len(uniqueErrors)) + batchErrors := make([]string, 0, len(uniqueErrors)) + for _, failure := range uniqueErrors { + errors = append(errors, failure) + batchErrors = append(batchErrors, failure.Error) + } + last := runs[len(runs)-1] + last.Passed = passed == count + last.Errors = errors + last.Duration = total / time.Duration(count) + last.BatchCount = count + last.BatchPassed = passed + last.BatchMinDur = min + last.BatchAvgDur = last.Duration + last.BatchMaxDur = max + last.BatchErrors = batchErrors + return last +} + +func bridgeMatrixFailure(result TestResult, start time.Time, assertion string, err error, context string) TestResult { + result.Duration = time.Since(start) + result.Errors = []AssertionError{{Assertion: assertion, Error: err.Error(), Context: context}} + return result +} + +func bridgeMatrixEndpoint(terminal stage.Endpoint, route bridgeMatrixRoute) (stage.Endpoint, *bridgeMatrixProbeStage, error) { + registry, err := stage.NewBridgeRegistry( + anthropicbridge.NewV1ToBeta(), + anthropicbridge.NewV1ToOpenAIChat(anthropicbridge.ChatOptions{}), + anthropicbridge.NewBetaToOpenAIChat(anthropicbridge.ChatOptions{}), + openaibridge.NewChatToAnthropicBeta(openaibridge.AnthropicOptions{}), + ) + if err != nil { + return nil, nil, err + } + if !route.isChain() && route.Pair.Source == route.Pair.Target { + identity, err := registry.Resolve(route.Pair.Source, route.Pair.Target, stage.AllBridgeCapabilities) + if err != nil { + return nil, nil, err + } + endpoint, err := stage.Adapt(terminal, identity) + return endpoint, nil, err + } + var ( + stages []stage.Stage + probe *bridgeMatrixProbeStage + ) + if route.isChain() { + probe = &bridgeMatrixProbeStage{api: route.Stage} + stages = []stage.Stage{probe} + } + endpoint, err := stage.BuildTopology(stage.TopologyConfig{ + Terminal: terminal, + Stages: stages, + ClientProtocol: route.Pair.Source, + Registry: registry, + RequiredCapabilities: stage.AllBridgeCapabilities, + }) + return endpoint, probe, err +} + +// bridgeMatrixProbeStage proves that the concrete middle level receives and +// returns its declared native protocol in both execution modes. +type bridgeMatrixProbeStage struct { + api protocol.APIType + + completeRequest any + completeResponse any + streamRequest any + streamEvents int + streamTypeError error +} + +func (*bridgeMatrixProbeStage) Name() string { return "bridge-matrix-probe" } + +func (s *bridgeMatrixProbeStage) Protocol() protocol.APIType { return s.api } + +func (s *bridgeMatrixProbeStage) Wrap(next stage.Endpoint) stage.Endpoint { + return &bridgeMatrixProbeEndpoint{stage: s, next: next} +} + +type bridgeMatrixProbeEndpoint struct { + stage *bridgeMatrixProbeStage + next stage.Endpoint +} + +func (e *bridgeMatrixProbeEndpoint) Protocol() protocol.APIType { return e.stage.api } + +func (e *bridgeMatrixProbeEndpoint) Complete(ctx context.Context, call stage.Call) (*stage.Response, error) { + e.stage.completeRequest = call.Request + response, err := e.next.Complete(ctx, call) + if response != nil { + e.stage.completeResponse = response.Value + } + return response, err +} + +func (e *bridgeMatrixProbeEndpoint) Stream(ctx context.Context, call stage.Call) (stage.EventStream, error) { + e.stage.streamRequest = call.Request + stream, err := e.next.Stream(ctx, call) + if err != nil { + return nil, err + } + return &bridgeMatrixProbeStream{stage: e.stage, next: stream}, nil +} + +type bridgeMatrixProbeStream struct { + stage *bridgeMatrixProbeStage + next stage.EventStream +} + +func (s *bridgeMatrixProbeStream) Next(ctx context.Context) (stage.Event, error) { + event, err := s.next.Next(ctx) + if err != nil { + return event, err + } + s.stage.streamEvents++ + switch event.Value.(type) { + case protocolstream.AnthropicEvent, anthropic.BetaRawMessageStreamEventUnion, *anthropic.BetaRawMessageStreamEventUnion: + default: + s.stage.streamTypeError = fmt.Errorf("middle Stage received %T, want Anthropic Beta event", event.Value) + } + return event, nil +} + +func (s *bridgeMatrixProbeStream) Close() error { return s.next.Close() } + +func (s *bridgeMatrixProbeStream) Result() stage.StreamResult { return s.next.Result() } + +func validateBridgeMatrixProbe(probe *bridgeMatrixProbeStage, streaming bool) []AssertionError { + if probe == nil { + return nil + } + var failures []AssertionError + if streaming { + if _, ok := probe.streamRequest.(*anthropic.BetaMessageNewParams); !ok { + failures = append(failures, AssertionError{Assertion: "chain/stage_request", Error: fmt.Sprintf("got %T, want *anthropic.BetaMessageNewParams", probe.streamRequest)}) + } + if probe.streamEvents == 0 { + failures = append(failures, AssertionError{Assertion: "chain/stage_events", Error: "middle Stage observed no response events"}) + } + if probe.streamTypeError != nil { + failures = append(failures, AssertionError{Assertion: "chain/stage_event_type", Error: probe.streamTypeError.Error()}) + } + return failures + } + if _, ok := probe.completeRequest.(*anthropic.BetaMessageNewParams); !ok { + failures = append(failures, AssertionError{Assertion: "chain/stage_request", Error: fmt.Sprintf("got %T, want *anthropic.BetaMessageNewParams", probe.completeRequest)}) + } + if _, ok := probe.completeResponse.(*anthropic.BetaMessage); !ok { + failures = append(failures, AssertionError{Assertion: "chain/stage_response", Error: fmt.Sprintf("got %T, want *anthropic.BetaMessage", probe.completeResponse)}) + } + return failures +} + +func bridgeMatrixRequest(source protocol.APIType, scenario string) (any, error) { + if source == protocol.TypeOpenAIChat { + return bridgeMatrixChatRequest(scenario) + } + messages := []any{ + map[string]any{"role": "user", "content": []any{map[string]any{"type": "text", "text": "What is the weather in Paris?"}}}, + } + request := map[string]any{ + "model": bridgeMatrixModel, + "max_tokens": 128, + "messages": messages, + } + switch scenario { + case "text": + case "tool_use": + request["tools"] = []any{bridgeMatrixToolDefinition()} + case "tool_result": + request["tools"] = []any{bridgeMatrixToolDefinition()} + request["messages"] = []any{ + messages[0], + map[string]any{"role": "assistant", "content": []any{map[string]any{ + "type": "tool_use", "id": "toolu_bridge_weather", "name": "get_weather", + "input": map[string]any{"location": "Paris"}, + }}}, + map[string]any{"role": "user", "content": []any{map[string]any{ + "type": "tool_result", "tool_use_id": "toolu_bridge_weather", "content": "18°C and sunny", "is_error": false, + }}}, + } + default: + return nil, fmt.Errorf("unsupported bridge scenario %q", scenario) + } + + switch source { + case protocol.TypeAnthropicV1: + value, err := decodeBridgeFixture[anthropic.MessageNewParams](request) + return &value, err + case protocol.TypeAnthropicBeta: + value, err := decodeBridgeFixture[anthropic.BetaMessageNewParams](request) + return &value, err + default: + return nil, fmt.Errorf("unsupported bridge source %q", source) + } +} + +func bridgeMatrixChatRequest(scenario string) (any, error) { + request := map[string]any{ + "model": bridgeMatrixModel, + "max_tokens": 128, + "messages": []any{ + map[string]any{"role": "user", "content": "What is the weather in Paris?"}, + }, + } + switch scenario { + case "text": + case "tool_use": + request["tools"] = []any{bridgeMatrixChatToolDefinition()} + case "tool_result": + request["tools"] = []any{bridgeMatrixChatToolDefinition()} + request["messages"] = []any{ + map[string]any{"role": "user", "content": "What is the weather in Paris?"}, + map[string]any{ + "role": "assistant", "content": nil, + "tool_calls": []any{map[string]any{ + "id": "toolu_bridge_weather", "type": "function", + "function": map[string]any{"name": "get_weather", "arguments": `{"location":"Paris"}`}, + }}, + }, + map[string]any{"role": "tool", "tool_call_id": "toolu_bridge_weather", "content": "18°C and sunny"}, + } + default: + return nil, fmt.Errorf("unsupported bridge scenario %q", scenario) + } + value, err := decodeBridgeFixture[openai.ChatCompletionNewParams](request) + return &value, err +} + +func bridgeMatrixToolDefinition() map[string]any { + return map[string]any{ + "name": "get_weather", + "description": "Get weather for a location", + "input_schema": map[string]any{ + "type": "object", + "properties": map[string]any{"location": map[string]any{"type": "string"}}, + "required": []string{"location"}, + }, + } +} + +func bridgeMatrixChatToolDefinition() map[string]any { + definition := bridgeMatrixToolDefinition() + return map[string]any{ + "type": "function", + "function": map[string]any{ + "name": definition["name"], + "description": definition["description"], + "parameters": definition["input_schema"], + }, + } +} + +type bridgeMatrixTerminal struct { + api protocol.APIType + response *stage.Response + events []stage.Event + result stage.StreamResult + lastCall stage.Call + lastStream *bridgeMatrixEventStream +} + +func newBridgeMatrixTerminal(route bridgeMatrixRoute, scenario string, streaming bool) (*bridgeMatrixTerminal, error) { + target := route.Pair.Target + terminal := &bridgeMatrixTerminal{api: target} + usage := protocol.NewTokenUsage(10, 8) + terminalModel := "bridge-provider-model" + if !route.isChain() && route.Pair.Source == route.Pair.Target { + terminalModel = bridgeMatrixModel + } + terminal.result = stage.StreamResult{Usage: usage, Model: terminalModel} + if target == protocol.TypeOpenAIChat { + if streaming { + events, err := bridgeMatrixChatEvents(scenario, terminalModel) + if err != nil { + return nil, err + } + terminal.events = events + return terminal, nil + } + response, err := bridgeMatrixChatResponse(scenario, terminalModel) + if err != nil { + return nil, err + } + terminal.response = &stage.Response{Value: response, Usage: usage, Model: terminalModel} + return terminal, nil + } + + if streaming { + events, err := bridgeMatrixAnthropicEvents(target, scenario, terminalModel) + if err != nil { + return nil, err + } + terminal.events = events + return terminal, nil + } + response, err := bridgeMatrixAnthropicResponse(target, scenario, terminalModel) + if err != nil { + return nil, err + } + terminal.response = &stage.Response{Value: response, Usage: usage, Model: terminalModel} + return terminal, nil +} + +func (e *bridgeMatrixTerminal) Protocol() protocol.APIType { return e.api } + +func (e *bridgeMatrixTerminal) Complete(_ context.Context, call stage.Call) (*stage.Response, error) { + e.lastCall = call + if e.response == nil { + return nil, fmt.Errorf("bridge matrix terminal %q has no complete fixture", e.api) + } + copy := *e.response + return ©, nil +} + +func (e *bridgeMatrixTerminal) Stream(_ context.Context, call stage.Call) (stage.EventStream, error) { + e.lastCall = call + e.lastStream = &bridgeMatrixEventStream{ + events: append([]stage.Event(nil), e.events...), + result: e.result, + } + return e.lastStream, nil +} + +func (e *bridgeMatrixTerminal) streamCloseCount() int { + if e.lastStream == nil { + return 0 + } + return e.lastStream.closeCount +} + +type bridgeMatrixEventStream struct { + events []stage.Event + result stage.StreamResult + closeCount int +} + +func (s *bridgeMatrixEventStream) Next(ctx context.Context) (stage.Event, error) { + if err := ctx.Err(); err != nil { + return stage.Event{}, err + } + if len(s.events) == 0 { + return stage.Event{}, io.EOF + } + event := s.events[0] + s.events = s.events[1:] + return event, nil +} + +func (s *bridgeMatrixEventStream) Close() error { + s.closeCount++ + return nil +} + +func (s *bridgeMatrixEventStream) Result() stage.StreamResult { return s.result } + +func bridgeMatrixAnthropicResponse(target protocol.APIType, scenario, model string) (any, error) { + body := bridgeMatrixAnthropicResponseBody(scenario, model) + switch target { + case protocol.TypeAnthropicV1: + value, err := decodeBridgeFixture[anthropic.Message](body) + return &value, err + case protocol.TypeAnthropicBeta: + value, err := decodeBridgeFixture[anthropic.BetaMessage](body) + return &value, err + default: + return nil, fmt.Errorf("unsupported Anthropic target %q", target) + } +} + +func bridgeMatrixAnthropicResponseBody(scenario, model string) map[string]any { + content := []any{map[string]any{"type": "text", "text": "The capital of France is Paris."}} + stopReason := "end_turn" + usage := map[string]any{"input_tokens": 10, "output_tokens": 8} + if scenario == "tool_use" { + content = []any{map[string]any{ + "type": "tool_use", "id": "toolu_bridge_weather", "name": "get_weather", + "input": map[string]any{"location": "Paris", "unit": "celsius"}, + }} + stopReason = "tool_use" + usage = map[string]any{"input_tokens": 15, "output_tokens": 20} + } + return map[string]any{ + "id": "msg_bridge_matrix", "type": "message", "role": "assistant", + "content": content, "model": model, "stop_reason": stopReason, + "stop_sequence": nil, "usage": usage, + } +} + +func bridgeMatrixChatResponse(scenario, model string) (*openai.ChatCompletion, error) { + message := map[string]any{"role": "assistant", "content": "The capital of France is Paris."} + finishReason := "stop" + usage := map[string]any{"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18} + if scenario == "tool_use" { + message = map[string]any{ + "role": "assistant", "content": nil, + "tool_calls": []any{map[string]any{ + "id": "call_bridge_weather", "type": "function", + "function": map[string]any{"name": "get_weather", "arguments": `{"location":"Paris","unit":"celsius"}`}, + }}, + } + finishReason = "tool_calls" + usage = map[string]any{"prompt_tokens": 15, "completion_tokens": 20, "total_tokens": 35} + } + body := map[string]any{ + "id": "chatcmpl_bridge_matrix", "object": "chat.completion", "created": 1, + "model": model, + "choices": []any{map[string]any{"index": 0, "message": message, "finish_reason": finishReason}}, + "usage": usage, + } + value, err := decodeBridgeFixture[openai.ChatCompletion](body) + return &value, err +} + +func bridgeMatrixChatEvents(scenario, model string) ([]stage.Event, error) { + var chunks []map[string]any + if scenario == "tool_use" { + chunks = []map[string]any{ + { + "id": "chatcmpl_bridge_stream", "object": "chat.completion.chunk", "created": 1, "model": model, + "choices": []any{map[string]any{ + "index": 0, "finish_reason": "", "delta": map[string]any{ + "role": "assistant", "tool_calls": []any{map[string]any{ + "index": 0, "id": "call_bridge_weather", "type": "function", + "function": map[string]any{"name": "get_weather", "arguments": `{"location":"Paris","unit":"celsius"}`}, + }}, + }, + }}, + }, + bridgeMatrixChatFinishChunk("tool_calls", model), + bridgeMatrixChatUsageChunk(15, 20, model), + } + } else { + chunks = []map[string]any{ + { + "id": "chatcmpl_bridge_stream", "object": "chat.completion.chunk", "created": 1, "model": model, + "choices": []any{map[string]any{ + "index": 0, "finish_reason": "", "delta": map[string]any{"role": "assistant", "content": "The capital of France is Paris."}, + }}, + }, + bridgeMatrixChatFinishChunk("stop", model), + bridgeMatrixChatUsageChunk(10, 8, model), + } + } + events := make([]stage.Event, 0, len(chunks)) + for _, chunk := range chunks { + value, err := decodeBridgeFixture[openai.ChatCompletionChunk](chunk) + if err != nil { + return nil, err + } + events = append(events, stage.Event{Value: value}) + } + return events, nil +} + +func bridgeMatrixChatFinishChunk(reason, model string) map[string]any { + return map[string]any{ + "id": "chatcmpl_bridge_stream", "object": "chat.completion.chunk", "created": 1, "model": model, + "choices": []any{map[string]any{"index": 0, "finish_reason": reason, "delta": map[string]any{}}}, + } +} + +func bridgeMatrixChatUsageChunk(input, output int, model string) map[string]any { + return map[string]any{ + "id": "chatcmpl_bridge_stream", "object": "chat.completion.chunk", "created": 1, "model": model, + "choices": []any{}, + "usage": map[string]any{"prompt_tokens": input, "completion_tokens": output, "total_tokens": input + output}, + } +} + +func bridgeMatrixAnthropicEvents(target protocol.APIType, scenario, model string) ([]stage.Event, error) { + inputTokens, outputTokens := 10, 8 + stopReason := "end_turn" + var contentEvents []map[string]any + if scenario == "tool_use" { + inputTokens, outputTokens = 15, 20 + stopReason = "tool_use" + contentEvents = []map[string]any{ + {"type": "content_block_start", "index": 0, "content_block": map[string]any{"type": "tool_use", "id": "toolu_bridge_weather", "name": "get_weather", "input": map[string]any{}}}, + {"type": "content_block_delta", "index": 0, "delta": map[string]any{"type": "input_json_delta", "partial_json": `{"location":"Paris","unit":"celsius"}`}}, + {"type": "content_block_stop", "index": 0}, + } + } else { + contentEvents = []map[string]any{ + {"type": "content_block_start", "index": 0, "content_block": map[string]any{"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": map[string]any{"type": "text_delta", "text": "The capital of France is Paris."}}, + {"type": "content_block_stop", "index": 0}, + } + } + events := []map[string]any{ + { + "type": "message_start", + "message": map[string]any{ + "id": "msg_bridge_stream", "type": "message", "role": "assistant", "content": []any{}, + "model": model, "stop_reason": nil, "stop_sequence": nil, + "usage": map[string]any{"input_tokens": inputTokens, "output_tokens": 0}, + }, + }, + } + events = append(events, contentEvents...) + events = append(events, + map[string]any{"type": "message_delta", "delta": map[string]any{"stop_reason": stopReason, "stop_sequence": nil}, "usage": map[string]any{"output_tokens": outputTokens}}, + map[string]any{"type": "message_stop"}, + ) + + result := make([]stage.Event, 0, len(events)) + for _, event := range events { + switch target { + case protocol.TypeAnthropicV1: + value, err := decodeBridgeFixture[anthropic.MessageStreamEventUnion](event) + if err != nil { + return nil, err + } + result = append(result, stage.Event{Value: value}) + case protocol.TypeAnthropicBeta: + value, err := decodeBridgeFixture[anthropic.BetaRawMessageStreamEventUnion](event) + if err != nil { + return nil, err + } + result = append(result, stage.Event{Value: value}) + default: + return nil, fmt.Errorf("unsupported Anthropic target %q", target) + } + } + return result, nil +} + +func parseBridgeMatrixResponse(response *stage.Response, route bridgeMatrixRoute, scenario string) (*RoundTripResult, []AssertionError) { + if response == nil { + return nil, []AssertionError{{Assertion: "response", Error: "response is nil"}} + } + pair := route.Pair + var failures []AssertionError + switch pair.Source { + case protocol.TypeAnthropicV1: + if _, ok := response.Value.(*anthropic.Message); !ok { + failures = append(failures, AssertionError{Assertion: "response/type", Error: fmt.Sprintf("got %T, want *anthropic.Message", response.Value)}) + } + case protocol.TypeAnthropicBeta: + if _, ok := response.Value.(*anthropic.BetaMessage); !ok { + failures = append(failures, AssertionError{Assertion: "response/type", Error: fmt.Sprintf("got %T, want *anthropic.BetaMessage", response.Value)}) + } + case protocol.TypeOpenAIChat: + switch response.Value.(type) { + case *openai.ChatCompletion, openai.ChatCompletion, wire.ChatCompletionWire, *wire.ChatCompletionWire: + default: + failures = append(failures, AssertionError{Assertion: "response/type", Error: fmt.Sprintf("got %T, want OpenAI Chat response value", response.Value)}) + } + } + raw, err := json.Marshal(response.Value) + if err != nil { + failures = append(failures, AssertionError{Assertion: "response/marshal", Error: err.Error()}) + return nil, failures + } + var body map[string]any + if err := json.Unmarshal(raw, &body); err != nil { + failures = append(failures, AssertionError{Assertion: "response/decode", Error: err.Error(), Context: string(raw)}) + return nil, failures + } + var parsed *sse.ParsedResult + if pair.Source == protocol.TypeOpenAIChat { + parsed = sse.ParseOpenAIChatResult(body) + } else { + parsed = sse.ParseAnthropicResult(body) + } + result := roundTripFromBridgeParsed(parsed, route, scenario, false, raw, nil) + if response.Model != bridgeMatrixModel { + failures = append(failures, AssertionError{Assertion: "response/model_fact", Error: fmt.Sprintf("got %q, want %q", response.Model, bridgeMatrixModel)}) + } + if response.Usage == nil || !response.Usage.HasUsage() { + failures = append(failures, AssertionError{Assertion: "response/usage_fact", Error: fmt.Sprintf("missing normalized usage: %+v", response.Usage)}) + } + return result, failures +} + +func consumeBridgeMatrixStream(stream stage.EventStream, route bridgeMatrixRoute, scenario string) (*RoundTripResult, []AssertionError) { + var failures []AssertionError + var eventLines []string + for { + event, err := stream.Next(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + failures = append(failures, AssertionError{Assertion: "stream/next", Error: err.Error()}) + break + } + if route.Pair.Source == protocol.TypeOpenAIChat { + data, err := normalizeBridgeMatrixChatEvent(event.Value) + if err != nil { + failures = append(failures, AssertionError{Assertion: "stream/event", Error: err.Error()}) + continue + } + eventLines = append(eventLines, "data: "+string(data)) + } else { + eventType, data, err := normalizeBridgeMatrixEvent(event.Value) + if err != nil { + failures = append(failures, AssertionError{Assertion: "stream/event", Error: err.Error()}) + continue + } + eventLines = append(eventLines, "event: "+eventType, "data: "+string(data)) + } + } + streamResult := stream.Result() + if streamResult.Model != bridgeMatrixModel { + failures = append(failures, AssertionError{Assertion: "stream/model_fact", Error: fmt.Sprintf("got %q, want %q", streamResult.Model, bridgeMatrixModel)}) + } + if streamResult.Usage == nil || !streamResult.Usage.HasUsage() { + failures = append(failures, AssertionError{Assertion: "stream/usage_fact", Error: fmt.Sprintf("missing normalized usage: %+v", streamResult.Usage)}) + } + if err := stream.Close(); err != nil { + failures = append(failures, AssertionError{Assertion: "stream/close", Error: err.Error()}) + } + var parsed *sse.ParsedResult + if route.Pair.Source == protocol.TypeOpenAIChat { + parsed = sse.AssembleOpenAIChatStream(eventLines) + } else { + parsed = sse.AssembleAnthropicStream(eventLines) + } + raw := []byte(strings.Join(eventLines, "\n") + "\n") + result := roundTripFromBridgeParsed(parsed, route, scenario, true, raw, eventLines) + if result.Usage == nil && streamResult.Usage != nil { + result.Usage = &TokenUsage{InputTokens: streamResult.Usage.InputTokens, OutputTokens: streamResult.Usage.OutputTokens} + } + return result, failures +} + +func normalizeBridgeMatrixChatEvent(value any) ([]byte, error) { + switch event := value.(type) { + case openai.ChatCompletionChunk: + return json.Marshal(event) + case *openai.ChatCompletionChunk: + if event == nil { + return nil, fmt.Errorf("nil OpenAI Chat chunk") + } + return json.Marshal(event) + case wire.ChatStreamChunk: + return json.Marshal(event) + case *wire.ChatStreamChunk: + if event == nil { + return nil, fmt.Errorf("nil OpenAI Chat wire chunk") + } + return json.Marshal(event) + default: + return nil, fmt.Errorf("unsupported OpenAI Chat stream event %T", value) + } +} + +func normalizeBridgeMatrixEvent(value any) (string, []byte, error) { + var eventType string + var data any + switch event := value.(type) { + case protocolstream.AnthropicEvent: + eventType, data = event.Type, event.Data + case anthropic.MessageStreamEventUnion: + eventType, data = event.Type, event + case *anthropic.MessageStreamEventUnion: + if event == nil { + return "", nil, fmt.Errorf("nil Anthropic v1 event") + } + eventType, data = event.Type, event + case anthropic.BetaRawMessageStreamEventUnion: + eventType, data = event.Type, event + case *anthropic.BetaRawMessageStreamEventUnion: + if event == nil { + return "", nil, fmt.Errorf("nil Anthropic beta event") + } + eventType, data = event.Type, event + default: + return "", nil, fmt.Errorf("unsupported Anthropic stream event %T", value) + } + if eventType == "" { + return "", nil, fmt.Errorf("Anthropic stream event %T has empty type", value) + } + raw, err := json.Marshal(data) + if err != nil { + return "", nil, err + } + return eventType, raw, nil +} + +func roundTripFromBridgeParsed(parsed *sse.ParsedResult, route bridgeMatrixRoute, scenario string, streaming bool, raw []byte, events []string) *RoundTripResult { + result := &RoundTripResult{ + SourceProtocol: route.Pair.Source, + TargetProtocol: route.Pair.Target, + ScenarioName: route.scenarioName(scenario), + IsStreaming: streaming, + RawBody: raw, + StreamEvents: events, + } + if parsed == nil { + return result + } + result.Content = parsed.Content + result.Role = parsed.Role + result.Model = parsed.Model + result.FinishReason = parsed.FinishReason + result.ThinkingContent = parsed.ThinkingContent + for _, toolCall := range parsed.ToolCalls { + result.ToolCalls = append(result.ToolCalls, ToolCallResult{ + ID: toolCall.ID, Name: toolCall.Name, Arguments: toolCall.Arguments, + }) + } + if parsed.Usage != nil { + result.Usage = &TokenUsage{InputTokens: parsed.Usage.InputTokens, OutputTokens: parsed.Usage.OutputTokens} + } + return result +} + +func runBridgeMatrixAssertions(result *RoundTripResult, scenario string) []AssertionError { + assertions := []Assertion{AssertUsageNonZero()} + switch scenario { + case "text", "tool_result": + assertions = append(assertions, + AssertRoleEquals("assistant"), + AssertContentContains("Paris"), + AssertContentNonEmpty(), + ) + case "tool_use": + assertions = append(assertions, + AssertHasToolCalls(1), + AssertToolCallName(0, "get_weather"), + AssertToolCallArgs(0, "location", "Paris"), + ) + } + var failures []AssertionError + for _, assertion := range assertions { + if err := assertion.Check(result); err != nil { + failures = append(failures, AssertionError{ + Assertion: assertion.Name, + Error: err.Error(), + Context: truncate(string(result.RawBody), 300), + }) + } + } + if result.Model != bridgeMatrixModel { + failures = append(failures, AssertionError{Assertion: "model/source_visible", Error: fmt.Sprintf("got %q, want %q", result.Model, bridgeMatrixModel)}) + } + return failures +} + +func validateBridgeMatrixTargetCall(call stage.Call, sourceRequest any, route bridgeMatrixRoute, scenario string, streaming bool) []AssertionError { + if call.Metadata.RequestID == "" { + return []AssertionError{{Assertion: "request/metadata", Error: "request ID was not preserved"}} + } + if !route.isChain() && route.Pair.Source == route.Pair.Target { + if call.Request != sourceRequest { + return []AssertionError{{Assertion: "request/identity", Error: fmt.Sprintf("identity request changed from %T to %T", sourceRequest, call.Request)}} + } + return nil + } + if route.Pair.Target == protocol.TypeAnthropicBeta { + return validateBridgeMatrixBetaTargetCall(call, scenario) + } + return validateBridgeMatrixChatTargetCall(call, scenario, streaming) +} + +func validateBridgeMatrixChatTargetCall(call stage.Call, scenario string, streaming bool) []AssertionError { + chatRequest, ok := call.Request.(*openai.ChatCompletionNewParams) + if !ok || chatRequest == nil { + return []AssertionError{{Assertion: "request/type", Error: fmt.Sprintf("got %T, want *openai.ChatCompletionNewParams", call.Request)}} + } + var failures []AssertionError + if string(chatRequest.Model) != bridgeMatrixModel { + failures = append(failures, AssertionError{Assertion: "request/model", Error: fmt.Sprintf("got %q, want %q", chatRequest.Model, bridgeMatrixModel)}) + } + if call.State.OpenAIChat == nil { + failures = append(failures, AssertionError{Assertion: "request/state", Error: "OpenAIConfig is nil"}) + } + includeUsage := chatRequest.StreamOptions.IncludeUsage.Valid() && chatRequest.StreamOptions.IncludeUsage.Value + if includeUsage != streaming { + failures = append(failures, AssertionError{Assertion: "request/stream_usage", Error: fmt.Sprintf("include_usage=%v, want %v", includeUsage, streaming)}) + } + raw, err := json.Marshal(chatRequest) + if err != nil { + failures = append(failures, AssertionError{Assertion: "request/marshal", Error: err.Error()}) + return failures + } + text := string(raw) + if !strings.Contains(text, "Paris") { + failures = append(failures, AssertionError{Assertion: "request/content", Error: "converted request lost Paris content", Context: truncate(text, 300)}) + } + if (scenario == "tool_use" || scenario == "tool_result") && !strings.Contains(text, "get_weather") { + failures = append(failures, AssertionError{Assertion: "request/tools", Error: "converted request lost get_weather", Context: truncate(text, 300)}) + } + if scenario == "tool_result" { + if !strings.Contains(text, `"role":"tool"`) || !strings.Contains(text, `"tool_call_id":"toolu_bridge_weather"`) { + failures = append(failures, AssertionError{Assertion: "request/tool_result", Error: "converted request lost tool result linkage", Context: truncate(text, 300)}) + } + } + return failures +} + +func validateBridgeMatrixBetaTargetCall(call stage.Call, scenario string) []AssertionError { + request, ok := call.Request.(*anthropic.BetaMessageNewParams) + if !ok || request == nil { + return []AssertionError{{Assertion: "request/type", Error: fmt.Sprintf("got %T, want *anthropic.BetaMessageNewParams", call.Request)}} + } + var failures []AssertionError + if string(request.Model) != bridgeMatrixModel { + failures = append(failures, AssertionError{Assertion: "request/model", Error: fmt.Sprintf("got %q, want %q", request.Model, bridgeMatrixModel)}) + } + if call.State.OpenAIChat != nil { + failures = append(failures, AssertionError{Assertion: "request/state", Error: "OpenAIConfig leaked into Anthropic target"}) + } + raw, err := json.Marshal(request) + if err != nil { + return append(failures, AssertionError{Assertion: "request/marshal", Error: err.Error()}) + } + text := string(raw) + if !strings.Contains(text, "Paris") { + failures = append(failures, AssertionError{Assertion: "request/content", Error: "converted request lost Paris content", Context: truncate(text, 300)}) + } + if (scenario == "tool_use" || scenario == "tool_result") && !strings.Contains(text, "get_weather") { + failures = append(failures, AssertionError{Assertion: "request/tools", Error: "converted request lost get_weather", Context: truncate(text, 300)}) + } + if scenario == "tool_result" && (!strings.Contains(text, `"type":"tool_result"`) || !strings.Contains(text, `"tool_use_id":"toolu_bridge_weather"`)) { + failures = append(failures, AssertionError{Assertion: "request/tool_result", Error: "converted request lost tool result linkage", Context: truncate(text, 300)}) + } + return failures +} + +func decodeBridgeFixture[T any](value any) (T, error) { + var result T + raw, err := json.Marshal(value) + if err != nil { + return result, err + } + if err := json.Unmarshal(raw, &result); err != nil { + return result, err + } + return result, nil +} diff --git a/internal/protocoltest/bridge_matrix_test.go b/internal/protocoltest/bridge_matrix_test.go new file mode 100644 index 000000000..9dee6eb48 --- /dev/null +++ b/internal/protocoltest/bridge_matrix_test.go @@ -0,0 +1,99 @@ +package protocoltest + +import ( + "strings" + "testing" +) + +func TestDefaultBridgeMatrix(t *testing.T) { + t.Parallel() + + results := DefaultBridgeMatrix().ExecuteAll() + if len(results) != 54 { + t.Fatalf("result count = %d, want 54", len(results)) + } + chainResults := 0 + for _, result := range results { + if strings.HasPrefix(result.Name, "bridges/chain/") { + chainResults++ + } + if !result.Passed { + t.Errorf("%s failed: %+v", result.Name, result.Errors) + } + if result.Skipped { + t.Errorf("%s unexpectedly skipped: %s", result.Name, result.SkipReason) + } + if result.Response == nil { + t.Errorf("%s has nil semantic response", result.Name) + } + } + if chainResults != 12 { + t.Fatalf("chain result count = %d, want 12", chainResults) + } +} + +func TestBridgeMatrixV1ToBeta(t *testing.T) { + t.Parallel() + + results := DefaultBridgeMatrix(). + OnlySources("anthropic_v1"). + OnlyTargets("anthropic_beta"). + ExecuteAll() + if len(results) != 6 { + t.Fatalf("result count = %d, want 6", len(results)) + } + for _, result := range results { + if !result.Passed || result.Skipped { + t.Fatalf("result = %+v", result) + } + } +} + +func TestBridgeMatrixConcreteChainFiltersAndBatch(t *testing.T) { + t.Parallel() + + results := DefaultBridgeMatrix(). + OnlyScenarios("tool_result"). + OnlySources("openai_chat"). + OnlyTargets("openai_chat"). + OnlyStreaming(true). + WithBatchCount(3). + ExecuteAll() + if len(results) != 2 { + t.Fatalf("result count = %d, want identity plus concrete chain", len(results)) + } + for _, result := range results { + if !result.Passed || result.BatchCount != 3 || result.BatchPassed != 3 { + t.Fatalf("batch result = %+v", result) + } + if strings.HasPrefix(result.Name, "bridges/chain/") { + if result.Name != "bridges/chain/chat_beta_stage_chat/tool_result/openai_chat/openai_chat/stream" { + t.Fatalf("chain result name = %q", result.Name) + } + return + } + } + t.Fatal("concrete chain result not found") +} + +func TestBridgeMatrixFiltersAndBatch(t *testing.T) { + t.Parallel() + + results := DefaultBridgeMatrix(). + OnlyScenarios("tool_result"). + OnlySources("anthropic_beta"). + OnlyTargets("openai_chat"). + OnlyStreaming(true). + WithBatchCount(3). + ExecuteAll() + if len(results) != 1 { + t.Fatalf("result count = %d, want 1", len(results)) + } + result := results[0] + if !result.Passed || result.BatchCount != 3 || result.BatchPassed != 3 { + t.Fatalf("batch result = %+v", result) + } + if result.Name != "bridges/tool_result/anthropic_beta/openai_chat/stream" { + t.Fatalf("result name = %q", result.Name) + } +} diff --git a/internal/protocoltest/failover_test.go b/internal/protocoltest/failover_test.go index f5db0f0ba..65591e7c5 100644 --- a/internal/protocoltest/failover_test.go +++ b/internal/protocoltest/failover_test.go @@ -77,6 +77,78 @@ func TestFailover_Stream_PreContent_500_RetriesAndSucceeds(t *testing.T) { assert.Equal(t, int64(1), route.PrimaryCallCount.Load()) } +func TestProtocolStageFailover_PreContent_RetriesAndSucceeds(t *testing.T) { + tests := []struct { + name string + source protocol.APIType + target protocol.APIType + endpoint pt.EndpointKind + }{ + {name: "chat_to_beta", source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta, endpoint: pt.EndpointAnthropic}, + {name: "v1_to_v1", source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicV1, endpoint: pt.EndpointAnthropic}, + {name: "v1_to_chat", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat, endpoint: pt.EndpointChat}, + {name: "beta_to_beta", source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta, endpoint: pt.EndpointAnthropic}, + {name: "beta_to_chat", source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIChat, endpoint: pt.EndpointChat}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + for _, streaming := range []bool{false, true} { + name := "nonstream" + scenario := pt.TextScenario() + if streaming { + name = "stream" + scenario = pt.StreamingTextScenario() + } + t.Run(name, func(t *testing.T) { + env := pt.NewTestEnv(t, pt.NewTestEnvOptionWithProtocolStage()) + defer env.Close() + + route := env.SetupFailoverRoute(t, test.source, test.target, scenario, pt.FailMockPreContent429) + result := env.SendWithModel(t, test.source, route.ModelName, streaming) + + require.Equal(t, 200, result.HTTPStatus) + assert.Equal(t, int64(1), route.PrimaryCallCount.Load()) + assert.Equal(t, 1, env.UpstreamEndpointHits(test.endpoint)) + assert.NotEmpty(t, result.Content) + }) + } + }) + } +} + +func TestProtocolStageFailover_MidStreamDoesNotRetry(t *testing.T) { + for _, test := range []struct { + name string + source protocol.APIType + target protocol.APIType + endpoint pt.EndpointKind + }{ + {name: "v1_to_v1", source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicV1, endpoint: pt.EndpointAnthropic}, + {name: "v1_to_chat", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat, endpoint: pt.EndpointChat}, + {name: "beta_to_beta", source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta, endpoint: pt.EndpointAnthropic}, + {name: "beta_to_chat", source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIChat, endpoint: pt.EndpointChat}, + } { + t.Run(test.name, func(t *testing.T) { + env := pt.NewTestEnv(t, pt.NewTestEnvOptionWithProtocolStage()) + defer env.Close() + + route := env.SetupFailoverRoute( + t, + test.source, + test.target, + pt.StreamingTextScenario(), + pt.FailMockMidStreamCut, + ) + result := env.SendWithModel(t, test.source, route.ModelName, true) + + require.Equal(t, 200, result.HTTPStatus, "first Anthropic event committed the attempt") + assert.Equal(t, int64(1), route.PrimaryCallCount.Load()) + assert.Equal(t, 0, env.UpstreamEndpointHits(test.endpoint), "fallback must not run after a committed Anthropic event") + require.NotEmpty(t, result.StreamEvents) + }) + } +} + // TestFailover_AllTiersFail_ClientSeesLastError — both services return 429. // After the loop exhausts the candidate pool, the deferred CommitIfBuffered // flushes the last buffered error to the wire. The client must see a non-200, diff --git a/internal/protocoltest/guardrails.go b/internal/protocoltest/guardrails.go new file mode 100644 index 000000000..cd4d782eb --- /dev/null +++ b/internal/protocoltest/guardrails.go @@ -0,0 +1,24 @@ +package protocoltest + +import ( + "context" + + "github.com/tingly-dev/tingly-box/internal/guardrails" + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" +) + +// NewAllowGuardrailsRuntime returns a deterministic active runtime for matrix +// compatibility checks. It evaluates every lifecycle phase but never changes +// the semantic response expected by the shared scenarios. +func NewAllowGuardrailsRuntime() *guardrails.Guardrails { + return &guardrails.Guardrails{ + Policy: allowGuardrailsPolicy{}, + HasActivePolicies: true, + } +} + +type allowGuardrailsPolicy struct{} + +func (allowGuardrailsPolicy) Evaluate(context.Context, guardrailscore.Input) (guardrailscore.Result, error) { + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil +} diff --git a/internal/protocoltest/matrix.go b/internal/protocoltest/matrix.go index 1e4d9cec2..7a254db21 100644 --- a/internal/protocoltest/matrix.go +++ b/internal/protocoltest/matrix.go @@ -30,13 +30,16 @@ func (p ProtocolPair) String() string { // Matrix defines the set of (source, target) pairs, scenarios, and // streaming modes to validate. type Matrix struct { - Pairs []ProtocolPair - Scenarios []Scenario - Streaming []bool - RecordDir string // Optional directory for recording requests/responses - BatchCount int // Number of times to run each test - MCPEnabled bool // Enable MCP feature flag in test env - Client Client // Client driver (nil = raw HTTP default) + Pairs []ProtocolPair + Scenarios []Scenario + Streaming []bool + RecordDir string // Optional directory for recording requests/responses + BatchCount int // Number of times to run each test + MCPEnabled bool // Enable MCP feature flag in test env + ProtocolStageEnabled bool // Enable the production Protocol Stage selector + MCPStageCoverage bool // Install the owned-tool fixture and servertool provider + GuardrailsEnabled bool // Enable an allow-only Guardrails runtime in the test env + Client Client // Client driver (nil = raw HTTP default) } // DefaultPairs is the canonical list of (source → target) conversion @@ -44,12 +47,9 @@ type Matrix struct { // to this list. // // Notes: -// - target=anthropic_v1 is intentionally absent. The harness picks -// providers by APIStyle and both Anthropic types map to the same -// style, so anthropic_beta as the target already exercises both -// Anthropic V1 passthrough (when source is V1) and the Beta -// conversions (when source is non-Anthropic). See -// internal/protocol/README.md. +// - V1 identity is explicit even though V1 and Beta providers share the +// Anthropic APIStyle. Protocol Stage selection is based on concrete API +// types, so the harness must preserve that distinction. // - Anthropic↔Anthropic cross-version (v1↔beta) is rejected by the // transform layer and not represented here. // - Google targets and the google→google passthrough are not yet @@ -57,6 +57,7 @@ type Matrix struct { func DefaultPairs() []ProtocolPair { return []ProtocolPair{ // Anthropic V1 source + {protocol.TypeAnthropicV1, protocol.TypeAnthropicV1}, // V1 passthrough {protocol.TypeAnthropicV1, protocol.TypeAnthropicBeta}, // V1 passthrough (provider APIStyle=Anthropic) {protocol.TypeAnthropicV1, protocol.TypeOpenAIChat}, // V1 → OpenAI Chat {protocol.TypeAnthropicV1, protocol.TypeOpenAIResponses}, // V1 → OpenAI Responses @@ -184,6 +185,37 @@ func (m *Matrix) WithMCPEnabled() *Matrix { return out } +// WithProtocolStage returns a copy that starts the real gateway with Stage +// selection enabled. This is production-path validation, unlike BridgeMatrix. +func (m *Matrix) WithProtocolStage() *Matrix { + out := m.clone() + out.ProtocolStageEnabled = true + return out +} + +// WithMCPStageCoverage adds the stateful owned-tool scenario and its local +// servertool provider. It is intentionally opt-in because ordinary protocol +// scenarios should not silently gain an executable server-owned tool. +func (m *Matrix) WithMCPStageCoverage() *Matrix { + out := m.clone() + out.MCPStageCoverage = true + for _, scenario := range out.Scenarios { + if scenario.Name == MCPStageOwnedToolScenarioName { + return out + } + } + out.Scenarios = append(out.Scenarios, newMCPStageOwnedToolScenario()) + return out +} + +// WithGuardrails enables an active allow-only Guardrails runtime. Matrix +// scenarios retain their normal semantics while exercising feature topology. +func (m *Matrix) WithGuardrails() *Matrix { + out := m.clone() + out.GuardrailsEnabled = true + return out +} + // WithClient returns a copy of the Matrix that drives requests through the // given client driver (official SDKs, subprocess drivers) instead of the // default raw HTTP client. @@ -200,6 +232,15 @@ func (m *Matrix) testEnvOpts() []TestEnvOption { if m.MCPEnabled { opts = append(opts, NewTestEnvOptionWithMCP()) } + if m.ProtocolStageEnabled { + opts = append(opts, NewTestEnvOptionWithProtocolStage()) + } + if m.MCPStageCoverage { + opts = append(opts, NewTestEnvOptionWithServertoolProviders(newMatrixEchoServertoolProvider())) + } + if m.GuardrailsEnabled { + opts = append(opts, NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())) + } if m.Client != nil { opts = append(opts, NewTestEnvOptionWithClient(m.Client)) } @@ -211,11 +252,7 @@ func (m *Matrix) testEnvOpts() []TestEnvOption { // test artifact; remove it when the defect is fixed. All tiers derive their // skips from this map (the matrix directly, replay via KnownDefectReason), so // closing a defect is a one-line deletion. -var skipSourceScenarios = map[string]string{ - // openai_responses source: tool_call conversion from provider back to Responses format loses tool calls - "openai_responses|tool_use": "Responses API source: tool_use conversion incomplete", - "openai_responses|streaming_tool_use": "Responses API source: streaming tool_use conversion incomplete", -} +var skipSourceScenarios = map[string]string{} // KnownDefectReason reports whether a (source protocol, scenario) combination // is in the known-defect registry, and why. Consumers outside the matrix @@ -439,6 +476,28 @@ func (m *Matrix) executeOneWithEnv(env *TestEnv, s Scenario, source, target prot base := m.newBaseResult(s.Name, source, target, streaming) env.SetupRoute(source, target, s) + requestModel := env.findRouteModel(source, target, s.Name) + verifyPersistedRecord := m.MCPStageCoverage && m.RecordDir != "" && s.Name == MCPStageOwnedToolScenarioName + var existingRecordIDs map[string]struct{} + if verifyPersistedRecord { + var snapshotErr error + existingRecordIDs, snapshotErr = persistedRequestRecordIDs(m.RecordDir) + if snapshotErr != nil { + return TestResult{ + Name: m.buildTestName(s.Name, source, target, streaming), + Scenario: s.Name, + Source: source, + Target: target, + Streaming: streaming, + Passed: false, + Errors: []AssertionError{{ + Assertion: "request_record_snapshot", + Error: snapshotErr.Error(), + }}, + Duration: time.Since(start), + } + } + } result, err := env.SendAsCLI(source, target, s, streaming) if err != nil { base.Errors = []AssertionError{{ @@ -462,6 +521,13 @@ func (m *Matrix) executeOneWithEnv(env *TestEnv, s Scenario, source, target prot }) } } + if verifyPersistedRecord { + recordErrors := verifyMCPStagePersistedRecord(env, m.RecordDir, requestModel, source, target, existingRecordIDs) + if len(recordErrors) > 0 { + passed = false + errors = append(errors, recordErrors...) + } + } base.Passed = passed base.Errors = errors diff --git a/internal/protocoltest/mcp_matrix.go b/internal/protocoltest/mcp_matrix.go new file mode 100644 index 000000000..23aa3d57a --- /dev/null +++ b/internal/protocoltest/mcp_matrix.go @@ -0,0 +1,231 @@ +package protocoltest + +import ( + "context" + "net/http" + "sync" + + "github.com/tingly-dev/tingly-box/internal/protocolserver/servertool" + coretool "github.com/tingly-dev/tingly-box/internal/tool" +) + +const ( + MCPStageOwnedToolScenarioName = "mcp_owned_tool" + matrixOwnedToolName = "tingly_box_mcp__builtin__echo" +) + +type matrixEchoServertoolProvider struct{} + +func newMatrixEchoServertoolProvider() servertool.ToolProvider { + return matrixEchoServertoolProvider{} +} + +func (matrixEchoServertoolProvider) Descriptor() coretool.VirtualTool { + return coretool.VirtualTool{ + Name: "echo", + Description: "Echo a value for protocol-stage matrix validation", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "q": map[string]any{"type": "string"}, + }, + "required": []string{"q"}, + }, + Handler: func(context.Context, coretool.ToolCall) (coretool.ToolResult, error) { + return coretool.TextToolResult("echo-result"), nil + }, + } +} + +func (matrixEchoServertoolProvider) Hook() servertool.Hook { return nil } + +// newMCPStageOwnedToolScenario returns alternating first-round tool calls and +// final text responses. Matrix execution is sequential per scenario, and each +// case makes exactly two calls to its target response format. +func newMCPStageOwnedToolScenario() Scenario { + var mu sync.Mutex + nonStreamCalls := make(map[ResponseFormat]int) + streamCalls := make(map[ResponseFormat]int) + + nonStream := func(format ResponseFormat, first, final any) func() (int, []byte) { + return func() (int, []byte) { + mu.Lock() + defer mu.Unlock() + nonStreamCalls[format]++ + if nonStreamCalls[format]%2 == 1 { + return http.StatusOK, mustMarshal(first) + } + return http.StatusOK, mustMarshal(final) + } + } + stream := func(format ResponseFormat, first, final []string) func() []string { + return func() []string { + mu.Lock() + defer mu.Unlock() + streamCalls[format]++ + if streamCalls[format]%2 == 1 { + return first + } + return final + } + } + + return Scenario{ + Name: MCPStageOwnedToolScenarioName, + Description: "Stage executes a server-owned tool and returns the provider's second-round answer", + Tags: []string{"mcp", "servertool", "stage"}, + MockResponses: map[ResponseFormat]MockResponseBuilder{ + FormatAnthropic: { + NonStream: nonStream(FormatAnthropic, matrixAnthropicOwnedTool(), matrixAnthropicOwnedToolFinal()), + Stream: stream(FormatAnthropic, matrixAnthropicOwnedToolStream(), matrixAnthropicOwnedToolFinalStream()), + }, + FormatOpenAIChat: { + NonStream: nonStream(FormatOpenAIChat, matrixChatOwnedTool(), matrixChatOwnedToolFinal()), + Stream: stream(FormatOpenAIChat, matrixChatOwnedToolStream(), matrixChatOwnedToolFinalStream()), + }, + FormatOpenAIResponses: { + NonStream: nonStream(FormatOpenAIResponses, matrixResponsesOwnedTool(), matrixResponsesOwnedToolFinal()), + Stream: stream(FormatOpenAIResponses, matrixResponsesOwnedToolStream(), matrixResponsesOwnedToolFinalStream()), + }, + }, + Assertions: []Assertion{ + AssertHTTPStatus(http.StatusOK), + AssertContentEquals("owned-tool-final"), + }, + } +} + +func matrixAnthropicOwnedTool() map[string]any { + return map[string]any{ + "id": "msg-owned-tool", "type": "message", "role": "assistant", "model": "worker-model", + "content": []map[string]any{{"type": "tool_use", "id": "toolu-owned-tool", "name": matrixOwnedToolName, "input": map[string]any{"q": "x"}}}, + "stop_reason": "tool_use", + "usage": map[string]any{"input_tokens": 8, "output_tokens": 3}, + } +} + +func matrixAnthropicOwnedToolFinal() map[string]any { + return map[string]any{ + "id": "msg-owned-tool-final", "type": "message", "role": "assistant", "model": "worker-model", + "content": []map[string]any{{"type": "text", "text": "owned-tool-final"}}, + "stop_reason": "end_turn", + "usage": map[string]any{"input_tokens": 12, "output_tokens": 5}, + } +} + +func matrixAnthropicOwnedToolStream() []string { + return []string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg-owned-tool","type":"message","role":"assistant","model":"worker-model","content":[],"stop_reason":null,"usage":{"input_tokens":8,"output_tokens":0}}}`, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu-owned-tool","name":"tingly_box_mcp__builtin__echo","input":{}}}`, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":\"x\"}"}}`, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":3}}`, + `event: message_stop`, + `data: {"type":"message_stop"}`, + } +} + +func matrixAnthropicOwnedToolFinalStream() []string { + return []string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg-owned-tool-final","type":"message","role":"assistant","model":"worker-model","content":[],"stop_reason":null,"usage":{"input_tokens":12,"output_tokens":0}}}`, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"owned-tool-final"}}`, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}`, + `event: message_stop`, + `data: {"type":"message_stop"}`, + } +} + +func matrixChatOwnedTool() map[string]any { + return map[string]any{ + "id": "chatcmpl-owned-tool", "object": "chat.completion", "created": 1, "model": "worker-model", + "choices": []map[string]any{{ + "index": 0, + "message": map[string]any{"role": "assistant", "content": "", "tool_calls": []map[string]any{{ + "id": "call-owned-tool", "type": "function", "function": map[string]any{"name": matrixOwnedToolName, "arguments": `{"q":"x"}`}, + }}}, + "finish_reason": "tool_calls", + }}, + "usage": map[string]any{"prompt_tokens": 8, "completion_tokens": 3, "total_tokens": 11}, + } +} + +func matrixChatOwnedToolFinal() map[string]any { + return map[string]any{ + "id": "chatcmpl-owned-tool-final", "object": "chat.completion", "created": 2, "model": "worker-model", + "choices": []map[string]any{{ + "index": 0, "message": map[string]any{"role": "assistant", "content": "owned-tool-final"}, "finish_reason": "stop", + }}, + "usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 5, "total_tokens": 17}, + } +} + +func matrixChatOwnedToolStream() []string { + return []string{ + `data: {"id":"chatcmpl-owned-tool","object":"chat.completion.chunk","created":1,"model":"worker-model","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call-owned-tool","type":"function","function":{"name":"tingly_box_mcp__builtin__echo","arguments":"{\"q\":\"x\"}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-owned-tool","object":"chat.completion.chunk","created":1,"model":"worker-model","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, + `data: [DONE]`, + } +} + +func matrixChatOwnedToolFinalStream() []string { + return []string{ + `data: {"id":"chatcmpl-owned-tool-final","object":"chat.completion.chunk","created":2,"model":"worker-model","choices":[{"index":0,"delta":{"role":"assistant","content":"owned-tool-final"},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-owned-tool-final","object":"chat.completion.chunk","created":2,"model":"worker-model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`, + `data: [DONE]`, + } +} + +func matrixResponsesOwnedTool() map[string]any { + return map[string]any{ + "id": "resp-owned-tool", "object": "response", "created_at": 1, "model": "worker-model", "status": "completed", + "output": []map[string]any{{ + "id": "fc-owned-tool", "type": "function_call", "call_id": "call-owned-tool", "name": matrixOwnedToolName, "arguments": `{"q":"x"}`, "status": "completed", + }}, + "usage": map[string]any{"input_tokens": 8, "output_tokens": 3, "total_tokens": 11}, + } +} + +func matrixResponsesOwnedToolFinal() map[string]any { + return map[string]any{ + "id": "resp-owned-tool-final", "object": "response", "created_at": 2, "model": "worker-model", "status": "completed", + "output": []map[string]any{{ + "id": "item-owned-tool-final", "type": "message", "role": "assistant", "status": "completed", + "content": []map[string]any{{"type": "output_text", "text": "owned-tool-final", "annotations": []any{}}}, + }}, + "usage": map[string]any{"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + } +} + +func matrixResponsesOwnedToolStream() []string { + return []string{ + `data: {"type":"response.created","response":{"id":"resp-owned-tool","object":"response","created_at":1,"model":"worker-model","status":"in_progress","output":[]}}`, + `data: {"type":"response.output_item.added","response_id":"resp-owned-tool","output_index":0,"item":{"id":"fc-owned-tool","type":"function_call","call_id":"call-owned-tool","name":"tingly_box_mcp__builtin__echo","status":"in_progress"}}`, + `data: {"type":"response.function_call_arguments.delta","response_id":"resp-owned-tool","item_id":"fc-owned-tool","output_index":0,"delta":"{\"q\":\"x\"}"}`, + `data: {"type":"response.function_call_arguments.done","response_id":"resp-owned-tool","item_id":"fc-owned-tool","output_index":0,"arguments":"{\"q\":\"x\"}"}`, + `data: {"type":"response.completed","response":{"id":"resp-owned-tool","object":"response","created_at":1,"model":"worker-model","status":"completed","output":[{"id":"fc-owned-tool","type":"function_call","call_id":"call-owned-tool","name":"tingly_box_mcp__builtin__echo","arguments":"{\"q\":\"x\"}","status":"completed"}],"usage":{"input_tokens":8,"output_tokens":3,"total_tokens":11}}}`, + `data: [DONE]`, + } +} + +func matrixResponsesOwnedToolFinalStream() []string { + return []string{ + `data: {"type":"response.created","response":{"id":"resp-owned-tool-final","object":"response","created_at":2,"model":"worker-model","status":"in_progress","output":[]}}`, + `data: {"type":"response.output_item.added","response_id":"resp-owned-tool-final","output_index":0,"item":{"id":"item-owned-tool-final","type":"message","role":"assistant","status":"in_progress","content":[]}}`, + `data: {"type":"response.output_text.delta","response_id":"resp-owned-tool-final","item_id":"item-owned-tool-final","output_index":0,"content_index":0,"delta":"owned-tool-final"}`, + `data: {"type":"response.output_text.done","response_id":"resp-owned-tool-final","item_id":"item-owned-tool-final","output_index":0,"content_index":0,"text":"owned-tool-final"}`, + `data: {"type":"response.completed","response":{"id":"resp-owned-tool-final","object":"response","created_at":2,"model":"worker-model","status":"completed","output":[{"id":"item-owned-tool-final","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"owned-tool-final","annotations":[]}]}],"usage":{"input_tokens":12,"output_tokens":5,"total_tokens":17}}}`, + `data: [DONE]`, + } +} diff --git a/internal/protocoltest/mcp_recording_matrix_test.go b/internal/protocoltest/mcp_recording_matrix_test.go new file mode 100644 index 000000000..a8f55ea35 --- /dev/null +++ b/internal/protocoltest/mcp_recording_matrix_test.go @@ -0,0 +1,46 @@ +package protocoltest + +import ( + "testing" + + requestrecord "github.com/tingly-dev/tingly-box/internal/record" +) + +func TestMCPStageRecordingMatrixPersistsStableBoundaries(t *testing.T) { + recordDir := t.TempDir() + matrix := DefaultMatrix(). + WithMCPEnabled(). + WithProtocolStage(). + WithMCPStageCoverage(). + WithRecordDir(recordDir). + OnlyScenarios(MCPStageOwnedToolScenarioName) + + results := matrix.ExecuteAll() + if len(results) != 26 { + t.Fatalf("matrix results = %d, want 26", len(results)) + } + for _, result := range results { + if result.Skipped || !result.Passed { + t.Fatalf("matrix case %s passed=%v skipped=%v errors=%#v", result.Name, result.Passed, result.Skipped, result.Errors) + } + } + + records, err := readPersistedRequestRecordArtifacts(recordDir) + if err != nil { + t.Fatal(err) + } + if len(records) != len(results) { + t.Fatalf("persisted records = %d, want %d", len(records), len(results)) + } + for _, record := range records { + if record == nil || record.Outcome != requestrecord.OutcomeSucceeded { + t.Fatalf("persisted record = %#v", record) + } + if len(record.ProviderExchanges) != 2 { + t.Fatalf("record %s provider exchanges = %d, want 2", record.RequestID, len(record.ProviderExchanges)) + } + if record.FinalResponse == nil { + t.Fatalf("record %s final response is missing", record.RequestID) + } + } +} diff --git a/internal/protocoltest/mcp_recording_validation.go b/internal/protocoltest/mcp_recording_validation.go new file mode 100644 index 000000000..16ecad33b --- /dev/null +++ b/internal/protocoltest/mcp_recording_validation.go @@ -0,0 +1,139 @@ +package protocoltest + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "time" + + "github.com/tingly-dev/tingly-box/internal/protocol" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" +) + +func verifyMCPStagePersistedRecord( + env *TestEnv, + recordDir string, + requestModel string, + source protocol.APIType, + target protocol.APIType, + existingRecordIDs map[string]struct{}, +) []AssertionError { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := env.ForceFlushRecordings(ctx); err != nil { + return []AssertionError{{Assertion: "request_record_flush", Error: err.Error()}} + } + + records, err := readPersistedRequestRecordArtifacts(recordDir) + if err != nil { + return []AssertionError{{Assertion: "request_record_read", Error: err.Error()}} + } + matched := make([]*requestrecord.RequestRecord, 0, 1) + for _, record := range records { + if record != nil && bytes.Contains(record.InputRequest.Body, []byte(requestModel)) { + if _, existed := existingRecordIDs[record.RequestID]; existed { + continue + } + matched = append(matched, record) + } + } + if len(matched) != 1 { + return []AssertionError{{ + Assertion: "request_record_identity", + Error: fmt.Sprintf("new persisted records for request model %q = %d, want 1", requestModel, len(matched)), + }} + } + + record := matched[0] + contextBody, _ := json.Marshal(record) + contextText := truncate(string(contextBody), 600) + var result []AssertionError + check := func(assertion string, condition bool, format string, args ...any) { + if condition { + return + } + result = append(result, AssertionError{ + Assertion: assertion, + Error: fmt.Sprintf(format, args...), + Context: contextText, + }) + } + + check("request_record_outcome", record.Outcome == requestrecord.OutcomeSucceeded, + "request outcome = %q, want %q", record.Outcome, requestrecord.OutcomeSucceeded) + check("request_record_input_protocol", record.InputRequest.Protocol == source, + "input protocol = %q, want %q", record.InputRequest.Protocol, source) + check("request_record_input_original", bytes.Contains(record.InputRequest.Body, []byte("capital of France")), + "input request does not contain the original user prompt") + check("request_record_input_unmodified", !bytes.Contains(record.InputRequest.Body, []byte(matrixOwnedToolName)), + "input request already contains injected server tool %q", matrixOwnedToolName) + check("request_record_exchange_count", len(record.ProviderExchanges) == 2, + "provider exchange count = %d, want 2", len(record.ProviderExchanges)) + + providerProtocol := target + if target == protocol.TypeAnthropicV1 { + // V1 MCP requests are promoted into the Beta-native Tool Loop and the + // provider boundary records that concrete protocol. + providerProtocol = protocol.TypeAnthropicBeta + } + for index, exchange := range record.ProviderExchanges { + check(fmt.Sprintf("request_record_exchange_%d_sequence", index+1), exchange.Sequence == index+1, + "exchange %d sequence = %d, want %d", index+1, exchange.Sequence, index+1) + check(fmt.Sprintf("request_record_exchange_%d_attempt", index+1), exchange.Attempt == 1, + "exchange %d attempt = %d, want 1", index+1, exchange.Attempt) + check(fmt.Sprintf("request_record_exchange_%d_protocol", index+1), exchange.Protocol == providerProtocol, + "exchange %d protocol = %q, want %q", index+1, exchange.Protocol, providerProtocol) + check(fmt.Sprintf("request_record_exchange_%d_request_protocol", index+1), exchange.Request.Protocol == providerProtocol, + "exchange %d request protocol = %q, want %q", index+1, exchange.Request.Protocol, providerProtocol) + check(fmt.Sprintf("request_record_exchange_%d_outcome", index+1), exchange.Outcome == requestrecord.OutcomeSucceeded, + "exchange %d outcome = %q, want %q", index+1, exchange.Outcome, requestrecord.OutcomeSucceeded) + check(fmt.Sprintf("request_record_exchange_%d_response", index+1), exchange.Response != nil, + "exchange %d provider response is missing", index+1) + if exchange.Response != nil { + check(fmt.Sprintf("request_record_exchange_%d_response_protocol", index+1), exchange.Response.Protocol == providerProtocol, + "exchange %d response protocol = %q, want %q", index+1, exchange.Response.Protocol, providerProtocol) + } + } + + if len(record.ProviderExchanges) == 2 { + first := record.ProviderExchanges[0] + second := record.ProviderExchanges[1] + check("request_record_first_provider_tool_injected", bytes.Contains(first.Request.Body, []byte(matrixOwnedToolName)), + "first provider request does not contain injected server tool %q", matrixOwnedToolName) + if first.Response != nil { + check("request_record_first_provider_tool_call", bytes.Contains(first.Response.Body, []byte(matrixOwnedToolName)), + "first provider response does not contain the owned tool call") + } + check("request_record_second_provider_tool_result", bytes.Contains(second.Request.Body, []byte("echo-result")), + "second provider request does not contain the local tool result") + if second.Response != nil { + check("request_record_second_provider_final", bytes.Contains(second.Response.Body, []byte("owned-tool-final")), + "second provider response does not contain the final answer") + } + } + + check("request_record_final_response", record.FinalResponse != nil, + "final client response is missing") + if record.FinalResponse != nil { + check("request_record_final_protocol", record.FinalResponse.Protocol == source, + "final response protocol = %q, want %q", record.FinalResponse.Protocol, source) + check("request_record_final_content", bytes.Contains(record.FinalResponse.Body, []byte("owned-tool-final")), + "final response does not contain the client-visible answer") + } + return result +} + +func persistedRequestRecordIDs(recordDir string) (map[string]struct{}, error) { + records, err := readPersistedRequestRecordArtifacts(recordDir) + if err != nil { + return nil, err + } + ids := make(map[string]struct{}, len(records)) + for _, record := range records { + if record != nil && record.RequestID != "" { + ids[record.RequestID] = struct{}{} + } + } + return ids, nil +} diff --git a/internal/protocoltest/protocol_stage_server_test.go b/internal/protocoltest/protocol_stage_server_test.go new file mode 100644 index 000000000..029d30db3 --- /dev/null +++ b/internal/protocoltest/protocol_stage_server_test.go @@ -0,0 +1,980 @@ +package protocoltest + +import ( + "bytes" + "compress/gzip" + "context" + "encoding/json" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/tingly-dev/tingly-box/internal/guardrails" + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + "github.com/tingly-dev/tingly-box/internal/protocol" + requestrecord "github.com/tingly-dev/tingly-box/internal/record" + "github.com/tingly-dev/tingly-box/internal/typ" +) + +func TestServerProtocolStageSelection(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts []TestEnvOption + source protocol.APIType + target protocol.APIType + streaming bool + wantHeader string + wantUpstream protocol.APIType + wantResponseModel bool + }{ + {name: "default chat route legacy", source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta, wantHeader: "legacy"}, + {name: "default responses route legacy", source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIResponses, wantHeader: "legacy"}, + { + name: "default chat Guardrail route remains legacy", + opts: []TestEnvOption{NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + wantHeader: "legacy", + }, + { + name: "default v1 MCP Guardrail route remains legacy", + opts: []TestEnvOption{NewTestEnvOptionWithMCP(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeAnthropicV1, + target: protocol.TypeOpenAIChat, + wantHeader: "legacy", + }, + { + name: "default v1 Guardrail route remains legacy", + opts: []TestEnvOption{NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeAnthropicV1, + target: protocol.TypeOpenAIChat, + wantHeader: "legacy", + }, + { + name: "default responses Guardrail route remains legacy", + opts: []TestEnvOption{NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIResponses, + target: protocol.TypeOpenAIChat, + wantHeader: "legacy", + }, + {name: "stage chat nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta, wantHeader: "stage"}, + {name: "stage chat stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta, streaming: true, wantHeader: "stage"}, + {name: "stage chat to responses nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIResponses, wantHeader: "stage", wantResponseModel: true}, + {name: "stage chat to responses stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIResponses, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage chat identity nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIChat, wantHeader: "stage", wantResponseModel: true}, + {name: "stage chat identity stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIChat, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage beta native nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta, wantHeader: "stage", wantResponseModel: true}, + {name: "stage beta native stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage beta to chat nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIChat, wantHeader: "stage", wantResponseModel: true}, + {name: "stage beta to chat stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIChat, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage beta to responses nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIResponses, wantHeader: "stage", wantResponseModel: true}, + {name: "stage beta to responses stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIResponses, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage v1 native nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicV1, wantHeader: "stage", wantResponseModel: true}, + {name: "stage v1 native stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicV1, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage v1 to chat nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat, wantHeader: "stage", wantResponseModel: true}, + {name: "stage v1 to chat stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage v1 to responses nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIResponses, wantHeader: "stage", wantResponseModel: true}, + {name: "stage v1 to responses stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIResponses, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage responses native nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIResponses, wantHeader: "stage", wantResponseModel: true}, + {name: "stage responses native stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIResponses, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage responses to beta nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIResponses, target: protocol.TypeAnthropicBeta, wantHeader: "stage", wantResponseModel: true}, + {name: "stage responses to beta stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIResponses, target: protocol.TypeAnthropicBeta, streaming: true, wantHeader: "stage", wantResponseModel: true}, + {name: "stage responses to chat nonstream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIChat, wantHeader: "stage", wantResponseModel: true}, + {name: "stage responses to chat stream", opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage()}, source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIChat, streaming: true, wantHeader: "stage", wantResponseModel: true}, + { + name: "stage beta runs MCP tool loop", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithMCP()}, + source: protocol.TypeAnthropicBeta, + target: protocol.TypeAnthropicBeta, + wantHeader: "stage", + }, + { + name: "stage v1 promotes MCP request to beta", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithMCP()}, + source: protocol.TypeAnthropicV1, + target: protocol.TypeAnthropicV1, + wantHeader: "stage", + wantUpstream: protocol.TypeAnthropicBeta, + }, + { + name: "stage v1 promotes Guardrail request to beta", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeAnthropicV1, + target: protocol.TypeAnthropicV1, + wantHeader: "stage", + wantUpstream: protocol.TypeAnthropicBeta, + }, + { + name: "stage v1 runs Guardrail through beta to chat", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeAnthropicV1, + target: protocol.TypeOpenAIChat, + wantHeader: "stage", + }, + { + name: "stage v1 runs Guardrail through beta to responses", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeAnthropicV1, + target: protocol.TypeOpenAIResponses, + wantHeader: "stage", + }, + { + name: "stage chat runs MCP tool loop", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithMCP()}, + source: protocol.TypeOpenAIChat, + target: protocol.TypeOpenAIChat, + wantHeader: "stage", + }, + { + name: "stage chat runs Guardrail through beta to chat", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIChat, + target: protocol.TypeOpenAIChat, + wantHeader: "stage", + }, + { + name: "stage chat runs Guardrail through beta", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIChat, + target: protocol.TypeAnthropicBeta, + wantHeader: "stage", + }, + { + name: "stage chat runs Guardrail through beta to responses", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIChat, + target: protocol.TypeOpenAIResponses, + wantHeader: "stage", + }, + { + name: "stage responses runs MCP tool loop", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithMCP()}, + source: protocol.TypeOpenAIResponses, + target: protocol.TypeOpenAIResponses, + wantHeader: "stage", + }, + { + name: "stage responses runs Guardrail through beta to responses", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIResponses, + target: protocol.TypeOpenAIResponses, + wantHeader: "stage", + }, + { + name: "stage responses runs Guardrail through beta", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIResponses, + target: protocol.TypeAnthropicBeta, + wantHeader: "stage", + }, + { + name: "stage responses runs Guardrail through beta to chat", + opts: []TestEnvOption{NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime())}, + source: protocol.TypeOpenAIResponses, + target: protocol.TypeOpenAIChat, + wantHeader: "stage", + }, + { + name: "stage v1 composes MCP and Guardrail through beta", + opts: []TestEnvOption{ + NewTestEnvOptionWithProtocolStage(), + NewTestEnvOptionWithMCP(), + NewTestEnvOptionWithGuardrails(NewAllowGuardrailsRuntime()), + }, + source: protocol.TypeAnthropicV1, + target: protocol.TypeOpenAIChat, + wantHeader: "stage", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + env := NewTestEnv(t, tt.opts...) + scenario := TextScenario() + env.SetupRoute(tt.source, tt.target, scenario) + model := env.findRouteModel(tt.source, tt.target, scenario.Name) + path, body := buildRequest(tt.source, model, tt.streaming) + req, err := http.NewRequest(http.MethodPost, env.GatewayURL()+path, bytes.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+env.ModelToken()) + req.Header.Set("X-Tingly-Debug-Routing", "1") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("do request: %v", err) + } + defer resp.Body.Close() + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read response: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d", resp.StatusCode) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != tt.wantHeader { + t.Fatalf("pipeline header = %q, want %q", got, tt.wantHeader) + } + wantUpstream := tt.wantUpstream + if wantUpstream == "" { + wantUpstream = tt.target + } + if got := resp.Header.Get("X-Tingly-Upstream-API"); got != string(wantUpstream) { + t.Fatalf("upstream API = %q, want %q", got, wantUpstream) + } + if tt.wantResponseModel && !strings.Contains(string(responseBody), `"model":"`+model+`"`) { + t.Fatalf("response does not expose request model %q: %s", model, responseBody) + } + }) + } +} + +func TestServerProtocolStageRecordingSelection(t *testing.T) { + t.Parallel() + + for _, tt := range []struct { + name string + source protocol.APIType + target protocol.APIType + wantHeader string + }{ + {name: "beta identity", source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta, wantHeader: "stage"}, + {name: "beta to chat", source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIChat, wantHeader: "stage"}, + {name: "beta to responses", source: protocol.TypeAnthropicBeta, target: protocol.TypeOpenAIResponses, wantHeader: "stage"}, + {name: "v1 identity", source: protocol.TypeAnthropicV1, target: protocol.TypeAnthropicV1, wantHeader: "stage"}, + {name: "v1 to chat", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat, wantHeader: "stage"}, + {name: "v1 to responses", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIResponses, wantHeader: "stage"}, + {name: "chat identity", source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIChat, wantHeader: "stage"}, + {name: "chat to beta", source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta, wantHeader: "stage"}, + {name: "chat to responses", source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIResponses, wantHeader: "stage"}, + {name: "responses identity", source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIResponses, wantHeader: "stage"}, + {name: "responses to beta", source: protocol.TypeOpenAIResponses, target: protocol.TypeAnthropicBeta, wantHeader: "stage"}, + {name: "responses to chat", source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIChat, wantHeader: "stage"}, + } { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + env := NewTestEnv(t, + NewTestEnvOptionWithProtocolStage(), + NewTestEnvOptionWithRecordDir(t.TempDir()), + ) + scenario := TextScenario() + env.SetupRoute(tt.source, tt.target, scenario) + model := env.findRouteModel(tt.source, tt.target, scenario.Name) + path, body := buildRequest(tt.source, model, false) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != tt.wantHeader { + t.Fatalf("pipeline header = %q, want %q", got, tt.wantHeader) + } + }) + } +} + +func TestServerProtocolStageRecordingPreservesOriginalAnthropicInput(t *testing.T) { + for _, source := range []protocol.APIType{protocol.TypeAnthropicV1, protocol.TypeAnthropicBeta} { + source := source + t.Run(string(source), func(t *testing.T) { + recordDir := t.TempDir() + env := NewTestEnv(t, + NewTestEnvOptionWithProtocolStage(), + NewTestEnvOptionWithRecordDir(recordDir), + ) + scenario := TextScenario() + env.SetupRoute(source, source, scenario) + model := env.findRouteModel(source, source, scenario.Name) + path, body := buildRequest(source, model, false) + + var input map[string]any + if err := json.Unmarshal(body, &input); err != nil { + t.Fatalf("decode request: %v", err) + } + input["client_extension"] = "preserve-me" + body, err := json.Marshal(input) + if err != nil { + t.Fatalf("encode request: %v", err) + } + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + env.Close() + + records := readPersistedRequestRecords(t, recordDir) + if len(records) != 1 { + t.Fatalf("RequestRecord count = %d, want 1", len(records)) + } + recordedBody := string(records[0].InputRequest.Body) + if !strings.Contains(recordedBody, `"client_extension":"preserve-me"`) { + t.Fatalf("input request lost original extension: %s", recordedBody) + } + }) + } +} + +func TestServerProtocolStageRecordingFailover(t *testing.T) { + routes := []struct { + name string + source protocol.APIType + primaryStyle protocol.APIStyle + fallbackTarget protocol.APIType + firstProtocol protocol.APIType + secondProtocol protocol.APIType + }{ + {name: "beta", source: protocol.TypeAnthropicBeta, primaryStyle: protocol.APIStyleAnthropic, fallbackTarget: protocol.TypeOpenAIChat, firstProtocol: protocol.TypeAnthropicBeta, secondProtocol: protocol.TypeOpenAIChat}, + {name: "v1", source: protocol.TypeAnthropicV1, primaryStyle: protocol.APIStyleAnthropic, fallbackTarget: protocol.TypeOpenAIChat, firstProtocol: protocol.TypeAnthropicV1, secondProtocol: protocol.TypeOpenAIChat}, + {name: "chat", source: protocol.TypeOpenAIChat, primaryStyle: protocol.APIStyleOpenAI, fallbackTarget: protocol.TypeAnthropicBeta, firstProtocol: protocol.TypeOpenAIChat, secondProtocol: protocol.TypeAnthropicBeta}, + {name: "responses", source: protocol.TypeOpenAIResponses, primaryStyle: protocol.APIStyleOpenAI, fallbackTarget: protocol.TypeAnthropicBeta, firstProtocol: protocol.TypeOpenAIChat, secondProtocol: protocol.TypeAnthropicBeta}, + } + for _, routeCase := range routes { + routeCase := routeCase + for _, streaming := range []bool{false, true} { + streaming := streaming + name := routeCase.name + "/complete" + if streaming { + name = routeCase.name + "/stream" + } + t.Run(name, func(t *testing.T) { + recordDir := t.TempDir() + env := NewTestEnv(t, + NewTestEnvOptionWithProtocolStage(), + NewTestEnvOptionWithRecordDir(recordDir), + ) + scenario := TextScenario() + if streaming { + scenario = StreamingTextScenario() + } + route := env.SetupCrossStyleFailoverRoute( + t, + routeCase.source, + routeCase.primaryStyle, + routeCase.fallbackTarget, + scenario, + FailMockPreContent500, + ) + + result := env.SendWithModel(t, routeCase.source, route.ModelName, streaming) + if result.HTTPStatus != http.StatusOK { + t.Fatalf("status = %d", result.HTTPStatus) + } + env.Close() + + records := readPersistedRequestRecords(t, recordDir) + if len(records) != 1 { + t.Fatalf("RequestRecord count = %d, want 1", len(records)) + } + record := records[0] + if record.Outcome != requestrecord.OutcomeSucceeded { + t.Fatalf("request outcome = %q, want succeeded", record.Outcome) + } + if len(record.ProviderExchanges) != 2 { + t.Fatalf("provider exchange count = %d, want 2", len(record.ProviderExchanges)) + } + first, second := record.ProviderExchanges[0], record.ProviderExchanges[1] + if first.Attempt != 1 || first.Protocol != routeCase.firstProtocol || first.Outcome != requestrecord.OutcomeFailed { + t.Fatalf("first exchange = attempt %d protocol %q outcome %q", first.Attempt, first.Protocol, first.Outcome) + } + if second.Attempt != 2 || second.Protocol != routeCase.secondProtocol || second.Outcome != requestrecord.OutcomeSucceeded { + t.Fatalf("second exchange = attempt %d protocol %q outcome %q", second.Attempt, second.Protocol, second.Outcome) + } + if record.FinalResponse == nil || record.FinalResponse.Protocol != routeCase.source { + t.Fatalf("final response = %#v, want %s", record.FinalResponse, routeCase.source) + } + }) + } + } +} + +func TestServerProtocolStageRecordingFailoverExhausted(t *testing.T) { + recordDir := t.TempDir() + env := NewTestEnv(t, + NewTestEnvOptionWithProtocolStage(), + NewTestEnvOptionWithRecordDir(recordDir), + ) + route := env.SetupBothFailingRoute(t, protocol.TypeOpenAIChat, protocol.TypeOpenAIChat, FailMockPreContent500) + + result := env.SendWithModel(t, protocol.TypeOpenAIChat, route.ModelName, false) + if result.HTTPStatus == http.StatusOK { + t.Fatal("exhausted failover unexpectedly returned 200") + } + env.Close() + + records := readPersistedRequestRecords(t, recordDir) + if len(records) != 1 { + t.Fatalf("RequestRecord count = %d, want 1", len(records)) + } + record := records[0] + if record.Outcome != requestrecord.OutcomeFailed { + t.Fatalf("request outcome = %q, want failed", record.Outcome) + } + if len(record.ProviderExchanges) != 2 { + t.Fatalf("provider exchange count = %d, want 2", len(record.ProviderExchanges)) + } + for index, exchange := range record.ProviderExchanges { + if exchange.Attempt != index+1 || exchange.Outcome != requestrecord.OutcomeFailed { + t.Fatalf("exchange %d = attempt %d outcome %q", index, exchange.Attempt, exchange.Outcome) + } + } + if record.FinalResponse != nil { + t.Fatalf("final response = %#v, want nil", record.FinalResponse) + } +} + +func readPersistedRequestRecords(t *testing.T, root string) []*requestrecord.RequestRecord { + t.Helper() + var records []*requestrecord.RequestRecord + err := filepath.WalkDir(root, func(path string, entry os.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if entry.IsDir() || !strings.HasSuffix(path, ".jsonl.gz") { + return nil + } + file, err := os.Open(path) + if err != nil { + return err + } + defer file.Close() + reader, err := gzip.NewReader(file) + if err != nil { + return err + } + defer reader.Close() + decoder := json.NewDecoder(reader) + for { + var envelope struct { + RequestRecord *requestrecord.RequestRecord `json:"request_record"` + } + if err := decoder.Decode(&envelope); err != nil { + if err == io.EOF { + break + } + return err + } + if envelope.RequestRecord != nil { + records = append(records, envelope.RequestRecord) + } + } + return nil + }) + if err != nil { + t.Fatalf("read persisted RequestRecords: %v", err) + } + return records +} + +func TestServerProtocolStageAnthropicBetaGuardrailComplete(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeAnthropicBeta, protocol.TypeOpenAIChat, protocol.TypeOpenAIResponses} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse { + return protocolStageBlockedResult("response denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := TextScenario() + env.SetupRoute(protocol.TypeAnthropicBeta, target, scenario) + model := env.findRouteModel(protocol.TypeAnthropicBeta, target, scenario.Name) + path, body := buildRequest(protocol.TypeAnthropicBeta, model, false) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("response was not blocked: %s", responseBody) + } + }) + } +} + +func TestServerProtocolStageAnthropicBetaGuardrailStream(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeOpenAIChat, protocol.TypeOpenAIResponses} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse && input.Content.Command != nil { + return protocolStageBlockedResult("command denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := StreamingToolUseScenario() + env.SetupRoute(protocol.TypeAnthropicBeta, target, scenario) + model := env.findRouteModel(protocol.TypeAnthropicBeta, target, scenario.Name) + path, body := buildRequest(protocol.TypeAnthropicBeta, model, true) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("stream was not blocked: %s", responseBody) + } + if strings.Contains(string(responseBody), `"type":"tool_use"`) { + t.Fatalf("blocked tool_use leaked to client: %s", responseBody) + } + }) + } +} + +func TestServerProtocolStageAnthropicV1GuardrailComplete(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeAnthropicV1, protocol.TypeOpenAIChat, protocol.TypeOpenAIResponses} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse { + return protocolStageBlockedResult("response denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := TextScenario() + env.SetupRoute(protocol.TypeAnthropicV1, target, scenario) + model := env.findRouteModel(protocol.TypeAnthropicV1, target, scenario.Name) + path, body := buildRequest(protocol.TypeAnthropicV1, model, false) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + wantUpstream := target + if target == protocol.TypeAnthropicV1 { + wantUpstream = protocol.TypeAnthropicBeta + } + if got := resp.Header.Get("X-Tingly-Upstream-API"); got != string(wantUpstream) { + t.Fatalf("upstream API = %q, want %q", got, wantUpstream) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("response was not blocked: %s", responseBody) + } + }) + } +} + +func TestServerProtocolStageAnthropicV1GuardrailPreservesCredentialMask(t *testing.T) { + t.Parallel() + + const ( + secret = "sk-stage-secret" + alias = "TINGLY_CRED_TOKEN_STAGE_TEST" + ) + runtime := newProtocolStageGuardrails(func(_ context.Context, _ guardrailscore.Input) (guardrailscore.Result, error) { + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + credentialCache := guardrails.BuildCredentialCache( + []guardrailscore.ProtectedCredential{{ + ID: "stage-credential", Name: "Stage credential", Type: guardrailscore.ProtectedCredentialTypeToken, + Secret: secret, AliasToken: alias, Enabled: true, + }}, + []string{string(typ.ScenarioAnthropic)}, + ) + runtime.SetCredentialCache(credentialCache) + + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + // Server initialization refreshes the runtime cache from the test database; + // install this request fixture after boot so the live handler sees it. + env.rootServer.CurrentGuardrailsRuntime().SetCredentialCache(credentialCache) + scenario := Scenario{ + Name: "credential_restore", + MockResponses: map[ResponseFormat]MockResponseBuilder{ + FormatAnthropic: { + NonStream: func() (int, []byte) { + return http.StatusOK, mustMarshal(map[string]any{ + "id": "msg-credential", "type": "message", "role": "assistant", + "content": []map[string]any{{"type": "text", "text": alias}}, + "model": "provider-model", "stop_reason": "end_turn", "stop_sequence": nil, + "usage": map[string]any{"input_tokens": 1, "output_tokens": 1}, + }) + }, + }, + }, + } + env.SetupRoute(protocol.TypeAnthropicV1, protocol.TypeAnthropicV1, scenario) + model := env.findRouteModel(protocol.TypeAnthropicV1, protocol.TypeAnthropicV1, scenario.Name) + path, body := buildRequest(protocol.TypeAnthropicV1, model, false) + var request map[string]any + if err := json.Unmarshal(body, &request); err != nil { + t.Fatalf("decode request: %v", err) + } + request["messages"] = []map[string]any{{"role": "user", "content": "use " + secret}} + body = mustMarshal(request) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if !strings.Contains(string(responseBody), secret) { + t.Fatalf("credential alias was not restored: %s", responseBody) + } + if strings.Contains(string(responseBody), alias) { + t.Fatalf("credential alias leaked to client: %s", responseBody) + } +} + +func TestServerProtocolStageAnthropicV1GuardrailStream(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeAnthropicV1, protocol.TypeOpenAIChat, protocol.TypeOpenAIResponses} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse && input.Content.Command != nil { + return protocolStageBlockedResult("command denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := StreamingToolUseScenario() + env.SetupRoute(protocol.TypeAnthropicV1, target, scenario) + model := env.findRouteModel(protocol.TypeAnthropicV1, target, scenario.Name) + path, body := buildRequest(protocol.TypeAnthropicV1, model, true) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + wantUpstream := target + if target == protocol.TypeAnthropicV1 { + wantUpstream = protocol.TypeAnthropicBeta + } + if got := resp.Header.Get("X-Tingly-Upstream-API"); got != string(wantUpstream) { + t.Fatalf("upstream API = %q, want %q", got, wantUpstream) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("stream was not blocked: %s", responseBody) + } + if strings.Contains(string(responseBody), `"type":"tool_use"`) { + t.Fatalf("blocked tool_use leaked to client: %s", responseBody) + } + }) + } +} + +func TestServerProtocolStageOpenAIResponsesGuardrailComplete(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeOpenAIResponses, protocol.TypeAnthropicBeta, protocol.TypeOpenAIChat} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse { + return protocolStageBlockedResult("response denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := TextScenario() + env.SetupRoute(protocol.TypeOpenAIResponses, target, scenario) + model := env.findRouteModel(protocol.TypeOpenAIResponses, target, scenario.Name) + path, body := buildRequest(protocol.TypeOpenAIResponses, model, false) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + if got := resp.Header.Get("X-Tingly-Upstream-API"); got != string(target) { + t.Fatalf("upstream API = %q, want %q", got, target) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("response was not blocked: %s", responseBody) + } + }) + } +} + +func TestServerProtocolStageOpenAIResponsesGuardrailStream(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeOpenAIResponses, protocol.TypeAnthropicBeta, protocol.TypeOpenAIChat} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse && input.Content.Command != nil { + return protocolStageBlockedResult("command denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := StreamingToolUseScenario() + env.SetupRoute(protocol.TypeOpenAIResponses, target, scenario) + model := env.findRouteModel(protocol.TypeOpenAIResponses, target, scenario.Name) + path, body := buildRequest(protocol.TypeOpenAIResponses, model, true) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + if got := resp.Header.Get("X-Tingly-Upstream-API"); got != string(target) { + t.Fatalf("upstream API = %q, want %q", got, target) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("stream was not blocked: %s", responseBody) + } + if strings.Contains(string(responseBody), `"type":"function_call"`) { + t.Fatalf("blocked function call leaked to client: %s", responseBody) + } + }) + } +} + +func TestServerProtocolStageOpenAIChatGuardrailComplete(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeOpenAIChat, protocol.TypeAnthropicBeta, protocol.TypeOpenAIResponses} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse { + return protocolStageBlockedResult("response denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := TextScenario() + env.SetupRoute(protocol.TypeOpenAIChat, target, scenario) + model := env.findRouteModel(protocol.TypeOpenAIChat, target, scenario.Name) + path, body := buildRequest(protocol.TypeOpenAIChat, model, false) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + if got := resp.Header.Get("X-Tingly-Upstream-API"); got != string(target) { + t.Fatalf("upstream API = %q, want %q", got, target) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("response was not blocked: %s", responseBody) + } + }) + } +} + +func TestServerProtocolStageOpenAIChatGuardrailStream(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeOpenAIChat, protocol.TypeAnthropicBeta, protocol.TypeOpenAIResponses} { + target := target + t.Run(string(target), func(t *testing.T) { + t.Parallel() + runtime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse && input.Content.Command != nil { + return protocolStageBlockedResult("command denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage(), NewTestEnvOptionWithGuardrails(runtime)) + scenario := StreamingToolUseScenario() + env.SetupRoute(protocol.TypeOpenAIChat, target, scenario) + model := env.findRouteModel(protocol.TypeOpenAIChat, target, scenario.Name) + path, body := buildRequest(protocol.TypeOpenAIChat, model, true) + + resp, responseBody := sendProtocolStageProbe(t, env, path, body) + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d: %s", resp.StatusCode, responseBody) + } + if got := resp.Header.Get("X-Tingly-Protocol-Pipeline"); got != "stage" { + t.Fatalf("pipeline header = %q, want stage", got) + } + if got := resp.Header.Get("X-Tingly-Upstream-API"); got != string(target) { + t.Fatalf("upstream API = %q, want %q", got, target) + } + if !strings.Contains(string(responseBody), "Blocked by guardrails") { + t.Fatalf("stream was not blocked: %s", responseBody) + } + if strings.Contains(string(responseBody), `"tool_calls"`) { + t.Fatalf("blocked tool call leaked to client: %s", responseBody) + } + }) + } +} + +type protocolStageGuardrailPolicy func(context.Context, guardrailscore.Input) (guardrailscore.Result, error) + +func (policy protocolStageGuardrailPolicy) Evaluate(ctx context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + return policy(ctx, input) +} + +func newProtocolStageGuardrails(policy protocolStageGuardrailPolicy) *guardrails.Guardrails { + return &guardrails.Guardrails{Policy: policy, HasActivePolicies: true} +} + +func protocolStageBlockedResult(reason string) guardrailscore.Result { + return guardrailscore.Result{ + Verdict: guardrailscore.VerdictBlock, + Reasons: []guardrailscore.PolicyResult{{ + PolicyID: "protocol-stage-test", + Verdict: guardrailscore.VerdictBlock, + Reason: reason, + }}, + } +} + +func sendProtocolStageProbe(t *testing.T, env *TestEnv, path string, body []byte) (*http.Response, []byte) { + t.Helper() + req, err := http.NewRequest(http.MethodPost, env.GatewayURL()+path, bytes.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+env.ModelToken()) + req.Header.Set("X-Tingly-Debug-Routing", "1") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("do request: %v", err) + } + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + resp.Body.Close() + t.Fatalf("read response: %v", err) + } + return resp, responseBody +} + +func TestServerProtocolStagePreservesSkipUsageFlag(t *testing.T) { + t.Parallel() + + for _, target := range []protocol.APIType{protocol.TypeAnthropicBeta, protocol.TypeOpenAIChat} { + target := target + for _, streaming := range []bool{false, true} { + streaming := streaming + name := string(target) + "/nonstream" + if streaming { + name = string(target) + "/stream" + } + t.Run(name, func(t *testing.T) { + t.Parallel() + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage()) + scenario := TextScenario() + model := env.SetupRouteWithFlags( + protocol.TypeOpenAIChat, + target, + scenario, + typ.RuleFlags{SkipUsage: true}, + ) + path, body := buildRequest(protocol.TypeOpenAIChat, model, streaming) + result, err := env.dispatch( + protocol.TypeOpenAIChat, + target, + scenario.Name, + path, + body, + map[string]string{"X-Tingly-Debug-Routing": "1"}, + streaming, + ) + if err != nil { + t.Fatalf("dispatch: %v", err) + } + if result.HTTPStatus != http.StatusOK { + t.Fatalf("status = %d", result.HTTPStatus) + } + if strings.Contains(string(result.RawBody), `"usage"`) { + t.Fatalf("response contains usage: %s", result.RawBody) + } + }) + } + } +} + +func TestServerProtocolStageAnthropicBetaPreservesRuleTransforms(t *testing.T) { + t.Parallel() + + t.Run("thinking effort", func(t *testing.T) { + t.Parallel() + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage()) + scenario := flagScenario() + model := env.SetupRouteWithFlags( + protocol.TypeAnthropicBeta, + protocol.TypeAnthropicBeta, + scenario, + typ.RuleFlags{ThinkingEffort: typ.ThinkingEffortHigh}, + ) + sendFlag(t, env, protocol.TypeAnthropicBeta, protocol.TypeAnthropicBeta, model, false, nil, nil) + thinking, ok := env.virtual.LastRequest(EndpointAnthropic).JSON()["thinking"].(map[string]any) + if !ok || thinking["type"] != "enabled" { + t.Fatalf("upstream thinking = %#v, want enabled", thinking) + } + }) + + t.Run("clean header", func(t *testing.T) { + t.Parallel() + env := NewTestEnv(t, NewTestEnvOptionWithProtocolStage()) + scenario := flagScenario() + model := env.SetupRouteWithFlags( + protocol.TypeAnthropicBeta, + protocol.TypeAnthropicBeta, + scenario, + typ.RuleFlags{CleanHeader: true}, + ) + sendFlag(t, env, protocol.TypeAnthropicBeta, protocol.TypeAnthropicBeta, model, false, func(request map[string]any) { + request["system"] = []map[string]any{ + {"type": "text", "text": "x-anthropic-billing-header: secret-token"}, + {"type": "text", "text": "You are a helpful assistant."}, + } + }, nil) + upstream := string(env.virtual.LastRequest(EndpointAnthropic).Body) + if strings.Contains(upstream, "x-anthropic-billing-header") { + t.Fatalf("billing header survived Stage transform: %s", truncate(upstream, 300)) + } + if !strings.Contains(upstream, "You are a helpful assistant.") { + t.Fatalf("normal system content was removed: %s", truncate(upstream, 300)) + } + }) +} diff --git a/internal/protocoltest/protocol_stage_tool_loop_test.go b/internal/protocoltest/protocol_stage_tool_loop_test.go new file mode 100644 index 000000000..ad3812565 --- /dev/null +++ b/internal/protocoltest/protocol_stage_tool_loop_test.go @@ -0,0 +1,408 @@ +package protocoltest + +import ( + "context" + "net/http" + "strings" + "sync" + "testing" + + guardrailscore "github.com/tingly-dev/tingly-box/internal/guardrails/core" + "github.com/tingly-dev/tingly-box/internal/protocol" + "github.com/tingly-dev/tingly-box/internal/protocolserver/servertool" + coretool "github.com/tingly-dev/tingly-box/internal/tool" +) + +const ownedToolName = "tingly_box_mcp__builtin__echo" + +type echoServertoolProvider struct { + mu sync.Mutex + calls int + arguments map[string]any +} + +func (p *echoServertoolProvider) Descriptor() coretool.VirtualTool { + return coretool.VirtualTool{ + Name: "echo", + Description: "Echo a value for protocol-stage integration tests", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "q": map[string]any{"type": "string"}, + }, + "required": []string{"q"}, + }, + Handler: func(_ context.Context, call coretool.ToolCall) (coretool.ToolResult, error) { + p.mu.Lock() + defer p.mu.Unlock() + p.calls++ + p.arguments = call.Arguments + return coretool.TextToolResult("echo-result"), nil + }, + } +} + +func (p *echoServertoolProvider) Hook() servertool.Hook { return nil } + +func (p *echoServertoolProvider) snapshot() (int, map[string]any) { + p.mu.Lock() + defer p.mu.Unlock() + return p.calls, p.arguments +} + +func TestServerProtocolStageOwnedToolLoopHTTP(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + source protocol.APIType + target protocol.APIType + }{ + {name: "beta_native", source: protocol.TypeAnthropicBeta, target: protocol.TypeAnthropicBeta}, + {name: "v1_promoted_to_beta_then_chat", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat}, + {name: "chat_through_beta_to_anthropic", source: protocol.TypeOpenAIChat, target: protocol.TypeAnthropicBeta}, + {name: "responses_through_beta_to_chat", source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIChat}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + clients := []Client{NewHTTPClient(), NewGoSDKClient()} + for _, client := range clients { + t.Run(client.Name(), func(t *testing.T) { + t.Parallel() + for _, streaming := range []bool{false, true} { + mode := "complete" + if streaming { + mode = "stream" + } + t.Run(mode, func(t *testing.T) { + t.Parallel() + provider := &echoServertoolProvider{} + env := NewTestEnv(t, + NewTestEnvOptionWithProtocolStage(), + NewTestEnvOptionWithMCP(), + NewTestEnvOptionWithServertoolProviders(provider), + NewTestEnvOptionWithClient(client), + ) + scenario := ownedToolLoopScenario() + env.SetupRoute(tt.source, tt.target, scenario) + + result := env.SendAs(t, tt.source, tt.target, scenario, streaming) + if result.HTTPStatus != http.StatusOK { + t.Fatalf("status = %d, body = %s", result.HTTPStatus, result.RawBody) + } + if result.Content != "owned-tool-final" { + t.Fatalf("content = %q, want owned-tool-final; body = %s", result.Content, result.RawBody) + } + if len(result.ToolCalls) != 0 { + t.Fatalf("final response leaked %d tool calls", len(result.ToolCalls)) + } + if env.VirtualCallCount() != 2 { + t.Fatalf("provider calls = %d, want 2", env.VirtualCallCount()) + } + calls, arguments := provider.snapshot() + if calls != 1 { + t.Fatalf("local tool executions = %d, want 1", calls) + } + if arguments["q"] != "x" { + t.Fatalf("local tool arguments = %#v, want q=x", arguments) + } + }) + } + }) + } + }) + } +} + +func TestServerProtocolStageMCPGuardrailComposition(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + source protocol.APIType + target protocol.APIType + }{ + {name: "v1_through_beta_to_chat", source: protocol.TypeAnthropicV1, target: protocol.TypeOpenAIChat}, + {name: "chat_through_beta_to_chat", source: protocol.TypeOpenAIChat, target: protocol.TypeOpenAIChat}, + {name: "responses_through_beta_to_chat", source: protocol.TypeOpenAIResponses, target: protocol.TypeOpenAIChat}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + for _, streaming := range []bool{false, true} { + mode := "complete" + if streaming { + mode = "stream" + } + t.Run(mode, func(t *testing.T) { + t.Parallel() + provider := &echoServertoolProvider{} + guardrailRuntime := newProtocolStageGuardrails(func(_ context.Context, input guardrailscore.Input) (guardrailscore.Result, error) { + if input.Direction == guardrailscore.DirectionResponse && input.Content.Command != nil { + return protocolStageBlockedResult("external tool denied"), nil + } + return guardrailscore.Result{Verdict: guardrailscore.VerdictAllow}, nil + }) + env := NewTestEnv(t, + NewTestEnvOptionWithProtocolStage(), + NewTestEnvOptionWithMCP(), + NewTestEnvOptionWithGuardrails(guardrailRuntime), + NewTestEnvOptionWithServertoolProviders(provider), + ) + scenario := ownedThenExternalToolScenario() + env.SetupRoute(tt.source, tt.target, scenario) + + result := env.SendAs(t, tt.source, tt.target, scenario, streaming) + if result.HTTPStatus != http.StatusOK { + t.Fatalf("status = %d, body = %s", result.HTTPStatus, result.RawBody) + } + if !strings.Contains(result.Content, "Blocked by guardrails") { + t.Fatalf("final response was not blocked: content=%q body=%s", result.Content, result.RawBody) + } + if len(result.ToolCalls) != 0 { + t.Fatalf("blocked external tool leaked in final response: %#v", result.ToolCalls) + } + if env.VirtualCallCount() != 2 { + t.Fatalf("provider calls = %d, want 2", env.VirtualCallCount()) + } + calls, _ := provider.snapshot() + if calls != 1 { + t.Fatalf("local tool executions = %d, want 1", calls) + } + }) + } + }) + } +} + +func ownedThenExternalToolScenario() Scenario { + var mu sync.Mutex + nonStreamCalls := map[ResponseFormat]int{} + streamCalls := map[ResponseFormat]int{} + + nextNonStream := func(format ResponseFormat, first, final any) func() (int, []byte) { + return func() (int, []byte) { + mu.Lock() + defer mu.Unlock() + nonStreamCalls[format]++ + if nonStreamCalls[format] == 1 { + return http.StatusOK, mustMarshal(first) + } + return http.StatusOK, mustMarshal(final) + } + } + nextStream := func(format ResponseFormat, first, final []string) func() []string { + return func() []string { + mu.Lock() + defer mu.Unlock() + streamCalls[format]++ + if streamCalls[format] == 1 { + return first + } + return final + } + } + + return Scenario{ + Name: "mcp_owned_then_external_tool", + Description: "Provider executes a server-owned tool, then requests a client-owned tool", + Tags: []string{"mcp", "servertool", "guardrail", "stage"}, + MockResponses: map[ResponseFormat]MockResponseBuilder{ + FormatAnthropic: { + NonStream: nextNonStream(FormatAnthropic, anthropicOwnedToolResponse(), anthropicExternalToolResponse()), + Stream: nextStream(FormatAnthropic, anthropicOwnedToolStream(), anthropicExternalToolStream()), + }, + FormatOpenAIChat: { + NonStream: nextNonStream(FormatOpenAIChat, openAIOwnedToolResponse(), openAIExternalToolResponse()), + Stream: nextStream(FormatOpenAIChat, openAIOwnedToolStream(), openAIExternalToolStream()), + }, + }, + } +} + +func ownedToolLoopScenario() Scenario { + var mu sync.Mutex + nonStreamCalls := map[ResponseFormat]int{} + streamCalls := map[ResponseFormat]int{} + + nextNonStream := func(format ResponseFormat, first, final any) func() (int, []byte) { + return func() (int, []byte) { + mu.Lock() + defer mu.Unlock() + nonStreamCalls[format]++ + if nonStreamCalls[format] == 1 { + return http.StatusOK, mustMarshal(first) + } + return http.StatusOK, mustMarshal(final) + } + } + nextStream := func(format ResponseFormat, first, final []string) func() []string { + return func() []string { + mu.Lock() + defer mu.Unlock() + streamCalls[format]++ + if streamCalls[format] == 1 { + return first + } + return final + } + } + + return Scenario{ + Name: "mcp_owned_tool", + Description: "Provider requests a server-owned tool, then returns a final answer", + Tags: []string{"mcp", "servertool", "stage"}, + MockResponses: map[ResponseFormat]MockResponseBuilder{ + FormatAnthropic: { + NonStream: nextNonStream(FormatAnthropic, anthropicOwnedToolResponse(), anthropicOwnedToolFinalResponse()), + Stream: nextStream(FormatAnthropic, anthropicOwnedToolStream(), anthropicOwnedToolFinalStream()), + }, + FormatOpenAIChat: { + NonStream: nextNonStream(FormatOpenAIChat, openAIOwnedToolResponse(), openAIOwnedToolFinalResponse()), + Stream: nextStream(FormatOpenAIChat, openAIOwnedToolStream(), openAIOwnedToolFinalStream()), + }, + }, + } +} + +func anthropicExternalToolResponse() map[string]any { + return map[string]any{ + "id": "msg-external-tool", "type": "message", "role": "assistant", "model": "worker-model", + "content": []map[string]any{{"type": "tool_use", "id": "toolu-external-tool", "name": "client_tool", "input": map[string]any{"q": "outside"}}}, + "stop_reason": "tool_use", + "usage": map[string]any{"input_tokens": 12, "output_tokens": 5}, + } +} + +func anthropicExternalToolStream() []string { + return []string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg-external-tool","type":"message","role":"assistant","model":"worker-model","content":[],"stop_reason":null,"usage":{"input_tokens":12,"output_tokens":0}}}`, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu-external-tool","name":"client_tool","input":{}}}`, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":\"outside\"}"}}`, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":5}}`, + `event: message_stop`, + `data: {"type":"message_stop"}`, + } +} + +func anthropicOwnedToolResponse() map[string]any { + return map[string]any{ + "id": "msg-owned-tool", "type": "message", "role": "assistant", "model": "worker-model", + "content": []map[string]any{{"type": "tool_use", "id": "toolu-owned-tool", "name": ownedToolName, "input": map[string]any{"q": "x"}}}, + "stop_reason": "tool_use", + "usage": map[string]any{"input_tokens": 8, "output_tokens": 3}, + } +} + +func anthropicOwnedToolFinalResponse() map[string]any { + return map[string]any{ + "id": "msg-owned-tool-final", "type": "message", "role": "assistant", "model": "worker-model", + "content": []map[string]any{{"type": "text", "text": "owned-tool-final"}}, + "stop_reason": "end_turn", + "usage": map[string]any{"input_tokens": 12, "output_tokens": 5}, + } +} + +func anthropicOwnedToolStream() []string { + return []string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg-owned-tool","type":"message","role":"assistant","model":"worker-model","content":[],"stop_reason":null,"usage":{"input_tokens":8,"output_tokens":0}}}`, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu-owned-tool","name":"tingly_box_mcp__builtin__echo","input":{}}}`, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":\"x\"}"}}`, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":3}}`, + `event: message_stop`, + `data: {"type":"message_stop"}`, + } +} + +func anthropicOwnedToolFinalStream() []string { + return []string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg-owned-tool-final","type":"message","role":"assistant","model":"worker-model","content":[],"stop_reason":null,"usage":{"input_tokens":12,"output_tokens":0}}}`, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"owned-tool-final"}}`, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}`, + `event: message_stop`, + `data: {"type":"message_stop"}`, + } +} + +func openAIOwnedToolResponse() map[string]any { + return map[string]any{ + "id": "chatcmpl-owned-tool", "object": "chat.completion", "created": 1, "model": "worker-model", + "choices": []map[string]any{{ + "index": 0, + "message": map[string]any{"role": "assistant", "content": "", "tool_calls": []map[string]any{{ + "id": "call-owned-tool", "type": "function", "function": map[string]any{"name": ownedToolName, "arguments": `{"q":"x"}`}, + }}}, + "finish_reason": "tool_calls", + }}, + "usage": map[string]any{"prompt_tokens": 8, "completion_tokens": 3, "total_tokens": 11}, + } +} + +func openAIOwnedToolFinalResponse() map[string]any { + return map[string]any{ + "id": "chatcmpl-owned-tool-final", "object": "chat.completion", "created": 2, "model": "worker-model", + "choices": []map[string]any{{ + "index": 0, "message": map[string]any{"role": "assistant", "content": "owned-tool-final"}, "finish_reason": "stop", + }}, + "usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 5, "total_tokens": 17}, + } +} + +func openAIExternalToolResponse() map[string]any { + return map[string]any{ + "id": "chatcmpl-external-tool", "object": "chat.completion", "created": 2, "model": "worker-model", + "choices": []map[string]any{{ + "index": 0, + "message": map[string]any{"role": "assistant", "content": "", "tool_calls": []map[string]any{{ + "id": "call-external-tool", "type": "function", "function": map[string]any{"name": "client_tool", "arguments": `{"q":"outside"}`}, + }}}, + "finish_reason": "tool_calls", + }}, + "usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 5, "total_tokens": 17}, + } +} + +func openAIOwnedToolStream() []string { + return []string{ + `data: {"id":"chatcmpl-owned-tool","object":"chat.completion.chunk","created":1,"model":"worker-model","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call-owned-tool","type":"function","function":{"name":"tingly_box_mcp__builtin__echo","arguments":"{\"q\":\"x\"}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-owned-tool","object":"chat.completion.chunk","created":1,"model":"worker-model","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, + `data: [DONE]`, + } +} + +func openAIOwnedToolFinalStream() []string { + return []string{ + `data: {"id":"chatcmpl-owned-tool-final","object":"chat.completion.chunk","created":2,"model":"worker-model","choices":[{"index":0,"delta":{"role":"assistant","content":"owned-tool-final"},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-owned-tool-final","object":"chat.completion.chunk","created":2,"model":"worker-model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`, + `data: [DONE]`, + } +} + +func openAIExternalToolStream() []string { + return []string{ + `data: {"id":"chatcmpl-external-tool","object":"chat.completion.chunk","created":2,"model":"worker-model","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call-external-tool","type":"function","function":{"name":"client_tool","arguments":"{\"q\":\"outside\"}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-external-tool","object":"chat.completion.chunk","created":2,"model":"worker-model","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, + `data: [DONE]`, + } +} diff --git a/internal/protocoltest/record_artifacts.go b/internal/protocoltest/record_artifacts.go new file mode 100644 index 000000000..8fff444a4 --- /dev/null +++ b/internal/protocoltest/record_artifacts.go @@ -0,0 +1,69 @@ +package protocoltest + +import ( + "compress/gzip" + "encoding/json" + "io" + "os" + "path/filepath" + "strings" + + requestrecord "github.com/tingly-dev/tingly-box/internal/record" +) + +// readPersistedRequestRecordArtifacts loads the additive request_record +// envelopes emitted by the real recording sink. Keeping this reader in the +// harness package lets CLI validation inspect the exact persisted artifact +// instead of relying on an in-memory recorder snapshot. +func readPersistedRequestRecordArtifacts(root string) ([]*requestrecord.RequestRecord, error) { + var records []*requestrecord.RequestRecord + err := filepath.WalkDir(root, func(path string, entry os.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if entry.IsDir() || !strings.HasSuffix(path, ".jsonl.gz") { + return nil + } + loaded, err := readRequestRecordArtifactFile(path) + if err != nil { + return err + } + records = append(records, loaded...) + return nil + }) + if os.IsNotExist(err) { + return nil, nil + } + return records, err +} + +func readRequestRecordArtifactFile(path string) ([]*requestrecord.RequestRecord, error) { + file, err := os.Open(path) + if err != nil { + return nil, err + } + defer file.Close() + + reader, err := gzip.NewReader(file) + if err != nil { + return nil, err + } + defer reader.Close() + + var records []*requestrecord.RequestRecord + decoder := json.NewDecoder(reader) + for { + var envelope struct { + RequestRecord *requestrecord.RequestRecord `json:"request_record"` + } + if err := decoder.Decode(&envelope); err != nil { + if err == io.EOF { + return records, nil + } + return nil, err + } + if envelope.RequestRecord != nil { + records = append(records, envelope.RequestRecord) + } + } +} diff --git a/internal/protocoltest/testenv.go b/internal/protocoltest/testenv.go index c52749f40..47072e417 100644 --- a/internal/protocoltest/testenv.go +++ b/internal/protocoltest/testenv.go @@ -2,6 +2,7 @@ package protocoltest import ( "bytes" + "context" "encoding/json" "fmt" "io" @@ -10,14 +11,17 @@ import ( "os" "sync" "testing" + "time" "github.com/tingly-dev/tingly-box/ai" "github.com/tingly-dev/tingly-box/internal/config" "github.com/tingly-dev/tingly-box/internal/constant" + "github.com/tingly-dev/tingly-box/internal/guardrails" "github.com/tingly-dev/tingly-box/internal/protocol" "github.com/tingly-dev/tingly-box/internal/protocol/sse" "github.com/tingly-dev/tingly-box/internal/server" serverconfig "github.com/tingly-dev/tingly-box/internal/server/config" + "github.com/tingly-dev/tingly-box/internal/protocolserver/servertool" "github.com/tingly-dev/tingly-box/internal/typ" ) @@ -37,6 +41,7 @@ import ( // native APIs (/v1/chat/completions, /v1/messages, etc.). type TestEnv struct { appConfig *config.AppConfig + rootServer *server.Server gatewayServer *httptest.Server // real HTTP server; every request traverses it virtual *VirtualServer modelToken string @@ -52,9 +57,12 @@ type TestEnv struct { type TestEnvOption func(*testEnvConfig) type testEnvConfig struct { - recordDir string - mcpEnabled bool - client Client + recordDir string + mcpEnabled bool + protocolStageEnabled bool + guardrailsRuntime *guardrails.Guardrails + client Client + servertoolProviders []servertool.ToolProvider } // NewTestEnvOptionWithRecordDir creates an option to set the record directory. @@ -72,6 +80,21 @@ func NewTestEnvOptionWithMCP() TestEnvOption { } } +// NewTestEnvOptionWithProtocolStage enables the real server Stage selector. +func NewTestEnvOptionWithProtocolStage() TestEnvOption { + return func(cfg *testEnvConfig) { + cfg.protocolStageEnabled = true + } +} + +// NewTestEnvOptionWithGuardrails enables Guardrails for the global scenario +// and injects the supplied runtime into the real gateway server. +func NewTestEnvOptionWithGuardrails(runtime *guardrails.Guardrails) TestEnvOption { + return func(cfg *testEnvConfig) { + cfg.guardrailsRuntime = runtime + } +} + // NewTestEnvOptionWithClient creates an option to set the client driver used // for sending requests through the gateway. Defaults to the raw HTTP client. func NewTestEnvOptionWithClient(c Client) TestEnvOption { @@ -127,6 +150,14 @@ func preseedEnterpriseContextKeys(configDir string) error { return serverconfig.WriteEnterpriseContextKeys(privatePath, publicPath, keys.private, keys.public) } +// NewTestEnvOptionWithServertoolProviders injects server-owned tools through +// the same Server option and startup path used by production embedders. +func NewTestEnvOptionWithServertoolProviders(providers ...servertool.ToolProvider) TestEnvOption { + return func(cfg *testEnvConfig) { + cfg.servertoolProviders = append(cfg.servertoolProviders, providers...) + } +} + // gatewayCore is the shared skeleton every single-process harness env builds // on: a temp config dir, an app config, a real gateway httptest.Server, and a // VirtualServer mock provider. TestEnv (matrix/flags) and AgentTestEnv @@ -134,6 +165,7 @@ func preseedEnterpriseContextKeys(configDir string) error { type gatewayCore struct { configDir string appConfig *config.AppConfig + rootServer *server.Server gateway *httptest.Server virtual *VirtualServer modelToken string @@ -166,12 +198,13 @@ func newGatewayCore(dirPattern string, configure func(*config.AppConfig), server configure(appConfig) } - gatewayServer := server.NewServer(appConfig.GetGlobalConfig(), serverOpts...) - ts := httptest.NewServer(gatewayServer.GetRouter()) + rootServer := server.NewServer(appConfig.GetGlobalConfig(), serverOpts...) + ts := httptest.NewServer(rootServer.GetRouter()) return &gatewayCore{ configDir: configDir, appConfig: appConfig, + rootServer: rootServer, gateway: ts, virtual: NewVirtualServerForCLI(), modelToken: appConfig.GetGlobalConfig().GetModelToken(), @@ -209,6 +242,11 @@ func (env *TestEnv) Close() { if env.gatewayServer != nil { env.gatewayServer.Close() } + if env.rootServer != nil { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + _ = env.rootServer.ForceFlushRecordings(ctx) + cancel() + } if env.virtual != nil { env.virtual.Close() } @@ -232,11 +270,32 @@ func NewTestEnvForCLI(opts ...TestEnvOption) (*TestEnv, error) { if cfg.recordDir != "" { serverOpts = append(serverOpts, server.WithRecordDir(cfg.recordDir)) } + if cfg.protocolStageEnabled { + serverOpts = append(serverOpts, server.WithProtocolStage(true)) + } + if cfg.guardrailsRuntime != nil { + serverOpts = append(serverOpts, server.WithGuardrails(cfg.guardrailsRuntime)) + } + if len(cfg.servertoolProviders) > 0 { + serverOpts = append(serverOpts, server.WithServertoolProviders(cfg.servertoolProviders...)) + } core, err := newGatewayCore("pv-env-*", func(ac *config.AppConfig) { + if cfg.recordDir != "" { + for _, scenario := range []typ.RuleScenario{typ.ScenarioAnthropic, typ.ScenarioOpenAI} { + _ = ac.GetGlobalConfig().SetScenarioStringFlag( + scenario, + serverconfig.FlagRecordingV2, + string(typ.RecordingModeStagedRequestResponse), + ) + } + } if cfg.mcpEnabled { _ = ac.GetGlobalConfig().SetScenarioFlag(typ.ScenarioGlobal, serverconfig.ExtensionMCP, true) } + if cfg.guardrailsRuntime != nil { + _ = ac.GetGlobalConfig().SetScenarioFlag(typ.ScenarioGlobal, serverconfig.ExtensionGuardrails, true) + } }, serverOpts...) if err != nil { return nil, err @@ -244,6 +303,7 @@ func NewTestEnvForCLI(opts ...TestEnvOption) (*TestEnv, error) { return &TestEnv{ appConfig: core.appConfig, + rootServer: core.rootServer, gatewayServer: core.gateway, virtual: core.virtual, modelToken: core.modelToken, @@ -267,6 +327,16 @@ func (env *TestEnv) VirtualURL() string { return env.virtual.URL() } // VirtualCallCount returns the number of requests received by the virtual server. func (env *TestEnv) VirtualCallCount() int { return env.virtual.CallCount() } +// ForceFlushRecordings waits until the real gateway has exported every queued +// recording artifact. It is safe to call repeatedly while the harness server +// remains active, allowing each matrix case to verify its own persisted record. +func (env *TestEnv) ForceFlushRecordings(ctx context.Context) error { + if env == nil || env.rootServer == nil { + return nil + } + return env.rootServer.ForceFlushRecordings(ctx) +} + // SetupRoute configures a gateway rule that routes source protocol requests // to the virtual server acting as a target protocol provider. // @@ -631,7 +701,7 @@ func assembleFromEvents(events []string, style protocol.APIStyle) sse.ParsedResu // Try to assemble as Responses API first r = sse.AssembleOpenAIResponsesStream(events) // If that failed, try Chat Completions - if r == nil || len(r.Content) == 0 { + if r == nil || (len(r.Content) == 0 && len(r.ToolCalls) == 0) { r = sse.AssembleOpenAIStream(events) } case protocol.APIStyleAnthropic: diff --git a/internal/protocoltest/virtual_client_test.go b/internal/protocoltest/virtual_client_test.go index 7542c9986..72ec0afd2 100644 --- a/internal/protocoltest/virtual_client_test.go +++ b/internal/protocoltest/virtual_client_test.go @@ -99,6 +99,17 @@ func TestVirtualClient_ToolUse_OpenAI(t *testing.T) { assert.Contains(t, result.ToolCalls[0].Arguments, "location") } +func TestVirtualClient_ToolUse_OpenAIResponsesStream(t *testing.T) { + vs := protocoltest.NewVirtualServer(t) + vc := vs.Client() + + result := vc.SendOpenAIResponses(t, protocoltest.ToolUseScenario(), true) + require.Equal(t, 200, result.HTTPStatus) + require.Len(t, result.ToolCalls, 1) + assert.Equal(t, "get_weather", result.ToolCalls[0].Name) + assert.Contains(t, result.ToolCalls[0].Arguments, "location") +} + func TestVirtualClient_ToolUse_Anthropic(t *testing.T) { vs := protocoltest.NewVirtualServer(t) vc := vs.Client()