Add progress-event callbacks for evaluation requests

.gitignore

@@ -23,4 +23,7 @@ lcov.info
 go.work
 
 # Local .env files
-*.local.idea/
+*.local
+
+# IDE / editor
+.idea/

app/lambda/handler.go

@@ -21,8 +21,10 @@ type LambdaHandlerParams struct {
 	// Config is the configuration for the Lambda handler.
 	Config Config
 
-	// Handlers is a slice of HTTP handlers grouped together.
-	Handlers []*server.HttpHandler `group:"handlers"`
+	// Mux is the shared, fully-wrapped application HTTP handler chain — the same
+	// one the standalone server serves, including OpenAPI request/response
+	// validation.
+	Mux *server.Mux
 
 	// Context is the context for the Lambda handler.
 	Context context.Context
@@ -32,30 +34,24 @@ type LambdaHandlerParams struct {
 }
 
 type LambdaHandler struct {
-	config Config
-	ctx    context.Context
-	cancel context.CancelFunc
-	mux    *http.ServeMux
-	log    *zap.Logger
+	config  Config
+	ctx     context.Context
+	cancel  context.CancelFunc
+	handler http.Handler
+	log     *zap.Logger
 }
 
 // NewLambdaHandler creates a new instance of LambdaHandler
 // with the given parameters.
 func NewLambdaHandler(params LambdaHandlerParams) *LambdaHandler {
 	ctx, cancel := context.WithCancel(params.Context)
 
-	mux := http.NewServeMux()
-
-	for _, handler := range params.Handlers {
-		mux.Handle(handler.Name, handler.Handler)
-	}
-
 	return &LambdaHandler{
-		config: params.Config,
-		ctx:    ctx,
-		cancel: cancel,
-		mux:    mux,
-		log:    params.Logger,
+		config:  params.Config,
+		ctx:     ctx,
+		cancel:  cancel,
+		handler: params.Mux,
+		log:     params.Logger,
 	}
 }
 
@@ -101,11 +97,11 @@ func (s *LambdaHandler) Shutdown() {
 func (s *LambdaHandler) getProxyFunction() (any, error) {
 	switch s.config.ProxySource {
 	case ProxySourceApiGatewayV1:
-		return httpadapter.New(server.NormalizePath(s.mux)).ProxyWithContext, nil
+		return httpadapter.New(s.handler).ProxyWithContext, nil
 	case ProxySourceApiGatewayV2:
-		return httpadapter.NewV2(server.NormalizePath(s.mux)).ProxyWithContext, nil
+		return httpadapter.NewV2(s.handler).ProxyWithContext, nil
 	case ProxySourceAlb:
-		return httpadapter.NewALB(server.NormalizePath(s.mux)).ProxyWithContext, nil
+		return httpadapter.NewALB(s.handler).ProxyWithContext, nil
 	default:
 		return nil, fmt.Errorf("invalid proxy source: %s", s.config.ProxySource)
 	}

app/lambda/module.go

@@ -4,6 +4,7 @@ import (
 	"go.uber.org/fx"
 
 	"github.com/lambda-feedback/shimmy/handler"
+	"github.com/lambda-feedback/shimmy/internal/server"
 	"github.com/lambda-feedback/shimmy/util/logging"
 )
 
@@ -14,8 +15,12 @@ func Module(config Config) fx.Option {
 		fx.Supply(config),
 		// rename logger for module
 		logging.DecorateLogger("lambda"),
+		// the Lambda proxy buffers the whole response — no incremental streaming
+		fx.Supply(handler.StreamingCapability{Enabled: false}),
 		// provide handlers
 		handler.Module(),
+		// provide the shared HTTP handler chain (specs + wrapped mux)
+		server.HandlerModule(),
 		// provide server
 		fx.Provide(NewLifecycleHandler),
 		// invoke server

app/lambda/module_test.go

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

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

cmd/root.go

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

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"`
 }