From 3a939ea0b2e08c6d32e8ae959b3e356e2f742bf8 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 14:30:38 +0100 Subject: [PATCH 1/9] 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 2/9] 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 3/9] =?UTF-8?q?Add=20callbackUrl=20support=20for=20progres?= =?UTF-8?q?s=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 4/9] 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 5/9] 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 6/9] 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 7/9] 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 8/9] 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 9/9] 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{}