agentplane: separate runner RPC adapters from sandbox reconciliation

agentplane/sandbox_controller/cmd/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
         "//agentplane/observability",
         "//agentplane/policy/api/v1alpha1",
         "//agentplane/sandbox_controller/internal/controller",
+        "//agentplane/sandbox_controller/internal/runnercontrol",
         "//agentplane/sandboxes/api",
         "@com_github_go_logr_logr//:logr",
         "@in_gopkg_yaml_v3//:yaml_v3",

agentplane/sandbox_controller/cmd/main.go

@@ -25,6 +25,7 @@ import (
 	"github.com/agentydragon/ducktape/agentplane/observability"
 	policyapi "github.com/agentydragon/ducktape/agentplane/policy/api/v1alpha1"
 	sandboxcontroller "github.com/agentydragon/ducktape/agentplane/sandbox_controller/internal/controller"
+	"github.com/agentydragon/ducktape/agentplane/sandbox_controller/internal/runnercontrol"
 	sandboxapi "github.com/agentydragon/ducktape/agentplane/sandboxes/api"
 )
 
@@ -92,7 +93,7 @@ func run(c config) error {
 	if err != nil {
 		return err
 	}
-	quiescer := sandboxcontroller.GatewayQuiescer{Target: c.RunnerGatewayTarget, TokenFile: c.RunnerGatewayTokenFile, TLS: c.RunnerGatewayTLS}
+	quiescer := runnercontrol.GatewayQuiescer{Target: c.RunnerGatewayTarget, TokenFile: c.RunnerGatewayTokenFile, TLS: c.RunnerGatewayTLS}
 	clientset, err := kubernetes.NewForConfig(restConfig)
 	if err != nil {
 		return err
@@ -105,7 +106,7 @@ func run(c config) error {
 		Client: mgr.GetClient(), Objects: mgr.GetAPIReader(), KubeVirt: kubeVirt,
 		Quiescer: quiescer, Leases: mgr.GetClient(), Credentials: credentials, Now: time.Now,
 		FirecrackerStatus: sandboxcontroller.PodProxyFirecrackerStatusReader{Client: clientset},
-		Suspender:         sandboxcontroller.GatewaySuspender{Target: c.RunnerGatewayTarget, TokenFile: c.RunnerGatewayTokenFile, TLS: c.RunnerGatewayTLS},
+		Suspender:         runnercontrol.GatewaySuspender{Target: c.RunnerGatewayTarget, TokenFile: c.RunnerGatewayTokenFile, TLS: c.RunnerGatewayTLS},
 	}
 	if err := reconciler.SetupWithManager(mgr); err != nil {
 		return err

agentplane/sandbox_controller/internal/controller/BUILD.bazel

@@ -21,17 +21,15 @@ go_library(
         "workspace.go",
     ],
     importpath = "github.com/agentydragon/ducktape/agentplane/sandbox_controller/internal/controller",
-    visibility = ["//agentplane/sandbox_controller/cmd:__pkg__"],
+    visibility = [
+        "//agentplane/sandbox_controller/cmd:__pkg__",
+        "//agentplane/sandbox_controller/internal/runnercontrol:__pkg__",
+    ],
     deps = [
-        "//agentplane/env/runtime/control:environment_runtime_go_proto",
         "//agentplane/identity/credentials/api/v1alpha1",
         "//agentplane/identity/credentials/requester",
-        "//agentplane/net/peertls",
-        "//agentplane/observability",
         "//agentplane/policy/api/v1alpha1",
-        "//agentplane/proto:session_go_proto",
         "//agentplane/sandboxes/api",
-        "//agentplane/sandboxes/environmentruntimeclient",
         "//agentplane/sandboxes/lease",
         "@io_k8s_api//coordination/v1:coordination",
         "@io_k8s_api//core/v1:core",
@@ -60,10 +58,8 @@ go_library(
         "@io_k8s_sigs_controller_runtime//pkg/reconcile",
         "@io_k8s_sigs_controller_runtime//pkg/scheme",
         "@io_kubevirt_api//core/v1:core",
-        "@io_opentelemetry_go_otel//attribute",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",
-        "@org_golang_google_protobuf//types/known/timestamppb",
     ],
 )
 
@@ -78,7 +74,6 @@ go_test(
         "identity_test.go",
         "idle_test.go",
         "launch_test.go",
-        "quiescer_test.go",
         "state_volume_test.go",
         "vm_test.go",
         "workspace_test.go",

agentplane/sandbox_controller/internal/controller/firecracker_suspend.go

@@ -6,13 +6,9 @@ import (
 	"math"
 	"time"
 
-	environmentruntimepb "github.com/agentydragon/ducktape/agentplane/env/runtime/control/environmentruntimepb"
-	"github.com/agentydragon/ducktape/agentplane/net/peertls"
 	sandboxapi "github.com/agentydragon/ducktape/agentplane/sandboxes/api"
-	"github.com/agentydragon/ducktape/agentplane/sandboxes/environmentruntimeclient"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
-	"google.golang.org/protobuf/types/known/timestamppb"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -25,36 +21,6 @@ type RuntimeSuspender interface {
 	Resume(context.Context, *sandboxapi.Sandbox, int64) error
 }
 
-type GatewaySuspender struct {
-	TLS       *peertls.ClientTLS
-	Target    string
-	TokenFile string
-}
-
-func (g GatewaySuspender) withClient(ctx context.Context, sandbox *sandboxapi.Sandbox, call func(context.Context, environmentruntimepb.EnvironmentRuntimeClient) error) error {
-	return environmentruntimeclient.WithClient(ctx, environmentruntimeclient.Config{
-		Target:    environmentruntimeclient.ExpandTarget(g.Target, sandbox.Name, sandbox.Namespace),
-		TokenFile: g.TokenFile, TLS: g.TLS, Timeout: quiesceTimeout,
-		SandboxName: sandbox.Name, SandboxUID: string(sandbox.UID), SandboxNamespace: sandbox.Namespace,
-	}, call)
-}
-
-func (g GatewaySuspender) Prepare(ctx context.Context, sandbox *sandboxapi.Sandbox, generation int64, idleSince time.Time) error {
-	return g.withClient(ctx, sandbox, func(ctx context.Context, client environmentruntimepb.EnvironmentRuntimeClient) error {
-		_, err := client.PrepareSuspend(ctx, &environmentruntimepb.PrepareSuspendRequest{
-			Generation: generation, IfIdleSince: timestamppb.New(idleSince),
-		})
-		return err
-	})
-}
-
-func (g GatewaySuspender) Resume(ctx context.Context, sandbox *sandboxapi.Sandbox, generation int64) error {
-	return g.withClient(ctx, sandbox, func(ctx context.Context, client environmentruntimepb.EnvironmentRuntimeClient) error {
-		_, err := client.ResumeAdmission(ctx, &environmentruntimepb.ResumeAdmissionRequest{Generation: generation})
-		return err
-	})
-}
-
 func firecrackerSuspending(sandbox *sandboxapi.Sandbox) bool {
 	attempt := sandbox.Status.FirecrackerSuspension
 	return attempt != nil && attempt.Phase != sandboxapi.FirecrackerSuspendReleased

agentplane/sandbox_controller/internal/controller/quiescer.go

@@ -2,27 +2,11 @@ package sandboxcontroller
 
 import (
 	"context"
-	"fmt"
 	"time"
 
-	"github.com/agentydragon/ducktape/agentplane/net/peertls"
-
-	"github.com/agentydragon/ducktape/agentplane/observability"
 	sandboxapi "github.com/agentydragon/ducktape/agentplane/sandboxes/api"
-	"github.com/agentydragon/ducktape/agentplane/sandboxes/environmentruntimeclient"
-	"go.opentelemetry.io/otel/attribute"
-
-	"google.golang.org/protobuf/types/known/timestamppb"
-
-	environmentruntimepb "github.com/agentydragon/ducktape/agentplane/env/runtime/control/environmentruntimepb"
-	sessionpb "github.com/agentydragon/ducktape/agentplane/proto/session"
-	"google.golang.org/grpc/status"
 )
 
-// quiesceTimeout bounds one Quiesce call; stopping harnesses can take a while, and Quiesce is idempotent, so
-// a call that runs out of time is repeated.
-const quiesceTimeout = 30 * time.Second
-
 // QuiesceResult is a runner's answer to a quiesce.
 type QuiesceResult struct {
 	// True unless the runner explicitly reports QUIESCE_STAGE_QUIESCED. An absent or unknown stage is
@@ -36,51 +20,3 @@ type QuiesceResult struct {
 type Quiescer interface {
 	Quiesce(ctx context.Context, sandbox *sandboxapi.Sandbox, ifIdleSince *time.Time) (QuiesceResult, error)
 }
-
-// GatewayQuiescer calls `Runner.Quiesce` through the Runner Gateway's proxy, at the Sandbox's own
-// control Service, which names the replica holding its runner connection.
-type GatewayQuiescer struct {
-	TLS *peertls.ClientTLS
-	// Target is a template of `{sandbox}` and `{namespace}`, typically
-	// sandboxapi.ControlServiceTargetTemplate.
-	Target string
-	// TokenFile is the projected ServiceAccount token the gateway accepts from this controller.
-	TokenFile string
-}
-
-func (g GatewayQuiescer) Quiesce(ctx context.Context, sandbox *sandboxapi.Sandbox, ifIdleSince *time.Time) (QuiesceResult, error) {
-	target := environmentruntimeclient.ExpandTarget(g.Target, sandbox.Name, sandbox.Namespace)
-	request := &sessionpb.QuiesceRequest{}
-	if ifIdleSince != nil {
-		request.IfIdleSince = timestamppb.New(*ifIdleSince)
-	}
-	var response *sessionpb.QuiesceResponse
-	called := false
-	err := environmentruntimeclient.WithClient(ctx, environmentruntimeclient.Config{
-		Target: target, TokenFile: g.TokenFile, TLS: g.TLS, Timeout: quiesceTimeout,
-		SandboxName: sandbox.Name, SandboxUID: string(sandbox.UID), SandboxNamespace: sandbox.Namespace,
-	}, func(ctx context.Context, client environmentruntimepb.EnvironmentRuntimeClient) error {
-		called = true
-		callCtx, span := observability.Start(ctx, "agentplane.sandbox_controller.runner_quiesce")
-		span.SetAttributes(
-			attribute.String("sandbox.name", sandbox.Name),
-			attribute.String("sandbox.uid", string(sandbox.UID)),
-		)
-		var callErr error
-		response, callErr = client.Quiesce(callCtx, request)
-		span.SetAttributes(attribute.Int("rpc.grpc.status_code", int(status.Code(callErr))))
-		observability.End(span, callErr)
-		return callErr
-	})
-	if err != nil {
-		if called {
-			return QuiesceResult{}, fmt.Errorf("quiescing %s through %s: %w", sandbox.Name, target, err)
-		}
-		return QuiesceResult{}, err
-	}
-	return quiesceResult(response), nil
-}
-
-func quiesceResult(response *sessionpb.QuiesceResponse) QuiesceResult {
-	return QuiesceResult{Quiescing: response.GetStage() != sessionpb.QuiesceStage_QUIESCE_STAGE_QUIESCED}
-}

agentplane/sandbox_controller/internal/runnercontrol/BUILD.bazel

@@ -0,0 +1,30 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "runnercontrol",
+    srcs = [
+        "quiescer.go",
+        "suspender.go",
+    ],
+    importpath = "github.com/agentydragon/ducktape/agentplane/sandbox_controller/internal/runnercontrol",
+    visibility = ["//agentplane/sandbox_controller/cmd:__pkg__"],
+    deps = [
+        "//agentplane/env/runtime/control:environment_runtime_go_proto",
+        "//agentplane/net/peertls",
+        "//agentplane/observability",
+        "//agentplane/proto:session_go_proto",
+        "//agentplane/sandbox_controller/internal/controller",
+        "//agentplane/sandboxes/api",
+        "//agentplane/sandboxes/environmentruntimeclient",
+        "@io_opentelemetry_go_otel//attribute",
+        "@org_golang_google_grpc//status",
+        "@org_golang_google_protobuf//types/known/timestamppb",
+    ],
+)
+
+go_test(
+    name = "runnercontrol_test",
+    srcs = ["quiescer_test.go"],
+    embed = [":runnercontrol"],
+    deps = ["//agentplane/proto:session_go_proto"],
+)

agentplane/sandbox_controller/internal/runnercontrol/quiescer.go

@@ -0,0 +1,73 @@
+package runnercontrol
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/agentydragon/ducktape/agentplane/net/peertls"
+
+	"github.com/agentydragon/ducktape/agentplane/observability"
+	sandboxapi "github.com/agentydragon/ducktape/agentplane/sandboxes/api"
+	"github.com/agentydragon/ducktape/agentplane/sandboxes/environmentruntimeclient"
+	"go.opentelemetry.io/otel/attribute"
+
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	environmentruntimepb "github.com/agentydragon/ducktape/agentplane/env/runtime/control/environmentruntimepb"
+	sessionpb "github.com/agentydragon/ducktape/agentplane/proto/session"
+	sandboxcontroller "github.com/agentydragon/ducktape/agentplane/sandbox_controller/internal/controller"
+	"google.golang.org/grpc/status"
+)
+
+// quiesceTimeout bounds one Runner Gateway call; stopping harnesses can take a while, and Quiesce is
+// idempotent, so a call that runs out of time is repeated.
+const quiesceTimeout = 30 * time.Second
+
+// GatewayQuiescer calls `Runner.Quiesce` through the Runner Gateway's proxy, at the Sandbox's own
+// control Service, which names the replica holding its runner connection.
+type GatewayQuiescer struct {
+	TLS *peertls.ClientTLS
+	// Target is a template of `{sandbox}` and `{namespace}`, typically
+	// sandboxapi.ControlServiceTargetTemplate.
+	Target string
+	// TokenFile is the projected ServiceAccount token the gateway accepts from this controller.
+	TokenFile string
+}
+
+func (g GatewayQuiescer) Quiesce(ctx context.Context, sandbox *sandboxapi.Sandbox, ifIdleSince *time.Time) (sandboxcontroller.QuiesceResult, error) {
+	target := environmentruntimeclient.ExpandTarget(g.Target, sandbox.Name, sandbox.Namespace)
+	request := &sessionpb.QuiesceRequest{}
+	if ifIdleSince != nil {
+		request.IfIdleSince = timestamppb.New(*ifIdleSince)
+	}
+	var response *sessionpb.QuiesceResponse
+	called := false
+	err := environmentruntimeclient.WithClient(ctx, environmentruntimeclient.Config{
+		Target: target, TokenFile: g.TokenFile, TLS: g.TLS, Timeout: quiesceTimeout,
+		SandboxName: sandbox.Name, SandboxUID: string(sandbox.UID), SandboxNamespace: sandbox.Namespace,
+	}, func(ctx context.Context, client environmentruntimepb.EnvironmentRuntimeClient) error {
+		called = true
+		callCtx, span := observability.Start(ctx, "agentplane.sandbox_controller.runner_quiesce")
+		span.SetAttributes(
+			attribute.String("sandbox.name", sandbox.Name),
+			attribute.String("sandbox.uid", string(sandbox.UID)),
+		)
+		var callErr error
+		response, callErr = client.Quiesce(callCtx, request)
+		span.SetAttributes(attribute.Int("rpc.grpc.status_code", int(status.Code(callErr))))
+		observability.End(span, callErr)
+		return callErr
+	})
+	if err != nil {
+		if called {
+			return sandboxcontroller.QuiesceResult{}, fmt.Errorf("quiescing %s through %s: %w", sandbox.Name, target, err)
+		}
+		return sandboxcontroller.QuiesceResult{}, err
+	}
+	return quiesceResult(response), nil
+}
+
+func quiesceResult(response *sessionpb.QuiesceResponse) sandboxcontroller.QuiesceResult {
+	return sandboxcontroller.QuiesceResult{Quiescing: response.GetStage() != sessionpb.QuiesceStage_QUIESCE_STAGE_QUIESCED}
+}

agentplane/sandbox_controller/internal/controller/quiescer_test.go → agentplane/sandbox_controller/internal/runnercontrol/quiescer_test.go

@@ -1,4 +1,4 @@
-package sandboxcontroller
+package runnercontrol
 
 import (
 	"testing"

agentplane/sandbox_controller/internal/runnercontrol/suspender.go

@@ -0,0 +1,44 @@
+package runnercontrol
+
+import (
+	"context"
+	"time"
+
+	environmentruntimepb "github.com/agentydragon/ducktape/agentplane/env/runtime/control/environmentruntimepb"
+	"github.com/agentydragon/ducktape/agentplane/net/peertls"
+	sandboxapi "github.com/agentydragon/ducktape/agentplane/sandboxes/api"
+	"github.com/agentydragon/ducktape/agentplane/sandboxes/environmentruntimeclient"
+	"google.golang.org/protobuf/types/known/timestamppb"
+)
+
+// GatewaySuspender fences Runtime admission through the Runner Gateway. Generations survive
+// controller restarts and distinguish delayed calls from a later attempt.
+type GatewaySuspender struct {
+	TLS       *peertls.ClientTLS
+	Target    string
+	TokenFile string
+}
+
+func (g GatewaySuspender) withClient(ctx context.Context, sandbox *sandboxapi.Sandbox, call func(context.Context, environmentruntimepb.EnvironmentRuntimeClient) error) error {
+	return environmentruntimeclient.WithClient(ctx, environmentruntimeclient.Config{
+		Target:    environmentruntimeclient.ExpandTarget(g.Target, sandbox.Name, sandbox.Namespace),
+		TokenFile: g.TokenFile, TLS: g.TLS, Timeout: quiesceTimeout,
+		SandboxName: sandbox.Name, SandboxUID: string(sandbox.UID), SandboxNamespace: sandbox.Namespace,
+	}, call)
+}
+
+func (g GatewaySuspender) Prepare(ctx context.Context, sandbox *sandboxapi.Sandbox, generation int64, idleSince time.Time) error {
+	return g.withClient(ctx, sandbox, func(ctx context.Context, client environmentruntimepb.EnvironmentRuntimeClient) error {
+		_, err := client.PrepareSuspend(ctx, &environmentruntimepb.PrepareSuspendRequest{
+			Generation: generation, IfIdleSince: timestamppb.New(idleSince),
+		})
+		return err
+	})
+}
+
+func (g GatewaySuspender) Resume(ctx context.Context, sandbox *sandboxapi.Sandbox, generation int64) error {
+	return g.withClient(ctx, sandbox, func(ctx context.Context, client environmentruntimepb.EnvironmentRuntimeClient) error {
+		_, err := client.ResumeAdmission(ctx, &environmentruntimepb.ResumeAdmissionRequest{Generation: generation})
+		return err
+	})
+}