diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..6be6d20e --- /dev/null +++ b/.dockerignore @@ -0,0 +1,20 @@ +.git +.github +.claude +.codex +.gstack +.worktrees +target +docs +scripts +*.db +*.db-shm +*.db-wal +# Keep local credentials and private keys out of the build context. +.env* +.npmrc +.cargo +*.pem +*.key +Dockerfile* +README.md diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..9f34e579 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,6 @@ +version: 2 +updates: + - package-ecosystem: docker + directory: / + schedule: + interval: weekly diff --git a/.github/workflows/container.yml b/.github/workflows/container.yml new file mode 100644 index 00000000..0d86567e --- /dev/null +++ b/.github/workflows/container.yml @@ -0,0 +1,180 @@ +name: Container + +on: + merge_group: + pull_request: + paths: + - ".dockerignore" + - ".github/workflows/container.yml" + - "Cargo.lock" + - "Cargo.toml" + - "Dockerfile" + - "docker-entrypoint.sh" + - "crates/**" + - "scripts/container-fixture-server.py" + - "scripts/fixtures/vllm-response-single-nonstreaming.json" + - "scripts/test-container-entrypoint.sh" + push: + branches: [main] + paths: + - ".dockerignore" + - ".github/workflows/container.yml" + - "Cargo.lock" + - "Cargo.toml" + - "Dockerfile" + - "docker-entrypoint.sh" + - "crates/**" + - "scripts/container-fixture-server.py" + - "scripts/fixtures/vllm-response-single-nonstreaming.json" + - "scripts/test-container-entrypoint.sh" + +permissions: + contents: read + +jobs: + build-and-smoke-test: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Check out repository + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Test container entrypoint + run: scripts/test-container-entrypoint.sh + + - name: Build runtime image + env: + DOCKER_BUILDKIT: "1" + run: | + docker build \ + --build-arg OCI_BUILD_PIPELINE="$GITHUB_WORKFLOW" \ + --build-arg OCI_BUILD_URL="$GITHUB_SERVER_URL/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID" \ + --build-arg OCI_CREATED="$(date -u +%Y-%m-%dT%H:%M:%SZ)" \ + --build-arg OCI_REVISION="$GITHUB_SHA" \ + --build-arg OCI_VERSION="${GITHUB_REF_NAME}" \ + --tag agentic-api:test \ + . + + - name: Verify runtime contents and identity + run: | + docker run --rm --entrypoint sh agentic-api:test -eu -c ' + for command in cargo rustc python vllm; do + ! command -v "$command" + done + test "$(id -u)" -ne 0 + test "$(id -g)" -eq 0 + test ! -w /usr/local/bin/agentic-server + test ! -w /usr/local/bin/docker-entrypoint.sh + ' + container_id=$(docker create agentic-api:test) + binary_copy=$(mktemp) + cleanup_binary() { + docker rm --force "$container_id" >/dev/null 2>&1 || true + rm -f "$binary_copy" + } + trap cleanup_binary EXIT + docker cp "$container_id:/usr/local/bin/agentic-server" "$binary_copy" + binary_description=$(file "$binary_copy") + echo "$binary_description" + case "$binary_description" in + *"not stripped"*) exit 1 ;; + *"stripped"*) ;; + *) echo "unable to verify that agentic-server is stripped" >&2; exit 1 ;; + esac + volume_dir=$(mktemp -d) + cleanup_volume() { + sudo rm -rf "$volume_dir" + } + trap 'cleanup_binary; cleanup_volume' EXIT + sudo chgrp 0 "$volume_dir" + chmod 2775 "$volume_dir" + docker run --rm --user 12345:0 \ + --volume "$volume_dir:/var/lib/agentic-api" \ + --entrypoint sh agentic-api:test \ + -eu -c 'touch /var/lib/agentic-api/write-check && test -w /var/lib/agentic-api/write-check' + + - name: Test Responses API end to end + run: | + python3 scripts/container-fixture-server.py \ + --fixture scripts/fixtures/vllm-response-single-nonstreaming.json \ + --port 18000 \ + >/tmp/agentic-api-upstream.log 2>&1 & + upstream_pid=$! + volume=agentic-api-smoke-data + upstream_response_id=$(jq --exit-status --raw-output \ + '.id | select(type == "string" and length > 0)' \ + scripts/fixtures/vllm-response-single-nonstreaming.json) + + cleanup() { + docker logs agentic-api-smoke 2>/dev/null || true + docker rm --force agentic-api-smoke >/dev/null 2>&1 || true + docker volume rm "$volume" >/dev/null 2>&1 || true + kill "$upstream_pid" 2>/dev/null || true + } + trap cleanup EXIT + + start_gateway() { + docker run --detach --name agentic-api-smoke --network host --user "$1:0" \ + --volume "$volume:/var/lib/agentic-api" \ + --env LLM_API_BASE=http://127.0.0.1:18000 \ + agentic-api:test + } + + wait_until_ready() { + for attempt in $(seq 1 30); do + if curl --connect-timeout 1 --max-time 3 --fail --silent http://127.0.0.1:9000/health >/dev/null && \ + curl --connect-timeout 1 --max-time 3 --fail --silent http://127.0.0.1:9000/ready >/dev/null; then + return 0 + fi + echo "gateway not ready (attempt $attempt/30)" + sleep 1 + done + return 1 + } + + post_response() { + jq --null-input --compact-output --arg previous_response_id "$1" ' + { + input: "Reply with exactly one word: HELLO", + model: "gpt-4o", + store: true, + stream: false + } + if $previous_response_id == "" then {} else { + previous_response_id: $previous_response_id + } end + ' >/tmp/agentic-api-request.json + curl --connect-timeout 2 --max-time 15 --fail --silent --show-error \ + --header 'Content-Type: application/json' \ + --data-binary @/tmp/agentic-api-request.json \ + --output "$2" \ + http://127.0.0.1:9000/v1/responses + jq --exit-status --arg upstream_response_id "$upstream_response_id" ' + .status == "completed" and + (.id | startswith("resp_")) and + .id != $upstream_response_id and + .model == "gpt-4o" and + .output[0].content[0].text == "HI" + ' "$2" + } + + docker volume create "$volume" + start_gateway 12345 + wait_until_ready + post_response "" /tmp/agentic-api-response-1.json + first_response_id=$(jq --exit-status --raw-output '.id' /tmp/agentic-api-response-1.json) + docker exec agentic-api-smoke test -f /var/lib/agentic-api/agentic_api.db + docker stop --timeout 10 agentic-api-smoke + test "$(docker inspect --format '{{.State.ExitCode}}' agentic-api-smoke)" -eq 0 + docker rm agentic-api-smoke + + start_gateway 23456 + wait_until_ready + post_response "$first_response_id" /tmp/agentic-api-response-2.json + second_response_id=$(jq --exit-status --raw-output '.id' /tmp/agentic-api-response-2.json) + post_response "$second_response_id" /tmp/agentic-api-response-3.json + if docker logs agentic-api-smoke 2>&1 | grep --quiet 'persist failed'; then + echo "SQLite persistence failed after arbitrary-UID rotation" >&2 + exit 1 + fi + docker stop --timeout 10 agentic-api-smoke + test "$(docker inspect --format '{{.State.ExitCode}}' agentic-api-smoke)" -eq 0 diff --git a/Cargo.lock b/Cargo.lock index 69a23195..7a7d66e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -48,6 +48,7 @@ dependencies = [ "http", "indexmap", "reqwest 0.12.28", + "reqwest 0.13.4", "rmcp", "serde", "serde_json", @@ -57,6 +58,7 @@ dependencies = [ "thiserror", "tokio", "tracing", + "url", "uuid", ] diff --git a/Cargo.toml b/Cargo.toml index ce3289ee..82f3388e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ futures = "0.3" indexmap = "2" http = "1" reqwest = { version = "0.12", default-features = false } +rmcp-reqwest = { package = "reqwest", version = "0.13.2", default-features = false, features = ["json", "stream", "rustls"] } rmcp = { version = "1.8", default-features = false } serde = { version = "1", features = ["derive"] } serde_json = "1" @@ -37,3 +38,4 @@ tower-http = { version = "0.6", features = ["cors"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } tokio-tungstenite = "0.29" +url = "2" diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 00000000..776bba54 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,65 @@ +# syntax=docker/dockerfile:1.7 + +ARG RUST_VERSION=1.96.0 +ARG DEBIAN_VERSION=bookworm +ARG RUST_IMAGE_DIGEST=sha256:5e2214abe154fe26e39f64488952e5c991eeed1d6d6da7cc8381ae83927f0cfc +ARG DEBIAN_IMAGE_DIGEST=sha256:7b140f374b289a7c2befc338f42ebe6441b7ea838a042bbd5acbfca6ec875818 + +FROM rust:${RUST_VERSION}-${DEBIAN_VERSION}@${RUST_IMAGE_DIGEST} AS rust-build + +ARG CARGO_BUILD_JOBS=4 +ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS} + +WORKDIR /workspace +COPY Cargo.toml Cargo.lock ./ +COPY crates ./crates + +RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=locked \ + --mount=type=cache,target=/usr/local/cargo/git,sharing=locked \ + --mount=type=cache,target=/workspace/target,sharing=locked \ + cargo build --locked --release -p agentic-server && \ + install -Dm755 -s target/release/agentic-server /out/agentic-server + +FROM debian:${DEBIAN_VERSION}-slim@${DEBIAN_IMAGE_DIGEST} AS runtime + +ARG RUNTIME_GID=0 +ARG RUNTIME_UID=10001 + +RUN apt-get update && \ + apt-get install --yes --no-install-recommends ca-certificates && \ + rm -rf /var/lib/apt/lists/* && \ + mkdir -p /var/lib/agentic-api && \ + chown "${RUNTIME_UID}:${RUNTIME_GID}" /var/lib/agentic-api && \ + chmod g=u,g+s /var/lib/agentic-api + +COPY --from=rust-build /out/agentic-server /usr/local/bin/agentic-server +COPY --chmod=0755 docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh + +ARG OCI_CREATED="" +ARG OCI_BUILD_PIPELINE=local +ARG OCI_BUILD_URL="" +ARG OCI_REVISION="" +ARG OCI_SOURCE="https://github.com/vllm-project/agentic-api" +ARG OCI_VERSION="" + +LABEL org.opencontainers.image.created="${OCI_CREATED}" \ + org.opencontainers.image.description="Rust gateway for stateful agentic APIs backed by vLLM" \ + org.opencontainers.image.licenses="Apache-2.0" \ + org.opencontainers.image.revision="${OCI_REVISION}" \ + org.opencontainers.image.source="${OCI_SOURCE}" \ + org.opencontainers.image.title="agentic-api" \ + org.opencontainers.image.url="${OCI_BUILD_URL}" \ + org.opencontainers.image.version="${OCI_VERSION}" \ + ai.vllm.build.commit="${OCI_REVISION}" \ + ai.vllm.build.pipeline="${OCI_BUILD_PIPELINE}" \ + ai.vllm.build.url="${OCI_BUILD_URL}" \ + ai.vllm.image.tag="${OCI_VERSION}" + +WORKDIR /var/lib/agentic-api +USER ${RUNTIME_UID}:${RUNTIME_GID} + +ENV GATEWAY_HOST=0.0.0.0 \ + GATEWAY_PORT=9000 + +EXPOSE 9000 +ENTRYPOINT ["docker-entrypoint.sh"] diff --git a/crates/agentic-server-core/Cargo.toml b/crates/agentic-server-core/Cargo.toml index b7919fcc..e84ec6ae 100644 --- a/crates/agentic-server-core/Cargo.toml +++ b/crates/agentic-server-core/Cargo.toml @@ -22,6 +22,7 @@ futures.workspace = true indexmap.workspace = true http.workspace = true reqwest = { workspace = true, features = ["default-tls", "stream"] } +rmcp-reqwest.workspace = true rmcp = { workspace = true, features = [ "client", "reqwest", @@ -33,6 +34,7 @@ serde_json.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["time"] } tracing.workspace = true +url.workspace = true sqlx = { version = "0.8", features = ["runtime-tokio-rustls", "any", "migrate", "sqlite", "postgres"] } chrono = { version = "0.4", features = ["serde"] } diff --git a/crates/agentic-server-core/src/tool/mcp/client.rs b/crates/agentic-server-core/src/tool/mcp/client.rs index 7abeaff4..c82a648a 100644 --- a/crates/agentic-server-core/src/tool/mcp/client.rs +++ b/crates/agentic-server-core/src/tool/mcp/client.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::fmt; +use std::net::SocketAddr; use std::process::Stdio; use std::sync::Arc; use std::time::Duration; @@ -15,6 +16,7 @@ use rmcp::service::{ClientInitializeError, PeerRequestOptions, RoleClient, Runni use rmcp::transport::StreamableHttpClientTransport; use rmcp::transport::child_process::TokioChildProcess; use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; +use rmcp_reqwest as http_client; use serde_json::Value; use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::process::Command; @@ -49,6 +51,15 @@ pub enum McpError { #[error("failed to connect to MCP server")] Connect(#[source] Box), + #[error("failed to resolve MCP server host")] + ResolveHost(#[source] std::io::Error), + + #[error("MCP server URL has no resolvable host")] + UnresolvableHost, + + #[error("failed to build MCP HTTP client")] + BuildHttpClient(#[source] http_client::Error), + #[error("invalid MCP HTTP header name")] InvalidHeaderName(#[source] http::header::InvalidHeaderName), @@ -95,8 +106,21 @@ impl McpClient { /// /// # Errors /// - /// Returns [`McpError::Connect`] if the MCP initialization handshake fails. + /// Returns an error if URL resolution, HTTP client or header construction, + /// the initialization timeout, or the MCP handshake fails. pub async fn connect(server_url: &str, headers: Option>) -> Result { + tokio::time::timeout(CONNECTION_TIMEOUT, Self::connect_streamable_http(server_url, headers)) + .await + .map_err(|_| McpError::Timeout { + operation: McpOperation::Connect, + })? + } + + async fn connect_streamable_http( + server_url: &str, + headers: Option>, + ) -> Result { + let http_client = pinned_http_client(server_url).await?; let mut config = StreamableHttpClientTransportConfig::with_uri(server_url.to_owned()); if let Some(headers) = headers.filter(|headers| !headers.is_empty()) { let mut custom_headers = HashMap::with_capacity(headers.len()); @@ -108,12 +132,10 @@ impl McpClient { } config = config.custom_headers(custom_headers); } - let transport = StreamableHttpClientTransport::from_config(config); - let service = tokio::time::timeout(CONNECTION_TIMEOUT, AgenticMcpClientHandler.serve(transport)) + let transport = StreamableHttpClientTransport::with_client(http_client, config); + let service = AgenticMcpClientHandler + .serve(transport) .await - .map_err(|_| McpError::Timeout { - operation: McpOperation::Connect, - })? .map_err(|error| McpError::Connect(Box::new(error)))?; Ok(Self { @@ -277,3 +299,199 @@ impl McpClient { }) } } + +/// Build the HTTP client used by an MCP connection with DNS pinned to the +/// addresses resolved during connection setup. This prevents a hostname from +/// resolving to different addresses between URL validation and later requests. +async fn pinned_http_client(server_url: &str) -> Result { + let url = http_client::Url::parse(server_url).map_err(|_| McpError::UnresolvableHost)?; + let port = url.port_or_known_default().ok_or(McpError::UnresolvableHost)?; + match url.host().ok_or(McpError::UnresolvableHost)? { + url::Host::Domain(host) => { + let addresses = tokio::net::lookup_host((host, port)) + .await + .map_err(McpError::ResolveHost)? + .collect::>(); + if addresses.is_empty() { + return Err(McpError::UnresolvableHost); + } + http_client_for_addresses(host, &addresses) + } + url::Host::Ipv4(_) | url::Host::Ipv6(_) => http_client_for_literal_address(), + } +} + +fn http_client_for_addresses(host: &str, addresses: &[SocketAddr]) -> Result { + http_client::Client::builder() + .no_proxy() + .redirect(http_client::redirect::Policy::none()) + .resolve_to_addrs(host, addresses) + .build() + .map_err(McpError::BuildHttpClient) +} + +fn http_client_for_literal_address() -> Result { + http_client::Client::builder() + .no_proxy() + .redirect(http_client::redirect::Policy::none()) + .build() + .map_err(McpError::BuildHttpClient) +} + +#[cfg(test)] +mod tests { + use super::{McpError, http_client_for_addresses, pinned_http_client}; + use std::io::{Read, Write}; + use std::net::{Ipv4Addr, SocketAddr, TcpListener as StdTcpListener}; + use std::process::Command; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::thread; + use std::time::Duration; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener as TokioTcpListener; + + const PROXY_CHILD_ENV: &str = "AGENTIC_MCP_PROXY_TEST_CHILD"; + const PROXY_TARGET_ENV: &str = "AGENTIC_MCP_PROXY_TEST_TARGET"; + + async fn spawn_http_server(response: &'static str) -> SocketAddr { + let listener = TokioTcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 1024]; + let _ = stream.read(&mut request).await.unwrap(); + stream.write_all(response.as_bytes()).await.unwrap(); + }); + address + } + + #[tokio::test] + async fn pinned_client_tries_every_resolved_address() { + let reachable = spawn_http_server("HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").await; + let unreachable = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], reachable.port())); + let client = http_client_for_addresses("mcp.test", &[unreachable, reachable]).unwrap(); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(5), + client.get(format!("http://mcp.test:{}/", reachable.port())).send(), + ) + .await + .expect("client must try the reachable address") + .unwrap(); + + assert_eq!(response.status(), http::StatusCode::OK); + } + + #[tokio::test] + async fn pinned_client_does_not_follow_redirects() { + let address = spawn_http_server( + "HTTP/1.1 302 Found\r\nLocation: http://127.0.0.1:9/\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await; + let client = http_client_for_addresses("mcp.test", &[address]).unwrap(); + + let response = client + .get(format!("http://mcp.test:{}/", address.port())) + .send() + .await + .unwrap(); + + assert!(response.status().is_redirection()); + } + + #[tokio::test] + async fn pinned_client_rejects_malformed_urls() { + assert!(matches!( + pinned_http_client("not a URL").await, + Err(McpError::UnresolvableHost) + )); + } + + #[tokio::test] + async fn pinned_client_accepts_ipv6_literal_urls_without_dns() { + pinned_http_client("http://[::1]:8000/mcp").await.unwrap(); + } + + fn serve_once(listener: &StdTcpListener, status: &str, accepted: &AtomicBool) { + listener.set_nonblocking(true).unwrap(); + for _ in 0..200 { + match listener.accept() { + Ok((mut stream, _)) => { + accepted.store(true, Ordering::SeqCst); + let mut request = [0_u8; 1024]; + let _ = stream.read(&mut request).unwrap(); + write!( + stream, + "HTTP/1.1 {status}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .unwrap(); + return; + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(5)); + } + Err(error) => panic!("test server failed: {error}"), + } + } + } + + #[test] + fn pinned_client_ignores_system_proxy() { + if std::env::var_os(PROXY_CHILD_ENV).is_some() { + let target = std::env::var(PROXY_TARGET_ENV).unwrap().parse::().unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async move { + let client = http_client_for_addresses("mcp.test", &[target]).unwrap(); + let response = client + .get(format!("http://mcp.test:{}/", target.port())) + .send() + .await + .unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + }); + return; + } + + let target_listener = StdTcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let target_address = target_listener.local_addr().unwrap(); + let target_accepted = Arc::new(AtomicBool::new(false)); + let target_thread = { + let accepted = Arc::clone(&target_accepted); + thread::spawn(move || serve_once(&target_listener, "200 OK", &accepted)) + }; + + let proxy_listener = StdTcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let proxy_address = proxy_listener.local_addr().unwrap(); + let proxy_accepted = Arc::new(AtomicBool::new(false)); + let proxy_thread = { + let accepted = Arc::clone(&proxy_accepted); + thread::spawn(move || serve_once(&proxy_listener, "418 I'm a teapot", &accepted)) + }; + + let output = Command::new(std::env::current_exe().unwrap()) + .arg("--exact") + .arg("tool::mcp::client::tests::pinned_client_ignores_system_proxy") + .arg("--nocapture") + .env(PROXY_CHILD_ENV, "1") + .env(PROXY_TARGET_ENV, target_address.to_string()) + .env("HTTP_PROXY", format!("http://{proxy_address}")) + .env("http_proxy", format!("http://{proxy_address}")) + .env("ALL_PROXY", format!("http://{proxy_address}")) + .env("all_proxy", format!("http://{proxy_address}")) + .env_remove("NO_PROXY") + .env_remove("no_proxy") + .output() + .unwrap(); + + target_thread.join().unwrap(); + proxy_thread.join().unwrap(); + assert!( + output.status.success(), + "child test failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(target_accepted.load(Ordering::SeqCst)); + assert!(!proxy_accepted.load(Ordering::SeqCst)); + } +} diff --git a/crates/agentic-server-core/src/tool/mcp/pool.rs b/crates/agentic-server-core/src/tool/mcp/pool.rs index c63e417e..33927bdf 100644 --- a/crates/agentic-server-core/src/tool/mcp/pool.rs +++ b/crates/agentic-server-core/src/tool/mcp/pool.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; -use std::net::IpAddr; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use reqwest::Url; use serde::{Deserialize, Serialize}; @@ -8,9 +7,10 @@ use serde::{Deserialize, Serialize}; use super::client::McpClient; use crate::types::tools::McpToolParam; -// Hostnames configured here are a trust boundary: validation compares the -// hostname string but does not pin its DNS resolution to the transport. Only -// add names whose DNS records are controlled by a trusted administrator. +// Hostnames configured here are a trust boundary. The HTTP client resolves a +// configured name once per connection and pins all returned addresses for that +// transport. Only add names whose DNS records are controlled by a trusted +// administrator. const MCP_ALLOWED_HOSTS_ENV: &str = "AGENTIC_MCP_ALLOWED_HOSTS"; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -177,8 +177,8 @@ fn validate_request_server_url(value: &str) -> Result { return Err("URL must not include credentials".to_owned()); } - let host = url.host_str().ok_or_else(|| "URL must include a host".to_owned())?; - if is_allowed_request_host(host) { + let host = url.host().ok_or_else(|| "URL must include a host".to_owned())?; + if is_allowed_request_host(&host) { return Ok(value.to_owned()); } @@ -187,19 +187,32 @@ fn validate_request_server_url(value: &str) -> Result { )) } -fn is_allowed_request_host(host: &str) -> bool { - host.eq_ignore_ascii_case("localhost") - || host.parse::().is_ok_and(|ip| ip.is_loopback()) - || host_allowed_by_env(host) +fn is_allowed_request_host(host: &url::Host<&str>) -> bool { + match host { + url::Host::Domain(host) => host.eq_ignore_ascii_case("localhost") || host_allowed_by_env(host), + url::Host::Ipv4(address) => address.is_loopback() || host_allowed_by_env(&address.to_string()), + url::Host::Ipv6(address) => address.is_loopback() || host_allowed_by_env(&address.to_string()), + } } fn host_allowed_by_env(host: &str) -> bool { - std::env::var(MCP_ALLOWED_HOSTS_ENV).is_ok_and(|allowed_hosts| { - allowed_hosts - .split(',') - .map(str::trim) - .any(|allowed_host| allowed_host.eq_ignore_ascii_case(host)) - }) + allowed_hosts() + .iter() + .any(|allowed_host| allowed_host.eq_ignore_ascii_case(host)) +} + +fn allowed_hosts() -> &'static [String] { + static ALLOWED_HOSTS: OnceLock> = OnceLock::new(); + ALLOWED_HOSTS.get_or_init(|| parse_allowed_hosts(&std::env::var(MCP_ALLOWED_HOSTS_ENV).unwrap_or_default())) +} + +fn parse_allowed_hosts(value: &str) -> Vec { + value + .split(',') + .map(str::trim) + .filter(|host| !host.is_empty()) + .map(str::to_owned) + .collect() } fn clean_string(value: Option<&str>) -> Option { @@ -211,7 +224,7 @@ fn clean_string(value: Option<&str>) -> Option { #[cfg(test)] mod tests { - use super::{McpServerEntry, server_entry_from_param, validate_request_server_url}; + use super::{McpServerEntry, parse_allowed_hosts, server_entry_from_param, validate_request_server_url}; use crate::types::tools::McpToolParam; #[test] @@ -263,6 +276,12 @@ mod tests { assert_eq!(url, "http://127.0.0.1:8000/mcp"); } + #[test] + fn request_server_url_allows_ipv6_loopback_http() { + let url = validate_request_server_url("http://[::1]:8000/mcp").unwrap(); + assert_eq!(url, "http://[::1]:8000/mcp"); + } + #[test] fn request_server_url_rejects_unallowlisted_host() { let error = validate_request_server_url("http://169.254.169.254/mcp").unwrap_err(); @@ -281,4 +300,12 @@ mod tests { assert!(server_entry_from_param(¶m).is_none()); } + + #[test] + fn allowed_host_parser_trims_and_discards_empty_entries() { + assert_eq!( + parse_allowed_hosts(" Example.COM, ,api.test "), + vec!["Example.COM", "api.test"] + ); + } } diff --git a/crates/agentic-server-core/tests/stateful_responses_integration.rs b/crates/agentic-server-core/tests/stateful_responses_integration.rs index 9c7a3b4e..040781cd 100644 --- a/crates/agentic-server-core/tests/stateful_responses_integration.rs +++ b/crates/agentic-server-core/tests/stateful_responses_integration.rs @@ -115,6 +115,51 @@ async fn test_single_turn_streaming_emits_response_completed_event() { ); } +#[tokio::test] +async fn test_stream_persists_when_client_disconnects_after_completion_event() { + let cassette = load_cassette(&format!("{DIR}/resp-single-gpt-4o-streaming.yaml")); + let t1 = &cassette.turns[0]; + let responses = vec![t1, t1]; + let fixture = TestFixture::new(&responses).await; + + let result = execute( + make_request(&t1.request.body.input, true, true, None, None), + Arc::clone(&fixture.exec_ctx), + ) + .await + .expect("execute"); + let Either::Right(stream) = result else { + panic!("expected streaming response"); + }; + let mut stream = Box::pin(stream); + let response_id = loop { + let Some(chunk) = stream.next().await else { + panic!("stream ended before response.completed"); + }; + let Some(data) = chunk.trim_end_matches('\n').strip_prefix("data: ") else { + continue; + }; + let Ok(event) = serde_json::from_str::(data) else { + continue; + }; + if event["type"] == "response.completed" { + break event["response"]["id"].as_str().expect("response id").to_owned(); + } + }; + + // Dropping before requesting the next chunk simulates a client disconnect + // immediately after receiving the terminal response event. + drop(stream); + + let follow_up = execute( + make_request("follow up", true, true, Some(response_id), None), + Arc::clone(&fixture.exec_ctx), + ) + .await + .expect("response must be persisted before response.completed is emitted"); + assert!(matches!(follow_up, Either::Right(_))); +} + /// Case 3 — two turns, non-streaming, chained via `previous_response_id`. #[tokio::test] async fn test_two_turn_nonstreaming_previous_response_id() { diff --git a/crates/agentic-server/benches/gateway_bench.rs b/crates/agentic-server/benches/gateway_bench.rs index 29881589..488e51b0 100644 --- a/crates/agentic-server/benches/gateway_bench.rs +++ b/crates/agentic-server/benches/gateway_bench.rs @@ -43,7 +43,7 @@ use agentic_core::config::Config; use agentic_core::executor::{ConversationHandler, ExecutionContext, ResponseHandler}; use agentic_core::proxy::ProxyState; use agentic_core::storage::{ConversationStore, ResponseStore, create_pool_with_schema}; -use agentic_server::app::{AppState, ServerConfig, build_router}; +use agentic_server::app::{AppState, ServerConfig, WebSocketTracker, build_router}; fn bench_turns() -> usize { std::env::var("BENCH_TURNS") @@ -173,6 +173,7 @@ async fn spawn_gateway(llm_url: &str) -> (Arc, String) { proxy_state, exec_ctx, shutdown_token: CancellationToken::new(), + websocket_tracker: WebSocketTracker::default(), llm_api_base: config.llm_api_base, openai_api_key: config.openai_api_key, }; diff --git a/crates/agentic-server/benches/proxy_bench.rs b/crates/agentic-server/benches/proxy_bench.rs index 4ec2a44f..4dce033b 100644 --- a/crates/agentic-server/benches/proxy_bench.rs +++ b/crates/agentic-server/benches/proxy_bench.rs @@ -18,7 +18,7 @@ use agentic_core::config::Config; use agentic_core::executor::{ConversationHandler, ExecutionContext, ResponseHandler}; use agentic_core::proxy::ProxyState; use agentic_core::storage::{ConversationStore, ResponseStore}; -use agentic_server::app::{AppState, ServerConfig, build_router}; +use agentic_server::app::{AppState, ServerConfig, WebSocketTracker, build_router}; const CONTENT_TYPE_JSON: &str = "application/json"; const PROMPT_SIZES: [usize; 3] = [1024, 10 * 1024, 100 * 1024]; @@ -93,6 +93,7 @@ async fn spawn_gateway(config: Config) -> String { proxy_state, exec_ctx, shutdown_token: CancellationToken::new(), + websocket_tracker: WebSocketTracker::default(), llm_api_base: config.llm_api_base, openai_api_key: config.openai_api_key, }; diff --git a/crates/agentic-server/src/app.rs b/crates/agentic-server/src/app.rs index 9e5baccf..c2771c57 100644 --- a/crates/agentic-server/src/app.rs +++ b/crates/agentic-server/src/app.rs @@ -1,8 +1,10 @@ use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use axum::Router; use axum::routing::{get, post}; use http::HeaderValue; +use tokio::sync::Notify; use tokio_util::sync::CancellationToken; use tower_http::cors::{AllowOrigin, Any, CorsLayer}; @@ -11,6 +13,51 @@ use agentic_core::proxy::ProxyState; use crate::handler::{conversations, count_tokens, health, messages, models, ready, responses, responses_ws}; +#[derive(Clone, Default)] +pub struct WebSocketTracker { + inner: Arc, +} + +#[derive(Default)] +struct WebSocketTrackerInner { + active: AtomicUsize, + idle: Notify, +} + +pub(crate) struct WebSocketGuard { + inner: Arc, +} + +impl WebSocketTracker { + pub(crate) fn track(&self) -> WebSocketGuard { + self.inner.active.fetch_add(1, Ordering::AcqRel); + WebSocketGuard { + inner: Arc::clone(&self.inner), + } + } + + /// Wait until every upgraded WebSocket task has finished. + pub async fn wait_until_idle(&self) { + loop { + let idle = self.inner.idle.notified(); + tokio::pin!(idle); + idle.as_mut().enable(); + if self.inner.active.load(Ordering::Acquire) == 0 { + return; + } + idle.await; + } + } +} + +impl Drop for WebSocketGuard { + fn drop(&mut self) { + if self.inner.active.fetch_sub(1, Ordering::AcqRel) == 1 { + self.inner.idle.notify_waiters(); + } + } +} + /// Server-level configuration read from environment variables. pub struct ServerConfig { pub cors_allowed_origins: Vec, @@ -62,6 +109,8 @@ pub struct AppState { pub exec_ctx: Arc, /// Shared cancellation signal used to drain long-lived handlers. pub shutdown_token: CancellationToken, + /// Tracks upgraded WebSocket tasks, which Axum does not await during HTTP drain. + pub websocket_tracker: WebSocketTracker, /// vLLM base URL — used by the `/ready` health probe. pub llm_api_base: String, /// Server-configured API key; used as fallback when the request carries no diff --git a/crates/agentic-server/src/handler/websocket/error.rs b/crates/agentic-server/src/handler/websocket/error.rs index 16b42f08..43b013d6 100644 --- a/crates/agentic-server/src/handler/websocket/error.rs +++ b/crates/agentic-server/src/handler/websocket/error.rs @@ -27,9 +27,6 @@ pub(super) enum WsError { #[error("websocket client disconnected")] ClientDisconnected, - #[error("websocket shutdown requested")] - Shutdown, - #[error("websocket receive failed: {0}")] Receive(String), } @@ -39,11 +36,9 @@ impl WsError { match self { Self::Executor(err) => err.http_status(), Self::InvalidJson(_) | Self::UnexpectedType | Self::BinaryFrame => StatusCode::BAD_REQUEST, - Self::SerializeJson(_) - | Self::SendFailed - | Self::ClientDisconnected - | Self::Shutdown - | Self::Receive(_) => StatusCode::INTERNAL_SERVER_ERROR, + Self::SerializeJson(_) | Self::SendFailed | Self::ClientDisconnected | Self::Receive(_) => { + StatusCode::INTERNAL_SERVER_ERROR + } } } @@ -52,18 +47,14 @@ impl WsError { Self::Executor(err) => err.error_code(), Self::InvalidJson(_) => "invalid_json", Self::UnexpectedType | Self::BinaryFrame => "invalid_request_error", - Self::SerializeJson(_) - | Self::SendFailed - | Self::ClientDisconnected - | Self::Shutdown - | Self::Receive(_) => "server_error", + Self::SerializeJson(_) | Self::SendFailed | Self::ClientDisconnected | Self::Receive(_) => "server_error", } } pub(super) fn to_ws_frame(&self) -> Option { if matches!( self, - Self::SerializeJson(_) | Self::SendFailed | Self::ClientDisconnected | Self::Shutdown | Self::Receive(_) + Self::SerializeJson(_) | Self::SendFailed | Self::ClientDisconnected | Self::Receive(_) ) { return None; } diff --git a/crates/agentic-server/src/handler/websocket/responses.rs b/crates/agentic-server/src/handler/websocket/responses.rs index 25ca6513..cf46efe8 100644 --- a/crates/agentic-server/src/handler/websocket/responses.rs +++ b/crates/agentic-server/src/handler/websocket/responses.rs @@ -7,7 +7,7 @@ use axum::http::HeaderMap; use axum::response::Response; use either::Either; use futures::stream::{SplitSink, SplitStream}; -use futures::{SinkExt, StreamExt}; +use futures::{SinkExt, Stream, StreamExt}; use serde_json::Value; use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; @@ -25,9 +25,13 @@ type WsSender = SplitSink; type WsReceiver = SplitStream; pub async fn responses_ws(State(state): State, headers: HeaderMap, ws: WebSocketUpgrade) -> Response { + let websocket_guard = state.websocket_tracker.track(); ws.max_message_size(MAX_BODY_SIZE) .max_frame_size(MAX_BODY_SIZE) - .on_upgrade(move |socket| responses_ws_loop(socket, state, headers)) + .on_upgrade(move |socket| async move { + let _websocket_guard = websocket_guard; + responses_ws_loop(socket, state, headers).await; + }) } async fn responses_ws_loop(socket: WebSocket, state: AppState, headers: HeaderMap) { @@ -39,6 +43,9 @@ async fn responses_ws_loop(socket: WebSocket, state: AppState, headers: HeaderMa let mut queue: VecDeque = VecDeque::new(); loop { + if shutdown_token.is_cancelled() { + break; + } let text = if let Some(buffered) = queue.pop_front() { buffered } else { @@ -93,6 +100,9 @@ async fn responses_ws_loop(socket: WebSocket, state: AppState, headers: HeaderMa } } } + if let Err(error) = sender.close().await { + debug!(%error, "failed to close responses websocket cleanly"); + } debug!("responses websocket session closed"); } @@ -203,6 +213,35 @@ fn empty_response_event( }) } +enum ShutdownInput { + Receiver(Option), + Upstream(Option), +} + +async fn next_shutdown_input( + receiver: &mut Receiver, + upstream: &mut Upstream, + prefer_receiver: bool, +) -> ShutdownInput +where + Receiver: Stream + Unpin, + Upstream: Stream + Unpin, +{ + if prefer_receiver { + tokio::select! { + biased; + message = receiver.next() => ShutdownInput::Receiver(message), + line = upstream.next() => ShutdownInput::Upstream(line), + } + } else { + tokio::select! { + biased; + line = upstream.next() => ShutdownInput::Upstream(line), + message = receiver.next() => ShutdownInput::Receiver(message), + } + } +} + /// Stream a response from the executor to the client. /// /// Requests arriving from the client while the stream is active are pushed @@ -214,9 +253,38 @@ async fn stream_ws_response( shutdown_token: &CancellationToken, queue: &mut VecDeque, ) -> Result<(), WsError> { + let mut prefer_shutdown_receiver = true; 'stream: loop { + if shutdown_token.is_cancelled() { + match next_shutdown_input(receiver, &mut stream, prefer_shutdown_receiver).await { + ShutdownInput::Receiver(message) => { + prefer_shutdown_receiver = false; + match message { + None | Some(Ok(Message::Close(_))) => return Err(WsError::ClientDisconnected), + Some(Ok(Message::Ping(payload))) => { + sender + .send(Message::Pong(payload)) + .await + .map_err(|_| WsError::SendFailed)?; + } + Some(Ok(Message::Text(_) | Message::Binary(_) | Message::Pong(_))) => {} + Some(Err(error)) => return Err(WsError::Receive(error.to_string())), + } + continue 'stream; + } + ShutdownInput::Upstream(line) => { + prefer_shutdown_receiver = true; + let Some(line) = line else { + break; + }; + forward_ws_stream_line(sender, &line).await?; + } + } + continue; + } + let next_line = tokio::select! { - () = shutdown_token.cancelled() => return Err(WsError::Shutdown), + () = shutdown_token.cancelled() => continue 'stream, message = receiver.next() => { match message { None | Some(Ok(Message::Close(_))) => return Err(WsError::ClientDisconnected), @@ -244,26 +312,30 @@ async fn stream_ws_response( let Some(line) = next_line else { break; }; - let Some(data) = line.strip_prefix("data: ") else { - continue; - }; - let data = data.trim(); - if data == "[DONE]" { - continue; - } - let value = match serde_json::from_str::(data) { - Ok(value) => value, - Err(e) => return Err(WsError::Executor(ExecutorError::from(e))), - }; - send_ws_json(sender, value).await?; + forward_ws_stream_line(sender, &line).await?; } Ok(()) } +async fn forward_ws_stream_line(sender: &mut WsSender, line: &str) -> Result<(), WsError> { + let Some(data) = line.strip_prefix("data: ") else { + return Ok(()); + }; + let data = data.trim(); + if data == "[DONE]" { + return Ok(()); + } + let value = match serde_json::from_str::(data) { + Ok(value) => value, + Err(e) => return Err(WsError::Executor(ExecutorError::from(e))), + }; + send_ws_json(sender, value).await +} + async fn handle_ws_error(sender: &mut WsSender, err: WsError) -> bool { match err { - WsError::Shutdown | WsError::ClientDisconnected | WsError::SendFailed => false, + WsError::ClientDisconnected | WsError::SendFailed => false, WsError::Receive(message) => { warn!("responses websocket receive error: {message}"); false @@ -286,3 +358,25 @@ async fn send_ws_json(sender: &mut WsSender, value: Value) -> Result<(), WsError .await .map_err(|_| WsError::SendFailed) } + +#[cfg(test)] +mod tests { + use futures::stream; + + use super::{ShutdownInput, next_shutdown_input}; + + #[tokio::test] + async fn shutdown_input_priority_alternates_when_both_streams_are_ready() { + let mut receiver = stream::repeat(()); + let mut upstream = stream::repeat(()); + + assert!(matches!( + next_shutdown_input(&mut receiver, &mut upstream, true).await, + ShutdownInput::Receiver(Some(())) + )); + assert!(matches!( + next_shutdown_input(&mut receiver, &mut upstream, false).await, + ShutdownInput::Upstream(Some(())) + )); + } +} diff --git a/crates/agentic-server/src/main.rs b/crates/agentic-server/src/main.rs index 827a05e3..890fa4d2 100644 --- a/crates/agentic-server/src/main.rs +++ b/crates/agentic-server/src/main.rs @@ -13,10 +13,10 @@ struct CommonArgs { #[arg(long, env = "OPENAI_API_KEY", hide_env_values = true, global = true)] openai_api_key: Option, - #[arg(long, default_value = "0.0.0.0", global = true)] + #[arg(long, env = "GATEWAY_HOST", default_value = "0.0.0.0", global = true)] gateway_host: String, - #[arg(long, default_value_t = 9000, global = true)] + #[arg(long, env = "GATEWAY_PORT", default_value_t = 9000, global = true)] gateway_port: u16, #[arg(long, default_value_t = 600.0, global = true)] @@ -46,7 +46,7 @@ struct Cli { #[command(subcommand)] command: Option, - #[arg(long)] + #[arg(long, env = "LLM_API_BASE")] llm_api_base: Option, #[command(flatten)] @@ -166,7 +166,7 @@ async fn main() -> Result<(), Error> { None => { let base = llm_api_base.ok_or_else(|| { Error::Config( - "standalone mode requires --llm-api-base; use `agentic-server serve ` for integrated mode" + "standalone mode requires LLM_API_BASE (or --llm-api-base); use `agentic-server serve ` for integrated mode" .to_owned(), ) })?; @@ -221,6 +221,25 @@ mod tests { assert!(cli.common.skip_llm_ready_check); } + #[test] + fn container_runtime_options_are_bound_to_environment_variables() { + let command = Cli::command(); + + for (argument, expected_env) in [ + ("llm_api_base", "LLM_API_BASE"), + ("gateway_host", "GATEWAY_HOST"), + ("gateway_port", "GATEWAY_PORT"), + ] { + let env = command + .get_arguments() + .find(|arg| arg.get_id() == argument) + .and_then(clap::Arg::get_env) + .unwrap_or_else(|| panic!("{argument} must be configurable through {expected_env}")); + + assert_eq!(env, expected_env); + } + } + #[test] fn sqlite_tuning_is_env_only_not_cli() { let mut help = Vec::new(); diff --git a/crates/agentic-server/src/server.rs b/crates/agentic-server/src/server.rs index 7240c79f..c9cd19f5 100644 --- a/crates/agentic-server/src/server.rs +++ b/crates/agentic-server/src/server.rs @@ -1,14 +1,19 @@ +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; +use std::time::Duration; use agentic_core::config::Config; use agentic_core::error::Error; use agentic_core::executor::ExecutionContext; use agentic_core::proxy::ProxyState; use agentic_core::readiness::wait_llm_ready; -use agentic_server::app::{AppState, ServerConfig, build_router}; +use agentic_server::app::{AppState, ServerConfig, WebSocketTracker, build_router}; use tokio::net::TcpListener; use tokio_util::sync::CancellationToken; -use tracing::info; +use tracing::{info, warn}; + +const GATEWAY_DRAIN_TIMEOUT: Duration = Duration::from_secs(8); async fn build_state(config: &Config, shutdown_token: CancellationToken) -> Result { let proxy_state = ProxyState::new(config.clone())?; @@ -18,6 +23,7 @@ async fn build_state(config: &Config, shutdown_token: CancellationToken) -> Resu proxy_state, exec_ctx, shutdown_token, + websocket_tracker: WebSocketTracker::default(), llm_api_base: config.llm_api_base.clone(), openai_api_key: config.openai_api_key.clone(), }) @@ -27,6 +33,7 @@ async fn serve_gateway(state: AppState, host: &str, port: u16) -> Result<(), Err let addr = format!("{host}:{port}"); let server_config = ServerConfig::from_env(); let shutdown_token = state.shutdown_token.clone(); + let websocket_tracker = state.websocket_tracker.clone(); let router = build_router(state, &server_config); let listener = TcpListener::bind(&addr).await?; info!("gateway listening on {addr}"); @@ -35,6 +42,7 @@ async fn serve_gateway(state: AppState, host: &str, port: u16) -> Result<(), Err shutdown_token.cancelled().await; }) .await?; + websocket_tracker.wait_until_idle().await; Ok(()) } @@ -45,15 +53,45 @@ async fn serve_gateway_until_signal(state: AppState, host: &str, port: u16) -> R tokio::select! { result = &mut gateway => result, - signal = tokio::signal::ctrl_c() => { + signal = shutdown_signal() => { signal?; info!("shutdown signal received"); shutdown_token.cancel(); - gateway.await + drain_gateway(gateway.as_mut()).await } } } +async fn drain_gateway(gateway: Pin<&mut F>) -> Result<(), Error> +where + F: Future>, +{ + if let Ok(result) = tokio::time::timeout(GATEWAY_DRAIN_TIMEOUT, gateway).await { + result + } else { + warn!( + timeout_seconds = GATEWAY_DRAIN_TIMEOUT.as_secs(), + "gateway drain timed out; closing remaining connections" + ); + Ok(()) + } +} + +#[cfg(unix)] +async fn shutdown_signal() -> Result<(), std::io::Error> { + let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; + + tokio::select! { + signal = tokio::signal::ctrl_c() => signal, + _ = terminate.recv() => Ok(()), + } +} + +#[cfg(not(unix))] +async fn shutdown_signal() -> Result<(), std::io::Error> { + tokio::signal::ctrl_c().await +} + async fn wait_until_llm_ready(config: &Config) -> Result<(), Error> { if config.skip_llm_ready_check { info!("skipping LLM readiness check: {}", config.llm_api_base); @@ -134,11 +172,15 @@ pub async fn run_with_llm(config: Config, host: &str, port: u16, llm_args: Vec { - signal?; - info!("shutdown signal received"); - shutdown_token.cancel(); - gateway.await + signal = shutdown_signal() => { + match signal { + Ok(()) => { + info!("shutdown signal received"); + shutdown_token.cancel(); + drain_gateway(gateway.as_mut()).await + } + Err(err) => Err(err.into()), + } } }; @@ -146,3 +188,27 @@ pub async fn run_with_llm(config: Config, host: &str, port: u16, llm_args: Vec>(); + tokio::pin!(gateway); + + drain_gateway(gateway.as_mut()).await.unwrap(); + } + + #[tokio::test] + async fn gateway_drain_preserves_server_errors() { + let gateway = std::future::ready(Err(Error::Config("gateway failed".to_owned()))); + tokio::pin!(gateway); + + let error = drain_gateway(gateway.as_mut()).await.unwrap_err(); + + assert_eq!(error.to_string(), "gateway failed"); + } +} diff --git a/crates/agentic-server/tests/cli_test.rs b/crates/agentic-server/tests/cli_test.rs new file mode 100644 index 00000000..9c85a1f0 --- /dev/null +++ b/crates/agentic-server/tests/cli_test.rs @@ -0,0 +1,16 @@ +use std::process::Command; + +#[test] +fn missing_llm_api_base_error_mentions_environment_and_flag() { + let output = Command::new(env!("CARGO_BIN_EXE_agentic-server")) + .env_remove("LLM_API_BASE") + .output() + .expect("agentic-server must run"); + + assert!(!output.status.success()); + let stderr = String::from_utf8(output.stderr).expect("stderr must be UTF-8"); + assert!( + stderr.contains("LLM_API_BASE (or --llm-api-base)"), + "unexpected error message: {stderr}" + ); +} diff --git a/crates/agentic-server/tests/common/mod.rs b/crates/agentic-server/tests/common/mod.rs index 929b8e02..983ae205 100644 --- a/crates/agentic-server/tests/common/mod.rs +++ b/crates/agentic-server/tests/common/mod.rs @@ -11,7 +11,7 @@ use agentic_core::config::Config; use agentic_core::executor::{ConversationHandler, ExecutionContext, ResponseHandler}; use agentic_core::proxy::ProxyState; use agentic_core::storage::{ConversationStore, ResponseStore}; -use agentic_server::app::{AppState, ServerConfig, build_router}; +use agentic_server::app::{AppState, ServerConfig, WebSocketTracker, build_router}; pub fn test_config(llm_url: &str) -> Config { Config { @@ -38,6 +38,7 @@ pub fn test_state(config: &Config) -> AppState { proxy_state, exec_ctx, shutdown_token: CancellationToken::new(), + websocket_tracker: WebSocketTracker::default(), llm_api_base: config.llm_api_base.clone(), openai_api_key: config.openai_api_key.clone(), } diff --git a/crates/agentic-server/tests/responses_websocket_test.rs b/crates/agentic-server/tests/responses_websocket_test.rs index 7781ca02..d8b80379 100644 --- a/crates/agentic-server/tests/responses_websocket_test.rs +++ b/crates/agentic-server/tests/responses_websocket_test.rs @@ -26,7 +26,7 @@ use agentic_core::executor::{ConversationHandler, ExecutionContext, ResponseHand use agentic_core::proxy::ProxyState; use agentic_core::storage::{ConversationStore, ResponseStore, create_pool_with_schema}; use agentic_core::tool::WebSearchHandler; -use agentic_server::app::AppState; +use agentic_server::app::{AppState, WebSocketTracker}; use common::{spawn_gateway, test_config}; @@ -93,6 +93,10 @@ impl Drop for MockYouSearchServer { enum MockResponse { Static(String), + Gated { + response: String, + release: oneshot::Receiver<()>, + }, Hanging { first_chunk: String, drop_tx: oneshot::Sender<()>, @@ -144,6 +148,16 @@ impl MockResponsesServer { (server, drop_rx) } + async fn start_gated(response: String) -> (Self, oneshot::Sender<()>) { + let (release_tx, release_rx) = oneshot::channel(); + let server = Self::start_with_responses(vec![MockResponse::Gated { + response, + release: release_rx, + }]) + .await; + (server, release_tx) + } + async fn start_with_responses(responses: Vec) -> Self { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -163,6 +177,10 @@ impl MockResponsesServer { let response = queue.lock().await.pop_front().expect("mock response queue exhausted"); let body = match response { MockResponse::Static(response) => axum::body::Body::from(response), + MockResponse::Gated { response, release } => { + let _ = release.await; + axum::body::Body::from(response) + } MockResponse::Hanging { first_chunk, drop_tx } => { axum::body::Body::from_stream(HangingSse::new(first_chunk, drop_tx)) } @@ -254,6 +272,7 @@ async fn storage_backed_state_with_web_search(llm_url: &str, web_search_base_url proxy_state, exec_ctx, shutdown_token: CancellationToken::new(), + websocket_tracker: WebSocketTracker::default(), llm_api_base: config.llm_api_base, openai_api_key: config.openai_api_key, }; @@ -1190,6 +1209,48 @@ async fn test_websocket_shutdown_token_closes_idle_connection() { assert!(mock.request_bodies().await.is_empty()); } +#[tokio::test] +async fn test_websocket_shutdown_drains_active_response_before_closing() { + let (mock, release) = + MockResponsesServer::start_gated(sse_response("resp_upstream_shutdown", "msg_upstream_shutdown", "DONE")).await; + let fixture = storage_backed_state(&mock.url).await; + let shutdown_token = fixture.state.shutdown_token.clone(); + let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await; + let mut ws = connect_responses_ws(&gateway_url).await; + + send_json( + &mut ws, + json!({ + "type": "response.create", + "model": "test-model", + "input": [{"type": "message", "role": "user", "content": "hi"}], + "store": true, + "stream": true + }), + ) + .await; + wait_for_request_count(&mock, 1).await; + + shutdown_token.cancel(); + send_json( + &mut ws, + json!({ + "type": "response.create", + "model": "test-model", + "input": "must not start during shutdown", + "store": true, + "stream": true + }), + ) + .await; + release.send(()).unwrap(); + + let events = recv_until_completed(&mut ws).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + recv_close_or_end(&mut ws).await; + assert_eq!(mock.request_bodies().await.len(), 1); +} + #[tokio::test] async fn test_websocket_client_close_cancels_hanging_upstream_stream() { let first_chunk = format!( diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh new file mode 100755 index 00000000..057b64e8 --- /dev/null +++ b/docker-entrypoint.sh @@ -0,0 +1,54 @@ +#!/bin/sh +set -eu + +# Keep runtime-created SQLite files writable when OpenShift rotates the +# arbitrary UID while retaining the image's root-group permission model. +umask 0002 + +# Mirror the server's default only for permission preparation. Keep DATABASE_URL +# unset so the Rust CLI remains authoritative for the actual connection URL. +database_url=${DATABASE_URL:-sqlite://./agentic_api.db} +case "$database_url" in + sqlite::memory:*) + ;; + sqlite://*) + case "$database_url" in + *%*) + echo "DATABASE_URL percent-encoded SQLite paths are not supported by this container entrypoint" >&2 + exit 64 + ;; + esac + sqlite_query= + case "$database_url" in + *\?*) + sqlite_query=${database_url#*\?} + sqlite_query=${sqlite_query%%\#*} + ;; + esac + prepare_sqlite=true + case "&$sqlite_query&" in + *"&mode=ro&"* | *"&mode=rw&"* | *"&mode=memory&"*) + prepare_sqlite=false + ;; + esac + + if [ "$prepare_sqlite" = true ]; then + # `?` and `#` are URI query/fragment delimiters, so they are not part + # of the filesystem path extracted here. Keep the parent directory + # group-writable so a rotated arbitrary UID can create SQLite sidecar files. + database_path=${database_url#sqlite://} + database_path=${database_path%%\?*} + database_path=${database_path%%\#*} + if [ -n "$database_path" ]; then + if [ ! -e "$database_path" ]; then + : >"$database_path" + chmod g+rw "$database_path" + fi + database_directory=$(dirname -- "$database_path") + chmod g+rwx "$database_directory" 2>/dev/null || true + fi + fi + ;; +esac + +exec agentic-server "$@" diff --git a/docs/deploying/container.md b/docs/deploying/container.md new file mode 100644 index 00000000..ecb736df --- /dev/null +++ b/docs/deploying/container.md @@ -0,0 +1,90 @@ +# Run agentic-api in a container + +The production image contains only the Rust gateway and its runtime libraries. It does not contain Python, vLLM, GPU libraries, model weights, Cargo, or the Rust toolchain. Run inference and PostgreSQL as external services. + +## Build the image + +The multi-stage build pins its Rust and Debian bases by digest, uses BuildKit caches, and copies only `agentic-server` into the runtime stage. Dependabot proposes weekly digest updates so base-image changes remain explicit and reviewable. + +```console +DOCKER_BUILDKIT=1 docker build \ + --build-arg OCI_CREATED="$(date -u +%Y-%m-%dT%H:%M:%SZ)" \ + --build-arg OCI_REVISION="$(git rev-parse HEAD)" \ + --build-arg OCI_VERSION="dev" \ + --tag agentic-api:dev \ + . +``` + +CI also records the workflow name and run URL in the vLLM-compatible image labels. Local builds use `local` as the pipeline label unless `OCI_BUILD_PIPELINE` is supplied as a build argument. + +Pass `--build-arg CARGO_BUILD_JOBS=` if the builder needs a different concurrency limit. For multi-platform publishing, use a Buildx builder with native workers where available; emulation is slower. + +## Configure the gateway + +The image starts `agentic-server` in standalone mode. At minimum, set `LLM_API_BASE` to an external OpenAI-compatible or vLLM endpoint. Production deployments should also set `DATABASE_URL` to an external PostgreSQL database. + +| Variable | Default | Purpose | +| --- | --- | --- | +| `LLM_API_BASE` | none | Required upstream inference URL | +| `GATEWAY_HOST` | `0.0.0.0` | Listen address | +| `GATEWAY_PORT` | `9000` | Listen port | +| `DATABASE_URL` | `sqlite://./agentic_api.db` | SQLite or PostgreSQL persistence URL | +| `OPENAI_API_KEY` | none | Credential sent to the upstream service when the client does not supply one | +| `SKIP_LLM_READY_CHECK` | `false` | Skip the startup probe for hosted providers without `/health` | +| `CORS_ALLOWED_ORIGINS` | none | Comma-separated browser origins | + +The container entrypoint rejects percent-encoded SQLite paths because SQLx decodes them before opening the database. Use a literal filesystem path or PostgreSQL instead. + +Do not put credentials into the image or Docker build arguments. Inject them at runtime through a secret manager. + +```console +docker run --rm --name agentic-api \ + --publish 127.0.0.1:9000:9000 \ + --env LLM_API_BASE=https://vllm.example.com \ + --env DATABASE_URL=postgresql://agentic-api@postgres.example.com/agentic_api \ + --env OPENAI_API_KEY \ + agentic-api:dev +``` + +The gateway does not provide inbound client authentication. `OPENAI_API_KEY` is an upstream credential, not a password for callers, so keep the port bound to loopback unless an authenticated ingress or proxy protects it. + +If the upstream is running on the Docker host, use `http://host.docker.internal:` on Docker Desktop. On Linux, add `--add-host host.docker.internal:host-gateway`. + +## Smoke test + +The liveness probe reports whether the gateway process is serving traffic. Readiness also checks the upstream inference service. + +```console +curl --fail http://127.0.0.1:9000/health +curl --fail http://127.0.0.1:9000/ready +``` + +The container CI workflow builds the image, verifies that build tools are absent, launches the gateway against a mock upstream, checks both probes, and exercises a stored Responses API request through SQLite persistence. HTTP streaming and WebSockets use the same gateway binary and exposed port; the image does not add a transport proxy. + +On `SIGTERM`, the gateway stops accepting connections and gives in-flight requests up to eight seconds to drain before closing the remaining connections. Set an orchestrator termination grace period longer than eight seconds; the default 30-second Kubernetes grace period and the documented 10-second Docker stop timeout both satisfy this requirement. + +## Kubernetes and OpenShift security context + +The image defaults to UID `10001` and GID `0`. Its working directory is setgid and the entrypoint uses a group-cooperative umask, so new SQLite files remain writable when OpenShift replaces the UID while retaining the group-0 permission model. Do not set a fixed `runAsUser` when the cluster assigns arbitrary UIDs. + +A volume mounted at `/var/lib/agentic-api` hides the ownership and mode stored in the image. For SQLite, configure the storage class or pod-level `fsGroup` so the mounted directory is writable by a supplemental group assigned to the container. The example below uses group 0 to match the image; if the cluster assigns a different permitted supplemental group, use that group and ensure the volume root is group-writable and setgid. PostgreSQL deployments do not need this pod-level filesystem setting. + +Volumes initialized by an older image may contain SQLite files without group-write permission. Before rotating to an arbitrary UID, repair those volumes once as an administrator with `chmod -R g+rwX /var/lib/agentic-api`. + +```yaml +spec: + securityContext: + fsGroup: 0 + fsGroupChangePolicy: OnRootMismatch + containers: + - name: agentic-api + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: ["ALL"] + runAsNonRoot: true + seccompProfile: + type: RuntimeDefault +``` + +Mount writable storage at `/var/lib/agentic-api` only when using SQLite. PostgreSQL deployments do not need a persistent filesystem for the gateway. diff --git a/docs/design/mcp-gateway-integration.md b/docs/design/mcp-gateway-integration.md index 0b92b967..d719d349 100644 --- a/docs/design/mcp-gateway-integration.md +++ b/docs/design/mcp-gateway-integration.md @@ -21,12 +21,13 @@ hostnames can be allowed with the comma-separated `AGENTIC_MCP_ALLOWED_HOSTS` environment variable. Treat every hostname added to `AGENTIC_MCP_ALLOWED_HOSTS` as fully trusted. The -allowlist validates the hostname text, while the HTTP transport performs DNS -resolution later and does not pin the resolved IP address. Consequently, this -setting is not an IP-level SSRF boundary and is unsafe for hostnames whose DNS -records could be changed by an untrusted party (including through DNS -rebinding). Only add hostnames whose DNS configuration is controlled by a -trusted administrator. +allowlist validates the hostname text. When it opens a connection, the HTTP +client resolves the hostname once, pins all returned addresses for that +transport, disables automatic proxy discovery, and disables redirects. This +prevents later DNS changes, proxy-side resolution, or redirect targets from +bypassing the validation, but it does not make an untrusted DNS record safe: a +poisoned initial resolution is still trusted. Only add hostnames whose DNS +configuration is controlled by a trusted administrator. ### First built-in tool: `read_mcp_resource` diff --git a/mkdocs.yaml b/mkdocs.yaml index 527c25e5..f50ee255 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -94,6 +94,8 @@ nav: - API Reference: api/index.md - Developing: - Getting Started: developing/getting-started.md + - Deploying: + - Container: deploying/container.md - Community: community/index.md extra: diff --git a/scripts/container-fixture-server.py b/scripts/container-fixture-server.py new file mode 100755 index 00000000..d58be6db --- /dev/null +++ b/scripts/container-fixture-server.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Serve a recorded vLLM Responses API fixture for container integration tests. + +The JSON response mirrors the existing agentic-server-core YAML cassette while +keeping the CI helper dependency-free. +""" + +import argparse +import json +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--fixture", required=True, type=Path) + parser.add_argument("--port", required=True, type=int) + return parser.parse_args() + + +def make_handler(response_body: bytes): + class FixtureHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + if self.path == "/health": + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + return + self.send_error(404) + + def do_POST(self) -> None: + if self.path != "/v1/responses": + self.send_error(404) + return + + try: + content_length = int(self.headers.get("Content-Length", "0")) + request = json.loads(self.rfile.read(content_length)) + except (ValueError, json.JSONDecodeError): + self.send_error(400, "request body must be valid JSON") + return + + input_items = request.get("input") + expected_prompt = { + "content": "Reply with exactly one word: HELLO", + "role": "user", + "type": "message", + } + if ( + not isinstance(input_items, list) + or not input_items + or input_items[-1] != expected_prompt + or request.get("model") != "gpt-4o" + or request.get("stream") is not False + ): + self.log_error("fixture request mismatch: %s", json.dumps(request, sort_keys=True)) + self.send_error(422, "request did not match the recorded fixture") + return + + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response_body))) + self.end_headers() + self.wfile.write(response_body) + + return FixtureHandler + + +def main() -> None: + args = parse_args() + response_body = json.dumps(json.loads(args.fixture.read_text())).encode() + server = ThreadingHTTPServer(("127.0.0.1", args.port), make_handler(response_body)) + server.serve_forever() + + +if __name__ == "__main__": + main() diff --git a/scripts/fixtures/vllm-response-single-nonstreaming.json b/scripts/fixtures/vllm-response-single-nonstreaming.json new file mode 100644 index 00000000..b2c3bf0f --- /dev/null +++ b/scripts/fixtures/vllm-response-single-nonstreaming.json @@ -0,0 +1,71 @@ +{ + "background": false, + "billing": { + "payer": "developer" + }, + "completed_at": 1776764143, + "created_at": 1776764142, + "error": null, + "frequency_penalty": 0.0, + "id": "resp_0508721937e20de90069e744ee9018819394c8e011dc6d7818", + "incomplete_details": null, + "instructions": null, + "max_output_tokens": null, + "max_tool_calls": null, + "metadata": {}, + "model": "gpt-4o-2024-08-06", + "object": "response", + "output": [ + { + "content": [ + { + "annotations": [], + "logprobs": [], + "text": "HI", + "type": "output_text" + } + ], + "id": "msg_0508721937e20de90069e744ef1b2881939596753c5951691e", + "role": "assistant", + "status": "completed", + "type": "message" + } + ], + "parallel_tool_calls": true, + "presence_penalty": 0.0, + "previous_response_id": null, + "prompt_cache_key": null, + "prompt_cache_retention": "in_memory", + "reasoning": { + "effort": null, + "summary": null + }, + "safety_identifier": null, + "service_tier": "default", + "status": "completed", + "store": true, + "temperature": 1.0, + "text": { + "format": { + "type": "text" + }, + "verbosity": "medium" + }, + "tool_choice": "auto", + "tools": [], + "top_logprobs": 0, + "top_p": 1.0, + "truncation": "disabled", + "usage": { + "input_tokens": 15, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens": 2, + "output_tokens_details": { + "reasoning_tokens": 0 + }, + "total_tokens": 17 + }, + "user": null +} diff --git a/scripts/test-container-entrypoint.sh b/scripts/test-container-entrypoint.sh new file mode 100755 index 00000000..c35379f9 --- /dev/null +++ b/scripts/test-container-entrypoint.sh @@ -0,0 +1,78 @@ +#!/bin/sh +set -eu + +repo_root=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd) +entrypoint=$repo_root/docker-entrypoint.sh +test_root=$(mktemp -d) +cleanup() { + status=$? + rm -rf "$test_root" + exit "$status" +} +trap cleanup EXIT + +mkdir -p "$test_root/bin" +printf '%s\n' \ + '#!/bin/sh' \ + 'if [ "${EXPECT_DATABASE_URL_UNSET:-}" = 1 ] && [ "${DATABASE_URL+x}" = x ]; then exit 1; fi' \ + 'exit 0' >"$test_root/bin/agentic-server" +chmod +x "$test_root/bin/agentic-server" +test_path=$test_root/bin:/usr/bin:/bin + +unset_dir=$test_root/unset +mkdir "$unset_dir" +( + cd "$unset_dir" + env -u DATABASE_URL EXPECT_DATABASE_URL_UNSET=1 PATH="$test_path" "$entrypoint" +) +test -f "$unset_dir/agentic_api.db" +test "$(stat -c '%a' "$unset_dir/agentic_api.db" 2>/dev/null || stat -f '%Lp' "$unset_dir/agentic_api.db")" = "664" + +sqlite_dir=$test_root/sqlite +mkdir "$sqlite_dir" +DATABASE_URL="sqlite://$sqlite_dir/state.db?mode=rwc&cache=shared#fragment" PATH="$test_path" "$entrypoint" +test -f "$sqlite_dir/state.db" +test -w "$sqlite_dir/state.db" + +printf '%s' preserved >"$sqlite_dir/existing.db" +DATABASE_URL="sqlite://$sqlite_dir/existing.db" PATH="$test_path" "$entrypoint" +test "$(cat "$sqlite_dir/existing.db")" = preserved + +for mode in ro rw memory; do + case_dir=$test_root/untouched-$mode + mkdir "$case_dir" + ( + cd "$case_dir" + DATABASE_URL="sqlite://$case_dir/state.db?mode=$mode" PATH="$test_path" "$entrypoint" + ) + test -z "$(find "$case_dir" -mindepth 1 -print -quit)" +done + +for database_url in 'sqlite::memory:' 'postgresql://agentic-api@postgres.example.com/agentic_api'; do + case_dir=$test_root/untouched-$(printf '%s' "$database_url" | tr -cd '[:alnum:]') + mkdir "$case_dir" + ( + cd "$case_dir" + DATABASE_URL=$database_url PATH="$test_path" "$entrypoint" + ) + test -z "$(find "$case_dir" -mindepth 1 -print -quit)" +done + +bare_dir=$test_root/bare +mkdir "$bare_dir" +chmod 700 "$bare_dir" +( + cd "$bare_dir" + DATABASE_URL=sqlite:// PATH="$test_path" "$entrypoint" +) +test -z "$(find "$bare_dir" -mindepth 1 -print -quit)" +test "$(ls -ld "$bare_dir" | cut -c5-7)" = "---" + +encoded_dir=$test_root/encoded +mkdir "$encoded_dir" +if DATABASE_URL="sqlite://$encoded_dir/state%23prod.db" PATH="$test_path" "$entrypoint" 2>"$encoded_dir/error.log"; then + echo "percent-encoded SQLite path unexpectedly succeeded" >&2 + exit 1 +fi +grep -q "percent-encoded SQLite paths are not supported" "$encoded_dir/error.log" +test ! -e "$encoded_dir/state%23prod.db"