Add progress-event callbacks for evaluation requests

app/lambda/module.go

@@ -14,6 +14,8 @@ func Module(config Config) fx.Option {
 		fx.Supply(config),
 		// rename logger for module
 		logging.DecorateLogger("lambda"),
+		// the Lambda proxy buffers the whole response — no incremental streaming
+		fx.Supply(handler.StreamingCapability{Enabled: false}),
 		// provide handlers
 		handler.Module(),
 		// provide server

app/lambda/module_test.go

@@ -0,0 +1,31 @@
+package lambda
+
+import (
+	"context"
+	"testing"
+
+	"go.uber.org/fx"
+	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/config"
+	"github.com/lambda-feedback/shimmy/runtime"
+)
+
+// TestModule_DependencyGraphResolves guards the fx wiring for Lambda mode
+// given the globals app.New supplies. StreamingCapability is supplied
+// here as {Enabled: false} — the Lambda proxy cannot stream.
+func TestModule_DependencyGraphResolves(t *testing.T) {
+	cfg := config.Config{}
+
+	err := fx.ValidateApp(
+		fx.NopLogger,
+		fx.Supply(fx.Annotate(context.Background(), fx.As(new(context.Context)))),
+		fx.Supply(zap.NewNop()),
+		fx.Supply(cfg),
+		runtime.Module(cfg.Runtime),
+		Module(Config{}),
+	)
+	if err != nil {
+		t.Fatalf("lambda fx graph failed validation: %v", err)
+	}
+}

app/standalone/module.go

@@ -13,6 +13,8 @@ func Module(config Config) fx.Option {
 		"serve",
 		// rename logger for module
 		logging.DecorateLogger("serve"),
+		// the standalone HTTP server can stream responses incrementally
+		fx.Supply(handler.StreamingCapability{Enabled: true}),
 		// provide handlers
 		handler.Module(),
 		// provide server

app/standalone/module_test.go

@@ -0,0 +1,33 @@
+package standalone
+
+import (
+	"context"
+	"testing"
+
+	"go.uber.org/fx"
+	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/config"
+	"github.com/lambda-feedback/shimmy/runtime"
+)
+
+// TestModule_DependencyGraphResolves guards the fx wiring: the standalone
+// module must be satisfiable given the globals app.New supplies (context,
+// logger, config.Config, runtime module). Regressions here — e.g. a
+// handler param with no provider — surface as a validation error rather
+// than a runtime panic on `shimmy serve`.
+func TestModule_DependencyGraphResolves(t *testing.T) {
+	cfg := config.Config{}
+
+	err := fx.ValidateApp(
+		fx.NopLogger,
+		fx.Supply(fx.Annotate(context.Background(), fx.As(new(context.Context)))),
+		fx.Supply(zap.NewNop()),
+		fx.Supply(cfg),
+		runtime.Module(cfg.Runtime),
+		Module(Config{}),
+	)
+	if err != nil {
+		t.Fatalf("standalone fx graph failed validation: %v", err)
+	}
+}

cmd/root.go

@@ -43,6 +43,76 @@ functions on arbitrary, serverless platforms.`
 				Category: "auth",
 				EnvVars:  []string{"AUTH_KEY"},
 			},
+			// progress flags
+			&cli.DurationFlag{
+				Name:     "progress-callback-timeout",
+				Usage:    "the timeout for a single progress callback delivery.",
+				Value:    time.Second,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_CALLBACK_TIMEOUT"},
+			},
+			&cli.StringSliceFlag{
+				Name:     "progress-allowed-hosts",
+				Usage:    "restrict progress callback URLs to these hosts. Entries may be an exact hostname or a \"*.example.com\" wildcard. Unset allows any host, subject to the private-network guard below.",
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_ALLOWED_HOSTS"},
+			},
+			&cli.BoolFlag{
+				Name:     "progress-allow-private-networks",
+				Usage:    "allow progress callback delivery to loopback, link-local, and private IP addresses. Leave disabled unless the callback target is known to live on a trusted private network.",
+				Value:    false,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_ALLOW_PRIVATE_NETWORKS"},
+			},
+			&cli.Int64Flag{
+				Name:     "progress-sidecar-max-body-bytes",
+				Usage:    "the maximum size, in bytes, of a single worker-authored progress event POST.",
+				Value:    16 * 1024,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_SIDECAR_MAX_BODY_BYTES"},
+			},
+			&cli.IntFlag{
+				Name:     "progress-sidecar-max-events",
+				Usage:    "the maximum number of worker-authored progress events relayed per evaluation.",
+				Value:    50,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_SIDECAR_MAX_EVENTS"},
+			},
+			&cli.IntFlag{
+				Name:     "progress-sidecar-burst-size",
+				Usage:    "how many worker-authored progress events at the start of an evaluation are exempt from the minimum spacing below, so a handful of legitimate back-to-back checkpoints aren't rate limited.",
+				Value:    5,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_SIDECAR_BURST_SIZE"},
+			},
+			&cli.DurationFlag{
+				Name:     "progress-sidecar-min-event-interval",
+				Usage:    "the minimum spacing between worker-authored progress events relayed per evaluation, once the burst allowance above is used up.",
+				Value:    10 * time.Millisecond,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_SIDECAR_MIN_EVENT_INTERVAL"},
+			},
+			&cli.DurationFlag{
+				Name:     "progress-sidecar-unbind-grace-period",
+				Usage:    "how long to keep relaying worker-authored progress events after a request returns, so a fire-and-forget POST dispatched just before the result can still land.",
+				Value:    250 * time.Millisecond,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD"},
+			},
+			&cli.BoolFlag{
+				Name:     "progress-stream-enabled",
+				Usage:    "stream progress back on the /evaluate and /chat responses as Server-Sent Events for requests that send 'Accept: text/event-stream'. Standalone/serve mode only; ignored under AWS Lambda.",
+				Value:    true,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_STREAM_ENABLED"},
+			},
+			&cli.IntFlag{
+				Name:     "progress-stream-heartbeat-seconds",
+				Usage:    "seconds between SSE heartbeat comments sent while an evaluation runs, so an idle streamed connection isn't dropped by an intermediary. 0 disables heartbeats.",
+				Value:    15,
+				Category: "progress",
+				EnvVars:  []string{"PROGRESS_STREAM_HEARTBEAT_SECONDS"},
+			},
 			// shim flags
 			&cli.StringFlag{
 				Name:     "interface",
@@ -324,21 +394,31 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) {
 
 	// map cli flags to config fields
 	cliMap := map[string]string{
-		"auth-key":                   "auth.key",
-		"max-workers":                "runtime.max_workers",
-		"command":                    "runtime.cmd",
-		"cwd":                        "runtime.cwd",
-		"arg":                        "runtime.arg",
-		"env":                        "runtime.env",
-		"interface":                  "runtime.io.interface",
-		"rpc-transport":              "runtime.io.rpc.transport",
-		"rpc-transport-ipc-endpoint": "runtime.io.rpc.ipc.endpoint",
-		"rpc-transport-http-url":     "runtime.io.rpc.http.url",
-		"rpc-transport-ws-url":       "runtime.io.rpc.ws.url",
-		"rpc-transport-tcp-address":  "runtime.io.rpc.tcp.address",
-		"worker-send-timeout":        "runtime.send.timeout",
-		"worker-stop-timeout":        "runtime.stop.timeout",
-		"worker-start-timeout":       "start_timeout",
+		"auth-key":                             "auth.key",
+		"progress-callback-timeout":            "progress.callback_timeout",
+		"progress-allowed-hosts":               "progress.allowed_hosts",
+		"progress-allow-private-networks":      "progress.allow_private_networks",
+		"progress-sidecar-max-body-bytes":      "progress.sidecar.max_body_bytes",
+		"progress-sidecar-max-events":          "progress.sidecar.max_events_per_span",
+		"progress-sidecar-burst-size":          "progress.sidecar.burst_size",
+		"progress-sidecar-min-event-interval":  "progress.sidecar.min_event_interval",
+		"progress-sidecar-unbind-grace-period": "progress.sidecar.unbind_grace_period",
+		"progress-stream-enabled":              "progress.stream.enabled",
+		"progress-stream-heartbeat-seconds":    "progress.stream.heartbeat_seconds",
+		"max-workers":                          "runtime.max_workers",
+		"command":                              "runtime.cmd",
+		"cwd":                                  "runtime.cwd",
+		"arg":                                  "runtime.arg",
+		"env":                                  "runtime.env",
+		"interface":                            "runtime.io.interface",
+		"rpc-transport":                        "runtime.io.rpc.transport",
+		"rpc-transport-ipc-endpoint":           "runtime.io.rpc.ipc.endpoint",
+		"rpc-transport-http-url":               "runtime.io.rpc.http.url",
+		"rpc-transport-ws-url":                 "runtime.io.rpc.ws.url",
+		"rpc-transport-tcp-address":            "runtime.io.rpc.tcp.address",
+		"worker-send-timeout":                  "runtime.send.timeout",
+		"worker-stop-timeout":                  "runtime.stop.timeout",
+		"worker-start-timeout":                 "start_timeout",
 		// sandbox
 		"sandbox":                 "runtime.sandbox.enabled",
 		"sandbox-nsjail-path":     "runtime.sandbox.nsjail_path",

config/config.go

@@ -3,6 +3,7 @@ package config
 import (
 	"time"
 
+	"github.com/lambda-feedback/shimmy/internal/progress"
 	"github.com/lambda-feedback/shimmy/runtime"
 )
 
@@ -30,6 +31,9 @@ type Config struct {
 	// Auth is the authentication configuration
 	Auth AuthConfig `conf:"auth"`
 
+	// Progress is the configuration for outbound progress-callback delivery
+	Progress progress.Config `conf:"progress"`
+
 	// StartTimeout is the duration to wait for the application to start.
 	StartTimeout time.Duration `conf:"start_timeout"`
 }

handler/chat.go

@@ -1,16 +1,24 @@
 package handler
 
 import (
+	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
 
+	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/internal/progress"
 	"github.com/lambda-feedback/shimmy/internal/server"
 	"github.com/lambda-feedback/shimmy/runtime"
 )
 
 // ServeChat handles POST /chat.
 func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) {
+	requestID := resolveRequestID(r)
+	w.Header().Set(muEdRequestIDHeader, requestID)
+
 	if !h.checkAuth(w, r) {
 		return
 	}
@@ -43,30 +51,129 @@ func (h *MuEdHandler) ServeChat(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	resp, err := h.runtime.Chat(r.Context(), runtime.ChatRequest{Data: reqData})
-	if err != nil {
-		h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", "chat failed", nil)
-		return
+	var callbackURL string
+	if chatReq.CallbackUrl != nil {
+		callbackURL = *chatReq.CallbackUrl
 	}
 
-	resultMap, ok := resp.Data["result"].(map[string]any)
-	if !ok {
-		h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", "invalid response from chat function", nil)
+	streaming := h.streamingCapable && h.config.Progress.Stream.Enabled && acceptsEventStream(r)
+	if streaming {
+		if _, ok := w.(http.Flusher); !ok {
+			h.log.Warn("response writer is not a flusher; serving buffered response")
+			streaming = false
+		}
+	}
+
+	ctx := r.Context()
+
+	if streaming {
+		h.serveChatStream(ctx, w, reqData, version, callbackURL, requestID)
 		return
 	}
 
-	chatResp, err := runtime.MuEdToChatResponse(resultMap)
-	if err != nil {
-		h.writeMuEdError(w, version, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error", err.Error(), nil)
+	if callbackURL != "" {
+		reporter, rerr := h.progressFactory.NewReporter(callbackURL, requestID)
+		if rerr != nil {
+			h.log.Warn("invalid callbackUrl, disabling progress reporting", zap.Error(rerr))
+		} else if reporter != nil {
+			ctx = progress.ContextWithReporter(ctx, reporter)
+		}
+	}
+
+	resp, err := h.runtime.Chat(ctx, runtime.ChatRequest{Data: reqData})
+	output, metadata, termErr := h.produceChatOutput(resp, err)
+	if termErr != nil {
+		progress.Emit(ctx, progress.Event{
+			Stage:   progress.StageFailed,
+			Command: string(runtime.CommandChat),
+			Message: termErr.userMessage,
+			Error:   termErr.rawError,
+		})
+		h.writeMuEdError(w, version, termErr.status, termErr.muEdCode, termErr.muEdTitle, termErr.muEdMessage, nil)
 		return
 	}
 
+	chatResp := map[string]any{"output": output}
+	if metadata != nil {
+		chatResp["metadata"] = metadata
+	}
+
+	progress.Emit(ctx, progress.Event{
+		Stage:   progress.StageCompleted,
+		Command: string(runtime.CommandChat),
+		Message: "Response is ready.",
+		Data:    map[string]any{"output": output, "metadata": metadata},
+	})
+
 	w.Header().Set("Content-Type", "application/json")
 	w.Header().Set(muEdVersionHeader, version)
 	w.WriteHeader(http.StatusOK)
 	json.NewEncoder(w).Encode(chatResp) //nolint:errcheck
 }
 
+// serveChatStream handles a POST /chat request that opted in to SSE
+// streaming. The streaming scaffold lives in streamProgress; this only
+// supplies the run step.
+func (h *MuEdHandler) serveChatStream(
+	ctx context.Context,
+	w http.ResponseWriter,
+	reqData map[string]any,
+	version string,
+	callbackURL string,
+	requestID string,
+) {
+	h.streamProgress(ctx, w, "chat", string(runtime.CommandChat), "Response is ready.", version, callbackURL, requestID,
+		func(ctx context.Context) (map[string]any, *terminalError) {
+			resp, err := h.runtime.Chat(ctx, runtime.ChatRequest{Data: reqData})
+			output, metadata, termErr := h.produceChatOutput(resp, err)
+			if termErr != nil {
+				return nil, termErr
+			}
+			data := map[string]any{"output": output}
+			if metadata != nil {
+				data["metadata"] = metadata
+			}
+			return data, nil
+		})
+}
+
+// produceChatOutput turns a runtime chat response into the µEd output
+// object (+ optional metadata), or a terminalError describing why it
+// couldn't. It is pure: no writes, no progress events. Unlike
+// produceFeedback there is no worker-non-200 passthrough — runtime.Chat
+// returns (response, error), not an HTTP status — so every failure is a
+// 500-class terminalError.
+func (h *MuEdHandler) produceChatOutput(resp runtime.ChatResponse, chatErr error) (output, metadata map[string]any, _ *terminalError) {
+	newErr := func(muEdMessage, rawError string) *terminalError {
+		return &terminalError{
+			status:      http.StatusInternalServerError,
+			muEdCode:    "INTERNAL_ERROR",
+			muEdTitle:   "Internal server error",
+			muEdMessage: muEdMessage,
+			userMessage: "We couldn't generate a response. Please try again.",
+			rawError:    rawError,
+		}
+	}
+
+	if chatErr != nil {
+		return nil, nil, newErr("chat failed", chatErr.Error())
+	}
+
+	resultMap, ok := resp.Data["result"].(map[string]any)
+	if !ok {
+		return nil, nil, newErr("invalid response from chat function", "invalid response from chat function")
+	}
+
+	chatResp, err := runtime.MuEdToChatResponse(resultMap)
+	if err != nil {
+		return nil, nil, newErr(err.Error(), fmt.Sprintf("invalid chat response: %v", err))
+	}
+
+	output, _ = chatResp["output"].(map[string]any)
+	metadata, _ = chatResp["metadata"].(map[string]any)
+	return output, metadata, nil
+}
+
 // ServeChatHealth handles GET /chat/health.
 func (h *MuEdHandler) ServeChatHealth(w http.ResponseWriter, r *http.Request) {
 	if !h.checkAuth(w, r) {