feat: Configurable timeouts + warmup for long-running / large-snapshot actors

cmd/atecontroller/internal/controllers/actortemplate_controller.go

@@ -17,6 +17,8 @@ package controllers
 import (
 	"context"
 	"fmt"
+	"os"
+	"strconv"
 	"time"
 
 	"github.com/agent-substrate/substrate/internal/resources"
@@ -42,9 +44,30 @@ const (
 	// a readiness probe. Templates whose containers all declare readyz skip
 	// this wait — ResumeActor only returns once readyz reports 200, so the
 	// workload is already initialized by the time we get here.
+	//
+	// The default can be overridden with the ATE_GOLDEN_WARMUP_SECONDS env
+	// var: workloads with a long, probe-less initialization (e.g. a heavy
+	// multi-process runtime that keeps warming up well past the 20s default)
+	// otherwise get checkpointed before they are ready.
 	goldenSnapshotWarmup = 20 * time.Second
+
+	// goldenWarmupEnv overrides goldenSnapshotWarmup with a whole number of
+	// seconds when set to a valid, non-negative integer.
+	goldenWarmupEnv = "ATE_GOLDEN_WARMUP_SECONDS"
 )
 
+// defaultGoldenWarmup returns the golden-snapshot warmup delay, honoring the
+// ATE_GOLDEN_WARMUP_SECONDS env override and falling back to goldenSnapshotWarmup
+// when the var is unset or invalid.
+func defaultGoldenWarmup() time.Duration {
+	if v := os.Getenv(goldenWarmupEnv); v != "" {
+		if secs, err := strconv.Atoi(v); err == nil && secs >= 0 {
+			return time.Duration(secs) * time.Second
+		}
+	}
+	return goldenSnapshotWarmup
+}
+
 type ActorTemplateReconciler struct {
 	client.Client
 	Scheme *runtime.Scheme
@@ -199,11 +222,11 @@ func (r *ActorTemplateReconciler) SetupWithManager(mgr ctrl.Manager) error {
 func goldenSnapshotWarmupFor(at *atev1alpha1.ActorTemplate) time.Duration {
 	containers := at.Spec.Containers
 	if len(containers) == 0 {
-		return goldenSnapshotWarmup
+		return defaultGoldenWarmup()
 	}
 	for i := range containers {
 		if containers[i].Readyz == nil {
-			return goldenSnapshotWarmup
+			return defaultGoldenWarmup()
 		}
 	}
 	return 0

cmd/atenet/internal/router/resumer.go

@@ -55,7 +55,12 @@ func (r *ActorResumer) ResumeActor(ctx context.Context, atespace, actorName stri
 		// We detach the context from the first caller using a fixed background timeout.
 		// This guarantees that if Caller 1 disconnects or times out, the underlying
 		// resume operation continues running for Caller 2 and Caller 3 without failing.
-		bgCtx, bgCancel := context.WithTimeout(context.Background(), 15*time.Second)
+		//
+		// The ceiling must cover a cold restore of a large snapshot: a ~60 MiB
+		// gVisor restore can exceed the default 15s, and a shorter timeout cancels
+		// the in-flight restore, surfacing as a 504. Configurable via
+		// ATE_RESUME_TIMEOUT_SECONDS (see timeouts.go).
+		bgCtx, bgCancel := context.WithTimeout(context.Background(), resumeTimeout())
 		defer bgCancel()
 
 		backoff := wait.Backoff{

cmd/atenet/internal/router/timeouts.go

@@ -0,0 +1,63 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package router
+
+import (
+	"os"
+	"strconv"
+	"time"
+)
+
+// Resume / forwarding timeouts and their environment overrides.
+//
+// Resume-on-demand of a suspended actor performs a cold gVisor restore, which
+// for a large snapshot (tens of MiB) can take tens of seconds. The default
+// route / ext_proc / background-resume timeouts are sized for steady-state
+// traffic and can cancel an in-flight restore, surfacing as a 504. These knobs
+// let an operator running heavy actors raise the ceilings without a code change.
+//
+// Defaults are unchanged from prior behavior, so this is purely additive.
+//
+// Longer term, suspend-safe actor networking (agent-substrate/substrate#465)
+// should remove most of the need to tune these.
+const (
+	resumeTimeoutEnv  = "ATE_RESUME_TIMEOUT_SECONDS"
+	routeTimeoutEnv   = "ATE_ROUTE_TIMEOUT_SECONDS"
+	extProcTimeoutEnv = "ATE_EXTPROC_TIMEOUT_SECONDS"
+
+	defaultResumeTimeout  = 15 * time.Second
+	defaultRouteTimeout   = 10 * time.Second
+	defaultExtProcTimeout = 5 * time.Second
+)
+
+// timeoutFromEnv returns the duration from a whole-seconds env var, falling back
+// to def when the var is unset or not a positive integer.
+func timeoutFromEnv(key string, def time.Duration) time.Duration {
+	if v := os.Getenv(key); v != "" {
+		if secs, err := strconv.Atoi(v); err == nil && secs > 0 {
+			return time.Duration(secs) * time.Second
+		}
+	}
+	return def
+}
+
+// resumeTimeout bounds a detached, background resume operation.
+func resumeTimeout() time.Duration { return timeoutFromEnv(resumeTimeoutEnv, defaultResumeTimeout) }
+
+// routeTimeout bounds a forwarded upstream request (e.g. a long LLM turn).
+func routeTimeout() time.Duration { return timeoutFromEnv(routeTimeoutEnv, defaultRouteTimeout) }
+
+// extProcTimeout bounds the ext_proc round-trip, which can drive a cold restore.
+func extProcTimeout() time.Duration { return timeoutFromEnv(extProcTimeoutEnv, defaultExtProcTimeout) }

cmd/atenet/internal/router/xds.go

@@ -373,7 +373,11 @@ func (x *XdsServer) buildRoutes() *routev3.RouteConfiguration {
 								ClusterSpecifier: &routev3.RouteAction_Cluster{
 									Cluster: "dynamic_forward_proxy_cluster",
 								},
-								Timeout: durationpb.New(10 * time.Second),
+								// Long-lived agent turns (e.g. an LLM streaming a
+								// response) can far exceed the default 10s route
+								// timeout and get cut off mid-turn. Configurable via
+								// ATE_ROUTE_TIMEOUT_SECONDS (see timeouts.go).
+								Timeout: durationpb.New(routeTimeout()),
 							},
 						},
 					},
@@ -391,13 +395,19 @@ func (x *XdsServer) buildHcm(statPrefix string) *anypb.Any {
 					ClusterName: ClusterName,
 				},
 			},
-			Timeout: durationpb.New(5 * time.Second),
+			// The ext_proc round-trip can drive a cold restore-on-demand of a
+			// suspended actor (a gVisor restore of a large snapshot), which may
+			// take tens of seconds; the default 5s timeout aborts the restore.
+			// Configurable via ATE_EXTPROC_TIMEOUT_SECONDS (see timeouts.go).
+			Timeout: durationpb.New(extProcTimeout()),
 		},
 		MutationRules: &mutationrulesv3.HeaderMutationRules{
 			AllowAllRouting: &wrapperspb.BoolValue{Value: true},
 		},
-		// Explicitly configure the message timeout to avoid the 200ms default
-		MessageTimeout: durationpb.New(5 * time.Second),
+		// Explicitly configure the message timeout to avoid the 200ms default.
+		// Sized to accommodate a cold restore-on-demand (see the GrpcService
+		// timeout above) rather than just the steady-state header exchange.
+		MessageTimeout: durationpb.New(extProcTimeout()),
 		ProcessingMode: &extprocv3filter.ProcessingMode{
 			RequestHeaderMode:   extprocv3filter.ProcessingMode_SEND,
 			ResponseHeaderMode:  extprocv3filter.ProcessingMode_SKIP,

cmd/ateom-gvisor/runsc.go

@@ -70,27 +70,45 @@ func (r *runsc) cmdCreate(ctx context.Context, out io.Writer, containerName stri
 	return nil
 }
 
+// allowConnectedOnSave reports whether `runsc start` should be passed
+// -allow-connected-on-save. It defaults to true (preserving prior behavior) and
+// can be disabled with ATEOM_RUNSC_ALLOW_CONNECTED_ON_SAVE=false for runsc
+// builds that do not define the flag — some GKE-Sandbox releases reject it on
+// `runsc start` with "flag provided but not defined" (they handle connected
+// sockets via SaveRestoreNetstack instead).
+func allowConnectedOnSave() bool {
+	return os.Getenv("ATEOM_RUNSC_ALLOW_CONNECTED_ON_SAVE") != "false"
+}
+
 func (r *runsc) cmdStart(ctx context.Context, out io.Writer, containerName string) error {
 	reapLock.RLock()
 	defer reapLock.RUnlock()
 
 	slog.InfoContext(ctx, "About to run runsc start", slog.String("container", containerName))
 
-	cmd := exec.CommandContext(
-		ctx,
-		r.path,
+	args := []string{
 		"-log-format", "json",
 		"--alsologtostderr",
 		// "-debug",
 		// "-debug-log", ateompath.RunscDebugLogDir(r.actorUID, containerName)+"/",
 		// "-debug-to-user-log",
 		// "-log-packets",
 		// "-strace",
-		"-allow-connected-on-save",
+	}
+	// -allow-connected-on-save lets runsc checkpoint a sandbox that still has
+	// open connected sockets. Some runsc builds (e.g. certain GKE-Sandbox
+	// releases) do not recognize the flag and reject `runsc start` with
+	// "flag provided but not defined", so it can be disabled by setting
+	// ATEOM_RUNSC_ALLOW_CONNECTED_ON_SAVE=false.
+	if allowConnectedOnSave() {
+		args = append(args, "-allow-connected-on-save")
+	}
+	args = append(args,
 		"-root", ateompath.RunSCStateDir(r.actorUID),
 		"start",
 		containerName, // Name of the container
 	)
+	cmd := exec.CommandContext(ctx, r.path, args...)
 	cmd.Stdout = out
 	cmd.Stderr = out