diff --git a/Cargo.lock b/Cargo.lock index 89f00d9c7..da40e70d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2575,6 +2575,27 @@ 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", + "uuid", +] + [[package]] name = "openab-gateway" version = "0.5.4" diff --git a/Cargo.toml b/Cargo.toml index 5a834a13f..191c2640f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["crates/openab-core", "crates/openab-gateway", "crates/openab-mcp"] +members = ["crates/openab-core", "crates/openab-gateway", "crates/openab-mcp", "crates/openab-cp"] exclude = ["openab-agent", "crates/platform-schema"] [package] diff --git a/Dockerfile b/Dockerfile index 36bf49ca5..540ba829f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,16 +11,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && \ if [ "$BUILD_MODE" = "unified" ]; then \ cargo build --release --features unified; \ elif [ -n "$FEATURES" ]; then \ diff --git a/Dockerfile.agentcore b/Dockerfile.agentcore index 7bdd6d70b..d2b153fc9 100644 --- a/Dockerfile.agentcore +++ b/Dockerfile.agentcore @@ -9,16 +9,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release --features agentcore \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release --features agentcore +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release --features agentcore # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.antigravity b/Dockerfile.antigravity index 77aef681e..2003d59f2 100644 --- a/Dockerfile.antigravity +++ b/Dockerfile.antigravity @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Build agy-acp adapter --- FROM rust:1-bookworm AS adapter-builder diff --git a/Dockerfile.builder b/Dockerfile.builder index ff344e5e4..c3122c707 100644 --- a/Dockerfile.builder +++ b/Dockerfile.builder @@ -23,20 +23,22 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml # 2. Dummy sources for dep-only build -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src # 3. Copy real sources and build COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && \ if [ "$BUILD_MODE" = "unified" ]; then \ cargo build --release --features unified; \ elif [ -n "$FEATURES" ]; then \ diff --git a/Dockerfile.claude b/Dockerfile.claude index aa68fe3a2..8b47f8975 100644 --- a/Dockerfile.claude +++ b/Dockerfile.claude @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.codex b/Dockerfile.codex index 4d0f35c44..512e31506 100644 --- a/Dockerfile.codex +++ b/Dockerfile.codex @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.copilot b/Dockerfile.copilot index e28a29089..bcef23aca 100644 --- a/Dockerfile.copilot +++ b/Dockerfile.copilot @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.cursor b/Dockerfile.cursor index 03e9db910..e82d71ed2 100644 --- a/Dockerfile.cursor +++ b/Dockerfile.cursor @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.devin b/Dockerfile.devin index e25dadb60..b1f6c0160 100644 --- a/Dockerfile.devin +++ b/Dockerfile.devin @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.gateway b/Dockerfile.gateway index 7b3741ff1..cce6dde33 100644 --- a/Dockerfile.gateway +++ b/Dockerfile.gateway @@ -17,17 +17,19 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo 'fn main() {}' > crates/openab-gateway/src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release -p openab-gateway \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch crates/openab-gateway/src/main.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && \ +RUN touch crates/openab-gateway/src/main.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && \ if [ -n "$FEATURES" ]; then \ cargo build --release -p openab-gateway --no-default-features --features "$FEATURES"; \ else \ diff --git a/Dockerfile.gemini b/Dockerfile.gemini index 3de113f06..6ed4c3bf1 100644 --- a/Dockerfile.gemini +++ b/Dockerfile.gemini @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.grok b/Dockerfile.grok index 0708c293f..ad12d6210 100644 --- a/Dockerfile.grok +++ b/Dockerfile.grok @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM debian:trixie-slim diff --git a/Dockerfile.hermes b/Dockerfile.hermes index 2217dbe65..bdf0b8eea 100644 --- a/Dockerfile.hermes +++ b/Dockerfile.hermes @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM python:3.12-slim-trixie diff --git a/Dockerfile.kimi b/Dockerfile.kimi index 85ffc7ca5..3bcd61f37 100644 --- a/Dockerfile.kimi +++ b/Dockerfile.kimi @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs \ && cargo build --release # --- Runtime stage --- diff --git a/Dockerfile.mimocode b/Dockerfile.mimocode index bdb50e952..df8331324 100644 --- a/Dockerfile.mimocode +++ b/Dockerfile.mimocode @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- # MiMo-Code (https://github.com/XiaomiMiMo/MiMo-Code) is a fork of OpenCode diff --git a/Dockerfile.native b/Dockerfile.native index f66b268ba..5de7f6ea7 100644 --- a/Dockerfile.native +++ b/Dockerfile.native @@ -5,17 +5,19 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml COPY openab-agent/ openab-agent/ -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release RUN cd openab-agent && cargo build --release # --- Runtime stage --- diff --git a/Dockerfile.opencode b/Dockerfile.opencode index c1f239f5e..7a7063575 100644 --- a/Dockerfile.opencode +++ b/Dockerfile.opencode @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- # node:22-trixie-slim mirrors the base image used by Dockerfile.claude, diff --git a/Dockerfile.pi b/Dockerfile.pi index 0d992f4af..ae5fb99cc 100644 --- a/Dockerfile.pi +++ b/Dockerfile.pi @@ -5,16 +5,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs && cargo build --release +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs && cargo build --release # --- Runtime stage --- FROM node:22-trixie-slim diff --git a/Dockerfile.unified b/Dockerfile.unified index 0c734ebe2..540786a97 100644 --- a/Dockerfile.unified +++ b/Dockerfile.unified @@ -31,16 +31,18 @@ COPY Cargo.toml Cargo.lock ./ COPY crates/openab-core/Cargo.toml crates/openab-core/Cargo.toml COPY crates/openab-gateway/Cargo.toml crates/openab-gateway/Cargo.toml COPY crates/openab-mcp/Cargo.toml crates/openab-mcp/Cargo.toml -RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src \ +COPY crates/openab-cp/Cargo.toml crates/openab-cp/Cargo.toml +RUN mkdir -p src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src \ && echo 'fn main() {}' > src/main.rs \ && echo '' > crates/openab-core/src/lib.rs \ && echo '' > crates/openab-gateway/src/lib.rs \ && echo '' > crates/openab-mcp/src/lib.rs \ + && echo '' > crates/openab-cp/src/lib.rs \ && cargo build --release --features unified \ - && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src + && rm -rf src crates/openab-core/src crates/openab-gateway/src crates/openab-mcp/src crates/openab-cp/src COPY crates/ crates/ COPY src/ src/ -RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs \ +RUN touch src/main.rs crates/openab-core/src/lib.rs crates/openab-gateway/src/lib.rs crates/openab-mcp/src/lib.rs crates/openab-cp/src/lib.rs \ && cargo build --release --features unified # Build openab-agent (used by native variant) — copied late to avoid cache busts COPY openab-agent/ openab-agent/ diff --git a/crates/openab-cp/Cargo.toml b/crates/openab-cp/Cargo.toml new file mode 100644 index 000000000..debff2e93 --- /dev/null +++ b/crates/openab-cp/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "openab-cp" +version = "0.1.0" +edition = "2021" +license = "MIT" +description = "OpenAB Agent Control Plane — registry, router, and policy for direct inter-agent delegation" + +[dependencies] +tokio = { version = "1", features = ["full"] } +axum = { version = "0.8", features = ["ws"] } +futures-util = "0.3" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +toml = "0.8" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +anyhow = "1" +uuid = { version = "1", features = ["v4"] } +chrono = { version = "0.4", features = ["serde"] } +parking_lot = "0.12" +clap = { version = "4", features = ["derive"] } +subtle = "2" + +[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..9b5df7943 --- /dev/null +++ b/crates/openab-cp/cp.toml.example @@ -0,0 +1,60 @@ +# 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 + +[[agents]] +key = "${CP_KEY_KOUDU}" # per-agent secret, never shared +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "${CP_KEY_WORKER1}" +namespace = "prod" +name = "worker-1" +type = "worker" +max_delegated_sessions_cap = 4 # CP-side clamp on advertised capacity + +# Per-namespace policy. Absent namespaces use the conservative defaults: +# max_depth = 1, allow_worker_initiation = false. +[namespaces.prod] +max_depth = 1 +allow_worker_initiation = false diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs new file mode 100644 index 000000000..21d672b3b --- /dev/null +++ b/crates/openab-cp/src/config.rs @@ -0,0 +1,395 @@ +//! CP-side configuration. +//! +//! Identity binding is the security core (review F1 on the ADR): every auth +//! key maps to **immutable claims** (`namespace`, `name`, `type`, optional +//! caps) owned by CP config. Registration frames are verified against these +//! claims — never the other way around. A compromised runtime cannot escalate +//! to another namespace or to `primary` by editing its own config. + +use anyhow::{bail, Context, Result}; +use serde::Deserialize; +use std::collections::BTreeMap; + +use crate::proto::AgentType; + +#[derive(Debug, Clone, Deserialize)] +pub struct CpConfig { + /// Bind address. 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 (review F5). + #[serde(default = "default_max_frame_bytes")] + pub max_frame_bytes: usize, + + /// Maximum `cp/delegate.prompt` size in bytes; oversized prompts are + /// rejected (unlike results, nothing has run yet). + #[serde(default = "default_max_prompt_bytes")] + pub max_prompt_bytes: usize, + + /// Deadline for the mandatory `cp/register` first frame, in seconds from + /// the completed WebSocket upgrade. A connection that authenticates but + /// never registers is closed when this elapses (review round-3 F4): + /// otherwise an authenticated peer could park unlimited sockets in the + /// pre-registration state, keeping them alive with pings forever. + #[serde(default = "default_register_timeout_secs")] + pub register_timeout_secs: u64, + + /// Maximum simultaneous connections per identity, counted from the + /// upgrade (so pre-registration sockets count too) and released on every + /// exit path (review round-3 F4). Replicas of one logical agent share one + /// identity, so this is the replica ceiling as well. + #[serde(default = "default_max_connections_per_identity")] + pub max_connections_per_identity: u32, + + /// Identity table: auth key → immutable claims. + /// Keyed by the key id (`kid`), with the secret alongside, so logs can + /// reference identities without printing secrets. + #[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 +} + +/// Immutable identity claims bound to one auth key. +#[derive(Debug, Clone, Deserialize)] +pub struct AgentIdentity { + /// The secret presented by the runtime (`OPENAB_CP_KEY`). Supports + /// `${ENV_VAR}` expansion so the config file itself holds no secrets. + pub key: String, + pub namespace: String, + pub name: String, + #[serde(rename = "type")] + pub agent_type: AgentType, + /// Optional CP-side clamp on the advertised concurrency budget. + #[serde(default)] + pub max_delegated_sessions_cap: Option, +} + +/// Per-namespace delegation policy. Defaults are the conservative ADR §5 +/// baseline; relaxation is CP-side config only. +#[derive(Debug, Clone, Deserialize)] +pub struct NamespacePolicy { + /// Maximum delegation chain depth (1 = primary → worker only). + #[serde(default = "default_depth")] + pub max_depth: u32, + /// Whether workers may initiate delegations (depth still applies). + #[serde(default)] + pub allow_worker_initiation: bool, +} + +fn default_depth() -> u32 { + 1 +} + +impl Default for NamespacePolicy { + fn default() -> Self { + Self { + max_depth: default_depth(), + allow_worker_initiation: false, + } + } +} + +impl CpConfig { + pub fn load(path: &str) -> Result { + let raw = + std::fs::read_to_string(path).with_context(|| format!("reading CP config {path}"))?; + let expanded = expand_env(&raw); + let cfg: CpConfig = toml::from_str(&expanded).context("parsing CP config")?; + cfg.validate()?; + Ok(cfg) + } + + pub fn validate(&self) -> Result<()> { + let mut seen_keys = std::collections::BTreeSet::new(); + let mut seen_names = std::collections::BTreeSet::new(); + for a in &self.agents { + if a.key.trim().is_empty() { + bail!("agent {}/{} has an empty key", a.namespace, a.name); + } + if !seen_keys.insert(a.key.as_str()) { + bail!( + "duplicate auth key (shared keys defeat per-agent revocation); \ + offending identity: {}/{}", + a.namespace, + a.name + ); + } + if !seen_names.insert((a.namespace.as_str(), a.name.as_str())) { + bail!( + "duplicate identity {}/{} — replicas share one identity (one key), \ + distinguished at registration by instance_id", + a.namespace, + a.name + ); + } + } + if self.lease_expiry_secs <= self.heartbeat_interval_secs { + bail!("lease_expiry_secs must exceed heartbeat_interval_secs"); + } + // Admission bounds must actually bound something (review round-3 F4): + // zero would mean "no registration deadline" / "no connection allowed". + if self.register_timeout_secs == 0 { + bail!("register_timeout_secs must be greater than 0"); + } + if self.max_connections_per_identity == 0 { + bail!("max_connections_per_identity must be at least 1"); + } + // Bearer keys over cleartext TCP must never reach an untrusted + // network: non-loopback binds require the explicit override + // (review round-2 F4). + 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() + } +} + +/// 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() { + // Review round-3 F4: absent fields keep working (serde defaults) and + // a zero bound is rejected rather than silently disabling the guard. + let cfg: CpConfig = toml::from_str(base_toml()).unwrap(); + cfg.validate().unwrap(); + assert_eq!(cfg.register_timeout_secs, 10); + assert_eq!(cfg.max_connections_per_identity, 8); + + let explicit: CpConfig = + toml::from_str("register_timeout_secs = 3\nmax_connections_per_identity = 2").unwrap(); + explicit.validate().unwrap(); + assert_eq!(explicit.register_timeout_secs, 3); + assert_eq!(explicit.max_connections_per_identity, 2); + + for bad in [ + "register_timeout_secs = 0", + "max_connections_per_identity = 0", + ] { + let cfg: CpConfig = toml::from_str(bad).unwrap(); + assert!(cfg.validate().is_err(), "{bad} must be rejected"); + } + } + + #[test] + fn env_expansion() { + std::env::set_var("CP_TEST_KEY_XYZ", "sekrit"); + assert_eq!( + expand_env("key = \"${CP_TEST_KEY_XYZ}\""), + "key = \"sekrit\"" + ); + assert_eq!(expand_env("no vars"), "no vars"); + assert_eq!(expand_env("${UNSET_VAR_ABC123}"), ""); + } +} diff --git a/crates/openab-cp/src/lib.rs b/crates/openab-cp/src/lib.rs new file mode 100644 index 000000000..f4a6ac14f --- /dev/null +++ b/crates/openab-cp/src/lib.rs @@ -0,0 +1,14 @@ +//! openab-cp — OpenAB Agent Control Plane. +//! +//! Hub-and-spoke registration and routing for direct agent-to-agent +//! delegation, per `docs/adr/agent-control-plane.md`. Runtimes dial out and +//! register over WebSocket; the CP authenticates them against config-bound +//! identities, enforces delegation policy authoritatively, and routes +//! `cp/delegate` / `cp/delegate_result` frames between them. + +pub mod config; +pub mod policy; +pub mod proto; +pub mod registry; +pub mod router; +pub mod server; diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs new file mode 100644 index 000000000..7a34753c3 --- /dev/null +++ b/crates/openab-cp/src/main.rs @@ -0,0 +1,56 @@ +//! Standalone control-plane binary. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use clap::Parser; +use tracing::info; + +use openab_cp::config::CpConfig; +use openab_cp::server::{app, run_sweeper, AppState}; + +#[derive(Parser)] +#[command(name = "openab-cp", about = "OpenAB Agent Control Plane")] +struct Cli { + /// Path to the CP config file (TOML). + #[arg(short, long, default_value = "cp.toml")] + config: String, +} + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), + ) + .init(); + + let cli = Cli::parse(); + let cfg = CpConfig::load(&cli.config)?; + let listen = cfg.listen.clone(); + if cfg.agents.is_empty() { + tracing::warn!("no [[agents]] identities configured — every connection will be rejected"); + } + info!( + listen = %listen, + identities = cfg.agents.len(), + namespaces = cfg.namespaces.len(), + "starting openab-cp" + ); + + let state = Arc::new(AppState::new(cfg)); + let sweeper = tokio::spawn(run_sweeper(state.clone())); + + let listener = tokio::net::TcpListener::bind(&listen) + .await + .with_context(|| format!("binding {listen}"))?; + axum::serve(listener, app(state)) + .with_graceful_shutdown(async { + let _ = tokio::signal::ctrl_c().await; + info!("shutdown signal received"); + }) + .await?; + + sweeper.abort(); + Ok(()) +} diff --git a/crates/openab-cp/src/policy.rs b/crates/openab-cp/src/policy.rs new file mode 100644 index 000000000..c8b485ca5 --- /dev/null +++ b/crates/openab-cp/src/policy.rs @@ -0,0 +1,246 @@ +//! CP-authoritative delegation policy (review F4 on the ADR). +//! +//! Every check here operates exclusively on CP-owned data: authenticated +//! identity claims, the CP-constructed ancestry chain, and CP config. Facade +//! checks in the runtime are defense-in-depth only; nothing in this module +//! trusts a client-supplied policy input. + +use chrono::{DateTime, Utc}; + +use crate::config::NamespacePolicy; +use crate::proto::AgentType; + +pub struct PolicyInput<'a> { + /// Authenticated initiator identity. + pub from_namespace: &'a str, + pub from_name: &'a str, + pub from_type: &'a AgentType, + /// Target namespace (v1: always the initiator's namespace; the router + /// resolves selectors within it. Cross-namespace grants are future work). + pub target_namespace: &'a str, + /// Logical name of the resolved target. + pub target_name: &'a str, + /// CP-constructed chain for the *parent* delegation (empty for a root + /// delegation). Elements are `namespace/name`. + pub parent_chain: &'a [String], + pub deadline: DateTime, + pub parent_deadline: Option>, + pub now: DateTime, + pub max_deadline_secs: u64, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum PolicyDenial { + WorkerInitiation, + DepthExceeded { max: u32, would_be: u32 }, + Cycle { target: String }, + CrossNamespace, + DeadlinePast, + DeadlineTooLong { max_secs: u64 }, + DeadlineExceedsParent, +} + +impl std::fmt::Display for PolicyDenial { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PolicyDenial::WorkerInitiation => { + write!(f, "workers may not initiate delegations in this namespace") + } + PolicyDenial::DepthExceeded { max, would_be } => write!( + f, + "delegation depth {would_be} exceeds namespace max_depth {max}" + ), + PolicyDenial::Cycle { target } => { + write!(f, "cycle: {target} is already in the delegation chain") + } + PolicyDenial::CrossNamespace => { + write!(f, "cross-namespace delegation is not granted") + } + PolicyDenial::DeadlinePast => write!(f, "deadline is in the past"), + PolicyDenial::DeadlineTooLong { max_secs } => { + write!(f, "deadline exceeds the CP cap of {max_secs}s") + } + PolicyDenial::DeadlineExceedsParent => { + write!(f, "child deadline exceeds the parent's remaining budget") + } + } + } +} + +/// Evaluate the full CP-side policy for one delegation attempt. +pub fn check(input: &PolicyInput<'_>, ns_policy: &NamespacePolicy) -> Result<(), PolicyDenial> { + // 1. Initiator role. + if *input.from_type == AgentType::Worker && !ns_policy.allow_worker_initiation { + return Err(PolicyDenial::WorkerInitiation); + } + + // 2. Namespace boundary (v1: strict). + if input.from_namespace != input.target_namespace { + return Err(PolicyDenial::CrossNamespace); + } + + // 3. Depth: the new chain would be parent_chain + initiator; its length + // equals the delegation depth (root delegation → depth 1). + let would_be = input.parent_chain.len() as u32 + 1; + if would_be > ns_policy.max_depth { + return Err(PolicyDenial::DepthExceeded { + max: ns_policy.max_depth, + would_be, + }); + } + + // 4. Cycle: target must not already be an ancestor (or the initiator). + let target_id = format!("{}/{}", input.target_namespace, input.target_name); + let from_id = format!("{}/{}", input.from_namespace, input.from_name); + if target_id == from_id || input.parent_chain.contains(&target_id) { + return Err(PolicyDenial::Cycle { target: target_id }); + } + + // 5. Deadline sanity and caps. + if input.deadline <= input.now { + return Err(PolicyDenial::DeadlinePast); + } + let remaining = (input.deadline - input.now).num_seconds(); + if remaining > input.max_deadline_secs as i64 { + return Err(PolicyDenial::DeadlineTooLong { + max_secs: input.max_deadline_secs, + }); + } + if let Some(parent) = input.parent_deadline { + if input.deadline > parent { + return Err(PolicyDenial::DeadlineExceedsParent); + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Duration; + + fn base<'a>(now: DateTime, chain: &'a [String]) -> PolicyInput<'a> { + PolicyInput { + from_namespace: "prod", + from_name: "koudu", + from_type: &AgentType::Primary, + target_namespace: "prod", + target_name: "worker-1", + parent_chain: chain, + deadline: now + Duration::seconds(60), + parent_deadline: None, + now, + max_deadline_secs: 1800, + } + } + + fn default_policy() -> NamespacePolicy { + NamespacePolicy::default() + } + + #[test] + fn root_primary_delegation_passes() { + let now = Utc::now(); + assert!(check(&base(now, &[]), &default_policy()).is_ok()); + } + + #[test] + fn worker_initiation_denied_by_default_allowed_by_config() { + let now = Utc::now(); + let chain: Vec = vec![]; + let mut input = base(now, &chain); + input.from_type = &AgentType::Worker; + assert_eq!( + check(&input, &default_policy()), + Err(PolicyDenial::WorkerInitiation) + ); + let relaxed = NamespacePolicy { + max_depth: 2, + allow_worker_initiation: true, + }; + assert!(check(&input, &relaxed).is_ok()); + } + + #[test] + fn depth_exceeded_at_default_depth_one() { + let now = Utc::now(); + let chain = vec!["prod/root".to_string()]; + let input = base(now, &chain); + assert_eq!( + check(&input, &default_policy()), + Err(PolicyDenial::DepthExceeded { + max: 1, + would_be: 2 + }) + ); + let relaxed = NamespacePolicy { + max_depth: 2, + allow_worker_initiation: true, + }; + assert!(check(&input, &relaxed).is_ok()); + } + + #[test] + fn cycle_rejected_including_self() { + let now = Utc::now(); + let chain = vec!["prod/worker-1".to_string()]; + let relaxed = NamespacePolicy { + max_depth: 5, + allow_worker_initiation: true, + }; + let input = base(now, &chain); + assert!(matches!( + check(&input, &relaxed), + Err(PolicyDenial::Cycle { .. }) + )); + // self-delegation + let chain2: Vec = vec![]; + let mut input2 = base(now, &chain2); + input2.target_name = "koudu"; + assert!(matches!( + check(&input2, &relaxed), + Err(PolicyDenial::Cycle { .. }) + )); + } + + #[test] + fn cross_namespace_denied() { + let now = Utc::now(); + let chain: Vec = vec![]; + let mut input = base(now, &chain); + input.target_namespace = "dev"; + assert_eq!( + check(&input, &default_policy()), + Err(PolicyDenial::CrossNamespace) + ); + } + + #[test] + fn deadline_rules() { + let now = Utc::now(); + let chain: Vec = vec![]; + + let mut past = base(now, &chain); + past.deadline = now - Duration::seconds(1); + assert_eq!( + check(&past, &default_policy()), + Err(PolicyDenial::DeadlinePast) + ); + + let mut long = base(now, &chain); + long.deadline = now + Duration::seconds(3600); + assert_eq!( + check(&long, &default_policy()), + Err(PolicyDenial::DeadlineTooLong { max_secs: 1800 }) + ); + + let mut over_parent = base(now, &chain); + over_parent.deadline = now + Duration::seconds(120); + over_parent.parent_deadline = Some(now + Duration::seconds(60)); + assert_eq!( + check(&over_parent, &default_policy()), + Err(PolicyDenial::DeadlineExceedsParent) + ); + } +} diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs new file mode 100644 index 000000000..91c4a27f4 --- /dev/null +++ b/crates/openab-cp/src/proto.rs @@ -0,0 +1,426 @@ +//! Control-plane wire protocol: JSON-RPC 2.0 envelopes and `cp/*` method +//! payloads, following the conventions of `openab-core/src/acp/protocol.rs`. +//! +//! ## Contract summary +//! +//! - Transport: one WebSocket per runtime, text frames, one JSON-RPC message +//! per frame. +//! - Every request carries `jsonrpc: "2.0"` and a `u64` id; responses echo the +//! id. Correlation of *delegations* (which span multiple request/response +//! pairs across two connections) uses `delegation_id`, never the JSON-RPC id. +//! - The first frame on a new connection MUST be `cp/register`. Anything else +//! is rejected with `NOT_REGISTERED` and the connection is closed. +//! - Delegation ancestry (`chain`) is **CP-constructed**: callers supply only +//! `parent_delegation_id`; the CP derives the chain from authenticated +//! identities and its in-flight table. A runtime cannot forge ancestry. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// Wire protocol version. Carried in `cp/register`; the CP rejects +/// registrations with a version it does not support. +pub const PROTOCOL_VERSION: u32 = 1; + +// --- JSON-RPC envelopes --- + +#[derive(Debug, Serialize)] +pub struct JsonRpcRequest { + pub jsonrpc: &'static str, + pub id: u64, + pub method: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub params: Option, +} + +impl JsonRpcRequest { + pub fn new(id: u64, method: impl Into, params: Option) -> Self { + Self { + jsonrpc: "2.0", + id, + method: method.into(), + params, + } + } +} + +#[derive(Debug, Serialize)] +pub struct JsonRpcResponse { + pub jsonrpc: &'static str, + pub id: u64, + pub result: Value, +} + +impl JsonRpcResponse { + pub fn new(id: u64, result: Value) -> Self { + Self { + jsonrpc: "2.0", + id, + result, + } + } +} + +#[derive(Debug, Serialize)] +pub struct JsonRpcErrorResponse { + pub jsonrpc: &'static str, + pub id: u64, + pub error: ErrorObject, +} + +impl JsonRpcErrorResponse { + pub fn new(id: u64, error: ErrorObject) -> Self { + Self { + jsonrpc: "2.0", + id, + error, + } + } +} + +/// Incoming message: request, response, or error — distinguished by fields. +#[derive(Debug, Deserialize)] +pub struct JsonRpcMessage { + pub jsonrpc: Option, + pub id: Option, + pub method: Option, + pub params: Option, + pub result: Option, + pub error: Option, +} + +impl JsonRpcMessage { + /// Validate this frame as a JSON-RPC 2.0 **request** (review F4): the + /// `jsonrpc` field must be exactly "2.0", and a request id must be + /// present (all `cp/*` client→CP methods are requests, not + /// notifications). Returns the request id. + pub fn require_request_envelope(&self) -> Result { + if self.jsonrpc.as_deref() != Some("2.0") { + return Err(ErrorObject::new( + codes::INVALID_REQUEST, + "jsonrpc must be \"2.0\"", + )); + } + match self.id { + Some(id) => Ok(id), + None => Err(ErrorObject::new( + codes::INVALID_REQUEST, + "cp/* methods are requests and require an id", + )), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ErrorObject { + pub code: i64, + pub message: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub data: Option, +} + +impl ErrorObject { + pub fn new(code: i64, message: impl Into) -> Self { + Self { + code, + message: message.into(), + data: None, + } + } +} + +// --- Error codes (application range, distinct and machine-actionable) --- + +pub mod codes { + /// Frame received before a successful `cp/register` on this connection. + pub const NOT_REGISTERED: i64 = -32001; + /// Auth key unknown, or registration claims do not match the identity + /// bound to the key. + pub const IDENTITY_MISMATCH: i64 = -32002; + /// Delegation denied by policy (initiator role, depth, cycle, namespace). + pub const POLICY_DENIED: i64 = -32003; + /// No registered, healthy runtime matches the target selector. + pub const NO_TARGET: i64 = -32004; + /// Matching targets exist but all are at their advertised capacity. + /// Explicit fast-fail: the CP never queues (v1 has no durable state). + pub const SATURATED: i64 = -32005; + /// Delegation deadline elapsed before a result frame arrived. + pub const DEADLINE_EXCEEDED: i64 = -32006; + /// Serving runtime disconnected while the delegation was in flight. + pub const TARGET_DISCONNECTED: i64 = -32007; + /// `delegation_id` already in flight (idempotency guard). + pub const DUPLICATE_DELEGATION: i64 = -32008; + /// `cp/register` carried an unsupported protocol version. + pub const UNSUPPORTED_VERSION: i64 = -32009; + /// Malformed params for an otherwise valid method. + pub const INVALID_PARAMS: i64 = -32602; + /// Invalid JSON-RPC 2.0 envelope (missing/wrong `jsonrpc`, missing id). + pub const INVALID_REQUEST: i64 = -32600; + /// Unknown method. + pub const METHOD_NOT_FOUND: i64 = -32601; +} + +// --- cp/register --- + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum AgentType { + Primary, + Worker, +} + +impl std::fmt::Display for AgentType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AgentType::Primary => write!(f, "primary"), + AgentType::Worker => write!(f, "worker"), + } + } +} + +/// Params of `cp/register`, the mandatory first frame. +/// +/// `namespace`, `name`, and `agent_type` are **assertions to be verified**, +/// not authorization inputs: the CP compares them against the immutable +/// claims bound to the presented auth key and rejects any mismatch with +/// `IDENTITY_MISMATCH`. They exist in the frame so a misconfigured runtime +/// fails loudly at registration instead of being silently re-identified. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterParams { + pub protocol_version: u32, + pub namespace: String, + pub name: String, + #[serde(rename = "type")] + pub agent_type: AgentType, + /// Runtime-generated per-process id; distinguishes replicas of the same + /// logical agent during rolling deploys. + pub instance_id: String, + #[serde(default)] + pub labels: std::collections::BTreeMap, + /// Advertised concurrency budget. The CP may clamp this to a + /// config-defined cap for the identity. + #[serde(default = "default_max_sessions")] + pub max_delegated_sessions: u32, +} + +fn default_max_sessions() -> u32 { + 1 +} + +/// Result of a successful `cp/register`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterAck { + pub protocol_version: u32, + /// Interval at which the runtime must send `cp/heartbeat`. + pub heartbeat_interval_secs: u64, + /// Lease duration; missing heartbeats past this window deregisters the + /// instance and fails its in-flight delegations. + pub lease_expiry_secs: u64, + /// The effective (possibly clamped) concurrency budget. + pub effective_max_delegated_sessions: u32, +} + +// --- cp/heartbeat --- + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HeartbeatParams { + pub instance_id: String, + /// Current number of delegated sessions the runtime is serving; lets the + /// CP correct drift in its own in-flight accounting. + #[serde(default)] + pub active_delegated_sessions: u32, +} + +// --- cp/delegate --- + +/// Target selector: exact logical name, or label match (all pairs must match). +/// Exactly one of the two must be set. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct TargetSelector { + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub labels: Option>, +} + +/// Params of `cp/delegate` as sent by the initiating runtime. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateParams { + /// Caller-generated unique id (idempotency key). The CP rejects a second + /// in-flight delegation with the same id. + pub delegation_id: String, + pub target: TargetSelector, + pub prompt: String, + /// Absolute RFC3339 deadline. Mandatory: the CP rejects missing, past, or + /// over-cap deadlines. A child deadline can never exceed the parent's + /// remaining budget. + pub deadline: chrono::DateTime, + /// If this delegation is issued while serving another delegation, the id + /// of that parent. The CP derives the ancestry chain from this — the + /// chain is never client-supplied. + #[serde(skip_serializing_if = "Option::is_none")] + pub parent_delegation_id: Option, +} + +/// Params of `cp/delegate` as forwarded to the serving runtime. The CP stamps +/// the authenticated origin and the CP-constructed chain. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateForward { + pub delegation_id: String, + pub prompt: String, + pub deadline: chrono::DateTime, + /// Authenticated identity of the initiating agent (`namespace/name`). + pub from: String, + /// CP-constructed delegation ancestry, root first. The serving runtime + /// can trust every element: each hop was authenticated by the CP. + pub chain: Vec, +} + +/// Immediate result of `cp/delegate` (routing acceptance, not completion). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateAck { + pub delegation_id: String, + /// The chosen serving instance's logical name (`namespace/name`). + pub assigned_to: String, +} + +// --- cp/delegate_result --- + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum DelegationStatus { + Completed, + Failed, + Timeout, + Cancelled, + TargetDisconnected, +} + +/// Params of `cp/delegate_result` — emitted by the serving **runtime** when +/// the agent's turn ends (protocol-mandatory; never depends on the model), +/// or synthesized by the CP on timeout/disconnect. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DelegateResultParams { + pub delegation_id: String, + pub status: DelegationStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +// --- cp/cancel --- + +/// Params of `cp/cancel`: from the initiator to abort an in-flight +/// delegation, or from the CP to the serving runtime (best effort) after a +/// timeout or initiator cancellation. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CancelParams { + pub delegation_id: String, + pub reason: String, +} + +// --- method names --- + +pub mod methods { + pub const REGISTER: &str = "cp/register"; + pub const HEARTBEAT: &str = "cp/heartbeat"; + pub const DELEGATE: &str = "cp/delegate"; + pub const DELEGATE_RESULT: &str = "cp/delegate_result"; + pub const CANCEL: &str = "cp/cancel"; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn register_params_roundtrip_with_type_rename() { + let json = serde_json::json!({ + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-abc", + "labels": {"backend": "kiro"}, + "max_delegated_sessions": 4 + }); + let p: RegisterParams = serde_json::from_value(json).unwrap(); + assert_eq!(p.agent_type, AgentType::Primary); + let back = serde_json::to_value(&p).unwrap(); + assert_eq!(back["type"], "primary"); + } + + #[test] + fn register_defaults_apply() { + let json = serde_json::json!({ + "protocol_version": 1, + "namespace": "prod", + "name": "w1", + "type": "worker", + "instance_id": "i-1" + }); + let p: RegisterParams = serde_json::from_value(json).unwrap(); + assert!(p.labels.is_empty()); + assert_eq!(p.max_delegated_sessions, 1); + } + + #[test] + fn delegate_params_require_deadline() { + let json = serde_json::json!({ + "delegation_id": "d-1", + "target": {"name": "w1"}, + "prompt": "hi" + }); + assert!(serde_json::from_value::(json).is_err()); + } + + #[test] + fn delegation_status_snake_case() { + assert_eq!( + serde_json::to_value(DelegationStatus::TargetDisconnected).unwrap(), + serde_json::json!("target_disconnected") + ); + } + + #[test] + fn incoming_message_distinguishes_request_and_response() { + let req: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","id":1,"method":"cp/heartbeat","params":{}}"#) + .unwrap(); + assert!(req.method.is_some() && req.result.is_none()); + let resp: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#).unwrap(); + assert!(resp.method.is_none() && resp.result.is_some()); + } + + #[test] + fn request_envelope_validation() { + let ok: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","id":7,"method":"cp/heartbeat"}"#).unwrap(); + assert_eq!(ok.require_request_envelope().unwrap(), 7); + + // Missing jsonrpc. + let no_ver: JsonRpcMessage = + serde_json::from_str(r#"{"id":7,"method":"cp/heartbeat"}"#).unwrap(); + assert_eq!( + no_ver.require_request_envelope().unwrap_err().code, + codes::INVALID_REQUEST + ); + + // Wrong version. + let bad_ver: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"1.0","id":7,"method":"cp/heartbeat"}"#).unwrap(); + assert_eq!( + bad_ver.require_request_envelope().unwrap_err().code, + codes::INVALID_REQUEST + ); + + // Notification shape (no id). + let no_id: JsonRpcMessage = + serde_json::from_str(r#"{"jsonrpc":"2.0","method":"cp/heartbeat"}"#).unwrap(); + assert_eq!( + no_id.require_request_envelope().unwrap_err().code, + codes::INVALID_REQUEST + ); + } +} diff --git a/crates/openab-cp/src/registry.rs b/crates/openab-cp/src/registry.rs new file mode 100644 index 000000000..dd7cd47cf --- /dev/null +++ b/crates/openab-cp/src/registry.rs @@ -0,0 +1,436 @@ +//! Instance registry: who is alive, with which claims, on which connection. +//! +//! Replica semantics (ADR §3): multiple instances may register under one +//! logical identity during rolling deploys. New delegations route to the +//! newest healthy instance; in-flight delegations complete on the instance +//! that accepted them. Lease expiry (missed heartbeats) deregisters an +//! instance and fails its in-flight delegations. + +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use parking_lot::RwLock; +use tokio::sync::{mpsc, watch}; + +use crate::proto::AgentType; + +/// Outbound frame sender for one WS connection (serialized JSON text). +/// Bounded: a peer that cannot drain its queue is disconnected rather than +/// growing CP memory (review F5). +pub type FrameTx = mpsc::Sender; + +/// Capacity of each per-connection outbound queue. +pub const OUTBOUND_QUEUE: usize = 256; + +/// Shutdown signal for one WS connection, held by the registry so the CP can +/// terminate a connection it no longer considers registered (review round-3 +/// F1: lease expiry must close the socket — otherwise the connection task +/// lives on with a registry entry that no longer exists, silently dropping +/// every subsequent frame and unable to re-register, since registration is +/// first-frame-only). +/// +/// `watch` (not `oneshot`) so the connection task can select on it repeatedly, +/// and wrapped in `Arc` so the registry entry and the connection task share +/// one signal without either side's drop cancelling it. +pub type ShutdownTx = Arc>; + +/// Create a fresh connection shutdown signal. The connection task keeps the +/// returned handle (to `subscribe()`), the registry keeps a clone. +pub fn shutdown_signal() -> ShutdownTx { + Arc::new(watch::channel(false).0) +} + +/// Registry slot: the public instance view plus CP-internal connection +/// control that is deliberately not part of `Instance` (nothing routed or +/// serialized should carry it). +struct Entry { + inst: Instance, + shutdown: ShutdownTx, +} + +/// A live, authenticated, registered runtime instance. +#[derive(Clone, Debug)] +pub struct Instance { + /// CP-generated registration handle — the registry key and the basis of + /// all ownership checks. Never client-supplied (review F1): a colliding + /// client `instance_id` cannot replace or tear down another identity's + /// registration. + pub handle: u64, + pub namespace: String, + pub name: String, + pub agent_type: AgentType, + /// Client-supplied replica discriminator (display/audit only; ownership + /// and teardown key on `handle`). + pub instance_id: String, + pub labels: BTreeMap, + pub max_delegated_sessions: u32, + /// Delegations currently routed to this instance. CP-owned and + /// authoritative — never merged from runtime reports (review F6). + pub active_sessions: u32, + pub registered_at: Instant, + pub last_heartbeat: Instant, + pub tx: FrameTx, +} + +impl Instance { + pub fn logical_id(&self) -> String { + format!("{}/{}", self.namespace, self.name) + } + + pub fn saturated(&self) -> bool { + self.active_sessions >= self.max_delegated_sessions + } + + fn matches_labels(&self, want: &BTreeMap) -> bool { + want.iter() + .all(|(k, v)| self.labels.get(k).map(|x| x == v).unwrap_or(false)) + } +} + +#[derive(Default)] +pub struct Registry { + /// Keyed by CP-generated registration handle. + inner: RwLock>, + next_handle: AtomicU64, +} + +impl Registry { + pub fn new() -> Self { + Self::default() + } + + /// Insert a newly registered instance under a fresh CP-generated handle + /// (returned). Re-registrations (reconnects) get a new handle; the stale + /// entry disappears when its socket closes or its lease expires — it can + /// never be replaced by another connection's registration. + /// + /// `shutdown` is the owning connection's termination signal: the CP + /// triggers it whenever it drops the registration on its own initiative + /// (lease expiry — review round-3 F1). + pub fn register_conn(&self, mut inst: Instance, shutdown: ShutdownTx) -> u64 { + let handle = self.next_handle.fetch_add(1, Ordering::Relaxed) + 1; + inst.handle = handle; + self.inner.write().insert(handle, Entry { inst, shutdown }); + handle + } + + /// Register an instance with a detached shutdown signal (no connection + /// task is listening). For tests and non-WS callers. + pub fn register(&self, inst: Instance) -> u64 { + self.register_conn(inst, shutdown_signal()) + } + + /// Ask the owning connection task to close. Returns whether a live + /// registration was signalled. Must be called BEFORE `deregister`, which + /// drops the registry's handle on the signal. + pub fn signal_shutdown(&self, handle: u64) -> bool { + match self.inner.read().get(&handle) { + Some(e) => { + // `send_replace` cannot fail even with no receivers left. + e.shutdown.send_replace(true); + true + } + None => false, + } + } + + /// Remove an instance by its registration handle (disconnect or lease + /// expiry). Only the owning connection or the sweeper knows the handle. + pub fn deregister(&self, handle: u64) -> Option { + self.inner.write().remove(&handle).map(|e| e.inst) + } + + /// Refresh the lease. The runtime-reported session count is intentionally + /// ignored: CP-owned in-flight accounting is authoritative (review F6 — + /// merging reports could pin an instance saturated forever). + pub fn heartbeat(&self, handle: u64) -> bool { + let mut g = self.inner.write(); + match g.get_mut(&handle) { + Some(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 (review F6): + /// - exact-name selection → replicas of one logical agent: newest + /// registration first (rolling-deploy rule), load as tie-breaker + /// - label selection → across logical agents: least loaded first, + /// registration recency as tie-breaker + pub fn select( + &self, + namespace: &str, + name: Option<&str>, + labels: Option<&BTreeMap>, + ) -> Result { + let g = self.inner.read(); + let mut matches: Vec<&Instance> = g + .values() + .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) = mpsc::channel(OUTBOUND_QUEUE); + Instance { + handle: 0, // assigned by register() + namespace: ns.into(), + name: name.into(), + agent_type: AgentType::Worker, + instance_id: id.into(), + labels: BTreeMap::new(), + max_delegated_sessions: max, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + } + } + + #[test] + fn select_by_name_and_namespace_isolation() { + let r = Registry::new(); + let h1 = r.register(inst("prod", "w1", "i-1", 2)); + r.register(inst("dev", "w1", "i-2", 2)); + let got = r.select("prod", Some("w1"), None).unwrap(); + assert_eq!(got.handle, h1); + assert!(matches!( + r.select("staging", Some("w1"), None), + Err(SelectError::NoTarget) + )); + } + + #[test] + fn replicas_route_to_newest() { + let r = Registry::new(); + r.register(inst("prod", "w1", "i-old", 2)); + std::thread::sleep(Duration::from_millis(5)); + let h_new = r.register(inst("prod", "w1", "i-new", 2)); + let got = r.select("prod", Some("w1"), None).unwrap(); + assert_eq!(got.handle, h_new); + } + + #[test] + fn saturation_is_distinct_from_no_target() { + let r = Registry::new(); + let mut i = inst("prod", "w1", "i-1", 1); + i.active_sessions = 1; + r.register(i); + assert!(matches!( + r.select("prod", Some("w1"), None), + Err(SelectError::Saturated) + )); + assert!(matches!( + r.select("prod", Some("nope"), None), + Err(SelectError::NoTarget) + )); + } + + #[test] + fn label_selection_least_loaded_first() { + let r = Registry::new(); + // Older but less loaded instance must win under label selection + // (inverse recency/load — review F6). + let mut a = inst("prod", "wa", "i-a", 4); + a.labels.insert("backend".into(), "kiro".into()); + a.active_sessions = 0; + let h_a = r.register(a); + std::thread::sleep(Duration::from_millis(5)); + let mut b = inst("prod", "wb", "i-b", 4); + b.labels.insert("backend".into(), "kiro".into()); + b.active_sessions = 3; + r.register(b); + + let mut want = BTreeMap::new(); + want.insert("backend".to_string(), "kiro".to_string()); + let got = r.select("prod", None, Some(&want)).unwrap(); + assert_eq!(got.handle, h_a, "least loaded wins despite being older"); + + // partial label mismatch -> NoTarget + want.insert("arch".to_string(), "x86".to_string()); + assert!(matches!( + r.select("prod", None, Some(&want)), + Err(SelectError::NoTarget) + )); + } + + #[test] + fn name_selection_newest_first_even_if_more_loaded() { + let r = Registry::new(); + let mut old = inst("prod", "w1", "i-old", 4); + old.active_sessions = 0; + r.register(old); + std::thread::sleep(Duration::from_millis(5)); + let mut new = inst("prod", "w1", "i-new", 4); + new.active_sessions = 2; + let h_new = r.register(new); + let got = r.select("prod", Some("w1"), None).unwrap(); + assert_eq!(got.handle, h_new, "replica rule: newest registration wins"); + } + + #[test] + fn lease_expiry_and_heartbeat() { + let r = Registry::new(); + let h = r.register(inst("prod", "w1", "i-1", 1)); + assert!(r.expired(Duration::from_secs(60)).is_empty()); + assert!(r.heartbeat(h)); + assert!(!r.heartbeat(h + 999)); + std::thread::sleep(Duration::from_millis(2)); + assert_eq!(r.expired(Duration::ZERO), vec![h]); + } + + #[test] + fn colliding_instance_id_cannot_replace_other_registration() { + // Review F1: a second connection registering the same client-supplied + // instance_id gets its own handle; the first registration survives + // and can only be torn down via its own handle. + let r = Registry::new(); + let h1 = r.register(inst("prod", "w1", "i-same", 1)); + let h2 = r.register(inst("prod", "w2", "i-same", 1)); + assert_ne!(h1, h2); + assert_eq!(r.list("prod").len(), 2); + // Tearing down the second leaves the first intact. + assert!(r.deregister(h2).is_some()); + assert!(r.get(h1).is_some()); + // Deregistering an already-gone handle is a no-op. + assert!(r.deregister(h2).is_none()); + } + + #[test] + fn heartbeat_does_not_mutate_session_count() { + let r = Registry::new(); + let h = r.register(inst("prod", "w1", "i-1", 2)); + r.adjust_sessions(h, 1); + assert!(r.heartbeat(h)); + assert_eq!( + r.get(h).unwrap().active_sessions, + 1, + "CP-owned count is authoritative; heartbeat never changes it" + ); + } + + #[test] + fn adjust_sessions_saturating() { + let r = Registry::new(); + let h = r.register(inst("prod", "w1", "i-1", 2)); + r.adjust_sessions(h, 1); + assert_eq!(r.get(h).unwrap().active_sessions, 1); + r.adjust_sessions(h, -5); + assert_eq!(r.get(h).unwrap().active_sessions, 0); + } + + #[tokio::test] + async fn signal_shutdown_reaches_the_owning_connection() { + // Review round-3 F1: the CP must be able to terminate a connection + // whose registration it drops on its own initiative. + let r = Registry::new(); + let sig = shutdown_signal(); + let mut rx = sig.subscribe(); + let h = r.register_conn(inst("prod", "w1", "i-1", 1), sig); + assert!(!*rx.borrow()); + + assert!(r.signal_shutdown(h)); + rx.changed().await.unwrap(); + assert!(*rx.borrow(), "connection task must observe the signal"); + + // After deregistration there is nothing left to signal. + r.deregister(h); + assert!(!r.signal_shutdown(h)); + } + + #[tokio::test] + async fn signal_shutdown_survives_registry_drop_of_the_entry() { + // The connection task keeps its own handle on the signal, so a + // signal delivered before deregistration is never lost. + let r = Registry::new(); + let sig = shutdown_signal(); + let mut rx = sig.subscribe(); + let h = r.register_conn(inst("prod", "w1", "i-1", 1), sig); + r.signal_shutdown(h); + r.deregister(h); + rx.changed().await.unwrap(); + assert!(*rx.borrow()); + } +} diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs new file mode 100644 index 000000000..57d77a05e --- /dev/null +++ b/crates/openab-cp/src/router.rs @@ -0,0 +1,1518 @@ +//! Delegation router: in-flight table, target selection, result routing, and +//! the failure semantics the ADR review required to be explicit: +//! +//! - **Deadline sweep** — an in-flight delegation whose deadline passes is +//! terminated: the initiator receives a synthesized `timeout` result and +//! the serving runtime receives a best-effort `cp/cancel` (stop burning +//! tokens). +//! - **Target disconnect / lease expiry** — in-flight delegations on that +//! instance fail immediately with `target_disconnected`. +//! - **Initiator disconnect** — its in-flight delegations are cancelled +//! downstream (best effort); nobody is left to receive the result. +//! - **CP restart** — the table is in-memory; all in-flight delegations +//! effectively end as initiator-side timeouts. Late `cp/delegate_result` +//! frames for unknown ids are acknowledged and dropped (logged), so +//! reconnecting runtimes do not error-loop. +//! - **Saturation** — routing never queues; `SATURATED` is returned +//! immediately (fast-fail, no hidden buffer). + +use std::collections::BTreeMap; + +use chrono::{DateTime, Utc}; +use parking_lot::Mutex; +use tracing::{info, warn}; + +use crate::config::CpConfig; +use crate::policy::{self, PolicyInput}; +use crate::proto::{ + codes, methods, AgentType, CancelParams, DelegateAck, DelegateForward, DelegateParams, + DelegateResultParams, DelegationStatus, ErrorObject, JsonRpcRequest, +}; +use crate::registry::{Instance, Registry, SelectError}; + +/// One in-flight delegation. Ownership is tracked by CP-generated +/// registration handles, never client-supplied ids (review F1). +#[derive(Clone)] +pub struct InFlight { + /// Namespace the delegation lives in — part of its identity (review + /// round-3 F3): `delegation_id` is client-supplied and only unique within + /// the namespace that produced it. + pub namespace: String, + pub delegation_id: String, + /// Authenticated initiator (`namespace/name`) and its registration handle. + pub from_logical: String, + 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, +} + +/// In-flight table key: `(namespace, delegation_id)` (review round-3 F3). +/// +/// Keying on the client-supplied `delegation_id` alone made one namespace's +/// ids observable from another: a colliding id was denied with +/// `DUPLICATE_DELEGATION`, and `cp/cancel` distinguished "no such id" from +/// "someone else's live id" — a cross-tenant existence oracle. The composite +/// key confines both to the namespace that owns the id. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct DelegationKey { + namespace: String, + delegation_id: String, +} + +impl DelegationKey { + fn new(namespace: &str, delegation_id: &str) -> Self { + Self { + namespace: namespace.to_string(), + delegation_id: delegation_id.to_string(), + } + } +} + +pub struct Router { + inflight: Mutex>, + /// Serializes the delegate admission sequence (duplicate check → target + /// selection → capacity reservation → in-flight insert) so concurrent + /// requests cannot double-admit one id or oversubscribe capacity + /// (review F2). Delegation rates are LLM-scale; a coarse admission lock + /// is simple and more than sufficient. + admission: Mutex<()>, +} + +pub enum DelegateOutcome { + /// Forwarded to the target; ack for the initiator. + Accepted(DelegateAck), + /// Rejected; error for the initiator. + Rejected(ErrorObject), +} + +impl Router { + pub fn new() -> Self { + Self { + inflight: Mutex::new(BTreeMap::new()), + admission: Mutex::new(()), + } + } + + /// Handle `cp/delegate` from an authenticated, registered initiator. + #[allow(clippy::too_many_arguments)] + pub fn delegate( + &self, + cfg: &CpConfig, + registry: &Registry, + from_namespace: &str, + from_name: &str, + from_type: &AgentType, + from_handle: u64, + params: DelegateParams, + next_rpc_id: u64, + ) -> DelegateOutcome { + let now = Utc::now(); + + // Admission is one atomic sequence (review F2): duplicate check, + // parent lookup, target selection, capacity reservation, and + // in-flight insertion all happen under this guard. + let _admission = self.admission.lock(); + + // Delegation identity is namespace-scoped (review round-3 F3): the + // same id in another namespace is a different delegation, so it + // neither collides here nor leaks its existence. + let key = DelegationKey::new(from_namespace, ¶ms.delegation_id); + if self.inflight.lock().contains_key(&key) { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::DUPLICATE_DELEGATION, + format!("delegation {} is already in flight", params.delegation_id), + )); + } + + // Selector sanity: exactly one of name/labels. + let (sel_name, sel_labels) = (params.target.name.as_deref(), params.target.labels.as_ref()); + if sel_name.is_some() == sel_labels.is_some() { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + "target must set exactly one of `name` or `labels`", + )); + } + + // Parent linkage: chain and deadline derive from the CP's own table, + // never from the client. The caller must BE the instance serving the + // parent delegation — otherwise any runtime knowing a live id could + // borrow its trusted chain and deadline budget (review F3). The + // lookup is namespace-scoped (review round-3 F3). Unknown and + // unauthorized parent ids return the same error (no enumeration). + let (parent_chain, parent_deadline) = match ¶ms.parent_delegation_id { + Some(pid) => { + let parent_key = DelegationKey::new(from_namespace, pid); + match self.inflight.lock().get(&parent_key) { + Some(p) if p.to_handle == from_handle => (p.chain.clone(), Some(p.deadline)), + _ => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::INVALID_PARAMS, + format!("parent delegation {pid} is not in flight for this instance"), + )) + } + } + } + None => (Vec::new(), None), + }; + + // Resolve target within the initiator's namespace (v1 boundary). + let target = match registry.select(from_namespace, sel_name, sel_labels) { + Ok(i) => i, + Err(SelectError::NoTarget) => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::NO_TARGET, + "no registered healthy runtime matches the target selector", + )) + } + Err(SelectError::Saturated) => { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::SATURATED, + "all matching runtimes are at capacity (CP does not queue; retry later)", + )) + } + }; + + // CP-authoritative policy. + let input = PolicyInput { + from_namespace, + from_name, + from_type, + target_namespace: &target.namespace, + target_name: &target.name, + parent_chain: &parent_chain, + deadline: params.deadline, + parent_deadline, + now, + max_deadline_secs: cfg.max_deadline_secs, + }; + if let Err(denial) = policy::check(&input, &cfg.policy_for(from_namespace)) { + return DelegateOutcome::Rejected(ErrorObject::new( + codes::POLICY_DENIED, + denial.to_string(), + )); + } + + // Build the forward frame with the CP-stamped chain. + let from_logical = format!("{from_namespace}/{from_name}"); + let mut chain = parent_chain; + chain.push(from_logical.clone()); + let forward = DelegateForward { + delegation_id: params.delegation_id.clone(), + prompt: params.prompt, + deadline: params.deadline, + from: from_logical.clone(), + chain: chain.clone(), + }; + let frame = JsonRpcRequest::new( + next_rpc_id, + methods::DELEGATE, + Some(serde_json::to_value(&forward).expect("serializable")), + ); + let text = serde_json::to_string(&frame).expect("serializable"); + + // Reserve capacity and record the in-flight entry BEFORE sending, so + // an immediately-arriving result finds it (review F2). Roll both + // back if the send fails. + registry.adjust_sessions(target.handle, 1); + let entry = InFlight { + 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, + }; + self.inflight.lock().insert(key.clone(), entry.clone()); + + if target.tx.try_send(text).is_err() { + // Disconnected or backpressured beyond its queue: roll back. + self.inflight.lock().remove(&key); + registry.adjust_sessions(target.handle, -1); + return DelegateOutcome::Rejected(ErrorObject::new( + codes::TARGET_DISCONNECTED, + "target disconnected or unresponsive during routing", + )); + } + + info!( + delegation = %entry.delegation_id, + from = %entry.from_logical, + to = %entry.to_logical, + chain = ?entry.chain, + deadline = %entry.deadline, + "delegation routed" + ); + + DelegateOutcome::Accepted(DelegateAck { + delegation_id: params.delegation_id, + assigned_to: target.logical_id(), + }) + } + + /// Handle `cp/delegate_result` from the serving runtime. Returns the + /// initiator-bound frame if the delegation is known; unknown ids (e.g. + /// results arriving after a CP restart) are dropped with a log. + /// + /// Ownership is validated under the SAME lock acquisition that removes + /// the entry (review round-3 F2): the previous remove-check-reinsert + /// dance opened a window in which a genuine result saw an empty table and + /// was dropped, and left the entry momentarily invisible to the deadline + /// sweep. + pub fn complete( + &self, + registry: &Registry, + serving_handle: u64, + mut params: DelegateResultParams, + max_result_bytes: usize, + next_rpc_id: u64, + ) -> Option<(Instance, String)> { + // The namespace comes from the authenticated sender's registration, + // never from the frame (review round-3 F3). + let namespace = match registry.get(serving_handle) { + Some(i) => i.namespace, + None => { + warn!( + handle = serving_handle, + delegation = %params.delegation_id, + "result from an unregistered connection — dropped" + ); + return None; + } + }; + let key = DelegationKey::new(&namespace, ¶ms.delegation_id); + let entry = { + let mut g = self.inflight.lock(); + match g.get(&key) { + Some(e) if e.to_handle == serving_handle => {} + Some(e) => { + // Only the instance the delegation was routed to may + // complete it. The entry stays exactly where it is. + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + expected = e.to_handle, + got = serving_handle, + "result from unexpected instance — dropped, delegation untouched" + ); + return None; + } + None => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "result for unknown delegation (late arrival or CP restart) — dropped" + ); + return None; + } + } + g.remove(&key).expect("present under the same lock") + }; + + registry.adjust_sessions(entry.to_handle, -1); + + // Truncate oversized results (keep the head; delegation already + // ran). The marker counts against the cap: the final value never + // exceeds max_result_bytes (review round-2 F5). + 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); + } + } + + info!( + delegation = %params.delegation_id, + status = ?params.status, + from = %entry.to_logical, + to = %entry.from_logical, + "delegation completed" + ); + + let initiator = registry.get(entry.from_handle)?; + let frame = JsonRpcRequest::new( + next_rpc_id, + methods::DELEGATE_RESULT, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + Some(( + initiator, + serde_json::to_string(&frame).expect("serializable"), + )) + } + + /// Handle `cp/cancel` from the initiator. Returns the frame to forward + /// to the serving runtime, if the delegation is in flight and owned by + /// the caller. + /// + /// Ownership is validated under the same lock acquisition that removes + /// the entry (review round-3 F2 — no remove/reinsert window), and every + /// refusal returns ONE byte-identical error (review round-3 F3): an + /// unknown id and another instance's live id are indistinguishable to the + /// caller, so `cp/cancel` cannot be used to probe for delegation ids. + /// The distinction is kept in the CP's own logs only. + pub fn cancel( + &self, + registry: &Registry, + 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 namespace = match registry.get(from_handle) { + Some(i) => i.namespace, + None => { + warn!( + handle = from_handle, + "cancel from an unregistered connection" + ); + return Err(refused()); + } + }; + let key = DelegationKey::new(&namespace, ¶ms.delegation_id); + let entry = { + let mut g = self.inflight.lock(); + match g.get(&key) { + Some(e) if e.from_handle == from_handle => {} + Some(_) => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + handle = from_handle, + "cancel refused: only the initiating instance may cancel" + ); + return Err(refused()); + } + None => { + warn!( + delegation = %params.delegation_id, + namespace = %namespace, + "cancel refused: delegation not in flight" + ); + return Err(refused()); + } + } + g.remove(&key).expect("present under the same lock") + }; + registry.adjust_sessions(entry.to_handle, -1); + info!(delegation = %params.delegation_id, "delegation cancelled by initiator"); + let target = registry.get(entry.to_handle); + Ok(target.map(|t| { + let frame = JsonRpcRequest::new( + next_rpc_id, + methods::CANCEL, + Some(serde_json::to_value(params).expect("serializable")), + ); + (t, serde_json::to_string(&frame).expect("serializable")) + })) + } + + /// Fail every in-flight delegation touching a deregistered instance. + /// Returns synthesized result/cancel frames to deliver: + /// - delegations SERVED by the instance → `target_disconnected` result to + /// the initiator + /// - delegations INITIATED by the instance → best-effort `cp/cancel` to + /// the serving runtime + pub fn fail_instance( + &self, + registry: &Registry, + handle: u64, + rpc_id: &mut impl FnMut() -> u64, + ) -> Vec<(Instance, String)> { + let mut affected = Vec::new(); + let entries: Vec = { + let mut g = self.inflight.lock(); + let 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(), + status: DelegationStatus::TargetDisconnected, + result: None, + error: Some(format!("{} disconnected", e.to_logical)), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::DELEGATE_RESULT, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + affected.push((init, serde_json::to_string(&frame).expect("serializable"))); + } + } else { + // Initiator died → cancel downstream, free worker capacity. + registry.adjust_sessions(e.to_handle, -1); + if let Some(target) = registry.get(e.to_handle) { + let params = CancelParams { + delegation_id: e.delegation_id.clone(), + reason: format!("initiator {} disconnected", e.from_logical), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::CANCEL, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + affected.push((target, serde_json::to_string(&frame).expect("serializable"))); + } + } + warn!(delegation = %e.delegation_id, handle, "in-flight delegation failed by disconnect"); + } + affected + } + + /// Deadline sweep: expire overdue delegations. Returns frames to deliver + /// (timeout result to the initiator, best-effort cancel to the server). + pub fn sweep_deadlines( + &self, + registry: &Registry, + now: DateTime, + rpc_id: &mut impl FnMut() -> u64, + ) -> Vec<(Instance, String)> { + let overdue: Vec = { + let mut g = self.inflight.lock(); + let 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(), + status: DelegationStatus::Timeout, + result: None, + error: Some("deadline exceeded".to_string()), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::DELEGATE_RESULT, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + frames.push((init, serde_json::to_string(&frame).expect("serializable"))); + } + if let Some(target) = registry.get(e.to_handle) { + let params = CancelParams { + delegation_id: e.delegation_id.clone(), + reason: "deadline exceeded".to_string(), + }; + let frame = JsonRpcRequest::new( + rpc_id(), + methods::CANCEL, + Some(serde_json::to_value(¶ms).expect("serializable")), + ); + frames.push((target, serde_json::to_string(&frame).expect("serializable"))); + } + } + frames + } + + /// Chain of an in-flight delegation (for tests/inspection). Delegation + /// ids are namespace-scoped (review round-3 F3), so the namespace is part + /// of the lookup. + pub fn chain_of(&self, namespace: &str, delegation_id: &str) -> Option> { + self.inflight + .lock() + .get(&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_QUEUE; + use chrono::Duration; + use std::time::Instant; + use tokio::sync::mpsc; + + fn cfg() -> CpConfig { + toml::from_str( + r#" +[[agents]] +key = "kp" +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "kw" +namespace = "prod" +name = "worker-1" +type = "worker" +"#, + ) + .unwrap() + } + + fn instance( + ns: &str, + name: &str, + ty: AgentType, + max: u32, + ) -> (Instance, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(OUTBOUND_QUEUE); + ( + Instance { + handle: 0, + namespace: ns.into(), + name: name.into(), + agent_type: ty, + instance_id: format!("i-{name}"), + labels: Default::default(), + max_delegated_sessions: max, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + }, + rx, + ) + } + + fn delegate_params(id: &str, target: &str, secs: i64) -> DelegateParams { + DelegateParams { + delegation_id: id.into(), + target: TargetSelector { + name: Some(target.into()), + labels: None, + }, + prompt: "do it".into(), + deadline: Utc::now() + Duration::seconds(secs), + parent_delegation_id: None, + } + } + + struct World { + cfg: CpConfig, + registry: Registry, + router: Router, + h_primary: u64, + h_worker: u64, + worker_rx: mpsc::Receiver, + primary_rx: mpsc::Receiver, + } + + fn world() -> World { + let registry = Registry::new(); + let (p, primary_rx) = instance("prod", "koudu", AgentType::Primary, 4); + let (w, worker_rx) = instance("prod", "worker-1", AgentType::Worker, 1); + let h_primary = registry.register(p); + let h_worker = registry.register(w); + World { + cfg: cfg(), + registry, + router: Router::new(), + h_primary, + h_worker, + worker_rx, + primary_rx, + } + } + + fn do_delegate(w: &World, params: DelegateParams) -> DelegateOutcome { + w.router.delegate( + &w.cfg, + &w.registry, + "prod", + "koudu", + &AgentType::Primary, + w.h_primary, + params, + 1, + ) + } + + #[test] + fn happy_path_roundtrip() { + let mut w = world(); + let out = do_delegate(&w, delegate_params("d-1", "worker-1", 60)); + let ack = match out { + DelegateOutcome::Accepted(a) => a, + DelegateOutcome::Rejected(e) => panic!("rejected: {}", e.message), + }; + assert_eq!(ack.assigned_to, "prod/worker-1"); + + let frame = w.worker_rx.try_recv().unwrap(); + let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); + assert_eq!(v["method"], "cp/delegate"); + assert_eq!(v["params"]["from"], "prod/koudu"); + assert_eq!(v["params"]["chain"], serde_json::json!(["prod/koudu"])); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 1); + + let result = DelegateResultParams { + delegation_id: "d-1".into(), + status: DelegationStatus::Completed, + result: Some("done".into()), + error: None, + }; + let (init, frame) = w + .router + .complete(&w.registry, w.h_worker, result, 1024, 2) + .unwrap(); + assert_eq!(init.handle, w.h_primary); + assert!(frame.contains("\"completed\"")); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + assert_eq!(w.router.inflight_count(), 0); + } + + #[test] + fn inflight_exists_before_target_receives_frame() { + // Review F2: an immediately-arriving result must find the entry. + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + // Complete BEFORE draining the worker's queue — entry must exist. + let result = DelegateResultParams { + delegation_id: "d-1".into(), + status: DelegationStatus::Completed, + result: Some("instant".into()), + error: None, + }; + assert!(w + .router + .complete(&w.registry, w.h_worker, result, 1024, 2) + .is_some()); + w.worker_rx.try_recv().unwrap(); + } + + #[test] + fn send_failure_rolls_back_reservation() { + // Close the worker's rx so try_send fails, then verify rollback. + let mut w = world(); + w.worker_rx.close(); + match do_delegate(&w, delegate_params("d-1", "worker-1", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::TARGET_DISCONNECTED), + _ => panic!("expected TARGET_DISCONNECTED"), + } + assert_eq!(w.router.inflight_count(), 0); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "capacity reservation must be rolled back" + ); + } + + #[test] + fn duplicate_delegation_id_rejected() { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + match do_delegate(&w, delegate_params("d-1", "worker-1", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::DUPLICATE_DELEGATION), + _ => panic!("expected rejection"), + } + } + + #[test] + fn saturation_fast_fails() { + let w = world(); // worker max = 1 + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + match do_delegate(&w, delegate_params("d-2", "worker-1", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::SATURATED), + _ => panic!("expected SATURATED"), + } + } + + #[test] + fn selector_must_be_exactly_one() { + let w = world(); + let mut p = delegate_params("d-1", "worker-1", 60); + p.target.labels = Some(Default::default()); + match do_delegate(&w, p) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::INVALID_PARAMS), + _ => panic!(), + } + let mut p2 = delegate_params("d-2", "worker-1", 60); + p2.target.name = None; + match do_delegate(&w, p2) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::INVALID_PARAMS), + _ => panic!(), + } + } + + #[test] + fn policy_denial_maps_to_error_code() { + let w = world(); + let out = w.router.delegate( + &w.cfg, + &w.registry, + "prod", + "worker-1", + &AgentType::Worker, + w.h_worker, + delegate_params("d-1", "koudu", 60), + 1, + ); + match out { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::POLICY_DENIED), + _ => panic!("expected POLICY_DENIED"), + } + } + + #[test] + fn unknown_target_is_no_target() { + let w = world(); + match do_delegate(&w, delegate_params("d-1", "ghost", 60)) { + DelegateOutcome::Rejected(e) => assert_eq!(e.code, codes::NO_TARGET), + _ => panic!(), + } + } + + #[test] + fn result_from_wrong_handle_dropped_and_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(), + status: DelegationStatus::Completed, + result: Some("spoofed".into()), + error: None, + }; + // h_primary is a valid handle but NOT the serving instance. + assert!(w + .router + .complete(&w.registry, w.h_primary, result, 1024, 2) + .is_none()); + assert_eq!(w.router.inflight_count(), 1); + } + + #[test] + fn late_result_after_restart_dropped() { + let w = world(); + let result = DelegateResultParams { + delegation_id: "d-unknown".into(), + status: DelegationStatus::Completed, + result: None, + error: None, + }; + assert!(w + .router + .complete(&w.registry, w.h_worker, result, 1024, 2) + .is_none()); + } + + #[test] + fn oversized_result_truncated() { + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + let result = DelegateResultParams { + delegation_id: "d-1".into(), + status: DelegationStatus::Completed, + result: Some("x".repeat(200)), + error: None, + }; + let cap = 96usize; + let (_, frame) = w + .router + .complete(&w.registry, w.h_worker, result, cap, 2) + .unwrap(); + let v: serde_json::Value = serde_json::from_str(&frame).unwrap(); + let out = v["params"]["result"].as_str().unwrap(); + assert!(out.contains("truncated by control plane")); + assert!( + out.len() <= cap, + "marker must count against the cap: {} > {}", + out.len(), + cap + ); + + // Degenerate tiny cap still never exceeds the cap. + assert!(matches!( + do_delegate(&w, delegate_params("d-2", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + let result2 = DelegateResultParams { + delegation_id: "d-2".into(), + status: DelegationStatus::Completed, + result: Some("y".repeat(100)), + error: None, + }; + let (_, frame2) = w + .router + .complete(&w.registry, w.h_worker, result2, 8, 3) + .unwrap(); + 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(), + reason: "changed my mind".into(), + }; + let err = w + .router + .cancel(&w.registry, w.h_worker, ¶ms, 5) + .unwrap_err(); + assert_eq!(err.code, codes::POLICY_DENIED); + assert_eq!(w.router.inflight_count(), 1); + let fwd = w + .router + .cancel(&w.registry, w.h_primary, ¶ms, 6) + .unwrap(); + let (inst, frame) = fwd.unwrap(); + assert_eq!(inst.handle, w.h_worker); + assert!(frame.contains("cp/cancel")); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + } + + #[test] + fn chain_extends_through_parent_and_foreign_parent_rejected() { + let w = world(); + let cfg: CpConfig = toml::from_str( + r#" +[[agents]] +key = "kp" +namespace = "prod" +name = "koudu" +type = "primary" + +[namespaces.prod] +max_depth = 5 +allow_worker_initiation = true +"#, + ) + .unwrap(); + let (w2, _rx2) = instance("prod", "worker-2", AgentType::Worker, 1); + let h_w2 = w.registry.register(w2); + + assert!(matches!( + w.router.delegate( + &cfg, + &w.registry, + "prod", + "koudu", + &AgentType::Primary, + w.h_primary, + delegate_params("d-root", "worker-1", 120), + 1, + ), + DelegateOutcome::Accepted(_) + )); + assert_eq!( + w.router.chain_of("prod", "d-root").unwrap(), + vec!["prod/koudu".to_string()] + ); + + // Review F3: worker-2 (NOT serving d-root) tries to borrow d-root + // as parent — rejected. + let mut foreign = delegate_params("d-foreign", "worker-2", 60); + foreign.parent_delegation_id = Some("d-root".into()); + match w.router.delegate( + &cfg, + &w.registry, + "prod", + "worker-2", + &AgentType::Worker, + h_w2, + foreign, + 2, + ) { + DelegateOutcome::Rejected(e) => { + assert_eq!(e.code, codes::INVALID_PARAMS); + assert!(e.message.contains("not in flight for this instance")); + } + _ => panic!("foreign parent must be rejected"), + } + + // worker-1 (serving d-root) delegates a legitimate child to worker-2. + let mut child = delegate_params("d-child", "worker-2", 60); + child.parent_delegation_id = Some("d-root".into()); + assert!(matches!( + w.router.delegate( + &cfg, + &w.registry, + "prod", + "worker-1", + &AgentType::Worker, + w.h_worker, + child, + 3, + ), + DelegateOutcome::Accepted(_) + )); + assert_eq!( + w.router.chain_of("prod", "d-child").unwrap(), + vec!["prod/koudu".to_string(), "prod/worker-1".to_string()] + ); + + // Cycle: worker-2 delegating back to koudu is rejected. + let mut cyc = delegate_params("d-cyc", "koudu", 30); + cyc.parent_delegation_id = Some("d-child".into()); + 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"), + } + } + + fn result_of(id: &str, body: &str) -> DelegateResultParams { + DelegateResultParams { + delegation_id: id.into(), + status: DelegationStatus::Completed, + result: Some(body.into()), + error: None, + } + } + + #[test] + fn wrong_handle_result_never_hides_the_genuine_one() { + // Review round-3 F2: ownership is validated under the same lock + // acquisition that removes the entry. The old remove → validate → + // reinsert sequence made the entry briefly invisible, so a genuine + // result arriving in that window was dropped as "unknown id". + for spoof_first in [true, false] { + let mut w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + w.worker_rx.try_recv().unwrap(); + + if spoof_first { + // h_primary is registered but is NOT the serving instance. + assert!(w + .router + .complete( + &w.registry, + w.h_primary, + result_of("d-1", "spoofed"), + 1024, + 2 + ) + .is_none()); + assert_eq!( + w.router.inflight_count(), + 1, + "a non-owner frame must not remove the entry" + ); + } + + let (init, frame) = w + .router + .complete( + &w.registry, + w.h_worker, + result_of("d-1", "genuine"), + 1024, + 3, + ) + .expect("genuine result must be delivered, never dropped"); + assert_eq!(init.handle, w.h_primary); + assert!(frame.contains("genuine")); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!(w.registry.get(w.h_worker).unwrap().active_sessions, 0); + + if !spoof_first { + // A late non-owner frame after completion is a plain no-op. + assert!(w + .router + .complete( + &w.registry, + w.h_primary, + result_of("d-1", "spoofed"), + 1024, + 4 + ) + .is_none()); + assert_eq!(w.router.inflight_count(), 0); + } + } + } + + #[test] + fn genuine_result_survives_concurrent_non_owner_frames() { + // Review round-3 F2, the racing case the sequential test above cannot + // observe: with remove → validate → reinsert, a genuine result that + // lands inside the window sees an empty table and is dropped, and the + // delegation then stalls to its deadline. Under a single lock + // acquisition the outcome is order-independent by construction. + for _ in 0..200 { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let gate = std::sync::Barrier::new(2); + let (spoofed, genuine) = std::thread::scope(|s| { + let spoof = s.spawn(|| { + gate.wait(); + // Registered, but not the serving instance. + w.router + .complete( + &w.registry, + w.h_primary, + result_of("d-1", "spoofed"), + 1024, + 2, + ) + .is_some() + }); + gate.wait(); + let genuine = w + .router + .complete( + &w.registry, + w.h_worker, + result_of("d-1", "genuine"), + 1024, + 3, + ) + .is_some(); + (spoof.join().unwrap(), genuine) + }); + assert!(!spoofed, "a non-owner must never complete a delegation"); + assert!(genuine, "the genuine result must never be dropped"); + assert_eq!(w.router.inflight_count(), 0); + } + } + + #[test] + fn genuine_cancel_survives_concurrent_non_owner_frames() { + // Review round-3 F2 (cancel side), racing case: with the old + // remove → validate → reinsert pattern, a genuine initiator cancel + // landing inside a non-owner cancel's window would see an empty table + // and be refused, leaving the delegation to stall to its deadline. + // Under a single lock acquisition the outcome is order-independent. + for _ in 0..200 { + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let params = CancelParams { + delegation_id: "d-1".into(), + reason: "race".into(), + }; + let gate = std::sync::Barrier::new(2); + let (spoofed, genuine) = std::thread::scope(|s| { + let spoof = s.spawn(|| { + gate.wait(); + // Registered, but not the initiator. + w.router.cancel(&w.registry, w.h_worker, ¶ms, 1).is_ok() + }); + gate.wait(); + let genuine = w + .router + .cancel(&w.registry, w.h_primary, ¶ms, 2) + .is_ok(); + (spoof.join().unwrap(), genuine) + }); + assert!(!spoofed, "a non-initiator must never cancel a delegation"); + assert!(genuine, "the genuine cancel must never be refused"); + assert_eq!(w.router.inflight_count(), 0); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "capacity must be released exactly once" + ); + } + } + + #[test] + fn refused_cancel_leaves_the_delegation_cancellable() { + // Review round-3 F2 (cancel side): a wrong-handle cancel must not + // remove-and-reinsert the entry, and must not disturb accounting. + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let params = CancelParams { + delegation_id: "d-1".into(), + reason: "not mine".into(), + }; + assert!(w + .router + .cancel(&w.registry, w.h_worker, ¶ms, 1) + .is_err()); + assert_eq!(w.router.inflight_count(), 1); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "a refused cancel must not release capacity" + ); + // The genuine initiator can still cancel. + let fwd = w + .router + .cancel(&w.registry, w.h_primary, ¶ms, 2) + .unwrap() + .unwrap(); + assert_eq!(fwd.0.handle, w.h_worker); + assert_eq!(w.router.inflight_count(), 0); + } + + #[test] + fn cancel_refusals_are_byte_identical() { + // Review round-3 F3: `cp/cancel` must not be an existence oracle — + // an unknown id and another instance's live id return the same error + // object, byte for byte. + let w = world(); + assert!(matches!( + do_delegate(&w, delegate_params("d-1", "worker-1", 60)), + DelegateOutcome::Accepted(_) + )); + let unknown = CancelParams { + delegation_id: "d-does-not-exist".into(), + reason: "probe".into(), + }; + let foreign = CancelParams { + delegation_id: "d-1".into(), + reason: "probe".into(), + }; + // Both probes come from the worker: it initiated neither. + let e_unknown = w + .router + .cancel(&w.registry, w.h_worker, &unknown, 1) + .unwrap_err(); + let e_foreign = w + .router + .cancel(&w.registry, w.h_worker, &foreign, 2) + .unwrap_err(); + assert_eq!( + serde_json::to_string(&e_unknown).unwrap(), + serde_json::to_string(&e_foreign).unwrap(), + "unknown and foreign delegation ids must be indistinguishable" + ); + assert_eq!(e_unknown.code, codes::POLICY_DENIED); + assert_eq!(w.router.inflight_count(), 1); + } + + #[test] + fn same_delegation_id_in_two_namespaces_is_independent() { + // Review round-3 F3: the in-flight table is keyed by + // (namespace, delegation_id). A client-supplied id in one namespace + // must neither collide with nor be observable from another. + let registry = Registry::new(); + let router = Router::new(); + let cfg = cfg(); + let (p_prod, _prod_init_rx) = instance("prod", "koudu", AgentType::Primary, 4); + let (w_prod, mut prod_rx) = instance("prod", "worker-1", AgentType::Worker, 2); + let (p_dev, _dev_init_rx) = instance("dev", "koudu", AgentType::Primary, 4); + let (w_dev, mut dev_rx) = instance("dev", "worker-1", AgentType::Worker, 2); + let hp_prod = registry.register(p_prod); + let hw_prod = registry.register(w_prod); + let hp_dev = registry.register(p_dev); + let hw_dev = registry.register(w_dev); + + for (ns, hp) in [("prod", hp_prod), ("dev", hp_dev)] { + match router.delegate( + &cfg, + ®istry, + ns, + "koudu", + &AgentType::Primary, + hp, + delegate_params("d-1", "worker-1", 60), + 1, + ) { + DelegateOutcome::Accepted(ack) => { + assert_eq!(ack.assigned_to, format!("{ns}/worker-1")) + } + DelegateOutcome::Rejected(e) => { + panic!("{ns} rejected ({}): {}", e.code, e.message) + } + } + } + assert_eq!( + router.inflight_count(), + 2, + "one `d-1` per namespace, both in flight" + ); + prod_rx.try_recv().unwrap(); + dev_rx.try_recv().unwrap(); + + // A dev instance cannot cancel prod's `d-1` — and cannot learn that + // it exists: same error as for an id that exists nowhere. + let probe = CancelParams { + delegation_id: "d-1".into(), + reason: "probe".into(), + }; + let nowhere = CancelParams { + delegation_id: "d-nowhere".into(), + reason: "probe".into(), + }; + let e_cross = router + .cancel(®istry, hw_dev, &probe, 10) + .unwrap_err() + .message; + let e_nowhere = router + .cancel(®istry, hw_dev, &nowhere, 11) + .unwrap_err() + .message; + assert_eq!(e_cross, e_nowhere); + assert_eq!(router.inflight_count(), 2); + + // Results route to the initiator of the SAME namespace only. + let (init, frame) = router + .complete(®istry, hw_dev, result_of("d-1", "dev-done"), 1024, 12) + .unwrap(); + assert_eq!(init.handle, hp_dev); + assert!(frame.contains("dev-done")); + assert!( + router.chain_of("prod", "d-1").is_some(), + "prod's delegation must be untouched" + ); + + let (init, frame) = router + .complete(®istry, hw_prod, result_of("d-1", "prod-done"), 1024, 13) + .unwrap(); + assert_eq!(init.handle, hp_prod); + assert!(frame.contains("prod-done")); + assert_eq!(router.inflight_count(), 0); + } + + #[test] + fn parent_lookup_is_namespace_scoped() { + // Review round-3 F3: parent-chain resolution must not reach into + // another namespace's in-flight table. + let cfg: CpConfig = toml::from_str( + r#" +[namespaces.prod] +max_depth = 5 +allow_worker_initiation = true + +[namespaces.dev] +max_depth = 5 +allow_worker_initiation = true +"#, + ) + .unwrap(); + let registry = Registry::new(); + let router = Router::new(); + let (p_prod, _rx1) = instance("prod", "koudu", AgentType::Primary, 4); + let (w_prod, mut rx2) = instance("prod", "worker-1", AgentType::Worker, 2); + let (t_prod, _rx3) = instance("prod", "worker-2", AgentType::Worker, 2); + let (w_dev, _rx4) = instance("dev", "worker-1", AgentType::Worker, 2); + let (t_dev, _rx5) = instance("dev", "worker-2", AgentType::Worker, 2); + let hp_prod = registry.register(p_prod); + let hw_prod = registry.register(w_prod); + registry.register(t_prod); + let hw_dev = registry.register(w_dev); + registry.register(t_dev); + + assert!(matches!( + router.delegate( + &cfg, + ®istry, + "prod", + "koudu", + &AgentType::Primary, + hp_prod, + delegate_params("d-root", "worker-1", 120), + 1, + ), + DelegateOutcome::Accepted(_) + )); + rx2.try_recv().unwrap(); + + // dev/worker-1 claims prod's `d-root` as its parent: invisible. + let mut child = delegate_params("d-child", "worker-2", 60); + child.parent_delegation_id = Some("d-root".into()); + match router.delegate( + &cfg, + ®istry, + "dev", + "worker-1", + &AgentType::Worker, + hw_dev, + child, + 2, + ) { + DelegateOutcome::Rejected(e) => { + assert_eq!(e.code, codes::INVALID_PARAMS); + assert!(e.message.contains("not in flight for this instance")); + } + _ => panic!("cross-namespace parent must be rejected"), + } + // ...and the legitimate in-namespace child still works. + let mut ok_child = delegate_params("d-child", "worker-2", 60); + ok_child.parent_delegation_id = Some("d-root".into()); + assert!(matches!( + router.delegate( + &cfg, + ®istry, + "prod", + "worker-1", + &AgentType::Worker, + hw_prod, + ok_child, + 3, + ), + DelegateOutcome::Accepted(_) + )); + assert_eq!( + router.chain_of("prod", "d-child").unwrap(), + vec!["prod/koudu".to_string(), "prod/worker-1".to_string()] + ); + } +} diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs new file mode 100644 index 000000000..ecf33b8b6 --- /dev/null +++ b/crates/openab-cp/src/server.rs @@ -0,0 +1,828 @@ +//! WebSocket server: authentication at upgrade, mandatory `cp/register` +//! first frame, then frame dispatch to registry/policy/router. +//! +//! Auth: the runtime presents its key as `Authorization: Bearer ` on the +//! upgrade request. Keys never appear in URLs (avoids access-log leakage). +//! +//! Resource bounds (review F5): the WS transport enforces +//! `max_frame_bytes` before parsing; each connection's outbound queue is +//! bounded — a peer that cannot drain it is treated as disconnected. +//! +//! Admission bounds (review round-3 F4): authentication alone is not a +//! bound. Every connection holds a per-identity slot from the upgrade until +//! it ends (`ConnPermit`, released on every exit path), and must complete +//! `cp/register` within `register_timeout_secs` or be closed. + +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use axum::extract::ws::{Message, WebSocket}; +use axum::extract::{State, WebSocketUpgrade}; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::IntoResponse; +use axum::routing::get; +use axum::Router as AxumRouter; +use futures_util::{SinkExt, StreamExt}; +use parking_lot::Mutex; +use tokio::sync::mpsc; +use tracing::{info, warn}; + +use crate::config::{AgentIdentity, CpConfig}; +use crate::proto::{ + codes, methods, CancelParams, DelegateParams, DelegateResultParams, ErrorObject, + JsonRpcErrorResponse, JsonRpcMessage, JsonRpcResponse, RegisterAck, RegisterParams, + PROTOCOL_VERSION, +}; +use crate::registry::{shutdown_signal, Instance, Registry, OUTBOUND_QUEUE}; +use crate::router::{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 (review round-3 + /// F4). + 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 (review round-3 F4). + pub fn try_acquire_conn(self: &Arc, identity: &AgentIdentity) -> Option { + let key = format!("{}/{}", identity.namespace, identity.name); + let mut g = self.conns.lock(); + let n = g.entry(key.clone()).or_insert(0); + if *n >= self.cfg.max_connections_per_identity { + return None; + } + *n += 1; + Some(ConnPermit { + state: Arc::clone(self), + key, + }) + } + + /// Live connection count for an identity (`namespace/name`). + pub fn conn_count(&self, logical_id: &str) -> u32 { + self.conns.lock().get(logical_id).copied().unwrap_or(0) + } +} + +/// RAII connection slot. Dropping it frees the identity's quota; it is never +/// released explicitly, so no early return can leak it (review round-3 F4). +pub struct ConnPermit { + state: Arc, + key: String, +} + +impl Drop for ConnPermit { + fn drop(&mut self) { + let mut g = self.state.conns.lock(); + if let Some(n) = g.get_mut(&self.key) { + *n = n.saturating_sub(1); + if *n == 0 { + g.remove(&self.key); + } + } + } +} + +pub fn app(state: Arc) -> AxumRouter { + AxumRouter::new() + .route("/cp", get(ws_handler)) + .route("/health", get(health)) + .with_state(state) +} + +async fn health() -> &'static str { + "ok" +} + +async fn ws_handler( + State(state): State>, + headers: HeaderMap, + ws: WebSocketUpgrade, +) -> axum::response::Response { + let key = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.strip_prefix("Bearer ")); + let identity = match key.and_then(|k| state.cfg.identity_for_key(k)) { + Some(id) => id.clone(), + None => { + warn!("WS rejected: missing or unknown auth key"); + return StatusCode::UNAUTHORIZED.into_response(); + } + }; + // Per-identity connection quota, taken before the upgrade so an + // over-quota peer is refused at the HTTP layer (review round-3 F4). + let permit = match state.try_acquire_conn(&identity) { + Some(p) => p, + None => { + warn!( + agent = %format!("{}/{}", identity.namespace, identity.name), + max = state.cfg.max_connections_per_identity, + "WS rejected: identity is at its connection quota" + ); + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + }; + let max_frame = state.cfg.max_frame_bytes; + ws.max_message_size(max_frame) + .max_frame_size(max_frame) + .on_upgrade(move |socket| handle_connection(state, socket, identity, 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 (review round-3 F4). + _permit: ConnPermit, +) { + let (mut sink, mut stream) = socket.split(); + + // --- Registration: mandatory first frame, within a deadline --- + // An authenticated peer must not be able to park idle sockets: pings keep + // the transport alive but do not extend this deadline (review round-3 F4). + let register = match tokio::time::timeout( + Duration::from_secs(state.cfg.register_timeout_secs), + async { + loop { + match stream.next().await { + Some(Ok(Message::Text(text))) => return Some(text), + Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue, + _ => return None, + } + } + }, + ) + .await + { + Ok(Some(text)) => text, + Ok(None) => { + warn!(agent = %identity.name, "connection closed before registration"); + return; + } + Err(_) => { + warn!( + agent = %format!("{}/{}", identity.namespace, identity.name), + timeout_secs = state.cfg.register_timeout_secs, + "no cp/register within the registration deadline — closing" + ); + let _ = sink.send(Message::Close(None)).await; + return; + } + }; + let (reg, reg_rpc_id) = match parse_register(®ister, &identity) { + Ok(ok) => ok, + Err((id, err)) => { + let resp = JsonRpcErrorResponse::new(id, err); + let _ = sink + .send(Message::Text( + serde_json::to_string(&resp).expect("serializable").into(), + )) + .await; + return; + } + }; + + // Outbound channel for this connection. Bounded (review F5): a peer that + // cannot drain OUTBOUND_QUEUE frames is disconnected, not buffered. + let (tx, mut rx) = mpsc::channel::(OUTBOUND_QUEUE); + + // Shutdown signal so the CP can close this socket when it drops the + // registration on its own initiative (lease expiry — review round-3 F1). + // Subscribed BEFORE registering so no signal can be missed, and kept + // alive here for the whole connection: closing is driven by an explicit + // signal, never by the registry happening to drop its side. + let shutdown = shutdown_signal(); + let mut shutdown_rx = shutdown.subscribe(); + + let effective_max = match identity.max_delegated_sessions_cap { + Some(cap) => reg.max_delegated_sessions.min(cap), + None => reg.max_delegated_sessions, + }; + // The registry assigns the CP-generated handle (review F1): ownership + // and teardown never key on the client-supplied instance_id. + let handle = state.registry.register_conn( + Instance { + handle: 0, + namespace: identity.namespace.clone(), + name: identity.name.clone(), + agent_type: identity.agent_type.clone(), + instance_id: reg.instance_id.clone(), + labels: reg.labels.clone(), + max_delegated_sessions: effective_max, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx: tx.clone(), + }, + Arc::clone(&shutdown), + ); + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + instance = %reg.instance_id, + handle, + r#type = %identity.agent_type, + max_sessions = effective_max, + "registered" + ); + + // Ack. The CP-generated handle is intentionally not disclosed. + let ack = RegisterAck { + protocol_version: PROTOCOL_VERSION, + heartbeat_interval_secs: state.cfg.heartbeat_interval_secs, + lease_expiry_secs: state.cfg.lease_expiry_secs, + effective_max_delegated_sessions: effective_max, + }; + let resp = JsonRpcResponse::new( + reg_rpc_id, + serde_json::to_value(&ack).expect("serializable"), + ); + if sink + .send(Message::Text( + serde_json::to_string(&resp).expect("serializable").into(), + )) + .await + .is_err() + { + teardown(&state, handle, &identity); + return; + } + + // --- Main loop: interleave inbound frames, outbound channel, shutdown --- + let mut cp_closed = false; + loop { + tokio::select! { + // The CP dropped this registration (lease expiry): the socket + // must go too (review round-3 F1). Keeping it open would leave a + // connection whose every frame hits an absent registry entry and + // which can never re-register, since registration is + // first-frame-only. Closing lets the client reconnect, + // re-authenticate, and register again. + _ = shutdown_rx.changed() => { + cp_closed = true; + break; + } + outbound = rx.recv() => { + match outbound { + Some(text) => { + if sink.send(Message::Text(text.into())).await.is_err() { + break; + } + } + None => break, + } + } + inbound = stream.next() => { + match inbound { + Some(Ok(Message::Text(text))) => { + if let Some(reply) = handle_frame(&state, handle, &text) { + if sink.send(Message::Text(reply.into())).await.is_err() { + break; + } + } + } + Some(Ok(Message::Ping(p))) => { + if sink.send(Message::Pong(p)).await.is_err() { + break; + } + } + Some(Ok(Message::Close(_))) | None => break, + Some(Ok(_)) => {} // binary/pong ignored + Some(Err(e)) => { + warn!(handle, err = %e, "WS error"); + break; + } + } + } + } + } + + if cp_closed { + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + handle, + "closing connection at the CP's request (registration dropped)" + ); + let _ = sink.send(Message::Close(None)).await; + } + + teardown(&state, handle, &identity); +} + +/// Deregister this connection's own registration (by handle — cannot touch +/// another connection's entry) and fail its in-flight delegations. +fn teardown(state: &Arc, handle: u64, identity: &AgentIdentity) { + state.registry.deregister(handle); + let mut next = || state.next_rpc_id(); + for (inst, frame) in state + .router + .fail_instance(&state.registry, handle, &mut next) + { + let _ = inst.tx.try_send(frame); + } + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + handle, + "disconnected" + ); +} + +/// Validate the registration frame against the authenticated identity. +/// Returns the parsed params and the request id, or an error payload. +fn parse_register( + text: &str, + identity: &AgentIdentity, +) -> Result<(RegisterParams, u64), (u64, ErrorObject)> { + let msg: JsonRpcMessage = match serde_json::from_str(text) { + Ok(m) => m, + Err(e) => { + return Err(( + 0, + ErrorObject::new(codes::INVALID_PARAMS, format!("malformed frame: {e}")), + )) + } + }; + let rpc_id = match msg.require_request_envelope() { + Ok(id) => id, + Err(err) => return Err((msg.id.unwrap_or(0), err)), + }; + if msg.method.as_deref() != Some(methods::REGISTER) { + return Err(( + rpc_id, + ErrorObject::new(codes::NOT_REGISTERED, "first frame must be cp/register"), + )); + } + let params: RegisterParams = match msg.params.and_then(|p| serde_json::from_value(p).ok()) { + Some(p) => p, + None => { + return Err(( + rpc_id, + ErrorObject::new(codes::INVALID_PARAMS, "invalid cp/register params"), + )) + } + }; + if params.protocol_version != PROTOCOL_VERSION { + return Err(( + rpc_id, + ErrorObject::new( + codes::UNSUPPORTED_VERSION, + format!( + "protocol version {} unsupported (CP speaks {})", + params.protocol_version, PROTOCOL_VERSION + ), + ), + )); + } + // Identity binding: claims must match the key's bound identity exactly. + if params.namespace != identity.namespace + || params.name != identity.name + || params.agent_type != identity.agent_type + { + return Err(( + rpc_id, + ErrorObject::new( + codes::IDENTITY_MISMATCH, + format!( + "registration claims {}/{} ({}) do not match the identity bound to this key", + params.namespace, params.name, params.agent_type + ), + ), + )); + } + if params.instance_id.trim().is_empty() { + return Err(( + rpc_id, + ErrorObject::new(codes::INVALID_PARAMS, "instance_id must be non-empty"), + )); + } + Ok((params, rpc_id)) +} + +/// Dispatch one post-registration frame. Returns an optional direct reply. +fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option { + let msg: JsonRpcMessage = match serde_json::from_str(text) { + Ok(m) => m, + Err(e) => { + let resp = JsonRpcErrorResponse::new( + 0, + ErrorObject::new(codes::INVALID_PARAMS, format!("malformed frame: {e}")), + ); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + }; + // Responses to CP-issued requests (forwarded delegates, cancels): v1 + // correlates by delegation_id inside result frames, so plain JSON-RPC + // acks are dropped. + let method = msg.method.as_deref()?.to_string(); + let rpc_id = match msg.require_request_envelope() { + Ok(id) => id, + Err(err) => { + let resp = JsonRpcErrorResponse::new(msg.id.unwrap_or(0), err); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + }; + // The sender's identity claims are never read from the frame: everything + // derives from the authenticated registration behind `handle`. + let me = state.registry.get(handle)?; + + macro_rules! params_or_err { + ($ty:ty) => { + match msg + .params + .clone() + .and_then(|p| serde_json::from_value::<$ty>(p).ok()) + { + Some(p) => p, + None => { + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new(codes::INVALID_PARAMS, "invalid params"), + ); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + } + }; + } + + match method.as_str() { + methods::HEARTBEAT => { + let _p = params_or_err!(crate::proto::HeartbeatParams); + state.registry.heartbeat(handle); + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } + methods::DELEGATE => { + let p = params_or_err!(DelegateParams); + if p.prompt.len() > state.cfg.max_prompt_bytes { + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new( + codes::INVALID_PARAMS, + format!( + "prompt exceeds max_prompt_bytes ({})", + state.cfg.max_prompt_bytes + ), + ), + ); + return Some(serde_json::to_string(&resp).expect("serializable")); + } + let outcome = state.router.delegate( + &state.cfg, + &state.registry, + &me.namespace, + &me.name, + &me.agent_type, + handle, + p, + state.next_rpc_id(), + ); + let reply = match outcome { + DelegateOutcome::Accepted(ack) => serde_json::to_string(&JsonRpcResponse::new( + rpc_id, + serde_json::to_value(&ack).expect("serializable"), + )), + DelegateOutcome::Rejected(err) => { + serde_json::to_string(&JsonRpcErrorResponse::new(rpc_id, err)) + } + }; + Some(reply.expect("serializable")) + } + methods::DELEGATE_RESULT => { + let p = params_or_err!(DelegateResultParams); + if let Some((initiator, frame)) = state.router.complete( + &state.registry, + handle, + p, + state.cfg.max_result_bytes, + state.next_rpc_id(), + ) { + let _ = initiator.tx.try_send(frame); + } + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } + methods::CANCEL => { + let p = params_or_err!(CancelParams); + match state + .router + .cancel(&state.registry, handle, &p, state.next_rpc_id()) + { + Ok(forward) => { + if let Some((target, frame)) = forward { + let _ = target.tx.try_send(frame); + } + let resp = JsonRpcResponse::new(rpc_id, serde_json::json!({"ok": true})); + Some(serde_json::to_string(&resp).expect("serializable")) + } + Err(err) => Some( + serde_json::to_string(&JsonRpcErrorResponse::new(rpc_id, err)) + .expect("serializable"), + ), + } + } + other => { + let resp = JsonRpcErrorResponse::new( + rpc_id, + ErrorObject::new(codes::METHOD_NOT_FOUND, format!("unknown method {other}")), + ); + Some(serde_json::to_string(&resp).expect("serializable")) + } + } +} + +/// 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 +/// (review round-3 F1): without it the connection task keeps running against +/// a registration that no longer exists — every later frame (heartbeats +/// included) finds no registry entry and gets no reply, 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); + 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 (review F4). + let no_ver = serde_json::json!({ + "id": 6, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-1" + } + }) + .to_string(); + let (_, err) = parse_register(&no_ver, &identity()).unwrap_err(); + assert_eq!(err.code, codes::INVALID_REQUEST); + + // Notification shape: no id. + let no_id = serde_json::json!({ + "jsonrpc": "2.0", "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-1" + } + }) + .to_string(); + let (_, err) = parse_register(&no_id, &identity()).unwrap_err(); + assert_eq!(err.code, codes::INVALID_REQUEST); + } + + fn state_with(cfg_toml: &str) -> Arc { + let cfg: CpConfig = toml::from_str(cfg_toml).unwrap(); + cfg.validate().unwrap(); + Arc::new(AppState::new(cfg)) + } + + #[test] + fn conn_quota_bounds_and_recycles_slots() { + // Review round-3 F4(b): the quota is a hard bound and the guard + // releases the slot on drop, so no exit path can leak it. + let state = state_with("max_connections_per_identity = 2"); + let id = identity(); + let p1 = state.try_acquire_conn(&id).expect("slot 1"); + let p2 = state.try_acquire_conn(&id).expect("slot 2"); + assert_eq!(state.conn_count("prod/koudu"), 2); + assert!( + state.try_acquire_conn(&id).is_none(), + "third concurrent connection must be refused" + ); + + drop(p1); + assert_eq!(state.conn_count("prod/koudu"), 1); + let p3 = state + .try_acquire_conn(&id) + .expect("released slot is reusable"); + drop(p2); + drop(p3); + assert_eq!(state.conn_count("prod/koudu"), 0); + assert!(state.try_acquire_conn(&id).is_some()); + } + + #[test] + fn conn_quota_is_per_identity() { + let state = state_with("max_connections_per_identity = 1"); + let a = identity(); + let mut b = identity(); + b.key = "k2".into(); + b.name = "worker-1".into(); + let _pa = state.try_acquire_conn(&a).expect("koudu slot"); + let _pb = state + .try_acquire_conn(&b) + .expect("worker-1 has its own quota"); + assert!( + state.try_acquire_conn(&a).is_none(), + "quota is per identity, not global" + ); + assert_eq!(state.conn_count("prod/koudu"), 1); + assert_eq!(state.conn_count("prod/worker-1"), 1); + } + + #[tokio::test] + async fn sweep_leases_signals_the_connection_before_dropping_it() { + // Review round-3 F1 at the sweeper level: the shutdown signal is + // delivered, not just the registry entry removed. (The end-to-end + // proof over a real socket lives in tests/ws_lifecycle.rs.) + let state = state_with("heartbeat_interval_secs = 1\nlease_expiry_secs = 2"); + let signal = crate::registry::shutdown_signal(); + let mut observer = signal.subscribe(); + let (tx, _rx) = mpsc::channel::(OUTBOUND_QUEUE); + let handle = state.registry.register_conn( + Instance { + handle: 0, + namespace: "prod".into(), + name: "koudu".into(), + agent_type: AgentType::Primary, + instance_id: "i-1".into(), + labels: Default::default(), + max_delegated_sessions: 1, + active_sessions: 0, + registered_at: Instant::now(), + last_heartbeat: Instant::now(), + tx, + }, + Arc::clone(&signal), + ); + + // A live lease is left alone. + sweep_leases(&state, Duration::from_secs(60)); + assert!(state.registry.get(handle).is_some()); + assert!(!*observer.borrow()); + + sweep_leases(&state, Duration::ZERO); + assert!(state.registry.get(handle).is_none()); + observer.changed().await.unwrap(); + assert!( + *observer.borrow(), + "the owning connection must be told to close" + ); + } +} diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs new file mode 100644 index 000000000..0447bf1df --- /dev/null +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -0,0 +1,254 @@ +//! End-to-end WebSocket lifecycle tests: connection termination on lease +//! expiry (review round-3 F1) and pre-registration admission bounds +//! (review round-3 F4). +//! +//! These drive a real CP over a loopback socket with a real WS client, which +//! is the only way to prove that the connection *task* reacts — the earlier +//! bug was invisible to unit tests of the registry/router alone. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use futures_util::{SinkExt, StreamExt}; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::StatusCode; +use tokio_tungstenite::tungstenite::{Error as WsError, Message}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use openab_cp::config::CpConfig; +use openab_cp::server::{app, sweep_leases, AppState}; + +const KEY: &str = "k-primary"; + +type Ws = WebSocketStream>; + +fn cfg(extra: &str) -> CpConfig { + let raw = format!( + r#" +{extra} + +[[agents]] +key = "{KEY}" +namespace = "prod" +name = "koudu" +type = "primary" +"# + ); + let cfg: CpConfig = toml::from_str(&raw).expect("test config parses"); + cfg.validate().expect("test config validates"); + cfg +} + +/// Start a CP on an ephemeral loopback port; returns its state and WS URL. +async fn spawn_cp(cfg: CpConfig) -> (Arc, String) { + let state = Arc::new(AppState::new(cfg)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let router = app(state.clone()); + tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + (state, format!("ws://{addr}/cp")) +} + +async fn connect(url: &str) -> Result { + let mut req = url.into_client_request().unwrap(); + req.headers_mut().insert( + "authorization", + format!("Bearer {KEY}").parse().expect("header value"), + ); + tokio_tungstenite::connect_async(req) + .await + .map(|(ws, _)| ws) +} + +/// Connect, retrying while the identity's quota slot is still being released +/// by the server task. +async fn connect_retry(url: &str) -> Ws { + for _ in 0..100 { + if let Ok(ws) = connect(url).await { + return ws; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("connection was never accepted"); +} + +fn register_frame(instance_id: &str) -> String { + serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": instance_id + } + }) + .to_string() +} + +/// Send `cp/register` and return the parsed reply. +async fn register(ws: &mut Ws, instance_id: &str) -> serde_json::Value { + ws.send(Message::Text(register_frame(instance_id).into())) + .await + .unwrap(); + let msg = tokio::time::timeout(Duration::from_secs(5), ws.next()) + .await + .expect("register must be answered") + .expect("stream open") + .expect("no ws error"); + serde_json::from_str(msg.to_text().unwrap()).unwrap() +} + +/// Wait until the peer closes the socket (Close frame, error, or EOF). +async fn wait_closed(ws: &mut Ws, within: Duration) -> bool { + let deadline = Instant::now() + within; + while Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(200), ws.next()).await { + Ok(None) | Ok(Some(Err(_))) => return true, + Ok(Some(Ok(Message::Close(_)))) => return true, + Ok(Some(Ok(_))) => continue, + Err(_) => continue, // read timeout: keep waiting + } + } + false +} + +#[tokio::test] +async fn lease_expiry_closes_the_connection_and_permits_reregistration() { + // Review round-3 F1: deregistering on lease expiry without terminating + // the connection left a live socket bound to a registration that no + // longer existed — heartbeats got no reply and re-registration was + // impossible (registration is first-frame-only). The CP must close it. + let (state, url) = spawn_cp(cfg("max_connections_per_identity = 1")).await; + + let mut ws = connect(&url).await.expect("first connection accepted"); + let ack = register(&mut ws, "i-1").await; + assert_eq!(ack["result"]["protocol_version"], 1, "registered"); + assert_eq!(state.registry.list("prod").len(), 1); + + // Zero lease: every registration is overdue on this pass. + sweep_leases(&state, Duration::ZERO); + assert!( + state.registry.list("prod").is_empty(), + "lease expiry deregisters" + ); + + assert!( + wait_closed(&mut ws, Duration::from_secs(5)).await, + "the connection task must observe the shutdown signal and close" + ); + drop(ws); + + // A reconnecting client re-authenticates and registers again. This also + // proves the connection slot was released (quota is 1 here). + let mut ws2 = connect_retry(&url).await; + let ack2 = register(&mut ws2, "i-2").await; + assert_eq!(ack2["result"]["protocol_version"], 1); + let live = state.registry.list("prod"); + assert_eq!(live.len(), 1); + assert_eq!(live[0].instance_id, "i-2"); +} + +#[tokio::test] +async fn ping_only_pre_registration_socket_is_closed_at_the_deadline() { + // Review round-3 F4(a): pings keep the transport alive but must not + // extend the registration deadline. + let (state, url) = spawn_cp(cfg("register_timeout_secs = 1")).await; + let mut ws = connect(&url).await.expect("connection accepted"); + + let started = Instant::now(); + let mut closed = false; + while started.elapsed() < Duration::from_secs(10) { + let _ = ws.send(Message::Ping(vec![7].into())).await; + match tokio::time::timeout(Duration::from_millis(250), ws.next()).await { + Ok(None) | Ok(Some(Err(_))) => { + closed = true; + break; + } + Ok(Some(Ok(Message::Close(_)))) => { + closed = true; + break; + } + _ => continue, + } + } + assert!( + closed, + "an authenticated socket that never registers must be closed" + ); + assert!(state.registry.list("prod").is_empty()); +} + +#[tokio::test] +async fn registration_after_the_deadline_is_not_accepted() { + // Review round-3 F4(a): the deadline is enforced, not merely advisory. + let (state, url) = spawn_cp(cfg("register_timeout_secs = 1")).await; + let mut ws = connect(&url).await.expect("connection accepted"); + tokio::time::sleep(Duration::from_millis(1_600)).await; + + let _ = ws + .send(Message::Text(register_frame("i-late").into())) + .await; + let acked = match tokio::time::timeout(Duration::from_secs(2), ws.next()).await { + Ok(Some(Ok(Message::Text(t)))) => t.contains("effective_max_delegated_sessions"), + _ => false, + }; + assert!(!acked, "a late cp/register must not be acked"); + assert!(state.registry.list("prod").is_empty()); +} + +#[tokio::test] +async fn connection_quota_rejects_over_limit_and_recycles_on_disconnect() { + // Review round-3 F4(b): the quota bounds concurrent sockets per identity + // and is released on every exit path (RAII), so connect → disconnect → + // connect always succeeds. + let (state, url) = spawn_cp(cfg( + "max_connections_per_identity = 1\nregister_timeout_secs = 30", + )) + .await; + + let mut ws = connect(&url).await.expect("first connection accepted"); + register(&mut ws, "i-1").await; + assert_eq!(state.conn_count("prod/koudu"), 1); + + match connect(&url).await { + Err(WsError::Http(resp)) => assert_eq!( + resp.status(), + StatusCode::SERVICE_UNAVAILABLE, + "over-quota upgrade must be refused before the WS handshake" + ), + Err(e) => panic!("unexpected error: {e}"), + Ok(_) => panic!("over-quota connection must be rejected"), + } + + let _ = ws.close(None).await; + drop(ws); + + let mut ws2 = connect_retry(&url).await; + register(&mut ws2, "i-2").await; + assert_eq!( + state.conn_count("prod/koudu"), + 1, + "the released slot was reused, not leaked" + ); +} + +#[tokio::test] +async fn pre_registration_sockets_count_against_the_quota() { + // Review round-3 F4(b): the quota is taken at the upgrade, so parked + // pre-registration sockets cannot be multiplied for free. + let (_state, url) = spawn_cp(cfg( + "max_connections_per_identity = 1\nregister_timeout_secs = 30", + )) + .await; + let _parked = connect(&url).await.expect("connection accepted"); + + match connect(&url).await { + Err(WsError::Http(resp)) => assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE), + Err(e) => panic!("unexpected error: {e}"), + Ok(_) => panic!("an unregistered socket must still occupy its quota slot"), + } +} diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 7a4a253d3..edf68c402 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -209,6 +209,69 @@ complete until the serving runtime returns a structured result frame The serving **runtime** emits this frame when the agent's turn ends — result delivery never depends on the sub-agent model "remembering" to report. +### v1 contract amendments (from PR #1465 review) + +The first implementation (`crates/openab-cp`) freezes the following +behaviors, resolving the review findings on identity, lifecycle, and +recovery semantics: + +- **Identity binding.** CP config owns an immutable identity table: auth key + → (`namespace`, `name`, `type`, optional capacity cap). The runtime's + registration claims are *verified against* the key's bound identity and + rejected on mismatch (`IDENTITY_MISMATCH`). Authorization never derives + from self-asserted registration fields. Keys are per-agent + (individually revocable) and presented as `Authorization: Bearer` on the + WebSocket upgrade — never in URLs. +- **CP-constructed chain.** `cp/delegate` carries only + `parent_delegation_id`; the CP derives the ancestry chain from its + in-flight table and the authenticated caller identity, then stamps it on + the forwarded frame. A runtime cannot forge ancestry, so depth/cycle + checks operate on trusted data. Policy (role, depth, cycle, namespace, + deadline caps) is enforced by the CP authoritatively; facade checks are + defense in depth only. +- **Registration lifecycle.** The first frame on a connection MUST be + `cp/register` (JSON-RPC 2.0 envelope validated — `jsonrpc: "2.0"` and a + request id are required; `protocol_version` field). Registrations are + keyed by a **CP-generated handle**, never the client-supplied + `instance_id`: a colliding `instance_id` cannot replace or tear down + another connection's registration, and all in-flight ownership checks + (completion, cancellation, parent linkage) compare handles. The ack + carries the heartbeat interval, lease window, and the effective (possibly + clamped) concurrency budget. Instances missing heartbeats past the lease + are deregistered; their in-flight delegations fail immediately with + `target_disconnected`. Heartbeats refresh the lease only — CP-owned + in-flight accounting is authoritative and never merged from runtime + reports. +- **Resource bounds.** The WS transport rejects messages over + `max_frame_bytes` before parsing; oversized `prompt`s are rejected + (`max_prompt_bytes`); per-connection outbound queues are bounded and a + peer that cannot drain its queue is treated as disconnected. Delegation + admission (duplicate check → target selection → capacity reservation → + in-flight insert) is one atomic sequence, and the in-flight entry exists + before the forward frame is sent. +- **Saturation = fast-fail.** When all matching targets are at capacity the + CP replies `SATURATED` immediately. The CP never queues — v1 has no + durable state, and a hidden in-memory queue would contradict that. + `NO_TARGET` (nothing matches) is a distinct error. +- **CP restart semantics.** The in-flight table is in-memory. After a CP + restart, in-flight delegations end as initiator-side timeouts (the + propagated deadline is the upper bound); late `cp/delegate_result` frames + for unknown ids are acknowledged, logged, and dropped so reconnecting + runtimes do not error-loop. +- **Timeout and disconnect synthesis.** A deadline sweep terminates overdue + delegations: the initiator receives a synthesized `timeout` result and the + serving runtime a best-effort `cp/cancel` (stop burning tokens). Worker + disconnect → `target_disconnected` to the initiator; initiator disconnect + → best-effort `cp/cancel` downstream. +- **Result size cap.** `cp/delegate_result.result` larger than the + configured `max_result_bytes` (default 256 KiB) is truncated head-first + with an explicit marker. +- **Idempotency.** `delegation_id` is the caller-generated idempotency key; + a duplicate in-flight id is rejected (`DUPLICATE_DELEGATION`). Only the + instance a delegation was routed to may complete it; only the initiating + instance may cancel it. + + --- ## 5. Delegation Policy @@ -365,14 +428,22 @@ of scope for v1. ## 11. Open Questions -1. **Streaming intermediate output** — should `cp/delegate` stream - `session/update`-style chunks back to the primary, or only the final - result frame? v1 leans final-only; streaming is additive. +1. ~~**Streaming intermediate output**~~ — *resolved (PR #1465 review): + committed scope as a fast-follow behind the same wire contract. Worker + runtimes will stream `session/update`-style chunks back through the CP. + Rationale: streaming is the observability substrate, not a feature — it + restores the free human visibility that Discord-mediated collaboration + provides today. It enables a read-only observer endpoint on the CP + (e.g. `wss://cp/.../observe?ns=prod`; separate read-only credential + class, namespace-scoped) so a human can tail all delegation traffic + across the fleet from one terminal. v1 ships final-result-only; the + stream frame shape is reserved in the wire contract. 2. **CP high availability** — single instance + fast re-registration is - acceptable for v1; is active/standby needed before multi-tenant use? -3. **Human-visibility directives** — should a primary be able to mirror - selected delegation traffic into a Discord thread (observability) via - existing output directives? + acceptable for v1 (restart semantics are now defined in §4); is + active/standby needed before multi-tenant use? +3. **Human-visibility directives** — Discord mirroring becomes a consumer + of the delegation stream (Q1) rather than a separate mechanism; exact + directive syntax TBD when streaming lands. 4. **AgentCore/remote runtimes** — an `agentcore-acp`-backed OAB registers like any other runtime; verify deadline propagation across the SDK boundary.