From 3a939ea0b2e08c6d32e8ae959b3e356e2f742bf8 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 14:30:38 +0100 Subject: [PATCH 01/28] Add progress reporting framework with HTTP callback support Implement a unified progress reporting framework that supports optional HTTP callbacks: - Define `Event` and `Reporter` abstractions in `internal/progress`. - Add `HTTPFactory` for building per-request HTTP callback reporters. - Introduce `Emit` convenience method to attach/report progress events via context. - Update supervisor and handler logic to emit lifecycle events. - Include comprehensive unit tests for reliability and correctness. --- config/config.go | 8 +- handler/module.go | 9 +- handler/mued.go | 80 ++++++-- handler/mued_test.go | 171 +++++++++++++++++- internal/execution/supervisor/supervisor.go | 17 ++ .../execution/supervisor/supervisor_test.go | 96 ++++++++++ internal/progress/event.go | 52 ++++++ internal/progress/factory.go | 81 +++++++++ internal/progress/factory_test.go | 67 +++++++ internal/progress/http_reporter.go | 115 ++++++++++++ internal/progress/http_reporter_test.go | 123 +++++++++++++ internal/progress/reporter.go | 44 +++++ internal/progress/reporter_test.go | 46 +++++ 13 files changed, 890 insertions(+), 19 deletions(-) create mode 100644 internal/progress/event.go create mode 100644 internal/progress/factory.go create mode 100644 internal/progress/factory_test.go create mode 100644 internal/progress/http_reporter.go create mode 100644 internal/progress/http_reporter_test.go create mode 100644 internal/progress/reporter.go create mode 100644 internal/progress/reporter_test.go diff --git a/config/config.go b/config/config.go index 020d38e..fc90d9b 100644 --- a/config/config.go +++ b/config/config.go @@ -1,6 +1,9 @@ package config -import "github.com/lambda-feedback/shimmy/runtime" +import ( + "github.com/lambda-feedback/shimmy/internal/progress" + "github.com/lambda-feedback/shimmy/runtime" +) type MessageEncoding string @@ -25,4 +28,7 @@ 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"` } diff --git a/handler/module.go b/handler/module.go index a58f29f..434dbda 100644 --- a/handler/module.go +++ b/handler/module.go @@ -1,6 +1,11 @@ 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" +) func Module() fx.Option { return fx.Module("common", @@ -10,5 +15,7 @@ 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), ) } diff --git a/handler/mued.go b/handler/mued.go index 53c4c73..0527366 100644 --- a/handler/mued.go +++ b/handler/mued.go @@ -10,33 +10,45 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" "github.com/lambda-feedback/shimmy/runtime" ) const muEdVersionHeader = "X-Api-Version" +// Progress-reporting headers. Deliberately distinct from the callbackUrl/ +// X-Request-Id pair documented (but not yet implemented) in the µEd schema +// for a different, unrelated feature (async whole-result delivery). +const ( + progressCallbackURLHeader = "X-Progress-Callback-Url" + progressCorrelationIDHeader = "X-Progress-Correlation-Id" +) + type MuEdHandlerParams struct { fx.In - Handler runtime.Handler - Runtime runtime.Runtime - Config config.Config - Log *zap.Logger + Handler runtime.Handler + Runtime runtime.Runtime + Config config.Config + Log *zap.Logger + ProgressFactory progress.Factory } type MuEdHandler struct { - handler runtime.Handler - runtime runtime.Runtime - config config.Config - log *zap.Logger + handler runtime.Handler + runtime runtime.Runtime + config config.Config + log *zap.Logger + progressFactory progress.Factory } 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, + config: params.Config, + log: params.Log, + progressFactory: params.ProgressFactory, } } @@ -157,9 +169,26 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { Header: header, } - resp := h.handler.Handle(r.Context(), req) + ctx := r.Context() + reporter, err := h.progressFactory.NewReporter( + r.Header.Get(progressCallbackURLHeader), + r.Header.Get(progressCorrelationIDHeader), + ) + if err != nil { + h.log.Warn("invalid progress callback header, disabling progress reporting", zap.Error(err)) + } else if reporter != nil { + ctx = progress.ContextWithReporter(ctx, reporter) + } + + resp := h.handler.Handle(ctx, req) if resp.StatusCode != http.StatusOK { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: string(command), + Message: muEdErrorMessageFromBody(resp.Body), + }) + for k, v := range resp.Header { for _, vv := range v { w.Header().Add(k, vv) @@ -190,12 +219,37 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { feedback = runtime.MuEdToEvaluateFeedback(result) } + progress.Emit(ctx, progress.Event{Stage: progress.StageFeedbackReady, Command: string(command)}) + w.Header().Set("Content-Type", "application/json") w.Header().Set(muEdVersionHeader, version) w.WriteHeader(http.StatusOK) json.NewEncoder(w).Encode(feedback) //nolint:errcheck } +// 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 = "evaluation failed" + + 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) { if !h.checkAuth(w, r) { diff --git a/handler/mued_test.go b/handler/mued_test.go index afb65af..a77285c 100644 --- a/handler/mued_test.go +++ b/handler/mued_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" @@ -39,12 +42,23 @@ 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 the X-Progress-Callback-Url header, +// 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, } } @@ -274,6 +288,155 @@ 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 +} + +func newProgressFactory(t *testing.T, timeout time.Duration) progress.Factory { + t.Helper() + return progress.NewHTTPFactory(progress.HTTPFactoryParams{ + Config: progress.Config{CallbackTimeout: timeout}, + 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(mathEvalBody(t))) + req.Header.Set(progressCallbackURLHeader, srv.URL) + req.Header.Set(progressCorrelationIDHeader, "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, "feedback_ready", evt["stage"]) + assert.Equal(t, "eval", evt["command"]) +} + +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(mathEvalBody(t))) + req.Header.Set(progressCallbackURLHeader, srv.URL) + req.Header.Set(progressCorrelationIDHeader, "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"]) +} + +func TestMuEdServeEvaluate_ProgressCallback_NoHeader_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 progress callback header should mean no callback requests") +} + +func TestMuEdServeEvaluate_ProgressCallback_InvalidURL_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(mathEvalBody(t))) + req.Header.Set(progressCallbackURLHeader, "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(mathEvalBody(t))) + req.Header.Set(progressCallbackURLHeader, srv.URL) + req.Header.Set(progressCorrelationIDHeader, "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) { diff --git a/internal/execution/supervisor/supervisor.go b/internal/execution/supervisor/supervisor.go index f3e6587..1028884 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 { @@ -165,12 +166,28 @@ func (s *WorkerSupervisor) Send( worker, err := s.acquireWorker(ctx) if err != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: method, + Message: "failed to acquire worker", + Error: err.Error(), + }) return nil, fmt.Errorf("failed to acquire worker: %w", err) } + progress.Emit(ctx, progress.Event{Stage: progress.StageWorkerAcquired, Command: method}) // NOTICE: unconventional error handling ahead, as we need // to release the worker before returning the error. + progress.Emit(ctx, progress.Event{Stage: progress.StageRunning, Command: method}) resData, err := worker.Send(ctx, method, data, s.sendParams.Timeout) + if err != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: method, + Message: "worker execution failed", + Error: 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..a52788b 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.StageWorkerAcquired, + progress.StageRunning, + }, 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.StageWorkerAcquired, + progress.StageRunning, + 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..8a77087 --- /dev/null +++ b/internal/progress/event.go @@ -0,0 +1,52 @@ +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 ( + // StageWorkerAcquired indicates a worker is ready to receive work, + // whether it was freshly booted or reused from a warm pool. + StageWorkerAcquired Stage = "worker_acquired" + + // StageRunning indicates the evaluation function is about to be invoked. + StageRunning Stage = "running" + + // StageFeedbackReady indicates feedback has been computed and is about + // to be returned to the caller. + StageFeedbackReady Stage = "feedback_ready" + + // StageFailed indicates a terminal failure at any layer of the pipeline. + StageFailed Stage = "failed" +) + +// 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 == StageFeedbackReady || s == StageFailed +} + +// 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 an optional human-readable note. + Message string + + // Error is populated only for StageFailed. + Error string + + // Data is a free-form extension point, reserved for future events + // (e.g. ones emitted by the evaluation function process itself). + 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..faaaf2b --- /dev/null +++ b/internal/progress/factory.go @@ -0,0 +1,81 @@ +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"` +} + +// 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 +} + +var _ Factory = (*HTTPFactory)(nil) + +// NewHTTPFactory builds a Factory that delivers progress events as +// outbound HTTP POST requests. +// +// The URL supplied to NewReporter is trusted as-is: today the only caller +// of shimmy's /evaluate endpoint is client-backend, already authenticated +// via the shared Auth.Key. If shimmy ever accepts callback URLs from less +// trusted callers, this is the place to add a host allowlist to close off +// the resulting SSRF surface. +func NewHTTPFactory(params HTTPFactoryParams) Factory { + timeout := params.Config.CallbackTimeout + if timeout <= 0 { + timeout = defaultCallbackTimeout + } + + return &HTTPFactory{ + client: &http.Client{}, + timeout: timeout, + log: params.Log, + } +} + +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) + } + + 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..d30599f --- /dev/null +++ b/internal/progress/http_reporter.go @@ -0,0 +1,115 @@ +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. +type payload struct { + CorrelationID string `json:"correlationId"` + Stage Stage `json:"stage"` + Command string `json:"command,omitempty"` + Message string `json:"message,omitempty"` + Error string `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 StageFeedbackReady) 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) { + body, err := json.Marshal(payload{ + CorrelationID: r.correlationID, + Stage: evt.Stage, + Command: evt.Command, + Message: evt.Message, + Error: evt.Error, + 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..7bf0731 --- /dev/null +++ b/internal/progress/http_reporter_test.go @@ -0,0 +1,123 @@ +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: StageRunning, 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 != StageRunning { + t.Errorf("expected stage %q, got %q", StageRunning, received[0].Stage) + } +} + +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: StageFeedbackReady}) + + 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: StageWorkerAcquired}) + r.Report(context.Background(), Event{Stage: StageRunning}) + + 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: StageRunning}) + 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: StageRunning}) +} 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..13fc9c5 --- /dev/null +++ b/internal/progress/reporter_test.go @@ -0,0 +1,46 @@ +package progress + +import ( + "context" + "testing" +) + +type recordingReporter struct { + events []Event +} + +func (r *recordingReporter) Report(_ context.Context, evt Event) { + r.events = append(r.events, evt) +} + +func TestEmit_NoReporterInContext_NoOp(t *testing.T) { + // must not panic, must not do anything observable + Emit(context.Background(), Event{Stage: StageRunning}) +} + +func TestEmit_WithReporter_DeliversEventAndSetsTimestamp(t *testing.T) { + r := &recordingReporter{} + ctx := ContextWithReporter(context.Background(), r) + + Emit(ctx, Event{Stage: StageWorkerAcquired, Command: "eval"}) + + if len(r.events) != 1 { + t.Fatalf("expected 1 event, got %d", len(r.events)) + } + evt := r.events[0] + if evt.Stage != StageWorkerAcquired { + t.Errorf("expected stage %q, got %q", StageWorkerAcquired, 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) + } +} From 4e29907dbf15cfb32d7f4f471c1c7e042e0f9c62 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 14:45:40 +0100 Subject: [PATCH 02/28] Refactor progress lifecycle stages and improve messaging clarity --- handler/mued.go | 8 ++++-- handler/mued_test.go | 2 +- internal/execution/supervisor/supervisor.go | 16 +++++++++--- .../execution/supervisor/supervisor_test.go | 8 +++--- internal/progress/event.go | 26 ++++++++++++------- internal/progress/http_reporter.go | 2 +- internal/progress/http_reporter_test.go | 16 ++++++------ internal/progress/reporter_test.go | 8 +++--- 8 files changed, 52 insertions(+), 34 deletions(-) diff --git a/handler/mued.go b/handler/mued.go index 0527366..f6f6893 100644 --- a/handler/mued.go +++ b/handler/mued.go @@ -219,7 +219,11 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { feedback = runtime.MuEdToEvaluateFeedback(result) } - progress.Emit(ctx, progress.Event{Stage: progress.StageFeedbackReady, Command: string(command)}) + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: string(command), + Message: "Feedback is ready.", + }) w.Header().Set("Content-Type", "application/json") w.Header().Set(muEdVersionHeader, version) @@ -230,7 +234,7 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { // 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 = "evaluation failed" + const fallback = "We couldn't evaluate your answer. Please try again." var errBody map[string]any if err := json.Unmarshal(body, &errBody); err != nil { diff --git a/handler/mued_test.go b/handler/mued_test.go index a77285c..2e23a7e 100644 --- a/handler/mued_test.go +++ b/handler/mued_test.go @@ -343,7 +343,7 @@ func TestMuEdServeEvaluate_ProgressCallback_Success(t *testing.T) { require.Len(t, *received, 1) evt := (*received)[0] assert.Equal(t, "corr-1", evt["correlationId"]) - assert.Equal(t, "feedback_ready", evt["stage"]) + assert.Equal(t, "completed", evt["stage"]) assert.Equal(t, "eval", evt["command"]) } diff --git a/internal/execution/supervisor/supervisor.go b/internal/execution/supervisor/supervisor.go index 1028884..9af5872 100644 --- a/internal/execution/supervisor/supervisor.go +++ b/internal/execution/supervisor/supervisor.go @@ -169,22 +169,30 @@ func (s *WorkerSupervisor) Send( progress.Emit(ctx, progress.Event{ Stage: progress.StageFailed, Command: method, - Message: "failed to acquire worker", + Message: "We couldn't start the evaluation. Please try again.", Error: err.Error(), }) return nil, fmt.Errorf("failed to acquire worker: %w", err) } - progress.Emit(ctx, progress.Event{Stage: progress.StageWorkerAcquired, Command: method}) + progress.Emit(ctx, progress.Event{ + Stage: progress.StagePreparing, + Command: method, + Message: "Preparing your evaluation…", + }) // NOTICE: unconventional error handling ahead, as we need // to release the worker before returning the error. - progress.Emit(ctx, progress.Event{Stage: progress.StageRunning, Command: method}) + progress.Emit(ctx, progress.Event{ + Stage: progress.StageEvaluating, + Command: method, + Message: "Evaluating your submission…", + }) resData, err := worker.Send(ctx, method, data, s.sendParams.Timeout) if err != nil { progress.Emit(ctx, progress.Event{ Stage: progress.StageFailed, Command: method, - Message: "worker execution failed", + Message: "Something went wrong while evaluating your answer. Please try again.", Error: err.Error(), }) } diff --git a/internal/execution/supervisor/supervisor_test.go b/internal/execution/supervisor/supervisor_test.go index a52788b..41630c1 100644 --- a/internal/execution/supervisor/supervisor_test.go +++ b/internal/execution/supervisor/supervisor_test.go @@ -300,8 +300,8 @@ func TestSupervisor_Send_EmitsWorkerAcquiredAndRunning(t *testing.T) { assert.NoError(t, err) assert.Equal(t, []progress.Stage{ - progress.StageWorkerAcquired, - progress.StageRunning, + progress.StagePreparing, + progress.StageEvaluating, }, r.stages()) } @@ -339,8 +339,8 @@ func TestSupervisor_Send_EmitsFailed_WhenWorkerSendFails(t *testing.T) { assert.ErrorIs(t, err, assert.AnError) assert.Equal(t, []progress.Stage{ - progress.StageWorkerAcquired, - progress.StageRunning, + progress.StagePreparing, + progress.StageEvaluating, progress.StageFailed, }, r.stages()) } diff --git a/internal/progress/event.go b/internal/progress/event.go index 8a77087..80814c7 100644 --- a/internal/progress/event.go +++ b/internal/progress/event.go @@ -7,16 +7,18 @@ import "time" type Stage string const ( - // StageWorkerAcquired indicates a worker is ready to receive work, - // whether it was freshly booted or reused from a warm pool. - StageWorkerAcquired Stage = "worker_acquired" + // 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. + StagePreparing Stage = "preparing" - // StageRunning indicates the evaluation function is about to be invoked. - StageRunning Stage = "running" + // StageEvaluating indicates the submission is being evaluated. + StageEvaluating Stage = "evaluating" - // StageFeedbackReady indicates feedback has been computed and is about + // StageCompleted indicates feedback has been computed and is about // to be returned to the caller. - StageFeedbackReady Stage = "feedback_ready" + StageCompleted Stage = "completed" // StageFailed indicates a terminal failure at any layer of the pipeline. StageFailed Stage = "failed" @@ -26,7 +28,7 @@ const ( // progress event stream. At most one terminal event is delivered per // Reporter instance. func (s Stage) terminal() bool { - return s == StageFeedbackReady || s == StageFailed + return s == StageCompleted || s == StageFailed } // Event describes a single progress update for an evaluation request. @@ -37,10 +39,14 @@ type Event struct { // Command is the µEd command being processed (e.g. "eval", "preview"). Command string - // Message is an optional human-readable note. + // 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 is populated only for StageFailed. + // 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 // Data is a free-form extension point, reserved for future events diff --git a/internal/progress/http_reporter.go b/internal/progress/http_reporter.go index d30599f..4d6fbb2 100644 --- a/internal/progress/http_reporter.go +++ b/internal/progress/http_reporter.go @@ -56,7 +56,7 @@ func newHTTPReporter( // 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 StageFeedbackReady) is delivered per reporter instance, since both +// 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() { diff --git a/internal/progress/http_reporter_test.go b/internal/progress/http_reporter_test.go index 7bf0731..2d7ee79 100644 --- a/internal/progress/http_reporter_test.go +++ b/internal/progress/http_reporter_test.go @@ -34,7 +34,7 @@ func TestHTTPCallbackReporter_Report_DeliversPayload(t *testing.T) { defer srv.Close() r := newTestReporter(t, srv.URL, time.Second) - r.Report(context.Background(), Event{Stage: StageRunning, Command: "eval"}) + r.Report(context.Background(), Event{Stage: StageEvaluating, Command: "eval"}) mu.Lock() defer mu.Unlock() @@ -44,8 +44,8 @@ func TestHTTPCallbackReporter_Report_DeliversPayload(t *testing.T) { if received[0].CorrelationID != "corr-1" { t.Errorf("expected correlationId %q, got %q", "corr-1", received[0].CorrelationID) } - if received[0].Stage != StageRunning { - t.Errorf("expected stage %q, got %q", StageRunning, received[0].Stage) + if received[0].Stage != StageEvaluating { + t.Errorf("expected stage %q, got %q", StageEvaluating, received[0].Stage) } } @@ -67,7 +67,7 @@ func TestHTTPCallbackReporter_Report_TerminalEventDeliveredOnlyOnce(t *testing.T // 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: StageFeedbackReady}) + r.Report(context.Background(), Event{Stage: StageCompleted}) mu.Lock() defer mu.Unlock() @@ -89,8 +89,8 @@ func TestHTTPCallbackReporter_Report_NonTerminalEventsAllDelivered(t *testing.T) defer srv.Close() r := newTestReporter(t, srv.URL, time.Second) - r.Report(context.Background(), Event{Stage: StageWorkerAcquired}) - r.Report(context.Background(), Event{Stage: StageRunning}) + r.Report(context.Background(), Event{Stage: StagePreparing}) + r.Report(context.Background(), Event{Stage: StageEvaluating}) mu.Lock() defer mu.Unlock() @@ -109,7 +109,7 @@ func TestHTTPCallbackReporter_Report_SlowReceiver_BoundedByTimeout(t *testing.T) r := newTestReporter(t, srv.URL, 20*time.Millisecond) start := time.Now() - r.Report(context.Background(), Event{Stage: StageRunning}) + r.Report(context.Background(), Event{Stage: StageEvaluating}) elapsed := time.Since(start) if elapsed > 150*time.Millisecond { @@ -119,5 +119,5 @@ func TestHTTPCallbackReporter_Report_SlowReceiver_BoundedByTimeout(t *testing.T) 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: StageRunning}) + r.Report(context.Background(), Event{Stage: StageEvaluating}) } diff --git a/internal/progress/reporter_test.go b/internal/progress/reporter_test.go index 13fc9c5..e5dd604 100644 --- a/internal/progress/reporter_test.go +++ b/internal/progress/reporter_test.go @@ -15,21 +15,21 @@ func (r *recordingReporter) Report(_ context.Context, evt Event) { func TestEmit_NoReporterInContext_NoOp(t *testing.T) { // must not panic, must not do anything observable - Emit(context.Background(), Event{Stage: StageRunning}) + Emit(context.Background(), Event{Stage: StageEvaluating}) } func TestEmit_WithReporter_DeliversEventAndSetsTimestamp(t *testing.T) { r := &recordingReporter{} ctx := ContextWithReporter(context.Background(), r) - Emit(ctx, Event{Stage: StageWorkerAcquired, Command: "eval"}) + Emit(ctx, Event{Stage: StagePreparing, Command: "eval"}) if len(r.events) != 1 { t.Fatalf("expected 1 event, got %d", len(r.events)) } evt := r.events[0] - if evt.Stage != StageWorkerAcquired { - t.Errorf("expected stage %q, got %q", StageWorkerAcquired, evt.Stage) + 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) From 3fe2dd4bfe91ac73187497dd0e2bc87ca0372760 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 15:22:37 +0100 Subject: [PATCH 03/28] =?UTF-8?q?Add=20callbackUrl=20support=20for=20progr?= =?UTF-8?q?ess=20events=20in=20=C2=B5Ed=20requests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Integrate callbackUrl field from µEd spec for progress reporting. - Replace progress callback headers with callbackUrl and request ID. - Include evaluation feedback payload in StageCompleted events. - Update tests to reflect callbackUrl usage and validation. --- handler/mued.go | 29 ++++++++++++---------- handler/mued_test.go | 49 +++++++++++++++++++++++++------------- internal/progress/event.go | 7 ++++-- runtime/mued.go | 10 ++++++++ 4 files changed, 64 insertions(+), 31 deletions(-) diff --git a/handler/mued.go b/handler/mued.go index f6f6893..2fa0ba4 100644 --- a/handler/mued.go +++ b/handler/mued.go @@ -16,13 +16,10 @@ import ( const muEdVersionHeader = "X-Api-Version" -// Progress-reporting headers. Deliberately distinct from the callbackUrl/ -// X-Request-Id pair documented (but not yet implemented) in the µEd schema -// for a different, unrelated feature (async whole-result delivery). -const ( - progressCallbackURLHeader = "X-Progress-Callback-Url" - progressCorrelationIDHeader = "X-Progress-Correlation-Id" -) +// muEdRequestIDHeader is the µEd spec's request-tracing header (see +// https://mued.org/spec, X-Request-Id parameter). Progress events reuse it +// as their correlation key, echoing back whatever the caller supplied. +const muEdRequestIDHeader = "X-Request-Id" type MuEdHandlerParams struct { fx.In @@ -169,13 +166,15 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { Header: header, } + var callbackURL string + if muEdReq.CallbackUrl != nil { + callbackURL = *muEdReq.CallbackUrl + } + ctx := r.Context() - reporter, err := h.progressFactory.NewReporter( - r.Header.Get(progressCallbackURLHeader), - r.Header.Get(progressCorrelationIDHeader), - ) + reporter, err := h.progressFactory.NewReporter(callbackURL, r.Header.Get(muEdRequestIDHeader)) if err != nil { - h.log.Warn("invalid progress callback header, disabling progress reporting", zap.Error(err)) + h.log.Warn("invalid callbackUrl, disabling progress reporting", zap.Error(err)) } else if reporter != nil { ctx = progress.ContextWithReporter(ctx, reporter) } @@ -219,10 +218,16 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { feedback = runtime.MuEdToEvaluateFeedback(result) } + // 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") diff --git a/handler/mued_test.go b/handler/mued_test.go index 2e23a7e..b60a50c 100644 --- a/handler/mued_test.go +++ b/handler/mued_test.go @@ -43,7 +43,7 @@ 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 the X-Progress-Callback-Url header, +// 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 { @@ -64,7 +64,12 @@ func newMuEdHandlerWithProgress(h runtime.Handler, r runtime.Runtime, key string 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"}, @@ -74,7 +79,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 } @@ -331,9 +340,8 @@ func TestMuEdServeEvaluate_ProgressCallback_Success(t *testing.T) { mockHandler.On("Handle", mock.Anything, mock.Anything). Return(evalHandlerResponse(true, "Well done")) - req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) - req.Header.Set(progressCallbackURLHeader, srv.URL) - req.Header.Set(progressCorrelationIDHeader, "corr-1") + 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) @@ -345,6 +353,16 @@ func TestMuEdServeEvaluate_ProgressCallback_Success(t *testing.T) { 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) { @@ -360,9 +378,8 @@ func TestMuEdServeEvaluate_ProgressCallback_Failure(t *testing.T) { Body: errorBody, }) - req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) - req.Header.Set(progressCallbackURLHeader, srv.URL) - req.Header.Set(progressCorrelationIDHeader, "corr-2") + 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) @@ -376,7 +393,7 @@ func TestMuEdServeEvaluate_ProgressCallback_Failure(t *testing.T) { assert.Equal(t, "boom", evt["message"]) } -func TestMuEdServeEvaluate_ProgressCallback_NoHeader_Unchanged(t *testing.T) { +func TestMuEdServeEvaluate_ProgressCallback_NoCallbackUrl_Unchanged(t *testing.T) { _, received := newProgressCallbackServer(t, nil) mockHandler := new(MockHandler) @@ -389,16 +406,15 @@ func TestMuEdServeEvaluate_ProgressCallback_NoHeader_Unchanged(t *testing.T) { newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) assert.Equal(t, http.StatusOK, w.Result().StatusCode) - assert.Empty(t, *received, "no progress callback header should mean no callback requests") + assert.Empty(t, *received, "no callbackUrl in the request body should mean no callback requests") } -func TestMuEdServeEvaluate_ProgressCallback_InvalidURL_EvaluationStillSucceeds(t *testing.T) { +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(mathEvalBody(t))) - req.Header.Set(progressCallbackURLHeader, "not-a-url") + 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) @@ -424,9 +440,8 @@ func TestMuEdServeEvaluate_ProgressCallback_SlowReceiver_DoesNotBlockResponse(t mockHandler.On("Handle", mock.Anything, mock.Anything). Return(evalHandlerResponse(true, "Well done")) - req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) - req.Header.Set(progressCallbackURLHeader, srv.URL) - req.Header.Set(progressCorrelationIDHeader, "corr-3") + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + req.Header.Set(muEdRequestIDHeader, "corr-3") w := httptest.NewRecorder() start := time.Now() diff --git a/internal/progress/event.go b/internal/progress/event.go index 80814c7..2629c8a 100644 --- a/internal/progress/event.go +++ b/internal/progress/event.go @@ -49,8 +49,11 @@ type Event struct { // students or teachers directly; show Message instead. Error string - // Data is a free-form extension point, reserved for future events - // (e.g. ones emitted by the evaluation function process itself). + // Data is a free-form extension point. On StageCompleted it carries + // the evaluation's feedback payload (so a callbackUrl-supplying + // caller gets the final result, not just a status ping). Otherwise + // it's reserved for future events, e.g. ones emitted by the + // evaluation function process itself. Data map[string]any // Timestamp is set by Emit, not by callers. diff --git a/runtime/mued.go b/runtime/mued.go index 24e8b2f..ca33c3a 100644 --- a/runtime/mued.go +++ b/runtime/mued.go @@ -34,6 +34,16 @@ 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"` } var SupportedMuEdVersions = []string{"0.1.0"} From 31b5e9bab2d6caf0975419be4e03667a98e3c9c1 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 15:31:02 +0100 Subject: [PATCH 04/28] Add `progress-callback-timeout` flag for configurable progress callback delivery timeout --- README.md | 4 ++++ cmd/root.go | 9 +++++++++ 2 files changed, 13 insertions(+) diff --git a/README.md b/README.md index 2cd6fda..f71eff1 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,10 @@ 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] + function --arg value, -a value [ --arg value, -a value ] additional arguments for to the worker process. [$FUNCTION_ARGS] diff --git a/cmd/root.go b/cmd/root.go index 275c31e..b719d0c 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -43,6 +43,14 @@ 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"}, + }, // shim flags &cli.StringFlag{ Name: "interface", @@ -318,6 +326,7 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { // map cli flags to config fields cliMap := map[string]string{ "auth-key": "auth.key", + "progress-callback-timeout": "progress.callback_timeout", "max-workers": "runtime.max_workers", "command": "runtime.cmd", "cwd": "runtime.cwd", From be31ea788a6d632228b44d7f9a7071c0936c9490 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 15:45:11 +0100 Subject: [PATCH 05/28] Add SSRF protections and enhanced safety controls to progress callback feature - Introduce `--progress-allowed-hosts` flag to restrict allowed callback hostnames. - Add `--progress-allow-private-networks` flag for optional private network access. - Implement automatic request ID generation for traceability and progress correlation. - Expand documentation with guidance on callback URL safety and SSRF prevention. - Update tests and internal logic for new configuration options and request IDs. --- README.md | 73 +++++++++++++++++++++++++++++++++++- cmd/root.go | 45 ++++++++++++++-------- handler/mued.go | 34 +++++++++++++++-- handler/mued_test.go | 66 +++++++++++++++++++++++++++++++- internal/progress/factory.go | 48 ++++++++++++++++++------ 5 files changed, 235 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index f71eff1..24b540f 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,9 @@ GLOBAL OPTIONS: progress - --progress-callback-timeout value the timeout for a single progress callback delivery. (default: 1s) [$PROGRESS_CALLBACK_TIMEOUT] + --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] function @@ -189,6 +191,75 @@ Example request using cases: } ``` +### Progress Events + +The shim also exposes a µEd-compatible endpoint at `POST /evaluate` (see the [µEd spec](https://mued.org/spec)), separate from the legacy `POST /` endpoint documented above. When a client calls `/evaluate` 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. "Evaluating your submission…") 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" +} +``` + +Four stages are emitted, in order, for a successful evaluation: + +| Stage | Meaning | +|-------|---------| +| `preparing` | The evaluation environment is being set up (a worker is ready — freshly booted or reused). | +| `evaluating` | The evaluation function is being invoked. | +| `completed` | Feedback has been computed. `data.feedback` carries the same array returned in the synchronous response body. | +| `failed` | A terminal failure occurred at some stage. `message` is safe to show to an end user; `error` carries raw technical detail for logs only. | + +`completed` and `failed` are terminal — at most one of them is delivered per request, whichever occurs first. + +Example event body: + +```json +{ + "correlationId": "req-7c193f38", + "stage": "evaluating", + "command": "eval", + "message": "Evaluating your submission…", + "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" +} +``` + +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. + ### Communication Channels The shim supports two interface modes, selected with `--interface`: diff --git a/cmd/root.go b/cmd/root.go index b719d0c..ececbc3 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -51,6 +51,19 @@ functions on arbitrary, serverless platforms.` 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"}, + }, // shim flags &cli.StringFlag{ Name: "interface", @@ -325,21 +338,23 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { // map cli flags to config fields cliMap := map[string]string{ - "auth-key": "auth.key", - "progress-callback-timeout": "progress.callback_timeout", - "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", + "auth-key": "auth.key", + "progress-callback-timeout": "progress.callback_timeout", + "progress-allowed-hosts": "progress.allowed_hosts", + "progress-allow-private-networks": "progress.allow_private_networks", + "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", // sandbox "sandbox": "runtime.sandbox.enabled", "sandbox-nsjail-path": "runtime.sandbox.nsjail_path", diff --git a/handler/mued.go b/handler/mued.go index 2fa0ba4..fbcdb03 100644 --- a/handler/mued.go +++ b/handler/mued.go @@ -1,10 +1,12 @@ package handler import ( + "crypto/rand" "encoding/json" "fmt" "io" "net/http" + "time" "go.uber.org/fx" "go.uber.org/zap" @@ -17,10 +19,31 @@ import ( const muEdVersionHeader = "X-Api-Version" // muEdRequestIDHeader is the µEd spec's request-tracing header (see -// https://mued.org/spec, X-Request-Id parameter). Progress events reuse it -// as their correlation key, echoing back whatever the caller supplied. +// 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 @@ -106,6 +129,9 @@ 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 } @@ -172,7 +198,7 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { } ctx := r.Context() - reporter, err := h.progressFactory.NewReporter(callbackURL, r.Header.Get(muEdRequestIDHeader)) + reporter, err := h.progressFactory.NewReporter(callbackURL, requestID) if err != nil { h.log.Warn("invalid callbackUrl, disabling progress reporting", zap.Error(err)) } else if reporter != nil { @@ -261,6 +287,8 @@ func muEdErrorMessageFromBody(body []byte) string { // 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 } diff --git a/handler/mued_test.go b/handler/mued_test.go index b60a50c..65388ac 100644 --- a/handler/mued_test.go +++ b/handler/mued_test.go @@ -325,10 +325,15 @@ func newProgressCallbackServer(t *testing.T, handlerFn http.HandlerFunc) (*httpt 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}, + Config: progress.Config{CallbackTimeout: timeout, AllowPrivateNetworks: true}, Log: zap.NewNop(), }) } @@ -674,3 +679,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/internal/progress/factory.go b/internal/progress/factory.go index faaaf2b..5ed1d18 100644 --- a/internal/progress/factory.go +++ b/internal/progress/factory.go @@ -18,6 +18,21 @@ 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"` } // Factory builds a per-request Reporter from caller-supplied callback @@ -39,9 +54,10 @@ type HTTPFactoryParams struct { } type HTTPFactory struct { - client *http.Client - timeout time.Duration - log *zap.Logger + client *http.Client + timeout time.Duration + log *zap.Logger + allowedHosts []string } var _ Factory = (*HTTPFactory)(nil) @@ -49,21 +65,27 @@ var _ Factory = (*HTTPFactory)(nil) // NewHTTPFactory builds a Factory that delivers progress events as // outbound HTTP POST requests. // -// The URL supplied to NewReporter is trusted as-is: today the only caller -// of shimmy's /evaluate endpoint is client-backend, already authenticated -// via the shared Auth.Key. If shimmy ever accepts callback URLs from less -// trusted callers, this is the place to add a host allowlist to close off -// the resulting SSRF surface. +// 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: &http.Client{}, - timeout: timeout, - log: params.Log, + client: client, + timeout: timeout, + log: params.Log, + allowedHosts: params.Config.AllowedHosts, } } @@ -77,5 +99,9 @@ func (f *HTTPFactory) NewReporter(callbackURL, correlationID string) (Reporter, 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 } From 76ad18bee6f3bcb83deaaf519a920cb3c23a4ae4 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 15:52:09 +0100 Subject: [PATCH 06/28] Add SSRF protections to HTTP progress callbacks - Introduce IP filtering to block private, loopback, and link-local addresses. - Add hostname wildcards for fine-grained allowed host configuration. - Implement custom HTTP transport with DNS-based IP validation. - Add comprehensive unit tests to cover SSRF scenarios and configuration options. --- internal/progress/ssrf.go | 74 +++++++++++++++++ internal/progress/ssrf_test.go | 146 +++++++++++++++++++++++++++++++++ 2 files changed, 220 insertions(+) create mode 100644 internal/progress/ssrf.go create mode 100644 internal/progress/ssrf_test.go 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") + } +} From 02dcc9e004c42acbe80ec124f73d0966b1ea50c3 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 17:47:33 +0100 Subject: [PATCH 07/28] Add SSRF protections to HTTP progress callbacks - Introduce IP filtering to block private, loopback, and link-local addresses. - Add hostname wildcards for fine-grained allowed host configuration. - Implement custom HTTP transport with DNS-based IP validation. - Add comprehensive unit tests to cover SSRF scenarios and configuration options. --- README.md | 38 +++ cmd/root.go | 58 +++-- internal/execution/dispatcher.go | 15 +- .../dispatcher/dispatcher_dedicated.go | 12 +- .../execution/dispatcher/dispatcher_pooled.go | 12 +- .../dispatcher/dispatcher_pooled_test.go | 5 +- internal/execution/supervisor/adapter.go | 35 +-- internal/execution/supervisor/adapter_file.go | 20 +- .../execution/supervisor/adapter_file_test.go | 87 +++++++ internal/execution/supervisor/adapter_rpc.go | 41 +++- .../execution/supervisor/adapter_rpc_test.go | 66 +++++ internal/execution/supervisor/adapter_test.go | 7 +- internal/execution/supervisor/supervisor.go | 7 +- internal/progress/event.go | 13 +- internal/progress/factory.go | 6 + internal/progress/reporter_test.go | 22 +- internal/progress/sidecar.go | 225 ++++++++++++++++++ internal/progress/sidecar_test.go | 195 +++++++++++++++ runtime/runtime.go | 13 +- 19 files changed, 818 insertions(+), 59 deletions(-) create mode 100644 internal/progress/sidecar.go create mode 100644 internal/progress/sidecar_test.go diff --git a/README.md b/README.md index 24b540f..8f905d4 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,9 @@ GLOBAL OPTIONS: --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-min-event-interval value the minimum spacing between worker-authored progress events relayed per evaluation. (default: 200ms) [$PROGRESS_SIDECAR_MIN_EVENT_INTERVAL] function @@ -260,6 +263,39 @@ A rejected callback URL behaves like any other delivery failure: it's logged and > **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. +#### Custom progress events from the evaluation function + +The four stages above are emitted by shimmy itself, around the evaluation function call as a whole — `evaluating` covers the entire invocation as one span. An evaluation 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` 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: an evaluation function can never claim `preparing`, `evaluating`, `completed`, or `failed` — those remain exclusively shim-authored. Custom events are always delivered with `"stage": "progress"`. + +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 after the request has already finished). + +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-min-event-interval` | `PROGRESS_SIDECAR_MIN_EVENT_INTERVAL` | `200ms` | Minimum spacing between relayed events. | + +> **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`: @@ -286,6 +322,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`) @@ -311,6 +348,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/cmd/root.go b/cmd/root.go index ececbc3..3ec0f4d 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -64,6 +64,27 @@ functions on arbitrary, serverless platforms.` 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.DurationFlag{ + Name: "progress-sidecar-min-event-interval", + Usage: "the minimum spacing between worker-authored progress events relayed per evaluation.", + Value: 200 * time.Millisecond, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_MIN_EVENT_INTERVAL"}, + }, // shim flags &cli.StringFlag{ Name: "interface", @@ -338,23 +359,26 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { // map cli flags to config fields cliMap := map[string]string{ - "auth-key": "auth.key", - "progress-callback-timeout": "progress.callback_timeout", - "progress-allowed-hosts": "progress.allowed_hosts", - "progress-allow-private-networks": "progress.allow_private_networks", - "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", + "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-min-event-interval": "progress.sidecar.min_event_interval", + "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", // sandbox "sandbox": "runtime.sandbox.enabled", "sandbox-nsjail-path": "runtime.sandbox.nsjail_path", 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..db6bc09 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.StageProgress, 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 ec40d5b..89374e0 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,15 @@ 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; Unbind closes the narrow window + // between this call returning and the next one starting, so a + // straggling POST from the worker can't be misattributed. + a.sidecar.Bind(method, progress.FromContext(ctx)) + defer a.sidecar.Unbind() + } + var result map[string]any ctx, cancel := context.WithTimeout(ctx, timeout) @@ -181,6 +208,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) } @@ -283,7 +316,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) } @@ -304,6 +337,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..0bfd2f2 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/Unbind 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.Unbind() + + 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.StageProgress, 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 9af5872..f3d6b65 100644 --- a/internal/execution/supervisor/supervisor.go +++ b/internal/execution/supervisor/supervisor.go @@ -75,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 } @@ -100,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) { diff --git a/internal/progress/event.go b/internal/progress/event.go index 2629c8a..c362e25 100644 --- a/internal/progress/event.go +++ b/internal/progress/event.go @@ -22,6 +22,13 @@ const ( // StageFailed indicates a terminal failure at any layer of the pipeline. StageFailed Stage = "failed" + + // StageProgress indicates a custom, evaluation-function-authored + // progress update. Unlike the other stages, these are never emitted + // by shimmy itself — only relayed from a worker's local progress + // side-channel (see Sidecar). A worker cannot claim any other stage; + // the wire contract for that side-channel has no way to set Stage. + StageProgress Stage = "progress" ) // terminal reports whether the stage marks the end of an evaluation's @@ -51,9 +58,9 @@ type Event struct { // Data is a free-form extension point. On StageCompleted it carries // the evaluation's feedback payload (so a callbackUrl-supplying - // caller gets the final result, not just a status ping). Otherwise - // it's reserved for future events, e.g. ones emitted by the - // evaluation function process itself. + // caller gets the final result, not just a status ping). On + // StageProgress 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. diff --git a/internal/progress/factory.go b/internal/progress/factory.go index 5ed1d18..de56695 100644 --- a/internal/progress/factory.go +++ b/internal/progress/factory.go @@ -33,6 +33,12 @@ type Config struct { // 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"` } // Factory builds a per-request Reporter from caller-supplied callback diff --git a/internal/progress/reporter_test.go b/internal/progress/reporter_test.go index e5dd604..3aee6b0 100644 --- a/internal/progress/reporter_test.go +++ b/internal/progress/reporter_test.go @@ -2,17 +2,32 @@ 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}) @@ -24,10 +39,11 @@ func TestEmit_WithReporter_DeliversEventAndSetsTimestamp(t *testing.T) { Emit(ctx, Event{Stage: StagePreparing, Command: "eval"}) - if len(r.events) != 1 { - t.Fatalf("expected 1 event, got %d", len(r.events)) + events := r.recorded() + if len(events) != 1 { + t.Fatalf("expected 1 event, got %d", len(events)) } - evt := r.events[0] + evt := events[0] if evt.Stage != StagePreparing { t.Errorf("expected stage %q, got %q", StagePreparing, evt.Stage) } diff --git a/internal/progress/sidecar.go b/internal/progress/sidecar.go new file mode 100644 index 0000000..2cfd31e --- /dev/null +++ b/internal/progress/sidecar.go @@ -0,0 +1,225 @@ +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 + defaultSidecarMinEventInterval = 200 * 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"` + + // MinEventInterval enforces a minimum spacing between accepted events + // within a span. If unset (<= 0), defaultSidecarMinEventInterval is used. + MinEventInterval time.Duration `conf:"min_event_interval"` +} + +func (c SidecarConfig) withDefaults() SidecarConfig { + if c.MaxBodyBytes <= 0 { + c.MaxBodyBytes = defaultSidecarMaxBodyBytes + } + if c.MaxEventsPerSpan <= 0 { + c.MaxEventsPerSpan = defaultSidecarMaxEventsPerSpan + } + if c.MinEventInterval <= 0 { + c.MinEventInterval = defaultSidecarMinEventInterval + } + 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 claim any stage other than StageProgress, which the sidecar +// hardcodes itself. 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"` +} + +// 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 +} + +// 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.command = command + s.reporter = reporter + s.count = 0 + s.lastSent = time.Time{} +} + +// Unbind detaches the current reporter, 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.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: StageProgress, + 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.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..1fad407 --- /dev/null +++ b/internal/progress/sidecar_test.go @@ -0,0 +1,195 @@ +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 != StageProgress { + t.Errorf("expected stage %q, got %q", StageProgress, 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 != StageProgress { + t.Errorf("worker-supplied stage must be ignored, got %q", events[0].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) { + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, 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_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/runtime/runtime.go b/runtime/runtime.go index a29ce92..8f8ee8f 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. @@ -46,6 +47,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 } @@ -53,9 +59,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 From 65b02c60a9448e97feef84a13d8d0e3d25f558e4 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Wed, 5 Aug 2026 14:16:21 +0100 Subject: [PATCH 08/28] Add unbind grace period to sidecar progress reporting - Introduce `--progress-sidecar-unbind-grace-period` flag with default value of 250ms. - Add `UnbindAfterGrace` method to allow delayed unbinding with generation-safe logic. - Update supervisor adapter to utilize `UnbindAfterGrace` for improved POST handling. --- cmd/root.go | 48 +++++++------ internal/execution/supervisor/adapter_rpc.go | 11 +-- internal/progress/sidecar.go | 72 +++++++++++++++++--- 3 files changed, 96 insertions(+), 35 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 3ec0f4d..7fbf9cc 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -85,6 +85,13 @@ functions on arbitrary, serverless platforms.` 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"}, + }, // shim flags &cli.StringFlag{ Name: "interface", @@ -359,26 +366,27 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { // map cli flags to config fields cliMap := map[string]string{ - "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-min-event-interval": "progress.sidecar.min_event_interval", - "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", + "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-min-event-interval": "progress.sidecar.min_event_interval", + "progress-sidecar-unbind-grace-period": "progress.sidecar.unbind_grace_period", + "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", // sandbox "sandbox": "runtime.sandbox.enabled", "sandbox-nsjail-path": "runtime.sandbox.nsjail_path", diff --git a/internal/execution/supervisor/adapter_rpc.go b/internal/execution/supervisor/adapter_rpc.go index 89374e0..193179b 100644 --- a/internal/execution/supervisor/adapter_rpc.go +++ b/internal/execution/supervisor/adapter_rpc.go @@ -184,11 +184,14 @@ func (a *rpcAdapter) Send( if a.sidecar != nil { // sendLock in the calling supervisor guarantees only one request - // is ever in flight per worker; Unbind closes the narrow window - // between this call returning and the next one starting, so a - // straggling POST from the worker can't be misattributed. + // 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.Unbind() + defer a.sidecar.UnbindAfterGrace() } var result map[string]any diff --git a/internal/progress/sidecar.go b/internal/progress/sidecar.go index 2cfd31e..f300d93 100644 --- a/internal/progress/sidecar.go +++ b/internal/progress/sidecar.go @@ -15,9 +15,10 @@ import ( ) const ( - defaultSidecarMaxBodyBytes int64 = 16 * 1024 - defaultSidecarMaxEventsPerSpan = 50 - defaultSidecarMinEventInterval = 200 * time.Millisecond + defaultSidecarMaxBodyBytes int64 = 16 * 1024 + defaultSidecarMaxEventsPerSpan = 50 + defaultSidecarMinEventInterval = 200 * time.Millisecond + defaultSidecarUnbindGracePeriod = 250 * time.Millisecond ) // SidecarConfig bounds abuse of the worker-authored progress side-channel. @@ -37,6 +38,14 @@ type SidecarConfig struct { // MinEventInterval enforces a minimum spacing between accepted events // within a span. 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 { @@ -49,6 +58,9 @@ func (c SidecarConfig) withDefaults() SidecarConfig { if c.MinEventInterval <= 0 { c.MinEventInterval = defaultSidecarMinEventInterval } + if c.UnbindGracePeriod <= 0 { + c.UnbindGracePeriod = defaultSidecarUnbindGracePeriod + } return c } @@ -80,11 +92,12 @@ type Sidecar struct { listener net.Listener server *http.Server - mu sync.Mutex - command string - reporter Reporter - count int - lastSent time.Time + 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. @@ -129,24 +142,61 @@ 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, 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 +// 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. From fdb7a564eb67afdb43337514f9c0833a4caab381 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Wed, 5 Aug 2026 17:38:45 +0100 Subject: [PATCH 09/28] Allow a burst of closely-spaced worker-authored progress events MinEventInterval's default (200ms) rate-limited a fast evaluation function reporting two checkpoints from compareSets' evaluation function to at most one event per span: even with delivery now serialized on the client side, two closely-spaced report_progress() calls could still both arrive well under any single fixed interval, since arrival timing is governed by local HTTP round-trip cost, not real application-level delay. Add BurstSize (default 5): the first N events in a span bypass MinEventInterval spacing entirely (still bounded by MaxEventsPerSpan), so a handful of legitimate back-to-back checkpoints go through, while MinEventInterval keeps guarding against sustained event spam once the burst is used up. Also lower the MinEventInterval default itself from 200ms to 10ms, since 200ms had no real abuse-prevention basis and was overly aggressive for normal use. --- README.md | 10 +++-- cmd/root.go | 12 +++++- .../execution/supervisor/adapter_rpc_test.go | 6 +-- internal/progress/sidecar.go | 24 ++++++++++-- internal/progress/sidecar_test.go | 39 ++++++++++++++++++- 5 files changed, 79 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 8f905d4..270d231 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,9 @@ GLOBAL OPTIONS: --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-min-event-interval value the minimum spacing between worker-authored progress events relayed per evaluation. (default: 200ms) [$PROGRESS_SIDECAR_MIN_EVENT_INTERVAL] + --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] function @@ -282,7 +284,7 @@ To emit a custom event, `POST` a small JSON body to `EVAL_PROGRESS_URL`: - `data` (object, optional): free-form, passed through as-is. - There is no `stage` field, by design: an evaluation function can never claim `preparing`, `evaluating`, `completed`, or `failed` — those remain exclusively shim-authored. Custom events are always delivered with `"stage": "progress"`. -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 after the request has already finished). +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: @@ -290,7 +292,9 @@ To bound how much an evaluation function (which may be running untrusted, sandbo |------|---------|---------|-------------| | `--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-min-event-interval` | `PROGRESS_SIDECAR_MIN_EVENT_INTERVAL` | `200ms` | Minimum spacing between relayed events. | +| `--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. diff --git a/cmd/root.go b/cmd/root.go index 7fbf9cc..427f48b 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -78,10 +78,17 @@ functions on arbitrary, serverless platforms.` 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.", - Value: 200 * time.Millisecond, + 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"}, }, @@ -372,6 +379,7 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { "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", "max-workers": "runtime.max_workers", diff --git a/internal/execution/supervisor/adapter_rpc_test.go b/internal/execution/supervisor/adapter_rpc_test.go index 0bfd2f2..1357fe6 100644 --- a/internal/execution/supervisor/adapter_rpc_test.go +++ b/internal/execution/supervisor/adapter_rpc_test.go @@ -174,8 +174,8 @@ func TestStdioAdapter_Start_InjectsProgressURL(t *testing.T) { } // TestStdioAdapter_Send_RelaysWorkerProgressEvents exercises the same -// Bind/Unbind path Send uses around the (separately, more fully) tested -// Sidecar, without needing a live RPC round trip - Send itself isn't +// 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) @@ -191,7 +191,7 @@ func TestStdioAdapter_Send_RelaysWorkerProgressEvents(t *testing.T) { // mirrors exactly what rpcAdapter.Send does with a.sidecar a.sidecar.Bind("eval", progress.FromContext(ctx)) - defer a.sidecar.Unbind() + defer a.sidecar.UnbindAfterGrace() resp, err := http.Post(a.sidecar.URL(), "application/json", strings.NewReader(`{"message":"checking correctness"}`)) assert.NoError(t, err) diff --git a/internal/progress/sidecar.go b/internal/progress/sidecar.go index f300d93..7d73f82 100644 --- a/internal/progress/sidecar.go +++ b/internal/progress/sidecar.go @@ -17,7 +17,8 @@ import ( const ( defaultSidecarMaxBodyBytes int64 = 16 * 1024 defaultSidecarMaxEventsPerSpan = 50 - defaultSidecarMinEventInterval = 200 * time.Millisecond + defaultSidecarBurstSize = 5 + defaultSidecarMinEventInterval = 10 * time.Millisecond defaultSidecarUnbindGracePeriod = 250 * time.Millisecond ) @@ -35,8 +36,20 @@ type SidecarConfig struct { // 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 - // within a span. If unset (<= 0), defaultSidecarMinEventInterval is used. + // 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 @@ -55,6 +68,11 @@ func (c SidecarConfig) withDefaults() SidecarConfig { 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 } @@ -264,7 +282,7 @@ func (s *Sidecar) accept() (command string, reporter Reporter, status int) { if s.count >= s.cfg.MaxEventsPerSpan { return "", nil, http.StatusTooManyRequests } - if !s.lastSent.IsZero() && now.Sub(s.lastSent) < s.cfg.MinEventInterval { + if s.count >= s.cfg.BurstSize && !s.lastSent.IsZero() && now.Sub(s.lastSent) < s.cfg.MinEventInterval { return "", nil, http.StatusTooManyRequests } diff --git a/internal/progress/sidecar_test.go b/internal/progress/sidecar_test.go index 1fad407..35548ac 100644 --- a/internal/progress/sidecar_test.go +++ b/internal/progress/sidecar_test.go @@ -139,7 +139,9 @@ func TestSidecar_RateLimit_MaxEventsPerSpan(t *testing.T) { } func TestSidecar_RateLimit_MinEventInterval(t *testing.T) { - s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, MinEventInterval: time.Hour}) + // 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) @@ -154,6 +156,41 @@ func TestSidecar_RateLimit_MinEventInterval(t *testing.T) { } } +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{} From bcfb3c0325954ea364a869f057b209886e96dceb Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Fri, 28 Aug 2026 17:57:59 +0100 Subject: [PATCH 10/28] Add unit tests for progress reporting, SSE handling, and runtime modules - Add test coverage for new SSEReporter behavior and intermediate step deduplication. - Include tests for dependency graph validation in Lambda and standalone runtime modules. - Add multi-reporter tests to verify fan-out behavior and isolated child panics. - Improve overall test reliability with enhanced mocks and structured assertions. --- app/lambda/module.go | 2 + app/lambda/module_test.go | 31 ++ app/standalone/module.go | 2 + app/standalone/module_test.go | 33 ++ cmd/root.go | 16 + handler/evaluate.go | 257 +++++++++++++--- handler/evaluate_stream_test.go | 371 +++++++++++++++++++++++ handler/module.go | 9 + internal/progress/factory.go | 18 ++ internal/progress/multi_reporter.go | 30 ++ internal/progress/multi_reporter_test.go | 45 +++ internal/progress/sse_reporter.go | 170 +++++++++++ internal/progress/sse_reporter_test.go | 258 ++++++++++++++++ internal/server/openapi.go | 22 +- internal/server/openapi_test.go | 100 +++++- internal/server/server.go | 9 +- 16 files changed, 1320 insertions(+), 53 deletions(-) create mode 100644 app/lambda/module_test.go create mode 100644 app/standalone/module_test.go create mode 100644 handler/evaluate_stream_test.go create mode 100644 internal/progress/multi_reporter.go create mode 100644 internal/progress/multi_reporter_test.go create mode 100644 internal/progress/sse_reporter.go create mode 100644 internal/progress/sse_reporter_test.go diff --git a/app/lambda/module.go b/app/lambda/module.go index 1ed820a..6f9bdb5 100644 --- a/app/lambda/module.go +++ b/app/lambda/module.go @@ -14,6 +14,8 @@ 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 server diff --git a/app/lambda/module_test.go b/app/lambda/module_test.go new file mode 100644 index 0000000..9c9f5ea --- /dev/null +++ b/app/lambda/module_test.go @@ -0,0 +1,31 @@ +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. +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..c95fdba --- /dev/null +++ b/app/standalone/module_test.go @@ -0,0 +1,33 @@ +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). 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 03ea365..3eb9712 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -99,6 +99,20 @@ functions on arbitrary, serverless platforms.` Category: "progress", EnvVars: []string{"PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD"}, }, + &cli.BoolFlag{ + Name: "progress-stream-enabled", + Usage: "stream progress back on the /evaluate response as Server-Sent Events for requests that send 'Accept: text/event-stream'. 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", @@ -389,6 +403,8 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { "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", diff --git a/handler/evaluate.go b/handler/evaluate.go index ffc7eda..c29095b 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -1,11 +1,14 @@ package handler import ( + "context" "crypto/rand" "encoding/json" "fmt" "io" "net/http" + "strings" + "sync" "time" "go.uber.org/fx" @@ -48,28 +51,31 @@ func generateRequestID() string { type MuEdHandlerParams struct { fx.In - Handler runtime.Handler - Runtime runtime.Runtime - Config config.Config - Log *zap.Logger - ProgressFactory progress.Factory + Handler runtime.Handler + Runtime runtime.Runtime + Config config.Config + Log *zap.Logger + ProgressFactory progress.Factory + StreamingCapability StreamingCapability } type MuEdHandler struct { - handler runtime.Handler - runtime runtime.Runtime - config config.Config - log *zap.Logger - progressFactory progress.Factory + handler runtime.Handler + runtime runtime.Runtime + config config.Config + log *zap.Logger + progressFactory progress.Factory + streamingCapable bool } func NewMuEdHandler(params MuEdHandlerParams) *MuEdHandler { return &MuEdHandler{ - handler: params.Handler, - runtime: params.Runtime, - config: params.Config, - log: params.Log, - progressFactory: params.ProgressFactory, + handler: params.Handler, + runtime: params.Runtime, + config: params.Config, + log: params.Log, + progressFactory: params.ProgressFactory, + streamingCapable: params.StreamingCapability.Enabled, } } @@ -198,7 +204,21 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { callbackURL = *muEdReq.CallbackUrl } + streaming := h.streamingCapable && h.config.Progress.Stream.Enabled && acceptsEventStream(r) + 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.serveEvaluateStream(ctx, w, req, command, isPreview, version, callbackURL, requestID) + return + } + reporter, err := h.progressFactory.NewReporter(callbackURL, requestID) if err != nil { h.log.Warn("invalid callbackUrl, disabling progress reporting", zap.Error(err)) @@ -208,43 +228,31 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { resp := h.handler.Handle(ctx, req) - if resp.StatusCode != http.StatusOK { + feedback, termErr := h.produceFeedback(resp, isPreview) + if termErr != nil { progress.Emit(ctx, progress.Event{ Stage: progress.StageFailed, Command: string(command), - Message: muEdErrorMessageFromBody(resp.Body), + Message: termErr.userMessage, + Error: termErr.rawError, }) - for k, v := range resp.Header { - for _, vv := range v { - w.Header().Add(k, vv) + 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 - return - } - 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) + h.writeMuEdError(w, version, termErr.status, termErr.muEdCode, termErr.muEdTitle, termErr.muEdMessage, nil) return } - 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 - } - - var feedback []map[string]any - if isPreview { - feedback = runtime.MuEdToPreviewFeedback(result) - } else { - feedback = runtime.MuEdToEvaluateFeedback(result) - } - // 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 @@ -263,6 +271,173 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { 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 for writeMuEdError (buffered + // path only). + muEdCode string + muEdTitle string + muEdMessage string + + // userMessage and rawError feed the StageFailed progress event and, + // on the streaming path, the "failed" SSE frame. + userMessage string + rawError string +} + +// produceFeedback turns a runtime response into muEd feedback, or a +// terminalError describing why it couldn't. It is pure: no writes, no +// progress events. +func (h *MuEdHandler) produceFeedback(resp runtime.Response, isPreview bool) ([]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 { + 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 { + 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", + } + } + + if isPreview { + return runtime.MuEdToPreviewFeedback(result), nil + } + return runtime.MuEdToEvaluateFeedback(result), nil +} + +// serveEvaluateStream handles a POST /evaluate request that opted in to +// SSE streaming. It commits a 200 + event-stream headers immediately, +// keeps the connection alive with heartbeats while the evaluation runs, +// and emits exactly one terminal frame (completed | failed) carrying the +// feedback plus every step that preceded it. Because the status is +// already committed, 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, + command runtime.Command, + isPreview bool, + version string, + callbackURL string, + requestID string, +) { + cmdLabel := "evaluate" + if isPreview { + cmdLabel = "preview" + } + + 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() + } + } + }() + } + + resp := h.handler.Handle(ctx, req) + + feedback, termErr := h.produceFeedback(resp, isPreview) + if termErr != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: string(command), + Message: termErr.userMessage, + Error: termErr.rawError, + }) + } else { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: string(command), + Message: "Feedback is ready.", + Data: map[string]any{"feedback": feedback}, + }) + } + + close(done) + hbWG.Wait() +} + // muEdErrorMessageFromBody best-effort extracts a human-readable message // from a JSON error body of the shape {"error": {"message": "..."}}. func muEdErrorMessageFromBody(body []byte) string { diff --git a/handler/evaluate_stream_test.go b/handler/evaluate_stream_test.go new file mode 100644 index 0000000..e734a06 --- /dev/null +++ b/handler/evaluate_stream_test.go @@ -0,0 +1,371 @@ +package handler + +import ( + "bufio" + "bytes" + "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") + return req +} + +// 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.Equal(t, "evaluate", data["command"]) + + 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_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.Equal(t, "preview", data["command"]) + 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"]) + assert.Equal(t, "boom", data["message"]) + assert.Contains(t, data["error"], "boom") +} + +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"]) +} + +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.Equal(t, "evaluate", data["command"]) + + // 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") + + 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) { + spec, err := server.LoadOpenAPISpec() + require.NoError(t, err) + mw, err := server.OpenAPIMiddleware(spec, zap.NewNop(), true) + 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") + + 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.Equal(t, "evaluate", data["command"]) +} + +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/module.go b/handler/module.go index fca3d4d..df10a89 100644 --- a/handler/module.go +++ b/handler/module.go @@ -7,6 +7,15 @@ import ( "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", fx.Provide(NewCommandHandler), diff --git a/internal/progress/factory.go b/internal/progress/factory.go index de56695..7e87579 100644 --- a/internal/progress/factory.go +++ b/internal/progress/factory.go @@ -39,6 +39,24 @@ type Config struct { // 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 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/sse_reporter.go b/internal/progress/sse_reporter.go new file mode 100644 index 0000000..3d9c827 --- /dev/null +++ b/internal/progress/sse_reporter.go @@ -0,0 +1,170 @@ +package progress + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + "time" + + "go.uber.org/zap" +) + +// sseStep is one accumulated progress step in the final SSE frame. Its +// shape is deliberately the same one a future mid-stream frame will use, +// so a client parses "a step" the same way whether it arrives inline or +// inside the terminal envelope. +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. The +// same shape is used for the "completed" and "failed" events: on failure +// Feedback is null and Error/Message carry the detail. +type sseEnvelope struct { + Command string `json:"command"` + Feedback []map[string]any `json:"feedback"` + Steps []sseStep `json:"steps"` + Error string `json:"error,omitempty"` + Message string `json:"message,omitempty"` +} + +// SSEReporter is a Reporter that streams progress back to the caller on +// the /evaluate response itself, as Server-Sent Events. In this phase it +// silently accumulates the intermediate steps and emits exactly one +// terminal frame (event: completed | failed) carrying the feedback 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 "progress" events — so all state and +// all writes to the ResponseWriter are guarded by mu. +type SSEReporter struct { + w http.ResponseWriter + flusher http.Flusher + command string + log *zap.Logger + + mu sync.Mutex + steps []sseStep + 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 accumulates a non-terminal event as a step, or writes the single +// terminal frame. Once the terminal frame is written, all further events +// (including a late worker "progress" 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 + } + + step := sseStep{ + Stage: string(evt.Stage), + Message: evt.Message, + Data: evt.Data, + Timestamp: evt.Timestamp, + } + if step.Timestamp.IsZero() { + // Worker-authored "progress" events bypass Emit and arrive + // without a timestamp. + step.Timestamp = time.Now().UTC() + } + + // Collapse a run of identical lifecycle stages — the per-case + // evaluation loop re-enters the supervisor and re-emits + // preparing/evaluating each time. "progress" steps are never + // collapsed. + if n := len(r.steps); n > 0 && r.steps[n-1].Stage == step.Stage && + (evt.Stage == StagePreparing || evt.Stage == StageEvaluating) { + return + } + + r.steps = append(r.steps, step) +} + +func (r *SSEReporter) writeEnvelopeLocked(evt Event) { + env := sseEnvelope{ + Command: r.command, + Steps: r.steps, + } + if env.Steps == nil { + env.Steps = []sseStep{} + } + + event := "completed" + if evt.Stage == StageFailed { + event = "failed" + env.Feedback = nil + env.Error = evt.Error + env.Message = evt.Message + } else { + feedback, ok := evt.Data["feedback"].([]map[string]any) + if !ok { + feedback = []map[string]any{} + } + env.Feedback = feedback + } + + body, err := json.Marshal(env) + 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() +} + +// 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..819314b --- /dev/null +++ b/internal/progress/sse_reporter_test.go @@ -0,0 +1,258 @@ +package progress + +import ( + "context" + "encoding/json" + "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 your evaluation…"}) + r.Report(context.Background(), Event{Stage: StageEvaluating, Message: "Evaluating your submission…"}) + 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) != 1 { + t.Fatalf("expected 1 frame, got %d: %q", len(frames), rec.Body.String()) + } + f := frames[0] + if f.event != "completed" { + t.Errorf("expected event 'completed', got %q", f.event) + } + if f.data["command"] != "evaluate" { + t.Errorf("expected command 'evaluate', got %v", f.data["command"]) + } + 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"] != "evaluating" { + 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.", + }) + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 1 { + t.Fatalf("expected 1 frame, got %d", len(frames)) + } + f := frames[0] + 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) + } + if f.data["error"] != "worker send: context deadline exceeded" { + t.Errorf("raw error not carried: %v", f.data["error"]) + } + if f.data["message"] != "We couldn't evaluate your answer. Please try again." { + t.Errorf("user message not carried: %v", f.data["message"]) + } + if steps, _ := f.data["steps"].([]any); len(steps) != 1 { + t.Errorf("expected 1 step, got %v", f.data["steps"]) + } +} + +func TestSSEReporter_PreviewCommandLabel(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()) + if frames[0].data["command"] != "preview" { + t.Errorf("expected command 'preview', got %v", frames[0].data["command"]) + } +} + +func TestSSEReporter_DedupConsecutivePreparingEvaluating(t *testing.T) { + rec, r := newRecorderReporter(t, "evaluate") + for _, s := range []Stage{StagePreparing, StagePreparing, StageEvaluating, StageEvaluating, StagePreparing} { + r.Report(context.Background(), Event{Stage: s}) + } + r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + + steps := parseSSEFrames(t, rec.Body.String())[0].data["steps"].([]any) + got := []string{} + for _, s := range steps { + got = append(got, s.(map[string]any)["stage"].(string)) + } + want := []string{"preparing", "evaluating", "preparing"} + 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: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) + + steps := parseSSEFrames(t, rec.Body.String())[0].data["steps"].([]any) + if len(steps) != 2 { + t.Fatalf("expected 2 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].data["message"] != "first" { + t.Errorf("expected the first terminal event to win, got %v", frames[0].data["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{}}}) + + frames := parseSSEFrames(t, rec.Body.String()) + if len(frames) != 1 || frames[0].event != "completed" { + t.Fatalf("expected exactly 1 completed frame, got %d: %q", len(frames), 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") + } +} diff --git a/internal/server/openapi.go b/internal/server/openapi.go index 3ef2727..dd86de2 100644 --- a/internal/server/openapi.go +++ b/internal/server/openapi.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "net/http/httptest" + "strings" "github.com/getkin/kin-openapi/openapi3" "github.com/getkin/kin-openapi/openapi3filter" @@ -26,7 +27,19 @@ func LoadOpenAPISpec() (*openapi3.T, error) { return spec, nil } -func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger) (func(http.Handler) http.Handler, error) { +// sseStreamsEvaluate reports whether this request is an SSE-streaming +// POST /evaluate: its response is written and flushed incrementally, so +// the middleware must not buffer it through httptest.NewRecorder (which +// also strips http.Flusher) or validate its non-JSON body against the +// spec. Request validation still runs. Matched with HasSuffix because the +// middleware runs before NormalizePath rewrites the path. +func sseStreamsEvaluate(r *http.Request) bool { + return r.Method == http.MethodPost && + strings.HasSuffix(r.URL.Path, "/evaluate") && + strings.Contains(strings.ToLower(r.Header.Get("Accept")), "text/event-stream") +} + +func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger, sseEnabled bool) (func(http.Handler) http.Handler, error) { router, err := legacy.NewRouter(spec, openapi3.IsOpenAPI31OrLater(), openapi3.AllowExtraSiblingFields("description", "summary"), @@ -57,6 +70,13 @@ func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger) (func(http.Handler) ht return } + // A streaming SSE response can't be buffered or JSON-validated; + // hand the real writer straight to the handler. + if sseEnabled && sseStreamsEvaluate(r) { + next.ServeHTTP(w, r) + return + } + // Capture response for validation rec := httptest.NewRecorder() next.ServeHTTP(rec, r) diff --git a/internal/server/openapi_test.go b/internal/server/openapi_test.go index 554f354..8bb8924 100644 --- a/internal/server/openapi_test.go +++ b/internal/server/openapi_test.go @@ -22,7 +22,7 @@ func TestOpenAPIMiddleware_Init(t *testing.T) { spec, err := LoadOpenAPISpec() require.NoError(t, err) - middleware, err := OpenAPIMiddleware(spec, zap.NewNop()) + middleware, err := OpenAPIMiddleware(spec, zap.NewNop(), true) require.NoError(t, err) assert.NotNil(t, middleware) } @@ -125,11 +125,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 +142,96 @@ func TestOpenAPIMiddleware_ValidHealthRequest_ReachesHandler(t *testing.T) { assert.Equal(t, http.StatusOK, w.Code) } +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_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) +} + +func TestOpenAPIMiddleware_SSEDisabled_StillBuffersAndValidates(t *testing.T) { + middleware := mustMiddleware(t, false) + + 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) +} + // mustMiddleware loads the real spec and returns the initialised middleware, failing the test on error. -func mustMiddleware(t *testing.T) func(http.Handler) http.Handler { +// SSE streaming bypass is enabled unless sseEnabled[0] is explicitly false. +func mustMiddleware(t *testing.T, sseEnabled ...bool) func(http.Handler) http.Handler { t.Helper() + enabled := true + if len(sseEnabled) > 0 { + enabled = sseEnabled[0] + } spec, err := LoadOpenAPISpec() require.NoError(t, err) - middleware, err := OpenAPIMiddleware(spec, zap.NewNop()) + middleware, err := OpenAPIMiddleware(spec, zap.NewNop(), enabled) require.NoError(t, err) return middleware } @@ -158,4 +242,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/server.go b/internal/server/server.go index 6a94ea5..e5088e7 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -11,6 +11,8 @@ import ( "go.uber.org/zap" "golang.org/x/net/http2" "golang.org/x/net/http2/h2c" + + "github.com/lambda-feedback/shimmy/config" ) type HttpServerParams struct { @@ -18,8 +20,9 @@ type HttpServerParams struct { Context context.Context - Config HttpConfig - Spec *openapi3.T + Config HttpConfig + AppConfig config.Config + Spec *openapi3.T Handlers []*HttpHandler `group:"handlers"` Logger *zap.Logger @@ -41,7 +44,7 @@ func NewHttpServer(params HttpServerParams) (*HttpServer, error) { } var handler http.Handler = NormalizePath(mux) - openAPIMiddleware, err := OpenAPIMiddleware(params.Spec, params.Logger) + openAPIMiddleware, err := OpenAPIMiddleware(params.Spec, params.Logger, params.AppConfig.Progress.Stream.Enabled) if err != nil { return nil, fmt.Errorf("initialising OpenAPI middleware: %w", err) } From 965150dde334c8500f2409ffaedcda4984b48fe6 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Mon, 31 Aug 2026 10:57:14 +0100 Subject: [PATCH 11/28] Add streaming progress updates via Server-Sent Events (SSE) - Implement SSE-based streaming for progress updates on `/evaluate` responses. - Add tests to validate SSE behavior, event streaming, and live frame correctness. - Collapse repeated lifecycle stages (`preparing`, `evaluating`) into single events per request. - Update README with detailed documentation on SSE usage, configuration, and behavior. - Introduce `--progress-stream-enabled` and `--progress-stream-heartbeat-seconds` flags for configuration. - Ensure terminal frames include accumulated steps and align with live frame data. --- README.md | 66 +++++++++++ handler/evaluate_stream_test.go | 68 +++++++++++ internal/progress/sse_reporter.go | 85 ++++++++++---- internal/progress/sse_reporter_test.go | 154 ++++++++++++++++++++++--- 4 files changed, 333 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 12f8b0a..910fb07 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,8 @@ GLOBAL OPTIONS: --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 response as Server-Sent Events for requests that send 'Accept: text/event-stream'. 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 @@ -266,6 +268,70 @@ A rejected callback URL behaves like any other delivery failure: it's logged and > **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` 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 evaluation runs, instead of the buffered `Feedback[]` JSON body. + +- **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 evaluation + is cancelled. + +Each non-terminal event is written as its own frame the moment it occurs, so the caller sees +progress live: + +``` +event: preparing +data: {"stage":"preparing","message":"Preparing your evaluation…","timestamp":"2026-08-31T09:16:29.474Z"} + +event: evaluating +data: {"stage":"evaluating","message":"Evaluating your submission…","timestamp":"2026-08-31T09:16:29.474Z"} + +event: progress +data: {"stage":"progress","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"2026-08-31T09:16:29.522Z"} +``` + +`preparing` and `evaluating` are streamed **once per request** even though a multi-case +evaluation re-enters those stages per case; worker-authored `progress` events 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 feedback plus every step that preceded it, and the connection closes: + +``` +event: completed +data: {"command":"evaluate", + "feedback":[{"awardedPoints":1,"message":"Well done"}], + "steps":[{"stage":"preparing","message":"Preparing your evaluation…","timestamp":"…"}, + {"stage":"evaluating","message":"Evaluating your submission…","timestamp":"…"}, + {"stage":"progress","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"…"}]} +``` + +``` +event: failed +data: {"command":"evaluate","feedback":null, + "steps":[ /* whatever streamed before the failure */ ], + "error":"worker send: context deadline exceeded", + "message":"We couldn't evaluate your answer. Please try again."} +``` + +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. `command` is `"evaluate"` or `"preview"`. 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 evaluation 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 four stages above are emitted by shimmy itself, around the evaluation function call as a whole — `evaluating` covers the entire invocation as one span. An evaluation 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` alongside shimmy's own events. diff --git a/handler/evaluate_stream_test.go b/handler/evaluate_stream_test.go index e734a06..c9ff822 100644 --- a/handler/evaluate_stream_test.go +++ b/handler/evaluate_stream_test.go @@ -3,6 +3,7 @@ package handler import ( "bufio" "bytes" + "context" "encoding/json" "io" "net/http" @@ -47,6 +48,35 @@ func sseRequest(t *testing.T, body []byte) *http.Request { 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() @@ -114,6 +144,44 @@ func TestServeEvaluate_SSE_Success(t *testing.T) { 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 events, emitted twice as the per-case loop + // would; only the first of each should reach the wire. + progress.Emit(ctx, progress.Event{Stage: progress.StagePreparing, Message: "Preparing your evaluation…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageEvaluating, Message: "Evaluating your submission…"}) + progress.Emit(ctx, progress.Event{Stage: progress.StageProgress, Message: "Parsing response and answer..."}) + progress.Emit(ctx, progress.Event{Stage: progress.StageProgress, Message: "Comparing sets for equivalence..."}) + progress.Emit(ctx, progress.Event{Stage: progress.StagePreparing, Message: "Preparing your evaluation…"}) + }). + 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", "evaluating", "progress", "progress", "completed"}, events) + + // The first 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, "progress", steps[3].(map[string]any)["stage"]) +} + func TestServeEvaluate_SSE_Preview(t *testing.T) { previewResp := runtime.Response{ StatusCode: http.StatusOK, diff --git a/internal/progress/sse_reporter.go b/internal/progress/sse_reporter.go index 3d9c827..aa049f2 100644 --- a/internal/progress/sse_reporter.go +++ b/internal/progress/sse_reporter.go @@ -11,10 +11,10 @@ import ( "go.uber.org/zap" ) -// sseStep is one accumulated progress step in the final SSE frame. Its -// shape is deliberately the same one a future mid-stream frame will use, -// so a client parses "a step" the same way whether it arrives inline or -// inside the terminal envelope. +// 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"` @@ -34,10 +34,12 @@ type sseEnvelope struct { } // SSEReporter is a Reporter that streams progress back to the caller on -// the /evaluate response itself, as Server-Sent Events. In this phase it -// silently accumulates the intermediate steps and emits exactly one -// terminal frame (event: completed | failed) carrying the feedback plus -// every step that preceded it, then the handler closes the connection. +// the /evaluate 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 feedback 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 @@ -49,10 +51,12 @@ type SSEReporter struct { command string log *zap.Logger - mu sync.Mutex - steps []sseStep - terminated bool - terminalOnce sync.Once + mu sync.Mutex + steps []sseStep + seenPreparing bool + seenEvaluating bool + terminated bool + terminalOnce sync.Once } var _ Reporter = (*SSEReporter)(nil) @@ -73,10 +77,11 @@ func NewSSEReporter(w http.ResponseWriter, command string, log *zap.Logger) (*SS }, nil } -// Report accumulates a non-terminal event as a step, or writes the single -// terminal frame. Once the terminal frame is written, all further events -// (including a late worker "progress" relayed after the request returned) -// are dropped without touching the ResponseWriter. +// 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 "progress" 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() @@ -93,6 +98,23 @@ func (r *SSEReporter) Report(_ context.Context, evt Event) { return } + // Collapse the lifecycle stages to their first occurrence for the + // whole request: the per-case evaluation loop re-enters the + // supervisor and re-emits preparing/evaluating once per case. + // Worker-authored "progress" events are never collapsed. + switch evt.Stage { + case StagePreparing: + if r.seenPreparing { + return + } + r.seenPreparing = true + case StageEvaluating: + if r.seenEvaluating { + return + } + r.seenEvaluating = true + } + step := sseStep{ Stage: string(evt.Stage), Message: evt.Message, @@ -105,16 +127,8 @@ func (r *SSEReporter) Report(_ context.Context, evt Event) { step.Timestamp = time.Now().UTC() } - // Collapse a run of identical lifecycle stages — the per-case - // evaluation loop re-enters the supervisor and re-emits - // preparing/evaluating each time. "progress" steps are never - // collapsed. - if n := len(r.steps); n > 0 && r.steps[n-1].Stage == step.Stage && - (evt.Stage == StagePreparing || evt.Stage == StageEvaluating) { - return - } - r.steps = append(r.steps, step) + r.writeStepLocked(step) } func (r *SSEReporter) writeEnvelopeLocked(evt Event) { @@ -153,6 +167,27 @@ func (r *SSEReporter) writeEnvelopeLocked(evt Event) { 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() { diff --git a/internal/progress/sse_reporter_test.go b/internal/progress/sse_reporter_test.go index 819314b..9441873 100644 --- a/internal/progress/sse_reporter_test.go +++ b/internal/progress/sse_reporter_test.go @@ -3,6 +3,7 @@ package progress import ( "context" "encoding/json" + "io" "net/http" "net/http/httptest" "strings" @@ -70,10 +71,13 @@ func TestSSEReporter_CompletedEnvelope(t *testing.T) { } frames := parseSSEFrames(t, rec.Body.String()) - if len(frames) != 1 { - t.Fatalf("expected 1 frame, got %d: %q", len(frames), rec.Body.String()) + if len(frames) != 3 { + t.Fatalf("expected 3 frames (preparing, evaluating, completed), got %d: %q", len(frames), rec.Body.String()) + } + if frames[0].event != "preparing" || frames[1].event != "evaluating" { + t.Errorf("unexpected live frame events: %q, %q", frames[0].event, frames[1].event) } - f := frames[0] + f := frames[2] if f.event != "completed" { t.Errorf("expected event 'completed', got %q", f.event) } @@ -110,10 +114,13 @@ func TestSSEReporter_FailedEnvelope(t *testing.T) { }) frames := parseSSEFrames(t, rec.Body.String()) - if len(frames) != 1 { - t.Fatalf("expected 1 frame, got %d", len(frames)) + 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[0] + f := frames[1] if f.event != "failed" { t.Errorf("expected event 'failed', got %q", f.event) } @@ -143,19 +150,32 @@ func TestSSEReporter_PreviewCommandLabel(t *testing.T) { } } -func TestSSEReporter_DedupConsecutivePreparingEvaluating(t *testing.T) { +func TestSSEReporter_DedupLifecycleStagesOncePerRequest(t *testing.T) { rec, r := newRecorderReporter(t, "evaluate") - for _, s := range []Stage{StagePreparing, StagePreparing, StageEvaluating, StageEvaluating, StagePreparing} { + // The per-case evaluation loop re-emits preparing/evaluating once per + // case; only the first of each for the whole request is kept. + for _, s := range []Stage{StagePreparing, StagePreparing, StageEvaluating, StageEvaluating, StageProgress, StagePreparing} { r.Report(context.Background(), Event{Stage: s}) } r.Report(context.Background(), Event{Stage: StageCompleted, Data: map[string]any{"feedback": []map[string]any{}}}) - steps := parseSSEFrames(t, rec.Body.String())[0].data["steps"].([]any) + frames := parseSSEFrames(t, rec.Body.String()) + + // Live frames: one preparing, one evaluating, one progress, then completed. + var liveEvents []string + for _, f := range frames[:len(frames)-1] { + liveEvents = append(liveEvents, f.event) + } + if strings.Join(liveEvents, ",") != "preparing,evaluating,progress" { + t.Errorf("expected live frames [preparing evaluating progress], 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", "evaluating", "preparing"} + want := []string{"preparing", "evaluating", "progress"} if strings.Join(got, ",") != strings.Join(want, ",") { t.Errorf("expected steps %v, got %v", want, got) } @@ -165,11 +185,22 @@ 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{}}}) - steps := parseSSEFrames(t, rec.Body.String())[0].data["steps"].([]any) - if len(steps) != 2 { - t.Fatalf("expected 2 progress steps, got %d", len(steps)) + 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) @@ -238,9 +269,22 @@ func TestSSEReporter_ConcurrentReport(t *testing.T) { 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()) - if len(frames) != 1 || frames[0].event != "completed" { - t.Fatalf("expected exactly 1 completed frame, got %d: %q", len(frames), 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()) } } @@ -256,3 +300,83 @@ func TestNewSSEReporter_NonFlusher_Error(t *testing.T) { 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") + } +} From 0c22ea89db97093df561fd72ce768781af4b5eec Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Mon, 31 Aug 2026 11:18:45 +0100 Subject: [PATCH 12/28] Refactor SSE scaffolding for progress streaming - Extract shared SSE streaming logic into `streamProgress` for reuse across `/evaluate` and `/chat`. - Add per-command terminal frame shapes (`sseEnvelope` for `/evaluate`, `sseChatEnvelope` for `/chat`). - Extend progress stages with `starting` and `thinking` for unified lifecycle reporting. - Refactor `/chat` to support streaming progress updates with callback compatibility. --- handler/chat.go | 127 ++++++++++++++++++-- handler/evaluate.go | 85 ++----------- handler/stream.go | 109 +++++++++++++++++ internal/execution/supervisor/supervisor.go | 10 +- internal/progress/event.go | 32 +++-- internal/progress/sidecar.go | 21 +++- internal/progress/sse_reporter.go | 113 ++++++++++------- internal/server/openapi.go | 28 +++-- runtime/chat.go | 5 + 9 files changed, 375 insertions(+), 155 deletions(-) create mode 100644 handler/stream.go diff --git a/handler/chat.go b/handler/chat.go index 24e9e26..4b40fe3 100644 --- a/handler/chat.go +++ b/handler/chat.go @@ -1,16 +1,24 @@ 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 } @@ -43,30 +51,129 @@ func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) { 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 + var callbackURL string + if 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.streamingCapable && h.config.Progress.Stream.Enabled && acceptsEventStream(r) + 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, 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 != "" { + 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}) + output, metadata, termErr := h.produceChatOutput(resp, err) + if termErr != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: string(runtime.CommandChat), + Message: termErr.userMessage, + Error: termErr.rawError, + }) + h.writeMuEdError(w, version, termErr.status, termErr.muEdCode, termErr.muEdTitle, termErr.muEdMessage, nil) return } + chatResp := map[string]any{"output": output} + if metadata != nil { + chatResp["metadata"] = metadata + } + + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: string(runtime.CommandChat), + Message: "Response is ready.", + Data: map[string]any{"output": output, "metadata": metadata}, + }) + 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, + 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}) + output, metadata, termErr := h.produceChatOutput(resp, err) + if termErr != nil { + return nil, termErr + } + data := map[string]any{"output": output} + if metadata != nil { + data["metadata"] = metadata + } + return data, nil + }) +} + +// produceChatOutput turns a runtime chat response into the µEd output +// object (+ optional metadata), 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) (output, metadata 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, nil, newErr("chat failed", chatErr.Error()) + } + + resultMap, ok := resp.Data["result"].(map[string]any) + if !ok { + return nil, nil, newErr("invalid response from chat function", "invalid response from chat function") + } + + chatResp, err := runtime.MuEdToChatResponse(resultMap) + if err != nil { + return nil, nil, newErr(err.Error(), fmt.Sprintf("invalid chat response: %v", err)) + } + + output, _ = chatResp["output"].(map[string]any) + metadata, _ = chatResp["metadata"].(map[string]any) + return output, metadata, nil +} + // ServeChatHealth handles GET /chat/health. func (h *MuEdHandler) ServeChatHealth(w http.ResponseWriter, r *http.Request) { if !h.checkAuth(w, r) { diff --git a/handler/evaluate.go b/handler/evaluate.go index c29095b..b995b6a 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -8,7 +8,6 @@ import ( "io" "net/http" "strings" - "sync" "time" "go.uber.org/fx" @@ -347,12 +346,11 @@ func (h *MuEdHandler) produceFeedback(resp runtime.Response, isPreview bool) ([] } // serveEvaluateStream handles a POST /evaluate request that opted in to -// SSE streaming. It commits a 200 + event-stream headers immediately, -// keeps the connection alive with heartbeats while the evaluation runs, -// and emits exactly one terminal frame (completed | failed) carrying the -// feedback plus every step that preceded it. Because the status is -// already committed, every post-Handle outcome — including an internal -// error — becomes a "failed" frame, never an HTTP error. +// 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, @@ -368,74 +366,15 @@ func (h *MuEdHandler) serveEvaluateStream( cmdLabel = "preview" } - 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() - } + 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, isPreview) + if termErr != nil { + return nil, termErr } - }() - } - - resp := h.handler.Handle(ctx, req) - - feedback, termErr := h.produceFeedback(resp, isPreview) - if termErr != nil { - progress.Emit(ctx, progress.Event{ - Stage: progress.StageFailed, - Command: string(command), - Message: termErr.userMessage, - Error: termErr.rawError, + return map[string]any{"feedback": feedback}, nil }) - } else { - progress.Emit(ctx, progress.Event{ - Stage: progress.StageCompleted, - Command: string(command), - Message: "Feedback is ready.", - Data: map[string]any{"feedback": feedback}, - }) - } - - close(done) - hbWG.Wait() } // muEdErrorMessageFromBody best-effort extracts a human-readable message diff --git a/handler/stream.go b/handler/stream.go new file mode 100644 index 0000000..8278142 --- /dev/null +++ b/handler/stream.go @@ -0,0 +1,109 @@ +package handler + +import ( + "context" + "net/http" + "sync" + "time" + + "go.uber.org/zap" + + "github.com/lambda-feedback/shimmy/internal/progress" +) + +// 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 envelope shape ("evaluate"/"preview" -> +// feedback[]; "chat" -> output/metadata) and is the frame's "command". +// 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() + } + } + }() + } + + data, termErr := run(ctx) + if termErr != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: command, + Message: termErr.userMessage, + Error: termErr.rawError, + }) + } else { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: command, + Message: doneMessage, + Data: data, + }) + } + + close(done) + hbWG.Wait() +} diff --git a/internal/execution/supervisor/supervisor.go b/internal/execution/supervisor/supervisor.go index f3d6b65..3adc424 100644 --- a/internal/execution/supervisor/supervisor.go +++ b/internal/execution/supervisor/supervisor.go @@ -174,7 +174,7 @@ func (s *WorkerSupervisor) Send( progress.Emit(ctx, progress.Event{ Stage: progress.StageFailed, Command: method, - Message: "We couldn't start the evaluation. Please try again.", + Message: "We couldn't start the request. Please try again.", Error: err.Error(), }) return nil, fmt.Errorf("failed to acquire worker: %w", err) @@ -182,22 +182,22 @@ func (s *WorkerSupervisor) Send( progress.Emit(ctx, progress.Event{ Stage: progress.StagePreparing, Command: method, - Message: "Preparing your evaluation…", + 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.StageEvaluating, + Stage: progress.StageStarting, Command: method, - Message: "Evaluating your submission…", + 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 while evaluating your answer. Please try again.", + Message: "Something went wrong. Please try again.", Error: err.Error(), }) } diff --git a/internal/progress/event.go b/internal/progress/event.go index c362e25..f2a5881 100644 --- a/internal/progress/event.go +++ b/internal/progress/event.go @@ -11,11 +11,25 @@ const ( // (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" - // StageEvaluating indicates the submission is being evaluated. + // 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" @@ -23,11 +37,9 @@ const ( // StageFailed indicates a terminal failure at any layer of the pipeline. StageFailed Stage = "failed" - // StageProgress indicates a custom, evaluation-function-authored - // progress update. Unlike the other stages, these are never emitted - // by shimmy itself — only relayed from a worker's local progress - // side-channel (see Sidecar). A worker cannot claim any other stage; - // the wire contract for that side-channel has no way to set Stage. + // 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" ) @@ -57,10 +69,10 @@ type Event struct { Error string // Data is a free-form extension point. On StageCompleted it carries - // the evaluation's feedback payload (so a callbackUrl-supplying - // caller gets the final result, not just a status ping). On - // StageProgress it carries whatever the evaluation function attached - // to its custom event (see Sidecar). + // 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. diff --git a/internal/progress/sidecar.go b/internal/progress/sidecar.go index 7d73f82..4319dc0 100644 --- a/internal/progress/sidecar.go +++ b/internal/progress/sidecar.go @@ -84,14 +84,27 @@ func (c SidecarConfig) withDefaults() SidecarConfig { // sidecarPayload is the JSON body a worker POSTs to report a custom // progress event. There is deliberately no "stage" field: a worker can -// never claim any stage other than StageProgress, which the sidecar -// hardcodes itself. Unknown fields (including a "stage" a worker might -// send anyway) are silently ignored by json.Decode, never merged in. +// 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, @@ -255,7 +268,7 @@ func (s *Sidecar) handle(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusAccepted) evt := Event{ - Stage: StageProgress, + Stage: stageForCommand(command), Command: command, Message: body.Message, Data: body.Data, diff --git a/internal/progress/sse_reporter.go b/internal/progress/sse_reporter.go index aa049f2..9580f9f 100644 --- a/internal/progress/sse_reporter.go +++ b/internal/progress/sse_reporter.go @@ -22,9 +22,10 @@ type sseStep struct { Timestamp time.Time `json:"timestamp"` } -// sseEnvelope is the JSON payload of the single terminal SSE frame. The -// same shape is used for the "completed" and "failed" events: on failure -// Feedback is null and Error/Message carry the detail. +// 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/Message carry the detail. type sseEnvelope struct { Command string `json:"command"` Feedback []map[string]any `json:"feedback"` @@ -33,30 +34,42 @@ type sseEnvelope struct { Message string `json:"message,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/Message carry the detail. +type sseChatEnvelope struct { + Command string `json:"command"` + Output map[string]any `json:"output"` + Metadata map[string]any `json:"metadata,omitempty"` + Steps []sseStep `json:"steps"` + Error string `json:"error,omitempty"` + Message string `json:"message,omitempty"` +} + // SSEReporter is a Reporter that streams progress back to the caller on -// the /evaluate 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 feedback plus every step that -// preceded it, then the handler closes the connection. +// 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 "progress" events — so all state and -// all writes to the ResponseWriter are guarded by mu. +// 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 string log *zap.Logger - mu sync.Mutex - steps []sseStep - seenPreparing bool - seenEvaluating bool - terminated bool - terminalOnce sync.Once + mu sync.Mutex + steps []sseStep + seenPreparing bool + seenStarting bool + terminated bool + terminalOnce sync.Once } var _ Reporter = (*SSEReporter)(nil) @@ -79,7 +92,7 @@ func NewSSEReporter(w http.ResponseWriter, command string, log *zap.Logger) (*SS // 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 "progress" relayed +// 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) { @@ -98,21 +111,22 @@ func (r *SSEReporter) Report(_ context.Context, evt Event) { return } - // Collapse the lifecycle stages to their first occurrence for the - // whole request: the per-case evaluation loop re-enters the - // supervisor and re-emits preparing/evaluating once per case. - // Worker-authored "progress" events are never collapsed. + // 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 StageEvaluating: - if r.seenEvaluating { + case StageStarting: + if r.seenStarting { return } - r.seenEvaluating = true + r.seenStarting = true } step := sseStep{ @@ -122,8 +136,8 @@ func (r *SSEReporter) Report(_ context.Context, evt Event) { Timestamp: evt.Timestamp, } if step.Timestamp.IsZero() { - // Worker-authored "progress" events bypass Emit and arrive - // without a timestamp. + // Worker-authored sub-steps bypass Emit (they come off the + // sidecar) and arrive without a timestamp. step.Timestamp = time.Now().UTC() } @@ -132,29 +146,44 @@ func (r *SSEReporter) Report(_ context.Context, evt Event) { } func (r *SSEReporter) writeEnvelopeLocked(evt Event) { - env := sseEnvelope{ - Command: r.command, - Steps: r.steps, - } - if env.Steps == nil { - env.Steps = []sseStep{} + steps := r.steps + if steps == nil { + steps = []sseStep{} } + failed := evt.Stage == StageFailed event := "completed" - if evt.Stage == StageFailed { + if failed { event = "failed" - env.Feedback = nil - env.Error = evt.Error - env.Message = evt.Message + } + + var payload any + if r.command == "chat" { + env := sseChatEnvelope{Command: r.command, Steps: steps} + if failed { + env.Error = evt.Error + env.Message = evt.Message + } else { + env.Output, _ = evt.Data["output"].(map[string]any) + env.Metadata, _ = evt.Data["metadata"].(map[string]any) + } + payload = env } else { - feedback, ok := evt.Data["feedback"].([]map[string]any) - if !ok { - feedback = []map[string]any{} + env := sseEnvelope{Command: r.command, Steps: steps} + if failed { + env.Error = evt.Error + env.Message = evt.Message + } else { + feedback, ok := evt.Data["feedback"].([]map[string]any) + if !ok { + feedback = []map[string]any{} + } + env.Feedback = feedback } - env.Feedback = feedback + payload = env } - body, err := json.Marshal(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 diff --git a/internal/server/openapi.go b/internal/server/openapi.go index dd86de2..add8be7 100644 --- a/internal/server/openapi.go +++ b/internal/server/openapi.go @@ -27,16 +27,22 @@ func LoadOpenAPISpec() (*openapi3.T, error) { return spec, nil } -// sseStreamsEvaluate reports whether this request is an SSE-streaming -// POST /evaluate: its response is written and flushed incrementally, so -// the middleware must not buffer it through httptest.NewRecorder (which -// also strips http.Flusher) or validate its non-JSON body against the -// spec. Request validation still runs. Matched with HasSuffix because the -// middleware runs before NormalizePath rewrites the path. -func sseStreamsEvaluate(r *http.Request) bool { - return r.Method == http.MethodPost && - strings.HasSuffix(r.URL.Path, "/evaluate") && - strings.Contains(strings.ToLower(r.Header.Get("Accept")), "text/event-stream") +// sseStreamsProgressRoute reports whether this request is an SSE-streaming +// POST to /evaluate or /chat: its response is written and flushed +// incrementally, so the middleware must not buffer it through +// httptest.NewRecorder (which also strips http.Flusher) or validate its +// non-JSON body against the spec. Request validation still runs. Matched +// with HasSuffix because the middleware runs before NormalizePath rewrites +// the path; "/chat/health" (GET) does not end with "/chat" and so is never +// matched. +func sseStreamsProgressRoute(r *http.Request) bool { + if r.Method != http.MethodPost { + return false + } + if !strings.Contains(strings.ToLower(r.Header.Get("Accept")), "text/event-stream") { + return false + } + return strings.HasSuffix(r.URL.Path, "/evaluate") || strings.HasSuffix(r.URL.Path, "/chat") } func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger, sseEnabled bool) (func(http.Handler) http.Handler, error) { @@ -72,7 +78,7 @@ func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger, sseEnabled bool) (func // A streaming SSE response can't be buffered or JSON-validated; // hand the real writer straight to the handler. - if sseEnabled && sseStreamsEvaluate(r) { + if sseEnabled && sseStreamsProgressRoute(r) { next.ServeHTTP(w, r) return } diff --git a/runtime/chat.go b/runtime/chat.go index 9c3c93f..4cebe92 100644 --- a/runtime/chat.go +++ b/runtime/chat.go @@ -44,6 +44,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 From 5387fc6557e582eafef65e311fafd575ebd242b6 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Mon, 31 Aug 2026 11:25:12 +0100 Subject: [PATCH 13/28] Update progress reporting with new `starting` stage and `/chat` support - Extend progress lifecycle to include `starting` stage for improved stage granularity. - Add `/chat` endpoint compatibility with SSE progress streaming and event validation. - Refactor evaluation tests to use `starting` in place of `evaluating` at appropriate stages. - Update README to document new stages and `/chat` progress models. - Add corresponding unit tests for new lifecycle reporting and endpoint behavior. --- README.md | 86 +++++++++++-------- handler/evaluate_stream_test.go | 25 +++--- .../execution/supervisor/adapter_file_test.go | 2 +- .../execution/supervisor/adapter_rpc_test.go | 2 +- .../execution/supervisor/supervisor_test.go | 4 +- internal/progress/sidecar_test.go | 32 ++++++- internal/progress/sse_reporter_test.go | 80 ++++++++++++++--- internal/server/openapi_test.go | 30 +++++++ 8 files changed, 196 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index 910fb07..10c9e46 100644 --- a/README.md +++ b/README.md @@ -201,9 +201,9 @@ Example request using cases: ### Progress Events -The shim also exposes a µEd-compatible endpoint at `POST /evaluate` (see the [µEd spec](https://mued.org/spec)), separate from the legacy `POST /` endpoint documented above. When a client calls `/evaluate` 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. +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. "Evaluating your submission…") 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. +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. @@ -215,25 +215,27 @@ To opt in, include `callbackUrl` in the request body and, optionally, an `X-Requ } ``` -Four stages are emitted, in order, for a successful evaluation: +Stages, in order: -| Stage | Meaning | -|-------|---------| -| `preparing` | The evaluation environment is being set up (a worker is ready — freshly booted or reused). | -| `evaluating` | The evaluation function is being invoked. | -| `completed` | Feedback has been computed. `data.feedback` carries the same array returned in the synchronous response body. | -| `failed` | A terminal failure occurred at some stage. `message` is safe to show to an end user; `error` carries raw technical detail for logs only. | +| 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 safe to show to an end user; `error` carries raw technical detail for logs only. | -`completed` and `failed` are terminal — at most one of them is delivered per request, whichever occurs first. +`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": "evaluating", + "stage": "starting", "command": "eval", - "message": "Evaluating your submission…", + "message": "Starting…", "timestamp": "2026-08-04T14:23:01.512Z" } ``` @@ -270,48 +272,50 @@ A rejected callback URL behaves like any other delivery failure: it's logged and #### Streaming progress on the response itself (Server-Sent Events) -A caller that would rather receive progress on the `/evaluate` 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 evaluation runs, instead of the buffered `Feedback[]` JSON body. +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. - **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 evaluation - is cancelled. +- 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: +progress live (`/evaluate` shown; `/chat` is identical but with `thinking` frames in place +of `evaluating`): ``` event: preparing -data: {"stage":"preparing","message":"Preparing your evaluation…","timestamp":"2026-08-31T09:16:29.474Z"} +data: {"stage":"preparing","message":"Preparing…","timestamp":"2026-08-31T09:16:29.474Z"} -event: evaluating -data: {"stage":"evaluating","message":"Evaluating your submission…","timestamp":"2026-08-31T09:16:29.474Z"} +event: starting +data: {"stage":"starting","message":"Starting…","timestamp":"2026-08-31T09:16:29.474Z"} -event: progress -data: {"stage":"progress","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"2026-08-31T09:16:29.522Z"} +event: evaluating +data: {"stage":"evaluating","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"2026-08-31T09:16:29.522Z"} ``` -`preparing` and `evaluating` are streamed **once per request** even though a multi-case -evaluation re-enters those stages per case; worker-authored `progress` events are streamed -every time. The `event:` line carries the stage; the `data` payload is a self-contained step -object (`stage`, `message`, optional `data`, `timestamp`). +`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 feedback plus every step that preceded it, and the connection closes: +— carrying the result plus every step that preceded it, and the connection closes: ``` event: completed data: {"command":"evaluate", "feedback":[{"awardedPoints":1,"message":"Well done"}], - "steps":[{"stage":"preparing","message":"Preparing your evaluation…","timestamp":"…"}, - {"stage":"evaluating","message":"Evaluating your submission…","timestamp":"…"}, - {"stage":"progress","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"…"}]} + "steps":[{"stage":"preparing","message":"Preparing…","timestamp":"…"}, + {"stage":"starting","message":"Starting…","timestamp":"…"}, + {"stage":"evaluating","message":"Ran 3/10 cases","data":{"completed":3,"total":10},"timestamp":"…"}]} ``` ``` @@ -322,19 +326,31 @@ data: {"command":"evaluate","feedback":null, "message":"We couldn't evaluate your answer. Please try again."} ``` +For `/chat` the terminal frame carries `output` (and optional `metadata`) instead of +`feedback`: + +``` +event: completed +data: {"command":"chat", + "output":{"role":"ASSISTANT","content":"…"}, + "metadata":{ /* optional, worker-supplied */ }, + "steps":[ /* preparing, starting, thinking… */ ]} +``` +A failed `/chat` frame has `"output":null` plus `error`/`message`. + 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. `command` is `"evaluate"` or `"preview"`. The correlation id is in the +is in-band. `command` is `"evaluate"`, `"preview"`, or `"chat"`. 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 evaluation runs, the shim also writes an SSE comment heartbeat (`: ping`) every +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 four stages above are emitted by shimmy itself, around the evaluation function call as a whole — `evaluating` covers the entire invocation as one span. An evaluation 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` alongside shimmy's own events. +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. diff --git a/handler/evaluate_stream_test.go b/handler/evaluate_stream_test.go index c9ff822..5ac9597 100644 --- a/handler/evaluate_stream_test.go +++ b/handler/evaluate_stream_test.go @@ -149,13 +149,15 @@ func TestServeEvaluate_SSE_StreamsLiveStepFrames(t *testing.T) { mockHandler.On("Handle", mock.Anything, mock.Anything). Run(func(args mock.Arguments) { ctx := args.Get(0).(context.Context) - // Shim lifecycle events, emitted twice as the per-case loop - // would; only the first of each should reach the wire. - progress.Emit(ctx, progress.Event{Stage: progress.StagePreparing, Message: "Preparing your evaluation…"}) - progress.Emit(ctx, progress.Event{Stage: progress.StageEvaluating, Message: "Evaluating your submission…"}) - progress.Emit(ctx, progress.Event{Stage: progress.StageProgress, Message: "Parsing response and answer..."}) - progress.Emit(ctx, progress.Event{Stage: progress.StageProgress, Message: "Comparing sets for equivalence..."}) - progress.Emit(ctx, progress.Event{Stage: progress.StagePreparing, Message: "Preparing your evaluation…"}) + // 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")) @@ -168,10 +170,10 @@ func TestServeEvaluate_SSE_StreamsLiveStepFrames(t *testing.T) { for _, f := range frames { events = append(events, f.event) } - assert.Equal(t, []string{"preparing", "evaluating", "progress", "progress", "completed"}, events) + assert.Equal(t, []string{"preparing", "starting", "evaluating", "evaluating", "completed"}, events) - // The first live frame's data is one step object, identical in shape - // to an element of the terminal frame's steps[]. + // 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"]) @@ -179,7 +181,8 @@ func TestServeEvaluate_SSE_StreamsLiveStepFrames(t *testing.T) { require.True(t, ok) require.Len(t, steps, 4) assert.Equal(t, "preparing", steps[0].(map[string]any)["stage"]) - assert.Equal(t, "progress", steps[3].(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) { diff --git a/internal/execution/supervisor/adapter_file_test.go b/internal/execution/supervisor/adapter_file_test.go index db6bc09..6d02991 100644 --- a/internal/execution/supervisor/adapter_file_test.go +++ b/internal/execution/supervisor/adapter_file_test.go @@ -214,7 +214,7 @@ func TestFileAdapter_Send_InjectsProgressURLAndRelaysWorkerEvents(t *testing.T) }, time.Second, 5*time.Millisecond, "expected the worker's progress event to be relayed") events := r.recorded() - assert.Equal(t, progress.StageProgress, events[0].Stage) + assert.Equal(t, progress.StageEvaluating, events[0].Stage) assert.Equal(t, "eval", events[0].Command) assert.Equal(t, "checking correctness", events[0].Message) } diff --git a/internal/execution/supervisor/adapter_rpc_test.go b/internal/execution/supervisor/adapter_rpc_test.go index 1357fe6..be1cbee 100644 --- a/internal/execution/supervisor/adapter_rpc_test.go +++ b/internal/execution/supervisor/adapter_rpc_test.go @@ -203,7 +203,7 @@ func TestStdioAdapter_Send_RelaysWorkerProgressEvents(t *testing.T) { }, time.Second, 5*time.Millisecond, "expected the worker's progress event to be relayed") events := r.recorded() - assert.Equal(t, progress.StageProgress, events[0].Stage) + assert.Equal(t, progress.StageEvaluating, events[0].Stage) assert.Equal(t, "eval", events[0].Command) assert.Equal(t, "checking correctness", events[0].Message) } diff --git a/internal/execution/supervisor/supervisor_test.go b/internal/execution/supervisor/supervisor_test.go index 41630c1..4357442 100644 --- a/internal/execution/supervisor/supervisor_test.go +++ b/internal/execution/supervisor/supervisor_test.go @@ -301,7 +301,7 @@ func TestSupervisor_Send_EmitsWorkerAcquiredAndRunning(t *testing.T) { assert.Equal(t, []progress.Stage{ progress.StagePreparing, - progress.StageEvaluating, + progress.StageStarting, }, r.stages()) } @@ -340,7 +340,7 @@ func TestSupervisor_Send_EmitsFailed_WhenWorkerSendFails(t *testing.T) { assert.Equal(t, []progress.Stage{ progress.StagePreparing, - progress.StageEvaluating, + progress.StageStarting, progress.StageFailed, }, r.stages()) } diff --git a/internal/progress/sidecar_test.go b/internal/progress/sidecar_test.go index 35548ac..b5d148b 100644 --- a/internal/progress/sidecar_test.go +++ b/internal/progress/sidecar_test.go @@ -57,8 +57,8 @@ func TestSidecar_Accept_RelaysEventThroughBoundReporter(t *testing.T) { events := waitForEvents(t, r, 1) evt := events[0] - if evt.Stage != StageProgress { - t.Errorf("expected stage %q, got %q", StageProgress, evt.Stage) + 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) @@ -82,11 +82,37 @@ func TestSidecar_IgnoresWorkerSuppliedStage(t *testing.T) { } events := waitForEvents(t, r, 1) - if events[0].Stage != StageProgress { + 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{}) diff --git a/internal/progress/sse_reporter_test.go b/internal/progress/sse_reporter_test.go index 9441873..a6b4d7d 100644 --- a/internal/progress/sse_reporter_test.go +++ b/internal/progress/sse_reporter_test.go @@ -59,8 +59,8 @@ func newRecorderReporter(t *testing.T, command string) (*httptest.ResponseRecord func TestSSEReporter_CompletedEnvelope(t *testing.T) { rec, r := newRecorderReporter(t, "evaluate") - r.Report(context.Background(), Event{Stage: StagePreparing, Message: "Preparing your evaluation…"}) - r.Report(context.Background(), Event{Stage: StageEvaluating, Message: "Evaluating your submission…"}) + 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"}}}, @@ -72,9 +72,9 @@ func TestSSEReporter_CompletedEnvelope(t *testing.T) { frames := parseSSEFrames(t, rec.Body.String()) if len(frames) != 3 { - t.Fatalf("expected 3 frames (preparing, evaluating, completed), got %d: %q", len(frames), rec.Body.String()) + t.Fatalf("expected 3 frames (preparing, starting, completed), got %d: %q", len(frames), rec.Body.String()) } - if frames[0].event != "preparing" || frames[1].event != "evaluating" { + 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] @@ -95,7 +95,7 @@ func TestSSEReporter_CompletedEnvelope(t *testing.T) { 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"] != "evaluating" { + 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"] == "" { @@ -150,24 +150,80 @@ func TestSSEReporter_PreviewCommandLabel(t *testing.T) { } } +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 f.data["command"] != "chat" { + t.Errorf("expected command 'chat', got %v", f.data["command"]) + } + 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.", + }) + + 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) + } + if f.data["error"] != "chat failed: worker exited" { + t.Errorf("raw error not carried: %v", f.data["error"]) + } +} + func TestSSEReporter_DedupLifecycleStagesOncePerRequest(t *testing.T) { rec, r := newRecorderReporter(t, "evaluate") - // The per-case evaluation loop re-emits preparing/evaluating once per - // case; only the first of each for the whole request is kept. - for _, s := range []Stage{StagePreparing, StagePreparing, StageEvaluating, StageEvaluating, StageProgress, StagePreparing} { + // 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()) - // Live frames: one preparing, one evaluating, one progress, then completed. var liveEvents []string for _, f := range frames[:len(frames)-1] { liveEvents = append(liveEvents, f.event) } - if strings.Join(liveEvents, ",") != "preparing,evaluating,progress" { - t.Errorf("expected live frames [preparing evaluating progress], got %v", liveEvents) + 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) @@ -175,7 +231,7 @@ func TestSSEReporter_DedupLifecycleStagesOncePerRequest(t *testing.T) { for _, s := range steps { got = append(got, s.(map[string]any)["stage"].(string)) } - want := []string{"preparing", "evaluating", "progress"} + want := []string{"preparing", "starting", "evaluating", "evaluating"} if strings.Join(got, ",") != strings.Join(want, ",") { t.Errorf("expected steps %v, got %v", want, got) } diff --git a/internal/server/openapi_test.go b/internal/server/openapi_test.go index 8bb8924..47ada1e 100644 --- a/internal/server/openapi_test.go +++ b/internal/server/openapi_test.go @@ -176,6 +176,36 @@ func TestOpenAPIMiddleware_SSEEvaluate_BypassesResponseValidation(t *testing.T) 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) From 3efc4f0206f0066755c99e610d47c3db6c9779e3 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Mon, 31 Aug 2026 11:25:21 +0100 Subject: [PATCH 14/28] Add unit tests for `/chat` SSE streaming behavior and edge cases - Introduce comprehensive test coverage for `/chat` endpoint's Server-Sent Events (SSE) streaming behavior. - Validate `thinking` and `starting` progress stages, terminal frame shapes, and fallback to JSON output. - Add tests for capability-disabled scenarios, callback handling, authentication errors, and heartbeat events. --- handler/chat_stream_test.go | 253 ++++++++++++++++++++++++++++++++++++ 1 file changed, 253 insertions(+) create mode 100644 handler/chat_stream_test.go diff --git a/handler/chat_stream_test.go b/handler/chat_stream_test.go new file mode 100644 index 0000000..0b0ef21 --- /dev/null +++ b/handler/chat_stream_test.go @@ -0,0 +1,253 @@ +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/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") + 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.Equal(t, "chat", data["command"]) + 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"]) + assert.Contains(t, data["error"], "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.Equal(t, "chat", data["command"]) + + 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") + + 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") +} + +// assertAnError is an error whose message contains "boom", for the failure-path test. +type assertAnError struct{} + +func (assertAnError) Error() string { return "chat failed: boom" } From add54b2379865e28bf6d0bfdc458327578dec465 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Mon, 31 Aug 2026 11:30:47 +0100 Subject: [PATCH 15/28] Update `/chat` endpoint to support SSE streaming - Extend `--progress-stream-enabled` flag and behavior to `/chat` responses. - Update README and CLI usage text to reflect new `/chat` SSE streaming support. --- README.md | 4 ++-- cmd/root.go | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 10c9e46..af12617 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ GLOBAL OPTIONS: --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 response as Server-Sent Events for requests that send 'Accept: text/event-stream'. Standalone/serve mode only; ignored under AWS Lambda. (default: true) [$PROGRESS_STREAM_ENABLED] + --progress-stream-enabled stream progress back on the /evaluate and /chat responses as Server-Sent Events for requests that send 'Accept: text/event-stream'. 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 @@ -365,7 +365,7 @@ To emit a custom event, `POST` a small JSON body to `EVAL_PROGRESS_URL`: - `message` (string, required): student/teacher-facing text. - `data` (object, optional): free-form, passed through as-is. -- There is no `stage` field, by design: an evaluation function can never claim `preparing`, `evaluating`, `completed`, or `failed` — those remain exclusively shim-authored. Custom events are always delivered with `"stage": "progress"`. +- 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). diff --git a/cmd/root.go b/cmd/root.go index 3eb9712..7cbeb63 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -101,7 +101,7 @@ functions on arbitrary, serverless platforms.` }, &cli.BoolFlag{ Name: "progress-stream-enabled", - Usage: "stream progress back on the /evaluate response as Server-Sent Events for requests that send 'Accept: text/event-stream'. Standalone/serve mode only; ignored under AWS Lambda.", + Usage: "stream progress back on the /evaluate and /chat responses as Server-Sent Events for requests that send 'Accept: text/event-stream'. Standalone/serve mode only; ignored under AWS Lambda.", Value: true, Category: "progress", EnvVars: []string{"PROGRESS_STREAM_ENABLED"}, From 1cd0e9c122428c784cf07f971e1a0ba5710d4ee1 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 1 Sep 2026 14:53:53 +0100 Subject: [PATCH 16/28] Add SSE terminal frame validation and refactor OpenAPI middleware - Validate terminal frame payloads against OpenAPI schemas for `/chat` and `/evaluate` SSE responses. - Refactor middleware to rely on runtime response sniffing instead of preflight checks. - Update tests for new validation logic and streaming behavior. - Document terminal frame structure in the OpenAPI schema. --- handler/chat_stream_test.go | 45 ++++++++++ handler/evaluate.go | 9 ++ handler/evaluate_stream_test.go | 2 +- handler/stream.go | 39 ++++++++- internal/server/openapi.go | 55 ++++--------- internal/server/openapi_test.go | 53 +++++++++--- internal/server/response_sniffer.go | 83 +++++++++++++++++++ internal/server/server.go | 9 +- internal/server/validate_body.go | 70 ++++++++++++++++ internal/server/validate_body_test.go | 52 ++++++++++++ runtime/schema/mued_v0.1.0.yml | 113 ++++++++++++++++++++++++++ 11 files changed, 473 insertions(+), 57 deletions(-) create mode 100644 internal/server/response_sniffer.go create mode 100644 internal/server/validate_body.go create mode 100644 internal/server/validate_body_test.go diff --git a/handler/chat_stream_test.go b/handler/chat_stream_test.go index 0b0ef21..6852910 100644 --- a/handler/chat_stream_test.go +++ b/handler/chat_stream_test.go @@ -14,6 +14,7 @@ import ( "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" @@ -247,6 +248,50 @@ func TestServeChat_SSE_Heartbeat(t *testing.T) { 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"]) + }) +} + // assertAnError is an error whose message contains "boom", for the failure-path test. type assertAnError struct{} diff --git a/handler/evaluate.go b/handler/evaluate.go index b995b6a..cc50689 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/getkin/kin-openapi/openapi3" "go.uber.org/fx" "go.uber.org/zap" @@ -56,6 +57,12 @@ type MuEdHandlerParams struct { 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 { @@ -65,6 +72,7 @@ type MuEdHandler struct { log *zap.Logger progressFactory progress.Factory streamingCapable bool + spec *openapi3.T } func NewMuEdHandler(params MuEdHandlerParams) *MuEdHandler { @@ -75,6 +83,7 @@ func NewMuEdHandler(params MuEdHandlerParams) *MuEdHandler { log: params.Log, progressFactory: params.ProgressFactory, streamingCapable: params.StreamingCapability.Enabled, + spec: params.Spec, } } diff --git a/handler/evaluate_stream_test.go b/handler/evaluate_stream_test.go index 5ac9597..4ac5479 100644 --- a/handler/evaluate_stream_test.go +++ b/handler/evaluate_stream_test.go @@ -389,7 +389,7 @@ func TestServeEvaluate_SSE_Heartbeat(t *testing.T) { func TestServeEvaluate_SSE_ThroughOpenAPIMiddleware(t *testing.T) { spec, err := server.LoadOpenAPISpec() require.NoError(t, err) - mw, err := server.OpenAPIMiddleware(spec, zap.NewNop(), true) + mw, err := server.OpenAPIMiddleware(spec, zap.NewNop()) require.NoError(t, err) mockHandler := new(MockHandler) diff --git a/handler/stream.go b/handler/stream.go index 8278142..e240950 100644 --- a/handler/stream.go +++ b/handler/stream.go @@ -9,6 +9,7 @@ import ( "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 @@ -88,14 +89,22 @@ func (h *MuEdHandler) streamProgress( } data, termErr := run(ctx) - if termErr != nil { + switch { + case termErr != nil: progress.Emit(ctx, progress.Event{ Stage: progress.StageFailed, Command: command, Message: termErr.userMessage, Error: termErr.rawError, }) - } else { + 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", + }) + default: progress.Emit(ctx, progress.Event{ Stage: progress.StageCompleted, Command: command, @@ -107,3 +116,29 @@ func (h *MuEdHandler) streamProgress( 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/server/openapi.go b/internal/server/openapi.go index add8be7..404fc6b 100644 --- a/internal/server/openapi.go +++ b/internal/server/openapi.go @@ -5,8 +5,6 @@ import ( "fmt" "io" "net/http" - "net/http/httptest" - "strings" "github.com/getkin/kin-openapi/openapi3" "github.com/getkin/kin-openapi/openapi3filter" @@ -27,25 +25,7 @@ func LoadOpenAPISpec() (*openapi3.T, error) { return spec, nil } -// sseStreamsProgressRoute reports whether this request is an SSE-streaming -// POST to /evaluate or /chat: its response is written and flushed -// incrementally, so the middleware must not buffer it through -// httptest.NewRecorder (which also strips http.Flusher) or validate its -// non-JSON body against the spec. Request validation still runs. Matched -// with HasSuffix because the middleware runs before NormalizePath rewrites -// the path; "/chat/health" (GET) does not end with "/chat" and so is never -// matched. -func sseStreamsProgressRoute(r *http.Request) bool { - if r.Method != http.MethodPost { - return false - } - if !strings.Contains(strings.ToLower(r.Header.Get("Accept")), "text/event-stream") { - return false - } - return strings.HasSuffix(r.URL.Path, "/evaluate") || strings.HasSuffix(r.URL.Path, "/chat") -} - -func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger, sseEnabled bool) (func(http.Handler) http.Handler, error) { +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"), @@ -76,25 +56,26 @@ func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger, sseEnabled bool) (func return } - // A streaming SSE response can't be buffered or JSON-validated; - // hand the real writer straight to the handler. - if sseEnabled && sseStreamsProgressRoute(r) { - next.ServeHTTP(w, r) + // 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 } - // Capture response for validation - rec := httptest.NewRecorder() - next.ServeHTTP(rec, r) - - // Snapshot body before validation — ValidateResponse drains the buffer. - bodyBytes := rec.Body.Bytes() + bodyBytes := sniffer.buf.Bytes() // Validate response (lenient — log only) 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, } @@ -104,11 +85,9 @@ func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger, sseEnabled bool) (func 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 diff --git a/internal/server/openapi_test.go b/internal/server/openapi_test.go index 47ada1e..6c3efec 100644 --- a/internal/server/openapi_test.go +++ b/internal/server/openapi_test.go @@ -22,7 +22,7 @@ func TestOpenAPIMiddleware_Init(t *testing.T) { spec, err := LoadOpenAPISpec() require.NoError(t, err) - middleware, err := OpenAPIMiddleware(spec, zap.NewNop(), true) + middleware, err := OpenAPIMiddleware(spec, zap.NewNop()) require.NoError(t, err) assert.NotNil(t, middleware) } @@ -225,8 +225,12 @@ func TestOpenAPIMiddleware_SSEEvaluate_RequestStillValidated(t *testing.T) { assert.Equal(t, http.StatusBadRequest, w.Code) } -func TestOpenAPIMiddleware_SSEDisabled_StillBuffersAndValidates(t *testing.T) { - middleware := mustMiddleware(t, false) +// 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{ @@ -251,17 +255,46 @@ func TestOpenAPIMiddleware_SSEDisabled_StillBuffersAndValidates(t *testing.T) { 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") +} + // mustMiddleware loads the real spec and returns the initialised middleware, failing the test on error. -// SSE streaming bypass is enabled unless sseEnabled[0] is explicitly false. -func mustMiddleware(t *testing.T, sseEnabled ...bool) func(http.Handler) http.Handler { +func mustMiddleware(t *testing.T) func(http.Handler) http.Handler { t.Helper() - enabled := true - if len(sseEnabled) > 0 { - enabled = sseEnabled[0] - } spec, err := LoadOpenAPISpec() require.NoError(t, err) - middleware, err := OpenAPIMiddleware(spec, zap.NewNop(), enabled) + middleware, err := OpenAPIMiddleware(spec, zap.NewNop()) require.NoError(t, err) return middleware } 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 e5088e7..6a94ea5 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -11,8 +11,6 @@ import ( "go.uber.org/zap" "golang.org/x/net/http2" "golang.org/x/net/http2/h2c" - - "github.com/lambda-feedback/shimmy/config" ) type HttpServerParams struct { @@ -20,9 +18,8 @@ type HttpServerParams struct { Context context.Context - Config HttpConfig - AppConfig config.Config - Spec *openapi3.T + Config HttpConfig + Spec *openapi3.T Handlers []*HttpHandler `group:"handlers"` Logger *zap.Logger @@ -44,7 +41,7 @@ func NewHttpServer(params HttpServerParams) (*HttpServer, error) { } var handler http.Handler = NormalizePath(mux) - openAPIMiddleware, err := OpenAPIMiddleware(params.Spec, params.Logger, params.AppConfig.Progress.Stream.Enabled) + openAPIMiddleware, err := OpenAPIMiddleware(params.Spec, params.Logger) if err != nil { return nil, fmt.Errorf("initialising OpenAPI middleware: %w", err) } diff --git a/internal/server/validate_body.go b/internal/server/validate_body.go new file mode 100644 index 0000000..fe2b932 --- /dev/null +++ b/internal/server/validate_body.go @@ -0,0 +1,70 @@ +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 + } + + // Round-trip 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. + 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/schema/mued_v0.1.0.yml b/runtime/schema/mued_v0.1.0.yml index c5ca3c8..f6ea8fc 100644 --- a/runtime/schema/mued_v0.1.0.yml +++ b/runtime/schema/mued_v0.1.0.yml @@ -375,6 +375,19 @@ paths: - feedbackId: fb-2 title: Overall structure message: The overall structure of your answer is clear and easy to follow. + text/event-stream: + schema: + type: string + description: | + Opt-in Server-Sent Events stream, selected with the request + header `Accept: text/event-stream` (standalone/serve mode + only; ignored under AWS Lambda). Intermediate frames: + `event: starting | preparing | evaluating`, `data` = a + `SseProgressStep` JSON object. Keep-alive `: ping` comment + lines. Exactly one terminal frame: `event: completed | failed`, + `data` = an `SseEvaluateTerminalFrame` JSON object. The HTTP + status is 200 even for a `failed` frame — the failure is + in-band. '202': $ref: '#/components/responses/202-Accepted' '400': @@ -638,6 +651,18 @@ paths: model: gpt-5.2 temperature: 0.5 outputTokens: 143 + text/event-stream: + schema: + type: string + description: | + Opt-in Server-Sent Events stream, selected with the request + header `Accept: text/event-stream` (standalone/serve mode + only; ignored under AWS Lambda). Intermediate frames: + `event: starting | preparing | thinking`, `data` = a + `SseProgressStep` JSON object. Keep-alive `: ping` comment + lines. Exactly one terminal frame: `event: completed | failed`, + `data` = a `SseChatTerminalFrame` JSON object. The HTTP status + is 200 even for a `failed` frame — the failure is in-band. '400': $ref: '#/components/responses/400-BadRequest-2' '403': @@ -1742,6 +1767,94 @@ components: description: Optional version of the chat service implementation. capabilities: $ref: '#/components/schemas/ChatCapabilities' + SseProgressStep: + type: object + description: | + Payload (`data`) of an intermediate SSE progress frame + (`event: starting | preparing | thinking | evaluating`). The same + shape also appears in `steps[]` of the terminal frame. + required: + - stage + - timestamp + properties: + stage: + type: string + description: Lifecycle stage this step reports. + message: + type: string + description: Short, learner/teacher-facing description of the step. + data: + type: object + additionalProperties: true + description: Free-form payload attached by a worker-authored sub-step. + timestamp: + type: string + format: date-time + SseChatTerminalFrame: + type: object + description: | + Payload (`data`) of the single terminal SSE frame for `POST /chat` + (`event: completed | failed`). On a `failed` frame `output` is null + and `error`/`message` carry the detail. + required: + - command + - steps + properties: + command: + type: string + description: The µEd command being processed (e.g. "chat"). + output: + type: + - object + - 'null' + additionalProperties: true + description: The generated assistant response; null on failure. + metadata: + type: + - object + - 'null' + additionalProperties: true + description: Optional metadata about response generation. + steps: + type: array + items: + $ref: '#/components/schemas/SseProgressStep' + error: + type: string + description: Raw technical detail (failure frames only; for logs/support). + message: + type: string + description: Human-facing status text. + SseEvaluateTerminalFrame: + type: object + description: | + Payload (`data`) of the single terminal SSE frame for + `POST /evaluate` (`event: completed | failed`). On a `failed` frame + `feedback` is null and `error`/`message` carry the detail. + required: + - command + - steps + properties: + command: + type: string + description: The µEd command being processed (e.g. "eval", "preview"). + feedback: + type: + - array + - 'null' + items: + $ref: '#/components/schemas/Feedback' + description: The generated feedback items; null on failure. + steps: + type: array + items: + $ref: '#/components/schemas/SseProgressStep' + error: + type: string + description: Raw technical detail (failure frames only; for logs/support). + message: + type: string + description: Human-facing status text. responses: 202-Accepted: description: Request accepted for asynchronous evaluation processing. From 64da428df40b6dbfbc1ede753278df30f7558d0a Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 1 Sep 2026 15:29:05 +0100 Subject: [PATCH 17/28] Add parity tests for SSE schemas and expand OpenAPI validation helpers - Add tests to enforce parity between SSE frame structs and OpenAPI schema definitions. - Introduce `ValidateComponentSchema` for validating payloads against specific OpenAPI component schemas. - Refactor `validate_body.go` to centralize schema validation logic and improve testability. --- internal/progress/sse_schema_parity_test.go | 173 ++++++++++++++++++++ internal/server/validate_body.go | 31 +++- 2 files changed, 200 insertions(+), 4 deletions(-) create mode 100644 internal/progress/sse_schema_parity_test.go diff --git a/internal/progress/sse_schema_parity_test.go b/internal/progress/sse_schema_parity_test.go new file mode 100644 index 0000000..052129e --- /dev/null +++ b/internal/progress/sse_schema_parity_test.go @@ -0,0 +1,173 @@ +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."}) + 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."}) + 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{ + Command: "chat", + Output: map[string]any{"role": "ASSISTANT", "content": "hi"}, + Metadata: map[string]any{"responseTimeMs": 12}, + Steps: []sseStep{step}, + }}, + {"SseChatTerminalFrame", sseChatEnvelope{ + Command: "chat", Steps: []sseStep{step}, Error: "boom", Message: "failed", + }}, + {"SseEvaluateTerminalFrame", sseEnvelope{ + Command: "evaluate", + Feedback: []map[string]any{{"feedbackId": "fb-1", "message": "ok"}}, + Steps: []sseStep{step}, + }}, + {"SseEvaluateTerminalFrame", sseEnvelope{ + Command: "evaluate", Steps: []sseStep{step}, Error: "boom", Message: "failed", + }}, + } + for _, c := range cases { + m := toMap(t, c.payload) + mustValidate(t, spec, c.schema, m) + assertAllFieldsDocumented(t, spec, c.schema, m) + } +} + +// --- helpers --- + +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) + } + for k := range m { + if ref.Value.Properties[k] == nil { + t.Errorf("%s: field %q is emitted by the struct but not a documented property", schemaName, k) + } + } +} + +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/server/validate_body.go b/internal/server/validate_body.go index fe2b932..1a4140d 100644 --- a/internal/server/validate_body.go +++ b/internal/server/validate_body.go @@ -26,9 +26,33 @@ func ValidateResponseBody(spec *openapi3.T, operationID string, payload any) err return err } - // Round-trip 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. + 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) @@ -37,7 +61,6 @@ func ValidateResponseBody(spec *openapi3.T, operationID string, payload any) err if err := json.Unmarshal(raw, &decoded); err != nil { return fmt.Errorf("decoding response payload: %w", err) } - return schema.VisitJSON(decoded) } From f0c1d08abca26c2daa5418f803cbadbe6acca9c6 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 09:17:30 +0100 Subject: [PATCH 18/28] Enhance terminal frame structure and refactor progress reporting - Standardize SSE terminal frames to exclude `command` and move failure details to `error` objects. - Add structured `ErrorInfo` for detailed failure representation, aligning with OpenAPI schemas. - Update `/chat` and `/evaluate` SSE handlers and tests to reflect the refined terminal frame structure. - Document streaming variant support and terminal frame changes in OpenAPI specifications. - Extend back-end validation to enforce parity between emitted events and schema definitions. --- handler/chat.go | 4 +- handler/chat_stream_test.go | 9 +- handler/evaluate.go | 25 +- handler/evaluate_stream_test.go | 19 +- handler/stream.go | 19 +- internal/execution/supervisor/supervisor.go | 12 + internal/progress/event.go | 19 ++ internal/progress/sse_reporter.go | 37 +-- internal/progress/sse_reporter_test.go | 73 ++++-- internal/progress/sse_schema_parity_test.go | 45 +++- runtime/chat_test.go | 19 +- runtime/evaluate.go | 26 ++- runtime/schema/mued_v0.1.0.yml | 240 +++++++++++++------- 13 files changed, 406 insertions(+), 141 deletions(-) diff --git a/handler/chat.go b/handler/chat.go index 4b40fe3..de53e0b 100644 --- a/handler/chat.go +++ b/handler/chat.go @@ -176,6 +176,8 @@ func (h *MuEdHandler) produceChatOutput(resp runtime.ChatResponse, chatErr error // 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 } @@ -202,7 +204,7 @@ func (h *MuEdHandler) ServeChatHealth(w http.ResponseWriter, r *http.Request) { return } - healthResp := runtime.MuEdToChatHealthResponse(resultMap) + healthResp := runtime.MuEdToChatHealthResponse(resultMap, h.streamingCapable && h.config.Progress.Stream.Enabled) 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 index 6852910..f8aaee0 100644 --- a/handler/chat_stream_test.go +++ b/handler/chat_stream_test.go @@ -72,7 +72,7 @@ func TestServeChat_SSE_Success(t *testing.T) { event, data := parseSSE(t, w.Body.String()) assert.Equal(t, "completed", event) - assert.Equal(t, "chat", data["command"]) + 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) } @@ -127,7 +127,10 @@ func TestServeChat_SSE_RuntimeError_BecomesFailedFrameAt200(t *testing.T) { event, data := parseSSE(t, w.Body.String()) assert.Equal(t, "failed", event) assert.Nil(t, data["output"]) - assert.Contains(t, data["error"], "boom") + 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) { @@ -191,7 +194,7 @@ func TestServeChat_SSE_WithCallbackUrl_BothDelivered(t *testing.T) { event, data := parseSSE(t, w.Body.String()) assert.Equal(t, "completed", event) - assert.Equal(t, "chat", data["command"]) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") require.Len(t, *received, 1) evt := (*received)[0] diff --git a/handler/evaluate.go b/handler/evaluate.go index cc50689..889b3f1 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -309,6 +309,29 @@ type terminalError struct { 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, or a // terminalError describing why it couldn't. It is pure: no writes, no // progress events. @@ -442,7 +465,7 @@ func (h *MuEdHandler) ServeHealth(w http.ResponseWriter, r *http.Request) { return } - result := runtime.MuEdToHealthResponse(legacyResult) + result := runtime.MuEdToHealthResponse(legacyResult, h.streamingCapable && h.config.Progress.Stream.Enabled) 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 index 4ac5479..da7ff6f 100644 --- a/handler/evaluate_stream_test.go +++ b/handler/evaluate_stream_test.go @@ -133,7 +133,7 @@ func TestServeEvaluate_SSE_Success(t *testing.T) { event, data := parseSSE(t, w.Body.String()) assert.Equal(t, "completed", event) - assert.Equal(t, "evaluate", data["command"]) + 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"]) @@ -202,7 +202,7 @@ func TestServeEvaluate_SSE_Preview(t *testing.T) { event, data := parseSSE(t, w.Body.String()) assert.Equal(t, "completed", event) - assert.Equal(t, "preview", data["command"]) + 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"] @@ -228,8 +228,12 @@ func TestServeEvaluate_SSE_WorkerNon200_BecomesFailedFrameAt200(t *testing.T) { event, data := parseSSE(t, w.Body.String()) assert.Equal(t, "failed", event) assert.Nil(t, data["feedback"]) - assert.Equal(t, "boom", data["message"]) - assert.Contains(t, data["error"], "boom") + 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) { @@ -247,6 +251,9 @@ func TestServeEvaluate_SSE_UnparseableWorkerResponse_FailedFrame(t *testing.T) { 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) { @@ -310,7 +317,7 @@ func TestServeEvaluate_SSE_WithCallbackUrl_BothDelivered(t *testing.T) { // SSE side event, data := parseSSE(t, w.Body.String()) assert.Equal(t, "completed", event) - assert.Equal(t, "evaluate", data["command"]) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") // callbackUrl side require.Len(t, *received, 1) @@ -431,7 +438,7 @@ func TestServeEvaluate_SSE_ThroughOpenAPIMiddleware(t *testing.T) { require.NoError(t, err) event, data := parseSSE(t, string(raw)) assert.Equal(t, "completed", event) - assert.Equal(t, "evaluate", data["command"]) + assert.NotContains(t, data, "command", "the standardised terminal frame drops the command key") } func mustMarshal(t *testing.T, v any) []byte { diff --git a/handler/stream.go b/handler/stream.go index e240950..338258b 100644 --- a/handler/stream.go +++ b/handler/stream.go @@ -88,14 +88,20 @@ func (h *MuEdHandler) streamProgress( }() } + 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, + 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{ @@ -103,6 +109,11 @@ func (h *MuEdHandler) streamProgress( 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{ diff --git a/internal/execution/supervisor/supervisor.go b/internal/execution/supervisor/supervisor.go index 3adc424..98a625d 100644 --- a/internal/execution/supervisor/supervisor.go +++ b/internal/execution/supervisor/supervisor.go @@ -176,6 +176,12 @@ func (s *WorkerSupervisor) Send( 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) } @@ -199,6 +205,12 @@ func (s *WorkerSupervisor) Send( 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(), + }, }) } diff --git a/internal/progress/event.go b/internal/progress/event.go index f2a5881..28064f8 100644 --- a/internal/progress/event.go +++ b/internal/progress/event.go @@ -50,6 +50,19 @@ 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. @@ -68,6 +81,12 @@ type Event struct { // 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 diff --git a/internal/progress/sse_reporter.go b/internal/progress/sse_reporter.go index 9580f9f..06a85e5 100644 --- a/internal/progress/sse_reporter.go +++ b/internal/progress/sse_reporter.go @@ -24,26 +24,24 @@ type sseStep struct { // 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/Message carry the detail. +// "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 { - Command string `json:"command"` Feedback []map[string]any `json:"feedback"` Steps []sseStep `json:"steps"` - Error string `json:"error,omitempty"` - Message string `json:"message,omitempty"` + 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/Message carry the detail. +// On failure Output is null and Error carries the detail. It matches the +// spec's SseChatTerminalFrame. type sseChatEnvelope struct { - Command string `json:"command"` Output map[string]any `json:"output"` Metadata map[string]any `json:"metadata,omitempty"` Steps []sseStep `json:"steps"` - Error string `json:"error,omitempty"` - Message string `json:"message,omitempty"` + Error *ErrorInfo `json:"error,omitempty"` } // SSEReporter is a Reporter that streams progress back to the caller on @@ -145,6 +143,17 @@ func (r *SSEReporter) Report(_ context.Context, evt Event) { 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 { @@ -159,20 +168,18 @@ func (r *SSEReporter) writeEnvelopeLocked(evt Event) { var payload any if r.command == "chat" { - env := sseChatEnvelope{Command: r.command, Steps: steps} + env := sseChatEnvelope{Steps: steps} if failed { - env.Error = evt.Error - env.Message = evt.Message + 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{Command: r.command, Steps: steps} + env := sseEnvelope{Steps: steps} if failed { - env.Error = evt.Error - env.Message = evt.Message + env.Error = failureErrorInfo(evt) } else { feedback, ok := evt.Data["feedback"].([]map[string]any) if !ok { diff --git a/internal/progress/sse_reporter_test.go b/internal/progress/sse_reporter_test.go index a6b4d7d..9d99867 100644 --- a/internal/progress/sse_reporter_test.go +++ b/internal/progress/sse_reporter_test.go @@ -81,8 +81,8 @@ func TestSSEReporter_CompletedEnvelope(t *testing.T) { if f.event != "completed" { t.Errorf("expected event 'completed', got %q", f.event) } - if f.data["command"] != "evaluate" { - t.Errorf("expected command 'evaluate', got %v", f.data["command"]) + 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 { @@ -111,6 +111,12 @@ func TestSSEReporter_FailedEnvelope(t *testing.T) { 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()) @@ -127,26 +133,47 @@ func TestSSEReporter_FailedEnvelope(t *testing.T) { if v, ok := f.data["feedback"]; !ok || v != nil { t.Errorf("expected feedback null, got %v (present=%v)", v, ok) } - if f.data["error"] != "worker send: context deadline exceeded" { - t.Errorf("raw error not carried: %v", f.data["error"]) + 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 f.data["message"] != "We couldn't evaluate your answer. Please try again." { - t.Errorf("user message not carried: %v", f.data["message"]) + 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_PreviewCommandLabel(t *testing.T) { +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()) - if frames[0].data["command"] != "preview" { - t.Errorf("expected command 'preview', got %v", frames[0].data["command"]) + 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"]) } } @@ -167,8 +194,8 @@ func TestSSEReporter_ChatEnvelope_Completed(t *testing.T) { t.Fatalf("expected [thinking, completed], got %q", rec.Body.String()) } f := frames[1] - if f.data["command"] != "chat" { - t.Errorf("expected command 'chat', got %v", f.data["command"]) + 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) @@ -192,6 +219,11 @@ func TestSSEReporter_ChatEnvelope_Failed(t *testing.T) { 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] @@ -201,8 +233,15 @@ func TestSSEReporter_ChatEnvelope_Failed(t *testing.T) { if v, ok := f.data["output"]; !ok || v != nil { t.Errorf("expected output null, got %v (present=%v)", v, ok) } - if f.data["error"] != "chat failed: worker exited" { - t.Errorf("raw error not carried: %v", f.data["error"]) + 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"]) } } @@ -277,8 +316,12 @@ func TestSSEReporter_TerminalOnce(t *testing.T) { if len(frames) != 1 { t.Fatalf("expected exactly 1 terminal frame, got %d", len(frames)) } - if frames[0].data["message"] != "first" { - t.Errorf("expected the first terminal event to win, got %v", frames[0].data["message"]) + 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"]) } } diff --git a/internal/progress/sse_schema_parity_test.go b/internal/progress/sse_schema_parity_test.go index 052129e..9fa876a 100644 --- a/internal/progress/sse_schema_parity_test.go +++ b/internal/progress/sse_schema_parity_test.go @@ -38,7 +38,12 @@ func TestSSEFrameSchemaParity(t *testing.T) { 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."}) + 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) }) @@ -57,7 +62,12 @@ func TestSSEFrameSchemaParity(t *testing.T) { 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."}) + 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) }) @@ -93,21 +103,19 @@ func TestSSEEnvelopeStructTagsMatchSchema(t *testing.T) { {"SseProgressStep", step}, {"SseProgressStep", sseStep{Stage: "thinking", Timestamp: now, Data: map[string]any{"k": "v"}}}, {"SseChatTerminalFrame", sseChatEnvelope{ - Command: "chat", Output: map[string]any{"role": "ASSISTANT", "content": "hi"}, Metadata: map[string]any{"responseTimeMs": 12}, Steps: []sseStep{step}, }}, {"SseChatTerminalFrame", sseChatEnvelope{ - Command: "chat", Steps: []sseStep{step}, Error: "boom", Message: "failed", + Steps: []sseStep{step}, Error: &ErrorInfo{Title: "Chat failed", Message: "failed", Trace: "boom"}, }}, {"SseEvaluateTerminalFrame", sseEnvelope{ - Command: "evaluate", Feedback: []map[string]any{{"feedbackId": "fb-1", "message": "ok"}}, Steps: []sseStep{step}, }}, {"SseEvaluateTerminalFrame", sseEnvelope{ - Command: "evaluate", Steps: []sseStep{step}, Error: "boom", Message: "failed", + Steps: []sseStep{step}, Error: &ErrorInfo{Title: "Evaluation failed", Message: "failed", Trace: "boom"}, }}, } for _, c := range cases { @@ -141,13 +149,36 @@ func assertAllFieldsDocumented(t *testing.T, spec *openapi3.T, schemaName string if ref == nil || ref.Value == nil { t.Fatalf("component schema %q not found", schemaName) } + documented := documentedProps(ref.Value) for k := range m { - if ref.Value.Properties[k] == nil { + 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 { diff --git a/runtime/chat_test.go b/runtime/chat_test.go index 0e6c79c..2b52bdf 100644 --- a/runtime/chat_test.go +++ b/runtime/chat_test.go @@ -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 b358a2b..0e6888d 100644 --- a/runtime/evaluate.go +++ b/runtime/evaluate.go @@ -1,6 +1,10 @@ package runtime -import "fmt" +import ( + "fmt" + + "github.com/lambda-feedback/shimmy/internal/progress" +) type MuEdSubmissionType string @@ -46,8 +50,11 @@ type MuEdEvaluateRequest struct { 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" @@ -61,10 +68,23 @@ func MuEdToHealthResponse(result map[string]any) map[string]any { "supportsSummativeFeedback": false, "supportsDataPolicy": "NOT_SUPPORTED", "supportedAPIVersions": SupportedMuEdVersions, + "supportsStreaming": streamingEnabled, + "supportedProgressStages": evaluateProgressStages, }, } } +// evaluateProgressStages is the set of SseProgressStep.stage values an +// /evaluate SSE stream can emit, advertised via +// capabilities.supportedProgressStages. +var evaluateProgressStages = []string{ + string(progress.StagePreparing), + string(progress.StageStarting), + string(progress.StageEvaluating), + string(progress.StageCompleted), + string(progress.StageFailed), +} + func muEdContentKey(t MuEdSubmissionType) string { switch t { case MuEdMath: diff --git a/runtime/schema/mued_v0.1.0.yml b/runtime/schema/mued_v0.1.0.yml index f6ea8fc..e617092 100644 --- a/runtime/schema/mued_v0.1.0.yml +++ b/runtime/schema/mued_v0.1.0.yml @@ -335,7 +335,23 @@ paths: version: 1 responses: '200': - description: Successfully generated feedback. + description: | + Successfully generated feedback. + + ### Streaming variant (opt-in) + + If the request sends `Accept: text/event-stream` (standalone/serve + mode only; ignored under AWS Lambda), the response is a Server-Sent + Events progress stream instead of a single JSON body. The HTTP + status stays `200` for the whole stream, including failures. It is: + zero or more progress frames whose SSE `event:` is the stage name + and `data:` is an `SseProgressStep`; optional `:`-prefixed + keep-alive comment lines; and exactly one terminal frame — + `event: completed` with `data:` an `SseEvaluateTerminalFrame` (the + `200` body under `feedback`, plus a `steps` replay), or + `event: failed` with `data:` an `SseEvaluateTerminalFrame` whose + `error` is an `ErrorResponse`. `X-Request-Id` / `X-Api-Version` are + sent once as response headers when the stream opens. headers: X-Request-Id: description: Request id for tracing this request across services. @@ -377,17 +393,7 @@ paths: message: The overall structure of your answer is clear and easy to follow. text/event-stream: schema: - type: string - description: | - Opt-in Server-Sent Events stream, selected with the request - header `Accept: text/event-stream` (standalone/serve mode - only; ignored under AWS Lambda). Intermediate frames: - `event: starting | preparing | evaluating`, `data` = a - `SseProgressStep` JSON object. Keep-alive `: ping` comment - lines. Exactly one terminal frame: `event: completed | failed`, - `data` = an `SseEvaluateTerminalFrame` JSON object. The HTTP - status is 200 even for a `failed` frame — the failure is - in-band. + $ref: '#/components/schemas/SseEvaluateTerminalFrame' '202': $ref: '#/components/responses/202-Accepted' '400': @@ -596,7 +602,23 @@ paths: temperature: 0.5 responses: '200': - description: Successful chat response. + description: | + Successful chat response. + + ### Streaming variant (opt-in) + + If the request sends `Accept: text/event-stream` (standalone/serve + mode only; ignored under AWS Lambda), the response is a Server-Sent + Events progress stream instead of a single JSON body. The HTTP + status stays `200` for the whole stream, including failures. It is: + zero or more progress frames whose SSE `event:` is the stage name + and `data:` is an `SseProgressStep`; optional `:`-prefixed + keep-alive comment lines; and exactly one terminal frame — + `event: completed` with `data:` an `SseChatTerminalFrame` (the + `200` body as `output` / `metadata`, plus a `steps` replay), or + `event: failed` with `data:` an `SseChatTerminalFrame` whose + `error` is an `ErrorResponse`. `X-Request-Id` / `X-Api-Version` are + sent once as response headers when the stream opens. headers: X-Request-Id: description: Request id for tracing this request across services. @@ -653,16 +675,7 @@ paths: outputTokens: 143 text/event-stream: schema: - type: string - description: | - Opt-in Server-Sent Events stream, selected with the request - header `Accept: text/event-stream` (standalone/serve mode - only; ignored under AWS Lambda). Intermediate frames: - `event: starting | preparing | thinking`, `data` = a - `SseProgressStep` JSON object. Keep-alive `: ping` comment - lines. Exactly one terminal frame: `event: completed | failed`, - `data` = a `SseChatTerminalFrame` JSON object. The HTTP status - is 200 even for a `failed` frame — the failure is in-band. + $ref: '#/components/schemas/SseChatTerminalFrame' '400': $ref: '#/components/responses/400-BadRequest-2' '403': @@ -1560,6 +1573,21 @@ components: supportsSummativeFeedback: type: boolean description: Indicates whether the service supports feedback with points / grading signals. + supportsStreaming: + type: boolean + description: | + Whether /evaluate supports opt-in SSE progress streaming via + `Accept: text/event-stream`. Distinct from + `configuration.llm.stream`. + supportedProgressStages: + type: + - array + - 'null' + description: | + Optional list of `SseProgressStep.stage` values an /evaluate SSE + stream may emit. Informative; clients must tolerate unlisted values. + items: + type: string supportsDataPolicy: $ref: '#/components/schemas/DataPolicySupport' supportedArtefactProfiles: @@ -1721,7 +1749,19 @@ components: description: Indicates whether the service supports adapting to user preferences. supportsStreaming: type: boolean - description: Indicates whether the service supports streaming responses. + description: | + Whether /chat supports opt-in SSE progress streaming via + `Accept: text/event-stream`. Distinct from + `configuration.llm.stream`. + supportedProgressStages: + type: + - array + - 'null' + description: | + Optional list of `SseProgressStep.stage` values a /chat SSE + stream may emit. Informative; clients must tolerate unlisted values. + items: + type: string supportsDataPolicy: $ref: '#/components/schemas/DataPolicySupport' supportedLanguages: @@ -1770,91 +1810,133 @@ components: SseProgressStep: type: object description: | - Payload (`data`) of an intermediate SSE progress frame - (`event: starting | preparing | thinking | evaluating`). The same - shape also appears in `steps[]` of the terminal frame. + A single progress step emitted while an operation runs. Carried as + the SSE `data:` payload of an intermediate progress frame (the SSE + `event:` field carries the stage name), and replayed in the + terminal frame's `steps` array. + additionalProperties: true required: - stage - timestamp properties: stage: type: string - description: Lifecycle stage this step reports. + description: | + Lifecycle stage this step reports. Informative, not a fixed + enum: implementations may add stages and clients must tolerate + unrecognised values. Common values: "preparing", "starting", + "evaluating" (evaluate), "thinking" (chat), and the terminal + "completed" / "failed". message: - type: string + type: + - string + - 'null' description: Short, learner/teacher-facing description of the step. data: type: object additionalProperties: true - description: Free-form payload attached by a worker-authored sub-step. + description: | + Free-form payload attached by a worker-authored sub-step + (shimmy extension; the canonical spec relies on + additionalProperties for this). timestamp: type: string format: date-time - SseChatTerminalFrame: + SseTerminalSteps: type: object description: | - Payload (`data`) of the single terminal SSE frame for `POST /chat` - (`event: completed | failed`). On a `failed` frame `output` is null - and `error`/`message` carry the detail. + Shared fragment of the terminal SSE frame: the ordered replay of + every progress step emitted during the stream, so a client that + connected late or dropped frames still receives the full trace. required: - - command - steps properties: - command: - type: string - description: The µEd command being processed (e.g. "chat"). - output: - type: - - object - - 'null' - additionalProperties: true - description: The generated assistant response; null on failure. - metadata: - type: - - object - - 'null' - additionalProperties: true - description: Optional metadata about response generation. steps: type: array + description: Ordered list of every SseProgressStep emitted during the stream. items: $ref: '#/components/schemas/SseProgressStep' - error: - type: string - description: Raw technical detail (failure frames only; for logs/support). - message: - type: string - description: Human-facing status text. - SseEvaluateTerminalFrame: + StreamingCapabilities: type: object description: | - Payload (`data`) of the single terminal SSE frame for - `POST /evaluate` (`event: completed | failed`). On a `failed` frame - `feedback` is null and `error`/`message` carry the detail. - required: - - command - - steps + Shared capability fragment describing an operation's support for + opt-in Server-Sent Events (SSE) progress streaming. + additionalProperties: true properties: - command: - type: string - description: The µEd command being processed (e.g. "eval", "preview"). - feedback: + supportsStreaming: + type: boolean + description: | + Whether this operation supports opt-in SSE progress streaming, + selected per request with `Accept: text/event-stream`. Distinct + from `configuration.llm.stream`, which governs token-level + streaming from the LLM provider. + supportedProgressStages: type: - array - 'null' + description: | + Optional list of `SseProgressStep.stage` values this service may + emit. Informative only; clients must tolerate unlisted stages. items: - $ref: '#/components/schemas/Feedback' - description: The generated feedback items; null on failure. - steps: - type: array - items: - $ref: '#/components/schemas/SseProgressStep' - error: - type: string - description: Raw technical detail (failure frames only; for logs/support). - message: - type: string - description: Human-facing status text. + type: string + SseChatTerminalFrame: + type: object + description: | + Payload (`data`) of the single terminal SSE frame for `POST /chat` + (`event: completed | failed`). On `completed`, `output`/`metadata` + hold the endpoint's normal 200 body and `error` is absent. On + `failed`, `output` is null and `error` holds an ErrorResponse. The + HTTP status stays 200 regardless. `steps` is always present. + allOf: + - $ref: '#/components/schemas/SseTerminalSteps' + - type: object + properties: + output: + type: + - object + - 'null' + additionalProperties: true + description: The generated assistant response; null on a failed frame. + metadata: + type: + - object + - 'null' + additionalProperties: true + description: Optional metadata about response generation. + error: + type: + - object + - 'null' + description: Present only on a failed frame. + allOf: + - $ref: '#/components/schemas/ErrorResponse' + SseEvaluateTerminalFrame: + type: object + description: | + Payload (`data`) of the single terminal SSE frame for + `POST /evaluate` (`event: completed | failed`). On `completed`, + `feedback` holds the endpoint's normal 200 body and `error` is + absent. On `failed`, `feedback` is null and `error` holds an + ErrorResponse. The HTTP status stays 200 regardless. `steps` is + always present. + allOf: + - $ref: '#/components/schemas/SseTerminalSteps' + - type: object + properties: + feedback: + type: + - array + - 'null' + items: + $ref: '#/components/schemas/Feedback' + description: The generated feedback items; null on a failed frame. + error: + type: + - object + - 'null' + description: Present only on a failed frame. + allOf: + - $ref: '#/components/schemas/ErrorResponse' responses: 202-Accepted: description: Request accepted for asynchronous evaluation processing. From 4a6650160402c0ba1749a6fb96f8f88ffb8f6759 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 09:17:41 +0100 Subject: [PATCH 19/28] Add SSE progress streaming support to `/chat` capabilities - Extend `MuEdToChatHealthResponse` to handle new `streamingEnabled` flag and expose `supportsStreaming` and `supportedProgressStages`. - Introduce `chatProgressStages` defining possible SSE progress stages for `/chat`. - Update `/chat` SSE capabilities to align with shimmy-layer features. --- runtime/chat.go | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/runtime/chat.go b/runtime/chat.go index 4cebe92..1dbb92d 100644 --- a/runtime/chat.go +++ b/runtime/chat.go @@ -3,6 +3,8 @@ package runtime import ( "encoding/json" "fmt" + + "github.com/lambda-feedback/shimmy/internal/progress" ) // ChatRequest is the dispatcher-level request for the chat command. @@ -110,7 +112,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) @@ -131,6 +138,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, @@ -144,3 +153,13 @@ 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. +var chatProgressStages = []string{ + string(progress.StagePreparing), + string(progress.StageStarting), + string(progress.StageThinking), + string(progress.StageCompleted), + string(progress.StageFailed), +} From ab0b034eb38c01c2b6fde28b4ed2e204a7c14fd6 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 09:21:44 +0100 Subject: [PATCH 20/28] Replace `MuEdChatRole` type with untyped string for `Role` field in chat messages - Simplify `MuEdChatMessage` by removing `MuEdChatRole` type in favor of untyped strings for role handling. - Update tests to align with the revised structure. --- runtime/chat.go | 17 ++++++----------- runtime/chat_test.go | 10 +++++----- 2 files changed, 11 insertions(+), 16 deletions(-) diff --git a/runtime/chat.go b/runtime/chat.go index 1dbb92d..b72d447 100644 --- a/runtime/chat.go +++ b/runtime/chat.go @@ -17,18 +17,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 diff --git a/runtime/chat_test.go b/runtime/chat_test.go index 2b52bdf..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) From 4ea03950e3f9291438c892629e679d184e3aafc2 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 09:31:43 +0100 Subject: [PATCH 21/28] Refactor progress reporting to include structured error details - Add `ErrorInfo` to payloads for `StageFailed` events, standardizing failure handling across callback and SSE responses. - Update `/chat` and `/evaluate` handlers and tests to validate structured error propagation. - Add unit tests ensuring error object correctness in HTTP callbacks and SSE terminal frames. --- handler/chat.go | 9 ++-- handler/evaluate.go | 9 ++-- handler/evaluate_test.go | 5 ++ internal/progress/http_reporter.go | 13 ++++-- internal/progress/http_reporter_test.go | 62 +++++++++++++++++++++++++ 5 files changed, 87 insertions(+), 11 deletions(-) diff --git a/handler/chat.go b/handler/chat.go index de53e0b..d070283 100644 --- a/handler/chat.go +++ b/handler/chat.go @@ -84,10 +84,11 @@ func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) { output, metadata, termErr := h.produceChatOutput(resp, err) if termErr != nil { progress.Emit(ctx, progress.Event{ - Stage: progress.StageFailed, - Command: string(runtime.CommandChat), - Message: termErr.userMessage, - Error: termErr.rawError, + 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 diff --git a/handler/evaluate.go b/handler/evaluate.go index 889b3f1..8969931 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -239,10 +239,11 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { feedback, termErr := h.produceFeedback(resp, isPreview) if termErr != nil { progress.Emit(ctx, progress.Event{ - Stage: progress.StageFailed, - Command: string(command), - Message: termErr.userMessage, - Error: termErr.rawError, + Stage: progress.StageFailed, + Command: string(command), + Message: termErr.userMessage, + Error: termErr.rawError, + ErrorInfo: termErr.progressErrorInfo("Evaluation failed"), }) if termErr.passthrough { diff --git a/handler/evaluate_test.go b/handler/evaluate_test.go index 9f83467..441c175 100644 --- a/handler/evaluate_test.go +++ b/handler/evaluate_test.go @@ -406,6 +406,11 @@ func TestMuEdServeEvaluate_ProgressCallback_Failure(t *testing.T) { 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) { diff --git a/internal/progress/http_reporter.go b/internal/progress/http_reporter.go index 4d6fbb2..5f47e2e 100644 --- a/internal/progress/http_reporter.go +++ b/internal/progress/http_reporter.go @@ -11,13 +11,16 @@ import ( "go.uber.org/zap" ) -// payload is the JSON body POSTed to the callback URL for each event. +// 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 string `json:"error,omitempty"` + Error *ErrorInfo `json:"error,omitempty"` Data map[string]any `json:"data,omitempty"` Timestamp time.Time `json:"timestamp"` } @@ -75,12 +78,16 @@ func (r *httpCallbackReporter) Report(ctx context.Context, evt Event) { } 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: evt.Error, + Error: errInfo, Data: evt.Data, Timestamp: evt.Timestamp, }) diff --git a/internal/progress/http_reporter_test.go b/internal/progress/http_reporter_test.go index 2d7ee79..efd07fa 100644 --- a/internal/progress/http_reporter_test.go +++ b/internal/progress/http_reporter_test.go @@ -49,6 +49,68 @@ func TestHTTPCallbackReporter_Report_DeliversPayload(t *testing.T) { } } +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 From f7a1622a65b79d29806e06fd413404fa549d7a45 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 09:31:57 +0100 Subject: [PATCH 22/28] Updated gitignore --- .gitignore | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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/ From 6a399419b3678bde073f80f3c215d1c888edcff8 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 09:40:37 +0100 Subject: [PATCH 23/28] Update README to document structured error object in `failed` events - Clarify the use of `ErrorResponse` in terminal `failed` stages for `/evaluate` and `/chat`. - Add examples showcasing the updated payload format and structured error propagation. - Refine descriptions of terminal frame handling and SSE response structure. --- README.md | 44 ++++++++++++++++++++++++++++++++------------ 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index af12617..046106a 100644 --- a/README.md +++ b/README.md @@ -224,7 +224,7 @@ Stages, in order: | `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 safe to show to an end user; `error` carries raw technical detail for logs only. | +| `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. @@ -257,6 +257,24 @@ Example terminal event, with the feedback payload attached: } ``` +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) @@ -307,12 +325,12 @@ though a multi-case evaluation re-enters them per case; worker-authored `evaluat `timestamp`). The stream then ends with exactly one terminal frame — `event: completed` or `event: failed` -— carrying the result plus every step that preceded it, and the connection closes: +— carrying the endpoint's normal `200` body plus every step that preceded it, and the +connection closes: ``` event: completed -data: {"command":"evaluate", - "feedback":[{"awardedPoints":1,"message":"Well done"}], +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":"…"}]} @@ -320,10 +338,12 @@ data: {"command":"evaluate", ``` event: failed -data: {"command":"evaluate","feedback":null, +data: {"feedback":null, "steps":[ /* whatever streamed before the failure */ ], - "error":"worker send: context deadline exceeded", - "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"}} ``` For `/chat` the terminal frame carries `output` (and optional `metadata`) instead of @@ -331,17 +351,17 @@ For `/chat` the terminal frame carries `output` (and optional `metadata`) instea ``` event: completed -data: {"command":"chat", - "output":{"role":"ASSISTANT","content":"…"}, +data: {"output":{"role":"ASSISTANT","content":"…"}, "metadata":{ /* optional, worker-supplied */ }, "steps":[ /* preparing, starting, thinking… */ ]} ``` -A failed `/chat` frame has `"output":null` plus `error`/`message`. +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. `command` is `"evaluate"`, `"preview"`, or `"chat"`. The correlation id is in the -`X-Request-Id` response header, not the body. Response headers: `Content-Type: +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 From 0ff55bc8ce0bbf109cbde63d41186cee6bcb4443 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 10:14:53 +0100 Subject: [PATCH 24/28] Refactor error handling and documentation in `/evaluate` and `/stream` - Clarify comments on error propagation and structured error details in `StageFailed` events. - Improve descriptions of terminal frame shapes, progress events, and fallback mechanisms. --- handler/evaluate.go | 10 ++++++---- handler/stream.go | 9 ++++----- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/handler/evaluate.go b/handler/evaluate.go index 8969931..74fcbff 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -298,14 +298,16 @@ type terminalError struct { status int body []byte - // muEd* describe a shimmy-internal error for writeMuEdError (buffered - // path only). + // 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 feed the StageFailed progress event and, - // on the streaming path, the "failed" SSE frame. + // 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 } diff --git a/handler/stream.go b/handler/stream.go index 338258b..b4fe283 100644 --- a/handler/stream.go +++ b/handler/stream.go @@ -21,11 +21,10 @@ import ( // already committed, every outcome of run — including an internal error — // becomes a "failed" frame, never an HTTP error. // -// cmdLabel selects the terminal envelope shape ("evaluate"/"preview" -> -// feedback[]; "chat" -> output/metadata) and is the frame's "command". -// command is the µEd command string carried on the emitted progress -// events. doneMessage is the human-facing text on the terminal completed -// event. +// 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 From b1798b05bbcf909dbd567f0b2c0edfb72368461a Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 10:20:00 +0100 Subject: [PATCH 25/28] Enhance error handling and test coverage for `/chat` SSE responses - Add validation for `ErrorResponse` objects in terminal `failed` events. - Improve structured error propagation by asserting `error.title` in SSE tests. - Clarify comments on `command` field usage in terminal frame shapes. --- handler/chat_stream_test.go | 3 +++ internal/execution/supervisor/supervisor.go | 1 + internal/progress/sse_reporter.go | 2 ++ 3 files changed, 6 insertions(+) diff --git a/handler/chat_stream_test.go b/handler/chat_stream_test.go index f8aaee0..58eadae 100644 --- a/handler/chat_stream_test.go +++ b/handler/chat_stream_test.go @@ -292,6 +292,9 @@ func TestServeChat_SSE_TerminalPayloadValidated(t *testing.T) { 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"]) }) } diff --git a/internal/execution/supervisor/supervisor.go b/internal/execution/supervisor/supervisor.go index 98a625d..b316026 100644 --- a/internal/execution/supervisor/supervisor.go +++ b/internal/execution/supervisor/supervisor.go @@ -184,6 +184,7 @@ func (s *WorkerSupervisor) Send( }, }) return nil, fmt.Errorf("failed to acquire worker: %w", err) + } progress.Emit(ctx, progress.Event{ Stage: progress.StagePreparing, diff --git a/internal/progress/sse_reporter.go b/internal/progress/sse_reporter.go index 06a85e5..00ceb07 100644 --- a/internal/progress/sse_reporter.go +++ b/internal/progress/sse_reporter.go @@ -59,6 +59,8 @@ type sseChatEnvelope struct { 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 From 39a8cc9ed12260c41b3017cca65b12f08523a338 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 12:20:11 +0100 Subject: [PATCH 26/28] =?UTF-8?q?Introduce=20versioned=20=C2=B5Ed=20adapte?= =?UTF-8?q?rs=20and=20middleware=20for=20multi-version=20API=20support?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/standalone/module_test.go | 28 +++++ handler/chat.go | 16 +-- handler/evaluate.go | 92 +++++++------- handler/mued_version_test.go | 145 +++++++++++++++++++++++ internal/server/module.go | 4 +- internal/server/openapi.go | 71 +++++++++-- internal/server/openapi_test.go | 135 +++++++++++++++++---- internal/server/server.go | 4 +- internal/server/testdata/mued_v0.2.0.yml | 43 +++++++ runtime/evaluate.go | 2 +- runtime/module.go | 3 + runtime/mued_adapter.go | 117 ++++++++++++++++++ runtime/mued_adapter_test.go | 94 +++++++++++++++ runtime/mued_v0_1_0.go | 61 ++++++++++ runtime/mued_v0_1_0_test.go | 118 ++++++++++++++++++ runtime/schema/openapi.go | 67 ++++++++++- runtime/version.go | 22 ++-- 17 files changed, 913 insertions(+), 109 deletions(-) create mode 100644 app/standalone/module_test.go create mode 100644 handler/mued_version_test.go create mode 100644 internal/server/testdata/mued_v0.2.0.yml create mode 100644 runtime/mued_adapter.go create mode 100644 runtime/mued_adapter_test.go create mode 100644 runtime/mued_v0_1_0.go create mode 100644 runtime/mued_v0_1_0_test.go diff --git a/app/standalone/module_test.go b/app/standalone/module_test.go new file mode 100644 index 0000000..08d9715 --- /dev/null +++ b/app/standalone/module_test.go @@ -0,0 +1,28 @@ +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" + "github.com/stretchr/testify/assert" +) + +// TestModule_DependencyGraphValid makes sure the standalone fx graph — the µEd +// handlers, the per-version OpenAPI specs, and the µEd version registry provided +// by runtime.Module — resolves without missing or cyclic dependencies. +func TestModule_DependencyGraphValid(t *testing.T) { + err := fx.ValidateApp( + fx.NopLogger, + fx.Supply(fx.Annotate(context.Background(), fx.As(new(context.Context)))), + fx.Supply(zap.NewNop()), + fx.Supply(config.Config{}), + runtime.Module(config.Config{}.Runtime), + Module(Config{}), + ) + assert.NoError(t, err) +} diff --git a/handler/chat.go b/handler/chat.go index 24e9e26..593c434 100644 --- a/handler/chat.go +++ b/handler/chat.go @@ -15,7 +15,7 @@ func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -31,13 +31,7 @@ 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 @@ -55,7 +49,7 @@ func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) { return } - chatResp, err := runtime.MuEdToChatResponse(resultMap) + chatResp, err := adapter.EncodeChat(resultMap) if err != nil { h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", err.Error(), nil) return @@ -73,7 +67,7 @@ func (h *MuEdHandler) ServeChatHealth(w http.ResponseWriter, r *http.Request) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -95,7 +89,7 @@ func (h *MuEdHandler) ServeChatHealth(w http.ResponseWriter, r *http.Request) { return } - healthResp := runtime.MuEdToChatHealthResponse(resultMap) + healthResp := adapter.EncodeChatHealth(resultMap) statusCode := http.StatusOK if status, ok := healthResp["status"].(string); ok && status == string(runtime.MuEdChatHealthStatusUnavailable) { diff --git a/handler/evaluate.go b/handler/evaluate.go index c462b17..1654dd7 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -19,58 +19,73 @@ const muEdVersionHeader = "X-Api-Version" 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 } 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 } 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, } } +// 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() +} + 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. @@ -102,7 +117,7 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -118,20 +133,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 +145,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)) @@ -184,11 +181,10 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { return } - var feedback []map[string]any - if isPreview { - feedback = runtime.MuEdToPreviewFeedback(result) - } else { - feedback = runtime.MuEdToEvaluateFeedback(result) + feedback, err := adapter.EncodeEvaluateFeedback(command, result) + if err != nil { + h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", "failed to build feedback", nil) + return } w.Header().Set("Content-Type", "application/json") @@ -203,7 +199,7 @@ func (h *MuEdHandler) ServeHealth(w http.ResponseWriter, r *http.Request) { return } - version, ok := h.checkMuEdVersion(w, r) + version, adapter, ok := h.checkMuEdVersion(w, r) if !ok { return } @@ -228,7 +224,7 @@ func (h *MuEdHandler) ServeHealth(w http.ResponseWriter, r *http.Request) { return } - result := runtime.MuEdToHealthResponse(legacyResult) + result := adapter.EncodeHealth(legacyResult) statusCode := http.StatusOK if s, ok := result["status"].(string); ok && s == "UNAVAILABLE" { diff --git a/handler/mued_version_test.go b/handler/mued_version_test.go new file mode 100644 index 0000000..b33eddc --- /dev/null +++ b/handler/mued_version_test.go @@ -0,0 +1,145 @@ +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) 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) map[string]any { + return map[string]any{"marker": "chat-health-9.9.9"} +} + +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/internal/server/module.go b/internal/server/module.go index 7644bed..630c3be 100644 --- a/internal/server/module.go +++ b/internal/server/module.go @@ -6,8 +6,8 @@ func Module(config HttpConfig) fx.Option { return fx.Module("server", // provide config fx.Supply(config), - // provide openapi spec - fx.Provide(LoadOpenAPISpec), + // provide openapi specs (one per supported µEd version) + fx.Provide(LoadOpenAPISpecs), // provide server fx.Provide(NewLifecycleServer), // invoke server diff --git a/internal/server/openapi.go b/internal/server/openapi.go index 8d8ed88..1fff705 100644 --- a/internal/server/openapi.go +++ b/internal/server/openapi.go @@ -9,35 +9,85 @@ import ( "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" "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 defaults to runtime.MuEdResolveVersion. 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 = runtime.MuEdResolveVersion + } + + 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 @@ -64,8 +114,7 @@ func OpenAPIMiddleware(spec *openapi3.T, log *zap.Logger) (func(http.Handler) ht // Snapshot body before validation — ValidateResponse drains the buffer. bodyBytes := rec.Body.Bytes() - - // Validate response (lenient — log only) + // Validate response respInput := &openapi3filter.ResponseValidationInput{ RequestValidationInput: reqInput, Status: rec.Code, diff --git a/internal/server/openapi_test.go b/internal/server/openapi_test.go index 554f354..eb67028 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,50 @@ 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") + for version, spec := range specs { + assert.NotNilf(t, spec, "spec for %s", version) + } +} + 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 +147,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 +164,85 @@ 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. +// 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 +253,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/server.go b/internal/server/server.go index 6a94ea5..426b8b1 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -19,7 +19,7 @@ type HttpServerParams struct { Context context.Context Config HttpConfig - Spec *openapi3.T + Specs map[string]*openapi3.T Handlers []*HttpHandler `group:"handlers"` Logger *zap.Logger @@ -41,7 +41,7 @@ func NewHttpServer(params HttpServerParams) (*HttpServer, error) { } var handler http.Handler = NormalizePath(mux) - openAPIMiddleware, err := OpenAPIMiddleware(params.Spec, params.Logger) + openAPIMiddleware, err := OpenAPIMiddleware(params.Specs, nil, params.Logger) if err != nil { return nil, fmt.Errorf("initialising OpenAPI middleware: %w", err) } 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/runtime/evaluate.go b/runtime/evaluate.go index bc9e868..1085363 100644 --- a/runtime/evaluate.go +++ b/runtime/evaluate.go @@ -50,7 +50,7 @@ func MuEdToHealthResponse(result map[string]any) map[string]any { "supportsFormativeFeedback": true, "supportsSummativeFeedback": false, "supportsDataPolicy": "NOT_SUPPORTED", - "supportedAPIVersions": SupportedMuEdVersions, + "supportedAPIVersions": SupportedMuEdVersions(), }, } } 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..b23f5c4 --- /dev/null +++ b/runtime/mued_adapter.go @@ -0,0 +1,117 @@ +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. + EncodeHealth(legacyResult map[string]any) 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. + EncodeChatHealth(result map[string]any) map[string]any +} + +// 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..054a206 --- /dev/null +++ b/runtime/mued_adapter_test.go @@ -0,0 +1,94 @@ +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) 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) map[string]any { + return map[string]any{"from": a.version} +} + +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_HasV010(t *testing.T) { + reg := runtime.DefaultMuEdRegistry() + + assert.Equal(t, []string{"0.1.0"}, reg.Versions()) + assert.Equal(t, []string{"0.1.0"}, runtime.SupportedMuEdVersions()) + assert.True(t, runtime.MuEdIsVersionSupported("0.1.0")) + assert.False(t, runtime.MuEdIsVersionSupported("99.0.0")) + + assert.Equal(t, "0.1.0", runtime.MuEdResolveVersion("")) + assert.Equal(t, "0.1.0", runtime.MuEdResolveVersion("0.1.0")) + assert.Equal(t, "0.1.0", runtime.MuEdResolveVersion("99.0.0")) + + require.NotNil(t, reg.Adapter("0.1.0")) + assert.Equal(t, "0.1.0", reg.Adapter("0.1.0").Version()) +} diff --git a/runtime/mued_v0_1_0.go b/runtime/mued_v0_1_0.go new file mode 100644 index 0000000..b16422d --- /dev/null +++ b/runtime/mued_v0_1_0.go @@ -0,0 +1,61 @@ +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) map[string]any { + return MuEdToHealthResponse(legacyResult) +} + +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) map[string]any { + return MuEdToChatHealthResponse(result) +} diff --git a/runtime/mued_v0_1_0_test.go b/runtime/mued_v0_1_0_test.go new file mode 100644 index 0000000..ad9087d --- /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), a.EncodeHealth(result)) + } +} + +func TestMuEdV010_Chat_MatchesFreeFunctions(t *testing.T) { + a := v010(t) + + chatReq := runtime.MuEdChatRequest{Messages: []runtime.MuEdChatMessage{{Role: runtime.MuEdChatRoleUser, 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), a.EncodeChatHealth(healthResult)) +} diff --git a/runtime/schema/openapi.go b/runtime/schema/openapi.go index b771890..e7548c3 100644 --- a/runtime/schema/openapi.go +++ b/runtime/schema/openapi.go @@ -1,6 +1,67 @@ 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. +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) } From 3571c8a45eb178a110d4e2a819f188b775634867 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 12:59:59 +0100 Subject: [PATCH 27/28] Serve one shared OpenAPI-validated handler chain on both deployments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The standalone server wrapped the route mux with NormalizePath + the OpenAPI request/response validation middleware; the Lambda adapter wrapped it with NormalizePath only. That left Lambda requests unvalidated against the spec and, more importantly, Lambda responses unchecked — so a non-conforming response would 500 on standalone but ship as-is on Lambda. Extract the wrapped chain into server.NewMux (mux + NormalizePath + per-version OpenAPI validation) and a server.HandlerModule fx module. Both app/standalone and app/lambda now depend on it and serve the exact same *server.Mux, so the two deployments validate identically. The only remaining deployment difference is transport: standalone runs an http.Server + listener (and optional h2c); Lambda hands the same handler to httpadapter. NewHttpServer / NewLifecycleServer no longer build the chain or return an error. Added fx.ValidateApp coverage for the Lambda graph and a NewMux test asserting validation + path normalisation are applied. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01YAusDUAMwEVN8hAV4N8qGk --- app/lambda/handler.go | 38 +++++++++---------- app/lambda/module.go | 3 ++ app/lambda/module_test.go | 28 ++++++++++++++ internal/server/module.go | 18 ++++++++- internal/server/mux.go | 44 ++++++++++++++++++++++ internal/server/mux_test.go | 73 +++++++++++++++++++++++++++++++++++++ internal/server/server.go | 32 ++++------------ 7 files changed, 189 insertions(+), 47 deletions(-) create mode 100644 app/lambda/module_test.go create mode 100644 internal/server/mux.go create mode 100644 internal/server/mux_test.go 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..1300a97 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" ) @@ -16,6 +17,8 @@ func Module(config Config) fx.Option { logging.DecorateLogger("lambda"), // 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..9b72dd0 --- /dev/null +++ b/app/lambda/module_test.go @@ -0,0 +1,28 @@ +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" + "github.com/stretchr/testify/assert" +) + +// TestModule_DependencyGraphValid makes sure the Lambda fx graph resolves — +// notably that it now pulls in server.HandlerModule so the Lambda adapter serves +// the same OpenAPI-validated handler chain as the standalone server. +func TestModule_DependencyGraphValid(t *testing.T) { + err := fx.ValidateApp( + fx.NopLogger, + fx.Supply(fx.Annotate(context.Background(), fx.As(new(context.Context)))), + fx.Supply(zap.NewNop()), + fx.Supply(config.Config{}), + runtime.Module(config.Config{}.Runtime), + Module(Config{}), + ) + assert.NoError(t, err) +} diff --git a/internal/server/module.go b/internal/server/module.go index 630c3be..96c503d 100644 --- a/internal/server/module.go +++ b/internal/server/module.go @@ -2,12 +2,26 @@ 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 specs (one per supported µEd version) - fx.Provide(LoadOpenAPISpecs), + // 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/server.go b/internal/server/server.go index 426b8b1..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 - Specs map[string]*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.Specs, nil, 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 { From 3cd9c295e33a38d4ed6fdf2c1af732901bb4bd9d Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Thu, 3 Sep 2026 15:27:01 +0100 Subject: [PATCH 28/28] =?UTF-8?q?Introduce=20`0.1.1-dev`=20=C2=B5Ed=20vers?= =?UTF-8?q?ion=20with=20SSE=20progress=20streaming=20support?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add versioned adapters and `SupportsStreaming` to conditionally enable SSE. - Update `/chat` and `/evaluate` handlers and tests to use `X-Api-Version`. - Refine OpenAPI specs and README to document `0.1.1-dev`-specific streaming behavior. - Enhance test coverage for version negotiation and SSE response validation. --- README.md | 7 +- cmd/root.go | 2 +- handler/chat.go | 2 +- handler/chat_stream_test.go | 3 + handler/chat_test.go | 4 +- handler/evaluate.go | 2 +- handler/evaluate_stream_test.go | 4 + handler/evaluate_test.go | 4 +- handler/mued_version_test.go | 1 + internal/progress/sse_schema_parity_test.go | 3 + internal/server/openapi_test.go | 12 ++ runtime/mued_adapter.go | 7 + runtime/mued_adapter_test.go | 21 +- runtime/mued_v0_1_0.go | 4 + runtime/schema/mued_v0.1.0.yml | 201 +------------------- runtime/schema/openapi.go | 6 + 16 files changed, 72 insertions(+), 211 deletions(-) diff --git a/README.md b/README.md index 7ac439e..624ba1c 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ GLOBAL OPTIONS: --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'. Standalone/serve mode only; ignored under AWS Lambda. (default: true) [$PROGRESS_STREAM_ENABLED] + --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 @@ -295,6 +295,11 @@ stand up a `callbackUrl` receiver can opt in with an `Accept: text/event-stream` 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`. diff --git a/cmd/root.go b/cmd/root.go index 56e64ff..11fcf7c 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -101,7 +101,7 @@ functions on arbitrary, serverless platforms.` }, &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'. Standalone/serve mode only; ignored under AWS Lambda.", + 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"}, diff --git a/handler/chat.go b/handler/chat.go index 9147a40..64b0b5e 100644 --- a/handler/chat.go +++ b/handler/chat.go @@ -55,7 +55,7 @@ func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) { callbackURL = *chatReq.CallbackUrl } - streaming := h.streamingEnabled() && acceptsEventStream(r) + 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") diff --git a/handler/chat_stream_test.go b/handler/chat_stream_test.go index 58eadae..7ec5eb4 100644 --- a/handler/chat_stream_test.go +++ b/handler/chat_stream_test.go @@ -40,6 +40,8 @@ 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 } @@ -228,6 +230,7 @@ func TestServeChat_SSE_Heartbeat(t *testing.T) { 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) 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 b3d144e..3c15e32 100644 --- a/handler/evaluate.go +++ b/handler/evaluate.go @@ -222,7 +222,7 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { callbackURL = *muEdReq.CallbackUrl } - streaming := h.streamingEnabled() && acceptsEventStream(r) + 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") diff --git a/handler/evaluate_stream_test.go b/handler/evaluate_stream_test.go index 5a73386..e94bbc3 100644 --- a/handler/evaluate_stream_test.go +++ b/handler/evaluate_stream_test.go @@ -45,6 +45,8 @@ 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 } @@ -366,6 +368,7 @@ func TestServeEvaluate_SSE_Heartbeat(t *testing.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) @@ -425,6 +428,7 @@ func TestServeEvaluate_SSE_ThroughOpenAPIMiddleware(t *testing.T) { 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) diff --git a/handler/evaluate_test.go b/handler/evaluate_test.go index 441c175..4274d28 100644 --- a/handler/evaluate_test.go +++ b/handler/evaluate_test.go @@ -636,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)) @@ -686,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)) diff --git a/handler/mued_version_test.go b/handler/mued_version_test.go index 433b27b..c332390 100644 --- a/handler/mued_version_test.go +++ b/handler/mued_version_test.go @@ -40,6 +40,7 @@ func (markerAdapter) EncodeChat(map[string]any) (map[string]any, error) { 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{ diff --git a/internal/progress/sse_schema_parity_test.go b/internal/progress/sse_schema_parity_test.go index 9fa876a..d1b6885 100644 --- a/internal/progress/sse_schema_parity_test.go +++ b/internal/progress/sse_schema_parity_test.go @@ -127,6 +127,9 @@ func TestSSEEnvelopeStructTagsMatchSchema(t *testing.T) { // --- 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() diff --git a/internal/server/openapi_test.go b/internal/server/openapi_test.go index b8a2652..4052b01 100644 --- a/internal/server/openapi_test.go +++ b/internal/server/openapi_test.go @@ -25,9 +25,21 @@ func TestLoadOpenAPISpecs(t *testing.T) { 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) { diff --git a/runtime/mued_adapter.go b/runtime/mued_adapter.go index e380710..8d3cd5a 100644 --- a/runtime/mued_adapter.go +++ b/runtime/mued_adapter.go @@ -37,6 +37,13 @@ type MuEdAdapter interface { // 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 diff --git a/runtime/mued_adapter_test.go b/runtime/mued_adapter_test.go index 2fb0cd5..9c25f49 100644 --- a/runtime/mued_adapter_test.go +++ b/runtime/mued_adapter_test.go @@ -32,6 +32,7 @@ func (a fakeAdapter) EncodeChat(map[string]any) (map[string]any, error) { 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() @@ -77,18 +78,28 @@ func TestMuEdRegistry_ReregisterKeepsPosition(t *testing.T) { assert.Equal(t, []string{"0.1.0", "0.2.0"}, reg.Versions()) } -func TestDefaultMuEdRegistry_HasV010(t *testing.T) { +func TestDefaultMuEdRegistry_RegisteredVersions(t *testing.T) { reg := runtime.DefaultMuEdRegistry() - assert.Equal(t, []string{"0.1.0"}, reg.Versions()) - assert.Equal(t, []string{"0.1.0"}, runtime.SupportedMuEdVersions()) + 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", runtime.MuEdResolveVersion("")) + 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.0", runtime.MuEdResolveVersion("99.0.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 index 0d76b39..f311f50 100644 --- a/runtime/mued_v0_1_0.go +++ b/runtime/mued_v0_1_0.go @@ -59,3 +59,7 @@ func (muEdV010) EncodeChat(result map[string]any) (map[string]any, error) { 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/schema/mued_v0.1.0.yml b/runtime/schema/mued_v0.1.0.yml index e617092..c5ca3c8 100644 --- a/runtime/schema/mued_v0.1.0.yml +++ b/runtime/schema/mued_v0.1.0.yml @@ -335,23 +335,7 @@ paths: version: 1 responses: '200': - description: | - Successfully generated feedback. - - ### Streaming variant (opt-in) - - If the request sends `Accept: text/event-stream` (standalone/serve - mode only; ignored under AWS Lambda), the response is a Server-Sent - Events progress stream instead of a single JSON body. The HTTP - status stays `200` for the whole stream, including failures. It is: - zero or more progress frames whose SSE `event:` is the stage name - and `data:` is an `SseProgressStep`; optional `:`-prefixed - keep-alive comment lines; and exactly one terminal frame — - `event: completed` with `data:` an `SseEvaluateTerminalFrame` (the - `200` body under `feedback`, plus a `steps` replay), or - `event: failed` with `data:` an `SseEvaluateTerminalFrame` whose - `error` is an `ErrorResponse`. `X-Request-Id` / `X-Api-Version` are - sent once as response headers when the stream opens. + description: Successfully generated feedback. headers: X-Request-Id: description: Request id for tracing this request across services. @@ -391,9 +375,6 @@ paths: - feedbackId: fb-2 title: Overall structure message: The overall structure of your answer is clear and easy to follow. - text/event-stream: - schema: - $ref: '#/components/schemas/SseEvaluateTerminalFrame' '202': $ref: '#/components/responses/202-Accepted' '400': @@ -602,23 +583,7 @@ paths: temperature: 0.5 responses: '200': - description: | - Successful chat response. - - ### Streaming variant (opt-in) - - If the request sends `Accept: text/event-stream` (standalone/serve - mode only; ignored under AWS Lambda), the response is a Server-Sent - Events progress stream instead of a single JSON body. The HTTP - status stays `200` for the whole stream, including failures. It is: - zero or more progress frames whose SSE `event:` is the stage name - and `data:` is an `SseProgressStep`; optional `:`-prefixed - keep-alive comment lines; and exactly one terminal frame — - `event: completed` with `data:` an `SseChatTerminalFrame` (the - `200` body as `output` / `metadata`, plus a `steps` replay), or - `event: failed` with `data:` an `SseChatTerminalFrame` whose - `error` is an `ErrorResponse`. `X-Request-Id` / `X-Api-Version` are - sent once as response headers when the stream opens. + description: Successful chat response. headers: X-Request-Id: description: Request id for tracing this request across services. @@ -673,9 +638,6 @@ paths: model: gpt-5.2 temperature: 0.5 outputTokens: 143 - text/event-stream: - schema: - $ref: '#/components/schemas/SseChatTerminalFrame' '400': $ref: '#/components/responses/400-BadRequest-2' '403': @@ -1573,21 +1535,6 @@ components: supportsSummativeFeedback: type: boolean description: Indicates whether the service supports feedback with points / grading signals. - supportsStreaming: - type: boolean - description: | - Whether /evaluate supports opt-in SSE progress streaming via - `Accept: text/event-stream`. Distinct from - `configuration.llm.stream`. - supportedProgressStages: - type: - - array - - 'null' - description: | - Optional list of `SseProgressStep.stage` values an /evaluate SSE - stream may emit. Informative; clients must tolerate unlisted values. - items: - type: string supportsDataPolicy: $ref: '#/components/schemas/DataPolicySupport' supportedArtefactProfiles: @@ -1749,19 +1696,7 @@ components: description: Indicates whether the service supports adapting to user preferences. supportsStreaming: type: boolean - description: | - Whether /chat supports opt-in SSE progress streaming via - `Accept: text/event-stream`. Distinct from - `configuration.llm.stream`. - supportedProgressStages: - type: - - array - - 'null' - description: | - Optional list of `SseProgressStep.stage` values a /chat SSE - stream may emit. Informative; clients must tolerate unlisted values. - items: - type: string + description: Indicates whether the service supports streaming responses. supportsDataPolicy: $ref: '#/components/schemas/DataPolicySupport' supportedLanguages: @@ -1807,136 +1742,6 @@ components: description: Optional version of the chat service implementation. capabilities: $ref: '#/components/schemas/ChatCapabilities' - SseProgressStep: - type: object - description: | - A single progress step emitted while an operation runs. Carried as - the SSE `data:` payload of an intermediate progress frame (the SSE - `event:` field carries the stage name), and replayed in the - terminal frame's `steps` array. - additionalProperties: true - required: - - stage - - timestamp - properties: - stage: - type: string - description: | - Lifecycle stage this step reports. Informative, not a fixed - enum: implementations may add stages and clients must tolerate - unrecognised values. Common values: "preparing", "starting", - "evaluating" (evaluate), "thinking" (chat), and the terminal - "completed" / "failed". - message: - type: - - string - - 'null' - description: Short, learner/teacher-facing description of the step. - data: - type: object - additionalProperties: true - description: | - Free-form payload attached by a worker-authored sub-step - (shimmy extension; the canonical spec relies on - additionalProperties for this). - timestamp: - type: string - format: date-time - SseTerminalSteps: - type: object - description: | - Shared fragment of the terminal SSE frame: the ordered replay of - every progress step emitted during the stream, so a client that - connected late or dropped frames still receives the full trace. - required: - - steps - properties: - steps: - type: array - description: Ordered list of every SseProgressStep emitted during the stream. - items: - $ref: '#/components/schemas/SseProgressStep' - StreamingCapabilities: - type: object - description: | - Shared capability fragment describing an operation's support for - opt-in Server-Sent Events (SSE) progress streaming. - additionalProperties: true - properties: - supportsStreaming: - type: boolean - description: | - Whether this operation supports opt-in SSE progress streaming, - selected per request with `Accept: text/event-stream`. Distinct - from `configuration.llm.stream`, which governs token-level - streaming from the LLM provider. - supportedProgressStages: - type: - - array - - 'null' - description: | - Optional list of `SseProgressStep.stage` values this service may - emit. Informative only; clients must tolerate unlisted stages. - items: - type: string - SseChatTerminalFrame: - type: object - description: | - Payload (`data`) of the single terminal SSE frame for `POST /chat` - (`event: completed | failed`). On `completed`, `output`/`metadata` - hold the endpoint's normal 200 body and `error` is absent. On - `failed`, `output` is null and `error` holds an ErrorResponse. The - HTTP status stays 200 regardless. `steps` is always present. - allOf: - - $ref: '#/components/schemas/SseTerminalSteps' - - type: object - properties: - output: - type: - - object - - 'null' - additionalProperties: true - description: The generated assistant response; null on a failed frame. - metadata: - type: - - object - - 'null' - additionalProperties: true - description: Optional metadata about response generation. - error: - type: - - object - - 'null' - description: Present only on a failed frame. - allOf: - - $ref: '#/components/schemas/ErrorResponse' - SseEvaluateTerminalFrame: - type: object - description: | - Payload (`data`) of the single terminal SSE frame for - `POST /evaluate` (`event: completed | failed`). On `completed`, - `feedback` holds the endpoint's normal 200 body and `error` is - absent. On `failed`, `feedback` is null and `error` holds an - ErrorResponse. The HTTP status stays 200 regardless. `steps` is - always present. - allOf: - - $ref: '#/components/schemas/SseTerminalSteps' - - type: object - properties: - feedback: - type: - - array - - 'null' - items: - $ref: '#/components/schemas/Feedback' - description: The generated feedback items; null on a failed frame. - error: - type: - - object - - 'null' - description: Present only on a failed frame. - allOf: - - $ref: '#/components/schemas/ErrorResponse' responses: 202-Accepted: description: Request accepted for asynchronous evaluation processing. diff --git a/runtime/schema/openapi.go b/runtime/schema/openapi.go index e7548c3..a81ec1b 100644 --- a/runtime/schema/openapi.go +++ b/runtime/schema/openapi.go @@ -51,6 +51,12 @@ func mustLoadMuEdSpecs() map[string][]byte { // 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 {