From e2853b4e13b3426aa032df8fd9e497d4f32144b8 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Mon, 10 Aug 2026 23:39:20 -0400 Subject: [PATCH 01/11] =?UTF-8?q?feat(cp):=20openab-cp=20control=20plane?= =?UTF-8?q?=20=E2=80=94=20registry,=20router,=20policy=20(ADR=20v1,=20PR?= =?UTF-8?q?=201/4)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the first slice of docs/adr/agent-control-plane.md: the standalone openab-cp binary with wire protocol, identity-bound registry, CP-authoritative policy engine, and delegation router. Addresses round-1 review findings: CP-generated registration handles for all ownership (F1), atomic delegation admission with insert-before-send (F2), parent delegation bound to its serving instance (F3), JSON-RPC 2.0 envelope validation (F4), transport/prompt/queue resource bounds (F5), lease-only heartbeats + least-loaded label scheduling (F6), and unrelated artifacts removed from the diff (F7). --- Cargo.lock | 20 + Cargo.toml | 2 +- Dockerfile | 8 +- Dockerfile.agentcore | 8 +- Dockerfile.antigravity | 8 +- Dockerfile.builder | 8 +- Dockerfile.claude | 8 +- Dockerfile.codex | 8 +- Dockerfile.copilot | 8 +- Dockerfile.cursor | 8 +- Dockerfile.devin | 8 +- Dockerfile.gateway | 8 +- Dockerfile.gemini | 8 +- Dockerfile.grok | 8 +- Dockerfile.hermes | 8 +- Dockerfile.kimi | 8 +- Dockerfile.mimocode | 8 +- Dockerfile.native | 8 +- Dockerfile.opencode | 8 +- Dockerfile.pi | 8 +- Dockerfile.unified | 8 +- crates/openab-cp/Cargo.toml | 22 + crates/openab-cp/cp.toml.example | 44 ++ crates/openab-cp/src/config.rs | 294 +++++++++ crates/openab-cp/src/lib.rs | 14 + crates/openab-cp/src/main.rs | 56 ++ crates/openab-cp/src/policy.rs | 246 ++++++++ crates/openab-cp/src/proto.rs | 426 +++++++++++++ crates/openab-cp/src/registry.rs | 350 +++++++++++ crates/openab-cp/src/router.rs | 992 +++++++++++++++++++++++++++++++ crates/openab-cp/src/server.rs | 597 +++++++++++++++++++ docs/adr/agent-control-plane.md | 85 ++- 32 files changed, 3235 insertions(+), 65 deletions(-) create mode 100644 crates/openab-cp/Cargo.toml create mode 100644 crates/openab-cp/cp.toml.example create mode 100644 crates/openab-cp/src/config.rs create mode 100644 crates/openab-cp/src/lib.rs create mode 100644 crates/openab-cp/src/main.rs create mode 100644 crates/openab-cp/src/policy.rs create mode 100644 crates/openab-cp/src/proto.rs create mode 100644 crates/openab-cp/src/registry.rs create mode 100644 crates/openab-cp/src/router.rs create mode 100644 crates/openab-cp/src/server.rs diff --git a/Cargo.lock b/Cargo.lock index 89f00d9c7..98dfb7d52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2575,6 +2575,26 @@ dependencies = [ "zip", ] +[[package]] +name = "openab-cp" +version = "0.1.0" +dependencies = [ + "anyhow", + "axum", + "chrono", + "clap", + "futures-util", + "parking_lot", + "serde", + "serde_json", + "subtle", + "tokio", + "toml", + "tracing", + "tracing-subscriber", + "uuid", +] + [[package]] name = "openab-gateway" version = "0.5.4" diff --git a/Cargo.toml b/Cargo.toml index 5a834a13f..191c2640f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["crates/openab-core", "crates/openab-gateway", "crates/openab-mcp"] +members = ["crates/openab-core", "crates/openab-gateway", "crates/openab-mcp", "crates/openab-cp"] exclude = ["openab-agent", "crates/platform-schema"] [package] diff --git a/Dockerfile b/Dockerfile index 36bf49ca5..540ba829f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,16 +11,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && \ if [ "$BUILD_MODE" = "unified" ]; then \ cargo build --release --features unified; \ elif [ -n "$FEATURES" ]; then \ diff --git a/Dockerfile.agentcore b/Dockerfile.agentcore index 7bdd6d70b..d2b153fc9 100644 --- a/Dockerfile.agentcore +++ b/Dockerfile.agentcore @@ -9,16 +9,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release --features agentcore \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release --features agentcore +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release --features agentcore # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.antigravity b/Dockerfile.antigravity index 77aef681e..2003d59f2 100644 --- a/Dockerfile.antigravity +++ b/Dockerfile.antigravity @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Build agy-acp adapter --- FROM rust:1-bookworm AS adapter-builder diff --git a/Dockerfile.builder b/Dockerfile.builder index ff344e5e4..c3122c707 100644 --- a/Dockerfile.builder +++ b/Dockerfile.builder @@ -23,20 +23,22 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml # 2. Dummy sources for dep-only build -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src # 3. Copy real sources and build COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && \ if [ "$BUILD_MODE" = "unified" ]; then \ cargo build --release --features unified; \ elif [ -n "$FEATURES" ]; then \ diff --git a/Dockerfile.claude b/Dockerfile.claude index aa68fe3a2..8b47f8975 100644 --- a/Dockerfile.claude +++ b/Dockerfile.claude @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.codex b/Dockerfile.codex index 4d0f35c44..512e31506 100644 --- a/Dockerfile.codex +++ b/Dockerfile.codex @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.copilot b/Dockerfile.copilot index e28a29089..bcef23aca 100644 --- a/Dockerfile.copilot +++ b/Dockerfile.copilot @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.cursor b/Dockerfile.cursor index 03e9db910..e82d71ed2 100644 --- a/Dockerfile.cursor +++ b/Dockerfile.cursor @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.devin b/Dockerfile.devin index e25dadb60..b1f6c0160 100644 --- a/Dockerfile.devin +++ b/Dockerfile.devin @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.gateway b/Dockerfile.gateway index 7b3741ff1..cce6dde33 100644 --- a/Dockerfile.gateway +++ b/Dockerfile.gateway @@ -17,17 +17,19 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo 'fn main() {}' > crates/openab-gateway/src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release -p openab-gateway \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch crates/openab-gateway/src/main.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && \ +RUN touch crates/openab-gateway/src/main.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && \ if [ -n "$FEATURES" ]; then \ cargo build --release -p openab-gateway --no-default-features --features "$FEATURES"; \ else \ diff --git a/Dockerfile.gemini b/Dockerfile.gemini index 3de113f06..6ed4c3bf1 100644 --- a/Dockerfile.gemini +++ b/Dockerfile.gemini @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.grok b/Dockerfile.grok index 0708c293f..ad12d6210 100644 --- a/Dockerfile.grok +++ b/Dockerfile.grok @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.hermes b/Dockerfile.hermes index 2217dbe65..bdf0b8eea 100644 --- a/Dockerfile.hermes +++ b/Dockerfile.hermes @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM python:3.12-slim-trixie diff --git a/Dockerfile.kimi b/Dockerfile.kimi index 85ffc7ca5..3bcd61f37 100644 --- a/Dockerfile.kimi +++ b/Dockerfile.kimi @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs \ && cargo build --release # --- Runtime stage --- diff --git a/Dockerfile.mimocode b/Dockerfile.mimocode index bdb50e952..df8331324 100644 --- a/Dockerfile.mimocode +++ b/Dockerfile.mimocode @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- # MiMo-Code (https://github.com/XiaomiMiMo/MiMo-Code) is a fork of OpenCode diff --git a/Dockerfile.native b/Dockerfile.native index f66b268ba..5de7f6ea7 100644 --- a/Dockerfile.native +++ b/Dockerfile.native @@ -5,17 +5,19 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml COPY openab-agent/ openab-agent/ -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release RUN cd openab-agent && cargo build --release # --- Runtime stage --- diff --git a/Dockerfile.opencode b/Dockerfile.opencode index c1f239f5e..7a7063575 100644 --- a/Dockerfile.opencode +++ b/Dockerfile.opencode @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- # node:22-trixie-slim mirrors the base image used by Dockerfile.claude, diff --git a/Dockerfile.pi b/Dockerfile.pi index 0d992f4af..ae5fb99cc 100644 --- a/Dockerfile.pi +++ b/Dockerfile.pi @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.unified b/Dockerfile.unified index 0c734ebe2..540786a97 100644 --- a/Dockerfile.unified +++ b/Dockerfile.unified @@ -31,16 +31,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release --features unified \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs \ && cargo build --release --features unified # Build openab-agent (used by native variant) — copied late to avoid cache busts COPY openab-agent/ openab-agent/ diff --git a/crates/openab-cp/Cargo.toml b/crates/openab-cp/Cargo.toml new file mode 100644 index 000000000..eff123bff --- /dev/null +++ b/crates/openab-cp/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "openab-cp" +version = "0.1.0" +edition = "2021" +license = "MIT" +description = "OpenAB Agent Control Plane — registry, router, and policy for direct inter-agent delegation" + +[dependencies] +tokio = { version = "1", features = ["full"] } +axum = { version = "0.8", features = ["ws"] } +futures-util = "0.3" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +toml = "0.8" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +anyhow = "1" +uuid = { version = "1", features = ["v4"] } +chrono = { version = "0.4", features = ["serde"] } +parking_lot = "0.12" +clap = { version = "4", features = ["derive"] } +subtle = "2" diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example new file mode 100644 index 000000000..8373d2839 --- /dev/null +++ b/crates/openab-cp/cp.toml.example @@ -0,0 +1,44 @@ +# openab-cp example configuration +# +# Identity table: every runtime authenticates with a per-agent key +# (Authorization: Bearer on the WebSocket upgrade). The claims below +# are IMMUTABLE and owned by this file — a runtime's own [control_plane] +# config is verified against them at registration and rejected on mismatch. + +listen = "0.0.0.0:9800" + +# Runtimes must heartbeat at this interval; missing heartbeats past the lease +# window deregisters the instance and fails its in-flight delegations. +heartbeat_interval_secs = 15 +lease_expiry_secs = 45 + +# Hard cap on delegation deadlines (seconds from now). +max_deadline_secs = 1800 + +# Results larger than this are truncated (head kept) with a marker. +max_result_bytes = 262144 + +# Transport-level cap on inbound WebSocket messages (enforced pre-parse). +max_frame_bytes = 1048576 + +# Delegation prompts larger than this are rejected. +max_prompt_bytes = 262144 + +[[agents]] +key = "${CP_KEY_KOUDU}" # per-agent secret, never shared +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "${CP_KEY_WORKER1}" +namespace = "prod" +name = "worker-1" +type = "worker" +max_delegated_sessions_cap = 4 # CP-side clamp on advertised capacity + +# Per-namespace policy. Absent namespaces use the conservative defaults: +# max_depth = 1, allow_worker_initiation = false. +[namespaces.prod] +max_depth = 1 +allow_worker_initiation = false diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs new file mode 100644 index 000000000..c92a5988f --- /dev/null +++ b/crates/openab-cp/src/config.rs @@ -0,0 +1,294 @@ +//! CP-side configuration. +//! +//! Identity binding is the security core (review F1 on the ADR): every auth +//! key maps to **immutable claims** (`namespace`, `name`, `type`, optional +//! caps) owned by CP config. Registration frames are verified against these +//! claims — never the other way around. A compromised runtime cannot escalate +//! to another namespace or to `primary` by editing its own config. + +use anyhow::{bail, Context, Result}; +use serde::Deserialize; +use std::collections::BTreeMap; + +use crate::proto::AgentType; + +#[derive(Debug, Clone, Deserialize)] +pub struct CpConfig { + /// Bind address, e.g. "0.0.0.0:9800". + #[serde(default = "default_listen")] + pub listen: String, + + /// Heartbeat interval communicated to runtimes. + #[serde(default = "default_heartbeat_secs")] + pub heartbeat_interval_secs: u64, + + /// Lease window: an instance missing heartbeats past this is deregistered + /// and its in-flight delegations fail with `TARGET_DISCONNECTED`. + #[serde(default = "default_lease_secs")] + pub lease_expiry_secs: u64, + + /// Hard cap on delegation deadline length (seconds from now). Deadlines + /// beyond this are rejected at `cp/delegate`. + #[serde(default = "default_max_deadline_secs")] + pub max_deadline_secs: u64, + + /// Maximum result payload size in bytes (`cp/delegate_result.result`). + /// Oversized results are truncated with a marker, not rejected — the + /// delegation already ran; losing the tail beats losing everything. + #[serde(default = "default_max_result_bytes")] + pub max_result_bytes: usize, + + /// Maximum WebSocket message size accepted from a runtime, enforced by + /// the transport before any parsing/allocation (review F5). + #[serde(default = "default_max_frame_bytes")] + pub max_frame_bytes: usize, + + /// Maximum `cp/delegate.prompt` size in bytes; oversized prompts are + /// rejected (unlike results, nothing has run yet). + #[serde(default = "default_max_prompt_bytes")] + pub max_prompt_bytes: usize, + + /// Identity table: auth key → immutable claims. + /// Keyed by the key id (`kid`), with the secret alongside, so logs can + /// reference identities without printing secrets. + #[serde(default)] + pub agents: Vec, + + /// Per-namespace policy overrides. + #[serde(default)] + pub namespaces: BTreeMap, +} + +fn default_listen() -> String { + "0.0.0.0:9800".to_string() +} +fn default_heartbeat_secs() -> u64 { + 15 +} +fn default_lease_secs() -> u64 { + 45 +} +fn default_max_deadline_secs() -> u64 { + 30 * 60 +} +fn default_max_result_bytes() -> usize { + 256 * 1024 +} +fn default_max_frame_bytes() -> usize { + 1024 * 1024 +} +fn default_max_prompt_bytes() -> usize { + 256 * 1024 +} + +/// Immutable identity claims bound to one auth key. +#[derive(Debug, Clone, Deserialize)] +pub struct AgentIdentity { + /// The secret presented by the runtime (`OPENAB_CP_KEY`). Supports + /// `${ENV_VAR}` expansion so the config file itself holds no secrets. + pub key: String, + pub namespace: String, + pub name: String, + #[serde(rename = "type")] + pub agent_type: AgentType, + /// Optional CP-side clamp on the advertised concurrency budget. + #[serde(default)] + pub max_delegated_sessions_cap: Option, +} + +/// Per-namespace delegation policy. Defaults are the conservative ADR §5 +/// baseline; relaxation is CP-side config only. +#[derive(Debug, Clone, Deserialize)] +pub struct NamespacePolicy { + /// Maximum delegation chain depth (1 = primary → worker only). + #[serde(default = "default_depth")] + pub max_depth: u32, + /// Whether workers may initiate delegations (depth still applies). + #[serde(default)] + pub allow_worker_initiation: bool, +} + +fn default_depth() -> u32 { + 1 +} + +impl Default for NamespacePolicy { + fn default() -> Self { + Self { + max_depth: default_depth(), + allow_worker_initiation: false, + } + } +} + +impl CpConfig { + pub fn load(path: &str) -> Result { + let raw = + std::fs::read_to_string(path).with_context(|| format!("reading CP config {path}"))?; + let expanded = expand_env(&raw); + let cfg: CpConfig = toml::from_str(&expanded).context("parsing CP config")?; + cfg.validate()?; + Ok(cfg) + } + + pub fn validate(&self) -> Result<()> { + let mut seen_keys = std::collections::BTreeSet::new(); + let mut seen_names = std::collections::BTreeSet::new(); + for a in &self.agents { + if a.key.trim().is_empty() { + bail!("agent {}/{} has an empty key", a.namespace, a.name); + } + if !seen_keys.insert(a.key.as_str()) { + bail!( + "duplicate auth key (shared keys defeat per-agent revocation); \ + offending identity: {}/{}", + a.namespace, + a.name + ); + } + if !seen_names.insert((a.namespace.as_str(), a.name.as_str())) { + bail!( + "duplicate identity {}/{} — replicas share one identity (one key), \ + distinguished at registration by instance_id", + a.namespace, + a.name + ); + } + } + if self.lease_expiry_secs <= self.heartbeat_interval_secs { + bail!("lease_expiry_secs must exceed heartbeat_interval_secs"); + } + Ok(()) + } + + /// Constant-time lookup of the identity bound to `key`. + pub fn identity_for_key(&self, key: &str) -> Option<&AgentIdentity> { + use subtle::ConstantTimeEq; + // Compare against every entry to avoid early-exit timing signal on + // which identity matched. + let mut found: Option<&AgentIdentity> = None; + for a in &self.agents { + let eq: bool = a.key.as_bytes().ct_eq(key.as_bytes()).into(); + if eq { + found = Some(a); + } + } + found + } + + pub fn policy_for(&self, namespace: &str) -> NamespacePolicy { + self.namespaces.get(namespace).cloned().unwrap_or_default() + } +} + +/// `${ENV_VAR}` expansion, mirroring openab-core config behavior. Unset vars +/// expand to the empty string (validation then rejects empty keys loudly). +fn expand_env(raw: &str) -> String { + let mut out = String::with_capacity(raw.len()); + let mut rest = raw; + while let Some(start) = rest.find("${") { + out.push_str(&rest[..start]); + match rest[start + 2..].find('}') { + Some(end) => { + let var = &rest[start + 2..start + 2 + end]; + out.push_str(&std::env::var(var).unwrap_or_default()); + rest = &rest[start + 2 + end + 1..]; + } + None => { + out.push_str(&rest[start..]); + rest = ""; + } + } + } + out.push_str(rest); + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn base_toml() -> &'static str { + r#" +listen = "127.0.0.1:9800" + +[[agents]] +key = "k-primary" +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "k-worker" +namespace = "prod" +name = "worker-1" +type = "worker" +max_delegated_sessions_cap = 2 + +[namespaces.prod] +max_depth = 2 +allow_worker_initiation = false +"# + } + + #[test] + fn parses_and_validates() { + let cfg: CpConfig = toml::from_str(base_toml()).unwrap(); + cfg.validate().unwrap(); + assert_eq!(cfg.agents.len(), 2); + assert_eq!(cfg.policy_for("prod").max_depth, 2); + // unknown namespace falls back to conservative defaults + let d = cfg.policy_for("dev"); + assert_eq!(d.max_depth, 1); + assert!(!d.allow_worker_initiation); + } + + #[test] + fn identity_lookup_binds_key_to_claims() { + let cfg: CpConfig = toml::from_str(base_toml()).unwrap(); + let id = cfg.identity_for_key("k-worker").unwrap(); + assert_eq!(id.name, "worker-1"); + assert_eq!(id.agent_type, AgentType::Worker); + assert!(cfg.identity_for_key("k-unknown").is_none()); + } + + #[test] + fn rejects_duplicate_keys() { + let toml_str = r#" +[[agents]] +key = "same" +namespace = "prod" +name = "a" +type = "primary" + +[[agents]] +key = "same" +namespace = "prod" +name = "b" +type = "worker" +"#; + let cfg: CpConfig = toml::from_str(toml_str).unwrap(); + assert!(cfg.validate().is_err()); + } + + #[test] + fn rejects_lease_not_exceeding_heartbeat() { + let toml_str = r#" +heartbeat_interval_secs = 30 +lease_expiry_secs = 30 +"#; + let cfg: CpConfig = toml::from_str(toml_str).unwrap(); + assert!(cfg.validate().is_err()); + } + + #[test] + fn env_expansion() { + std::env::set_var("CP_TEST_KEY_XYZ", "sekrit"); + assert_eq!( + expand_env("key = \"${CP_TEST_KEY_XYZ}\""), + "key = \"sekrit\"" + ); + assert_eq!(expand_env("no vars"), "no vars"); + assert_eq!(expand_env("${UNSET_VAR_ABC123}"), ""); + } +} diff --git a/crates/openab-cp/src/lib.rs b/crates/openab-cp/src/lib.rs new file mode 100644 index 000000000..f4a6ac14f --- /dev/null +++ b/crates/openab-cp/src/lib.rs @@ -0,0 +1,14 @@ +//! openab-cp — OpenAB Agent Control Plane. +//! +//! Hub-and-spoke registration and routing for direct agent-to-agent +//! delegation, per `docs/adr/agent-control-plane.md`. Runtimes dial out and +//! register over WebSocket; the CP authenticates them against config-bound +//! identities, enforces delegation policy authoritatively, and routes +//! `cp/delegate` / `cp/delegate_result` frames between them. + +pub mod config; +pub mod policy; +pub mod proto; +pub mod registry; +pub mod router; +pub mod server; diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs new file mode 100644 index 000000000..7a34753c3 --- /dev/null +++ b/crates/openab-cp/src/main.rs @@ -0,0 +1,56 @@ +//! Standalone control-plane binary. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use clap::Parser; +use tracing::info; + +use openab_cp::config::CpConfig; +use openab_cp::server::{app, run_sweeper, AppState}; + +#[derive(Parser)] +#[command(name = "openab-cp", about = "OpenAB Agent Control Plane")] +struct Cli { + /// Path to the CP config file (TOML). + #[arg(short, long, default_value = "cp.toml")] + config: String, +} + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), + ) + .init(); + + let cli = Cli::parse(); + let cfg = CpConfig::load(&cli.config)?; + let listen = cfg.listen.clone(); + if cfg.agents.is_empty() { + tracing::warn!("no [[agents]] identities configured — every connection will be rejected"); + } + info!( + listen = %listen, + identities = cfg.agents.len(), + namespaces = cfg.namespaces.len(), + "starting openab-cp" + ); + + let state = Arc::new(AppState::new(cfg)); + let sweeper = tokio::spawn(run_sweeper(state.clone())); + + let listener = tokio::net::TcpListener::bind(&listen) + .await + .with_context(|| format!("binding {listen}"))?; + axum::serve(listener, app(state)) + .with_graceful_shutdown(async { + let _ = tokio::signal::ctrl_c().await; + info!("shutdown signal received"); + }) + .await?; + + sweeper.abort(); + Ok(()) +} diff --git a/crates/openab-cp/src/policy.rs b/crates/openab-cp/src/policy.rs new file mode 100644 index 000000000..c8b485ca5 --- /dev/null +++ b/crates/openab-cp/src/policy.rs @@ -0,0 +1,246 @@ +//! CP-authoritative delegation policy (review F4 on the ADR). +//! +//! Every check here operates exclusively on CP-owned data: authenticated +//! identity claims, the CP-constructed ancestry chain, and CP config. Facade +//! checks in the runtime are defense-in-depth only; nothing in this module +//! trusts a client-supplied policy input. + +use chrono::{DateTime, Utc}; + +use crate::config::NamespacePolicy; +use crate::proto::AgentType; + +pub struct PolicyInput<'a> { + /// Authenticated initiator identity. + pub from_namespace: &'a str, + pub from_name: &'a str, + pub from_type: &'a AgentType, + /// Target namespace (v1: always the initiator's namespace; the router + /// resolves selectors within it. Cross-namespace grants are future work). + pub target_namespace: &'a str, + /// Logical name of the resolved target. + pub target_name: &'a str, + /// CP-constructed chain for the *parent* delegation (empty for a root + /// delegation). Elements are `namespace/name`. + pub parent_chain: &'a [String], + pub deadline: DateTime, + pub parent_deadline: Option>, + pub now: DateTime, + pub max_deadline_secs: u64, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum PolicyDenial { + WorkerInitiation, + DepthExceeded { max: u32, would_be: u32 }, + Cycle { target: String }, + CrossNamespace, + DeadlinePast, + DeadlineTooLong { max_secs: u64 }, + DeadlineExceedsParent, +} + +impl std::fmt::Display for PolicyDenial { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PolicyDenial::WorkerInitiation => { + write!(f, "workers may not initiate delegations in this namespace") + } + PolicyDenial::DepthExceeded { max, would_be } => write!( + f, + "delegation depth {would_be} exceeds namespace max_depth {max}" + ), + PolicyDenial::Cycle { target } => { + write!(f, "cycle: {target} is already in the delegation chain") + } + PolicyDenial::CrossNamespace => { + write!(f, "cross-namespace delegation is not granted") + } + PolicyDenial::DeadlinePast => write!(f, "deadline is in the past"), + PolicyDenial::DeadlineTooLong { max_secs } => { + write!(f, "deadline exceeds the CP cap of {max_secs}s") + } + PolicyDenial::DeadlineExceedsParent => { + write!(f, "child deadline exceeds the parent's remaining budget") + } + } + } +} + +/// Evaluate the full CP-side policy for one delegation attempt. +pub fn check(input: &PolicyInput<'_>, ns_policy: &NamespacePolicy) -> Result<(), PolicyDenial> { + // 1. Initiator role. + if *input.from_type == AgentType::Worker && !ns_policy.allow_worker_initiation { + return Err(PolicyDenial::WorkerInitiation); + } + + // 2. Namespace boundary (v1: strict). + if input.from_namespace != input.target_namespace { + return Err(PolicyDenial::CrossNamespace); + } + + // 3. Depth: the new chain would be parent_chain + initiator; its length + // equals the delegation depth (root delegation → depth 1). + let would_be = input.parent_chain.len() as u32 + 1; + if would_be > ns_policy.max_depth { + return Err(PolicyDenial::DepthExceeded { + max: ns_policy.max_depth, + would_be, + }); + } + + // 4. Cycle: target must not already be an ancestor (or the initiator). + let target_id = format!("{}/{}", input.target_namespace, input.target_name); + let from_id = format!("{}/{}", input.from_namespace, input.from_name); + if target_id == from_id || input.parent_chain.contains(&target_id) { + return Err(PolicyDenial::Cycle { target: target_id }); + } + + // 5. Deadline sanity and caps. + if input.deadline <= input.now { + return Err(PolicyDenial::DeadlinePast); + } + let remaining = (input.deadline - input.now).num_seconds(); + if remaining > input.max_deadline_secs as i64 { + return Err(PolicyDenial::DeadlineTooLong { + max_secs: input.max_deadline_secs, + }); + } + if let Some(parent) = input.parent_deadline { + if input.deadline > parent { + return Err(PolicyDenial::DeadlineExceedsParent); + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Duration; + + fn base<'a>(now: DateTime, chain: &'a [String]) -> PolicyInput<'a> { + PolicyInput { + from_namespace: "prod", + from_name: "koudu", + from_type: &AgentType::Primary, + target_namespace: "prod", + target_name: "worker-1", + parent_chain: chain, + deadline: now + Duration::seconds(60), + parent_deadline: None, + now, + max_deadline_secs: 1800, + } + } + + fn default_policy() -> NamespacePolicy { + NamespacePolicy::default() + } + + #[test] + fn root_primary_delegation_passes() { + let now = Utc::now(); + assert!(check(&base(now, &[]), &default_policy()).is_ok()); + } + + #[test] + fn worker_initiation_denied_by_default_allowed_by_config() { + let now = Utc::now(); + let chain: Vec = vec![]; + let mut input = base(now, &chain); + input.from_type = &AgentType::Worker; + assert_eq!( + check(&input, &default_policy()), + Err(PolicyDenial::WorkerInitiation) + ); + let relaxed = NamespacePolicy { + max_depth: 2, + allow_worker_initiation: true, + }; + assert!(check(&input, &relaxed).is_ok()); + } + + #[test] + fn depth_exceeded_at_default_depth_one() { + let now = Utc::now(); + let chain = vec!["prod/root".to_string()]; + let input = base(now, &chain); + assert_eq!( + check(&input, &default_policy()), + Err(PolicyDenial::DepthExceeded { + max: 1, + would_be: 2 + }) + ); + let relaxed = NamespacePolicy { + max_depth: 2, + allow_worker_initiation: true, + }; + assert!(check(&input, &relaxed).is_ok()); + } + + #[test] + fn cycle_rejected_including_self() { + let now = Utc::now(); + let chain = vec!["prod/worker-1".to_string()]; + let relaxed = NamespacePolicy { + max_depth: 5, + allow_worker_initiation: true, + }; + let input = base(now, &chain); + assert!(matches!( + check(&input, &relaxed), + Err(PolicyDenial::Cycle { .. }) + )); + // self-delegation + let chain2: Vec = vec![]; + let mut input2 = base(now, &chain2); + input2.target_name = "koudu"; + assert!(matches!( + check(&input2, &relaxed), + Err(PolicyDenial::Cycle { .. }) + )); + } + + #[test] + fn cross_namespace_denied() { + let now = Utc::now(); + let chain: Vec = vec![]; + let mut input = base(now, &chain); + input.target_namespace = "dev"; + assert_eq!( + check(&input, &default_policy()), + Err(PolicyDenial::CrossNamespace) + ); + } + + #[test] + fn deadline_rules() { + let now = Utc::now(); + let chain: Vec = vec![]; + + let mut past = base(now, &chain); + past.deadline = now - Duration::seconds(1); + assert_eq!( + check(&past, &default_policy()), + Err(PolicyDenial::DeadlinePast) + ); + + let mut long = base(now, &chain); + long.deadline = now + Duration::seconds(3600); + assert_eq!( + check(&long, &default_policy()), + Err(PolicyDenial::DeadlineTooLong { max_secs: 1800 }) + ); + + let mut over_parent = base(now, &chain); + over_parent.deadline = now + Duration::seconds(120); + over_parent.parent_deadline = Some(now + Duration::seconds(60)); + assert_eq!( + check(&over_parent, &default_policy()), + Err(PolicyDenial::DeadlineExceedsParent) + ); + } +} diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs new file mode 100644 index 000000000..91c4a27f4 --- /dev/null +++ b/crates/openab-cp/src/proto.rs @@ -0,0 +1,426 @@ +//! Control-plane wire protocol: JSON-RPC 2.0 envelopes and `cp/*` method +//! payloads, following the conventions of `openab-core/src/acp/protocol.rs`. +//! +//! ## Contract summary +//! +//! - Transport: one WebSocket per runtime, text frames, one JSON-RPC message +//! per frame. +//! - Every request carries `jsonrpc: "2.0"` and a `u64` id; responses echo the +//! id. Correlation of *delegations* (which span multiple request/response +//! pairs across two connections) uses `delegation_id`, never the JSON-RPC id. +//! - The first frame on a new connection MUST be `cp/register`. Anything else +//! is rejected with `NOT_REGISTERED` and the connection is closed. +//! - Delegation ancestry (`chain`) is **CP-constructed**: callers supply only +//! `parent_delegation_id`; the CP derives the chain from authenticated +//! identities and its in-flight table. A runtime cannot forge ancestry. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// Wire protocol version. Carried in `cp/register`; the CP rejects +/// registrations with a version it does not support. +pub const PROTOCOL_VERSION: u32 = 1; + +// --- JSON-RPC envelopes --- + +#[derive(Debug, Serialize)] +pub struct JsonRpcRequest { + pub jsonrpc: &'static str, + pub id: u64, + pub method: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub params: Option, +} + +impl JsonRpcRequest { + pub fn new(id: u64, method: impl Into, params: Option) -> Self { + Self { + jsonrpc: "2.0", + id, + method: method.into(), + params, + } + } +} + +#[derive(Debug, Serialize)] +pub struct JsonRpcResponse { + pub jsonrpc: &'static str, + pub id: u64, + pub result: Value, +} + +impl JsonRpcResponse { + pub fn new(id: u64, result: Value) -> Self { + Self { + jsonrpc: "2.0", + id, + result, + } + } +} + +#[derive(Debug, Serialize)] +pub struct JsonRpcErrorResponse { + pub jsonrpc: &'static str, + pub id: u64, + pub error: ErrorObject, +} + +impl JsonRpcErrorResponse { + pub fn new(id: u64, error: ErrorObject) -> Self { + Self { + jsonrpc: "2.0", + id, + error, + } + } +} + +/// Incoming message: request, response, or error — distinguished by fields. +#[derive(Debug, Deserialize)] +pub struct JsonRpcMessage { + pub jsonrpc: Option, + pub id: Option, + pub method: Option, + pub params: Option, + pub result: Option, + pub error: Option, +} + +impl JsonRpcMessage { + /// Validate this frame as a JSON-RPC 2.0 **request** (review F4): the + /// `jsonrpc` field must be exactly "2.0", and a request id must be + /// present (all `cp/*` client→CP methods are requests, not + /// notifications). Returns the request id. + pub fn require_request_envelope(&self) -> Result { + if self.jsonrpc.as_deref() != Some("2.0") { + return Err(ErrorObject::new( + codes::INVALID_REQUEST, + "jsonrpc must be \"2.0\"", + )); + } + match self.id { + Some(id) => Ok(id), + None => Err(ErrorObject::new( + codes::INVALID_REQUEST, + "cp/* methods are requests and require an id", + )), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ErrorObject { + pub code: i64, + pub message: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub data: Option, +} + +impl ErrorObject { + pub fn new(code: i64, message: impl Into) -> Self { + Self { + code, + message: message.into(), + data: None, + } + } +} + +// --- Error codes (application range, distinct and machine-actionable) --- + +pub mod codes { + /// Frame received before a successful `cp/register` on this connection. + pub const NOT_REGISTERED: i64 = -32001; + /// Auth key unknown, or registration claims do not match the identity + /// bound to the key. + pub const IDENTITY_MISMATCH: i64 = -32002; + /// Delegation denied by policy (initiator role, depth, cycle, namespace). + pub const POLICY_DENIED: i64 = -32003; + /// No registered, healthy runtime matches the target selector. + pub const NO_TARGET: i64 = -32004; + /// Matching targets exist but all are at their advertised capacity. + /// Explicit fast-fail: the CP never queues (v1 has no durable state). + pub const SATURATED: i64 = -32005; + /// Delegation deadline elapsed before a result frame arrived. + pub const DEADLINE_EXCEEDED: i64 = -32006; + /// Serving runtime disconnected while the delegation was in flight. + pub const TARGET_DISCONNECTED: i64 = -32007; + /// `delegation_id` already in flight (idempotency guard). + pub const DUPLICATE_DELEGATION: i64 = -32008; + /// `cp/register` carried an unsupported protocol version. + pub const UNSUPPORTED_VERSION: i64 = -32009; + /// Malformed params for an otherwise valid method. + pub const INVALID_PARAMS: i64 = -32602; + /// Invalid JSON-RPC 2.0 envelope (missing/wrong `jsonrpc`, missing id). + pub const INVALID_REQUEST: i64 = -32600; + /// Unknown method. + pub const METHOD_NOT_FOUND: i64 = -32601; +} + +// --- cp/register --- + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum AgentType { + Primary, + Worker, +} + +impl std::fmt::Display for AgentType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AgentType::Primary => write!(f, "primary"), + AgentType::Worker => write!(f, "worker"), + } + } +} + +/// Params of `cp/register`, the mandatory first frame. +/// +/// `namespace`, `name`, and `agent_type` are **assertions to be verified**, +/// not authorization inputs: the CP compares them against the immutable +/// claims bound to the presented auth key and rejects any mismatch with +/// `IDENTITY_MISMATCH`. They exist in the frame so a misconfigured runtime +/// fails loudly at registration instead of being silently re-identified. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterParams { + pub protocol_version: u32, + pub namespace: String, + pub name: String, + #[serde(rename = "type")] + pub agent_type: AgentType, + /// Runtime-generated per-process id; distinguishes replicas of the same + /// logical agent during rolling deploys. + pub instance_id: String, + #[serde(default)] + pub labels: std::collections::BTreeMap, + /// Advertised concurrency budget. The CP may clamp this to a + /// config-defined cap for the identity. + #[serde(default = "default_max_sessions")] + pub max_delegated_sessions: u32, +} + +fn default_max_sessions() -> u32 { + 1 +} + +/// Result of a successful `cp/register`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterAck { + pub protocol_version: u32, + /// Interval at which the runtime must send `cp/heartbeat`. + pub heartbeat_interval_secs: u64, + /// Lease duration; missing heartbeats past this window deregisters the + /// instance and fails its in-flight delegations. + pub lease_expiry_secs: u64, + /// The effective (possibly clamped) concurrency budget. + pub effective_max_delegated_sessions: u32, +} + +// --- cp/heartbeat --- + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HeartbeatParams { + pub instance_id: String, + /// Current number of delegated sessions the runtime is serving; lets the + /// CP correct drift in its own in-flight accounting. + #[serde(default)] + pub active_delegated_sessions: u32, +} + +// --- cp/delegate --- + +/// Target selector: exact logical name, or label match (all pairs must match). +/// Exactly one of the two must be set. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct TargetSelector { + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub labels: Option>, +} + +/// Params of `cp/delegate` as sent by the initiating runtime. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateParams { + /// Caller-generated unique id (idempotency key). The CP rejects a second + /// in-flight delegation with the same id. + pub delegation_id: String, + pub target: TargetSelector, + pub prompt: String, + /// Absolute RFC3339 deadline. Mandatory: the CP rejects missing, past, or + /// over-cap deadlines. A child deadline can never exceed the parent's + /// remaining budget. + pub deadline: chrono::DateTime, + /// If this delegation is issued while serving another delegation, the id + /// of that parent. The CP derives the ancestry chain from this — the + /// chain is never client-supplied. + #[serde(skip_serializing_if = "Option::is_none")] + pub parent_delegation_id: Option, +} + +/// Params of `cp/delegate` as forwarded to the serving runtime. The CP stamps +/// the authenticated origin and the CP-constructed chain. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateForward { + pub delegation_id: String, + pub prompt: String, + pub deadline: chrono::DateTime, + /// Authenticated identity of the initiating agent (`namespace/name`). + pub from: String, + /// CP-constructed delegation ancestry, root first. The serving runtime + /// can trust every element: each hop was authenticated by the CP. + pub chain: Vec, +} + +/// Immediate result of `cp/delegate` (routing acceptance, not completion). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateAck { + pub delegation_id: String, + /// The chosen serving instance's logical name (`namespace/name`). + pub assigned_to: String, +} + +// --- cp/delegate_result --- + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum DelegationStatus { + Completed, + Failed, + Timeout, + Cancelled, + TargetDisconnected, +} + +/// Params of `cp/delegate_result` — emitted by the serving **runtime** when +/// the agent's turn ends (protocol-mandatory; never depends on the model), +/// or synthesized by the CP on timeout/disconnect. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateResultParams { + pub delegation_id: String, + pub status: DelegationStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +// --- cp/cancel --- + +/// Params of `cp/cancel`: from the initiator to abort an in-flight +/// delegation, or from the CP to the serving runtime (best effort) after a +/// timeout or initiator cancellation. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CancelParams { + pub delegation_id: String, + pub reason: String, +} + +// --- method names --- + +pub mod methods { + pub const REGISTER: &str = "cp/register"; + pub const HEARTBEAT: &str = "cp/heartbeat"; + pub const DELEGATE: &str = "cp/delegate"; + pub const DELEGATE_RESULT: &str = "cp/delegate_result"; + pub const CANCEL: &str = "cp/cancel"; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn register_params_roundtrip_with_type_rename() { + let json = serde_json::json!({ + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-abc", + "labels": {"backend": "kiro"}, + "max_delegated_sessions": 4 + }); + let p: RegisterParams = serde_json::from_value(json).unwrap(); + assert_eq!(p.agent_type, AgentType::Primary); + let back = serde_json::to_value(&p).unwrap(); + assert_eq!(back["type"], "primary"); + } + + #[test] + fn register_defaults_apply() { + let json = serde_json::json!({ + "protocol_version": 1, + "namespace": "prod", + "name": "w1", + "type": "worker", + "instance_id": "i-1" + }); + let p: RegisterParams = serde_json::from_value(json).unwrap(); + assert!(p.labels.is_empty()); + assert_eq!(p.max_delegated_sessions, 1); + } + + #[test] + fn delegate_params_require_deadline() { + let json = serde_json::json!({ + "delegation_id": "d-1", + "target": {"name": "w1"}, + "prompt": "hi" + }); + assert!(serde_json::from_value::(json).is_err()); + } + + #[test] + fn delegation_status_snake_case() { + assert_eq!( + serde_json::to_value(DelegationStatus::TargetDisconnected).unwrap(), + serde_json::json!("target_disconnected") + ); + } + + #[test] + fn incoming_message_distinguishes_request_and_response() { + let req: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","id":1,"method":"cp/heartbeat","params":{}}"#) + .unwrap(); + assert!(req.method.is_some() && req.result.is_none()); + let resp: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#).unwrap(); + assert!(resp.method.is_none() && resp.result.is_some()); + } + + #[test] + fn request_envelope_validation() { + let ok: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","id":7,"method":"cp/heartbeat"}"#).unwrap(); + assert_eq!(ok.require_request_envelope().unwrap(), 7); + + // Missing jsonrpc. + let no_ver: JsonRpcMessage = + serde_json::from_str(r#"{"id":7,"method":"cp/heartbeat"}"#).unwrap(); + assert_eq!( + no_ver.require_request_envelope().unwrap_err().code, + codes::INVALID_REQUEST + ); + + // Wrong version. + let bad_ver: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"1.0","id":7,"method":"cp/heartbeat"}"#).unwrap(); + assert_eq!( + bad_ver.require_request_envelope().unwrap_err().code, + codes::INVALID_REQUEST + ); + + // Notification shape (no id). + let no_id: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","method":"cp/heartbeat"}"#).unwrap(); + assert_eq!( + no_id.require_request_envelope().unwrap_err().code, + codes::INVALID_REQUEST + ); + } +} diff --git a/crates/openab-cp/src/registry.rs b/crates/openab-cp/src/registry.rs new file mode 100644 index 000000000..d97663bad --- /dev/null +++ b/crates/openab-cp/src/registry.rs @@ -0,0 +1,350 @@ +//! Instance registry: who is alive, with which claims, on which connection. +//! +//! Replica semantics (ADR §3): multiple instances may register under one +//! logical identity during rolling deploys. New delegations route to the +//! newest healthy instance; in-flight delegations complete on the instance +//! that accepted them. Lease expiry (missed heartbeats) deregisters an +//! instance and fails its in-flight delegations. + +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use parking_lot::RwLock; +use tokio::sync::mpsc; + +use crate::proto::AgentType; + +/// Outbound frame sender for one WS connection (serialized JSON text). +/// Bounded: a peer that cannot drain its queue is disconnected rather than +/// growing CP memory (review F5). +pub type FrameTx = mpsc::Sender; + +/// Capacity of each per-connection outbound queue. +pub const OUTBOUND_QUEUE: usize = 256; + +/// A live, authenticated, registered runtime instance. +#[derive(Clone, Debug)] +pub struct Instance { + /// CP-generated registration handle — the registry key and the basis of + /// all ownership checks. Never client-supplied (review F1): a colliding + /// client `instance_id` cannot replace or tear down another identity's + /// registration. + pub handle: u64, + pub namespace: String, + pub name: String, + pub agent_type: AgentType, + /// Client-supplied replica discriminator (display/audit only; ownership + /// and teardown key on `handle`). + pub instance_id: String, + pub labels: BTreeMap, + pub max_delegated_sessions: u32, + /// Delegations currently routed to this instance. CP-owned and + /// authoritative — never merged from runtime reports (review F6). + pub active_sessions: u32, + pub registered_at: Instant, + pub last_heartbeat: Instant, + pub tx: FrameTx, +} + +impl Instance { + pub fn logical_id(&self) -> String { + format!("{}/{}", self.namespace, self.name) + } + + pub fn saturated(&self) -> bool { + self.active_sessions >= self.max_delegated_sessions + } + + fn matches_labels(&self, want: &BTreeMap) -> bool { + want.iter() + .all(|(k, v)| self.labels.get(k).map(|x| x == v).unwrap_or(false)) + } +} + +#[derive(Default)] +pub struct Registry { + /// Keyed by CP-generated registration handle. + inner: RwLock>, + next_handle: AtomicU64, +} + +impl Registry { + pub fn new() -> Self { + Self::default() + } + + /// Insert a newly registered instance under a fresh CP-generated handle + /// (returned). Re-registrations (reconnects) get a new handle; the stale + /// entry disappears when its socket closes or its lease expires — it can + /// never be replaced by another connection's registration. + pub fn register(&self, mut inst: Instance) -> u64 { + let handle = self.next_handle.fetch_add(1, Ordering::Relaxed) + 1; + inst.handle = handle; + self.inner.write().insert(handle, inst); + handle + } + + /// Remove an instance by its registration handle (disconnect or lease + /// expiry). Only the owning connection or the sweeper knows the handle. + pub fn deregister(&self, handle: u64) -> Option { + self.inner.write().remove(&handle) + } + + /// Refresh the lease. The runtime-reported session count is intentionally + /// ignored: CP-owned in-flight accounting is authoritative (review F6 — + /// merging reports could pin an instance saturated forever). + pub fn heartbeat(&self, handle: u64) -> bool { + let mut g = self.inner.write(); + match g.get_mut(&handle) { + Some(i) => { + i.last_heartbeat = Instant::now(); + true + } + None => false, + } + } + + /// Handles whose lease has expired. + pub fn expired(&self, lease: Duration) -> Vec { + let now = Instant::now(); + self.inner + .read() + .values() + .filter(|i| now.duration_since(i.last_heartbeat) > lease) + .map(|i| i.handle) + .collect() + } + + pub fn get(&self, handle: u64) -> Option { + self.inner.read().get(&handle).cloned() + } + + /// Select a serving instance within `namespace` by exact name or labels. + /// + /// Unsaturated matches only. Ordering (review F6): + /// - exact-name selection → replicas of one logical agent: newest + /// registration first (rolling-deploy rule), load as tie-breaker + /// - label selection → across logical agents: least loaded first, + /// registration recency as tie-breaker + pub fn select( + &self, + namespace: &str, + name: Option<&str>, + labels: Option<&BTreeMap>, + ) -> Result { + let g = self.inner.read(); + let mut matches: Vec<&Instance> = g + .values() + .filter(|i| i.namespace == namespace) + .filter(|i| match name { + Some(n) => i.name == n, + None => true, + }) + .filter(|i| match labels { + Some(want) => i.matches_labels(want), + None => true, + }) + .collect(); + + if matches.is_empty() { + return Err(SelectError::NoTarget); + } + matches.retain(|i| !i.saturated()); + if matches.is_empty() { + return Err(SelectError::Saturated); + } + if name.is_some() { + matches.sort_by(|a, b| { + b.registered_at + .cmp(&a.registered_at) + .then(a.active_sessions.cmp(&b.active_sessions)) + }); + } else { + matches.sort_by(|a, b| { + a.active_sessions + .cmp(&b.active_sessions) + .then(b.registered_at.cmp(&a.registered_at)) + }); + } + Ok(matches[0].clone()) + } + + /// Adjust the CP-owned in-flight count for an instance. + pub fn adjust_sessions(&self, handle: u64, delta: i32) { + let mut g = self.inner.write(); + if let Some(i) = g.get_mut(&handle) { + i.active_sessions = i.active_sessions.saturating_add_signed(delta); + } + } + + /// Registry snapshot for one namespace (basis for a future `list_agents`). + pub fn list(&self, namespace: &str) -> Vec { + self.inner + .read() + .values() + .filter(|i| i.namespace == namespace) + .cloned() + .collect() + } +} + +#[derive(Debug, PartialEq, Eq)] +pub enum SelectError { + NoTarget, + Saturated, +} + +#[cfg(test)] +mod tests { + use super::*; + + fn inst(ns: &str, name: &str, id: &str, max: u32) -> Instance { + let (tx, _rx) = mpsc::channel(OUTBOUND_QUEUE); + Instance { + handle: 0, // assigned by register() + namespace: ns.into(), + name: name.into(), + agent_type: AgentType::Worker, + instance_id: id.into(), + labels: BTreeMap::new(), + max_delegated_sessions: max, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + } + } + + #[test] + fn select_by_name_and_namespace_isolation() { + let r = Registry::new(); + let h1 = r.register(inst("prod", "w1", "i-1", 2)); + r.register(inst("dev", "w1", "i-2", 2)); + let got = r.select("prod", Some("w1"), None).unwrap(); + assert_eq!(got.handle, h1); + assert!(matches!( + r.select("staging", Some("w1"), None), + Err(SelectError::NoTarget) + )); + } + + #[test] + fn replicas_route_to_newest() { + let r = Registry::new(); + r.register(inst("prod", "w1", "i-old", 2)); + std::thread::sleep(Duration::from_millis(5)); + let h_new = r.register(inst("prod", "w1", "i-new", 2)); + let got = r.select("prod", Some("w1"), None).unwrap(); + assert_eq!(got.handle, h_new); + } + + #[test] + fn saturation_is_distinct_from_no_target() { + let r = Registry::new(); + let mut i = inst("prod", "w1", "i-1", 1); + i.active_sessions = 1; + r.register(i); + assert!(matches!( + r.select("prod", Some("w1"), None), + Err(SelectError::Saturated) + )); + assert!(matches!( + r.select("prod", Some("nope"), None), + Err(SelectError::NoTarget) + )); + } + + #[test] + fn label_selection_least_loaded_first() { + let r = Registry::new(); + // Older but less loaded instance must win under label selection + // (inverse recency/load — review F6). + let mut a = inst("prod", "wa", "i-a", 4); + a.labels.insert("backend".into(), "kiro".into()); + a.active_sessions = 0; + let h_a = r.register(a); + std::thread::sleep(Duration::from_millis(5)); + let mut b = inst("prod", "wb", "i-b", 4); + b.labels.insert("backend".into(), "kiro".into()); + b.active_sessions = 3; + r.register(b); + + let mut want = BTreeMap::new(); + want.insert("backend".to_string(), "kiro".to_string()); + let got = r.select("prod", None, Some(&want)).unwrap(); + assert_eq!(got.handle, h_a, "least loaded wins despite being older"); + + // partial label mismatch -> NoTarget + want.insert("arch".to_string(), "x86".to_string()); + assert!(matches!( + r.select("prod", None, Some(&want)), + Err(SelectError::NoTarget) + )); + } + + #[test] + fn name_selection_newest_first_even_if_more_loaded() { + let r = Registry::new(); + let mut old = inst("prod", "w1", "i-old", 4); + old.active_sessions = 0; + r.register(old); + std::thread::sleep(Duration::from_millis(5)); + let mut new = inst("prod", "w1", "i-new", 4); + new.active_sessions = 2; + let h_new = r.register(new); + let got = r.select("prod", Some("w1"), None).unwrap(); + assert_eq!(got.handle, h_new, "replica rule: newest registration wins"); + } + + #[test] + fn lease_expiry_and_heartbeat() { + let r = Registry::new(); + let h = r.register(inst("prod", "w1", "i-1", 1)); + assert!(r.expired(Duration::from_secs(60)).is_empty()); + assert!(r.heartbeat(h)); + assert!(!r.heartbeat(h + 999)); + std::thread::sleep(Duration::from_millis(2)); + assert_eq!(r.expired(Duration::ZERO), vec![h]); + } + + #[test] + fn colliding_instance_id_cannot_replace_other_registration() { + // Review F1: a second connection registering the same client-supplied + // instance_id gets its own handle; the first registration survives + // and can only be torn down via its own handle. + let r = Registry::new(); + let h1 = r.register(inst("prod", "w1", "i-same", 1)); + let h2 = r.register(inst("prod", "w2", "i-same", 1)); + assert_ne!(h1, h2); + assert_eq!(r.list("prod").len(), 2); + // Tearing down the second leaves the first intact. + assert!(r.deregister(h2).is_some()); + assert!(r.get(h1).is_some()); + // Deregistering an already-gone handle is a no-op. + assert!(r.deregister(h2).is_none()); + } + + #[test] + fn heartbeat_does_not_mutate_session_count() { + let r = Registry::new(); + let h = r.register(inst("prod", "w1", "i-1", 2)); + r.adjust_sessions(h, 1); + assert!(r.heartbeat(h)); + assert_eq!( + r.get(h).unwrap().active_sessions, + 1, + "CP-owned count is authoritative; heartbeat never changes it" + ); + } + + #[test] + fn adjust_sessions_saturating() { + let r = Registry::new(); + let h = r.register(inst("prod", "w1", "i-1", 2)); + r.adjust_sessions(h, 1); + assert_eq!(r.get(h).unwrap().active_sessions, 1); + r.adjust_sessions(h, -5); + assert_eq!(r.get(h).unwrap().active_sessions, 0); + } +} diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs new file mode 100644 index 000000000..4aae8ccaf --- /dev/null +++ b/crates/openab-cp/src/router.rs @@ -0,0 +1,992 @@ +//! Delegation router: in-flight table, target selection, result routing, and +//! the failure semantics the ADR review required to be explicit: +//! +//! - **Deadline sweep** — an in-flight delegation whose deadline passes is +//! terminated: the initiator receives a synthesized `timeout` result and +//! the serving runtime receives a best-effort `cp/cancel` (stop burning +//! tokens). +//! - **Target disconnect / lease expiry** — in-flight delegations on that +//! instance fail immediately with `target_disconnected`. +//! - **Initiator disconnect** — its in-flight delegations are cancelled +//! downstream (best effort); nobody is left to receive the result. +//! - **CP restart** — the table is in-memory; all in-flight delegations +//! effectively end as initiator-side timeouts. Late `cp/delegate_result` +//! frames for unknown ids are acknowledged and dropped (logged), so +//! reconnecting runtimes do not error-loop. +//! - **Saturation** — routing never queues; `SATURATED` is returned +//! immediately (fast-fail, no hidden buffer). + +use std::collections::BTreeMap; + +use chrono::{DateTime, Utc}; +use parking_lot::Mutex; +use tracing::{info, warn}; + +use crate::config::CpConfig; +use crate::policy::{self, PolicyInput}; +use crate::proto::{ + codes, methods, AgentType, CancelParams, DelegateAck, DelegateForward, DelegateParams, + DelegateResultParams, DelegationStatus, ErrorObject, JsonRpcRequest, +}; +use crate::registry::{Instance, Registry, SelectError}; + +/// One in-flight delegation. Ownership is tracked by CP-generated +/// registration handles, never client-supplied ids (review F1). +#[derive(Clone)] +pub struct InFlight { + pub delegation_id: String, + /// Authenticated initiator (`namespace/name`) and its registration handle. + pub from_logical: String, + pub from_handle: u64, + /// Chosen serving instance. + pub to_logical: String, + pub to_handle: u64, + pub deadline: DateTime, + /// CP-constructed chain for THIS delegation (root first, ends with the + /// initiator). Children extend it. + pub chain: Vec, +} + +pub struct Router { + inflight: Mutex>, + /// Serializes the delegate admission sequence (duplicate check → target + /// selection → capacity reservation → in-flight insert) so concurrent + /// requests cannot double-admit one id or oversubscribe capacity + /// (review F2). Delegation rates are LLM-scale; a coarse admission lock + /// is simple and more than sufficient. + admission: Mutex<()>, +} + +pub enum DelegateOutcome { + /// Forwarded to the target; ack for the initiator. + Accepted(DelegateAck), + /// Rejected; error for the initiator. + Rejected(ErrorObject), +} + +impl Router { + pub fn new() -> Self { + Self { + inflight: Mutex::new(BTreeMap::new()), + admission: Mutex::new(()), + } + } + + /// Handle `cp/delegate` from an authenticated, registered initiator. + #[allow(clippy::too_many_arguments)] + pub fn delegate( + &self, + cfg: &CpConfig, + registry: &Registry, + from_namespace: &str, + from_name: &str, + from_type: &AgentType, + from_handle: u64, + params: DelegateParams, + next_rpc_id: u64, + ) -> DelegateOutcome { + let now = Utc::now(); + + // Admission is one atomic sequence (review F2): duplicate check, + // parent lookup, target selection, capacity reservation, and + // in-flight insertion all happen under this guard. + let _admission = self.admission.lock(); + + if self.inflight.lock().contains_key(¶ms.delegation_id) { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::DUPLICATE_DELEGATION, + format!("delegation {} is already in flight", params.delegation_id), + )); + } + + // Selector sanity: exactly one of name/labels. + let (sel_name, sel_labels) = (params.target.name.as_deref(), params.target.labels.as_ref()); + if sel_name.is_some() == sel_labels.is_some() { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + "target must set exactly one of `name` or `labels`", + )); + } + + // Parent linkage: chain and deadline derive from the CP's own table, + // never from the client. The caller must BE the instance serving the + // parent delegation — otherwise any runtime knowing a live id could + // borrow its trusted chain and deadline budget (review F3). Unknown + // and unauthorized parent ids return the same error (no enumeration). + let (parent_chain, parent_deadline) = match ¶ms.parent_delegation_id { + Some(pid) => match self.inflight.lock().get(pid) { + Some(p) if p.to_handle == from_handle => (p.chain.clone(), Some(p.deadline)), + _ => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + format!("parent delegation {pid} is not in flight for this instance"), + )) + } + }, + None => (Vec::new(), None), + }; + + // Resolve target within the initiator's namespace (v1 boundary). + let target = match registry.select(from_namespace, sel_name, sel_labels) { + Ok(i) => i, + Err(SelectError::NoTarget) => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::NO_TARGET, + "no registered healthy runtime matches the target selector", + )) + } + Err(SelectError::Saturated) => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::SATURATED, + "all matching runtimes are at capacity (CP does not queue; retry later)", + )) + } + }; + + // CP-authoritative policy. + let input = PolicyInput { + from_namespace, + from_name, + from_type, + target_namespace: &target.namespace, + target_name: &target.name, + parent_chain: &parent_chain, + deadline: params.deadline, + parent_deadline, + now, + max_deadline_secs: cfg.max_deadline_secs, + }; + if let Err(denial) = policy::check(&input, &cfg.policy_for(from_namespace)) { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::POLICY_DENIED, + denial.to_string(), + )); + } + + // Build the forward frame with the CP-stamped chain. + let from_logical = format!("{from_namespace}/{from_name}"); + let mut chain = parent_chain; + chain.push(from_logical.clone()); + let forward = DelegateForward { + delegation_id: params.delegation_id.clone(), + prompt: params.prompt, + deadline: params.deadline, + from: from_logical.clone(), + chain: chain.clone(), + }; + let frame = JsonRpcRequest::new( + next_rpc_id, + methods::DELEGATE, + Some(serde_json::to_value(&forward).expect("serializable")), + ); + let text = serde_json::to_string(&frame).expect("serializable"); + + // Reserve capacity and record the in-flight entry BEFORE sending, so + // an immediately-arriving result finds it (review F2). Roll both + // back if the send fails. + registry.adjust_sessions(target.handle, 1); + let entry = InFlight { + delegation_id: params.delegation_id.clone(), + from_logical, + from_handle, + to_logical: target.logical_id(), + to_handle: target.handle, + deadline: params.deadline, + chain, + }; + self.inflight + .lock() + .insert(params.delegation_id.clone(), entry.clone()); + + if target.tx.try_send(text).is_err() { + // Disconnected or backpressured beyond its queue: roll back. + self.inflight.lock().remove(¶ms.delegation_id); + registry.adjust_sessions(target.handle, -1); + return DelegateOutcome::Rejected(ErrorObject::new( + codes::TARGET_DISCONNECTED, + "target disconnected or unresponsive during routing", + )); + } + + info!( + delegation = %entry.delegation_id, + from = %entry.from_logical, + to = %entry.to_logical, + chain = ?entry.chain, + deadline = %entry.deadline, + "delegation routed" + ); + + DelegateOutcome::Accepted(DelegateAck { + delegation_id: params.delegation_id, + assigned_to: target.logical_id(), + }) + } + + /// Handle `cp/delegate_result` from the serving runtime. Returns the + /// initiator-bound frame if the delegation is known; unknown ids (e.g. + /// results arriving after a CP restart) are dropped with a log. + pub fn complete( + &self, + registry: &Registry, + serving_handle: u64, + mut params: DelegateResultParams, + max_result_bytes: usize, + next_rpc_id: u64, + ) -> Option<(Instance, String)> { + let entry = { self.inflight.lock().remove(¶ms.delegation_id) }; + let entry = match entry { + Some(e) => e, + None => { + warn!( + delegation = %params.delegation_id, + "result for unknown delegation (late arrival or CP restart) — dropped" + ); + return None; + } + }; + if entry.to_handle != serving_handle { + // Only the instance the delegation was routed to may complete it. + warn!( + delegation = %params.delegation_id, + expected = entry.to_handle, + got = serving_handle, + "result from unexpected instance — dropped, delegation restored" + ); + self.inflight + .lock() + .insert(params.delegation_id.clone(), entry); + return None; + } + + registry.adjust_sessions(entry.to_handle, -1); + + // Truncate oversized results (keep the head; delegation already ran). + if let Some(r) = ¶ms.result { + if r.len() > max_result_bytes { + let mut cut = max_result_bytes; + while !r.is_char_boundary(cut) { + cut -= 1; + } + params.result = Some(format!( + "{}\n…[truncated by control plane: {} of {} bytes]", + &r[..cut], + cut, + r.len() + )); + } + } + + info!( + delegation = %params.delegation_id, + status = ?params.status, + from = %entry.to_logical, + to = %entry.from_logical, + "delegation completed" + ); + + let initiator = registry.get(entry.from_handle)?; + let frame = JsonRpcRequest::new( + next_rpc_id, + methods::DELEGATE_RESULT, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + Some(( + initiator, + serde_json::to_string(&frame).expect("serializable"), + )) + } + + /// Handle `cp/cancel` from the initiator. Returns the frame to forward + /// to the serving runtime, if the delegation is in flight and owned by + /// the caller. + pub fn cancel( + &self, + registry: &Registry, + from_handle: u64, + params: &CancelParams, + next_rpc_id: u64, + ) -> Result, ErrorObject> { + let entry = { self.inflight.lock().remove(¶ms.delegation_id) }; + let entry = match entry { + Some(e) => e, + None => { + return Err(ErrorObject::new( + codes::INVALID_PARAMS, + format!("delegation {} is not in flight", params.delegation_id), + )) + } + }; + if entry.from_handle != from_handle { + self.inflight + .lock() + .insert(params.delegation_id.clone(), entry); + return Err(ErrorObject::new( + codes::POLICY_DENIED, + "only the initiating instance may cancel a delegation", + )); + } + registry.adjust_sessions(entry.to_handle, -1); + info!(delegation = %params.delegation_id, "delegation cancelled by initiator"); + let target = registry.get(entry.to_handle); + Ok(target.map(|t| { + let frame = JsonRpcRequest::new( + next_rpc_id, + methods::CANCEL, + Some(serde_json::to_value(params).expect("serializable")), + ); + (t, serde_json::to_string(&frame).expect("serializable")) + })) + } + + /// Fail every in-flight delegation touching a deregistered instance. + /// Returns synthesized result/cancel frames to deliver: + /// - delegations SERVED by the instance → `target_disconnected` result to + /// the initiator + /// - delegations INITIATED by the instance → best-effort `cp/cancel` to + /// the serving runtime + pub fn fail_instance( + &self, + registry: &Registry, + handle: u64, + rpc_id: &mut impl FnMut() -> u64, + ) -> Vec<(Instance, String)> { + let mut affected = Vec::new(); + let entries: Vec = { + let mut g = self.inflight.lock(); + let ids: Vec = g + .values() + .filter(|e| e.to_handle == handle || e.from_handle == handle) + .map(|e| e.delegation_id.clone()) + .collect(); + ids.iter().filter_map(|id| g.remove(id)).collect() + }; + for e in entries { + if e.to_handle == handle { + // Serving side died → tell the initiator. + if let Some(init) = registry.get(e.from_handle) { + let params = DelegateResultParams { + delegation_id: e.delegation_id.clone(), + status: DelegationStatus::TargetDisconnected, + result: None, + error: Some(format!("{} disconnected", e.to_logical)), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::DELEGATE_RESULT, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + affected.push((init, serde_json::to_string(&frame).expect("serializable"))); + } + } else { + // Initiator died → cancel downstream, free worker capacity. + registry.adjust_sessions(e.to_handle, -1); + if let Some(target) = registry.get(e.to_handle) { + let params = CancelParams { + delegation_id: e.delegation_id.clone(), + reason: format!("initiator {} disconnected", e.from_logical), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::CANCEL, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + affected.push((target, serde_json::to_string(&frame).expect("serializable"))); + } + } + warn!(delegation = %e.delegation_id, handle, "in-flight delegation failed by disconnect"); + } + affected + } + + /// Deadline sweep: expire overdue delegations. Returns frames to deliver + /// (timeout result to the initiator, best-effort cancel to the server). + pub fn sweep_deadlines( + &self, + registry: &Registry, + now: DateTime, + rpc_id: &mut impl FnMut() -> u64, + ) -> Vec<(Instance, String)> { + let overdue: Vec = { + let mut g = self.inflight.lock(); + let ids: Vec = g + .values() + .filter(|e| e.deadline <= now) + .map(|e| e.delegation_id.clone()) + .collect(); + ids.iter().filter_map(|id| g.remove(id)).collect() + }; + let mut frames = Vec::new(); + for e in overdue { + registry.adjust_sessions(e.to_handle, -1); + warn!(delegation = %e.delegation_id, deadline = %e.deadline, "delegation deadline exceeded"); + if let Some(init) = registry.get(e.from_handle) { + let params = DelegateResultParams { + delegation_id: e.delegation_id.clone(), + status: DelegationStatus::Timeout, + result: None, + error: Some("deadline exceeded".to_string()), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::DELEGATE_RESULT, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + frames.push((init, serde_json::to_string(&frame).expect("serializable"))); + } + if let Some(target) = registry.get(e.to_handle) { + let params = CancelParams { + delegation_id: e.delegation_id.clone(), + reason: "deadline exceeded".to_string(), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::CANCEL, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + frames.push((target, serde_json::to_string(&frame).expect("serializable"))); + } + } + frames + } + + /// Chain of an in-flight delegation (for tests/inspection). + pub fn chain_of(&self, delegation_id: &str) -> Option> { + self.inflight + .lock() + .get(delegation_id) + .map(|e| e.chain.clone()) + } + + pub fn inflight_count(&self) -> usize { + self.inflight.lock().len() + } +} + +impl Default for Router { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::proto::TargetSelector; + use crate::registry::OUTBOUND_QUEUE; + use chrono::Duration; + use std::time::Instant; + use tokio::sync::mpsc; + + fn cfg() -> CpConfig { + toml::from_str( + r#" +[[agents]] +key = "kp" +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "kw" +namespace = "prod" +name = "worker-1" +type = "worker" +"#, + ) + .unwrap() + } + + fn instance( + ns: &str, + name: &str, + ty: AgentType, + max: u32, + ) -> (Instance, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(OUTBOUND_QUEUE); + ( + Instance { + handle: 0, + namespace: ns.into(), + name: name.into(), + agent_type: ty, + instance_id: format!("i-{name}"), + labels: Default::default(), + max_delegated_sessions: max, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + }, + rx, + ) + } + + fn delegate_params(id: &str, target: &str, secs: i64) -> DelegateParams { + DelegateParams { + delegation_id: id.into(), + target: TargetSelector { + name: Some(target.into()), + labels: None, + }, + prompt: "do it".into(), + deadline: Utc::now() + Duration::seconds(secs), + parent_delegation_id: None, + } + } + + struct World { + cfg: CpConfig, + registry: Registry, + router: Router, + h_primary: u64, + h_worker: u64, + worker_rx: mpsc::Receiver, + primary_rx: mpsc::Receiver, + } + + fn world() -> World { + let registry = Registry::new(); + let (p, primary_rx) = instance("prod", "koudu", AgentType::Primary, 4); + let (w, worker_rx) = instance("prod", "worker-1", AgentType::Worker, 1); + let h_primary = registry.register(p); + let h_worker = registry.register(w); + World { + cfg: cfg(), + registry, + router: Router::new(), + h_primary, + h_worker, + worker_rx, + primary_rx, + } + } + + fn do_delegate(w: &World, params: DelegateParams) -> DelegateOutcome { + w.router.delegate( + &w.cfg, + &w.registry, + "prod", + "koudu", + &AgentType::Primary, + w.h_primary, + params, + 1, + ) + } + + #[test] + fn happy_path_roundtrip() { + let mut w = world(); + let out = do_delegate(&w, delegate_params("d-1", "worker-1", 60)); + let ack = match out { + DelegateOutcome::Accepted(a) => a, + DelegateOutcome::Rejected(e) => panic!("rejected: {}", e.message), + }; + assert_eq!(ack.assigned_to, "prod/worker-1"); + + let frame = w.worker_rx.try_recv().unwrap(); + let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); + assert_eq!(v["method"], "cp/delegate"); + assert_eq!(v["params"]["from"], "prod/koudu"); + assert_eq!(v["params"]["chain"], serde_json::json!(["prod/koudu"])); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 1); + + let result = DelegateResultParams { + delegation_id: "d-1".into(), + status: DelegationStatus::Completed, + result: Some("done".into()), + error: None, + }; + let (init, frame) = w + .router + .complete(&w.registry, w.h_worker, result, 1024, 2) + .unwrap(); + assert_eq!(init.handle, w.h_primary); + assert!(frame.contains("\"completed\"")); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + assert_eq!(w.router.inflight_count(), 0); + } + + #[test] + fn inflight_exists_before_target_receives_frame() { + // Review F2: an immediately-arriving result must find the entry. + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + // Complete BEFORE draining the worker's queue — entry must exist. + let result = DelegateResultParams { + delegation_id: "d-1".into(), + status: DelegationStatus::Completed, + result: Some("instant".into()), + error: None, + }; + assert!(w + .router + .complete(&w.registry, w.h_worker, result, 1024, 2) + .is_some()); + w.worker_rx.try_recv().unwrap(); + } + + #[test] + fn send_failure_rolls_back_reservation() { + // Close the worker's rx so try_send fails, then verify rollback. + let mut w = world(); + w.worker_rx.close(); + match do_delegate(&w, delegate_params("d-1", "worker-1", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::TARGET_DISCONNECTED), + _ => panic!("expected TARGET_DISCONNECTED"), + } + assert_eq!(w.router.inflight_count(), 0); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "capacity reservation must be rolled back" + ); + } + + #[test] + fn duplicate_delegation_id_rejected() { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + match do_delegate(&w, delegate_params("d-1", "worker-1", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::DUPLICATE_DELEGATION), + _ => panic!("expected rejection"), + } + } + + #[test] + fn saturation_fast_fails() { + let w = world(); // worker max = 1 + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + match do_delegate(&w, delegate_params("d-2", "worker-1", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::SATURATED), + _ => panic!("expected SATURATED"), + } + } + + #[test] + fn selector_must_be_exactly_one() { + let w = world(); + let mut p = delegate_params("d-1", "worker-1", 60); + p.target.labels = Some(Default::default()); + match do_delegate(&w, p) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::INVALID_PARAMS), + _ => panic!(), + } + let mut p2 = delegate_params("d-2", "worker-1", 60); + p2.target.name = None; + match do_delegate(&w, p2) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::INVALID_PARAMS), + _ => panic!(), + } + } + + #[test] + fn policy_denial_maps_to_error_code() { + let w = world(); + let out = w.router.delegate( + &w.cfg, + &w.registry, + "prod", + "worker-1", + &AgentType::Worker, + w.h_worker, + delegate_params("d-1", "koudu", 60), + 1, + ); + match out { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::POLICY_DENIED), + _ => panic!("expected POLICY_DENIED"), + } + } + + #[test] + fn unknown_target_is_no_target() { + let w = world(); + match do_delegate(&w, delegate_params("d-1", "ghost", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::NO_TARGET), + _ => panic!(), + } + } + + #[test] + fn result_from_wrong_handle_dropped_and_restored() { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let result = DelegateResultParams { + delegation_id: "d-1".into(), + status: DelegationStatus::Completed, + result: Some("spoofed".into()), + error: None, + }; + // h_primary is a valid handle but NOT the serving instance. + assert!(w + .router + .complete(&w.registry, w.h_primary, result, 1024, 2) + .is_none()); + assert_eq!(w.router.inflight_count(), 1); + } + + #[test] + fn late_result_after_restart_dropped() { + let w = world(); + let result = DelegateResultParams { + delegation_id: "d-unknown".into(), + status: DelegationStatus::Completed, + result: None, + error: None, + }; + assert!(w + .router + .complete(&w.registry, w.h_worker, result, 1024, 2) + .is_none()); + } + + #[test] + fn oversized_result_truncated() { + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + let result = DelegateResultParams { + delegation_id: "d-1".into(), + status: DelegationStatus::Completed, + result: Some("x".repeat(100)), + error: None, + }; + let (_, frame) = w + .router + .complete(&w.registry, w.h_worker, result, 10, 2) + .unwrap(); + assert!(frame.contains("truncated by control plane")); + } + + #[test] + fn deadline_sweep_times_out_and_cancels() { + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + + let mut id = 100u64; + let mut next = || { + id += 1; + id + }; + assert!(w + .router + .sweep_deadlines(&w.registry, Utc::now(), &mut next) + .is_empty()); + let frames = + w.router + .sweep_deadlines(&w.registry, Utc::now() + Duration::seconds(120), &mut next); + assert_eq!(frames.len(), 2); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + + for (inst, frame) in frames { + let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); + match v["method"].as_str().unwrap() { + "cp/delegate_result" => { + assert_eq!(inst.handle, w.h_primary); + assert_eq!(v["params"]["status"], "timeout"); + } + "cp/cancel" => assert_eq!(inst.handle, w.h_worker), + m => panic!("unexpected method {m}"), + } + } + } + + #[test] + fn worker_disconnect_fails_delegation_to_initiator() { + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + w.registry.deregister(w.h_worker); + + let mut id = 0u64; + let mut next = || { + id += 1; + id + }; + let frames = w.router.fail_instance(&w.registry, w.h_worker, &mut next); + assert_eq!(frames.len(), 1); + let (inst, frame) = &frames[0]; + assert_eq!(inst.handle, w.h_primary); + assert!(frame.contains("target_disconnected")); + assert_eq!(w.router.inflight_count(), 0); + inst.tx.try_send(frame.clone()).unwrap(); + assert!(w.primary_rx.try_recv().unwrap().contains("d-1")); + } + + #[test] + fn initiator_disconnect_cancels_downstream() { + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + w.registry.deregister(w.h_primary); + + let mut id = 0u64; + let mut next = || { + id += 1; + id + }; + let frames = w.router.fail_instance(&w.registry, w.h_primary, &mut next); + assert_eq!(frames.len(), 1); + let (inst, frame) = &frames[0]; + assert_eq!(inst.handle, w.h_worker); + assert!(frame.contains("cp/cancel")); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + } + + #[test] + fn cancel_only_by_initiator() { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let params = CancelParams { + delegation_id: "d-1".into(), + reason: "changed my mind".into(), + }; + let err = w + .router + .cancel(&w.registry, w.h_worker, ¶ms, 5) + .unwrap_err(); + assert_eq!(err.code, codes::POLICY_DENIED); + assert_eq!(w.router.inflight_count(), 1); + let fwd = w + .router + .cancel(&w.registry, w.h_primary, ¶ms, 6) + .unwrap(); + let (inst, frame) = fwd.unwrap(); + assert_eq!(inst.handle, w.h_worker); + assert!(frame.contains("cp/cancel")); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + } + + #[test] + fn chain_extends_through_parent_and_foreign_parent_rejected() { + let w = world(); + let cfg: CpConfig = toml::from_str( + r#" +[[agents]] +key = "kp" +namespace = "prod" +name = "koudu" +type = "primary" + +[namespaces.prod] +max_depth = 5 +allow_worker_initiation = true +"#, + ) + .unwrap(); + let (w2, _rx2) = instance("prod", "worker-2", AgentType::Worker, 1); + let h_w2 = w.registry.register(w2); + + assert!(matches!( + w.router.delegate( + &cfg, + &w.registry, + "prod", + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-root", "worker-1", 120), + 1, + ), + DelegateOutcome::Accepted(_) + )); + assert_eq!( + w.router.chain_of("d-root").unwrap(), + vec!["prod/koudu".to_string()] + ); + + // Review F3: worker-2 (NOT serving d-root) tries to borrow d-root + // as parent — rejected. + let mut foreign = delegate_params("d-foreign", "worker-2", 60); + foreign.parent_delegation_id = Some("d-root".into()); + match w.router.delegate( + &cfg, + &w.registry, + "prod", + "worker-2", + &AgentType::Worker, + h_w2, + foreign, + 2, + ) { + DelegateOutcome::Rejected(e) => { + assert_eq!(e.code, codes::INVALID_PARAMS); + assert!(e.message.contains("not in flight for this instance")); + } + _ => panic!("foreign parent must be rejected"), + } + + // worker-1 (serving d-root) delegates a legitimate child to worker-2. + let mut child = delegate_params("d-child", "worker-2", 60); + child.parent_delegation_id = Some("d-root".into()); + assert!(matches!( + w.router.delegate( + &cfg, + &w.registry, + "prod", + "worker-1", + &AgentType::Worker, + w.h_worker, + child, + 3, + ), + DelegateOutcome::Accepted(_) + )); + assert_eq!( + w.router.chain_of("d-child").unwrap(), + vec!["prod/koudu".to_string(), "prod/worker-1".to_string()] + ); + + // Cycle: worker-2 delegating back to koudu is rejected. + let mut cyc = delegate_params("d-cyc", "koudu", 30); + cyc.parent_delegation_id = Some("d-child".into()); + match w.router.delegate( + &cfg, + &w.registry, + "prod", + "worker-2", + &AgentType::Worker, + h_w2, + cyc, + 4, + ) { + DelegateOutcome::Rejected(e) => { + assert_eq!(e.code, codes::POLICY_DENIED); + assert!(e.message.contains("cycle"), "{}", e.message); + } + _ => panic!("expected cycle rejection"), + } + } +} diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs new file mode 100644 index 000000000..9504b5c42 --- /dev/null +++ b/crates/openab-cp/src/server.rs @@ -0,0 +1,597 @@ +//! WebSocket server: authentication at upgrade, mandatory `cp/register` +//! first frame, then frame dispatch to registry/policy/router. +//! +//! Auth: the runtime presents its key as `Authorization: Bearer ` on the +//! upgrade request. Keys never appear in URLs (avoids access-log leakage). +//! +//! Resource bounds (review F5): the WS transport enforces +//! `max_frame_bytes` before parsing; each connection's outbound queue is +//! bounded — a peer that cannot drain it is treated as disconnected. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Instant; + +use axum::extract::ws::{Message, WebSocket}; +use axum::extract::{State, WebSocketUpgrade}; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::IntoResponse; +use axum::routing::get; +use axum::Router as AxumRouter; +use futures_util::{SinkExt, StreamExt}; +use tokio::sync::mpsc; +use tracing::{info, warn}; + +use crate::config::{AgentIdentity, CpConfig}; +use crate::proto::{ + codes, methods, CancelParams, DelegateParams, DelegateResultParams, ErrorObject, + JsonRpcErrorResponse, JsonRpcMessage, JsonRpcResponse, RegisterAck, RegisterParams, + PROTOCOL_VERSION, +}; +use crate::registry::{Instance, Registry, OUTBOUND_QUEUE}; +use crate::router::{DelegateOutcome, Router}; + +pub struct AppState { + pub cfg: CpConfig, + pub registry: Registry, + pub router: Router, + rpc_id: AtomicU64, +} + +impl AppState { + pub fn new(cfg: CpConfig) -> Self { + Self { + cfg, + registry: Registry::new(), + router: Router::new(), + rpc_id: AtomicU64::new(1), + } + } + + pub fn next_rpc_id(&self) -> u64 { + self.rpc_id.fetch_add(1, Ordering::Relaxed) + } +} + +pub fn app(state: Arc) -> AxumRouter { + AxumRouter::new() + .route("/cp", get(ws_handler)) + .route("/health", get(health)) + .with_state(state) +} + +async fn health() -> &'static str { + "ok" +} + +async fn ws_handler( + State(state): State>, + headers: HeaderMap, + ws: WebSocketUpgrade, +) -> axum::response::Response { + let key = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.strip_prefix("Bearer ")); + let identity = match key.and_then(|k| state.cfg.identity_for_key(k)) { + Some(id) => id.clone(), + None => { + warn!("WS rejected: missing or unknown auth key"); + return StatusCode::UNAUTHORIZED.into_response(); + } + }; + let max_frame = state.cfg.max_frame_bytes; + ws.max_message_size(max_frame) + .max_frame_size(max_frame) + .on_upgrade(move |socket| handle_connection(state, socket, identity)) +} + +async fn handle_connection(state: Arc, socket: WebSocket, identity: AgentIdentity) { + let (mut sink, mut stream) = socket.split(); + + // --- Registration: mandatory first frame --- + let register = loop { + match stream.next().await { + Some(Ok(Message::Text(text))) => break text, + Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue, + _ => { + warn!(agent = %identity.name, "connection closed before registration"); + return; + } + } + }; + let (reg, reg_rpc_id) = match parse_register(®ister, &identity) { + Ok(ok) => ok, + Err((id, err)) => { + let resp = JsonRpcErrorResponse::new(id, err); + let _ = sink + .send(Message::Text( + serde_json::to_string(&resp).expect("serializable").into(), + )) + .await; + return; + } + }; + + // Outbound channel for this connection. Bounded (review F5): a peer that + // cannot drain OUTBOUND_QUEUE frames is disconnected, not buffered. + let (tx, mut rx) = mpsc::channel::(OUTBOUND_QUEUE); + + let effective_max = match identity.max_delegated_sessions_cap { + Some(cap) => reg.max_delegated_sessions.min(cap), + None => reg.max_delegated_sessions, + }; + // The registry assigns the CP-generated handle (review F1): ownership + // and teardown never key on the client-supplied instance_id. + let handle = state.registry.register(Instance { + handle: 0, + namespace: identity.namespace.clone(), + name: identity.name.clone(), + agent_type: identity.agent_type.clone(), + instance_id: reg.instance_id.clone(), + labels: reg.labels.clone(), + max_delegated_sessions: effective_max, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx: tx.clone(), + }); + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + instance = %reg.instance_id, + handle, + r#type = %identity.agent_type, + max_sessions = effective_max, + "registered" + ); + + // Ack. The CP-generated handle is intentionally not disclosed. + let ack = RegisterAck { + protocol_version: PROTOCOL_VERSION, + heartbeat_interval_secs: state.cfg.heartbeat_interval_secs, + lease_expiry_secs: state.cfg.lease_expiry_secs, + effective_max_delegated_sessions: effective_max, + }; + let resp = JsonRpcResponse::new( + reg_rpc_id, + serde_json::to_value(&ack).expect("serializable"), + ); + if sink + .send(Message::Text( + serde_json::to_string(&resp).expect("serializable").into(), + )) + .await + .is_err() + { + teardown(&state, handle, &identity); + return; + } + + // --- Main loop: interleave inbound frames and outbound channel --- + loop { + tokio::select! { + outbound = rx.recv() => { + match outbound { + Some(text) => { + if sink.send(Message::Text(text.into())).await.is_err() { + break; + } + } + None => break, + } + } + inbound = stream.next() => { + match inbound { + Some(Ok(Message::Text(text))) => { + if let Some(reply) = handle_frame(&state, handle, &text) { + if sink.send(Message::Text(reply.into())).await.is_err() { + break; + } + } + } + Some(Ok(Message::Ping(p))) => { + if sink.send(Message::Pong(p)).await.is_err() { + break; + } + } + Some(Ok(Message::Close(_))) | None => break, + Some(Ok(_)) => {} // binary/pong ignored + Some(Err(e)) => { + warn!(handle, err = %e, "WS error"); + break; + } + } + } + } + } + + teardown(&state, handle, &identity); +} + +/// Deregister this connection's own registration (by handle — cannot touch +/// another connection's entry) and fail its in-flight delegations. +fn teardown(state: &Arc, handle: u64, identity: &AgentIdentity) { + state.registry.deregister(handle); + let mut next = || state.next_rpc_id(); + for (inst, frame) in state + .router + .fail_instance(&state.registry, handle, &mut next) + { + let _ = inst.tx.try_send(frame); + } + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + handle, + "disconnected" + ); +} + +/// Validate the registration frame against the authenticated identity. +/// Returns the parsed params and the request id, or an error payload. +fn parse_register( + text: &str, + identity: &AgentIdentity, +) -> Result<(RegisterParams, u64), (u64, ErrorObject)> { + let msg: JsonRpcMessage = match serde_json::from_str(text) { + Ok(m) => m, + Err(e) => { + return Err(( + 0, + ErrorObject::new(codes::INVALID_PARAMS, format!("malformed frame: {e}")), + )) + } + }; + let rpc_id = match msg.require_request_envelope() { + Ok(id) => id, + Err(err) => return Err((msg.id.unwrap_or(0), err)), + }; + if msg.method.as_deref() != Some(methods::REGISTER) { + return Err(( + rpc_id, + ErrorObject::new(codes::NOT_REGISTERED, "first frame must be cp/register"), + )); + } + let params: RegisterParams = match msg.params.and_then(|p| serde_json::from_value(p).ok()) { + Some(p) => p, + None => { + return Err(( + rpc_id, + ErrorObject::new(codes::INVALID_PARAMS, "invalid cp/register params"), + )) + } + }; + if params.protocol_version != PROTOCOL_VERSION { + return Err(( + rpc_id, + ErrorObject::new( + codes::UNSUPPORTED_VERSION, + format!( + "protocol version {} unsupported (CP speaks {})", + params.protocol_version, PROTOCOL_VERSION + ), + ), + )); + } + // Identity binding: claims must match the key's bound identity exactly. + if params.namespace != identity.namespace + || params.name != identity.name + || params.agent_type != identity.agent_type + { + return Err(( + rpc_id, + ErrorObject::new( + codes::IDENTITY_MISMATCH, + format!( + "registration claims {}/{} ({}) do not match the identity bound to this key", + params.namespace, params.name, params.agent_type + ), + ), + )); + } + if params.instance_id.trim().is_empty() { + return Err(( + rpc_id, + ErrorObject::new(codes::INVALID_PARAMS, "instance_id must be non-empty"), + )); + } + Ok((params, rpc_id)) +} + +/// Dispatch one post-registration frame. Returns an optional direct reply. +fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option { + let msg: JsonRpcMessage = match serde_json::from_str(text) { + Ok(m) => m, + Err(e) => { + let resp = JsonRpcErrorResponse::new( + 0, + ErrorObject::new(codes::INVALID_PARAMS, format!("malformed frame: {e}")), + ); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + }; + // Responses to CP-issued requests (forwarded delegates, cancels): v1 + // correlates by delegation_id inside result frames, so plain JSON-RPC + // acks are dropped. + let method = msg.method.as_deref()?.to_string(); + let rpc_id = match msg.require_request_envelope() { + Ok(id) => id, + Err(err) => { + let resp = JsonRpcErrorResponse::new(msg.id.unwrap_or(0), err); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + }; + // The sender's identity claims are never read from the frame: everything + // derives from the authenticated registration behind `handle`. + let me = state.registry.get(handle)?; + + macro_rules! params_or_err { + ($ty:ty) => { + match msg + .params + .clone() + .and_then(|p| serde_json::from_value::<$ty>(p).ok()) + { + Some(p) => p, + None => { + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new(codes::INVALID_PARAMS, "invalid params"), + ); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + } + }; + } + + match method.as_str() { + methods::HEARTBEAT => { + let _p = params_or_err!(crate::proto::HeartbeatParams); + state.registry.heartbeat(handle); + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } + methods::DELEGATE => { + let p = params_or_err!(DelegateParams); + if p.prompt.len() > state.cfg.max_prompt_bytes { + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new( + codes::INVALID_PARAMS, + format!( + "prompt exceeds max_prompt_bytes ({})", + state.cfg.max_prompt_bytes + ), + ), + ); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + let outcome = state.router.delegate( + &state.cfg, + &state.registry, + &me.namespace, + &me.name, + &me.agent_type, + handle, + p, + state.next_rpc_id(), + ); + let reply = match outcome { + DelegateOutcome::Accepted(ack) => serde_json::to_string(&JsonRpcResponse::new( + rpc_id, + serde_json::to_value(&ack).expect("serializable"), + )), + DelegateOutcome::Rejected(err) => { + serde_json::to_string(&JsonRpcErrorResponse::new(rpc_id, err)) + } + }; + Some(reply.expect("serializable")) + } + methods::DELEGATE_RESULT => { + let p = params_or_err!(DelegateResultParams); + if let Some((initiator, frame)) = state.router.complete( + &state.registry, + handle, + p, + state.cfg.max_result_bytes, + state.next_rpc_id(), + ) { + let _ = initiator.tx.try_send(frame); + } + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } + methods::CANCEL => { + let p = params_or_err!(CancelParams); + match state + .router + .cancel(&state.registry, handle, &p, state.next_rpc_id()) + { + Ok(forward) => { + if let Some((target, frame)) = forward { + let _ = target.tx.try_send(frame); + } + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } + Err(err) => Some( + serde_json::to_string(&JsonRpcErrorResponse::new(rpc_id, err)) + .expect("serializable"), + ), + } + } + other => { + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new(codes::METHOD_NOT_FOUND, format!("unknown method {other}")), + ); + Some(serde_json::to_string(&resp).expect("serializable")) + } + } +} + +/// Background sweeps: lease expiry and delegation deadlines. +pub async fn run_sweeper(state: Arc) { + let mut tick = tokio::time::interval(std::time::Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tick.tick().await; + + // Lease expiry → deregister + fail in-flight. + let lease = std::time::Duration::from_secs(state.cfg.lease_expiry_secs); + for handle in state.registry.expired(lease) { + warn!(handle, "lease expired — deregistering"); + state.registry.deregister(handle); + let mut next = || state.next_rpc_id(); + for (inst, frame) in state + .router + .fail_instance(&state.registry, handle, &mut next) + { + let _ = inst.tx.try_send(frame); + } + } + + // Deadline sweep. + let mut next = || state.next_rpc_id(); + for (inst, frame) in + state + .router + .sweep_deadlines(&state.registry, chrono::Utc::now(), &mut next) + { + let _ = inst.tx.try_send(frame); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::proto::AgentType; + + fn identity() -> AgentIdentity { + AgentIdentity { + key: "k".into(), + namespace: "prod".into(), + name: "koudu".into(), + agent_type: AgentType::Primary, + max_delegated_sessions_cap: None, + } + } + + #[test] + fn register_valid() { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-1" + } + }) + .to_string(); + let (params, rpc) = parse_register(&frame, &identity()).unwrap(); + assert_eq!(params.instance_id, "i-1"); + assert_eq!(rpc, 1); + } + + #[test] + fn register_identity_mismatch_rejected() { + for (ns, name, ty) in [ + ("dev", "koudu", "primary"), + ("prod", "other", "primary"), + ("prod", "koudu", "worker"), + ] { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 2, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": ns, + "name": name, + "type": ty, + "instance_id": "i-1" + } + }) + .to_string(); + let (_, err) = parse_register(&frame, &identity()).unwrap_err(); + assert_eq!(err.code, codes::IDENTITY_MISMATCH, "{ns}/{name}/{ty}"); + } + } + + #[test] + fn register_wrong_first_method_rejected() { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 3, "method": "cp/heartbeat", "params": {"instance_id": "i-1"} + }) + .to_string(); + let (_, err) = parse_register(&frame, &identity()).unwrap_err(); + assert_eq!(err.code, codes::NOT_REGISTERED); + } + + #[test] + fn register_unsupported_version_rejected() { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 4, "method": "cp/register", + "params": { + "protocol_version": 99, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-1" + } + }) + .to_string(); + let (_, err) = parse_register(&frame, &identity()).unwrap_err(); + assert_eq!(err.code, codes::UNSUPPORTED_VERSION); + } + + #[test] + fn register_empty_instance_id_rejected() { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 5, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": " " + } + }) + .to_string(); + let (_, err) = parse_register(&frame, &identity()).unwrap_err(); + assert_eq!(err.code, codes::INVALID_PARAMS); + } + + #[test] + fn register_invalid_envelope_rejected() { + // Missing jsonrpc field (review F4). + let no_ver = serde_json::json!({ + "id": 6, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-1" + } + }) + .to_string(); + let (_, err) = parse_register(&no_ver, &identity()).unwrap_err(); + assert_eq!(err.code, codes::INVALID_REQUEST); + + // Notification shape: no id. + let no_id = serde_json::json!({ + "jsonrpc": "2.0", "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-1" + } + }) + .to_string(); + let (_, err) = parse_register(&no_id, &identity()).unwrap_err(); + assert_eq!(err.code, codes::INVALID_REQUEST); + } +} diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 7a4a253d3..edf68c402 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -209,6 +209,69 @@ complete until the serving runtime returns a structured result frame The serving **runtime** emits this frame when the agent's turn ends — result delivery never depends on the sub-agent model "remembering" to report. +### v1 contract amendments (from PR #1465 review) + +The first implementation (`crates/openab-cp`) freezes the following +behaviors, resolving the review findings on identity, lifecycle, and +recovery semantics: + +- **Identity binding.** CP config owns an immutable identity table: auth key + → (`namespace`, `name`, `type`, optional capacity cap). The runtime's + registration claims are *verified against* the key's bound identity and + rejected on mismatch (`IDENTITY_MISMATCH`). Authorization never derives + from self-asserted registration fields. Keys are per-agent + (individually revocable) and presented as `Authorization: Bearer` on the + WebSocket upgrade — never in URLs. +- **CP-constructed chain.** `cp/delegate` carries only + `parent_delegation_id`; the CP derives the ancestry chain from its + in-flight table and the authenticated caller identity, then stamps it on + the forwarded frame. A runtime cannot forge ancestry, so depth/cycle + checks operate on trusted data. Policy (role, depth, cycle, namespace, + deadline caps) is enforced by the CP authoritatively; facade checks are + defense in depth only. +- **Registration lifecycle.** The first frame on a connection MUST be + `cp/register` (JSON-RPC 2.0 envelope validated — `jsonrpc: "2.0"` and a + request id are required; `protocol_version` field). Registrations are + keyed by a **CP-generated handle**, never the client-supplied + `instance_id`: a colliding `instance_id` cannot replace or tear down + another connection's registration, and all in-flight ownership checks + (completion, cancellation, parent linkage) compare handles. The ack + carries the heartbeat interval, lease window, and the effective (possibly + clamped) concurrency budget. Instances missing heartbeats past the lease + are deregistered; their in-flight delegations fail immediately with + `target_disconnected`. Heartbeats refresh the lease only — CP-owned + in-flight accounting is authoritative and never merged from runtime + reports. +- **Resource bounds.** The WS transport rejects messages over + `max_frame_bytes` before parsing; oversized `prompt`s are rejected + (`max_prompt_bytes`); per-connection outbound queues are bounded and a + peer that cannot drain its queue is treated as disconnected. Delegation + admission (duplicate check → target selection → capacity reservation → + in-flight insert) is one atomic sequence, and the in-flight entry exists + before the forward frame is sent. +- **Saturation = fast-fail.** When all matching targets are at capacity the + CP replies `SATURATED` immediately. The CP never queues — v1 has no + durable state, and a hidden in-memory queue would contradict that. + `NO_TARGET` (nothing matches) is a distinct error. +- **CP restart semantics.** The in-flight table is in-memory. After a CP + restart, in-flight delegations end as initiator-side timeouts (the + propagated deadline is the upper bound); late `cp/delegate_result` frames + for unknown ids are acknowledged, logged, and dropped so reconnecting + runtimes do not error-loop. +- **Timeout and disconnect synthesis.** A deadline sweep terminates overdue + delegations: the initiator receives a synthesized `timeout` result and the + serving runtime a best-effort `cp/cancel` (stop burning tokens). Worker + disconnect → `target_disconnected` to the initiator; initiator disconnect + → best-effort `cp/cancel` downstream. +- **Result size cap.** `cp/delegate_result.result` larger than the + configured `max_result_bytes` (default 256 KiB) is truncated head-first + with an explicit marker. +- **Idempotency.** `delegation_id` is the caller-generated idempotency key; + a duplicate in-flight id is rejected (`DUPLICATE_DELEGATION`). Only the + instance a delegation was routed to may complete it; only the initiating + instance may cancel it. + + --- ## 5. Delegation Policy @@ -365,14 +428,22 @@ of scope for v1. ## 11. Open Questions -1. **Streaming intermediate output** — should `cp/delegate` stream - `session/update`-style chunks back to the primary, or only the final - result frame? v1 leans final-only; streaming is additive. +1. ~~**Streaming intermediate output**~~ — *resolved (PR #1465 review): + committed scope as a fast-follow behind the same wire contract. Worker + runtimes will stream `session/update`-style chunks back through the CP. + Rationale: streaming is the observability substrate, not a feature — it + restores the free human visibility that Discord-mediated collaboration + provides today. It enables a read-only observer endpoint on the CP + (e.g. `wss://cp/.../observe?ns=prod`; separate read-only credential + class, namespace-scoped) so a human can tail all delegation traffic + across the fleet from one terminal. v1 ships final-result-only; the + stream frame shape is reserved in the wire contract. 2. **CP high availability** — single instance + fast re-registration is - acceptable for v1; is active/standby needed before multi-tenant use? -3. **Human-visibility directives** — should a primary be able to mirror - selected delegation traffic into a Discord thread (observability) via - existing output directives? + acceptable for v1 (restart semantics are now defined in §4); is + active/standby needed before multi-tenant use? +3. **Human-visibility directives** — Discord mirroring becomes a consumer + of the delegation stream (Q1) rather than a separate mechanism; exact + directive syntax TBD when streaming lands. 4. **AgentCore/remote runtimes** — an `agentcore-acp`-backed OAB registers like any other runtime; verify deadline propagation across the SDK boundary. From 33ed835f67396007d994878b89935a39638a95e6 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Tue, 11 Aug 2026 00:10:17 -0400 Subject: [PATCH 02/11] fix(cp): loopback-default bind guard + truncation marker respects result cap Round-2 review F4/F5: bearer keys never cross cleartext non-loopback TCP without an explicit allow_insecure_bind acknowledgment (TLS proxy or private network required), and result truncation counts the marker against max_result_bytes. --- crates/openab-cp/cp.toml.example | 6 ++- crates/openab-cp/src/config.rs | 52 +++++++++++++++++++++++++- crates/openab-cp/src/router.rs | 64 +++++++++++++++++++++++++------- 3 files changed, 106 insertions(+), 16 deletions(-) diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index 8373d2839..21cedd4f6 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -5,7 +5,11 @@ # are IMMUTABLE and owned by this file — a runtime's own [control_plane] # config is verified against them at registration and rejected on mismatch. -listen = "0.0.0.0:9800" +# The CP terminates no TLS itself. The safe default is loopback; to bind a +# non-loopback address you must front it with a TLS proxy (wss://) or a +# private network (e.g. tailnet) AND set allow_insecure_bind = true. +listen = "127.0.0.1:9800" +# allow_insecure_bind = true # Runtimes must heartbeat at this interval; missing heartbeats past the lease # window deregisters the instance and fails its in-flight delegations. diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index c92a5988f..ce3c2751b 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -14,10 +14,19 @@ use crate::proto::AgentType; #[derive(Debug, Clone, Deserialize)] pub struct CpConfig { - /// Bind address, e.g. "0.0.0.0:9800". + /// Bind address. Defaults to loopback: the CP carries bearer + /// credentials and terminates no TLS itself, so non-loopback binds + /// require `allow_insecure_bind = true` and a TLS-terminating proxy + /// (or a private overlay network) in front. #[serde(default = "default_listen")] pub listen: String, + /// Explicit opt-in to bind a non-loopback address WITHOUT in-process + /// TLS. Only set this when a trusted TLS proxy terminates wss:// in + /// front of the CP, or the network is private (e.g. a tailnet). + #[serde(default)] + pub allow_insecure_bind: bool, + /// Heartbeat interval communicated to runtimes. #[serde(default = "default_heartbeat_secs")] pub heartbeat_interval_secs: u64, @@ -60,7 +69,7 @@ pub struct CpConfig { } fn default_listen() -> String { - "0.0.0.0:9800".to_string() + "127.0.0.1:9800".to_string() } fn default_heartbeat_secs() -> u64 { 15 @@ -158,6 +167,17 @@ impl CpConfig { if self.lease_expiry_secs <= self.heartbeat_interval_secs { bail!("lease_expiry_secs must exceed heartbeat_interval_secs"); } + // Bearer keys over cleartext TCP must never reach an untrusted + // network: non-loopback binds require the explicit override + // (review round-2 F4). + if !self.allow_insecure_bind && !is_loopback(&self.listen) { + bail!( + "listen = \"{}\" is not loopback and the CP terminates no TLS. \ + Put a TLS proxy (wss://) or a private network in front and set \ + allow_insecure_bind = true to acknowledge this", + self.listen + ); + } Ok(()) } @@ -181,6 +201,20 @@ impl CpConfig { } } +/// Whether a `host:port` bind address is loopback. +fn is_loopback(listen: &str) -> bool { + let host = match listen.rsplit_once(':') { + Some((h, _)) => h.trim_start_matches('[').trim_end_matches(']'), + None => listen, + }; + if host == "localhost" { + return true; + } + host.parse::() + .map(|ip| ip.is_loopback()) + .unwrap_or(false) +} + /// `${ENV_VAR}` expansion, mirroring openab-core config behavior. Unset vars /// expand to the empty string (validation then rejects empty keys loudly). fn expand_env(raw: &str) -> String { @@ -281,6 +315,20 @@ lease_expiry_secs = 30 assert!(cfg.validate().is_err()); } + #[test] + fn non_loopback_bind_requires_override() { + let cfg: CpConfig = toml::from_str("listen = \"0.0.0.0:9800\"").unwrap(); + assert!(cfg.validate().is_err()); + let cfg: CpConfig = + toml::from_str("listen = \"0.0.0.0:9800\"\nallow_insecure_bind = true").unwrap(); + cfg.validate().unwrap(); + // Loopback variants pass without the override. + for l in ["127.0.0.1:9800", "localhost:9800", "[::1]:9800"] { + let cfg: CpConfig = toml::from_str(&format!("listen = \"{l}\"")).unwrap(); + cfg.validate().unwrap(); + } + } + #[test] fn env_expansion() { std::env::set_var("CP_TEST_KEY_XYZ", "sekrit"); diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index 4aae8ccaf..b314b62d1 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -261,19 +261,20 @@ impl Router { registry.adjust_sessions(entry.to_handle, -1); - // Truncate oversized results (keep the head; delegation already ran). + // Truncate oversized results (keep the head; delegation already + // ran). The marker counts against the cap: the final value never + // exceeds max_result_bytes (review round-2 F5). if let Some(r) = ¶ms.result { if r.len() > max_result_bytes { - let mut cut = max_result_bytes; - while !r.is_char_boundary(cut) { - cut -= 1; + let marker = format!("\n…[truncated by control plane: {} bytes total]", r.len()); + let budget = max_result_bytes.saturating_sub(marker.len()); + let cut = floor_char_boundary(r, budget); + let mut out = format!("{}{}", &r[..cut], marker); + if out.len() > max_result_bytes { + // Degenerate tiny cap: keep whatever fits. + out.truncate(floor_char_boundary(&out, max_result_bytes)); } - params.result = Some(format!( - "{}\n…[truncated by control plane: {} of {} bytes]", - &r[..cut], - cut, - r.len() - )); + params.result = Some(out); } } @@ -463,6 +464,15 @@ impl Router { } } +/// Largest index `<= max` that lands on a char boundary of `s`. +fn floor_char_boundary(s: &str, max: usize) -> usize { + let mut cut = max.min(s.len()); + while cut > 0 && !s.is_char_boundary(cut) { + cut -= 1; + } + cut +} + impl Default for Router { fn default() -> Self { Self::new() @@ -765,14 +775,42 @@ type = "worker" let result = DelegateResultParams { delegation_id: "d-1".into(), status: DelegationStatus::Completed, - result: Some("x".repeat(100)), + result: Some("x".repeat(200)), error: None, }; + let cap = 96usize; let (_, frame) = w .router - .complete(&w.registry, w.h_worker, result, 10, 2) + .complete(&w.registry, w.h_worker, result, cap, 2) + .unwrap(); + let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); + let out = v["params"]["result"].as_str().unwrap(); + assert!(out.contains("truncated by control plane")); + assert!( + out.len() <= cap, + "marker must count against the cap: {} > {}", + out.len(), + cap + ); + + // Degenerate tiny cap still never exceeds the cap. + assert!(matches!( + do_delegate(&w, delegate_params("d-2", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + let result2 = DelegateResultParams { + delegation_id: "d-2".into(), + status: DelegationStatus::Completed, + result: Some("y".repeat(100)), + error: None, + }; + let (_, frame2) = w + .router + .complete(&w.registry, w.h_worker, result2, 8, 3) .unwrap(); - assert!(frame.contains("truncated by control plane")); + let v2: serde_json::Value = serde_json::from_str(&frame2).unwrap(); + assert!(v2["params"]["result"].as_str().unwrap().len() <= 8); } #[test] From b9850dae9344ba0f3773df4af234f0764d453753 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Wed, 12 Aug 2026 07:28:40 -0400 Subject: [PATCH 03/11] fix(cp): close lease-expired connections, atomic ownership checks, namespace-scoped delegation ids, admission bounds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 — lease expiry left a zombie connection. The sweeper deregistered the instance and failed its in-flight work but never told the connection task to stop, and the task holds its own clone of the outbound sender, so `rx.recv()` never ended. The socket stayed open bound to a registration that no longer existed: every later frame (heartbeats included) hit `registry.get(handle) -> None` and got no reply, and the client could not recover because registration is first-frame-only. Each connection now creates a `watch` shutdown signal (`registry::shutdown_signal`), subscribes to it before registering, and hands a clone to the registry, which keeps it in a CP-internal `Entry` alongside the public `Instance` (nothing routed or serialized carries it). `sweep_leases` calls `signal_shutdown` before `deregister`, and the main loop `select!`s on the signal, sends a Close frame, and tears down; the client then reconnects, re-authenticates, and registers again. Proven end-to-end over a real loopback socket by tests/ws_lifecycle.rs `lease_expiry_closes_the_connection_and_permits_reregistration`, and at the sweeper level by `sweep_leases_signals_the_connection_before_dropping_it` / registry `signal_shutdown_reaches_the_owning_connection`. F2 — remove-then-check race in `complete` and `cancel`. Both removed the in-flight entry, released the lock, validated the caller's handle, and reinserted on mismatch. A genuine result landing in that window saw an unknown id and was silently dropped, the entry was briefly invisible to the deadline sweep, and the wrong-handle frame's reinsert left the delegation to stall to its deadline. Ownership is now validated under the SAME lock acquisition that removes: `get()`, check the handle, `remove()` only on success, never reinsert. Proven by `wrong_handle_result_never_hides_the_genuine_one` (both frame orders), `genuine_result_survives_concurrent_non_owner_frames` (200 barrier-synchronised races), and `refused_cancel_leaves_the_delegation_cancellable` for the cancel path. The pre-existing test `result_from_wrong_handle_dropped_and_restored` was renamed to `result_from_wrong_handle_dropped_and_entry_untouched`: its name asserted the remove-and-reinsert behaviour this commit removes. Its body and assertions are unchanged. F3 — global delegation_id keyspace crossed namespaces. The in-flight table was keyed on the client-supplied `delegation_id` alone, so one namespace's ids were observable from another: a colliding id was denied with `DUPLICATE_DELEGATION`, and `cp/cancel` distinguished "no such id" (`INVALID_PARAMS`) from "someone else's live id" (`POLICY_DENIED`) — a cross-tenant existence oracle. The table is now keyed by a composite `DelegationKey { namespace, delegation_id }`, and parent-chain resolution is scoped to the caller's namespace too, so a foreign parent id cannot be borrowed for its trusted chain and deadline budget. The namespace always comes from the authenticated registration behind the handle, never from the frame. `cp/cancel` returns one indistinguishable error — same code and same message — for an unknown id and for not-the-initiator; only the CP's own logs keep the distinction. Proven by `same_delegation_id_in_two_namespaces_is_independent` (one `d-1` live in prod and dev at once, each completing to its own initiator), `parent_lookup_is_namespace_scoped`, and `cancel_refusals_are_byte_identical` (the two error objects compared as serialized JSON). F4 — unbounded authenticated pre-registration sockets. A valid key could open unlimited sockets and park them in the registration loop forever, kept alive by pings. Two bounds, both new config fields with serde defaults so absent fields keep working (documented in cp.toml.example, rejected when zero by `CpConfig::validate`): (a) `register_timeout_secs` (default 10) is a deadline on the mandatory `cp/register` first frame, measured from the completed upgrade — pings are skipped inside the timeout rather than extending it; (b) `max_connections_per_identity` (default 8) is acquired in the upgrade handler, before `on_upgrade`, so pre-registration sockets count too and an over-quota peer is refused at the HTTP layer with 503. The slot is held by an RAII `ConnPermit` whose `Drop` releases it, so no early return — including a failed ack write, a rejected registration, or an upgrade that never completes — can leak it. Proven by ws_lifecycle `ping_only_pre_registration_socket_is_closed_at_the_deadline`, `registration_after_the_deadline_is_not_accepted`, `connection_quota_rejects_over_limit_and_recycles_on_disconnect`, `pre_registration_sockets_count_against_the_quota`, and the unit tests `conn_quota_bounds_and_recycles_slots` / `conn_quota_is_per_identity` / `admission_bounds_default_and_are_validated`. The end-to-end tests need a WebSocket client, added as a dev-dependency on tokio-tungstenite 0.29 — the version axum 0.8.9 already resolves to, so Cargo.lock gains only openab-cp's dependency edge and no new package version. Audit follow-up: added genuine_cancel_survives_concurrent_non_owner_frames — the cancel-side racing regression mirroring the complete-side one (200 barrier-synchronised rounds, non-initiator cancel racing the genuine initiator cancel; genuine must always win, capacity released exactly once). --- Cargo.lock | 1 + crates/openab-cp/Cargo.toml | 5 + crates/openab-cp/cp.toml.example | 12 + crates/openab-cp/src/config.rs | 53 +++ crates/openab-cp/src/registry.rs | 110 ++++- crates/openab-cp/src/router.rs | 618 ++++++++++++++++++++++--- crates/openab-cp/src/server.rs | 311 +++++++++++-- crates/openab-cp/tests/ws_lifecycle.rs | 254 ++++++++++ 8 files changed, 1247 insertions(+), 117 deletions(-) create mode 100644 crates/openab-cp/tests/ws_lifecycle.rs diff --git a/Cargo.lock b/Cargo.lock index 98dfb7d52..da40e70d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2589,6 +2589,7 @@ dependencies = [ "serde_json", "subtle", "tokio", + "tokio-tungstenite 0.29.0", "toml", "tracing", "tracing-subscriber", diff --git a/crates/openab-cp/Cargo.toml b/crates/openab-cp/Cargo.toml index eff123bff..debff2e93 100644 --- a/crates/openab-cp/Cargo.toml +++ b/crates/openab-cp/Cargo.toml @@ -20,3 +20,8 @@ chrono = { version = "0.4", features = ["serde"] } parking_lot = "0.12" clap = { version = "4", features = ["derive"] } subtle = "2" + +[dev-dependencies] +# WebSocket client for the end-to-end admission/lifecycle tests. Same version +# axum 0.8 already uses, so it adds no new dependency to the workspace. +tokio-tungstenite = "0.29" diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index 21cedd4f6..9b5df7943 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -28,6 +28,18 @@ max_frame_bytes = 1048576 # Delegation prompts larger than this are rejected. max_prompt_bytes = 262144 +# A connection must send its cp/register first frame within this many seconds +# of the WebSocket upgrade, or it is closed. Authentication alone is not a +# bound: without this, an authenticated peer could park idle sockets +# indefinitely (pings do not extend the deadline). +register_timeout_secs = 10 + +# Maximum simultaneous connections per identity, counted from the upgrade — +# so sockets that have not registered yet count too — and released as soon as +# a connection ends. Replicas share one identity, so this is also the ceiling +# on concurrent replicas of one logical agent. +max_connections_per_identity = 8 + [[agents]] key = "${CP_KEY_KOUDU}" # per-agent secret, never shared namespace = "prod" diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index ce3c2751b..21d672b3b 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -57,6 +57,21 @@ pub struct CpConfig { #[serde(default = "default_max_prompt_bytes")] pub max_prompt_bytes: usize, + /// Deadline for the mandatory `cp/register` first frame, in seconds from + /// the completed WebSocket upgrade. A connection that authenticates but + /// never registers is closed when this elapses (review round-3 F4): + /// otherwise an authenticated peer could park unlimited sockets in the + /// pre-registration state, keeping them alive with pings forever. + #[serde(default = "default_register_timeout_secs")] + pub register_timeout_secs: u64, + + /// Maximum simultaneous connections per identity, counted from the + /// upgrade (so pre-registration sockets count too) and released on every + /// exit path (review round-3 F4). Replicas of one logical agent share one + /// identity, so this is the replica ceiling as well. + #[serde(default = "default_max_connections_per_identity")] + pub max_connections_per_identity: u32, + /// Identity table: auth key → immutable claims. /// Keyed by the key id (`kid`), with the secret alongside, so logs can /// reference identities without printing secrets. @@ -89,6 +104,12 @@ fn default_max_frame_bytes() -> usize { fn default_max_prompt_bytes() -> usize { 256 * 1024 } +fn default_register_timeout_secs() -> u64 { + 10 +} +fn default_max_connections_per_identity() -> u32 { + 8 +} /// Immutable identity claims bound to one auth key. #[derive(Debug, Clone, Deserialize)] @@ -167,6 +188,14 @@ impl CpConfig { if self.lease_expiry_secs <= self.heartbeat_interval_secs { bail!("lease_expiry_secs must exceed heartbeat_interval_secs"); } + // Admission bounds must actually bound something (review round-3 F4): + // zero would mean "no registration deadline" / "no connection allowed". + if self.register_timeout_secs == 0 { + bail!("register_timeout_secs must be greater than 0"); + } + if self.max_connections_per_identity == 0 { + bail!("max_connections_per_identity must be at least 1"); + } // Bearer keys over cleartext TCP must never reach an untrusted // network: non-loopback binds require the explicit override // (review round-2 F4). @@ -329,6 +358,30 @@ lease_expiry_secs = 30 } } + #[test] + fn admission_bounds_default_and_are_validated() { + // Review round-3 F4: absent fields keep working (serde defaults) and + // a zero bound is rejected rather than silently disabling the guard. + let cfg: CpConfig = toml::from_str(base_toml()).unwrap(); + cfg.validate().unwrap(); + assert_eq!(cfg.register_timeout_secs, 10); + assert_eq!(cfg.max_connections_per_identity, 8); + + let explicit: CpConfig = + toml::from_str("register_timeout_secs = 3\nmax_connections_per_identity = 2").unwrap(); + explicit.validate().unwrap(); + assert_eq!(explicit.register_timeout_secs, 3); + assert_eq!(explicit.max_connections_per_identity, 2); + + for bad in [ + "register_timeout_secs = 0", + "max_connections_per_identity = 0", + ] { + let cfg: CpConfig = toml::from_str(bad).unwrap(); + assert!(cfg.validate().is_err(), "{bad} must be rejected"); + } + } + #[test] fn env_expansion() { std::env::set_var("CP_TEST_KEY_XYZ", "sekrit"); diff --git a/crates/openab-cp/src/registry.rs b/crates/openab-cp/src/registry.rs index d97663bad..dd7cd47cf 100644 --- a/crates/openab-cp/src/registry.rs +++ b/crates/openab-cp/src/registry.rs @@ -8,10 +8,11 @@ use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; use std::time::{Duration, Instant}; use parking_lot::RwLock; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use crate::proto::AgentType; @@ -23,6 +24,32 @@ pub type FrameTx = mpsc::Sender; /// Capacity of each per-connection outbound queue. pub const OUTBOUND_QUEUE: usize = 256; +/// Shutdown signal for one WS connection, held by the registry so the CP can +/// terminate a connection it no longer considers registered (review round-3 +/// F1: lease expiry must close the socket — otherwise the connection task +/// lives on with a registry entry that no longer exists, silently dropping +/// every subsequent frame and unable to re-register, since registration is +/// first-frame-only). +/// +/// `watch` (not `oneshot`) so the connection task can select on it repeatedly, +/// and wrapped in `Arc` so the registry entry and the connection task share +/// one signal without either side's drop cancelling it. +pub type ShutdownTx = Arc>; + +/// Create a fresh connection shutdown signal. The connection task keeps the +/// returned handle (to `subscribe()`), the registry keeps a clone. +pub fn shutdown_signal() -> ShutdownTx { + Arc::new(watch::channel(false).0) +} + +/// Registry slot: the public instance view plus CP-internal connection +/// control that is deliberately not part of `Instance` (nothing routed or +/// serialized should carry it). +struct Entry { + inst: Instance, + shutdown: ShutdownTx, +} + /// A live, authenticated, registered runtime instance. #[derive(Clone, Debug)] pub struct Instance { @@ -65,7 +92,7 @@ impl Instance { #[derive(Default)] pub struct Registry { /// Keyed by CP-generated registration handle. - inner: RwLock>, + inner: RwLock>, next_handle: AtomicU64, } @@ -78,17 +105,41 @@ impl Registry { /// (returned). Re-registrations (reconnects) get a new handle; the stale /// entry disappears when its socket closes or its lease expires — it can /// never be replaced by another connection's registration. - pub fn register(&self, mut inst: Instance) -> u64 { + /// + /// `shutdown` is the owning connection's termination signal: the CP + /// triggers it whenever it drops the registration on its own initiative + /// (lease expiry — review round-3 F1). + pub fn register_conn(&self, mut inst: Instance, shutdown: ShutdownTx) -> u64 { let handle = self.next_handle.fetch_add(1, Ordering::Relaxed) + 1; inst.handle = handle; - self.inner.write().insert(handle, inst); + self.inner.write().insert(handle, Entry { inst, shutdown }); handle } + /// Register an instance with a detached shutdown signal (no connection + /// task is listening). For tests and non-WS callers. + pub fn register(&self, inst: Instance) -> u64 { + self.register_conn(inst, shutdown_signal()) + } + + /// Ask the owning connection task to close. Returns whether a live + /// registration was signalled. Must be called BEFORE `deregister`, which + /// drops the registry's handle on the signal. + pub fn signal_shutdown(&self, handle: u64) -> bool { + match self.inner.read().get(&handle) { + Some(e) => { + // `send_replace` cannot fail even with no receivers left. + e.shutdown.send_replace(true); + true + } + None => false, + } + } + /// Remove an instance by its registration handle (disconnect or lease /// expiry). Only the owning connection or the sweeper knows the handle. pub fn deregister(&self, handle: u64) -> Option { - self.inner.write().remove(&handle) + self.inner.write().remove(&handle).map(|e| e.inst) } /// Refresh the lease. The runtime-reported session count is intentionally @@ -97,8 +148,8 @@ impl Registry { pub fn heartbeat(&self, handle: u64) -> bool { let mut g = self.inner.write(); match g.get_mut(&handle) { - Some(i) => { - i.last_heartbeat = Instant::now(); + Some(e) => { + e.inst.last_heartbeat = Instant::now(); true } None => false, @@ -111,13 +162,13 @@ impl Registry { self.inner .read() .values() - .filter(|i| now.duration_since(i.last_heartbeat) > lease) - .map(|i| i.handle) + .filter(|e| now.duration_since(e.inst.last_heartbeat) > lease) + .map(|e| e.inst.handle) .collect() } pub fn get(&self, handle: u64) -> Option { - self.inner.read().get(&handle).cloned() + self.inner.read().get(&handle).map(|e| e.inst.clone()) } /// Select a serving instance within `namespace` by exact name or labels. @@ -136,6 +187,7 @@ impl Registry { let g = self.inner.read(); let mut matches: Vec<&Instance> = g .values() + .map(|e| &e.inst) .filter(|i| i.namespace == namespace) .filter(|i| match name { Some(n) => i.name == n, @@ -173,8 +225,8 @@ impl Registry { /// Adjust the CP-owned in-flight count for an instance. pub fn adjust_sessions(&self, handle: u64, delta: i32) { let mut g = self.inner.write(); - if let Some(i) = g.get_mut(&handle) { - i.active_sessions = i.active_sessions.saturating_add_signed(delta); + if let Some(e) = g.get_mut(&handle) { + e.inst.active_sessions = e.inst.active_sessions.saturating_add_signed(delta); } } @@ -183,6 +235,7 @@ impl Registry { self.inner .read() .values() + .map(|e| &e.inst) .filter(|i| i.namespace == namespace) .cloned() .collect() @@ -347,4 +400,37 @@ mod tests { r.adjust_sessions(h, -5); assert_eq!(r.get(h).unwrap().active_sessions, 0); } + + #[tokio::test] + async fn signal_shutdown_reaches_the_owning_connection() { + // Review round-3 F1: the CP must be able to terminate a connection + // whose registration it drops on its own initiative. + let r = Registry::new(); + let sig = shutdown_signal(); + let mut rx = sig.subscribe(); + let h = r.register_conn(inst("prod", "w1", "i-1", 1), sig); + assert!(!*rx.borrow()); + + assert!(r.signal_shutdown(h)); + rx.changed().await.unwrap(); + assert!(*rx.borrow(), "connection task must observe the signal"); + + // After deregistration there is nothing left to signal. + r.deregister(h); + assert!(!r.signal_shutdown(h)); + } + + #[tokio::test] + async fn signal_shutdown_survives_registry_drop_of_the_entry() { + // The connection task keeps its own handle on the signal, so a + // signal delivered before deregistration is never lost. + let r = Registry::new(); + let sig = shutdown_signal(); + let mut rx = sig.subscribe(); + let h = r.register_conn(inst("prod", "w1", "i-1", 1), sig); + r.signal_shutdown(h); + r.deregister(h); + rx.changed().await.unwrap(); + assert!(*rx.borrow()); + } } diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index b314b62d1..57d77a05e 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -34,6 +34,10 @@ use crate::registry::{Instance, Registry, SelectError}; /// registration handles, never client-supplied ids (review F1). #[derive(Clone)] pub struct InFlight { + /// Namespace the delegation lives in — part of its identity (review + /// round-3 F3): `delegation_id` is client-supplied and only unique within + /// the namespace that produced it. + pub namespace: String, pub delegation_id: String, /// Authenticated initiator (`namespace/name`) and its registration handle. pub from_logical: String, @@ -47,8 +51,30 @@ pub struct InFlight { pub chain: Vec, } +/// In-flight table key: `(namespace, delegation_id)` (review round-3 F3). +/// +/// Keying on the client-supplied `delegation_id` alone made one namespace's +/// ids observable from another: a colliding id was denied with +/// `DUPLICATE_DELEGATION`, and `cp/cancel` distinguished "no such id" from +/// "someone else's live id" — a cross-tenant existence oracle. The composite +/// key confines both to the namespace that owns the id. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct DelegationKey { + namespace: String, + delegation_id: String, +} + +impl DelegationKey { + fn new(namespace: &str, delegation_id: &str) -> Self { + Self { + namespace: namespace.to_string(), + delegation_id: delegation_id.to_string(), + } + } +} + pub struct Router { - inflight: Mutex>, + inflight: Mutex>, /// Serializes the delegate admission sequence (duplicate check → target /// selection → capacity reservation → in-flight insert) so concurrent /// requests cannot double-admit one id or oversubscribe capacity @@ -92,7 +118,11 @@ impl Router { // in-flight insertion all happen under this guard. let _admission = self.admission.lock(); - if self.inflight.lock().contains_key(¶ms.delegation_id) { + // Delegation identity is namespace-scoped (review round-3 F3): the + // same id in another namespace is a different delegation, so it + // neither collides here nor leaks its existence. + let key = DelegationKey::new(from_namespace, ¶ms.delegation_id); + if self.inflight.lock().contains_key(&key) { return DelegateOutcome::Rejected(ErrorObject::new( codes::DUPLICATE_DELEGATION, format!("delegation {} is already in flight", params.delegation_id), @@ -111,18 +141,22 @@ impl Router { // Parent linkage: chain and deadline derive from the CP's own table, // never from the client. The caller must BE the instance serving the // parent delegation — otherwise any runtime knowing a live id could - // borrow its trusted chain and deadline budget (review F3). Unknown - // and unauthorized parent ids return the same error (no enumeration). + // borrow its trusted chain and deadline budget (review F3). The + // lookup is namespace-scoped (review round-3 F3). Unknown and + // unauthorized parent ids return the same error (no enumeration). let (parent_chain, parent_deadline) = match ¶ms.parent_delegation_id { - Some(pid) => match self.inflight.lock().get(pid) { - Some(p) if p.to_handle == from_handle => (p.chain.clone(), Some(p.deadline)), - _ => { - return DelegateOutcome::Rejected(ErrorObject::new( - codes::INVALID_PARAMS, - format!("parent delegation {pid} is not in flight for this instance"), - )) + Some(pid) => { + let parent_key = DelegationKey::new(from_namespace, pid); + match self.inflight.lock().get(&parent_key) { + Some(p) if p.to_handle == from_handle => (p.chain.clone(), Some(p.deadline)), + _ => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + format!("parent delegation {pid} is not in flight for this instance"), + )) + } } - }, + } None => (Vec::new(), None), }; @@ -186,6 +220,7 @@ impl Router { // back if the send fails. registry.adjust_sessions(target.handle, 1); let entry = InFlight { + namespace: from_namespace.to_string(), delegation_id: params.delegation_id.clone(), from_logical, from_handle, @@ -194,13 +229,11 @@ impl Router { deadline: params.deadline, chain, }; - self.inflight - .lock() - .insert(params.delegation_id.clone(), entry.clone()); + self.inflight.lock().insert(key.clone(), entry.clone()); if target.tx.try_send(text).is_err() { // Disconnected or backpressured beyond its queue: roll back. - self.inflight.lock().remove(¶ms.delegation_id); + self.inflight.lock().remove(&key); registry.adjust_sessions(target.handle, -1); return DelegateOutcome::Rejected(ErrorObject::new( codes::TARGET_DISCONNECTED, @@ -226,6 +259,12 @@ impl Router { /// Handle `cp/delegate_result` from the serving runtime. Returns the /// initiator-bound frame if the delegation is known; unknown ids (e.g. /// results arriving after a CP restart) are dropped with a log. + /// + /// Ownership is validated under the SAME lock acquisition that removes + /// the entry (review round-3 F2): the previous remove-check-reinsert + /// dance opened a window in which a genuine result saw an empty table and + /// was dropped, and left the entry momentarily invisible to the deadline + /// sweep. pub fn complete( &self, registry: &Registry, @@ -234,30 +273,47 @@ impl Router { max_result_bytes: usize, next_rpc_id: u64, ) -> Option<(Instance, String)> { - let entry = { self.inflight.lock().remove(¶ms.delegation_id) }; - let entry = match entry { - Some(e) => e, + // The namespace comes from the authenticated sender's registration, + // never from the frame (review round-3 F3). + let namespace = match registry.get(serving_handle) { + Some(i) => i.namespace, None => { warn!( + handle = serving_handle, delegation = %params.delegation_id, - "result for unknown delegation (late arrival or CP restart) — dropped" + "result from an unregistered connection — dropped" ); return None; } }; - if entry.to_handle != serving_handle { - // Only the instance the delegation was routed to may complete it. - warn!( - delegation = %params.delegation_id, - expected = entry.to_handle, - got = serving_handle, - "result from unexpected instance — dropped, delegation restored" - ); - self.inflight - .lock() - .insert(params.delegation_id.clone(), entry); - return None; - } + let key = DelegationKey::new(&namespace, ¶ms.delegation_id); + let entry = { + let mut g = self.inflight.lock(); + match g.get(&key) { + Some(e) if e.to_handle == serving_handle => {} + Some(e) => { + // Only the instance the delegation was routed to may + // complete it. The entry stays exactly where it is. + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + expected = e.to_handle, + got = serving_handle, + "result from unexpected instance — dropped, delegation untouched" + ); + return None; + } + None => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "result for unknown delegation (late arrival or CP restart) — dropped" + ); + return None; + } + } + g.remove(&key).expect("present under the same lock") + }; registry.adjust_sessions(entry.to_handle, -1); @@ -301,6 +357,13 @@ impl Router { /// Handle `cp/cancel` from the initiator. Returns the frame to forward /// to the serving runtime, if the delegation is in flight and owned by /// the caller. + /// + /// Ownership is validated under the same lock acquisition that removes + /// the entry (review round-3 F2 — no remove/reinsert window), and every + /// refusal returns ONE byte-identical error (review round-3 F3): an + /// unknown id and another instance's live id are indistinguishable to the + /// caller, so `cp/cancel` cannot be used to probe for delegation ids. + /// The distinction is kept in the CP's own logs only. pub fn cancel( &self, registry: &Registry, @@ -308,25 +371,47 @@ impl Router { params: &CancelParams, next_rpc_id: u64, ) -> Result, ErrorObject> { - let entry = { self.inflight.lock().remove(¶ms.delegation_id) }; - let entry = match entry { - Some(e) => e, + let refused = || { + ErrorObject::new( + codes::POLICY_DENIED, + "delegation is not in flight for this instance", + ) + }; + let namespace = match registry.get(from_handle) { + Some(i) => i.namespace, None => { - return Err(ErrorObject::new( - codes::INVALID_PARAMS, - format!("delegation {} is not in flight", params.delegation_id), - )) + warn!( + handle = from_handle, + "cancel from an unregistered connection" + ); + return Err(refused()); } }; - if entry.from_handle != from_handle { - self.inflight - .lock() - .insert(params.delegation_id.clone(), entry); - return Err(ErrorObject::new( - codes::POLICY_DENIED, - "only the initiating instance may cancel a delegation", - )); - } + let key = DelegationKey::new(&namespace, ¶ms.delegation_id); + let entry = { + let mut g = self.inflight.lock(); + match g.get(&key) { + Some(e) if e.from_handle == from_handle => {} + Some(_) => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + handle = from_handle, + "cancel refused: only the initiating instance may cancel" + ); + return Err(refused()); + } + None => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "cancel refused: delegation not in flight" + ); + return Err(refused()); + } + } + g.remove(&key).expect("present under the same lock") + }; registry.adjust_sessions(entry.to_handle, -1); info!(delegation = %params.delegation_id, "delegation cancelled by initiator"); let target = registry.get(entry.to_handle); @@ -355,12 +440,12 @@ impl Router { let mut affected = Vec::new(); let entries: Vec = { let mut g = self.inflight.lock(); - let ids: Vec = g - .values() - .filter(|e| e.to_handle == handle || e.from_handle == handle) - .map(|e| e.delegation_id.clone()) + let keys: Vec = g + .iter() + .filter(|(_, e)| e.to_handle == handle || e.from_handle == handle) + .map(|(k, _)| k.clone()) .collect(); - ids.iter().filter_map(|id| g.remove(id)).collect() + keys.iter().filter_map(|k| g.remove(k)).collect() }; for e in entries { if e.to_handle == handle { @@ -410,12 +495,12 @@ impl Router { ) -> Vec<(Instance, String)> { let overdue: Vec = { let mut g = self.inflight.lock(); - let ids: Vec = g - .values() - .filter(|e| e.deadline <= now) - .map(|e| e.delegation_id.clone()) + let keys: Vec = g + .iter() + .filter(|(_, e)| e.deadline <= now) + .map(|(k, _)| k.clone()) .collect(); - ids.iter().filter_map(|id| g.remove(id)).collect() + keys.iter().filter_map(|k| g.remove(k)).collect() }; let mut frames = Vec::new(); for e in overdue { @@ -451,11 +536,13 @@ impl Router { frames } - /// Chain of an in-flight delegation (for tests/inspection). - pub fn chain_of(&self, delegation_id: &str) -> Option> { + /// Chain of an in-flight delegation (for tests/inspection). Delegation + /// ids are namespace-scoped (review round-3 F3), so the namespace is part + /// of the lookup. + pub fn chain_of(&self, namespace: &str, delegation_id: &str) -> Option> { self.inflight .lock() - .get(delegation_id) + .get(&DelegationKey::new(namespace, delegation_id)) .map(|e| e.chain.clone()) } @@ -729,7 +816,7 @@ type = "worker" } #[test] - fn result_from_wrong_handle_dropped_and_restored() { + fn result_from_wrong_handle_dropped_and_entry_untouched() { let w = world(); assert!(matches!( do_delegate(&w, delegate_params("d-1", "worker-1", 60)), @@ -961,7 +1048,7 @@ allow_worker_initiation = true DelegateOutcome::Accepted(_) )); assert_eq!( - w.router.chain_of("d-root").unwrap(), + w.router.chain_of("prod", "d-root").unwrap(), vec!["prod/koudu".to_string()] ); @@ -1003,7 +1090,7 @@ allow_worker_initiation = true DelegateOutcome::Accepted(_) )); assert_eq!( - w.router.chain_of("d-child").unwrap(), + w.router.chain_of("prod", "d-child").unwrap(), vec!["prod/koudu".to_string(), "prod/worker-1".to_string()] ); @@ -1027,4 +1114,405 @@ allow_worker_initiation = true _ => panic!("expected cycle rejection"), } } + + fn result_of(id: &str, body: &str) -> DelegateResultParams { + DelegateResultParams { + delegation_id: id.into(), + status: DelegationStatus::Completed, + result: Some(body.into()), + error: None, + } + } + + #[test] + fn wrong_handle_result_never_hides_the_genuine_one() { + // Review round-3 F2: ownership is validated under the same lock + // acquisition that removes the entry. The old remove → validate → + // reinsert sequence made the entry briefly invisible, so a genuine + // result arriving in that window was dropped as "unknown id". + for spoof_first in [true, false] { + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + + if spoof_first { + // h_primary is registered but is NOT the serving instance. + assert!(w + .router + .complete( + &w.registry, + w.h_primary, + result_of("d-1", "spoofed"), + 1024, + 2 + ) + .is_none()); + assert_eq!( + w.router.inflight_count(), + 1, + "a non-owner frame must not remove the entry" + ); + } + + let (init, frame) = w + .router + .complete( + &w.registry, + w.h_worker, + result_of("d-1", "genuine"), + 1024, + 3, + ) + .expect("genuine result must be delivered, never dropped"); + assert_eq!(init.handle, w.h_primary); + assert!(frame.contains("genuine")); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + + if !spoof_first { + // A late non-owner frame after completion is a plain no-op. + assert!(w + .router + .complete( + &w.registry, + w.h_primary, + result_of("d-1", "spoofed"), + 1024, + 4 + ) + .is_none()); + assert_eq!(w.router.inflight_count(), 0); + } + } + } + + #[test] + fn genuine_result_survives_concurrent_non_owner_frames() { + // Review round-3 F2, the racing case the sequential test above cannot + // observe: with remove → validate → reinsert, a genuine result that + // lands inside the window sees an empty table and is dropped, and the + // delegation then stalls to its deadline. Under a single lock + // acquisition the outcome is order-independent by construction. + for _ in 0..200 { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let gate = std::sync::Barrier::new(2); + let (spoofed, genuine) = std::thread::scope(|s| { + let spoof = s.spawn(|| { + gate.wait(); + // Registered, but not the serving instance. + w.router + .complete( + &w.registry, + w.h_primary, + result_of("d-1", "spoofed"), + 1024, + 2, + ) + .is_some() + }); + gate.wait(); + let genuine = w + .router + .complete( + &w.registry, + w.h_worker, + result_of("d-1", "genuine"), + 1024, + 3, + ) + .is_some(); + (spoof.join().unwrap(), genuine) + }); + assert!(!spoofed, "a non-owner must never complete a delegation"); + assert!(genuine, "the genuine result must never be dropped"); + assert_eq!(w.router.inflight_count(), 0); + } + } + + #[test] + fn genuine_cancel_survives_concurrent_non_owner_frames() { + // Review round-3 F2 (cancel side), racing case: with the old + // remove → validate → reinsert pattern, a genuine initiator cancel + // landing inside a non-owner cancel's window would see an empty table + // and be refused, leaving the delegation to stall to its deadline. + // Under a single lock acquisition the outcome is order-independent. + for _ in 0..200 { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let params = CancelParams { + delegation_id: "d-1".into(), + reason: "race".into(), + }; + let gate = std::sync::Barrier::new(2); + let (spoofed, genuine) = std::thread::scope(|s| { + let spoof = s.spawn(|| { + gate.wait(); + // Registered, but not the initiator. + w.router.cancel(&w.registry, w.h_worker, ¶ms, 1).is_ok() + }); + gate.wait(); + let genuine = w + .router + .cancel(&w.registry, w.h_primary, ¶ms, 2) + .is_ok(); + (spoof.join().unwrap(), genuine) + }); + assert!(!spoofed, "a non-initiator must never cancel a delegation"); + assert!(genuine, "the genuine cancel must never be refused"); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "capacity must be released exactly once" + ); + } + } + + #[test] + fn refused_cancel_leaves_the_delegation_cancellable() { + // Review round-3 F2 (cancel side): a wrong-handle cancel must not + // remove-and-reinsert the entry, and must not disturb accounting. + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let params = CancelParams { + delegation_id: "d-1".into(), + reason: "not mine".into(), + }; + assert!(w + .router + .cancel(&w.registry, w.h_worker, ¶ms, 1) + .is_err()); + assert_eq!(w.router.inflight_count(), 1); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "a refused cancel must not release capacity" + ); + // The genuine initiator can still cancel. + let fwd = w + .router + .cancel(&w.registry, w.h_primary, ¶ms, 2) + .unwrap() + .unwrap(); + assert_eq!(fwd.0.handle, w.h_worker); + assert_eq!(w.router.inflight_count(), 0); + } + + #[test] + fn cancel_refusals_are_byte_identical() { + // Review round-3 F3: `cp/cancel` must not be an existence oracle — + // an unknown id and another instance's live id return the same error + // object, byte for byte. + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let unknown = CancelParams { + delegation_id: "d-does-not-exist".into(), + reason: "probe".into(), + }; + let foreign = CancelParams { + delegation_id: "d-1".into(), + reason: "probe".into(), + }; + // Both probes come from the worker: it initiated neither. + let e_unknown = w + .router + .cancel(&w.registry, w.h_worker, &unknown, 1) + .unwrap_err(); + let e_foreign = w + .router + .cancel(&w.registry, w.h_worker, &foreign, 2) + .unwrap_err(); + assert_eq!( + serde_json::to_string(&e_unknown).unwrap(), + serde_json::to_string(&e_foreign).unwrap(), + "unknown and foreign delegation ids must be indistinguishable" + ); + assert_eq!(e_unknown.code, codes::POLICY_DENIED); + assert_eq!(w.router.inflight_count(), 1); + } + + #[test] + fn same_delegation_id_in_two_namespaces_is_independent() { + // Review round-3 F3: the in-flight table is keyed by + // (namespace, delegation_id). A client-supplied id in one namespace + // must neither collide with nor be observable from another. + let registry = Registry::new(); + let router = Router::new(); + let cfg = cfg(); + let (p_prod, _prod_init_rx) = instance("prod", "koudu", AgentType::Primary, 4); + let (w_prod, mut prod_rx) = instance("prod", "worker-1", AgentType::Worker, 2); + let (p_dev, _dev_init_rx) = instance("dev", "koudu", AgentType::Primary, 4); + let (w_dev, mut dev_rx) = instance("dev", "worker-1", AgentType::Worker, 2); + let hp_prod = registry.register(p_prod); + let hw_prod = registry.register(w_prod); + let hp_dev = registry.register(p_dev); + let hw_dev = registry.register(w_dev); + + for (ns, hp) in [("prod", hp_prod), ("dev", hp_dev)] { + match router.delegate( + &cfg, + ®istry, + ns, + "koudu", + &AgentType::Primary, + hp, + delegate_params("d-1", "worker-1", 60), + 1, + ) { + DelegateOutcome::Accepted(ack) => { + assert_eq!(ack.assigned_to, format!("{ns}/worker-1")) + } + DelegateOutcome::Rejected(e) => { + panic!("{ns} rejected ({}): {}", e.code, e.message) + } + } + } + assert_eq!( + router.inflight_count(), + 2, + "one `d-1` per namespace, both in flight" + ); + prod_rx.try_recv().unwrap(); + dev_rx.try_recv().unwrap(); + + // A dev instance cannot cancel prod's `d-1` — and cannot learn that + // it exists: same error as for an id that exists nowhere. + let probe = CancelParams { + delegation_id: "d-1".into(), + reason: "probe".into(), + }; + let nowhere = CancelParams { + delegation_id: "d-nowhere".into(), + reason: "probe".into(), + }; + let e_cross = router + .cancel(®istry, hw_dev, &probe, 10) + .unwrap_err() + .message; + let e_nowhere = router + .cancel(®istry, hw_dev, &nowhere, 11) + .unwrap_err() + .message; + assert_eq!(e_cross, e_nowhere); + assert_eq!(router.inflight_count(), 2); + + // Results route to the initiator of the SAME namespace only. + let (init, frame) = router + .complete(®istry, hw_dev, result_of("d-1", "dev-done"), 1024, 12) + .unwrap(); + assert_eq!(init.handle, hp_dev); + assert!(frame.contains("dev-done")); + assert!( + router.chain_of("prod", "d-1").is_some(), + "prod's delegation must be untouched" + ); + + let (init, frame) = router + .complete(®istry, hw_prod, result_of("d-1", "prod-done"), 1024, 13) + .unwrap(); + assert_eq!(init.handle, hp_prod); + assert!(frame.contains("prod-done")); + assert_eq!(router.inflight_count(), 0); + } + + #[test] + fn parent_lookup_is_namespace_scoped() { + // Review round-3 F3: parent-chain resolution must not reach into + // another namespace's in-flight table. + let cfg: CpConfig = toml::from_str( + r#" +[namespaces.prod] +max_depth = 5 +allow_worker_initiation = true + +[namespaces.dev] +max_depth = 5 +allow_worker_initiation = true +"#, + ) + .unwrap(); + let registry = Registry::new(); + let router = Router::new(); + let (p_prod, _rx1) = instance("prod", "koudu", AgentType::Primary, 4); + let (w_prod, mut rx2) = instance("prod", "worker-1", AgentType::Worker, 2); + let (t_prod, _rx3) = instance("prod", "worker-2", AgentType::Worker, 2); + let (w_dev, _rx4) = instance("dev", "worker-1", AgentType::Worker, 2); + let (t_dev, _rx5) = instance("dev", "worker-2", AgentType::Worker, 2); + let hp_prod = registry.register(p_prod); + let hw_prod = registry.register(w_prod); + registry.register(t_prod); + let hw_dev = registry.register(w_dev); + registry.register(t_dev); + + assert!(matches!( + router.delegate( + &cfg, + ®istry, + "prod", + "koudu", + &AgentType::Primary, + hp_prod, + delegate_params("d-root", "worker-1", 120), + 1, + ), + DelegateOutcome::Accepted(_) + )); + rx2.try_recv().unwrap(); + + // dev/worker-1 claims prod's `d-root` as its parent: invisible. + let mut child = delegate_params("d-child", "worker-2", 60); + child.parent_delegation_id = Some("d-root".into()); + match router.delegate( + &cfg, + ®istry, + "dev", + "worker-1", + &AgentType::Worker, + hw_dev, + child, + 2, + ) { + DelegateOutcome::Rejected(e) => { + assert_eq!(e.code, codes::INVALID_PARAMS); + assert!(e.message.contains("not in flight for this instance")); + } + _ => panic!("cross-namespace parent must be rejected"), + } + // ...and the legitimate in-namespace child still works. + let mut ok_child = delegate_params("d-child", "worker-2", 60); + ok_child.parent_delegation_id = Some("d-root".into()); + assert!(matches!( + router.delegate( + &cfg, + ®istry, + "prod", + "worker-1", + &AgentType::Worker, + hw_prod, + ok_child, + 3, + ), + DelegateOutcome::Accepted(_) + )); + assert_eq!( + router.chain_of("prod", "d-child").unwrap(), + vec!["prod/koudu".to_string(), "prod/worker-1".to_string()] + ); + } } diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index 9504b5c42..ecf33b8b6 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -7,10 +7,16 @@ //! Resource bounds (review F5): the WS transport enforces //! `max_frame_bytes` before parsing; each connection's outbound queue is //! bounded — a peer that cannot drain it is treated as disconnected. +//! +//! Admission bounds (review round-3 F4): authentication alone is not a +//! bound. Every connection holds a per-identity slot from the upgrade until +//! it ends (`ConnPermit`, released on every exit path), and must complete +//! `cp/register` within `register_timeout_secs` or be closed. +use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use axum::extract::ws::{Message, WebSocket}; use axum::extract::{State, WebSocketUpgrade}; @@ -19,6 +25,7 @@ use axum::response::IntoResponse; use axum::routing::get; use axum::Router as AxumRouter; use futures_util::{SinkExt, StreamExt}; +use parking_lot::Mutex; use tokio::sync::mpsc; use tracing::{info, warn}; @@ -28,7 +35,7 @@ use crate::proto::{ JsonRpcErrorResponse, JsonRpcMessage, JsonRpcResponse, RegisterAck, RegisterParams, PROTOCOL_VERSION, }; -use crate::registry::{Instance, Registry, OUTBOUND_QUEUE}; +use crate::registry::{shutdown_signal, Instance, Registry, OUTBOUND_QUEUE}; use crate::router::{DelegateOutcome, Router}; pub struct AppState { @@ -36,6 +43,10 @@ pub struct AppState { pub registry: Registry, pub router: Router, rpc_id: AtomicU64, + /// Live connections per identity (`namespace/name`), counted from the + /// upgrade so pre-registration sockets are bounded too (review round-3 + /// F4). + conns: Mutex>, } impl AppState { @@ -45,12 +56,55 @@ impl AppState { registry: Registry::new(), router: Router::new(), rpc_id: AtomicU64::new(1), + conns: Mutex::new(BTreeMap::new()), } } pub fn next_rpc_id(&self) -> u64 { self.rpc_id.fetch_add(1, Ordering::Relaxed) } + + /// Take a connection slot for `identity`, or `None` when the identity is + /// already at `max_connections_per_identity`. The returned guard releases + /// the slot on drop — including on every early return and on an upgrade + /// that never completes (review round-3 F4). + pub fn try_acquire_conn(self: &Arc, identity: &AgentIdentity) -> Option { + let key = format!("{}/{}", identity.namespace, identity.name); + let mut g = self.conns.lock(); + let n = g.entry(key.clone()).or_insert(0); + if *n >= self.cfg.max_connections_per_identity { + return None; + } + *n += 1; + Some(ConnPermit { + state: Arc::clone(self), + key, + }) + } + + /// Live connection count for an identity (`namespace/name`). + pub fn conn_count(&self, logical_id: &str) -> u32 { + self.conns.lock().get(logical_id).copied().unwrap_or(0) + } +} + +/// RAII connection slot. Dropping it frees the identity's quota; it is never +/// released explicitly, so no early return can leak it (review round-3 F4). +pub struct ConnPermit { + state: Arc, + key: String, +} + +impl Drop for ConnPermit { + fn drop(&mut self) { + let mut g = self.state.conns.lock(); + if let Some(n) = g.get_mut(&self.key) { + *n = n.saturating_sub(1); + if *n == 0 { + g.remove(&self.key); + } + } + } } pub fn app(state: Arc) -> AxumRouter { @@ -80,24 +134,65 @@ async fn ws_handler( return StatusCode::UNAUTHORIZED.into_response(); } }; + // Per-identity connection quota, taken before the upgrade so an + // over-quota peer is refused at the HTTP layer (review round-3 F4). + let permit = match state.try_acquire_conn(&identity) { + Some(p) => p, + None => { + warn!( + agent = %format!("{}/{}", identity.namespace, identity.name), + max = state.cfg.max_connections_per_identity, + "WS rejected: identity is at its connection quota" + ); + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + }; let max_frame = state.cfg.max_frame_bytes; ws.max_message_size(max_frame) .max_frame_size(max_frame) - .on_upgrade(move |socket| handle_connection(state, socket, identity)) + .on_upgrade(move |socket| handle_connection(state, socket, identity, permit)) } -async fn handle_connection(state: Arc, socket: WebSocket, identity: AgentIdentity) { +async fn handle_connection( + state: Arc, + socket: WebSocket, + identity: AgentIdentity, + // Held for the connection's whole lifetime; dropped here on every exit + // path, including the early returns below (review round-3 F4). + _permit: ConnPermit, +) { let (mut sink, mut stream) = socket.split(); - // --- Registration: mandatory first frame --- - let register = loop { - match stream.next().await { - Some(Ok(Message::Text(text))) => break text, - Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue, - _ => { - warn!(agent = %identity.name, "connection closed before registration"); - return; + // --- Registration: mandatory first frame, within a deadline --- + // An authenticated peer must not be able to park idle sockets: pings keep + // the transport alive but do not extend this deadline (review round-3 F4). + let register = match tokio::time::timeout( + Duration::from_secs(state.cfg.register_timeout_secs), + async { + loop { + match stream.next().await { + Some(Ok(Message::Text(text))) => return Some(text), + Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue, + _ => return None, + } } + }, + ) + .await + { + Ok(Some(text)) => text, + Ok(None) => { + warn!(agent = %identity.name, "connection closed before registration"); + return; + } + Err(_) => { + warn!( + agent = %format!("{}/{}", identity.namespace, identity.name), + timeout_secs = state.cfg.register_timeout_secs, + "no cp/register within the registration deadline — closing" + ); + let _ = sink.send(Message::Close(None)).await; + return; } }; let (reg, reg_rpc_id) = match parse_register(®ister, &identity) { @@ -117,25 +212,36 @@ async fn handle_connection(state: Arc, socket: WebSocket, identity: Ag // cannot drain OUTBOUND_QUEUE frames is disconnected, not buffered. let (tx, mut rx) = mpsc::channel::(OUTBOUND_QUEUE); + // Shutdown signal so the CP can close this socket when it drops the + // registration on its own initiative (lease expiry — review round-3 F1). + // Subscribed BEFORE registering so no signal can be missed, and kept + // alive here for the whole connection: closing is driven by an explicit + // signal, never by the registry happening to drop its side. + let shutdown = shutdown_signal(); + let mut shutdown_rx = shutdown.subscribe(); + let effective_max = match identity.max_delegated_sessions_cap { Some(cap) => reg.max_delegated_sessions.min(cap), None => reg.max_delegated_sessions, }; // The registry assigns the CP-generated handle (review F1): ownership // and teardown never key on the client-supplied instance_id. - let handle = state.registry.register(Instance { - handle: 0, - namespace: identity.namespace.clone(), - name: identity.name.clone(), - agent_type: identity.agent_type.clone(), - instance_id: reg.instance_id.clone(), - labels: reg.labels.clone(), - max_delegated_sessions: effective_max, - active_sessions: 0, - registered_at: Instant::now(), - last_heartbeat: Instant::now(), - tx: tx.clone(), - }); + let handle = state.registry.register_conn( + Instance { + handle: 0, + namespace: identity.namespace.clone(), + name: identity.name.clone(), + agent_type: identity.agent_type.clone(), + instance_id: reg.instance_id.clone(), + labels: reg.labels.clone(), + max_delegated_sessions: effective_max, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx: tx.clone(), + }, + Arc::clone(&shutdown), + ); info!( agent = %format!("{}/{}", identity.namespace, identity.name), instance = %reg.instance_id, @@ -167,9 +273,20 @@ async fn handle_connection(state: Arc, socket: WebSocket, identity: Ag return; } - // --- Main loop: interleave inbound frames and outbound channel --- + // --- Main loop: interleave inbound frames, outbound channel, shutdown --- + let mut cp_closed = false; loop { tokio::select! { + // The CP dropped this registration (lease expiry): the socket + // must go too (review round-3 F1). Keeping it open would leave a + // connection whose every frame hits an absent registry entry and + // which can never re-register, since registration is + // first-frame-only. Closing lets the client reconnect, + // re-authenticate, and register again. + _ = shutdown_rx.changed() => { + cp_closed = true; + break; + } outbound = rx.recv() => { match outbound { Some(text) => { @@ -205,6 +322,15 @@ async fn handle_connection(state: Arc, socket: WebSocket, identity: Ag } } + if cp_closed { + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + handle, + "closing connection at the CP's request (registration dropped)" + ); + let _ = sink.send(Message::Close(None)).await; + } + teardown(&state, handle, &identity); } @@ -429,26 +555,41 @@ fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option, lease: Duration) { + for handle in state.registry.expired(lease) { + warn!( + handle, + "lease expired — deregistering and closing connection" + ); + // Signal first: `deregister` drops the registry's side of the signal. + state.registry.signal_shutdown(handle); + state.registry.deregister(handle); + let mut next = || state.next_rpc_id(); + for (inst, frame) in state + .router + .fail_instance(&state.registry, handle, &mut next) + { + let _ = inst.tx.try_send(frame); + } + } +} + /// Background sweeps: lease expiry and delegation deadlines. pub async fn run_sweeper(state: Arc) { - let mut tick = tokio::time::interval(std::time::Duration::from_secs(1)); + let mut tick = tokio::time::interval(Duration::from_secs(1)); tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { tick.tick().await; - // Lease expiry → deregister + fail in-flight. - let lease = std::time::Duration::from_secs(state.cfg.lease_expiry_secs); - for handle in state.registry.expired(lease) { - warn!(handle, "lease expired — deregistering"); - state.registry.deregister(handle); - let mut next = || state.next_rpc_id(); - for (inst, frame) in state - .router - .fail_instance(&state.registry, handle, &mut next) - { - let _ = inst.tx.try_send(frame); - } - } + sweep_leases(&state, Duration::from_secs(state.cfg.lease_expiry_secs)); // Deadline sweep. let mut next = || state.next_rpc_id(); @@ -594,4 +735,94 @@ mod tests { let (_, err) = parse_register(&no_id, &identity()).unwrap_err(); assert_eq!(err.code, codes::INVALID_REQUEST); } + + fn state_with(cfg_toml: &str) -> Arc { + let cfg: CpConfig = toml::from_str(cfg_toml).unwrap(); + cfg.validate().unwrap(); + Arc::new(AppState::new(cfg)) + } + + #[test] + fn conn_quota_bounds_and_recycles_slots() { + // Review round-3 F4(b): the quota is a hard bound and the guard + // releases the slot on drop, so no exit path can leak it. + let state = state_with("max_connections_per_identity = 2"); + let id = identity(); + let p1 = state.try_acquire_conn(&id).expect("slot 1"); + let p2 = state.try_acquire_conn(&id).expect("slot 2"); + assert_eq!(state.conn_count("prod/koudu"), 2); + assert!( + state.try_acquire_conn(&id).is_none(), + "third concurrent connection must be refused" + ); + + drop(p1); + assert_eq!(state.conn_count("prod/koudu"), 1); + let p3 = state + .try_acquire_conn(&id) + .expect("released slot is reusable"); + drop(p2); + drop(p3); + assert_eq!(state.conn_count("prod/koudu"), 0); + assert!(state.try_acquire_conn(&id).is_some()); + } + + #[test] + fn conn_quota_is_per_identity() { + let state = state_with("max_connections_per_identity = 1"); + let a = identity(); + let mut b = identity(); + b.key = "k2".into(); + b.name = "worker-1".into(); + let _pa = state.try_acquire_conn(&a).expect("koudu slot"); + let _pb = state + .try_acquire_conn(&b) + .expect("worker-1 has its own quota"); + assert!( + state.try_acquire_conn(&a).is_none(), + "quota is per identity, not global" + ); + assert_eq!(state.conn_count("prod/koudu"), 1); + assert_eq!(state.conn_count("prod/worker-1"), 1); + } + + #[tokio::test] + async fn sweep_leases_signals_the_connection_before_dropping_it() { + // Review round-3 F1 at the sweeper level: the shutdown signal is + // delivered, not just the registry entry removed. (The end-to-end + // proof over a real socket lives in tests/ws_lifecycle.rs.) + let state = state_with("heartbeat_interval_secs = 1\nlease_expiry_secs = 2"); + let signal = crate::registry::shutdown_signal(); + let mut observer = signal.subscribe(); + let (tx, _rx) = mpsc::channel::(OUTBOUND_QUEUE); + let handle = state.registry.register_conn( + Instance { + handle: 0, + namespace: "prod".into(), + name: "koudu".into(), + agent_type: AgentType::Primary, + instance_id: "i-1".into(), + labels: Default::default(), + max_delegated_sessions: 1, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + }, + Arc::clone(&signal), + ); + + // A live lease is left alone. + sweep_leases(&state, Duration::from_secs(60)); + assert!(state.registry.get(handle).is_some()); + assert!(!*observer.borrow()); + + sweep_leases(&state, Duration::ZERO); + assert!(state.registry.get(handle).is_none()); + observer.changed().await.unwrap(); + assert!( + *observer.borrow(), + "the owning connection must be told to close" + ); + } } diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs new file mode 100644 index 000000000..0447bf1df --- /dev/null +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -0,0 +1,254 @@ +//! End-to-end WebSocket lifecycle tests: connection termination on lease +//! expiry (review round-3 F1) and pre-registration admission bounds +//! (review round-3 F4). +//! +//! These drive a real CP over a loopback socket with a real WS client, which +//! is the only way to prove that the connection *task* reacts — the earlier +//! bug was invisible to unit tests of the registry/router alone. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use futures_util::{SinkExt, StreamExt}; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::StatusCode; +use tokio_tungstenite::tungstenite::{Error as WsError, Message}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use openab_cp::config::CpConfig; +use openab_cp::server::{app, sweep_leases, AppState}; + +const KEY: &str = "k-primary"; + +type Ws = WebSocketStream>; + +fn cfg(extra: &str) -> CpConfig { + let raw = format!( + r#" +{extra} + +[[agents]] +key = "{KEY}" +namespace = "prod" +name = "koudu" +type = "primary" +"# + ); + let cfg: CpConfig = toml::from_str(&raw).expect("test config parses"); + cfg.validate().expect("test config validates"); + cfg +} + +/// Start a CP on an ephemeral loopback port; returns its state and WS URL. +async fn spawn_cp(cfg: CpConfig) -> (Arc, String) { + let state = Arc::new(AppState::new(cfg)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let router = app(state.clone()); + tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + (state, format!("ws://{addr}/cp")) +} + +async fn connect(url: &str) -> Result { + let mut req = url.into_client_request().unwrap(); + req.headers_mut().insert( + "authorization", + format!("Bearer {KEY}").parse().expect("header value"), + ); + tokio_tungstenite::connect_async(req) + .await + .map(|(ws, _)| ws) +} + +/// Connect, retrying while the identity's quota slot is still being released +/// by the server task. +async fn connect_retry(url: &str) -> Ws { + for _ in 0..100 { + if let Ok(ws) = connect(url).await { + return ws; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("connection was never accepted"); +} + +fn register_frame(instance_id: &str) -> String { + serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": instance_id + } + }) + .to_string() +} + +/// Send `cp/register` and return the parsed reply. +async fn register(ws: &mut Ws, instance_id: &str) -> serde_json::Value { + ws.send(Message::Text(register_frame(instance_id).into())) + .await + .unwrap(); + let msg = tokio::time::timeout(Duration::from_secs(5), ws.next()) + .await + .expect("register must be answered") + .expect("stream open") + .expect("no ws error"); + serde_json::from_str(msg.to_text().unwrap()).unwrap() +} + +/// Wait until the peer closes the socket (Close frame, error, or EOF). +async fn wait_closed(ws: &mut Ws, within: Duration) -> bool { + let deadline = Instant::now() + within; + while Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(200), ws.next()).await { + Ok(None) | Ok(Some(Err(_))) => return true, + Ok(Some(Ok(Message::Close(_)))) => return true, + Ok(Some(Ok(_))) => continue, + Err(_) => continue, // read timeout: keep waiting + } + } + false +} + +#[tokio::test] +async fn lease_expiry_closes_the_connection_and_permits_reregistration() { + // Review round-3 F1: deregistering on lease expiry without terminating + // the connection left a live socket bound to a registration that no + // longer existed — heartbeats got no reply and re-registration was + // impossible (registration is first-frame-only). The CP must close it. + let (state, url) = spawn_cp(cfg("max_connections_per_identity = 1")).await; + + let mut ws = connect(&url).await.expect("first connection accepted"); + let ack = register(&mut ws, "i-1").await; + assert_eq!(ack["result"]["protocol_version"], 1, "registered"); + assert_eq!(state.registry.list("prod").len(), 1); + + // Zero lease: every registration is overdue on this pass. + sweep_leases(&state, Duration::ZERO); + assert!( + state.registry.list("prod").is_empty(), + "lease expiry deregisters" + ); + + assert!( + wait_closed(&mut ws, Duration::from_secs(5)).await, + "the connection task must observe the shutdown signal and close" + ); + drop(ws); + + // A reconnecting client re-authenticates and registers again. This also + // proves the connection slot was released (quota is 1 here). + let mut ws2 = connect_retry(&url).await; + let ack2 = register(&mut ws2, "i-2").await; + assert_eq!(ack2["result"]["protocol_version"], 1); + let live = state.registry.list("prod"); + assert_eq!(live.len(), 1); + assert_eq!(live[0].instance_id, "i-2"); +} + +#[tokio::test] +async fn ping_only_pre_registration_socket_is_closed_at_the_deadline() { + // Review round-3 F4(a): pings keep the transport alive but must not + // extend the registration deadline. + let (state, url) = spawn_cp(cfg("register_timeout_secs = 1")).await; + let mut ws = connect(&url).await.expect("connection accepted"); + + let started = Instant::now(); + let mut closed = false; + while started.elapsed() < Duration::from_secs(10) { + let _ = ws.send(Message::Ping(vec![7].into())).await; + match tokio::time::timeout(Duration::from_millis(250), ws.next()).await { + Ok(None) | Ok(Some(Err(_))) => { + closed = true; + break; + } + Ok(Some(Ok(Message::Close(_)))) => { + closed = true; + break; + } + _ => continue, + } + } + assert!( + closed, + "an authenticated socket that never registers must be closed" + ); + assert!(state.registry.list("prod").is_empty()); +} + +#[tokio::test] +async fn registration_after_the_deadline_is_not_accepted() { + // Review round-3 F4(a): the deadline is enforced, not merely advisory. + let (state, url) = spawn_cp(cfg("register_timeout_secs = 1")).await; + let mut ws = connect(&url).await.expect("connection accepted"); + tokio::time::sleep(Duration::from_millis(1_600)).await; + + let _ = ws + .send(Message::Text(register_frame("i-late").into())) + .await; + let acked = match tokio::time::timeout(Duration::from_secs(2), ws.next()).await { + Ok(Some(Ok(Message::Text(t)))) => t.contains("effective_max_delegated_sessions"), + _ => false, + }; + assert!(!acked, "a late cp/register must not be acked"); + assert!(state.registry.list("prod").is_empty()); +} + +#[tokio::test] +async fn connection_quota_rejects_over_limit_and_recycles_on_disconnect() { + // Review round-3 F4(b): the quota bounds concurrent sockets per identity + // and is released on every exit path (RAII), so connect → disconnect → + // connect always succeeds. + let (state, url) = spawn_cp(cfg( + "max_connections_per_identity = 1\nregister_timeout_secs = 30", + )) + .await; + + let mut ws = connect(&url).await.expect("first connection accepted"); + register(&mut ws, "i-1").await; + assert_eq!(state.conn_count("prod/koudu"), 1); + + match connect(&url).await { + Err(WsError::Http(resp)) => assert_eq!( + resp.status(), + StatusCode::SERVICE_UNAVAILABLE, + "over-quota upgrade must be refused before the WS handshake" + ), + Err(e) => panic!("unexpected error: {e}"), + Ok(_) => panic!("over-quota connection must be rejected"), + } + + let _ = ws.close(None).await; + drop(ws); + + let mut ws2 = connect_retry(&url).await; + register(&mut ws2, "i-2").await; + assert_eq!( + state.conn_count("prod/koudu"), + 1, + "the released slot was reused, not leaked" + ); +} + +#[tokio::test] +async fn pre_registration_sockets_count_against_the_quota() { + // Review round-3 F4(b): the quota is taken at the upgrade, so parked + // pre-registration sockets cannot be multiplied for free. + let (_state, url) = spawn_cp(cfg( + "max_connections_per_identity = 1\nregister_timeout_secs = 30", + )) + .await; + let _parked = connect(&url).await.expect("connection accepted"); + + match connect(&url).await { + Err(WsError::Http(resp)) => assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE), + Err(e) => panic!("unexpected error: {e}"), + Ok(_) => panic!("an unregistered socket must still occupy its quota slot"), + } +} From 09968d19689bdca972e7a3e9380b83d449509875 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Thu, 13 Aug 2026 08:46:27 -0400 Subject: [PATCH 04/11] =?UTF-8?q?fix(cp):=20round-4=20review=20=E2=80=94?= =?UTF-8?q?=20close=20codes,=20swept-frame=20errors,=20ownership=20helper,?= =?UTF-8?q?=20ADR=20contract=20sync?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 InFlight.namespace documented, not removed: the field is read by the observer event layer one PR up this stack (per-namespace cp/event fan-out in fail_instance/sweep_deadlines), and it is part of the entry's identity because delegation ids are only unique within their namespace. Removing it here would only force the stack to re-add it. F2 Dropped the unused `uuid` dependency from crates/openab-cp/Cargo.toml — registration handles come from an AtomicU64, nothing generated a UUID. Cargo.lock regenerated (openab-cp's dep list loses `uuid`; the package itself stays, other crates still use it). F3 Removed `codes::DEADLINE_EXCEEDED` instead of reserving it. Verified that admission already rejects an already-expired deadline: policy::check returns PolicyDenial::DeadlinePast, which router::delegate maps to POLICY_DENIED — so the reserved-for-admission rationale does not hold. A deadline that elapses in flight is not an error response either; the sweeper synthesizes a cp/delegate_result with status `timeout`. -32006 is left unassigned with a comment so the published numbering stays stable. F4 Extracted the duplicated ownership-check-and-remove sequence from complete() and cancel() into one private helper, Router::claim(), returning Claim { Owned(InFlight), WrongOwner{..}, NotFound{..}, Unregistered }. The single-lock-acquisition property is preserved inside the helper: the owner comparison and the removal happen under one `inflight` lock, so a non-owner frame never makes the entry momentarily invisible. Callers keep their distinct logging and return shapes (None vs. byte-identical POLICY_DENIED). All pre-existing tests pass with zero non-comment changes inside the tests module — that is the behavior-preservation proof. F5 Rewrote every `(review round-N FX)` comment tag across the crate into self-standing rationale, keeping the WHY prose and dropping the round index (src/**, plus tests/ws_lifecycle.rs, which a reviewer greps too). `grep -rniE 'review|round-[0-9]' crates/openab-cp/` is now empty. F6 docs/adr/agent-control-plane.md §4 amended to the enforced contract: registration deadline (pings do not extend it; 1008 close), lease expiry actively closing the connection with reconnect + re-register as the only recovery, per-identity connection quota counted from the upgrade and RAII-released, and (namespace, delegation_id) scoping including the indistinguishable cancel refusals. The "CP restart semantics" section was corrected against the code rather than to the reviewer's phrasing: a restart is all leases expiring at once, but the CP cannot synthesize any failure because the table and the sockets die with the process — initiators reconcile against the deadline they already propagated. Synthesized failures are the live-CP path only, and they are **side-specific**: Router::fail_instance is a mutually exclusive branch — a dead serving instance sends `target_disconnected` to the initiator, a dead initiator sends a best-effort `cp/cancel` downstream (and releases the worker's reserved capacity), never both for one delegation. Only sweep_deadlines emits both frames, because there both peers are still connected. Two sentences that claimed lease expiry/disconnect of either side produces both outcomes were rewritten to match. F7 CP-initiated closes now carry meaning: registration timeout and lease expiry both send WS close 1008 (policy violation) with reason "registration timeout" / "lease expired"; the over-quota upgrade's HTTP 503 gets a body naming max_connections_per_identity. ws_lifecycle assertions extended (not weakened) to check code + reason and the 503 body; verified they fail when the code/reason/body is mutated. F13 handle_frame() now answers NOT_REGISTERED when registry.get(handle) is None — frames landing between a lease sweep and the connection close were dropped silently, indistinguishable from a hung CP. New unit test: heartbeat answered while registered, then the same frame on the swept handle returns error code NOT_REGISTERED correlated with the request id; a cp/delegate on a swept handle likewise never reaches the router. F15 Added a module-level lock-hierarchy doc to router.rs: admission → inflight, never the reverse, with the registry lock called out as an independent lock never held together with `inflight`. Verification (macmini, isolated worktree matching this branch so openab-core resolves): 63 lib + 5 ws_lifecycle tests pass (baseline 62 + 5; +1 is the new F13 test), cargo clippy -p openab-cp --all-targets -- -D warnings clean, cargo fmt --check 0 diffs — all equal to or better than the b9850da baseline. --- Cargo.lock | 1 - crates/openab-cp/Cargo.toml | 1 - crates/openab-cp/src/config.rs | 21 +- crates/openab-cp/src/policy.rs | 2 +- crates/openab-cp/src/proto.rs | 11 +- crates/openab-cp/src/registry.rs | 32 +-- crates/openab-cp/src/router.rs | 319 ++++++++++++++++--------- crates/openab-cp/src/server.rs | 178 +++++++++++--- crates/openab-cp/tests/ws_lifecycle.rs | 121 +++++++--- docs/adr/agent-control-plane.md | 74 +++++- 10 files changed, 543 insertions(+), 217 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index da40e70d2..fc11ec7e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2593,7 +2593,6 @@ dependencies = [ "toml", "tracing", "tracing-subscriber", - "uuid", ] [[package]] diff --git a/crates/openab-cp/Cargo.toml b/crates/openab-cp/Cargo.toml index debff2e93..86ca5e30c 100644 --- a/crates/openab-cp/Cargo.toml +++ b/crates/openab-cp/Cargo.toml @@ -15,7 +15,6 @@ toml = "0.8" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } anyhow = "1" -uuid = { version = "1", features = ["v4"] } chrono = { version = "0.4", features = ["serde"] } parking_lot = "0.12" clap = { version = "4", features = ["derive"] } diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index 21d672b3b..75e27ffc6 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -1,6 +1,6 @@ //! CP-side configuration. //! -//! Identity binding is the security core (review F1 on the ADR): every auth +//! Identity binding is the security core: every auth //! key maps to **immutable claims** (`namespace`, `name`, `type`, optional //! caps) owned by CP config. Registration frames are verified against these //! claims — never the other way around. A compromised runtime cannot escalate @@ -48,7 +48,8 @@ pub struct CpConfig { pub max_result_bytes: usize, /// Maximum WebSocket message size accepted from a runtime, enforced by - /// the transport before any parsing/allocation (review F5). + /// the transport before any parsing/allocation, so an oversized frame is + /// never buffered in full. #[serde(default = "default_max_frame_bytes")] pub max_frame_bytes: usize, @@ -59,16 +60,17 @@ pub struct CpConfig { /// Deadline for the mandatory `cp/register` first frame, in seconds from /// the completed WebSocket upgrade. A connection that authenticates but - /// never registers is closed when this elapses (review round-3 F4): - /// otherwise an authenticated peer could park unlimited sockets in the + /// never registers is closed when this elapses: otherwise an + /// authenticated peer could park unlimited sockets in the /// pre-registration state, keeping them alive with pings forever. #[serde(default = "default_register_timeout_secs")] pub register_timeout_secs: u64, /// Maximum simultaneous connections per identity, counted from the /// upgrade (so pre-registration sockets count too) and released on every - /// exit path (review round-3 F4). Replicas of one logical agent share one - /// identity, so this is the replica ceiling as well. + /// exit path. Authentication alone is not a bound: one leaked key could + /// otherwise open unlimited sockets. Replicas of one logical agent share + /// one identity, so this is the replica ceiling as well. #[serde(default = "default_max_connections_per_identity")] pub max_connections_per_identity: u32, @@ -188,7 +190,7 @@ impl CpConfig { if self.lease_expiry_secs <= self.heartbeat_interval_secs { bail!("lease_expiry_secs must exceed heartbeat_interval_secs"); } - // Admission bounds must actually bound something (review round-3 F4): + // Admission bounds must actually bound something: // zero would mean "no registration deadline" / "no connection allowed". if self.register_timeout_secs == 0 { bail!("register_timeout_secs must be greater than 0"); @@ -197,8 +199,7 @@ impl CpConfig { bail!("max_connections_per_identity must be at least 1"); } // Bearer keys over cleartext TCP must never reach an untrusted - // network: non-loopback binds require the explicit override - // (review round-2 F4). + // network: non-loopback binds require the explicit override. if !self.allow_insecure_bind && !is_loopback(&self.listen) { bail!( "listen = \"{}\" is not loopback and the CP terminates no TLS. \ @@ -360,7 +361,7 @@ lease_expiry_secs = 30 #[test] fn admission_bounds_default_and_are_validated() { - // Review round-3 F4: absent fields keep working (serde defaults) and + // Absent fields keep working (serde defaults) and // a zero bound is rejected rather than silently disabling the guard. let cfg: CpConfig = toml::from_str(base_toml()).unwrap(); cfg.validate().unwrap(); diff --git a/crates/openab-cp/src/policy.rs b/crates/openab-cp/src/policy.rs index c8b485ca5..6792de59b 100644 --- a/crates/openab-cp/src/policy.rs +++ b/crates/openab-cp/src/policy.rs @@ -1,4 +1,4 @@ -//! CP-authoritative delegation policy (review F4 on the ADR). +//! CP-authoritative delegation policy. //! //! Every check here operates exclusively on CP-owned data: authenticated //! identity claims, the CP-constructed ancestry chain, and CP config. Facade diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs index 91c4a27f4..89dce87f9 100644 --- a/crates/openab-cp/src/proto.rs +++ b/crates/openab-cp/src/proto.rs @@ -89,7 +89,7 @@ pub struct JsonRpcMessage { } impl JsonRpcMessage { - /// Validate this frame as a JSON-RPC 2.0 **request** (review F4): the + /// Validate this frame as a JSON-RPC 2.0 **request**: the /// `jsonrpc` field must be exactly "2.0", and a request id must be /// present (all `cp/*` client→CP methods are requests, not /// notifications). Returns the request id. @@ -143,8 +143,13 @@ pub mod codes { /// Matching targets exist but all are at their advertised capacity. /// Explicit fast-fail: the CP never queues (v1 has no durable state). pub const SATURATED: i64 = -32005; - /// Delegation deadline elapsed before a result frame arrived. - pub const DEADLINE_EXCEEDED: i64 = -32006; + // -32006 is deliberately unassigned. It held a `DEADLINE_EXCEEDED` code + // that nothing could ever emit: a deadline already in the past is + // rejected at admission by the policy check (`POLICY_DENIED`, "deadline + // is in the past"), and a deadline that elapses while the delegation runs + // is not an error response at all — the sweeper synthesizes a + // `cp/delegate_result` with status `timeout`. Leaving the slot empty keeps + // the published numbering stable for anyone who saw the earlier constant. /// Serving runtime disconnected while the delegation was in flight. pub const TARGET_DISCONNECTED: i64 = -32007; /// `delegation_id` already in flight (idempotency guard). diff --git a/crates/openab-cp/src/registry.rs b/crates/openab-cp/src/registry.rs index dd7cd47cf..7a3c27a5d 100644 --- a/crates/openab-cp/src/registry.rs +++ b/crates/openab-cp/src/registry.rs @@ -18,18 +18,17 @@ use crate::proto::AgentType; /// Outbound frame sender for one WS connection (serialized JSON text). /// Bounded: a peer that cannot drain its queue is disconnected rather than -/// growing CP memory (review F5). +/// growing CP memory. pub type FrameTx = mpsc::Sender; /// Capacity of each per-connection outbound queue. pub const OUTBOUND_QUEUE: usize = 256; /// Shutdown signal for one WS connection, held by the registry so the CP can -/// terminate a connection it no longer considers registered (review round-3 -/// F1: lease expiry must close the socket — otherwise the connection task -/// lives on with a registry entry that no longer exists, silently dropping -/// every subsequent frame and unable to re-register, since registration is -/// first-frame-only). +/// terminate a connection it no longer considers registered. Lease expiry must +/// close the socket — otherwise the connection task lives on with a registry +/// entry that no longer exists, and the client cannot re-register because +/// registration is first-frame-only. /// /// `watch` (not `oneshot`) so the connection task can select on it repeatedly, /// and wrapped in `Arc` so the registry entry and the connection task share @@ -54,8 +53,8 @@ struct Entry { #[derive(Clone, Debug)] pub struct Instance { /// CP-generated registration handle — the registry key and the basis of - /// all ownership checks. Never client-supplied (review F1): a colliding - /// client `instance_id` cannot replace or tear down another identity's + /// all ownership checks. Never client-supplied: a colliding client + /// `instance_id` cannot replace or tear down another identity's /// registration. pub handle: u64, pub namespace: String, @@ -67,7 +66,7 @@ pub struct Instance { pub labels: BTreeMap, pub max_delegated_sessions: u32, /// Delegations currently routed to this instance. CP-owned and - /// authoritative — never merged from runtime reports (review F6). + /// authoritative — never merged from runtime reports. pub active_sessions: u32, pub registered_at: Instant, pub last_heartbeat: Instant, @@ -108,7 +107,7 @@ impl Registry { /// /// `shutdown` is the owning connection's termination signal: the CP /// triggers it whenever it drops the registration on its own initiative - /// (lease expiry — review round-3 F1). + /// (lease expiry). pub fn register_conn(&self, mut inst: Instance, shutdown: ShutdownTx) -> u64 { let handle = self.next_handle.fetch_add(1, Ordering::Relaxed) + 1; inst.handle = handle; @@ -143,8 +142,8 @@ impl Registry { } /// Refresh the lease. The runtime-reported session count is intentionally - /// ignored: CP-owned in-flight accounting is authoritative (review F6 — - /// merging reports could pin an instance saturated forever). + /// ignored: CP-owned in-flight accounting is authoritative, and merging + /// runtime reports could pin an instance saturated forever. pub fn heartbeat(&self, handle: u64) -> bool { let mut g = self.inner.write(); match g.get_mut(&handle) { @@ -173,7 +172,7 @@ impl Registry { /// Select a serving instance within `namespace` by exact name or labels. /// - /// Unsaturated matches only. Ordering (review F6): + /// Unsaturated matches only. Ordering: /// - exact-name selection → replicas of one logical agent: newest /// registration first (rolling-deploy rule), load as tie-breaker /// - label selection → across logical agents: least loaded first, @@ -312,7 +311,8 @@ mod tests { fn label_selection_least_loaded_first() { let r = Registry::new(); // Older but less loaded instance must win under label selection - // (inverse recency/load — review F6). + // (load first, recency only as tie-breaker — the inverse of the + // exact-name replica rule). let mut a = inst("prod", "wa", "i-a", 4); a.labels.insert("backend".into(), "kiro".into()); a.active_sessions = 0; @@ -363,7 +363,7 @@ mod tests { #[test] fn colliding_instance_id_cannot_replace_other_registration() { - // Review F1: a second connection registering the same client-supplied + // A second connection registering the same client-supplied // instance_id gets its own handle; the first registration survives // and can only be torn down via its own handle. let r = Registry::new(); @@ -403,7 +403,7 @@ mod tests { #[tokio::test] async fn signal_shutdown_reaches_the_owning_connection() { - // Review round-3 F1: the CP must be able to terminate a connection + // The CP must be able to terminate a connection // whose registration it drops on its own initiative. let r = Registry::new(); let sig = shutdown_signal(); diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index 57d77a05e..2d10d31ce 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -1,5 +1,5 @@ //! Delegation router: in-flight table, target selection, result routing, and -//! the failure semantics the ADR review required to be explicit: +//! the failure semantics the ADR requires to be explicit rather than implied: //! //! - **Deadline sweep** — an in-flight delegation whose deadline passes is //! terminated: the initiator receives a synthesized `timeout` result and @@ -9,12 +9,36 @@ //! instance fail immediately with `target_disconnected`. //! - **Initiator disconnect** — its in-flight delegations are cancelled //! downstream (best effort); nobody is left to receive the result. -//! - **CP restart** — the table is in-memory; all in-flight delegations -//! effectively end as initiator-side timeouts. Late `cp/delegate_result` -//! frames for unknown ids are acknowledged and dropped (logged), so -//! reconnecting runtimes do not error-loop. +//! - **CP restart** — the table dies with the process, and the connections +//! die with it, so the CP cannot synthesize anything: in-flight delegations +//! end as initiator-side timeouts against the already-propagated deadline. +//! Late `cp/delegate_result` frames for unknown ids are acknowledged and +//! dropped (logged), so reconnecting runtimes do not error-loop. //! - **Saturation** — routing never queues; `SATURATED` is returned //! immediately (fast-fail, no hidden buffer). +//! +//! # Lock hierarchy +//! +//! The router holds two locks and acquires them in ONE order only: +//! +//! ```text +//! admission → inflight (never the reverse) +//! ``` +//! +//! `admission` serializes the whole delegate admission sequence; `inflight` +//! guards the in-flight table itself and is taken for short, self-contained +//! critical sections. Every path that needs both — only `delegate` does — +//! takes `admission` first and then `inflight`, possibly several times. +//! No path may take `inflight` and then reach for `admission`: because +//! `delegate` holds `admission` across `inflight` acquisitions, doing so +//! would close a deadlock cycle. Paths that need only the table +//! (`complete`, `cancel`, `fail_instance`, `sweep_deadlines`, `chain_of`) +//! take `inflight` alone and never touch `admission`. +//! +//! Registry access is a third, independent lock owned by [`Registry`]. It is +//! always acquired and released *outside* an `inflight` critical section +//! (e.g. `registry.get(...)` completes before the table is locked), so it +//! does not participate in this hierarchy. use std::collections::BTreeMap; @@ -31,12 +55,22 @@ use crate::proto::{ use crate::registry::{Instance, Registry, SelectError}; /// One in-flight delegation. Ownership is tracked by CP-generated -/// registration handles, never client-supplied ids (review F1). +/// registration handles, never client-supplied ids: a colliding +/// `instance_id` on another connection can neither complete nor cancel this +/// delegation. #[derive(Clone)] pub struct InFlight { - /// Namespace the delegation lives in — part of its identity (review - /// round-3 F3): `delegation_id` is client-supplied and only unique within + /// Namespace that owns this delegation — part of its identity, not just a + /// lookup key: `delegation_id` is client-supplied and only unique within /// the namespace that produced it. + /// + /// Stored on the entry because the delegation outlives the request that + /// created it, and the paths that end it without a client request — + /// `fail_instance` and `sweep_deadlines` — have no namespace of their own + /// to work from. Its consumer is the observer event layer added by the + /// next PR in this stack (per-namespace `cp/event` fan-out reads + /// `e.namespace` in exactly those two paths), so the field is part of the + /// entry's contract rather than an unused remnant. pub namespace: String, pub delegation_id: String, /// Authenticated initiator (`namespace/name`) and its registration handle. @@ -51,7 +85,7 @@ pub struct InFlight { pub chain: Vec, } -/// In-flight table key: `(namespace, delegation_id)` (review round-3 F3). +/// In-flight table key: `(namespace, delegation_id)`. /// /// Keying on the client-supplied `delegation_id` alone made one namespace's /// ids observable from another: a colliding id was denied with @@ -77,9 +111,10 @@ pub struct Router { inflight: Mutex>, /// Serializes the delegate admission sequence (duplicate check → target /// selection → capacity reservation → in-flight insert) so concurrent - /// requests cannot double-admit one id or oversubscribe capacity - /// (review F2). Delegation rates are LLM-scale; a coarse admission lock - /// is simple and more than sufficient. + /// requests cannot double-admit one id or oversubscribe capacity: without + /// it, two racing delegates both see a free slot and both reserve it. + /// Delegation rates are LLM-scale; a coarse admission lock is simple and + /// more than sufficient. admission: Mutex<()>, } @@ -90,6 +125,43 @@ pub enum DelegateOutcome { Rejected(ErrorObject), } +/// Which side of a delegation a caller must be to act on it. +#[derive(Clone, Copy)] +enum Owner { + /// The instance the delegation was routed to — the only one that may + /// complete it. + Server, + /// The instance that initiated the delegation — the only one that may + /// cancel it. + Initiator, +} + +/// Result of looking up an in-flight delegation on a caller's behalf and +/// asserting the caller owns it. +/// +/// The whole check happens under ONE acquisition of the in-flight lock, which +/// is the property that matters: an earlier version removed the entry, +/// validated ownership, then reinserted it on refusal, and a genuine frame +/// landing in that window saw an empty table and was dropped as "unknown id", +/// leaving the delegation to stall until its deadline. Here the entry is +/// either removed because the caller owns it, or never touched at all. +enum Claim { + /// The caller owns it; the entry has already been removed from the table. + Owned(InFlight), + /// The entry exists but belongs to another instance. Left in place. + WrongOwner { + namespace: String, + /// Handle of the instance that does own it (CP-side logs only — it is + /// never disclosed to the caller). + owner_handle: u64, + }, + /// No entry for `(namespace, delegation_id)`. + NotFound { namespace: String }, + /// The calling connection has no registration: it was swept (lease + /// expiry) or never registered, so it has no namespace to look in. + Unregistered, +} + impl Router { pub fn new() -> Self { Self { @@ -113,14 +185,15 @@ impl Router { ) -> DelegateOutcome { let now = Utc::now(); - // Admission is one atomic sequence (review F2): duplicate check, - // parent lookup, target selection, capacity reservation, and - // in-flight insertion all happen under this guard. + // Admission is one atomic sequence: duplicate check, parent lookup, + // target selection, capacity reservation, and in-flight insertion all + // happen under this guard, so two racing delegates can neither + // double-admit an id nor both claim the last free slot. let _admission = self.admission.lock(); - // Delegation identity is namespace-scoped (review round-3 F3): the - // same id in another namespace is a different delegation, so it - // neither collides here nor leaks its existence. + // Delegation identity is namespace-scoped: the same id in another + // namespace is a different delegation, so it neither collides here + // nor leaks its existence. let key = DelegationKey::new(from_namespace, ¶ms.delegation_id); if self.inflight.lock().contains_key(&key) { return DelegateOutcome::Rejected(ErrorObject::new( @@ -141,9 +214,9 @@ impl Router { // Parent linkage: chain and deadline derive from the CP's own table, // never from the client. The caller must BE the instance serving the // parent delegation — otherwise any runtime knowing a live id could - // borrow its trusted chain and deadline budget (review F3). The - // lookup is namespace-scoped (review round-3 F3). Unknown and - // unauthorized parent ids return the same error (no enumeration). + // borrow its trusted chain and deadline budget. The lookup is + // namespace-scoped. Unknown and unauthorized parent ids return the + // same error (no enumeration). let (parent_chain, parent_deadline) = match ¶ms.parent_delegation_id { Some(pid) => { let parent_key = DelegationKey::new(from_namespace, pid); @@ -216,8 +289,8 @@ impl Router { let text = serde_json::to_string(&frame).expect("serializable"); // Reserve capacity and record the in-flight entry BEFORE sending, so - // an immediately-arriving result finds it (review F2). Roll both - // back if the send fails. + // an immediately-arriving result finds it. Roll both back if the send + // fails. registry.adjust_sessions(target.handle, 1); let entry = InFlight { namespace: from_namespace.to_string(), @@ -256,15 +329,47 @@ impl Router { }) } + /// Look up `delegation_id` in the caller's namespace, assert the caller is + /// the delegation's `owner` side, and remove the entry if so — all under + /// one acquisition of the in-flight lock (see [`Claim`]). + /// + /// The namespace is taken from the caller's authenticated registration, + /// never from the frame, so a delegation id can only ever be resolved + /// inside the namespace of the connection that named it. + fn claim(&self, registry: &Registry, handle: u64, delegation_id: &str, owner: Owner) -> Claim { + // Registry lookup completes before the in-flight lock is taken; the + // two locks are never held together (see the lock hierarchy above). + let namespace = match registry.get(handle) { + Some(i) => i.namespace, + None => return Claim::Unregistered, + }; + let key = DelegationKey::new(&namespace, delegation_id); + let mut g = self.inflight.lock(); + let owner_handle = match g.get(&key) { + Some(e) => match owner { + Owner::Server => e.to_handle, + Owner::Initiator => e.from_handle, + }, + None => return Claim::NotFound { namespace }, + }; + if owner_handle != handle { + return Claim::WrongOwner { + namespace, + owner_handle, + }; + } + let entry = g.remove(&key).expect("present under the same lock"); + Claim::Owned(entry) + } + /// Handle `cp/delegate_result` from the serving runtime. Returns the /// initiator-bound frame if the delegation is known; unknown ids (e.g. /// results arriving after a CP restart) are dropped with a log. /// - /// Ownership is validated under the SAME lock acquisition that removes - /// the entry (review round-3 F2): the previous remove-check-reinsert - /// dance opened a window in which a genuine result saw an empty table and - /// was dropped, and left the entry momentarily invisible to the deadline - /// sweep. + /// Only the instance the delegation was routed to may complete it, and + /// that check shares the lock acquisition that removes the entry (see + /// [`Claim`]), so a non-owner frame can never make the delegation + /// momentarily invisible to a genuine result or to the deadline sweep. pub fn complete( &self, registry: &Registry, @@ -273,11 +378,37 @@ impl Router { max_result_bytes: usize, next_rpc_id: u64, ) -> Option<(Instance, String)> { - // The namespace comes from the authenticated sender's registration, - // never from the frame (review round-3 F3). - let namespace = match registry.get(serving_handle) { - Some(i) => i.namespace, - None => { + let entry = match self.claim( + registry, + serving_handle, + ¶ms.delegation_id, + Owner::Server, + ) { + Claim::Owned(entry) => entry, + Claim::WrongOwner { + namespace, + owner_handle, + } => { + // Only the instance the delegation was routed to may + // complete it. The entry stays exactly where it is. + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + expected = owner_handle, + got = serving_handle, + "result from unexpected instance — dropped, delegation untouched" + ); + return None; + } + Claim::NotFound { namespace } => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "result for unknown delegation (late arrival or CP restart) — dropped" + ); + return None; + } + Claim::Unregistered => { warn!( handle = serving_handle, delegation = %params.delegation_id, @@ -286,40 +417,12 @@ impl Router { return None; } }; - let key = DelegationKey::new(&namespace, ¶ms.delegation_id); - let entry = { - let mut g = self.inflight.lock(); - match g.get(&key) { - Some(e) if e.to_handle == serving_handle => {} - Some(e) => { - // Only the instance the delegation was routed to may - // complete it. The entry stays exactly where it is. - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - expected = e.to_handle, - got = serving_handle, - "result from unexpected instance — dropped, delegation untouched" - ); - return None; - } - None => { - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - "result for unknown delegation (late arrival or CP restart) — dropped" - ); - return None; - } - } - g.remove(&key).expect("present under the same lock") - }; registry.adjust_sessions(entry.to_handle, -1); // Truncate oversized results (keep the head; delegation already // ran). The marker counts against the cap: the final value never - // exceeds max_result_bytes (review round-2 F5). + // exceeds max_result_bytes. if let Some(r) = ¶ms.result { if r.len() > max_result_bytes { let marker = format!("\n…[truncated by control plane: {} bytes total]", r.len()); @@ -358,12 +461,12 @@ impl Router { /// to the serving runtime, if the delegation is in flight and owned by /// the caller. /// - /// Ownership is validated under the same lock acquisition that removes - /// the entry (review round-3 F2 — no remove/reinsert window), and every - /// refusal returns ONE byte-identical error (review round-3 F3): an - /// unknown id and another instance's live id are indistinguishable to the - /// caller, so `cp/cancel` cannot be used to probe for delegation ids. - /// The distinction is kept in the CP's own logs only. + /// Ownership is validated under the same lock acquisition that removes the + /// entry (see [`Claim`] — no remove/reinsert window), and every refusal + /// returns ONE byte-identical error: an unknown id and another instance's + /// live id are indistinguishable to the caller, so `cp/cancel` cannot be + /// used to probe for delegation ids. The distinction is kept in the CP's + /// own logs only. pub fn cancel( &self, registry: &Registry, @@ -377,40 +480,37 @@ impl Router { "delegation is not in flight for this instance", ) }; - let namespace = match registry.get(from_handle) { - Some(i) => i.namespace, - None => { + let entry = match self.claim( + registry, + from_handle, + ¶ms.delegation_id, + Owner::Initiator, + ) { + Claim::Owned(entry) => entry, + Claim::WrongOwner { namespace, .. } => { warn!( + delegation = %params.delegation_id, + namespace = %namespace, handle = from_handle, - "cancel from an unregistered connection" + "cancel refused: only the initiating instance may cancel" ); return Err(refused()); } - }; - let key = DelegationKey::new(&namespace, ¶ms.delegation_id); - let entry = { - let mut g = self.inflight.lock(); - match g.get(&key) { - Some(e) if e.from_handle == from_handle => {} - Some(_) => { - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - handle = from_handle, - "cancel refused: only the initiating instance may cancel" - ); - return Err(refused()); - } - None => { - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - "cancel refused: delegation not in flight" - ); - return Err(refused()); - } + Claim::NotFound { namespace } => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "cancel refused: delegation not in flight" + ); + return Err(refused()); + } + Claim::Unregistered => { + warn!( + handle = from_handle, + "cancel from an unregistered connection" + ); + return Err(refused()); } - g.remove(&key).expect("present under the same lock") }; registry.adjust_sessions(entry.to_handle, -1); info!(delegation = %params.delegation_id, "delegation cancelled by initiator"); @@ -537,8 +637,7 @@ impl Router { } /// Chain of an in-flight delegation (for tests/inspection). Delegation - /// ids are namespace-scoped (review round-3 F3), so the namespace is part - /// of the lookup. + /// ids are namespace-scoped, so the namespace is part of the lookup. pub fn chain_of(&self, namespace: &str, delegation_id: &str) -> Option> { self.inflight .lock() @@ -707,7 +806,8 @@ type = "worker" #[test] fn inflight_exists_before_target_receives_frame() { - // Review F2: an immediately-arriving result must find the entry. + // An immediately-arriving result must find the entry: it is inserted + // before the forward frame is sent. let mut w = world(); assert!(matches!( do_delegate(&w, delegate_params("d-1", "worker-1", 60)), @@ -1052,8 +1152,9 @@ allow_worker_initiation = true vec!["prod/koudu".to_string()] ); - // Review F3: worker-2 (NOT serving d-root) tries to borrow d-root - // as parent — rejected. + // Borrowed ancestry: worker-2 (NOT serving d-root) tries to use + // d-root as its parent — rejected, so a trusted chain and deadline + // budget cannot be inherited by a stranger. let mut foreign = delegate_params("d-foreign", "worker-2", 60); foreign.parent_delegation_id = Some("d-root".into()); match w.router.delegate( @@ -1126,8 +1227,8 @@ allow_worker_initiation = true #[test] fn wrong_handle_result_never_hides_the_genuine_one() { - // Review round-3 F2: ownership is validated under the same lock - // acquisition that removes the entry. The old remove → validate → + // Ownership is validated under the same lock acquisition that removes + // the entry. The old remove → validate → // reinsert sequence made the entry briefly invisible, so a genuine // result arriving in that window was dropped as "unknown id". for spoof_first in [true, false] { @@ -1191,7 +1292,7 @@ allow_worker_initiation = true #[test] fn genuine_result_survives_concurrent_non_owner_frames() { - // Review round-3 F2, the racing case the sequential test above cannot + // The racing case the sequential test above cannot // observe: with remove → validate → reinsert, a genuine result that // lands inside the window sees an empty table and is dropped, and the // delegation then stalls to its deadline. Under a single lock @@ -1238,7 +1339,7 @@ allow_worker_initiation = true #[test] fn genuine_cancel_survives_concurrent_non_owner_frames() { - // Review round-3 F2 (cancel side), racing case: with the old + // Cancel side, racing case: with the old // remove → validate → reinsert pattern, a genuine initiator cancel // landing inside a non-owner cancel's window would see an empty table // and be refused, leaving the delegation to stall to its deadline. @@ -1280,7 +1381,7 @@ allow_worker_initiation = true #[test] fn refused_cancel_leaves_the_delegation_cancellable() { - // Review round-3 F2 (cancel side): a wrong-handle cancel must not + // Cancel side: a wrong-handle cancel must not // remove-and-reinsert the entry, and must not disturb accounting. let w = world(); assert!(matches!( @@ -1313,7 +1414,7 @@ allow_worker_initiation = true #[test] fn cancel_refusals_are_byte_identical() { - // Review round-3 F3: `cp/cancel` must not be an existence oracle — + // `cp/cancel` must not be an existence oracle — // an unknown id and another instance's live id return the same error // object, byte for byte. let w = world(); @@ -1349,7 +1450,7 @@ allow_worker_initiation = true #[test] fn same_delegation_id_in_two_namespaces_is_independent() { - // Review round-3 F3: the in-flight table is keyed by + // The in-flight table is keyed by // (namespace, delegation_id). A client-supplied id in one namespace // must neither collide with nor be observable from another. let registry = Registry::new(); @@ -1433,7 +1534,7 @@ allow_worker_initiation = true #[test] fn parent_lookup_is_namespace_scoped() { - // Review round-3 F3: parent-chain resolution must not reach into + // Parent-chain resolution must not reach into // another namespace's in-flight table. let cfg: CpConfig = toml::from_str( r#" diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index ecf33b8b6..08db9c41e 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -4,21 +4,28 @@ //! Auth: the runtime presents its key as `Authorization: Bearer ` on the //! upgrade request. Keys never appear in URLs (avoids access-log leakage). //! -//! Resource bounds (review F5): the WS transport enforces -//! `max_frame_bytes` before parsing; each connection's outbound queue is -//! bounded — a peer that cannot drain it is treated as disconnected. +//! Resource bounds: the WS transport enforces `max_frame_bytes` before +//! parsing; each connection's outbound queue is bounded — a peer that cannot +//! drain it is treated as disconnected. //! -//! Admission bounds (review round-3 F4): authentication alone is not a -//! bound. Every connection holds a per-identity slot from the upgrade until -//! it ends (`ConnPermit`, released on every exit path), and must complete -//! `cp/register` within `register_timeout_secs` or be closed. +//! Admission bounds: authentication alone is not a bound. Every connection +//! holds a per-identity slot from the upgrade until it ends (`ConnPermit`, +//! released on every exit path), and must complete `cp/register` within +//! `register_timeout_secs` or be closed. +//! +//! CP-initiated closes carry meaning: both the registration-timeout close and +//! the lease-expiry close are WS code 1008 (policy violation) with a short +//! reason (`registration timeout` / `lease expired`), and an over-quota +//! upgrade is refused with HTTP 503 naming the quota. A client can therefore +//! tell "I misbehaved / my lease lapsed" from a transport-level drop without +//! consulting CP logs. use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use axum::extract::ws::{Message, WebSocket}; +use axum::extract::ws::{close_code, CloseFrame, Message, WebSocket}; use axum::extract::{State, WebSocketUpgrade}; use axum::http::{HeaderMap, StatusCode}; use axum::response::IntoResponse; @@ -44,8 +51,7 @@ pub struct AppState { pub router: Router, rpc_id: AtomicU64, /// Live connections per identity (`namespace/name`), counted from the - /// upgrade so pre-registration sockets are bounded too (review round-3 - /// F4). + /// upgrade so pre-registration sockets are bounded too. conns: Mutex>, } @@ -67,7 +73,7 @@ impl AppState { /// Take a connection slot for `identity`, or `None` when the identity is /// already at `max_connections_per_identity`. The returned guard releases /// the slot on drop — including on every early return and on an upgrade - /// that never completes (review round-3 F4). + /// that never completes. pub fn try_acquire_conn(self: &Arc, identity: &AgentIdentity) -> Option { let key = format!("{}/{}", identity.namespace, identity.name); let mut g = self.conns.lock(); @@ -89,7 +95,7 @@ impl AppState { } /// RAII connection slot. Dropping it frees the identity's quota; it is never -/// released explicitly, so no early return can leak it (review round-3 F4). +/// released explicitly, so no early return can leak it. pub struct ConnPermit { state: Arc, key: String, @@ -118,6 +124,22 @@ async fn health() -> &'static str { "ok" } +/// Reason string on the close frame sent when `cp/register` never arrived. +pub const REASON_REGISTER_TIMEOUT: &str = "registration timeout"; +/// Reason string on the close frame sent when the CP drops a registration +/// because its lease elapsed. +pub const REASON_LEASE_EXPIRED: &str = "lease expired"; + +/// A CP-initiated close that states why. Code 1008 (policy violation) plus a +/// short reason, so a client can distinguish "the CP closed me on purpose" +/// from a transport-level drop and act on it (re-register vs. plain retry). +fn policy_close(reason: &'static str) -> Message { + Message::Close(Some(CloseFrame { + code: close_code::POLICY, + reason: reason.into(), + })) +} + async fn ws_handler( State(state): State>, headers: HeaderMap, @@ -135,7 +157,7 @@ async fn ws_handler( } }; // Per-identity connection quota, taken before the upgrade so an - // over-quota peer is refused at the HTTP layer (review round-3 F4). + // over-quota peer is refused at the HTTP layer. let permit = match state.try_acquire_conn(&identity) { Some(p) => p, None => { @@ -144,7 +166,16 @@ async fn ws_handler( max = state.cfg.max_connections_per_identity, "WS rejected: identity is at its connection quota" ); - return StatusCode::SERVICE_UNAVAILABLE.into_response(); + // Name the quota in the body: without it a client cannot tell an + // exhausted quota from an overloaded CP, and both are 503. + return ( + StatusCode::SERVICE_UNAVAILABLE, + format!( + "identity is at its connection quota (max_connections_per_identity = {})\n", + state.cfg.max_connections_per_identity + ), + ) + .into_response(); } }; let max_frame = state.cfg.max_frame_bytes; @@ -158,14 +189,14 @@ async fn handle_connection( socket: WebSocket, identity: AgentIdentity, // Held for the connection's whole lifetime; dropped here on every exit - // path, including the early returns below (review round-3 F4). + // path, including the early returns below. _permit: ConnPermit, ) { let (mut sink, mut stream) = socket.split(); // --- Registration: mandatory first frame, within a deadline --- // An authenticated peer must not be able to park idle sockets: pings keep - // the transport alive but do not extend this deadline (review round-3 F4). + // the transport alive but do not extend this deadline. let register = match tokio::time::timeout( Duration::from_secs(state.cfg.register_timeout_secs), async { @@ -191,7 +222,7 @@ async fn handle_connection( timeout_secs = state.cfg.register_timeout_secs, "no cp/register within the registration deadline — closing" ); - let _ = sink.send(Message::Close(None)).await; + let _ = sink.send(policy_close(REASON_REGISTER_TIMEOUT)).await; return; } }; @@ -208,12 +239,12 @@ async fn handle_connection( } }; - // Outbound channel for this connection. Bounded (review F5): a peer that + // Outbound channel for this connection. Bounded: a peer that // cannot drain OUTBOUND_QUEUE frames is disconnected, not buffered. let (tx, mut rx) = mpsc::channel::(OUTBOUND_QUEUE); // Shutdown signal so the CP can close this socket when it drops the - // registration on its own initiative (lease expiry — review round-3 F1). + // registration on its own initiative (lease expiry). // Subscribed BEFORE registering so no signal can be missed, and kept // alive here for the whole connection: closing is driven by an explicit // signal, never by the registry happening to drop its side. @@ -224,7 +255,7 @@ async fn handle_connection( Some(cap) => reg.max_delegated_sessions.min(cap), None => reg.max_delegated_sessions, }; - // The registry assigns the CP-generated handle (review F1): ownership + // The registry assigns the CP-generated handle: ownership // and teardown never key on the client-supplied instance_id. let handle = state.registry.register_conn( Instance { @@ -278,7 +309,7 @@ async fn handle_connection( loop { tokio::select! { // The CP dropped this registration (lease expiry): the socket - // must go too (review round-3 F1). Keeping it open would leave a + // must go too. Keeping it open would leave a // connection whose every frame hits an absent registry entry and // which can never re-register, since registration is // first-frame-only. Closing lets the client reconnect, @@ -328,7 +359,7 @@ async fn handle_connection( handle, "closing connection at the CP's request (registration dropped)" ); - let _ = sink.send(Message::Close(None)).await; + let _ = sink.send(policy_close(REASON_LEASE_EXPIRED)).await; } teardown(&state, handle, &identity); @@ -448,7 +479,32 @@ fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option i, + None => { + warn!( + handle, + method = %method, + "frame on a connection whose registration is gone (swept lease) — NOT_REGISTERED" + ); + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new( + codes::NOT_REGISTERED, + "connection is no longer registered (lease expired); reconnect and re-register", + ), + ); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + }; macro_rules! params_or_err { ($ty:ty) => { @@ -558,10 +614,10 @@ fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option, lease: Duration) { for handle in state.registry.expired(lease) { @@ -705,7 +761,7 @@ mod tests { #[test] fn register_invalid_envelope_rejected() { - // Missing jsonrpc field (review F4). + // Missing jsonrpc field: the envelope is validated, not assumed. let no_ver = serde_json::json!({ "id": 6, "method": "cp/register", "params": { @@ -744,7 +800,7 @@ mod tests { #[test] fn conn_quota_bounds_and_recycles_slots() { - // Review round-3 F4(b): the quota is a hard bound and the guard + // The quota is a hard bound and the guard // releases the slot on drop, so no exit path can leak it. let state = state_with("max_connections_per_identity = 2"); let id = identity(); @@ -788,7 +844,7 @@ mod tests { #[tokio::test] async fn sweep_leases_signals_the_connection_before_dropping_it() { - // Review round-3 F1 at the sweeper level: the shutdown signal is + // At the sweeper level: the shutdown signal is // delivered, not just the registry entry removed. (The end-to-end // proof over a real socket lives in tests/ws_lifecycle.rs.) let state = state_with("heartbeat_interval_secs = 1\nlease_expiry_secs = 2"); @@ -825,4 +881,68 @@ mod tests { "the owning connection must be told to close" ); } + + #[tokio::test] + async fn frame_on_a_swept_handle_is_answered_not_registered() { + // Frames can arrive between the sweeper dropping a registration and + // the connection task observing the close signal. They must be + // answered: silence is indistinguishable from a hung CP, and the + // client needs to know it has to reconnect and register again. + let state = state_with("heartbeat_interval_secs = 1\nlease_expiry_secs = 2"); + let signal = crate::registry::shutdown_signal(); + let (tx, _rx) = mpsc::channel::(OUTBOUND_QUEUE); + let handle = state.registry.register_conn( + Instance { + handle: 0, + namespace: "prod".into(), + name: "koudu".into(), + agent_type: AgentType::Primary, + instance_id: "i-1".into(), + labels: Default::default(), + max_delegated_sessions: 1, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + }, + Arc::clone(&signal), + ); + + // While registered, a heartbeat is answered normally. + let hb = serde_json::json!({ + "jsonrpc": "2.0", "id": 7, "method": "cp/heartbeat", + "params": {"instance_id": "i-1"} + }) + .to_string(); + let ok: serde_json::Value = + serde_json::from_str(&handle_frame(&state, handle, &hb).expect("answered")).unwrap(); + assert_eq!(ok["result"]["ok"], true); + + // Sweep the lease, then replay the same frame on the same handle. + sweep_leases(&state, Duration::ZERO); + assert!(state.registry.get(handle).is_none(), "handle was swept"); + + let reply = handle_frame(&state, handle, &hb) + .expect("a frame on a swept handle must be answered, not dropped"); + let v: serde_json::Value = serde_json::from_str(&reply).unwrap(); + assert_eq!(v["id"], 7, "the error must correlate with the request"); + assert_eq!(v["error"]["code"], codes::NOT_REGISTERED); + + // Same for a delegate attempt: no method reaches the router with an + // absent registration. + let del = serde_json::json!({ + "jsonrpc": "2.0", "id": 8, "method": "cp/delegate", + "params": { + "delegation_id": "d-1", + "target": {"name": "worker-1"}, + "prompt": "do it", + "deadline": "2999-01-01T00:00:00Z" + } + }) + .to_string(); + let v2: serde_json::Value = + serde_json::from_str(&handle_frame(&state, handle, &del).expect("answered")).unwrap(); + assert_eq!(v2["error"]["code"], codes::NOT_REGISTERED); + assert_eq!(state.router.inflight_count(), 0); + } } diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index 0447bf1df..b3df9025a 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -1,6 +1,6 @@ //! End-to-end WebSocket lifecycle tests: connection termination on lease -//! expiry (review round-3 F1) and pre-registration admission bounds -//! (review round-3 F4). +//! expiry, pre-registration admission bounds, and the meaning the CP attaches +//! to its own closes (WS 1008 + reason, HTTP 503 naming the quota). //! //! These drive a real CP over a loopback socket with a real WS client, which //! is the only way to prove that the connection *task* reacts — the earlier @@ -102,26 +102,63 @@ async fn register(ws: &mut Ws, instance_id: &str) -> serde_json::Value { serde_json::from_str(msg.to_text().unwrap()).unwrap() } -/// Wait until the peer closes the socket (Close frame, error, or EOF). -async fn wait_closed(ws: &mut Ws, within: Duration) -> bool { +/// How a connection ended, as observed by the client. +#[derive(Debug, PartialEq, Eq)] +enum Closed { + /// A WS Close frame carrying a code and a reason. + Frame { code: u16, reason: String }, + /// A Close frame with no payload — the CP never sends these on purpose. + Bare, + /// EOF or transport error with no Close frame at all. + Dropped, +} + +/// Wait until the peer closes the socket, and report how. `None` means the +/// socket was still open when `within` elapsed. +async fn wait_closed(ws: &mut Ws, within: Duration) -> Option { let deadline = Instant::now() + within; while Instant::now() < deadline { match tokio::time::timeout(Duration::from_millis(200), ws.next()).await { - Ok(None) | Ok(Some(Err(_))) => return true, - Ok(Some(Ok(Message::Close(_)))) => return true, + Ok(None) | Ok(Some(Err(_))) => return Some(Closed::Dropped), + Ok(Some(Ok(Message::Close(Some(cf))))) => { + return Some(Closed::Frame { + code: cf.code.into(), + reason: cf.reason.to_string(), + }) + } + Ok(Some(Ok(Message::Close(None)))) => return Some(Closed::Bare), Ok(Some(Ok(_))) => continue, Err(_) => continue, // read timeout: keep waiting } } - false + None +} + +/// The close a CP-initiated termination must carry: 1008 (policy violation) +/// plus a short reason the client can act on. +fn policy_close(reason: &str) -> Closed { + Closed::Frame { + code: 1008, + reason: reason.to_string(), + } +} + +/// Body of a refused upgrade, as text. +fn http_body(resp: &tokio_tungstenite::tungstenite::http::Response>>) -> String { + resp.body() + .as_ref() + .map(|b| String::from_utf8_lossy(b).to_string()) + .unwrap_or_default() } #[tokio::test] async fn lease_expiry_closes_the_connection_and_permits_reregistration() { - // Review round-3 F1: deregistering on lease expiry without terminating + // Deregistering on lease expiry without terminating // the connection left a live socket bound to a registration that no // longer existed — heartbeats got no reply and re-registration was - // impossible (registration is first-frame-only). The CP must close it. + // impossible (registration is first-frame-only). The CP must close it, + // and the close must say why so the client can distinguish it from a + // network drop. let (state, url) = spawn_cp(cfg("max_connections_per_identity = 1")).await; let mut ws = connect(&url).await.expect("first connection accepted"); @@ -136,9 +173,13 @@ async fn lease_expiry_closes_the_connection_and_permits_reregistration() { "lease expiry deregisters" ); - assert!( - wait_closed(&mut ws, Duration::from_secs(5)).await, - "the connection task must observe the shutdown signal and close" + let closed = wait_closed(&mut ws, Duration::from_secs(5)) + .await + .expect("the connection task must observe the shutdown signal and close"); + assert_eq!( + closed, + policy_close("lease expired"), + "a lease-expiry close must be 1008 with a reason, not an anonymous close" ); drop(ws); @@ -154,37 +195,32 @@ async fn lease_expiry_closes_the_connection_and_permits_reregistration() { #[tokio::test] async fn ping_only_pre_registration_socket_is_closed_at_the_deadline() { - // Review round-3 F4(a): pings keep the transport alive but must not - // extend the registration deadline. + // Pings keep the transport alive but must not + // extend the registration deadline — and the close must name the reason. let (state, url) = spawn_cp(cfg("register_timeout_secs = 1")).await; let mut ws = connect(&url).await.expect("connection accepted"); let started = Instant::now(); - let mut closed = false; + let mut closed = None; while started.elapsed() < Duration::from_secs(10) { let _ = ws.send(Message::Ping(vec![7].into())).await; - match tokio::time::timeout(Duration::from_millis(250), ws.next()).await { - Ok(None) | Ok(Some(Err(_))) => { - closed = true; - break; - } - Ok(Some(Ok(Message::Close(_)))) => { - closed = true; - break; - } - _ => continue, + if let Some(how) = wait_closed(&mut ws, Duration::from_millis(250)).await { + closed = Some(how); + break; } } - assert!( + let closed = closed.expect("an authenticated socket that never registers must be closed"); + assert_eq!( closed, - "an authenticated socket that never registers must be closed" + policy_close("registration timeout"), + "a registration-timeout close must be 1008 with a reason" ); assert!(state.registry.list("prod").is_empty()); } #[tokio::test] async fn registration_after_the_deadline_is_not_accepted() { - // Review round-3 F4(a): the deadline is enforced, not merely advisory. + // The deadline is enforced, not merely advisory. let (state, url) = spawn_cp(cfg("register_timeout_secs = 1")).await; let mut ws = connect(&url).await.expect("connection accepted"); tokio::time::sleep(Duration::from_millis(1_600)).await; @@ -202,9 +238,10 @@ async fn registration_after_the_deadline_is_not_accepted() { #[tokio::test] async fn connection_quota_rejects_over_limit_and_recycles_on_disconnect() { - // Review round-3 F4(b): the quota bounds concurrent sockets per identity + // The quota bounds concurrent sockets per identity // and is released on every exit path (RAII), so connect → disconnect → - // connect always succeeds. + // connect always succeeds. The refusal names the quota: 503 alone cannot + // be told apart from an overloaded CP. let (state, url) = spawn_cp(cfg( "max_connections_per_identity = 1\nregister_timeout_secs = 30", )) @@ -215,11 +252,18 @@ async fn connection_quota_rejects_over_limit_and_recycles_on_disconnect() { assert_eq!(state.conn_count("prod/koudu"), 1); match connect(&url).await { - Err(WsError::Http(resp)) => assert_eq!( - resp.status(), - StatusCode::SERVICE_UNAVAILABLE, - "over-quota upgrade must be refused before the WS handshake" - ), + Err(WsError::Http(resp)) => { + assert_eq!( + resp.status(), + StatusCode::SERVICE_UNAVAILABLE, + "over-quota upgrade must be refused before the WS handshake" + ); + let body = http_body(&resp); + assert!( + body.contains("max_connections_per_identity"), + "the 503 body must name the quota, got {body:?}" + ); + } Err(e) => panic!("unexpected error: {e}"), Ok(_) => panic!("over-quota connection must be rejected"), } @@ -238,7 +282,7 @@ async fn connection_quota_rejects_over_limit_and_recycles_on_disconnect() { #[tokio::test] async fn pre_registration_sockets_count_against_the_quota() { - // Review round-3 F4(b): the quota is taken at the upgrade, so parked + // The quota is taken at the upgrade, so parked // pre-registration sockets cannot be multiplied for free. let (_state, url) = spawn_cp(cfg( "max_connections_per_identity = 1\nregister_timeout_secs = 30", @@ -247,7 +291,10 @@ async fn pre_registration_sockets_count_against_the_quota() { let _parked = connect(&url).await.expect("connection accepted"); match connect(&url).await { - Err(WsError::Http(resp)) => assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE), + Err(WsError::Http(resp)) => { + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); + assert!(http_body(&resp).contains("max_connections_per_identity")); + } Err(e) => panic!("unexpected error: {e}"), Ok(_) => panic!("an unregistered socket must still occupy its quota slot"), } diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index edf68c402..b347288e2 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -100,8 +100,8 @@ Key properties: contract (§9). 4. **CP connectivity is strictly additive.** Loss of the CP link never affects normal platform (Discord/Slack) operation; the runtime reconnects - with backoff. The CP itself is stateless enough that a restart only means - re-registration. + with backoff. The CP holds no durable state, so a restart costs + re-registration plus whatever delegations were in flight (§4). ### Naming @@ -238,10 +238,40 @@ recovery semantics: (completion, cancellation, parent linkage) compare handles. The ack carries the heartbeat interval, lease window, and the effective (possibly clamped) concurrency budget. Instances missing heartbeats past the lease - are deregistered; their in-flight delegations fail immediately with - `target_disconnected`. Heartbeats refresh the lease only — CP-owned + are deregistered, and their in-flight delegations fail immediately with a + **side-specific** outcome: delegations the expired instance was *serving* + send `target_disconnected` to the initiator, while delegations it had + *initiated* send a best-effort `cp/cancel` to the still-live server (and + release its reserved capacity). The two are mutually exclusive — one dead + instance never produces both frames for the same delegation. Heartbeats + refresh the lease only — CP-owned in-flight accounting is authoritative and never merged from runtime reports. +- **Registration deadline.** `cp/register` must arrive within + `register_timeout_secs` of the completed WebSocket upgrade. WS Ping/Pong + keeps the transport alive but does **not** extend the deadline, and a + `cp/register` that arrives after it is never acked. On expiry the CP closes + the socket with WS code **1008** (policy violation) and reason + `registration timeout`. Rationale: authentication alone bounds nothing — an + authenticated peer could otherwise park sockets indefinitely in the + pre-registration state. +- **Registration is per-connection and first-frame-only, so lease expiry ends + the connection.** When the CP drops a registration on its own initiative it + also closes the socket — WS code **1008**, reason `lease expired`. Leaving + it open would strand a connection whose every subsequent frame hits an + absent registry entry and which can never re-register. Recovery is + therefore always the same shape: **reconnect, re-authenticate, re-register** + (a new handle; the old one is gone for good). Frames that arrive in the + window between the sweep and the close are answered `NOT_REGISTERED` rather + than dropped silently, so a client can tell a swept lease from a hung CP. +- **Per-identity connection quota.** `max_connections_per_identity` bounds + concurrent sockets per identity and is counted **from the upgrade**, so + pre-registration sockets occupy a slot too. An over-quota upgrade is refused + at the HTTP layer with **503** and a body naming the quota (a bare 503 is + indistinguishable from an overloaded CP). The slot is held by an RAII guard + released on every exit path, so no early return, failed handshake, or panic + can leak it. Replicas of one logical agent share one identity, so this is + also the replica ceiling. - **Resource bounds.** The WS transport rejects messages over `max_frame_bytes` before parsing; oversized `prompt`s are rejected (`max_prompt_bytes`); per-connection outbound queues are bounded and a @@ -253,11 +283,34 @@ recovery semantics: CP replies `SATURATED` immediately. The CP never queues — v1 has no durable state, and a hidden in-memory queue would contradict that. `NO_TARGET` (nothing matches) is a distinct error. -- **CP restart semantics.** The in-flight table is in-memory. After a CP - restart, in-flight delegations end as initiator-side timeouts (the - propagated deadline is the upper bound); late `cp/delegate_result` frames - for unknown ids are acknowledged, logged, and dropped so reconnecting - runtimes do not error-loop. +- **Delegation ids are scoped to `(namespace, delegation_id)`.** The id is + client-supplied, so it is only unique within the namespace that produced + it. Two namespaces may hold the same id concurrently, legally and + invisibly: `DUPLICATE_DELEGATION` only ever refers to the caller's own + namespace, parent-chain lookup never reaches across namespaces, and result + routing resolves the id inside the sender's registered namespace. `cp/cancel` + refusals are deliberately **indistinguishable**: an unknown id and another + instance's live id return the same `POLICY_DENIED` error object, byte for + byte, so cancel cannot be used as an existence oracle for other tenants' + delegation ids. The CP's own logs keep the distinction. +- **CP restart semantics.** A CP restart is equivalent to every lease + expiring at once *with* the connection closure that implies — except that + the CP is not there to send it: the in-flight table and the sockets die + together with the process, so no synthesized `timeout` or + `target_disconnected` frame can be emitted for delegations that were in + flight. Runtimes observe the transport drop, reconnect with backoff, and + re-register (new handles, empty in-flight table). Initiators reconcile + against the deadline they already propagated, which is the upper bound on + every orphaned delegation. Once the CP is back, late + `cp/delegate_result` frames for unknown ids are acknowledged, logged, and + dropped so reconnecting runtimes do not error-loop, and a frame from a + connection that has not re-registered is answered `NOT_REGISTERED`. Within a + *live* CP the synthesized failures do happen, but per side, not both at + once: lease expiry or disconnect sends `target_disconnected` to the + initiator when the *serving* instance died, and a best-effort `cp/cancel` + downstream when the *initiating* instance died. Only a deadline sweep emits + both frames for one delegation (see below), because there both peers are + still connected. - **Timeout and disconnect synthesis.** A deadline sweep terminates overdue delegations: the initiator receives a synthesized `timeout` result and the serving runtime a best-effort `cp/cancel` (stop burning tokens). Worker @@ -267,7 +320,8 @@ recovery semantics: configured `max_result_bytes` (default 256 KiB) is truncated head-first with an explicit marker. - **Idempotency.** `delegation_id` is the caller-generated idempotency key; - a duplicate in-flight id is rejected (`DUPLICATE_DELEGATION`). Only the + a duplicate id already in flight **in the caller's own namespace** is + rejected (`DUPLICATE_DELEGATION`). Only the instance a delegation was routed to may complete it; only the initiating instance may cancel it. From 1de59fb3ba6ac0e57a4714f10e2e7169384e4e03 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Thu, 13 Aug 2026 16:16:35 +0000 Subject: [PATCH 05/11] =?UTF-8?q?fix(cp):=20round-5=20review=20=E2=80=94?= =?UTF-8?q?=20exactly-once=20capacity=20release,=20terminal-result=20deliv?= =?UTF-8?q?ery=20commit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 (critical): delegate's send-failure rollback decremented the target's session count even when a concurrent fail_instance/sweep_deadlines had already removed the entry and released the reservation. The saturating math hid the double release, so saturated() could admit work to a full instance. The decrement is now conditional on this call actually removing the entry: whoever removes an in-flight entry releases its capacity, exactly once. Regression: concurrent_fail_instance_never_double_releases_capacity (200-iteration barrier race with a live baseline delegation). F2 (critical): cp/delegate_result delivery was fire-and-forget after the entry was irreversibly removed — a full initiator queue silently lost the result while the serving side was acked ok:true, contradicting the documented cannot-drain-means-disconnected contract. complete() is now two-phase: peek (validate ownership without removing), send, then commit (claim + capacity release) only after the initiator's queue accepted the frame. On refusal the entry stays in flight, the serving side gets a TARGET_DISCONNECTED error instead of a false ack, and the initiator is closed (WS 1008 "outbound queue overflow") so the delegation resolves through the fail_instance path. The shutdown signal now carries the close reason. Regressions: stalled_initiator_result_is_never_silently_lost, stalled_initiator_is_disconnected_and_serving_side_not_falsely_acked. F3: teardown documents its idempotency contract with sweep_leases. F4: ADR section 7 now distinguishes the runtime-to-CP Bearer/WS boundary (shipped) from the agent-to-facade UDS boundary (PR 3/4); the never-on-TCP claim is scoped to the local facade. The v1 contract amendments record the terminal-result and capacity-release semantics. F5: docs/control-plane.md quickstart + README feature entry. --- README.md | 1 + crates/openab-cp/src/registry.rs | 35 ++- crates/openab-cp/src/router.rs | 453 ++++++++++++++++++++++--------- crates/openab-cp/src/server.rs | 164 +++++++++-- docs/adr/agent-control-plane.md | 40 ++- docs/control-plane.md | 54 ++++ 6 files changed, 587 insertions(+), 160 deletions(-) create mode 100644 docs/control-plane.md diff --git a/README.md b/README.md index ae73e02e9..a9628a824 100644 --- a/README.md +++ b/README.md @@ -58,6 +58,7 @@ webhook platforms and `WS/webhook` for Feishu/Lark. - **@mention trigger** — mention the bot in an allowed channel to start a conversation - **Thread-based multi-turn** — auto-creates threads; no @mention needed for follow-ups - **Multi-agent collaboration** — bot-to-bot messaging for coordinated workflows ([docs/multi-agent.md](docs/multi-agent.md)) +- **Agent control plane (preview)** — standalone `openab-cp` service for direct agent-to-agent delegation over WebSocket, with identity-bound registration and CP-authoritative policy; runtime client and facade land in follow-up releases ([docs/control-plane.md](docs/control-plane.md)) - **Agent-controlled reply-to** — agents choose which message to reply to via `[[reply_to:id]]` directive, enabling clear conversation threads in multi-bot channels ([docs/output-directives.md](docs/output-directives.md)) - **Edit-streaming** — live-updates the Discord message every 1.5s as tokens arrive - **Emoji status reactions** — 👀→🤔→🔥/👨‍💻/⚡→👍+random mood face diff --git a/crates/openab-cp/src/registry.rs b/crates/openab-cp/src/registry.rs index 7a3c27a5d..3b0518460 100644 --- a/crates/openab-cp/src/registry.rs +++ b/crates/openab-cp/src/registry.rs @@ -33,12 +33,14 @@ pub const OUTBOUND_QUEUE: usize = 256; /// `watch` (not `oneshot`) so the connection task can select on it repeatedly, /// and wrapped in `Arc` so the registry entry and the connection task share /// one signal without either side's drop cancelling it. -pub type ShutdownTx = Arc>; +pub type ShutdownTx = Arc>>; /// Create a fresh connection shutdown signal. The connection task keeps the -/// returned handle (to `subscribe()`), the registry keeps a clone. +/// returned handle (to `subscribe()`), the registry keeps a clone. The value +/// is `None` until the CP requests the close, then the close reason (sent on +/// the WS close frame so the client knows why it was terminated). pub fn shutdown_signal() -> ShutdownTx { - Arc::new(watch::channel(false).0) + Arc::new(watch::channel(None).0) } /// Registry slot: the public instance view plus CP-internal connection @@ -121,14 +123,15 @@ impl Registry { self.register_conn(inst, shutdown_signal()) } - /// Ask the owning connection task to close. Returns whether a live - /// registration was signalled. Must be called BEFORE `deregister`, which - /// drops the registry's handle on the signal. - pub fn signal_shutdown(&self, handle: u64) -> bool { + /// Ask the owning connection task to close, with the reason to put on + /// the close frame. Returns whether a live registration was signalled. + /// Must be called BEFORE `deregister`, which drops the registry's handle + /// on the signal. + pub fn signal_shutdown(&self, handle: u64, reason: &'static str) -> bool { match self.inner.read().get(&handle) { Some(e) => { // `send_replace` cannot fail even with no receivers left. - e.shutdown.send_replace(true); + e.shutdown.send_replace(Some(reason)); true } None => false, @@ -409,15 +412,19 @@ mod tests { let sig = shutdown_signal(); let mut rx = sig.subscribe(); let h = r.register_conn(inst("prod", "w1", "i-1", 1), sig); - assert!(!*rx.borrow()); + assert!(rx.borrow().is_none()); - assert!(r.signal_shutdown(h)); + assert!(r.signal_shutdown(h, "lease expired")); rx.changed().await.unwrap(); - assert!(*rx.borrow(), "connection task must observe the signal"); + assert_eq!( + *rx.borrow(), + Some("lease expired"), + "connection task must observe the signal and its reason" + ); // After deregistration there is nothing left to signal. r.deregister(h); - assert!(!r.signal_shutdown(h)); + assert!(!r.signal_shutdown(h, "lease expired")); } #[tokio::test] @@ -428,9 +435,9 @@ mod tests { let sig = shutdown_signal(); let mut rx = sig.subscribe(); let h = r.register_conn(inst("prod", "w1", "i-1", 1), sig); - r.signal_shutdown(h); + r.signal_shutdown(h, "lease expired"); r.deregister(h); rx.changed().await.unwrap(); - assert!(*rx.borrow()); + assert_eq!(*rx.borrow(), Some("lease expired")); } } diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index 2d10d31ce..bf1f8b22c 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -162,6 +162,27 @@ enum Claim { Unregistered, } +/// Outcome of a `cp/delegate_result` frame (see [`Router::complete`]). +#[derive(Debug, PartialEq, Eq)] +pub enum CompleteOutcome { + /// The result was accepted by the initiator's queue; the delegation is + /// finished and the serving instance's capacity was released. + Delivered, + /// The frame was refused or the delegation is unknown (wrong owner, + /// unknown id, unregistered caller, or the initiator is gone). Nothing + /// changed; each case is logged. + Dropped, + /// The initiator's bounded outbound queue refused the terminal result. + /// The entry is still in flight: the caller must treat the initiator as + /// disconnected (close its connection), whose teardown then fails the + /// delegation through `fail_instance` — capacity is released exactly + /// once and the serving runtime receives `cp/cancel`. + InitiatorStalled { + /// Registration handle of the stalled initiator. + initiator_handle: u64, + }, +} + impl Router { pub fn new() -> Self { Self { @@ -306,8 +327,18 @@ impl Router { if target.tx.try_send(text).is_err() { // Disconnected or backpressured beyond its queue: roll back. - self.inflight.lock().remove(&key); - registry.adjust_sessions(target.handle, -1); + // + // Roll back only what this call still owns. `fail_instance` and + // `sweep_deadlines` take the in-flight lock without the admission + // lock, so they can remove this very entry between the insert + // above and this branch — and whoever removes an entry also + // releases its capacity reservation. Decrementing here after a + // concurrent removal would double-release: the saturating math + // hides the underflow and `saturated()` then admits new work to + // an instance that is actually full. + if self.inflight.lock().remove(&key).is_some() { + registry.adjust_sessions(target.handle, -1); + } return DelegateOutcome::Rejected(ErrorObject::new( codes::TARGET_DISCONNECTED, "target disconnected or unresponsive during routing", @@ -362,14 +393,33 @@ impl Router { Claim::Owned(entry) } - /// Handle `cp/delegate_result` from the serving runtime. Returns the - /// initiator-bound frame if the delegation is known; unknown ids (e.g. - /// results arriving after a CP restart) are dropped with a log. + /// Handle `cp/delegate_result` from the serving runtime. + /// + /// The terminal result is the one frame that must never be silently + /// dropped, so delivery happens in two phases: /// - /// Only the instance the delegation was routed to may complete it, and - /// that check shares the lock acquisition that removes the entry (see - /// [`Claim`]), so a non-owner frame can never make the delegation - /// momentarily invisible to a genuine result or to the deadline sweep. + /// 1. **Peek** — validate ownership under one in-flight lock acquisition + /// without removing the entry, then build and `try_send` the + /// initiator-bound frame. + /// 2. **Commit** — only after the initiator's queue accepted the frame, + /// remove the entry (via [`Claim`], same single-lock property) and + /// release the serving instance's capacity. + /// + /// If the initiator's bounded queue refuses the frame, the entry stays + /// in flight and [`CompleteOutcome::InitiatorStalled`] tells the caller + /// to treat the initiator as disconnected (per the bounded-queue + /// contract): its teardown runs `fail_instance`, which releases capacity + /// exactly once and sends `cp/cancel` to the serving runtime. + /// + /// Peek-then-commit admits one benign race: two concurrent genuine + /// results for the same id can both pass the peek and both be delivered, + /// but only the first commit releases capacity (the second finds the + /// entry gone and does nothing). Duplicate `cp/delegate_result` frames + /// are correlated by `delegation_id` and idempotent for the initiator. + /// + /// Only the instance the delegation was routed to may complete it; a + /// non-owner frame can never make the delegation momentarily invisible + /// to a genuine result or to the deadline sweep. pub fn complete( &self, registry: &Registry, @@ -377,48 +427,48 @@ impl Router { mut params: DelegateResultParams, max_result_bytes: usize, next_rpc_id: u64, - ) -> Option<(Instance, String)> { - let entry = match self.claim( - registry, - serving_handle, - ¶ms.delegation_id, - Owner::Server, - ) { - Claim::Owned(entry) => entry, - Claim::WrongOwner { - namespace, - owner_handle, - } => { - // Only the instance the delegation was routed to may - // complete it. The entry stays exactly where it is. - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - expected = owner_handle, - got = serving_handle, - "result from unexpected instance — dropped, delegation untouched" - ); - return None; - } - Claim::NotFound { namespace } => { - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - "result for unknown delegation (late arrival or CP restart) — dropped" - ); - return None; - } - Claim::Unregistered => { + ) -> CompleteOutcome { + // Phase 1 — peek: validate without removing. Removing before the + // send would make a refused send unrecoverable (silent loss of a + // computed result while the serving side is acked as delivered). + let namespace = match registry.get(serving_handle) { + Some(i) => i.namespace, + None => { warn!( handle = serving_handle, delegation = %params.delegation_id, "result from an unregistered connection — dropped" ); - return None; + return CompleteOutcome::Dropped; + } + }; + let key = DelegationKey::new(&namespace, ¶ms.delegation_id); + let entry = { + let g = self.inflight.lock(); + match g.get(&key) { + Some(e) if e.to_handle == serving_handle => e.clone(), + Some(e) => { + // Only the instance the delegation was routed to may + // complete it. The entry stays exactly where it is. + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + expected = e.to_handle, + got = serving_handle, + "result from unexpected instance — dropped, delegation untouched" + ); + return CompleteOutcome::Dropped; + } + None => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "result for unknown delegation (late arrival or CP restart) — dropped" + ); + return CompleteOutcome::Dropped; + } } }; - - registry.adjust_sessions(entry.to_handle, -1); // Truncate oversized results (keep the head; delegation already // ran). The marker counts against the cap: the final value never @@ -437,24 +487,57 @@ impl Router { } } - info!( - delegation = %params.delegation_id, - status = ?params.status, - from = %entry.to_logical, - to = %entry.from_logical, - "delegation completed" - ); - - let initiator = registry.get(entry.from_handle)?; + let Some(initiator) = registry.get(entry.from_handle) else { + // The initiator deregistered concurrently: its `fail_instance` + // pass removes this entry, releases capacity, and cancels the + // serving side — nothing to do here. + warn!( + delegation = %params.delegation_id, + "result for a delegation whose initiator is gone — dropped" + ); + return CompleteOutcome::Dropped; + }; let frame = JsonRpcRequest::new( next_rpc_id, methods::DELEGATE_RESULT, Some(serde_json::to_value(¶ms).expect("serializable")), ); - Some(( - initiator, - serde_json::to_string(&frame).expect("serializable"), - )) + let text = serde_json::to_string(&frame).expect("serializable"); + + if initiator.tx.try_send(text).is_err() { + // Bounded-queue contract: a peer that cannot drain its queue is + // treated as disconnected, never silently skipped. The entry + // stays in flight; the caller closes the initiator, whose + // teardown fails the delegation over the `fail_instance` path. + warn!( + delegation = %params.delegation_id, + initiator = %entry.from_logical, + "initiator queue full — terminal result refused, treating initiator as disconnected" + ); + return CompleteOutcome::InitiatorStalled { + initiator_handle: entry.from_handle, + }; + } + + // Phase 2 — commit. If a concurrent path (duplicate result, cancel, + // sweep, fail_instance) removed the entry between peek and now, that + // path also released the capacity — do not decrement twice. + if let Claim::Owned(e) = self.claim( + registry, + serving_handle, + ¶ms.delegation_id, + Owner::Server, + ) { + registry.adjust_sessions(e.to_handle, -1); + info!( + delegation = %params.delegation_id, + status = ?params.status, + from = %e.to_logical, + to = %e.from_logical, + "delegation completed" + ); + } + CompleteOutcome::Delivered } /// Handle `cp/cancel` from the initiator. Returns the frame to forward @@ -794,11 +877,11 @@ type = "worker" result: Some("done".into()), error: None, }; - let (init, frame) = w - .router - .complete(&w.registry, w.h_worker, result, 1024, 2) - .unwrap(); - assert_eq!(init.handle, w.h_primary); + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, 1024, 2), + CompleteOutcome::Delivered + ); + let frame = w.primary_rx.try_recv().unwrap(); assert!(frame.contains("\"completed\"")); assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); assert_eq!(w.router.inflight_count(), 0); @@ -820,10 +903,10 @@ type = "worker" result: Some("instant".into()), error: None, }; - assert!(w - .router - .complete(&w.registry, w.h_worker, result, 1024, 2) - .is_some()); + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, 1024, 2), + CompleteOutcome::Delivered + ); w.worker_rx.try_recv().unwrap(); } @@ -844,6 +927,118 @@ type = "worker" ); } + #[test] + fn concurrent_fail_instance_never_double_releases_capacity() { + // delegate's rollback and fail_instance can race on the same entry: + // fail_instance takes the in-flight lock without the admission lock, + // so it can remove the entry (and release its reservation) between + // delegate's insert and a failing try_send. Whoever removes the + // entry releases the capacity — exactly once. A double release + // silently undercounts the target (saturating math) and lets + // `saturated()` admit work to an instance that is actually full. + for _ in 0..200 { + let registry = Registry::new(); + let router = Router::new(); + let cfg = cfg(); + let (p1, _p1_rx) = instance("prod", "koudu", AgentType::Primary, 4); + let (p2, _p2_rx) = instance("prod", "koudu-2", AgentType::Primary, 4); + let (wk, mut worker_rx) = instance("prod", "worker-1", AgentType::Worker, 4); + let hp1 = registry.register(p1); + let hp2 = registry.register(p2); + let hw = registry.register(wk); + + // Baseline: a live delegation from p1 keeps the true count at 1. + match router.delegate( + &cfg, + ®istry, + "prod", + "koudu", + &AgentType::Primary, + hp1, + delegate_params("d-0", "worker-1", 60), + 1, + ) { + DelegateOutcome::Accepted(_) => {} + DelegateOutcome::Rejected(e) => panic!("baseline rejected: {}", e.message), + } + worker_rx.try_recv().unwrap(); + // Close the worker's queue so p2's forward frame is refused and + // its delegate call takes the rollback path. + worker_rx.close(); + + let gate = std::sync::Barrier::new(2); + std::thread::scope(|s| { + s.spawn(|| { + gate.wait(); + // p2 dies while its delegate call is in flight. + let mut next = || 99; + router.fail_instance(®istry, hp2, &mut next); + }); + gate.wait(); + let _ = router.delegate( + &cfg, + ®istry, + "prod", + "koudu-2", + &AgentType::Primary, + hp2, + delegate_params("d-1", "worker-1", 60), + 2, + ); + }); + + assert_eq!( + registry.get(hw).unwrap().active_sessions, + 1, + "exactly the baseline delegation must stay reserved" + ); + assert_eq!(router.inflight_count(), 1); + } + } + + #[test] + fn stalled_initiator_result_is_never_silently_lost() { + // Terminal results honor the bounded-queue contract: if the + // initiator cannot drain its queue, the entry stays in flight and + // the caller is told to treat the initiator as disconnected. The + // delegation then resolves through fail_instance (cp/cancel to the + // serving side, capacity released once) — never by silently + // dropping a computed result while acking the serving side. + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + + // Fill the initiator's bounded queue so the result frame is refused. + let initiator_tx = w.registry.get(w.h_primary).unwrap().tx; + while initiator_tx.try_send("filler".into()).is_ok() {} + + assert_eq!( + w.router + .complete(&w.registry, w.h_worker, result_of("d-1", "late"), 1024, 2), + CompleteOutcome::InitiatorStalled { + initiator_handle: w.h_primary + } + ); + assert_eq!(w.router.inflight_count(), 1, "entry must stay in flight"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "capacity must not be released while the delegation is unresolved" + ); + + // The stalled initiator is then failed (disconnect path): capacity + // is released exactly once and the serving side is told to cancel. + let mut next = || 3; + let frames = w.router.fail_instance(&w.registry, w.h_primary, &mut next); + assert_eq!(frames.len(), 1); + assert!(frames[0].1.contains("cp/cancel")); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + } + #[test] fn duplicate_delegation_id_rejected() { let w = world(); @@ -929,10 +1124,10 @@ type = "worker" error: None, }; // h_primary is a valid handle but NOT the serving instance. - assert!(w - .router - .complete(&w.registry, w.h_primary, result, 1024, 2) - .is_none()); + assert_eq!( + w.router.complete(&w.registry, w.h_primary, result, 1024, 2), + CompleteOutcome::Dropped + ); assert_eq!(w.router.inflight_count(), 1); } @@ -945,10 +1140,10 @@ type = "worker" result: None, error: None, }; - assert!(w - .router - .complete(&w.registry, w.h_worker, result, 1024, 2) - .is_none()); + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, 1024, 2), + CompleteOutcome::Dropped + ); } #[test] @@ -966,10 +1161,11 @@ type = "worker" error: None, }; let cap = 96usize; - let (_, frame) = w - .router - .complete(&w.registry, w.h_worker, result, cap, 2) - .unwrap(); + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, cap, 2), + CompleteOutcome::Delivered + ); + let frame = w.primary_rx.try_recv().unwrap(); let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); let out = v["params"]["result"].as_str().unwrap(); assert!(out.contains("truncated by control plane")); @@ -992,10 +1188,11 @@ type = "worker" result: Some("y".repeat(100)), error: None, }; - let (_, frame2) = w - .router - .complete(&w.registry, w.h_worker, result2, 8, 3) - .unwrap(); + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result2, 8, 3), + CompleteOutcome::Delivered + ); + let frame2 = w.primary_rx.try_recv().unwrap(); let v2: serde_json::Value = serde_json::from_str(&frame2).unwrap(); assert!(v2["params"]["result"].as_str().unwrap().len() <= 8); } @@ -1241,16 +1438,16 @@ allow_worker_initiation = true if spoof_first { // h_primary is registered but is NOT the serving instance. - assert!(w - .router - .complete( + assert_eq!( + w.router.complete( &w.registry, w.h_primary, result_of("d-1", "spoofed"), 1024, 2 - ) - .is_none()); + ), + CompleteOutcome::Dropped + ); assert_eq!( w.router.inflight_count(), 1, @@ -1258,33 +1455,34 @@ allow_worker_initiation = true ); } - let (init, frame) = w - .router - .complete( + assert_eq!( + w.router.complete( &w.registry, w.h_worker, result_of("d-1", "genuine"), 1024, 3, - ) - .expect("genuine result must be delivered, never dropped"); - assert_eq!(init.handle, w.h_primary); + ), + CompleteOutcome::Delivered, + "genuine result must be delivered, never dropped" + ); + let frame = w.primary_rx.try_recv().unwrap(); assert!(frame.contains("genuine")); assert_eq!(w.router.inflight_count(), 0); assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); if !spoof_first { // A late non-owner frame after completion is a plain no-op. - assert!(w - .router - .complete( + assert_eq!( + w.router.complete( &w.registry, w.h_primary, result_of("d-1", "spoofed"), 1024, 4 - ) - .is_none()); + ), + CompleteOutcome::Dropped + ); assert_eq!(w.router.inflight_count(), 0); } } @@ -1308,27 +1506,22 @@ allow_worker_initiation = true let spoof = s.spawn(|| { gate.wait(); // Registered, but not the serving instance. - w.router - .complete( - &w.registry, - w.h_primary, - result_of("d-1", "spoofed"), - 1024, - 2, - ) - .is_some() - }); - gate.wait(); - let genuine = w - .router - .complete( + w.router.complete( &w.registry, - w.h_worker, - result_of("d-1", "genuine"), + w.h_primary, + result_of("d-1", "spoofed"), 1024, - 3, - ) - .is_some(); + 2, + ) == CompleteOutcome::Delivered + }); + gate.wait(); + let genuine = w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", "genuine"), + 1024, + 3, + ) == CompleteOutcome::Delivered; (spoof.join().unwrap(), genuine) }); assert!(!spoofed, "a non-owner must never complete a delegation"); @@ -1456,9 +1649,9 @@ allow_worker_initiation = true let registry = Registry::new(); let router = Router::new(); let cfg = cfg(); - let (p_prod, _prod_init_rx) = instance("prod", "koudu", AgentType::Primary, 4); + let (p_prod, mut prod_init_rx) = instance("prod", "koudu", AgentType::Primary, 4); let (w_prod, mut prod_rx) = instance("prod", "worker-1", AgentType::Worker, 2); - let (p_dev, _dev_init_rx) = instance("dev", "koudu", AgentType::Primary, 4); + let (p_dev, mut dev_init_rx) = instance("dev", "koudu", AgentType::Primary, 4); let (w_dev, mut dev_rx) = instance("dev", "worker-1", AgentType::Worker, 2); let hp_prod = registry.register(p_prod); let hw_prod = registry.register(w_prod); @@ -1514,20 +1707,26 @@ allow_worker_initiation = true assert_eq!(router.inflight_count(), 2); // Results route to the initiator of the SAME namespace only. - let (init, frame) = router - .complete(®istry, hw_dev, result_of("d-1", "dev-done"), 1024, 12) - .unwrap(); - assert_eq!(init.handle, hp_dev); + assert_eq!( + router.complete(®istry, hw_dev, result_of("d-1", "dev-done"), 1024, 12), + CompleteOutcome::Delivered + ); + let frame = dev_init_rx.try_recv().unwrap(); assert!(frame.contains("dev-done")); + assert!( + prod_init_rx.try_recv().is_err(), + "prod's initiator must not receive dev's result" + ); assert!( router.chain_of("prod", "d-1").is_some(), "prod's delegation must be untouched" ); - let (init, frame) = router - .complete(®istry, hw_prod, result_of("d-1", "prod-done"), 1024, 13) - .unwrap(); - assert_eq!(init.handle, hp_prod); + assert_eq!( + router.complete(®istry, hw_prod, result_of("d-1", "prod-done"), 1024, 13), + CompleteOutcome::Delivered + ); + let frame = prod_init_rx.try_recv().unwrap(); assert!(frame.contains("prod-done")); assert_eq!(router.inflight_count(), 0); } diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index 08db9c41e..38b65ec1a 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -43,7 +43,7 @@ use crate::proto::{ PROTOCOL_VERSION, }; use crate::registry::{shutdown_signal, Instance, Registry, OUTBOUND_QUEUE}; -use crate::router::{DelegateOutcome, Router}; +use crate::router::{CompleteOutcome, DelegateOutcome, Router}; pub struct AppState { pub cfg: CpConfig, @@ -129,6 +129,10 @@ pub const REASON_REGISTER_TIMEOUT: &str = "registration timeout"; /// Reason string on the close frame sent when the CP drops a registration /// because its lease elapsed. pub const REASON_LEASE_EXPIRED: &str = "lease expired"; +/// Reason string on the close frame sent when a peer's bounded outbound +/// queue refused a terminal frame: per the queue contract the peer is +/// treated as disconnected, not buffered. +pub const REASON_BACKPRESSURE: &str = "outbound queue overflow"; /// A CP-initiated close that states why. Code 1008 (policy violation) plus a /// short reason, so a client can distinguish "the CP closed me on purpose" @@ -305,17 +309,18 @@ async fn handle_connection( } // --- Main loop: interleave inbound frames, outbound channel, shutdown --- - let mut cp_closed = false; + let mut cp_close_reason: Option<&'static str> = None; loop { tokio::select! { - // The CP dropped this registration (lease expiry): the socket - // must go too. Keeping it open would leave a + // The CP dropped this registration (lease expiry) or must + // terminate the connection (terminal-frame backpressure): the + // socket must go too. Keeping it open would leave a // connection whose every frame hits an absent registry entry and // which can never re-register, since registration is // first-frame-only. Closing lets the client reconnect, // re-authenticate, and register again. _ = shutdown_rx.changed() => { - cp_closed = true; + cp_close_reason = *shutdown_rx.borrow_and_update(); break; } outbound = rx.recv() => { @@ -353,13 +358,14 @@ async fn handle_connection( } } - if cp_closed { + if let Some(reason) = cp_close_reason { info!( agent = %format!("{}/{}", identity.namespace, identity.name), handle, - "closing connection at the CP's request (registration dropped)" + reason, + "closing connection at the CP's request" ); - let _ = sink.send(policy_close(REASON_LEASE_EXPIRED)).await; + let _ = sink.send(policy_close(reason)).await; } teardown(&state, handle, &identity); @@ -367,6 +373,13 @@ async fn handle_connection( /// Deregister this connection's own registration (by handle — cannot touch /// another connection's entry) and fail its in-flight delegations. +/// +/// Deliberately idempotent with the sweeper: when `sweep_leases` already ran +/// `deregister` + `fail_instance` for this handle, both calls here find +/// nothing (the registry entry and the in-flight entries are gone) and are +/// no-ops. That idempotency is a contract — `fail_instance` releases +/// capacity only for entries it actually removes, so a second pass can never +/// double-release (see the capacity note in `Router::delegate`'s rollback). fn teardown(state: &Arc, handle: u64, identity: &AgentIdentity) { state.registry.deregister(handle); let mut next = || state.next_rpc_id(); @@ -570,17 +583,40 @@ fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option { let p = params_or_err!(DelegateResultParams); - if let Some((initiator, frame)) = state.router.complete( + match state.router.complete( &state.registry, handle, p, state.cfg.max_result_bytes, state.next_rpc_id(), ) { - let _ = initiator.tx.try_send(frame); + CompleteOutcome::InitiatorStalled { initiator_handle } => { + // The initiator cannot drain its bounded queue: per the + // queue contract it is treated as disconnected, never + // silently skipped. Its teardown fails the delegation + // over the fail_instance path (capacity released once, + // cp/cancel to this serving runtime). Do NOT ack the + // result as delivered — the serving side must know its + // result did not reach the initiator. + state + .registry + .signal_shutdown(initiator_handle, REASON_BACKPRESSURE); + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new( + codes::TARGET_DISCONNECTED, + "initiator cannot receive the result; the delegation will be cancelled", + ), + ); + Some(serde_json::to_string(&resp).expect("serializable")) + } + // Delivered, or dropped as unknown/foreign (each case is + // logged; late results after a CP restart are expected). + _ => { + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } } - let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); - Some(serde_json::to_string(&resp).expect("serializable")) } methods::CANCEL => { let p = params_or_err!(CancelParams); @@ -626,7 +662,7 @@ pub fn sweep_leases(state: &Arc, lease: Duration) { "lease expired — deregistering and closing connection" ); // Signal first: `deregister` drops the registry's side of the signal. - state.registry.signal_shutdown(handle); + state.registry.signal_shutdown(handle, REASON_LEASE_EXPIRED); state.registry.deregister(handle); let mut next = || state.next_rpc_id(); for (inst, frame) in state @@ -871,14 +907,15 @@ mod tests { // A live lease is left alone. sweep_leases(&state, Duration::from_secs(60)); assert!(state.registry.get(handle).is_some()); - assert!(!*observer.borrow()); + assert!(observer.borrow().is_none()); sweep_leases(&state, Duration::ZERO); assert!(state.registry.get(handle).is_none()); observer.changed().await.unwrap(); - assert!( + assert_eq!( *observer.borrow(), - "the owning connection must be told to close" + Some(REASON_LEASE_EXPIRED), + "the owning connection must be told to close, and why" ); } @@ -945,4 +982,99 @@ mod tests { assert_eq!(v2["error"]["code"], codes::NOT_REGISTERED); assert_eq!(state.router.inflight_count(), 0); } + + #[tokio::test] + async fn stalled_initiator_is_disconnected_and_serving_side_not_falsely_acked() { + // The bounded-queue contract for the one frame that matters most: + // when the initiator's queue refuses the terminal result, the + // serving side must NOT receive `ok: true`, and the initiator must + // be closed (treated as disconnected) rather than silently skipped. + let state = state_with(""); + + // Initiator with a full single-slot queue. + let signal_i = crate::registry::shutdown_signal(); + let mut observer = signal_i.subscribe(); + let (tx_i, _rx_i) = mpsc::channel::(1); + let h_i = state.registry.register_conn( + Instance { + handle: 0, + namespace: "prod".into(), + name: "koudu".into(), + agent_type: AgentType::Primary, + instance_id: "i-1".into(), + labels: Default::default(), + max_delegated_sessions: 4, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx: tx_i.clone(), + }, + Arc::clone(&signal_i), + ); + let (tx_w, mut rx_w) = mpsc::channel::(OUTBOUND_QUEUE); + let h_w = state.registry.register_conn( + Instance { + handle: 0, + namespace: "prod".into(), + name: "worker-1".into(), + agent_type: AgentType::Worker, + instance_id: "i-2".into(), + labels: Default::default(), + max_delegated_sessions: 1, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx: tx_w, + }, + crate::registry::shutdown_signal(), + ); + + // Route one delegation through the wire-facing handler. + let deadline = (chrono::Utc::now() + chrono::Duration::seconds(60)).to_rfc3339(); + let del = serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/delegate", + "params": { + "delegation_id": "d-1", + "target": {"name": "worker-1"}, + "prompt": "do it", + "deadline": deadline + } + }) + .to_string(); + let ack: serde_json::Value = + serde_json::from_str(&handle_frame(&state, h_i, &del).expect("answered")).unwrap(); + assert!(ack.get("error").is_none(), "delegation must be accepted"); + rx_w.try_recv().expect("worker received the forward frame"); + + // Fill the initiator's queue, then complete. + tx_i.try_send("filler".into()).unwrap(); + let res = serde_json::json!({ + "jsonrpc": "2.0", "id": 2, "method": "cp/delegate_result", + "params": {"delegation_id": "d-1", "status": "completed", "result": "done"} + }) + .to_string(); + let reply: serde_json::Value = + serde_json::from_str(&handle_frame(&state, h_w, &res).expect("answered")).unwrap(); + assert_eq!( + reply["error"]["code"], + codes::TARGET_DISCONNECTED, + "the serving side must not be acked as delivered" + ); + assert_eq!(reply["id"], 2, "the error must correlate with the request"); + + // The initiator is told to close, with the backpressure reason. + observer.changed().await.unwrap(); + assert_eq!(*observer.borrow(), Some(REASON_BACKPRESSURE)); + + // The delegation is still in flight: teardown of the stalled + // initiator resolves it through fail_instance (capacity released + // once, cp/cancel to the serving runtime). + assert_eq!(state.router.inflight_count(), 1); + let mut next = || 9; + let frames = state.router.fail_instance(&state.registry, h_i, &mut next); + assert_eq!(frames.len(), 1); + assert!(frames[0].1.contains("cp/cancel")); + assert_eq!(state.registry.get(h_w).unwrap().active_sessions, 0); + assert_eq!(state.router.inflight_count(), 0); + } } diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index b347288e2..36fd8d501 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -279,6 +279,27 @@ recovery semantics: admission (duplicate check → target selection → capacity reservation → in-flight insert) is one atomic sequence, and the in-flight entry exists before the forward frame is sent. +- **Terminal results are never silently dropped.** `cp/delegate_result` + delivery commits only after the initiator's queue accepts the frame: the + CP validates ownership, sends, and only then removes the in-flight entry + and releases the serving instance's capacity. If the initiator's bounded + queue refuses the frame, the serving runtime receives an error (not + `ok: true`), the initiator is closed (WS 1008, reason + `outbound queue overflow` — the "cannot drain → disconnected" rule applied + to the frame where it matters most), and the delegation resolves through + the disconnect path: `cp/cancel` to the serving runtime and exactly one + capacity release. Two concurrent duplicate results may both be delivered + (correlation is by `delegation_id`; delivery is idempotent for the + initiator), but capacity is released exactly once. Best-effort frames + (`cp/cancel`, sweep-synthesized `timeout`) remain fire-and-forget: the + propagated deadline is their backstop. +- **Capacity release follows entry removal.** Whichever path removes an + in-flight entry (result commit, cancel, deadline sweep, instance failure, + or a failed forward's rollback) releases its capacity reservation — and + only that path does, exactly once. A rollback that finds its entry already + removed by a concurrent sweep or disconnect must not decrement again: + session counts are saturating, so a double release is silent and would + let `saturated()` admit work to a full instance. - **Saturation = fast-fail.** When all matching targets are at capacity the CP replies `SATURATED` immediately. The CP never queues — v1 has no durable state, and a hidden in-memory queue would contradict that. @@ -413,11 +434,24 @@ facade without changing anything shipped in v1. ## 7. Security +Two distinct auth boundaries exist, and they must not be conflated: + +1. **Runtime ↔ CP (shipped in PR 1/4):** the OAB runtime authenticates to the + CP with `Authorization: Bearer ` on the WebSocket upgrade, over TCP. + The CP binds loopback by default; any non-loopback bind requires the + explicit `allow_insecure_bind` override and a TLS-terminating proxy (or a + private network) in front — bearer keys must never cross untrusted + cleartext TCP. See the "v1 contract amendments" in §4 for the enforced + registration semantics. +2. **Agent subprocess ↔ local facade (PR 3/4, not yet shipped):** the UDS + path is the only thing the child needs; filesystem permissions on the + socket are the local auth boundary. The *local facade* is never exposed + on TCP — this claim is about the UDS facade, not about the CP itself, + which is a TCP service by design. + - **No CP credentials in the agent process.** `OPENAB_CP_KEY` lives in the OAB runtime env; agent subprocesses keep the existing `env_clear` - whitelist. The UDS path is the only thing the child needs; filesystem - permissions on the socket are the local auth boundary. The local API is - never exposed on TCP. + whitelist. - **Per-agent auth keys** to the CP (not one shared fleet key), so a single compromised runtime is individually revocable. - **Per-peer identity.** Delegated prompts arrive attributed to the sending diff --git a/docs/control-plane.md b/docs/control-plane.md new file mode 100644 index 000000000..fae85aed6 --- /dev/null +++ b/docs/control-plane.md @@ -0,0 +1,54 @@ +# Agent Control Plane (`openab-cp`) + +Standalone control-plane service for direct agent-to-agent delegation over +WebSocket JSON-RPC, so agents delegate work to each other without +round-tripping through a chat platform. Design and wire contract: +[ADR: Agent Control Plane](adr/agent-control-plane.md). + +> **Status: PR 1/4 of the control-plane stack.** This slice ships the CP +> server binary (registry, policy, router, wire protocol). The OAB-runtime +> client (`[control_plane]` config + registration), the MCP facade/CLI, and +> streaming land in the follow-up slices — until then nothing connects to +> this server in a stock deployment, and there is no packaged container +> image yet. + +## Run + +```bash +cargo run -p openab-cp -- --config cp.toml +``` + +Start from the annotated example config: + +```bash +cp crates/openab-cp/cp.toml.example cp.toml +``` + +Every field is documented in the example file, including the security +rationale. The essentials: + +- `listen` — defaults to loopback (`127.0.0.1:9800`). A non-loopback bind is + refused unless `allow_insecure_bind = true` is set explicitly, and then a + TLS-terminating proxy (`wss://`) or a private network in front is + required: runtimes authenticate with bearer keys that must never cross + untrusted cleartext TCP. +- `[[identity]]` — one entry per agent identity: the auth key (supports + `${ENV_VAR}` expansion) and its immutable `namespace`/`name`/`type` + claims. A connecting runtime must register as exactly the identity its + key is bound to. +- Heartbeats, lease expiry, registration deadline, per-identity connection + quotas, and frame/prompt/result size caps are all configurable with safe + defaults. + +## Health + +`GET /health` answers `ok` (liveness only; deeper checks are tracked in +issue #1474). + +## Client behavior to expect + +- CP-initiated closes use WS code 1008 with a reason: `registration + timeout`, `lease expired`, or `outbound queue overflow`. On any of these, + reconnect, re-authenticate, and re-register. +- After a lease expires or the CP restarts, in-flight delegations are gone: + initiators reconcile against their own deadlines and re-delegate. From e313045a3704c32df57032df33e3c2b32f283104 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Thu, 13 Aug 2026 16:33:32 +0000 Subject: [PATCH 06/11] fix(cp): read all Claim fields in the completion commit phase The two-phase complete() rewrite removed the only reader of Claim::WrongOwner's owner_handle, tripping dead_code under -D warnings. The commit phase now logs every non-Owned claim outcome, which also makes the concurrent-removal and re-owned cases observable in CP logs. --- crates/openab-cp/src/router.rs | 43 +++++++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index bf1f8b22c..1472bc257 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -522,20 +522,45 @@ impl Router { // Phase 2 — commit. If a concurrent path (duplicate result, cancel, // sweep, fail_instance) removed the entry between peek and now, that // path also released the capacity — do not decrement twice. - if let Claim::Owned(e) = self.claim( + match self.claim( registry, serving_handle, ¶ms.delegation_id, Owner::Server, ) { - registry.adjust_sessions(e.to_handle, -1); - info!( - delegation = %params.delegation_id, - status = ?params.status, - from = %e.to_logical, - to = %e.from_logical, - "delegation completed" - ); + Claim::Owned(e) => { + registry.adjust_sessions(e.to_handle, -1); + info!( + delegation = %params.delegation_id, + status = ?params.status, + from = %e.to_logical, + to = %e.from_logical, + "delegation completed" + ); + } + Claim::WrongOwner { + namespace, + owner_handle, + } => { + // Only possible if the id was removed and re-admitted between + // peek and commit. The new delegation is not ours to touch. + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + owner = owner_handle, + "delegation re-owned between delivery and commit — entry left untouched" + ); + } + Claim::NotFound { namespace } => { + // Concurrent removal (duplicate result, cancel, sweep, or + // fail_instance): whoever removed it released the capacity. + info!( + delegation = %params.delegation_id, + namespace = %namespace, + "entry removed concurrently after delivery — capacity already released" + ); + } + Claim::Unregistered => {} } CompleteOutcome::Delivered } From bde414a9e3c34436c7639cf043177001e98817ee Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Thu, 13 Aug 2026 20:12:33 -0400 Subject: [PATCH 07/11] =?UTF-8?q?fix(cp):=20round-6=20review=20=E2=80=94?= =?UTF-8?q?=20generation-stamped=20commits,=20bounded=20writes,=20terminal?= =?UTF-8?q?-frame=20contract?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 (same-owner ABA in two-phase completion). Every admission is now stamped with a CP-generated, never-reused generation (AtomicU64 on Router, stamped into InFlight under the admission lock). The commit phase is extracted into `Router::commit_completion`, which under ONE inflight lock acquisition removes the entry only when key AND serving handle AND generation all match, and otherwise leaves the live entry — and its capacity reservation — strictly untouched. Previously the commit claimed by (namespace, delegation_id) + serving handle only: peek d-1, concurrent cancel/sweep removes it and frees the slot, the initiator legitimately reuses d-1 (cancel-then-retry) and re-admission routes to the SAME worker with one replica, and the stale commit then removed the NEW entry and released its capacity — the live delegation became invisible to its own genuine result and to the sweep. Phase 1 is likewise extracted (`peek_for_completion`) so both steps are directly callable and deterministically testable. `delegate`'s rollback now matches on the generation too, so it removes the exact entry it inserted without relying on the admission lock for that property. `Owner` is gone: `claim` is initiator-only (cancel's single-phase path) and logs the owning handle. CompleteOutcome::Delivered becomes Delivered { committed: bool }, modelling wire delivery separately from state commit; server.rs matches it explicitly and still acks the serving side (the result did reach the initiator) while logging the non-committing case. F2 (blocked writer pins quota slots). All outbound writes go through `send_bounded`, which races the shutdown watch against a `write_timeout_secs`-bounded send; timeout or transport error is a disconnect (break → teardown → ConnPermit drop → downstream cancellation). Covers the registration-phase writes (the shutdown signal is now created before the registration read, so those writes are shutdown-aware too), the ack, both select arms, and the post-loop close. A close signal landing on the ack write still delivers its 1008 reason. New config field `write_timeout_secs` (serde default 30, validated > 0, documented in cp.toml.example): a `select!` arm body is not cancelled by its siblings, so an unbounded `sink.send().await` parked the whole connection task — a peer with a closed receive window or a half-open connection pinned its identity's quota indefinitely and defeated lease-expiry recovery, which works by signalling that very task. F3 (contradictory terminal frames). Documented rather than suppressed: "first terminal frame wins" is now part of the ADR's v1 contract amendments (initiators MUST treat the first terminal frame for a delegation_id as authoritative and ignore later ones) and restated in docs/control-plane.md's client-behavior section. The router module doc and `complete`'s doc no longer describe the peek-send window as benign duplicate delivery; they reference the contract and state that CP state is exact under the generation rule. No suppression machinery added this round. F4: docs/control-plane.md called the identity table `[[identity]]`; the key is `[[agents]]`, cross-checked against cp.toml.example and config.rs. F5: README.zh-TW.md gains the control-plane feature entry mirroring README.md's placement (after 多 agent 協作) and content. F6: deterministic interleaving tests driving the extracted commit directly, no barrier timing: - stale_commit_never_claims_a_reused_delegation_id — the F1 regression, run for both removal kinds (cancel, sweep): peek A → remove A → re-admit B with the same id and same worker → commit A ⇒ Superseded, B still in flight and visible, capacity intact, B then completes normally. - commit_after_cancel_releases_capacity_exactly_once, commit_after_deadline_sweep_releases_capacity_exactly_once — commit-vs-cancel and commit-vs-sweep with no re-admission ⇒ Vanished, no double release (checked against a second target and by re-saturating the worker). - delivered_reports_whether_it_committed — commit is claimed exactly once; a repeated commit is a no-op. - tests/ws_lifecycle.rs: a_peer_that_stops_reading_is_disconnected_and_frees_ its_quota — real socket, real client; the peer registers then stops reading while the CP forwards ~48 MiB of results, and the test asserts the task terminates, conn_count returns to 0, a new connection for the same identity registers, and the remaining delegation is cancelled downstream. The existing concurrency race tests are kept unchanged. Assertion extensions (extended, never weakened): the 9 existing `CompleteOutcome::Delivered` assertions in router.rs now assert `Delivered { committed: true }` — strictly stronger, they additionally pin that the frame performed the state commit. config.rs's admission_bounds_default_and_are_validated extends to cover write_timeout_secs' default, explicit value, and zero rejection. Verified on an isolated checkout: 70 unit + 6 integration tests pass (was 66 + 5), clippy --all-targets -D warnings clean, rustfmt clean. Both new regressions were negative-controlled: reverting the generation check makes stale_commit_never_claims_a_reused_delegation_id fail (Claimed vs Superseded), and restoring the unbounded outbound write makes the ws_lifecycle test fail with the connection still pinned after 30s. Audit round-2 fixes: - Generation minting uses fetch_update, not fetch_add: an exhausted counter is never written, so it can neither wrap the atomic nor re-issue a generation — admission fails closed permanently at the ceiling. Boundary test pins the refusal, the parked counter, and that no capacity is reserved by a refused admission. (Unreachable in a realistic process lifetime; insurance against a refactor downgrading it to wrapping math.) - The stalled-peer regression asserts permit release within the CONFIGURED write timeout plus scheduling allowance, not an unrelated 30s poll — a 20s pin would previously have passed. --- README.zh-TW.md | 1 + crates/openab-cp/cp.toml.example | 10 + crates/openab-cp/src/config.rs | 34 +- crates/openab-cp/src/router.rs | 685 +++++++++++++++++++++---- crates/openab-cp/src/server.rs | 193 +++++-- crates/openab-cp/tests/ws_lifecycle.rs | 200 +++++++- docs/adr/agent-control-plane.md | 44 +- docs/control-plane.md | 13 +- 8 files changed, 1032 insertions(+), 148 deletions(-) diff --git a/README.zh-TW.md b/README.zh-TW.md index 619c2d1e4..c429cdb42 100644 --- a/README.zh-TW.md +++ b/README.zh-TW.md @@ -58,6 +58,7 @@ platforms 使用 `webhook/API`，Feishu/Lark 則使用 `WS/webhook`。 - **@mention 觸發** — 在允許的頻道中 mention bot，即可開始對話 - **以討論串進行多輪對話** — 自動建立討論串；後續訊息不需再次 @mention - **多 agent 協作** — 支援 bot-to-bot 訊息，實現協調式工作流程（[docs/multi-agent.md](docs/multi-agent.md)） +- **Agent control plane（預覽）** — 獨立的 `openab-cp` 服務，讓 agent 之間可透過 WebSocket 直接委派任務，具備綁定身分的註冊機制與由 CP 主導的政策控管；runtime client 與 facade 將於後續版本推出（[docs/control-plane.md](docs/control-plane.md)） - **由 agent 控制回覆對象** — agent 可透過 `[[reply_to:id]]` 指令選擇要回覆的訊息，讓多 bot 頻道中的對話脈絡更清楚（[docs/output-directives.md](docs/output-directives.md)） - **編輯式串流輸出** — token 產生時每 1.5 秒即時更新 Discord 訊息 - **Emoji 狀態反應** — 👀→🤔→🔥/👨‍💻/⚡→👍+隨機情緒表情 diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index 9b5df7943..92dec8df5 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -40,6 +40,16 @@ register_timeout_secs = 10 # on concurrent replicas of one logical agent. max_connections_per_identity = 8 +# How long a single outbound WebSocket write may block before the peer is +# treated as disconnected. A bounded queue does not bound the writer: a peer +# that stops reading (closed TCP receive window, half-open connection) would +# otherwise park the CP inside one send, pinning that identity's connection +# slots and its in-flight delegations for as long as the socket survives — and +# lease expiry cannot recover it, because the connection task is what acts on +# the CP's close signal. Keep it generous enough for a slow-but-alive peer +# draining a large result, short enough that a dead one frees its quota. +write_timeout_secs = 30 + [[agents]] key = "${CP_KEY_KOUDU}" # per-agent secret, never shared namespace = "prod" diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index 75e27ffc6..02b492273 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -74,6 +74,23 @@ pub struct CpConfig { #[serde(default = "default_max_connections_per_identity")] pub max_connections_per_identity: u32, + /// How long one outbound WebSocket write may block before the peer is + /// treated as disconnected. + /// + /// A bounded outbound queue does not bound the *writer*: a peer whose TCP + /// receive window is closed (or whose connection is half-open) leaves the + /// CP parked inside a single `send`, holding that identity's connection + /// slot and its in-flight delegations for as long as the socket survives — + /// which defeats lease-expiry recovery, because the connection task is + /// what reacts to the CP's own close signal. Every write is therefore + /// bounded by this timeout and raced against the shutdown signal; a + /// timeout is a disconnect. + /// + /// Generous enough for a slow-but-alive peer draining a large result, + /// short enough that a dead one frees its quota promptly. + #[serde(default = "default_write_timeout_secs")] + pub write_timeout_secs: u64, + /// Identity table: auth key → immutable claims. /// Keyed by the key id (`kid`), with the secret alongside, so logs can /// reference identities without printing secrets. @@ -112,6 +129,9 @@ fn default_register_timeout_secs() -> u64 { fn default_max_connections_per_identity() -> u32 { 8 } +fn default_write_timeout_secs() -> u64 { + 30 +} /// Immutable identity claims bound to one auth key. #[derive(Debug, Clone, Deserialize)] @@ -198,6 +218,11 @@ impl CpConfig { if self.max_connections_per_identity == 0 { bail!("max_connections_per_identity must be at least 1"); } + // Zero would mean "every write times out instantly", i.e. no peer + // could ever be written to. + if self.write_timeout_secs == 0 { + bail!("write_timeout_secs must be greater than 0"); + } // Bearer keys over cleartext TCP must never reach an untrusted // network: non-loopback binds require the explicit override. if !self.allow_insecure_bind && !is_loopback(&self.listen) { @@ -367,16 +392,21 @@ lease_expiry_secs = 30 cfg.validate().unwrap(); assert_eq!(cfg.register_timeout_secs, 10); assert_eq!(cfg.max_connections_per_identity, 8); + assert_eq!(cfg.write_timeout_secs, 30); - let explicit: CpConfig = - toml::from_str("register_timeout_secs = 3\nmax_connections_per_identity = 2").unwrap(); + let explicit: CpConfig = toml::from_str( + "register_timeout_secs = 3\nmax_connections_per_identity = 2\nwrite_timeout_secs = 5", + ) + .unwrap(); explicit.validate().unwrap(); assert_eq!(explicit.register_timeout_secs, 3); assert_eq!(explicit.max_connections_per_identity, 2); + assert_eq!(explicit.write_timeout_secs, 5); for bad in [ "register_timeout_secs = 0", "max_connections_per_identity = 0", + "write_timeout_secs = 0", ] { let cfg: CpConfig = toml::from_str(bad).unwrap(); assert!(cfg.validate().is_err(), "{bad} must be rejected"); diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index 1472bc257..4e45e0289 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -17,6 +17,21 @@ //! - **Saturation** — routing never queues; `SATURATED` is returned //! immediately (fast-fail, no hidden buffer). //! +//! # Terminal frames +//! +//! Ending a delegation is a two-sided event: a frame goes out on the wire and +//! CP state is committed. The commit is exact — it claims the one admission it +//! delivered a result for (key + serving handle + [`InFlight::generation`]) or +//! nothing at all — so CP state stays consistent under any interleaving. +//! +//! The wire is a different matter: a `completed` result racing the deadline +//! sweep's synthesized `timeout` can put TWO terminal frames on the wire for +//! one `delegation_id`. v1 resolves that by contract instead of CP-side +//! suppression (which would need per-id terminal state the CP deliberately +//! does not keep): **the first terminal frame for a `delegation_id` wins**, and +//! initiators MUST ignore later ones. See the v1 contract amendments in +//! `docs/adr/agent-control-plane.md`. +//! //! # Lock hierarchy //! //! The router holds two locks and acquires them in ONE order only: @@ -41,6 +56,7 @@ //! does not participate in this hierarchy. use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU64, Ordering}; use chrono::{DateTime, Utc}; use parking_lot::Mutex; @@ -83,6 +99,18 @@ pub struct InFlight { /// CP-constructed chain for THIS delegation (root first, ends with the /// initiator). Children extend it. pub chain: Vec, + /// CP-generated, never-reused admission stamp. `(namespace, + /// delegation_id)` is NOT a stable identity over time: the id is + /// client-supplied, and cancel-then-retry — a natural client pattern — + /// legitimately re-admits the same id, which with a single replica routes + /// to the same serving instance again. Key plus serving handle therefore + /// cannot distinguish the entry a two-phase completion peeked from a + /// later, unrelated admission wearing the same clothes (an ABA race). + /// + /// The generation makes that distinction total: it is minted once per + /// admission from a monotonic counter and never reused, so the commit + /// step claims the entry it actually delivered a result for, or nothing. + pub generation: u64, } /// In-flight table key: `(namespace, delegation_id)`. @@ -116,6 +144,11 @@ pub struct Router { /// Delegation rates are LLM-scale; a coarse admission lock is simple and /// more than sufficient. admission: Mutex<()>, + /// Source of [`InFlight::generation`] stamps. Monotonic and never reset + /// (the table dies with the process, so a restart cannot collide with + /// anything still in flight): every admission gets a value no earlier or + /// later admission has worn. + next_generation: AtomicU64, } pub enum DelegateOutcome { @@ -125,19 +158,8 @@ pub enum DelegateOutcome { Rejected(ErrorObject), } -/// Which side of a delegation a caller must be to act on it. -#[derive(Clone, Copy)] -enum Owner { - /// The instance the delegation was routed to — the only one that may - /// complete it. - Server, - /// The instance that initiated the delegation — the only one that may - /// cancel it. - Initiator, -} - -/// Result of looking up an in-flight delegation on a caller's behalf and -/// asserting the caller owns it. +/// Result of looking up an in-flight delegation on behalf of its claimed +/// initiator and removing it if the claim holds. /// /// The whole check happens under ONE acquisition of the in-flight lock, which /// is the property that matters: an earlier version removed the entry, @@ -145,10 +167,16 @@ enum Owner { /// landing in that window saw an empty table and was dropped as "unknown id", /// leaving the delegation to stall until its deadline. Here the entry is /// either removed because the caller owns it, or never touched at all. +/// +/// Single-phase by construction — the caller acts on the returned entry +/// without going back to the table — so no window exists in which the id +/// could be re-admitted under the caller's feet. Contrast the two-phase +/// completion path, which needs [`InFlight::generation`] for exactly that +/// reason. `cp/cancel` is this helper's only caller. enum Claim { - /// The caller owns it; the entry has already been removed from the table. + /// The caller initiated it; the entry has already been removed. Owned(InFlight), - /// The entry exists but belongs to another instance. Left in place. + /// The entry exists but was initiated by another instance. Left in place. WrongOwner { namespace: String, /// Handle of the instance that does own it (CP-side logs only — it is @@ -162,12 +190,64 @@ enum Claim { Unregistered, } +/// Phase-1 snapshot of a completion: the entry as it existed at peek time, +/// including its [`InFlight::generation`] stamp, or why no result can be +/// delivered for it. A peek never removes anything. +enum Peek { + /// The caller is the instance the delegation was routed to. + Serving(InFlight), + /// The entry exists but another instance serves it. Left in place — a + /// non-owner frame must never make the delegation momentarily invisible + /// to a genuine result or to the deadline sweep. + Foreign { + /// CP-side logs only; never disclosed to the caller. + owner_handle: u64, + }, + /// No entry for `(namespace, delegation_id)`. + Unknown, +} + +/// Phase-2 result of committing a delivered completion (see +/// [`Router::commit_completion`]). +#[derive(Debug, PartialEq, Eq)] +enum Commit { + /// The peeked admission was still the live one: entry removed and the + /// serving instance's capacity released, exactly once. + Claimed, + /// The entry is gone — a concurrent cancel, sweep, disconnect, or a + /// duplicate result's commit removed it. Whichever path removed it + /// released the capacity; this one must not decrement again. + Vanished, + /// A DIFFERENT admission holds `(namespace, delegation_id)` now: the id + /// was removed and re-admitted between peek and commit (cancel-then-retry + /// routed back to the same worker is the ordinary way this happens, and + /// with a single replica it is the *only* way it happens). The live entry + /// and its capacity are left untouched: claiming it would erase a + /// delegation that is genuinely running and silently drop its real result + /// later. + Superseded { + /// Generation now holding the id (CP-side logs only). + generation: u64, + }, +} + /// Outcome of a `cp/delegate_result` frame (see [`Router::complete`]). +/// +/// Wire delivery and state commit are distinct events: the initiator can have +/// received the result while the CP's own bookkeeping was concluded by +/// somebody else (a concurrent cancel, sweep, or disconnect). Collapsing the +/// two hid whether this frame is the one that ended the delegation. #[derive(Debug, PartialEq, Eq)] pub enum CompleteOutcome { - /// The result was accepted by the initiator's queue; the delegation is - /// finished and the serving instance's capacity was released. - Delivered, + /// The result reached the initiator's queue. + Delivered { + /// Whether THIS frame also committed the state transition — removed + /// the in-flight entry it peeked and released the serving instance's + /// capacity. `false` means a concurrent path had already ended the + /// delegation (or its id was re-admitted), so nothing was changed + /// here; the frame was still delivered. + committed: bool, + }, /// The frame was refused or the delegation is unknown (wrong owner, /// unknown id, unregistered caller, or the initiator is gone). Nothing /// changed; each case is logged. @@ -188,6 +268,7 @@ impl Router { Self { inflight: Mutex::new(BTreeMap::new()), admission: Mutex::new(()), + next_generation: AtomicU64::new(0), } } @@ -312,6 +393,40 @@ impl Router { // Reserve capacity and record the in-flight entry BEFORE sending, so // an immediately-arriving result finds it. Roll both back if the send // fails. + // + // Minted BEFORE the capacity reservation so exhaustion cannot leave a + // reserved slot behind. `fetch_add` on an exhausted counter would wrap + // and re-issue generation values, silently recreating the ABA this + // stamp exists to prevent — so exhaustion fails the admission instead. + // (Unreachable in practice: one admission per nanosecond exhausts a + // u64 after ~584 years; this is fail-closed insurance, not a path.) + // `fetch_update` rather than `fetch_add`: an exhausted counter must + // NOT be written (fetch_add would wrap the atomic itself and re-issue + // generations from 0). Refusal leaves the counter parked at the + // ceiling, so every later admission is refused too — fail closed, + // permanently, with no wrapping path. + let generation = + match self + .next_generation + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |g| { + if g >= u64::MAX - 1 { + None + } else { + Some(g + 1) + } + }) { + Ok(prev) => prev + 1, + Err(_) => { + tracing::error!("delegation generation space exhausted — refusing admission"); + // -32603 = JSON-RPC internal error; no protocol-specific code + // is warranted for a condition that cannot occur in a + // process's realistic lifetime. + return DelegateOutcome::Rejected(ErrorObject::new( + -32603, + "control plane generation space exhausted; restart the CP", + )); + } + }; registry.adjust_sessions(target.handle, 1); let entry = InFlight { namespace: from_namespace.to_string(), @@ -322,6 +437,10 @@ impl Router { to_handle: target.handle, deadline: params.deadline, chain, + // Stamped under the admission lock, so every admission — including + // a re-admission of an id that was just cancelled — is + // distinguishable from every other for the life of the process. + generation, }; self.inflight.lock().insert(key.clone(), entry.clone()); @@ -336,7 +455,12 @@ impl Router { // concurrent removal would double-release: the saturating math // hides the underflow and `saturated()` then admits new work to // an instance that is actually full. - if self.inflight.lock().remove(&key).is_some() { + // + // Matched on the generation, not just the key: the admission lock + // happens to rule out a re-admission of this id while we are + // here, but the rollback does not need that argument to be + // correct — it removes the exact entry it inserted or nothing. + if self.remove_generation(&key, entry.generation).is_some() { registry.adjust_sessions(target.handle, -1); } return DelegateOutcome::Rejected(ErrorObject::new( @@ -360,14 +484,25 @@ impl Router { }) } - /// Look up `delegation_id` in the caller's namespace, assert the caller is - /// the delegation's `owner` side, and remove the entry if so — all under - /// one acquisition of the in-flight lock (see [`Claim`]). + /// Remove `key` only if it still holds `generation` — the exact admission + /// the caller is acting for — under one lock acquisition. Any other entry + /// (or none) is left untouched. + fn remove_generation(&self, key: &DelegationKey, generation: u64) -> Option { + let mut g = self.inflight.lock(); + match g.get(key) { + Some(e) if e.generation == generation => g.remove(key), + _ => None, + } + } + + /// Look up `delegation_id` in the caller's namespace, assert the caller + /// initiated it, and remove the entry if so — all under one acquisition of + /// the in-flight lock (see [`Claim`]). /// /// The namespace is taken from the caller's authenticated registration, /// never from the frame, so a delegation id can only ever be resolved /// inside the namespace of the connection that named it. - fn claim(&self, registry: &Registry, handle: u64, delegation_id: &str, owner: Owner) -> Claim { + fn claim(&self, registry: &Registry, handle: u64, delegation_id: &str) -> Claim { // Registry lookup completes before the in-flight lock is taken; the // two locks are never held together (see the lock hierarchy above). let namespace = match registry.get(handle) { @@ -377,10 +512,7 @@ impl Router { let key = DelegationKey::new(&namespace, delegation_id); let mut g = self.inflight.lock(); let owner_handle = match g.get(&key) { - Some(e) => match owner { - Owner::Server => e.to_handle, - Owner::Initiator => e.from_handle, - }, + Some(e) => e.from_handle, None => return Claim::NotFound { namespace }, }; if owner_handle != handle { @@ -393,17 +525,80 @@ impl Router { Claim::Owned(entry) } + /// Phase 1 of a completion: snapshot the entry for `(namespace, + /// delegation_id)` and assert `serving_handle` is the instance it was + /// routed to — under one in-flight lock acquisition, removing nothing. + /// + /// The returned [`InFlight`] carries the [`InFlight::generation`] the + /// commit step must match, so delivery can happen outside the lock without + /// the commit ever being able to claim a different admission. + fn peek_for_completion( + &self, + namespace: &str, + delegation_id: &str, + serving_handle: u64, + ) -> Peek { + let key = DelegationKey::new(namespace, delegation_id); + let g = self.inflight.lock(); + match g.get(&key) { + Some(e) if e.to_handle == serving_handle => Peek::Serving(e.clone()), + Some(e) => Peek::Foreign { + owner_handle: e.to_handle, + }, + None => Peek::Unknown, + } + } + + /// Phase 2 of a completion: end the delegation `peeked` describes. + /// + /// Under ONE in-flight lock acquisition, the entry is removed and the + /// serving instance's capacity released only if the live entry is still + /// the same admission — key, serving handle, AND generation all match. + /// Anything else is left strictly untouched, including its capacity: + /// + /// - a concurrent cancel/sweep/disconnect already ended it → [`Commit::Vanished`], + /// and that path already released the capacity (releasing it here too + /// would let `saturated()` admit work to a full instance); + /// - the id was re-admitted in the meantime → [`Commit::Superseded`]. This + /// is the ABA case that made a stale commit destructive: matching on + /// key + serving handle alone, a commit for delegation *n* would remove + /// the live entry of delegation *n+1* (cancel-then-retry re-admits the + /// same id, and with one replica it routes to the same worker), making a + /// running delegation invisible to its own genuine result and to the + /// sweep, and wrongly freeing its slot. + fn commit_completion(&self, registry: &Registry, peeked: &InFlight) -> Commit { + let key = DelegationKey::new(&peeked.namespace, &peeked.delegation_id); + let removed = { + let mut g = self.inflight.lock(); + match g.get(&key) { + Some(e) if e.to_handle == peeked.to_handle && e.generation == peeked.generation => { + g.remove(&key).expect("present under the same lock") + } + Some(e) => { + return Commit::Superseded { + generation: e.generation, + } + } + None => return Commit::Vanished, + } + }; + registry.adjust_sessions(removed.to_handle, -1); + Commit::Claimed + } + /// Handle `cp/delegate_result` from the serving runtime. /// /// The terminal result is the one frame that must never be silently /// dropped, so delivery happens in two phases: /// /// 1. **Peek** — validate ownership under one in-flight lock acquisition - /// without removing the entry, then build and `try_send` the - /// initiator-bound frame. + /// without removing the entry ([`Router::peek_for_completion`]), then + /// build and `try_send` the initiator-bound frame. /// 2. **Commit** — only after the initiator's queue accepted the frame, - /// remove the entry (via [`Claim`], same single-lock property) and - /// release the serving instance's capacity. + /// end the delegation ([`Router::commit_completion`]): remove the + /// entry and release the serving instance's capacity, but only if the + /// live entry is still the very admission that was peeked (key + + /// serving handle + [`InFlight::generation`]). /// /// If the initiator's bounded queue refuses the frame, the entry stays /// in flight and [`CompleteOutcome::InitiatorStalled`] tells the caller @@ -411,11 +606,19 @@ impl Router { /// contract): its teardown runs `fail_instance`, which releases capacity /// exactly once and sends `cp/cancel` to the serving runtime. /// - /// Peek-then-commit admits one benign race: two concurrent genuine - /// results for the same id can both pass the peek and both be delivered, - /// but only the first commit releases capacity (the second finds the - /// entry gone and does nothing). Duplicate `cp/delegate_result` frames - /// are correlated by `delegation_id` and idempotent for the initiator. + /// Nothing outside the commit's exact-match window is touched, so the + /// peek-send window cannot corrupt CP state: a concurrent cancel, sweep, + /// or disconnect that already ended the delegation leaves this frame with + /// `Delivered { committed: false }`, and an id re-admitted in the window + /// keeps its own live entry and capacity. + /// + /// What the window *can* still produce is more than one terminal frame on + /// the wire for one `delegation_id` — a `completed` result racing the + /// sweep's synthesized `timeout`, or two duplicate results both passing + /// the peek. That is resolved by contract, not by CP-side suppression: + /// initiators MUST treat the FIRST terminal frame for a `delegation_id` as + /// authoritative and ignore later ones (see "first terminal frame wins" + /// in the ADR's v1 contract amendments). /// /// Only the instance the delegation was routed to may complete it; a /// non-owner frame can never make the delegation momentarily invisible @@ -442,24 +645,22 @@ impl Router { return CompleteOutcome::Dropped; } }; - let key = DelegationKey::new(&namespace, ¶ms.delegation_id); - let entry = { - let g = self.inflight.lock(); - match g.get(&key) { - Some(e) if e.to_handle == serving_handle => e.clone(), - Some(e) => { + let entry = + match self.peek_for_completion(&namespace, ¶ms.delegation_id, serving_handle) { + Peek::Serving(e) => e, + Peek::Foreign { owner_handle } => { // Only the instance the delegation was routed to may // complete it. The entry stays exactly where it is. warn!( delegation = %params.delegation_id, namespace = %namespace, - expected = e.to_handle, + expected = owner_handle, got = serving_handle, "result from unexpected instance — dropped, delegation untouched" ); return CompleteOutcome::Dropped; } - None => { + Peek::Unknown => { warn!( delegation = %params.delegation_id, namespace = %namespace, @@ -467,8 +668,7 @@ impl Router { ); return CompleteOutcome::Dropped; } - } - }; + }; // Truncate oversized results (keep the head; delegation already // ran). The marker counts against the cap: the final value never @@ -519,50 +719,45 @@ impl Router { }; } - // Phase 2 — commit. If a concurrent path (duplicate result, cancel, - // sweep, fail_instance) removed the entry between peek and now, that - // path also released the capacity — do not decrement twice. - match self.claim( - registry, - serving_handle, - ¶ms.delegation_id, - Owner::Server, - ) { - Claim::Owned(e) => { - registry.adjust_sessions(e.to_handle, -1); + // Phase 2 — commit. Claims ONLY the admission that was peeked; see + // `commit_completion` for why key + serving handle is not enough. + let committed = match self.commit_completion(registry, &entry) { + Commit::Claimed => { info!( delegation = %params.delegation_id, status = ?params.status, - from = %e.to_logical, - to = %e.from_logical, + from = %entry.to_logical, + to = %entry.from_logical, "delegation completed" ); + true } - Claim::WrongOwner { - namespace, - owner_handle, - } => { - // Only possible if the id was removed and re-admitted between - // peek and commit. The new delegation is not ours to touch. - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - owner = owner_handle, - "delegation re-owned between delivery and commit — entry left untouched" - ); - } - Claim::NotFound { namespace } => { + Commit::Vanished => { // Concurrent removal (duplicate result, cancel, sweep, or // fail_instance): whoever removed it released the capacity. info!( delegation = %params.delegation_id, - namespace = %namespace, + namespace = %entry.namespace, "entry removed concurrently after delivery — capacity already released" ); + false } - Claim::Unregistered => {} - } - CompleteOutcome::Delivered + Commit::Superseded { generation } => { + // The id was cancelled/expired and re-admitted between peek + // and commit. That new delegation is live and not ours to + // touch: removing it would strand a running delegation and + // free a slot it still occupies. + warn!( + delegation = %params.delegation_id, + namespace = %entry.namespace, + peeked_generation = entry.generation, + live_generation = generation, + "delegation id re-admitted between delivery and commit — live entry left untouched" + ); + false + } + }; + CompleteOutcome::Delivered { committed } } /// Handle `cp/cancel` from the initiator. Returns the frame to forward @@ -588,18 +783,17 @@ impl Router { "delegation is not in flight for this instance", ) }; - let entry = match self.claim( - registry, - from_handle, - ¶ms.delegation_id, - Owner::Initiator, - ) { + let entry = match self.claim(registry, from_handle, ¶ms.delegation_id) { Claim::Owned(entry) => entry, - Claim::WrongOwner { namespace, .. } => { + Claim::WrongOwner { + namespace, + owner_handle, + } => { warn!( delegation = %params.delegation_id, namespace = %namespace, handle = from_handle, + initiator = owner_handle, "cancel refused: only the initiating instance may cancel" ); return Err(refused()); @@ -904,7 +1098,7 @@ type = "worker" }; assert_eq!( w.router.complete(&w.registry, w.h_worker, result, 1024, 2), - CompleteOutcome::Delivered + CompleteOutcome::Delivered { committed: true } ); let frame = w.primary_rx.try_recv().unwrap(); assert!(frame.contains("\"completed\"")); @@ -930,7 +1124,7 @@ type = "worker" }; assert_eq!( w.router.complete(&w.registry, w.h_worker, result, 1024, 2), - CompleteOutcome::Delivered + CompleteOutcome::Delivered { committed: true } ); w.worker_rx.try_recv().unwrap(); } @@ -1188,7 +1382,7 @@ type = "worker" let cap = 96usize; assert_eq!( w.router.complete(&w.registry, w.h_worker, result, cap, 2), - CompleteOutcome::Delivered + CompleteOutcome::Delivered { committed: true } ); let frame = w.primary_rx.try_recv().unwrap(); let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); @@ -1215,7 +1409,7 @@ type = "worker" }; assert_eq!( w.router.complete(&w.registry, w.h_worker, result2, 8, 3), - CompleteOutcome::Delivered + CompleteOutcome::Delivered { committed: true } ); let frame2 = w.primary_rx.try_recv().unwrap(); let v2: serde_json::Value = serde_json::from_str(&frame2).unwrap(); @@ -1488,7 +1682,7 @@ allow_worker_initiation = true 1024, 3, ), - CompleteOutcome::Delivered, + CompleteOutcome::Delivered { committed: true }, "genuine result must be delivered, never dropped" ); let frame = w.primary_rx.try_recv().unwrap(); @@ -1537,7 +1731,7 @@ allow_worker_initiation = true result_of("d-1", "spoofed"), 1024, 2, - ) == CompleteOutcome::Delivered + ) == CompleteOutcome::Delivered { committed: true } }); gate.wait(); let genuine = w.router.complete( @@ -1546,7 +1740,7 @@ allow_worker_initiation = true result_of("d-1", "genuine"), 1024, 3, - ) == CompleteOutcome::Delivered; + ) == CompleteOutcome::Delivered { committed: true }; (spoof.join().unwrap(), genuine) }); assert!(!spoofed, "a non-owner must never complete a delegation"); @@ -1734,7 +1928,7 @@ allow_worker_initiation = true // Results route to the initiator of the SAME namespace only. assert_eq!( router.complete(®istry, hw_dev, result_of("d-1", "dev-done"), 1024, 12), - CompleteOutcome::Delivered + CompleteOutcome::Delivered { committed: true } ); let frame = dev_init_rx.try_recv().unwrap(); assert!(frame.contains("dev-done")); @@ -1749,7 +1943,7 @@ allow_worker_initiation = true assert_eq!( router.complete(®istry, hw_prod, result_of("d-1", "prod-done"), 1024, 13), - CompleteOutcome::Delivered + CompleteOutcome::Delivered { committed: true } ); let frame = prod_init_rx.try_recv().unwrap(); assert!(frame.contains("prod-done")); @@ -1840,4 +2034,309 @@ allow_worker_initiation = true vec!["prod/koudu".to_string(), "prod/worker-1".to_string()] ); } + + /// Phase-1 snapshot exactly as `complete` takes it, so a test can hold a + /// real pre-delivery entry (generation included) across an interleaved + /// operation and then drive the commit step directly — deterministic, + /// no barrier timing. + fn peek(w: &World, id: &str) -> InFlight { + match w.router.peek_for_completion("prod", id, w.h_worker) { + Peek::Serving(e) => e, + _ => panic!("{id} must be in flight and served by the worker"), + } + } + + /// How the in-flight entry is removed between peek and commit. + #[derive(Debug, Clone, Copy)] + enum Interleaved { + /// The initiator cancels (`cp/cancel`). + Cancel, + /// The deadline sweep expires it. + Sweep, + } + + /// Perform the interleaved removal and return the frames it synthesized + /// for delivery (the router builds them; the caller is what sends them). + fn interleave(w: &World, how: Interleaved, id: &str) -> Vec { + match how { + Interleaved::Cancel => { + let params = CancelParams { + delegation_id: id.into(), + reason: "changed my mind".into(), + }; + w.router + .cancel(&w.registry, w.h_primary, ¶ms, 90) + .expect("the initiator may cancel") + .map(|(_, frame)| frame) + .into_iter() + .collect() + } + Interleaved::Sweep => { + let mut id_seq = 900u64; + let mut next = || { + id_seq += 1; + id_seq + }; + let frames = w.router.sweep_deadlines( + &w.registry, + Utc::now() + Duration::seconds(3600), + &mut next, + ); + assert!(!frames.is_empty(), "the sweep must have expired something"); + frames.into_iter().map(|(_, frame)| frame).collect() + } + } + } + + fn drain(w: &mut World) { + while w.worker_rx.try_recv().is_ok() {} + while w.primary_rx.try_recv().is_ok() {} + } + + #[test] + fn stale_commit_never_claims_a_reused_delegation_id() { + // The ABA case. `(namespace, delegation_id)` + serving handle is not a + // stable identity: cancel-then-retry is an ordinary client pattern, the + // id is client-supplied, and with a single replica the retry routes to + // the SAME worker. A commit that matched on those alone would remove + // the RETRY's live entry and release its capacity — the running + // delegation becomes invisible to its own genuine result and to the + // sweep, and its slot is handed out while still occupied. + // + // The generation stamp makes the commit claim the admission it + // actually delivered for, or nothing. + for how in [Interleaved::Cancel, Interleaved::Sweep] { + let mut w = world(); // worker max_delegated_sessions = 1 + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + // A: peeked, its initiator-bound frame notionally sent. + let a = peek(&w, "d-1"); + + // A is removed and its capacity released by another path. + interleave(&w, how, "d-1"); + assert_eq!(w.router.inflight_count(), 0, "{how:?}"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "{how:?}: the removing path releases the capacity" + ); + drain(&mut w); + + // B: the initiator retries the same id; the only replica is the + // same worker, so key AND serving handle repeat exactly. + assert!( + matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + ), + "{how:?}: the freed slot must admit the retry" + ); + let b = peek(&w, "d-1"); + assert_eq!(a.to_handle, b.to_handle, "{how:?}: same worker"); + assert_eq!(a.delegation_id, b.delegation_id); + assert_ne!( + a.generation, b.generation, + "{how:?}: generations are never reused" + ); + + // The stale commit for A lands. It must claim nothing. + assert_eq!( + w.router.commit_completion(&w.registry, &a), + Commit::Superseded { + generation: b.generation + }, + "{how:?}: a stale commit must not claim the re-admitted entry" + ); + assert_eq!( + w.router.inflight_count(), + 1, + "{how:?}: B must remain in flight" + ); + assert_eq!( + w.router.chain_of("prod", "d-1").as_deref(), + Some(&["prod/koudu".to_string()][..]), + "{how:?}: B must still be visible to results and to the sweep" + ); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "{how:?}: B's capacity reservation must be intact" + ); + + // B then completes normally — its genuine result is delivered and + // commits, releasing the capacity exactly once. + assert_eq!( + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", "genuine"), + 1024, + 7 + ), + CompleteOutcome::Delivered { committed: true }, + "{how:?}" + ); + let frame = w.primary_rx.try_recv().expect("initiator got the result"); + assert!(frame.contains("genuine"), "{how:?}"); + assert_eq!(w.router.inflight_count(), 0, "{how:?}"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "{how:?}" + ); + } + } + + #[test] + fn commit_after_cancel_releases_capacity_exactly_once() { + // commit-vs-cancel, driven directly: the initiator cancels between the + // peek and the commit and no retry follows. The commit finds its + // admission gone and must not decrement again — session counts are + // saturating, so a double release is silent and `saturated()` would + // then admit work to a full instance. + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + // Second delegation on a 4-slot target, to make an underflow visible + // as a wrong count rather than a saturating clamp at zero. + let (extra, _extra_rx) = instance("prod", "worker-2", AgentType::Worker, 4); + let h_extra = w.registry.register(extra); + assert!(matches!( + do_delegate(&w, delegate_params("d-2", "worker-2", 60)), + DelegateOutcome::Accepted(_) + )); + let a = peek(&w, "d-1"); + + interleave(&w, Interleaved::Cancel, "d-1"); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + drain(&mut w); + + assert_eq!( + w.router.commit_completion(&w.registry, &a), + Commit::Vanished + ); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "capacity must be released exactly once, by the cancel" + ); + // The unrelated delegation is untouched by any of this. + assert_eq!(w.router.inflight_count(), 1); + assert_eq!(w.registry.get(h_extra).unwrap().active_sessions, 1); + } + + #[test] + fn commit_after_deadline_sweep_releases_capacity_exactly_once() { + // commit-vs-sweep, driven directly: the deadline sweep expires the + // delegation between the peek and the commit. Same requirement as the + // cancel case, and the initiator has already been sent the sweep's + // `timeout` — the wire then carries two terminal frames for one id, + // which the ADR resolves with "first terminal frame wins". + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let a = peek(&w, "d-1"); + + // The sweep's `timeout` is the first terminal frame for this id; the + // late `completed` result below is the second, which the initiator + // ignores per the ADR's "first terminal frame wins". + let swept = interleave(&w, Interleaved::Sweep, "d-1"); + assert!( + swept.iter().any(|f| f.contains("\"timeout\"")), + "the sweep must synthesize a timeout result for the initiator" + ); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + drain(&mut w); + + assert_eq!( + w.router.commit_completion(&w.registry, &a), + Commit::Vanished + ); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "capacity must be released exactly once, by the sweep" + ); + assert_eq!(w.router.inflight_count(), 0); + // The freed slot is genuinely free (a double release would have made + // the count underflow and this admission could exceed max=1). + assert!(matches!( + do_delegate(&w, delegate_params("d-2", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + match do_delegate(&w, delegate_params("d-3", "worker-1", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!( + e.code, + codes::SATURATED, + "the worker's single slot must still bound admission" + ), + _ => panic!("expected SATURATED — capacity accounting drifted"), + } + } + + #[test] + fn delivered_reports_whether_it_committed() { + // Wire delivery and state commit are separate events, and the outcome + // says which happened: an unconditional `Delivered` could not + // distinguish "this frame ended the delegation" from "somebody else + // already had". + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let a = peek(&w, "d-1"); + assert_eq!( + w.router.commit_completion(&w.registry, &a), + Commit::Claimed, + "the live admission is claimed exactly once" + ); + assert_eq!( + w.router.commit_completion(&w.registry, &a), + Commit::Vanished, + "a repeated commit of the same admission is a no-op" + ); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + drain(&mut w); + + // Through the public path: the entry is gone, so the frame is not + // even delivered — a peek that finds nothing is a plain drop. + assert_eq!( + w.router + .complete(&w.registry, w.h_worker, result_of("d-1", "late"), 1024, 8), + CompleteOutcome::Dropped + ); + assert!(w.primary_rx.try_recv().is_err()); + } + + #[test] + fn generation_exhaustion_fails_closed_instead_of_wrapping() { + // fetch_add on an exhausted counter would wrap and re-issue + // generations, silently recreating the ABA the stamp prevents. + // Unreachable in a realistic process lifetime; pinned here so a + // refactor cannot quietly downgrade it to wrapping arithmetic. + let w = world(); + w.router + .next_generation + .store(u64::MAX - 1, Ordering::Relaxed); + // At the ceiling, fetch_update declines to store: the admission is + // refused and the counter stays parked at MAX-1 forever. + let out = do_delegate(&w, delegate_params("d-last", "worker-1", 60)); + assert!( + matches!(out, DelegateOutcome::Rejected(ref e) if e.code == -32603), + "exhausted generation space must refuse admission" + ); + // No capacity was reserved by the refused admission. + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + // And it stays parked: the next attempt is refused too. + let out2 = do_delegate(&w, delegate_params("d-next", "worker-1", 60)); + assert!(matches!(out2, DelegateOutcome::Rejected(ref e) if e.code == -32603)); + } } diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index 38b65ec1a..77cbcbe22 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -6,7 +6,10 @@ //! //! Resource bounds: the WS transport enforces `max_frame_bytes` before //! parsing; each connection's outbound queue is bounded — a peer that cannot -//! drain it is treated as disconnected. +//! drain it is treated as disconnected — and every outbound write is itself +//! bounded by `write_timeout_secs` and raced against the CP's close signal, so +//! a peer that stops reading cannot park the connection task (and with it the +//! identity's connection quota and its in-flight delegations). //! //! Admission bounds: authentication alone is not a bound. Every connection //! holds a per-identity slot from the upgrade until it ends (`ConnPermit`, @@ -33,7 +36,7 @@ use axum::routing::get; use axum::Router as AxumRouter; use futures_util::{SinkExt, StreamExt}; use parking_lot::Mutex; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tracing::{info, warn}; use crate::config::{AgentIdentity, CpConfig}; @@ -144,6 +147,49 @@ fn policy_close(reason: &'static str) -> Message { })) } +/// Why a bounded write did not complete. Either way the connection ends. +enum WriteStop { + /// Transport error, or the peer did not accept the frame within + /// `write_timeout_secs` — a peer that cannot be written to is treated as + /// disconnected, exactly like one that cannot drain its queue. + Disconnected, + /// The CP asked this connection to close while the write was pending, + /// with the reason to put on the close frame (`None` if the signal + /// carried none). + Shutdown(Option<&'static str>), +} + +/// Send one frame with a bound on how long it may block, while remaining +/// responsive to the CP's own close signal. +/// +/// Both properties are load-bearing. `sink.send().await` inside a `select!` +/// arm body is NOT cancelled by the other arms, so an unbounded write parks +/// the whole connection task: it stops reading inbound frames, stops +/// observing the shutdown watch, and — because [`ConnPermit`] is released +/// only when the task returns — pins its identity's connection quota. A peer +/// with a closed TCP receive window or a half-open connection could hold +/// those slots indefinitely and no lease expiry could reclaim them, since +/// lease expiry works by signalling this very task. +/// +/// A cancelled or timed-out write can leave a partially written frame on the +/// wire; that is acceptable precisely because both outcomes end the +/// connection (the caller breaks to teardown, dropping the socket). +async fn send_bounded( + sink: &mut futures_util::stream::SplitSink, + msg: Message, + write_timeout: Duration, + shutdown_rx: &mut watch::Receiver>, +) -> Result<(), WriteStop> { + tokio::select! { + _ = shutdown_rx.changed() => Err(WriteStop::Shutdown(*shutdown_rx.borrow_and_update())), + sent = tokio::time::timeout(write_timeout, sink.send(msg)) => match sent { + Ok(Ok(())) => Ok(()), + Ok(Err(_)) => Err(WriteStop::Disconnected), + Err(_) => Err(WriteStop::Disconnected), + }, + } +} + async fn ws_handler( State(state): State>, headers: HeaderMap, @@ -197,6 +243,19 @@ async fn handle_connection( _permit: ConnPermit, ) { let (mut sink, mut stream) = socket.split(); + let write_timeout = Duration::from_secs(state.cfg.write_timeout_secs); + + // Shutdown signal so the CP can close this socket when it drops the + // registration on its own initiative (lease expiry) or must terminate the + // connection (terminal-frame backpressure). + // + // Created BEFORE the registration read — and therefore before the registry + // ever holds a clone — so no signal can be missed, and every write in this + // task, registration-phase writes included, can be raced against it. Kept + // alive here for the whole connection: closing is driven by an explicit + // signal, never by the registry happening to drop its side. + let shutdown = shutdown_signal(); + let mut shutdown_rx = shutdown.subscribe(); // --- Registration: mandatory first frame, within a deadline --- // An authenticated peer must not be able to park idle sockets: pings keep @@ -226,7 +285,13 @@ async fn handle_connection( timeout_secs = state.cfg.register_timeout_secs, "no cp/register within the registration deadline — closing" ); - let _ = sink.send(policy_close(REASON_REGISTER_TIMEOUT)).await; + let _ = send_bounded( + &mut sink, + policy_close(REASON_REGISTER_TIMEOUT), + write_timeout, + &mut shutdown_rx, + ) + .await; return; } }; @@ -234,11 +299,13 @@ async fn handle_connection( Ok(ok) => ok, Err((id, err)) => { let resp = JsonRpcErrorResponse::new(id, err); - let _ = sink - .send(Message::Text( - serde_json::to_string(&resp).expect("serializable").into(), - )) - .await; + let _ = send_bounded( + &mut sink, + Message::Text(serde_json::to_string(&resp).expect("serializable").into()), + write_timeout, + &mut shutdown_rx, + ) + .await; return; } }; @@ -247,14 +314,6 @@ async fn handle_connection( // cannot drain OUTBOUND_QUEUE frames is disconnected, not buffered. let (tx, mut rx) = mpsc::channel::(OUTBOUND_QUEUE); - // Shutdown signal so the CP can close this socket when it drops the - // registration on its own initiative (lease expiry). - // Subscribed BEFORE registering so no signal can be missed, and kept - // alive here for the whole connection: closing is driven by an explicit - // signal, never by the registry happening to drop its side. - let shutdown = shutdown_signal(); - let mut shutdown_rx = shutdown.subscribe(); - let effective_max = match identity.max_delegated_sessions_cap { Some(cap) => reg.max_delegated_sessions.min(cap), None => reg.max_delegated_sessions, @@ -297,19 +356,59 @@ async fn handle_connection( reg_rpc_id, serde_json::to_value(&ack).expect("serializable"), ); - if sink - .send(Message::Text( - serde_json::to_string(&resp).expect("serializable").into(), - )) - .await - .is_err() + if let Err(stop) = send_bounded( + &mut sink, + Message::Text(serde_json::to_string(&resp).expect("serializable").into()), + write_timeout, + &mut shutdown_rx, + ) + .await { + // CP-initiated closes carry meaning even here: if the CP signalled + // this connection while the ack was in flight, still tell the client + // why before tearing down. + if let WriteStop::Shutdown(Some(reason)) = stop { + let _ = send_bounded( + &mut sink, + policy_close(reason), + write_timeout, + &mut shutdown_rx, + ) + .await; + } teardown(&state, handle, &identity); return; } // --- Main loop: interleave inbound frames, outbound channel, shutdown --- + // + // Every write goes through `send_bounded`: a `select!` arm body is not + // cancelled by the other arms, so an unbounded write here would stop this + // task from observing the shutdown watch and from releasing its + // `ConnPermit` — see `send_bounded`. let mut cp_close_reason: Option<&'static str> = None; + // A macro, not a closure: each expansion borrows `sink` only for the + // duration of its own arm body, and `break` acts on the loop below. + macro_rules! write_or_break { + ($msg:expr) => { + match send_bounded(&mut sink, $msg, write_timeout, &mut shutdown_rx).await { + Ok(()) => {} + Err(WriteStop::Shutdown(reason)) => { + cp_close_reason = reason; + break; + } + Err(WriteStop::Disconnected) => { + warn!( + handle, + timeout_secs = state.cfg.write_timeout_secs, + "outbound write failed or exceeded write_timeout_secs — \ + treating the peer as disconnected" + ); + break; + } + } + }; + } loop { tokio::select! { // The CP dropped this registration (lease expiry) or must @@ -325,11 +424,7 @@ async fn handle_connection( } outbound = rx.recv() => { match outbound { - Some(text) => { - if sink.send(Message::Text(text.into())).await.is_err() { - break; - } - } + Some(text) => write_or_break!(Message::Text(text.into())), None => break, } } @@ -337,16 +432,10 @@ async fn handle_connection( match inbound { Some(Ok(Message::Text(text))) => { if let Some(reply) = handle_frame(&state, handle, &text) { - if sink.send(Message::Text(reply.into())).await.is_err() { - break; - } - } - } - Some(Ok(Message::Ping(p))) => { - if sink.send(Message::Pong(p)).await.is_err() { - break; + write_or_break!(Message::Text(reply.into())); } } + Some(Ok(Message::Ping(p))) => write_or_break!(Message::Pong(p)), Some(Ok(Message::Close(_))) | None => break, Some(Ok(_)) => {} // binary/pong ignored Some(Err(e)) => { @@ -365,7 +454,16 @@ async fn handle_connection( reason, "closing connection at the CP's request" ); - let _ = sink.send(policy_close(reason)).await; + // Bounded like every other write: a peer that has stopped reading must + // not be able to delay teardown (and its quota slot) by refusing to + // accept the close frame. + let _ = send_bounded( + &mut sink, + policy_close(reason), + write_timeout, + &mut shutdown_rx, + ) + .await; } teardown(&state, handle, &identity); @@ -610,9 +708,26 @@ fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option { + // The result reached the initiator, which is what the serving + // side is being acked for. Whether THIS frame also committed + // the state transition is a CP-internal matter: a concurrent + // cancel/sweep/disconnect may have ended the delegation + // first, and the initiator resolves competing terminal frames + // by "first one wins" (see the ADR wire contract). + CompleteOutcome::Delivered { committed } => { + if !committed { + info!( + handle, + "terminal result delivered, but the delegation had already been \ + ended (or its id re-admitted) — no state change" + ); + } + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } + // Dropped as unknown/foreign; each case is logged in the + // router (late results after a CP restart are expected). + CompleteOutcome::Dropped => { let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); Some(serde_json::to_string(&resp).expect("serializable")) } diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index b3df9025a..5d24cacd1 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -20,6 +20,7 @@ use openab_cp::config::CpConfig; use openab_cp::server::{app, sweep_leases, AppState}; const KEY: &str = "k-primary"; +const KEY_WORKER: &str = "k-worker"; type Ws = WebSocketStream>; @@ -33,6 +34,12 @@ key = "{KEY}" namespace = "prod" name = "koudu" type = "primary" + +[[agents]] +key = "{KEY_WORKER}" +namespace = "prod" +name = "worker-1" +type = "worker" "# ); let cfg: CpConfig = toml::from_str(&raw).expect("test config parses"); @@ -52,17 +59,21 @@ async fn spawn_cp(cfg: CpConfig) -> (Arc, String) { (state, format!("ws://{addr}/cp")) } -async fn connect(url: &str) -> Result { +async fn connect_as(url: &str, key: &str) -> Result { let mut req = url.into_client_request().unwrap(); req.headers_mut().insert( "authorization", - format!("Bearer {KEY}").parse().expect("header value"), + format!("Bearer {key}").parse().expect("header value"), ); tokio_tungstenite::connect_async(req) .await .map(|(ws, _)| ws) } +async fn connect(url: &str) -> Result { + connect_as(url, KEY).await +} + /// Connect, retrying while the identity's quota slot is still being released /// by the server task. async fn connect_retry(url: &str) -> Ws { @@ -299,3 +310,188 @@ async fn pre_registration_sockets_count_against_the_quota() { Ok(_) => panic!("an unregistered socket must still occupy its quota slot"), } } + +/// Register a worker that advertises `max_sessions` concurrent delegations. +async fn register_worker(ws: &mut Ws, instance_id: &str, max_sessions: u32) -> serde_json::Value { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "worker-1", + "type": "worker", + "instance_id": instance_id, + "max_delegated_sessions": max_sessions + } + }) + .to_string(); + ws.send(Message::Text(frame.into())).await.unwrap(); + let msg = tokio::time::timeout(Duration::from_secs(5), ws.next()) + .await + .expect("register must be answered") + .expect("stream open") + .expect("no ws error"); + serde_json::from_str(msg.to_text().unwrap()).unwrap() +} + +#[tokio::test] +async fn a_peer_that_stops_reading_is_disconnected_and_frees_its_quota() { + // A bounded outbound queue does not bound the WRITER. `sink.send().await` + // inside a `select!` arm body is not cancelled by the other arms, so a peer + // that stops reading (closed TCP receive window / half-open connection) + // parks the connection task inside one write: it no longer observes the + // shutdown watch, and its `ConnPermit` — released only when the task + // returns — pins the identity's quota slots. Lease expiry cannot recover + // that, because lease expiry works by signalling this very task. + // + // Every write is therefore bounded by `write_timeout_secs` and raced + // against the shutdown signal: a timeout is a disconnect, which runs + // teardown, drops the permit, and cancels the delegations downstream. + const WRITE_TIMEOUT_SECS: u64 = 1; + // Enough result payload to overrun any socket buffer on the way to a peer + // that never reads (8 MiB inbound cap, ~4 MiB per result, 12 of them). + const BIG: usize = 4 * 1024 * 1024; + const RESULTS: usize = 12; + + let (state, url) = spawn_cp(cfg(&format!( + "max_connections_per_identity = 1 +register_timeout_secs = 30 +write_timeout_secs = {WRITE_TIMEOUT_SECS} +max_frame_bytes = 8388608 +max_result_bytes = 8388608" + ))) + .await; + + // The initiator registers and then never reads again. + let mut stalled = connect(&url).await.expect("initiator accepted"); + assert_eq!( + register(&mut stalled, "i-stalled").await["result"]["protocol_version"], + 1 + ); + assert_eq!(state.conn_count("prod/koudu"), 1); + + // A worker that behaves normally: it serves the delegations and answers. + let mut worker = connect_as(&url, KEY_WORKER).await.expect("worker accepted"); + assert_eq!( + register_worker(&mut worker, "i-worker", (RESULTS + 1) as u32).await["result"] + ["protocol_version"], + 1 + ); + + // The initiator delegates RESULTS + 1 times. `d-keep` is never completed, + // so a live delegation remains for the disconnect path to cancel. + let deadline = (chrono::Utc::now() + chrono::Duration::seconds(600)).to_rfc3339(); + let mut ids: Vec = vec!["d-keep".to_string()]; + ids.extend((0..RESULTS).map(|i| format!("d-{i}"))); + for (n, id) in ids.iter().enumerate() { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 100 + n, "method": "cp/delegate", + "params": { + "delegation_id": id, + "target": {"name": "worker-1"}, + "prompt": "do it", + "deadline": deadline + } + }) + .to_string(); + stalled.send(Message::Text(frame.into())).await.unwrap(); + } + + // Collect ALL forwards on the worker side first, so every delegation is + // admitted before any large write can block the initiator's task. + let mut pending = Vec::new(); + while pending.len() < ids.len() { + let msg = tokio::time::timeout(Duration::from_secs(20), worker.next()) + .await + .expect("worker must receive its forwards") + .expect("stream open") + .expect("no ws error"); + let v: serde_json::Value = match msg { + Message::Text(t) => serde_json::from_str(&t).unwrap(), + _ => continue, + }; + if v["method"] != "cp/delegate" { + continue; + } + pending.push(v["params"]["delegation_id"].as_str().unwrap().to_string()); + } + assert_eq!( + state.router.inflight_count(), + ids.len(), + "every delegation must be admitted before the writer is stalled" + ); + + // Answer all but `d-keep` with a large result. Each is forwarded to the + // stalled initiator, whose socket buffers fill and whose write then blocks. + for (n, id) in pending.iter().enumerate() { + if id == "d-keep" { + continue; + } + let result = serde_json::json!({ + "jsonrpc": "2.0", "id": 900 + n, "method": "cp/delegate_result", + "params": { + "delegation_id": id, + "status": "completed", + "result": "x".repeat(BIG) + } + }) + .to_string(); + worker.send(Message::Text(result.into())).await.unwrap(); + } + + // The stalled peer's connection task must terminate on its own — the write + // bound is the only thing that can make it happen — releasing the quota + // slot. Before the fix this never occurred and the assertion timed out. + let started = Instant::now(); + let mut released = false; + // The bound is the point of the regression: the permit must free within + // the configured write timeout plus scheduling allowance — NOT eventually. + // (A 30s poll here would pass even if the permit were pinned for 20s.) + let bound = Duration::from_secs(WRITE_TIMEOUT_SECS + 4); + while started.elapsed() < bound { + if state.conn_count("prod/koudu") == 0 { + released = true; + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!( + released, + "a peer that stopped reading must be disconnected within the write bound; \ + conn_count was still {} after {:?}", + state.conn_count("prod/koudu"), + started.elapsed() + ); + + // Downstream cancellation fires: teardown fails the initiator's remaining + // delegation, which sends `cp/cancel` to the serving worker. + let mut cancelled = false; + let cancel_deadline = Instant::now() + Duration::from_secs(20); + while Instant::now() < cancel_deadline { + match tokio::time::timeout(Duration::from_millis(500), worker.next()).await { + Ok(Some(Ok(Message::Text(t)))) => { + if t.contains("cp/cancel") && t.contains("d-keep") { + cancelled = true; + break; + } + } + Ok(None) | Ok(Some(Err(_))) => break, + _ => continue, + } + } + assert!( + cancelled, + "the stalled initiator's teardown must cancel its in-flight delegation downstream" + ); + assert_eq!(state.router.inflight_count(), 0); + + // The quota slot is genuinely free: a fresh connection for the SAME + // identity is accepted and can register (quota is 1 here). + drop(stalled); + let mut fresh = connect_retry(&url).await; + assert_eq!( + register(&mut fresh, "i-fresh").await["result"]["protocol_version"], + 1, + "the released slot must be reusable by the same identity" + ); +} diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 36fd8d501..877e5625b 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -275,10 +275,17 @@ recovery semantics: - **Resource bounds.** The WS transport rejects messages over `max_frame_bytes` before parsing; oversized `prompt`s are rejected (`max_prompt_bytes`); per-connection outbound queues are bounded and a - peer that cannot drain its queue is treated as disconnected. Delegation - admission (duplicate check → target selection → capacity reservation → - in-flight insert) is one atomic sequence, and the in-flight entry exists - before the forward frame is sent. + peer that cannot drain its queue is treated as disconnected. Every outbound + write is additionally bounded by `write_timeout_secs` and raced against the + CP's own close signal, because a bounded queue does not bound the *writer*: + a peer that stops reading (closed TCP receive window, half-open connection) + would otherwise park the connection task inside one `send`, pinning that + identity's connection quota and its in-flight delegations for as long as the + socket survives — and lease expiry could not reclaim them, since lease + expiry works by signalling that very task. A write that times out is a + disconnect. Delegation admission (duplicate check → target selection → + capacity reservation → in-flight insert) is one atomic sequence, and the + in-flight entry exists before the forward frame is sent. - **Terminal results are never silently dropped.** `cp/delegate_result` delivery commits only after the initiator's queue accepts the frame: the CP validates ownership, sends, and only then removes the in-flight entry @@ -288,11 +295,30 @@ recovery semantics: `outbound queue overflow` — the "cannot drain → disconnected" rule applied to the frame where it matters most), and the delegation resolves through the disconnect path: `cp/cancel` to the serving runtime and exactly one - capacity release. Two concurrent duplicate results may both be delivered - (correlation is by `delegation_id`; delivery is idempotent for the - initiator), but capacity is released exactly once. Best-effort frames - (`cp/cancel`, sweep-synthesized `timeout`) remain fire-and-forget: the - propagated deadline is their backstop. + capacity release. Best-effort frames (`cp/cancel`, sweep-synthesized + `timeout`) remain fire-and-forget: the propagated deadline is their backstop. +- **Admissions carry a generation; commits are exact.** Every admission is + stamped with a CP-generated, never-reused generation. The commit phase of a + completion removes an in-flight entry only when key, serving handle, AND + generation all match the admission it delivered a result for; anything else + is left strictly untouched, capacity included. `(namespace, delegation_id)` + is deliberately not a stable identity over time — the id is client-supplied + and cancel-then-retry is an ordinary client pattern, which with a single + replica re-admits the same id to the same worker — so without the generation + a stale commit could remove a *live* delegation's entry, making it invisible + to its own genuine result and to the deadline sweep while freeing a slot it + still occupies. +- **First terminal frame wins.** More than one terminal frame may reach an + initiator for one `delegation_id`: a `completed` result can race the deadline + sweep's synthesized `timeout`, and duplicate results are possible in the + window between delivery and commit. **Initiators MUST treat the first + terminal frame (`completed`, `failed`, `timeout`, `target_disconnected`) for + a `delegation_id` as authoritative and ignore every later terminal frame for + that id.** The CP does not suppress the later frames: doing so would require + per-id terminal state that a CP with no durable state deliberately does not + keep, and the initiator already correlates by `delegation_id`. CP-side state + is unaffected either way — the generation rule above makes the commit exact, + so exactly one path ever releases the capacity. - **Capacity release follows entry removal.** Whichever path removes an in-flight entry (result commit, cancel, deadline sweep, instance failure, or a failed forward's rollback) releases its capacity reservation — and diff --git a/docs/control-plane.md b/docs/control-plane.md index fae85aed6..31c211c70 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -32,13 +32,13 @@ rationale. The essentials: TLS-terminating proxy (`wss://`) or a private network in front is required: runtimes authenticate with bearer keys that must never cross untrusted cleartext TCP. -- `[[identity]]` — one entry per agent identity: the auth key (supports +- `[[agents]]` — one entry per agent identity: the auth key (supports `${ENV_VAR}` expansion) and its immutable `namespace`/`name`/`type` claims. A connecting runtime must register as exactly the identity its key is bound to. - Heartbeats, lease expiry, registration deadline, per-identity connection - quotas, and frame/prompt/result size caps are all configurable with safe - defaults. + quotas, the outbound write timeout, and frame/prompt/result size caps are + all configurable with safe defaults. ## Health @@ -50,5 +50,12 @@ issue #1474). - CP-initiated closes use WS code 1008 with a reason: `registration timeout`, `lease expired`, or `outbound queue overflow`. On any of these, reconnect, re-authenticate, and re-register. +- A peer that stops reading is disconnected: any single outbound write that + blocks longer than `write_timeout_secs` is treated as a dead peer, so keep + draining the socket even while busy. +- **The first terminal frame for a `delegation_id` wins.** A `completed` + result can race the CP's synthesized `timeout`, so an initiator may receive + more than one terminal frame for the same delegation. Treat the first as + authoritative and ignore later ones; the CP does not suppress them. - After a lease expires or the CP restarts, in-flight delegations are gone: initiators reconcile against their own deadlines and re-delegate. From d1b51048f73e452bca085bfcc07038276c1129a4 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Thu, 13 Aug 2026 23:43:37 -0400 Subject: [PATCH 08/11] =?UTF-8?q?fix(cp):=20round-8=20review=20=E2=80=94?= =?UTF-8?q?=20wire=20admission=20tokens,=20byte/inflight/capacity=20bounds?= =?UTF-8?q?,=20panic-safe=20teardown?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 (blocking) — admission identity is now protocol-visible. The round-6 generation stopped at the CP boundary: the wire carried only the reusable delegation_id, so after cancel-then-retry re-admitted the same id to the same worker, a late cp/delegate_result for the old admission A peeked the live B, was delivered to the initiator as B's terminal frame, and committed B (peek and commit both saw B, so B's own generation matched) — freeing capacity B still occupied and leaving B's genuine result to be dropped as unknown. - proto.rs: `admission` (AdmissionToken) added to DelegateAck (initiator learns it), DelegateForward (worker learns it), and DelegateResultParams as a REQUIRED field (worker MUST echo it; absent = INVALID_PARAMS, never a wildcard). CP-synthesized terminal frames build DelegateResultParams from the in-flight entry, so sweep `timeout` and `target_disconnected` carry the token of the admission they end. - router.rs: peek_for_completion matches the echoed token against the live entry BEFORE the initiator-bound frame is built (delivery is the irreversible half). Mismatch → new Peek::StaleAdmission → Dropped, logged distinctly with both tokens, but answered with the same generic ack as any other drop: a distinguishable reply would tell the serving side whether an id is currently re-admitted (the oracle class round-3 F3 removed). - Token design: the generation counter was ONE GLOBAL monotonic counter; putting that raw value on the wire would leak cross-namespace delegation volume. Commit matching only needs never-reuse per (namespace, delegation_id), so counters are now per namespace, held as the value guarded by the admission lock (BTreeMap) — minting already happened there, and this makes minting outside the admission sequence structurally impossible. Fail-closed exhaustion is preserved (a counter at the ceiling is never bumped, so it can never wrap and re-issue tokens) and is now also proven to be per-namespace: prod exhaustion does not refuse dev. - Contract: ADR §4 amendment rewritten — the token's full round trip, the echo requirement, the stale-token drop, the namespace-scoping rationale — and "first terminal frame wins" is re-keyed per admission token (correlating on the reusable id lets a late frame for a superseded admission permanently mask the live one). docs/control-plane.md carries the same for runtime authors. cp/cancel deliberately does NOT carry the token yet; it lands with the runtime-client slice. F2 — per-connection outbound BYTE budget. 256 entries is not a memory bound when frame sizes are configured independently (~2 GiB at 8 MiB results). registry.rs now owns OutboundBudget: bytes are reserved before enqueue, released on dequeue and on teardown (FrameRx::drop drains and releases, so a surviving sender clone never sees budget consumed by a dead connection). FrameTx/FrameRx replace the bare mpsc alias; refusal semantics are unchanged (SendRefused is still just "this peer cannot take the frame → treat as disconnected"). Config `max_outbound_queue_bytes` (default 16 MiB, validated > max_frame_bytes so one maximum-size frame always fits), documented in cp.toml.example. F3 — global in-flight bound. Config `max_inflight_delegations` (default 4096, validated > 0), enforced at admission before any capacity is reserved, refused with SATURATED naming the bound. DEVIATION with rationale: the bound is the in-flight table's own length rather than a parallel counter, so no removal path has to remember to decrement — commit, cancel, sweep, fail_instance and the send-failure rollback all remove the entry and the count follows by construction. A parallel counter is one refactor away from drifting, and a drifted global bound wedges the whole CP. The regression walks all five removal paths and re-admits after each. F4 — default clamp for advertised capacity. Config `default_max_delegated_sessions_cap` (default 16, validated > 0) applied via CpConfig::effective_max_sessions to every identity with no max_delegated_sessions_cap; per-identity value still overrides in either direction. There is no longer an uncapped path, so a worker cannot advertise its way out of saturation-based backpressure. F5 — panic-safe teardown. Teardown ran only on handle_connection's normal return path; a panic (reachable from expect("serializable") on production paths) left the registry entry, the in-flight rows, and the capacity they reserve on OTHER instances for the lease sweeper to reclaim up to lease_expiry_secs later. A RegistrationGuard scoped to the registered lifetime now owns teardown and is its ONLY caller here, so the normal and unwind paths cannot diverge. Idempotency (relied on because the sweeper runs the same deregister + fail_instance pair) is verified, not assumed: deregister is handle-keyed and returns None for an absent entry, and fail_instance releases capacity only for entries it actually removes. Tests: 71 lib + 6 e2e → 82 lib + 9 e2e; clippy -D warnings clean, fmt 0. Touched existing assertions, each an extension rather than a weakening: - every DelegateResultParams literal and the result_of helper gained the now required `admission` field, sourced from the live entry via a new `token()` test helper or from the ack — mechanical consequence of the required field. - happy_path_roundtrip additionally asserts the ack, the forwarded frame and the terminal frame all carry the same token (added coverage). - peek_for_completion's test helper passes the live token (new parameter). - generation_exhaustion_fails_closed_instead_of_wrapping → admission_token_exhaustion_fails_closed_instead_of_wrapping: same fail-closed property, now set on a per-namespace counter, plus a new assertion that another namespace is unaffected and that the refused admission never advances the counter. - the stalled-initiator server test now exhausts the initiator's BYTE budget instead of a 1-entry channel — same contract, exercising the new bound. - the e2e stalled-writer test echoes the token from each forward (required field) and asserts the forward carries one. Audit follow-up: the stalled-writer e2e regression sets its byte budget above the ~48 MiB it enqueues, so byte-refusal backpressure cannot preempt the write-timeout path the test exists to prove (budget refusal has dedicated regressions of its own). --- crates/openab-cp/cp.toml.example | 26 +- crates/openab-cp/src/config.rs | 168 +++++- crates/openab-cp/src/proto.rs | 98 +++ crates/openab-cp/src/registry.rs | 240 +++++++- crates/openab-cp/src/router.rs | 789 ++++++++++++++++++++----- crates/openab-cp/src/server.rs | 261 +++++++- crates/openab-cp/tests/ws_lifecycle.rs | 257 +++++++- docs/adr/agent-control-plane.md | 109 +++- docs/control-plane.md | 30 +- 9 files changed, 1786 insertions(+), 192 deletions(-) diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index 92dec8df5..9ae1c1642 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -50,6 +50,30 @@ max_connections_per_identity = 8 # draining a large result, short enough that a dead one frees its quota. write_timeout_secs = 30 +# Memory ceiling for ONE connection's outbound queue, in bytes. The queue is +# also bounded in entries, but entries are not a memory bound: frame sizes are +# configured independently, so a queue full of maximum-size frames is +# entry-legal and megabytes-to-gigabytes large. Bytes are reserved before a +# frame is enqueued and released when it is dequeued or when the connection is +# torn down; a frame that would exceed the budget is refused exactly like a full +# queue (the peer is treated as disconnected, never buffered). Must exceed +# max_frame_bytes so one maximum-size frame always fits. +max_outbound_queue_bytes = 16777216 + +# Process-wide ceiling on simultaneously in-flight delegations, enforced at +# admission (before any capacity is reserved) with SATURATED. Per-target +# max_delegated_sessions is advertised by the runtime and bounds one target's +# concurrency; it does not bound the CP's own state, and the in-flight table +# retains identity and ancestry per live admission. +max_inflight_delegations = 4096 + +# Default clamp on a runtime's advertised max_delegated_sessions, applied to +# every identity that does not set its own max_delegated_sessions_cap below. +# The advertised value is self-asserted and saturation is the CP's only +# backpressure signal, so there is no uncapped path: a worker advertising an +# enormous budget would otherwise disable saturation entirely. +default_max_delegated_sessions_cap = 16 + [[agents]] key = "${CP_KEY_KOUDU}" # per-agent secret, never shared namespace = "prod" @@ -61,7 +85,7 @@ key = "${CP_KEY_WORKER1}" namespace = "prod" name = "worker-1" type = "worker" -max_delegated_sessions_cap = 4 # CP-side clamp on advertised capacity +max_delegated_sessions_cap = 4 # overrides default_max_delegated_sessions_cap # Per-namespace policy. Absent namespaces use the conservative defaults: # max_depth = 1, allow_worker_initiation = false. diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index 02b492273..4385eaeb6 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -91,6 +91,44 @@ pub struct CpConfig { #[serde(default = "default_write_timeout_secs")] pub write_timeout_secs: u64, + /// Memory ceiling for ONE connection's outbound queue, in bytes. + /// + /// The queue is bounded in entries as well, but entries are not a memory + /// bound: frame sizes are independently configurable, so 256 queued frames + /// of `max_frame_bytes` each is entry-legal and megabytes-to-gigabytes + /// large. Bytes are reserved before a frame is enqueued and released when + /// it is dequeued (or when the connection is torn down); a frame that would + /// exceed the budget is refused exactly like a full queue — the peer is + /// treated as disconnected, never buffered. + /// + /// `write_timeout_secs` bounds how long a stalled writer *lives*; this + /// bounds how much it can have accumulated in the meantime. + #[serde(default = "default_max_outbound_queue_bytes")] + pub max_outbound_queue_bytes: usize, + + /// Process-wide ceiling on simultaneously in-flight delegations, enforced + /// at admission (before any capacity is reserved) with `SATURATED`. + /// + /// Per-target `max_delegated_sessions` is advertised by the runtime and + /// bounds one target's concurrency; it does not bound the CP's own state. + /// The in-flight table retains identity and ancestry per live admission, so + /// without a global ceiling an authenticated initiator with enough targets + /// can grow control-plane memory without limit. + #[serde(default = "default_max_inflight_delegations")] + pub max_inflight_delegations: usize, + + /// Default clamp on a runtime's advertised `max_delegated_sessions`, + /// applied to every identity that does not set its own + /// `max_delegated_sessions_cap`. + /// + /// The advertised value is self-asserted, and saturation is the CP's only + /// backpressure signal: an uncapped identity advertising a huge budget + /// disables it entirely. A per-identity cap overrides this default (up or + /// down), so raising a specific worker's ceiling stays a deliberate, + /// CP-side act. + #[serde(default = "default_max_delegated_sessions_cap")] + pub default_max_delegated_sessions_cap: u32, + /// Identity table: auth key → immutable claims. /// Keyed by the key id (`kid`), with the secret alongside, so logs can /// reference identities without printing secrets. @@ -132,6 +170,15 @@ fn default_max_connections_per_identity() -> u32 { fn default_write_timeout_secs() -> u64 { 30 } +fn default_max_outbound_queue_bytes() -> usize { + 16 * 1024 * 1024 +} +fn default_max_inflight_delegations() -> usize { + 4096 +} +fn default_max_delegated_sessions_cap() -> u32 { + 16 +} /// Immutable identity claims bound to one auth key. #[derive(Debug, Clone, Deserialize)] @@ -143,7 +190,11 @@ pub struct AgentIdentity { pub name: String, #[serde(rename = "type")] pub agent_type: AgentType, - /// Optional CP-side clamp on the advertised concurrency budget. + /// Optional CP-side clamp on the advertised concurrency budget. When + /// absent, `default_max_delegated_sessions_cap` applies — a runtime is + /// never trusted with an unclamped budget, because saturation is the CP's + /// only backpressure signal. Setting it here overrides the default in + /// either direction. #[serde(default)] pub max_delegated_sessions_cap: Option, } @@ -223,6 +274,26 @@ impl CpConfig { if self.write_timeout_secs == 0 { bail!("write_timeout_secs must be greater than 0"); } + // A budget that cannot hold one maximum-size frame would refuse every + // enqueue and disconnect every peer the moment a large frame is routed + // to it. + if self.max_outbound_queue_bytes <= self.max_frame_bytes { + bail!( + "max_outbound_queue_bytes ({}) must exceed max_frame_bytes ({}) — \ + a queue that cannot hold one maximum-size frame refuses every peer", + self.max_outbound_queue_bytes, + self.max_frame_bytes + ); + } + // Zero would refuse every delegation the CP could ever route. + if self.max_inflight_delegations == 0 { + bail!("max_inflight_delegations must be at least 1"); + } + // Zero would clamp every uncapped identity to no capacity at all, so + // every delegation to it would be SATURATED. + if self.default_max_delegated_sessions_cap == 0 { + bail!("default_max_delegated_sessions_cap must be at least 1"); + } // Bearer keys over cleartext TCP must never reach an untrusted // network: non-loopback binds require the explicit override. if !self.allow_insecure_bind && !is_loopback(&self.listen) { @@ -254,6 +325,22 @@ impl CpConfig { pub fn policy_for(&self, namespace: &str) -> NamespacePolicy { self.namespaces.get(namespace).cloned().unwrap_or_default() } + + /// Effective concurrency budget for `identity` given what its runtime + /// advertised at registration. + /// + /// The advertised value is self-asserted, so it is always clamped: by the + /// identity's own `max_delegated_sessions_cap` when set, otherwise by + /// `default_max_delegated_sessions_cap`. There is no uncapped path — an + /// omitted per-identity cap used to mean "trust the runtime", which let a + /// buggy or compromised worker advertise a budget large enough to nullify + /// saturation as a backpressure signal. + pub fn effective_max_sessions(&self, identity: &AgentIdentity, advertised: u32) -> u32 { + let cap = identity + .max_delegated_sessions_cap + .unwrap_or(self.default_max_delegated_sessions_cap); + advertised.min(cap) + } } /// Whether a `host:port` bind address is loopback. @@ -413,6 +500,85 @@ lease_expiry_secs = 30 } } + #[test] + fn resource_bounds_default_and_are_validated() { + // The byte budget, the global in-flight ceiling, and the default + // session clamp all default to safe values, and every degenerate + // setting is rejected rather than silently disabling the guard. + let cfg: CpConfig = toml::from_str(base_toml()).unwrap(); + cfg.validate().unwrap(); + assert_eq!(cfg.max_outbound_queue_bytes, 16 * 1024 * 1024); + assert_eq!(cfg.max_inflight_delegations, 4096); + assert_eq!(cfg.default_max_delegated_sessions_cap, 16); + + let explicit: CpConfig = toml::from_str( + "max_outbound_queue_bytes = 2097152 +max_inflight_delegations = 32 +default_max_delegated_sessions_cap = 3", + ) + .unwrap(); + explicit.validate().unwrap(); + assert_eq!(explicit.max_outbound_queue_bytes, 2 * 1024 * 1024); + assert_eq!(explicit.max_inflight_delegations, 32); + assert_eq!(explicit.default_max_delegated_sessions_cap, 3); + + for bad in [ + "max_inflight_delegations = 0", + "default_max_delegated_sessions_cap = 0", + // Equal to max_frame_bytes is not enough: one maximum-size frame + // must fit, or the connection is unusable. + "max_outbound_queue_bytes = 1048576", + "max_outbound_queue_bytes = 4096\nmax_frame_bytes = 8192", + ] { + let cfg: CpConfig = toml::from_str(bad).unwrap(); + assert!(cfg.validate().is_err(), "{bad} must be rejected"); + } + } + + #[test] + fn advertised_capacity_is_always_clamped() { + // An identity WITHOUT its own cap is clamped by the global default — + // there is no uncapped path, so a worker advertising u32::MAX cannot + // nullify saturation as a backpressure signal. + let cfg: CpConfig = toml::from_str(base_toml()).unwrap(); + let uncapped = cfg.identity_for_key("k-primary").unwrap(); + assert!(uncapped.max_delegated_sessions_cap.is_none()); + assert_eq!( + cfg.effective_max_sessions(uncapped, u32::MAX), + 16, + "the global default clamps an identity with no cap of its own" + ); + // Below the clamp, the advertised value stands. + assert_eq!(cfg.effective_max_sessions(uncapped, 2), 2); + + // A per-identity cap overrides the default in either direction. + let capped = cfg.identity_for_key("k-worker").unwrap(); + assert_eq!(capped.max_delegated_sessions_cap, Some(2)); + assert_eq!(cfg.effective_max_sessions(capped, u32::MAX), 2); + + let raised: CpConfig = toml::from_str( + r#" +default_max_delegated_sessions_cap = 4 + +[[agents]] +key = "k-big" +namespace = "prod" +name = "fat-worker" +type = "worker" +max_delegated_sessions_cap = 64 +"#, + ) + .unwrap(); + raised.validate().unwrap(); + let big = raised.identity_for_key("k-big").unwrap(); + assert_eq!( + raised.effective_max_sessions(big, 32), + 32, + "a per-identity cap above the default raises the ceiling deliberately" + ); + assert_eq!(raised.effective_max_sessions(big, 128), 64); + } + #[test] fn env_expansion() { std::env::set_var("CP_TEST_KEY_XYZ", "sekrit"); diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs index 89dce87f9..6ebade73d 100644 --- a/crates/openab-cp/src/proto.rs +++ b/crates/openab-cp/src/proto.rs @@ -8,6 +8,12 @@ //! - Every request carries `jsonrpc: "2.0"` and a `u64` id; responses echo the //! id. Correlation of *delegations* (which span multiple request/response //! pairs across two connections) uses `delegation_id`, never the JSON-RPC id. +//! - `delegation_id` is client-supplied and REUSABLE (cancel-then-retry), so +//! correlation of one *admission* of that id — which terminal frame belongs +//! to which routing decision — uses the CP-minted +//! [`AdmissionToken`]: carried on the `cp/delegate` ack and forwarded frame, +//! echoed by the serving runtime in `cp/delegate_result` (required), and +//! stamped on every initiator-bound terminal frame. //! - The first frame on a new connection MUST be `cp/register`. Anything else //! is rejected with `NOT_REGISTERED` and the connection is closed. //! - Delegation ancestry (`chain`) is **CP-constructed**: callers supply only @@ -271,6 +277,10 @@ pub struct DelegateParams { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DelegateForward { pub delegation_id: String, + /// The CP-minted admission token for THIS admission of + /// `delegation_id` — see [`AdmissionToken`]. The serving runtime MUST + /// echo it in the matching `cp/delegate_result`. + pub admission: AdmissionToken, pub prompt: String, pub deadline: chrono::DateTime, /// Authenticated identity of the initiating agent (`namespace/name`). @@ -284,10 +294,44 @@ pub struct DelegateForward { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DelegateAck { pub delegation_id: String, + /// The CP-minted admission token for this admission — see + /// [`AdmissionToken`]. The initiator correlates terminal frames on it, + /// because `delegation_id` alone is reusable. + pub admission: AdmissionToken, /// The chosen serving instance's logical name (`namespace/name`). pub assigned_to: String, } +/// Protocol-visible identity of ONE admission of a `delegation_id`. +/// +/// `delegation_id` is client-supplied and deliberately reusable: +/// cancel-then-retry is an ordinary client pattern, and with a single replica +/// the retry routes to the same worker, so `(namespace, delegation_id)` + +/// serving instance is not a stable identity over time. Every frame that +/// belongs to a specific admission therefore carries this token: +/// +/// - `cp/delegate` ack → the initiator learns it; +/// - forwarded `cp/delegate` → the serving runtime learns it; +/// - `cp/delegate_result` → the serving runtime MUST echo it (required field); +/// the CP drops a result whose token is not the live admission's, so a late +/// result for a cancelled admission can never be delivered as, or commit, +/// the delegation that reused the id; +/// - every initiator-bound terminal frame, CP-synthesized `timeout` and +/// `target_disconnected` included, carries the token of the admission it +/// ends, so "first terminal frame wins" is keyed per admission rather than +/// per reusable id. +/// +/// The value is a per-namespace monotonic counter. Namespace-scoped on +/// purpose: a single global counter placed on the wire would disclose +/// cross-namespace delegation volume, the class of oracle the namespace-scoped +/// in-flight key exists to remove. Within a namespace the number is no more +/// than what `cp/list_agents` already shows that namespace about itself. +/// +/// `cp/cancel` does not carry a token yet; it lands with the runtime-client +/// slice (PR 3/4), where the serving side gains the state to disambiguate +/// cancellation targets. +pub type AdmissionToken = u64; + // --- cp/delegate_result --- #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -303,9 +347,19 @@ pub enum DelegationStatus { /// Params of `cp/delegate_result` — emitted by the serving **runtime** when /// the agent's turn ends (protocol-mandatory; never depends on the model), /// or synthesized by the CP on timeout/disconnect. +/// +/// `admission` is REQUIRED, not optional: it is the only thing that ties the +/// frame to one admission of a reusable `delegation_id`. A missing token is a +/// malformed frame (`INVALID_PARAMS`), never a wildcard — an optional token +/// would leave exactly the misdelivery path it exists to close. The wire is +/// pre-1.0 and every serving runtime in this stack learns the token from the +/// forwarded `cp/delegate`, so echoing it costs nothing. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DelegateResultParams { pub delegation_id: String, + /// Echo of [`DelegateForward::admission`]. On CP-synthesized terminal + /// frames the CP fills in the token of the admission it is ending. + pub admission: AdmissionToken, pub status: DelegationStatus, #[serde(skip_serializing_if = "Option::is_none")] pub result: Option, @@ -387,6 +441,50 @@ mod tests { ); } + #[test] + fn delegate_result_requires_the_admission_token() { + // The token is the only thing tying a result frame to ONE admission of + // a reusable delegation_id. A frame without it must be malformed, not + // a wildcard that matches whatever admission is live. + let without = serde_json::json!({ + "delegation_id": "d-1", + "status": "completed", + "result": "done" + }); + assert!(serde_json::from_value::(without).is_err()); + + let with = serde_json::json!({ + "delegation_id": "d-1", + "admission": 7, + "status": "completed", + "result": "done" + }); + let p: DelegateResultParams = serde_json::from_value(with).unwrap(); + assert_eq!(p.admission, 7); + // And it round-trips onto the wire (initiator-bound frames carry it). + let back = serde_json::to_value(&p).unwrap(); + assert_eq!(back["admission"], 7); + } + + #[test] + fn ack_and_forward_carry_the_admission_token() { + let ack = DelegateAck { + delegation_id: "d-1".into(), + admission: 3, + assigned_to: "prod/worker-1".into(), + }; + assert_eq!(serde_json::to_value(&ack).unwrap()["admission"], 3); + let fwd = DelegateForward { + delegation_id: "d-1".into(), + admission: 3, + prompt: "hi".into(), + deadline: chrono::Utc::now(), + from: "prod/koudu".into(), + chain: vec!["prod/koudu".into()], + }; + assert_eq!(serde_json::to_value(&fwd).unwrap()["admission"], 3); + } + #[test] fn incoming_message_distinguishes_request_and_response() { let req: JsonRpcMessage = diff --git a/crates/openab-cp/src/registry.rs b/crates/openab-cp/src/registry.rs index 3b0518460..7404bf0d4 100644 --- a/crates/openab-cp/src/registry.rs +++ b/crates/openab-cp/src/registry.rs @@ -7,7 +7,7 @@ //! instance and fails its in-flight delegations. use std::collections::BTreeMap; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -16,13 +16,165 @@ use tokio::sync::{mpsc, watch}; use crate::proto::AgentType; +/// Capacity of each per-connection outbound queue, in entries. Entries alone +/// are not a memory bound — see [`OutboundBudget`], which bounds the bytes. +pub const OUTBOUND_QUEUE: usize = 256; + +/// Byte budget shared by one connection's [`FrameTx`] clones and its +/// [`FrameRx`]. +/// +/// The entry-bounded queue does not bound memory: frame sizes are +/// independently configurable, so 256 queued frames is entry-legal and +/// arbitrarily large in bytes (a stalled initiator with an 8 MiB result budget +/// could hold ~2 GiB). Bytes are reserved before an enqueue and released on +/// dequeue or teardown, so a stalled peer is refused on whichever bound it hits +/// first. +#[derive(Debug)] +pub struct OutboundBudget { + reserved: AtomicUsize, + max: usize, +} + +impl OutboundBudget { + fn new(max: usize) -> Self { + Self { + reserved: AtomicUsize::new(0), + max, + } + } + + /// Reserve `n` bytes, or refuse. Refusal is not an error state: the caller + /// treats it exactly like a full queue (the peer is disconnected). + fn reserve(&self, n: usize) -> bool { + self.reserved + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |cur| { + match cur.checked_add(n) { + Some(next) if next <= self.max => Some(next), + _ => None, + } + }) + .is_ok() + } + + fn release(&self, n: usize) { + self.reserved.fetch_sub(n, Ordering::AcqRel); + } + + /// Bytes currently reserved (queued but not yet dequeued). + pub fn reserved_bytes(&self) -> usize { + self.reserved.load(Ordering::Acquire) + } + + pub fn max_bytes(&self) -> usize { + self.max + } +} + +/// Why an outbound frame was refused. Both outcomes mean the same thing to +/// callers — this peer cannot take the frame, so treat it as disconnected — +/// and are distinguished only for logs and tests. +#[derive(Debug, PartialEq, Eq)] +pub enum SendRefused { + /// The connection is gone (receiver dropped). + Closed, + /// The queue is full, in entries or in bytes. + Full, +} + /// Outbound frame sender for one WS connection (serialized JSON text). -/// Bounded: a peer that cannot drain its queue is disconnected rather than -/// growing CP memory. -pub type FrameTx = mpsc::Sender; +/// +/// Bounded twice — in entries by the channel and in bytes by +/// [`OutboundBudget`] — because a peer that cannot drain its queue must be +/// disconnected rather than grow CP memory. Cloneable: the registry hands +/// clones to every path that routes a frame to this connection. +#[derive(Clone, Debug)] +pub struct FrameTx { + tx: mpsc::Sender, + budget: Arc, +} -/// Capacity of each per-connection outbound queue. -pub const OUTBOUND_QUEUE: usize = 256; +impl FrameTx { + /// Enqueue one frame without blocking. Bytes are reserved BEFORE the + /// enqueue (and released again if the enqueue fails), so the budget can + /// never be observed lower than what the queue actually holds. + pub fn try_send(&self, text: String) -> Result<(), SendRefused> { + let n = text.len(); + if !self.budget.reserve(n) { + return Err(SendRefused::Full); + } + match self.tx.try_send(text) { + Ok(()) => Ok(()), + Err(mpsc::error::TrySendError::Full(_)) => { + self.budget.release(n); + Err(SendRefused::Full) + } + Err(mpsc::error::TrySendError::Closed(_)) => { + self.budget.release(n); + Err(SendRefused::Closed) + } + } + } + + /// The connection's byte budget (for observability and tests). + pub fn budget(&self) -> &Arc { + &self.budget + } +} + +/// Receiving half of one connection's outbound queue, owned by that +/// connection's task. Dequeueing releases the frame's byte reservation; +/// dropping it (connection teardown) releases whatever is still queued, so a +/// torn-down connection never leaves reservations behind. +pub struct FrameRx { + rx: mpsc::Receiver, + budget: Arc, +} + +impl FrameRx { + pub async fn recv(&mut self) -> Option { + let text = self.rx.recv().await?; + self.budget.release(text.len()); + Some(text) + } + + /// Non-blocking dequeue (tests and drain loops). + pub fn try_recv(&mut self) -> Result { + let text = self.rx.try_recv()?; + self.budget.release(text.len()); + Ok(text) + } + + /// Stop accepting frames (tests simulate a dead peer with this). + pub fn close(&mut self) { + self.rx.close(); + } +} + +impl Drop for FrameRx { + fn drop(&mut self) { + // Teardown release: drain what is still queued so any surviving + // sender clone sees the budget freed rather than permanently consumed + // by a connection that no longer exists. + self.rx.close(); + while let Ok(text) = self.rx.try_recv() { + self.budget.release(text.len()); + } + } +} + +/// Create one connection's outbound queue: bounded in entries by +/// [`OUTBOUND_QUEUE`] and in bytes by `max_bytes`. +pub fn outbound_channel(max_bytes: usize) -> (FrameTx, FrameRx) { + let (tx, rx) = mpsc::channel(OUTBOUND_QUEUE); + let budget = Arc::new(OutboundBudget::new(max_bytes)); + ( + FrameTx { + tx, + budget: Arc::clone(&budget), + }, + FrameRx { rx, budget }, + ) +} /// Shutdown signal for one WS connection, held by the registry so the CP can /// terminate a connection it no longer considers registered. Lease expiry must @@ -255,7 +407,7 @@ mod tests { use super::*; fn inst(ns: &str, name: &str, id: &str, max: u32) -> Instance { - let (tx, _rx) = mpsc::channel(OUTBOUND_QUEUE); + let (tx, _rx) = outbound_channel(1024 * 1024); Instance { handle: 0, // assigned by register() namespace: ns.into(), @@ -271,6 +423,80 @@ mod tests { } } + #[test] + fn outbound_queue_is_bounded_in_bytes_not_only_entries() { + // 256 entries is not a memory bound: with independently configurable + // frame sizes a stalled peer can hold entries × max-frame bytes. The + // byte budget must refuse well before the entry count is exhausted. + const BUDGET: usize = 64 * 1024; + const FRAME: usize = 8 * 1024; + let (tx, rx) = outbound_channel(BUDGET); + + let mut accepted = 0; + loop { + match tx.try_send("x".repeat(FRAME)) { + Ok(()) => accepted += 1, + Err(e) => { + assert_eq!(e, SendRefused::Full); + break; + } + } + assert!(accepted <= OUTBOUND_QUEUE, "entry bound would have won"); + } + assert_eq!( + accepted, + BUDGET / FRAME, + "the byte budget, not the entry count, must be the binding limit" + ); + assert!( + accepted < OUTBOUND_QUEUE, + "refusal must happen well before entry-count exhaustion ({accepted} < {OUTBOUND_QUEUE})" + ); + assert_eq!(tx.budget().reserved_bytes(), BUDGET); + + // A frame that cannot fit is refused without consuming budget, and a + // smaller one that does fit is still accepted (no sticky refusal). + drop(rx); + let (tx2, mut rx2) = outbound_channel(BUDGET); + assert!(tx2.try_send("y".repeat(BUDGET + 1)).is_err()); + assert_eq!( + tx2.budget().reserved_bytes(), + 0, + "a refused reservation must not leak budget" + ); + assert!(tx2.try_send("y".repeat(FRAME)).is_ok()); + assert_eq!(tx2.budget().reserved_bytes(), FRAME); + + // Dequeueing releases the reservation. + assert_eq!(rx2.try_recv().unwrap().len(), FRAME); + assert_eq!(tx2.budget().reserved_bytes(), 0); + } + + #[test] + fn teardown_releases_the_whole_outbound_reservation() { + // Dropping the receiving half is connection teardown: whatever is + // still queued must be released, or a surviving sender clone would see + // the budget permanently consumed by a connection that is gone. + const BUDGET: usize = 32 * 1024; + let (tx, rx) = outbound_channel(BUDGET); + let observer = Arc::clone(tx.budget()); + for _ in 0..4 { + tx.try_send("z".repeat(4 * 1024)).unwrap(); + } + assert_eq!(observer.reserved_bytes(), 16 * 1024); + + drop(rx); + assert_eq!( + observer.reserved_bytes(), + 0, + "teardown must release every queued frame's reservation" + ); + // The connection is gone, so further sends are refused as closed — + // and still leak no budget. + assert_eq!(tx.try_send("more".into()), Err(SendRefused::Closed)); + assert_eq!(observer.reserved_bytes(), 0); + } + #[test] fn select_by_name_and_namespace_isolation() { let r = Registry::new(); diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index 4e45e0289..1bf74afb6 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -22,15 +22,21 @@ //! Ending a delegation is a two-sided event: a frame goes out on the wire and //! CP state is committed. The commit is exact — it claims the one admission it //! delivered a result for (key + serving handle + [`InFlight::generation`]) or -//! nothing at all — so CP state stays consistent under any interleaving. +//! nothing at all — so CP state stays consistent under any interleaving. The +//! same admission stamp is protocol-visible as +//! [`crate::proto::AdmissionToken`]: the serving runtime echoes it in +//! `cp/delegate_result` and a frame naming a stale admission is dropped before +//! anything is delivered, so exactness does not stop at the CP boundary. //! //! The wire is a different matter: a `completed` result racing the deadline //! sweep's synthesized `timeout` can put TWO terminal frames on the wire for -//! one `delegation_id`. v1 resolves that by contract instead of CP-side -//! suppression (which would need per-id terminal state the CP deliberately -//! does not keep): **the first terminal frame for a `delegation_id` wins**, and -//! initiators MUST ignore later ones. See the v1 contract amendments in -//! `docs/adr/agent-control-plane.md`. +//! one admission. v1 resolves that by contract instead of CP-side suppression +//! (which would need per-id terminal state the CP deliberately does not +//! keep): **the first terminal frame for an admission token wins**, and +//! initiators MUST ignore later ones for that token. Because every +//! initiator-bound terminal frame carries the token, a frame for a superseded +//! admission is distinguishable from — and cannot mask — the live one. See the +//! v1 contract amendments in `docs/adr/agent-control-plane.md`. //! //! # Lock hierarchy //! @@ -56,7 +62,6 @@ //! does not participate in this hierarchy. use std::collections::BTreeMap; -use std::sync::atomic::{AtomicU64, Ordering}; use chrono::{DateTime, Utc}; use parking_lot::Mutex; @@ -99,7 +104,8 @@ pub struct InFlight { /// CP-constructed chain for THIS delegation (root first, ends with the /// initiator). Children extend it. pub chain: Vec, - /// CP-generated, never-reused admission stamp. `(namespace, + /// CP-generated, never-reused admission stamp — the value carried on the + /// wire as [`crate::proto::AdmissionToken`]. `(namespace, /// delegation_id)` is NOT a stable identity over time: the id is /// client-supplied, and cancel-then-retry — a natural client pattern — /// legitimately re-admits the same id, which with a single replica routes @@ -108,8 +114,17 @@ pub struct InFlight { /// later, unrelated admission wearing the same clothes (an ABA race). /// /// The generation makes that distinction total: it is minted once per - /// admission from a monotonic counter and never reused, so the commit - /// step claims the entry it actually delivered a result for, or nothing. + /// admission from a **per-namespace** monotonic counter and never reused + /// for that namespace, so the commit step claims the entry it actually + /// delivered a result for, or nothing. Because the value is on the wire, + /// the serving runtime echoes it in `cp/delegate_result` and the CP + /// refuses a result that does not name the live admission — the same + /// exactness, extended past the CP boundary. + /// + /// Per-namespace rather than global precisely because it is + /// protocol-visible: a global counter would let one namespace observe + /// another's delegation volume. Commit matching only needs never-reuse per + /// `(namespace, delegation_id)` key, which per-namespace counters give. pub generation: u64, } @@ -137,18 +152,26 @@ impl DelegationKey { pub struct Router { inflight: Mutex>, - /// Serializes the delegate admission sequence (duplicate check → target - /// selection → capacity reservation → in-flight insert) so concurrent - /// requests cannot double-admit one id or oversubscribe capacity: without - /// it, two racing delegates both see a free slot and both reserve it. - /// Delegation rates are LLM-scale; a coarse admission lock is simple and - /// more than sufficient. - admission: Mutex<()>, - /// Source of [`InFlight::generation`] stamps. Monotonic and never reset - /// (the table dies with the process, so a restart cannot collide with - /// anything still in flight): every admission gets a value no earlier or - /// later admission has worn. - next_generation: AtomicU64, + /// Serializes the delegate admission sequence (duplicate check → global + /// in-flight bound → target selection → capacity reservation → in-flight + /// insert) so concurrent requests cannot double-admit one id or + /// oversubscribe capacity: without it, two racing delegates both see a + /// free slot and both reserve it. Delegation rates are LLM-scale; a coarse + /// admission lock is simple and more than sufficient. + /// + /// The guarded value is the source of [`InFlight::generation`] stamps: + /// namespace → last minted generation. Keeping the counters *inside* the + /// admission lock makes it structurally impossible to mint a generation + /// outside the sequence that consumes it. Never reset (the table dies with + /// the process, so a restart cannot collide with anything still in + /// flight): every admission in a namespace gets a value no earlier or + /// later admission in that namespace has worn. + /// + /// Per-namespace, not global, because the value is now protocol-visible + /// (see [`crate::proto::AdmissionToken`]): a global counter on the wire + /// would disclose other namespaces' delegation volume. Commit matching + /// only requires never-reuse per `(namespace, delegation_id)` key. + admission: Mutex>, } pub enum DelegateOutcome { @@ -194,7 +217,8 @@ enum Claim { /// including its [`InFlight::generation`] stamp, or why no result can be /// delivered for it. A peek never removes anything. enum Peek { - /// The caller is the instance the delegation was routed to. + /// The caller is the instance the delegation was routed to, and the frame + /// names that instance's live admission. Serving(InFlight), /// The entry exists but another instance serves it. Left in place — a /// non-owner frame must never make the delegation momentarily invisible @@ -203,6 +227,18 @@ enum Peek { /// CP-side logs only; never disclosed to the caller. owner_handle: u64, }, + /// The entry exists and the caller serves it, but the frame echoes a + /// DIFFERENT admission token than the live one: a late result for an + /// admission that was cancelled/expired before the same `delegation_id` + /// was re-admitted. Left strictly in place — delivering this payload would + /// hand the initiator one admission's result as another's, and the commit + /// that followed would remove a delegation that is genuinely running. + StaleAdmission { + /// Token the frame echoed (CP-side logs only). + echoed: u64, + /// Token of the live admission (CP-side logs only). + live: u64, + }, /// No entry for `(namespace, delegation_id)`. Unknown, } @@ -267,8 +303,7 @@ impl Router { pub fn new() -> Self { Self { inflight: Mutex::new(BTreeMap::new()), - admission: Mutex::new(()), - next_generation: AtomicU64::new(0), + admission: Mutex::new(BTreeMap::new()), } } @@ -287,11 +322,13 @@ impl Router { ) -> DelegateOutcome { let now = Utc::now(); - // Admission is one atomic sequence: duplicate check, parent lookup, - // target selection, capacity reservation, and in-flight insertion all - // happen under this guard, so two racing delegates can neither - // double-admit an id nor both claim the last free slot. - let _admission = self.admission.lock(); + // Admission is one atomic sequence: duplicate check, global in-flight + // bound, parent lookup, target selection, capacity reservation, and + // in-flight insertion all happen under this guard, so two racing + // delegates can neither double-admit an id nor both claim the last + // free slot. The guard also owns the per-namespace generation + // counters, so a token cannot be minted outside this sequence. + let mut admission = self.admission.lock(); // Delegation identity is namespace-scoped: the same id in another // namespace is a different delegation, so it neither collides here @@ -304,6 +341,35 @@ impl Router { )); } + // Process-wide bound on live admissions, checked before any capacity + // is reserved and before a target is even selected (nothing to roll + // back on refusal). Per-target `max_delegated_sessions` is + // runtime-advertised, so it bounds one target's concurrency, not the + // CP's own memory: the in-flight table retains the chain and identity + // of every live admission. The bound is the table's own length rather + // than a separate counter, which is why no removal path has to + // remember to decrement it — commit, cancel, sweep, fail_instance and + // the send-failure rollback all remove the entry, and the count + // follows by construction. A parallel counter would be one refactor + // away from drifting, and a drifted global bound wedges the whole CP. + let live = self.inflight.lock().len(); + if live >= cfg.max_inflight_delegations { + warn!( + live, + max = cfg.max_inflight_delegations, + namespace = %from_namespace, + "global in-flight delegation bound reached — refusing admission" + ); + return DelegateOutcome::Rejected(ErrorObject::new( + codes::SATURATED, + format!( + "control plane is at its global in-flight limit \ + (max_inflight_delegations = {}); retry later", + cfg.max_inflight_delegations + ), + )); + } + // Selector sanity: exactly one of name/labels. let (sel_name, sel_labels) = (params.target.name.as_deref(), params.target.labels.as_ref()); if sel_name.is_some() == sel_labels.is_some() { @@ -372,12 +438,44 @@ impl Router { )); } - // Build the forward frame with the CP-stamped chain. + // Mint the admission token BEFORE anything is reserved or sent, so + // exhaustion cannot leave a reserved slot behind — and because the + // forwarded frame carries the token, the serving runtime learns which + // admission it is serving and can echo it back. + // + // Counters are per namespace and never wrap: an exhausted counter must + // NOT be bumped (wrapping would re-issue tokens from 0 and silently + // recreate the ABA this stamp exists to prevent), so exhaustion fails + // the admission and leaves the counter parked at the ceiling — every + // later admission in that namespace is refused too. Fail closed, + // permanently, with no wrapping path. (Unreachable in practice: one + // admission per nanosecond exhausts a u64 after ~584 years; this is + // insurance, not a path.) + let counter = admission.entry(from_namespace.to_string()).or_insert(0); + if *counter >= u64::MAX - 1 { + tracing::error!( + namespace = %from_namespace, + "admission token space exhausted for this namespace — refusing admission" + ); + // -32603 = JSON-RPC internal error; no protocol-specific code + // is warranted for a condition that cannot occur in a + // process's realistic lifetime. + return DelegateOutcome::Rejected(ErrorObject::new( + -32603, + "control plane admission token space exhausted; restart the CP", + )); + } + *counter += 1; + let generation = *counter; + + // Build the forward frame with the CP-stamped chain and admission + // token. let from_logical = format!("{from_namespace}/{from_name}"); let mut chain = parent_chain; chain.push(from_logical.clone()); let forward = DelegateForward { delegation_id: params.delegation_id.clone(), + admission: generation, prompt: params.prompt, deadline: params.deadline, from: from_logical.clone(), @@ -393,40 +491,6 @@ impl Router { // Reserve capacity and record the in-flight entry BEFORE sending, so // an immediately-arriving result finds it. Roll both back if the send // fails. - // - // Minted BEFORE the capacity reservation so exhaustion cannot leave a - // reserved slot behind. `fetch_add` on an exhausted counter would wrap - // and re-issue generation values, silently recreating the ABA this - // stamp exists to prevent — so exhaustion fails the admission instead. - // (Unreachable in practice: one admission per nanosecond exhausts a - // u64 after ~584 years; this is fail-closed insurance, not a path.) - // `fetch_update` rather than `fetch_add`: an exhausted counter must - // NOT be written (fetch_add would wrap the atomic itself and re-issue - // generations from 0). Refusal leaves the counter parked at the - // ceiling, so every later admission is refused too — fail closed, - // permanently, with no wrapping path. - let generation = - match self - .next_generation - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |g| { - if g >= u64::MAX - 1 { - None - } else { - Some(g + 1) - } - }) { - Ok(prev) => prev + 1, - Err(_) => { - tracing::error!("delegation generation space exhausted — refusing admission"); - // -32603 = JSON-RPC internal error; no protocol-specific code - // is warranted for a condition that cannot occur in a - // process's realistic lifetime. - return DelegateOutcome::Rejected(ErrorObject::new( - -32603, - "control plane generation space exhausted; restart the CP", - )); - } - }; registry.adjust_sessions(target.handle, 1); let entry = InFlight { namespace: from_namespace.to_string(), @@ -439,7 +503,8 @@ impl Router { chain, // Stamped under the admission lock, so every admission — including // a re-admission of an id that was just cancelled — is - // distinguishable from every other for the life of the process. + // distinguishable from every other in this namespace for the life + // of the process. generation, }; self.inflight.lock().insert(key.clone(), entry.clone()); @@ -471,6 +536,7 @@ impl Router { info!( delegation = %entry.delegation_id, + admission = entry.generation, from = %entry.from_logical, to = %entry.to_logical, chain = ?entry.chain, @@ -480,6 +546,7 @@ impl Router { DelegateOutcome::Accepted(DelegateAck { delegation_id: params.delegation_id, + admission: generation, assigned_to: target.logical_id(), }) } @@ -526,8 +593,16 @@ impl Router { } /// Phase 1 of a completion: snapshot the entry for `(namespace, - /// delegation_id)` and assert `serving_handle` is the instance it was - /// routed to — under one in-flight lock acquisition, removing nothing. + /// delegation_id)`, assert `serving_handle` is the instance it was routed + /// to, and assert the frame's echoed admission token names that live + /// admission — under one in-flight lock acquisition, removing nothing. + /// + /// The token check happens HERE, before the initiator-bound frame is + /// built, because delivery is the irreversible half: a stale result that + /// passed the peek would be handed to the initiator as the live + /// admission's terminal frame (and, being the first terminal frame for + /// that id, would mask the genuine one) even if the commit later declined + /// to touch anything. /// /// The returned [`InFlight`] carries the [`InFlight::generation`] the /// commit step must match, so delivery can happen outside the lock without @@ -537,14 +612,19 @@ impl Router { namespace: &str, delegation_id: &str, serving_handle: u64, + admission: u64, ) -> Peek { let key = DelegationKey::new(namespace, delegation_id); let g = self.inflight.lock(); match g.get(&key) { - Some(e) if e.to_handle == serving_handle => Peek::Serving(e.clone()), - Some(e) => Peek::Foreign { + Some(e) if e.to_handle != serving_handle => Peek::Foreign { owner_handle: e.to_handle, }, + Some(e) if e.generation != admission => Peek::StaleAdmission { + echoed: admission, + live: e.generation, + }, + Some(e) => Peek::Serving(e.clone()), None => Peek::Unknown, } } @@ -591,15 +671,24 @@ impl Router { /// The terminal result is the one frame that must never be silently /// dropped, so delivery happens in two phases: /// - /// 1. **Peek** — validate ownership under one in-flight lock acquisition - /// without removing the entry ([`Router::peek_for_completion`]), then - /// build and `try_send` the initiator-bound frame. + /// 1. **Peek** — validate ownership AND the echoed admission token under + /// one in-flight lock acquisition without removing the entry + /// ([`Router::peek_for_completion`]), then build and `try_send` the + /// initiator-bound frame. /// 2. **Commit** — only after the initiator's queue accepted the frame, /// end the delegation ([`Router::commit_completion`]): remove the /// entry and release the serving instance's capacity, but only if the /// live entry is still the very admission that was peeked (key + /// serving handle + [`InFlight::generation`]). /// + /// The token check in phase 1 is what extends admission exactness past the + /// CP boundary. Without it a late result for a cancelled admission A, + /// arriving after the same `delegation_id` was re-admitted as B to the same + /// worker, would peek B, be delivered to the initiator as B's terminal + /// frame, and then commit B (peek and commit both saw B, so B's own + /// generation matched) — releasing capacity B still occupies and leaving + /// B's genuine result to be dropped later as unknown. + /// /// If the initiator's bounded queue refuses the frame, the entry stays /// in flight and [`CompleteOutcome::InitiatorStalled`] tells the caller /// to treat the initiator as disconnected (per the bounded-queue @@ -616,9 +705,12 @@ impl Router { /// the wire for one `delegation_id` — a `completed` result racing the /// sweep's synthesized `timeout`, or two duplicate results both passing /// the peek. That is resolved by contract, not by CP-side suppression: - /// initiators MUST treat the FIRST terminal frame for a `delegation_id` as - /// authoritative and ignore later ones (see "first terminal frame wins" - /// in the ADR's v1 contract amendments). + /// initiators MUST treat the FIRST terminal frame for a given **admission + /// token** as authoritative and ignore later ones for that token (see + /// "first terminal frame wins" in the ADR's v1 contract amendments). Every + /// initiator-bound terminal frame carries the token of the admission it + /// ends — CP-synthesized `timeout` and `target_disconnected` included — so + /// a late frame for a superseded admission can never mask the live one. /// /// Only the instance the delegation was routed to may complete it; a /// non-owner frame can never make the delegation momentarily invisible @@ -645,30 +737,52 @@ impl Router { return CompleteOutcome::Dropped; } }; - let entry = - match self.peek_for_completion(&namespace, ¶ms.delegation_id, serving_handle) { - Peek::Serving(e) => e, - Peek::Foreign { owner_handle } => { - // Only the instance the delegation was routed to may - // complete it. The entry stays exactly where it is. - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - expected = owner_handle, - got = serving_handle, - "result from unexpected instance — dropped, delegation untouched" - ); - return CompleteOutcome::Dropped; - } - Peek::Unknown => { - warn!( - delegation = %params.delegation_id, - namespace = %namespace, - "result for unknown delegation (late arrival or CP restart) — dropped" - ); - return CompleteOutcome::Dropped; - } - }; + let entry = match self.peek_for_completion( + &namespace, + ¶ms.delegation_id, + serving_handle, + params.admission, + ) { + Peek::Serving(e) => e, + Peek::Foreign { owner_handle } => { + // Only the instance the delegation was routed to may + // complete it. The entry stays exactly where it is. + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + expected = owner_handle, + got = serving_handle, + "result from unexpected instance — dropped, delegation untouched" + ); + return CompleteOutcome::Dropped; + } + Peek::StaleAdmission { echoed, live } => { + // A late result for an admission that no longer exists, while + // the same id is live under a new one (cancel-then-retry). + // Logged distinctly for operators, but answered with the same + // generic ack as any other drop: a distinguishable reply would + // tell the serving side whether the id is currently + // re-admitted, which is exactly the kind of existence oracle + // the namespace-scoped keys and byte-identical cancel + // refusals removed. + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + echoed_admission = echoed, + live_admission = live, + "result echoes a stale admission token — dropped, live delegation untouched" + ); + return CompleteOutcome::Dropped; + } + Peek::Unknown => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "result for unknown delegation (late arrival or CP restart) — dropped" + ); + return CompleteOutcome::Dropped; + } + }; // Truncate oversized results (keep the head; delegation already // ran). The marker counts against the cap: the final value never @@ -855,6 +969,10 @@ impl Router { if let Some(init) = registry.get(e.from_handle) { let params = DelegateResultParams { delegation_id: e.delegation_id.clone(), + // The admission this frame ends — the initiator + // correlates terminal frames on the token, not on the + // reusable id. + admission: e.generation, status: DelegationStatus::TargetDisconnected, result: None, error: Some(format!("{} disconnected", e.to_logical)), @@ -911,6 +1029,10 @@ impl Router { if let Some(init) = registry.get(e.from_handle) { let params = DelegateResultParams { delegation_id: e.delegation_id.clone(), + // The admission that expired. A retry of the same id gets + // its own token, so this timeout can never be mistaken for + // the retry's terminal frame. + admission: e.generation, status: DelegationStatus::Timeout, result: None, error: Some("deadline exceeded".to_string()), @@ -971,10 +1093,9 @@ impl Default for Router { mod tests { use super::*; use crate::proto::TargetSelector; - use crate::registry::OUTBOUND_QUEUE; + use crate::registry::{outbound_channel, FrameRx}; use chrono::Duration; use std::time::Instant; - use tokio::sync::mpsc; fn cfg() -> CpConfig { toml::from_str( @@ -995,13 +1116,8 @@ type = "worker" .unwrap() } - fn instance( - ns: &str, - name: &str, - ty: AgentType, - max: u32, - ) -> (Instance, mpsc::Receiver) { - let (tx, rx) = mpsc::channel(OUTBOUND_QUEUE); + fn instance(ns: &str, name: &str, ty: AgentType, max: u32) -> (Instance, FrameRx) { + let (tx, rx) = outbound_channel(1024 * 1024); ( Instance { handle: 0, @@ -1020,6 +1136,18 @@ type = "worker" ) } + /// The live admission token for `(namespace, delegation_id)` — what a + /// well-behaved serving runtime echoes, since it learned it from the + /// forwarded `cp/delegate`. + fn token(router: &Router, namespace: &str, delegation_id: &str) -> u64 { + router + .inflight + .lock() + .get(&DelegationKey::new(namespace, delegation_id)) + .expect("delegation must be in flight to have a token") + .generation + } + fn delegate_params(id: &str, target: &str, secs: i64) -> DelegateParams { DelegateParams { delegation_id: id.into(), @@ -1039,8 +1167,8 @@ type = "worker" router: Router, h_primary: u64, h_worker: u64, - worker_rx: mpsc::Receiver, - primary_rx: mpsc::Receiver, + worker_rx: FrameRx, + primary_rx: FrameRx, } fn world() -> World { @@ -1073,6 +1201,320 @@ type = "worker" ) } + /// Unwrap an accepted admission, keeping its ack (and with it the + /// admission token the initiator learns). + fn accept(out: DelegateOutcome) -> DelegateAck { + match out { + DelegateOutcome::Accepted(a) => a, + DelegateOutcome::Rejected(e) => panic!("rejected ({}): {}", e.code, e.message), + } + } + + /// The `params.admission` of a JSON-RPC frame the CP put on the wire. + fn frame_admission(frame: &str) -> u64 { + let v: serde_json::Value = serde_json::from_str(frame).unwrap(); + v["params"]["admission"] + .as_u64() + .unwrap_or_else(|| panic!("frame carries no admission token: {frame}")) + } + + #[test] + fn late_result_for_a_superseded_admission_is_dropped() { + // The protocol-level half of the ABA. Cancel-then-retry re-admits the + // same client-supplied id and, with a single replica, to the SAME + // worker. A late `cp/delegate_result` for the cancelled admission A + // therefore matches the live admission B on everything the CP used to + // check — namespace, delegation_id, serving handle — so before the + // token existed A's payload was delivered to the initiator as B's + // terminal frame, and the commit then removed B (peek and commit both + // saw B, so B's own generation matched), releasing capacity B still + // occupied and leaving B's genuine result to be dropped as unknown. + let mut w = world(); // worker max_delegated_sessions = 1 + let a = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + assert_eq!( + frame_admission(&w.worker_rx.try_recv().unwrap()), + a.admission, + "the worker learns the token it must echo" + ); + + // A is cancelled; the id and the worker's slot are free again. + let cancel = CancelParams { + delegation_id: "d-1".into(), + reason: "changed my mind".into(), + }; + w.router + .cancel(&w.registry, w.h_primary, &cancel, 2) + .expect("the initiator may cancel"); + drain(&mut w); + + // B: the same id, re-admitted, routed to the same worker. + let b = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + assert_ne!( + a.admission, b.admission, + "each admission of a reusable id gets its own token" + ); + assert_eq!( + frame_admission(&w.worker_rx.try_recv().unwrap()), + b.admission + ); + + // The late result for A arrives, echoing A's token. + assert_eq!( + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", a.admission, "A's stale payload"), + 1024, + 3 + ), + CompleteOutcome::Dropped, + "a result naming a superseded admission must be dropped" + ); + assert!( + w.primary_rx.try_recv().is_err(), + "A's payload must never be delivered as B's result" + ); + assert_eq!(w.router.inflight_count(), 1, "B must remain in flight"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "B's capacity reservation must be intact" + ); + assert_eq!( + w.router.chain_of("prod", "d-1").as_deref(), + Some(&["prod/koudu".to_string()][..]), + "B must still be visible to its own result and to the sweep" + ); + + // B then completes normally with its own token. + assert_eq!( + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", b.admission, "B's genuine result"), + 1024, + 4 + ), + CompleteOutcome::Delivered { committed: true } + ); + let frame = w.primary_rx.try_recv().expect("initiator got B's result"); + assert!(frame.contains("B's genuine result")); + assert_eq!( + frame_admission(&frame), + b.admission, + "the terminal frame names the admission it ends" + ); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + } + + #[test] + fn terminal_frames_carry_the_token_of_the_admission_they_end() { + // "First terminal frame wins" is only usable if the initiator can tell + // WHICH admission a terminal frame ends: for a reusable id, a late + // frame for a superseded admission would otherwise mask the live one + // permanently. Every initiator-bound terminal frame therefore carries + // the token — CP-synthesized `timeout` and `target_disconnected` + // included, since both are built from the in-flight entry. + let mut w = world(); + let a = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + drain(&mut w); + + // A's terminal frame is the sweep's synthesized timeout. + let mut seq = 900u64; + let mut next = || { + seq += 1; + seq + }; + let swept = + w.router + .sweep_deadlines(&w.registry, Utc::now() + Duration::seconds(3600), &mut next); + let timeout = swept + .iter() + .map(|(_, f)| f.clone()) + .find(|f| f.contains("\"timeout\"")) + .expect("the sweep must synthesize a timeout for the initiator"); + assert_eq!( + frame_admission(&timeout), + a.admission, + "the synthesized timeout must name the admission that expired" + ); + drain(&mut w); + + // B reuses the id and gets its own token; its terminal frame is + // therefore distinguishable from A's. + let b = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + assert_ne!(a.admission, b.admission); + drain(&mut w); + assert_eq!( + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", b.admission, "B done"), + 1024, + 5 + ), + CompleteOutcome::Delivered { committed: true } + ); + let terminal_b = w.primary_rx.try_recv().unwrap(); + assert_eq!(frame_admission(&terminal_b), b.admission); + assert_ne!( + frame_admission(&timeout), + frame_admission(&terminal_b), + "A's and B's terminal frames must be distinguishable on the wire" + ); + + // The disconnect-synthesized terminal frame carries it too. + let c = accept(do_delegate(&w, delegate_params("d-2", "worker-1", 60))); + drain(&mut w); + w.registry.deregister(w.h_worker); + let frames = w.router.fail_instance(&w.registry, w.h_worker, &mut next); + assert_eq!(frames.len(), 1); + assert!(frames[0].1.contains("target_disconnected")); + assert_eq!( + frame_admission(&frames[0].1), + c.admission, + "target_disconnected must name the admission it ends" + ); + } + + #[test] + fn global_inflight_bound_refuses_and_is_released_by_every_removal_path() { + // Per-target `max_delegated_sessions` is runtime-advertised and bounds + // one target's concurrency, not the CP's own state: the in-flight table + // retains identity and ancestry per live admission. The global bound is + // the ceiling on that table — and because the bound IS the table's + // length, every path that removes an entry releases it by construction. + // Each removal path below is followed by a fresh admission that only + // succeeds if the slot came back. + let cfg: CpConfig = toml::from_str( + r#" +max_inflight_delegations = 1 + +[[agents]] +key = "kp" +namespace = "prod" +name = "koudu" +type = "primary" +"#, + ) + .unwrap(); + cfg.validate().unwrap(); + let registry = Registry::new(); + let router = Router::new(); + // Capacity 8 on the target, so the GLOBAL bound is the binding limit. + let (p, mut primary_rx) = instance("prod", "koudu", AgentType::Primary, 8); + let (wk, mut worker_rx) = instance("prod", "worker-1", AgentType::Worker, 8); + let hp = registry.register(p); + let hw = registry.register(wk); + let go = |id: &str| { + router.delegate( + &cfg, + ®istry, + "prod", + "koudu", + &AgentType::Primary, + hp, + delegate_params(id, "worker-1", 60), + 1, + ) + }; + let mut drain_all = || { + while worker_rx.try_recv().is_ok() {} + while primary_rx.try_recv().is_ok() {} + }; + + let first = accept(go("d-1")); + match go("d-2") { + DelegateOutcome::Rejected(e) => { + assert_eq!(e.code, codes::SATURATED); + assert!( + e.message.contains("max_inflight_delegations"), + "the refusal must name the bound it hit: {}", + e.message + ); + } + _ => panic!("the global in-flight bound must refuse the second admission"), + } + // Refusal reserves nothing. + assert_eq!(registry.get(hw).unwrap().active_sessions, 1); + drain_all(); + + // 1. commit (a delivered terminal result). + assert_eq!( + router.complete( + ®istry, + hw, + result_of("d-1", first.admission, "done"), + 1024, + 2 + ), + CompleteOutcome::Delivered { committed: true } + ); + assert_eq!(router.inflight_count(), 0); + let second = accept(go("d-2")); + drain_all(); + + // 2. cancel. + let cancel = CancelParams { + delegation_id: "d-2".into(), + reason: "no longer needed".into(), + }; + router.cancel(®istry, hp, &cancel, 3).expect("owned"); + assert_eq!(router.inflight_count(), 0); + assert_ne!(second.admission, accept(go("d-3")).admission); + drain_all(); + + // 3. deadline sweep. + let mut seq = 500u64; + let mut next = || { + seq += 1; + seq + }; + assert!(!router + .sweep_deadlines(®istry, Utc::now() + Duration::seconds(3600), &mut next) + .is_empty()); + assert_eq!(router.inflight_count(), 0); + accept(go("d-4")); + drain_all(); + + // 4. fail_instance (the initiator's own disconnect). + router.fail_instance(®istry, hp, &mut next); + assert_eq!(router.inflight_count(), 0); + accept(go("d-5")); + drain_all(); + + // 5. send-failure rollback. The worker's queue is closed, so the + // forward is refused and the admission rolls back — the global slot + // must come back with it, which a second target proves. + let (wk2, mut worker2_rx) = instance("prod", "worker-2", AgentType::Worker, 8); + registry.register(wk2); + router.fail_instance(®istry, hp, &mut next); + assert_eq!(router.inflight_count(), 0); + worker_rx.close(); + match go("d-6") { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::TARGET_DISCONNECTED), + _ => panic!("a closed target queue must fail the forward"), + } + assert_eq!(router.inflight_count(), 0, "the rollback removed the entry"); + match router.delegate( + &cfg, + ®istry, + "prod", + "koudu", + &AgentType::Primary, + hp, + delegate_params("d-7", "worker-2", 60), + 9, + ) { + DelegateOutcome::Accepted(_) => {} + DelegateOutcome::Rejected(e) => { + panic!("the rolled-back admission must free the global slot: {e:?}") + } + } + worker2_rx.try_recv().expect("worker-2 got the forward"); + } + #[test] fn happy_path_roundtrip() { let mut w = world(); @@ -1088,10 +1530,15 @@ type = "worker" assert_eq!(v["method"], "cp/delegate"); assert_eq!(v["params"]["from"], "prod/koudu"); assert_eq!(v["params"]["chain"], serde_json::json!(["prod/koudu"])); + // The serving runtime learns the admission token it must echo, and it + // is the same one the initiator was acked with. + assert_eq!(v["params"]["admission"], ack.admission); + assert_eq!(ack.admission, token(&w.router, "prod", "d-1")); assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 1); let result = DelegateResultParams { delegation_id: "d-1".into(), + admission: ack.admission, status: DelegationStatus::Completed, result: Some("done".into()), error: None, @@ -1102,6 +1549,9 @@ type = "worker" ); let frame = w.primary_rx.try_recv().unwrap(); assert!(frame.contains("\"completed\"")); + // The initiator-bound terminal frame names the admission it ends. + let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); + assert_eq!(v["params"]["admission"], ack.admission); assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); assert_eq!(w.router.inflight_count(), 0); } @@ -1118,6 +1568,7 @@ type = "worker" // Complete BEFORE draining the worker's queue — entry must exist. let result = DelegateResultParams { delegation_id: "d-1".into(), + admission: token(&w.router, "prod", "d-1"), status: DelegationStatus::Completed, result: Some("instant".into()), error: None, @@ -1229,14 +1680,20 @@ type = "worker" DelegateOutcome::Accepted(_) )); w.worker_rx.try_recv().unwrap(); + let tok = token(&w.router, "prod", "d-1"); // Fill the initiator's bounded queue so the result frame is refused. let initiator_tx = w.registry.get(w.h_primary).unwrap().tx; while initiator_tx.try_send("filler".into()).is_ok() {} assert_eq!( - w.router - .complete(&w.registry, w.h_worker, result_of("d-1", "late"), 1024, 2), + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", tok, "late"), + 1024, + 2 + ), CompleteOutcome::InitiatorStalled { initiator_handle: w.h_primary } @@ -1338,6 +1795,7 @@ type = "worker" )); let result = DelegateResultParams { delegation_id: "d-1".into(), + admission: token(&w.router, "prod", "d-1"), status: DelegationStatus::Completed, result: Some("spoofed".into()), error: None, @@ -1355,6 +1813,9 @@ type = "worker" let w = world(); let result = DelegateResultParams { delegation_id: "d-unknown".into(), + // Any token at all: with no entry for the id there is nothing to + // match against, so the frame is dropped as unknown. + admission: 1, status: DelegationStatus::Completed, result: None, error: None, @@ -1375,6 +1836,7 @@ type = "worker" w.worker_rx.try_recv().unwrap(); let result = DelegateResultParams { delegation_id: "d-1".into(), + admission: token(&w.router, "prod", "d-1"), status: DelegationStatus::Completed, result: Some("x".repeat(200)), error: None, @@ -1403,6 +1865,7 @@ type = "worker" w.worker_rx.try_recv().unwrap(); let result2 = DelegateResultParams { delegation_id: "d-2".into(), + admission: token(&w.router, "prod", "d-2"), status: DelegationStatus::Completed, result: Some("y".repeat(100)), error: None, @@ -1632,9 +2095,12 @@ allow_worker_initiation = true } } - fn result_of(id: &str, body: &str) -> DelegateResultParams { + /// A `cp/delegate_result` frame as a well-behaved serving runtime builds + /// it: echoing the admission token it was forwarded. + fn result_of(id: &str, admission: u64, body: &str) -> DelegateResultParams { DelegateResultParams { delegation_id: id.into(), + admission, status: DelegationStatus::Completed, result: Some(body.into()), error: None, @@ -1654,6 +2120,7 @@ allow_worker_initiation = true DelegateOutcome::Accepted(_) )); w.worker_rx.try_recv().unwrap(); + let tok = token(&w.router, "prod", "d-1"); if spoof_first { // h_primary is registered but is NOT the serving instance. @@ -1661,7 +2128,7 @@ allow_worker_initiation = true w.router.complete( &w.registry, w.h_primary, - result_of("d-1", "spoofed"), + result_of("d-1", tok, "spoofed"), 1024, 2 ), @@ -1678,7 +2145,7 @@ allow_worker_initiation = true w.router.complete( &w.registry, w.h_worker, - result_of("d-1", "genuine"), + result_of("d-1", tok, "genuine"), 1024, 3, ), @@ -1696,7 +2163,7 @@ allow_worker_initiation = true w.router.complete( &w.registry, w.h_primary, - result_of("d-1", "spoofed"), + result_of("d-1", tok, "spoofed"), 1024, 4 ), @@ -1720,6 +2187,7 @@ allow_worker_initiation = true do_delegate(&w, delegate_params("d-1", "worker-1", 60)), DelegateOutcome::Accepted(_) )); + let tok = token(&w.router, "prod", "d-1"); let gate = std::sync::Barrier::new(2); let (spoofed, genuine) = std::thread::scope(|s| { let spoof = s.spawn(|| { @@ -1728,7 +2196,7 @@ allow_worker_initiation = true w.router.complete( &w.registry, w.h_primary, - result_of("d-1", "spoofed"), + result_of("d-1", tok, "spoofed"), 1024, 2, ) == CompleteOutcome::Delivered { committed: true } @@ -1737,7 +2205,7 @@ allow_worker_initiation = true let genuine = w.router.complete( &w.registry, w.h_worker, - result_of("d-1", "genuine"), + result_of("d-1", tok, "genuine"), 1024, 3, ) == CompleteOutcome::Delivered { committed: true }; @@ -1927,7 +2395,13 @@ allow_worker_initiation = true // Results route to the initiator of the SAME namespace only. assert_eq!( - router.complete(®istry, hw_dev, result_of("d-1", "dev-done"), 1024, 12), + router.complete( + ®istry, + hw_dev, + result_of("d-1", token(&router, "dev", "d-1"), "dev-done"), + 1024, + 12 + ), CompleteOutcome::Delivered { committed: true } ); let frame = dev_init_rx.try_recv().unwrap(); @@ -1942,7 +2416,13 @@ allow_worker_initiation = true ); assert_eq!( - router.complete(®istry, hw_prod, result_of("d-1", "prod-done"), 1024, 13), + router.complete( + ®istry, + hw_prod, + result_of("d-1", token(&router, "prod", "d-1"), "prod-done"), + 1024, + 13 + ), CompleteOutcome::Delivered { committed: true } ); let frame = prod_init_rx.try_recv().unwrap(); @@ -2040,7 +2520,8 @@ allow_worker_initiation = true /// operation and then drive the commit step directly — deterministic, /// no barrier timing. fn peek(w: &World, id: &str) -> InFlight { - match w.router.peek_for_completion("prod", id, w.h_worker) { + let live = token(&w.router, "prod", id); + match w.router.peek_for_completion("prod", id, w.h_worker, live) { Peek::Serving(e) => e, _ => panic!("{id} must be in flight and served by the worker"), } @@ -2171,7 +2652,7 @@ allow_worker_initiation = true w.router.complete( &w.registry, w.h_worker, - result_of("d-1", "genuine"), + result_of("d-1", b.generation, "genuine"), 1024, 7 ), @@ -2309,34 +2790,72 @@ allow_worker_initiation = true // Through the public path: the entry is gone, so the frame is not // even delivered — a peek that finds nothing is a plain drop. assert_eq!( - w.router - .complete(&w.registry, w.h_worker, result_of("d-1", "late"), 1024, 8), + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", a.generation, "late"), + 1024, + 8 + ), CompleteOutcome::Dropped ); assert!(w.primary_rx.try_recv().is_err()); } #[test] - fn generation_exhaustion_fails_closed_instead_of_wrapping() { - // fetch_add on an exhausted counter would wrap and re-issue - // generations, silently recreating the ABA the stamp prevents. - // Unreachable in a realistic process lifetime; pinned here so a - // refactor cannot quietly downgrade it to wrapping arithmetic. + fn admission_token_exhaustion_fails_closed_instead_of_wrapping() { + // Wrapping an exhausted counter would re-issue tokens and silently + // recreate the ABA the stamp prevents. Unreachable in a realistic + // process lifetime; pinned here so a refactor cannot quietly downgrade + // it to wrapping arithmetic. The counters are per namespace now, so + // the ceiling is set on the namespace under test — and exhaustion in + // one namespace must not refuse another's admissions. let w = world(); w.router - .next_generation - .store(u64::MAX - 1, Ordering::Relaxed); - // At the ceiling, fetch_update declines to store: the admission is - // refused and the counter stays parked at MAX-1 forever. + .admission + .lock() + .insert("prod".to_string(), u64::MAX - 1); + // At the ceiling the counter is NOT bumped: the admission is refused + // and the counter stays parked at MAX-1 forever. let out = do_delegate(&w, delegate_params("d-last", "worker-1", 60)); assert!( matches!(out, DelegateOutcome::Rejected(ref e) if e.code == -32603), - "exhausted generation space must refuse admission" + "exhausted admission token space must refuse admission" ); // No capacity was reserved by the refused admission. assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); // And it stays parked: the next attempt is refused too. let out2 = do_delegate(&w, delegate_params("d-next", "worker-1", 60)); assert!(matches!(out2, DelegateOutcome::Rejected(ref e) if e.code == -32603)); + assert_eq!( + *w.router.admission.lock().get("prod").unwrap(), + u64::MAX - 1, + "a refused admission must not advance (or wrap) the counter" + ); + + // A different namespace is unaffected — counters are independent. + let (p_dev, _dev_init_rx) = instance("dev", "koudu", AgentType::Primary, 4); + let (w_dev, mut dev_rx) = instance("dev", "worker-1", AgentType::Worker, 1); + let hp_dev = w.registry.register(p_dev); + w.registry.register(w_dev); + match w.router.delegate( + &w.cfg, + &w.registry, + "dev", + "koudu", + &AgentType::Primary, + hp_dev, + delegate_params("d-dev", "worker-1", 60), + 9, + ) { + DelegateOutcome::Accepted(ack) => assert_eq!( + ack.admission, 1, + "each namespace starts its own token sequence at 1" + ), + DelegateOutcome::Rejected(e) => { + panic!("dev must be unaffected by prod exhaustion: {}", e.message) + } + } + dev_rx.try_recv().unwrap(); } } diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index 77cbcbe22..c4ef4f76c 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -36,7 +36,7 @@ use axum::routing::get; use axum::Router as AxumRouter; use futures_util::{SinkExt, StreamExt}; use parking_lot::Mutex; -use tokio::sync::{mpsc, watch}; +use tokio::sync::watch; use tracing::{info, warn}; use crate::config::{AgentIdentity, CpConfig}; @@ -45,7 +45,7 @@ use crate::proto::{ JsonRpcErrorResponse, JsonRpcMessage, JsonRpcResponse, RegisterAck, RegisterParams, PROTOCOL_VERSION, }; -use crate::registry::{shutdown_signal, Instance, Registry, OUTBOUND_QUEUE}; +use crate::registry::{outbound_channel, shutdown_signal, Instance, Registry}; use crate::router::{CompleteOutcome, DelegateOutcome, Router}; pub struct AppState { @@ -310,14 +310,16 @@ async fn handle_connection( } }; - // Outbound channel for this connection. Bounded: a peer that - // cannot drain OUTBOUND_QUEUE frames is disconnected, not buffered. - let (tx, mut rx) = mpsc::channel::(OUTBOUND_QUEUE); + // Outbound channel for this connection. Bounded twice: a peer that cannot + // drain OUTBOUND_QUEUE frames — or that has accumulated + // max_outbound_queue_bytes of them — is disconnected, not buffered. + let (tx, mut rx) = outbound_channel(state.cfg.max_outbound_queue_bytes); - let effective_max = match identity.max_delegated_sessions_cap { - Some(cap) => reg.max_delegated_sessions.min(cap), - None => reg.max_delegated_sessions, - }; + // The advertised budget is self-asserted and therefore always clamped: + // by this identity's own cap when set, otherwise by the global default. + let effective_max = state + .cfg + .effective_max_sessions(&identity, reg.max_delegated_sessions); // The registry assigns the CP-generated handle: ownership // and teardown never key on the client-supplied instance_id. let handle = state.registry.register_conn( @@ -336,6 +338,18 @@ async fn handle_connection( }, Arc::clone(&shutdown), ); + // From here on, teardown is owned by an RAII guard rather than the return + // path: a panic anywhere below (the `expect("serializable")` sites are on + // production paths) would otherwise skip deregistration and leave this + // instance's in-flight delegations — and the capacity they reserve on + // OTHER instances — pinned until the lease expires. The guard runs on both + // the normal return and an unwind, and is the ONLY caller of `teardown` + // here, so the two paths cannot diverge. + let _registered = RegistrationGuard { + state: Arc::clone(&state), + handle, + identity: identity.clone(), + }; info!( agent = %format!("{}/{}", identity.namespace, identity.name), instance = %reg.instance_id, @@ -376,7 +390,7 @@ async fn handle_connection( ) .await; } - teardown(&state, handle, &identity); + // `_registered` runs teardown on the way out. return; } @@ -466,12 +480,49 @@ async fn handle_connection( .await; } - teardown(&state, handle, &identity); + // Teardown runs here, when `_registered` drops — on this path and on an + // unwind alike. +} + +/// RAII owner of a *registered* connection's CP-side state. +/// +/// Scoped to the registered lifetime: constructed immediately after +/// `register_conn`, dropped when the connection task returns **or unwinds**. +/// [`ConnPermit`] already made the identity's connection quota panic-safe; +/// this does the same for the registry entry and the in-flight rows, which a +/// panic between registration and return would otherwise leave for the lease +/// sweeper — up to `lease_expiry_secs` of capacity reserved on *other* +/// instances for delegations nobody is serving any more. +/// +/// [`teardown`] is idempotent by construction, which is what makes a guard +/// safe here: `deregister` is keyed by handle and returns `None` for an +/// already-removed entry, and `fail_instance` releases capacity only for +/// entries it actually removes. A guard that fires after `sweep_leases` +/// already tore this handle down therefore finds nothing and changes nothing +/// (proved by `teardown_runs_twice_without_double_releasing`). +struct RegistrationGuard { + state: Arc, + handle: u64, + identity: AgentIdentity, +} + +impl Drop for RegistrationGuard { + fn drop(&mut self) { + // Must not panic: a panic here during an unwind aborts the process. + // Everything it touches is lock-guarded map mutation and non-blocking + // sends — no `expect`, no allocation-dependent invariants. parking_lot + // locks are not poisoned and are released by the unwind itself, so a + // panic taken while holding one cannot deadlock this call. + teardown(&self.state, self.handle, &self.identity); + } } /// Deregister this connection's own registration (by handle — cannot touch /// another connection's entry) and fail its in-flight delegations. /// +/// Invoked from [`RegistrationGuard::drop`], so it runs on the normal return +/// path and on an unwind alike. +/// /// Deliberately idempotent with the sweeper: when `sweep_leases` already ran /// `deregister` + `fail_instance` for this handle, both calls here find /// nothing (the registry entry and the in-flight entries are gone) and are @@ -993,6 +1044,167 @@ mod tests { assert_eq!(state.conn_count("prod/worker-1"), 1); } + /// Register one instance on `state` with a fresh outbound queue. + fn register_test_instance( + state: &Arc, + name: &str, + agent_type: AgentType, + max_sessions: u32, + ) -> (u64, crate::registry::FrameRx) { + let (tx, rx) = outbound_channel(1024 * 1024); + let handle = state.registry.register_conn( + Instance { + handle: 0, + namespace: "prod".into(), + name: name.into(), + agent_type, + instance_id: format!("i-{name}"), + labels: Default::default(), + max_delegated_sessions: max_sessions, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + }, + crate::registry::shutdown_signal(), + ); + (handle, rx) + } + + /// Route one delegation through the wire-facing handler and return the + /// admission token from the ack. + fn delegate_through_handler(state: &Arc, from: u64, id: &str, target: &str) -> u64 { + let deadline = (chrono::Utc::now() + chrono::Duration::seconds(60)).to_rfc3339(); + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/delegate", + "params": { + "delegation_id": id, + "target": {"name": target}, + "prompt": "do it", + "deadline": deadline + } + }) + .to_string(); + let ack: serde_json::Value = + serde_json::from_str(&handle_frame(state, from, &frame).expect("answered")).unwrap(); + assert!( + ack.get("error").is_none(), + "delegation must be accepted: {ack}" + ); + ack["result"]["admission"] + .as_u64() + .expect("token on the ack") + } + + #[test] + fn registered_teardown_survives_a_panic() { + // Teardown used to run only on the connection task's normal return + // path. A panic after successful registration — reachable from the + // `expect("serializable")` sites on production paths — skipped it: the + // RAII `ConnPermit` still freed the identity's quota, but the registry + // entry and the in-flight rows survived until the lease swept them, + // keeping capacity reserved on OTHER instances for up to + // `lease_expiry_secs`. The Drop guard makes the unwind path identical + // to the normal one. + let state = state_with(""); + let (h_i, _rx_i) = register_test_instance(&state, "koudu", AgentType::Primary, 4); + let (h_w, mut rx_w) = register_test_instance(&state, "worker-1", AgentType::Worker, 1); + delegate_through_handler(&state, h_i, "d-1", "worker-1"); + rx_w.try_recv().expect("worker received the forward"); + assert_eq!(state.registry.get(h_w).unwrap().active_sessions, 1); + assert_eq!(state.router.inflight_count(), 1); + + // Panic inside the registered lifetime of the INITIATOR's connection. + let silent = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); + let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _registered = RegistrationGuard { + state: Arc::clone(&state), + handle: h_i, + identity: identity(), + }; + panic!("expect(\"serializable\") on a production path"); + })); + std::panic::set_hook(silent); + assert!(panicked.is_err(), "the test must actually unwind"); + + // Everything the guard owns is reclaimed immediately — not at lease + // expiry. + assert!( + state.registry.get(h_i).is_none(), + "the registry entry must be gone" + ); + assert_eq!( + state.router.inflight_count(), + 0, + "in-flight rows must be reclaimed" + ); + assert_eq!( + state.registry.get(h_w).unwrap().active_sessions, + 0, + "capacity reserved on the serving instance must be released" + ); + // ...and the serving runtime is told to stop working. + let cancel = rx_w.try_recv().expect("downstream cancel was queued"); + assert!(cancel.contains("cp/cancel") && cancel.contains("d-1")); + } + + #[test] + fn teardown_runs_twice_without_double_releasing() { + // The guard can fire on a handle that was already torn down — the lease + // sweeper runs the same `deregister` + `fail_instance` pair on its own + // initiative. A second pass must change nothing: `deregister` is keyed + // by handle and returns `None` for an absent entry, and `fail_instance` + // releases capacity only for entries it actually removes. A double + // release would be silent (session counts saturate) and would let + // `saturated()` admit work to a full instance. + let state = state_with(""); + let (h_i, _rx_i) = register_test_instance(&state, "koudu", AgentType::Primary, 4); + let (h_w, mut rx_w) = register_test_instance(&state, "worker-1", AgentType::Worker, 4); + // A second initiator keeps one delegation alive on the same worker, so + // an over-release shows up as a wrong count rather than a clamp at zero. + let (h_i2, _rx_i2) = register_test_instance(&state, "koudu-2", AgentType::Primary, 4); + delegate_through_handler(&state, h_i, "d-1", "worker-1"); + delegate_through_handler(&state, h_i2, "d-2", "worker-1"); + assert_eq!(state.registry.get(h_w).unwrap().active_sessions, 2); + while rx_w.try_recv().is_ok() {} + + let guard = || RegistrationGuard { + state: Arc::clone(&state), + handle: h_i, + identity: identity(), + }; + + // First pass. + drop(guard()); + assert!(state.registry.get(h_i).is_none()); + assert_eq!(state.router.inflight_count(), 1, "only d-1 was failed"); + assert_eq!( + state.registry.get(h_w).unwrap().active_sessions, + 1, + "exactly d-1's reservation was released" + ); + + // Second pass on the same handle: a no-op. + drop(guard()); + assert!(state.registry.get(h_i).is_none()); + assert_eq!(state.router.inflight_count(), 1); + assert_eq!( + state.registry.get(h_w).unwrap().active_sessions, + 1, + "a repeated teardown must not release capacity a second time" + ); + // And the sweeper's own pass over the same handle is equally inert. + sweep_leases(&state, Duration::ZERO); + assert_eq!( + state.router.inflight_count(), + 0, + "the sweeper expired the remaining leases" + ); + drop(guard()); + assert_eq!(state.router.inflight_count(), 0); + } + #[tokio::test] async fn sweep_leases_signals_the_connection_before_dropping_it() { // At the sweeper level: the shutdown signal is @@ -1001,7 +1213,7 @@ mod tests { let state = state_with("heartbeat_interval_secs = 1\nlease_expiry_secs = 2"); let signal = crate::registry::shutdown_signal(); let mut observer = signal.subscribe(); - let (tx, _rx) = mpsc::channel::(OUTBOUND_QUEUE); + let (tx, _rx) = outbound_channel(1024 * 1024); let handle = state.registry.register_conn( Instance { handle: 0, @@ -1042,7 +1254,7 @@ mod tests { // client needs to know it has to reconnect and register again. let state = state_with("heartbeat_interval_secs = 1\nlease_expiry_secs = 2"); let signal = crate::registry::shutdown_signal(); - let (tx, _rx) = mpsc::channel::(OUTBOUND_QUEUE); + let (tx, _rx) = outbound_channel(1024 * 1024); let handle = state.registry.register_conn( Instance { handle: 0, @@ -1106,10 +1318,13 @@ mod tests { // be closed (treated as disconnected) rather than silently skipped. let state = state_with(""); - // Initiator with a full single-slot queue. + // Initiator whose outbound byte budget one filler frame exhausts — + // the queue refuses on bytes here, which is the same + // "cannot drain → disconnected" contract as an exhausted entry count. + const FILLER: &str = "filler"; let signal_i = crate::registry::shutdown_signal(); let mut observer = signal_i.subscribe(); - let (tx_i, _rx_i) = mpsc::channel::(1); + let (tx_i, _rx_i) = outbound_channel(FILLER.len()); let h_i = state.registry.register_conn( Instance { handle: 0, @@ -1126,7 +1341,7 @@ mod tests { }, Arc::clone(&signal_i), ); - let (tx_w, mut rx_w) = mpsc::channel::(OUTBOUND_QUEUE); + let (tx_w, mut rx_w) = outbound_channel(1024 * 1024); let h_w = state.registry.register_conn( Instance { handle: 0, @@ -1159,13 +1374,21 @@ mod tests { let ack: serde_json::Value = serde_json::from_str(&handle_frame(&state, h_i, &del).expect("answered")).unwrap(); assert!(ack.get("error").is_none(), "delegation must be accepted"); + let admission = ack["result"]["admission"] + .as_u64() + .expect("the ack must carry the admission token"); rx_w.try_recv().expect("worker received the forward frame"); - // Fill the initiator's queue, then complete. - tx_i.try_send("filler".into()).unwrap(); + // Fill the initiator's queue to its byte budget, then complete. + tx_i.try_send(FILLER.into()).unwrap(); let res = serde_json::json!({ "jsonrpc": "2.0", "id": 2, "method": "cp/delegate_result", - "params": {"delegation_id": "d-1", "status": "completed", "result": "done"} + "params": { + "delegation_id": "d-1", + "admission": admission, + "status": "completed", + "result": "done" + } }) .to_string(); let reply: serde_json::Value = diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index 5d24cacd1..3ab800268 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -334,6 +334,241 @@ async fn register_worker(ws: &mut Ws, instance_id: &str, max_sessions: u32) -> s serde_json::from_str(msg.to_text().unwrap()).unwrap() } +/// Read the next JSON frame, or `None` if none arrives within `within`. +async fn next_json(ws: &mut Ws, within: Duration) -> Option { + let deadline = Instant::now() + within; + while Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(100), ws.next()).await { + Ok(Some(Ok(Message::Text(t)))) => return Some(serde_json::from_str(&t).unwrap()), + Ok(None) | Ok(Some(Err(_))) => return None, + Ok(Some(Ok(_))) => continue, + Err(_) => continue, + } + } + None +} + +/// Read frames until one satisfies `want`; panics on timeout. +async fn await_frame( + ws: &mut Ws, + what: &str, + want: impl Fn(&serde_json::Value) -> bool, +) -> serde_json::Value { + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline { + if let Some(v) = next_json(ws, Duration::from_millis(500)).await { + if want(&v) { + return v; + } + } + } + panic!("never received {what}"); +} + +#[tokio::test] +async fn advertised_capacity_is_clamped_by_the_global_default() { + // The advertised budget is self-asserted and saturation is the CP's only + // backpressure signal, so an identity with no cap of its own is clamped by + // `default_max_delegated_sessions_cap` — there is no uncapped path. + let (state, url) = spawn_cp(cfg( + "register_timeout_secs = 30\ndefault_max_delegated_sessions_cap = 5", + )) + .await; + let mut worker = connect_as(&url, KEY_WORKER).await.expect("worker accepted"); + let ack = register_worker(&mut worker, "i-greedy", u32::MAX).await; + assert_eq!( + ack["result"]["effective_max_delegated_sessions"], 5, + "an uncapped identity advertising u32::MAX must be clamped to the global default" + ); + assert_eq!( + state.registry.list("prod")[0].max_delegated_sessions, + 5, + "the clamped value is what the registry routes on" + ); +} + +#[tokio::test] +async fn a_late_result_for_a_superseded_admission_is_not_delivered() { + // Over a real pair of sockets: cancel-then-retry reuses the + // client-supplied `delegation_id` and, with a single replica, routes to the + // SAME worker. Before the admission token was on the wire, a late + // `cp/delegate_result` for the cancelled admission matched the live one on + // everything the CP checked — so its payload was delivered to the initiator + // as the live delegation's terminal frame (masking the genuine one, per + // "first terminal frame wins") and its commit removed the live entry. + let (state, url) = spawn_cp(cfg( + "register_timeout_secs = 30\nmax_connections_per_identity = 2", + )) + .await; + + let mut initiator = connect(&url).await.expect("initiator accepted"); + assert_eq!( + register(&mut initiator, "i-1").await["result"]["protocol_version"], + 1 + ); + let mut worker = connect_as(&url, KEY_WORKER).await.expect("worker accepted"); + assert_eq!( + register_worker(&mut worker, "i-w", 2).await["result"]["protocol_version"], + 1 + ); + + let deadline = (chrono::Utc::now() + chrono::Duration::seconds(600)).to_rfc3339(); + let delegate = |rpc_id: u64| { + serde_json::json!({ + "jsonrpc": "2.0", "id": rpc_id, "method": "cp/delegate", + "params": { + "delegation_id": "d-1", + "target": {"name": "worker-1"}, + "prompt": "do it", + "deadline": deadline + } + }) + .to_string() + }; + + // Admission A. + initiator + .send(Message::Text(delegate(10).into())) + .await + .unwrap(); + let ack_a = await_frame(&mut initiator, "ack for A", |v| v["id"] == 10).await; + let a = ack_a["result"]["admission"] + .as_u64() + .expect("the ack must carry the admission token"); + let fwd_a = await_frame(&mut worker, "forward for A", |v| { + v["method"] == "cp/delegate" + }) + .await; + assert_eq!( + fwd_a["params"]["admission"].as_u64(), + Some(a), + "the worker learns the token it must echo" + ); + + // The initiator cancels A, then retries the SAME id: admission B. + let cancel = serde_json::json!({ + "jsonrpc": "2.0", "id": 11, "method": "cp/cancel", + "params": {"delegation_id": "d-1", "reason": "changed my mind"} + }) + .to_string(); + initiator.send(Message::Text(cancel.into())).await.unwrap(); + assert_eq!( + await_frame(&mut initiator, "cancel ack", |v| v["id"] == 11).await["result"]["ok"], + true + ); + + initiator + .send(Message::Text(delegate(12).into())) + .await + .unwrap(); + let ack_b = await_frame(&mut initiator, "ack for B", |v| v["id"] == 12).await; + let b = ack_b["result"]["admission"].as_u64().expect("token on B"); + assert_ne!(a, b, "each admission of the reused id gets its own token"); + let fwd_b = await_frame(&mut worker, "forward for B", |v| { + v["method"] == "cp/delegate" && v["params"]["admission"].as_u64() == Some(b) + }) + .await; + assert_eq!(fwd_b["params"]["delegation_id"], "d-1"); + assert_eq!(state.router.inflight_count(), 1, "only B is in flight"); + + // The stale result for A arrives (a worker that had already computed it). + let stale = serde_json::json!({ + "jsonrpc": "2.0", "id": 20, "method": "cp/delegate_result", + "params": { + "delegation_id": "d-1", + "admission": a, + "status": "completed", + "result": "A's stale payload" + } + }) + .to_string(); + worker.send(Message::Text(stale.into())).await.unwrap(); + let ack_stale = await_frame(&mut worker, "ack for the stale result", |v| v["id"] == 20).await; + assert_eq!( + ack_stale["result"]["ok"], true, + "the drop must look exactly like any other generic ack — not a distinguishable \ + reply that would reveal whether the id is currently re-admitted" + ); + + // The initiator must see nothing: B is untouched and its capacity intact. + if let Some(v) = next_json(&mut initiator, Duration::from_millis(750)).await { + assert_ne!( + v["method"], "cp/delegate_result", + "A's payload must never be delivered as B's terminal frame: {v}" + ); + } + assert_eq!(state.router.inflight_count(), 1, "B must remain in flight"); + + // B's genuine result is then delivered, carrying B's token. + let genuine = serde_json::json!({ + "jsonrpc": "2.0", "id": 21, "method": "cp/delegate_result", + "params": { + "delegation_id": "d-1", + "admission": b, + "status": "completed", + "result": "B's genuine result" + } + }) + .to_string(); + worker.send(Message::Text(genuine.into())).await.unwrap(); + let terminal = await_frame(&mut initiator, "B's terminal frame", |v| { + v["method"] == "cp/delegate_result" + }) + .await; + assert_eq!(terminal["params"]["admission"].as_u64(), Some(b)); + assert_eq!(terminal["params"]["result"], "B's genuine result"); + assert_eq!(state.router.inflight_count(), 0); +} + +#[tokio::test] +async fn a_result_without_an_admission_token_is_rejected_as_malformed() { + // The token is required, not optional: an absent token must be + // INVALID_PARAMS rather than a wildcard matching whatever is live. + let (state, url) = spawn_cp(cfg( + "register_timeout_secs = 30\nmax_connections_per_identity = 2", + )) + .await; + let mut initiator = connect(&url).await.expect("initiator accepted"); + register(&mut initiator, "i-1").await; + let mut worker = connect_as(&url, KEY_WORKER).await.expect("worker accepted"); + register_worker(&mut worker, "i-w", 2).await; + + let deadline = (chrono::Utc::now() + chrono::Duration::seconds(600)).to_rfc3339(); + let del = serde_json::json!({ + "jsonrpc": "2.0", "id": 10, "method": "cp/delegate", + "params": { + "delegation_id": "d-1", + "target": {"name": "worker-1"}, + "prompt": "do it", + "deadline": deadline + } + }) + .to_string(); + initiator.send(Message::Text(del.into())).await.unwrap(); + await_frame(&mut initiator, "ack", |v| v["id"] == 10).await; + await_frame(&mut worker, "forward", |v| v["method"] == "cp/delegate").await; + + let untokened = serde_json::json!({ + "jsonrpc": "2.0", "id": 30, "method": "cp/delegate_result", + "params": {"delegation_id": "d-1", "status": "completed", "result": "no token"} + }) + .to_string(); + worker.send(Message::Text(untokened.into())).await.unwrap(); + let reply = await_frame(&mut worker, "error for the untokened result", |v| { + v["id"] == 30 + }) + .await; + assert_eq!(reply["error"]["code"], -32602, "INVALID_PARAMS"); + assert_eq!( + state.router.inflight_count(), + 1, + "the delegation is untouched by a malformed result" + ); + if let Some(v) = next_json(&mut initiator, Duration::from_millis(500)).await { + assert_ne!(v["method"], "cp/delegate_result"); + } +} + #[tokio::test] async fn a_peer_that_stops_reading_is_disconnected_and_frees_its_quota() { // A bounded outbound queue does not bound the WRITER. `sink.send().await` @@ -358,7 +593,12 @@ async fn a_peer_that_stops_reading_is_disconnected_and_frees_its_quota() { register_timeout_secs = 30 write_timeout_secs = {WRITE_TIMEOUT_SECS} max_frame_bytes = 8388608 -max_result_bytes = 8388608" +max_result_bytes = 8388608 +# Byte budget deliberately ABOVE the ~48 MiB this test enqueues: this +# regression proves the WRITE-TIMEOUT recovery path, and a budget below the +# payload would trip byte-refusal backpressure first, silently changing what +# is being tested (byte-budget refusal has its own dedicated regressions). +max_outbound_queue_bytes = 67108864" ))) .await; @@ -398,8 +638,9 @@ max_result_bytes = 8388608" } // Collect ALL forwards on the worker side first, so every delegation is - // admitted before any large write can block the initiator's task. - let mut pending = Vec::new(); + // admitted before any large write can block the initiator's task. Each + // forward carries the admission token the result MUST echo. + let mut pending: Vec<(String, u64)> = Vec::new(); while pending.len() < ids.len() { let msg = tokio::time::timeout(Duration::from_secs(20), worker.next()) .await @@ -413,7 +654,12 @@ max_result_bytes = 8388608" if v["method"] != "cp/delegate" { continue; } - pending.push(v["params"]["delegation_id"].as_str().unwrap().to_string()); + pending.push(( + v["params"]["delegation_id"].as_str().unwrap().to_string(), + v["params"]["admission"] + .as_u64() + .expect("the forwarded frame must carry the admission token"), + )); } assert_eq!( state.router.inflight_count(), @@ -423,7 +669,7 @@ max_result_bytes = 8388608" // Answer all but `d-keep` with a large result. Each is forwarded to the // stalled initiator, whose socket buffers fill and whose write then blocks. - for (n, id) in pending.iter().enumerate() { + for (n, (id, admission)) in pending.iter().enumerate() { if id == "d-keep" { continue; } @@ -431,6 +677,7 @@ max_result_bytes = 8388608" "jsonrpc": "2.0", "id": 900 + n, "method": "cp/delegate_result", "params": { "delegation_id": id, + "admission": admission, "status": "completed", "result": "x".repeat(BIG) } diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 877e5625b..0eba0f0d0 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -209,6 +209,25 @@ complete until the serving runtime returns a structured result frame The serving **runtime** emits this frame when the agent's turn ends — result delivery never depends on the sub-agent model "remembering" to report. +The frame MUST echo the `admission` token the CP stamped on the forwarded +`cp/delegate` (see "Admissions carry a protocol-visible token" below): + +```json +{ + "method": "cp/delegate_result", + "params": { + "delegation_id": "d-01J...", + "admission": 42, + "status": "completed", + "result": "…" + } +} +``` + +`delegation_id` says *which delegation*; `admission` says *which admission of +it*, and the id is reusable. A frame without the token is rejected as +malformed, and one naming a superseded admission is dropped. + ### v1 contract amendments (from PR #1465 review) The first implementation (`crates/openab-cp`) freezes the following @@ -297,28 +316,76 @@ recovery semantics: the disconnect path: `cp/cancel` to the serving runtime and exactly one capacity release. Best-effort frames (`cp/cancel`, sweep-synthesized `timeout`) remain fire-and-forget: the propagated deadline is their backstop. -- **Admissions carry a generation; commits are exact.** Every admission is - stamped with a CP-generated, never-reused generation. The commit phase of a - completion removes an in-flight entry only when key, serving handle, AND - generation all match the admission it delivered a result for; anything else - is left strictly untouched, capacity included. `(namespace, delegation_id)` - is deliberately not a stable identity over time — the id is client-supplied - and cancel-then-retry is an ordinary client pattern, which with a single - replica re-admits the same id to the same worker — so without the generation - a stale commit could remove a *live* delegation's entry, making it invisible - to its own genuine result and to the deadline sweep while freeing a slot it - still occupies. -- **First terminal frame wins.** More than one terminal frame may reach an - initiator for one `delegation_id`: a `completed` result can race the deadline - sweep's synthesized `timeout`, and duplicate results are possible in the - window between delivery and commit. **Initiators MUST treat the first +- **Admissions carry a protocol-visible token; commits are exact.** Every + admission is stamped with a CP-minted, never-reused **admission token** + (`admission`, a per-namespace monotonic counter). `(namespace, + delegation_id)` is deliberately not a stable identity over time — the id is + client-supplied and cancel-then-retry is an ordinary client pattern, which + with a single replica re-admits the same id to the same worker — so the token + is what identifies one admission, and it travels the whole round trip: + - the `cp/delegate` ack carries it, so the initiator can correlate; + - the forwarded `cp/delegate` carries it, so the serving runtime learns it; + - `cp/delegate_result` MUST echo it. It is a **required** field: a missing + token is `INVALID_PARAMS`, never a wildcard; + - every initiator-bound terminal frame carries it, CP-synthesized `timeout` + and `target_disconnected` included (both are built from the in-flight + entry). + + The CP checks the echoed token *before* building the initiator-bound frame, + and the commit phase removes an in-flight entry only when key, serving + handle, AND token all match the admission it delivered a result for; anything + else is left strictly untouched, capacity included. Without the token on the + wire, a late result for a cancelled admission A — arriving after the same id + was re-admitted as B to the same worker — would be delivered to the initiator + as B's terminal frame and would then commit B (peek and commit both saw B), + releasing capacity B still occupies and leaving B's genuine result to be + dropped later as unknown. A stale-token result is dropped and answered with + the same generic ack as any other drop: a distinguishable reply would tell + the serving side whether an id is currently re-admitted, the class of oracle + namespace-scoped keys and byte-identical `cp/cancel` refusals removed. The + counter is per namespace for the same reason: a single global counter on the + wire would disclose other namespaces' delegation volume, while commit + matching only needs never-reuse per `(namespace, delegation_id)`. Exhaustion + fails closed (the admission is refused; the counter never wraps). + `cp/cancel` does not carry the token yet — it lands with the runtime-client + slice, where the serving side gains the state to disambiguate cancellation + targets. +- **First terminal frame per admission wins.** More than one terminal frame may + reach an initiator for one admission: a `completed` result can race the + deadline sweep's synthesized `timeout`, and duplicate results are possible in + the window between delivery and commit. **Initiators MUST treat the first terminal frame (`completed`, `failed`, `timeout`, `target_disconnected`) for - a `delegation_id` as authoritative and ignore every later terminal frame for - that id.** The CP does not suppress the later frames: doing so would require - per-id terminal state that a CP with no durable state deliberately does not - keep, and the initiator already correlates by `delegation_id`. CP-side state - is unaffected either way — the generation rule above makes the commit exact, - so exactly one path ever releases the capacity. + a given `admission` token as authoritative and ignore every later terminal + frame for that token.** Correlation is per admission, not per + `delegation_id`: keyed on the reusable id, a late frame for a superseded + admission would permanently mask the live admission's genuine terminal frame. + The CP does not suppress the later frames: doing so would require per-id + terminal state that a CP with no durable state deliberately does not keep. + CP-side state is unaffected either way — the token rule above makes the + commit exact, so exactly one path ever releases the capacity. +- **Global admission and memory bounds.** Beyond per-frame and per-connection + limits the CP bounds its own aggregate state: + `max_inflight_delegations` (default 4096) caps simultaneously in-flight + delegations process-wide, enforced at admission before any capacity is + reserved and refused with `SATURATED` — per-target + `max_delegated_sessions` is runtime-advertised and bounds one target's + concurrency, not the CP's memory; `max_outbound_queue_bytes` (default 16 MiB) + caps each connection's outbound queue in **bytes** as well as entries, since + 256 queued frames of a configurable frame size is not a memory bound; and + `default_max_delegated_sessions_cap` (default 16) clamps every advertised + capacity that has no per-identity cap, so a runtime can never advertise its + way out of saturation-based backpressure. +- **Teardown of a registered connection is panic-safe.** Deregistration, + failing the connection's in-flight delegations, and downstream cancellation + run from an RAII guard scoped to the registered lifetime, so they happen on + the normal return path and on an unwind alike (the `expect("serializable")` + sites on production paths make a panic reachable). Without it a panicking + connection task left its registry entry and in-flight rows — and the capacity + they reserve on *other* instances — for the lease sweeper to reclaim up to + `lease_expiry_secs` later. Teardown is idempotent by construction: + deregistration is keyed by handle and capacity is released only for in-flight + entries actually removed, so the guard and the sweeper can both run over the + same handle without double-releasing. - **Capacity release follows entry removal.** Whichever path removes an in-flight entry (result commit, cancel, deadline sweep, instance failure, or a failed forward's rollback) releases its capacity reservation — and diff --git a/docs/control-plane.md b/docs/control-plane.md index 31c211c70..704913624 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -39,6 +39,10 @@ rationale. The essentials: - Heartbeats, lease expiry, registration deadline, per-identity connection quotas, the outbound write timeout, and frame/prompt/result size caps are all configurable with safe defaults. +- Aggregate bounds keep the CP itself bounded: `max_inflight_delegations` + (global live-admission ceiling), `max_outbound_queue_bytes` (per-connection + outbound memory), and `default_max_delegated_sessions_cap` (clamp on + runtime-advertised capacity for identities with no cap of their own). ## Health @@ -52,10 +56,30 @@ issue #1474). reconnect, re-authenticate, and re-register. - A peer that stops reading is disconnected: any single outbound write that blocks longer than `write_timeout_secs` is treated as a dead peer, so keep - draining the socket even while busy. -- **The first terminal frame for a `delegation_id` wins.** A `completed` + draining the socket even while busy. The same rule applies to the queue + behind it: a connection whose outbound queue exceeds + `max_outbound_queue_bytes` (or its entry count) is disconnected rather than + buffered. +- **Echo the `admission` token.** The `cp/delegate` ack and the forwarded + `cp/delegate` both carry an `admission` token identifying that one admission + of a `delegation_id`. A serving runtime MUST copy it into the matching + `cp/delegate_result`; the field is required, and a result naming a superseded + admission is dropped (the ack looks the same as any other, so do not treat + `ok: true` as proof of delivery — that is what the initiator's terminal frame + is for). +- **The first terminal frame for an `admission` token wins.** A `completed` result can race the CP's synthesized `timeout`, so an initiator may receive - more than one terminal frame for the same delegation. Treat the first as + more than one terminal frame for the same admission. Treat the first as authoritative and ignore later ones; the CP does not suppress them. + Correlate on `admission`, not on `delegation_id`: the id is yours to reuse + (cancel-then-retry is legal), and a late frame for the cancelled admission + would otherwise mask the retry's genuine result. Every terminal frame carries + the token, including CP-synthesized `timeout` and `target_disconnected`. +- A delegation may be refused with `SATURATED` because the target is at + capacity *or* because the CP is at `max_inflight_delegations`; the error + message says which. The CP never queues — retry later. +- The capacity a runtime advertises in `max_delegated_sessions` is clamped by + the CP (`default_max_delegated_sessions_cap`, or a per-identity override). + The ack's `effective_max_delegated_sessions` is the value that counts. - After a lease expires or the CP restarts, in-flight delegations are gone: initiators reconcile against their own deadlines and re-delegate. From fbc8f967dffe2dfb2249b73575549eb665b95d91 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Fri, 14 Aug 2026 00:53:31 -0400 Subject: [PATCH 09/11] =?UTF-8?q?fix(cp):=20round-9=20review=20=E2=80=94?= =?UTF-8?q?=20admission=20tokens=20on=20cancellation,=20ADR=20example=20fi?= =?UTF-8?q?x,=20wire-break=20callout?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 (blocking) — the admission token now covers cancellation, symmetrically with the result path. Round 8 stopped at results, leaving `cp/cancel` keyed on the reusable `delegation_id` alone, and that gap was reachable from two directions: - CP-synthesized, no client error required. `sweep_deadlines` removes an expired entry under the in-flight lock and builds its best-effort `cp/cancel` AFTER releasing it. In that gap the initiator can legitimately re-admit the same id; with a single replica admission B routes to the SAME worker and its forward is enqueued first. The worker then sees `forward(B)` followed by an id-only cancel for A — different producers into one queue, so per-connection ordering guarantees say nothing — and B's work is aborted at the source while the CP's table still shows B live, so B can only resolve by deadline. Result-token validation cannot repair this: the work is already gone. - Initiator-facing. An application-level retry of `cancel(A)` arriving after the re-admission matched B (`claim` compared key + `from_handle` only) and removed it: B's capacity was released while its work continued, and B's genuine result was later dropped as unknown with no synthesized terminal frame. An ordinary retry became a silent kill of the retry it was cleaning up after. Changes: - proto.rs: `admission: AdmissionToken` added to `CancelParams` as a REQUIRED field. Optional would be a "cancel whatever holds this id now" wildcard — exactly the misdelivery the token exists to close. `AdmissionToken`'s doc loses the "cp/cancel does not carry a token yet" deferral and documents the both-directions round trip instead. - router.rs: `claim` takes the token and matches (from_handle, namespace + delegation_id, admission) — all three under the ONE lock acquisition that removes the entry. New `Claim::StaleAdmission` leaves the live entry and its capacity strictly untouched and is refused BYTE-IDENTICALLY with the unknown-id and wrong-owner refusals: a distinguishable reply would tell a caller whether the id it reused is currently re-admitted, the oracle class round-3 F3 removed. - Every CP-synthesized cancel stamps the ended admission's generation. Grepped all 15 `CancelParams` construction sites: the two production ones are `fail_instance` (initiator-disconnect downstream cancels) and `sweep_deadlines`, and both build from the entry they removed. The stalled-initiator/backpressure teardown needs no separate change because it synthesizes nothing of its own — `server::teardown`, which the `RegistrationGuard` drives, and `sweep_leases` both route through `fail_instance`, and `grep methods::CANCEL src/` confirms those two plus the initiator forward are the only frame producers. The initiator forward passes the caller's params verbatim, which is correct: the token was just matched against the live entry. - `Claim`'s doc comment no longer lets its single-phase framing imply the cancel path needs no admission identity. The atomicity claim is kept and scoped: it is about this operation's consistency, not about the frame having been composed against state that may since have moved on. Both properties are required; the two-phase completion path needs the token for the additional reason that it has a peek-to-commit window of its own. Deterministic regressions (no barrier timing), both negative-controlled: - `swept_admissions_cancel_frame_cannot_target_a_reused_id` — expire A via the sweep, admit same-id/same-worker B in the gap, then assert A's outbound cancel AND timeout frames carry A's token and not B's, and that B is still in flight with its capacity intact. Mis-stamping the sweep's token fails it. - `retried_cancel_for_a_superseded_admission_never_removes_the_live_one` — the initiator retries cancel with A's token after B's re-admission: refused, byte-identical to the unknown-id refusal, B untouched (in flight, capacity held, live token unchanged, nothing forwarded downstream), and B then completes normally with its own token. Disabling the generation check fails it. - `cancel_refusals_are_byte_identical` extended (not weakened) with the stale-token case and with post-refusal assertions that the live token, the in-flight count, and the capacity reservation are all unchanged. Disabling the generation check fails it too. F2 — ADR §4's `cp/delegate` example showed a client-sent `chain`, a field `DelegateParams` has never had; serde ignored it silently, which is why the drift survived. The example now shows what an initiator actually sends (`parent_delegation_id`, no chain, no admission) and the CP-forwarded frame is documented separately with `chain` + `admission` + the authenticated `from`. A `cp/cancel` example is added, since that frame is now wire-relevant. BONUS: `adr_wire_examples_round_trip_through_the_serde_structs` makes the ADR examples part of the compiled contract. `include_str!` pulls the markdown in at compile time (no runtime fs, so the crate's test discipline holds), every ` ```json ` block must parse into the struct its `method` names, and the round trip must reproduce the exact key set — so a field the example has and the struct does not fails here. The initiator-sent arm asserts the absence of `chain` and `admission` explicitly, with the message a future drift would need. An unrecognized example shape panics rather than being skipped. Verified by reintroducing the drift: the test fails with "the initiator-sent cp/delegate example must not show a client-supplied `chain`". The `../../../docs` path does couple the test to the repo layout; that is the trade, and the doc comment says so — a fixture copied into the file would drift from the ADR the same way the ADR drifted from the structs. F3 — the required-field additions are wire-breaking. ADR §4's amendment block and docs/control-plane.md now carry an explicit ⚠️ callout: `admission` is required on `cp/delegate_result` (round 8) and on `cp/cancel` (here), there is no backward-compatible optional spelling, and a runtime built against the pre-token contract has EVERY result and EVERY cancel refused with INVALID_PARAMS after the CP is upgraded. The runtime-author guidance covers both roles and states that refusals are indistinguishable by design. The release-note paragraph for the PR description is at /tmp/1469-wire-break-note.md. Verification (isolated worktree on the build host): 85 lib + 9 e2e pass (baseline 82 + 9; +3 = two F1 regressions and the F2 drift test), clippy --all-targets -D warnings clean, fmt 0 diffs. Assertion changes are extensions only; the 13 test `CancelParams` literals gained the now-required field, sourced from the ack or the live entry via the existing `token()` helper. --- crates/openab-cp/src/proto.rs | 149 ++++++++++- crates/openab-cp/src/router.rs | 326 ++++++++++++++++++++++--- crates/openab-cp/tests/ws_lifecycle.rs | 4 +- docs/adr/agent-control-plane.md | 102 +++++++- docs/control-plane.md | 15 ++ 5 files changed, 552 insertions(+), 44 deletions(-) diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs index 6ebade73d..28bfdd2fa 100644 --- a/crates/openab-cp/src/proto.rs +++ b/crates/openab-cp/src/proto.rs @@ -12,8 +12,9 @@ //! correlation of one *admission* of that id — which terminal frame belongs //! to which routing decision — uses the CP-minted //! [`AdmissionToken`]: carried on the `cp/delegate` ack and forwarded frame, -//! echoed by the serving runtime in `cp/delegate_result` (required), and -//! stamped on every initiator-bound terminal frame. +//! echoed by the serving runtime in `cp/delegate_result` (required), carried +//! on `cp/cancel` in both directions (required), and stamped on every +//! initiator-bound terminal frame. //! - The first frame on a new connection MUST be `cp/register`. Anything else //! is rejected with `NOT_REGISTERED` and the connection is closed. //! - Delegation ancestry (`chain`) is **CP-constructed**: callers supply only @@ -319,17 +320,20 @@ pub struct DelegateAck { /// - every initiator-bound terminal frame, CP-synthesized `timeout` and /// `target_disconnected` included, carries the token of the admission it /// ends, so "first terminal frame wins" is keyed per admission rather than -/// per reusable id. +/// per reusable id; +/// - `cp/cancel` carries it in BOTH directions (required field). From the +/// initiator it names the admission to abort, so a retried cancel cannot +/// remove the re-admission that replaced its target; on every CP-synthesized +/// cancel the CP stamps the token of the admission it is ending, so a +/// best-effort cancel that overtakes a same-id re-admission's forward is +/// identifiable at the worker as belonging to the admission that is already +/// over. /// /// The value is a per-namespace monotonic counter. Namespace-scoped on /// purpose: a single global counter placed on the wire would disclose /// cross-namespace delegation volume, the class of oracle the namespace-scoped /// in-flight key exists to remove. Within a namespace the number is no more /// than what `cp/list_agents` already shows that namespace about itself. -/// -/// `cp/cancel` does not carry a token yet; it lands with the runtime-client -/// slice (PR 3/4), where the serving side gains the state to disambiguate -/// cancellation targets. pub type AdmissionToken = u64; // --- cp/delegate_result --- @@ -372,9 +376,24 @@ pub struct DelegateResultParams { /// Params of `cp/cancel`: from the initiator to abort an in-flight /// delegation, or from the CP to the serving runtime (best effort) after a /// timeout or initiator cancellation. +/// +/// `admission` is REQUIRED for the same reason it is required on +/// [`DelegateResultParams`], and the asymmetry would have been the hole: +/// `delegation_id` alone names a *slot* that cancel-then-retry legitimately +/// reuses, not the work being aborted. From the initiator the token is the +/// abort target, so a retried cancel is refused instead of removing the +/// re-admission that replaced its target. On CP-synthesized cancels the CP +/// stamps the token of the admission it is ending, so a best-effort cancel for +/// an expired admission that reaches the worker *after* a same-id +/// re-admission's forward names the finished admission rather than the live +/// one. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CancelParams { pub delegation_id: String, + /// The admission being cancelled — see [`AdmissionToken`]. Required: a + /// missing token is a malformed frame (`INVALID_PARAMS`), never a + /// "cancel whatever holds this id now" wildcard. + pub admission: AdmissionToken, pub reason: String, } @@ -392,6 +411,122 @@ pub mod methods { mod tests { use super::*; + /// Every ` ```json ` block in the ADR, in document order. + /// + /// `include_str!` is a compile-time read: the examples become part of this + /// crate's source, so the check costs no runtime filesystem access and + /// obeys the crate's no-fs-in-unit-tests rule. It does couple the test to + /// the repository layout — deliberately. `openab-cp` is a workspace binary + /// crate, never vendored on its own, and a fixture copied into this file + /// would drift from the ADR exactly the way the ADR drifted from the + /// structs. + fn adr_json_examples() -> Vec { + const ADR: &str = include_str!("../../../docs/adr/agent-control-plane.md"); + let mut out = Vec::new(); + let mut open: Option = None; + for line in ADR.lines() { + let line = line.trim_end(); + match (&mut open, line) { + (None, "```json") => open = Some(String::new()), + (Some(_), "```") => { + let body = open.take().expect("block is open"); + out.push(serde_json::from_str(&body).unwrap_or_else(|e| { + panic!("an ADR ```json block is not valid JSON: {e}\n{body}") + })); + } + (Some(buf), l) => { + buf.push_str(l); + buf.push('\n'); + } + (None, _) => {} + } + } + assert!(open.is_none(), "unterminated ```json block in the ADR"); + out + } + + fn sorted_keys(v: &Value) -> Vec { + let mut k: Vec = v + .as_object() + .unwrap_or_else(|| panic!("expected a JSON object, got {v}")) + .keys() + .cloned() + .collect(); + k.sort(); + k + } + + #[test] + fn adr_wire_examples_round_trip_through_the_serde_structs() { + // `docs/adr/agent-control-plane.md` is the authoritative wire contract + // and its examples are what a client implementer copies. They drifted + // once: the `cp/delegate` example showed a client-sent `chain`, a field + // `DelegateParams` has never had, and serde silently ignored it, so + // nothing failed. This test makes the examples part of the compiled + // contract — each block must parse into the struct its `method` names, + // and the round trip must reproduce the exact key set, so a field the + // example has and the struct does not (or the reverse) fails here. + // + // A new ADR example with a shape this dispatch does not recognize + // fails rather than being skipped: extend both together. + let examples = adr_json_examples(); + assert!( + examples.len() >= 4, + "expected the delegate, forwarded-delegate, result, and cancel \ + examples; found {}", + examples.len() + ); + for ex in &examples { + let method = ex["method"] + .as_str() + .unwrap_or_else(|| panic!("ADR example carries no method: {ex}")); + let params = &ex["params"]; + // `from` is stamped by the CP and appears only on the forwarded + // frame, so it distinguishes the two `cp/delegate` shapes. + let forwarded = params.get("from").is_some(); + let round_tripped = match (method, forwarded) { + (methods::DELEGATE, false) => { + // The finding, pinned: what an initiator sends carries + // neither the CP-constructed chain nor the CP-minted token. + assert!( + params.get("chain").is_none(), + "the initiator-sent cp/delegate example must not show a \ + client-supplied `chain` — the CP constructs it" + ); + assert!( + params.get("admission").is_none(), + "the initiator-sent cp/delegate example must not show an \ + `admission` — the CP mints it and returns it in the ack" + ); + let p: DelegateParams = serde_json::from_value(params.clone()) + .expect("cp/delegate example must parse as DelegateParams"); + serde_json::to_value(&p).expect("serializable") + } + (methods::DELEGATE, true) => { + let p: DelegateForward = serde_json::from_value(params.clone()) + .expect("forwarded cp/delegate example must parse as DelegateForward"); + serde_json::to_value(&p).expect("serializable") + } + (methods::DELEGATE_RESULT, _) => { + let p: DelegateResultParams = serde_json::from_value(params.clone()) + .expect("cp/delegate_result example must parse as DelegateResultParams"); + serde_json::to_value(&p).expect("serializable") + } + (methods::CANCEL, _) => { + let p: CancelParams = serde_json::from_value(params.clone()) + .expect("cp/cancel example must parse as CancelParams"); + serde_json::to_value(&p).expect("serializable") + } + (m, _) => panic!("the ADR documents a `{m}` example this test does not check"), + }; + assert_eq!( + sorted_keys(params), + sorted_keys(&round_tripped), + "the `{method}` example in the ADR has drifted from its struct" + ); + } + } + #[test] fn register_params_roundtrip_with_type_rename() { let json = serde_json::json!({ diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index 1bf74afb6..4567c5660 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -192,12 +192,21 @@ pub enum DelegateOutcome { /// either removed because the caller owns it, or never touched at all. /// /// Single-phase by construction — the caller acts on the returned entry -/// without going back to the table — so no window exists in which the id -/// could be re-admitted under the caller's feet. Contrast the two-phase -/// completion path, which needs [`InFlight::generation`] for exactly that -/// reason. `cp/cancel` is this helper's only caller. +/// without going back to the table — so no window exists in which the id could +/// be re-admitted under the caller's feet. That is a statement about *this +/// operation's atomicity*, and it is deliberately NOT a claim that the cancel +/// path needs no admission identity: the frame the caller sent was composed +/// against whatever it last observed, and by the time it arrives the id may +/// legitimately hold a different admission. Atomicity keeps the CP's table +/// consistent; the [`AdmissionToken`] the caller must name keeps the operation +/// aimed at the admission it meant. Both are required, and the two-phase +/// completion path needs the token for the additional reason that it has a +/// peek-to-commit window of its own. +/// +/// `cp/cancel` is this helper's only caller. enum Claim { - /// The caller initiated it; the entry has already been removed. + /// The caller initiated it and named its live admission; the entry has + /// already been removed. Owned(InFlight), /// The entry exists but was initiated by another instance. Left in place. WrongOwner { @@ -206,6 +215,19 @@ enum Claim { /// never disclosed to the caller). owner_handle: u64, }, + /// The caller initiated the entry that holds the id, but named a + /// DIFFERENT admission: its target was already ended (cancelled, swept, + /// or completed) and the id re-admitted. Left strictly in place — + /// removing it would abort work the caller never asked to abort, release + /// capacity the live admission still occupies, and orphan its result with + /// no synthesized terminal frame. + StaleAdmission { + namespace: String, + /// Token the frame named (CP-side logs only). + named: u64, + /// Token of the live admission (CP-side logs only). + live: u64, + }, /// No entry for `(namespace, delegation_id)`. NotFound { namespace: String }, /// The calling connection has no registration: it was swept (lease @@ -563,13 +585,22 @@ impl Router { } /// Look up `delegation_id` in the caller's namespace, assert the caller - /// initiated it, and remove the entry if so — all under one acquisition of - /// the in-flight lock (see [`Claim`]). + /// initiated it AND that `admission` names its live admission, and remove + /// the entry if so — all under one acquisition of the in-flight lock (see + /// [`Claim`]). /// /// The namespace is taken from the caller's authenticated registration, /// never from the frame, so a delegation id can only ever be resolved - /// inside the namespace of the connection that named it. - fn claim(&self, registry: &Registry, handle: u64, delegation_id: &str) -> Claim { + /// inside the namespace of the connection that named it. The token, by + /// contrast, is the caller's statement of intent: it says *which admission* + /// of that id is meant to end, which the reusable id cannot. + fn claim( + &self, + registry: &Registry, + handle: u64, + delegation_id: &str, + admission: u64, + ) -> Claim { // Registry lookup completes before the in-flight lock is taken; the // two locks are never held together (see the lock hierarchy above). let namespace = match registry.get(handle) { @@ -578,8 +609,8 @@ impl Router { }; let key = DelegationKey::new(&namespace, delegation_id); let mut g = self.inflight.lock(); - let owner_handle = match g.get(&key) { - Some(e) => e.from_handle, + let (owner_handle, live) = match g.get(&key) { + Some(e) => (e.from_handle, e.generation), None => return Claim::NotFound { namespace }, }; if owner_handle != handle { @@ -588,6 +619,13 @@ impl Router { owner_handle, }; } + if live != admission { + return Claim::StaleAdmission { + namespace, + named: admission, + live, + }; + } let entry = g.remove(&key).expect("present under the same lock"); Claim::Owned(entry) } @@ -875,15 +913,16 @@ impl Router { } /// Handle `cp/cancel` from the initiator. Returns the frame to forward - /// to the serving runtime, if the delegation is in flight and owned by - /// the caller. + /// to the serving runtime, if the delegation is in flight, owned by the + /// caller, and the caller named its live admission. /// - /// Ownership is validated under the same lock acquisition that removes the - /// entry (see [`Claim`] — no remove/reinsert window), and every refusal - /// returns ONE byte-identical error: an unknown id and another instance's - /// live id are indistinguishable to the caller, so `cp/cancel` cannot be - /// used to probe for delegation ids. The distinction is kept in the CP's - /// own logs only. + /// All three facts are established under the same lock acquisition that + /// removes the entry (see [`Claim`] — no remove/reinsert window), and every + /// refusal returns ONE byte-identical error: an unknown id, another + /// instance's live id, and the caller's own id under a superseded admission + /// are indistinguishable to the caller, so `cp/cancel` cannot be used to + /// probe for delegation ids or for whether an id is currently re-admitted. + /// The distinctions are kept in the CP's own logs only. pub fn cancel( &self, registry: &Registry, @@ -897,7 +936,12 @@ impl Router { "delegation is not in flight for this instance", ) }; - let entry = match self.claim(registry, from_handle, ¶ms.delegation_id) { + let entry = match self.claim( + registry, + from_handle, + ¶ms.delegation_id, + params.admission, + ) { Claim::Owned(entry) => entry, Claim::WrongOwner { namespace, @@ -912,6 +956,20 @@ impl Router { ); return Err(refused()); } + Claim::StaleAdmission { + namespace, + named, + live, + } => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + named_admission = named, + live_admission = live, + "cancel refused: names a superseded admission; the live one is untouched" + ); + return Err(refused()); + } Claim::NotFound { namespace } => { warn!( delegation = %params.delegation_id, @@ -929,9 +987,15 @@ impl Router { } }; registry.adjust_sessions(entry.to_handle, -1); - info!(delegation = %params.delegation_id, "delegation cancelled by initiator"); + info!( + delegation = %params.delegation_id, + admission = entry.generation, + "delegation cancelled by initiator" + ); let target = registry.get(entry.to_handle); Ok(target.map(|t| { + // Forwarded verbatim: the token was just matched against the live + // entry, so it names exactly the admission the worker is serving. let frame = JsonRpcRequest::new( next_rpc_id, methods::CANCEL, @@ -990,6 +1054,11 @@ impl Router { if let Some(target) = registry.get(e.to_handle) { let params = CancelParams { delegation_id: e.delegation_id.clone(), + // The admission this cancel ends. Built from the entry + // this loop removed, so if the id is re-admitted before + // this best-effort frame reaches the worker, the frame + // still names the admission that is over. + admission: e.generation, reason: format!("initiator {} disconnected", e.from_logical), }; let frame = JsonRpcRequest::new( @@ -1047,6 +1116,13 @@ impl Router { if let Some(target) = registry.get(e.to_handle) { let params = CancelParams { delegation_id: e.delegation_id.clone(), + // The admission that expired. The frame is built from the + // entry this sweep removed, and the removal happened under + // the in-flight lock while the send happens after it is + // released: a same-id retry can be admitted and forwarded + // in that gap, so the token is what keeps this cancel aimed + // at the expired admission instead of the live one. + admission: e.generation, reason: "deadline exceeded".to_string(), }; let frame = JsonRpcRequest::new( @@ -1240,6 +1316,7 @@ type = "worker" // A is cancelled; the id and the worker's slot are free again. let cancel = CancelParams { delegation_id: "d-1".into(), + admission: a.admission, reason: "changed my mind".into(), }; w.router @@ -1308,6 +1385,164 @@ type = "worker" assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); } + #[test] + fn swept_admissions_cancel_frame_cannot_target_a_reused_id() { + // The cancel-direction ABA, and the one that needs no client error at + // all. `sweep_deadlines` removes the expired entry under the in-flight + // lock and builds its best-effort `cp/cancel` AFTER releasing it. In + // that gap the initiator can legitimately re-admit the same id, and + // with a single replica admission B routes to the SAME worker and its + // forward is enqueued first. The worker then sees `forward(B)` followed + // by a cancel for the id — two different producers into one queue, so + // per-connection ordering guarantees say nothing. Keyed only on the + // reusable `delegation_id` that cancel terminates B at the source: the + // CP's table still shows B live, but its work is aborted and it can + // only resolve by deadline. The token is what makes the frame name the + // admission that is over. + let mut w = world(); // worker max_delegated_sessions = 1 + let a = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 1))); + assert_eq!( + frame_admission(&w.worker_rx.try_recv().unwrap()), + a.admission + ); + + let mut id = 500u64; + let mut next = || { + id += 1; + id + }; + let swept = + w.router + .sweep_deadlines(&w.registry, Utc::now() + Duration::seconds(120), &mut next); + assert_eq!(w.router.inflight_count(), 0, "A expired"); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + + // The gap: B is admitted and forwarded before the sweep's frames are + // sent. This is the interleaving the sweep cannot prevent, only survive. + let b = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + assert_ne!(a.admission, b.admission); + assert_eq!( + frame_admission(&w.worker_rx.try_recv().unwrap()), + b.admission + ); + + let cancel = swept + .iter() + .map(|(_, f)| f) + .find(|f| f.contains("cp/cancel")) + .expect("the sweep synthesizes a best-effort cancel"); + assert_eq!( + frame_admission(cancel), + a.admission, + "the swept cancel names the admission it ended, not the id" + ); + assert_ne!( + frame_admission(cancel), + b.admission, + "a worker matching on the token cannot mistake it for B" + ); + // The timeout frame is A's too, so the initiator does not mistake it + // for B's terminal frame either. + let timeout = swept + .iter() + .map(|(_, f)| f) + .find(|f| f.contains("cp/delegate_result")) + .expect("the sweep synthesizes a timeout result"); + assert_eq!(frame_admission(timeout), a.admission); + + // B is untouched by A's expiry: still in flight, capacity still held. + assert_eq!(w.router.inflight_count(), 1, "B must remain in flight"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "B's capacity reservation must be intact" + ); + assert_eq!(token(&w.router, "prod", "d-1"), b.admission); + } + + #[test] + fn retried_cancel_for_a_superseded_admission_never_removes_the_live_one() { + // The initiator-facing direction. An application-level retry of + // cancel(A) — an ordinary thing to do when the first attempt's ack was + // not observed — arrives after the same id was re-admitted as B. Keyed + // on `(namespace, delegation_id)` + `from_handle` alone it matched B + // and removed it: B's capacity was released while its work continued, + // and B's genuine result was later dropped as unknown, with no + // synthesized terminal frame because the entry was gone. So an ordinary + // retry became a silent kill of the retry it was cleaning up after. + let mut w = world(); + let a = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + let first = CancelParams { + delegation_id: "d-1".into(), + admission: a.admission, + reason: "changed my mind".into(), + }; + w.router + .cancel(&w.registry, w.h_primary, &first, 2) + .expect("the initiator may cancel its live admission"); + drain(&mut w); + + let b = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + assert_ne!(a.admission, b.admission); + assert_eq!( + frame_admission(&w.worker_rx.try_recv().unwrap()), + b.admission + ); + + // The retry: same initiator, same id, A's token. + let err = w + .router + .cancel(&w.registry, w.h_primary, &first, 3) + .expect_err("a cancel naming a superseded admission must be refused"); + assert_eq!(err.code, codes::POLICY_DENIED); + // Byte-identical to the unknown-id refusal: the retry learns nothing + // about whether its id was re-admitted. + let unknown = CancelParams { + delegation_id: "d-nope".into(), + admission: a.admission, + reason: "probe".into(), + }; + assert_eq!( + serde_json::to_string(&err).unwrap(), + serde_json::to_string( + &w.router + .cancel(&w.registry, w.h_primary, &unknown, 4) + .unwrap_err() + ) + .unwrap(), + ); + + // B untouched in every respect the refusal could have disturbed. + assert_eq!(w.router.inflight_count(), 1, "B must remain in flight"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "B's capacity must not be released by A's cancel" + ); + assert_eq!(token(&w.router, "prod", "d-1"), b.admission); + assert!( + w.worker_rx.try_recv().is_err(), + "no cancel is forwarded downstream for a refused cancel" + ); + + // And B still completes normally with its own token. + assert_eq!( + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", b.admission, "B's genuine result"), + 1024, + 5 + ), + CompleteOutcome::Delivered { committed: true } + ); + let frame = w.primary_rx.try_recv().expect("initiator got B's result"); + assert!(frame.contains("B's genuine result")); + assert_eq!(frame_admission(&frame), b.admission); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + } + #[test] fn terminal_frames_carry_the_token_of_the_admission_they_end() { // "First terminal frame wins" is only usable if the initiator can tell @@ -1458,6 +1693,7 @@ type = "primary" // 2. cancel. let cancel = CancelParams { delegation_id: "d-2".into(), + admission: second.admission, reason: "no longer needed".into(), }; router.cancel(®istry, hp, &cancel, 3).expect("owned"); @@ -1974,6 +2210,9 @@ type = "primary" )); let params = CancelParams { delegation_id: "d-1".into(), + // The live token: the refusal must be about who is asking, not + // about which admission was named. + admission: token(&w.router, "prod", "d-1"), reason: "changed my mind".into(), }; let err = w @@ -2232,6 +2471,7 @@ allow_worker_initiation = true )); let params = CancelParams { delegation_id: "d-1".into(), + admission: token(&w.router, "prod", "d-1"), reason: "race".into(), }; let gate = std::sync::Barrier::new(2); @@ -2270,6 +2510,7 @@ allow_worker_initiation = true )); let params = CancelParams { delegation_id: "d-1".into(), + admission: token(&w.router, "prod", "d-1"), reason: "not mine".into(), }; assert!(w @@ -2295,19 +2536,22 @@ allow_worker_initiation = true #[test] fn cancel_refusals_are_byte_identical() { // `cp/cancel` must not be an existence oracle — - // an unknown id and another instance's live id return the same error - // object, byte for byte. + // an unknown id, another instance's live id, and the caller's OWN id + // under a superseded admission token all return the same error object, + // byte for byte. The last case matters as much as the first two: a + // distinguishable refusal would tell an initiator whether the id it + // reused is currently re-admitted, which is a fact about the CP's + // scheduling state it has no frame that reports. let w = world(); - assert!(matches!( - do_delegate(&w, delegate_params("d-1", "worker-1", 60)), - DelegateOutcome::Accepted(_) - )); + let live = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))).admission; let unknown = CancelParams { delegation_id: "d-does-not-exist".into(), + admission: live, reason: "probe".into(), }; let foreign = CancelParams { delegation_id: "d-1".into(), + admission: live, reason: "probe".into(), }; // Both probes come from the worker: it initiated neither. @@ -2319,13 +2563,34 @@ allow_worker_initiation = true .router .cancel(&w.registry, w.h_worker, &foreign, 2) .unwrap_err(); + // This one comes from the genuine initiator, naming a token that is + // not the live admission's. + let stale = CancelParams { + delegation_id: "d-1".into(), + admission: live.wrapping_add(1), + reason: "probe".into(), + }; + let e_stale = w + .router + .cancel(&w.registry, w.h_primary, &stale, 3) + .unwrap_err(); + let as_json = |e: &ErrorObject| serde_json::to_string(e).unwrap(); assert_eq!( - serde_json::to_string(&e_unknown).unwrap(), - serde_json::to_string(&e_foreign).unwrap(), + as_json(&e_unknown), + as_json(&e_foreign), "unknown and foreign delegation ids must be indistinguishable" ); + assert_eq!( + as_json(&e_unknown), + as_json(&e_stale), + "a stale admission token must be indistinguishable from an unknown id" + ); assert_eq!(e_unknown.code, codes::POLICY_DENIED); + assert_eq!(e_stale.code, codes::POLICY_DENIED); + // Every refusal left the delegation exactly as it was. assert_eq!(w.router.inflight_count(), 1); + assert_eq!(token(&w.router, "prod", "d-1"), live); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 1); } #[test] @@ -2376,10 +2641,12 @@ allow_worker_initiation = true // it exists: same error as for an id that exists nowhere. let probe = CancelParams { delegation_id: "d-1".into(), + admission: token(&router, "prod", "d-1"), reason: "probe".into(), }; let nowhere = CancelParams { delegation_id: "d-nowhere".into(), + admission: 1, reason: "probe".into(), }; let e_cross = router @@ -2543,6 +2810,7 @@ allow_worker_initiation = true Interleaved::Cancel => { let params = CancelParams { delegation_id: id.into(), + admission: token(&w.router, "prod", id), reason: "changed my mind".into(), }; w.router diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index 3ab800268..d2ebabe74 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -446,9 +446,11 @@ async fn a_late_result_for_a_superseded_admission_is_not_delivered() { ); // The initiator cancels A, then retries the SAME id: admission B. + // The cancel names admission A explicitly — `delegation_id` alone would + // name whatever holds the id when the frame lands, not the work to abort. let cancel = serde_json::json!({ "jsonrpc": "2.0", "id": 11, "method": "cp/cancel", - "params": {"delegation_id": "d-1", "reason": "changed my mind"} + "params": {"delegation_id": "d-1", "admission": a, "reason": "changed my mind"} }) .to_string(); initiator.send(Message::Text(cancel.into())).await.unwrap(); diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 0eba0f0d0..1df2d0331 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -180,6 +180,11 @@ Agent A ◄──── result ◄──────── OAB-A ◄──── ### Delegate frame +What the initiator sends. Note what is *absent*: there is no `chain` field. +Callers supply at most `parent_delegation_id`; the CP constructs the ancestry +from authenticated identities and its own in-flight table, so a runtime cannot +forge it. + ```json { "method": "cp/delegate", @@ -187,20 +192,45 @@ Agent A ◄──── result ◄──────── OAB-A ◄──── "delegation_id": "d-01J...", "target": { "name": "worker-1" }, "prompt": "…", - "chain": ["koudu"], + "parent_delegation_id": "d-01H...", "deadline": "2026-08-06T22:45:00Z" } } ``` - `target` — exact `name` or a `labels` selector (CP schedules among matches) -- `chain` — the full delegation ancestry, appended at every hop. Enables - cycle rejection (target already in chain), depth enforcement, fan-out - budgets, and audit tracing back to the human-facing root. +- `parent_delegation_id` — omitted for a root delegation. If present, the + caller must be the instance currently *serving* that parent, in the caller's + own namespace; the CP appends to that parent's chain. - `deadline` — propagated absolute deadline. A child's timeout can never exceed its parent's remaining budget, so orphaned workers cannot keep consuming tokens after the root gave up. +What the serving runtime receives is a different frame: the CP adds the +authenticated `from`, the constructed `chain`, and the `admission` token for +this admission. + +```json +{ + "method": "cp/delegate", + "params": { + "delegation_id": "d-01J...", + "admission": 42, + "prompt": "…", + "deadline": "2026-08-06T22:45:00Z", + "from": "prod/koudu", + "chain": ["prod/koudu", "prod/worker-2"] + } +} +``` + +- `chain` — the full delegation ancestry, root first, appended at every hop. + Enables cycle rejection (target already in chain), depth enforcement, and + audit tracing back to the human-facing root. Every element was authenticated + by the CP, so the serving runtime can trust it. +- `admission` — the token this runtime must echo on `cp/delegate_result` and + match on `cp/cancel` (see "Admissions carry a protocol-visible token"). + ### Result delivery is protocol-mandatory Adopting Kiro's `summary` lesson at the protocol level: a delegation is not @@ -228,6 +258,29 @@ The frame MUST echo the `admission` token the CP stamped on the forwarded it*, and the id is reusable. A frame without the token is rejected as malformed, and one naming a superseded admission is dropped. +### Cancel frame + +`cp/cancel` travels both ways — initiator → CP → serving runtime — and carries +the token in both, for the same reason results do: + +```json +{ + "method": "cp/cancel", + "params": { + "delegation_id": "d-01J...", + "admission": 42, + "reason": "changed my mind" + } +} +``` + +From the initiator, `admission` is the abort target: a cancel naming a +superseded admission is refused rather than removing whatever holds the id now. +CP-synthesized cancels (deadline sweep, initiator disconnect, +stalled-initiator teardown) stamp the token of the admission they are ending, so +a best-effort cancel that overtakes a same-id re-admission's forward is +identifiable at the worker as belonging to work that is already over. + ### v1 contract amendments (from PR #1465 review) The first implementation (`crates/openab-cp`) freezes the following @@ -327,6 +380,11 @@ recovery semantics: - the forwarded `cp/delegate` carries it, so the serving runtime learns it; - `cp/delegate_result` MUST echo it. It is a **required** field: a missing token is `INVALID_PARAMS`, never a wildcard; + - `cp/cancel` carries it in **both** directions, also required. From the + initiator it names the admission to abort; on every CP-synthesized cancel + (deadline sweep, initiator disconnect, stalled-initiator teardown) the CP + stamps the token of the admission it is ending, built from the in-flight + entry it removed; - every initiator-bound terminal frame carries it, CP-synthesized `timeout` and `target_disconnected` included (both are built from the in-flight entry). @@ -347,9 +405,39 @@ recovery semantics: wire would disclose other namespaces' delegation volume, while commit matching only needs never-reuse per `(namespace, delegation_id)`. Exhaustion fails closed (the admission is refused; the counter never wraps). - `cp/cancel` does not carry the token yet — it lands with the runtime-client - slice, where the serving side gains the state to disambiguate cancellation - targets. + + Cancellation needs the token for the same reason results do, in both + directions. `cp/cancel` from the initiator is matched on `(from_handle, + namespace + delegation_id, admission)` — all three under the one lock + acquisition that removes the entry — and any mismatch is refused with the + same byte-identical `POLICY_DENIED` as an unknown id or another instance's + live id, so a caller cannot learn whether the id it reused is currently + re-admitted. Without the token, an ordinary application-level *retry* of + `cancel(A)` landing after the same id was re-admitted as B matched B and + removed it: B's capacity was released while its work continued, and B's + genuine result was later dropped as unknown with no synthesized terminal + frame, because the entry was gone. In the CP-synthesized direction the gap is + structural rather than client-dependent: the deadline sweep removes an + expired entry under the in-flight lock and builds its best-effort + `cp/cancel` after releasing it, so a same-id re-admission can be admitted and + its forward enqueued to the same worker first. The worker would then receive + `forward(B)` followed by an id-only cancel for A — different producers into + one queue, so connection ordering does not help — and B's work would be + aborted at the source while the CP still shows it live. Stamping the ended + admission's token makes that frame identifiable as belonging to work that is + already over. + + ⚠️ **Wire-breaking (pre-1.0).** `admission` is a **required** field on both + `cp/delegate_result` (added in the round-8 revision of this contract) and + `cp/cancel` (added here). Optional would be a wildcard, which is exactly the + misdelivery path the token exists to close, so there is no + backward-compatible spelling of it. A runtime built against the pre-token + contract will have **every** result and **every** cancel refused with + `INVALID_PARAMS` after the CP is upgraded. Runtimes must echo + `DelegateForward::admission` on results and name the target admission on + cancels. There are no shipped clients at this point in the stack — every + serving runtime learns the token from the forwarded `cp/delegate` — so the + migration is mechanical, but it is not silent and it is not optional. - **First terminal frame per admission wins.** More than one terminal frame may reach an initiator for one admission: a `completed` result can race the deadline sweep's synthesized `timeout`, and duplicate results are possible in diff --git a/docs/control-plane.md b/docs/control-plane.md index 704913624..07a5a5c0e 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -67,6 +67,21 @@ issue #1474). admission is dropped (the ack looks the same as any other, so do not treat `ok: true` as proof of delivery — that is what the initiator's terminal frame is for). +- **Name the admission on `cp/cancel` too.** `admission` is required there as + well, in both directions. As an initiator, send the token of the admission + you mean to abort: a cancel naming a superseded admission is refused, which + is what stops a retried cancel from killing the re-admission that replaced + its target. Refusals are deliberately indistinguishable from an unknown id, + so treat one as "not mine / not live" and reconcile against your own state + rather than inferring anything about the CP's. As a serving runtime, match + incoming cancels on the token: a CP-synthesized cancel carries the token of + the admission it ends, and it can arrive *after* a forward that reused the + same `delegation_id` — cancelling on the id alone would abort the wrong work. +- ⚠️ **Wire-breaking change (pre-1.0).** `admission` is required on + `cp/delegate_result` and on `cp/cancel`. A runtime built against the earlier + contract has every result and every cancel refused with `INVALID_PARAMS` + after upgrading the CP; there is no compatible optional spelling, because an + absent token would be the wildcard the field exists to remove. - **The first terminal frame for an `admission` token wins.** A `completed` result can race the CP's synthesized `timeout`, so an initiator may receive more than one terminal frame for the same admission. Treat the first as From cadca62773764cb8cf087440e6bd7e401762a709 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Fri, 14 Aug 2026 08:14:24 -0400 Subject: [PATCH 10/11] =?UTF-8?q?fix(cp):=20round-10=20review=20=E2=80=94?= =?UTF-8?q?=20admission=20token=20on=20parent=20linkage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 🔴 Parent linkage was the third surface binding CP-authoritative state to the reusable `delegation_id` without an admission token, after results (round 8) and cancels (round 9). `delegate()` resolved a parent by (namespace, parent_delegation_id) plus serving handle only, so once parent admission A ended (cancel/completion/sweep) and the id was re-admitted as B to the same worker, a residual child request composed against A satisfied every check against B: the child inherited B's CP-constructed chain and B's remaining deadline budget, and depth, cycle and parent-budget were evaluated for the wrong admission. Unlike the earlier two this is a policy-envelope hijack feeding authorization decisions, and a best-effort `cp/cancel` cannot make the worker police it. - proto.rs: `DelegateParams` gains `parent_admission: Option`. The parent reference is a coupled pair — a parented request must carry both halves; an id without a token is INVALID_PARAMS, never a wildcard. A root delegation carries neither and is byte-unchanged on the wire. A token without an id is refused rather than ignored (documented choice), so a client that drops the id sees its bug instead of silently getting an unparented delegation. - router.rs: the parent branch matches (namespace + id) + serving handle + generation under the same single in-flight lock acquisition it used before. Chain and deadline are still read from the `p` binding that acquisition validated — no re-lookup after validation. Stale admissions join the existing single refusal shape for unknown/unauthorized parents, so the reply is not an oracle for whether an id is currently re-admitted. - Docs: ADR §4 redefines "currently serving that parent" as the specific forwarded admission (serving handle + admission), not handle + reusable id; the parented `cp/delegate` example gains `parent_admission`; the wire-breaking callout now covers parented delegations while stating roots are unaffected. docs/control-plane.md quickstart mirrors it. The ADR round-trip test additionally asserts the example shows both parent halves or neither, since key-set equality alone cannot catch a missing `None` half. Tests (85 → 91 lib, 9 e2e unchanged): - child_naming_a_cancelled_parent_admission_cannot_hijack_the_re_admission - child_naming_a_swept_parent_admission_cannot_hijack_the_re_admission Both assert the byte-identical refusal and that B is untouched — chain, deadline, admission stamp, endpoints, capacity, and no child admitted. - child_naming_the_live_parent_admission_extends_it_normally — positive control: extends B's chain and is clamped to B's budget. A 90s child is denied because it fits A's old 120s but not B's 60s, so the clamp is provably B's. - parent_reference_requires_both_the_id_and_the_admission — both coupling directions, plus the root shape still admitted. - parent_refusals_are_byte_identical — one parent id, three CP states (unknown, foreign handle, stale admission), identical error bytes. - proto: parent_reference_is_a_pair_and_root_delegations_are_unchanged_on_the_wire Existing parent tests were threaded with the token mechanically: each names its parent's real admission so the refusal under test stays attributable to its original cause (foreign serving handle, cross-namespace lookup, policy cycle) rather than to a missing token. Negative control: reverting only the generation comparison fails both ABA regressions and the oracle test while the positive control still passes; stripping `parent_admission` from the ADR example fails the round-trip test. Scope: crates/openab-cp + the two CP docs. openab-core untouched (PR 1471 constructs root delegations only). No dependency changes. Round-10 loop follow-up: the proto.rs module-level contract summary was still pre-round-10 — its AdmissionToken usage list omitted parent_admission and the chain bullet said callers supply only parent_delegation_id, contradicting DelegateParams, the router enforcement, the ADR and the quickstart. The summary now states that parent linkage is the coupled parent_delegation_id + parent_admission pair (both required on a parented request, INVALID_PARAMS otherwise, neither on a root), that AdmissionToken is carried there, and that no surface accepts a bare reusable id as a reference to an existing delegation. --- crates/openab-cp/src/proto.rs | 121 +++++- crates/openab-cp/src/router.rs | 628 ++++++++++++++++++++++++++++---- docs/adr/agent-control-plane.md | 72 +++- docs/control-plane.md | 22 +- 4 files changed, 755 insertions(+), 88 deletions(-) diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs index 28bfdd2fa..e8ec2bc37 100644 --- a/crates/openab-cp/src/proto.rs +++ b/crates/openab-cp/src/proto.rs @@ -13,13 +13,23 @@ //! to which routing decision — uses the CP-minted //! [`AdmissionToken`]: carried on the `cp/delegate` ack and forwarded frame, //! echoed by the serving runtime in `cp/delegate_result` (required), carried -//! on `cp/cancel` in both directions (required), and stamped on every -//! initiator-bound terminal frame. +//! on `cp/cancel` in both directions (required), carried on `cp/delegate` as +//! `parent_admission` whenever the request names a parent (required), and +//! stamped on every initiator-bound terminal frame. Every surface that +//! references an *existing* delegation names it by the (id, admission) pair; +//! a bare reusable id is never accepted as a reference. //! - The first frame on a new connection MUST be `cp/register`. Anything else //! is rejected with `NOT_REGISTERED` and the connection is closed. -//! - Delegation ancestry (`chain`) is **CP-constructed**: callers supply only -//! `parent_delegation_id`; the CP derives the chain from authenticated -//! identities and its in-flight table. A runtime cannot forge ancestry. +//! - Delegation ancestry (`chain`) is **CP-constructed**: callers never supply +//! a chain. They supply a parent *reference*, and that reference is the +//! coupled pair `parent_delegation_id` + `parent_admission` — a parented +//! `cp/delegate` MUST carry both (an id without a token is +//! `INVALID_PARAMS`), so the parent is named by a specific admission rather +//! than by a reusable id. The CP derives the chain, depth, cycle and +//! deadline-budget inputs from the entry that reference resolves to, using +//! authenticated identities and its in-flight table. A root delegation +//! carries neither field. A runtime cannot forge ancestry, nor inherit the +//! chain and budget of an admission it was not serving. use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -255,6 +265,15 @@ pub struct TargetSelector { } /// Params of `cp/delegate` as sent by the initiating runtime. +/// +/// The parent reference is a **pair**: `parent_delegation_id` names the +/// delegation, `parent_admission` names which admission of it (see +/// [`AdmissionToken`]). Both are `Option` on the struct because a root +/// delegation carries neither — that shape is unchanged on the wire — but the +/// CP enforces their coupling at admission: an id without a token is +/// `INVALID_PARAMS`, never "whatever admission holds that id now". An optional +/// token on a parented request would be exactly the wildcard that lets stale +/// work inherit a live re-admission's chain and deadline budget. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DelegateParams { /// Caller-generated unique id (idempotency key). The CP rejects a second @@ -269,8 +288,28 @@ pub struct DelegateParams { /// If this delegation is issued while serving another delegation, the id /// of that parent. The CP derives the ancestry chain from this — the /// chain is never client-supplied. + /// + /// Present ⇒ `parent_admission` MUST also be present. #[serde(skip_serializing_if = "Option::is_none")] pub parent_delegation_id: Option, + /// Which admission of `parent_delegation_id` this delegation is a child + /// of — the token the serving runtime learned from + /// [`DelegateForward::admission`] for the parent it is currently serving. + /// + /// Required whenever `parent_delegation_id` is present, and meaningless + /// without it (the CP refuses a token-without-id as malformed rather than + /// ignoring it, so a client bug surfaces instead of silently producing a + /// root delegation). `delegation_id` alone names a *slot* that + /// cancel-then-retry legitimately reuses; the parent lookup feeds the + /// CP's own policy inputs — ancestry chain, depth, cycle, and the + /// deadline budget a child is clamped to — so it must name the admission + /// those inputs were computed for. Without the token, a residual task + /// from a parent admission A that has already ended could submit a child + /// against the same id after it was re-admitted as B to the same worker, + /// and the child would inherit B's trusted chain and B's remaining + /// budget. + #[serde(skip_serializing_if = "Option::is_none")] + pub parent_admission: Option, } /// Params of `cp/delegate` as forwarded to the serving runtime. The CP stamps @@ -328,6 +367,14 @@ pub struct DelegateAck { /// best-effort cancel that overtakes a same-id re-admission's forward is /// identifiable at the worker as belonging to the admission that is already /// over. +/// - the parent reference on `cp/delegate` carries it +/// (`parent_admission`, required whenever `parent_delegation_id` is +/// present). The parent lookup is what lets a child inherit a +/// CP-constructed chain and a parent's remaining deadline budget, so it +/// must name the admission those were derived for: after a parent +/// admission ends and its id is re-admitted, a child composed against the +/// ended admission is refused instead of adopting the live one's policy +/// envelope. /// /// The value is a per-namespace monotonic counter. Namespace-scoped on /// purpose: a single global counter placed on the wire would disclose @@ -498,6 +545,20 @@ mod tests { "the initiator-sent cp/delegate example must not show an \ `admission` — the CP mints it and returns it in the ack" ); + // The parent reference is a pair, and key-set equality + // alone cannot catch a missing half: an example with + // `parent_delegation_id` and no `parent_admission` parses + // into `None`, which is then skipped on serialization, so + // the key sets would still match. Assert the coupling + // directly — an example showing a bare parent id would + // document exactly the wildcard the CP refuses, and it is + // what a client implementer copies. + assert_eq!( + params.get("parent_delegation_id").is_some(), + params.get("parent_admission").is_some(), + "the cp/delegate example must show `parent_delegation_id` \ + and `parent_admission` together, or neither (root)" + ); let p: DelegateParams = serde_json::from_value(params.clone()) .expect("cp/delegate example must parse as DelegateParams"); serde_json::to_value(&p).expect("serializable") @@ -576,6 +637,56 @@ mod tests { ); } + #[test] + fn parent_reference_is_a_pair_and_root_delegations_are_unchanged_on_the_wire() { + // A root delegation carries neither field, and neither appears when it + // is serialized — the pre-token root shape is byte-compatible, which is + // why the wire break is confined to parented requests. + let root = serde_json::json!({ + "delegation_id": "d-1", + "target": {"name": "w1"}, + "prompt": "hi", + "deadline": "2026-08-06T22:45:00Z" + }); + let p: DelegateParams = serde_json::from_value(root).unwrap(); + assert!(p.parent_delegation_id.is_none()); + assert!(p.parent_admission.is_none()); + let back = serde_json::to_value(&p).unwrap(); + assert!(back.get("parent_delegation_id").is_none()); + assert!( + back.get("parent_admission").is_none(), + "a root delegation must not put an empty parent reference on the wire" + ); + + // A parented request carries both halves, and both round-trip. + let parented = serde_json::json!({ + "delegation_id": "d-2", + "target": {"name": "w1"}, + "prompt": "hi", + "deadline": "2026-08-06T22:45:00Z", + "parent_delegation_id": "d-1", + "parent_admission": 42 + }); + let p: DelegateParams = serde_json::from_value(parented).unwrap(); + assert_eq!(p.parent_delegation_id.as_deref(), Some("d-1")); + assert_eq!(p.parent_admission, Some(42)); + let back = serde_json::to_value(&p).unwrap(); + assert_eq!(back["parent_admission"], 42); + // Deserialization deliberately accepts a half-filled pair: the coupling + // is a CP admission rule (`INVALID_PARAMS`, with a message naming the + // missing half) rather than a serde error, so the refusal is a protocol + // error the caller can act on instead of a parse failure. + let half = serde_json::json!({ + "delegation_id": "d-3", + "target": {"name": "w1"}, + "prompt": "hi", + "deadline": "2026-08-06T22:45:00Z", + "parent_delegation_id": "d-1" + }); + let p: DelegateParams = serde_json::from_value(half).unwrap(); + assert!(p.parent_admission.is_none()); + } + #[test] fn delegate_result_requires_the_admission_token() { // The token is the only thing tying a result frame to ONE admission of diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index 4567c5660..de86ece26 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -402,26 +402,77 @@ impl Router { } // Parent linkage: chain and deadline derive from the CP's own table, - // never from the client. The caller must BE the instance serving the - // parent delegation — otherwise any runtime knowing a live id could - // borrow its trusted chain and deadline budget. The lookup is - // namespace-scoped. Unknown and unauthorized parent ids return the - // same error (no enumeration). - let (parent_chain, parent_deadline) = match ¶ms.parent_delegation_id { - Some(pid) => { - let parent_key = DelegationKey::new(from_namespace, pid); - match self.inflight.lock().get(&parent_key) { - Some(p) if p.to_handle == from_handle => (p.chain.clone(), Some(p.deadline)), - _ => { - return DelegateOutcome::Rejected(ErrorObject::new( - codes::INVALID_PARAMS, - format!("parent delegation {pid} is not in flight for this instance"), - )) + // never from the client. Three things must hold together, and the + // admission token is the one that makes the other two sufficient: + // + // 1. the parent lives in the CALLER's namespace (scoped lookup); + // 2. the caller IS the instance serving it (`to_handle`) — otherwise + // any runtime knowing a live id could borrow its trusted chain and + // deadline budget; + // 3. the caller names the SPECIFIC admission it is serving. Without + // this, "currently serving that parent" degrades to "holds the + // connection that serves whatever wears this id now": once parent + // admission A ends (cancel, completion, or sweep) and the id is + // re-admitted as B — with a single replica, to the same worker — a + // residual child request composed against A satisfies 1 and 2 + // against B and inherits B's CP-constructed chain and B's remaining + // deadline budget, so depth, cycle, and parent-budget are evaluated + // for the wrong admission. Cancels are best effort, so the CP + // cannot delegate policing of this to the worker: a buggy or + // malicious runtime holding the serving connection could trigger it + // deliberately. + // + // The id and the token are a coupled pair, enforced here because this + // is where parent presence is decided. An id without a token is + // malformed, never a wildcard; a token without an id is malformed too, + // rather than silently ignored, so a client that drops the id sees its + // bug instead of getting an unintended root delegation. + // + // Unknown, unauthorized (wrong serving handle) and stale (superseded + // admission) parents all return the SAME error — one refusal shape, no + // enumeration. A distinguishable stale refusal would be an oracle + // telling the caller whether the id it references is currently + // re-admitted, which is CP scheduling state no frame reports. + let (parent_chain, parent_deadline) = + match (¶ms.parent_delegation_id, params.parent_admission) { + (Some(pid), Some(padmission)) => { + let parent_key = DelegationKey::new(from_namespace, pid); + // ONE acquisition of the in-flight lock resolves the parent and + // validates all three conditions, and the chain and deadline + // below are read from the very entry that was validated — the + // `p` binding, not a second lookup. A re-lookup after + // validation would reintroduce the race the token closes. + match self.inflight.lock().get(&parent_key) { + Some(p) if p.to_handle == from_handle && p.generation == padmission => { + (p.chain.clone(), Some(p.deadline)) + } + _ => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + format!( + "parent delegation {pid} is not in flight for this instance" + ), + )) + } } } - } - None => (Vec::new(), None), - }; + (Some(_), None) => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + "parent_delegation_id requires parent_admission: name the \ + admission token this instance was forwarded for that parent \ + (a delegation id alone is reusable and cannot identify it)", + )) + } + (None, Some(_)) => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + "parent_admission is meaningless without parent_delegation_id: \ + omit both for a root delegation, or send both", + )) + } + (None, None) => (Vec::new(), None), + }; // Resolve target within the initiator's namespace (v1 boundary). let target = match registry.select(from_namespace, sel_name, sel_labels) { @@ -1234,6 +1285,24 @@ type = "worker" prompt: "do it".into(), deadline: Utc::now() + Duration::seconds(secs), parent_delegation_id: None, + parent_admission: None, + } + } + + /// A child request naming a parent as a well-behaved serving runtime + /// composes it: the parent id AND the admission token that runtime was + /// forwarded for it. The pair is what the CP requires. + fn child_params( + id: &str, + target: &str, + secs: i64, + parent_id: &str, + parent_admission: u64, + ) -> DelegateParams { + DelegateParams { + parent_delegation_id: Some(parent_id.into()), + parent_admission: Some(parent_admission), + ..delegate_params(id, target, secs) } } @@ -2252,18 +2321,18 @@ allow_worker_initiation = true let (w2, _rx2) = instance("prod", "worker-2", AgentType::Worker, 1); let h_w2 = w.registry.register(w2); - assert!(matches!( - w.router.delegate( - &cfg, - &w.registry, - "prod", - "koudu", - &AgentType::Primary, - w.h_primary, - delegate_params("d-root", "worker-1", 120), - 1, - ), - DelegateOutcome::Accepted(_) + // Capturing the ack (rather than asserting the shape) because every + // parented request below must name the admission token of the parent it + // references — the ack is where a real initiator learns it. + let root = accept(w.router.delegate( + &cfg, + &w.registry, + "prod", + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-root", "worker-1", 120), + 1, )); assert_eq!( w.router.chain_of("prod", "d-root").unwrap(), @@ -2272,9 +2341,10 @@ allow_worker_initiation = true // Borrowed ancestry: worker-2 (NOT serving d-root) tries to use // d-root as its parent — rejected, so a trusted chain and deadline - // budget cannot be inherited by a stranger. - let mut foreign = delegate_params("d-foreign", "worker-2", 60); - foreign.parent_delegation_id = Some("d-root".into()); + // budget cannot be inherited by a stranger. It names the CORRECT + // admission token, so the refusal is still attributable to the wrong + // serving handle rather than to a bad token. + let foreign = child_params("d-foreign", "worker-2", 60, "d-root", root.admission); match w.router.delegate( &cfg, &w.registry, @@ -2293,29 +2363,26 @@ allow_worker_initiation = true } // worker-1 (serving d-root) delegates a legitimate child to worker-2. - let mut child = delegate_params("d-child", "worker-2", 60); - child.parent_delegation_id = Some("d-root".into()); - assert!(matches!( - w.router.delegate( - &cfg, - &w.registry, - "prod", - "worker-1", - &AgentType::Worker, - w.h_worker, - child, - 3, - ), - DelegateOutcome::Accepted(_) + let child = child_params("d-child", "worker-2", 60, "d-root", root.admission); + let child_ack = accept(w.router.delegate( + &cfg, + &w.registry, + "prod", + "worker-1", + &AgentType::Worker, + w.h_worker, + child, + 3, )); assert_eq!( w.router.chain_of("prod", "d-child").unwrap(), vec!["prod/koudu".to_string(), "prod/worker-1".to_string()] ); - // Cycle: worker-2 delegating back to koudu is rejected. - let mut cyc = delegate_params("d-cyc", "koudu", 30); - cyc.parent_delegation_id = Some("d-child".into()); + // Cycle: worker-2 delegating back to koudu is rejected. Naming + // d-child's real admission token, so the request clears parent + // validation and the refusal comes from the policy engine. + let cyc = child_params("d-cyc", "koudu", 30, "d-child", child_ack.admission); match w.router.delegate( &cfg, &w.registry, @@ -2334,6 +2401,434 @@ allow_worker_initiation = true } } + /// The full in-flight entry for `(namespace, delegation_id)` — so a test + /// can assert that a refused operation left a live admission's + /// CP-authoritative state (chain, deadline, stamp, endpoints) untouched, + /// not merely that the entry still exists. + fn entry(router: &Router, namespace: &str, delegation_id: &str) -> InFlight { + router + .inflight + .lock() + .get(&DelegationKey::new(namespace, delegation_id)) + .expect("delegation must be in flight") + .clone() + } + + fn rejected(out: DelegateOutcome) -> ErrorObject { + match out { + DelegateOutcome::Rejected(e) => e, + DelegateOutcome::Accepted(a) => { + panic!("expected a refusal, got admission {}", a.admission) + } + } + } + + /// `world()`'s identity table plus a depth budget and worker initiation, + /// so an instance serving a parent may legitimately delegate a child. + fn deep_cfg() -> CpConfig { + toml::from_str( + r#" +[[agents]] +key = "kp" +namespace = "prod" +name = "koudu" +type = "primary" + +[namespaces.prod] +max_depth = 5 +allow_worker_initiation = true +"#, + ) + .unwrap() + } + + /// A world with a second worker, so a child can be routed somewhere other + /// than the instance serving its parent. The returned `FrameRx` must stay + /// alive: dropping it closes worker-2's outbound channel and every send to + /// it would fail as a disconnect. + fn parent_world() -> (World, CpConfig, u64, FrameRx) { + let w = world(); + let (w2, rx2) = instance("prod", "worker-2", AgentType::Worker, 2); + let h_w2 = w.registry.register(w2); + (w, deep_cfg(), h_w2, rx2) + } + + fn delegate_as( + w: &World, + cfg: &CpConfig, + name: &str, + ty: &AgentType, + handle: u64, + params: DelegateParams, + rpc: u64, + ) -> DelegateOutcome { + w.router + .delegate(cfg, &w.registry, "prod", name, ty, handle, params, rpc) + } + + fn cancel_admission(w: &World, delegation_id: &str, admission: u64, rpc: u64) { + w.router + .cancel( + &w.registry, + w.h_primary, + &CancelParams { + delegation_id: delegation_id.into(), + admission, + reason: "changed my mind".into(), + }, + rpc, + ) + .expect("the initiator may cancel its own live admission"); + } + + /// Admit parent id `d-parent` on worker-1 as admission A (120s of budget), + /// end A, then re-admit the SAME id as B (60s) — which, with a single + /// replica, routes to the same worker. Returns (A's token, B's token). + /// + /// The asymmetric budgets are load-bearing: 60s is strictly inside A's + /// 120s, so a child clamped against the wrong admission is observable. + fn parent_aba_via_cancel(w: &World, cfg: &CpConfig) -> (u64, u64) { + let a = accept(delegate_as( + w, + cfg, + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-parent", "worker-1", 120), + 1, + )); + cancel_admission(w, "d-parent", a.admission, 2); + assert_eq!(w.router.inflight_count(), 0, "A is over"); + let b = accept(delegate_as( + w, + cfg, + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-parent", "worker-1", 60), + 3, + )); + assert_ne!(a.admission, b.admission, "a re-admission gets a new token"); + (a.admission, b.admission) + } + + /// Assert that the refusal is the single parent refusal shape, and that a + /// stale child changed nothing: no child exists, and B is exactly the + /// entry it was — chain, deadline, admission stamp, endpoints, capacity. + fn assert_stale_child_refused_and_b_untouched( + w: &World, + b: u64, + before: &InFlight, + out: DelegateOutcome, + h_w2: u64, + ) { + let e = rejected(out); + assert_eq!(e.code, codes::INVALID_PARAMS); + assert!( + e.message.contains("not in flight for this instance"), + "{}", + e.message + ); + assert!( + w.router.chain_of("prod", "d-child").is_none(), + "no child may have been admitted" + ); + let after = entry(&w.router, "prod", "d-parent"); + assert_eq!(after.generation, b, "B's admission stamp is unchanged"); + assert_eq!(after.chain, before.chain, "B's CP-constructed chain"); + assert_eq!(after.deadline, before.deadline, "B's deadline budget"); + assert_eq!(after.to_handle, before.to_handle, "B's serving instance"); + assert_eq!(after.from_handle, before.from_handle, "B's initiator"); + assert_eq!(w.router.inflight_count(), 1, "B and nothing else"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "B's capacity reservation is intact" + ); + assert_eq!( + w.registry.get(h_w2).unwrap().active_sessions, + 0, + "the refused child reserved nothing on its target" + ); + } + + #[test] + fn child_naming_a_cancelled_parent_admission_cannot_hijack_the_re_admission() { + // The parent-linkage ABA, cancel flavour — the third reusable-id + // surface after results and cancels. worker-1 serves admission A of + // parent id P. A is cancelled and the initiator retries (an ordinary + // pattern), so P is re-admitted as B to the SAME worker. A residual + // task from A then submits a child against P over that same + // connection. + // + // Every check that predates the token still passes: the namespace is + // right and the caller IS the instance serving P. So the child would + // have inherited B's CP-constructed chain and B's remaining deadline + // budget, and depth, cycle and parent-budget would all have been + // evaluated for an admission it has nothing to do with. Worker-side + // discipline cannot substitute: cancels are best effort, and the + // runtime holding the serving connection is exactly who would exploit + // this deliberately. + let (w, cfg, h_w2, _rx2) = parent_world(); + let (a, b) = parent_aba_via_cancel(&w, &cfg); + let before = entry(&w.router, "prod", "d-parent"); + assert_eq!(before.generation, b); + + let out = delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + child_params("d-child", "worker-2", 30, "d-parent", a), + 4, + ); + assert_stale_child_refused_and_b_untouched(&w, b, &before, out, h_w2); + } + + #[test] + fn child_naming_a_swept_parent_admission_cannot_hijack_the_re_admission() { + // Same defect, reached without any client error at all. The deadline + // sweep removes the expired parent under the in-flight lock, and the + // initiator legitimately re-admits the id afterwards; a task the sweep + // could only *best-effort* cancel is still running and still composes + // children against the admission it was forwarded. This is the path + // the CP cannot delegate to client discipline, only survive by + // identity. + let (w, cfg, h_w2, _rx2) = parent_world(); + let a = accept(delegate_as( + &w, + &cfg, + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-parent", "worker-1", 1), + 1, + )); + let mut id = 900u64; + let mut next = || { + id += 1; + id + }; + let swept = + w.router + .sweep_deadlines(&w.registry, Utc::now() + Duration::seconds(60), &mut next); + assert!(!swept.is_empty(), "A expired and was swept"); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + + // The gap the sweep leaves open: B is admitted to the same worker. + let b = accept(delegate_as( + &w, + &cfg, + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-parent", "worker-1", 60), + 2, + )); + assert_ne!(a.admission, b.admission); + let before = entry(&w.router, "prod", "d-parent"); + + let out = delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + child_params("d-child", "worker-2", 30, "d-parent", a.admission), + 3, + ); + assert_stale_child_refused_and_b_untouched(&w, b.admission, &before, out, h_w2); + } + + #[test] + fn child_naming_the_live_parent_admission_extends_it_normally() { + // The positive control for both regressions above: the token is an + // identity check, not a new obstacle. A child naming the admission its + // worker is actually serving is admitted, extends THAT admission's + // chain, and is clamped to THAT admission's remaining budget. + let (w, cfg, _h_w2, _rx2) = parent_world(); + let (_a, b) = parent_aba_via_cancel(&w, &cfg); + let b_deadline = entry(&w.router, "prod", "d-parent").deadline; + + accept(delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + child_params("d-child", "worker-2", 30, "d-parent", b), + 4, + )); + assert_eq!( + w.router.chain_of("prod", "d-child").unwrap(), + vec!["prod/koudu".to_string(), "prod/worker-1".to_string()], + "the child extends the live parent admission's chain" + ); + let child_deadline = entry(&w.router, "prod", "d-child").deadline; + assert!( + child_deadline <= b_deadline, + "child deadline {child_deadline} must sit inside B's budget {b_deadline}" + ); + + // And the clamp is B's, not the id's. 90s is comfortably inside the + // 120s that admission A carried and outside B's 60s: if the derivation + // ever read the wrong admission for this id, this child would be + // admitted. + let e = rejected(delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + child_params("d-child-2", "worker-2", 90, "d-parent", b), + 5, + )); + assert_eq!(e.code, codes::POLICY_DENIED); + assert!( + e.message.contains("parent's remaining budget"), + "{}", + e.message + ); + } + + #[test] + fn parent_reference_requires_both_the_id_and_the_admission() { + // The pair is coupled at admission, where parent presence is decided. + let (w, cfg, _h_w2, _rx2) = parent_world(); + let a = accept(delegate_as( + &w, + &cfg, + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-parent", "worker-1", 120), + 1, + )); + + // (1) An id with no token is the wildcard shape, and it is refused — + // this request names a real, live parent from its genuine serving + // instance, so the ONLY thing wrong with it is the missing token. + // Accepting it would restore the whole defect under a different name. + let mut no_token = child_params("d-child", "worker-2", 30, "d-parent", a.admission); + no_token.parent_admission = None; + let e = rejected(delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + no_token, + 2, + )); + assert_eq!(e.code, codes::INVALID_PARAMS); + assert!( + e.message.contains("requires parent_admission"), + "{}", + e.message + ); + + // (2) A token with no id is refused rather than ignored — the + // documented choice. Silently treating it as a root delegation would + // hide a client bug and hand the caller an unintended delegation with + // no ancestry, no depth accounting and no parent budget; a request + // that names an admission plainly means to be parented. + let mut no_id = delegate_params("d-orphan", "worker-2", 30); + no_id.parent_admission = Some(a.admission); + let e = rejected(delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + no_id, + 3, + )); + assert_eq!(e.code, codes::INVALID_PARAMS); + assert!( + e.message.contains("without parent_delegation_id"), + "{}", + e.message + ); + + // Neither refusal admitted anything, and the ROOT shape — neither + // field present — is unchanged on the wire and still accepted. + assert_eq!(w.router.inflight_count(), 1, "only the parent"); + accept(delegate_as( + &w, + &cfg, + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-root-2", "worker-2", 30), + 4, + )); + } + + #[test] + fn parent_refusals_are_byte_identical() { + // The parent lookup must not become an existence oracle. For ONE + // parent id, three different CP states return the same error object + // byte for byte: no such parent, a live parent this caller does not + // serve, and a parent whose live admission is newer than the one the + // caller names. The third is the case added with the token, and it + // matters as much as the others — a distinguishable stale refusal + // would tell the caller that the id it references has been + // re-admitted, which is CP scheduling state no frame reports. The id + // itself appears in the message, but that is the caller's own input. + let (w, cfg, h_w2, _rx2) = parent_world(); + let (a, b) = parent_aba_via_cancel(&w, &cfg); + + // (i) Right serving instance, superseded admission. + let e_stale = rejected(delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + child_params("d-child", "worker-2", 30, "d-parent", a), + 4, + )); + // (ii) Live parent, correct token, caller does not serve it. + let e_foreign = rejected(delegate_as( + &w, + &cfg, + "worker-2", + &AgentType::Worker, + h_w2, + child_params("d-child", "worker-1", 30, "d-parent", b), + 5, + )); + // (iii) No such parent at all, with a token that was live a moment ago. + cancel_admission(&w, "d-parent", b, 6); + assert_eq!(w.router.inflight_count(), 0); + let e_unknown = rejected(delegate_as( + &w, + &cfg, + "worker-1", + &AgentType::Worker, + w.h_worker, + child_params("d-child", "worker-2", 30, "d-parent", b), + 7, + )); + + let as_json = |e: &ErrorObject| serde_json::to_string(e).unwrap(); + assert_eq!( + as_json(&e_stale), + as_json(&e_foreign), + "a stale admission must be indistinguishable from a parent this \ + caller does not serve" + ); + assert_eq!( + as_json(&e_stale), + as_json(&e_unknown), + "a stale admission must be indistinguishable from no such parent — \ + otherwise the refusal reports whether the id was re-admitted" + ); + assert_eq!(e_stale.code, codes::INVALID_PARAMS); + } + /// A `cp/delegate_result` frame as a well-behaved serving runtime builds /// it: echoing the admission token it was forwarded. fn result_of(id: &str, admission: u64, body: &str) -> DelegateResultParams { @@ -2726,24 +3221,24 @@ allow_worker_initiation = true let hw_dev = registry.register(w_dev); registry.register(t_dev); - assert!(matches!( - router.delegate( - &cfg, - ®istry, - "prod", - "koudu", - &AgentType::Primary, - hp_prod, - delegate_params("d-root", "worker-1", 120), - 1, - ), - DelegateOutcome::Accepted(_) + // Captured for its admission token: both children below must name the + // admission of the parent they reference. + let root = accept(router.delegate( + &cfg, + ®istry, + "prod", + "koudu", + &AgentType::Primary, + hp_prod, + delegate_params("d-root", "worker-1", 120), + 1, )); rx2.try_recv().unwrap(); - // dev/worker-1 claims prod's `d-root` as its parent: invisible. - let mut child = delegate_params("d-child", "worker-2", 60); - child.parent_delegation_id = Some("d-root".into()); + // dev/worker-1 claims prod's `d-root` as its parent: invisible. It + // names prod's real admission token, so only the namespace scope can + // account for the refusal. + let child = child_params("d-child", "worker-2", 60, "d-root", root.admission); match router.delegate( &cfg, ®istry, @@ -2761,8 +3256,7 @@ allow_worker_initiation = true _ => panic!("cross-namespace parent must be rejected"), } // ...and the legitimate in-namespace child still works. - let mut ok_child = delegate_params("d-child", "worker-2", 60); - ok_child.parent_delegation_id = Some("d-root".into()); + let ok_child = child_params("d-child", "worker-2", 60, "d-root", root.admission); assert!(matches!( router.delegate( &cfg, diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 1df2d0331..7b5d91ba7 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -181,9 +181,10 @@ Agent A ◄──── result ◄──────── OAB-A ◄──── ### Delegate frame What the initiator sends. Note what is *absent*: there is no `chain` field. -Callers supply at most `parent_delegation_id`; the CP constructs the ancestry +Callers supply at most a parent reference; the CP constructs the ancestry from authenticated identities and its own in-flight table, so a runtime cannot -forge it. +forge it. This example is a *parented* delegation — a root one simply omits +both parent fields. ```json { @@ -193,15 +194,26 @@ forge it. "target": { "name": "worker-1" }, "prompt": "…", "parent_delegation_id": "d-01H...", + "parent_admission": 41, "deadline": "2026-08-06T22:45:00Z" } } ``` - `target` — exact `name` or a `labels` selector (CP schedules among matches) -- `parent_delegation_id` — omitted for a root delegation. If present, the - caller must be the instance currently *serving* that parent, in the caller's - own namespace; the CP appends to that parent's chain. +- `parent_delegation_id` + `parent_admission` — a **pair**, both omitted for a + root delegation and both required together otherwise. If present, the caller + must be the instance currently serving *that specific admission* of the + parent, in the caller's own namespace; the CP appends to that admission's + chain. "Currently serving that parent" means the exact forwarded admission — + the serving handle plus the `admission` token the CP stamped on the parent's + forwarded `cp/delegate` — **not** the serving handle plus the reusable + `delegation_id`. The id alone names a slot that cancel-then-retry + legitimately reuses, so handle + id would let a task from an admission that + has already ended inherit the chain and deadline budget of whatever was + re-admitted under the same id (see "Admissions carry a protocol-visible + token"). An id without a token is `INVALID_PARAMS`; so is a token without an + id, rather than being ignored as an accidental root delegation. - `deadline` — propagated absolute deadline. A child's timeout can never exceed its parent's remaining budget, so orphaned workers cannot keep consuming tokens after the root gave up. @@ -294,8 +306,9 @@ recovery semantics: from self-asserted registration fields. Keys are per-agent (individually revocable) and presented as `Authorization: Bearer` on the WebSocket upgrade — never in URLs. -- **CP-constructed chain.** `cp/delegate` carries only - `parent_delegation_id`; the CP derives the ancestry chain from its +- **CP-constructed chain.** `cp/delegate` carries only a parent *reference* + (`parent_delegation_id` + `parent_admission`); the CP derives the ancestry + chain from its in-flight table and the authenticated caller identity, then stamps it on the forwarded frame. A runtime cannot forge ancestry, so depth/cycle checks operate on trusted data. Policy (role, depth, cycle, namespace, @@ -307,7 +320,9 @@ recovery semantics: keyed by a **CP-generated handle**, never the client-supplied `instance_id`: a colliding `instance_id` cannot replace or tear down another connection's registration, and all in-flight ownership checks - (completion, cancellation, parent linkage) compare handles. The ack + (completion, cancellation, parent linkage) compare handles — each paired + with the `admission` token, since the handle alone cannot say *which* + admission of a reusable id is meant. The ack carries the heartbeat interval, lease window, and the effective (possibly clamped) concurrency budget. Instances missing heartbeats past the lease are deregistered, and their in-flight delegations fail immediately with a @@ -387,7 +402,9 @@ recovery semantics: entry it removed; - every initiator-bound terminal frame carries it, CP-synthesized `timeout` and `target_disconnected` included (both are built from the in-flight - entry). + entry); + - a **parent reference** on `cp/delegate` carries it as `parent_admission`, + required whenever `parent_delegation_id` is present. The CP checks the echoed token *before* building the initiator-bound frame, and the commit phase removes an in-flight entry only when key, serving @@ -427,15 +444,44 @@ recovery semantics: admission's token makes that frame identifiable as belonging to work that is already over. + Parent linkage is the third surface, and the one that feeds the CP's own + authorization decisions rather than frame delivery. The parent lookup exists + to stop any runtime that knows a live id from borrowing its trusted chain and + deadline budget, and it checks that the caller is the instance *serving* that + parent — but "serving" resolved through `(namespace, delegation_id)` plus the + serving handle is not an admission. Once parent admission A ends (cancel, + completion, or sweep) and the id is re-admitted as B — with a single replica, + to the same worker — a residual task from A submitting + `cp/delegate { parent_delegation_id: P }` satisfies every one of those checks + against B. The child then inherits B's CP-constructed chain and B's remaining + deadline budget: depth, cycle, and parent-budget are evaluated for the wrong + admission, and the audit chain attributes A's work to B's root. Unlike the + result and cancel cases this is not misdelivery but a policy-envelope hijack, + and the CP cannot push it onto the worker: `cp/cancel` is best effort, and the + runtime holding the serving connection is precisely who would trigger it + deliberately. So the reference is a pair — `parent_delegation_id` plus + `parent_admission` — matched together with the serving handle under the one + in-flight lock acquisition the parent branch already takes, and the chain and + deadline are read from the entry that acquisition validated rather than from a + second lookup. Unknown, unauthorized and stale parents share one refusal + shape, so the reply is not an oracle for whether an id is currently + re-admitted. + ⚠️ **Wire-breaking (pre-1.0).** `admission` is a **required** field on both `cp/delegate_result` (added in the round-8 revision of this contract) and - `cp/cancel` (added here). Optional would be a wildcard, which is exactly the + `cp/cancel` (added in round 9), and a **parented** `cp/delegate` must carry + `parent_admission` alongside `parent_delegation_id` (added here). Optional + would be a wildcard, which is exactly the misdelivery path the token exists to close, so there is no backward-compatible spelling of it. A runtime built against the pre-token contract will have **every** result and **every** cancel refused with - `INVALID_PARAMS` after the CP is upgraded. Runtimes must echo - `DelegateForward::admission` on results and name the target admission on - cancels. There are no shipped clients at this point in the stack — every + `INVALID_PARAMS` after the CP is upgraded, and every *parented* delegation + refused with it too. Runtimes must echo + `DelegateForward::admission` on results, name the target admission on + cancels, and name the parent's admission when delegating a child while + serving that parent. Root delegations are unaffected — they carry neither + parent field, so their wire shape is unchanged. There are no shipped clients + at this point in the stack — every serving runtime learns the token from the forwarded `cp/delegate` — so the migration is mechanical, but it is not silent and it is not optional. - **First terminal frame per admission wins.** More than one terminal frame may diff --git a/docs/control-plane.md b/docs/control-plane.md index 07a5a5c0e..c55b7600f 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -77,11 +77,27 @@ issue #1474). incoming cancels on the token: a CP-synthesized cancel carries the token of the admission it ends, and it can arrive *after* a forward that reused the same `delegation_id` — cancelling on the id alone would abort the wrong work. +- **Name the parent's admission when you delegate a child.** If you issue + `cp/delegate` *while serving* another delegation, send + `parent_delegation_id` **and** `parent_admission` — the token you were + forwarded for that parent. Both or neither: an id without a token is + refused, and so is a token without an id. "The instance currently serving + that parent" means the specific admission you were forwarded, not your + connection plus the parent's `delegation_id`, because that id is reusable — + otherwise a task whose parent has already ended could inherit the chain and + deadline budget of whatever was re-admitted under the same id. Refusals here + use the same shape for an unknown parent, a parent you do not serve, and a + superseded admission, so treat one as "that parent admission is over" and + stop fanning out rather than retrying with a different token. A **root** + delegation omits both fields and is unchanged. - ⚠️ **Wire-breaking change (pre-1.0).** `admission` is required on - `cp/delegate_result` and on `cp/cancel`. A runtime built against the earlier - contract has every result and every cancel refused with `INVALID_PARAMS` + `cp/delegate_result` and on `cp/cancel`, and `parent_admission` is required + on any `cp/delegate` that names a parent. A runtime built against the earlier + contract has every result, every cancel, and every parented delegation + refused with `INVALID_PARAMS` after upgrading the CP; there is no compatible optional spelling, because an - absent token would be the wildcard the field exists to remove. + absent token would be the wildcard the field exists to remove. Root + delegations are unaffected. - **The first terminal frame for an `admission` token wins.** A `completed` result can race the CP's synthesized `timeout`, so an initiator may receive more than one terminal frame for the same admission. Treat the first as From 33e99d2be991efefb23581fe7bd7fca88a7e13d2 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Fri, 14 Aug 2026 16:58:27 +0000 Subject: [PATCH 11/11] docs(adr): extend the (id, admission) invariant to the facade/CLI contract The v1 facade tool table and CLI examples still referenced existing delegations by bare reusable delegation_id, contradicting the invariant established for the wire protocol (result r8, cancel r9, parent r10): every surface that references an existing delegation names it by the (id, admission) pair. Async spawn_agent now returns an opaque delegation handle encapsulating (delegation_id, admission); check_delegation, cancel_delegation, and 'openab agent status ' take that handle, so a delayed local API operation can never observe or abort a same-id re-admission. This is a contract-only amendment directing PR 3; no shipped code changes. --- docs/adr/agent-control-plane.md | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 7b5d91ba7..6efbbf146 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -625,10 +625,19 @@ integration. v1 tool surface, intentionally minimal: | Tool | Behavior | |------|----------| -| `spawn_agent` | Delegate a task. Blocking (waits up to deadline) or async (returns `delegation_id` immediately). | -| `check_delegation` | Status / result by `delegation_id`. | +| `spawn_agent` | Delegate a task. Blocking (waits up to deadline) or async (returns a `delegation` handle immediately). | +| `check_delegation` | Status / result by `delegation` handle. | | `list_agents` | Registry view for the caller's namespace (names, types, labels, availability) — lets the model discover targets by label. | -| `cancel_delegation` | Cancel an in-flight delegation. | +| `cancel_delegation` | Cancel an in-flight delegation by `delegation` handle. | + +The `delegation` handle returned by `spawn_agent` is opaque to callers and +encapsulates the `(delegation_id, admission)` pair. This extends the wire +invariant ("every surface that references an existing delegation names it by +the (id, admission) pair; a bare reusable id is never accepted as a +reference") to the local API: `check_delegation` and `cancel_delegation` +resolve the handle to the exact admission it was minted for, so a delayed +check or cancel can never observe or abort a same-id re-admission. Callers +never see or construct the two halves separately. The facade is where policy is enforced *before* frames leave the box: schema validation, chain/depth checks, deadline clamping, audit logging. A @@ -646,7 +655,10 @@ evolves underneath. ### CLI (`openab agent `) — secondary client, same socket - **Ops/debugging:** exec into a task and run `openab agent list` / - `openab agent status ` when a delegation hangs + `openab agent status ` when a delegation hangs — the handle is the + same opaque `(delegation_id, admission)` value `spawn` printed, so a status + or cancel typed minutes later still names the exact admission, never a + same-id re-admission - **Hooks & cron:** lifecycle hooks and cron jobs can fire `openab agent spawn …` without new plumbing - **Escape hatch** for backends where MCP injection proves awkward