diff --git a/Cargo.lock b/Cargo.lock index 89f00d9c7..fc11ec7e5 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", + "tokio-tungstenite 0.29.0", + "toml", + "tracing", + "tracing-subscriber", +] + [[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/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/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/Cargo.toml b/crates/openab-cp/Cargo.toml new file mode 100644 index 000000000..86ca5e30c --- /dev/null +++ b/crates/openab-cp/Cargo.toml @@ -0,0 +1,26 @@ +[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" +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 new file mode 100644 index 000000000..9ae1c1642 --- /dev/null +++ b/crates/openab-cp/cp.toml.example @@ -0,0 +1,94 @@ +# 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. + +# 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. +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 + +# 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 + +# 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 + +# 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" +name = "koudu" +type = "primary" + +[[agents]] +key = "${CP_KEY_WORKER1}" +namespace = "prod" +name = "worker-1" +type = "worker" +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. +[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..4385eaeb6 --- /dev/null +++ b/crates/openab-cp/src/config.rs @@ -0,0 +1,592 @@ +//! CP-side configuration. +//! +//! 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 +//! 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. 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, + + /// 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, so an oversized frame is + /// never buffered in full. + #[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, + + /// 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: 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. 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, + + /// 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, + + /// 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. + #[serde(default)] + pub agents: Vec, + + /// Per-namespace policy overrides. + #[serde(default)] + pub namespaces: BTreeMap, +} + +fn default_listen() -> String { + "127.0.0.1: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 +} +fn default_register_timeout_secs() -> u64 { + 10 +} +fn default_max_connections_per_identity() -> u32 { + 8 +} +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)] +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. 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, +} + +/// 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"); + } + // 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"); + } + 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"); + } + // 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) { + 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(()) + } + + /// 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() + } + + /// 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. +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 { + 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 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 admission_bounds_default_and_are_validated() { + // 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); + assert_eq!(cfg.write_timeout_secs, 30); + + 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"); + } + } + + #[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"); + 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..6792de59b --- /dev/null +++ b/crates/openab-cp/src/policy.rs @@ -0,0 +1,246 @@ +//! 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 +//! 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..e8ec2bc37 --- /dev/null +++ b/crates/openab-cp/src/proto.rs @@ -0,0 +1,775 @@ +//! 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. +//! - `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), carried +//! 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 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; + +/// 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**: 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; + // -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). + 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. +/// +/// 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 + /// 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. + /// + /// 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 +/// the authenticated origin and the CP-constructed chain. +#[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`). + 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 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; +/// - `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 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 +/// 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. +pub type AdmissionToken = u64; + +// --- 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. +/// +/// `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, + #[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. +/// +/// `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, +} + +// --- 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::*; + + /// 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" + ); + // 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") + } + (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!({ + "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 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 + // 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 = + 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..7404bf0d4 --- /dev/null +++ b/crates/openab-cp/src/registry.rs @@ -0,0 +1,669 @@ +//! 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, AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use parking_lot::RwLock; +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 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, +} + +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 +/// 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 +/// 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. 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(None).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 { + /// CP-generated registration handle — the registry key and the basis of + /// 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, + 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. + 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. + /// + /// `shutdown` is the owning connection's termination signal: the CP + /// triggers it whenever it drops the registration on its own initiative + /// (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; + 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, 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(Some(reason)); + 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).map(|e| e.inst) + } + + /// Refresh the lease. The runtime-reported session count is intentionally + /// 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) { + Some(e) => { + e.inst.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(|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).map(|e| e.inst.clone()) + } + + /// Select a serving instance within `namespace` by exact name or labels. + /// + /// 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, + /// 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() + .map(|e| &e.inst) + .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(e) = g.get_mut(&handle) { + e.inst.active_sessions = e.inst.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() + .map(|e| &e.inst) + .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) = outbound_channel(1024 * 1024); + 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 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(); + 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 + // (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; + 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() { + // 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); + } + + #[tokio::test] + async fn signal_shutdown_reaches_the_owning_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(); + let mut rx = sig.subscribe(); + let h = r.register_conn(inst("prod", "w1", "i-1", 1), sig); + assert!(rx.borrow().is_none()); + + assert!(r.signal_shutdown(h, "lease expired")); + rx.changed().await.unwrap(); + 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, "lease expired")); + } + + #[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, "lease expired"); + r.deregister(h); + rx.changed().await.unwrap(); + assert_eq!(*rx.borrow(), Some("lease expired")); + } +} diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs new file mode 100644 index 000000000..de86ece26 --- /dev/null +++ b/crates/openab-cp/src/router.rs @@ -0,0 +1,3623 @@ +//! Delegation router: in-flight table, target selection, result routing, and +//! 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 +//! 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 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). +//! +//! # 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 +//! 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 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 +//! +//! 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; + +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: a colliding +/// `instance_id` on another connection can neither complete nor cancel this +/// delegation. +#[derive(Clone)] +pub struct InFlight { + /// 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. + 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, + /// 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 + /// 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 **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, +} + +/// 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 +/// `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>, + /// 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 { + /// Forwarded to the target; ack for the initiator. + Accepted(DelegateAck), + /// Rejected; error for the initiator. + Rejected(ErrorObject), +} + +/// 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, +/// 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. +/// +/// 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. 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 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 { + namespace: String, + /// Handle of the instance that does own it (CP-side logs only — it is + /// 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 + /// expiry) or never registered, so it has no namespace to look in. + 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, 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 + /// to a genuine result or to the deadline sweep. + Foreign { + /// 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, +} + +/// 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 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. + 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 { + inflight: Mutex::new(BTreeMap::new()), + admission: Mutex::new(BTreeMap::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: 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 + // 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), + )); + } + + // 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() { + 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. 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" + ), + )) + } + } + } + (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) { + 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(), + )); + } + + // 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(), + 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. Roll both 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, + to_logical: target.logical_id(), + 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 in this namespace for the life + // of the process. + generation, + }; + self.inflight.lock().insert(key.clone(), entry.clone()); + + if target.tx.try_send(text).is_err() { + // Disconnected or backpressured beyond its queue: roll back. + // + // 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. + // + // 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( + codes::TARGET_DISCONNECTED, + "target disconnected or unresponsive during routing", + )); + } + + info!( + delegation = %entry.delegation_id, + admission = entry.generation, + from = %entry.from_logical, + to = %entry.to_logical, + chain = ?entry.chain, + deadline = %entry.deadline, + "delegation routed" + ); + + DelegateOutcome::Accepted(DelegateAck { + delegation_id: params.delegation_id, + admission: generation, + assigned_to: target.logical_id(), + }) + } + + /// 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 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. 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) { + Some(i) => i.namespace, + None => return Claim::Unregistered, + }; + let key = DelegationKey::new(&namespace, delegation_id); + let mut g = self.inflight.lock(); + let (owner_handle, live) = match g.get(&key) { + Some(e) => (e.from_handle, e.generation), + None => return Claim::NotFound { namespace }, + }; + if owner_handle != handle { + return Claim::WrongOwner { + namespace, + 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) + } + + /// Phase 1 of a completion: snapshot the entry for `(namespace, + /// 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 + /// the commit ever being able to claim a different admission. + fn peek_for_completion( + &self, + 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::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, + } + } + + /// 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 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 + /// contract): its teardown runs `fail_instance`, which releases capacity + /// exactly once and sends `cp/cancel` to the serving runtime. + /// + /// 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 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 + /// to a genuine result or to the deadline sweep. + pub fn complete( + &self, + registry: &Registry, + serving_handle: u64, + mut params: DelegateResultParams, + max_result_bytes: usize, + next_rpc_id: u64, + ) -> 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 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 + // 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()); + 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(out); + } + } + + 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")), + ); + 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. 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 = %entry.to_logical, + to = %entry.from_logical, + "delegation completed" + ); + true + } + Commit::Vanished => { + // Concurrent removal (duplicate result, cancel, sweep, or + // fail_instance): whoever removed it released the capacity. + info!( + delegation = %params.delegation_id, + namespace = %entry.namespace, + "entry removed concurrently after delivery — capacity already released" + ); + false + } + 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 + /// to the serving runtime, if the delegation is in flight, owned by the + /// caller, and the caller named its live admission. + /// + /// 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, + from_handle: u64, + params: &CancelParams, + next_rpc_id: u64, + ) -> Result, ErrorObject> { + let refused = || { + ErrorObject::new( + codes::POLICY_DENIED, + "delegation is not in flight for this instance", + ) + }; + let entry = match self.claim( + registry, + from_handle, + ¶ms.delegation_id, + params.admission, + ) { + Claim::Owned(entry) => entry, + 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()); + } + 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, + 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()); + } + }; + registry.adjust_sessions(entry.to_handle, -1); + 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, + 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 keys: Vec = g + .iter() + .filter(|(_, e)| e.to_handle == handle || e.from_handle == handle) + .map(|(k, _)| k.clone()) + .collect(); + keys.iter().filter_map(|k| g.remove(k)).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(), + // 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)), + }; + 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(), + // 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( + 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 keys: Vec = g + .iter() + .filter(|(_, e)| e.deadline <= now) + .map(|(k, _)| k.clone()) + .collect(); + keys.iter().filter_map(|k| g.remove(k)).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(), + // 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()), + }; + 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(), + // 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( + 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). Delegation + /// 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() + .get(&DelegationKey::new(namespace, delegation_id)) + .map(|e| e.chain.clone()) + } + + pub fn inflight_count(&self) -> usize { + self.inflight.lock().len() + } +} + +/// 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() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::proto::TargetSelector; + use crate::registry::{outbound_channel, FrameRx}; + use chrono::Duration; + use std::time::Instant; + + 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, FrameRx) { + let (tx, rx) = outbound_channel(1024 * 1024); + ( + 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, + ) + } + + /// 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(), + target: TargetSelector { + name: Some(target.into()), + labels: None, + }, + 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) + } + } + + struct World { + cfg: CpConfig, + registry: Registry, + router: Router, + h_primary: u64, + h_worker: u64, + worker_rx: FrameRx, + primary_rx: FrameRx, + } + + 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, + ) + } + + /// 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(), + admission: a.admission, + 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 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 + // 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(), + admission: second.admission, + 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(); + 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"])); + // 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, + }; + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, 1024, 2), + CompleteOutcome::Delivered { committed: true } + ); + 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); + } + + #[test] + fn inflight_exists_before_target_receives_frame() { + // 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)), + DelegateOutcome::Accepted(_) + )); + // 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, + }; + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, 1024, 2), + CompleteOutcome::Delivered { committed: true } + ); + 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 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(); + 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", tok, "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(); + 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_entry_untouched() { + 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(), + admission: token(&w.router, "prod", "d-1"), + status: DelegationStatus::Completed, + result: Some("spoofed".into()), + error: None, + }; + // h_primary is a valid handle but NOT the serving instance. + assert_eq!( + w.router.complete(&w.registry, w.h_primary, result, 1024, 2), + CompleteOutcome::Dropped + ); + 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(), + // 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, + }; + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, 1024, 2), + CompleteOutcome::Dropped + ); + } + + #[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(), + admission: token(&w.router, "prod", "d-1"), + status: DelegationStatus::Completed, + result: Some("x".repeat(200)), + error: None, + }; + let cap = 96usize; + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result, cap, 2), + CompleteOutcome::Delivered { committed: true } + ); + 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")); + 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(), + admission: token(&w.router, "prod", "d-2"), + status: DelegationStatus::Completed, + result: Some("y".repeat(100)), + error: None, + }; + assert_eq!( + w.router.complete(&w.registry, w.h_worker, result2, 8, 3), + CompleteOutcome::Delivered { committed: true } + ); + 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); + } + + #[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(), + // 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 + .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); + + // 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(), + vec!["prod/koudu".to_string()] + ); + + // 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. 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, + "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 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. 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, + "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"), + } + } + + /// 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 { + DelegateResultParams { + delegation_id: id.into(), + admission, + status: DelegationStatus::Completed, + result: Some(body.into()), + error: None, + } + } + + #[test] + fn wrong_handle_result_never_hides_the_genuine_one() { + // 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(); + let tok = token(&w.router, "prod", "d-1"); + + if spoof_first { + // h_primary is registered but is NOT the serving instance. + assert_eq!( + w.router.complete( + &w.registry, + w.h_primary, + result_of("d-1", tok, "spoofed"), + 1024, + 2 + ), + CompleteOutcome::Dropped + ); + assert_eq!( + w.router.inflight_count(), + 1, + "a non-owner frame must not remove the entry" + ); + } + + assert_eq!( + w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", tok, "genuine"), + 1024, + 3, + ), + CompleteOutcome::Delivered { committed: true }, + "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_eq!( + w.router.complete( + &w.registry, + w.h_primary, + result_of("d-1", tok, "spoofed"), + 1024, + 4 + ), + CompleteOutcome::Dropped + ); + assert_eq!(w.router.inflight_count(), 0); + } + } + } + + #[test] + fn genuine_result_survives_concurrent_non_owner_frames() { + // 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 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(|| { + gate.wait(); + // Registered, but not the serving instance. + w.router.complete( + &w.registry, + w.h_primary, + result_of("d-1", tok, "spoofed"), + 1024, + 2, + ) == CompleteOutcome::Delivered { committed: true } + }); + gate.wait(); + let genuine = w.router.complete( + &w.registry, + w.h_worker, + result_of("d-1", tok, "genuine"), + 1024, + 3, + ) == CompleteOutcome::Delivered { committed: true }; + (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() { + // 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(), + admission: token(&w.router, "prod", "d-1"), + 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() { + // 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(), + admission: token(&w.router, "prod", "d-1"), + 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() { + // `cp/cancel` must not be an existence oracle — + // 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(); + 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. + 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(); + // 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!( + 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] + fn same_delegation_id_in_two_namespaces_is_independent() { + // 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, 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, 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); + 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(), + 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 + .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. + assert_eq!( + 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(); + 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" + ); + + assert_eq!( + 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(); + assert!(frame.contains("prod-done")); + assert_eq!(router.inflight_count(), 0); + } + + #[test] + fn parent_lookup_is_namespace_scoped() { + // 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); + + // 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. 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, + "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 ok_child = child_params("d-child", "worker-2", 60, "d-root", root.admission); + 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()] + ); + } + + /// 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 { + 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"), + } + } + + /// 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(), + admission: token(&w.router, "prod", id), + 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", b.generation, "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", a.generation, "late"), + 1024, + 8 + ), + CompleteOutcome::Dropped + ); + assert!(w.primary_rx.try_recv().is_err()); + } + + #[test] + 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 + .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 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 new file mode 100644 index 000000000..c4ef4f76c --- /dev/null +++ b/crates/openab-cp/src/server.rs @@ -0,0 +1,1418 @@ +//! 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: 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 — 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`, +//! 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::{close_code, CloseFrame, 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 parking_lot::Mutex; +use tokio::sync::watch; +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::{outbound_channel, shutdown_signal, Instance, Registry}; +use crate::router::{CompleteOutcome, DelegateOutcome, Router}; + +pub struct AppState { + pub cfg: CpConfig, + 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. + conns: Mutex>, +} + +impl AppState { + pub fn new(cfg: CpConfig) -> Self { + Self { + cfg, + 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. + 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. +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 { + AxumRouter::new() + .route("/cp", get(ws_handler)) + .route("/health", get(health)) + .with_state(state) +} + +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"; +/// 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" +/// 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(), + })) +} + +/// 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, + 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(); + } + }; + // Per-identity connection quota, taken before the upgrade so an + // over-quota peer is refused at the HTTP layer. + 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" + ); + // 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; + ws.max_message_size(max_frame) + .max_frame_size(max_frame) + .on_upgrade(move |socket| handle_connection(state, socket, identity, permit)) +} + +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. + _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 + // the transport alive but do not extend this deadline. + 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 _ = send_bounded( + &mut sink, + policy_close(REASON_REGISTER_TIMEOUT), + write_timeout, + &mut shutdown_rx, + ) + .await; + return; + } + }; + let (reg, reg_rpc_id) = match parse_register(®ister, &identity) { + Ok(ok) => ok, + Err((id, err)) => { + let resp = JsonRpcErrorResponse::new(id, err); + let _ = send_bounded( + &mut sink, + Message::Text(serde_json::to_string(&resp).expect("serializable").into()), + write_timeout, + &mut shutdown_rx, + ) + .await; + return; + } + }; + + // 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); + + // 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( + 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), + ); + // 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, + 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 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; + } + // `_registered` runs teardown on the way out. + 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 + // 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_close_reason = *shutdown_rx.borrow_and_update(); + break; + } + outbound = rx.recv() => { + match outbound { + Some(text) => write_or_break!(Message::Text(text.into())), + None => break, + } + } + inbound = stream.next() => { + match inbound { + Some(Ok(Message::Text(text))) => { + if let Some(reply) = handle_frame(&state, handle, &text) { + 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)) => { + warn!(handle, err = %e, "WS error"); + break; + } + } + } + } + } + + if let Some(reason) = cp_close_reason { + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + handle, + reason, + "closing connection at the CP's request" + ); + // 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 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 +/// 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(); + 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`. + // + // The registration can be gone while this task is still running: the + // sweeper deregisters an expired lease and signals the connection, but the + // signal is observed asynchronously, so frames already in flight land + // here first. Answer them instead of dropping them silently — a client + // whose heartbeat vanished into nothing cannot tell a swept lease from a + // hung CP, whereas NOT_REGISTERED tells it exactly what to do (reconnect + // and register again; registration is first-frame-only). + let me = match state.registry.get(handle) { + Some(i) => 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) => { + 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); + match state.router.complete( + &state.registry, + handle, + p, + state.cfg.max_result_bytes, + state.next_rpc_id(), + ) { + 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")) + } + // 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")) + } + } + } + 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")) + } + } +} + +/// One lease-expiry pass: drop registrations whose lease elapsed, close their +/// connections, and fail their in-flight delegations. +/// +/// Signalling the connection is what makes the deregistration complete: +/// without it the connection task keeps running against +/// a registration that no longer exists — every later frame (heartbeats +/// included) is answered `NOT_REGISTERED` at best, and the client cannot +/// re-register because registration is first-frame-only. +pub fn sweep_leases(state: &Arc, 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, REASON_LEASE_EXPIRED); + 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(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tick.tick().await; + + sweep_leases(&state, Duration::from_secs(state.cfg.lease_expiry_secs)); + + // 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: the envelope is validated, not assumed. + 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); + } + + 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() { + // 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); + } + + /// 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 + // 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) = outbound_channel(1024 * 1024); + 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().is_none()); + + sweep_leases(&state, Duration::ZERO); + assert!(state.registry.get(handle).is_none()); + observer.changed().await.unwrap(); + assert_eq!( + *observer.borrow(), + Some(REASON_LEASE_EXPIRED), + "the owning connection must be told to close, and why" + ); + } + + #[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) = outbound_channel(1024 * 1024); + 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); + } + + #[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 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) = outbound_channel(FILLER.len()); + 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) = outbound_channel(1024 * 1024); + 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"); + 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 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", + "admission": admission, + "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/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs new file mode 100644 index 000000000..d2ebabe74 --- /dev/null +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -0,0 +1,746 @@ +//! End-to-end WebSocket lifecycle tests: connection termination on lease +//! 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 +//! 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"; +const KEY_WORKER: &str = "k-worker"; + +type Ws = WebSocketStream>; + +fn cfg(extra: &str) -> CpConfig { + let raw = format!( + r#" +{extra} + +[[agents]] +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"); + 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_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"), + ); + 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 { + 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() +} + +/// 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 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 + } + } + 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() { + // 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, + // 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"); + 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" + ); + + 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); + + // 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() { + // 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 = None; + while started.elapsed() < Duration::from_secs(10) { + let _ = ws.send(Message::Ping(vec![7].into())).await; + if let Some(how) = wait_closed(&mut ws, Duration::from_millis(250)).await { + closed = Some(how); + break; + } + } + let closed = closed.expect("an authenticated socket that never registers must be closed"); + assert_eq!( + 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() { + // 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() { + // The quota bounds concurrent sockets per identity + // and is released on every exit path (RAII), so connect → disconnect → + // 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", + )) + .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" + ); + 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"), + } + + 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() { + // 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); + 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"), + } +} + +/// 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() +} + +/// 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. + // 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", "admission": a, "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` + // 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 +# 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; + + // 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. 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 + .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(), + v["params"]["admission"] + .as_u64() + .expect("the forwarded frame must carry the admission token"), + )); + } + 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, admission)) 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, + "admission": admission, + "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 7a4a253d3..6efbbf146 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 @@ -180,6 +180,12 @@ Agent A ◄──── result ◄──────── OAB-A ◄──── ### Delegate frame +What the initiator sends. Note what is *absent*: there is no `chain` field. +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. This example is a *parented* delegation — a root one simply omits +both parent fields. + ```json { "method": "cp/delegate", @@ -187,20 +193,56 @@ Agent A ◄──── result ◄──────── OAB-A ◄──── "delegation_id": "d-01J...", "target": { "name": "worker-1" }, "prompt": "…", - "chain": ["koudu"], + "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) -- `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` + `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. +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 @@ -209,6 +251,329 @@ 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. + +### 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 +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 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, + 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 — 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 + **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 + 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 + 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. Best-effort frames (`cp/cancel`, sweep-synthesized + `timeout`) remain fire-and-forget: the propagated deadline is their backstop. +- **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; + - `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); + - 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 + 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). + + 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. + + 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 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, 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 + 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 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 + 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. + `NO_TARGET` (nothing matches) is a distinct error. +- **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 + 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 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. + + --- ## 5. Delegation Policy @@ -260,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 @@ -281,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 @@ -296,11 +673,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 @@ -365,14 +755,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. diff --git a/docs/control-plane.md b/docs/control-plane.md new file mode 100644 index 000000000..c55b7600f --- /dev/null +++ b/docs/control-plane.md @@ -0,0 +1,116 @@ +# 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. +- `[[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, 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 + +`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. +- 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 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). +- **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. +- **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`, 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. 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 + 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.