feat(tool-broker): extract connector execution into standalone aegis-tool-broker binary - #1856
Conversation
…tool-broker binary Phase 1 of the tool-broker extraction: the gateway process no longer links aegis-tool-broker-connectors (removed from src/Cargo.toml production deps, kept as a dev-dependency for the one deliberate mock-server test exception). It never resolves a real provider credential or runs a shell command itself anymore. POST /v1/broker/execute now calls the new standalone aegis-tool-broker binary over HTTP with a required service-to-service bearer token (AEGIS_TOOL_BROKER_API_TOKEN); with no broker configured it fails closed with 501, matching the OIDC/signed-policy-bundle precedent. The gateway is unchanged for authorization, hash-bound approval consumption, and receipt emission -- only the connector-execution step moved process. Ships with a Helm chart (helm/aegis-tool-broker/) and a docker-compose.full.yml service, both wired for local/dev use. Known Phase-1 gap (documented in Tool_Broker.md's Honest Scope, not glossed over): the gateway consumes the approval before the broker HTTP call, so a broker-unreachable failure burns the approval without executing the action. Covered by execute_fails_closed_when_the_broker_is_unreachable. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
ECC bundle files are already tracked in this repository. Skipping generation of another bundle PR. |
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (28)
📝 WalkthroughWalkthroughChangesStandalone tool broker
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Gateway
participant ToolBrokerClient
participant ToolBroker
participant Connector
Gateway->>ToolBrokerClient: Submit authenticated execute request
ToolBrokerClient->>ToolBroker: POST /v1/execute
ToolBroker->>Connector: Resolve credentials and execute action
Connector-->>ToolBroker: Return sanitized output
ToolBroker-->>ToolBrokerClient: Return JSON response
ToolBrokerClient-->>Gateway: Return output or mapped error
Possibly related PRs
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ast-grep (0.44.1)src/src/main.rsast-grep timed out on this file Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request extracts the tool broker execution engine from the gateway into a standalone binary (aegis-tool-broker) to isolate credentials and connector execution. The gateway now communicates with the broker over HTTP using a shared service-to-service bearer token. The feedback highlights several critical issues and improvements: a non-existent Rust version (1.96) in the Dockerfile that will break the build, a required CLI token argument that breaks the container health check, missing HTTP timeouts in both the health check and the gateway client that could cause hangs, a need for defense-in-depth checks against empty tokens, and a Helm chart configuration that should fail fast if the API token secret is missing.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # extraction). A separate binary/image from the gateway — the gateway no | ||
| # longer links aegis-tool-broker-connectors at all, so this is the only | ||
| # process that resolves a real provider credential or executes a connector. | ||
| FROM rust:1.96-bookworm AS builder |
There was a problem hiding this comment.
The Rust version 1.96 specified in the builder stage does not exist (the current stable version is in the 1.8x range). This will cause the Docker build to fail immediately. Please use a valid stable Rust version, such as 1.82-bookworm.
FROM rust:1.82-bookworm AS builder
| #[arg(long, env = "AEGIS_TOOL_BROKER_API_TOKEN", hide_env_values = true)] | ||
| api_token: String, |
There was a problem hiding this comment.
Making api_token a required non-optional string in clap will cause the --healthcheck command to fail if the AEGIS_TOOL_BROKER_API_TOKEN environment variable is not set during the health check. Additionally, we should explicitly validate that the token is not empty or whitespace-only to prevent starting the privileged service with an insecure empty token. Changing api_token to Option and validating it in main resolves both issues.
| #[arg(long, env = "AEGIS_TOOL_BROKER_API_TOKEN", hide_env_values = true)] | |
| api_token: String, | |
| #[arg(long, env = "AEGIS_TOOL_BROKER_API_TOKEN", hide_env_values = true)] | |
| api_token: Option<String>, |
| let executor = build_broker_executor(); | ||
| let app = router(executor, cli.api_token.clone()); |
There was a problem hiding this comment.
Validate that the API token is present and non-empty before starting the server, and pass the validated token to the router.
| let executor = build_broker_executor(); | |
| let app = router(executor, cli.api_token.clone()); | |
| let Some(api_token) = cli.api_token.filter(|t| !t.trim().is_empty()) else { | |
| return Err("AEGIS_TOOL_BROKER_API_TOKEN is required and cannot be empty".into()); | |
| }; | |
| let executor = build_broker_executor(); | |
| let app = router(executor, api_token); |
| pub fn new(base_url: String, api_token: String, timeout_secs: u64) -> Self { | ||
| Self { | ||
| execute_url: format!("{base_url}/v1/execute"), | ||
| api_token, | ||
| http: reqwest::Client::new(), | ||
| timeout: Duration::from_secs(timeout_secs), | ||
| } | ||
| } |
There was a problem hiding this comment.
The reqwest::Client is created without a timeout. While tokio::time::timeout is used around the .send() call, it only protects the initial request sending and header receipt. The subsequent response body reading (via .json().await) is executed outside the timeout and can hang indefinitely if the connection stalls during streaming. Configuring the timeout on the reqwest::Client itself ensures that both the request and the body read are safely bounded.
pub fn new(base_url: String, api_token: String, timeout_secs: u64) -> Self {
let timeout = Duration::from_secs(timeout_secs);
let http = reqwest::Client::builder()
.timeout(timeout)
.build()
.unwrap_or_else(|_| reqwest::Client::new());
Self {
execute_url: format!("{base_url}/v1/execute"),
api_token,
http,
timeout,
}
}| if cli.healthcheck { | ||
| let url = format!("http://{}/livez", cli.listen); | ||
| return match reqwest::get(&url).await { | ||
| Ok(resp) if resp.status().is_success() => Ok(()), | ||
| _ => std::process::exit(1), | ||
| }; | ||
| } |
There was a problem hiding this comment.
The healthcheck request uses reqwest::get which has no timeout configured. If the server is deadlocked or hanging, the healthcheck process will hang indefinitely, potentially leading to resource exhaustion. Configuring a client with a reasonable timeout (e.g., 5 seconds) prevents this.
| if cli.healthcheck { | |
| let url = format!("http://{}/livez", cli.listen); | |
| return match reqwest::get(&url).await { | |
| Ok(resp) if resp.status().is_success() => Ok(()), | |
| _ => std::process::exit(1), | |
| }; | |
| } | |
| if cli.healthcheck { | |
| let url = format!("http://{}/livez", cli.listen); | |
| let client = reqwest::Client::builder() | |
| .timeout(std::time::Duration::from_secs(5)) | |
| .build() | |
| .unwrap_or_else(|_| reqwest::Client::new()); | |
| return match client.get(&url).send().await { | |
| Ok(resp) if resp.status().is_success() => Ok(()), | |
| _ => std::process::exit(1), | |
| }; | |
| } |
| pub fn tokens_match(configured: &str, provided: &str) -> bool { | ||
| fn hash(token: &str) -> String { | ||
| use sha2::{Digest, Sha256}; | ||
| hex::encode(Sha256::digest(token.as_bytes())) | ||
| } | ||
| hash(configured) == hash(provided) | ||
| } |
There was a problem hiding this comment.
As a defense-in-depth measure, tokens_match should explicitly return false if the configured token is empty. This prevents any potential authentication bypass if the server is misconfigured with an empty token.
pub fn tokens_match(configured: &str, provided: &str) -> bool {
if configured.is_empty() {
return false;
}
fn hash(token: &str) -> String {
use sha2::{Digest, Sha256};
hex::encode(Sha256::digest(token.as_bytes()))
}
hash(configured) == hash(provided)
}| fn two_empty_tokens_are_equal_but_the_server_never_configures_an_empty_one() { | ||
| // Documents the invariant this relies on: Cli::api_token has no | ||
| // default_value, so an empty configured token is unreachable in | ||
| // production -- clap itself refuses to start the process. | ||
| assert!(tokens_match("", "")); | ||
| } |
There was a problem hiding this comment.
Update the test to assert that an empty configured token never matches, reflecting the new defense-in-depth behavior.
#[test]
fn empty_configured_token_never_matches() {
assert!(!tokens_match("", ""));
}| - name: AEGIS_TOOL_BROKER_API_TOKEN | ||
| valueFrom: | ||
| secretKeyRef: | ||
| name: {{ include "aegis-tool-broker.secretName" . }} | ||
| key: AEGIS_TOOL_BROKER_API_TOKEN | ||
| optional: true |
There was a problem hiding this comment.
Setting optional: true for the AEGIS_TOOL_BROKER_API_TOKEN secret reference allows the pod to start even if the secret or the key is missing, which will result in the container running without a token (or failing to start at runtime). Setting optional: false (or omitting it) ensures Kubernetes fails fast with a clear CreateContainerConfigError if the required secret configuration is missing.
- name: AEGIS_TOOL_BROKER_API_TOKEN
valueFrom:
secretKeyRef:
name: {{ include "aegis-tool-broker.secretName" . }}
key: AEGIS_TOOL_BROKER_API_TOKEN
optional: false
Summary
Phase 1 of the tool-broker extraction (roadmap:
docs/current-vs-roadmap.md"Tool broker" partial item; design rationale inAegisAgent_Runtime_Data_Plane.md/AegisAgent_Agent_Cage.md: "sensor/cage/egress/broker must ship as separate binaries... never inside the gateway process").aegis-tool-brokerbinary (bins/aegis-tool-broker/) owns connector execution:EnvCredentialResolver,ShellConnector/GithubConnector/etc, and the actual provider API calls. It exposesPOST /v1/executebehind a required bearer token (AEGIS_TOOL_BROKER_API_TOKEN, no default — this binary's entire surface is the privileged action) plus/livez//readyz.src/) no longer linksaegis-tool-broker-connectorsin production (removed fromsrc/Cargo.toml's[dependencies]; kept as a[dev-dependencies]-only exception soroutes/broker.rs's tests can spin up a realBrokerExecutorbehind a mock HTTP server). It never resolves a real credential or shells out itself anymore.ToolBrokerClient(src/src/tool_broker_client.rs) calls the new binary over HTTP;AppState.tool_broker: Option<Arc<ToolBrokerClient>>replaces the old in-processbroker_executor. With no broker configured,POST /v1/broker/executefails closed with 501 — same precedent as OIDC and signed policy bundles, no in-process fallback.helm/aegis-tool-broker/, hardenedsecurityContextcopied fromaegis-egress-proxy,httpGetprobes against/livez//readyz) and adocker-compose.full.ymlservice wiringAEGIS_TOOL_BROKER_URL/AEGIS_TOOL_BROKER_API_TOKENinto the gateway service.Tool_Broker.md,Implementation_Status.md,current-vs-roadmap.md) corrected to describe the actual Phase 1 topology instead of the aspirational reversed-topology/scoped-token end-state, which is not built here and remains a named follow-up.Known Phase-1 gap, documented not glossed over (
Tool_Broker.md§6 Honest Scope): the gateway consumes the approval before the broker HTTP call. If the broker is unreachable at that point, the approval is burned without the action executing — fail-closed (not fail-open), but a real UX gap. Regression test:execute_fails_closed_when_the_broker_is_unreachable.Test plan
cargo check --workspacecargo fmt --all -- --checkcargo clippy --workspace --all-targets -- -D warningscargo test --workspace -- --test-threads=1(new binary's own unit tests: auth middleware, executor config, handlers incl. credential-never-leaks; adaptedroutes/broker.rsexecute-path tests against a mock HTTP broker; newexecute_returns_501_when_tool_broker_is_unconfigured/execute_fails_closed_when_the_broker_is_unreachable)helm lint helm/aegis-tool-broker/helm template helm/aegis-tool-broker/renders correctlycargo tree -p gateway -i aegis-tool-broker-connectorsconfirms the crate is dev-dependency-only nowbins/aegis-tool-brokerstandalone against a real gateway +docker-compose.full.ymltool-broker service🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Deployment
Bug Fixes