diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..151a4ad --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,16 @@ +# Agent instructions for RAX + +RAX is a Rust port of AX (https://github.com/google/ax). These rules apply to +every agent working in this repository. + +## Hard rules + +- **Never mention or reference upstream AX pull requests or issues directly.** + That covers PR and issue numbers (`#123`, `google/ax#123`, "AX PR 123") and + links to them (`github.com/google/ax/pull/...`, `.../issues/...`), in RAX + pull requests, commit messages, code comments, docs, review replies and any + other message. Mentions like these notify the upstream project and create + cross-references there. Describe the upstream change in your own words + instead, for example "as AX now does" or "AX's CreateTask now rejects + missing workspaces". Naming an AX commit hash to pin a parity point (as the + README does) is fine. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..43c994c --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/Cargo.lock b/Cargo.lock index a622dfe..0c9b0e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -76,6 +76,45 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +[[package]] +name = "asn1-rs" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5493c3bedbacf7fd7382c6346bbd66687d12bbaad3a89a2d2c303ee6cf20b048" +dependencies = [ + "asn1-rs-derive", + "asn1-rs-impl", + "displaydoc", + "nom", + "num-traits", + "rusticata-macros", + "thiserror 1.0.69", + "time", +] + +[[package]] +name = "asn1-rs-derive" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "965c2d33e53cb6b267e148a4cb0760bc01f4904c1cd4bb4002a085bb016d1490" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "asn1-rs-impl" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -347,6 +386,12 @@ dependencies = [ "cmov", ] +[[package]] +name = "data-encoding" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" + [[package]] name = "deadpool" version = "0.13.1" @@ -380,6 +425,20 @@ dependencies = [ "tokio", ] +[[package]] +name = "der-parser" +version = "9.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cd0a5c643689626bec213c4d8bd4d96acc8ffdb4ad4bb6bc16abf27d5f4b553" +dependencies = [ + "asn1-rs", + "displaydoc", + "nom", + "num-bigint", + "num-traits", + "rusticata-macros", +] + [[package]] name = "deranged" version = "0.5.9" @@ -408,6 +467,17 @@ dependencies = [ "ctutils", ] +[[package]] +name = "displaydoc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + [[package]] name = "either" version = "1.19.0" @@ -933,6 +1003,12 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "mio" version = "1.2.4" @@ -950,6 +1026,16 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1022,6 +1108,15 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "oid-registry" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8d8034d9489cdaf79228eb9f6a3b8d7bb32ba00d6645ebd48eef4077ceb5bd9" +dependencies = [ + "asn1-rs", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1308,7 +1403,7 @@ dependencies = [ "prost-reflect", "prost-types", "protox-parse", - "thiserror", + "thiserror 2.0.21", ] [[package]] @@ -1320,7 +1415,7 @@ dependencies = [ "logos", "miette", "prost-types", - "thiserror", + "thiserror 2.0.21", ] [[package]] @@ -1429,7 +1524,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml_ng", - "thiserror", + "thiserror 2.0.21", "tonic", "tonic-prost", "tonic-prost-build", @@ -1473,8 +1568,10 @@ dependencies = [ "serde_json", "serde_yaml_ng", "tempfile", + "time", "tokio", "tonic", + "x509-parser", ] [[package]] @@ -1535,7 +1632,7 @@ dependencies = [ "rand 0.9.5", "rax-api", "serde_json", - "thiserror", + "thiserror 2.0.21", "tokio", "tokio-postgres", "tracing", @@ -1653,6 +1750,15 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rusticata-macros" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632" +dependencies = [ + "nom", +] + [[package]] name = "rustix" version = "1.1.6" @@ -1926,6 +2032,17 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "tempfile" version = "3.27.0" @@ -1939,13 +2056,33 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.21", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -1979,6 +2116,7 @@ dependencies = [ "powerfmt", "serde_core", "time-core", + "time-macros", ] [[package]] @@ -1987,6 +2125,16 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" +[[package]] +name = "time-macros" +version = "0.2.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e689342a48d2ea927c87ea50cabf8594854bf940e9310208848d680d668ed85" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinyvec" version = "1.13.3" @@ -2604,6 +2752,23 @@ version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" +[[package]] +name = "x509-parser" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcbc162f30700d6f3f82a24bf7cc62ffe7caea42c0b2cba8bf7f3ae50cf51f69" +dependencies = [ + "asn1-rs", + "data-encoding", + "der-parser", + "lazy_static", + "nom", + "oid-registry", + "rusticata-macros", + "thiserror 1.0.69", + "time", +] + [[package]] name = "yasna" version = "0.5.2" diff --git a/Cargo.toml b/Cargo.toml index 67e9ce3..be00bb3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,3 +60,4 @@ rcgen = { version = "0.13", default-features = false, features = ["crypto", "pem tonic-reflection = "0.14" rax-guest = { path = "crates/rax-guest" } rax-task-runner = { path = "crates/rax-task-runner" } +x509-parser = { version = "0.16", default-features = false } diff --git a/NOTICE b/NOTICE index 5c9348c..6d1d442 100644 --- a/NOTICE +++ b/NOTICE @@ -6,8 +6,9 @@ proto/ax/v1alpha1/ax.proto, keep their original copyright and license headers. The Rust code is a translation of AX's Go implementation and is offered under the terms in LICENSE, with the Apache 2.0 notices above preserved. -proto/ateapi/ateapi.proto is copied from Agent Substrate -(https://github.com/agent-substrate/substrate, go.ate.dev/substrate v0.6.0), +proto/ateapi/v1/ateapi.proto is copied from Agent Substrate +(https://github.com/agent-substrate/substrate, api/ateapi/v1 at +go.ate.dev/substrate v1.0.0-rc.1), Copyright 2026 Google LLC, licensed under the Apache License, Version 2.0. proto/ateenv/v1alpha/guest.proto is copied from Agent Substrate's env module diff --git a/README.md b/README.md index 0c1e7b8..735855f 100644 --- a/README.md +++ b/README.md @@ -160,7 +160,7 @@ refuses leaves the task `Pending`, to retry. ## Parity with AX **RAX covers everything AX ships**, as of AX's `main` at -[`f0228f1`](https://github.com/google/ax/commit/f0228f1) (2026-10-30): the +[`6e970c5`](https://github.com/google/ax/commit/6e970c5) (2026-11-03): the control plane (`ax-server`), the process inside each task container (`ax-task-runner`, including the guest services `ax ssh` talks to), the Claude Code runner and its image, and the CLI (`ax`, with its tunnels). The @@ -183,12 +183,12 @@ and runner, and so on (see [What is tested](#what-is-tested)). | `pkg/apis/v1alpha1` (types, resources) | Task, Workspace and Model types, YAML and JSON encoding, name and spec validation, CPU and memory quantities | `rax-api` | done | | `internal/store`, `store/memory`, `store/postgres` | Resource store; PostgreSQL with AX's schema | `rax-store` | done (no TLS to PostgreSQL yet) | | `internal/lock`, `store/postgres/lock.go` | Per-resource locks; PostgreSQL leases with renewal and takeover | `rax-store` | done | -| `internal/server` | The `ax.v1alpha1.AX` gRPC API (all 14 RPCs), actor-owned status, `/healthz` | `rax-server` | done | +| `internal/server` | The `ax.v1alpha1.AX` gRPC API (all 14 RPCs), actor-owned status, bound workspaces checked before a task is stored, `/healthz` | `rax-server` | done | | `cmd/ax-server` | Server binary, flags and environment | `rax-server` binary | done (same flags) | -| `internal/substrate` | Substrate control API client: TLS, token, retries, atespaces, templates, actors | `rax-substrate` | done | +| `internal/substrate` | Substrate control API client (`substrate.ateapi.v1`, Substrate v1.0.0-rc.1): TLS, token, retries, atespaces, templates with the cluster-wide snapshot location, actors | `rax-substrate` | done | | `internal/controller` | Task reconciler: per-task templates, explicit `Suspend`, `Resume` and `Delete`, workspace readiness, delete cleanup | `rax-server::controller` | done | | `cmd/ax` | CLI: `apply`, `get`, `describe`, `suspend`, `resume`, `delete`, `ssh`, `ctx`/`context`, `tunnel list`/`stop`, `version` | `rax-cli` (`rax`) | done (same flags, output and exit codes) | -| `internal/tunnel` | CLI tunnels: per-context state, `kubectl port-forward`, router TLS through a PodCertificateRequest | `rax-cli::tunnel`, `rax-cli::router_tls` | done (same state and cache files) | +| `internal/tunnel` | CLI tunnels: per-context state, `kubectl port-forward`, router TLS through a PodCertificateRequest, cached no longer than the certificate is valid | `rax-cli::tunnel`, `rax-cli::router_tls` | done (same state and cache files) | | `internal/guest` | Client for the in-actor guest daemon (`ax ssh`: exec, stdin, signals, Ctrl-C twice to leave) | `rax-guest::client` | done | | Substrate `env/guest` (embedded by AX's runner) | Guest daemon: process service with spooled output, stdin, signals, limits; filesystem service; reflection | `rax-guest` | done | | `cmd/ax-task-runner`, `runner` | The process inside each task container: runs setup, then supervises the task's command | `rax-task-runner` | done (same flags, `-flag` spellings too) | @@ -227,7 +227,7 @@ and runner, and so on (see [What is tested](#what-is-tested)). ### What is tested -167 tests, all run in CI (`cargo test --all`, with a PostgreSQL 16 service). +172 tests, all run in CI (`cargo test --all`, with a PostgreSQL 16 service). They port the Go tests for the packages above and add some of their own. | Suite | Tests | Covers | @@ -235,16 +235,16 @@ They port the Go tests for the packages above and add some of their own. | `rax-api` unit | 21 | YAML round trips and strict decoding (unknown fields rejected), lenient decoding, Workspace MCP and Model secret-key spellings, task config, workspace binding, name/meta/task/file validation, resource validation, CPU and memory parsing | | `rax-store` unit | 12 | One conformance suite run on both the memory and the PostgreSQL store (CRUD, defaults, paging, not-found), concurrent schema creation, records with unknown fields, lock serialization, cancellation, lease renewal and takeover of expired leases, a crashed server's tasks kept and its locks freed after a restart, URL redaction | | `rax-server` unit | 15 | Actor state to task phase and Ready, stored status keeps only AX's fields, `/healthz`, template naming and digest, the template name pattern, worker and router addresses, Go-style `:port` listen addresses; the pod runner's manifests, pod phases, retryable errors, one run per resume and label-scoped listing | -| `rax-server` `tests/service.rs` | 17 | gRPC round trip over the real router, invalid names, workspace binding, suspend and resume lifecycle, delete (terminating, cleanup failure, status write failure), suspend success and failure, actor state read from Substrate, new tasks start suspended, failed and transiently failed provisioning, a later success clearing a stored failure, a suspend cancelling a failed delete, creation time kept, concurrent creates serialized | +| `rax-server` `tests/service.rs` | 18 | gRPC round trip over the real router, invalid names, workspace binding, bound workspace required (missing, other atespace, store error), suspend and resume lifecycle, delete (terminating, cleanup failure, status write failure), suspend success and failure, actor state read from Substrate, new tasks start suspended, failed and transiently failed provisioning, a later success clearing a stored failure, a suspend cancelling a failed delete, creation time kept, concurrent creates serialized | | `rax-server` `tests/controller.rs` | 17 | Ports of AX's reconciler tests against an in-memory Substrate: actor named after the task, worker IP, default image, suspend and failed suspend, rejected template (bare Substrate message, no fallback), transient template failure, workspace readiness polled directly and through the router and kept across suspend and resume, no new template on status changes, resource limits and their failure paths, delete removes only the task's templates, waits for the actor, times out, transient failures and recovery; plus a full create, resume, get, suspend and delete through the service | -| `rax-substrate` unit | 12 | Retry and backoff rules, transient errors, template building, snapshot location, resource limits, TLS config, HTTP readiness probe | +| `rax-substrate` unit | 11 | Retry and backoff rules, transient errors, template building (no snapshot location unless given), resource limits, TLS config, HTTP readiness probe | | `rax-substrate` `tests/client.rs` | 13 | Client against an in-memory Substrate: atespace, template and actor creation, revert or recreate of crashed actors, waiting out deleting actors, resume retries (`ResourceExhausted` not retried), delete polling, missing actors and timeout, template deletion retries, pagination, token re-read on every call | | `rax-substrate` `tests/tls.rs` | 1 | A real TLS connection: CA file, authority override, wrong server name rejected, unreadable CA, skip-verify | | `rax-guest` unit | 9 | Output spooling and replay from offsets, exit codes and signal exits (128+n), stdin writes and close, signals to the process group, the concurrency cap, log truncation, the default timeout, pruning, start failures, timeouts, path cleaning and confinement, the signal table | | `rax-task-runner` unit | 8 | First-run marker, git clone into a subdirectory with depth, git errors withholding the marker, inlined files, marker names, clone destinations, spec loading with unknown fields and empty specs | | `rax-task-runner` `tests/runner.rs` | 8 | Ports of AX's runner, metadata and guest client tests: command env, working directory and exit code, serving after the command exits, stopping the command's process group, waiting without a command, bound workspace setup and readiness, metadata only on loopback, guest services only with `spec.debug`, exec with stdin, empty stdin, exit codes and forwarded signals | | `rax-claude-code-runner` unit | 7 | Ports of AX's claude-code-runner tests: rejected configs, the minimal invocation, flags by value type in sorted order, the Workspace's MCP servers as `--mcp-config` byte for byte, Go-style JSON escaping and numbers, `spec.command` left alone | -| `rax-cli` unit | 15 | Ports of AX's argument parsing, `ssh` argument and Ctrl-C escalation tests, tunnel context names and the port-forward pattern, AX's tunnel state files, tabwriter-style tables, Go `%v` values, document splitting, kinds, worker addresses, router TLS configs | +| `rax-cli` unit | 20 | Ports of AX's argument parsing (with `-n` over `$AX_SYSTEM_NAMESPACE`), `ssh` argument and Ctrl-C escalation tests, tunnel context names and the port-forward pattern, AX's tunnel state files, tabwriter-style tables, Go `%v` values, document splitting, kinds, worker addresses, router TLS configs, router certificate expiry in the cache | | `rax-cli` `tests/cli.rs` | 5 | Ports of AX's CLI tests, running the `rax` binary against rax-server: `apply` atespaces for every kind, immutable tasks, `unchanged` and `configured`, empty and invalid documents, `get` aliases and missing resources, exact `describe` output, `delete`, and the offline commands (`version`, `ctx`, `tunnel list`/`stop`, usage) | | `rax-cli` `tests/e2e_local.rs` | 7 | AX's seven E2E scenarios, replayed without Kubernetes through the CLI binary against rax-server, an in-memory Substrate and real task-runners: CLI version, workspace and model CRUD, task lifecycle and Substrate sync, metadata server and guest exec, workspace setup with files and git, suspend and resume persistence, describe and atespace isolation | diff --git a/crates/rax-cli/Cargo.toml b/crates/rax-cli/Cargo.toml index 39c5caf..3775721 100644 --- a/crates/rax-cli/Cargo.toml +++ b/crates/rax-cli/Cargo.toml @@ -26,6 +26,7 @@ serde_json.workspace = true serde_yaml_ng.workspace = true tokio = { workspace = true, features = ["io-std"] } tonic.workspace = true +x509-parser.workspace = true [dev-dependencies] axum.workspace = true @@ -34,4 +35,5 @@ rax-store.workspace = true rax-substrate = { workspace = true, features = ["testing"] } rax-task-runner.workspace = true tempfile.workspace = true +time = "0.3" tokio = { workspace = true, features = ["io-std", "rt-multi-thread"] } diff --git a/crates/rax-cli/src/args.rs b/crates/rax-cli/src/args.rs index f1d387e..883c8d5 100644 --- a/crates/rax-cli/src/args.rs +++ b/crates/rax-cli/src/args.rs @@ -32,9 +32,18 @@ impl Default for GlobalArgs { /// Splits out the command name and the global flags. Flags are not /// interpreted after `--`, so a command run by `rax ssh -- ...` is -/// passed through intact. +/// passed through intact. The system namespace defaults to +/// `$AX_SYSTEM_NAMESPACE`, then `ax-system`, and `-n`/`--namespace` +/// overrides both. pub fn parse_global_args(args: &[String]) -> GlobalArgs { + parse_global_args_with_env(args, std::env::var("AX_SYSTEM_NAMESPACE").ok().as_deref()) +} + +fn parse_global_args_with_env(args: &[String], env_namespace: Option<&str>) -> GlobalArgs { let mut g = GlobalArgs::default(); + if let Some(ns) = env_namespace.filter(|n| !n.is_empty()) { + g.namespace = ns.into(); + } let mut i = 0; while i < args.len() { let arg = args[i].as_str(); @@ -242,6 +251,39 @@ mod tests { } } + #[test] + fn global_args_namespace_precedence() { + for (env, args, want) in [ + // Defaults to ax-system. + (None, v(&["ssh", "task123"]), "ax-system"), + (Some(""), v(&["ssh", "task123"]), "ax-system"), + // Falls back to AX_SYSTEM_NAMESPACE when -n is omitted. + ( + Some("env-ax-system"), + v(&["ssh", "task123"]), + "env-ax-system", + ), + // An explicit -n overrides AX_SYSTEM_NAMESPACE. + ( + Some("env-ax-system"), + v(&["-n", "flag-ax-system", "ssh", "task123"]), + "flag-ax-system", + ), + // Even when it names the default. + ( + Some("env-ax-system"), + v(&["--namespace=ax-system", "ssh", "task123"]), + "ax-system", + ), + ] { + assert_eq!( + parse_global_args_with_env(&args, env).namespace, + want, + "{args:?} with env={env:?}" + ); + } + } + #[test] fn ssh_args() { let ok = |args: &[&str], task: &str, cmd: &[&str]| { diff --git a/crates/rax-cli/src/commands.rs b/crates/rax-cli/src/commands.rs index 1fe3046..645d595 100644 --- a/crates/rax-cli/src/commands.rs +++ b/crates/rax-cli/src/commands.rs @@ -860,6 +860,76 @@ pub fn escalate_interrupts( rx } +/// Connects to a task's guest services: directly when the worker is +/// reachable, else through Substrate's atenet router (mTLS on 443, falling +/// back to plaintext on 80). Returns the port-forward the connection runs +/// over, if any. +fn dial_guest( + host: &str, + port: u16, + target_actor: &str, + kube_context: &str, + ax_namespace: &str, +) -> Result<(GuestClient, Option), Error> { + // 1. The worker directly, when it is reachable (in-cluster or a mesh). + let endpoint = join_host_port(host, port); + let direct = endpoint + .to_socket_addrs() + .ok() + .and_then(|mut a| a.next()) + .is_some_and(|a| TcpStream::connect_timeout(&a, Duration::from_millis(500)).is_ok()); + if direct { + let c = GuestClient::connect(&endpoint, None, None) + .map_err(|e| format!("connecting to guest at {endpoint}: {e}"))?; + return Ok((c, None)); + } + + // 2. Through Substrate's atenet router: mTLS on 443, else plaintext 80. + let ax_namespace = if ax_namespace.is_empty() { + rax_api::DEFAULT_SYSTEM_NAMESPACE + } else { + ax_namespace + }; + let setup_err = match dial_router_tls(target_actor, kube_context, ax_namespace)? { + Ok(connected) => return Ok(connected), + Err(setup_err) => setup_err, + }; + let pf = tunnel::port_forward(kube_context, "ate-system", "svc/atenet-router", 80) + .map_err(|e| format!("establishing port-forward to atenet-router: {e} ({setup_err})"))?; + let router = format!("127.0.0.1:{}", pf.port); + let c = GuestClient::connect(&router, Some(target_actor), None) + .map_err(|e| format!("connecting to guest via atenet-router at {router}: {e}"))?; + Ok((c, Some(pf))) +} + +/// Connects to the guest through the atenet router over mTLS on 443. The +/// outer error is a failed connection, which ends `rax ssh`; the inner one is +/// a failure to set up mTLS or the port-forward, after which the plaintext +/// route is tried. +#[allow(clippy::type_complexity)] +fn dial_router_tls( + target_actor: &str, + kube_context: &str, + ax_namespace: &str, +) -> Result), String>, Error> { + let tls = match router_tls_config(kube_context, ax_namespace, "ax-server") { + Ok(tls) => tls, + Err(e) => return Ok(Err(format!("mTLS setup error: {e}"))), + }; + let pf = match tunnel::port_forward(kube_context, "ate-system", "svc/atenet-router", 443) { + Ok(pf) => pf, + Err(e) => return Ok(Err(format!("port 443 forward error: {e}"))), + }; + let router = format!("127.0.0.1:{}", pf.port); + let c = GuestClient::connect( + &router, + Some(target_actor), + Some((tls, ROUTER_SERVER_NAME.to_owned())), + ) + .map_err(|e| format!("connecting to guest via atenet-router at {router}: {e}"))?; + Ok(Ok((c, Some(pf)))) +} + fn stdin_is_terminal() -> bool { // SAFETY: isatty only inspects the descriptor. unsafe { libc::isatty(0) == 1 } @@ -870,6 +940,7 @@ pub async fn run_ssh( server: &str, atespace: &str, kube_context: &str, + ax_namespace: &str, args: &[String], ) -> Result { let parsed = parse_ssh_args(args)?; @@ -906,54 +977,8 @@ pub async fn run_ssh( let (host, port) = guest_endpoint(&status.worker_ip); let target_actor = format!("{}/{}", meta_of(&task.metadata).1, status.actor); - let endpoint = join_host_port(&host, port); - // Kept alive for the session. - let mut _forward = None; - - // 1. The worker directly, when it is reachable (in-cluster or a mesh). - let direct = endpoint - .to_socket_addrs() - .ok() - .and_then(|mut a| a.next()) - .is_some_and(|a| TcpStream::connect_timeout(&a, Duration::from_millis(500)).is_ok()); - let guest = if direct { - GuestClient::connect(&endpoint, None, None) - .map_err(|e| format!("connecting to guest at {endpoint}: {e}"))? - } else { - // 2. Through Substrate's atenet router: mTLS on 443, else plaintext 80. - let ax_namespace = std::env::var("AX_SYSTEM_NAMESPACE") - .ok() - .filter(|n| !n.is_empty()) - .unwrap_or_else(|| rax_api::DEFAULT_SYSTEM_NAMESPACE.into()); - let mut via_tls = None; - if let Ok(tls) = router_tls_config(kube_context, &ax_namespace, "ax-server") - && let Ok(pf) = - tunnel::port_forward(kube_context, "ate-system", "svc/atenet-router", 443) - { - let router = format!("127.0.0.1:{}", pf.port); - let c = GuestClient::connect( - &router, - Some(&target_actor), - Some((tls, ROUTER_SERVER_NAME.to_owned())), - ) - .map_err(|e| format!("connecting to guest via atenet-router at {router}: {e}"))?; - _forward = Some(pf); - via_tls = Some(c); - } - match via_tls { - Some(c) => c, - None => { - let pf = tunnel::port_forward(kube_context, "ate-system", "svc/atenet-router", 80) - .map_err(|e| format!("establishing port-forward to atenet-router: {e}"))?; - let router = format!("127.0.0.1:{}", pf.port); - let c = GuestClient::connect(&router, Some(&target_actor), None).map_err(|e| { - format!("connecting to guest via atenet-router at {router}: {e}") - })?; - _forward = Some(pf); - c - } - } - }; + // The port-forward, if any, is kept alive for the session. + let (guest, _forward) = dial_guest(&host, port, &target_actor, kube_context, ax_namespace)?; // A bare shell on a terminal runs interactively, so it prompts and Ctrl-C // stops what it runs rather than the shell. diff --git a/crates/rax-cli/src/lib.rs b/crates/rax-cli/src/lib.rs index 439e104..6bc4a4a 100644 --- a/crates/rax-cli/src/lib.rs +++ b/crates/rax-cli/src/lib.rs @@ -131,13 +131,18 @@ pub fn run(argv: Vec) -> i32 { "ssh" => { // ssh leaves through process::exit so a blocked stdin read // can't hold the runtime open. - let code = - match commands::run_ssh(&server, &atespace, &opts.kube_context, &opts.args) - .await - { - Ok(code) => code, - Err(e) => fail(e), - }; + let code = match commands::run_ssh( + &server, + &atespace, + &opts.kube_context, + &opts.namespace, + &opts.args, + ) + .await + { + Ok(code) => code, + Err(e) => fail(e), + }; std::process::exit(code); } _ => unreachable!(), diff --git a/crates/rax-cli/src/router_tls.rs b/crates/rax-cli/src/router_tls.rs index b76b018..4992a33 100644 --- a/crates/rax-cli/src/router_tls.rs +++ b/crates/rax-cli/src/router_tls.rs @@ -3,8 +3,9 @@ //! The router's serving CA comes from the live servicedns //! ClusterTrustBundle; the client certificate is a short-lived one from the //! podidentity signer, requested with a PodCertificateRequest for the AX -//! server's service account. Credentials are cached for 30 minutes in the -//! tunnel directory, in the same format AX uses. +//! server's service account. Credentials are cached in the tunnel directory, +//! in the same format AX uses, for 30 minutes or until a minute before the +//! client certificate expires, whichever is sooner. use std::io::Write; use std::process::{Command, Stdio}; @@ -25,6 +26,9 @@ const LIVE_BUNDLE_SELECTOR: &str = "podcert.ate.dev/canarying=live"; /// The name the router's certificate is verified against. pub const ROUTER_SERVER_NAME: &str = "atenet-router.ate-system.svc"; const CACHE_TTL: Duration = Duration::from_secs(30 * 60); +/// How long before its NotAfter a client certificate is treated as expired, +/// so it is not used for a connection that outlives it. +const ROUTER_CERT_EXPIRY: Duration = Duration::from_secs(60); #[derive(Debug, Serialize, Deserialize)] struct CachedRouterTls { @@ -51,7 +55,15 @@ pub fn router_tls_config( let (chain_pem, key_pem) = request_pod_identity_cert(&context, namespace, service_account) .map_err(|e| format!("requesting podidentity client certificate: {e}"))?; let config = client_config(&trust_pem, &chain_pem, &key_pem)?; - let expires = chrono::DateTime::::from(SystemTime::now() + CACHE_TTL); + let not_after = leaf_not_after(&chain_pem)?; + if SystemTime::now() >= not_after { + return Err(format!( + "issued podidentity client certificate already expired at {}", + chrono::DateTime::::from(not_after) + .to_rfc3339_opts(chrono::SecondsFormat::Secs, true) + )); + } + let expires = chrono::DateTime::::from(cache_expiry(SystemTime::now(), not_after)); let _ = save_cached( &context, namespace, @@ -96,6 +108,38 @@ pub fn client_config( Ok(Arc::new(config)) } +/// When credentials written at `now` for a certificate expiring at +/// `not_after` stop being reused: 30 minutes on, but never past a minute +/// before the certificate expires. +fn cache_expiry(now: SystemTime, not_after: SystemTime) -> SystemTime { + let ttl = now + CACHE_TTL; + let cert = not_after + .checked_sub(ROUTER_CERT_EXPIRY) + .unwrap_or(SystemTime::UNIX_EPOCH); + ttl.min(cert) +} + +/// Whether a cached certificate expiring at `not_after` may still be used. +fn usable(now: SystemTime, not_after: SystemTime) -> bool { + now + ROUTER_CERT_EXPIRY <= not_after +} + +/// The expiry of the leaf (first) certificate in a PEM chain. +fn leaf_not_after(chain_pem: &str) -> Result { + let leaf = CertificateDer::pem_slice_iter(chain_pem.as_bytes()) + .next() + .ok_or("certificate chain is empty")? + .map_err(|e| format!("parsing leaf certificate: {e}"))?; + let (_, cert) = x509_parser::parse_x509_certificate(&leaf) + .map_err(|e| format!("parsing leaf certificate: {e}"))?; + let secs = cert.validity().not_after.timestamp(); + Ok(if secs >= 0 { + SystemTime::UNIX_EPOCH + Duration::from_secs(secs as u64) + } else { + SystemTime::UNIX_EPOCH + }) +} + fn cache_path(context: &str, namespace: &str, sa: &str) -> Option { let key = sanitize_context(&format!("{context}-{namespace}-{sa}")); tunnel_dir() @@ -104,12 +148,22 @@ fn cache_path(context: &str, namespace: &str, sa: &str) -> Option Option> { - let data = std::fs::read(cache_path(context, namespace, sa)?).ok()?; + load_cached_from(&cache_path(context, namespace, sa)?) +} + +fn load_cached_from(path: &std::path::Path) -> Option> { + let data = std::fs::read(path).ok()?; let cached: CachedRouterTls = serde_json::from_slice(&data).ok()?; let expires = chrono::DateTime::parse_from_rfc3339(&cached.expires_at).ok()?; if SystemTime::now() > SystemTime::from(expires) { return None; } + if !usable( + SystemTime::now(), + leaf_not_after(&cached.cert_chain_pem).ok()?, + ) { + return None; + } client_config( &cached.trust_bundle_pem, &cached.cert_chain_pem, @@ -390,6 +444,80 @@ mod tests { assert!("89ab".contains(&u[19..20])); } + fn cert_expiring_at(not_after: SystemTime) -> String { + cert_for(&rcgen::KeyPair::generate().unwrap(), not_after) + } + + fn cert_for(key: &rcgen::KeyPair, not_after: SystemTime) -> String { + let mut params = rcgen::CertificateParams::new(vec!["client".into()]).unwrap(); + let not_after = chrono::DateTime::::from(not_after).timestamp(); + params.not_before = rcgen::date_time_ymd(2000, 1, 1); + params.not_after = time::OffsetDateTime::from_unix_timestamp(not_after).unwrap(); + params.self_signed(key).unwrap().pem() + } + + #[test] + fn reads_leaf_expiry() { + let at = SystemTime::UNIX_EPOCH + Duration::from_secs(2_000_000_000); + assert_eq!(leaf_not_after(&cert_expiring_at(at)).unwrap(), at); + assert!(leaf_not_after("").is_err()); + assert!( + leaf_not_after("-----BEGIN CERTIFICATE-----\nAAAA\n-----END CERTIFICATE-----\n") + .is_err() + ); + } + + #[test] + fn caches_until_the_certificate_expires() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000); + let min = Duration::from_secs(60); + // A long-lived certificate is cached for 30 minutes. + assert_eq!(cache_expiry(now, now + 60 * min), now + CACHE_TTL); + // A short-lived one until a minute before it expires. + assert_eq!(cache_expiry(now, now + 10 * min), now + 9 * min); + assert!(cache_expiry(now, now + min / 2) < now); + } + + #[test] + fn ignores_cached_certificates_near_expiry() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000_000); + let min = Duration::from_secs(60); + assert!(usable(now, now + 10 * min)); + assert!(usable(now, now + min)); + // Within the margin, or expired. + assert!(!usable(now, now + min / 2)); + assert!(!usable(now, now - min)); + } + + #[test] + fn rejects_expired_cache_entries() { + let dir = tempfile::tempdir().unwrap(); + let ca_key = rcgen::KeyPair::generate().unwrap(); + let ca = rcgen::CertificateParams::new(vec![ROUTER_SERVER_NAME.into()]) + .unwrap() + .self_signed(&ca_key) + .unwrap(); + let key = rcgen::KeyPair::generate().unwrap(); + for (not_after, want) in [ + (SystemTime::now() + Duration::from_secs(3600), true), + (SystemTime::now() + Duration::from_secs(30), false), + (SystemTime::now() - Duration::from_secs(3600), false), + ] { + let cached = CachedRouterTls { + trust_bundle_pem: ca.pem(), + cert_chain_pem: cert_for(&key, not_after), + key_pem: key.serialize_pem(), + // The cache entry itself has not expired; only the + // certificate may have. + expires_at: chrono::DateTime::::from(SystemTime::now() + CACHE_TTL) + .to_rfc3339_opts(chrono::SecondsFormat::AutoSi, true), + }; + let path = dir.path().join("cache.json"); + std::fs::write(&path, serde_json::to_vec(&cached).unwrap()).unwrap(); + assert_eq!(load_cached_from(&path).is_some(), want, "{not_after:?}"); + } + } + #[test] fn builds_client_configs() { let ca_key = rcgen::KeyPair::generate().unwrap(); diff --git a/crates/rax-cli/tests/e2e_local.rs b/crates/rax-cli/tests/e2e_local.rs index 29516f9..74fa945 100644 --- a/crates/rax-cli/tests/e2e_local.rs +++ b/crates/rax-cli/tests/e2e_local.rs @@ -566,14 +566,14 @@ async fn task_lifecycle_and_substrate_sync() { .as_ref() .unwrap(); assert_eq!((probe.path.as_str(), probe.port), ("/readyz", 80)); - assert!( - !tmpl - .snapshot_config - .as_ref() - .unwrap() - .storage_location - .is_empty() + // A SnapshotConfig at volume fidelity, with no template-level + // location so Substrate 1.0 uses its cluster-wide one. + let snap = tmpl.snapshot_config.as_ref().unwrap(); + assert_eq!( + snap.preferred_fidelity, + rax_substrate::ateapi::SnapshotFidelity::Volumes as i32 ); + assert!(snap.storage_location.is_empty()); tmpl_name }; diff --git a/crates/rax-server/src/service.rs b/crates/rax-server/src/service.rs index 88b1705..dcbe42e 100644 --- a/crates/rax-server/src/service.rs +++ b/crates/rax-server/src/service.rs @@ -159,6 +159,31 @@ impl AxService { Err(e) => return Err(Status::internal(format!("checking existing task: {e}"))), } + // Tasks are immutable, so a task created before its Workspace would + // keep an empty workspace for its whole life. Reject it instead. + let workspace = match task + .spec + .as_ref() + .and_then(|s| s.workspace.as_ref()) + .map(|w| w.name.as_str()) + .filter(|n| !n.is_empty()) + { + None => None, + Some(ws_name) => match self.resources.get_workspace(atespace, ws_name).await { + Ok(ws) => Some(ws), + Err(StoreError::NotFound) => { + return Err(Status::failed_precondition(format!( + "workspace {ws_name:?} not found in atespace {atespace:?}" + ))); + } + Err(e) => { + return Err(Status::internal(format!( + "getting workspace {ws_name:?}: {e}" + ))); + } + }, + }; + let meta = task.metadata.as_mut().expect("metadata was defaulted"); if meta.create_time.is_none() { meta.create_time = Some(now()); @@ -173,8 +198,7 @@ impl AxService { // New tasks start suspended. match &self.reconciler { Some(r) => { - let ws = self.fetch_workspace(atespace, &task).await; - if let Err(e) = r.suspend(&mut task, ws).await { + if let Err(e) = r.suspend(&mut task, workspace).await { return Err(self.record_failure(&mut task, "provisioning", &e).await); } // The actor exists, so reads take the task's phase from it diff --git a/crates/rax-server/tests/service.rs b/crates/rax-server/tests/service.rs index ce524fb..0caf71f 100644 --- a/crates/rax-server/tests/service.rs +++ b/crates/rax-server/tests/service.rs @@ -69,6 +69,8 @@ struct FakeState { on_delete: Option>, /// The state of each task's actor, by task name. actors: HashMap, + /// The workspace the last suspend or resume was given. + workspace: Option, } impl FakeReconciler { @@ -122,11 +124,13 @@ impl FakeReconciler { #[async_trait] impl Reconciler for FakeReconciler { - async fn suspend(&self, task: &mut Task, _ws: Option) -> Result<(), Status> { + async fn suspend(&self, task: &mut Task, ws: Option) -> Result<(), Status> { + self.state.lock().unwrap().workspace = ws; self.drive(task, ActorState::Suspended) } - async fn resume(&self, task: &mut Task, _ws: Option) -> Result<(), Status> { + async fn resume(&self, task: &mut Task, ws: Option) -> Result<(), Status> { + self.state.lock().unwrap().workspace = ws; self.drive(task, ActorState::Running) } @@ -493,6 +497,14 @@ async fn validates_workspace_binding() { .await .unwrap_err(); assert_eq!(err.code(), Code::InvalidArgument); + svc.update_workspace(UpdateWorkspaceRequest { + workspace: Some(Workspace { + metadata: meta("a"), + ..Default::default() + }), + }) + .await + .unwrap(); svc.create_task(bind( "good", WorkspaceRef { @@ -610,6 +622,8 @@ async fn delete_keeps_task_when_cleanup_fails() { struct FailingStore { inner: MemoryStore, fail_updates: AtomicBool, + /// What every workspace lookup fails with, if anything. + get_workspace_err: Mutex>, } #[async_trait] @@ -640,6 +654,9 @@ impl ResourceStore for FailingStore { self.inner.save_workspace(ws).await } async fn get_workspace(&self, a: &str, n: &str) -> rax_store::Result { + if let Some(e) = self.get_workspace_err.lock().unwrap().take() { + return Err(e); + } self.inner.get_workspace(a, n).await } async fn list_workspaces(&self, a: &str) -> rax_store::Result> { @@ -662,6 +679,151 @@ impl ResourceStore for FailingStore { } } +/// A task is immutable, so CreateTask must not save or provision a task whose +/// Workspace does not exist yet: it would run with an empty workspace forever. +#[tokio::test] +async fn create_task_requires_bound_workspace() { + struct Case { + name: &'static str, + task_atespace: &'static str, + /// Where the "data" Workspace exists; empty for nowhere. + ws_atespace: &'static str, + binding: &'static str, + get_err: Option, + want_code: Code, + want_msg: &'static str, + } + let base = || Case { + name: "", + task_atespace: "", + ws_atespace: "", + binding: "data", + get_err: None, + want_code: Code::Ok, + want_msg: "", + }; + for tc in [ + Case { + name: "missing", + want_code: Code::FailedPrecondition, + want_msg: r#"workspace "data" not found in atespace "default""#, + ..base() + }, + Case { + name: "other atespace", + task_atespace: "team-a", + ws_atespace: "team-b", + want_code: Code::FailedPrecondition, + want_msg: r#"workspace "data" not found in atespace "team-a""#, + ..base() + }, + Case { + name: "store error", + get_err: Some(StoreError::Backend("store unavailable".into())), + want_code: Code::Internal, + want_msg: r#"getting workspace "data": store unavailable"#, + ..base() + }, + Case { + name: "present", + ws_atespace: "default", + ..base() + }, + Case { + name: "present in task atespace", + task_atespace: "team-a", + ws_atespace: "team-a", + ..base() + }, + Case { + name: "no binding", + binding: "", + ..base() + }, + ] { + let store = Arc::new(FailingStore::default()); + *store.get_workspace_err.lock().unwrap() = tc.get_err; + if !tc.ws_atespace.is_empty() { + store + .inner + .save_workspace(&mut Workspace { + metadata: Some(ObjectMeta { + atespace: tc.ws_atespace.into(), + name: "data".into(), + ..Default::default() + }), + spec: Some(WorkspaceSpec::default()), + ..Default::default() + }) + .await + .unwrap(); + } + let rec = Arc::new(FakeReconciler::default()); + let svc = AxService::new(store.clone(), on(rec.clone())); + let task = Task { + metadata: Some(ObjectMeta { + atespace: tc.task_atespace.into(), + name: "task".into(), + ..Default::default() + }), + spec: Some(TaskSpec { + image: "alpine".into(), + workspace: (!tc.binding.is_empty()).then(|| WorkspaceRef { + name: tc.binding.into(), + path: String::new(), + }), + ..Default::default() + }), + ..Default::default() + }; + let result = svc + .create_task(CreateTaskRequest { task: Some(task) }) + .await; + let (code, msg) = match &result { + Ok(_) => (Code::Ok, String::new()), + Err(e) => (e.code(), e.message().to_owned()), + }; + assert_eq!(code, tc.want_code, "{}: {msg}", tc.name); + assert!(msg.contains(tc.want_msg), "{}: {msg:?}", tc.name); + + let atespace = if tc.task_atespace.is_empty() { + "default" + } else { + tc.task_atespace + }; + let stored = store.inner.get_task(atespace, "task").await; + if tc.want_code != Code::Ok { + assert!( + matches!(stored, Err(StoreError::NotFound)), + "{}: rejected task was stored", + tc.name + ); + assert_eq!( + rec.counts(), + (0, 0), + "{}: rejected task was provisioned", + tc.name + ); + continue; + } + stored.unwrap_or_else(|e| panic!("{}: stored task: {e}", tc.name)); + assert_eq!( + rec.counts(), + (1, 0), + "{}: accepted task not provisioned", + tc.name + ); + // The reconciler gets the workspace the server resolved. + let ws = rec.state.lock().unwrap().workspace.clone(); + assert_eq!( + ws.and_then(|w| w.metadata).map(|m| m.name), + (!tc.binding.is_empty()).then(|| "data".to_owned()), + "{}", + tc.name + ); + } +} + #[tokio::test] async fn delete_stops_when_status_update_fails() { let store = Arc::new(FailingStore::default()); diff --git a/crates/rax-substrate/build.rs b/crates/rax-substrate/build.rs index c66c807..048009d 100644 --- a/crates/rax-substrate/build.rs +++ b/crates/rax-substrate/build.rs @@ -1,10 +1,10 @@ -//! Generates the Agent Substrate control API client from proto/ateapi/ateapi.proto. +//! Generates the Agent Substrate control API client from proto/ateapi/v1/ateapi.proto. use std::path::PathBuf; fn main() -> Result<(), Box> { let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../proto"); - let proto = root.join("ateapi/ateapi.proto"); + let proto = root.join("ateapi/v1/ateapi.proto"); println!("cargo:rerun-if-changed={}", proto.display()); let fds = protox::compile([&proto], [&root])?; diff --git a/crates/rax-substrate/src/lib.rs b/crates/rax-substrate/src/lib.rs index 8b36cdd..3dd5c94 100644 --- a/crates/rax-substrate/src/lib.rs +++ b/crates/rax-substrate/src/lib.rs @@ -10,16 +10,17 @@ pub mod testing; pub mod tls; pub mod ateapi { - //! Types and services generated from Substrate's `ateapi.proto`. + //! Types and services generated from Substrate's `api/ateapi/v1/ateapi.proto` + //! (package `substrate.ateapi.v1`). #![allow(clippy::all)] - tonic::include_proto!("ateapi"); + tonic::include_proto!("substrate.ateapi.v1"); } pub use client::{Client, ClientOptions, DEFAULT_TARGET, TokenInterceptor, worker_ip}; pub use retry::RetryPolicy; pub use template::{ - DEFAULT_CPU_MILLICORES, DEFAULT_GUEST_COMMAND, DEFAULT_MEMORY_BYTES, DEFAULT_SNAPSHOTS_BUCKET, - build_actor_template, resource_limits, + DEFAULT_CPU_MILLICORES, DEFAULT_GUEST_COMMAND, DEFAULT_MEMORY_BYTES, build_actor_template, + resource_limits, }; use tonic::{Code, Status}; diff --git a/crates/rax-substrate/src/template.rs b/crates/rax-substrate/src/template.rs index f3b6758..501cd61 100644 --- a/crates/rax-substrate/src/template.rs +++ b/crates/rax-substrate/src/template.rs @@ -13,9 +13,6 @@ use crate::ateapi::{ /// The command the guest container runs when a template names none. pub const DEFAULT_GUEST_COMMAND: &str = "/usr/local/bin/ax-task-runner"; -/// Where snapshots go when neither the caller nor `AX_SNAPSHOTS_BUCKET` names -/// a location. -pub const DEFAULT_SNAPSHOTS_BUCKET: &str = "gs://dberkov-gke-dev3/ate-env/"; /// The CPU limit (1 core) for a task that sets none. pub const DEFAULT_CPU_MILLICORES: i64 = 1000; /// The memory limit (512 MiB) for a task that sets none. @@ -43,28 +40,23 @@ pub fn resource_limits(reqs: Option<&ResourcesSpec>) -> Option) -> String { - if !explicit.is_empty() { - return explicit.to_owned(); - } - env.filter(|b| !b.is_empty()) - .unwrap_or_else(|| DEFAULT_SNAPSHOTS_BUCKET.to_owned()) -} - /// Builds the template every task actor runs from: one `guest` container /// with a `/readyz` wakeup probe and a durable `/workspace` volume, /// snapshotted at volume fidelity. Empty arguments take their defaults, and /// missing resource limits take [`DEFAULT_CPU_MILLICORES`] and /// [`DEFAULT_MEMORY_BYTES`]. +/// +/// `snapshots_bucket` sets the template's own snapshot storage location. +/// Substrate rejects a non-empty one (`storage_location: Forbidden`) unless +/// the cluster enables the `PerTemplateSnapshotLocation` preview feature +/// gate; leaving it empty uses Substrate's cluster-wide snapshot location. pub fn build_actor_template( atespace: &str, name: &str, image: &str, env: &BTreeMap, command: &[String], - snapshots_bucket_location: &str, + snapshots_bucket: &str, resources: Option<&ResourceRequirements>, ) -> ActorTemplate { let atespace = if atespace.is_empty() { @@ -82,11 +74,6 @@ pub fn build_actor_template( } else { command.to_vec() }; - let bucket = snapshots_bucket( - snapshots_bucket_location, - std::env::var("AX_SNAPSHOTS_BUCKET").ok(), - ); - let limits = resources.and_then(|r| r.limits.as_ref()); let cpu_millicores = limits .map(|l| l.cpu_millicores) @@ -133,7 +120,7 @@ pub fn build_actor_template( ..Default::default() }], snapshot_config: Some(SnapshotConfig { - storage_location: bucket, + storage_location: snapshots_bucket.to_owned(), preferred_fidelity: SnapshotFidelity::Volumes as i32, ..Default::default() }), @@ -155,15 +142,7 @@ mod tests { #[test] fn builds_the_task_template() { let env = BTreeMap::from([("FOO".to_owned(), "BAR".to_owned())]); - let tmpl = build_actor_template( - "", - "test-tmpl", - "", - &env, - &[], - "gs://ate-snapshots/ax/", - None, - ); + let tmpl = build_actor_template("", "test-tmpl", "", &env, &[], "", None); let meta = tmpl.metadata.as_ref().unwrap(); assert_eq!(meta.atespace, "default"); assert_eq!(meta.name, "test-tmpl"); @@ -184,22 +163,24 @@ mod tests { assert_eq!(c.volume_mounts[0].mount_path, "/workspace"); assert!(tmpl.volumes[0].durable_dir.is_some()); let snap = tmpl.snapshot_config.as_ref().unwrap(); - assert_eq!(snap.storage_location, "gs://ate-snapshots/ax/"); + // Empty by default, so Substrate 1.0 uses its cluster-wide location. + assert_eq!(snap.storage_location, ""); assert_eq!(snap.preferred_fidelity, SnapshotFidelity::Volumes as i32); - } - #[test] - fn snapshot_location_falls_back_to_env_then_default() { - assert_eq!( - snapshots_bucket("gs://x/", Some("gs://env/".into())), - "gs://x/" + // An explicit bucket is the template's location. + let tmpl = build_actor_template( + "", + "test-tmpl", + "", + &BTreeMap::new(), + &[], + "gs://ate-snapshots/ax/", + None, ); - assert_eq!(snapshots_bucket("", Some("gs://env/".into())), "gs://env/"); assert_eq!( - snapshots_bucket("", Some(String::new())), - DEFAULT_SNAPSHOTS_BUCKET + tmpl.snapshot_config.unwrap().storage_location, + "gs://ate-snapshots/ax/" ); - assert_eq!(snapshots_bucket("", None), DEFAULT_SNAPSHOTS_BUCKET); } fn limits(cpu: i64, memory: i64) -> ResourceRequirements { diff --git a/docs/local-kubernetes.md b/docs/local-kubernetes.md index 571a706..3273663 100644 --- a/docs/local-kubernetes.md +++ b/docs/local-kubernetes.md @@ -11,7 +11,8 @@ the control plane and the CLI are RAX's, from this checkout; `RUNNER=ax`, `ax`, one at a time or together. > **Status: not yet run end to end.** The scripts were written against AX's -> kind setup and Substrate v0.6.0's installer and are statically checked +> kind setup and Substrate's installer (v0.6.0's, then v1.0.0-rc.1's +> `ate-setup`) and are statically checked > (shellcheck, `just --dry-run`), but no cluster has been brought up with > them yet: the environment they were written in cannot give a nested > cluster network access. Expect to hit rough edges on the first run, and see @@ -122,8 +123,9 @@ manifests from `test/e2e/manifests` unchanged: the port (`--ports=127.0.0.1:5001:5001`). A `kind-registry` Service in `ate-system` stands in for kind's registry container, which Substrate's atelet pulls task images from. -2. **Substrate.** Substrate's `hack/install-ate-kind.sh deploy ate-system` - with AX's flags, pointed at this cluster's context. Substrate's kind +2. **Substrate.** Substrate's `ate-setup --kind deploy ate-system` (or + `hack/install-ate-kind.sh` on versions before 1.0) with AX's flags, + pointed at this cluster's context. Substrate's kind profile fits any single local node: in-cluster Postgres, RustFS for snapshots, and no cloud credential provider. The installed commit is recorded on the `ate-system` namespace, and a different one is uninstalled @@ -181,9 +183,11 @@ Use `SERVER=ax`, `RUNNER=ax` or `CLI=ax` to tell a RAX bug from a setup problem. cluster without them. `ax-server`'s manifest mounts a `podCertificate` and a `clusterTrustBundle` projected volume, so the server cannot start without them either. Kubernetes 1.34 or newer is needed; v1.37.0 is the default. -- **A Substrate version with `hack/install-ate-kind.sh`.** The version AX - pins (v0.6.0 at the time of writing) has it; Substrate's `main` branch no - longer does, so `SUBSTRATE_REF=main` does not work here. +- **A Substrate version AX's kind setup supports.** Substrate v1.0 installs + with `go run ./cmd/ate-setup --kind`, which AX's `hack/setup-e2e-kind.sh` + now uses and so does `deploy` (v1.0.0-rc.1 is the version AX pins at the + time of writing); older versions with `hack/install-ate-kind.sh` still + work. Either way Go is needed on the host to run the installer. - **Docker on the host**, even for MicroK8s: the images are built with it, and Substrate's installer checks `docker info` in its kind profile. - **gVisor only.** Micro-VM sandbox classes (kata with cloud-hypervisor) need diff --git a/hack/local/local.sh b/hack/local/local.sh index 0ee7663..c50abbf 100755 --- a/hack/local/local.sh +++ b/hack/local/local.sh @@ -64,8 +64,6 @@ AX_REF="${AX_REF:-main}" # AX's ax-server manifest and Substrate's bootstrap owners both assume this. export AX_SYSTEM_NAMESPACE=ax-system export AX_WORKER_REPLICAS="${AX_WORKER_REPLICAS:-2}" -# Substrate's kind profile serves snapshots from in-cluster RustFS at this bucket. -export AX_SNAPSHOTS_BUCKET="${AX_SNAPSHOTS_BUCKET:-gs://ate-snapshots/ax/}" # The registry in registry.yaml, as the node and the host both see it. REGISTRY_ADDR=localhost:5001 @@ -349,14 +347,14 @@ resolve_substrate_dir() { git -C "${SUBSTRATE_DIR}" checkout --quiet --detach "${ref}" fi fi - [[ -x "${SUBSTRATE_DIR}/hack/install-ate-kind.sh" ]] || - die "${SUBSTRATE_DIR} has no hack/install-ate-kind.sh; use a Substrate version AX supports" + [[ -x "${SUBSTRATE_DIR}/hack/install-ate-kind.sh" || -d "${SUBSTRATE_DIR}/cmd/ate-setup" ]] || + die "${SUBSTRATE_DIR} has neither hack/install-ate-kind.sh nor cmd/ate-setup; use a Substrate version AX supports" SUBSTRATE_COMMIT="$(git -C "${SUBSTRATE_DIR}" rev-parse HEAD)" log "Using Substrate at ${SUBSTRATE_DIR} (${SUBSTRATE_COMMIT:0:12})" } # render substitutes only the AX_* variables in one of AX's E2E manifests. -MANIFEST_VARS=(AX_SYSTEM_NAMESPACE AX_WORKER_REPLICAS AX_SERVER_IMAGE AX_DEFAULT_TASK_IMAGE AX_SNAPSHOTS_BUCKET) +MANIFEST_VARS=(AX_SYSTEM_NAMESPACE AX_WORKER_REPLICAS AX_SERVER_IMAGE AX_DEFAULT_TASK_IMAGE) render() { local manifest="${AX_DIR}/test/e2e/manifests/$1" sed_args=() var for var in "${MANIFEST_VARS[@]}"; do @@ -436,9 +434,16 @@ deploy_substrate() { # RustFS for snapshots, no cloud credential provider. Its only kind-specific # assumptions are the kind-registry name, which registry.yaml provides, and # a reachable Docker daemon, which builds the images anyway. + # Substrate 1.0 dropped hack/install-ate-kind.sh for `ate-setup --kind`, + # which AX's kind setup now runs; older versions still have the script. + local install_cmd=(./hack/install-ate-kind.sh) + if [[ ! -f "${SUBSTRATE_DIR}/hack/install-ate-kind.sh" ]]; then + install_cmd=(go run ./cmd/ate-setup --kind --context="${KUBE_CONTEXT}") + fi ( cd "${SUBSTRATE_DIR}" - KUBECTL_CONTEXT="${KUBE_CONTEXT}" KO_DOCKER_REPO="${REGISTRY_ADDR}" ./hack/install-ate-kind.sh \ + KUBECTL_CONTEXT="${KUBE_CONTEXT}" KO_DOCKER_REPO="${REGISTRY_ADDR}" \ + KO_DEFAULTPLATFORMS="linux/${NODE_ARCH}" "${install_cmd[@]}" \ deploy ate-system \ --rollout-timeout="${SUBSTRATE_ROLLOUT_TIMEOUT:-5m}" \ --credential-provider='{"enabled":false}' \ diff --git a/proto/ateapi/ateapi.proto b/proto/ateapi/v1/ateapi.proto similarity index 98% rename from proto/ateapi/ateapi.proto rename to proto/ateapi/v1/ateapi.proto index e453a31..d997c8d 100644 --- a/proto/ateapi/ateapi.proto +++ b/proto/ateapi/v1/ateapi.proto @@ -14,12 +14,12 @@ syntax = "proto3"; -package ateapi; +package substrate.ateapi.v1; import "google/protobuf/descriptor.proto"; import "google/protobuf/timestamp.proto"; -option go_package = "go.ate.dev/substrate/pkg/proto/ateapipb"; +option go_package = "go.ate.dev/substrate/api/ateapi/v1;ateapi"; // Control is the primary RPC interface for Agentic Substrate. service Control { @@ -90,16 +90,6 @@ service Control { option (method).resource = "Actor"; } - // Create a Substrate-issued SPIFFE certificate asserting the actor identity. - // - // SPIFFE URI: spiffe://${trustdomain}/actor/${atespace}/${actor} - // - // * Called by the egress gateway when actor client certificate injection is - // configured for outbound requests. - rpc MintActorCertificate(MintActorCertificateRequest) returns (MintActorCertificateResponse) { - option (method).resource = "Actor"; - } - // Tag the external snapshot a suspended Actor holds. The tag gets its own // copy of that snapshot, so suspending or deleting the Actor afterwards // cannot collect it. When the Actor's latest snapshot is not uploaded yet, @@ -291,24 +281,24 @@ enum SnapshotOwner { // Keep this in sync with Snapshot.owner's maximum. } -// SnapshotStorageStatus indicates the lifecycle stage of a durable snapshot. -enum SnapshotStorageStatus { - SNAPSHOT_STORAGE_STATUS_UNSPECIFIED = 0; +// SnapshotStorageState indicates the lifecycle stage of a durable snapshot. +enum SnapshotStorageState { + SNAPSHOT_STORAGE_STATE_UNSPECIFIED = 0; // Snapshot creation in durable storage is in progress. - SNAPSHOT_STORAGE_STATUS_PENDING = 1; + SNAPSHOT_STORAGE_STATE_PENDING = 1; // Snapshot is complete and ready for use in durable storage. - SNAPSHOT_STORAGE_STATUS_COMPLETED = 2; - // Keep this in sync with SnapshotStorage.status's maximum. + SNAPSHOT_STORAGE_STATE_COMPLETED = 2; + // Keep this in sync with SnapshotStorage.state's maximum. } // SnapshotStorage records the durable storage representation of a Snapshot. message SnapshotStorage { - // status is the lifecycle status of the durable snapshot. + // state is the lifecycle state of the durable snapshot. // // +k8s:required // +k8s:minimum=1 - // +k8s:maximum=2 # keep this in sync with the SnapshotStorageStatus enum - SnapshotStorageStatus status = 1; + // +k8s:maximum=2 # keep this in sync with the SnapshotStorageState enum + SnapshotStorageState state = 1; // object snapshot is stored in object storage. // @@ -1279,9 +1269,10 @@ message SnapshotConfig { // storage_location is the base object-storage URI snapshots of actors on // this version are stored under. If unset, ate-api-server's // --snapshot-location is used. Requires the PerTemplateSnapshotLocation preview - // gate; ignored without it. + // gate. // // +ate:preview=PerTemplateSnapshotLocation + // +k8s:ifDisabled(PerTemplateSnapshotLocation)=+k8s:forbidden // +k8s:optional // +k8s:maxLength=1024 // +k8s:customValidation # Validate URI @@ -2012,39 +2003,6 @@ message MintActorJwtResponse { google.protobuf.Timestamp expire_time = 2; } -// Request message for MintActorCertificate. -message MintActorCertificateRequest { - // The actor for which the certificate should be issued. - // - // Must be a valid actor that currently exists according to the actor store. - // - // +k8s:required - ObjectRef actor = 1; - - // The UID of the actor --- used to guard against deletion and recreation of - // an actor with the same name. - // - // +k8s:required - // +k8s:format=k8s-uuid - string actor_uid = 2; - - // Request contains DER encoded bytes of a x509 certificate signing request. - // The signer will ignore the contents of the CSR except to extract the - // subject public key. - // - // +k8s:required - // +k8s:maxBytes=16384 - bytes certificate_signing_request = 3; -} - -// Response message for MintActorCertificate. -message MintActorCertificateResponse { - // Response contains a list of DER encoded certificates. The first entry is the - // leaf certificate, and any remaining entries are intermediates in - // leaf-to-root order. - repeated bytes actor_certificates = 1; -} - message ListTagsRequest { // The atespace to list tags from. Empty lists across all atespaces. // @@ -2409,7 +2367,8 @@ message WorkerStatus { WorkerResources capacity = 2; // The allocated resources of the Worker. This is the sum of all the - // resources consumed by the Actors currently assigned to this Worker. + // resources consumed by the Actors currently assigned to this Worker, + // including their runtime's actor_overhead. // // +k8s:optional WorkerResources allocated = 3; @@ -2505,10 +2464,39 @@ message WorkerSandboxRuntime { // +k8s:required VersionedSandboxCompat version = 1; + // The resource overhead incurred by a Worker for each Actor running under + // this runtime, on top of the Actor's stated limits. Placement considers + // this when deciding if a Worker has enough capacity for an Actor, and it is + // accounted for in allocated. If not specified, Actors under this runtime + // are assumed to have negligible overhead. + // + // +k8s:optional + RuntimeOverhead actor_overhead = 2; + // TODO: Add an informational SandboxConfig version name (e.g. "gvisor-2") // once workers report one. } +// RuntimeOverhead is the memory an Actor uses on its Worker outside its own +// limits: a fixed amount, plus an amount in proportion to its memory limit. +message RuntimeOverhead { + // memory_bytes is a fixed amount of memory for each Actor. + // + // +k8s:optional + // +k8s:minimum=0 + // +k8s:maximum=1099511627776 # 1TiB + int64 memory_bytes = 1; + + // memory_bytes_per_gib is memory per GiB of the Actor's memory limit, e.g. + // host page tables for a microVM's guest RAM. It scales in proportion and + // rounds down to a byte, so an Actor with a 512MiB limit adds half of it. + // + // +k8s:optional + // +k8s:minimum=0 + // +k8s:maximum=1073741824 # 1GiB + int64 memory_bytes_per_gib = 2; +} + // ActorSandboxRuntime is the runtime an actor runs under, as recorded on its // assignment and on the snapshots taken from it: one class and one of the // compatibility versions a Worker reports for it. @@ -2630,7 +2618,8 @@ message ActorAssignment { // +k8s:format=k8s-uuid string actor_uid = 3; - // The resources allocated to the Actor when it was assigned. + // The resources allocated to the Actor when it was assigned: its limits, + // plus the actor_overhead of the runtime it runs under on the Worker. // // +k8s:optional Resources resources = 4;