diff --git a/.gitignore b/.gitignore index 83e00e5..d0dcd7b 100644 --- a/.gitignore +++ b/.gitignore @@ -23,4 +23,7 @@ lcov.info go.work # Local .env files -*.local.idea/ +*.local + +# IDE / editor +.idea/ diff --git a/README.md b/README.md index 896085e..624ba1c 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,19 @@ GLOBAL OPTIONS: --auth-key value, -k value the authentication key to use for incoming requests. [$AUTH_KEY] + progress + + --progress-callback-timeout value the timeout for a single progress callback delivery. (default: 1s) [$PROGRESS_CALLBACK_TIMEOUT] + --progress-allowed-hosts value [ --progress-allowed-hosts value ] restrict progress callback URLs to these hosts. Entries may be an exact hostname or a "*.example.com" wildcard. Unset allows any host, subject to the private-network guard below. [$PROGRESS_ALLOWED_HOSTS] + --progress-allow-private-networks allow progress callback delivery to loopback, link-local, and private IP addresses. Leave disabled unless the callback target is known to live on a trusted private network. (default: false) [$PROGRESS_ALLOW_PRIVATE_NETWORKS] + --progress-sidecar-max-body-bytes value the maximum size, in bytes, of a single worker-authored progress event POST. (default: 16384) [$PROGRESS_SIDECAR_MAX_BODY_BYTES] + --progress-sidecar-max-events value the maximum number of worker-authored progress events relayed per evaluation. (default: 50) [$PROGRESS_SIDECAR_MAX_EVENTS] + --progress-sidecar-burst-size value how many worker-authored progress events at the start of an evaluation are exempt from the minimum spacing below, so a handful of legitimate back-to-back checkpoints aren't rate limited. (default: 5) [$PROGRESS_SIDECAR_BURST_SIZE] + --progress-sidecar-min-event-interval value the minimum spacing between worker-authored progress events relayed per evaluation, once the burst allowance above is used up. (default: 10ms) [$PROGRESS_SIDECAR_MIN_EVENT_INTERVAL] + --progress-sidecar-unbind-grace-period value how long to keep relaying worker-authored progress events after a request returns, so a fire-and-forget POST dispatched just before the result can still land. (default: 250ms) [$PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD] + --progress-stream-enabled stream progress back on the /evaluate and /chat responses as Server-Sent Events for requests that send 'Accept: text/event-stream' and negotiate 'X-Api-Version: 0.1.1-dev'. Standalone/serve mode only; ignored under AWS Lambda. (default: true) [$PROGRESS_STREAM_ENABLED] + --progress-stream-heartbeat-seconds value seconds between SSE heartbeat comments sent while an evaluation runs, so an idle streamed connection isn't dropped by an intermediary. 0 disables heartbeats. (default: 15) [$PROGRESS_STREAM_HEARTBEAT_SECONDS] + function --arg value, -a value [ --arg value, -a value ] additional arguments for to the worker process. [$FUNCTION_ARGS] @@ -186,6 +199,215 @@ Example request using cases: } ``` +### Progress Events + +The shim also exposes µEd-compatible endpoints at `POST /evaluate` and `POST /chat` (see the [µEd spec](https://mued.org/spec)), separate from the legacy `POST /` endpoint documented above. When a client calls either with a `callbackUrl` in the request body, the shim POSTs a small JSON event to that URL at each stage of processing — in addition to, not instead of, the normal synchronous HTTP response. + +This lets a caller show progress to the end user (e.g. "Starting…") without polling, and without the shim needing to hold a connection open. It works identically whether the shim is deployed standalone or on AWS Lambda. + +To opt in, include `callbackUrl` in the request body and, optionally, an `X-Request-Id` header — both are part of the µEd spec's own request contract, not shim-specific additions. Every event echoes back the `X-Request-Id` value verbatim so the caller can correlate it with the original request. + +```json +{ + "submission": { "type": "TEXT", "content": { "text": "..." } }, + "task": { "referenceSolution": { "text": "..." } }, + "callbackUrl": "https://your-service.example.com/hooks/shimmy-progress" +} +``` + +Stages, in order: + +| Stage | Producer | Meaning | +|-------|----------|---------| +| `preparing` | shim | A worker is being made ready (freshly booted or reused from the pool). Emitted once per request. | +| `starting` | shim | The worker is about to be invoked. Emitted once per request. | +| `evaluating` | worker | A progress checkpoint the evaluation function reported during an `/evaluate` (or `/preview`) call. Zero or more, in the function's own order. | +| `thinking` | worker | The `/chat` equivalent of `evaluating` — a checkpoint the chat function reported. | +| `completed` | shim | The result has been computed. For `/evaluate`, `data.feedback` carries the same array as the synchronous body; for `/chat`, `data.output` carries the message. | +| `failed` | shim | A terminal failure occurred. `message` is a short end-user-safe line; `error` is an `ErrorResponse` object (`title`, optional `message`/`code`/`trace`/`details`) for programmatic handling and logs. | + +`completed` and `failed` are terminal — at most one of them is delivered per request, whichever occurs first. `preparing` and `starting` are each delivered at most once even for a multi-case evaluation that internally re-enters those stages per case. + +Example event body: + +```json +{ + "correlationId": "req-7c193f38", + "stage": "starting", + "command": "eval", + "message": "Starting…", + "timestamp": "2026-08-04T14:23:01.512Z" +} +``` + +Example terminal event, with the feedback payload attached: + +```json +{ + "correlationId": "req-7c193f38", + "stage": "completed", + "command": "eval", + "message": "Feedback is ready.", + "data": { + "feedback": [ + { "awardedPoints": 1, "message": "Well done" } + ] + }, + "timestamp": "2026-08-04T14:23:02.310Z" +} +``` + +A `failed` terminal event carries an `error` object instead of `data`: + +```json +{ + "correlationId": "req-7c193f38", + "stage": "failed", + "command": "eval", + "message": "We couldn't evaluate your answer. Please try again.", + "error": { + "title": "Evaluation failed", + "message": "We couldn't evaluate your answer. Please try again.", + "code": "INTERNAL_ERROR", + "trace": "worker send: context deadline exceeded" + }, + "timestamp": "2026-08-04T14:23:02.310Z" +} +``` + +Delivery is best-effort and never blocks or fails the evaluation itself: each callback POST is bounded by `--progress-callback-timeout` (default `1s`, see [Usage](#usage)); a slow, unreachable, or erroring receiver is logged and skipped, never surfaced to the caller as an evaluation failure. + +#### Callback URL safety (SSRF protection) + +Since `callbackUrl` is caller-supplied, the shim guards against it being used to reach services it shouldn't be able to reach: + +- **By default**, callback delivery refuses to dial loopback, link-local (this includes cloud metadata endpoints like `169.254.169.254`), and private (RFC1918/RFC4193) IP addresses — checked against the address actually resolved and dialed, not just the URL's literal hostname, so a public-looking domain that resolves to a private address is still blocked. Set `--progress-allow-private-networks` only if the callback target is known to live on a private network you trust (e.g. a same-VPC service). +- **`--progress-allowed-hosts`** optionally restricts callback URLs to an explicit list of hostnames (exact match, or `*.example.com` wildcards). Unset means any (non-private) host is accepted. + +A rejected callback URL behaves like any other delivery failure: it's logged and skipped, never surfaced to the caller as an evaluation failure. + +> **Note:** the µEd spec describes `callbackUrl` for asynchronous *final-result* delivery — the service may return `202 Accepted` immediately and POST the result later. The shim doesn't implement that 202 flow; it always responds synchronously with `200 OK` and the feedback body as normal. It reuses the same `callbackUrl` field to additionally deliver progress events — including the final feedback, via the `completed` event's `data` field — rather than requiring a shim-specific header for the same concept. + +#### Streaming progress on the response itself (Server-Sent Events) + +A caller that would rather receive progress on the `/evaluate` or `/chat` response than +stand up a `callbackUrl` receiver can opt in with an `Accept: text/event-stream` request +header. The shim then keeps the response open and streams [Server-Sent Events](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events) +as the request runs, instead of the buffered JSON body. + +- **Requires the `0.1.1-dev` µEd API version.** SSE progress streaming is a shimmy + extension not yet in the published µEd `0.1.0` contract, pending upstream standardisation, + so it lives in a separate `0.1.1-dev` version. Select it with an `X-Api-Version: 0.1.1-dev` + request header; a request that resolves to `0.1.0` (the pinned default for header-less + clients) always gets the buffered JSON body even with `Accept: text/event-stream`. +- **Standalone / `serve` mode only.** Under AWS Lambda the proxy buffers the whole response, + so the `Accept` header is ignored and the normal buffered JSON body is returned. Disable it + everywhere with `--progress-stream-enabled=false`. +- Works alongside `callbackUrl`: if a request carries both, every event is delivered to the + stream **and** POSTed to the callback. +- The connection is bound to the request context — if the caller disconnects, the work is + cancelled. + +Each non-terminal event is written as its own frame the moment it occurs, so the caller sees +progress live (`/evaluate` shown; `/chat` is identical but with `thinking` frames in place +of `evaluating`): + +``` +event: preparing +data: {"stage":"preparing","message":"Preparing…","timestamp":"2026-08-31T09:16:29.474Z"} + +event: starting +data: {"stage":"starting","message":"Starting…","timestamp":"2026-08-31T09:16:29.474Z"} + +event: evaluating +data: {"stage":"evaluating","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"2026-08-31T09:16:29.522Z"} +``` + +`preparing` and `starting` (the shim's own markers) are streamed **once per request** even +though a multi-case evaluation re-enters them per case; worker-authored `evaluating` / +`thinking` sub-steps are streamed every time. The `event:` line carries the stage; the +`data` payload is a self-contained step object (`stage`, `message`, optional `data`, +`timestamp`). + +The stream then ends with exactly one terminal frame — `event: completed` or `event: failed` +— carrying the endpoint's normal `200` body plus every step that preceded it, and the +connection closes: + +``` +event: completed +data: {"feedback":[{"awardedPoints":1,"message":"Well done"}], + "steps":[{"stage":"preparing","message":"Preparing…","timestamp":"…"}, + {"stage":"starting","message":"Starting…","timestamp":"…"}, + {"stage":"evaluating","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"…"}]} +``` + +``` +event: failed +data: {"feedback":null, + "steps":[ /* whatever streamed before the failure */ ], + "error":{"title":"Evaluation failed", + "message":"We couldn't evaluate your answer. Please try again.", + "code":"INTERNAL_ERROR", + "trace":"worker send: context deadline exceeded"}} +``` + +For `/chat` the terminal frame carries `output` (and optional `metadata`) instead of +`feedback`: + +``` +event: completed +data: {"output":{"role":"ASSISTANT","content":"…"}, + "metadata":{ /* optional, worker-supplied */ }, + "steps":[ /* preparing, starting, thinking… */ ]} +``` +A failed `/chat` frame has `"output":null` plus the same `error` object. + +Each element of the terminal frame's `steps[]` is byte-identical to the `data` payload of the +live frame that carried it. The HTTP status is `200` even for a `failed` frame — the failure +is in-band. The terminal frame's `data` is the µEd spec's `SseEvaluateTerminalFrame` / +`SseChatTerminalFrame`; on failure its `error` is a standard `ErrorResponse`. The correlation +id is in the `X-Request-Id` response header, not the body. Response headers: `Content-Type: +text/event-stream`, `Cache-Control: no-cache`, `X-Accel-Buffering: no`, no `Content-Length`. + +While the request runs, the shim also writes an SSE comment heartbeat (`: ping`) every +`--progress-stream-heartbeat-seconds` seconds (default `15`; `0` disables) so an idle +connection isn't dropped by an intermediary. + +#### Custom progress events from the evaluation function + +The `preparing` and `starting` stages are emitted by shimmy itself, around the worker call as a whole. An evaluation or chat function that does multiple steps internally (e.g. several model calls) can emit its own progress events *during* that span, which are relayed through the same `callbackUrl` (and SSE stream) alongside shimmy's own events. + +When a request opts in to progress reporting (via `callbackUrl`), shimmy starts a loopback-only HTTP listener and passes its address to the evaluation function process as the `EVAL_PROGRESS_URL` environment variable, the same way it passes `EVAL_RPC_TRANSPORT`, `EVAL_FILE_NAME_REQUEST`, etc. (see [Communication Channels](#communication-channels) below). This works identically regardless of interface (`rpc` or `file`) or RPC transport, and regardless of the evaluation function's language — it only needs to be able to make an HTTP POST. + +To emit a custom event, `POST` a small JSON body to `EVAL_PROGRESS_URL`: + +```json +{ + "message": "Checking correctness…", + "data": { "step": 2, "of": 3 } +} +``` + +- `message` (string, required): student/teacher-facing text. +- `data` (object, optional): free-form, passed through as-is. +- There is no `stage` field, by design: a worker can never choose its own stage. The shim assigns one from the command in flight — `evaluating` for an `/evaluate` (or `/preview`) request, `thinking` for `/chat` — and the shim-only stages `preparing`, `starting`, `completed`, and `failed` are never available to a worker. + +The response status is informational only — the evaluation function should treat every response as fire-and-forget and never fail on a non-2xx status. Delivery is best-effort, same as outbound callback delivery: `202` accepted (delivery to `callbackUrl` is then attempted in the background), `400` malformed body or empty `message`, `413` body too large, `429` rate limited, `503` no request currently associated with the listener (e.g. a stray POST arriving after both the request has finished and the grace period below has elapsed). + +To bound how much an evaluation function (which may be running untrusted, sandboxed code) can push through this channel, events are capped before relay: + +| Flag | Env var | Default | Description | +|------|---------|---------|-------------| +| `--progress-sidecar-max-body-bytes` | `PROGRESS_SIDECAR_MAX_BODY_BYTES` | `16384` | Maximum size, in bytes, of a single event POST. | +| `--progress-sidecar-max-events` | `PROGRESS_SIDECAR_MAX_EVENTS` | `50` | Maximum number of events relayed per evaluation. | +| `--progress-sidecar-burst-size` | `PROGRESS_SIDECAR_BURST_SIZE` | `5` | Events at the start of a span exempt from the minimum spacing below, so a handful of legitimate back-to-back checkpoints aren't rate limited. | +| `--progress-sidecar-min-event-interval` | `PROGRESS_SIDECAR_MIN_EVENT_INTERVAL` | `10ms` | Minimum spacing between relayed events, once the burst allowance is used up. | +| `--progress-sidecar-unbind-grace-period` | `PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD` | `250ms` | How long the listener keeps relaying after a request returns, so a fire-and-forget event POST dispatched by the worker just before returning its result still has a window to land. | + +> **Sandboxing note:** under `--sandbox` alone, the worker keeps the host network namespace and can reach the loopback listener normally. Only the separate, explicit `--sandbox-disable-network` flag isolates networking (and loopback specifically) — under that flag, custom progress events are silently dropped, the same as any other best-effort delivery failure. + +This is a shim-side contract only; no client library ships in this repo. Evaluation function libraries (e.g. per-language toolkits) can build a thin wrapper around reading `EVAL_PROGRESS_URL` and POSTing to it. + ### Communication Channels The shim supports two interface modes, selected with `--interface`: @@ -212,6 +434,7 @@ The shim injects the following environment variables into the evaluation functio | `EVAL_RPC_HTTP_URL` | HTTP URL (HTTP transport only) | | `EVAL_RPC_WS_URL` | WebSocket URL (WS transport only) | | `EVAL_RPC_TCP_ADDRESS` | TCP address (TCP transport only) | +| `EVAL_PROGRESS_URL` | Local URL to POST [custom progress events](#custom-progress-events-from-the-evaluation-function) to (only set when the request opted in via `callbackUrl`) | #### File System (`--interface file`) @@ -237,6 +460,7 @@ The shim also sets the following environment variables: | `EVAL_IO` | `FILE` | | `EVAL_FILE_NAME_REQUEST` | Path to the input file | | `EVAL_FILE_NAME_RESPONSE` | Path to the output file | +| `EVAL_PROGRESS_URL` | Local URL to POST [custom progress events](#custom-progress-events-from-the-evaluation-function) to (only set when the request opted in via `callbackUrl`) | > Using the file interface is recommended for large payloads such as base64-encoded images. diff --git a/app/lambda/handler.go b/app/lambda/handler.go index c230576..45139c0 100644 --- a/app/lambda/handler.go +++ b/app/lambda/handler.go @@ -21,8 +21,10 @@ type LambdaHandlerParams struct { // Config is the configuration for the Lambda handler. Config Config - // Handlers is a slice of HTTP handlers grouped together. - Handlers []*server.HttpHandler `group:"handlers"` + // Mux is the shared, fully-wrapped application HTTP handler chain — the same + // one the standalone server serves, including OpenAPI request/response + // validation. + Mux *server.Mux // Context is the context for the Lambda handler. Context context.Context @@ -32,11 +34,11 @@ type LambdaHandlerParams struct { } type LambdaHandler struct { - config Config - ctx context.Context - cancel context.CancelFunc - mux *http.ServeMux - log *zap.Logger + config Config + ctx context.Context + cancel context.CancelFunc + handler http.Handler + log *zap.Logger } // NewLambdaHandler creates a new instance of LambdaHandler @@ -44,18 +46,12 @@ type LambdaHandler struct { func NewLambdaHandler(params LambdaHandlerParams) *LambdaHandler { ctx, cancel := context.WithCancel(params.Context) - mux := http.NewServeMux() - - for _, handler := range params.Handlers { - mux.Handle(handler.Name, handler.Handler) - } - return &LambdaHandler{ - config: params.Config, - ctx: ctx, - cancel: cancel, - mux: mux, - log: params.Logger, + config: params.Config, + ctx: ctx, + cancel: cancel, + handler: params.Mux, + log: params.Logger, } } @@ -101,11 +97,11 @@ func (s *LambdaHandler) Shutdown() { func (s *LambdaHandler) getProxyFunction() (any, error) { switch s.config.ProxySource { case ProxySourceApiGatewayV1: - return httpadapter.New(server.NormalizePath(s.mux)).ProxyWithContext, nil + return httpadapter.New(s.handler).ProxyWithContext, nil case ProxySourceApiGatewayV2: - return httpadapter.NewV2(server.NormalizePath(s.mux)).ProxyWithContext, nil + return httpadapter.NewV2(s.handler).ProxyWithContext, nil case ProxySourceAlb: - return httpadapter.NewALB(server.NormalizePath(s.mux)).ProxyWithContext, nil + return httpadapter.NewALB(s.handler).ProxyWithContext, nil default: return nil, fmt.Errorf("invalid proxy source: %s", s.config.ProxySource) } diff --git a/app/lambda/module.go b/app/lambda/module.go index 1ed820a..0f87592 100644 --- a/app/lambda/module.go +++ b/app/lambda/module.go @@ -4,6 +4,7 @@ import ( "go.uber.org/fx" "github.com/lambda-feedback/shimmy/handler" + "github.com/lambda-feedback/shimmy/internal/server" "github.com/lambda-feedback/shimmy/util/logging" ) @@ -14,8 +15,12 @@ func Module(config Config) fx.Option { fx.Supply(config), // rename logger for module logging.DecorateLogger("lambda"), + // the Lambda proxy buffers the whole response — no incremental streaming + fx.Supply(handler.StreamingCapability{Enabled: false}), // provide handlers handler.Module(), + // provide the shared HTTP handler chain (specs + wrapped mux) + server.HandlerModule(), // provide server fx.Provide(NewLifecycleHandler), // invoke server diff --git a/app/lambda/module_test.go b/app/lambda/module_test.go new file mode 100644 index 0000000..ecec613 --- /dev/null +++ b/app/lambda/module_test.go @@ -0,0 +1,33 @@ +package lambda + +import ( + "context" + "testing" + + "go.uber.org/fx" + "go.uber.org/zap" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/runtime" +) + +// TestModule_DependencyGraphResolves guards the fx wiring for Lambda mode +// given the globals app.New supplies. StreamingCapability is supplied +// here as {Enabled: false} — the Lambda proxy cannot stream. It now also +// pulls in server.HandlerModule so the Lambda adapter serves the same +// OpenAPI-validated handler chain as the standalone server. +func TestModule_DependencyGraphResolves(t *testing.T) { + cfg := config.Config{} + + err := fx.ValidateApp( + fx.NopLogger, + fx.Supply(fx.Annotate(context.Background(), fx.As(new(context.Context)))), + fx.Supply(zap.NewNop()), + fx.Supply(cfg), + runtime.Module(cfg.Runtime), + Module(Config{}), + ) + if err != nil { + t.Fatalf("lambda fx graph failed validation: %v", err) + } +} diff --git a/app/standalone/module.go b/app/standalone/module.go index e4be75c..02a4c30 100644 --- a/app/standalone/module.go +++ b/app/standalone/module.go @@ -13,6 +13,8 @@ func Module(config Config) fx.Option { "serve", // rename logger for module logging.DecorateLogger("serve"), + // the standalone HTTP server can stream responses incrementally + fx.Supply(handler.StreamingCapability{Enabled: true}), // provide handlers handler.Module(), // provide server diff --git a/app/standalone/module_test.go b/app/standalone/module_test.go new file mode 100644 index 0000000..5678946 --- /dev/null +++ b/app/standalone/module_test.go @@ -0,0 +1,34 @@ +package standalone + +import ( + "context" + "testing" + + "go.uber.org/fx" + "go.uber.org/zap" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/runtime" +) + +// TestModule_DependencyGraphResolves guards the fx wiring: the standalone +// module must be satisfiable given the globals app.New supplies (context, +// logger, config.Config, runtime module — which provides the µEd version +// registry). Regressions here — e.g. a handler param with no provider — +// surface as a validation error rather than a runtime panic on +// `shimmy serve`. +func TestModule_DependencyGraphResolves(t *testing.T) { + cfg := config.Config{} + + err := fx.ValidateApp( + fx.NopLogger, + fx.Supply(fx.Annotate(context.Background(), fx.As(new(context.Context)))), + fx.Supply(zap.NewNop()), + fx.Supply(cfg), + runtime.Module(cfg.Runtime), + Module(Config{}), + ) + if err != nil { + t.Fatalf("standalone fx graph failed validation: %v", err) + } +} diff --git a/cmd/root.go b/cmd/root.go index 7086d4e..11fcf7c 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -43,6 +43,76 @@ functions on arbitrary, serverless platforms.` Category: "auth", EnvVars: []string{"AUTH_KEY"}, }, + // progress flags + &cli.DurationFlag{ + Name: "progress-callback-timeout", + Usage: "the timeout for a single progress callback delivery.", + Value: time.Second, + Category: "progress", + EnvVars: []string{"PROGRESS_CALLBACK_TIMEOUT"}, + }, + &cli.StringSliceFlag{ + Name: "progress-allowed-hosts", + Usage: "restrict progress callback URLs to these hosts. Entries may be an exact hostname or a \"*.example.com\" wildcard. Unset allows any host, subject to the private-network guard below.", + Category: "progress", + EnvVars: []string{"PROGRESS_ALLOWED_HOSTS"}, + }, + &cli.BoolFlag{ + Name: "progress-allow-private-networks", + Usage: "allow progress callback delivery to loopback, link-local, and private IP addresses. Leave disabled unless the callback target is known to live on a trusted private network.", + Value: false, + Category: "progress", + EnvVars: []string{"PROGRESS_ALLOW_PRIVATE_NETWORKS"}, + }, + &cli.Int64Flag{ + Name: "progress-sidecar-max-body-bytes", + Usage: "the maximum size, in bytes, of a single worker-authored progress event POST.", + Value: 16 * 1024, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_MAX_BODY_BYTES"}, + }, + &cli.IntFlag{ + Name: "progress-sidecar-max-events", + Usage: "the maximum number of worker-authored progress events relayed per evaluation.", + Value: 50, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_MAX_EVENTS"}, + }, + &cli.IntFlag{ + Name: "progress-sidecar-burst-size", + Usage: "how many worker-authored progress events at the start of an evaluation are exempt from the minimum spacing below, so a handful of legitimate back-to-back checkpoints aren't rate limited.", + Value: 5, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_BURST_SIZE"}, + }, + &cli.DurationFlag{ + Name: "progress-sidecar-min-event-interval", + Usage: "the minimum spacing between worker-authored progress events relayed per evaluation, once the burst allowance above is used up.", + Value: 10 * time.Millisecond, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_MIN_EVENT_INTERVAL"}, + }, + &cli.DurationFlag{ + Name: "progress-sidecar-unbind-grace-period", + Usage: "how long to keep relaying worker-authored progress events after a request returns, so a fire-and-forget POST dispatched just before the result can still land.", + Value: 250 * time.Millisecond, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD"}, + }, + &cli.BoolFlag{ + Name: "progress-stream-enabled", + Usage: "stream progress back on the /evaluate and /chat responses as Server-Sent Events for requests that send 'Accept: text/event-stream' and negotiate 'X-Api-Version: 0.1.1-dev'. Standalone/serve mode only; ignored under AWS Lambda.", + Value: true, + Category: "progress", + EnvVars: []string{"PROGRESS_STREAM_ENABLED"}, + }, + &cli.IntFlag{ + Name: "progress-stream-heartbeat-seconds", + Usage: "seconds between SSE heartbeat comments sent while an evaluation runs, so an idle streamed connection isn't dropped by an intermediary. 0 disables heartbeats.", + Value: 15, + Category: "progress", + EnvVars: []string{"PROGRESS_STREAM_HEARTBEAT_SECONDS"}, + }, // shim flags &cli.StringFlag{ Name: "interface", @@ -371,21 +441,31 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { // map cli flags to config fields cliMap := map[string]string{ - "auth-key": "auth.key", - "max-workers": "runtime.max_workers", - "command": "runtime.cmd", - "cwd": "runtime.cwd", - "arg": "runtime.arg", - "env": "runtime.env", - "interface": "runtime.io.interface", - "rpc-transport": "runtime.io.rpc.transport", - "rpc-transport-ipc-endpoint": "runtime.io.rpc.ipc.endpoint", - "rpc-transport-http-url": "runtime.io.rpc.http.url", - "rpc-transport-ws-url": "runtime.io.rpc.ws.url", - "rpc-transport-tcp-address": "runtime.io.rpc.tcp.address", - "worker-send-timeout": "runtime.send.timeout", - "worker-stop-timeout": "runtime.stop.timeout", - "worker-start-timeout": "start_timeout", + "auth-key": "auth.key", + "progress-callback-timeout": "progress.callback_timeout", + "progress-allowed-hosts": "progress.allowed_hosts", + "progress-allow-private-networks": "progress.allow_private_networks", + "progress-sidecar-max-body-bytes": "progress.sidecar.max_body_bytes", + "progress-sidecar-max-events": "progress.sidecar.max_events_per_span", + "progress-sidecar-burst-size": "progress.sidecar.burst_size", + "progress-sidecar-min-event-interval": "progress.sidecar.min_event_interval", + "progress-sidecar-unbind-grace-period": "progress.sidecar.unbind_grace_period", + "progress-stream-enabled": "progress.stream.enabled", + "progress-stream-heartbeat-seconds": "progress.stream.heartbeat_seconds", + "max-workers": "runtime.max_workers", + "command": "runtime.cmd", + "cwd": "runtime.cwd", + "arg": "runtime.arg", + "env": "runtime.env", + "interface": "runtime.io.interface", + "rpc-transport": "runtime.io.rpc.transport", + "rpc-transport-ipc-endpoint": "runtime.io.rpc.ipc.endpoint", + "rpc-transport-http-url": "runtime.io.rpc.http.url", + "rpc-transport-ws-url": "runtime.io.rpc.ws.url", + "rpc-transport-tcp-address": "runtime.io.rpc.tcp.address", + "worker-send-timeout": "runtime.send.timeout", + "worker-stop-timeout": "runtime.stop.timeout", + "worker-start-timeout": "start_timeout", // sandbox "sandbox": "runtime.sandbox.enabled", "sandbox-nsjail-path": "runtime.sandbox.nsjail_path", diff --git a/config/config.go b/config/config.go index 59daf40..3a18018 100644 --- a/config/config.go +++ b/config/config.go @@ -3,6 +3,7 @@ package config import ( "time" + "github.com/lambda-feedback/shimmy/internal/progress" "github.com/lambda-feedback/shimmy/runtime" ) @@ -30,6 +31,9 @@ type Config struct { // Auth is the authentication configuration Auth AuthConfig `conf:"auth"` + // Progress is the configuration for outbound progress-callback delivery + Progress progress.Config `conf:"progress"` + // StartTimeout is the duration to wait for the application to start. StartTimeout time.Duration `conf:"start_timeout"` } diff --git a/handler/chat.go b/handler/chat.go index 24e9e26..64b0b5e 100644 --- a/handler/chat.go +++ b/handler/chat.go @@ -1,21 +1,29 @@ package handler import ( + "context" "encoding/json" + "fmt" "io" "net/http" + "go.uber.org/zap" + + "github.com/lambda-feedback/shimmy/internal/progress" "github.com/lambda-feedback/shimmy/internal/server" "github.com/lambda-feedback/shimmy/runtime" ) // ServeChat handles POST /chat. func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) { + requestID := resolveRequestID(r) + w.Header().Set(muEdRequestIDHeader, requestID) + if !h.checkAuth(w, r) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -31,49 +39,139 @@ func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) { return } - var chatReq runtime.MuEdChatRequest - if err := json.Unmarshal(body, &chatReq); err != nil { - h.writeMuEdError(w, version, http.StatusBadRequest, "VALIDATION_ERROR", "Bad request", "invalid request body", nil) - return - } - - reqData, err := runtime.MuEdBuildChatRequest(chatReq) + reqData, err := adapter.DecodeChat(body) if err != nil { h.writeMuEdError(w, version, http.StatusBadRequest, "VALIDATION_ERROR", "Bad request", err.Error(), nil) return } - resp, err := h.runtime.Chat(r.Context(), runtime.ChatRequest{Data: reqData}) - if err != nil { - h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", "chat failed", nil) - return + // The adapter's DecodeChat deliberately does not surface transport + // concerns like callbackUrl, so pull it from the raw body here. A parse + // failure is impossible in practice — DecodeChat already parsed the + // same bytes — so a miss just means no out-of-band progress delivery. + var callbackURL string + var chatReq runtime.MuEdChatRequest + if json.Unmarshal(body, &chatReq) == nil && chatReq.CallbackUrl != nil { + callbackURL = *chatReq.CallbackUrl } - resultMap, ok := resp.Data["result"].(map[string]any) - if !ok { - h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", "invalid response from chat function", nil) + streaming := h.streamingEnabled() && acceptsEventStream(r) && adapter.SupportsStreaming() + if streaming { + if _, ok := w.(http.Flusher); !ok { + h.log.Warn("response writer is not a flusher; serving buffered response") + streaming = false + } + } + + ctx := r.Context() + + if streaming { + h.serveChatStream(ctx, w, reqData, adapter, version, callbackURL, requestID) return } - chatResp, err := runtime.MuEdToChatResponse(resultMap) - if err != nil { - h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", err.Error(), nil) + if callbackURL != "" && h.progressFactory != nil { + reporter, rerr := h.progressFactory.NewReporter(callbackURL, requestID) + if rerr != nil { + h.log.Warn("invalid callbackUrl, disabling progress reporting", zap.Error(rerr)) + } else if reporter != nil { + ctx = progress.ContextWithReporter(ctx, reporter) + } + } + + resp, err := h.runtime.Chat(ctx, runtime.ChatRequest{Data: reqData}) + chatResp, termErr := h.produceChatOutput(resp, err, adapter) + if termErr != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: string(runtime.CommandChat), + Message: termErr.userMessage, + Error: termErr.rawError, + ErrorInfo: termErr.progressErrorInfo("Chat failed"), + }) + h.writeMuEdError(w, version, termErr.status, termErr.muEdCode, termErr.muEdTitle, termErr.muEdMessage, nil) return } + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: string(runtime.CommandChat), + Message: "Response is ready.", + Data: chatResp, + }) + w.Header().Set("Content-Type", "application/json") w.Header().Set(muEdVersionHeader, version) w.WriteHeader(http.StatusOK) json.NewEncoder(w).Encode(chatResp) //nolint:errcheck } +// serveChatStream handles a POST /chat request that opted in to SSE +// streaming. The streaming scaffold lives in streamProgress; this only +// supplies the run step. +func (h *MuEdHandler) serveChatStream( + ctx context.Context, + w http.ResponseWriter, + reqData map[string]any, + adapter runtime.MuEdAdapter, + version string, + callbackURL string, + requestID string, +) { + h.streamProgress(ctx, w, "chat", string(runtime.CommandChat), "Response is ready.", version, callbackURL, requestID, + func(ctx context.Context) (map[string]any, *terminalError) { + resp, err := h.runtime.Chat(ctx, runtime.ChatRequest{Data: reqData}) + chatResp, termErr := h.produceChatOutput(resp, err, adapter) + if termErr != nil { + return nil, termErr + } + return chatResp, nil + }) +} + +// produceChatOutput turns a runtime chat response into the µEd chat +// response object via the resolved version adapter, or a terminalError +// describing why it couldn't. It is pure: no writes, no progress events. +// Unlike produceFeedback there is no worker-non-200 passthrough — +// runtime.Chat returns (response, error), not an HTTP status — so every +// failure is a 500-class terminalError. +func (h *MuEdHandler) produceChatOutput(resp runtime.ChatResponse, chatErr error, adapter runtime.MuEdAdapter) (map[string]any, *terminalError) { + newErr := func(muEdMessage, rawError string) *terminalError { + return &terminalError{ + status: http.StatusInternalServerError, + muEdCode: "INTERNAL_ERROR", + muEdTitle: "Internal server error", + muEdMessage: muEdMessage, + userMessage: "We couldn't generate a response. Please try again.", + rawError: rawError, + } + } + + if chatErr != nil { + return nil, newErr("chat failed", chatErr.Error()) + } + + resultMap, ok := resp.Data["result"].(map[string]any) + if !ok { + return nil, newErr("invalid response from chat function", "invalid response from chat function") + } + + chatResp, err := adapter.EncodeChat(resultMap) + if err != nil { + return nil, newErr(err.Error(), fmt.Sprintf("invalid chat response: %v", err)) + } + return chatResp, nil +} + // ServeChatHealth handles GET /chat/health. func (h *MuEdHandler) ServeChatHealth(w http.ResponseWriter, r *http.Request) { + w.Header().Set(muEdRequestIDHeader, resolveRequestID(r)) + if !h.checkAuth(w, r) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -95,7 +193,7 @@ func (h *MuEdHandler) ServeChatHealth(w http.ResponseWriter, r *http.Request) { return } - healthResp := runtime.MuEdToChatHealthResponse(resultMap) + healthResp := adapter.EncodeChatHealth(resultMap, h.streamingEnabled()) statusCode := http.StatusOK if status, ok := healthResp["status"].(string); ok && status == string(runtime.MuEdChatHealthStatusUnavailable) { diff --git a/handler/chat_stream_test.go b/handler/chat_stream_test.go new file mode 100644 index 0000000..7ec5eb4 --- /dev/null +++ b/handler/chat_stream_test.go @@ -0,0 +1,307 @@ +package handler + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" + "github.com/lambda-feedback/shimmy/internal/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +// --- helpers --- + +func newChatStreamHandler(rt *MockRuntime, pf progress.Factory, opts progress.StreamConfig) *MuEdHandler { + if pf == nil { + pf = inertFactory() + } + return &MuEdHandler{ + runtime: rt, + config: config.Config{Progress: progress.Config{Stream: opts}}, + log: zap.NewNop(), + progressFactory: pf, + streamingCapable: true, + } +} + +func chatSSERequest(t *testing.T, body []byte) *http.Request { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/chat", bytes.NewReader(body)) + req.Header.Set("Accept", "text/event-stream") + // SSE streaming is only offered on µEd versions whose contract declares it. + req.Header.Set("X-Api-Version", "0.1.1-dev") + return req +} + +func chatBodyWithCallback(t *testing.T, callbackURL string) []byte { + t.Helper() + return mustMarshal(t, map[string]any{ + "messages": []map[string]any{{"role": "USER", "content": "hello"}}, + "callbackUrl": callbackURL, + }) +} + +// --- tests --- + +func TestServeChat_SSE_Success(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "Here you go"), nil) + + req := chatSSERequest(t, chatRequestBody(t)) + req.Header.Set(muEdRequestIDHeader, "corr-chat") + w := httptest.NewRecorder() + + newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}).ServeChat(w, req) + + res := w.Result() + assert.Equal(t, http.StatusOK, res.StatusCode) + assert.Equal(t, "text/event-stream", res.Header.Get("Content-Type")) + assert.Equal(t, "corr-chat", res.Header.Get(muEdRequestIDHeader)) + assert.Empty(t, res.Header.Get("Content-Length")) + + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "completed", event) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") + if _, hasFeedback := data["feedback"]; hasFeedback { + t.Errorf("chat frame must not carry a feedback key: %v", data) + } + out, ok := data["output"].(map[string]any) + require.True(t, ok, "output should be an object: %v", data["output"]) + assert.Equal(t, "Here you go", out["content"]) + _, ok = data["steps"].([]any) + assert.True(t, ok, "steps should always be present as an array") +} + +func TestServeChat_SSE_StreamsThinkingFrames(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + ctx := args.Get(0).(context.Context) + progress.Emit(ctx, progress.Event{Stage: progress.StagePreparing, Message: "Preparing…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageStarting, Message: "Starting…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageThinking, Message: "Searching your notes…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageThinking, Message: "Drafting a reply…"}) + }). + Return(chatRuntimeResponse("ASSISTANT", "Done"), nil) + + w := httptest.NewRecorder() + newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}). + ServeChat(w, chatSSERequest(t, chatRequestBody(t))) + + frames := parseSSEAll(t, w.Body.String()) + var events []string + for _, f := range frames { + events = append(events, f.event) + } + assert.Equal(t, []string{"preparing", "starting", "thinking", "thinking", "completed"}, events) + assert.Equal(t, "Searching your notes…", frames[2].data["message"]) + + steps := frames[4].data["steps"].([]any) + require.Len(t, steps, 4) + assert.Equal(t, "thinking", steps[3].(map[string]any)["stage"]) +} + +func TestServeChat_SSE_RuntimeError_BecomesFailedFrameAt200(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", ""), assertAnError{}) + + w := httptest.NewRecorder() + newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}). + ServeChat(w, chatSSERequest(t, chatRequestBody(t))) + + res := w.Result() + assert.Equal(t, http.StatusOK, res.StatusCode, "the stream stays 200; failure is in-band") + + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "failed", event) + assert.Nil(t, data["output"]) + errObj, ok := data["error"].(map[string]any) + require.True(t, ok, "error should be an ErrorResponse object, got %T", data["error"]) + assert.NotEmpty(t, errObj["title"]) + assert.Contains(t, errObj["trace"], "boom") +} + +func TestServeChat_SSE_CapabilityDisabled_FallsBackToJSON(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "hi"), nil) + + h := newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}) + h.streamingCapable = false + + w := httptest.NewRecorder() + h.ServeChat(w, chatSSERequest(t, chatRequestBody(t))) + + res := w.Result() + assert.Equal(t, http.StatusOK, res.StatusCode) + assert.Equal(t, "application/json", res.Header.Get("Content-Type")) + + var chatResp map[string]any + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &chatResp)) + out := chatResp["output"].(map[string]any) + assert.Equal(t, "hi", out["content"]) +} + +func TestServeChat_SSE_StreamConfigDisabled_FallsBackToJSON(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "hi"), nil) + + w := httptest.NewRecorder() + newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: false}). + ServeChat(w, chatSSERequest(t, chatRequestBody(t))) + + assert.Equal(t, "application/json", w.Result().Header.Get("Content-Type")) +} + +func TestServeChat_SSE_NoAcceptHeader_Unchanged(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "hi"), nil) + + req := httptest.NewRequest(http.MethodPost, "/chat", bytes.NewReader(chatRequestBody(t))) + w := httptest.NewRecorder() + newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}).ServeChat(w, req) + + assert.Equal(t, "application/json", w.Result().Header.Get("Content-Type")) +} + +func TestServeChat_SSE_WithCallbackUrl_BothDelivered(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "Here you go"), nil) + + req := chatSSERequest(t, chatBodyWithCallback(t, srv.URL)) + req.Header.Set(muEdRequestIDHeader, "corr-chat-both") + w := httptest.NewRecorder() + + newChatStreamHandler(rt, newProgressFactory(t, time.Second), progress.StreamConfig{Enabled: true}). + ServeChat(w, req) + + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "completed", event) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") + + require.Len(t, *received, 1) + evt := (*received)[0] + assert.Equal(t, "corr-chat-both", evt["correlationId"]) + assert.Equal(t, "completed", evt["stage"]) +} + +func TestServeChat_SSE_AuthFailure_StillHTTPError(t *testing.T) { + rt := new(MockRuntime) + h := newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}) + h.config.Auth.Key = "secret" + + w := httptest.NewRecorder() + h.ServeChat(w, chatSSERequest(t, chatRequestBody(t))) + + assert.Equal(t, http.StatusUnauthorized, w.Result().StatusCode) + assert.NotEqual(t, "text/event-stream", w.Result().Header.Get("Content-Type")) + rt.AssertNotCalled(t, "Chat", mock.Anything, mock.Anything) +} + +func TestServeChat_SSE_Heartbeat(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "hi"), nil). + After(1200 * time.Millisecond) + + h := newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true, HeartbeatSeconds: 1}) + srv := httptest.NewServer(http.HandlerFunc(h.ServeChat)) + defer srv.Close() + + req, err := http.NewRequest(http.MethodPost, srv.URL+"/chat", bytes.NewReader(chatRequestBody(t))) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("X-Api-Version", "0.1.1-dev") + + resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + var sawPing bool + reader := bufio.NewReader(resp.Body) + for { + line, err := reader.ReadString('\n') + if strings.HasPrefix(line, ": ping") { + sawPing = true + } + if strings.HasPrefix(line, "event: completed") { + break + } + if err == io.EOF { + break + } + require.NoError(t, err) + } + assert.True(t, sawPing, "expected at least one heartbeat before the completed frame") +} + +// TestServeChat_SSE_TerminalPayloadValidated covers the terminal-frame +// schema check: with the spec wired in, a worker response that would +// violate the µEd ChatResponse schema is turned into a "failed" frame +// rather than shipped as "completed". +func TestServeChat_SSE_TerminalPayloadValidated(t *testing.T) { + spec, err := server.LoadOpenAPISpec() + require.NoError(t, err) + + t.Run("valid payload still completes", func(t *testing.T) { + rt := new(MockRuntime) + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "all good"), nil) + + h := newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}) + h.spec = spec + + w := httptest.NewRecorder() + h.ServeChat(w, chatSSERequest(t, chatRequestBody(t))) + + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "completed", event) + out := data["output"].(map[string]any) + assert.Equal(t, "all good", out["content"]) + }) + + t.Run("schema-invalid payload becomes a failed frame", func(t *testing.T) { + rt := new(MockRuntime) + // "ROBOT" is not in the Message.role enum. + rt.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ROBOT", "hello"), nil) + + h := newChatStreamHandler(rt, nil, progress.StreamConfig{Enabled: true}) + h.spec = spec + + w := httptest.NewRecorder() + h.ServeChat(w, chatSSERequest(t, chatRequestBody(t))) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode, "failure is in-band") + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "failed", event) + assert.Nil(t, data["output"]) + errObj, ok := data["error"].(map[string]any) + require.True(t, ok, "error should be an ErrorResponse object, got %T", data["error"]) + assert.Equal(t, "Invalid response", errObj["title"]) + }) +} + +// assertAnError is an error whose message contains "boom", for the failure-path test. +type assertAnError struct{} + +func (assertAnError) Error() string { return "chat failed: boom" } diff --git a/handler/chat_test.go b/handler/chat_test.go index d6e717b..58f0ab8 100644 --- a/handler/chat_test.go +++ b/handler/chat_test.go @@ -266,7 +266,7 @@ func TestServeChat_UnsupportedVersionHeader(t *testing.T) { raw, _ := io.ReadAll(res.Body) assert.Equal(t, http.StatusNotAcceptable, res.StatusCode) - assert.Equal(t, "0.1.0", res.Header.Get("X-Api-Version")) + assert.Equal(t, "0.1.1-dev", res.Header.Get("X-Api-Version"), "406 stamps the latest supported version") var body map[string]any require.NoError(t, json.Unmarshal(raw, &body)) @@ -320,7 +320,7 @@ func TestServeChatHealth_UnsupportedVersionHeader(t *testing.T) { raw, _ := io.ReadAll(res.Body) assert.Equal(t, http.StatusNotAcceptable, res.StatusCode) - assert.Equal(t, "0.1.0", res.Header.Get("X-Api-Version")) + assert.Equal(t, "0.1.1-dev", res.Header.Get("X-Api-Version"), "406 stamps the latest supported version") var body map[string]any require.NoError(t, json.Unmarshal(raw, &body)) diff --git a/handler/evaluate.go b/handler/evaluate.go index c462b17..3c15e32 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -1,76 +1,145 @@ package handler import ( + "context" + "crypto/rand" "encoding/json" "fmt" "io" "net/http" + "strings" + "time" + "github.com/getkin/kin-openapi/openapi3" "go.uber.org/fx" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" "github.com/lambda-feedback/shimmy/internal/server" "github.com/lambda-feedback/shimmy/runtime" ) const muEdVersionHeader = "X-Api-Version" +// muEdRequestIDHeader is the µEd spec's request-tracing header (see +// https://mued.org/spec, X-Request-Id parameter). It's echoed back on every +// response, generating one if the caller didn't supply it, and progress +// events reuse the resolved value as their correlation key. +const muEdRequestIDHeader = "X-Request-Id" + +// resolveRequestID returns the caller-supplied X-Request-Id, or generates +// one if absent, so every request is traceable and correlatable even when +// the caller doesn't participate in tracing itself. +func resolveRequestID(r *http.Request) string { + if id := r.Header.Get(muEdRequestIDHeader); id != "" { + return id + } + return generateRequestID() +} + +func generateRequestID() string { + b := make([]byte, 4) + if _, err := rand.Read(b); err != nil { + // crypto/rand.Read on a real OS essentially never fails; fall back + // to a timestamp-based id rather than leaving the request untraceable. + return fmt.Sprintf("req-%08x", time.Now().UnixNano()) + } + return fmt.Sprintf("req-%x", b) +} + type MuEdHandlerParams struct { fx.In - Handler runtime.Handler - Runtime runtime.Runtime - Config config.Config - Log *zap.Logger + Handler runtime.Handler + Runtime runtime.Runtime + Registry *runtime.MuEdRegistry + Config config.Config + Log *zap.Logger + ProgressFactory progress.Factory + StreamingCapability StreamingCapability + + // Spec is the µEd OpenAPI spec, used to validate the SSE terminal + // frame payload (the streamed analogue of the buffered path's + // response validation). Optional: absent under AWS Lambda, which + // cannot stream anyway. + Spec *openapi3.T `optional:"true"` } type MuEdHandler struct { - handler runtime.Handler - runtime runtime.Runtime - config config.Config - log *zap.Logger + handler runtime.Handler + runtime runtime.Runtime + registry *runtime.MuEdRegistry + config config.Config + log *zap.Logger + progressFactory progress.Factory + streamingCapable bool + spec *openapi3.T } func NewMuEdHandler(params MuEdHandlerParams) *MuEdHandler { return &MuEdHandler{ - handler: params.Handler, - runtime: params.Runtime, - config: params.Config, - log: params.Log, + handler: params.Handler, + runtime: params.Runtime, + registry: params.Registry, + config: params.Config, + log: params.Log, + progressFactory: params.ProgressFactory, + streamingCapable: params.StreamingCapability.Enabled, + spec: params.Spec, } } +// muEdRegistry returns the handler's version registry, falling back to the +// process-wide default when none was injected (e.g. in unit tests). +func (h *MuEdHandler) muEdRegistry() *runtime.MuEdRegistry { + if h.registry != nil { + return h.registry + } + return runtime.DefaultMuEdRegistry() +} + +// streamingEnabled reports whether this deployment can and should stream +// SSE progress: a streaming-capable environment with streaming turned on +// in config. It gates both the runtime decision to stream a response and +// the capability shimmy advertises on its health endpoints. +func (h *MuEdHandler) streamingEnabled() bool { + return h.streamingCapable && h.config.Progress.Stream.Enabled +} + func writeJSONError(w http.ResponseWriter, msg string, status int) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(status) json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message": msg}}) //nolint:errcheck } -// checkMuEdVersion validates the X-Api-Version request header. -// Returns (resolvedVersion, true) on success, or writes a 406 and returns ("", false). -func (h *MuEdHandler) checkMuEdVersion(w http.ResponseWriter, r *http.Request) (string, bool) { +// checkMuEdVersion validates the X-Api-Version request header and resolves it to +// a concrete version adapter. Returns (resolvedVersion, adapter, true) on +// success, or writes a 406 and returns ("", nil, false). +func (h *MuEdHandler) checkMuEdVersion(w http.ResponseWriter, r *http.Request) (string, runtime.MuEdAdapter, bool) { + reg := h.muEdRegistry() requested := r.Header.Get(muEdVersionHeader) - if requested != "" && !runtime.MuEdIsVersionSupported(requested) { + if requested != "" && !reg.Supports(requested) { body, _ := json.Marshal(map[string]any{ "title": "API version not supported", "message": fmt.Sprintf( "The requested API version '%s' is not supported. Supported versions are: %v.", - requested, runtime.SupportedMuEdVersions, + requested, reg.Versions(), ), "code": "VERSION_NOT_SUPPORTED", "details": map[string]any{ "requestedVersion": requested, - "supportedVersions": runtime.SupportedMuEdVersions, + "supportedVersions": reg.Versions(), }, }) - w.Header().Set(muEdVersionHeader, runtime.MuEdResolveVersion(requested)) + w.Header().Set(muEdVersionHeader, reg.Resolve(requested)) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusNotAcceptable) w.Write(body) //nolint:errcheck - return "", false + return "", nil, false } - return runtime.MuEdResolveVersion(requested), true + version := reg.Resolve(requested) + return version, reg.Adapter(version), true } // writeMuEdError writes a structured muEd JSON error response with X-Api-Version header. @@ -98,11 +167,14 @@ func (h *MuEdHandler) checkAuth(w http.ResponseWriter, r *http.Request) bool { // ServeEvaluate handles POST /evaluate. func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { + requestID := resolveRequestID(r) + w.Header().Set(muEdRequestIDHeader, requestID) + if !h.checkAuth(w, r) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -118,20 +190,7 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { return } - var muEdReq runtime.MuEdEvaluateRequest - if err := json.Unmarshal(body, &muEdReq); err != nil { - h.writeMuEdError(w, version, http.StatusBadRequest, "VALIDATION_ERROR", "Bad request", "invalid request body", nil) - return - } - - isPreview := muEdReq.PreSubmissionFeedback != nil && muEdReq.PreSubmissionFeedback.Enabled - - var legacyBody map[string]any - if isPreview { - legacyBody, err = runtime.MuEdBuildLegacyPreviewRequest(muEdReq) - } else { - legacyBody, err = runtime.MuEdBuildLegacyEvaluateRequest(muEdReq) - } + legacyBody, command, err := adapter.DecodeEvaluate(body) if err != nil { h.writeMuEdError(w, version, http.StatusBadRequest, "VALIDATION_ERROR", "Bad request", err.Error(), nil) return @@ -143,11 +202,6 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { return } - command := runtime.CommandEvaluate - if isPreview { - command = runtime.CommandPreview - } - header := http.Header{} header.Set("Command", string(command)) @@ -158,52 +212,258 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { Header: header, } - resp := h.handler.Handle(r.Context(), req) + // The adapter's DecodeEvaluate deliberately does not surface transport + // concerns like callbackUrl, so pull it from the raw body here. A parse + // failure is impossible in practice — DecodeEvaluate already parsed the + // same bytes — so a miss just means no out-of-band progress delivery. + var callbackURL string + var muEdReq runtime.MuEdEvaluateRequest + if json.Unmarshal(body, &muEdReq) == nil && muEdReq.CallbackUrl != nil { + callbackURL = *muEdReq.CallbackUrl + } + + streaming := h.streamingEnabled() && acceptsEventStream(r) && adapter.SupportsStreaming() + if streaming { + if _, ok := w.(http.Flusher); !ok { + h.log.Warn("response writer is not a flusher; serving buffered response") + streaming = false + } + } + + ctx := r.Context() - if resp.StatusCode != http.StatusOK { - for k, v := range resp.Header { - for _, vv := range v { - w.Header().Add(k, vv) + if streaming { + h.serveEvaluateStream(ctx, w, req, adapter, command, version, callbackURL, requestID) + return + } + + if h.progressFactory != nil { + reporter, rerr := h.progressFactory.NewReporter(callbackURL, requestID) + if rerr != nil { + h.log.Warn("invalid callbackUrl, disabling progress reporting", zap.Error(rerr)) + } else if reporter != nil { + ctx = progress.ContextWithReporter(ctx, reporter) + } + } + + resp := h.handler.Handle(ctx, req) + + feedback, termErr := h.produceFeedback(resp, adapter, command) + if termErr != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: string(command), + Message: termErr.userMessage, + Error: termErr.rawError, + ErrorInfo: termErr.progressErrorInfo("Evaluation failed"), + }) + + if termErr.passthrough { + for k, v := range termErr.header { + for _, vv := range v { + w.Header().Add(k, vv) + } } + w.Header().Set(muEdVersionHeader, version) + w.WriteHeader(termErr.status) + w.Write(termErr.body) //nolint:errcheck + return } - w.Header().Set(muEdVersionHeader, version) - w.WriteHeader(resp.StatusCode) - w.Write(resp.Body) //nolint:errcheck + + h.writeMuEdError(w, version, termErr.status, termErr.muEdCode, termErr.muEdTitle, termErr.muEdMessage, nil) return } + // Carry the feedback itself on the completed event so that, when a + // caller supplies callbackUrl, that callback genuinely fulfils the + // µEd spec's "deliver feedback results to this URL" wording — even + // though shimmy always takes the synchronous 200 path rather than + // the spec's 202-Accepted deferred-delivery flow. + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: string(command), + Message: "Feedback is ready.", + Data: map[string]any{"feedback": feedback}, + }) + + w.Header().Set("Content-Type", "application/json") + w.Header().Set(muEdVersionHeader, version) + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(feedback) //nolint:errcheck +} + +// acceptsEventStream reports whether the caller opted in to an SSE +// streaming response via the Accept header. +func acceptsEventStream(r *http.Request) bool { + return strings.Contains(strings.ToLower(r.Header.Get("Accept")), "text/event-stream") +} + +// terminalError is the outcome of produceFeedback when feedback can't be +// produced. It carries everything the buffered and streaming paths each +// need to report the failure, so produceFeedback itself performs no +// writes and emits no events. +type terminalError struct { + // passthrough replays the evaluation function's own non-2xx response + // verbatim (buffered path only). + passthrough bool + header http.Header + status int + body []byte + + // muEd* describe a shimmy-internal error. The buffered path passes + // them straight to writeMuEdError; both paths also feed them into the + // StageFailed event's ErrorInfo via progressErrorInfo. + muEdCode string + muEdTitle string + muEdMessage string + + // userMessage and rawError are the human-facing line and raw detail + // for the StageFailed progress event (Message / Error), and the + // fallbacks progressErrorInfo uses for the ErrorInfo message / trace. + userMessage string + rawError string +} + +// progressErrorInfo maps the terminalError to the structured error object +// carried on the StageFailed progress event and emitted as the SSE +// "failed" frame's `error` (shaped like the spec's ErrorResponse). title +// falls back to fallbackTitle when the error has no µEd title (e.g. a +// worker-response passthrough), so the object always satisfies +// ErrorResponse, whose only required field is title. +func (e *terminalError) progressErrorInfo(fallbackTitle string) *progress.ErrorInfo { + title := e.muEdTitle + if title == "" { + title = fallbackTitle + } + msg := e.muEdMessage + if msg == "" { + msg = e.userMessage + } + return &progress.ErrorInfo{ + Title: title, + Message: msg, + Code: e.muEdCode, + Trace: e.rawError, + } +} + +// produceFeedback turns a runtime response into muEd feedback via the +// resolved version adapter, or a terminalError describing why it couldn't. +// It is pure: no writes, no progress events. +func (h *MuEdHandler) produceFeedback(resp runtime.Response, adapter runtime.MuEdAdapter, command runtime.Command) ([]map[string]any, *terminalError) { + if resp.StatusCode != http.StatusOK { + return nil, &terminalError{ + passthrough: true, + header: resp.Header, + status: resp.StatusCode, + body: resp.Body, + userMessage: muEdErrorMessageFromBody(resp.Body), + rawError: string(resp.Body), + } + } + var respBody map[string]any if err := json.Unmarshal(resp.Body, &respBody); err != nil { - h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", "failed to parse response", nil) - return + return nil, &terminalError{ + status: http.StatusInternalServerError, + muEdCode: "INTERNAL_ERROR", + muEdTitle: "Internal server error", + muEdMessage: "failed to parse response", + userMessage: "We couldn't evaluate your answer. Please try again.", + rawError: fmt.Sprintf("failed to parse response: %v", err), + } } result, ok := respBody["result"].(map[string]any) if !ok { - h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", "invalid response from evaluation function", nil) - return + return nil, &terminalError{ + status: http.StatusInternalServerError, + muEdCode: "INTERNAL_ERROR", + muEdTitle: "Internal server error", + muEdMessage: "invalid response from evaluation function", + userMessage: "We couldn't evaluate your answer. Please try again.", + rawError: "invalid response from evaluation function", + } + } + + feedback, err := adapter.EncodeEvaluateFeedback(command, result) + if err != nil { + return nil, &terminalError{ + status: http.StatusInternalServerError, + muEdCode: "INTERNAL_ERROR", + muEdTitle: "Internal server error", + muEdMessage: "failed to build feedback", + userMessage: "We couldn't evaluate your answer. Please try again.", + rawError: fmt.Sprintf("failed to build feedback: %v", err), + } } + return feedback, nil +} - var feedback []map[string]any - if isPreview { - feedback = runtime.MuEdToPreviewFeedback(result) - } else { - feedback = runtime.MuEdToEvaluateFeedback(result) +// serveEvaluateStream handles a POST /evaluate request that opted in to +// SSE streaming. The streaming scaffold (headers, reporter, heartbeats, +// terminal frame) lives in streamProgress; this only supplies the run +// step. Because the 200 is committed before the worker runs, every +// post-Handle outcome — including an internal error — becomes a "failed" +// frame, never an HTTP error. +func (h *MuEdHandler) serveEvaluateStream( + ctx context.Context, + w http.ResponseWriter, + req runtime.Request, + adapter runtime.MuEdAdapter, + command runtime.Command, + version string, + callbackURL string, + requestID string, +) { + cmdLabel := "evaluate" + if command == runtime.CommandPreview { + cmdLabel = "preview" } - w.Header().Set("Content-Type", "application/json") - w.Header().Set(muEdVersionHeader, version) - w.WriteHeader(http.StatusOK) - json.NewEncoder(w).Encode(feedback) //nolint:errcheck + h.streamProgress(ctx, w, cmdLabel, string(command), "Feedback is ready.", version, callbackURL, requestID, + func(ctx context.Context) (map[string]any, *terminalError) { + resp := h.handler.Handle(ctx, req) + feedback, termErr := h.produceFeedback(resp, adapter, command) + if termErr != nil { + return nil, termErr + } + return map[string]any{"feedback": feedback}, nil + }) +} + +// muEdErrorMessageFromBody best-effort extracts a human-readable message +// from a JSON error body of the shape {"error": {"message": "..."}}. +func muEdErrorMessageFromBody(body []byte) string { + const fallback = "We couldn't evaluate your answer. Please try again." + + var errBody map[string]any + if err := json.Unmarshal(body, &errBody); err != nil { + return fallback + } + + errObj, ok := errBody["error"].(map[string]any) + if !ok { + return fallback + } + + msg, ok := errObj["message"].(string) + if !ok || msg == "" { + return fallback + } + + return msg } // ServeHealth handles GET /evaluate/health. func (h *MuEdHandler) ServeHealth(w http.ResponseWriter, r *http.Request) { + w.Header().Set(muEdRequestIDHeader, resolveRequestID(r)) + if !h.checkAuth(w, r) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -228,7 +488,7 @@ func (h *MuEdHandler) ServeHealth(w http.ResponseWriter, r *http.Request) { return } - result := runtime.MuEdToHealthResponse(legacyResult) + result := adapter.EncodeHealth(legacyResult, h.streamingEnabled()) statusCode := http.StatusOK if s, ok := result["status"].(string); ok && s == "UNAVAILABLE" { diff --git a/handler/evaluate_stream_test.go b/handler/evaluate_stream_test.go new file mode 100644 index 0000000..e94bbc3 --- /dev/null +++ b/handler/evaluate_stream_test.go @@ -0,0 +1,453 @@ +package handler + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" + "github.com/lambda-feedback/shimmy/internal/server" + "github.com/lambda-feedback/shimmy/runtime" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +// --- helpers --- + +func inertFactory() progress.Factory { + return progress.NewHTTPFactory(progress.HTTPFactoryParams{Log: zap.NewNop()}) +} + +func newStreamHandler(h runtime.Handler, pf progress.Factory, opts progress.StreamConfig) *MuEdHandler { + if pf == nil { + pf = inertFactory() + } + return &MuEdHandler{ + handler: h, + config: config.Config{Progress: progress.Config{Stream: opts}}, + log: zap.NewNop(), + progressFactory: pf, + streamingCapable: true, + } +} + +func sseRequest(t *testing.T, body []byte) *http.Request { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(body)) + req.Header.Set("Accept", "text/event-stream") + // SSE streaming is only offered on µEd versions whose contract declares it. + req.Header.Set("X-Api-Version", "0.1.1-dev") + return req +} + +type sseFrame struct { + event string + data map[string]any +} + +// parseSSEAll returns every non-comment frame in order. +func parseSSEAll(t *testing.T, raw string) []sseFrame { + t.Helper() + var frames []sseFrame + for _, block := range strings.Split(strings.TrimSpace(raw), "\n\n") { + block = strings.TrimSpace(block) + if block == "" || strings.HasPrefix(block, ":") { + continue + } + var f sseFrame + for _, line := range strings.Split(block, "\n") { + switch { + case strings.HasPrefix(line, "event: "): + f.event = strings.TrimPrefix(line, "event: ") + case strings.HasPrefix(line, "data: "): + f.data = map[string]any{} + require.NoError(t, json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &f.data)) + } + } + frames = append(frames, f) + } + return frames +} + +// parseSSE returns (eventName, decoded data) of the single terminal frame. +func parseSSE(t *testing.T, raw string) (string, map[string]any) { + t.Helper() + var event string + var data map[string]any + for _, block := range strings.Split(strings.TrimSpace(raw), "\n\n") { + block = strings.TrimSpace(block) + if block == "" || strings.HasPrefix(block, ":") { + continue + } + for _, line := range strings.Split(block, "\n") { + switch { + case strings.HasPrefix(line, "event: "): + event = strings.TrimPrefix(line, "event: ") + case strings.HasPrefix(line, "data: "): + require.NoError(t, json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &data)) + } + } + } + return event, data +} + +func previewBody(t *testing.T) []byte { + t.Helper() + b, err := json.Marshal(map[string]any{ + "submission": map[string]any{ + "type": "MATH", + "content": map[string]any{"expression": "x^2"}, + }, + "preSubmissionFeedback": map[string]any{"enabled": true}, + }) + require.NoError(t, err) + return b +} + +// --- tests --- + +func TestServeEvaluate_SSE_Success(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := sseRequest(t, mathEvalBody(t)) + req.Header.Set(muEdRequestIDHeader, "corr-sse") + w := httptest.NewRecorder() + + newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}).ServeEvaluate(w, req) + + res := w.Result() + assert.Equal(t, http.StatusOK, res.StatusCode) + assert.Equal(t, "text/event-stream", res.Header.Get("Content-Type")) + assert.Equal(t, "corr-sse", res.Header.Get(muEdRequestIDHeader)) + assert.Empty(t, res.Header.Get("Content-Length")) + + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "completed", event) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") + + fb, ok := data["feedback"].([]any) + require.True(t, ok, "feedback should be an array: %v", data["feedback"]) + require.Len(t, fb, 1) + assert.Equal(t, "Well done", fb[0].(map[string]any)["message"]) + + _, ok = data["steps"].([]any) + assert.True(t, ok, "steps should always be present as an array") +} + +func TestServeEvaluate_SSE_StreamsLiveStepFrames(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + ctx := args.Get(0).(context.Context) + // Shim lifecycle markers, emitted twice as the per-case loop + // would; only the first of each reaches the wire. Worker + // sub-steps come in as StageEvaluating and are never collapsed. + progress.Emit(ctx, progress.Event{Stage: progress.StagePreparing, Message: "Preparing…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageStarting, Message: "Starting…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageEvaluating, Message: "Parsing response and answer..."}) + progress.Emit(ctx, progress.Event{Stage: progress.StageEvaluating, Message: "Comparing sets for equivalence..."}) + progress.Emit(ctx, progress.Event{Stage: progress.StagePreparing, Message: "Preparing…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageStarting, Message: "Starting…"}) + }). + Return(evalHandlerResponse(true, "Well done")) + + w := httptest.NewRecorder() + newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}). + ServeEvaluate(w, sseRequest(t, mathEvalBody(t))) + + frames := parseSSEAll(t, w.Body.String()) + var events []string + for _, f := range frames { + events = append(events, f.event) + } + assert.Equal(t, []string{"preparing", "starting", "evaluating", "evaluating", "completed"}, events) + + // A live frame's data is one step object, identical in shape to an + // element of the terminal frame's steps[]. + assert.Equal(t, "preparing", frames[0].data["stage"]) + assert.Equal(t, "Parsing response and answer...", frames[2].data["message"]) + + steps, ok := frames[4].data["steps"].([]any) + require.True(t, ok) + require.Len(t, steps, 4) + assert.Equal(t, "preparing", steps[0].(map[string]any)["stage"]) + assert.Equal(t, "starting", steps[1].(map[string]any)["stage"]) + assert.Equal(t, "evaluating", steps[3].(map[string]any)["stage"]) +} + +func TestServeEvaluate_SSE_Preview(t *testing.T) { + previewResp := runtime.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: mustMarshal(t, map[string]any{ + "command": "preview", + "result": map[string]any{"preview": map[string]any{"feedback": "looks right"}}, + }), + } + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything).Return(previewResp) + + w := httptest.NewRecorder() + newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}).ServeEvaluate(w, sseRequest(t, previewBody(t))) + + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "completed", event) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") + fb := data["feedback"].([]any) + require.Len(t, fb, 1) + _, ok := fb[0].(map[string]any)["preSubmissionFeedback"] + assert.True(t, ok, "expected preSubmissionFeedback wrapper, got %v", fb[0]) +} + +func TestServeEvaluate_SSE_WorkerNon200_BecomesFailedFrameAt200(t *testing.T) { + errorBody := mustMarshal(t, map[string]any{"error": map[string]any{"message": "boom"}}) + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything).Return(runtime.Response{ + StatusCode: http.StatusInternalServerError, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: errorBody, + }) + + w := httptest.NewRecorder() + newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}).ServeEvaluate(w, sseRequest(t, mathEvalBody(t))) + + res := w.Result() + assert.Equal(t, http.StatusOK, res.StatusCode, "the stream stays 200; failure is in-band") + assert.Equal(t, "text/event-stream", res.Header.Get("Content-Type")) + + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "failed", event) + assert.Nil(t, data["feedback"]) + errObj, ok := data["error"].(map[string]any) + require.True(t, ok, "error should be an ErrorResponse object, got %T", data["error"]) + assert.NotEmpty(t, errObj["title"]) + assert.Equal(t, "boom", errObj["message"]) + assert.Contains(t, errObj["trace"], "boom") + assert.NotContains(t, data, "message", "failure detail now lives in the error object, not a top-level message") +} + +func TestServeEvaluate_SSE_UnparseableWorkerResponse_FailedFrame(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything).Return(runtime.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: []byte("not json"), + }) + + w := httptest.NewRecorder() + newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}).ServeEvaluate(w, sseRequest(t, mathEvalBody(t))) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode) + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "failed", event) + assert.Nil(t, data["feedback"]) + errObj, ok := data["error"].(map[string]any) + require.True(t, ok, "error should be an ErrorResponse object, got %T", data["error"]) + assert.NotEmpty(t, errObj["title"]) +} + +func TestServeEvaluate_SSE_CapabilityDisabled_FallsBackToJSON(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + h := newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}) + h.streamingCapable = false + + w := httptest.NewRecorder() + h.ServeEvaluate(w, sseRequest(t, mathEvalBody(t))) + + res := w.Result() + assert.Equal(t, http.StatusOK, res.StatusCode) + assert.Equal(t, "application/json", res.Header.Get("Content-Type")) + + var feedback []map[string]any + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &feedback)) + require.Len(t, feedback, 1) + assert.Equal(t, "Well done", feedback[0]["message"]) +} + +func TestServeEvaluate_SSE_StreamConfigDisabled_FallsBackToJSON(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + w := httptest.NewRecorder() + newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: false}).ServeEvaluate(w, sseRequest(t, mathEvalBody(t))) + + assert.Equal(t, "application/json", w.Result().Header.Get("Content-Type")) +} + +func TestServeEvaluate_SSE_NoAcceptHeader_Unchanged(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + w := httptest.NewRecorder() + newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}).ServeEvaluate(w, req) + + assert.Equal(t, "application/json", w.Result().Header.Get("Content-Type")) +} + +func TestServeEvaluate_SSE_WithCallbackUrl_BothDelivered(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := sseRequest(t, mathEvalBodyWithCallback(t, srv.URL)) + req.Header.Set(muEdRequestIDHeader, "corr-both") + w := httptest.NewRecorder() + + newStreamHandler(mockHandler, newProgressFactory(t, time.Second), progress.StreamConfig{Enabled: true}). + ServeEvaluate(w, req) + + // SSE side + event, data := parseSSE(t, w.Body.String()) + assert.Equal(t, "completed", event) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") + + // callbackUrl side + require.Len(t, *received, 1) + evt := (*received)[0] + assert.Equal(t, "corr-both", evt["correlationId"]) + assert.Equal(t, "completed", evt["stage"]) +} + +func TestServeEvaluate_SSE_AuthFailure_StillHTTPError(t *testing.T) { + mockHandler := new(MockHandler) + h := newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}) + h.config.Auth.Key = "secret" + + w := httptest.NewRecorder() + h.ServeEvaluate(w, sseRequest(t, mathEvalBody(t))) + + assert.Equal(t, http.StatusUnauthorized, w.Result().StatusCode) + assert.NotEqual(t, "text/event-stream", w.Result().Header.Get("Content-Type")) + mockHandler.AssertNotCalled(t, "Handle", mock.Anything, mock.Anything) +} + +func TestServeEvaluate_SSE_UnsupportedVersion_StillHTTPError(t *testing.T) { + mockHandler := new(MockHandler) + h := newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}) + + req := sseRequest(t, mathEvalBody(t)) + req.Header.Set(muEdVersionHeader, "99.0.0") + w := httptest.NewRecorder() + h.ServeEvaluate(w, req) + + assert.Equal(t, http.StatusNotAcceptable, w.Result().StatusCode) + mockHandler.AssertNotCalled(t, "Handle", mock.Anything, mock.Anything) +} + +func TestServeEvaluate_SSE_Heartbeat(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")). + After(1200 * time.Millisecond) + + h := newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true, HeartbeatSeconds: 1}) + srv := httptest.NewServer(http.HandlerFunc(h.ServeEvaluate)) + defer srv.Close() + + reqBody := bytes.NewReader(mathEvalBody(t)) + req, err := http.NewRequest(http.MethodPost, srv.URL+"/evaluate", reqBody) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("X-Api-Version", "0.1.1-dev") + + resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + var sawPing bool + reader := bufio.NewReader(resp.Body) + for { + line, err := reader.ReadString('\n') + if strings.HasPrefix(line, ": ping") { + sawPing = true + } + if strings.HasPrefix(line, "event: completed") { + break + } + if err == io.EOF { + break + } + require.NoError(t, err) + } + assert.True(t, sawPing, "expected at least one heartbeat before the completed frame") +} + +// TestServeEvaluate_SSE_ThroughOpenAPIMiddleware exercises the real serve-mode +// chain end to end: a live socket, the OpenAPI middleware (which must NOT +// buffer the stream), NormalizePath, and the streaming handler with a real +// flushable ResponseWriter. +func TestServeEvaluate_SSE_ThroughOpenAPIMiddleware(t *testing.T) { + specs, err := server.LoadOpenAPISpecs() + require.NoError(t, err) + mw, err := server.OpenAPIMiddleware(specs, nil, zap.NewNop()) + require.NoError(t, err) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + h := newStreamHandler(mockHandler, nil, progress.StreamConfig{Enabled: true}) + + mux := http.NewServeMux() + mux.HandleFunc("/evaluate", h.ServeEvaluate) + srv := httptest.NewServer(mw(server.NormalizePath(mux))) + defer srv.Close() + + // Spec-valid body: the OpenAPI middleware validates the request before + // the streaming bypass, and the spec requires task.title. + body := mustMarshal(t, map[string]any{ + "submission": map[string]any{ + "type": "MATH", + "content": map[string]any{"expression": "x^2"}, + }, + "task": map[string]any{ + "title": "t", + "referenceSolution": map[string]any{"expression": "x^2"}, + }, + }) + req, err := http.NewRequest(http.MethodPost, srv.URL+"/evaluate", bytes.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Api-Version", "0.1.1-dev") + + resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) + assert.Empty(t, resp.Header.Get("Content-Length")) + + raw, err := io.ReadAll(resp.Body) + require.NoError(t, err) + event, data := parseSSE(t, string(raw)) + assert.Equal(t, "completed", event) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") +} + +func mustMarshal(t *testing.T, v any) []byte { + t.Helper() + b, err := json.Marshal(v) + require.NoError(t, err) + return b +} diff --git a/handler/evaluate_test.go b/handler/evaluate_test.go index 5fa03b6..4274d28 100644 --- a/handler/evaluate_test.go +++ b/handler/evaluate_test.go @@ -8,9 +8,12 @@ import ( "io" "net/http" "net/http/httptest" + "sync" "testing" + "time" "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" "github.com/lambda-feedback/shimmy/runtime" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" @@ -49,18 +52,34 @@ func (m *MockRuntime) Shutdown(ctx context.Context) error { // --- Helpers --- +// newMuEdHandler builds a handler with a default, inert progress factory: +// since none of the existing tests set callbackUrl in the request body, +// NewReporter always returns (nil, nil) and behavior is unchanged. Tests +// that exercise progress reporting itself use newMuEdHandlerWithProgress. func newMuEdHandler(h runtime.Handler, r runtime.Runtime, key string) *MuEdHandler { + return newMuEdHandlerWithProgress(h, r, key, progress.NewHTTPFactory(progress.HTTPFactoryParams{ + Log: zap.NewNop(), + })) +} + +func newMuEdHandlerWithProgress(h runtime.Handler, r runtime.Runtime, key string, pf progress.Factory) *MuEdHandler { return &MuEdHandler{ - handler: h, - runtime: r, - config: config.Config{Auth: config.AuthConfig{Key: key}}, - log: zap.NewNop(), + handler: h, + runtime: r, + config: config.Config{Auth: config.AuthConfig{Key: key}}, + log: zap.NewNop(), + progressFactory: pf, } } func mathEvalBody(t *testing.T) []byte { t.Helper() - b, err := json.Marshal(map[string]any{ + return mathEvalBodyWithCallback(t, "") +} + +func mathEvalBodyWithCallback(t *testing.T, callbackURL string) []byte { + t.Helper() + body := map[string]any{ "submission": map[string]any{ "type": "MATH", "content": map[string]any{"expression": "x^2"}, @@ -70,7 +89,11 @@ func mathEvalBody(t *testing.T) []byte { "expression": "x^2", }, }, - }) + } + if callbackURL != "" { + body["callbackUrl"] = callbackURL + } + b, err := json.Marshal(body) require.NoError(t, err) return b } @@ -284,6 +307,171 @@ func TestMuEdServeEvaluate_WorkerErrorForwarded(t *testing.T) { assert.Equal(t, errorBody, bytes.TrimRight(raw, "\n")) } +// --- Progress callback tests (ServeEvaluate) --- + +// newProgressCallbackServer spins up a fake progress-callback receiver +// that records every decoded request body it receives. +func newProgressCallbackServer(t *testing.T, handlerFn http.HandlerFunc) (*httptest.Server, *[]map[string]any) { + t.Helper() + + var mu sync.Mutex + var received []map[string]any + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + mu.Lock() + received = append(received, body) + mu.Unlock() + + if handlerFn != nil { + handlerFn(w, r) + return + } + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + + return srv, &received +} + +// newProgressFactory builds a factory with SSRF protection relaxed: these +// tests use httptest.NewServer (a loopback address) to stand in for the +// caller's real, non-loopback callback receiver, so the default +// private-network guard would otherwise reject every delivery here. The +// guard itself is covered directly in internal/progress. +func newProgressFactory(t *testing.T, timeout time.Duration) progress.Factory { + t.Helper() + return progress.NewHTTPFactory(progress.HTTPFactoryParams{ + Config: progress.Config{CallbackTimeout: timeout, AllowPrivateNetworks: true}, + Log: zap.NewNop(), + }) +} + +func TestMuEdServeEvaluate_ProgressCallback_Success(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + req.Header.Set(muEdRequestIDHeader, "corr-1") + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode) + + require.Len(t, *received, 1) + evt := (*received)[0] + assert.Equal(t, "corr-1", evt["correlationId"]) + assert.Equal(t, "completed", evt["stage"]) + assert.Equal(t, "eval", evt["command"]) + + data, ok := evt["data"].(map[string]any) + require.True(t, ok, "expected data field on the completed event") + feedback, ok := data["feedback"].([]any) + require.True(t, ok, "expected data.feedback array") + require.Len(t, feedback, 1) + item, ok := feedback[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "Well done", item["message"]) + assert.Equal(t, 1.0, item["awardedPoints"]) +} + +func TestMuEdServeEvaluate_ProgressCallback_Failure(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + errorBody, _ := json.Marshal(map[string]any{ + "error": map[string]any{"message": "boom"}, + }) + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything).Return(runtime.Response{ + StatusCode: http.StatusInternalServerError, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: errorBody, + }) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + req.Header.Set(muEdRequestIDHeader, "corr-2") + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + assert.Equal(t, http.StatusInternalServerError, w.Result().StatusCode) + + require.Len(t, *received, 1) + evt := (*received)[0] + assert.Equal(t, "corr-2", evt["correlationId"]) + assert.Equal(t, "failed", evt["stage"]) + assert.Equal(t, "boom", evt["message"]) + // error is the same ErrorResponse-shaped object as on the SSE frame. + errObj, ok := evt["error"].(map[string]any) + require.True(t, ok, "callback error should be an ErrorResponse object, got %T", evt["error"]) + assert.NotEmpty(t, errObj["title"]) + assert.Equal(t, "boom", errObj["message"]) +} + +func TestMuEdServeEvaluate_ProgressCallback_NoCallbackUrl_Unchanged(t *testing.T) { + _, received := newProgressCallbackServer(t, nil) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode) + assert.Empty(t, *received, "no callbackUrl in the request body should mean no callback requests") +} + +func TestMuEdServeEvaluate_ProgressCallback_InvalidCallbackUrl_EvaluationStillSucceeds(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, "not-a-url"))) + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + res := w.Result() + defer res.Body.Close() + body, _ := io.ReadAll(res.Body) + + assert.Equal(t, http.StatusOK, res.StatusCode) + + var feedback []map[string]any + require.NoError(t, json.Unmarshal(body, &feedback)) + require.Len(t, feedback, 1) +} + +func TestMuEdServeEvaluate_ProgressCallback_SlowReceiver_DoesNotBlockResponse(t *testing.T) { + srv, _ := newProgressCallbackServer(t, func(w http.ResponseWriter, r *http.Request) { + time.Sleep(200 * time.Millisecond) + w.WriteHeader(http.StatusOK) + }) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + req.Header.Set(muEdRequestIDHeader, "corr-3") + w := httptest.NewRecorder() + + start := time.Now() + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, 20*time.Millisecond)).ServeEvaluate(w, req) + elapsed := time.Since(start) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode) + assert.Less(t, elapsed, 150*time.Millisecond, "ServeEvaluate should return promptly, bounded by CallbackTimeout") +} + // --- ServeHealth tests --- func TestMuEdServeHealth_Success(t *testing.T) { @@ -448,7 +636,7 @@ func TestMuEdServeEvaluate_UnsupportedVersionHeader(t *testing.T) { raw, _ := io.ReadAll(res.Body) assert.Equal(t, http.StatusNotAcceptable, res.StatusCode) - assert.Equal(t, "0.1.0", res.Header.Get("X-Api-Version")) + assert.Equal(t, "0.1.1-dev", res.Header.Get("X-Api-Version"), "406 stamps the latest supported version") var body map[string]any require.NoError(t, json.Unmarshal(raw, &body)) @@ -498,7 +686,7 @@ func TestMuEdServeHealth_UnsupportedVersionHeader(t *testing.T) { raw, _ := io.ReadAll(res.Body) assert.Equal(t, http.StatusNotAcceptable, res.StatusCode) - assert.Equal(t, "0.1.0", res.Header.Get("X-Api-Version")) + assert.Equal(t, "0.1.1-dev", res.Header.Get("X-Api-Version"), "406 stamps the latest supported version") var body map[string]any require.NoError(t, json.Unmarshal(raw, &body)) @@ -506,3 +694,62 @@ func TestMuEdServeHealth_UnsupportedVersionHeader(t *testing.T) { mockRuntime.AssertNotCalled(t, "Handle", mock.Anything, mock.Anything) } + +// --- Request ID tests (ServeEvaluate) --- + +func TestMuEdServeEvaluate_RequestID_EchoedWhenSupplied(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "ok")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + req.Header.Set(muEdRequestIDHeader, "caller-supplied-id") + w := httptest.NewRecorder() + + newMuEdHandler(mockHandler, nil, "").ServeEvaluate(w, req) + + assert.Equal(t, "caller-supplied-id", w.Result().Header.Get(muEdRequestIDHeader)) +} + +func TestMuEdServeEvaluate_RequestID_GeneratedWhenAbsent(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "ok")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + w := httptest.NewRecorder() + + newMuEdHandler(mockHandler, nil, "").ServeEvaluate(w, req) + + assert.NotEmpty(t, w.Result().Header.Get(muEdRequestIDHeader)) +} + +func TestMuEdServeEvaluate_RequestID_EchoedOnErrorResponses(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader([]byte("not json"))) + req.Header.Set(muEdRequestIDHeader, "caller-supplied-id") + w := httptest.NewRecorder() + + newMuEdHandler(new(MockHandler), nil, "").ServeEvaluate(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Result().StatusCode) + assert.Equal(t, "caller-supplied-id", w.Result().Header.Get(muEdRequestIDHeader)) +} + +func TestMuEdServeEvaluate_ProgressCallback_GeneratedRequestIDUsedAsCorrelation(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + respRequestID := w.Result().Header.Get(muEdRequestIDHeader) + require.NotEmpty(t, respRequestID) + + require.Len(t, *received, 1) + assert.Equal(t, respRequestID, (*received)[0]["correlationId"]) +} diff --git a/handler/module.go b/handler/module.go index c16e381..df10a89 100644 --- a/handler/module.go +++ b/handler/module.go @@ -1,6 +1,20 @@ package handler -import "go.uber.org/fx" +import ( + "go.uber.org/fx" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" +) + +// StreamingCapability tells the muEd handler whether the current +// execution environment can stream an HTTP response incrementally. +// It is true under the standalone HTTP server and false under the AWS +// Lambda proxy (which buffers the whole response). Each app module +// supplies its own value — it is deliberately not provided here. +type StreamingCapability struct { + Enabled bool +} func Module() fx.Option { return fx.Module("common", @@ -10,6 +24,8 @@ func Module() fx.Option { fx.Provide(NewHealthRoute), fx.Provide(NewMuEdEvaluateRoute), fx.Provide(NewMuEdEvaluateHealthRoute), + fx.Provide(func(cfg config.Config) progress.Config { return cfg.Progress }), + fx.Provide(progress.NewHTTPFactory), fx.Provide(NewMuEdChatRoute), fx.Provide(NewMuEdChatHealthRoute), ) diff --git a/handler/mued_version_test.go b/handler/mued_version_test.go new file mode 100644 index 0000000..c332390 --- /dev/null +++ b/handler/mued_version_test.go @@ -0,0 +1,146 @@ +package handler + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/runtime" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +// markerAdapter is a synthetic MuEdAdapter whose outputs are easy to recognise, +// used to prove the handler dispatches to the resolved version's adapter rather +// than to hard-coded v0.1.0 logic. +type markerAdapter struct{} + +func (markerAdapter) Version() string { return "9.9.9" } +func (markerAdapter) DecodeEvaluate([]byte) (map[string]any, runtime.Command, error) { + return map[string]any{"marker": "decoded"}, runtime.CommandEvaluate, nil +} +func (markerAdapter) EncodeEvaluateFeedback(runtime.Command, map[string]any) ([]map[string]any, error) { + return []map[string]any{{"marker": "feedback-9.9.9"}}, nil +} +func (markerAdapter) EncodeHealth(map[string]any, bool) map[string]any { + return map[string]any{"marker": "health-9.9.9"} +} +func (markerAdapter) DecodeChat([]byte) (map[string]any, error) { + return map[string]any{"marker": "chat"}, nil +} +func (markerAdapter) EncodeChat(map[string]any) (map[string]any, error) { + return map[string]any{"marker": "chat-9.9.9"}, nil +} +func (markerAdapter) EncodeChatHealth(map[string]any, bool) map[string]any { + return map[string]any{"marker": "chat-health-9.9.9"} +} +func (markerAdapter) SupportsStreaming() bool { return false } + +func newMuEdHandlerWithRegistry(h runtime.Handler, r runtime.Runtime, reg *runtime.MuEdRegistry) *MuEdHandler { + return &MuEdHandler{ + handler: h, + runtime: r, + registry: reg, + config: config.Config{}, + log: zap.NewNop(), + } +} + +// TestMuEdServeEvaluate_DispatchesToResolvedAdapter proves version dispatch: +// an X-Api-Version the injected registry supports is routed to that version's +// adapter, and the response echoes the resolved version. +func TestMuEdServeEvaluate_DispatchesToResolvedAdapter(t *testing.T) { + reg := runtime.NewMuEdRegistry() + reg.Register(markerAdapter{}) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything).Return(evalHandlerResponse(true, "ignored")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + req.Header.Set("X-Api-Version", "9.9.9") + w := httptest.NewRecorder() + + newMuEdHandlerWithRegistry(mockHandler, nil, reg).ServeEvaluate(w, req) + + res := w.Result() + defer res.Body.Close() + raw, _ := io.ReadAll(res.Body) + + assert.Equal(t, http.StatusOK, res.StatusCode) + assert.Equal(t, "9.9.9", res.Header.Get("X-Api-Version")) + + var feedback []map[string]any + require.NoError(t, json.Unmarshal(raw, &feedback)) + require.Len(t, feedback, 1) + assert.Equal(t, "feedback-9.9.9", feedback[0]["marker"]) + + mockHandler.AssertExpectations(t) +} + +// TestMuEdServeEvaluate_VersionParity runs the happy path with both an absent +// header and an explicit supported header and asserts identical output. +func TestMuEdServeEvaluate_VersionParity(t *testing.T) { + run := func(setHeader bool) (int, string, []map[string]any) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + if setHeader { + req.Header.Set("X-Api-Version", "0.1.0") + } + w := httptest.NewRecorder() + newMuEdHandler(mockHandler, nil, "").ServeEvaluate(w, req) + + res := w.Result() + defer res.Body.Close() + raw, _ := io.ReadAll(res.Body) + var fb []map[string]any + require.NoError(t, json.Unmarshal(raw, &fb)) + return res.StatusCode, res.Header.Get("X-Api-Version"), fb + } + + absentCode, absentVer, absentFb := run(false) + explicitCode, explicitVer, explicitFb := run(true) + + assert.Equal(t, http.StatusOK, absentCode) + assert.Equal(t, absentCode, explicitCode) + assert.Equal(t, "0.1.0", absentVer) + assert.Equal(t, absentVer, explicitVer) + assert.Equal(t, absentFb, explicitFb) +} + +// TestMuEdServeChat_DispatchesToResolvedAdapter is the chat-side counterpart. +func TestMuEdServeChat_DispatchesToResolvedAdapter(t *testing.T) { + reg := runtime.NewMuEdRegistry() + reg.Register(markerAdapter{}) + + mockRuntime := new(MockRuntime) + mockRuntime.On("Chat", mock.Anything, mock.Anything). + Return(chatRuntimeResponse("ASSISTANT", "ignored"), nil) + + req := httptest.NewRequest(http.MethodPost, "/chat", bytes.NewReader(chatRequestBody(t))) + req.Header.Set("X-Api-Version", "9.9.9") + w := httptest.NewRecorder() + + newMuEdHandlerWithRegistry(nil, mockRuntime, reg).ServeChat(w, req) + + res := w.Result() + defer res.Body.Close() + raw, _ := io.ReadAll(res.Body) + + assert.Equal(t, http.StatusOK, res.StatusCode) + assert.Equal(t, "9.9.9", res.Header.Get("X-Api-Version")) + + var resp map[string]any + require.NoError(t, json.Unmarshal(raw, &resp)) + assert.Equal(t, "chat-9.9.9", resp["marker"]) + + mockRuntime.AssertExpectations(t) +} diff --git a/handler/stream.go b/handler/stream.go new file mode 100644 index 0000000..b4fe283 --- /dev/null +++ b/handler/stream.go @@ -0,0 +1,154 @@ +package handler + +import ( + "context" + "net/http" + "sync" + "time" + + "go.uber.org/zap" + + "github.com/lambda-feedback/shimmy/internal/progress" + "github.com/lambda-feedback/shimmy/internal/server" +) + +// streamProgress runs a request whose progress is streamed back on the +// response as Server-Sent Events. It commits a 200 + event-stream headers +// immediately, wires an SSE reporter (fanned out to a callbackURL reporter +// too, if callbackURL is set) into ctx, keeps the connection alive with +// heartbeats while run executes, then emits exactly one terminal frame +// (completed | failed) built from run's result. Because the status is +// already committed, every outcome of run — including an internal error — +// becomes a "failed" frame, never an HTTP error. +// +// cmdLabel selects the terminal frame shape ("evaluate"/"preview" -> +// feedback[]; "chat" -> output/metadata). command is the µEd command +// string carried on the emitted progress events. doneMessage is the +// human-facing text on the terminal completed event. +// +// run returns the terminal event's Data payload (e.g. {"feedback": …} or +// {"output": …, "metadata": …}) on success, or a *terminalError. It must +// not write to w or emit progress events itself. +func (h *MuEdHandler) streamProgress( + ctx context.Context, + w http.ResponseWriter, + cmdLabel string, + command string, + doneMessage string, + version string, + callbackURL string, + requestID string, + run func(ctx context.Context) (map[string]any, *terminalError), +) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.Header().Set(muEdVersionHeader, version) + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + + sseReporter, err := progress.NewSSEReporter(w, cmdLabel, h.log) + if err != nil { + // Guarded against by the caller; don't panic if it slips through. + h.log.Error("failed to create SSE reporter", zap.Error(err)) + return + } + + var reporter progress.Reporter = sseReporter + if callbackURL != "" { + cbReporter, cbErr := h.progressFactory.NewReporter(callbackURL, requestID) + if cbErr != nil { + h.log.Warn("invalid callbackUrl, disabling callback delivery", zap.Error(cbErr)) + } else if cbReporter != nil { + reporter = progress.NewMultiReporter(sseReporter, cbReporter) + } + } + ctx = progress.ContextWithReporter(ctx, reporter) + + done := make(chan struct{}) + var hbWG sync.WaitGroup + if secs := h.config.Progress.Stream.HeartbeatSeconds; secs > 0 { + hbWG.Add(1) + go func() { + defer hbWG.Done() + ticker := time.NewTicker(time.Duration(secs) * time.Second) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ctx.Done(): + return + case <-ticker.C: + sseReporter.Heartbeat() + } + } + }() + } + + failTitle := "Evaluation failed" + if cmdLabel == "chat" { + failTitle = "Chat failed" + } + + data, termErr := run(ctx) + switch { + case termErr != nil: + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: command, + Message: termErr.userMessage, + Error: termErr.rawError, + ErrorInfo: termErr.progressErrorInfo(failTitle), + }) + case h.terminalFrameInvalid(cmdLabel, data): + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: command, + Message: "We couldn't produce a valid response. Please try again.", + Error: "SSE terminal payload failed OpenAPI validation", + ErrorInfo: &progress.ErrorInfo{ + Title: "Invalid response", + Message: "We couldn't produce a valid response. Please try again.", + Code: "INTERNAL_ERROR", + }, + }) + default: + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: command, + Message: doneMessage, + Data: data, + }) + } + + close(done) + hbWG.Wait() +} + +// terminalFrameInvalid reports whether the terminal frame's data payload +// fails the µEd response schema for the equivalent non-streaming body. It +// gives the streamed path the schema guarantee the buffered path gets +// from the OpenAPI response filter. A nil spec (e.g. under Lambda, which +// never streams) or an unmapped command reports valid. +func (h *MuEdHandler) terminalFrameInvalid(cmdLabel string, data map[string]any) bool { + var operationID string + var payload any + switch cmdLabel { + case "chat": + operationID, payload = "chat", data + case "evaluate": + operationID, payload = "evaluateSubmission", data["feedback"] + default: + // "preview" has no dedicated spec operation of its own. + return false + } + + if err := server.ValidateResponseBody(h.spec, operationID, payload); err != nil { + h.log.Error("SSE terminal payload failed OpenAPI validation", + zap.String("command", cmdLabel), zap.Error(err)) + return true + } + return false +} diff --git a/internal/execution/dispatcher.go b/internal/execution/dispatcher.go index 300ca3f..f1a562f 100644 --- a/internal/execution/dispatcher.go +++ b/internal/execution/dispatcher.go @@ -7,6 +7,7 @@ import ( "github.com/lambda-feedback/shimmy/internal/execution/dispatcher" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) type Dispatcher dispatcher.Dispatcher @@ -27,6 +28,10 @@ type Params struct { // Config is the config for the dispatcher and the underlying supervisors Config Config + // Progress configures worker-authored progress event delivery, + // passed through to the underlying supervisor(s). + Progress progress.Config + // Log is the logger to use for the dispatcher Log *zap.Logger } @@ -38,8 +43,9 @@ func NewDispatcher(params Params) (dispatcher.Dispatcher, error) { Config: dispatcher.DedicatedDispatcherConfig{ Supervisor: params.Config.Supervisor, }, - Context: params.Context, - Log: params.Log, + Context: params.Context, + Progress: params.Progress, + Log: params.Log, }, ) } else { @@ -49,8 +55,9 @@ func NewDispatcher(params Params) (dispatcher.Dispatcher, error) { Supervisor: params.Config.Supervisor, MaxWorkers: params.Config.MaxWorkers, }, - Context: params.Context, - Log: params.Log, + Context: params.Context, + Progress: params.Progress, + Log: params.Log, }, ) } diff --git a/internal/execution/dispatcher/dispatcher_dedicated.go b/internal/execution/dispatcher/dispatcher_dedicated.go index 2cb5223..842e647 100644 --- a/internal/execution/dispatcher/dispatcher_dedicated.go +++ b/internal/execution/dispatcher/dispatcher_dedicated.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) type DedicatedDispatcher struct { @@ -31,6 +32,10 @@ type DedicatedDispatcherParams struct { // SupervisorFactory is the factory function to create a new supervisor SupervisorFactory SupervisorFactory + // Progress configures worker-authored progress event delivery, + // passed through to the underlying supervisor. + Progress progress.Config + // Log is the logger to use for the dispatcher Log *zap.Logger } @@ -110,8 +115,9 @@ func createSupervisor( params DedicatedDispatcherParams, ) (supervisor.Supervisor, error) { return params.SupervisorFactory(supervisor.Params{ - Context: params.Context, - Config: params.Config.Supervisor, - Log: params.Log, + Context: params.Context, + Config: params.Config.Supervisor, + Progress: params.Progress, + Log: params.Log, }) } diff --git a/internal/execution/dispatcher/dispatcher_pooled.go b/internal/execution/dispatcher/dispatcher_pooled.go index 7a49429..967e26c 100644 --- a/internal/execution/dispatcher/dispatcher_pooled.go +++ b/internal/execution/dispatcher/dispatcher_pooled.go @@ -9,6 +9,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) type PooledDispatcher struct { @@ -37,6 +38,10 @@ type PooledDispatcherParams struct { // SupervisorFactory is the factory function to create a new supervisor SupervisorFactory SupervisorFactory + // Progress configures worker-authored progress event delivery, + // passed through to each pooled supervisor. + Progress progress.Config + // Log is the logger to use for the dispatcher Log *zap.Logger } @@ -157,9 +162,10 @@ func createPool( constructor := func(ctx context.Context) (supervisor.Supervisor, error) { sv, err := params.SupervisorFactory(supervisor.Params{ - Context: ctx, - Config: params.Config.Supervisor, - Log: params.Log, + Context: ctx, + Config: params.Config.Supervisor, + Progress: params.Progress, + Log: params.Log, }) if err != nil { return nil, err diff --git a/internal/execution/dispatcher/dispatcher_pooled_test.go b/internal/execution/dispatcher/dispatcher_pooled_test.go index 0fee760..723ec9a 100644 --- a/internal/execution/dispatcher/dispatcher_pooled_test.go +++ b/internal/execution/dispatcher/dispatcher_pooled_test.go @@ -3,7 +3,6 @@ package dispatcher_test import ( "context" "testing" - "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" @@ -210,8 +209,8 @@ func TestPooledDispatcher_Send_ReleaseSupervisorWaitErrorShutdown(t *testing.T) _, err := m.Send(context.Background(), "test", data) assert.NoError(t, err) - // wait for the release to happen in a goroutine - <-time.After(1 * time.Millisecond) + // wait for the background goroutine to finish by draining the pool + m.Shutdown(context.Background()) assert.Equal(t, 1, waited) } diff --git a/internal/execution/supervisor/adapter.go b/internal/execution/supervisor/adapter.go index e31eb13..7ec6803 100644 --- a/internal/execution/supervisor/adapter.go +++ b/internal/execution/supervisor/adapter.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) // AdapterWorkerFactoryFn is a type alias for a function that creates a worker @@ -42,20 +43,26 @@ type Adapter interface { // MARK: - factory -// defaultAdapterFactory is the default adapter factory -// that creates an adapter based on the given IO mode. -func defaultAdapterFactory( - workerFactory AdapterWorkerFactoryFn, - config IOConfig, - log *zap.Logger, -) (Adapter, error) { - switch config.Interface { - case FileIO: - return newFileAdapter(workerFactory, log), nil - case RpcIO: - return newRpcAdapter(workerFactory, config.Rpc, log), nil - default: - return nil, ErrUnsupportedIOInterface +// newDefaultAdapterFactory returns the default AdapterFactoryFn, wiring +// each created adapter's worker-authored progress side-channel (see +// internal/progress.Sidecar) with the given limits. It's a closure rather +// than a plain function so that AdapterFactoryFn's signature - and every +// test double built against it - doesn't need to carry progress.Config +// through every caller. +func newDefaultAdapterFactory(progressCfg progress.Config) AdapterFactoryFn { + return func( + workerFactory AdapterWorkerFactoryFn, + config IOConfig, + log *zap.Logger, + ) (Adapter, error) { + switch config.Interface { + case FileIO: + return newFileAdapter(workerFactory, progressCfg.Sidecar, log), nil + case RpcIO: + return newRpcAdapter(workerFactory, config.Rpc, progressCfg.Sidecar, log), nil + default: + return nil, ErrUnsupportedIOInterface + } } } diff --git a/internal/execution/supervisor/adapter_file.go b/internal/execution/supervisor/adapter_file.go index 7917f47..1bd2838 100644 --- a/internal/execution/supervisor/adapter_file.go +++ b/internal/execution/supervisor/adapter_file.go @@ -15,6 +15,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) // fileAdapter is an adapter that allows supervisors to use files to @@ -32,6 +33,9 @@ type fileAdapter struct { // worker is the worker that is managed by the adapter. worker worker.Worker + // sidecarCfg configures the worker-authored progress side-channel. + sidecarCfg progress.SidecarConfig + log *zap.Logger } @@ -39,10 +43,12 @@ var _ Adapter = (*fileAdapter)(nil) func newFileAdapter( workerFactory AdapterWorkerFactoryFn, + sidecarCfg progress.SidecarConfig, log *zap.Logger, ) *fileAdapter { return &fileAdapter{ workerFactory: workerFactory, + sidecarCfg: sidecarCfg, log: log.Named("adapter_file"), } } @@ -153,14 +159,26 @@ func (a *fileAdapter) Send( // ensure env is not nil if startParams.Env == nil { - startParams.Env = make([]string, 0, 3) + startParams.Env = make([]string, 0, 4) } + // the file interface is one process per request, so the sidecar is + // scoped entirely to this call - no Bind/Unbind swap needed, unlike + // the persistent rpcAdapter. + sidecar, err := progress.NewSidecar(a.sidecarCfg, a.log) + if err != nil { + return nil, fmt.Errorf("error starting progress sidecar: %w", err) + } + defer sidecar.Close() + + sidecar.Bind(method, progress.FromContext(ctx)) + // append req and res file names to worker env startParams.Env = append(startParams.Env, "EVAL_IO=FILE", "EVAL_FILE_NAME_REQUEST="+reqFile.Name(), "EVAL_FILE_NAME_RESPONSE="+resFile.Name(), + "EVAL_PROGRESS_URL="+sidecar.URL(), ) // create the worker with modified args and env diff --git a/internal/execution/supervisor/adapter_file_test.go b/internal/execution/supervisor/adapter_file_test.go index 380238b..6d02991 100644 --- a/internal/execution/supervisor/adapter_file_test.go +++ b/internal/execution/supervisor/adapter_file_test.go @@ -3,17 +3,54 @@ package supervisor import ( "context" "io" + "net/http" "os" "strings" + "sync" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) +// recordingReporter is a minimal progress.Reporter test double, local to +// this package since progress.Reporter's own test double is unexported +// in a different package. +type recordingReporter struct { + mu sync.Mutex + events []progress.Event +} + +func (r *recordingReporter) Report(_ context.Context, evt progress.Event) { + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, evt) +} + +func (r *recordingReporter) recorded() []progress.Event { + r.mu.Lock() + defer r.mu.Unlock() + return append([]progress.Event(nil), r.events...) +} + +// envValue returns the value of the first "key=value" entry in env, or "" +// if key isn't present. +func envValue(env []string, key string) string { + prefix := key + "=" + for _, e := range env { + if strings.HasPrefix(e, prefix) { + return strings.TrimPrefix(e, prefix) + } + } + return "" +} + func TestFileAdapter_Start_DoesNotStartWorker(t *testing.T) { a, w := createFileAdapter(t) @@ -132,6 +169,56 @@ func TestFileAdapter_Send_ReturnsInvalidDataError(t *testing.T) { w.AssertNotCalled(t, "Start") } +func TestFileAdapter_Send_InjectsProgressURLAndRelaysWorkerEvents(t *testing.T) { + w := worker.NewMockWorker(t) + + var sp *worker.StartConfig + workerFactory := func(params worker.StartConfig) (worker.Worker, error) { + sp = ¶ms + return w, nil + } + + a := &fileAdapter{ + workerFactory: workerFactory, + log: zap.NewNop(), + } + + r := &recordingReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + data := map[string]any{"foo": "bar"} + + w.EXPECT().Start(mock.Anything).RunAndReturn(func(ctx context.Context) error { + progressURL := envValue(sp.Env, "EVAL_PROGRESS_URL") + require.NotEmpty(t, progressURL, "expected EVAL_PROGRESS_URL in worker env") + + resp, err := http.Post(progressURL, "application/json", strings.NewReader(`{"message":"checking correctness"}`)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusAccepted, resp.StatusCode) + + requestFileName := sp.Args[len(sp.Args)-2] + responseFileName := sp.Args[len(sp.Args)-1] + reqData, _ := os.ReadFile(requestFileName) + _ = os.WriteFile(responseFileName, reqData, os.ModeAppend) + return nil + }) + w.EXPECT().ReadPipe().Return(io.NopCloser(strings.NewReader("")), nil) + var cell int + w.EXPECT().Wait(mock.Anything).Return(worker.ExitEvent{Code: &cell}, nil) + + _, err := a.Send(ctx, "eval", data, 10) + require.NoError(t, err) + + assert.Eventually(t, func() bool { + return len(r.recorded()) == 1 + }, time.Second, 5*time.Millisecond, "expected the worker's progress event to be relayed") + + events := r.recorded() + assert.Equal(t, progress.StageEvaluating, events[0].Stage) + assert.Equal(t, "eval", events[0].Command) + assert.Equal(t, "checking correctness", events[0].Message) +} + func createFileAdapter(t *testing.T) (*fileAdapter, *worker.MockWorker) { w := worker.NewMockWorker(t) diff --git a/internal/execution/supervisor/adapter_rpc.go b/internal/execution/supervisor/adapter_rpc.go index 837c5c3..184e22b 100644 --- a/internal/execution/supervisor/adapter_rpc.go +++ b/internal/execution/supervisor/adapter_rpc.go @@ -14,6 +14,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) // RpcConfig describes the configuration for the rpc interface. @@ -91,16 +92,27 @@ type rpcAdapter struct { config RpcConfig log *zap.Logger + + // sidecarCfg configures the worker-authored progress side-channel. + sidecarCfg progress.SidecarConfig + + // sidecar is the loopback HTTP listener for worker-authored progress + // events, injected into the worker's env as EVAL_PROGRESS_URL. It + // lives for this adapter's whole lifetime (one persistent worker can + // serve many requests), and is Bind/Unbind-ed around each Send call. + sidecar *progress.Sidecar } func newRpcAdapter( workerFactory AdapterWorkerFactoryFn, config RpcConfig, + sidecarCfg progress.SidecarConfig, log *zap.Logger, ) *rpcAdapter { return &rpcAdapter{ workerFactory: workerFactory, config: config, + sidecarCfg: sidecarCfg, log: log.Named("adapter_rpc"), } } @@ -113,7 +125,13 @@ func (a *rpcAdapter) Start( return errors.New("no worker factory provided") } - params.Env = buildEnv(params.Env, a.config) + sidecar, err := progress.NewSidecar(a.sidecarCfg, a.log) + if err != nil { + return fmt.Errorf("error starting progress sidecar: %w", err) + } + a.sidecar = sidecar + + params.Env = buildEnv(params.Env, a.config, sidecar.URL()) // create the worker worker, err := a.workerFactory(params) @@ -164,6 +182,18 @@ func (a *rpcAdapter) Send( return nil, errors.New("rpc client not available") } + if a.sidecar != nil { + // sendLock in the calling supervisor guarantees only one request + // is ever in flight per worker; UnbindAfterGrace closes the window + // between this call returning and the next one starting (after a + // short grace period, to give a fire-and-forget progress POST the + // worker dispatched just before returning its result a chance to + // still land), so a straggling POST can't be misattributed to an + // unrelated future request. + a.sidecar.Bind(method, progress.FromContext(ctx)) + defer a.sidecar.UnbindAfterGrace() + } + var result map[string]any ctx, cancel := context.WithTimeout(ctx, timeout) @@ -181,6 +211,12 @@ func (a *rpcAdapter) Stop() (ReleaseFunc, error) { return nil, errors.New("no worker provided") } + if a.sidecar != nil { + if err := a.sidecar.Close(); err != nil { + a.log.Warn("error closing progress sidecar", zap.Error(err)) + } + } + return stopWorker(a.worker) } @@ -285,7 +321,7 @@ func getIPCEndpoint(config IpcTransportConfig) string { } } -func buildEnv(env []string, config RpcConfig) []string { +func buildEnv(env []string, config RpcConfig, progressURL string) []string { if env == nil { env = make([]string, 0) } @@ -306,6 +342,10 @@ func buildEnv(env []string, config RpcConfig) []string { env = append(env, "EVAL_RPC_TCP_ADDRESS="+config.Tcp.Address) } + if progressURL != "" { + env = append(env, "EVAL_PROGRESS_URL="+progressURL) + } + return env } diff --git a/internal/execution/supervisor/adapter_rpc_test.go b/internal/execution/supervisor/adapter_rpc_test.go index 2ac8860..be1cbee 100644 --- a/internal/execution/supervisor/adapter_rpc_test.go +++ b/internal/execution/supervisor/adapter_rpc_test.go @@ -4,13 +4,17 @@ import ( "bytes" "context" "io" + "net/http" + "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) type rwc struct { @@ -38,6 +42,14 @@ func createRpcAdapter(t *testing.T) (*rpcAdapter, *worker.MockWorker) { config: RpcConfig{Transport: StdioTransport}, } + // Start (called by most tests using this helper) always spins up a + // real progress sidecar listener; close it so tests don't leak ports. + t.Cleanup(func() { + if adapter.sidecar != nil { + adapter.sidecar.Close() + } + }) + return adapter, w } @@ -142,6 +154,60 @@ func TestStdioAdapter_Stop_WaitForError(t *testing.T) { assert.ErrorIs(t, err, assert.AnError) } +func TestStdioAdapter_Start_InjectsProgressURL(t *testing.T) { + a, w := createRpcAdapter(t) + + var sp *worker.StartConfig + baseFactory := a.workerFactory + a.workerFactory = func(params worker.StartConfig) (worker.Worker, error) { + sp = ¶ms + return baseFactory(params) + } + + w.EXPECT().DuplexPipe().Return(newRwc(), nil) + w.EXPECT().Start(mock.Anything).Return(nil) + + err := a.Start(context.Background(), worker.StartConfig{}) + assert.NoError(t, err) + + assert.Contains(t, sp.Env, "EVAL_PROGRESS_URL="+a.sidecar.URL()) +} + +// TestStdioAdapter_Send_RelaysWorkerProgressEvents exercises the same +// Bind/UnbindAfterGrace path Send uses around the (separately, more fully) +// tested Sidecar, without needing a live RPC round trip - Send itself isn't +// otherwise exercised in this file (see the disabled tests below). +func TestStdioAdapter_Send_RelaysWorkerProgressEvents(t *testing.T) { + a, w := createRpcAdapter(t) + + w.EXPECT().DuplexPipe().Return(newRwc(), nil) + w.EXPECT().Start(mock.Anything).Return(nil) + + err := a.Start(context.Background(), worker.StartConfig{}) + assert.NoError(t, err) + + r := &recordingReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + // mirrors exactly what rpcAdapter.Send does with a.sidecar + a.sidecar.Bind("eval", progress.FromContext(ctx)) + defer a.sidecar.UnbindAfterGrace() + + resp, err := http.Post(a.sidecar.URL(), "application/json", strings.NewReader(`{"message":"checking correctness"}`)) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusAccepted, resp.StatusCode) + + assert.Eventually(t, func() bool { + return len(r.recorded()) == 1 + }, time.Second, 5*time.Millisecond, "expected the worker's progress event to be relayed") + + events := r.recorded() + assert.Equal(t, progress.StageEvaluating, events[0].Stage) + assert.Equal(t, "eval", events[0].Command) + assert.Equal(t, "checking correctness", events[0].Message) +} + // func TestStdioAdapter_Send(t *testing.T) { // a, w := createStdioAdapter(t) diff --git a/internal/execution/supervisor/adapter_test.go b/internal/execution/supervisor/adapter_test.go index c426f7a..1374d1b 100644 --- a/internal/execution/supervisor/adapter_test.go +++ b/internal/execution/supervisor/adapter_test.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) func TestDefaultAdapterFactory(t *testing.T) { @@ -16,9 +17,11 @@ func TestDefaultAdapterFactory(t *testing.T) { return w, nil } + factory := newDefaultAdapterFactory(progress.Config{}) + cases := []IOConfig{{Interface: FileIO}, {Interface: RpcIO}} for _, mode := range cases { - _, err := defaultAdapterFactory(workerFactory, mode, zap.NewNop()) + _, err := factory(workerFactory, mode, zap.NewNop()) assert.NoError(t, err) } @@ -31,7 +34,7 @@ func TestDefaultAdapterFactory_Fails(t *testing.T) { return w, nil } - _, err := defaultAdapterFactory(workerFactory, IOConfig{Interface: ""}, zap.NewNop()) + _, err := newDefaultAdapterFactory(progress.Config{})(workerFactory, IOConfig{Interface: ""}, zap.NewNop()) assert.ErrorIs(t, err, ErrUnsupportedIOInterface) } diff --git a/internal/execution/supervisor/supervisor.go b/internal/execution/supervisor/supervisor.go index f3e6587..b316026 100644 --- a/internal/execution/supervisor/supervisor.go +++ b/internal/execution/supervisor/supervisor.go @@ -9,6 +9,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) type Supervisor interface { @@ -74,6 +75,11 @@ type Params struct { // is called when the supervisor needs to create a new worker. WorkerFactory WorkerFactoryFn + // Progress configures worker-authored progress event delivery (the + // EVAL_PROGRESS_URL side-channel). Only used when AdapterFactory is + // nil, since the default adapter factory is what wires it up. + Progress progress.Config + // Log is the logger to use for the supervisor Log *zap.Logger } @@ -99,7 +105,7 @@ func New(params Params) (Supervisor, error) { } if params.AdapterFactory == nil { - params.AdapterFactory = defaultAdapterFactory + params.AdapterFactory = newDefaultAdapterFactory(params.Progress) } createAdapter := func() (*workerRef, error) { @@ -165,12 +171,49 @@ func (s *WorkerSupervisor) Send( worker, err := s.acquireWorker(ctx) if err != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: method, + Message: "We couldn't start the request. Please try again.", + Error: err.Error(), + ErrorInfo: &progress.ErrorInfo{ + Title: "Request failed", + Message: "We couldn't start the request. Please try again.", + Code: "INTERNAL_ERROR", + Trace: err.Error(), + }, + }) return nil, fmt.Errorf("failed to acquire worker: %w", err) + } + progress.Emit(ctx, progress.Event{ + Stage: progress.StagePreparing, + Command: method, + Message: "Preparing…", + }) // NOTICE: unconventional error handling ahead, as we need // to release the worker before returning the error. + progress.Emit(ctx, progress.Event{ + Stage: progress.StageStarting, + Command: method, + Message: "Starting…", + }) resData, err := worker.Send(ctx, method, data, s.sendParams.Timeout) + if err != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: method, + Message: "Something went wrong. Please try again.", + Error: err.Error(), + ErrorInfo: &progress.ErrorInfo{ + Title: "Request failed", + Message: "Something went wrong. Please try again.", + Code: "INTERNAL_ERROR", + Trace: err.Error(), + }, + }) + } release, releaseErr := s.releaseWorker() if releaseErr != nil { diff --git a/internal/execution/supervisor/supervisor_test.go b/internal/execution/supervisor/supervisor_test.go index 82bdb95..4357442 100644 --- a/internal/execution/supervisor/supervisor_test.go +++ b/internal/execution/supervisor/supervisor_test.go @@ -9,6 +9,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) func TestSupervisor_New_DefaultWorkerFactory(t *testing.T) { @@ -280,6 +281,101 @@ func TestSupervisor_Send_Fails(t *testing.T) { assert.NotNil(t, res) } +// MARK: - progress + +func TestSupervisor_Send_EmitsWorkerAcquiredAndRunning(t *testing.T) { + s, a, err := createSupervisor(t, supervisor.RpcIO) + assert.NoError(t, err) + + data := map[string]any{"data": "data"} + resData := map[string]any{"result": "result"} + + a.EXPECT().Start(mock.Anything, mock.Anything).Return(nil) + a.EXPECT().Send(mock.Anything, "test", data, mock.Anything).Return(resData, nil) + + r := &fakeReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + _, err = s.Send(ctx, "test", data) + assert.NoError(t, err) + + assert.Equal(t, []progress.Stage{ + progress.StagePreparing, + progress.StageStarting, + }, r.stages()) +} + +func TestSupervisor_Send_EmitsFailed_WhenAcquireFails(t *testing.T) { + mockFactory := func(supervisor.AdapterWorkerFactoryFn, supervisor.IOConfig, *zap.Logger) (supervisor.Adapter, error) { + return nil, assert.AnError + } + + s, err := createSupervisorWithFactory(supervisor.RpcIO, mockFactory) + assert.NoError(t, err) + + r := &fakeReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + data := map[string]any{"data": "data"} + _, err = s.Send(ctx, "test", data) + assert.ErrorIs(t, err, assert.AnError) + + assert.Equal(t, []progress.Stage{progress.StageFailed}, r.stages()) +} + +func TestSupervisor_Send_EmitsFailed_WhenWorkerSendFails(t *testing.T) { + s, a, err := createSupervisor(t, supervisor.RpcIO) + assert.NoError(t, err) + + data := map[string]any{"data": "data"} + + a.EXPECT().Start(mock.Anything, mock.Anything).Return(nil) + a.EXPECT().Send(mock.Anything, "test", data, mock.Anything).Return(nil, assert.AnError) + + r := &fakeReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + _, err = s.Send(ctx, "test", data) + assert.ErrorIs(t, err, assert.AnError) + + assert.Equal(t, []progress.Stage{ + progress.StagePreparing, + progress.StageStarting, + progress.StageFailed, + }, r.stages()) +} + +func TestSupervisor_Send_NoReporterInContext_BehavesUnchanged(t *testing.T) { + s, a, err := createSupervisor(t, supervisor.RpcIO) + assert.NoError(t, err) + + data := map[string]any{"data": "data"} + resData := map[string]any{"result": "result"} + + a.EXPECT().Start(mock.Anything, mock.Anything).Return(nil) + a.EXPECT().Send(mock.Anything, "test", data, mock.Anything).Return(resData, nil) + + res, err := s.Send(context.Background(), "test", data) + assert.NoError(t, err) + assert.Equal(t, resData, res.Data) +} + +type fakeReporter struct { + events []progress.Event +} + +func (r *fakeReporter) Report(_ context.Context, evt progress.Event) { + r.events = append(r.events, evt) +} + +func (r *fakeReporter) stages() []progress.Stage { + stages := make([]progress.Stage, len(r.events)) + for i, evt := range r.events { + stages[i] = evt.Stage + } + return stages +} + // MARK: - mocks func createSupervisor(t *testing.T, mode supervisor.IOInterface) ( diff --git a/internal/progress/event.go b/internal/progress/event.go new file mode 100644 index 0000000..28064f8 --- /dev/null +++ b/internal/progress/event.go @@ -0,0 +1,99 @@ +package progress + +import "time" + +// Stage identifies a point in the lifecycle of an evaluation request that +// progress events can be emitted for. +type Stage string + +const ( + // StagePreparing indicates the evaluation environment is being set up + // (a worker is ready to receive work, whether freshly booted or reused + // from a warm pool). Deliberately named around what a student or + // teacher would recognise, not shimmy's internal "worker" concept. + // Emitted by shimmy itself, once per request. + StagePreparing Stage = "preparing" + + // StageStarting indicates the worker is about to be invoked. Emitted by + // shimmy itself, once per request, for both /evaluate and /chat. + StageStarting Stage = "starting" + + // StageEvaluating indicates a worker-authored sub-step during an + // /evaluate (or /preview) request, relayed from the worker's local + // progress side-channel (see Sidecar). A worker cannot claim any stage; + // the sidecar assigns this based on the command in flight. + StageEvaluating Stage = "evaluating" + + // StageThinking indicates a worker-authored sub-step during a /chat + // request, relayed from the worker's local progress side-channel (see + // Sidecar). Like StageEvaluating, the sidecar assigns it by command; + // the worker cannot set it. + StageThinking Stage = "thinking" + + // StageCompleted indicates feedback has been computed and is about + // to be returned to the caller. + StageCompleted Stage = "completed" + + // StageFailed indicates a terminal failure at any layer of the pipeline. + StageFailed Stage = "failed" + + // StageProgress is retained for compatibility but is no longer emitted: + // worker-authored sub-steps are now relayed as StageEvaluating or + // StageThinking depending on the command in flight (see Sidecar). + StageProgress Stage = "progress" +) + +// terminal reports whether the stage marks the end of an evaluation's +// progress event stream. At most one terminal event is delivered per +// Reporter instance. +func (s Stage) terminal() bool { + return s == StageCompleted || s == StageFailed +} + +// ErrorInfo is structured failure detail for a StageFailed event. On the +// SSE terminal "failed" frame it is emitted as the frame's `error` +// object, shaped like the µEd spec's ErrorResponse (title is required; +// the rest are optional). It carries no student/teacher-facing copy — +// that stays on Event.Message. +type ErrorInfo struct { + Title string `json:"title"` + Message string `json:"message,omitempty"` + Code string `json:"code,omitempty"` + Trace string `json:"trace,omitempty"` + Details map[string]any `json:"details,omitempty"` +} + +// Event describes a single progress update for an evaluation request. +type Event struct { + // Stage is the lifecycle point this event describes. + Stage Stage + + // Command is the µEd command being processed (e.g. "eval", "preview"). + Command string + + // Message is a short, student/teacher-facing description of this + // event, safe to display as-is (e.g. "Evaluating your submission…"). + // It must never contain raw technical error detail — see Error. + Message string + + // Error carries raw technical error detail for StageFailed events, + // intended for logs/support diagnostics. Never display this to + // students or teachers directly; show Message instead. + Error string + + // ErrorInfo is the structured failure detail for a StageFailed event. + // The SSE reporter emits it as the terminal "failed" frame's `error` + // object; when nil it falls back to a minimal object built from + // Message/Error. Ignored by non-SSE reporters. + ErrorInfo *ErrorInfo + + // Data is a free-form extension point. On StageCompleted it carries + // the final result payload (so a callbackUrl-supplying caller gets the + // result, not just a status ping). On a worker-authored sub-step + // (StageEvaluating / StageThinking) it carries whatever the evaluation + // function attached to its custom event (see Sidecar). + Data map[string]any + + // Timestamp is set by Emit, not by callers. + Timestamp time.Time +} diff --git a/internal/progress/factory.go b/internal/progress/factory.go new file mode 100644 index 0000000..7e87579 --- /dev/null +++ b/internal/progress/factory.go @@ -0,0 +1,131 @@ +package progress + +import ( + "fmt" + "net/http" + "net/url" + "time" + + "go.uber.org/fx" + "go.uber.org/zap" +) + +// defaultCallbackTimeout is used when Config.CallbackTimeout is unset. +const defaultCallbackTimeout = time.Second + +// Config is the configuration for outbound progress-callback delivery. +type Config struct { + // CallbackTimeout bounds a single progress callback POST. If unset + // (or <= 0), defaultCallbackTimeout is used. + CallbackTimeout time.Duration `conf:"callback_timeout"` + + // AllowedHosts, if non-empty, restricts callback URLs to these hosts. + // Entries may be an exact hostname (e.g. "api.example.com") or a + // "*.example.com" wildcard matching any subdomain. Empty means no + // host restriction — callback delivery is still subject to the + // private-network protection below. + AllowedHosts []string `conf:"allowed_hosts"` + + // AllowPrivateNetworks disables the default SSRF protection that + // refuses to dial loopback, link-local (including cloud metadata + // endpoints such as 169.254.169.254), and private (RFC1918/RFC4193) + // IP addresses, however the URL's hostname resolves. Only enable + // this if shimmy's callback targets are known to live on a private + // network you trust (e.g. a same-VPC service). + AllowPrivateNetworks bool `conf:"allow_private_networks"` + + // Sidecar bounds worker-authored progress events delivered via the + // EVAL_PROGRESS_URL side-channel (see sidecar.go), before they're + // relayed through the same outbound delivery path as shim-authored + // events. + Sidecar SidecarConfig `conf:"sidecar"` + + // Stream configures in-band SSE delivery of progress for /evaluate + // requests that send "Accept: text/event-stream" (see sse_reporter.go). + // Only effective in standalone/serve mode; Lambda cannot stream. + Stream StreamConfig `conf:"stream"` +} + +// StreamConfig configures in-band Server-Sent Events progress delivery. +type StreamConfig struct { + // Enabled turns SSE streaming on. When false, the "Accept: + // text/event-stream" request header is ignored and /evaluate serves + // its normal buffered JSON response. + Enabled bool `conf:"enabled"` + + // HeartbeatSeconds is the spacing between SSE heartbeat comments sent + // while an evaluation runs, so an idle connection isn't dropped by an + // intermediary. 0 disables heartbeats. + HeartbeatSeconds int `conf:"heartbeat_seconds"` +} + +// Factory builds a per-request Reporter from caller-supplied callback +// coordinates. +type Factory interface { + // NewReporter returns a Reporter that delivers events to callbackURL, + // tagging each with correlationID. If callbackURL is empty, it returns + // (nil, nil) — the signal that progress reporting is disabled for this + // request. An error is returned only when callbackURL is non-empty but + // invalid. + NewReporter(callbackURL, correlationID string) (Reporter, error) +} + +type HTTPFactoryParams struct { + fx.In + + Config Config + Log *zap.Logger +} + +type HTTPFactory struct { + client *http.Client + timeout time.Duration + log *zap.Logger + allowedHosts []string +} + +var _ Factory = (*HTTPFactory)(nil) + +// NewHTTPFactory builds a Factory that delivers progress events as +// outbound HTTP POST requests. +// +// Since the callback URL is caller-supplied, delivery is guarded against +// SSRF by default: the underlying transport refuses to dial loopback, +// link-local, or private IP addresses (see Config.AllowPrivateNetworks), +// and Config.AllowedHosts can further restrict which hostnames are +// accepted at all. +func NewHTTPFactory(params HTTPFactoryParams) Factory { + timeout := params.Config.CallbackTimeout + if timeout <= 0 { + timeout = defaultCallbackTimeout + } + + client := &http.Client{} + if !params.Config.AllowPrivateNetworks { + client.Transport = newSSRFGuardedTransport() + } + + return &HTTPFactory{ + client: client, + timeout: timeout, + log: params.Log, + allowedHosts: params.Config.AllowedHosts, + } +} + +func (f *HTTPFactory) NewReporter(callbackURL, correlationID string) (Reporter, error) { + if callbackURL == "" { + return nil, nil + } + + u, err := url.ParseRequestURI(callbackURL) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") { + return nil, fmt.Errorf("invalid progress callback url: %q", callbackURL) + } + + if len(f.allowedHosts) > 0 && !hostAllowed(u.Hostname(), f.allowedHosts) { + return nil, fmt.Errorf("progress callback host %q is not in the allowed hosts list", u.Hostname()) + } + + return newHTTPReporter(f.client, callbackURL, correlationID, f.timeout, f.log.Named("progress")), nil +} diff --git a/internal/progress/factory_test.go b/internal/progress/factory_test.go new file mode 100644 index 0000000..3e6e095 --- /dev/null +++ b/internal/progress/factory_test.go @@ -0,0 +1,67 @@ +package progress + +import ( + "testing" + "time" + + "go.uber.org/zap" +) + +func newTestFactory() *HTTPFactory { + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{CallbackTimeout: time.Second}, + Log: zap.NewNop(), + }) + return f.(*HTTPFactory) +} + +func TestHTTPFactory_NewReporter_EmptyURL_ReturnsNilReporterNoError(t *testing.T) { + f := newTestFactory() + + r, err := f.NewReporter("", "corr-1") + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if r != nil { + t.Fatalf("expected nil reporter for empty callback url, got %v", r) + } +} + +func TestHTTPFactory_NewReporter_ValidURL_ReturnsReporter(t *testing.T) { + f := newTestFactory() + + r, err := f.NewReporter("https://example.com/callback", "corr-1") + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if r == nil { + t.Fatalf("expected non-nil reporter for valid url") + } +} + +func TestHTTPFactory_NewReporter_InvalidURL_ReturnsError(t *testing.T) { + f := newTestFactory() + + cases := []string{ + "not-a-url", + "ftp://example.com/callback", + "://broken", + } + + for _, c := range cases { + if _, err := f.NewReporter(c, "corr-1"); err == nil { + t.Errorf("expected error for callback url %q, got nil", c) + } + } +} + +func TestNewHTTPFactory_DefaultsTimeoutWhenUnset(t *testing.T) { + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{}, + Log: zap.NewNop(), + }).(*HTTPFactory) + + if f.timeout != defaultCallbackTimeout { + t.Errorf("expected default timeout %v, got %v", defaultCallbackTimeout, f.timeout) + } +} diff --git a/internal/progress/http_reporter.go b/internal/progress/http_reporter.go new file mode 100644 index 0000000..5f47e2e --- /dev/null +++ b/internal/progress/http_reporter.go @@ -0,0 +1,122 @@ +package progress + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "sync" + "time" + + "go.uber.org/zap" +) + +// payload is the JSON body POSTed to the callback URL for each event. On +// a StageFailed event, error is an ErrorResponse-shaped object identical +// to the one on the SSE terminal "failed" frame, so a callbackUrl +// consumer and an SSE consumer handle failure the same way. +type payload struct { + CorrelationID string `json:"correlationId"` + Stage Stage `json:"stage"` + Command string `json:"command,omitempty"` + Message string `json:"message,omitempty"` + Error *ErrorInfo `json:"error,omitempty"` + Data map[string]any `json:"data,omitempty"` + Timestamp time.Time `json:"timestamp"` +} + +// httpCallbackReporter delivers progress events as outbound HTTP POST +// requests to a caller-supplied URL. +type httpCallbackReporter struct { + client *http.Client + url string + correlationID string + timeout time.Duration + log *zap.Logger + + terminalOnce sync.Once +} + +var _ Reporter = (*httpCallbackReporter)(nil) + +func newHTTPReporter( + client *http.Client, + url string, + correlationID string, + timeout time.Duration, + log *zap.Logger, +) Reporter { + return &httpCallbackReporter{ + client: client, + url: url, + correlationID: correlationID, + timeout: timeout, + log: log, + } +} + +// Report POSTs evt to the configured callback URL. Delivery is best-effort: +// any error (invalid payload, dial failure, timeout, non-2xx response) is +// logged and swallowed — it must never fail or slow down the evaluation +// beyond the configured timeout. At most one terminal event (StageFailed +// or StageCompleted) is delivered per reporter instance, since both +// the supervisor and handler layers can independently detect failure. +func (r *httpCallbackReporter) Report(ctx context.Context, evt Event) { + if evt.Stage.terminal() { + sent := false + r.terminalOnce.Do(func() { + r.send(ctx, evt) + sent = true + }) + if !sent { + r.log.Debug("dropping duplicate terminal progress event", zap.String("stage", string(evt.Stage))) + } + return + } + + r.send(ctx, evt) +} + +func (r *httpCallbackReporter) send(ctx context.Context, evt Event) { + var errInfo *ErrorInfo + if evt.Stage == StageFailed { + errInfo = failureErrorInfo(evt) + } + body, err := json.Marshal(payload{ + CorrelationID: r.correlationID, + Stage: evt.Stage, + Command: evt.Command, + Message: evt.Message, + Error: errInfo, + Data: evt.Data, + Timestamp: evt.Timestamp, + }) + if err != nil { + r.log.Warn("failed to marshal progress event", zap.String("stage", string(evt.Stage)), zap.Error(err)) + return + } + + ctx, cancel := context.WithTimeout(ctx, r.timeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.url, bytes.NewReader(body)) + if err != nil { + r.log.Warn("failed to build progress callback request", zap.String("stage", string(evt.Stage)), zap.Error(err)) + return + } + req.Header.Set("Content-Type", "application/json") + + resp, err := r.client.Do(req) + if err != nil { + r.log.Warn("progress callback delivery failed", zap.String("stage", string(evt.Stage)), zap.Error(err)) + return + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + r.log.Warn("progress callback returned non-2xx status", + zap.String("stage", string(evt.Stage)), + zap.Int("status", resp.StatusCode), + ) + } +} diff --git a/internal/progress/http_reporter_test.go b/internal/progress/http_reporter_test.go new file mode 100644 index 0000000..efd07fa --- /dev/null +++ b/internal/progress/http_reporter_test.go @@ -0,0 +1,185 @@ +package progress + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "go.uber.org/zap" +) + +func newTestReporter(t *testing.T, url string, timeout time.Duration) *httpCallbackReporter { + t.Helper() + return newHTTPReporter(&http.Client{}, url, "corr-1", timeout, zap.NewNop()).(*httpCallbackReporter) +} + +func TestHTTPCallbackReporter_Report_DeliversPayload(t *testing.T) { + var mu sync.Mutex + var received []payload + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var p payload + if err := json.NewDecoder(r.Body).Decode(&p); err != nil { + t.Errorf("failed to decode payload: %v", err) + } + mu.Lock() + received = append(received, p) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + r.Report(context.Background(), Event{Stage: StageEvaluating, Command: "eval"}) + + mu.Lock() + defer mu.Unlock() + if len(received) != 1 { + t.Fatalf("expected 1 request, got %d", len(received)) + } + if received[0].CorrelationID != "corr-1" { + t.Errorf("expected correlationId %q, got %q", "corr-1", received[0].CorrelationID) + } + if received[0].Stage != StageEvaluating { + t.Errorf("expected stage %q, got %q", StageEvaluating, received[0].Stage) + } +} + +func TestHTTPCallbackReporter_Report_FailedEventCarriesStructuredError(t *testing.T) { + var mu sync.Mutex + var received []payload + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var p payload + if err := json.NewDecoder(r.Body).Decode(&p); err != nil { + t.Errorf("failed to decode payload: %v", err) + } + mu.Lock() + received = append(received, p) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + r.Report(context.Background(), Event{ + Stage: StageFailed, + Message: "We couldn't evaluate your answer.", + Error: "worker exited 1", + ErrorInfo: &ErrorInfo{Title: "Evaluation failed", Message: "We couldn't evaluate your answer.", Code: "INTERNAL_ERROR", Trace: "worker exited 1"}, + }) + + mu.Lock() + defer mu.Unlock() + if len(received) != 1 { + t.Fatalf("expected 1 request, got %d", len(received)) + } + got := received[0].Error + if got == nil { + t.Fatalf("expected a structured error object on the failed callback payload") + } + if got.Title != "Evaluation failed" || got.Code != "INTERNAL_ERROR" || got.Trace != "worker exited 1" { + t.Errorf("error object not carried through: %+v", got) + } +} + +func TestHTTPCallbackReporter_Report_NonFailedEventHasNoError(t *testing.T) { + var mu sync.Mutex + var received []payload + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var p payload + _ = json.NewDecoder(r.Body).Decode(&p) + mu.Lock() + received = append(received, p) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + r.Report(context.Background(), Event{Stage: StageCompleted, Message: "done"}) + + mu.Lock() + defer mu.Unlock() + if len(received) != 1 || received[0].Error != nil { + t.Fatalf("expected no error object on a non-failed event, got %+v", received) + } +} + +func TestHTTPCallbackReporter_Report_TerminalEventDeliveredOnlyOnce(t *testing.T) { + var mu sync.Mutex + var count int + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + count++ + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + + // simulate both the supervisor and handler layers independently + // detecting failure and trying to emit a terminal event + r.Report(context.Background(), Event{Stage: StageFailed, Message: "boot failed"}) + r.Report(context.Background(), Event{Stage: StageFailed, Message: "handler backstop"}) + r.Report(context.Background(), Event{Stage: StageCompleted}) + + mu.Lock() + defer mu.Unlock() + if count != 1 { + t.Fatalf("expected exactly 1 terminal event delivered, got %d", count) + } +} + +func TestHTTPCallbackReporter_Report_NonTerminalEventsAllDelivered(t *testing.T) { + var mu sync.Mutex + var count int + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + count++ + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + r.Report(context.Background(), Event{Stage: StagePreparing}) + r.Report(context.Background(), Event{Stage: StageEvaluating}) + + mu.Lock() + defer mu.Unlock() + if count != 2 { + t.Fatalf("expected 2 non-terminal events delivered, got %d", count) + } +} + +func TestHTTPCallbackReporter_Report_SlowReceiver_BoundedByTimeout(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(200 * time.Millisecond) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, 20*time.Millisecond) + + start := time.Now() + r.Report(context.Background(), Event{Stage: StageEvaluating}) + elapsed := time.Since(start) + + if elapsed > 150*time.Millisecond { + t.Errorf("expected Report to return promptly bounded by timeout, took %v", elapsed) + } +} + +func TestHTTPCallbackReporter_Report_UnreachableURL_DoesNotPanic(t *testing.T) { + r := newTestReporter(t, "http://127.0.0.1:0", 50*time.Millisecond) + r.Report(context.Background(), Event{Stage: StageEvaluating}) +} diff --git a/internal/progress/multi_reporter.go b/internal/progress/multi_reporter.go new file mode 100644 index 0000000..8ae0046 --- /dev/null +++ b/internal/progress/multi_reporter.go @@ -0,0 +1,30 @@ +package progress + +import "context" + +// multiReporter fans a single event out to several reporters. It's used +// when a request both opens an SSE stream and supplies a callbackUrl: +// each underlying reporter keeps its own terminal-once guard, so the +// fan-out needs no extra state. +type multiReporter struct { + reporters []Reporter +} + +var _ Reporter = (*multiReporter)(nil) + +// NewMultiReporter returns a Reporter that delivers each event to every +// reporter in rs, in order. A reporter that panics or blocks must not +// prevent the others from receiving the event, nor propagate out to the +// evaluation goroutine. +func NewMultiReporter(rs ...Reporter) Reporter { + return &multiReporter{reporters: rs} +} + +func (m *multiReporter) Report(ctx context.Context, evt Event) { + for _, r := range m.reporters { + func() { + defer func() { _ = recover() }() + r.Report(ctx, evt) + }() + } +} diff --git a/internal/progress/multi_reporter_test.go b/internal/progress/multi_reporter_test.go new file mode 100644 index 0000000..72519b9 --- /dev/null +++ b/internal/progress/multi_reporter_test.go @@ -0,0 +1,45 @@ +package progress + +import ( + "context" + "testing" +) + +type panicReporter struct{ called bool } + +func (p *panicReporter) Report(context.Context, Event) { + p.called = true + panic("boom") +} + +func TestMultiReporter_FansOutInOrder(t *testing.T) { + a := &recordingReporter{} + b := &recordingReporter{} + m := NewMultiReporter(a, b) + + m.Report(context.Background(), Event{Stage: StagePreparing}) + m.Report(context.Background(), Event{Stage: StageCompleted}) + + for name, r := range map[string]*recordingReporter{"a": a, "b": b} { + evts := r.recorded() + if len(evts) != 2 || evts[0].Stage != StagePreparing || evts[1].Stage != StageCompleted { + t.Errorf("reporter %s: expected both events in order, got %v", name, evts) + } + } +} + +func TestMultiReporter_ChildPanicIsolated(t *testing.T) { + p := &panicReporter{} + b := &recordingReporter{} + m := NewMultiReporter(p, b) + + // must not panic out to the caller + m.Report(context.Background(), Event{Stage: StageEvaluating}) + + if !p.called { + t.Error("expected the panicking reporter to have been called") + } + if evts := b.recorded(); len(evts) != 1 || evts[0].Stage != StageEvaluating { + t.Errorf("expected the second reporter to still receive the event, got %v", evts) + } +} diff --git a/internal/progress/reporter.go b/internal/progress/reporter.go new file mode 100644 index 0000000..e21f051 --- /dev/null +++ b/internal/progress/reporter.go @@ -0,0 +1,44 @@ +package progress + +import ( + "context" + "time" +) + +// Reporter delivers progress events for a single evaluation request. +type Reporter interface { + // Report emits a single event. Implementations MUST NOT return an + // error to the caller and MUST apply their own bounded timeout — + // progress delivery must never fail or slow down the evaluation. + Report(ctx context.Context, evt Event) +} + +type contextKey int + +var reporterKey = contextKey(0) + +// ContextWithReporter returns a copy of ctx carrying the given Reporter. +func ContextWithReporter(ctx context.Context, r Reporter) context.Context { + return context.WithValue(ctx, reporterKey, r) +} + +// FromContext returns the Reporter attached to ctx, or nil if none is +// attached. A nil Reporter is the expected, common case: most requests +// don't opt in to progress reporting. +func FromContext(ctx context.Context) Reporter { + r, _ := ctx.Value(reporterKey).(Reporter) + return r +} + +// Emit is the call-site convenience for reporting a progress event. It is +// a silent no-op when no Reporter is attached to ctx, which is what makes +// progress reporting purely opt-in/additive. +func Emit(ctx context.Context, evt Event) { + r := FromContext(ctx) + if r == nil { + return + } + + evt.Timestamp = time.Now().UTC() + r.Report(ctx, evt) +} diff --git a/internal/progress/reporter_test.go b/internal/progress/reporter_test.go new file mode 100644 index 0000000..3aee6b0 --- /dev/null +++ b/internal/progress/reporter_test.go @@ -0,0 +1,62 @@ +package progress + +import ( + "context" + "sync" + "testing" +) + +// recordingReporter is a test double shared across this package's test +// files. It's safe for concurrent use since sidecar_test.go exercises it +// from the sidecar's detached relay goroutine as well as the test +// goroutine polling for results. +type recordingReporter struct { + mu sync.Mutex + events []Event +} + +func (r *recordingReporter) Report(_ context.Context, evt Event) { + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, evt) +} + +// recorded returns a snapshot of the events received so far. +func (r *recordingReporter) recorded() []Event { + r.mu.Lock() + defer r.mu.Unlock() + return append([]Event(nil), r.events...) +} + +func TestEmit_NoReporterInContext_NoOp(t *testing.T) { + // must not panic, must not do anything observable + Emit(context.Background(), Event{Stage: StageEvaluating}) +} + +func TestEmit_WithReporter_DeliversEventAndSetsTimestamp(t *testing.T) { + r := &recordingReporter{} + ctx := ContextWithReporter(context.Background(), r) + + Emit(ctx, Event{Stage: StagePreparing, Command: "eval"}) + + events := r.recorded() + if len(events) != 1 { + t.Fatalf("expected 1 event, got %d", len(events)) + } + evt := events[0] + if evt.Stage != StagePreparing { + t.Errorf("expected stage %q, got %q", StagePreparing, evt.Stage) + } + if evt.Command != "eval" { + t.Errorf("expected command %q, got %q", "eval", evt.Command) + } + if evt.Timestamp.IsZero() { + t.Errorf("expected Emit to set a non-zero timestamp") + } +} + +func TestFromContext_NoReporter_ReturnsNil(t *testing.T) { + if r := FromContext(context.Background()); r != nil { + t.Errorf("expected nil reporter, got %v", r) + } +} diff --git a/internal/progress/sidecar.go b/internal/progress/sidecar.go new file mode 100644 index 0000000..4319dc0 --- /dev/null +++ b/internal/progress/sidecar.go @@ -0,0 +1,306 @@ +package progress + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "strings" + "sync" + "time" + + "go.uber.org/zap" +) + +const ( + defaultSidecarMaxBodyBytes int64 = 16 * 1024 + defaultSidecarMaxEventsPerSpan = 50 + defaultSidecarBurstSize = 5 + defaultSidecarMinEventInterval = 10 * time.Millisecond + defaultSidecarUnbindGracePeriod = 250 * time.Millisecond +) + +// SidecarConfig bounds abuse of the worker-authored progress side-channel. +// Since EVAL_PROGRESS_URL is reachable by arbitrary (and, under sandboxing, +// untrusted) worker code, delivery to the real callbackUrl must stay bounded +// regardless of how the worker behaves. +type SidecarConfig struct { + // MaxBodyBytes caps the size of a single progress event POST body. + // If unset (<= 0), defaultSidecarMaxBodyBytes is used. + MaxBodyBytes int64 `conf:"max_body_bytes"` + + // MaxEventsPerSpan caps how many progress events a single evaluation + // span (the window between Bind and the next Bind/Unbind) may relay. + // If unset (<= 0), defaultSidecarMaxEventsPerSpan is used. + MaxEventsPerSpan int `conf:"max_events_per_span"` + + // BurstSize is how many events at the start of a span are exempt from + // MinEventInterval spacing, so a handful of legitimate back-to-back + // checkpoints (e.g. a fast evaluation reporting progress at several + // points microseconds to a few ms apart) aren't rate-limited just + // because they arrive faster than any fixed spacing could accommodate. + // MinEventInterval spacing only applies once the burst is used up. + // Still bounded by MaxEventsPerSpan. If unset (== 0), + // defaultSidecarBurstSize is used; a negative value explicitly + // disables the burst allowance (spacing applies from the first event). + BurstSize int `conf:"burst_size"` + + // MinEventInterval enforces a minimum spacing between accepted events + // once a span's BurstSize allowance is used up. If unset (<= 0), + // defaultSidecarMinEventInterval is used. + MinEventInterval time.Duration `conf:"min_event_interval"` + + // UnbindGracePeriod delays detaching the bound reporter after a span + // ends, so a worker-authored progress POST that was already in flight + // (e.g. dispatched fire-and-forget just before the worker returned its + // result) still has a window to arrive and be relayed, instead of + // racing the RPC response back to shim. If unset (<= 0), + // defaultSidecarUnbindGracePeriod is used. + UnbindGracePeriod time.Duration `conf:"unbind_grace_period"` +} + +func (c SidecarConfig) withDefaults() SidecarConfig { + if c.MaxBodyBytes <= 0 { + c.MaxBodyBytes = defaultSidecarMaxBodyBytes + } + if c.MaxEventsPerSpan <= 0 { + c.MaxEventsPerSpan = defaultSidecarMaxEventsPerSpan + } + if c.BurstSize < 0 { + c.BurstSize = 0 + } else if c.BurstSize == 0 { + c.BurstSize = defaultSidecarBurstSize + } + if c.MinEventInterval <= 0 { + c.MinEventInterval = defaultSidecarMinEventInterval + } + if c.UnbindGracePeriod <= 0 { + c.UnbindGracePeriod = defaultSidecarUnbindGracePeriod + } + return c +} + +// sidecarPayload is the JSON body a worker POSTs to report a custom +// progress event. There is deliberately no "stage" field: a worker can +// never choose its own stage. The sidecar assigns one from the command in +// flight (see stageForCommand). Unknown fields (including a "stage" a +// worker might send anyway) are silently ignored by json.Decode, never +// merged in. +type sidecarPayload struct { + Message string `json:"message"` + Data map[string]any `json:"data,omitempty"` +} + +// stageForCommand maps the command bound to the sidecar onto the stage a +// worker-authored sub-step is relayed under: chat commands report +// "thinking", everything else (eval, preview, …) reports "evaluating". +func stageForCommand(command string) Stage { + switch command { + case "chat", "chat/health": + return StageThinking + default: + return StageEvaluating + } +} + +// Sidecar is a loopback-only HTTP listener that accepts worker-authored +// progress events and relays them, best-effort, through whichever Reporter +// is currently Bind-ed to it. It is the counterpart, on the inbound side, +// to the outbound delivery in http_reporter.go: since it only ever binds +// to 127.0.0.1, it needs no SSRF guarding, but it does need its own abuse +// limits, since the worker producing events may be untrusted. +// +// Its lifetime differs by adapter: for a persistent RPC worker, one Sidecar +// lives for the worker's whole lifetime and is Bind/Unbind-ed around each +// request; for the transient file interface, one Sidecar is created and +// Closed per request. +type Sidecar struct { + cfg SidecarConfig + log *zap.Logger + + listener net.Listener + server *http.Server + + mu sync.Mutex + command string + reporter Reporter + count int + lastSent time.Time + generation uint64 +} + +// NewSidecar starts a loopback HTTP listener on an OS-assigned port. +func NewSidecar(cfg SidecarConfig, log *zap.Logger) (*Sidecar, error) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("failed to start progress sidecar listener: %w", err) + } + + s := &Sidecar{ + cfg: cfg.withDefaults(), + log: log.Named("progress_sidecar"), + listener: ln, + } + + s.server = &http.Server{ + Handler: http.HandlerFunc(s.handle), + ReadHeaderTimeout: 2 * time.Second, + ReadTimeout: 2 * time.Second, + WriteTimeout: 2 * time.Second, + } + + go func() { + if err := s.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { + s.log.Warn("progress sidecar listener stopped unexpectedly", zap.Error(err)) + } + }() + + return s, nil +} + +// URL returns the sidecar's loopback address, suitable for EVAL_PROGRESS_URL. +func (s *Sidecar) URL() string { + return "http://" + s.listener.Addr().String() +} + +// Bind associates command/reporter with the sidecar for the duration of one +// evaluation span, resetting rate-limit state so a fresh span isn't +// poisoned by the previous request's usage. Call at the start of an +// adapter's Send. A nil reporter behaves like Unbind. +func (s *Sidecar) Bind(command string, reporter Reporter) { + s.mu.Lock() + defer s.mu.Unlock() + + s.generation++ + s.command = command + s.reporter = reporter + s.count = 0 + s.lastSent = time.Time{} +} + +// Unbind detaches the current reporter immediately, so any subsequent POST +// (e.g. a straggler arriving after the bound request has already returned) +// is rejected with 503 rather than misattributed to a future, unrelated +// request. +func (s *Sidecar) Unbind() { + s.mu.Lock() + defer s.mu.Unlock() + + s.generation++ + s.command = "" + s.reporter = nil +} + +// UnbindAfterGrace schedules the detach for after cfg.UnbindGracePeriod +// instead of doing it immediately, without blocking the caller. This gives +// a worker-authored progress POST dispatched fire-and-forget just before +// the RPC response reached shim a window to still arrive and be relayed, +// rather than losing the race against Unbind and being rejected with 503. +// +// If a new span is Bind-ed (or explicitly Unbind-ed) before the grace +// period elapses, this is a no-op: the generation captured at schedule time +// will no longer match, so the stale detach never fires and never clobbers +// the newer span. +func (s *Sidecar) UnbindAfterGrace() { + s.mu.Lock() + gen := s.generation + grace := s.cfg.UnbindGracePeriod + s.mu.Unlock() + + if grace <= 0 { + s.Unbind() + return + } + + time.AfterFunc(grace, func() { + s.mu.Lock() + defer s.mu.Unlock() + + if s.generation != gen { + return + } + + s.generation++ + s.command = "" + s.reporter = nil + }) +} + +// Close shuts down the sidecar's listener. It does not wait for any +// in-flight relayed events (those run detached from the listener, see +// handle) — consistent with progress delivery never blocking shutdown. +func (s *Sidecar) Close() error { + return s.server.Close() +} + +func (s *Sidecar) handle(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + + r.Body = http.MaxBytesReader(w, r.Body, s.cfg.MaxBodyBytes) + + var body sidecarPayload + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + w.WriteHeader(http.StatusRequestEntityTooLarge) + return + } + w.WriteHeader(http.StatusBadRequest) + return + } + + if strings.TrimSpace(body.Message) == "" { + w.WriteHeader(http.StatusBadRequest) + return + } + + command, reporter, status := s.accept() + if status != 0 { + w.WriteHeader(status) + return + } + + w.WriteHeader(http.StatusAccepted) + + evt := Event{ + Stage: stageForCommand(command), + Command: command, + Message: body.Message, + Data: body.Data, + } + + // Relay detached from the inbound request: the worker's POST must + // never be held open for the outbound callbackUrl delivery, which has + // its own bounded timeout inside Report. + go reporter.Report(context.Background(), evt) +} + +// accept reports whether a new event may be relayed right now, applying +// the bound reporter check and the abuse limits. status is 0 on success, +// or the HTTP status to reject the request with otherwise. +func (s *Sidecar) accept() (command string, reporter Reporter, status int) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.reporter == nil { + return "", nil, http.StatusServiceUnavailable + } + + now := time.Now() + if s.count >= s.cfg.MaxEventsPerSpan { + return "", nil, http.StatusTooManyRequests + } + if s.count >= s.cfg.BurstSize && !s.lastSent.IsZero() && now.Sub(s.lastSent) < s.cfg.MinEventInterval { + return "", nil, http.StatusTooManyRequests + } + + s.count++ + s.lastSent = now + + return s.command, s.reporter, 0 +} diff --git a/internal/progress/sidecar_test.go b/internal/progress/sidecar_test.go new file mode 100644 index 0000000..b5d148b --- /dev/null +++ b/internal/progress/sidecar_test.go @@ -0,0 +1,258 @@ +package progress + +import ( + "bytes" + "net/http" + "strings" + "testing" + "time" + + "go.uber.org/zap" +) + +func newTestSidecar(t *testing.T, cfg SidecarConfig) *Sidecar { + t.Helper() + s, err := NewSidecar(cfg, zap.NewNop()) + if err != nil { + t.Fatalf("failed to start sidecar: %v", err) + } + t.Cleanup(func() { s.Close() }) + return s +} + +func postSidecar(t *testing.T, s *Sidecar, body string) *http.Response { + t.Helper() + resp, err := http.Post(s.URL(), "application/json", bytes.NewBufferString(body)) + if err != nil { + t.Fatalf("failed to POST to sidecar: %v", err) + } + defer resp.Body.Close() + return resp +} + +// waitForEvents polls until r has at least n events or the timeout expires, +// since the sidecar relays events in a detached goroutine. +func waitForEvents(t *testing.T, r *recordingReporter, n int) []Event { + t.Helper() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if events := r.recorded(); len(events) >= n { + return events + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("timed out waiting for %d events, got %d", n, len(r.recorded())) + return nil +} + +func TestSidecar_Accept_RelaysEventThroughBoundReporter(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + r := &recordingReporter{} + s.Bind("eval", r) + + resp := postSidecar(t, s, `{"message":"checking correctness…","data":{"step":2}}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected 202, got %d", resp.StatusCode) + } + + events := waitForEvents(t, r, 1) + evt := events[0] + if evt.Stage != StageEvaluating { + t.Errorf("expected stage %q, got %q", StageEvaluating, evt.Stage) + } + if evt.Command != "eval" { + t.Errorf("expected command %q, got %q", "eval", evt.Command) + } + if evt.Message != "checking correctness…" { + t.Errorf("unexpected message %q", evt.Message) + } + if evt.Data["step"] != float64(2) { + t.Errorf("expected data.step=2, got %v", evt.Data["step"]) + } +} + +func TestSidecar_IgnoresWorkerSuppliedStage(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + r := &recordingReporter{} + s.Bind("eval", r) + + resp := postSidecar(t, s, `{"message":"trying to spoof","stage":"completed"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected 202, got %d", resp.StatusCode) + } + + events := waitForEvents(t, r, 1) + if events[0].Stage != StageEvaluating { + t.Errorf("worker-supplied stage must be ignored, got %q", events[0].Stage) + } +} + +func TestSidecar_StageFollowsBoundCommand(t *testing.T) { + cases := map[string]Stage{ + "eval": StageEvaluating, + "preview": StageEvaluating, + "chat": StageThinking, + "chat/health": StageThinking, + } + for command, wantStage := range cases { + t.Run(command, func(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + r := &recordingReporter{} + s.Bind(command, r) + + resp := postSidecar(t, s, `{"message":"working…"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected 202, got %d", resp.StatusCode) + } + + evt := waitForEvents(t, r, 1)[0] + if evt.Stage != wantStage { + t.Errorf("command %q: expected stage %q, got %q", command, wantStage, evt.Stage) + } + }) + } +} + +func TestSidecar_RejectsEmptyMessage(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + s.Bind("eval", &recordingReporter{}) + + resp := postSidecar(t, s, `{"message":""}`) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", resp.StatusCode) + } +} + +func TestSidecar_RejectsMalformedJSON(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + s.Bind("eval", &recordingReporter{}) + + resp := postSidecar(t, s, `not json`) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", resp.StatusCode) + } +} + +func TestSidecar_RejectsOversizedBody(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{MaxBodyBytes: 16}) + s.Bind("eval", &recordingReporter{}) + + body := `{"message":"` + strings.Repeat("x", 100) + `"}` + resp := postSidecar(t, s, body) + if resp.StatusCode != http.StatusRequestEntityTooLarge { + t.Fatalf("expected 413, got %d", resp.StatusCode) + } +} + +func TestSidecar_RateLimit_MaxEventsPerSpan(t *testing.T) { + // MinEventInterval is small (not disabled - 0 means "use the default") + // and slept past between POSTs, so only MaxEventsPerSpan is under test. + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 1, MinEventInterval: time.Millisecond}) + r := &recordingReporter{} + s.Bind("eval", r) + + first := postSidecar(t, s, `{"message":"one"}`) + if first.StatusCode != http.StatusAccepted { + t.Fatalf("expected first event accepted (202), got %d", first.StatusCode) + } + + time.Sleep(5 * time.Millisecond) + + second := postSidecar(t, s, `{"message":"two"}`) + if second.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected second event rate limited (429), got %d", second.StatusCode) + } +} + +func TestSidecar_RateLimit_MinEventInterval(t *testing.T) { + // BurstSize disabled so the very first event is already subject to + // interval spacing, isolating what this test exercises. + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, BurstSize: -1, MinEventInterval: time.Hour}) + r := &recordingReporter{} + s.Bind("eval", r) + + first := postSidecar(t, s, `{"message":"one"}`) + if first.StatusCode != http.StatusAccepted { + t.Fatalf("expected first event accepted (202), got %d", first.StatusCode) + } + + second := postSidecar(t, s, `{"message":"two"}`) + if second.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected second event rate limited (429) by min interval, got %d", second.StatusCode) + } +} + +func TestSidecar_Burst_AllowsCloselySpacedEventsWithinBurst(t *testing.T) { + // A large MinEventInterval would reject any second event immediately - + // unless it falls within the burst allowance, which is what this + // exercises: events 2 and 3 land inside BurstSize and must be accepted + // even though far less than MinEventInterval separates them. + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, BurstSize: 3, MinEventInterval: time.Hour}) + r := &recordingReporter{} + s.Bind("eval", r) + + for i, msg := range []string{"one", "two", "three"} { + resp := postSidecar(t, s, `{"message":"`+msg+`"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected burst event %d accepted (202), got %d", i+1, resp.StatusCode) + } + } +} + +func TestSidecar_Burst_ThenEnforcesMinEventInterval(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, BurstSize: 2, MinEventInterval: time.Hour}) + r := &recordingReporter{} + s.Bind("eval", r) + + for i, msg := range []string{"one", "two"} { + resp := postSidecar(t, s, `{"message":"`+msg+`"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected burst event %d accepted (202), got %d", i+1, resp.StatusCode) + } + } + + third := postSidecar(t, s, `{"message":"three"}`) + if third.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected event past burst allowance rate limited (429), got %d", third.StatusCode) + } +} + +func TestSidecar_Bind_ResetsRateLimitState(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 1, MinEventInterval: time.Millisecond}) + r1 := &recordingReporter{} + s.Bind("eval", r1) + + postSidecar(t, s, `{"message":"one"}`) + exhausted := postSidecar(t, s, `{"message":"two"}`) + if exhausted.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected span to be exhausted (429), got %d", exhausted.StatusCode) + } + + r2 := &recordingReporter{} + s.Bind("eval", r2) + + resp := postSidecar(t, s, `{"message":"fresh span"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected fresh span to accept after re-Bind (202), got %d", resp.StatusCode) + } +} + +func TestSidecar_Unbound_Returns503(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + + resp := postSidecar(t, s, `{"message":"nobody home"}`) + if resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("expected 503 with no bound reporter, got %d", resp.StatusCode) + } +} + +func TestSidecar_Unbind_Returns503(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + s.Bind("eval", &recordingReporter{}) + s.Unbind() + + resp := postSidecar(t, s, `{"message":"straggler"}`) + if resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("expected 503 after Unbind, got %d", resp.StatusCode) + } +} diff --git a/internal/progress/sse_reporter.go b/internal/progress/sse_reporter.go new file mode 100644 index 0000000..00ceb07 --- /dev/null +++ b/internal/progress/sse_reporter.go @@ -0,0 +1,243 @@ +package progress + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + "time" + + "go.uber.org/zap" +) + +// sseStep is one progress step. The same shape is written both as its own +// live frame (event: ) the moment the event arrives and as an +// element of the terminal envelope's steps[], so a client parses "a step" +// the same way whether it arrives inline or inside the terminal frame. +type sseStep struct { + Stage string `json:"stage"` + Message string `json:"message,omitempty"` + Data map[string]any `json:"data,omitempty"` + Timestamp time.Time `json:"timestamp"` +} + +// sseEnvelope is the JSON payload of the single terminal SSE frame for an +// /evaluate (or /preview) request. The same shape is used for the +// "completed" and "failed" events: on failure Feedback is null and Error +// (an ErrorResponse-shaped object) carries the detail. It matches the +// spec's SseEvaluateTerminalFrame. +type sseEnvelope struct { + Feedback []map[string]any `json:"feedback"` + Steps []sseStep `json:"steps"` + Error *ErrorInfo `json:"error,omitempty"` +} + +// sseChatEnvelope is the terminal-frame payload for a /chat request. Chat +// has no feedback[]; it returns an output object plus optional metadata. +// On failure Output is null and Error carries the detail. It matches the +// spec's SseChatTerminalFrame. +type sseChatEnvelope struct { + Output map[string]any `json:"output"` + Metadata map[string]any `json:"metadata,omitempty"` + Steps []sseStep `json:"steps"` + Error *ErrorInfo `json:"error,omitempty"` +} + +// SSEReporter is a Reporter that streams progress back to the caller on +// the /evaluate or /chat response itself, as Server-Sent Events. Each +// non-terminal event is written immediately as its own frame +// (event: , data = the step object) so the caller sees progress as +// it happens, and is also accumulated; on completion/failure a single +// terminal frame (event: completed | failed) carries the result plus +// every step that preceded it, then the handler closes the connection. +// +// Report is called concurrently — synchronously from the request +// goroutine for shim-authored events, and from detached sidecar +// goroutines for worker-authored sub-steps — so all state and all writes +// to the ResponseWriter are guarded by mu. +type SSEReporter struct { + w http.ResponseWriter + flusher http.Flusher + // command ("evaluate" | "preview" | "chat") only selects the terminal + // frame shape (feedback[] vs output/metadata); it is not serialised. + command string + log *zap.Logger + + mu sync.Mutex + steps []sseStep + seenPreparing bool + seenStarting bool + terminated bool + terminalOnce sync.Once +} + +var _ Reporter = (*SSEReporter)(nil) + +// NewSSEReporter returns a reporter that writes SSE frames to w. It +// returns an error if w cannot be flushed incrementally, so the caller +// can fall back to a buffered response. +func NewSSEReporter(w http.ResponseWriter, command string, log *zap.Logger) (*SSEReporter, error) { + flusher, ok := w.(http.Flusher) + if !ok { + return nil, fmt.Errorf("response writer does not support flushing") + } + return &SSEReporter{ + w: w, + flusher: flusher, + command: command, + log: log, + }, nil +} + +// Report streams a non-terminal event as its own frame (and accumulates +// it), or writes the single terminal frame. Once the terminal frame is +// written, all further events (including a late worker sub-step relayed +// after the request returned) are dropped without touching the +// ResponseWriter. +func (r *SSEReporter) Report(_ context.Context, evt Event) { + r.mu.Lock() + defer r.mu.Unlock() + + if r.terminated { + return + } + + if evt.Stage.terminal() { + r.terminalOnce.Do(func() { + r.terminated = true + r.writeEnvelopeLocked(evt) + }) + return + } + + // Collapse the shim's lifecycle markers to their first occurrence for + // the whole request: the per-case evaluation loop re-enters the + // supervisor and re-emits preparing/starting once per case. + // Worker-authored sub-steps (evaluating / thinking) are never + // collapsed — they sit on their own stages. + switch evt.Stage { + case StagePreparing: + if r.seenPreparing { + return + } + r.seenPreparing = true + case StageStarting: + if r.seenStarting { + return + } + r.seenStarting = true + } + + step := sseStep{ + Stage: string(evt.Stage), + Message: evt.Message, + Data: evt.Data, + Timestamp: evt.Timestamp, + } + if step.Timestamp.IsZero() { + // Worker-authored sub-steps bypass Emit (they come off the + // sidecar) and arrive without a timestamp. + step.Timestamp = time.Now().UTC() + } + + r.steps = append(r.steps, step) + r.writeStepLocked(step) +} + +// failureErrorInfo returns the ErrorResponse-shaped object for a "failed" +// terminal frame. It prefers the structured ErrorInfo the handler +// attached; failing that it synthesises a minimal object from the event's +// human-facing Message and raw Error so `title` is never empty. +func failureErrorInfo(evt Event) *ErrorInfo { + if evt.ErrorInfo != nil { + return evt.ErrorInfo + } + return &ErrorInfo{Title: "Error", Message: evt.Message, Trace: evt.Error} +} + +func (r *SSEReporter) writeEnvelopeLocked(evt Event) { + steps := r.steps + if steps == nil { + steps = []sseStep{} + } + failed := evt.Stage == StageFailed + + event := "completed" + if failed { + event = "failed" + } + + var payload any + if r.command == "chat" { + env := sseChatEnvelope{Steps: steps} + if failed { + env.Error = failureErrorInfo(evt) + } else { + env.Output, _ = evt.Data["output"].(map[string]any) + env.Metadata, _ = evt.Data["metadata"].(map[string]any) + } + payload = env + } else { + env := sseEnvelope{Steps: steps} + if failed { + env.Error = failureErrorInfo(evt) + } else { + feedback, ok := evt.Data["feedback"].([]map[string]any) + if !ok { + feedback = []map[string]any{} + } + env.Feedback = feedback + } + payload = env + } + + body, err := json.Marshal(payload) + if err != nil { + r.log.Warn("failed to marshal SSE envelope", zap.String("event", event), zap.Error(err)) + return + } + + if _, err := fmt.Fprintf(r.w, "event: %s\ndata: %s\n\n", event, body); err != nil { + r.log.Debug("failed to write SSE terminal frame", zap.Error(err)) + return + } + r.flusher.Flush() +} + +// writeStepLocked streams a single intermediate progress step as its own +// SSE frame (event: ), so the caller sees progress as it happens +// rather than only in the terminal frame. The step is already recorded in +// r.steps for the terminal envelope, so a marshal or write failure here +// only costs the live frame. A write failure does not set terminated: the +// terminal-frame attempt and further accumulation continue. Callers hold +// r.mu. +func (r *SSEReporter) writeStepLocked(step sseStep) { + body, err := json.Marshal(step) + if err != nil { + r.log.Warn("failed to marshal SSE step", zap.String("stage", step.Stage), zap.Error(err)) + return + } + + if _, err := fmt.Fprintf(r.w, "event: %s\ndata: %s\n\n", step.Stage, body); err != nil { + r.log.Debug("failed to write SSE step frame", zap.Error(err)) + return + } + r.flusher.Flush() +} + +// Heartbeat writes an SSE comment line to keep the connection alive. It +// is a no-op once the terminal frame has been written. +func (r *SSEReporter) Heartbeat() { + r.mu.Lock() + defer r.mu.Unlock() + + if r.terminated { + return + } + if _, err := r.w.Write([]byte(": ping\n\n")); err != nil { + r.log.Debug("failed to write SSE heartbeat", zap.Error(err)) + return + } + r.flusher.Flush() +} diff --git a/internal/progress/sse_reporter_test.go b/internal/progress/sse_reporter_test.go new file mode 100644 index 0000000..9d99867 --- /dev/null +++ b/internal/progress/sse_reporter_test.go @@ -0,0 +1,481 @@ +package progress + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "go.uber.org/zap" +) + +type sseFrame struct { + event string + data map[string]any + raw string +} + +func parseSSEFrames(t *testing.T, raw string) []sseFrame { + t.Helper() + var frames []sseFrame + for _, block := range strings.Split(strings.TrimSpace(raw), "\n\n") { + block = strings.TrimSpace(block) + if block == "" || strings.HasPrefix(block, ":") { + continue // heartbeat / comment + } + var f sseFrame + f.raw = block + for _, line := range strings.Split(block, "\n") { + switch { + case strings.HasPrefix(line, "event: "): + f.event = strings.TrimPrefix(line, "event: ") + case strings.HasPrefix(line, "data: "): + payload := strings.TrimPrefix(line, "data: ") + if err := json.Unmarshal([]byte(payload), &f.data); err != nil { + t.Fatalf("frame data is not valid JSON: %v\n%s", err, payload) + } + } + } + frames = append(frames, f) + } + return frames +} + +func newRecorderReporter(t *testing.T, command string) (*httptest.ResponseRecorder, *SSEReporter) { + t.Helper() + rec := httptest.NewRecorder() + r, err := NewSSEReporter(rec, command, zap.NewNop()) + if err != nil { + t.Fatalf("NewSSEReporter: %v", err) + } + return rec, r +} + +func TestSSEReporter_CompletedEnvelope(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + + r.Report(context.Background(), Event{Stage: StagePreparing, Message: "Preparing…"}) + r.Report(context.Background(), Event{Stage: StageStarting, Message: "Starting…"}) + r.Report(context.Background(), Event{ + Stage: StageCompleted, + Data: map[string]any{"feedback": []map[string]any{{"message": "Well done"}}}, + }) + + if !rec.Flushed { + t.Error("expected the response to be flushed") + } + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 3 { + t.Fatalf("expected 3 frames (preparing, starting, completed), got %d: %q", len(frames), rec.Body.String()) + } + if frames[0].event != "preparing" || frames[1].event != "starting" { + t.Errorf("unexpected live frame events: %q, %q", frames[0].event, frames[1].event) + } + f := frames[2] + if f.event != "completed" { + t.Errorf("expected event 'completed', got %q", f.event) + } + if _, hasCommand := f.data["command"]; hasCommand { + t.Errorf("terminal frame must not carry a command key: %v", f.data) + } + fb, ok := f.data["feedback"].([]any) + if !ok || len(fb) != 1 { + t.Fatalf("expected feedback array of 1, got %v", f.data["feedback"]) + } + if fb[0].(map[string]any)["message"] != "Well done" { + t.Errorf("feedback item not carried through: %v", fb[0]) + } + steps, ok := f.data["steps"].([]any) + if !ok || len(steps) != 2 { + t.Fatalf("expected 2 steps, got %v", f.data["steps"]) + } + if steps[0].(map[string]any)["stage"] != "preparing" || steps[1].(map[string]any)["stage"] != "starting" { + t.Errorf("unexpected step stages: %v", steps) + } + if steps[0].(map[string]any)["timestamp"] == "" { + t.Errorf("expected step timestamp to be set") + } +} + +func TestSSEReporter_FailedEnvelope(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + + r.Report(context.Background(), Event{Stage: StagePreparing, Message: "Preparing…"}) + r.Report(context.Background(), Event{ + Stage: StageFailed, + Error: "worker send: context deadline exceeded", + Message: "We couldn't evaluate your answer. Please try again.", + ErrorInfo: &ErrorInfo{ + Title: "Evaluation failed", + Message: "We couldn't evaluate your answer. Please try again.", + Code: "INTERNAL_ERROR", + Trace: "worker send: context deadline exceeded", + }, + }) + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 2 { + t.Fatalf("expected 2 frames (preparing, failed), got %d", len(frames)) + } + if frames[0].event != "preparing" { + t.Errorf("expected first frame 'preparing', got %q", frames[0].event) + } + f := frames[1] + if f.event != "failed" { + t.Errorf("expected event 'failed', got %q", f.event) + } + if v, ok := f.data["feedback"]; !ok || v != nil { + t.Errorf("expected feedback null, got %v (present=%v)", v, ok) + } + errObj, ok := f.data["error"].(map[string]any) + if !ok { + t.Fatalf("expected error to be an ErrorResponse object, got %T: %v", f.data["error"], f.data["error"]) + } + if errObj["title"] != "Evaluation failed" { + t.Errorf("error title not carried: %v", errObj["title"]) + } + if errObj["message"] != "We couldn't evaluate your answer. Please try again." { + t.Errorf("error message not carried: %v", errObj["message"]) + } + if errObj["trace"] != "worker send: context deadline exceeded" { + t.Errorf("error trace not carried: %v", errObj["trace"]) + } + if _, hasMessage := f.data["message"]; hasMessage { + t.Errorf("terminal frame must not carry a top-level message key: %v", f.data) + } + if steps, _ := f.data["steps"].([]any); len(steps) != 1 { + t.Errorf("expected 1 step, got %v", f.data["steps"]) + } +} + +func TestSSEReporter_PreviewUsesFeedbackEnvelope(t *testing.T) { + rec, r := newRecorderReporter(t, "preview") + r.Report(context.Background(), Event{ + Stage: StageCompleted, + Data: map[string]any{"feedback": []map[string]any{{"preSubmissionFeedback": map[string]any{}}}}, + }) + frames := parseSSEFrames(t, rec.Body.String()) + f := frames[0] + if f.event != "completed" { + t.Fatalf("expected 'completed', got %q", f.event) + } + if _, hasCommand := f.data["command"]; hasCommand { + t.Errorf("terminal frame must not carry a command key: %v", f.data) + } + fb, ok := f.data["feedback"].([]any) + if !ok || len(fb) != 1 { + t.Fatalf("preview should use the feedback envelope, got %v", f.data["feedback"]) + } + if _, ok := f.data["steps"].([]any); !ok { + t.Errorf("steps should always be present as an array, got %v", f.data["steps"]) + } +} + +func TestSSEReporter_ChatEnvelope_Completed(t *testing.T) { + rec, r := newRecorderReporter(t, "chat") + + r.Report(context.Background(), Event{Stage: StageThinking, Message: "Searching your notes…"}) + r.Report(context.Background(), Event{ + Stage: StageCompleted, + Data: map[string]any{ + "output": map[string]any{"role": "ASSISTANT", "content": "Here you go"}, + "metadata": map[string]any{"model": "x"}, + }, + }) + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 2 || frames[0].event != "thinking" || frames[1].event != "completed" { + t.Fatalf("expected [thinking, completed], got %q", rec.Body.String()) + } + f := frames[1] + if _, hasCommand := f.data["command"]; hasCommand { + t.Errorf("terminal frame must not carry a command key: %v", f.data) + } + if _, hasFeedback := f.data["feedback"]; hasFeedback { + t.Errorf("chat envelope must not carry a feedback key: %v", f.data) + } + out, ok := f.data["output"].(map[string]any) + if !ok || out["content"] != "Here you go" { + t.Fatalf("expected output object, got %v", f.data["output"]) + } + if md, ok := f.data["metadata"].(map[string]any); !ok || md["model"] != "x" { + t.Errorf("expected metadata carried, got %v", f.data["metadata"]) + } + if steps, _ := f.data["steps"].([]any); len(steps) != 1 { + t.Errorf("expected 1 step, got %v", f.data["steps"]) + } +} + +func TestSSEReporter_ChatEnvelope_Failed(t *testing.T) { + rec, r := newRecorderReporter(t, "chat") + + r.Report(context.Background(), Event{ + Stage: StageFailed, + Error: "chat failed: worker exited", + Message: "We couldn't generate a response. Please try again.", + ErrorInfo: &ErrorInfo{ + Title: "Chat failed", + Message: "We couldn't generate a response. Please try again.", + Trace: "chat failed: worker exited", + }, + }) + + f := parseSSEFrames(t, rec.Body.String())[0] + if f.event != "failed" { + t.Fatalf("expected 'failed', got %q", f.event) + } + if v, ok := f.data["output"]; !ok || v != nil { + t.Errorf("expected output null, got %v (present=%v)", v, ok) + } + errObj, ok := f.data["error"].(map[string]any) + if !ok { + t.Fatalf("expected error to be an ErrorResponse object, got %T: %v", f.data["error"], f.data["error"]) + } + if errObj["title"] != "Chat failed" { + t.Errorf("error title not carried: %v", errObj["title"]) + } + if errObj["trace"] != "chat failed: worker exited" { + t.Errorf("error trace not carried: %v", errObj["trace"]) + } +} + +func TestSSEReporter_DedupLifecycleStagesOncePerRequest(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + // The per-case evaluation loop re-emits the shim's preparing/starting + // markers once per case; only the first of each for the whole request + // is kept. Worker-authored evaluating sub-steps are never collapsed. + for _, s := range []Stage{StagePreparing, StagePreparing, StageStarting, StageStarting, StageEvaluating, StageEvaluating, StageStarting} { + r.Report(context.Background(), Event{Stage: s}) + } + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + + frames := parseSSEFrames(t, rec.Body.String()) + + var liveEvents []string + for _, f := range frames[:len(frames)-1] { + liveEvents = append(liveEvents, f.event) + } + if strings.Join(liveEvents, ",") != "preparing,starting,evaluating,evaluating" { + t.Errorf("expected live frames [preparing starting evaluating evaluating], got %v", liveEvents) + } + + steps := frames[len(frames)-1].data["steps"].([]any) + got := []string{} + for _, s := range steps { + got = append(got, s.(map[string]any)["stage"].(string)) + } + want := []string{"preparing", "starting", "evaluating", "evaluating"} + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Errorf("expected steps %v, got %v", want, got) + } +} + +func TestSSEReporter_ProgressStepsNotDedupedAndTimestamped(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + r.Report(context.Background(), Event{Stage: StageProgress, Message: "same"}) + r.Report(context.Background(), Event{Stage: StageProgress, Message: "same"}) + r.Report(context.Background(), Event{Stage: StageProgress, Message: "same"}) + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 4 { + t.Fatalf("expected 4 frames (3 progress + completed), got %d: %q", len(frames), rec.Body.String()) + } + for _, f := range frames[:3] { + if f.event != "progress" { + t.Errorf("expected a 'progress' live frame, got %q", f.event) + } + } + + steps := frames[3].data["steps"].([]any) + if len(steps) != 3 { + t.Fatalf("expected 3 progress steps, got %d", len(steps)) + } + for _, s := range steps { + ts, _ := s.(map[string]any)["timestamp"].(string) + parsed, err := time.Parse(time.RFC3339Nano, ts) + if err != nil || parsed.IsZero() { + t.Errorf("expected a non-zero RFC3339 timestamp, got %q (err=%v)", ts, err) + } + } +} + +func TestSSEReporter_TerminalOnce(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + r.Report(context.Background(), Event{Stage: StageFailed, Message: "first"}) + r.Report(context.Background(), Event{Stage: StageFailed, Message: "second"}) + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 1 { + t.Fatalf("expected exactly 1 terminal frame, got %d", len(frames)) + } + if frames[0].event != "failed" { + t.Errorf("expected the first terminal event ('failed') to win, got %q", frames[0].event) + } + errObj, _ := frames[0].data["error"].(map[string]any) + if errObj["message"] != "first" { + t.Errorf("expected the first terminal event to win, got %v", errObj["message"]) + } +} + +func TestSSEReporter_LateProgressAfterTerminalDropped(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + before := rec.Body.String() + + r.Report(context.Background(), Event{Stage: StageProgress, Message: "too late"}) + r.Heartbeat() + + if rec.Body.String() != before { + t.Errorf("expected no output after terminal frame, got extra: %q", strings.TrimPrefix(rec.Body.String(), before)) + } +} + +func TestSSEReporter_Heartbeat(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + r.Report(context.Background(), Event{Stage: StagePreparing}) + r.Heartbeat() + r.Heartbeat() + if got := strings.Count(rec.Body.String(), ": ping\n\n"); got != 2 { + t.Errorf("expected 2 heartbeat comments, got %d", got) + } + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + r.Heartbeat() + if got := strings.Count(rec.Body.String(), ": ping\n\n"); got != 2 { + t.Errorf("expected heartbeat to be a no-op after terminal, got %d", got) + } +} + +func TestSSEReporter_ConcurrentReport(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + + var wg sync.WaitGroup + for i := 0; i < 20; i++ { + wg.Add(1) + go func() { + defer wg.Done() + r.Report(context.Background(), Event{Stage: StageProgress, Message: "p"}) + r.Heartbeat() + }() + } + wg.Wait() + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + + // Every frame must be a well-formed event/data pair (parseSSEFrames + // fails the test otherwise). Exactly one terminal frame, and it's last. + frames := parseSSEFrames(t, rec.Body.String()) + terminals := 0 + for i, f := range frames { + if f.event == "completed" || f.event == "failed" { + terminals++ + if i != len(frames)-1 { + t.Errorf("terminal frame at index %d is not last of %d", i, len(frames)) + } + } else if f.event != "progress" { + t.Errorf("unexpected intermediate frame event %q", f.event) + } + } + if terminals != 1 { + t.Fatalf("expected exactly 1 terminal frame, got %d: %q", terminals, rec.Body.String()) + } +} + +type nonFlusherWriter struct{ h http.Header } + +func (n nonFlusherWriter) Header() http.Header { return n.h } +func (n nonFlusherWriter) Write(b []byte) (int, error) { return len(b), nil } +func (n nonFlusherWriter) WriteHeader(int) {} + +func TestNewSSEReporter_NonFlusher_Error(t *testing.T) { + _, err := NewSSEReporter(nonFlusherWriter{h: http.Header{}}, "evaluate", zap.NewNop()) + if err == nil { + t.Fatal("expected an error for a non-flushable writer") + } +} + +func TestSSEReporter_LiveFrameMatchesTerminalStep(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + + r.Report(context.Background(), Event{ + Stage: StageProgress, + Message: "Parsing response and answer...", + Data: map[string]any{"step": float64(1), "of": float64(4)}, + }) + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 2 { + t.Fatalf("expected 2 frames, got %d: %q", len(frames), rec.Body.String()) + } + + live := frames[0] + if live.event != "progress" { + t.Fatalf("expected 'progress' live frame, got %q", live.event) + } + + steps := frames[1].data["steps"].([]any) + if len(steps) != 1 { + t.Fatalf("expected 1 terminal step, got %d", len(steps)) + } + + // The live frame's data payload must be byte-identical to the matching + // terminal steps[] element. + wantJSON, _ := json.Marshal(steps[0]) + gotJSON, _ := json.Marshal(live.data) + if string(wantJSON) != string(gotJSON) { + t.Errorf("live frame data != terminal step:\n live: %s\n step: %s", gotJSON, wantJSON) + } +} + +// failingAfterNWriter is an http.Flusher whose Write starts returning an +// error after okWrites successful writes. +type failingAfterNWriter struct { + h http.Header + okWrites int + writes int + flushed int +} + +func (w *failingAfterNWriter) Header() http.Header { return w.h } +func (w *failingAfterNWriter) WriteHeader(int) {} +func (w *failingAfterNWriter) Flush() { w.flushed++ } +func (w *failingAfterNWriter) Write(b []byte) (int, error) { + w.writes++ + if w.writes > w.okWrites { + return 0, io.ErrClosedPipe + } + return len(b), nil +} + +func TestSSEReporter_LiveFrameWriteErrorDoesNotTerminate(t *testing.T) { + w := &failingAfterNWriter{h: http.Header{}, okWrites: 1} + r, err := NewSSEReporter(w, "evaluate", zap.NewNop()) + if err != nil { + t.Fatalf("NewSSEReporter: %v", err) + } + + // First live frame writes OK; the second fails at the writer. + r.Report(context.Background(), Event{Stage: StageProgress, Message: "one"}) + r.Report(context.Background(), Event{Stage: StageProgress, Message: "two"}) + + if r.terminated { + t.Fatal("a live-frame write error must not set terminated") + } + + // The terminal frame is still attempted (Write is called again). + writesBefore := w.writes + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + if w.writes == writesBefore { + t.Error("expected the terminal frame to still attempt a write after a live-frame write error") + } + if !r.terminated { + t.Error("expected terminated to be set once the terminal frame ran") + } +} diff --git a/internal/progress/sse_schema_parity_test.go b/internal/progress/sse_schema_parity_test.go new file mode 100644 index 0000000..d1b6885 --- /dev/null +++ b/internal/progress/sse_schema_parity_test.go @@ -0,0 +1,207 @@ +package progress + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/getkin/kin-openapi/openapi3" + + "github.com/lambda-feedback/shimmy/internal/server" +) + +// These tests keep the hand-written Sse* component schemas in +// runtime/schema/mued_v0.1.0.yml in step with the structs this package +// actually serialises onto the SSE stream. They are pure parity checks — +// production does not validate frames per request (see handler/stream.go: +// only the terminal frame's data payload is checked, against the +// endpoint's own response schema). + +func TestSSEFrameSchemaParity(t *testing.T) { + spec := mustSpec(t) + ctx := context.Background() + now := time.Now().UTC() + + t.Run("chat step + terminal frames from a live reporter", func(t *testing.T) { + rec, r := newRecorderReporter(t, "chat") + r.Report(ctx, Event{Stage: StageThinking, Message: "Drafting a reply…", Timestamp: now}) + r.Report(ctx, Event{Stage: StageCompleted, Data: map[string]any{ + "output": map[string]any{"role": "ASSISTANT", "content": "hi"}, + "metadata": map[string]any{"responseTimeMs": 12}, + }}) + + frames := parseSSEFrames(t, rec.Body.String()) + mustValidate(t, spec, "SseProgressStep", frameByEvent(t, frames, "thinking").data) + mustValidate(t, spec, "SseChatTerminalFrame", frameByEvent(t, frames, "completed").data) + }) + + t.Run("chat failed terminal frame", func(t *testing.T) { + rec, r := newRecorderReporter(t, "chat") + r.Report(ctx, Event{ + Stage: StageFailed, + Error: "boom", + Message: "We couldn't generate a response.", + ErrorInfo: &ErrorInfo{Title: "Chat failed", Message: "We couldn't generate a response.", Trace: "boom"}, + }) + frames := parseSSEFrames(t, rec.Body.String()) + mustValidate(t, spec, "SseChatTerminalFrame", frameByEvent(t, frames, "failed").data) + }) + + t.Run("evaluate step + terminal frames from a live reporter", func(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + r.Report(ctx, Event{Stage: StageEvaluating, Message: "Checking your working…", Timestamp: now}) + r.Report(ctx, Event{Stage: StageCompleted, Data: map[string]any{ + "feedback": []map[string]any{{"feedbackId": "fb-1", "message": "ok"}}, + }}) + + frames := parseSSEFrames(t, rec.Body.String()) + mustValidate(t, spec, "SseProgressStep", frameByEvent(t, frames, "evaluating").data) + mustValidate(t, spec, "SseEvaluateTerminalFrame", frameByEvent(t, frames, "completed").data) + }) + + t.Run("evaluate failed terminal frame", func(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + r.Report(ctx, Event{ + Stage: StageFailed, + Error: "boom", + Message: "We couldn't evaluate your answer.", + ErrorInfo: &ErrorInfo{Title: "Evaluation failed", Message: "We couldn't evaluate your answer.", Trace: "boom"}, + }) + frames := parseSSEFrames(t, rec.Body.String()) + mustValidate(t, spec, "SseEvaluateTerminalFrame", frameByEvent(t, frames, "failed").data) + }) +} + +// TestSSEStepSchema_CoversEveryStage asserts every Stage this package can +// put on a step frame validates against SseProgressStep — a canary if a +// stage is added, or an enum is later added to the schema without it. +func TestSSEStepSchema_CoversEveryStage(t *testing.T) { + spec := mustSpec(t) + for _, stage := range []Stage{ + StagePreparing, StageStarting, StageEvaluating, StageThinking, + StageCompleted, StageFailed, StageProgress, + } { + step := sseStep{Stage: string(stage), Message: "x", Timestamp: time.Now().UTC()} + mustValidate(t, spec, "SseProgressStep", toMap(t, step)) + } +} + +// TestSSEEnvelopeStructTagsMatchSchema builds the envelope structs +// directly — real JSON tags, real field set — and asserts each (a) +// satisfies its schema and (b) emits no field the schema doesn't +// document. Catches a struct field rename/retag that skips the spec. +func TestSSEEnvelopeStructTagsMatchSchema(t *testing.T) { + spec := mustSpec(t) + now := time.Now().UTC() + step := sseStep{Stage: "starting", Message: "Starting…", Timestamp: now} + + cases := []struct { + schema string + payload any + }{ + {"SseProgressStep", step}, + {"SseProgressStep", sseStep{Stage: "thinking", Timestamp: now, Data: map[string]any{"k": "v"}}}, + {"SseChatTerminalFrame", sseChatEnvelope{ + Output: map[string]any{"role": "ASSISTANT", "content": "hi"}, + Metadata: map[string]any{"responseTimeMs": 12}, + Steps: []sseStep{step}, + }}, + {"SseChatTerminalFrame", sseChatEnvelope{ + Steps: []sseStep{step}, Error: &ErrorInfo{Title: "Chat failed", Message: "failed", Trace: "boom"}, + }}, + {"SseEvaluateTerminalFrame", sseEnvelope{ + Feedback: []map[string]any{{"feedbackId": "fb-1", "message": "ok"}}, + Steps: []sseStep{step}, + }}, + {"SseEvaluateTerminalFrame", sseEnvelope{ + Steps: []sseStep{step}, Error: &ErrorInfo{Title: "Evaluation failed", Message: "failed", Trace: "boom"}, + }}, + } + for _, c := range cases { + m := toMap(t, c.payload) + mustValidate(t, spec, c.schema, m) + assertAllFieldsDocumented(t, spec, c.schema, m) + } +} + +// --- helpers --- + +// mustSpec loads the latest embedded µEd spec, which is 0.1.1-dev — the version +// that carries the Sse* schemas these tests validate against. Canonical 0.1.0 +// deliberately does not define them. +func mustSpec(t *testing.T) *openapi3.T { + t.Helper() + spec, err := server.LoadOpenAPISpec() + if err != nil { + t.Fatalf("LoadOpenAPISpec: %v", err) + } + return spec +} + +func mustValidate(t *testing.T, spec *openapi3.T, schemaName string, payload any) { + t.Helper() + if err := server.ValidateComponentSchema(spec, schemaName, payload); err != nil { + t.Errorf("payload does not satisfy %s: %v\npayload: %+v", schemaName, err, payload) + } +} + +func assertAllFieldsDocumented(t *testing.T, spec *openapi3.T, schemaName string, m map[string]any) { + t.Helper() + ref := spec.Components.Schemas[schemaName] + if ref == nil || ref.Value == nil { + t.Fatalf("component schema %q not found", schemaName) + } + documented := documentedProps(ref.Value) + for k := range m { + if !documented[k] { + t.Errorf("%s: field %q is emitted by the struct but not a documented property", schemaName, k) + } + } +} + +// documentedProps collects every property name a schema documents, +// following allOf composition (the Sse*TerminalFrame schemas merge a +// shared SseTerminalSteps fragment with an inline branch). +func documentedProps(schema *openapi3.Schema) map[string]bool { + out := map[string]bool{} + if schema == nil { + return out + } + for k := range schema.Properties { + out[k] = true + } + for _, sub := range schema.AllOf { + if sub == nil || sub.Value == nil { + continue + } + for k := range documentedProps(sub.Value) { + out[k] = true + } + } + return out +} + +func frameByEvent(t *testing.T, frames []sseFrame, event string) sseFrame { + t.Helper() + for _, f := range frames { + if f.event == event { + return f + } + } + t.Fatalf("no %q frame among %d frames", event, len(frames)) + return sseFrame{} +} + +func toMap(t *testing.T, v any) map[string]any { + t.Helper() + raw, err := json.Marshal(v) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var m map[string]any + if err := json.Unmarshal(raw, &m); err != nil { + t.Fatalf("unmarshal: %v", err) + } + return m +} diff --git a/internal/progress/ssrf.go b/internal/progress/ssrf.go new file mode 100644 index 0000000..3104dc3 --- /dev/null +++ b/internal/progress/ssrf.go @@ -0,0 +1,74 @@ +package progress + +import ( + "context" + "fmt" + "net" + "net/http" + "strings" +) + +// isDisallowedIP reports whether ip must never be a target for an outbound +// progress callback: loopback, link-local (this also covers cloud metadata +// endpoints such as AWS's 169.254.169.254), private (RFC1918/RFC4193), +// unspecified, and multicast addresses. +func isDisallowedIP(ip net.IP) bool { + return ip.IsLoopback() || + ip.IsLinkLocalUnicast() || + ip.IsLinkLocalMulticast() || + ip.IsInterfaceLocalMulticast() || + ip.IsMulticast() || + ip.IsUnspecified() || + ip.IsPrivate() +} + +// hostAllowed reports whether host matches one of the allowed patterns. +// A pattern is either an exact hostname (e.g. "api.example.com") or a +// "*.example.com" wildcard matching any subdomain of example.com (but not +// example.com itself, which must be listed separately if intended). +func hostAllowed(host string, allowed []string) bool { + host = strings.ToLower(strings.TrimSuffix(host, ".")) + for _, pattern := range allowed { + pattern = strings.ToLower(pattern) + if pattern == host { + return true + } + if suffix, ok := strings.CutPrefix(pattern, "*."); ok && strings.HasSuffix(host, "."+suffix) { + return true + } + } + return false +} + +// newSSRFGuardedTransport returns an http.Transport that resolves DNS +// itself and refuses to dial any IP address isDisallowedIP flags, rather +// than trusting the request's literal hostname string. Checking the +// hostname alone would miss the common bypass of pointing an +// innocent-looking domain at a private or link-local address. +func newSSRFGuardedTransport() *http.Transport { + transport := http.DefaultTransport.(*http.Transport).Clone() + + dialer := &net.Dialer{} + transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) { + host, port, err := net.SplitHostPort(addr) + if err != nil { + return nil, err + } + + ips, err := net.DefaultResolver.LookupIP(ctx, "ip", host) + if err != nil { + return nil, err + } + + for _, ip := range ips { + if isDisallowedIP(ip) { + continue + } + return dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port)) + } + + return nil, fmt.Errorf("host %q resolves only to disallowed private/loopback/link-local addresses", host) + } + + return transport +} diff --git a/internal/progress/ssrf_test.go b/internal/progress/ssrf_test.go new file mode 100644 index 0000000..2ef5ce0 --- /dev/null +++ b/internal/progress/ssrf_test.go @@ -0,0 +1,146 @@ +package progress + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "go.uber.org/zap" +) + +func TestIsDisallowedIP(t *testing.T) { + disallowed := []string{ + "127.0.0.1", // loopback + "::1", // loopback (v6) + "169.254.169.254", // link-local: cloud metadata endpoint + "fe80::1", // link-local (v6) + "10.0.0.1", // private RFC1918 + "172.16.0.1", // private RFC1918 + "192.168.1.1", // private RFC1918 + "fc00::1", // private RFC4193 + "0.0.0.0", // unspecified + "224.0.0.1", // multicast + } + for _, s := range disallowed { + ip := net.ParseIP(s) + if ip == nil { + t.Fatalf("failed to parse test IP %q", s) + } + if !isDisallowedIP(ip) { + t.Errorf("expected %q to be disallowed", s) + } + } + + allowed := []string{ + "8.8.8.8", + "1.1.1.1", + "93.184.216.34", + } + for _, s := range allowed { + ip := net.ParseIP(s) + if ip == nil { + t.Fatalf("failed to parse test IP %q", s) + } + if isDisallowedIP(ip) { + t.Errorf("expected %q to be allowed", s) + } + } +} + +func TestHostAllowed(t *testing.T) { + allowed := []string{"api.example.com", "*.example.org"} + + cases := []struct { + host string + want bool + }{ + {"api.example.com", true}, + {"API.EXAMPLE.COM", true}, + {"other.example.com", false}, + {"foo.example.org", true}, + {"a.b.example.org", true}, + {"example.org", false}, // bare domain not covered by wildcard + {"evil.com", false}, + } + + for _, c := range cases { + if got := hostAllowed(c.host, allowed); got != c.want { + t.Errorf("hostAllowed(%q, %v) = %v, want %v", c.host, allowed, got, c.want) + } + } +} + +func TestHTTPFactory_DefaultBlocksLoopbackDelivery(t *testing.T) { + var received bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + received = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{CallbackTimeout: 500 * time.Millisecond}, + Log: zap.NewNop(), + }) + + r, err := f.NewReporter(srv.URL, "corr-1") + if err != nil { + t.Fatalf("expected NewReporter to succeed (block happens at delivery time), got %v", err) + } + + r.Report(context.Background(), Event{Stage: StageEvaluating}) + + if received { + t.Errorf("expected delivery to a loopback address to be blocked by default") + } +} + +func TestHTTPFactory_AllowPrivateNetworks_PermitsLoopbackDelivery(t *testing.T) { + var received bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + received = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{CallbackTimeout: time.Second, AllowPrivateNetworks: true}, + Log: zap.NewNop(), + }) + + r, err := f.NewReporter(srv.URL, "corr-1") + if err != nil { + t.Fatalf("expected NewReporter to succeed, got %v", err) + } + + r.Report(context.Background(), Event{Stage: StageEvaluating}) + + if !received { + t.Errorf("expected delivery to succeed with AllowPrivateNetworks: true") + } +} + +func TestHTTPFactory_AllowedHosts_RejectsUnlistedHost(t *testing.T) { + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{ + CallbackTimeout: time.Second, + AllowedHosts: []string{"good.example.com"}, + }, + Log: zap.NewNop(), + }) + + if _, err := f.NewReporter("https://evil.example.com/hook", "corr-1"); err == nil { + t.Errorf("expected an error for a host not in AllowedHosts") + } + + r, err := f.NewReporter("https://good.example.com/hook", "corr-1") + if err != nil { + t.Fatalf("expected no error for an allowed host, got %v", err) + } + if r == nil { + t.Fatalf("expected a non-nil reporter for an allowed host") + } +} diff --git a/internal/server/module.go b/internal/server/module.go index 7644bed..becb770 100644 --- a/internal/server/module.go +++ b/internal/server/module.go @@ -2,12 +2,30 @@ package server import "go.uber.org/fx" +// HandlerModule provides the shared application HTTP handler chain — the +// per-version OpenAPI specs and the wrapped Mux. Both the standalone server and +// the Lambda adapter depend on it so they serve an identical, identically +// validated handler. +func HandlerModule() fx.Option { + return fx.Module("http-handler", + // provide openapi specs (one per supported µEd version) + fx.Provide(LoadOpenAPISpecs), + // provide the wrapped handler chain + fx.Provide(NewMux), + ) +} + +// Module provides the standalone HTTP server on top of HandlerModule. func Module(config HttpConfig) fx.Option { return fx.Module("server", // provide config fx.Supply(config), - // provide openapi spec + // provide the single latest spec, used by the streaming handler to + // validate SSE terminal frames (Lambda never streams, so its module + // deliberately omits this) fx.Provide(LoadOpenAPISpec), + // provide the shared handler chain + HandlerModule(), // provide server fx.Provide(NewLifecycleServer), // invoke server diff --git a/internal/server/mux.go b/internal/server/mux.go new file mode 100644 index 0000000..20bd110 --- /dev/null +++ b/internal/server/mux.go @@ -0,0 +1,44 @@ +package server + +import ( + "fmt" + "net/http" + + "github.com/getkin/kin-openapi/openapi3" + "go.uber.org/fx" + "go.uber.org/zap" +) + +// MuxParams are the dependency-injected pieces the shared HTTP handler chain is +// built from. +type MuxParams struct { + fx.In + + Specs map[string]*openapi3.T + Handlers []*HttpHandler `group:"handlers"` + Logger *zap.Logger +} + +// Mux is the fully-wrapped application HTTP handler: the route mux, path +// normalisation, and per-version OpenAPI request/response validation. Both the +// standalone server and the Lambda adapter serve this exact chain, so the two +// deployments validate requests and responses identically. +type Mux struct { + http.Handler +} + +// NewMux assembles the shared HTTP handler chain from the registered route +// handlers and the embedded OpenAPI specs. +func NewMux(params MuxParams) (*Mux, error) { + mux := http.NewServeMux() + for _, h := range params.Handlers { + mux.Handle(h.Name, h.Handler) + } + + openAPIMiddleware, err := OpenAPIMiddleware(params.Specs, nil, params.Logger) + if err != nil { + return nil, fmt.Errorf("initialising OpenAPI middleware: %w", err) + } + + return &Mux{Handler: openAPIMiddleware(NormalizePath(mux))}, nil +} diff --git a/internal/server/mux_test.go b/internal/server/mux_test.go new file mode 100644 index 0000000..7d664e6 --- /dev/null +++ b/internal/server/mux_test.go @@ -0,0 +1,73 @@ +package server + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +// TestNewMux_AppliesValidationAndNormalisation proves the shared chain both +// deployments serve actually wraps the route handlers with OpenAPI validation +// and path normalisation. +func TestNewMux_AppliesValidationAndNormalisation(t *testing.T) { + specs, err := LoadOpenAPISpecs() + require.NoError(t, err) + + var gotPath string + evaluate := &HttpHandler{ + Name: "/evaluate", + Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`[]`)) //nolint:errcheck + }), + } + + mux, err := NewMux(MuxParams{ + Specs: specs, + Handlers: []*HttpHandler{evaluate}, + Logger: zap.NewNop(), + }) + require.NoError(t, err) + + t.Run("valid request is normalised and reaches the handler", func(t *testing.T) { + gotPath = "" + body := mustJSON(t, map[string]any{ + "submission": map[string]any{"type": "TEXT", "content": map[string]any{"text": "hi"}}, + }) + req := httptest.NewRequest(http.MethodPost, "/myFunction/evaluate", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + mux.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "/evaluate", gotPath, "NormalizePath should rewrite the prefixed path") + }) + + t.Run("spec-violating request is rejected before the handler", func(t *testing.T) { + gotPath = "" + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader([]byte(`{}`))) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + mux.ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Empty(t, gotPath, "handler must not be reached") + }) + + t.Run("unknown route passes through", func(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, "/not-a-mued-route", nil) + w := httptest.NewRecorder() + mux.ServeHTTP(w, req) + + // no handler registered for this path -> mux 404, but the middleware + // must not have turned it into a 400/500 + assert.Equal(t, http.StatusNotFound, w.Code) + }) +} diff --git a/internal/server/openapi.go b/internal/server/openapi.go index 8d8ed88..0b3ea24 100644 --- a/internal/server/openapi.go +++ b/internal/server/openapi.go @@ -5,39 +5,90 @@ import ( "fmt" "io" "net/http" - "net/http/httptest" + "sort" "github.com/getkin/kin-openapi/openapi3" "github.com/getkin/kin-openapi/openapi3filter" + "github.com/getkin/kin-openapi/routers" "github.com/getkin/kin-openapi/routers/legacy" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/runtime/schema" ) -func LoadOpenAPISpec() (*openapi3.T, error) { +func loadSpec(data []byte) (*openapi3.T, error) { loader := openapi3.NewLoader() loader.IsExternalRefsAllowed = true - spec, err := loader.LoadFromData(schema.OpenAPISpec) + spec, err := loader.LoadFromData(data) if err != nil { - return nil, fmt.Errorf("loading OpenAPI spec: %w", err) + return nil, err } // Skip validation for OpenAPI 3.1.0 — the legacy router validates on NewRouter. return spec, nil } -func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger) (func(http.Handler) http.Handler, error) { - router, err := legacy.NewRouter(spec, - openapi3.IsOpenAPI31OrLater(), - openapi3.AllowExtraSiblingFields("description", "summary"), - ) +// LoadOpenAPISpec loads the latest embedded µEd OpenAPI spec. +func LoadOpenAPISpec() (*openapi3.T, error) { + spec, err := loadSpec(schema.OpenAPISpec) if err != nil { - return nil, fmt.Errorf("creating OpenAPI router: %w", err) + return nil, fmt.Errorf("loading OpenAPI spec: %w", err) + } + return spec, nil +} + +// LoadOpenAPISpecs loads every embedded µEd OpenAPI spec, keyed by version. +func LoadOpenAPISpecs() (map[string]*openapi3.T, error) { + out := make(map[string]*openapi3.T, len(schema.MuEdOpenAPISpecs)) + for version, data := range schema.MuEdOpenAPISpecs { + spec, err := loadSpec(data) + if err != nil { + return nil, fmt.Errorf("loading OpenAPI spec %s: %w", version, err) + } + out[version] = spec + } + return out, nil +} + +// OpenAPIMiddleware validates µEd requests and responses against the OpenAPI +// spec for the version the client is targeting. The spec is selected from the +// X-Api-Version header via resolveVersion — the same resolver the handlers use — +// so a request is validated against exactly the version that will serve it. A +// nil resolveVersion falls back to defaultSpecVersionResolver, which routes off +// the loaded spec versions alone (this package can't import runtime without an +// import cycle via internal/progress). Routes that no selected spec defines +// (e.g. the legacy "/" route) pass through unvalidated. +func OpenAPIMiddleware(specs map[string]*openapi3.T, resolveVersion func(string) string, log *zap.Logger) (func(http.Handler) http.Handler, error) { + if len(specs) == 0 { + return nil, fmt.Errorf("no OpenAPI specs provided") + } + if resolveVersion == nil { + resolveVersion = defaultSpecVersionResolver(specs) + } + + routerByVersion := make(map[string]routers.Router, len(specs)) + for version, spec := range specs { + router, err := legacy.NewRouter(spec, + openapi3.IsOpenAPI31OrLater(), + openapi3.AllowExtraSiblingFields("description", "summary"), + ) + if err != nil { + return nil, fmt.Errorf("creating OpenAPI router for %s: %w", version, err) + } + routerByVersion[version] = router } opts := &openapi3filter.Options{AuthenticationFunc: openapi3filter.NoopAuthenticationFunc} return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + version := resolveVersion(r.Header.Get("X-Api-Version")) + router, ok := routerByVersion[version] + if !ok { + // No spec for the resolved version — cannot validate, pass through. + // The handler still rejects genuinely unsupported versions with a 406. + next.ServeHTTP(w, r) + return + } + route, pathParams, err := router.FindRoute(r) if err != nil { // Not a µEd route — pass through unvalidated @@ -57,19 +108,26 @@ func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger) (func(http.Handler) ht return } - // Capture response for validation - rec := httptest.NewRecorder() - next.ServeHTTP(rec, r) - - // Snapshot body before validation — ValidateResponse drains the buffer. - bodyBytes := rec.Body.Bytes() + // Buffer the response so it can be validated — unless the + // handler streams it (Content-Type: text/event-stream), in + // which case the sniffer has already committed it to the + // client and there is nothing to validate: the filter has no + // model for a frame sequence and buffering would defeat the + // stream. The decision follows what the handler actually did, + // so it can't disagree with the handler's own streaming check. + sniffer := newResponseSniffer(w) + next.ServeHTTP(sniffer, r) + if sniffer.streamed() { + return + } + bodyBytes := sniffer.buf.Bytes() - // Validate response (lenient — log only) + // Validate response respInput := &openapi3filter.ResponseValidationInput{ RequestValidationInput: reqInput, - Status: rec.Code, - Header: rec.Header(), + Status: sniffer.status, + Header: sniffer.Header(), Body: io.NopCloser(bytes.NewReader(bodyBytes)), Options: opts, } @@ -79,12 +137,37 @@ func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger) (func(http.Handler) ht return } - // Forward captured response - for k, v := range rec.Header() { - w.Header()[k] = v - } - w.WriteHeader(rec.Code) + // Forward the buffered response. Headers set by the handler are + // already on w — the sniffer passed w's header map through. + w.WriteHeader(sniffer.status) w.Write(bodyBytes) //nolint:errcheck }) }, nil } + +// defaultSpecVersionResolver resolves an X-Api-Version header against the set +// of loaded spec versions when the caller passes no resolver: an exact match +// wins, an empty header selects the lowest version, and anything else selects +// the highest. With a single embedded spec (the common case) every input +// resolves to that one version. It mirrors runtime.MuEdRegistry.Resolve +// closely enough for validation routing, without importing runtime. +func defaultSpecVersionResolver(specs map[string]*openapi3.T) func(string) string { + versions := make([]string, 0, len(specs)) + for v := range specs { + versions = append(versions, v) + } + sort.Strings(versions) + + return func(requested string) string { + if _, ok := specs[requested]; ok { + return requested + } + if len(versions) == 0 { + return requested + } + if requested == "" { + return versions[0] + } + return versions[len(versions)-1] + } +} diff --git a/internal/server/openapi_test.go b/internal/server/openapi_test.go index 554f354..4052b01 100644 --- a/internal/server/openapi_test.go +++ b/internal/server/openapi_test.go @@ -5,8 +5,10 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "os" "testing" + "github.com/getkin/kin-openapi/openapi3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/zap" @@ -18,30 +20,62 @@ func TestLoadOpenAPISpec(t *testing.T) { assert.NotNil(t, spec) } +func TestLoadOpenAPISpecs(t *testing.T) { + specs, err := LoadOpenAPISpecs() + require.NoError(t, err) + require.NotEmpty(t, specs) + assert.Contains(t, specs, "0.1.0") + assert.Contains(t, specs, "0.1.1-dev") + for version, spec := range specs { + assert.NotNilf(t, spec, "spec for %s", version) + } + + // The opt-in SSE progress-streaming surface is a shimmy extension pending + // upstream µEd PRs: it lives in 0.1.1-dev only, never in canonical 0.1.0. + sseSchemas := []string{ + "SseProgressStep", "SseTerminalSteps", "StreamingCapabilities", + "SseChatTerminalFrame", "SseEvaluateTerminalFrame", + } + for _, name := range sseSchemas { + assert.Containsf(t, specs["0.1.1-dev"].Components.Schemas, name, "0.1.1-dev should define %s", name) + assert.NotContainsf(t, specs["0.1.0"].Components.Schemas, name, "canonical 0.1.0 must not define %s", name) + } +} + func TestOpenAPIMiddleware_Init(t *testing.T) { - spec, err := LoadOpenAPISpec() + specs, err := LoadOpenAPISpecs() require.NoError(t, err) - middleware, err := OpenAPIMiddleware(spec, zap.NewNop()) + middleware, err := OpenAPIMiddleware(specs, nil, zap.NewNop()) require.NoError(t, err) assert.NotNil(t, middleware) } +func TestOpenAPIMiddleware_NoSpecs_Errors(t *testing.T) { + _, err := OpenAPIMiddleware(map[string]*openapi3.T{}, nil, zap.NewNop()) + assert.Error(t, err) +} + func TestOpenAPIMiddleware_UnknownRoute_PassesThrough(t *testing.T) { middleware := mustMiddleware(t) - called := false - next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - called = true - w.WriteHeader(http.StatusOK) - }) + for _, version := range []string{"", "0.1.0", "0.2.0", "9.9.9"} { + called := false + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + }) - req := httptest.NewRequest(http.MethodGet, "/not-a-mued-route", nil) - w := httptest.NewRecorder() - middleware(next).ServeHTTP(w, req) + req := httptest.NewRequest(http.MethodGet, "/not-a-mued-route", nil) + if version != "" { + req.Header.Set("X-Api-Version", version) + } + w := httptest.NewRecorder() + middleware(next).ServeHTTP(w, req) - assert.True(t, called, "next handler should be called for unknown route") - assert.Equal(t, http.StatusOK, w.Code) + assert.Truef(t, called, "next handler should be called for unknown route (version %q)", version) + assert.Equal(t, http.StatusOK, w.Code) + } } func TestOpenAPIMiddleware_ValidRequest_ReachesHandler(t *testing.T) { @@ -125,11 +159,11 @@ func TestOpenAPIMiddleware_ValidHealthRequest_ReachesHandler(t *testing.T) { w.Write(mustJSON(t, map[string]any{ //nolint:errcheck "status": "OK", "capabilities": map[string]any{ - "supportsEvaluate": true, + "supportsEvaluate": true, "supportsPreSubmissionFeedback": false, - "supportsFormativeFeedback": true, - "supportsSummativeFeedback": true, - "supportsDataPolicy": "NOT_SUPPORTED", + "supportsFormativeFeedback": true, + "supportsSummativeFeedback": true, + "supportsDataPolicy": "NOT_SUPPORTED", }, })) }) @@ -142,12 +176,232 @@ func TestOpenAPIMiddleware_ValidHealthRequest_ReachesHandler(t *testing.T) { assert.Equal(t, http.StatusOK, w.Code) } -// mustMiddleware loads the real spec and returns the initialised middleware, failing the test on error. +func TestOpenAPIMiddleware_SSEEvaluate_BypassesResponseValidation(t *testing.T) { + middleware := mustMiddleware(t) + + body := mustJSON(t, map[string]any{ + "submission": map[string]any{ + "type": "TEXT", + "content": map[string]any{"text": "hello"}, + }, + }) + + var flushed bool + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // A non-JSON, non-spec body that the buffered path would 500. + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.Write([]byte("event: completed\ndata: {}\n\n")) //nolint:errcheck + if f, ok := w.(http.Flusher); ok { + f.Flush() + flushed = true + } + }) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + w := httptest.NewRecorder() + middleware(next).ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type")) + assert.Equal(t, "event: completed\ndata: {}\n\n", w.Body.String()) + assert.True(t, flushed, "handler should receive a flushable writer") +} + +func TestOpenAPIMiddleware_SSEChat_BypassesResponseValidation(t *testing.T) { + middleware := mustMiddleware(t) + + body := mustJSON(t, map[string]any{ + "messages": []map[string]any{{"role": "USER", "content": "hello"}}, + }) + + var flushed bool + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.Write([]byte("event: completed\ndata: {}\n\n")) //nolint:errcheck + if f, ok := w.(http.Flusher); ok { + f.Flush() + flushed = true + } + }) + + req := httptest.NewRequest(http.MethodPost, "/chat", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + w := httptest.NewRecorder() + middleware(next).ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type")) + assert.Equal(t, "event: completed\ndata: {}\n\n", w.Body.String()) + assert.True(t, flushed, "handler should receive a flushable writer") +} + +func TestOpenAPIMiddleware_SSEEvaluate_RequestStillValidated(t *testing.T) { + middleware := mustMiddleware(t) + + // missing required "submission" + body := mustJSON(t, map[string]any{}) + + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Error("next handler must not be called for invalid request") + }) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + w := httptest.NewRecorder() + middleware(next).ServeHTTP(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) +} + +// A request that asks for an SSE stream but whose handler falls back to a +// buffered JSON response (e.g. streaming not supported in this runtime) +// must still be response-validated — the bypass keys on what the handler +// wrote, not on the request's Accept header. +func TestOpenAPIMiddleware_AcceptSSEButJSONResponse_StillValidated(t *testing.T) { + middleware := mustMiddleware(t) + + body := mustJSON(t, map[string]any{ + "submission": map[string]any{ + "type": "TEXT", + "content": map[string]any{"text": "hello"}, + }, + }) + + // object body: valid JSON but spec requires an array for POST /evaluate 200 + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"unexpected": "object"}`)) //nolint:errcheck + }) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + w := httptest.NewRecorder() + middleware(next).ServeHTTP(w, req) + + assert.Equal(t, http.StatusInternalServerError, w.Code) +} + +// A streamed response is forwarded verbatim and left unvalidated even +// when its body would fail the spec, and the handler still gets a +// flushable writer. +func TestOpenAPIMiddleware_StreamedResponse_ForwardedUnvalidated(t *testing.T) { + middleware := mustMiddleware(t) + + body := mustJSON(t, map[string]any{ + "submission": map[string]any{ + "type": "TEXT", + "content": map[string]any{"text": "hello"}, + }, + }) + + var flushed bool + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + // no explicit WriteHeader — the sniffer must decide on first Write + w.Write([]byte("event: completed\ndata: {\"not\":\"an array\"}\n\n")) //nolint:errcheck + if f, ok := w.(http.Flusher); ok { + f.Flush() + flushed = true + } + }) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + middleware(next).ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "event: completed\ndata: {\"not\":\"an array\"}\n\n", w.Body.String()) + assert.True(t, flushed, "handler should receive a flushable writer") +} + +// TestOpenAPIMiddleware_VersionSelectsSpec proves the middleware validates a +// request against the spec for the version the client targets: the same +// /evaluate body is accepted under v0.1.0 but rejected under the synthetic +// v0.2.0 spec, which additionally requires "extraField". +func TestOpenAPIMiddleware_VersionSelectsSpec(t *testing.T) { + bodyNoExtra := map[string]any{ + "submission": map[string]any{"type": "TEXT", "content": map[string]any{"text": "hi"}}, + } + bodyWithExtra := map[string]any{ + "submission": map[string]any{"type": "TEXT", "content": map[string]any{"text": "hi"}}, + "extraField": "present", + } + + tests := []struct { + name string + version string + body map[string]any + wantCode int + wantHandler bool + }{ + {"v0.1.0 accepts body without extraField", "0.1.0", bodyNoExtra, http.StatusOK, true}, + {"no header resolves to v0.1.0 and accepts", "", bodyNoExtra, http.StatusOK, true}, + {"v0.2.0 rejects body without extraField", "0.2.0", bodyNoExtra, http.StatusBadRequest, false}, + {"v0.2.0 accepts body with extraField", "0.2.0", bodyWithExtra, http.StatusOK, true}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + middleware := mustMiddleware(t) + + called := false + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`[]`)) //nolint:errcheck + }) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mustJSON(t, tc.body))) + req.Header.Set("Content-Type", "application/json") + if tc.version != "" { + req.Header.Set("X-Api-Version", tc.version) + } + w := httptest.NewRecorder() + middleware(next).ServeHTTP(w, req) + + assert.Equal(t, tc.wantCode, w.Code) + assert.Equal(t, tc.wantHandler, called) + }) + } +} + +// mustMiddleware loads the real v0.1.0 spec plus the synthetic v0.2.0 testdata +// spec and returns the initialised middleware, with a resolver that mirrors +// runtime.MuEdRegistry.Resolve for the order [0.1.0, 0.2.0]. func mustMiddleware(t *testing.T) func(http.Handler) http.Handler { t.Helper() - spec, err := LoadOpenAPISpec() + + specs, err := LoadOpenAPISpecs() + require.NoError(t, err) + + data, err := os.ReadFile("testdata/mued_v0.2.0.yml") require.NoError(t, err) - middleware, err := OpenAPIMiddleware(spec, zap.NewNop()) + v020, err := loadSpec(data) + require.NoError(t, err) + specs["0.2.0"] = v020 + + resolve := func(v string) string { + switch v { + case "": + return "0.1.0" // Default(): first registered, pinned + case "0.1.0", "0.2.0": + return v + default: + return "0.2.0" // Latest() + } + } + + middleware, err := OpenAPIMiddleware(specs, resolve, zap.NewNop()) require.NoError(t, err) return middleware } @@ -158,4 +412,4 @@ func mustJSON(t *testing.T, v any) []byte { b, err := json.Marshal(v) require.NoError(t, err) return b -} \ No newline at end of file +} diff --git a/internal/server/response_sniffer.go b/internal/server/response_sniffer.go new file mode 100644 index 0000000..f1a462c --- /dev/null +++ b/internal/server/response_sniffer.go @@ -0,0 +1,83 @@ +package server + +import ( + "bytes" + "net/http" + "strings" +) + +// responseSniffer wraps the real http.ResponseWriter and decides, on the +// handler's first write, whether the response is a Server-Sent Events +// stream (Content-Type: text/event-stream) or a normal buffered response: +// +// - streaming: the status and headers are committed to the real writer +// immediately and every subsequent Write is forwarded straight +// through; Flush delegates to the real writer so frames reach the +// client incrementally. The OpenAPI response filter is skipped — it +// has no model for a frame sequence and buffering would defeat the +// stream. +// - buffered: the body is accumulated in memory so the middleware can +// run ValidateResponse against it before anything is sent. +// +// The choice is driven by what the handler actually did, not by a +// pre-flight guess from the request, so the middleware and the handler +// can never disagree about whether a response is streamed. +type responseSniffer struct { + real http.ResponseWriter + status int + decided bool + stream bool + buf bytes.Buffer +} + +func newResponseSniffer(real http.ResponseWriter) *responseSniffer { + return &responseSniffer{real: real, status: http.StatusOK} +} + +func (s *responseSniffer) Header() http.Header { return s.real.Header() } + +func (s *responseSniffer) WriteHeader(code int) { + if s.decided { + return + } + s.status = code + s.decide() +} + +func (s *responseSniffer) Write(p []byte) (int, error) { + if !s.decided { + s.decide() + } + if s.stream { + return s.real.Write(p) + } + return s.buf.Write(p) +} + +// Flush forwards to the real writer only once the response has been +// identified as a stream; for a buffered response it is a no-op — the +// body is still being collected for validation. +func (s *responseSniffer) Flush() { + if !s.stream { + return + } + if f, ok := s.real.(http.Flusher); ok { + f.Flush() + } +} + +func (s *responseSniffer) decide() { + s.decided = true + s.stream = strings.Contains( + strings.ToLower(s.real.Header().Get("Content-Type")), + "text/event-stream", + ) + if s.stream { + s.real.WriteHeader(s.status) + } +} + +// streamed reports whether the handler wrote a Server-Sent Events +// response that has already been committed to the client, so the +// middleware has nothing left to validate or forward. +func (s *responseSniffer) streamed() bool { return s.decided && s.stream } diff --git a/internal/server/server.go b/internal/server/server.go index 6a94ea5..bbea100 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -6,7 +6,6 @@ import ( "net" "net/http" - "github.com/getkin/kin-openapi/openapi3" "go.uber.org/fx" "go.uber.org/zap" "golang.org/x/net/http2" @@ -19,10 +18,9 @@ type HttpServerParams struct { Context context.Context Config HttpConfig - Spec *openapi3.T + Mux *Mux - Handlers []*HttpHandler `group:"handlers"` - Logger *zap.Logger + Logger *zap.Logger } type HttpServer struct { @@ -33,19 +31,8 @@ type HttpServer struct { log *zap.Logger } -func NewHttpServer(params HttpServerParams) (*HttpServer, error) { - mux := http.NewServeMux() - - for _, handler := range params.Handlers { - mux.Handle(handler.Name, handler.Handler) - } - - var handler http.Handler = NormalizePath(mux) - openAPIMiddleware, err := OpenAPIMiddleware(params.Spec, params.Logger) - if err != nil { - return nil, fmt.Errorf("initialising OpenAPI middleware: %w", err) - } - handler = openAPIMiddleware(handler) +func NewHttpServer(params HttpServerParams) *HttpServer { + var handler http.Handler = params.Mux if params.Config.H2c { handler = h2c.NewHandler(handler, &http2.Server{}) } @@ -61,14 +48,11 @@ func NewHttpServer(params HttpServerParams) (*HttpServer, error) { port: params.Config.Port, server: server, log: params.Logger, - }, nil + } } -func NewLifecycleServer(params HttpServerParams, lc fx.Lifecycle) (*HttpServer, error) { - server, err := NewHttpServer(params) - if err != nil { - return nil, err - } +func NewLifecycleServer(params HttpServerParams, lc fx.Lifecycle) *HttpServer { + server := NewHttpServer(params) lc.Append(fx.Hook{ OnStart: func(ctx context.Context) error { go server.Serve(ctx) @@ -78,7 +62,7 @@ func NewLifecycleServer(params HttpServerParams, lc fx.Lifecycle) (*HttpServer, return server.Shutdown(ctx) }, }) - return server, nil + return server } func (s *HttpServer) Serve(context.Context) error { diff --git a/internal/server/testdata/mued_v0.2.0.yml b/internal/server/testdata/mued_v0.2.0.yml new file mode 100644 index 0000000..e16112f --- /dev/null +++ b/internal/server/testdata/mued_v0.2.0.yml @@ -0,0 +1,43 @@ +# Synthetic µEd spec used only by openapi_test.go to exercise per-version spec +# selection. It is a deliberately trimmed OpenAPI 3.1.0 document whose only +# meaningful difference from v0.1.0 is that POST /evaluate additionally requires +# an "extraField" property. +openapi: 3.1.0 +info: + title: Synthetic µEd test spec + version: 0.2.0 +paths: + /evaluate: + post: + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - submission + - extraField + properties: + submission: + type: object + extraField: + type: string + responses: + '200': + description: OK + content: + application/json: + schema: + type: array + items: + type: object + /evaluate/health: + get: + responses: + '200': + description: OK + content: + application/json: + schema: + type: object diff --git a/internal/server/validate_body.go b/internal/server/validate_body.go new file mode 100644 index 0000000..1a4140d --- /dev/null +++ b/internal/server/validate_body.go @@ -0,0 +1,93 @@ +package server + +import ( + "encoding/json" + "fmt" + + "github.com/getkin/kin-openapi/openapi3" +) + +// ValidateResponseBody checks payload against the 200 application/json +// schema of the given operation in the spec. It gives the SSE terminal +// frame — whose payload mirrors the non-streaming response body — the +// same schema guarantee the buffered path gets from the OpenAPI response +// filter (which can't run on a streamed response). +// +// A nil spec means "no schema available" and returns nil, so callers that +// may run without the spec loaded (e.g. under AWS Lambda) need no extra +// guard. +func ValidateResponseBody(spec *openapi3.T, operationID string, payload any) error { + if spec == nil { + return nil + } + + schema, err := responseSchemaFor(spec, operationID) + if err != nil { + return err + } + + return validateAgainstSchema(schema, payload) +} + +// ValidateComponentSchema checks payload against the named schema in +// components/schemas. Like ValidateResponseBody, a nil spec is a no-op. +// It is used by the SSE frame parity test to keep the hand-written +// Sse* schemas in step with the structs progress emits. +func ValidateComponentSchema(spec *openapi3.T, schemaName string, payload any) error { + if spec == nil { + return nil + } + + if spec.Components == nil || spec.Components.Schemas == nil { + return fmt.Errorf("spec has no component schemas") + } + ref := spec.Components.Schemas[schemaName] + if ref == nil || ref.Value == nil { + return fmt.Errorf("component schema %q not found", schemaName) + } + + return validateAgainstSchema(ref.Value, payload) +} + +// validateAgainstSchema round-trips payload through JSON so VisitJSON +// sees the generic shapes it expects (map[string]any, []any, float64) +// rather than concrete Go types such as []map[string]any. +func validateAgainstSchema(schema *openapi3.Schema, payload any) error { + raw, err := json.Marshal(payload) + if err != nil { + return fmt.Errorf("encoding response payload: %w", err) + } + var decoded any + if err := json.Unmarshal(raw, &decoded); err != nil { + return fmt.Errorf("decoding response payload: %w", err) + } + return schema.VisitJSON(decoded) +} + +// responseSchemaFor returns the 200 application/json schema for the +// operation with the given operationId. +func responseSchemaFor(spec *openapi3.T, operationID string) (*openapi3.Schema, error) { + if spec.Paths == nil { + return nil, fmt.Errorf("spec has no paths") + } + for _, item := range spec.Paths.Map() { + for _, op := range item.Operations() { + if op == nil || op.OperationID != operationID { + continue + } + if op.Responses == nil { + return nil, fmt.Errorf("operation %q has no responses", operationID) + } + resp := op.Responses.Status(200) + if resp == nil || resp.Value == nil { + return nil, fmt.Errorf("operation %q has no 200 response", operationID) + } + mt := resp.Value.Content.Get("application/json") + if mt == nil || mt.Schema == nil || mt.Schema.Value == nil { + return nil, fmt.Errorf("operation %q 200 response has no application/json schema", operationID) + } + return mt.Schema.Value, nil + } + } + return nil, fmt.Errorf("operation %q not found in spec", operationID) +} diff --git a/internal/server/validate_body_test.go b/internal/server/validate_body_test.go new file mode 100644 index 0000000..82fbe79 --- /dev/null +++ b/internal/server/validate_body_test.go @@ -0,0 +1,52 @@ +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateResponseBody_NilSpec(t *testing.T) { + assert.NoError(t, ValidateResponseBody(nil, "chat", map[string]any{"anything": true})) +} + +func TestValidateResponseBody_Chat(t *testing.T) { + spec, err := LoadOpenAPISpec() + require.NoError(t, err) + + valid := map[string]any{ + "output": map[string]any{"role": "ASSISTANT", "content": "hello"}, + "metadata": nil, + } + assert.NoError(t, ValidateResponseBody(spec, "chat", valid)) + + // role not in the Message enum + badRole := map[string]any{ + "output": map[string]any{"role": "ROBOT", "content": "hello"}, + } + assert.Error(t, ValidateResponseBody(spec, "chat", badRole)) + + // missing required "output" + assert.Error(t, ValidateResponseBody(spec, "chat", map[string]any{"metadata": nil})) +} + +func TestValidateResponseBody_EvaluateSubmission(t *testing.T) { + spec, err := LoadOpenAPISpec() + require.NoError(t, err) + + feedback := []map[string]any{ + {"feedbackId": "fb-1", "message": "looks good"}, + } + assert.NoError(t, ValidateResponseBody(spec, "evaluateSubmission", feedback)) + + // 200 schema is an array, not an object + assert.Error(t, ValidateResponseBody(spec, "evaluateSubmission", map[string]any{"nope": true})) +} + +func TestValidateResponseBody_UnknownOperation(t *testing.T) { + spec, err := LoadOpenAPISpec() + require.NoError(t, err) + + assert.Error(t, ValidateResponseBody(spec, "noSuchOperation", map[string]any{})) +} diff --git a/runtime/chat.go b/runtime/chat.go index 9c3c93f..09cbb0d 100644 --- a/runtime/chat.go +++ b/runtime/chat.go @@ -15,18 +15,13 @@ type ChatResponse struct { Data map[string]any } -type MuEdChatRole string - -const ( - MuEdChatRoleUser MuEdChatRole = "USER" - MuEdChatRoleAssistant MuEdChatRole = "ASSISTANT" - MuEdChatRoleSystem MuEdChatRole = "SYSTEM" - MuEdChatRoleTool MuEdChatRole = "TOOL" -) - +// MuEdChatMessage is one entry in a chat request's messages array. Role +// (USER / ASSISTANT / SYSTEM / TOOL per the µEd spec) is passed straight +// through to the worker, never inspected by shimmy, so it stays an +// untyped string like the other freeform chat fields. type MuEdChatMessage struct { - Role MuEdChatRole `json:"role"` - Content string `json:"content"` + Role string `json:"role"` + Content string `json:"content"` } // MuEdChatRequest is the request body for the chat endpoint. Only messages @@ -44,6 +39,11 @@ type MuEdChatRequest struct { User map[string]any `json:"user,omitempty"` Context map[string]any `json:"context,omitempty"` Configuration map[string]any `json:"configuration,omitempty"` + + // CallbackUrl, when set, receives out-of-band progress events for this + // chat request, exactly as on /evaluate. Part of the µEd request + // contract, not a shim-specific field. + CallbackUrl *string `json:"callbackUrl,omitempty"` } type MuEdChatHealthStatus string @@ -105,7 +105,12 @@ func MuEdToChatResponse(result map[string]any) (map[string]any, error) { // this passes the worker's capabilities through largely as-is — it only // fills in the spec's required keys/defaults and normalises nil slices to // empty ones so they serialise as [] not null. -func MuEdToChatHealthResponse(result map[string]any) map[string]any { +// +// SSE progress streaming is the exception: it is a shimmy-layer capability, +// not the worker's, so supportsStreaming/supportedProgressStages are set +// from streamingEnabled (shimmy's streaming build + config), overriding +// anything the worker reported. +func MuEdToChatHealthResponse(result map[string]any, streamingEnabled bool) map[string]any { status, _ := result["status"].(string) if status == "" { status = string(MuEdChatHealthStatusOK) @@ -126,6 +131,8 @@ func MuEdToChatHealthResponse(result map[string]any) map[string]any { capabilities[key] = []string{} } } + capabilities["supportsStreaming"] = streamingEnabled + capabilities["supportedProgressStages"] = chatProgressStages resp := map[string]any{ "status": status, @@ -139,3 +146,15 @@ func MuEdToChatHealthResponse(result map[string]any) map[string]any { } return resp } + +// chatProgressStages is the set of SseProgressStep.stage values a /chat +// SSE stream can emit, advertised via capabilities.supportedProgressStages. +// The literals mirror the progress.Stage* constants; they are inlined here +// to keep the runtime package free of a dependency on internal/progress. +var chatProgressStages = []string{ + "preparing", + "starting", + "thinking", + "completed", + "failed", +} diff --git a/runtime/chat_test.go b/runtime/chat_test.go index 0e6c79c..47ab05f 100644 --- a/runtime/chat_test.go +++ b/runtime/chat_test.go @@ -15,7 +15,7 @@ import ( func TestMuEdBuildChatRequest_Valid(t *testing.T) { req := runtime.MuEdChatRequest{ Messages: []runtime.MuEdChatMessage{ - {Role: runtime.MuEdChatRoleUser, Content: "hello"}, + {Role: "USER", Content: "hello"}, }, } body, err := runtime.MuEdBuildChatRequest(req) @@ -44,7 +44,7 @@ func TestMuEdBuildChatRequest_NilMessages(t *testing.T) { func TestMuEdBuildChatRequest_OptionalFieldsOmitted(t *testing.T) { req := runtime.MuEdChatRequest{ Messages: []runtime.MuEdChatMessage{ - {Role: runtime.MuEdChatRoleUser, Content: "hi"}, + {Role: "USER", Content: "hi"}, }, } body, err := runtime.MuEdBuildChatRequest(req) @@ -60,7 +60,7 @@ func TestMuEdBuildChatRequest_OptionalFieldsOmitted(t *testing.T) { func TestMuEdBuildChatRequest_ConversationIDIncluded(t *testing.T) { req := runtime.MuEdChatRequest{ - Messages: []runtime.MuEdChatMessage{{Role: runtime.MuEdChatRoleUser, Content: "hi"}}, + Messages: []runtime.MuEdChatMessage{{Role: "USER", Content: "hi"}}, ConversationID: "abc-123", } body, err := runtime.MuEdBuildChatRequest(req) @@ -85,7 +85,7 @@ func TestMuEdBuildChatRequest_UserPassedThroughIntact(t *testing.T) { }, } req := runtime.MuEdChatRequest{ - Messages: []runtime.MuEdChatMessage{{Role: runtime.MuEdChatRoleUser, Content: "hi"}}, + Messages: []runtime.MuEdChatMessage{{Role: "USER", Content: "hi"}}, User: user, } body, err := runtime.MuEdBuildChatRequest(req) @@ -120,7 +120,7 @@ func TestMuEdBuildChatRequest_ContextPassedThroughIntact(t *testing.T) { }, } req := runtime.MuEdChatRequest{ - Messages: []runtime.MuEdChatMessage{{Role: runtime.MuEdChatRoleUser, Content: "hi"}}, + Messages: []runtime.MuEdChatMessage{{Role: "USER", Content: "hi"}}, Context: context, } body, err := runtime.MuEdBuildChatRequest(req) @@ -208,7 +208,7 @@ func TestMuEdToChatHealthResponse_Valid(t *testing.T) { "statusMessage": "partially degraded", "version": "1.2.3", } - resp := runtime.MuEdToChatHealthResponse(result) + resp := runtime.MuEdToChatHealthResponse(result, true) assert.Equal(t, "DEGRADED", resp["status"]) assert.Equal(t, "partially degraded", resp["statusMessage"]) assert.Equal(t, "1.2.3", resp["version"]) @@ -221,29 +221,34 @@ func TestMuEdToChatHealthResponse_Valid(t *testing.T) { assert.Equal(t, []string{}, capabilities["supportedLanguages"]) assert.Equal(t, []string{}, capabilities["supportedModels"]) assert.Equal(t, []string{}, capabilities["supportedAPIVersions"]) + // SSE progress streaming is shimmy-authoritative, driven by the arg. + assert.Equal(t, true, capabilities["supportsStreaming"]) + assert.Contains(t, capabilities["supportedProgressStages"], "thinking") } func TestMuEdToChatHealthResponse_CapabilitiesPassedThroughIntact(t *testing.T) { // The worker is authoritative on its own capabilities (unlike evaluate, // which hardcodes them) — arbitrary worker-supplied keys must survive. + // SSE progress streaming is the exception: it is a shimmy-layer + // capability, so the worker's supportsStreaming is overridden. result := map[string]any{ "status": "OK", "capabilities": map[string]any{ "supportsChat": true, "supportsUserPreferences": true, - "supportsStreaming": false, + "supportsStreaming": true, "supportsDataPolicy": "PARTIAL", "supportedLanguages": []any{"en", "de"}, "supportedModels": []any{"gpt-4o"}, "supportedAPIVersions": []any{"0.1.0"}, }, } - resp := runtime.MuEdToChatHealthResponse(result) + resp := runtime.MuEdToChatHealthResponse(result, false) capabilities, ok := resp["capabilities"].(map[string]any) require.True(t, ok) assert.Equal(t, true, capabilities["supportsChat"]) assert.Equal(t, true, capabilities["supportsUserPreferences"]) - assert.Equal(t, false, capabilities["supportsStreaming"]) + assert.Equal(t, false, capabilities["supportsStreaming"], "shimmy overrides the worker's streaming flag") assert.Equal(t, "PARTIAL", capabilities["supportsDataPolicy"]) assert.Equal(t, []any{"en", "de"}, capabilities["supportedLanguages"]) assert.Equal(t, []any{"gpt-4o"}, capabilities["supportedModels"]) @@ -251,12 +256,12 @@ func TestMuEdToChatHealthResponse_CapabilitiesPassedThroughIntact(t *testing.T) } func TestMuEdToChatHealthResponse_DefaultsStatusOK(t *testing.T) { - resp := runtime.MuEdToChatHealthResponse(map[string]any{}) + resp := runtime.MuEdToChatHealthResponse(map[string]any{}, false) assert.Equal(t, "OK", resp["status"]) } func TestMuEdToChatHealthResponse_DefaultsMissingCapabilities(t *testing.T) { - resp := runtime.MuEdToChatHealthResponse(map[string]any{}) + resp := runtime.MuEdToChatHealthResponse(map[string]any{}, false) capabilities, ok := resp["capabilities"].(map[string]any) require.True(t, ok) assert.Equal(t, false, capabilities["supportsChat"]) @@ -264,7 +269,7 @@ func TestMuEdToChatHealthResponse_DefaultsMissingCapabilities(t *testing.T) { } func TestMuEdToChatHealthResponse_NilSlicesDefaultToEmpty(t *testing.T) { - resp := runtime.MuEdToChatHealthResponse(map[string]any{}) + resp := runtime.MuEdToChatHealthResponse(map[string]any{}, false) raw, err := json.Marshal(resp) require.NoError(t, err) diff --git a/runtime/evaluate.go b/runtime/evaluate.go index bc9e868..101fa9b 100644 --- a/runtime/evaluate.go +++ b/runtime/evaluate.go @@ -34,10 +34,23 @@ type MuEdEvaluateRequest struct { Task *MuEdTask `json:"task"` Configuration *MuEdConfiguration `json:"configuration"` PreSubmissionFeedback *MuEdPreSubmissionFeedback `json:"preSubmissionFeedback"` + + // CallbackUrl is the µEd spec's optional HTTPS callback URL (see + // https://mued.org/spec, EvaluateRequest.callbackUrl). The spec + // describes it for asynchronous final-result delivery (the service + // may return 202 Accepted and POST the result here later); shimmy + // doesn't implement that 202 flow, but reuses this same field as the + // target for progress events, since both describe "send updates + // about this request to this URL" and a caller shouldn't need a + // shimmy-specific header for something the spec already defines. + CallbackUrl *string `json:"callbackUrl"` } -// MuEdToHealthResponse converts a legacy runtime health result to muEd format. -func MuEdToHealthResponse(result map[string]any) map[string]any { +// MuEdToHealthResponse converts a legacy runtime health result to muEd +// format. streamingEnabled is shimmy's own opt-in SSE progress-streaming +// capability for this deployment (streaming build + config enabled); it +// is advertised verbatim as capabilities.supportsStreaming. +func MuEdToHealthResponse(result map[string]any, streamingEnabled bool) map[string]any { status := "DEGRADED" if passed, ok := result["tests_passed"].(bool); ok && passed { status = "OK" @@ -50,11 +63,26 @@ func MuEdToHealthResponse(result map[string]any) map[string]any { "supportsFormativeFeedback": true, "supportsSummativeFeedback": false, "supportsDataPolicy": "NOT_SUPPORTED", - "supportedAPIVersions": SupportedMuEdVersions, + "supportedAPIVersions": SupportedMuEdVersions(), + "supportsStreaming": streamingEnabled, + "supportedProgressStages": evaluateProgressStages, }, } } +// evaluateProgressStages is the set of SseProgressStep.stage values an +// /evaluate SSE stream can emit, advertised via +// capabilities.supportedProgressStages. The literals mirror the +// progress.Stage* constants; they are inlined here to keep the runtime +// package free of a dependency on internal/progress. +var evaluateProgressStages = []string{ + "preparing", + "starting", + "evaluating", + "completed", + "failed", +} + func muEdContentKey(t MuEdSubmissionType) string { switch t { case MuEdMath: diff --git a/runtime/module.go b/runtime/module.go index fa46f32..d6dbea7 100644 --- a/runtime/module.go +++ b/runtime/module.go @@ -10,6 +10,9 @@ func Module(config Config) fx.Option { // provide runtime config fx.Supply(config), + // provide the µEd version adapter registry + fx.Provide(DefaultMuEdRegistry), + // provide runtime fx.Provide(NewLifecycleRuntime), diff --git a/runtime/mued_adapter.go b/runtime/mued_adapter.go new file mode 100644 index 0000000..8d3cd5a --- /dev/null +++ b/runtime/mued_adapter.go @@ -0,0 +1,128 @@ +package runtime + +// MuEdAdapter translates between one specific µEd API version's wire format and +// the legacy worker protocol. Exactly one implementation is registered per +// supported version; the HTTP handlers resolve the client's X-Api-Version to an +// adapter and drive it, staying version-agnostic themselves. +// +// Decode (not just transform) sits behind this interface on purpose: a future +// version with a different request shape owns its own json.Unmarshal target and +// its own preview-detection rule without the handlers changing. +type MuEdAdapter interface { + // Version is the µEd API version this adapter implements, e.g. "0.1.0". + Version() string + + // DecodeEvaluate parses a POST /evaluate request body into the legacy + // worker request map and the command to run — CommandEvaluate or + // CommandPreview. The adapter owns preview detection. + DecodeEvaluate(body []byte) (legacy map[string]any, command Command, err error) + + // EncodeEvaluateFeedback converts a legacy worker result into the µEd + // feedback array for the command DecodeEvaluate returned. + EncodeEvaluateFeedback(command Command, result map[string]any) ([]map[string]any, error) + + // EncodeHealth converts a legacy health result into the µEd health response. + // streamingEnabled is shimmy's own SSE progress-streaming capability for + // this deployment; the adapter folds it into the advertised capabilities. + EncodeHealth(legacyResult map[string]any, streamingEnabled bool) map[string]any + + // DecodeChat parses a POST /chat request body into the worker request map. + DecodeChat(body []byte) (map[string]any, error) + + // EncodeChat converts a worker chat result into the µEd chat response. + EncodeChat(result map[string]any) (map[string]any, error) + + // EncodeChatHealth converts a worker chat health result into the µEd chat + // health response. streamingEnabled is shimmy's own SSE progress-streaming + // capability for this deployment, overlaid on the worker's reported + // capabilities. + EncodeChatHealth(result map[string]any, streamingEnabled bool) map[string]any + + // SupportsStreaming reports whether this µEd version's contract defines the + // opt-in SSE progress-streaming response surface (the text/event-stream + // media type on /evaluate and /chat). Handlers only stream when this is + // true, so a client negotiating a version without that surface always gets + // the buffered JSON body. + SupportsStreaming() bool +} + +// MuEdRegistry holds the µEd version adapters known to the process, in +// registration order (oldest first). +type MuEdRegistry struct { + order []string + byVersion map[string]MuEdAdapter +} + +// NewMuEdRegistry returns an empty registry. +func NewMuEdRegistry() *MuEdRegistry { + return &MuEdRegistry{byVersion: map[string]MuEdAdapter{}} +} + +// Register adds an adapter. Registering a version again replaces the earlier +// adapter but keeps its position in the order. +func (r *MuEdRegistry) Register(a MuEdAdapter) { + v := a.Version() + if _, seen := r.byVersion[v]; !seen { + r.order = append(r.order, v) + } + r.byVersion[v] = a +} + +// Versions returns the supported versions in registration order. +func (r *MuEdRegistry) Versions() []string { + out := make([]string, len(r.order)) + copy(out, r.order) + return out +} + +// Supports reports whether version is registered. +func (r *MuEdRegistry) Supports(version string) bool { + _, ok := r.byVersion[version] + return ok +} + +// Latest is the most recently registered version, or "" when the registry is +// empty. Used only as the value stamped on a 406 response for an unsupported +// version. +func (r *MuEdRegistry) Latest() string { + if len(r.order) == 0 { + return "" + } + return r.order[len(r.order)-1] +} + +// Default is the version used when a client sends no X-Api-Version header: the +// first registered version. Pinned deliberately — it does not track Latest, so +// registering a newer version never silently moves header-less clients onto new +// semantics. Bump it in its own change. +func (r *MuEdRegistry) Default() string { + if len(r.order) == 0 { + return "" + } + return r.order[0] +} + +// Resolve maps a requested version to a concrete supported one: the default +// version when requested is empty, the request itself when supported, otherwise +// the latest supported version. +func (r *MuEdRegistry) Resolve(requested string) string { + if requested == "" { + return r.Default() + } + if r.Supports(requested) { + return requested + } + return r.Latest() +} + +// Adapter returns the adapter for an already-resolved version, or nil. +func (r *MuEdRegistry) Adapter(version string) MuEdAdapter { + return r.byVersion[version] +} + +// defaultMuEdRegistry is the process-wide registry. Version adapter files +// register into it from their init(); see mued_v0_1_0.go. +var defaultMuEdRegistry = NewMuEdRegistry() + +// DefaultMuEdRegistry returns the process-wide µEd version registry. +func DefaultMuEdRegistry() *MuEdRegistry { return defaultMuEdRegistry } diff --git a/runtime/mued_adapter_test.go b/runtime/mued_adapter_test.go new file mode 100644 index 0000000..9c25f49 --- /dev/null +++ b/runtime/mued_adapter_test.go @@ -0,0 +1,105 @@ +package runtime_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/lambda-feedback/shimmy/runtime" +) + +// fakeAdapter is a minimal MuEdAdapter used to exercise multi-version registry +// behaviour without a second real µEd version. +type fakeAdapter struct{ version string } + +func (a fakeAdapter) Version() string { return a.version } +func (a fakeAdapter) DecodeEvaluate([]byte) (map[string]any, runtime.Command, error) { + return map[string]any{"from": a.version}, runtime.CommandEvaluate, nil +} +func (a fakeAdapter) EncodeEvaluateFeedback(runtime.Command, map[string]any) ([]map[string]any, error) { + return []map[string]any{{"from": a.version}}, nil +} +func (a fakeAdapter) EncodeHealth(map[string]any, bool) map[string]any { + return map[string]any{"from": a.version} +} +func (a fakeAdapter) DecodeChat([]byte) (map[string]any, error) { + return map[string]any{"from": a.version}, nil +} +func (a fakeAdapter) EncodeChat(map[string]any) (map[string]any, error) { + return map[string]any{"from": a.version}, nil +} +func (a fakeAdapter) EncodeChatHealth(map[string]any, bool) map[string]any { + return map[string]any{"from": a.version} +} +func (a fakeAdapter) SupportsStreaming() bool { return false } + +func TestMuEdRegistry_OrderAndResolution(t *testing.T) { + reg := runtime.NewMuEdRegistry() + reg.Register(fakeAdapter{version: "0.1.0"}) + reg.Register(fakeAdapter{version: "0.2.0"}) + reg.Register(fakeAdapter{version: "0.3.0"}) + + assert.Equal(t, []string{"0.1.0", "0.2.0", "0.3.0"}, reg.Versions()) + assert.Equal(t, "0.1.0", reg.Default(), "default is the first registered version, pinned") + assert.Equal(t, "0.3.0", reg.Latest()) + + assert.True(t, reg.Supports("0.2.0")) + assert.False(t, reg.Supports("9.9.9")) + + assert.Equal(t, "0.1.0", reg.Resolve(""), "empty request resolves to the pinned default") + assert.Equal(t, "0.2.0", reg.Resolve("0.2.0"), "supported request resolves to itself") + assert.Equal(t, "0.3.0", reg.Resolve("9.9.9"), "unsupported request resolves to latest") +} + +func TestMuEdRegistry_Adapter(t *testing.T) { + reg := runtime.NewMuEdRegistry() + reg.Register(fakeAdapter{version: "0.1.0"}) + reg.Register(fakeAdapter{version: "9.9.9"}) + + got := reg.Adapter("9.9.9") + require.NotNil(t, got) + assert.Equal(t, "9.9.9", got.Version()) + + feedback, err := got.EncodeEvaluateFeedback(runtime.CommandEvaluate, nil) + require.NoError(t, err) + require.Len(t, feedback, 1) + assert.Equal(t, "9.9.9", feedback[0]["from"]) + + assert.Nil(t, reg.Adapter("0.5.0"), "unknown version has no adapter") +} + +func TestMuEdRegistry_ReregisterKeepsPosition(t *testing.T) { + reg := runtime.NewMuEdRegistry() + reg.Register(fakeAdapter{version: "0.1.0"}) + reg.Register(fakeAdapter{version: "0.2.0"}) + reg.Register(fakeAdapter{version: "0.1.0"}) // replace, don't reorder + + assert.Equal(t, []string{"0.1.0", "0.2.0"}, reg.Versions()) +} + +func TestDefaultMuEdRegistry_RegisteredVersions(t *testing.T) { + reg := runtime.DefaultMuEdRegistry() + + assert.Equal(t, []string{"0.1.0", "0.1.1-dev"}, reg.Versions()) + assert.Equal(t, []string{"0.1.0", "0.1.1-dev"}, runtime.SupportedMuEdVersions()) + assert.True(t, runtime.MuEdIsVersionSupported("0.1.0")) + assert.True(t, runtime.MuEdIsVersionSupported("0.1.1-dev")) + assert.False(t, runtime.MuEdIsVersionSupported("99.0.0")) + + assert.Equal(t, "0.1.0", reg.Default(), "0.1.0 stays the pinned default") + assert.Equal(t, "0.1.1-dev", reg.Latest()) + + assert.Equal(t, "0.1.0", runtime.MuEdResolveVersion(""), "header-less clients stay on 0.1.0") + assert.Equal(t, "0.1.0", runtime.MuEdResolveVersion("0.1.0")) + assert.Equal(t, "0.1.1-dev", runtime.MuEdResolveVersion("0.1.1-dev")) + assert.Equal(t, "0.1.1-dev", runtime.MuEdResolveVersion("99.0.0"), "unsupported resolves to latest") + + require.NotNil(t, reg.Adapter("0.1.0")) + assert.Equal(t, "0.1.0", reg.Adapter("0.1.0").Version()) + require.NotNil(t, reg.Adapter("0.1.1-dev")) + assert.Equal(t, "0.1.1-dev", reg.Adapter("0.1.1-dev").Version()) + + assert.False(t, reg.Adapter("0.1.0").SupportsStreaming()) + assert.True(t, reg.Adapter("0.1.1-dev").SupportsStreaming()) +} diff --git a/runtime/mued_v0_1_0.go b/runtime/mued_v0_1_0.go new file mode 100644 index 0000000..f311f50 --- /dev/null +++ b/runtime/mued_v0_1_0.go @@ -0,0 +1,65 @@ +package runtime + +import ( + "encoding/json" + "fmt" +) + +// muEdV010 is the MuEdAdapter for µEd API version 0.1.0. Every method delegates +// to the package-level transform functions in evaluate.go / chat.go, so 0.1.0 +// behaviour is exactly what it was before the adapter layer existed. +type muEdV010 struct{} + +var _ MuEdAdapter = muEdV010{} + +func init() { + defaultMuEdRegistry.Register(muEdV010{}) +} + +func (muEdV010) Version() string { return "0.1.0" } + +func (muEdV010) DecodeEvaluate(body []byte) (map[string]any, Command, error) { + var req MuEdEvaluateRequest + if err := json.Unmarshal(body, &req); err != nil { + return nil, "", fmt.Errorf("invalid request body") + } + + if req.PreSubmissionFeedback != nil && req.PreSubmissionFeedback.Enabled { + legacy, err := MuEdBuildLegacyPreviewRequest(req) + return legacy, CommandPreview, err + } + + legacy, err := MuEdBuildLegacyEvaluateRequest(req) + return legacy, CommandEvaluate, err +} + +func (muEdV010) EncodeEvaluateFeedback(command Command, result map[string]any) ([]map[string]any, error) { + if command == CommandPreview { + return MuEdToPreviewFeedback(result), nil + } + return MuEdToEvaluateFeedback(result), nil +} + +func (muEdV010) EncodeHealth(legacyResult map[string]any, streamingEnabled bool) map[string]any { + return MuEdToHealthResponse(legacyResult, streamingEnabled) +} + +func (muEdV010) DecodeChat(body []byte) (map[string]any, error) { + var req MuEdChatRequest + if err := json.Unmarshal(body, &req); err != nil { + return nil, fmt.Errorf("invalid request body") + } + return MuEdBuildChatRequest(req) +} + +func (muEdV010) EncodeChat(result map[string]any) (map[string]any, error) { + return MuEdToChatResponse(result) +} + +func (muEdV010) EncodeChatHealth(result map[string]any, streamingEnabled bool) map[string]any { + return MuEdToChatHealthResponse(result, streamingEnabled) +} + +// SupportsStreaming is false: canonical µEd 0.1.0 has no text/event-stream +// response surface. shimmy's SSE progress streaming is offered from 0.1.1-dev on. +func (muEdV010) SupportsStreaming() bool { return false } diff --git a/runtime/mued_v0_1_0_test.go b/runtime/mued_v0_1_0_test.go new file mode 100644 index 0000000..5d52561 --- /dev/null +++ b/runtime/mued_v0_1_0_test.go @@ -0,0 +1,118 @@ +package runtime_test + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/lambda-feedback/shimmy/runtime" +) + +// The v0.1.0 adapter must be a pure delegation to the package-level transform +// functions — these tests are the byte-for-byte regression guard for that. + +func v010(t *testing.T) runtime.MuEdAdapter { + t.Helper() + a := runtime.DefaultMuEdRegistry().Adapter("0.1.0") + require.NotNil(t, a) + return a +} + +func mustJSON(t *testing.T, v any) string { + t.Helper() + b, err := json.Marshal(v) + require.NoError(t, err) + return string(b) +} + +func TestMuEdV010_DecodeEvaluate_MatchesFreeFunctions(t *testing.T) { + a := v010(t) + + evalReq := runtime.MuEdEvaluateRequest{ + Submission: runtime.MuEdSubmission{Type: runtime.MuEdMath, Content: map[string]any{"expression": "x^2"}}, + Task: &runtime.MuEdTask{ReferenceSolution: map[string]any{"expression": "x^2"}}, + } + evalBody := mustJSON(t, evalReq) + + gotLegacy, gotCmd, err := a.DecodeEvaluate([]byte(evalBody)) + require.NoError(t, err) + assert.Equal(t, runtime.CommandEvaluate, gotCmd) + wantLegacy, err := runtime.MuEdBuildLegacyEvaluateRequest(evalReq) + require.NoError(t, err) + assert.Equal(t, wantLegacy, gotLegacy) + + previewReq := runtime.MuEdEvaluateRequest{ + Submission: runtime.MuEdSubmission{Type: runtime.MuEdMath, Content: map[string]any{"expression": "x^2"}}, + PreSubmissionFeedback: &runtime.MuEdPreSubmissionFeedback{Enabled: true}, + } + previewBody := mustJSON(t, previewReq) + + gotLegacy, gotCmd, err = a.DecodeEvaluate([]byte(previewBody)) + require.NoError(t, err) + assert.Equal(t, runtime.CommandPreview, gotCmd) + wantLegacy, err = runtime.MuEdBuildLegacyPreviewRequest(previewReq) + require.NoError(t, err) + assert.Equal(t, wantLegacy, gotLegacy) +} + +func TestMuEdV010_DecodeEvaluate_Errors(t *testing.T) { + a := v010(t) + + _, _, err := a.DecodeEvaluate([]byte("not json")) + assert.Error(t, err) + + missingRef := mustJSON(t, runtime.MuEdEvaluateRequest{ + Submission: runtime.MuEdSubmission{Type: runtime.MuEdMath, Content: map[string]any{"expression": "x^2"}}, + }) + _, _, err = a.DecodeEvaluate([]byte(missingRef)) + assert.Error(t, err, "missing task.referenceSolution is a decode error") +} + +func TestMuEdV010_EncodeEvaluateFeedback_MatchesFreeFunctions(t *testing.T) { + a := v010(t) + + evalResult := map[string]any{"is_correct": true, "feedback": "Well done"} + gotFb, err := a.EncodeEvaluateFeedback(runtime.CommandEvaluate, evalResult) + require.NoError(t, err) + assert.Equal(t, runtime.MuEdToEvaluateFeedback(evalResult), gotFb) + + previewResult := map[string]any{"preview": map[string]any{"latex": "x^{2}"}} + gotFb, err = a.EncodeEvaluateFeedback(runtime.CommandPreview, previewResult) + require.NoError(t, err) + assert.Equal(t, runtime.MuEdToPreviewFeedback(previewResult), gotFb) +} + +func TestMuEdV010_EncodeHealth_MatchesFreeFunction(t *testing.T) { + a := v010(t) + + for _, passed := range []bool{true, false} { + result := map[string]any{"tests_passed": passed} + assert.Equal(t, runtime.MuEdToHealthResponse(result, false), a.EncodeHealth(result, false)) + } +} + +func TestMuEdV010_Chat_MatchesFreeFunctions(t *testing.T) { + a := v010(t) + + chatReq := runtime.MuEdChatRequest{Messages: []runtime.MuEdChatMessage{{Role: "USER", Content: "hi"}}} + gotData, err := a.DecodeChat([]byte(mustJSON(t, chatReq))) + require.NoError(t, err) + wantData, err := runtime.MuEdBuildChatRequest(chatReq) + require.NoError(t, err) + assert.Equal(t, wantData, gotData) + + _, err = a.DecodeChat([]byte(`{"messages":[]}`)) + assert.Error(t, err, "empty messages is a decode error") + + chatResult := map[string]any{"output": map[string]any{"role": "ASSISTANT", "content": "hello"}} + gotResp, err := a.EncodeChat(chatResult) + require.NoError(t, err) + wantResp, err := runtime.MuEdToChatResponse(chatResult) + require.NoError(t, err) + assert.Equal(t, wantResp, gotResp) + + healthResult := map[string]any{} + assert.Equal(t, runtime.MuEdToChatHealthResponse(healthResult, false), a.EncodeChatHealth(healthResult, false)) +} diff --git a/runtime/runtime.go b/runtime/runtime.go index 742ee38..2d25afa 100644 --- a/runtime/runtime.go +++ b/runtime/runtime.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution" + "github.com/lambda-feedback/shimmy/internal/progress" ) // Runtime is the interface for a runtime. @@ -49,6 +50,11 @@ type RuntimeParams struct { // Config is the config for the underlying runtime manager Config Config + // Progress configures worker-authored progress event delivery (the + // EVAL_PROGRESS_URL side-channel). Provided by handler.Module, shared + // with the outbound callbackUrl delivery configuration. + Progress progress.Config + // Log is the logger to use for the runtime Log *zap.Logger } @@ -56,9 +62,10 @@ type RuntimeParams struct { // NewRuntime creates a new runtime. func NewRuntime(params RuntimeParams) (Runtime, error) { dispatcher, err := execution.NewDispatcher(Params{ - Context: params.Context, - Config: params.Config, - Log: params.Log, + Context: params.Context, + Config: params.Config, + Progress: params.Progress, + Log: params.Log, }) if err != nil { return nil, err diff --git a/runtime/schema/openapi.go b/runtime/schema/openapi.go index b771890..a81ec1b 100644 --- a/runtime/schema/openapi.go +++ b/runtime/schema/openapi.go @@ -1,6 +1,73 @@ package schema -import _ "embed" +import ( + "embed" + "fmt" + "sort" + "strings" +) -//go:embed mued_v0.1.0.yml -var OpenAPISpec []byte +// muEdSpecFS holds every embedded µEd OpenAPI spec. Files are named +// mued_v.yml; adding a new version is a matter of dropping in another +// such file — no Go change is required here. +// +//go:embed mued_v*.yml +var muEdSpecFS embed.FS + +// MuEdOpenAPISpecs maps µEd API version -> raw OpenAPI spec bytes, discovered +// from the embedded mued_v.yml files at package load. +var MuEdOpenAPISpecs = mustLoadMuEdSpecs() + +// OpenAPISpec is the latest embedded µEd OpenAPI spec, retained for callers that +// still expect a single spec blob. +var OpenAPISpec = MuEdOpenAPISpecs[LatestMuEdSpecVersion()] + +func mustLoadMuEdSpecs() map[string][]byte { + entries, err := muEdSpecFS.ReadDir(".") + if err != nil { + panic(fmt.Sprintf("reading embedded µEd specs: %v", err)) + } + + out := make(map[string][]byte) + for _, e := range entries { + name := e.Name() + if !strings.HasPrefix(name, "mued_v") || !strings.HasSuffix(name, ".yml") { + continue + } + version := strings.TrimSuffix(strings.TrimPrefix(name, "mued_v"), ".yml") + data, err := muEdSpecFS.ReadFile(name) + if err != nil { + panic(fmt.Sprintf("reading embedded µEd spec %s: %v", name, err)) + } + out[version] = data + } + + if len(out) == 0 { + panic("no embedded µEd OpenAPI specs found") + } + return out +} + +// MuEdSpecVersions returns the embedded spec versions in ascending order. +// Ordering is lexical, which is sufficient while versions stay single-digit; +// revisit if a component ever reaches double digits. +// +// The lexical sort also treats a pre-release tag as newer than its base +// release: "0.1.1-dev" sorts after "0.1.0", so the dev spec is the "latest" +// and drives the single-spec callers (LoadOpenAPISpec, MuEdHandler.Spec) that +// need its SSE schemas. Note a future real "0.1.1" would sort *before* +// "0.1.1-dev" — revisit this ordering when the dev tag is promoted. +func MuEdSpecVersions() []string { + versions := make([]string, 0, len(MuEdOpenAPISpecs)) + for v := range MuEdOpenAPISpecs { + versions = append(versions, v) + } + sort.Strings(versions) + return versions +} + +// LatestMuEdSpecVersion returns the highest embedded spec version. +func LatestMuEdSpecVersion() string { + versions := MuEdSpecVersions() + return versions[len(versions)-1] +} diff --git a/runtime/version.go b/runtime/version.go index 1763f37..161103f 100644 --- a/runtime/version.go +++ b/runtime/version.go @@ -1,19 +1,19 @@ package runtime -var SupportedMuEdVersions = []string{"0.1.0"} +// SupportedMuEdVersions returns the µEd API versions this build supports, in +// registration order (oldest first). Backed by DefaultMuEdRegistry. +func SupportedMuEdVersions() []string { + return defaultMuEdRegistry.Versions() +} +// MuEdIsVersionSupported reports whether the given µEd API version is supported. func MuEdIsVersionSupported(version string) bool { - for _, v := range SupportedMuEdVersions { - if v == version { - return true - } - } - return false + return defaultMuEdRegistry.Supports(version) } +// MuEdResolveVersion maps a requested µEd API version to a concrete supported +// one: the default version when requested is empty, the request itself when +// supported, otherwise the latest supported version. func MuEdResolveVersion(requested string) string { - if MuEdIsVersionSupported(requested) { - return requested - } - return SupportedMuEdVersions[len(SupportedMuEdVersions)-1] + return defaultMuEdRegistry.Resolve(requested) }