From fbb6f02f04f385233d36fc293037a5604825d550 Mon Sep 17 00:00:00 2001 From: bjo4 <9113867+bjo4@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:19:55 +0800 Subject: [PATCH] feat(runtime): pause on Kubernetes, unbounded checkpoints, cluster CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the three remaining Kubernetes differences. The fourth item on the list, hosted tenant provisioning, is a product decision rather than a gap and is not touched here. Pause. I had called this unsolvable because Kubernetes has no cgroup freezer. It does not, but signalling is the equivalent from the Agent's point of view: SIGSTOP to every process in the container's PID namespace, SIGCONT to resume, over the exec subresource. It walks /proc rather than using kill -1 so the helper shell does not stop itself before finishing. The cost is a dependency on the Agent image having a shell and /proc, which a distroless image does not, and a failed pause is reported rather than leaving the Session in `pausing`. Checkpoints. They rode in a per-Session Secret, capped at Kubernetes' 1 MB object limit. An init container now fetches from a new agent-scoped endpoint and re-checks the SHA-256, so a truncated transfer fails the Pod instead of handing the Agent a partial file. The endpoint exposes nothing new — the same bytes are mounted into the workspace either way — and serves only the checkpoint the calling Session was created from. CI. The cluster-backed suites now run on the default branch and on demand rather than every push, since building a cluster costs minutes. The egress policy job builds its own Calico cluster instead of using kind-action, whose kindnet accepts NetworkPolicy and ignores it. The pause test watches a counter inside the container stop and start again, so a pause that reported success without freezing anything fails it. The checkpoint change is covered by raising the e2e checkpoint to 1.5 MB, above what a Secret could ever have carried, and asserting the restored size — the old path could not have passed that. Writing the pause test I first shipped helpers that returned a constant, so the comparison was always equal and the test would have passed vacuously. exec now captures stdout and the test compares real counts. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci.yml | 60 +++++++++ README.md | 6 +- api/openapi.yaml | 21 ++++ deploy/helm/README.md | 10 +- .../templates/runtime-worker.yaml | 5 + docs/agent-protocol.md | 6 + docs/examples/kubernetes-rbac.yaml | 5 + docs/kubernetes.md | 47 +++++-- examples/research-report-agent/main.py | 9 +- internal/httpapi/server.go | 40 ++++++ internal/worker/kubernetes.go | 116 +++++++++++------- internal/worker/kubernetes_attach.go | 79 ++++++++++++ .../worker/kubernetes_integration_test.go | 98 +++++++++++++-- internal/worker/kubernetes_test.go | 15 +-- test/e2e/runtime_policy_flow.py | 26 +++- 15 files changed, 460 insertions(+), 83 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f09e95d..3df1759 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,6 +3,9 @@ name: ci on: push: pull_request: + # The cluster-backed jobs are gated on this so they can be run against a + # branch without paying for a cluster on every push. + workflow_dispatch: permissions: contents: read @@ -172,3 +175,60 @@ jobs: - name: Tear down if: always() run: docker compose --profile sample --profile e2e --profile e2e-images --profile ha down -v + + # Kubernetes coverage needs a cluster, so it is separated from the fast jobs. + # It runs on the default branch and on demand rather than on every push, + # because building a cluster and loading images costs several minutes. + kubernetes: + runs-on: ubuntu-latest + if: github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' + needs: [go, helm] + timeout-minutes: 45 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version-file: go.mod + cache: true + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - uses: helm/kind-action@v1 + with: + cluster_name: agent-platform + - uses: azure/setup-helm@v4 + # Deploys the platform with the chart and drives the suite through it, so + # Agent Sessions run as Pods and reach the Control Plane as a Service. + - name: End-to-end flows on Kubernetes + run: bash test/e2e/kubernetes/run.sh + - name: Cluster state on failure + if: failure() + run: | + kubectl -n agent-platform get pods -o wide || true + kubectl -n agent-platform logs -l app.kubernetes.io/name=agent-platform-control-plane --tail 200 || true + kubectl -n agent-platform logs -l app.kubernetes.io/name=agent-platform-runtime-worker --tail 200 || true + kubectl -n agent-platform-runtime get pods -o wide || true + + # The Agent egress policy is inert on a CNI that ignores NetworkPolicy, so + # this builds its own cluster with Calico rather than reusing the kind default. + egress-policy: + runs-on: ubuntu-latest + if: github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' + needs: [go] + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version-file: go.mod + cache: true + # kind-action is not used here: it would create a cluster with kindnet, + # which accepts NetworkPolicy and ignores it. policy-check.sh builds its + # own cluster with Calico, so only the binary is needed. + - name: Install kind + run: | + curl -fsSLo /usr/local/bin/kind \ + https://kind.sigs.k8s.io/dl/v0.27.0/kind-linux-amd64 + chmod +x /usr/local/bin/kind + - name: Verify the egress policy is enforced + run: bash test/e2e/kubernetes/policy-check.sh diff --git a/README.md b/README.md index f08c4a9..31d4368 100644 --- a/README.md +++ b/README.md @@ -195,7 +195,7 @@ launched by the runtime worker, never by the Control Plane. OIDC single sign-on is implemented; SAML, SCIM, and per-organization identity providers are not. Runtime HTTP/HTTPS egress enforcement and checkpoint-based child Session resume are implemented. Several Control Plane instances can run over one PostgreSQL, NATS, and object store; see [High availability](docs/high-availability.md). Agent Sessions run as -Docker containers or Kubernetes Pods, including stdio MCP servers on both; see -[Kubernetes](docs/kubernetes.md) for the remaining backend differences. A Helm -chart deploys the platform itself onto a cluster. Hosted tenant provisioning remains a future +Docker containers or Kubernetes Pods, with the same controls on both; see +[Kubernetes](docs/kubernetes.md) for how pause and checkpoint restore differ +underneath. A Helm chart deploys the platform itself onto a cluster. Hosted tenant provisioning remains a future milestone. diff --git a/api/openapi.yaml b/api/openapi.yaml index cbfe489..c24a2b5 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -372,6 +372,27 @@ paths: responses: "201": { description: Artifact stored } "413": { description: Artifact exceeds the configured size limit } + /agent/v1/sessions/{sessionId}/checkpoint: + parameters: + - $ref: "#/components/parameters/SessionId" + get: + tags: [agent] + security: + - agentToken: [] + summary: Download the checkpoint this Session was created from + description: >- + Lets a runtime stream a restored checkpoint into the workspace instead of + carrying the bytes itself, which the Kubernetes runtime needs because a + Secret caps at 1 MB. Serves only the one checkpoint named by this + Session. The SHA-256 is returned in X-Checkpoint-Sha256 so the runtime + can verify the transfer. + responses: + "200": + description: Checkpoint content + headers: + X-Checkpoint-Sha256: + schema: { type: string } + "404": { description: "This Session was not created from a checkpoint" } /agent/v1/sessions/{sessionId}/checkpoints: parameters: - $ref: "#/components/parameters/SessionId" diff --git a/deploy/helm/README.md b/deploy/helm/README.md index b508d94..36650a7 100644 --- a/deploy/helm/README.md +++ b/deploy/helm/README.md @@ -54,8 +54,11 @@ updates stop crossing instances. See [high availability](../../docs/high-availab ## What the worker gets A ServiceAccount with a namespace-scoped Role over Pods, Secrets, -`pods/attach`, and NetworkPolicies in `runtimeWorker.runtimeNamespace`. Nothing -cluster-scoped. +`pods/attach`, `pods/exec`, and NetworkPolicies in +`runtimeWorker.runtimeNamespace`. Nothing cluster-scoped. + +`pods/attach` drives stdio MCP servers; `pods/exec` is how pause and resume +signal the Agent's processes, since Kubernetes has no cgroup freezer. The Control Plane Pod label `app.kubernetes.io/name: agent-platform-control-plane` is not cosmetic: the egress NetworkPolicy the worker writes selects Control Plane @@ -82,3 +85,6 @@ Three flows stay Compose-only because they drive Docker itself — `durable_command_flow` detaches a container network and `recovery_flow` restarts Compose services — and `oidc_flow` needs an identity provider configured at install time. + +CI runs this on the default branch and on demand rather than on every push, +since building a cluster and loading images costs several minutes. diff --git a/deploy/helm/agent-platform/templates/runtime-worker.yaml b/deploy/helm/agent-platform/templates/runtime-worker.yaml index e3f8c75..f8dbea4 100644 --- a/deploy/helm/agent-platform/templates/runtime-worker.yaml +++ b/deploy/helm/agent-platform/templates/runtime-worker.yaml @@ -34,6 +34,11 @@ rules: - apiGroups: [""] resources: ["pods/attach"] verbs: ["create", "get"] + # Pause and resume signal the Agent's processes, since Kubernetes has no + # equivalent of the cgroup freezer. + - apiGroups: [""] + resources: ["pods/exec"] + verbs: ["create", "get"] # The Agent egress policy the worker writes on start-up. - apiGroups: ["networking.k8s.io"] resources: ["networkpolicies"] diff --git a/docs/agent-protocol.md b/docs/agent-protocol.md index 4dc350b..571c4df 100644 --- a/docs/agent-protocol.md +++ b/docs/agent-protocol.md @@ -131,6 +131,12 @@ variables: Python Agents can call `session.restored_checkpoint()` to read the bytes. It returns `None` for a Session that was not started from a checkpoint. +How the bytes arrive depends on the runtime. Docker streams them through an +init container; Kubernetes fetches them from +`GET /agent/v1/sessions/{sessionId}/checkpoint` with the Session's Agent token +and verifies the digest before the Agent starts. Either way the Agent sees the +same file at the same path. + ## Runtime network egress The worker configures standard `HTTP_PROXY`, `HTTPS_PROXY`, `http_proxy`, and diff --git a/docs/examples/kubernetes-rbac.yaml b/docs/examples/kubernetes-rbac.yaml index fc04c57..587f1c0 100644 --- a/docs/examples/kubernetes-rbac.yaml +++ b/docs/examples/kubernetes-rbac.yaml @@ -27,6 +27,11 @@ rules: - apiGroups: [""] resources: ["pods/attach"] verbs: ["create", "get"] + # Pause and resume signal the Agent's processes, since Kubernetes has no + # equivalent of the cgroup freezer. + - apiGroups: [""] + resources: ["pods/exec"] + verbs: ["create", "get"] # The egress policy the worker writes on start-up. Effective only on a CNI # that enforces NetworkPolicy; see docs/kubernetes.md. - apiGroups: ["networking.k8s.io"] diff --git a/docs/kubernetes.md b/docs/kubernetes.md index 9666f3e..0227cf2 100644 --- a/docs/kubernetes.md +++ b/docs/kubernetes.md @@ -16,6 +16,9 @@ Select it with `RUNTIME_PROVIDER=kubernetes` on the worker. | `KUBERNETES_CA_CERT` | PEM CA bundle when running outside the cluster | | `KUBERNETES_INSECURE_SKIP_VERIFY` | Local testing only; makes the connection interceptable | +The worker needs `create` on `pods`, `secrets`, `pods/attach`, `pods/exec`, and +`networkpolicies` in its runtime namespace. Nothing cluster-scoped. + In-cluster the worker reads its credentials from the projected service account and needs no configuration beyond `RUNTIME_PROVIDER` and the namespace. @@ -81,19 +84,33 @@ MCP servers, so there is no declared user to impose. This needs `create` on `pods/attach` in the runtime namespace. -## What Kubernetes cannot do +## Pause and resume + +Kubernetes has no equivalent of Docker's cgroup freezer, so a pause is signalled +instead: `SIGSTOP` to every process in the Agent container's PID namespace, +`SIGCONT` to resume. From the Agent's point of view that is what a freeze +amounts to. + +The signalling walks `/proc` rather than using `kill -1`, so the helper shell +does not stop itself before it has finished with the rest. It needs `create` on +`pods/exec`, and it needs the Agent image to have a shell and `/proc` — a +distroless image cannot be paused this way. A pause that fails is reported +rather than leaving the Session stuck in `pausing`. -The worker asks the runtime what it supports rather than assuming, so these are -refused visibly instead of hanging. +## Checkpoints -**Pause and resume.** There is no API to freeze a running Pod; Docker uses the -cgroup freezer. A pause request against a Kubernetes Session returns it to -`running` with a message saying so, rather than leaving it stuck in `pausing`. +A restored checkpoint is fetched by an init container from +`GET /agent/v1/sessions/{id}/checkpoint`, authenticated with the Session's own +Agent token, and its SHA-256 is re-checked before the Agent starts so a +truncated transfer fails the Pod rather than handing over a partial file. -**Checkpoints above 900 KB.** They travel in a per-Session Secret, which -Kubernetes caps at 1 MB. Docker streams them through an init container's stdin -and has no such limit. An oversized checkpoint is refused at Start rather than -failing partway through. +Carrying the bytes in a per-Session Secret would cap them at Kubernetes' 1 MB +object limit. The endpoint exposes nothing new: the same bytes are mounted into +the Agent's workspace either way, and it serves only the one checkpoint the +Session was created from. + +The init container image is `busybox:1.37` by default and needs a shell, `wget`, +and `sha256sum`. ## RBAC @@ -104,6 +121,9 @@ rules: - apiGroups: [""] resources: ["pods", "secrets"] verbs: ["create", "get", "list", "delete"] + - apiGroups: [""] + resources: ["pods/attach", "pods/exec"] + verbs: ["create", "get"] - apiGroups: ["networking.k8s.io"] resources: ["networkpolicies"] verbs: ["create", "get", "list"] @@ -128,9 +148,10 @@ go test ./internal/worker -run Integration -v Those cover starting a Session to completion, reporting a non-zero exit, distinguishing a missing workload from a failed one, mounting Session secrets -read-only, restoring a checkpoint into the workspace, refusing an oversized one, -stopping a running Session, writing the egress policy idempotently, and running -a stdio MCP tool through attach. The secret and checkpoint cases assert from +read-only, stopping a running Session, writing the egress policy idempotently, +running a stdio MCP tool through attach, and pausing and resuming an Agent — +that last one by watching a counter inside the container stop and start again, +so a pause that reported success without freezing anything would fail. The secret and checkpoint cases assert from inside the container, so a mount that silently did not arrive fails the test. The stdio case uses the same fixture server as the Docker end-to-end suite, diff --git a/examples/research-report-agent/main.py b/examples/research-report-agent/main.py index 0c8a1de..739c82f 100644 --- a/examples/research-report-agent/main.py +++ b/examples/research-report-agent/main.py @@ -56,14 +56,21 @@ async def main() -> None: if isinstance(checkpoint_test, dict): restored = session.restored_checkpoint() if restored is not None: + # Report the size rather than megabytes of padding. + head = restored.decode(errors="replace").split("\n", 1)[0] await session.message( - "Restored checkpoint: " + restored.decode(errors="replace") + f"Restored checkpoint: {head} ({len(restored)} bytes)" ) await session.complete( {"checkpointRestored": True, "checkpointId": session.checkpoint_id} ) return content = str(checkpoint_test.get("content", "checkpoint state")) + # Padding lets a test request a checkpoint larger than a Kubernetes + # Secret can carry, which is what the fetch path has to handle. + padding = int(checkpoint_test.get("sizeBytes", 0)) + if padding > len(content): + content += "\n" + "x" * (padding - len(content) - 1) checkpoint_path = Path("/workspace/validation.checkpoint") checkpoint_path.write_text(content, encoding="utf-8") checkpoint = await session.upload_checkpoint( diff --git a/internal/httpapi/server.go b/internal/httpapi/server.go index c7b7aa7..c0d433f 100644 --- a/internal/httpapi/server.go +++ b/internal/httpapi/server.go @@ -135,6 +135,7 @@ func (s *Server) routes() { s.mux.HandleFunc("GET /agent/v1/sessions/{sessionID}/connect", s.agentConnect) s.mux.HandleFunc("POST /agent/v1/sessions/{sessionID}/artifacts", s.agentArtifactUpload) s.mux.HandleFunc("POST /agent/v1/sessions/{sessionID}/checkpoints", s.agentCheckpointUpload) + s.mux.HandleFunc("GET /agent/v1/sessions/{sessionID}/checkpoint", s.agentCheckpointDownload) s.mux.HandleFunc("POST /agent/v1/sessions/{sessionID}/models/{binding}/chat/completions", s.agentModelCompletion) s.mux.HandleFunc("POST /agent/v1/sessions/{sessionID}/mcp/{serverID}/tools/{toolName}", s.agentMCPToolCall) @@ -1032,6 +1033,45 @@ func (s *Server) downloadArtifact(w http.ResponseWriter, r *http.Request) { _, _ = io.Copy(w, content) } +// agentCheckpointDownload serves the checkpoint a Session was created from, +// authenticated by that Session's own Agent token. +// +// It exists so a runtime can stream a restored checkpoint into the workspace +// instead of carrying the bytes itself. The Kubernetes runtime used a per- +// Session Secret, which caps at 1 MB; fetching removes that ceiling. +// +// This exposes nothing new to the Agent: the same bytes are mounted into its +// workspace either way. It serves only the one checkpoint named by this +// Session, so a compromised Agent cannot walk the organization's checkpoints. +func (s *Server) agentCheckpointDownload(w http.ResponseWriter, r *http.Request) { + session, ok := s.authenticateAgentRequest(w, r) + if !ok { + return + } + if session.CheckpointID == "" { + writeError(w, http.StatusNotFound, "no_checkpoint", "This Session was not created from a checkpoint.") + return + } + checkpoint, err := s.store.Checkpoint(r.Context(), session.CheckpointID) + if err != nil || checkpoint.AgentVersionID != session.AgentVersionID { + writeError(w, http.StatusNotFound, "checkpoint_not_found", "Checkpoint not found.") + return + } + content, err := s.objects.Get(r.Context(), checkpoint.ObjectKey) + if err != nil { + writeError(w, http.StatusNotFound, "checkpoint_content_missing", "Checkpoint content is unavailable.") + return + } + defer content.Close() + w.Header().Set("Content-Type", "application/octet-stream") + // The runtime verifies this before handing the bytes to the Agent. + w.Header().Set("X-Checkpoint-Sha256", checkpoint.SHA256) + w.Header().Set("Cache-Control", "private, max-age=0") + if _, err := io.Copy(w, io.LimitReader(content, s.config.MaxArtifactBytes)); err != nil { + s.log.Warn("stream checkpoint to agent", "sessionId", session.ID, "error", err) + } +} + func (s *Server) downloadCheckpoint(w http.ResponseWriter, r *http.Request) { checkpoint, err := s.store.Checkpoint(r.Context(), r.PathValue("checkpointID")) if err != nil { diff --git a/internal/worker/kubernetes.go b/internal/worker/kubernetes.go index b98a76d..1f7f3c4 100644 --- a/internal/worker/kubernetes.go +++ b/internal/worker/kubernetes.go @@ -17,10 +17,8 @@ import ( const ( agentContainerName = "agent" toolContainerName = "mcp-server" - // checkpointLimit is what a Kubernetes Secret can carry. Docker streams the - // checkpoint through an init container's stdin and has no such ceiling. - checkpointLimit = 900 << 10 - secretMountPath = "/run/secrets/agent-platform" + secretMountPath = "/run/secrets/agent-platform" + checkpointPath = "/workspace/.agent-platform/checkpoint" ) // Kubernetes runs each Agent Session as a Pod. @@ -33,6 +31,9 @@ const ( type Kubernetes struct { client *kubeClient ControlURL string + // CheckpointImage runs the init container that fetches a restored + // checkpoint. It needs a shell, wget, and sha256sum. + CheckpointImage string // PolicyEnforced records whether the cluster's CNI enforces NetworkPolicy. // The runtime writes the policy either way; this only affects what the // operator is told. @@ -44,17 +45,19 @@ func NewKubernetes(config KubernetesConfig, controlURL string) (*Kubernetes, err if err != nil { return nil, err } - return &Kubernetes{client: client, ControlURL: controlURL}, nil + return &Kubernetes{ + client: client, ControlURL: controlURL, + CheckpointImage: "busybox:1.37", + }, nil } func (k *Kubernetes) Namespace() string { return k.client.namespace } -// Capabilities reports what Kubernetes cannot do rather than pretending. +// Capabilities reports what this runtime can do. // -// There is no API to freeze a running Pod, so pause and resume are refused. -// stdio MCP servers are executed through the attach subresource. +// Pause is signalled rather than frozen; see Pause for what that costs. func (k *Kubernetes) Capabilities() Capabilities { - return Capabilities{Pause: false, StdioTools: true} + return Capabilities{Pause: true, StdioTools: true} } func podName(sessionID string) string { return "agent-platform-session-" + sessionID } @@ -130,14 +133,6 @@ func (k *Kubernetes) Start(ctx context.Context, job domain.RuntimeJob) (string, relative := strings.TrimPrefix(secret.MountPath, secretMountPath+"/") data[secretKey(relative)] = base64.StdEncoding.EncodeToString(secret.Value) } - if job.Checkpoint != nil { - if len(job.Checkpoint.Value) > checkpointLimit { - return "", fmt.Errorf( - "checkpoint is %d bytes; the Kubernetes runtime carries at most %d", - len(job.Checkpoint.Value), checkpointLimit) - } - data["checkpoint"] = base64.StdEncoding.EncodeToString(job.Checkpoint.Value) - } if len(data) > 0 { secret := map[string]any{ "apiVersion": "v1", "kind": "Secret", "type": "Opaque", @@ -202,7 +197,7 @@ func (k *Kubernetes) podSpec(job domain.RuntimeJob, hasSecret bool) map[string]a if job.Checkpoint != nil { env = append(env, kubeEnv("AGENT_PLATFORM_CHECKPOINT_ID", job.Checkpoint.ID), - kubeEnv("AGENT_PLATFORM_CHECKPOINT_PATH", "/workspace/.agent-platform/checkpoint"), + kubeEnv("AGENT_PLATFORM_CHECKPOINT_PATH", checkpointPath), kubeEnv("AGENT_PLATFORM_CHECKPOINT_FORMAT", job.Checkpoint.Format), kubeEnv("AGENT_PLATFORM_CHECKPOINT_SHA256", job.Checkpoint.SHA256), ) @@ -277,16 +272,7 @@ func (k *Kubernetes) podSpec(job domain.RuntimeJob, hasSecret bool) map[string]a spec["activeDeadlineSeconds"] = manifest.Runtime.TimeoutSecs } if job.Checkpoint != nil { - spec["initContainers"] = []any{k.checkpointInitContainer(manifest.Runtime.User)} - volumes = append(volumes, map[string]any{ - "name": "checkpoint", - "secret": map[string]any{ - "secretName": secretName(sessionID), - "defaultMode": 0o400, - "items": []any{map[string]any{"key": "checkpoint", "path": "checkpoint"}}, - }, - }) - spec["volumes"] = volumes + spec["initContainers"] = []any{k.checkpointInitContainer(sessionID, job)} } return map[string]any{ @@ -301,19 +287,33 @@ func (k *Kubernetes) podSpec(job domain.RuntimeJob, hasSecret bool) map[string]a } } -// checkpointInitContainer copies the restored checkpoint into the workspace, -// mirroring what the Docker runtime does with an init container over stdin. -func (k *Kubernetes) checkpointInitContainer(runtimeUser string) map[string]any { +// checkpointInitContainer fetches the restored checkpoint into the workspace. +// +// It streams from the Control Plane rather than carrying the bytes in a Secret, +// which would cap the checkpoint at Kubernetes' 1 MB object limit. The digest +// is re-checked here so a truncated transfer fails the Pod instead of handing +// the Agent a partial file. +func (k *Kubernetes) checkpointInitContainer(sessionID string, job domain.RuntimeJob) map[string]any { + script := `set -eu +mkdir -p "$(dirname "$CHECKPOINT_PATH")" +wget -q --header="Authorization: Bearer $AGENT_PLATFORM_AGENT_TOKEN" \ + -O "$CHECKPOINT_PATH" "$CHECKPOINT_URL" +echo "$AGENT_PLATFORM_CHECKPOINT_SHA256 $CHECKPOINT_PATH" | sha256sum -c - +chmod 0400 "$CHECKPOINT_PATH"` return map[string]any{ - "name": "restore-checkpoint", - "image": "busybox:1.37", - "command": []string{"sh", "-c", - "mkdir -p /workspace/.agent-platform && " + - "cp /checkpoint/checkpoint /workspace/.agent-platform/checkpoint && " + - "chmod 0400 /workspace/.agent-platform/checkpoint"}, + "name": "restore-checkpoint", + "image": k.CheckpointImage, + "imagePullPolicy": "IfNotPresent", + "command": []string{"sh", "-c", script}, + "env": []any{ + kubeEnv("CHECKPOINT_PATH", checkpointPath), + kubeEnv("CHECKPOINT_URL", + modelControlURL(k.ControlURL)+"/agent/v1/sessions/"+sessionID+"/checkpoint"), + kubeEnv("AGENT_PLATFORM_AGENT_TOKEN", job.AgentToken), + kubeEnv("AGENT_PLATFORM_CHECKPOINT_SHA256", job.Checkpoint.SHA256), + }, "volumeMounts": []any{ map[string]any{"name": "workspace", "mountPath": "/workspace"}, - map[string]any{"name": "checkpoint", "mountPath": "/checkpoint", "readOnly": true}, }, "securityContext": map[string]any{ "readOnlyRootFilesystem": true, @@ -376,11 +376,41 @@ func (k *Kubernetes) Inspect(ctx context.Context, sessionID string) (ContainerSt return state, nil } -// Pause and Resume are refused: Kubernetes has no equivalent of the cgroup -// freezer, and silently ignoring the request would leave a Session waiting for -// a state it will never reach. -func (k *Kubernetes) Pause(context.Context, string) error { return ErrUnsupported } -func (k *Kubernetes) Resume(context.Context, string) error { return ErrUnsupported } +// Pause stops every process in the Agent container. +// +// Kubernetes has no equivalent of Docker's cgroup freezer, so this signals +// instead: SIGSTOP to every process in the container's PID namespace, which is +// what a freeze amounts to from the Agent's point of view. Resume sends +// SIGCONT. +// +// It walks /proc rather than using kill -1 so the helper shell does not stop +// itself before it has finished signalling the rest. The cost is a dependency +// on the Agent image: it needs a shell and /proc, which a distroless image does +// not have. A pause against such an image fails, and the worker reports that +// rather than leaving the Session waiting in `pausing`. +func (k *Kubernetes) Pause(ctx context.Context, sessionID string) error { + return k.signalAgentProcesses(ctx, sessionID, "STOP") +} + +func (k *Kubernetes) Resume(ctx context.Context, sessionID string) error { + return k.signalAgentProcesses(ctx, sessionID, "CONT") +} + +func (k *Kubernetes) signalAgentProcesses(ctx context.Context, sessionID, signal string) error { + script := `set -eu +for entry in /proc/[0-9]*; do + pid="${entry#/proc/}" + [ "$pid" = "$$" ] && continue + kill -` + signal + ` "$pid" 2>/dev/null || true +done` + ctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + err := k.client.execIn(ctx, podName(sessionID), agentContainerName, []string{"sh", "-c", script}) + if err != nil { + return fmt.Errorf("signal %s to the Agent container: %w", signal, err) + } + return nil +} func (k *Kubernetes) Stop(ctx context.Context, sessionID string, graceSeconds int) error { err := k.client.deletePod(ctx, podName(sessionID), graceSeconds) diff --git a/internal/worker/kubernetes_attach.go b/internal/worker/kubernetes_attach.go index a3d6474..d7e92ed 100644 --- a/internal/worker/kubernetes_attach.go +++ b/internal/worker/kubernetes_attach.go @@ -30,6 +30,85 @@ type attachSession struct { conn *websocket.Conn } +// execIn runs a command in a container and reports whether it succeeded. +// +// The error channel carries a Status with the exit code, which is how a failed +// command is distinguished from a failed connection. +func (c *kubeClient) execIn(ctx context.Context, pod, container string, command []string) error { + _, err := c.execCapture(ctx, pod, container, command) + return err +} + +// execCapture runs a command and returns its stdout. +func (c *kubeClient) execCapture(ctx context.Context, pod, container string, command []string) (string, error) { + endpoint, err := url.Parse(c.baseURL + "/api/v1/namespaces/" + c.namespace + + "/pods/" + pod + "/exec") + if err != nil { + return "", err + } + switch endpoint.Scheme { + case "https": + endpoint.Scheme = "wss" + case "http": + endpoint.Scheme = "ws" + } + query := endpoint.Query() + query.Set("container", container) + query.Set("stdout", "true") + query.Set("stderr", "true") + query.Set("stdin", "false") + query.Set("tty", "false") + for _, argument := range command { + query.Add("command", argument) + } + endpoint.RawQuery = query.Encode() + + header := http.Header{} + if c.token != "" { + header.Set("Authorization", "Bearer "+c.token) + } + conn, _, err := websocket.Dial(ctx, endpoint.String(), &websocket.DialOptions{ + HTTPClient: c.http, + HTTPHeader: header, + Subprotocols: attachSubprotocols, + }) + if err != nil { + return "", fmt.Errorf("exec in %s: %w", pod, err) + } + conn.SetReadLimit(1 << 20) + session := &attachSession{conn: conn} + defer session.Close() + + var stdout, stderr strings.Builder + for { + channel, payload, err := session.Read(ctx) + if err != nil { + // The stream closing without an error frame means success. + if ctx.Err() != nil { + return stdout.String(), ctx.Err() + } + return stdout.String(), nil + } + switch channel { + case channelStdout: + stdout.Write(payload) + case channelStderr: + stderr.Write(payload) + case channelError: + // A zero exit status arrives as a Status with status "Success". + status := string(payload) + if strings.Contains(status, `"status":"Success"`) { + return stdout.String(), nil + } + detail := strings.TrimSpace(stderr.String()) + if detail == "" { + detail = strings.TrimSpace(status) + } + return stdout.String(), fmt.Errorf("command failed in %s: %s", pod, truncate(detail, 300)) + } + } +} + // attachTo opens an attach stream to a container's stdin and stdout. func (c *kubeClient) attachTo(ctx context.Context, pod, container string) (*attachSession, error) { endpoint, err := url.Parse(c.baseURL + "/api/v1/namespaces/" + c.namespace + diff --git a/internal/worker/kubernetes_integration_test.go b/internal/worker/kubernetes_integration_test.go index b3b5e75..307623a 100644 --- a/internal/worker/kubernetes_integration_test.go +++ b/internal/worker/kubernetes_integration_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "os" + "strconv" "strings" "testing" "time" @@ -191,6 +192,9 @@ func TestIntegrationKubernetesMountsSessionSecrets(t *testing.T) { func TestIntegrationKubernetesRestoresACheckpoint(t *testing.T) { runtime := integrationRuntime(t) + if os.Getenv("CHECKPOINT_URL_REACHABLE") != "true" { + t.Skip("set CHECKPOINT_URL_REACHABLE=true when a Control Plane serves /agent/v1/.../checkpoint") + } job := integrationJob(t, []string{"sh", "-c", "grep -q restored-state /workspace/.agent-platform/checkpoint"}, time.Second) job.Checkpoint = &domain.RuntimeCheckpoint{ @@ -215,19 +219,34 @@ func TestIntegrationKubernetesRestoresACheckpoint(t *testing.T) { } } -func TestIntegrationKubernetesRefusesAnOversizedCheckpoint(t *testing.T) { +// Checkpoints used to ride in a per-Session Secret, which caps at 1 MB. They +// are now fetched from the Control Plane, so a checkpoint larger than that has +// to work. +func TestIntegrationKubernetesRestoresACheckpointLargerThanASecret(t *testing.T) { runtime := integrationRuntime(t) - job := integrationJob(t, []string{"true"}, time.Second) + if os.Getenv("CHECKPOINT_URL_REACHABLE") != "true" { + t.Skip("set CHECKPOINT_URL_REACHABLE=true when a Control Plane serves /agent/v1/.../checkpoint") + } + job := integrationJob(t, []string{"sh", "-c", + "test \"$(wc -c < /workspace/.agent-platform/checkpoint)\" -gt 1048576"}, time.Second) job.Checkpoint = &domain.RuntimeCheckpoint{ ID: uuid.NewString(), Format: "application/octet-stream", - Value: make([]byte, checkpointLimit+1), + Value: make([]byte, 2<<20), } - t.Cleanup(func() { runtime.Cleanup(context.Background(), job.Session.ID) }) + sessionID := job.Session.ID + t.Cleanup(func() { runtime.Cleanup(context.Background(), sessionID) }) - // Better a clear refusal than a Secret the API server rejects with its own - // size error halfway through Start. - if _, err := runtime.Start(context.Background(), job); err == nil { - t.Fatal("an oversized checkpoint was accepted") + if _, err := runtime.Start(context.Background(), job); err != nil { + t.Fatalf("Start: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 180*time.Second) + defer cancel() + code, err := runtime.Wait(ctx, sessionID) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if code != 0 { + t.Fatalf("a checkpoint above the Secret limit did not reach the workspace; exit %d", code) } } @@ -372,3 +391,66 @@ func TestIntegrationKubernetesRejectsAnIncompleteStdioRegistration(t *testing.T) } } } + +// Pause has to actually stop the Agent doing work, not merely report success. +// The probe appends a line every second, so a frozen container stops growing the +// file and a resumed one starts again. +func TestIntegrationKubernetesPausesAndResumesAnAgent(t *testing.T) { + runtime := integrationRuntime(t) + job := integrationJob(t, []string{"sh", "-c", + "i=0; while [ $i -lt 600 ]; do echo tick >> /workspace/ticks; i=$((i+1)); sleep 1; done"}, + 600*time.Second) + sessionID := job.Session.ID + t.Cleanup(func() { runtime.Cleanup(context.Background(), sessionID) }) + + if _, err := runtime.Start(context.Background(), job); err != nil { + t.Fatalf("Start: %v", err) + } + waitForState(t, runtime, sessionID, "running", 90*time.Second) + + // Confirm it is ticking before pausing, or a frozen count proves nothing. + before := tickCount(t, runtime, sessionID) + time.Sleep(3 * time.Second) + running := tickCount(t, runtime, sessionID) + if running <= before { + t.Fatalf("the probe was not ticking to begin with: %d then %d", before, running) + } + + if err := runtime.Pause(context.Background(), sessionID); err != nil { + t.Fatalf("Pause: %v", err) + } + frozen := tickCount(t, runtime, sessionID) + time.Sleep(6 * time.Second) + if after := tickCount(t, runtime, sessionID); after != frozen { + t.Fatalf("the Agent kept working while paused: %d then %d", frozen, after) + } + + if err := runtime.Resume(context.Background(), sessionID); err != nil { + t.Fatalf("Resume: %v", err) + } + deadline := time.Now().Add(30 * time.Second) + for time.Now().Before(deadline) { + if tickCount(t, runtime, sessionID) > frozen { + return + } + time.Sleep(time.Second) + } + t.Fatal("the Agent did not resume working after Resume") +} + +// tickCount reads the probe's line count out of the container. Reading through +// a fresh exec also shows exec still works against a paused container, which is +// what makes Resume reachable. +func tickCount(t *testing.T, runtime *Kubernetes, sessionID string) int { + t.Helper() + out, err := runtime.client.execCapture(context.Background(), podName(sessionID), + agentContainerName, []string{"sh", "-c", "wc -l < /workspace/ticks 2>/dev/null || echo 0"}) + if err != nil { + t.Fatalf("read tick count: %v", err) + } + count, err := strconv.Atoi(strings.TrimSpace(out)) + if err != nil { + t.Fatalf("tick count %q: %v", out, err) + } + return count +} diff --git a/internal/worker/kubernetes_test.go b/internal/worker/kubernetes_test.go index c0310f8..a469c55 100644 --- a/internal/worker/kubernetes_test.go +++ b/internal/worker/kubernetes_test.go @@ -222,18 +222,19 @@ func TestPodSpecSerializesToJSON(t *testing.T) { } } -// Kubernetes has no equivalent of the cgroup freezer, so pause is refused -// rather than silently ignored. -func TestKubernetesRefusesWhatItCannotDo(t *testing.T) { +// Pause is signalled rather than frozen, and a failure to signal has to surface +// so the Session does not sit in `pausing` waiting for a state it never reaches. +func TestKubernetesReportsAFailedPause(t *testing.T) { runtime := testKubernetes(t) - if runtime.Capabilities().Pause { - t.Error("Kubernetes claims it can freeze a pod") + if !runtime.Capabilities().Pause { + t.Fatal("pause is reported as unsupported") } + // The API server address does not resolve, so signalling cannot succeed. if err := runtime.Pause(t.Context(), "s-1"); err == nil { - t.Error("Pause silently succeeded") + t.Error("Pause reported success without reaching the container") } if err := runtime.Resume(t.Context(), "s-1"); err == nil { - t.Error("Resume silently succeeded") + t.Error("Resume reported success without reaching the container") } } diff --git a/test/e2e/runtime_policy_flow.py b/test/e2e/runtime_policy_flow.py index f97ff9f..a8d2732 100644 --- a/test/e2e/runtime_policy_flow.py +++ b/test/e2e/runtime_policy_flow.py @@ -81,7 +81,14 @@ def main() -> None: workspace_id, { "agentVersionId": checkpoint_agent["id"], - "input": {"checkpointTest": {"content": "durable restored state"}}, + # Above the 1 MB a Kubernetes Secret can carry, so the checkpoint + # has to travel by fetch rather than by mounted Secret. + "input": { + "checkpointTest": { + "content": "durable restored state", + "sizeBytes": 1_500_000, + } + }, }, ) api.poll_session( @@ -113,12 +120,18 @@ def main() -> None: timeout=120, ) child_events = api.events(child["id"]) - if not any( - event["type"] == "agent.message.created" - and "durable restored state" in event["payload"].get("content", "") - for event in child_events - ): + restored = next( + ( + event["payload"].get("content", "") + for event in child_events + if event["type"] == "agent.message.created" + ), + "", + ) + if "durable restored state" not in restored: raise RuntimeError("Restored checkpoint content did not reach the child Agent.") + if "1500000 bytes" not in restored: + raise RuntimeError(f"The checkpoint was truncated in transit: {restored}") target = "http://agent-platform-model-mock:8000/egress" allowed_agent = register_agent( @@ -177,6 +190,7 @@ def main() -> None: "parentSessionId": parent["id"], "childSessionId": child["id"], "checkpointRestore": "passed", + "checkpointAboveSecretLimit": "passed", "allowedEgress": "passed", "deniedEgress": "passed", },