diff --git a/Cargo.lock b/Cargo.lock index 8d1f205..2e0ded3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -41,6 +41,15 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "arc-swap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +dependencies = [ + "rustversion", +] + [[package]] name = "assert-json-diff" version = "2.0.2" @@ -150,6 +159,15 @@ dependencies = [ "tower-service", ] +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", +] + [[package]] name = "backtrace" version = "0.3.76" @@ -267,7 +285,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -421,6 +443,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + [[package]] name = "encoding_rs" version = "0.8.35" @@ -804,7 +832,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2", + "socket2 0.6.3", "system-configuration", "tokio", "tower-service", @@ -973,6 +1001,15 @@ dependencies = [ "serde", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1179,12 +1216,31 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967" +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1481,10 +1537,12 @@ dependencies = [ "futures-util", "http-body-util", "jsonwebtoken", + "redis", "serde", "serde_json", "thiserror 2.0.18", "tokio", + "tokio-stream", "tokio-tungstenite 0.29.0", "tower", "tracing", @@ -1611,7 +1669,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2", + "socket2 0.6.3", "thiserror 2.0.18", "tokio", "tracing", @@ -1649,7 +1707,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.6.3", "tracing", "windows-sys 0.60.2", ] @@ -1713,6 +1771,32 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "redis" +version = "0.27.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09d8f99a4090c89cc489a94833c901ead69bfbf3877b4867d5482e321ee875bc" +dependencies = [ + "arc-swap", + "async-trait", + "backon", + "bytes", + "combine", + "futures", + "futures-util", + "itertools", + "itoa", + "num-bigint", + "percent-encoding", + "pin-project-lite", + "ryu", + "sha1_smol", + "socket2 0.5.10", + "tokio", + "tokio-util", + "url", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2198,6 +2282,12 @@ dependencies = [ "digest", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sharded-slab" version = "0.1.7" @@ -2244,6 +2334,16 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.3" @@ -2448,7 +2548,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.6.3", "tokio-macros", "windows-sys 0.61.2", ] diff --git a/SPEC.md b/SPEC.md index 746326e..db08fd8 100644 --- a/SPEC.md +++ b/SPEC.md @@ -482,16 +482,20 @@ frontmatter and drop the SPEC link in the same change. ## Open questions -- **Reconnect & redelivery semantics.** The current design holds a - bounded in-memory `MessageQueue` per agent for pushes that arrive - while the tunnel is disconnected. Lost on broker restart; operator - re-POSTs are the recovery path. A persistent queue (Redis stream, - NATS jetstream) is a future option, decided once a real - multi-replica deployment lands. -- **Multi-replica routing.** Single-replica works out of the box. - Multi-replica needs either sticky LB routing on `agent_id` or a - shared registry. Either is feasible behind the existing - `AgentRegistry` interface; the call gets made when needed. +- **Reconnect & redelivery semantics.** The in-memory + `MessageQueue` is the default for single-instance deployments + (lossy across restarts). The optional Redis backend + (`features = ["redis"]`, env `OVERACP_REDIS_URL`) provides + durable per-agent inbox streams (`XADD` / `XREADGROUP`) with + XAUTOCLAIM orphan recovery and a dead-letter queue. See + [`docs/design/ha.md`](./docs/design/ha.md). +- **Multi-replica routing.** Resolved. The Redis backend uses + ownership leases (`SET EX` + Lua CAS heartbeat) keyed on the + tunnel holder. `deliver()` writes to the per-agent inbox stream + unconditionally; the owner's consumer drains it. SSE fan-out + uses Redis pub/sub. No sticky LB required — consistent-hash on + `agent_id` is an optional optimization. See + [`docs/design/ha.md`](./docs/design/ha.md). - **`turn/end` vs webhook.** Currently a notification fanned out over SSE for the operator to consume. Some operators may prefer a push webhook (HTTP POST from the broker to a configured URL on diff --git a/TODO.md b/TODO.md index ca9536c..f508c95 100644 --- a/TODO.md +++ b/TODO.md @@ -153,9 +153,12 @@ broker itself. - [ ] Agent JWT rotation strategy. Currently mints once with a 30-day TTL per `docs/design/protocol.md` § 2.4 and never refreshes. -- [ ] Persistent `MessageQueue` for multi-replica deployments - (Redis stream, NATS jetstream). Current in-memory queue - survives reconnects but not broker restarts. +- [x] Persistent `MessageQueue` for multi-replica deployments. + Landed as the `redis` feature: `RedisMessageQueue` (inbox + streams + buffer streams), `RedisAgentRegistry` (ownership + leases + XREADGROUP consumer), `RedisStreamBroker` (pub/sub + SSE fan-out). See [`docs/design/ha.md`](./docs/design/ha.md) + and `tests/smoke-ha.sh`. - [ ] Streaming completions (multiple `stream/textDelta` per turn, proper turn terminator framing). diff --git a/codecov.yml b/codecov.yml index af1a077..e35b60e 100644 --- a/codecov.yml +++ b/codecov.yml @@ -10,6 +10,11 @@ coverage: target: 70% if_ci_failed: error +ignore: + # Redis backend requires a live Redis/Valkey instance for + # integration tests — excluded from patch coverage checks. + - "server/src/redis_backend/**" + comment: layout: "reach, diff, flags, files" require_changes: false diff --git a/docs/design/INDEX.md b/docs/design/INDEX.md index c726d82..1443420 100644 --- a/docs/design/INDEX.md +++ b/docs/design/INDEX.md @@ -22,3 +22,7 @@ For the high-level architecture and milestone roadmap see hydration and persistence: where it lives in the architecture (agent supervisor, not the broker), the `WorkspaceSync` trait, the per-backend crate convention, and the configuration model. +- [`ha.md`](./ha.md) — multi-instance HA via Redis/Valkey: ownership + leases, inbox streams, pub/sub SSE fan-out, and the `redis` feature + gate. Optional backend behind `AgentRegistryProvider`, + `MessageQueueProvider`, and `StreamBrokerProvider` traits. diff --git a/docs/design/ha.md b/docs/design/ha.md new file mode 100644 index 0000000..2a72b06 --- /dev/null +++ b/docs/design/ha.md @@ -0,0 +1,106 @@ +--- +status: Active +--- + +# Multi-Instance HA — Redis/Valkey Backend + +Design doc for the optional Redis/Valkey-backed implementations of +`AgentRegistryProvider`, `MessageQueueProvider`, and +`StreamBrokerProvider` that enable multi-instance overacp-server +deployments behind a plain (non-sticky) load balancer. + +## Motivation + +The in-memory defaults (`InMemoryAgentRegistry`, `InMemoryMessageQueue`, +`InMemoryStreamBroker`) keep all routing state in-process. A second +server instance behind the same LB would not know about agents +connected to the first. REST pushes that land on the wrong instance +silently fail. This module solves that by externalizing routing state +to Redis/Valkey. + +## Architecture + +### Design principles + +1. **Symmetric instances.** No leader election. Any instance serves + any REST request; whichever holds the tunnel is the owner. +2. **Tunnel-holder is owner.** The ownership lock follows the + WebSocket tunnel, not a pre-computed hash. +3. **Produce unconditionally.** `deliver()` always writes to the + per-agent inbox stream. The owner's XREADGROUP consumer drains it. +4. **Pub/sub for SSE fan-out.** Broadcast semantics; no durability. +5. **Natural pinning.** The agent_id → instance_id lease is the pin. + Consistent-hash LB is an optional optimization. + +### Keyspace + +| Key | Type | TTL | Purpose | +|-----|------|-----|---------| +| `overacp:owner:{agent_id}` | STRING | 30s | Ownership lease. Value = instance_id. | +| `overacp:tunnel:{agent_id}` | HASH | 60s | Tunnel directory (instance_id, connected_at). | +| `overacp:agents:connected` | SET | — | Live agent_id set. | +| `overacp:agents:recent` | ZSET | — | Recently-disconnected log (score = ts_ms, capped at 64). | +| `overacp:agents:claims:{agent_id}` | STRING | 60s | Serialized claims snapshot. | +| `overacp:inbox:{agent_id}` | STREAM | — | Per-agent delivery stream. Consumer group `owners`. | +| `overacp:inbox:dlq` | STREAM | — | Dead-letter queue. | +| `overacp:buffer:{agent_id}` | STREAM | — | Offline message buffer (drained on reconnect). | +| `overacp:stream:{agent_id}` | PUBSUB | — | SSE fan-out channel. | +| `overacp:control:{agent_id}` | PUBSUB | — | Control signals (takeover, disconnect). | + +### Ownership lease lifecycle + +1. **Acquire** (`SET EX 30`, force — no NX). Publishes `takeover` + on `overacp:control:{agent_id}` if a different instance holds it. +2. **Heartbeat** (10s interval). Lua CAS `EXPIRE` on the owner key. + Also refreshes tunnel directory and claims TTL. +3. **Release** (Drop / explicit). Lua CAS `DEL`. Updates connected + set and recently-disconnected ZSET. + +Lua scripts ported from +`overfolder/agent-runner/src/session_lock.rs`. + +### Message delivery + +`deliver(agent_id, frame)`: +1. `EXISTS overacp:owner:{agent_id}` +2. Owner present → `XADD overacp:inbox:{agent_id} * frame {frame}` → `Live` +3. Owner absent → `NoTunnel(frame)` — caller buffers via `MessageQueueProvider` + +Owner's inbox consumer: `XREADGROUP GROUP owners {consumer_id} COUNT 1 STREAMS overacp:inbox:{agent_id} >`. +ACK on successful local delivery. XAUTOCLAIM every ~5 min for orphan +recovery. DLQ on `MAX_DELIVERIES=3` or age > 4h. + +### SSE fan-out + +`publish(agent_id, frame)` → `PUBLISH overacp:stream:{agent_id} {frame}`. +`subscribe(agent_id)` → dedicated pub/sub connection, returns `BoxStream`. + +## Configuration + +| Env var | Required | Default | Notes | +|---------|----------|---------|-------| +| `OVERACP_REDIS_URL` | No | — | Enables Redis backend when set. | +| `OVERACP_INSTANCE_ID` | No | hostname / random | Instance identifier for leases. | +| `PORT` | No | `8080` | HTTP listen port. | + +## Feature gate + +The Redis backend is behind `features = ["redis"]` in +`server/Cargo.toml`. The `redis` crate is not compiled without +the feature; all in-memory tests run without it. + +## Relationship to overfolder-on-overacp + +Per `overfolder/docs/design/overfolder-on-overacp.md` § "Storage, +Valkey, and Locks", the lock primitives are designed to be +extractable to a shared `valkey-session` crate. For now they live +inside `overacp-server` behind the `redis` feature gate. Overfolder's +controlplane embeds overacp-server as a library and gets HA for free. + +## Verification + +- `tests/smoke-ha.sh` — two-instance end-to-end test against + dockerized Valkey. Covers cross-instance push, cross-instance SSE, + and failover after instance kill. +- Unit tests of in-memory implementations remain the primary safety + net (150 tests, all pass without the `redis` feature). diff --git a/server/Cargo.toml b/server/Cargo.toml index 9fb5acf..63cd60b 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -26,9 +26,16 @@ tracing-subscriber = { workspace = true, features = ["env-filter"] } futures = "0.3" futures-util = { workspace = true } dashmap = { workspace = true } +tokio-stream = { version = "0.1", features = ["sync"] } jsonwebtoken = { workspace = true } dotenvy = { workspace = true } +# Optional Redis/Valkey backend for multi-instance HA. +redis = { version = "0.27", features = ["tokio-comp", "connection-manager", "streams"], optional = true } + +[features] +redis = ["dep:redis"] + [lints] workspace = true diff --git a/server/src/api/agents.rs b/server/src/api/agents.rs index 96666a8..f891cd4 100644 --- a/server/src/api/agents.rs +++ b/server/src/api/agents.rs @@ -33,14 +33,14 @@ use axum::http::StatusCode; use axum::response::sse::{Event, KeepAlive, Sse}; use axum::routing::{get, post}; use axum::{Json, Router}; -use futures::stream::{self, Stream}; +use futures::stream::Stream; +use futures::StreamExt; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use tokio::sync::broadcast::error::RecvError; use uuid::Uuid; use crate::api::error::ApiError; -use crate::registry::{AgentDescription, QueueError}; +use crate::registry::{AgentDescription, DeliveryOutcome, QueueError}; use crate::state::AppState; /// Admin-only routes: list/describe/force-disconnect agents. Per @@ -107,7 +107,7 @@ pub struct AgentListResponse { /// wired up in [`crate::routes`]). async fn list_agents(State(state): State) -> Result, ApiError> { Ok(Json(AgentListResponse { - agents: state.registry.list(), + agents: state.registry.list_agents().await, })) } @@ -118,7 +118,8 @@ async fn describe_agent( ) -> Result, ApiError> { state .registry - .describe(id) + .describe_agent(id) + .await .map(Json) .ok_or_else(|| ApiError::NotFound(format!("agent {id}"))) } @@ -131,10 +132,10 @@ async fn disconnect_agent( State(state): State, Path(id): Path, ) -> Result { - if !state.registry.is_connected(id) { + if !state.registry.is_connected(id).await { return Err(ApiError::NotFound(format!("agent {id}"))); } - state.registry.unregister(id); + state.registry.disconnect(id).await; Ok(StatusCode::ACCEPTED) } @@ -165,36 +166,14 @@ async fn send_message( }) .to_string(); - let delivery = if let Some(entry) = state.registry.get(id) { - // Live tunnel — deliver inline. If the channel has been - // closed between our lookup and the send (a disconnect - // race), `SendError` hands us the original frame back, so - // we fall through to the queue without rebuilding it. We - // deliberately do NOT call `registry.unregister(id)` here: - // `run_tunnel`'s own cleanup owns that, and reaching in - // from the REST handler risks racing a fresh reconnect - // that slotted in at the same `agent_id` between our - // failed send and the cleanup call. - match entry.tx.send(notif) { - Ok(()) => Delivery::Live, - Err(send_err) => { - let recovered = send_err.0; - match state.message_queue.push(id, recovered) { - Ok(()) => Delivery::Queued, - Err(QueueError::Full { capacity, .. }) => { - return Err(queue_full(id, capacity)); - } - } - } - } - } else { - // No live tunnel — buffer for the next reconnect. - match state.message_queue.push(id, notif) { + let delivery = match state.registry.deliver(id, notif).await { + DeliveryOutcome::Live => Delivery::Live, + DeliveryOutcome::NoTunnel(frame) => match state.message_queue.push(id, frame).await { Ok(()) => Delivery::Queued, Err(QueueError::Full { capacity, .. }) => { return Err(queue_full(id, capacity)); } - } + }, }; Ok((StatusCode::ACCEPTED, Json(SendMessageResponse { delivery }))) @@ -215,20 +194,9 @@ async fn stream_events( State(state): State, Path(id): Path, ) -> Result>>, ApiError> { - let rx = state.stream_broker.subscribe(id); - let stream = stream::unfold(rx, |mut rx| async move { - loop { - match rx.recv().await { - Ok(text) => { - return Some((Ok::(Event::default().data(text)), rx)); - } - // Slow consumer — skip the missed frames and keep going. - Err(RecvError::Lagged(_)) => continue, - Err(RecvError::Closed) => return None, - } - } - }); - Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))) + let stream = state.stream_broker.subscribe(id); + let sse_stream = stream.map(|text| Ok::(Event::default().data(text))); + Ok(Sse::new(sse_stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))) } /// `POST /agents/{id}/cancel` — inject a cancel notification down @@ -238,14 +206,13 @@ async fn cancel_turn( State(state): State, Path(id): Path, ) -> Result { - if let Some(entry) = state.registry.get(id) { - let notif = json!({ - "jsonrpc": "2.0", - "method": "session/cancel", - "params": {}, - }); - let _ = entry.tx.send(notif.to_string()); - } + let notif = json!({ + "jsonrpc": "2.0", + "method": "session/cancel", + "params": {}, + }) + .to_string(); + let _ = state.registry.deliver(id, notif).await; Ok(StatusCode::ACCEPTED) } @@ -255,7 +222,7 @@ mod tests { use super::*; use crate::auth::{Claims, StaticJwtAuthenticator}; - use crate::registry::{AgentEntry, MessageQueue}; + use crate::registry::{MessageQueue, TunnelLease}; use crate::state::AppState; use tokio::sync::mpsc; @@ -266,20 +233,21 @@ mod tests { fn state_with_queue_capacity(cap: usize) -> AppState { let base = fresh_state(); AppState { - message_queue: MessageQueue::new(cap), + message_queue: Arc::new(MessageQueue::new(cap)), ..base } } /// Register a fake agent in the new registry and return the - /// receiving side of its tunnel channel. - fn register_fake(state: &AppState, agent_id: Uuid) -> mpsc::UnboundedReceiver { + /// receiving side of its tunnel channel plus the lease guard. + async fn register_fake( + state: &AppState, + agent_id: Uuid, + ) -> (mpsc::UnboundedReceiver, TunnelLease) { let (tx, rx) = mpsc::unbounded_channel::(); let claims = Claims::agent(agent_id, Some(Uuid::new_v4()), 60, "test"); - state - .registry - .register(agent_id, AgentEntry::new(tx, claims)); - rx + let lease = state.registry.acquire(agent_id, tx, claims).await.unwrap(); + (rx, lease) } fn parse(s: &str) -> Value { @@ -292,7 +260,7 @@ mod tests { async fn send_message_delivers_inline_when_connected() { let state = fresh_state(); let agent_id = Uuid::new_v4(); - let mut rx = register_fake(&state, agent_id); + let (mut rx, _lease) = register_fake(&state, agent_id).await; let (status, Json(resp)) = send_message( State(state), @@ -333,9 +301,9 @@ mod tests { assert_eq!(status, StatusCode::ACCEPTED); assert_eq!(resp.delivery, Delivery::Queued); - assert_eq!(state.message_queue.len(agent_id), 1); + assert_eq!(state.message_queue.len(agent_id).await, 1); - let drained = state.message_queue.drain(agent_id); + let drained = state.message_queue.drain(agent_id).await; assert_eq!(parse(&drained[0])["params"]["content"], "pending"); } @@ -361,11 +329,12 @@ mod tests { state .message_queue .push(agent_id, "already-buffered".to_string()) + .await .unwrap(); // Register a fake entry and immediately drop its receiver - // so `tx.send` fails. - let rx = register_fake(&state, agent_id); + // so `deliver()` returns `NoTunnel`. + let (rx, _lease) = register_fake(&state, agent_id).await; drop(rx); let err = send_message( @@ -381,7 +350,7 @@ mod tests { assert!(matches!(err, ApiError::ServiceUnavailable(_))); // The queue still holds exactly one frame (the pre-seed). - assert_eq!(state.message_queue.len(agent_id), 1); + assert_eq!(state.message_queue.len(agent_id).await, 1); } #[tokio::test] @@ -397,7 +366,7 @@ mod tests { // same agent_id. let state = fresh_state(); let agent_id = Uuid::new_v4(); - let rx = register_fake(&state, agent_id); + let (rx, _lease) = register_fake(&state, agent_id).await; drop(rx); // close the receiver without unregistering let (status, Json(resp)) = send_message( @@ -415,11 +384,11 @@ mod tests { assert_eq!(resp.delivery, Delivery::Queued); // The frame landed in the queue (not silently dropped). - assert_eq!(state.message_queue.len(agent_id), 1); + assert_eq!(state.message_queue.len(agent_id).await, 1); // The registry entry is still present — the handler did // not reach in and unregister it. - assert!(state.registry.is_connected(agent_id)); + assert!(state.registry.is_connected(agent_id).await); } #[tokio::test] @@ -470,7 +439,7 @@ mod tests { // And we still have exactly one frame buffered — no ghost // state from the rejected push. - assert_eq!(state.message_queue.len(agent_id), 1); + assert_eq!(state.message_queue.len(agent_id).await, 1); } // ── GET /agents, GET /agents/{id}, DELETE /agents/{id} ── @@ -487,9 +456,9 @@ mod tests { let state = fresh_state(); let live = Uuid::new_v4(); let dead = Uuid::new_v4(); - let _live_rx = register_fake(&state, live); - let _dead_rx = register_fake(&state, dead); - state.registry.unregister(dead); + let (_live_rx, _live_lease) = register_fake(&state, live).await; + let (_dead_rx, _dead_lease) = register_fake(&state, dead).await; + drop(_dead_lease); // triggers disconnect via TunnelLease drop let Json(resp) = list_agents(State(state)).await.unwrap(); assert_eq!(resp.agents.len(), 2); @@ -507,7 +476,7 @@ mod tests { async fn describe_agent_returns_connected_shape() { let state = fresh_state(); let id = Uuid::new_v4(); - let _rx = register_fake(&state, id); + let (_rx, _lease) = register_fake(&state, id).await; let Json(desc) = describe_agent(State(state), Path(id)).await.unwrap(); assert_eq!(desc.agent_id, id); @@ -528,13 +497,13 @@ mod tests { async fn disconnect_unregisters_and_returns_202() { let state = fresh_state(); let id = Uuid::new_v4(); - let _rx = register_fake(&state, id); + let (_rx, _lease) = register_fake(&state, id).await; let status = disconnect_agent(State(state.clone()), Path(id)) .await .unwrap(); assert_eq!(status, StatusCode::ACCEPTED); - assert!(!state.registry.is_connected(id)); + assert!(!state.registry.is_connected(id).await); } #[tokio::test] @@ -552,7 +521,7 @@ mod tests { async fn cancel_emits_notification_when_connected() { let state = fresh_state(); let id = Uuid::new_v4(); - let mut rx = register_fake(&state, id); + let (mut rx, _lease) = register_fake(&state, id).await; let status = cancel_turn(State(state), Path(id)).await.unwrap(); assert_eq!(status, StatusCode::ACCEPTED); @@ -585,27 +554,25 @@ mod tests { #[tokio::test] async fn stream_handler_forwards_broadcast_frames() { - // Subscribe via the handler and drive a frame through the - // underlying broadcast channel. The `Stream::unfold` body - // inside `stream_events` pumps the frame out as an SSE - // event. use axum::response::IntoResponse as _; use http_body_util::BodyExt as _; let state = fresh_state(); let agent_id = Uuid::new_v4(); - // Warm up the broker channel before the handler subscribes, - // so the stream_events call finds a sender already. - let sender = state.stream_broker.sender_for(agent_id); - - let sse = stream_events(State(state), Path(agent_id)) + // Subscribe via the handler first (internally creates the channel). + let sse = stream_events(State(state.clone()), Path(agent_id)) .await .expect("sse handle"); - sender - .send(r#"{"jsonrpc":"2.0","method":"stream/textDelta","params":{}}"#.to_string()) - .expect("send to broker"); + // Publish a frame through the broker. + state + .stream_broker + .publish( + agent_id, + r#"{"jsonrpc":"2.0","method":"stream/textDelta","params":{}}"#.to_string(), + ) + .await; let resp = sse.into_response(); let mut body = resp.into_body(); diff --git a/server/src/lib.rs b/server/src/lib.rs index 2b1fa01..257ae67 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -5,6 +5,8 @@ pub use startup::{build_state_from_env, StartupError}; pub mod api; pub mod auth; pub mod hooks; +#[cfg(feature = "redis")] +pub mod redis_backend; pub mod registry; pub mod routes; pub mod state; @@ -15,7 +17,10 @@ pub use hooks::{ BootError, BootProvider, DefaultBootProvider, DefaultQuotaPolicy, DefaultToolHost, QuotaError, QuotaPolicy, ToolError, ToolHost, }; -pub use registry::{AgentDescription, AgentEntry, AgentRegistry, MessageQueue, QueueError}; +pub use registry::{ + AgentDescription, AgentEntry, AgentRegistry, AgentRegistryProvider, DeliveryOutcome, + MessageQueue, MessageQueueProvider, QueueError, RegistryError, TunnelLease, +}; pub use routes::router; pub use state::AppState; -pub use tunnel::StreamBroker; +pub use tunnel::{InMemoryStreamBroker, StreamBrokerProvider}; diff --git a/server/src/main.rs b/server/src/main.rs index ffe7799..caf2143 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -1,3 +1,4 @@ +use std::env; use std::error::Error as StdError; use overacp_server::build_state_from_env; @@ -15,11 +16,14 @@ async fn main() -> Result<(), Box> { ) .init(); - let state = build_state_from_env()?; + let state = build_state_from_env().await?; let app = overacp_server::router(state); - let listener = TcpListener::bind("0.0.0.0:8080").await?; - tracing::info!("overacp-server listening on 0.0.0.0:8080"); + let port = env::var("PORT").unwrap_or_else(|_| "8080".to_string()); + let addr = format!("0.0.0.0:{port}"); + + let listener = TcpListener::bind(&addr).await?; + tracing::info!("overacp-server listening on {addr}"); axum::serve(listener, app).await?; Ok(()) } diff --git a/server/src/redis_backend/broker.rs b/server/src/redis_backend/broker.rs new file mode 100644 index 0000000..4ded802 --- /dev/null +++ b/server/src/redis_backend/broker.rs @@ -0,0 +1,101 @@ +//! `RedisStreamBroker` — Redis pub/sub-backed fan-out for `stream/*`, +//! `turn/end`, and `heartbeat` notifications to SSE subscribers. +//! +//! Uses `PUBLISH` / `SUBSCRIBE` on channel `overacp:stream:{agent_id}`. +//! Each subscriber gets a dedicated pub/sub connection (Redis +//! multiplexed connections can't mix SUBSCRIBE with other commands). + +use async_trait::async_trait; +use futures::stream::BoxStream; +use futures::StreamExt; +use redis::aio::ConnectionManager; +use tokio::sync::mpsc; +use tokio_stream::wrappers::UnboundedReceiverStream; +use tracing::warn; +use uuid::Uuid; + +use crate::tunnel::broker::StreamBrokerProvider; + +use super::keys::stream_channel; + +pub struct RedisStreamBroker { + conn: ConnectionManager, + /// Redis URL for creating dedicated subscriber connections. + redis_url: String, +} + +impl RedisStreamBroker { + pub fn new(conn: ConnectionManager, redis_url: String) -> Self { + Self { conn, redis_url } + } +} + +#[async_trait] +impl StreamBrokerProvider for RedisStreamBroker { + async fn publish(&self, agent_id: Uuid, frame: String) { + let channel = stream_channel(agent_id); + let mut conn = self.conn.clone(); + let _: Result<(), _> = redis::cmd("PUBLISH") + .arg(&channel) + .arg(&frame) + .query_async(&mut conn) + .await; + } + + fn subscribe(&self, agent_id: Uuid) -> BoxStream<'static, String> { + let channel = stream_channel(agent_id); + let redis_url = self.redis_url.clone(); + + let (tx, rx) = mpsc::unbounded_channel::(); + + // Spawn a dedicated pub/sub listener. It runs until the + // receiver is dropped (which closes the channel via tx). + tokio::spawn(async move { + let client = match redis::Client::open(redis_url.as_str()) { + Ok(c) => c, + Err(e) => { + warn!("redis subscribe client error: {e}"); + return; + } + }; + let mut pubsub = match client.get_async_pubsub().await { + Ok(ps) => ps, + Err(e) => { + warn!("redis pubsub connect error: {e}"); + return; + } + }; + if let Err(e) = pubsub.subscribe(&channel).await { + warn!("redis subscribe error: {e}"); + return; + } + + let mut stream = pubsub.on_message(); + loop { + tokio::select! { + msg = stream.next() => { + match msg { + Some(msg) => { + let payload: String = match msg.get_payload() { + Ok(p) => p, + Err(_) => continue, + }; + if tx.send(payload).is_err() { + // Receiver dropped — SSE client disconnected. + break; + } + } + None => break, + } + } + _ = tx.closed() => { + // Receiver dropped — SSE client disconnected. + break; + } + } + } + }); + + UnboundedReceiverStream::new(rx).boxed() + } +} diff --git a/server/src/redis_backend/inbox.rs b/server/src/redis_backend/inbox.rs new file mode 100644 index 0000000..2e66dd9 --- /dev/null +++ b/server/src/redis_backend/inbox.rs @@ -0,0 +1,255 @@ +//! Per-agent inbox consumer — drains `overacp:inbox:{agent_id}` via +//! XREADGROUP and forwards frames to the local mpsc sender. +//! +//! One consumer task is spawned per locally-owned agent (started by +//! `RedisAgentRegistry::acquire`, aborted when the `TunnelLease` is +//! dropped). Mirrors the consumer-group pattern from +//! `overfolder/agent-runner/src/triggers/redis_consumer.rs`. + +use std::time::Duration; + +use redis::aio::ConnectionManager; +use redis::streams::{StreamReadOptions, StreamReadReply}; +use redis::AsyncCommands; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tokio::time; +use tracing::{debug, trace, warn}; +use uuid::Uuid; + +use super::keys::{dlq_key, inbox_key}; + +/// Consumer group name shared by all instances. +const CONSUMER_GROUP: &str = "owners"; + +/// How long to sleep between polls when no messages are available. +const IDLE_POLL_MS: u64 = 200; + +/// How many poll iterations between autoclaim sweeps (~5 min). +const AUTOCLAIM_INTERVAL: u64 = 1500; // 1500 * 200ms = 300s + +/// Minimum idle time before an entry can be auto-claimed (5 min). +const AUTOCLAIM_MIN_IDLE_MS: u64 = 300_000; + +/// Max delivery attempts before dead-lettering. +const MAX_DELIVERIES: u64 = 3; + +/// Max wall-clock age of a stream entry before dead-lettering (4h). +const MAX_RECLAIM_AGE_MS: u64 = 4 * 3600 * 1000; + +/// Ensure the consumer group exists on the inbox stream. +pub async fn ensure_consumer_group(conn: &ConnectionManager, agent_id: Uuid) { + let key = inbox_key(agent_id); + let mut c = conn.clone(); + // XGROUP CREATE ... 0 MKSTREAM — idempotent (BUSYGROUP error if + // the group already exists, which we ignore). + let _: Result<(), _> = redis::cmd("XGROUP") + .arg("CREATE") + .arg(&key) + .arg(CONSUMER_GROUP) + .arg("0") + .arg("MKSTREAM") + .query_async(&mut c) + .await; +} + +/// Spawn the inbox consumer for `agent_id`. Returns a `JoinHandle` +/// that the caller should abort when the tunnel lease is released. +pub fn spawn_consumer( + conn: ConnectionManager, + agent_id: Uuid, + consumer_id: String, + local_tx: mpsc::UnboundedSender, +) -> JoinHandle<()> { + tokio::spawn(async move { + let key = inbox_key(agent_id); + let mut c = conn; + let mut poll_count: u64 = 0; + + // First pass: drain pending entries (auto-claimed from dead consumers). + // After draining pending, switch to reading new messages. + let mut read_id = "0".to_string(); + + loop { + poll_count += 1; + + // ── XREADGROUP ──────────────────────────────────── + let opts = StreamReadOptions::default() + .group(CONSUMER_GROUP, &consumer_id) + .count(1); + + let reply: Result = + c.xread_options(&[&key], &[read_id.as_str()], &opts).await; + + match reply { + Ok(reply) if !reply.keys.is_empty() => { + let mut had_entries = false; + for stream_key in &reply.keys { + for entry in &stream_key.ids { + had_entries = true; + let frame = entry.map.get("frame").and_then(|v| { + if let redis::Value::BulkString(bytes) = v { + String::from_utf8(bytes.clone()).ok() + } else { + None + } + }); + + if let Some(frame) = frame { + if local_tx.send(frame).is_err() { + // Receiver dropped — tunnel closed. Do NOT ack; + // the autoclaim will handle redelivery or DLQ. + debug!(%agent_id, "inbox consumer: local tx closed"); + return; + } + // ACK on successful delivery. + let _: Result<(), _> = redis::cmd("XACK") + .arg(&key) + .arg(CONSUMER_GROUP) + .arg(&entry.id) + .query_async(&mut c) + .await; + trace!(%agent_id, id = %entry.id, "inbox: delivered + acked"); + } + } + } + if !had_entries && read_id == "0" { + // No more pending entries; switch to reading new. + read_id = ">".to_string(); + } + } + Ok(_) => { + if read_id == "0" { + // No pending entries; switch to reading new. + read_id = ">".to_string(); + } + // No messages — idle poll. + time::sleep(Duration::from_millis(IDLE_POLL_MS)).await; + } + Err(e) => { + // Transient Redis error — back off. + warn!(%agent_id, "inbox XREADGROUP error: {e}"); + time::sleep(Duration::from_millis(IDLE_POLL_MS)).await; + } + } + + // ── Periodic XAUTOCLAIM ────────────────────────── + if poll_count.is_multiple_of(AUTOCLAIM_INTERVAL) { + autoclaim_orphaned(&mut c, agent_id, &consumer_id, &key).await; + } + } + }) +} + +/// Return type of the XAUTOCLAIM Redis command: `(cursor, entries, deleted_ids)`. +/// Redis 7.0+ returns a 3-element array. +type AutoclaimResult = ( + String, + Vec<(String, Vec<(String, redis::Value)>)>, + Vec, +); + +/// Reclaim orphaned entries from dead consumers and either redeliver +/// or dead-letter them. +async fn autoclaim_orphaned( + conn: &mut ConnectionManager, + agent_id: Uuid, + consumer_id: &str, + stream_key: &str, +) { + // XAUTOCLAIM COUNT 10 + let result: Result = redis::cmd("XAUTOCLAIM") + .arg(stream_key) + .arg(CONSUMER_GROUP) + .arg(consumer_id) + .arg(AUTOCLAIM_MIN_IDLE_MS) + .arg("0-0") + .arg("COUNT") + .arg(10) + .query_async(conn) + .await; + + let (_cursor, entries, _deleted) = match result { + Ok(r) => r, + Err(e) => { + debug!(%agent_id, "xautoclaim error (may be normal): {e}"); + return; + } + }; + + for (entry_id, fields) in entries { + let delivery_count = get_delivery_count(conn, stream_key, &entry_id).await; + let age_ms = stream_id_age_ms(&entry_id); + + let should_dlq = + delivery_count >= MAX_DELIVERIES || age_ms.is_some_and(|a| a > MAX_RECLAIM_AGE_MS); + + if should_dlq { + // Move to DLQ. + let mut args: Vec<(String, String)> = fields + .into_iter() + .filter_map(|(k, v)| { + if let redis::Value::BulkString(bytes) = v { + Some((k, String::from_utf8_lossy(&bytes).to_string())) + } else { + None + } + }) + .collect(); + args.push(("dlq_original_id".to_string(), entry_id.clone())); + args.push(("dlq_delivery_count".to_string(), delivery_count.to_string())); + args.push(("dlq_agent_id".to_string(), agent_id.to_string())); + + let dlq = dlq_key(); + let mut cmd = redis::cmd("XADD"); + cmd.arg(dlq).arg("*"); + for (k, v) in &args { + cmd.arg(k).arg(v); + } + let _: Result = cmd.query_async(conn).await; + + // ACK the original entry. + let _: Result<(), _> = redis::cmd("XACK") + .arg(stream_key) + .arg(CONSUMER_GROUP) + .arg(&entry_id) + .query_async(conn) + .await; + + warn!( + %agent_id, + id = %entry_id, + delivery_count, + "inbox entry dead-lettered" + ); + } + // If not DLQ'd, the entry was auto-claimed to us and will + // be picked up on the next XREADGROUP with ID ">" or "0". + } +} + +/// Get the delivery count for a specific entry via XPENDING. +async fn get_delivery_count(conn: &mut ConnectionManager, stream_key: &str, entry_id: &str) -> u64 { + // XPENDING 1 + let result: Result, _> = redis::cmd("XPENDING") + .arg(stream_key) + .arg(CONSUMER_GROUP) + .arg(entry_id) + .arg(entry_id) + .arg(1) + .query_async(conn) + .await; + + match result { + Ok(entries) if !entries.is_empty() => entries[0].3, + _ => 1, + } +} + +/// Parse the millisecond timestamp from a stream ID (`{ms}-{seq}`). +fn stream_id_age_ms(id: &str) -> Option { + let ms_str = id.split('-').next()?; + let ms: u64 = ms_str.parse().ok()?; + let now_ms = chrono::Utc::now().timestamp_millis() as u64; + Some(now_ms.saturating_sub(ms)) +} diff --git a/server/src/redis_backend/keys.rs b/server/src/redis_backend/keys.rs new file mode 100644 index 0000000..5268e81 --- /dev/null +++ b/server/src/redis_backend/keys.rs @@ -0,0 +1,57 @@ +//! Centralized Redis key helpers for the `redis_backend` module. +//! +//! Every Redis key used by over/ACP lives here so naming conventions +//! stay consistent and duplicates are impossible. + +use uuid::Uuid; + +/// Ownership lease key — `SET EX` with the owning instance ID. +pub fn owner_key(agent_id: Uuid) -> String { + format!("overacp:owner:{agent_id}") +} + +/// Tunnel directory hash — `HSET` with `instance_id`, `connected_at`. +pub fn tunnel_key(agent_id: Uuid) -> String { + format!("overacp:tunnel:{agent_id}") +} + +/// Set of currently-connected agent IDs. +pub fn connected_set_key() -> &'static str { + "overacp:agents:connected" +} + +/// Sorted set of recently-disconnected agent IDs (score = disconnect +/// timestamp in ms). +pub fn recent_zset_key() -> &'static str { + "overacp:agents:recent" +} + +/// Serialised JWT claims snapshot for an agent. +pub fn claims_key(agent_id: Uuid) -> String { + format!("overacp:agents:claims:{agent_id}") +} + +/// Pub/sub control channel for takeover / disconnect signals. +pub fn control_channel(agent_id: Uuid) -> String { + format!("overacp:control:{agent_id}") +} + +/// Per-agent inbox stream consumed via XREADGROUP. +pub fn inbox_key(agent_id: Uuid) -> String { + format!("overacp:inbox:{agent_id}") +} + +/// Dead-letter queue for undeliverable inbox entries. +pub fn dlq_key() -> &'static str { + "overacp:inbox:dlq" +} + +/// Per-agent offline message buffer (separate from the live inbox). +pub fn buffer_key(agent_id: Uuid) -> String { + format!("overacp:buffer:{agent_id}") +} + +/// Pub/sub channel for stream fan-out (SSE subscribers). +pub fn stream_channel(agent_id: Uuid) -> String { + format!("overacp:stream:{agent_id}") +} diff --git a/server/src/redis_backend/lease.rs b/server/src/redis_backend/lease.rs new file mode 100644 index 0000000..513d3bc --- /dev/null +++ b/server/src/redis_backend/lease.rs @@ -0,0 +1,260 @@ +//! Ownership lease for a tunnel — Redis `SET EX` + Lua CAS heartbeat. +//! +//! Ported from `overfolder/agent-runner/src/session_lock.rs`. Key +//! differences from the overfolder version: +//! +//! - **Force-takeover, not NX.** over/ACP allows reconnection from a +//! different instance (matches current in-memory behavior where +//! `register()` overwrites). On acquire we `SET EX` unconditionally +//! and notify the previous owner via pub/sub. +//! - **30s TTL, 10s heartbeat** (vs 5s/2s in overfolder). Tunnel +//! sessions are longer-lived than agentic loop locks. +//! - **Tunnel directory.** We also maintain `overacp:tunnel:{agent_id}` +//! (HASH with instance_id, connected_at) and the connected-agents +//! set + claims record. + +use std::sync::Arc; +use std::time::Duration; + +use redis::aio::ConnectionManager; +use redis::Script; +use tokio::sync::Notify; +use tokio::task::JoinHandle; +use tokio::time; +use tracing::{debug, info, warn}; +use uuid::Uuid; + +use super::keys::{ + claims_key, connected_set_key, control_channel, owner_key, recent_zset_key, tunnel_key, +}; + +/// How long the ownership lease lives without a heartbeat renewal. +pub const LEASE_TTL_SECS: u64 = 30; + +/// How often the heartbeat renews the lease (must be < TTL). +const HEARTBEAT_INTERVAL_SECS: u64 = 10; + +/// Lua script: atomic compare-and-expire (heartbeat). +const HEARTBEAT_SCRIPT: &str = r#" +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("EXPIRE", KEYS[1], ARGV[2]) +else + return 0 +end +"#; + +/// Lua script: atomic compare-and-delete (release). +const RELEASE_SCRIPT: &str = r#" +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) +else + return 0 +end +"#; + +/// Maximum entries in the recently-disconnected sorted set. +const RECENT_CAPACITY: i64 = 64; + +/// A held ownership lease with background heartbeat. +pub struct OwnershipLease { + agent_id: Uuid, + instance_id: String, + conn: ConnectionManager, + stop: Arc, + heartbeat_handle: JoinHandle<()>, +} + +impl OwnershipLease { + /// Release the lease explicitly (async). Also called best-effort + /// from Drop. + pub async fn release(self) { + self.stop.notify_one(); + self.heartbeat_handle.abort(); + release_lease(&self.conn, self.agent_id, &self.instance_id).await; + } +} + +impl Drop for OwnershipLease { + fn drop(&mut self) { + self.stop.notify_one(); + self.heartbeat_handle.abort(); + // Spawn a fire-and-forget cleanup. If the runtime is shutting + // down this may not complete — the TTL will clean up. + let conn = self.conn.clone(); + let agent_id = self.agent_id; + let instance_id = self.instance_id.clone(); + tokio::spawn(async move { + release_lease(&conn, agent_id, &instance_id).await; + }); + } +} + +/// Acquire the ownership lease for `agent_id`. Force-takes over from +/// any existing owner (publishes a `takeover` signal on the control +/// channel so the previous owner can close its WebSocket). +pub async fn acquire_lease( + conn: &ConnectionManager, + agent_id: Uuid, + instance_id: &str, + claims_json: &str, +) -> Result { + let key = owner_key(agent_id); + let mut c = conn.clone(); + + // Check for existing owner and notify takeover if different. + let existing: Option = redis::cmd("GET").arg(&key).query_async(&mut c).await?; + + if let Some(ref prev) = existing { + if prev != instance_id { + info!( + %agent_id, + prev_owner = %prev, + new_owner = %instance_id, + "taking over tunnel ownership" + ); + let _: Result<(), _> = redis::cmd("PUBLISH") + .arg(control_channel(agent_id)) + .arg("takeover") + .query_async(&mut c) + .await; + } + } + + // Force-set the ownership lease (SET EX, no NX). + let _: () = redis::cmd("SET") + .arg(&key) + .arg(instance_id) + .arg("EX") + .arg(LEASE_TTL_SECS) + .query_async(&mut c) + .await?; + + // Populate the tunnel directory. + let tunnel = tunnel_key(agent_id); + let now = chrono::Utc::now().to_rfc3339(); + redis::pipe() + .cmd("HSET") + .arg(&tunnel) + .arg("instance_id") + .arg(instance_id) + .arg("connected_at") + .arg(&now) + .ignore() + .cmd("EXPIRE") + .arg(&tunnel) + .arg(LEASE_TTL_SECS * 2) + .ignore() + // Add to connected set. + .cmd("SADD") + .arg(connected_set_key()) + .arg(agent_id.to_string()) + .ignore() + // Remove from recently-disconnected. + .cmd("ZREM") + .arg(recent_zset_key()) + .arg(agent_id.to_string()) + .ignore() + // Store serialized claims. + .cmd("SET") + .arg(claims_key(agent_id)) + .arg(claims_json) + .arg("EX") + .arg(LEASE_TTL_SECS * 2) + .ignore() + .query_async::<()>(&mut c) + .await?; + + // Start heartbeat. + let stop = Arc::new(Notify::new()); + let stop_clone = stop.clone(); + let mut hb_conn = conn.clone(); + let hb_key = key.clone(); + let hb_tunnel = tunnel_key(agent_id); + let hb_claims_key = claims_key(agent_id); + let hb_value = instance_id.to_string(); + + let heartbeat_handle = tokio::spawn(async move { + let mut interval = time::interval(Duration::from_secs(HEARTBEAT_INTERVAL_SECS)); + interval.tick().await; // skip immediate first tick + loop { + tokio::select! { + () = stop_clone.notified() => break, + _ = interval.tick() => { + // CAS-renew the ownership lease. + let renewed: i32 = Script::new(HEARTBEAT_SCRIPT) + .key(&hb_key) + .arg(&hb_value) + .arg(LEASE_TTL_SECS) + .invoke_async(&mut hb_conn) + .await + .unwrap_or(0); + if renewed == 0 { + warn!(%agent_id, "ownership heartbeat: key missing or stolen"); + break; + } + // Refresh tunnel directory + claims TTL. + let _: Result<(), _> = redis::pipe() + .cmd("EXPIRE").arg(&hb_tunnel).arg(LEASE_TTL_SECS * 2).ignore() + .cmd("EXPIRE").arg(&hb_claims_key).arg(LEASE_TTL_SECS * 2).ignore() + .query_async::<()>(&mut hb_conn) + .await; + } + } + } + }); + + debug!(%agent_id, %instance_id, ttl = LEASE_TTL_SECS, "ownership lease acquired"); + + Ok(OwnershipLease { + agent_id, + instance_id: instance_id.to_string(), + conn: conn.clone(), + stop, + heartbeat_handle, + }) +} + +/// Release ownership — CAS-DEL the lease key, update connected set +/// and recently-disconnected log. +async fn release_lease(conn: &ConnectionManager, agent_id: Uuid, instance_id: &str) { + let key = owner_key(agent_id); + let mut c = conn.clone(); + + // Atomic compare-and-delete. + let deleted: i32 = Script::new(RELEASE_SCRIPT) + .key(&key) + .arg(instance_id) + .invoke_async(&mut c) + .await + .unwrap_or(0); + + if deleted == 1 { + let now_ms = chrono::Utc::now().timestamp_millis(); + let _: Result<(), _> = redis::pipe() + // Remove from connected set. + .cmd("SREM") + .arg(connected_set_key()) + .arg(agent_id.to_string()) + .ignore() + // Add to recently-disconnected sorted set. + .cmd("ZADD") + .arg(recent_zset_key()) + .arg(now_ms) + .arg(agent_id.to_string()) + .ignore() + // Trim recent set to capacity. + .cmd("ZREMRANGEBYRANK") + .arg(recent_zset_key()) + .arg(0) + .arg(-(RECENT_CAPACITY + 1)) + .ignore() + // Clean up tunnel directory. + .cmd("DEL") + .arg(tunnel_key(agent_id)) + .ignore() + .query_async::<()>(&mut c) + .await; + + debug!(%agent_id, "ownership lease released"); + } +} diff --git a/server/src/redis_backend/mod.rs b/server/src/redis_backend/mod.rs new file mode 100644 index 0000000..fd63eb6 --- /dev/null +++ b/server/src/redis_backend/mod.rs @@ -0,0 +1,66 @@ +//! Redis/Valkey backend for multi-instance HA. +//! +//! Behind `#[cfg(feature = "redis")]`. Provides distributed +//! implementations of `AgentRegistryProvider`, `MessageQueueProvider`, +//! and `StreamBrokerProvider` using Redis ownership leases, inbox +//! streams, and pub/sub. + +pub mod broker; +pub mod inbox; +pub mod keys; +pub mod lease; +pub mod queue; +pub mod registry; + +use std::env; +use std::sync::{Arc, OnceLock}; + +use redis::aio::ConnectionManager; + +use self::broker::RedisStreamBroker; +use self::queue::RedisMessageQueue; +use self::registry::RedisAgentRegistry; +use crate::registry::agent::AgentRegistryProvider; +use crate::registry::queue::MessageQueueProvider; +use crate::tunnel::broker::StreamBrokerProvider; + +/// Default per-agent message buffer capacity. +const DEFAULT_QUEUE_CAPACITY: usize = 64; + +/// Lazily resolved instance identifier (hostname or random fallback). +/// Mirrors `overfolder/agent-runner/src/session_lock.rs:43-51`. +pub fn instance_id() -> &'static str { + static ID: OnceLock = OnceLock::new(); + ID.get_or_init(|| { + env::var("OVERACP_INSTANCE_ID") + .or_else(|_| env::var("HOSTNAME")) + .or_else(|_| env::var("K_REVISION")) + .unwrap_or_else(|_| uuid::Uuid::new_v4().to_string()[..8].to_string()) + }) +} + +/// The three Redis-backed providers, ready to plug into `AppState`. +pub struct RedisProviders { + pub registry: Arc, + pub message_queue: Arc, + pub stream_broker: Arc, +} + +impl RedisProviders { + /// Connect to Redis and build all three providers. + pub async fn from_url(redis_url: &str) -> Result { + let client = redis::Client::open(redis_url)?; + let conn = ConnectionManager::new(client).await?; + let id = instance_id().to_string(); + + Ok(Self { + registry: Arc::new(RedisAgentRegistry::new( + conn.clone(), + redis_url.to_string(), + id, + )), + message_queue: Arc::new(RedisMessageQueue::new(conn.clone(), DEFAULT_QUEUE_CAPACITY)), + stream_broker: Arc::new(RedisStreamBroker::new(conn, redis_url.to_string())), + }) + } +} diff --git a/server/src/redis_backend/queue.rs b/server/src/redis_backend/queue.rs new file mode 100644 index 0000000..73c18c6 --- /dev/null +++ b/server/src/redis_backend/queue.rs @@ -0,0 +1,124 @@ +//! `RedisMessageQueue` — Redis-backed per-agent buffer for +//! `session/message` pushes that arrive while the agent is offline. +//! +//! Uses a separate key `overacp:buffer:{agent_id}` from the inbox +//! stream (which is for live delivery to a connected agent). The +//! buffer is drained atomically on reconnect. + +use async_trait::async_trait; +use redis::aio::ConnectionManager; +use redis::Script; +use uuid::Uuid; + +use crate::registry::queue::{MessageQueueProvider, QueueError}; + +use super::keys::buffer_key; + +/// Lua script: atomic XRANGE + DEL. Reads all entries and deletes the +/// stream in one round-trip so no push can slip in between. +const DRAIN_SCRIPT: &str = r#" +local entries = redis.call("XRANGE", KEYS[1], "-", "+", "COUNT", ARGV[1]) +if #entries > 0 then + redis.call("DEL", KEYS[1]) +end +return entries +"#; + +/// Lua script: atomic check-length-then-add. Returns 1 on success, +/// 0 if the stream already has >= capacity entries. +const PUSH_SCRIPT: &str = r#" +local len = redis.call("XLEN", KEYS[1]) +if len >= tonumber(ARGV[1]) then + return 0 +end +redis.call("XADD", KEYS[1], "*", "frame", ARGV[2]) +return 1 +"#; + +pub struct RedisMessageQueue { + conn: ConnectionManager, + per_agent_capacity: usize, +} + +impl RedisMessageQueue { + pub fn new(conn: ConnectionManager, per_agent_capacity: usize) -> Self { + Self { + conn, + per_agent_capacity, + } + } +} + +#[async_trait] +impl MessageQueueProvider for RedisMessageQueue { + async fn push(&self, agent_id: Uuid, frame: String) -> Result<(), QueueError> { + let key = buffer_key(agent_id); + let mut conn = self.conn.clone(); + + // Atomic check-and-add via Lua to avoid TOCTOU race between + // XLEN and XADD. + let added: i32 = Script::new(PUSH_SCRIPT) + .key(&key) + .arg(self.per_agent_capacity) + .arg(&frame) + .invoke_async(&mut conn) + .await + .map_err(|_| QueueError::Full { + agent_id, + capacity: self.per_agent_capacity, + })?; + + if added == 0 { + return Err(QueueError::Full { + agent_id, + capacity: self.per_agent_capacity, + }); + } + + Ok(()) + } + + async fn drain(&self, agent_id: Uuid) -> Vec { + let key = buffer_key(agent_id); + let mut conn = self.conn.clone(); + + // Atomic XRANGE + DEL via Lua to avoid losing messages pushed + // between the two commands. + let result: Result)>, _> = Script::new(DRAIN_SCRIPT) + .key(&key) + .arg(self.per_agent_capacity) + .invoke_async(&mut conn) + .await; + + match result { + Ok(entries) => entries + .into_iter() + .filter_map(|(_id, fields)| { + fields + .into_iter() + .find(|(k, _)| k == "frame") + .map(|(_, v)| v) + }) + .collect(), + Err(_) => Vec::new(), + } + } + + async fn len(&self, agent_id: Uuid) -> usize { + let key = buffer_key(agent_id); + let mut conn = self.conn.clone(); + redis::cmd("XLEN") + .arg(&key) + .query_async(&mut conn) + .await + .unwrap_or(0) + } + + async fn is_empty(&self, agent_id: Uuid) -> bool { + self.len(agent_id).await == 0 + } + + fn capacity(&self) -> usize { + self.per_agent_capacity + } +} diff --git a/server/src/redis_backend/registry.rs b/server/src/redis_backend/registry.rs new file mode 100644 index 0000000..963a296 --- /dev/null +++ b/server/src/redis_backend/registry.rs @@ -0,0 +1,340 @@ +//! `RedisAgentRegistry` — Redis-backed agent registry implementing +//! `AgentRegistryProvider`. +//! +//! Combines ownership leases, inbox streams, and the connected/recent +//! metadata to provide multi-instance agent routing. The "produce +//! unconditionally" pattern means `deliver()` always writes to the +//! per-agent inbox stream; the owner's XREADGROUP consumer drains it. + +use std::sync::Arc; + +use async_trait::async_trait; +use dashmap::DashMap; +use futures::StreamExt; +use redis::aio::ConnectionManager; +use redis::{AsyncCommands, Script}; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tracing::{info, warn}; +use uuid::Uuid; + +use crate::auth::Claims; +use crate::registry::agent::{ + AgentDescription, AgentRegistryProvider, DeliveryOutcome, RegistryError, TunnelLease, +}; + +use super::inbox; +use super::keys::{ + claims_key, connected_set_key, control_channel, inbox_key, owner_key, recent_zset_key, +}; +use super::lease; + +/// Serializable claims snapshot stored in Redis for `list` / `describe`. +#[derive(Debug, Clone, Serialize, Deserialize)] +struct ClaimsSnapshot { + agent_id: Uuid, + user: Option, + connected_at: String, +} + +/// Lua script: atomic EXISTS + XADD. Checks the owner key exists before +/// writing to the inbox, preventing orphaned messages when the owner +/// disconnects between the two operations. +const DELIVER_SCRIPT: &str = r#" +if redis.call("EXISTS", KEYS[1]) == 1 then + redis.call("XADD", KEYS[2], "*", "frame", ARGV[1]) + return 1 +end +return 0 +"#; + +pub struct RedisAgentRegistry { + conn: ConnectionManager, + /// Redis URL for creating dedicated pub/sub connections. + redis_url: String, + instance_id: String, + /// Local senders for agents owned by this instance. The inbox + /// consumer reads from the stream and forwards to these senders. + local_senders: Arc>>, + /// Inbox consumer handles — aborted when the lease drops. + consumer_handles: Arc>>, +} + +impl RedisAgentRegistry { + pub fn new(conn: ConnectionManager, redis_url: String, instance_id: String) -> Self { + Self { + conn, + redis_url, + instance_id, + local_senders: Arc::new(DashMap::new()), + consumer_handles: Arc::new(DashMap::new()), + } + } +} + +#[async_trait] +impl AgentRegistryProvider for RedisAgentRegistry { + async fn acquire( + &self, + agent_id: Uuid, + local_tx: mpsc::UnboundedSender, + claims: Claims, + ) -> Result { + // Serialize claims for the Redis metadata store. + let snapshot = ClaimsSnapshot { + agent_id, + user: claims.user, + connected_at: chrono::Utc::now().to_rfc3339(), + }; + let claims_json = serde_json::to_string(&snapshot).unwrap_or_default(); + + // Acquire the ownership lease. + let ownership = lease::acquire_lease(&self.conn, agent_id, &self.instance_id, &claims_json) + .await + .map_err(|e| RegistryError::AcquireFailed { + agent_id, + reason: e.to_string(), + })?; + + // Ensure the consumer group exists on the inbox stream. + inbox::ensure_consumer_group(&self.conn, agent_id).await; + + // Store the local sender so the inbox consumer can forward. + self.local_senders.insert(agent_id, local_tx.clone()); + + // Spawn the inbox consumer. + let consumer_id = format!("{}-{agent_id}", self.instance_id); + let handle = inbox::spawn_consumer(self.conn.clone(), agent_id, consumer_id, local_tx); + self.consumer_handles.insert(agent_id, handle); + + info!(%agent_id, instance = %self.instance_id, "redis tunnel lease acquired"); + + // Build the TunnelLease RAII guard. + let senders = self.local_senders.clone(); + let handles = self.consumer_handles.clone(); + Ok(TunnelLease::new(move || { + // Abort the inbox consumer. + if let Some((_, handle)) = handles.remove(&agent_id) { + handle.abort(); + } + // Remove local sender. + senders.remove(&agent_id); + // The OwnershipLease drop handles Redis cleanup. + drop(ownership); + })) + } + + async fn deliver(&self, agent_id: Uuid, frame: String) -> DeliveryOutcome { + let owner = owner_key(agent_id); + let inbox = inbox_key(agent_id); + let mut conn = self.conn.clone(); + + // Atomic EXISTS + XADD via Lua to prevent orphaned messages + // when the owner disconnects between the two operations. + let delivered: i32 = Script::new(DELIVER_SCRIPT) + .key(&owner) + .key(&inbox) + .arg(&frame) + .invoke_async(&mut conn) + .await + .unwrap_or(0); + + if delivered == 1 { + DeliveryOutcome::Live + } else { + DeliveryOutcome::NoTunnel(frame) + } + } + + async fn is_connected(&self, agent_id: Uuid) -> bool { + let key = owner_key(agent_id); + let mut conn = self.conn.clone(); + redis::cmd("EXISTS") + .arg(&key) + .query_async(&mut conn) + .await + .unwrap_or(false) + } + + async fn list_agents(&self) -> Vec { + let mut conn = self.conn.clone(); + let mut out = Vec::new(); + + // Connected agents. + let connected: Vec = conn.smembers(connected_set_key()).await.unwrap_or_default(); + + for id_str in &connected { + if let Ok(agent_id) = Uuid::parse_str(id_str) { + if let Some(desc) = hydrate_description(&mut conn, agent_id, true).await { + out.push(desc); + } + } + } + + // Recently-disconnected agents. + let recent: Vec<(String, f64)> = redis::cmd("ZRANGE") + .arg(recent_zset_key()) + .arg(0) + .arg(-1) + .arg("WITHSCORES") + .query_async(&mut conn) + .await + .unwrap_or_default(); + + for (id_str, disconnect_ts) in &recent { + if let Ok(agent_id) = Uuid::parse_str(id_str) { + // Skip if also in connected (race). + if connected.contains(id_str) { + continue; + } + let idle_secs = { + let now_ms = chrono::Utc::now().timestamp_millis() as f64; + ((now_ms - disconnect_ts) / 1000.0).max(0.0) as u64 + }; + out.push(AgentDescription { + agent_id, + connected: false, + uptime_secs: None, + idle_secs: Some(idle_secs), + user: None, + }); + } + } + + out + } + + async fn describe_agent(&self, agent_id: Uuid) -> Option { + let mut conn = self.conn.clone(); + + // Check connected first. + let is_member: bool = conn + .sismember(connected_set_key(), agent_id.to_string()) + .await + .unwrap_or(false); + + if is_member { + return hydrate_description(&mut conn, agent_id, true).await; + } + + // Check recently-disconnected. + let score: Option = redis::cmd("ZSCORE") + .arg(recent_zset_key()) + .arg(agent_id.to_string()) + .query_async(&mut conn) + .await + .ok(); + + score.map(|disconnect_ts| { + let now_ms = chrono::Utc::now().timestamp_millis() as f64; + let idle_secs = ((now_ms - disconnect_ts) / 1000.0).max(0.0) as u64; + AgentDescription { + agent_id, + connected: false, + uptime_secs: None, + idle_secs: Some(idle_secs), + user: None, + } + }) + } + + async fn disconnect(&self, agent_id: Uuid) { + let mut conn = self.conn.clone(); + let channel = control_channel(agent_id); + let _: Result<(), _> = redis::cmd("PUBLISH") + .arg(&channel) + .arg("disconnect") + .query_async(&mut conn) + .await; + } + + async fn touch(&self, _agent_id: Uuid) { + // No-op in Redis mode. Activity is tracked via the lease + // heartbeat timestamp. The `idle_secs` field in + // `AgentDescription` is derived from the connected_at + // or heartbeat time. + } + + fn control_receiver(&self, agent_id: Uuid) -> Option> { + let channel = control_channel(agent_id); + let redis_url = self.redis_url.clone(); + let (tx, rx) = mpsc::unbounded_channel(); + + tokio::spawn(async move { + let client = match redis::Client::open(redis_url.as_str()) { + Ok(c) => c, + Err(e) => { + warn!("control subscriber: redis client error: {e}"); + return; + } + }; + let mut pubsub = match client.get_async_pubsub().await { + Ok(ps) => ps, + Err(e) => { + warn!("control subscriber: pubsub connect error: {e}"); + return; + } + }; + if let Err(e) = pubsub.subscribe(&channel).await { + warn!("control subscriber: subscribe error: {e}"); + return; + } + + let mut stream = pubsub.on_message(); + loop { + tokio::select! { + msg = stream.next() => { + match msg { + Some(msg) => { + let payload: String = match msg.get_payload() { + Ok(p) => p, + Err(_) => continue, + }; + if tx.send(payload).is_err() { + break; + } + } + None => break, + } + } + _ = tx.closed() => { + break; + } + } + } + }); + + Some(rx) + } +} + +/// Hydrate an `AgentDescription` from the claims snapshot in Redis. +async fn hydrate_description( + conn: &mut ConnectionManager, + agent_id: Uuid, + connected: bool, +) -> Option { + let claims_json: Option = conn.get(claims_key(agent_id)).await.ok()?; + + let snapshot: Option = claims_json.and_then(|j| serde_json::from_str(&j).ok()); + + let (uptime_secs, idle_secs, user) = if let Some(snap) = snapshot { + let connected_at = chrono::DateTime::parse_from_rfc3339(&snap.connected_at) + .map(|dt| dt.timestamp()) + .unwrap_or(0); + let now = chrono::Utc::now().timestamp(); + let uptime = (now - connected_at).max(0) as u64; + (Some(uptime), Some(0), snap.user) + } else { + (None, None, None) + }; + + Some(AgentDescription { + agent_id, + connected, + uptime_secs, + idle_secs, + user, + }) +} diff --git a/server/src/registry/agent.rs b/server/src/registry/agent.rs index 461f84d..cb1b3f9 100644 --- a/server/src/registry/agent.rs +++ b/server/src/registry/agent.rs @@ -1,4 +1,4 @@ -//! `AgentRegistry` — in-memory routing table of currently-connected +//! `InMemoryAgentRegistry` — in-memory routing table of currently-connected //! agents. //! //! The registry stores one [`AgentEntry`] per active tunnel, keyed by @@ -16,12 +16,100 @@ use serde::Serialize; use tokio::sync::mpsc; use uuid::Uuid; +use async_trait::async_trait; +use thiserror::Error; + use crate::auth::Claims; /// How many recently-disconnected agents to remember. Bounded so a /// flapping client can't grow the registry without limit. const RECENT_CAPACITY: usize = 64; +/// Outcome of a `deliver` call on the registry. +pub enum DeliveryOutcome { + /// Frame was delivered to the agent's tunnel (locally or via stream). + Live, + /// No active tunnel for this agent. The original frame is returned + /// so the caller can buffer it in the `MessageQueue`. + NoTunnel(String), +} + +/// Errors from tunnel lease acquisition. +#[derive(Debug, Error)] +pub enum RegistryError { + #[error("failed to acquire tunnel for agent {agent_id}: {reason}")] + AcquireFailed { agent_id: Uuid, reason: String }, +} + +/// RAII guard for a tunnel lease. Dropping the guard releases the +/// lease (unregisters the agent in the in-memory case, releases the +/// Redis ownership lock in the distributed case). +pub struct TunnelLease { + cleanup: Option>, +} + +impl TunnelLease { + pub fn new(cleanup: impl FnOnce() + Send + 'static) -> Self { + Self { + cleanup: Some(Box::new(cleanup)), + } + } +} + +impl Drop for TunnelLease { + fn drop(&mut self) { + if let Some(f) = self.cleanup.take() { + f(); + } + } +} + +/// Trait for agent registry implementations. The in-memory default +/// stores entries in a `DashMap`; the Redis backend (behind the +/// `redis` feature) uses ownership leases and inbox streams. +#[async_trait] +pub trait AgentRegistryProvider: Send + Sync { + /// Register a freshly-connected agent and acquire a tunnel lease. + /// The returned `TunnelLease` is an RAII guard: dropping it + /// releases the lease and unregisters the agent. + async fn acquire( + &self, + agent_id: Uuid, + local_tx: mpsc::UnboundedSender, + claims: Claims, + ) -> Result; + + /// Deliver a frame to the agent's tunnel. Returns `Live` if the + /// frame was accepted for delivery (locally or via a stream), + /// `NoTunnel(frame)` if no tunnel exists (caller should buffer). + async fn deliver(&self, agent_id: Uuid, frame: String) -> DeliveryOutcome; + + /// Whether `agent_id` currently has a live tunnel anywhere in + /// the cluster. + async fn is_connected(&self, agent_id: Uuid) -> bool; + + /// Snapshot of all connected + recently-disconnected agents. + async fn list_agents(&self) -> Vec; + + /// Describe a single agent by ID. + async fn describe_agent(&self, agent_id: Uuid) -> Option; + + /// Force-disconnect the tunnel for `agent_id`. + async fn disconnect(&self, agent_id: Uuid); + + /// Record activity from the agent (bumps last-activity timestamp). + async fn touch(&self, agent_id: Uuid); + + /// Subscribe to control signals (e.g. `takeover`, `disconnect`) + /// for this agent. Returns `None` for backends that do not support + /// cross-instance control (the in-memory default). When `Some`, + /// the receiver yields signal strings; the tunnel read loop should + /// break when one arrives. + fn control_receiver(&self, _agent_id: Uuid) -> Option> { + None + } +} + /// Routing entry for one connected agent. pub struct AgentEntry { /// Send messages to the connected agent over its tunnel. The @@ -120,18 +208,18 @@ struct RecentEntry { /// Per-agent routing table. Cheap to clone — internally an `Arc`- /// equipped DashMap and a Mutex-guarded VecDeque. #[derive(Clone)] -pub struct AgentRegistry { +pub struct InMemoryAgentRegistry { connected: Arc>>, recent: Arc>>, } -impl Default for AgentRegistry { +impl Default for InMemoryAgentRegistry { fn default() -> Self { Self::new() } } -impl AgentRegistry { +impl InMemoryAgentRegistry { pub fn new() -> Self { Self { connected: Arc::new(DashMap::new()), @@ -232,6 +320,55 @@ impl AgentRegistry { } } +#[async_trait] +impl AgentRegistryProvider for InMemoryAgentRegistry { + async fn acquire( + &self, + agent_id: Uuid, + local_tx: mpsc::UnboundedSender, + claims: Claims, + ) -> Result { + self.register(agent_id, AgentEntry::new(local_tx, claims)); + let registry = self.clone(); + Ok(TunnelLease::new(move || { + registry.unregister(agent_id); + })) + } + + async fn deliver(&self, agent_id: Uuid, frame: String) -> DeliveryOutcome { + if let Some(entry) = self.get(agent_id) { + match entry.tx.send(frame) { + Ok(()) => DeliveryOutcome::Live, + Err(err) => DeliveryOutcome::NoTunnel(err.0), + } + } else { + DeliveryOutcome::NoTunnel(frame) + } + } + + async fn is_connected(&self, agent_id: Uuid) -> bool { + self.connected.contains_key(&agent_id) + } + + async fn list_agents(&self) -> Vec { + self.list() + } + + async fn describe_agent(&self, agent_id: Uuid) -> Option { + self.describe(agent_id) + } + + async fn disconnect(&self, agent_id: Uuid) { + self.unregister(agent_id); + } + + async fn touch(&self, agent_id: Uuid) { + if let Some(entry) = self.get(agent_id) { + entry.touch(); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -244,7 +381,7 @@ mod tests { #[test] fn register_and_get() { - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); let id = Uuid::new_v4(); let (_rx, entry) = entry_for(id); reg.register(id, entry); @@ -254,16 +391,16 @@ mod tests { #[test] fn default_matches_new() { - // `AgentRegistry::default()` should produce an empty + // `InMemoryAgentRegistry::default()` should produce an empty // registry equivalent to `::new()`. - let reg: AgentRegistry = Default::default(); + let reg: InMemoryAgentRegistry = Default::default(); assert!(reg.list().is_empty()); assert!(!reg.is_connected(Uuid::new_v4())); } #[test] fn unregister_moves_to_recent_log() { - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); let id = Uuid::new_v4(); let (_rx, entry) = entry_for(id); reg.register(id, entry); @@ -277,7 +414,7 @@ mod tests { #[test] fn describe_finds_recent_entry_after_disconnect() { - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); let id = Uuid::new_v4(); let (_rx, entry) = entry_for(id); reg.register(id, entry); @@ -289,7 +426,7 @@ mod tests { #[test] fn reconnect_clears_recent_entry() { - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); let id = Uuid::new_v4(); let (_rx, e1) = entry_for(id); reg.register(id, e1); @@ -317,13 +454,13 @@ mod tests { #[test] fn describe_unknown_agent_is_none() { - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); assert!(reg.describe(Uuid::new_v4()).is_none()); } #[test] fn recent_log_is_bounded() { - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); // Push more than RECENT_CAPACITY agents through the // disconnect path; the log should never exceed the cap. for _ in 0..(RECENT_CAPACITY + 32) { @@ -338,7 +475,7 @@ mod tests { #[test] fn list_orders_connected_first() { - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); // disconnected agent let dead_id = Uuid::new_v4(); let (_rx, dead) = entry_for(dead_id); @@ -368,7 +505,7 @@ mod tests { // the race, but we pin the end state after an unregister // + register cycle: the agent appears in the list exactly // once and is marked connected. - let reg = AgentRegistry::new(); + let reg = InMemoryAgentRegistry::new(); let id = Uuid::new_v4(); let (_rx1, e1) = entry_for(id); diff --git a/server/src/registry/mod.rs b/server/src/registry/mod.rs index 32544bb..4922858 100644 --- a/server/src/registry/mod.rs +++ b/server/src/registry/mod.rs @@ -17,5 +17,8 @@ pub mod agent; pub mod queue; -pub use agent::{AgentDescription, AgentEntry, AgentRegistry}; -pub use queue::{MessageQueue, QueueError}; +pub use agent::{ + AgentDescription, AgentEntry, AgentRegistryProvider, DeliveryOutcome, + InMemoryAgentRegistry as AgentRegistry, RegistryError, TunnelLease, +}; +pub use queue::{InMemoryMessageQueue as MessageQueue, MessageQueueProvider, QueueError}; diff --git a/server/src/registry/queue.rs b/server/src/registry/queue.rs index d2fe361..625f4f5 100644 --- a/server/src/registry/queue.rs +++ b/server/src/registry/queue.rs @@ -1,4 +1,4 @@ -//! `MessageQueue` — bounded per-agent buffer for `session/message` +//! `InMemoryMessageQueue` — bounded per-agent buffer for `session/message` //! pushes that arrive while the agent's tunnel is disconnected. //! //! When `POST /agents/{id}/messages` runs against an agent that has @@ -15,6 +15,7 @@ use std::collections::VecDeque; use std::sync::{Arc, Mutex}; +use async_trait::async_trait; use dashmap::DashMap; use thiserror::Error; use uuid::Uuid; @@ -34,20 +35,43 @@ pub enum QueueError { Full { agent_id: Uuid, capacity: usize }, } +/// Trait for message queue implementations. The in-memory default +/// uses a bounded `VecDeque` per agent; the Redis backend uses a +/// capped stream. +#[async_trait] +pub trait MessageQueueProvider: Send + Sync { + /// Append `frame` to `agent_id`'s buffer. Returns `Err(Full)` + /// if the buffer is at capacity. + async fn push(&self, agent_id: Uuid, frame: String) -> Result<(), QueueError>; + + /// Drain `agent_id`'s buffer in FIFO order. Returns an empty + /// vec if there is nothing to drain. + async fn drain(&self, agent_id: Uuid) -> Vec; + + /// Number of buffered frames for `agent_id`. + async fn len(&self, agent_id: Uuid) -> usize; + + /// Whether `agent_id`'s buffer is empty (or absent). + async fn is_empty(&self, agent_id: Uuid) -> bool; + + /// Per-agent capacity ceiling. + fn capacity(&self) -> usize; +} + /// Per-agent bounded notification buffer. Cheap to clone. #[derive(Clone)] -pub struct MessageQueue { +pub struct InMemoryMessageQueue { inner: Arc>>>, per_agent_capacity: usize, } -impl Default for MessageQueue { +impl Default for InMemoryMessageQueue { fn default() -> Self { Self::new(DEFAULT_PER_AGENT_CAPACITY) } } -impl MessageQueue { +impl InMemoryMessageQueue { /// Build a queue with the given per-agent capacity. pub fn new(per_agent_capacity: usize) -> Self { Self { @@ -114,6 +138,29 @@ impl MessageQueue { } } +#[async_trait] +impl MessageQueueProvider for InMemoryMessageQueue { + async fn push(&self, agent_id: Uuid, frame: String) -> Result<(), QueueError> { + InMemoryMessageQueue::push(self, agent_id, frame) + } + + async fn drain(&self, agent_id: Uuid) -> Vec { + InMemoryMessageQueue::drain(self, agent_id) + } + + async fn len(&self, agent_id: Uuid) -> usize { + InMemoryMessageQueue::len(self, agent_id) + } + + async fn is_empty(&self, agent_id: Uuid) -> bool { + InMemoryMessageQueue::is_empty(self, agent_id) + } + + fn capacity(&self) -> usize { + self.per_agent_capacity + } +} + #[cfg(test)] mod tests { use super::*; @@ -124,13 +171,13 @@ mod tests { #[test] fn empty_drain_returns_empty() { - let q = MessageQueue::default(); + let q = InMemoryMessageQueue::default(); assert!(q.drain(Uuid::new_v4()).is_empty()); } #[test] fn push_then_drain_preserves_order() { - let q = MessageQueue::default(); + let q = InMemoryMessageQueue::default(); let id = Uuid::new_v4(); q.push(id, frame(1)).unwrap(); q.push(id, frame(2)).unwrap(); @@ -144,7 +191,7 @@ mod tests { #[test] fn drain_removes_slot() { - let q = MessageQueue::default(); + let q = InMemoryMessageQueue::default(); let id = Uuid::new_v4(); q.push(id, frame(1)).unwrap(); let _ = q.drain(id); @@ -156,7 +203,7 @@ mod tests { #[test] fn capacity_overflow_returns_full() { - let q = MessageQueue::new(3); + let q = InMemoryMessageQueue::new(3); let id = Uuid::new_v4(); q.push(id, frame(1)).unwrap(); q.push(id, frame(2)).unwrap(); @@ -169,7 +216,7 @@ mod tests { #[test] fn agents_are_isolated() { - let q = MessageQueue::default(); + let q = InMemoryMessageQueue::default(); let a = Uuid::new_v4(); let b = Uuid::new_v4(); q.push(a, frame(1)).unwrap(); @@ -183,7 +230,7 @@ mod tests { #[test] fn len_and_is_empty_track_buffer() { - let q = MessageQueue::default(); + let q = InMemoryMessageQueue::default(); let id = Uuid::new_v4(); assert!(q.is_empty(id)); q.push(id, frame(1)).unwrap(); @@ -193,9 +240,9 @@ mod tests { #[test] fn capacity_returns_configured_value() { - assert_eq!(MessageQueue::new(7).capacity(), 7); + assert_eq!(InMemoryMessageQueue::new(7).capacity(), 7); assert_eq!( - MessageQueue::default().capacity(), + InMemoryMessageQueue::default().capacity(), DEFAULT_PER_AGENT_CAPACITY ); } diff --git a/server/src/startup.rs b/server/src/startup.rs index c7ae28f..4d0b7aa 100644 --- a/server/src/startup.rs +++ b/server/src/startup.rs @@ -14,6 +14,10 @@ use crate::{AppState, StaticJwtAuthenticator}; pub enum StartupError { #[error("OVERACP_JWT_SIGNING_KEY is required (set it in .env or the environment)")] MissingSigningKey, + + #[cfg(feature = "redis")] + #[error("redis connection failed: {0}")] + Redis(#[from] redis::RedisError), } /// Build the server's [`AppState`] from process environment variables. @@ -22,13 +26,27 @@ pub enum StartupError { /// /// - `OVERACP_JWT_SIGNING_KEY` (required) /// - `OVERACP_JWT_ISSUER` (optional, defaults to `"overacp"`) -pub fn build_state_from_env() -> Result { +/// - `OVERACP_REDIS_URL` (optional, enables Redis backend for HA) +/// - `OVERACP_INSTANCE_ID` (optional, defaults to hostname or random) +pub async fn build_state_from_env() -> Result { let signing_key = env::var("OVERACP_JWT_SIGNING_KEY").map_err(|_| StartupError::MissingSigningKey)?; let issuer = env::var("OVERACP_JWT_ISSUER").unwrap_or_else(|_| "overacp".to_string()); - Ok(AppState::new(Arc::new(StaticJwtAuthenticator::new( - signing_key, - issuer, - )))) + #[allow(unused_mut)] + let mut state = AppState::new(Arc::new(StaticJwtAuthenticator::new(signing_key, issuer))); + + #[cfg(feature = "redis")] + if let Ok(redis_url) = env::var("OVERACP_REDIS_URL") { + use crate::redis_backend; + tracing::info!( + redis_url = %redis_url, + instance_id = %redis_backend::instance_id(), + "enabling redis backend for multi-instance HA" + ); + let providers = redis_backend::RedisProviders::from_url(&redis_url).await?; + state = state.with_redis_providers(providers); + } + + Ok(state) } diff --git a/server/src/state.rs b/server/src/state.rs index 315a476..841c00c 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -4,19 +4,23 @@ use crate::auth::Authenticator; use crate::hooks::{ BootProvider, DefaultBootProvider, DefaultQuotaPolicy, DefaultToolHost, QuotaPolicy, ToolHost, }; -use crate::registry::{AgentRegistry, MessageQueue}; -use crate::tunnel::broker::StreamBroker; +use crate::registry::agent::{AgentRegistryProvider, InMemoryAgentRegistry}; +use crate::registry::queue::{InMemoryMessageQueue, MessageQueueProvider}; +use crate::tunnel::broker::{InMemoryStreamBroker, StreamBrokerProvider}; + +#[cfg(feature = "redis")] +use crate::redis_backend; #[derive(Clone)] pub struct AppState { pub authenticator: Arc, /// Per-agent routing table. The source of truth for every REST /// handler in `api/agents.rs` and the tunnel write path. - pub registry: AgentRegistry, + pub registry: Arc, /// Bounded per-agent buffer for `session/message` pushes that /// arrive while the agent's tunnel is disconnected. - pub message_queue: MessageQueue, - pub stream_broker: Arc, + pub message_queue: Arc, + pub stream_broker: Arc, /// Operator hook for `initialize` — see `hooks::BootProvider`. pub boot_provider: Arc, /// Operator hook for `tools/list` and `tools/call`. @@ -29,9 +33,9 @@ impl AppState { pub fn new(authenticator: Arc) -> Self { Self { authenticator, - registry: AgentRegistry::new(), - message_queue: MessageQueue::default(), - stream_broker: StreamBroker::new(), + registry: Arc::new(InMemoryAgentRegistry::new()), + message_queue: Arc::new(InMemoryMessageQueue::default()), + stream_broker: InMemoryStreamBroker::new(), boot_provider: Arc::new(DefaultBootProvider), tool_host: Arc::new(DefaultToolHost), quota_policy: Arc::new(DefaultQuotaPolicy), @@ -58,6 +62,16 @@ impl AppState { self.quota_policy = policy; self } + + /// Builder: swap all three routing providers to Redis-backed + /// implementations for multi-instance HA. + #[cfg(feature = "redis")] + pub fn with_redis_providers(mut self, providers: redis_backend::RedisProviders) -> Self { + self.registry = providers.registry; + self.message_queue = providers.message_queue; + self.stream_broker = providers.stream_broker; + self + } } #[cfg(test)] diff --git a/server/src/tunnel/broker.rs b/server/src/tunnel/broker.rs index 8fbd989..cad6af6 100644 --- a/server/src/tunnel/broker.rs +++ b/server/src/tunnel/broker.rs @@ -9,22 +9,41 @@ //! //! For now this lives in process only; multi-replica deployments //! will swap in a Redis/Valkey-backed impl behind the same -//! `StreamBroker` interface. +//! `StreamBrokerProvider` interface. use std::sync::Arc; +use async_trait::async_trait; use dashmap::DashMap; +use futures::stream::BoxStream; +use futures::StreamExt; use tokio::sync::broadcast; +use tokio_stream::wrappers::BroadcastStream; use uuid::Uuid; const STREAM_BUFFER: usize = 128; +/// Trait for stream broker implementations. The in-memory default +/// uses `tokio::sync::broadcast` channels; the Redis backend uses +/// pub/sub for cross-instance fan-out. +#[async_trait] +pub trait StreamBrokerProvider: Send + Sync { + /// Publish a frame for fan-out to all SSE subscribers of + /// `agent_id`, wherever they live in the cluster. + async fn publish(&self, agent_id: Uuid, frame: String); + + /// Subscribe to an agent's stream fan-out. Returns a boxed + /// stream that yields frames as they arrive. Dropping the + /// stream unsubscribes. + fn subscribe(&self, agent_id: Uuid) -> BoxStream<'static, String>; +} + #[derive(Default)] -pub struct StreamBroker { +pub struct InMemoryStreamBroker { bus: DashMap>, } -impl StreamBroker { +impl InMemoryStreamBroker { pub fn new() -> Arc { Arc::new(Self::default()) } @@ -43,13 +62,29 @@ impl StreamBroker { } } +#[async_trait] +impl StreamBrokerProvider for InMemoryStreamBroker { + async fn publish(&self, agent_id: Uuid, frame: String) { + let sender = self.sender_for(agent_id); + // Best-effort: if no subscribers, the send silently fails. + let _ = sender.send(frame); + } + + fn subscribe(&self, agent_id: Uuid) -> BoxStream<'static, String> { + let rx = InMemoryStreamBroker::subscribe(self, agent_id); + BroadcastStream::new(rx) + .filter_map(|r: Result| async { r.ok() }) + .boxed() + } +} + #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn fan_out_to_subscriber() { - let b = StreamBroker::new(); + let b = InMemoryStreamBroker::new(); let id = Uuid::new_v4(); let mut rx = b.subscribe(id); b.sender_for(id).send("hello".into()).unwrap(); diff --git a/server/src/tunnel/dispatch.rs b/server/src/tunnel/dispatch.rs index 2252d99..54686ed 100644 --- a/server/src/tunnel/dispatch.rs +++ b/server/src/tunnel/dispatch.rs @@ -172,16 +172,17 @@ mod tests { BootProvider, DefaultBootProvider, DefaultQuotaPolicy, DefaultToolHost, QuotaPolicy, ToolHost, }; - use crate::registry::{AgentRegistry, MessageQueue}; - use crate::tunnel::broker::StreamBroker; + use crate::registry::agent::InMemoryAgentRegistry; + use crate::registry::queue::InMemoryMessageQueue; + use crate::tunnel::broker::InMemoryStreamBroker; use crate::tunnel::run::TunnelContext; fn ctx_default(agent_id: Uuid) -> TunnelContext { TunnelContext { claims: Claims::agent(agent_id, Some(Uuid::new_v4()), 60, "test"), - registry: AgentRegistry::new(), - message_queue: MessageQueue::default(), - stream_broker: StreamBroker::new(), + registry: Arc::new(InMemoryAgentRegistry::new()), + message_queue: Arc::new(InMemoryMessageQueue::default()), + stream_broker: InMemoryStreamBroker::new(), boot_provider: Arc::new(DefaultBootProvider), tool_host: Arc::new(DefaultToolHost), quota_policy: Arc::new(DefaultQuotaPolicy), @@ -427,9 +428,9 @@ mod tests { ) -> TunnelContext { TunnelContext { claims: Claims::agent(Uuid::new_v4(), None, 60, "test"), - registry: AgentRegistry::new(), - message_queue: MessageQueue::default(), - stream_broker: StreamBroker::new(), + registry: Arc::new(InMemoryAgentRegistry::new()), + message_queue: Arc::new(InMemoryMessageQueue::default()), + stream_broker: InMemoryStreamBroker::new(), boot_provider: boot, tool_host: tools, quota_policy: quota, diff --git a/server/src/tunnel/mod.rs b/server/src/tunnel/mod.rs index 67fcd23..587567b 100644 --- a/server/src/tunnel/mod.rs +++ b/server/src/tunnel/mod.rs @@ -3,12 +3,12 @@ //! [`run::run_tunnel`] against the operator hooks on //! [`crate::AppState`]; incoming `stream/*`, `turn/end`, and //! `heartbeat` notifications are fanned out through -//! [`broker::StreamBroker`] to SSE subscribers. +//! [`broker::InMemoryStreamBroker`] to SSE subscribers. pub mod broker; pub mod dispatch; pub mod run; -pub use broker::StreamBroker; +pub use broker::{InMemoryStreamBroker, StreamBrokerProvider}; pub use dispatch::handle_message; pub use run::{run_tunnel, TunnelContext}; diff --git a/server/src/tunnel/run.rs b/server/src/tunnel/run.rs index bb5e511..7b0c02e 100644 --- a/server/src/tunnel/run.rs +++ b/server/src/tunnel/run.rs @@ -3,22 +3,25 @@ use std::sync::Arc; use std::time::Duration; +use std::future::pending; + use axum::extract::ws::{Message, WebSocket}; use futures_util::{SinkExt, StreamExt}; use tokio::sync::mpsc; use tokio::time; -use tracing::{info, trace}; +use tracing::{info, trace, warn}; use crate::auth::Claims; use crate::hooks::{BootProvider, QuotaPolicy, ToolHost}; -use crate::registry::{AgentEntry, AgentRegistry, MessageQueue}; -use crate::tunnel::broker::StreamBroker; +use crate::registry::agent::AgentRegistryProvider; +use crate::registry::queue::MessageQueueProvider; +use crate::tunnel::broker::StreamBrokerProvider; use crate::tunnel::dispatch::handle_message; const PING_INTERVAL: Duration = Duration::from_secs(20); /// Context passed to message handlers. Carries the agent's claims, -/// the in-memory routing state (`registry`, `stream_broker`, +/// the routing state (`registry`, `stream_broker`, /// `message_queue`), and the three operator hooks the dispatch /// table delegates to (`BootProvider`, `ToolHost`, `QuotaPolicy`). /// @@ -28,9 +31,9 @@ const PING_INTERVAL: Duration = Duration::from_secs(20); /// dispatch runs the JWT has already been validated. pub struct TunnelContext { pub claims: Claims, - pub registry: AgentRegistry, - pub message_queue: MessageQueue, - pub stream_broker: Arc, + pub registry: Arc, + pub message_queue: Arc, + pub stream_broker: Arc, pub boot_provider: Arc, pub tool_host: Arc, pub quota_policy: Arc, @@ -45,28 +48,42 @@ pub async fn run_tunnel(ws: WebSocket, claims: Claims, ctx: Arc) let (tx, mut rx) = mpsc::unbounded_channel::(); - // Register the agent in the routing table. - ctx.registry - .register(agent_id, AgentEntry::new(tx, claims.clone())); + // Keep a local clone of the sender for dispatch responses and + // buffered-frame delivery. The original `tx` is moved into the + // registry via `acquire`. + let local_tx = tx.clone(); + + // Acquire a tunnel lease (registers the agent in the routing + // table; in Redis mode also takes the ownership lock and starts + // the heartbeat + inbox consumer). + let lease = match ctx.registry.acquire(agent_id, tx, claims.clone()).await { + Ok(lease) => lease, + Err(e) => { + warn!(%agent_id, "failed to acquire tunnel lease: {e}"); + return; + } + }; // Drain any session/message pushes that arrived while this // agent's tunnel was disconnected. The drain happens before we // yield to the read loop so the agent sees the buffered frames // first. - let buffered = ctx.message_queue.drain(agent_id); + let buffered = ctx.message_queue.drain(agent_id).await; if !buffered.is_empty() { info!( %agent_id, count = buffered.len(), "draining buffered session/message frames on reconnect" ); - if let Some(entry) = ctx.registry.get(agent_id) { - for frame in buffered { - let _ = entry.tx.send(frame); - } + for frame in buffered { + let _ = local_tx.send(frame); } } + // Subscribe to control signals (takeover / disconnect) so + // cross-instance events can close this tunnel. + let mut control_rx = ctx.registry.control_receiver(agent_id); + info!(%agent_id, role = %claims.role, "tunnel connected"); // Periodic WS ping. Long-poll proxies (e.g. cloudflared) close @@ -102,44 +119,62 @@ pub async fn run_tunnel(ws: WebSocket, claims: Claims, ctx: Arc) } }); - // Read loop. - while let Some(Ok(msg)) = ws_rx.next().await { - match msg { - Message::Text(text) => { - trace!(%agent_id, payload = %text, "agent→server"); - if let Some(entry) = ctx.registry.get(agent_id) { - entry.touch(); - } - - // Best-effort fan-out of stream/* and turn/end - // notifications to the in-memory broker so SSE - // subscribers receive them. Cheap string sniff to - // avoid parsing every frame twice. - if text.contains("\"stream/") - || text.contains("\"turn/end\"") - || text.contains("\"heartbeat\"") - { - let sender = ctx.stream_broker.sender_for(agent_id); - let _ = sender.send(text.to_string()); - } - - if let Some(response) = handle_message(&text, &ctx).await { - trace!(%agent_id, payload = %response, "server→agent"); - if let Some(entry) = ctx.registry.get(agent_id) { - let _ = entry.tx.send(response); + // Read loop — also listens for control signals so cross-instance + // takeover / disconnect events can close the tunnel promptly. + loop { + tokio::select! { + ws_msg = ws_rx.next() => { + let Some(Ok(msg)) = ws_msg else { break; }; + match msg { + Message::Text(text) => { + trace!(%agent_id, payload = %text, "agent→server"); + ctx.registry.touch(agent_id).await; + + // Best-effort fan-out of stream/* and turn/end + // notifications to the broker so SSE subscribers + // receive them. Cheap string sniff to avoid + // parsing every frame twice. + if text.contains("\"stream/") + || text.contains("\"turn/end\"") + || text.contains("\"heartbeat\"") + { + ctx.stream_broker + .publish(agent_id, text.to_string()) + .await; + } + + if let Some(response) = handle_message(&text, &ctx).await { + trace!(%agent_id, payload = %response, "server→agent"); + let _ = local_tx.send(response); + } + } + Message::Close(_) => { + info!(%agent_id, "tunnel closed by client"); + break; } + _ => {} } } - Message::Close(_) => { - info!(%agent_id, "tunnel closed by client"); + Some(signal) = recv_control(&mut control_rx) => { + info!(%agent_id, signal = %signal, "control signal, closing tunnel"); break; } - _ => {} } } - ctx.registry.unregister(agent_id); + // Dropping the lease unregisters the agent (in-memory) or + // releases the ownership lock (Redis). + drop(lease); ping_task.abort(); write_task.abort(); info!(%agent_id, "tunnel disconnected"); } + +/// Helper for `tokio::select!`: yields the next control signal, or +/// pends forever when no control receiver is present (in-memory mode). +async fn recv_control(rx: &mut Option>) -> Option { + match rx { + Some(rx) => rx.recv().await, + None => pending().await, + } +} diff --git a/server/tests/dispatch_wireup.rs b/server/tests/dispatch_wireup.rs index ef324b6..a56b33c 100644 --- a/server/tests/dispatch_wireup.rs +++ b/server/tests/dispatch_wireup.rs @@ -20,7 +20,7 @@ use overacp_server::hooks::{ }; use overacp_server::tunnel::dispatch::handle_message; use overacp_server::tunnel::run::TunnelContext; -use overacp_server::tunnel::StreamBroker; +use overacp_server::tunnel::InMemoryStreamBroker; use overacp_server::{AppState, StaticJwtAuthenticator}; use uuid::Uuid; @@ -196,7 +196,7 @@ async fn turn_end_is_fanned_out_to_stream_broker_subscribers() { // forwards them to the per-agent broadcast channel. This test // exercises the broker fan-out directly: if you subscribe to an // agent before pushing a frame, you receive it. - let stream_broker = StreamBroker::new(); + let stream_broker = InMemoryStreamBroker::new(); let agent_id = Uuid::new_v4(); let mut rx = stream_broker.subscribe(agent_id); diff --git a/server/tests/main_env.rs b/server/tests/main_env.rs index cb433e3..f118d4f 100644 --- a/server/tests/main_env.rs +++ b/server/tests/main_env.rs @@ -65,6 +65,7 @@ async fn healthz_status(state: overacp_server::AppState) -> StatusCode { async fn missing_signing_key_errors() { let _g = EnvGuard::new(); let err = build_state_from_env() + .await .err() .expect("should fail without signing key"); assert!(matches!(err, StartupError::MissingSigningKey)); @@ -75,7 +76,7 @@ async fn missing_signing_key_errors() { async fn signing_key_only_uses_default_issuer() { let g = EnvGuard::new(); g.set("OVERACP_JWT_SIGNING_KEY", "k"); - let state = build_state_from_env().expect("should build"); + let state = build_state_from_env().await.expect("should build"); assert_eq!(state.authenticator.issuer(), "overacp"); // And the state produces a live router with a responsive healthz. assert_eq!(healthz_status(state).await, StatusCode::OK); @@ -86,6 +87,6 @@ async fn custom_issuer_is_honoured() { let g = EnvGuard::new(); g.set("OVERACP_JWT_SIGNING_KEY", "k"); g.set("OVERACP_JWT_ISSUER", "custom-issuer"); - let state = build_state_from_env().expect("should build"); + let state = build_state_from_env().await.expect("should build"); assert_eq!(state.authenticator.issuer(), "custom-issuer"); } diff --git a/server/tests/registry.rs b/server/tests/registry.rs index 1b90298..2faf637 100644 --- a/server/tests/registry.rs +++ b/server/tests/registry.rs @@ -24,10 +24,10 @@ fn live_entry() -> (mpsc::UnboundedReceiver, AgentEntry, Claims) { (rx, AgentEntry::new(tx, claims.clone()), claims) } -#[test] -fn registry_default_is_exposed_via_appstate() { +#[tokio::test] +async fn registry_default_is_exposed_via_appstate() { let state = fresh_state(); - assert!(state.registry.list().is_empty()); + assert!(state.registry.list_agents().await.is_empty()); assert_eq!(state.message_queue.capacity(), 64); } @@ -100,37 +100,40 @@ fn message_queue_full_error_is_re_exported_from_crate_root() { assert_eq!(capacity, 1); } -#[test] -fn drain_on_reconnect_via_appstate() { +#[tokio::test] +async fn drain_on_reconnect_via_appstate() { // Smoke: simulate the "agent disconnected, REST pushes, agent // reconnects, drain happens" loop using the public surface. let state = fresh_state(); let agent_id = Uuid::new_v4(); // 1. Push two notifications while disconnected. - state.message_queue.push(agent_id, "first".into()).unwrap(); - state.message_queue.push(agent_id, "second".into()).unwrap(); - assert_eq!(state.message_queue.len(agent_id), 2); - - // 2. Reconnect: register an entry, then drain the queue and - // forward to the entry. - let (mut rx, entry, claims) = live_entry(); - // Use the test agent_id, not the live_entry random one. - let claims = Claims::agent(agent_id, claims.user, 60, "overacp"); - let entry = AgentEntry::new(entry.tx.clone(), claims); - state.registry.register(agent_id, entry); - - let buffered = state.message_queue.drain(agent_id); - let registered = state.registry.get(agent_id).expect("registered"); + state + .message_queue + .push(agent_id, "first".into()) + .await + .unwrap(); + state + .message_queue + .push(agent_id, "second".into()) + .await + .unwrap(); + assert_eq!(state.message_queue.len(agent_id).await, 2); + + // 2. Reconnect: register an entry via acquire, then drain the + // queue and deliver to the agent. + let (tx, mut rx) = mpsc::unbounded_channel::(); + let claims = Claims::agent(agent_id, Some(Uuid::new_v4()), 60, "overacp"); + let _lease = state.registry.acquire(agent_id, tx, claims).await.unwrap(); + + let buffered = state.message_queue.drain(agent_id).await; for frame in buffered { - registered.tx.send(frame).unwrap(); + let _ = state.registry.deliver(agent_id, frame).await; } // 3. The reconnecting agent receives both buffered frames in - // order before any live traffic. The channel is already - // populated synchronously by the for-loop above, so - // `try_recv` is sufficient — no need to spin up tokio. + // order before any live traffic. assert_eq!(rx.try_recv().ok(), Some("first".to_string())); assert_eq!(rx.try_recv().ok(), Some("second".to_string())); - assert!(state.message_queue.is_empty(agent_id)); + assert!(state.message_queue.is_empty(agent_id).await); } diff --git a/server/tests/rest_surface.rs b/server/tests/rest_surface.rs index 2c6dbab..b1ea772 100644 --- a/server/tests/rest_surface.rs +++ b/server/tests/rest_surface.rs @@ -25,8 +25,8 @@ use axum::body::Body; use axum::http::{Method, Request, StatusCode}; use http_body_util::BodyExt; use overacp_server::auth::{Authenticator, Claims}; -use overacp_server::registry::{AgentEntry, MessageQueue}; -use overacp_server::{router, AppState, StaticJwtAuthenticator}; +use overacp_server::registry::MessageQueue; +use overacp_server::{router, AppState, StaticJwtAuthenticator, TunnelLease}; use serde_json::{json, Value}; use tokio::sync::mpsc; use tower::ServiceExt; @@ -44,7 +44,7 @@ fn fresh_state() -> (AppState, Arc) { fn state_with_queue_capacity(cap: usize) -> (AppState, Arc) { let (state, auth) = fresh_state(); let state = AppState { - message_queue: MessageQueue::new(cap), + message_queue: Arc::new(MessageQueue::new(cap)), ..state }; (state, auth) @@ -61,14 +61,15 @@ fn mint_agent_for(auth: &Arc, agent_id: Uuid) -> String { } /// Register a fake tunnel entry in the registry and return its -/// receiving channel. Mirrors the `run_tunnel` registration path. -fn register_fake(state: &AppState, agent_id: Uuid) -> mpsc::UnboundedReceiver { +/// receiving channel plus the lease guard. +async fn register_fake( + state: &AppState, + agent_id: Uuid, +) -> (mpsc::UnboundedReceiver, TunnelLease) { let (tx, rx) = mpsc::unbounded_channel::(); let claims = Claims::agent(agent_id, None, 300, ISSUER); - state - .registry - .register(agent_id, AgentEntry::new(tx, claims)); - rx + let lease = state.registry.acquire(agent_id, tx, claims).await.unwrap(); + (rx, lease) } async fn oneshot(app: axum::Router, req: Request) -> (StatusCode, Value) { @@ -218,7 +219,7 @@ async fn describe_agent_is_admin_only_per_spec() { // hit the agent-scoped streaming endpoints instead. let (state, auth) = fresh_state(); let agent_id = Uuid::new_v4(); - let _rx = register_fake(&state, agent_id); + let (_rx, _lease) = register_fake(&state, agent_id).await; let app = router(state); // Admin can describe any agent. @@ -271,7 +272,7 @@ async fn describe_unknown_agent_returns_404_for_admin() { async fn delete_agent_admin_only_and_disconnects() { let (state, auth) = fresh_state(); let agent_id = Uuid::new_v4(); - let _rx = register_fake(&state, agent_id); + let (_rx, _lease) = register_fake(&state, agent_id).await; let admin_jwt = mint_admin(&auth); let app = router(state.clone()); @@ -282,7 +283,7 @@ async fn delete_agent_admin_only_and_disconnects() { ) .await; assert_eq!(status, StatusCode::ACCEPTED); - assert!(!state.registry.is_connected(agent_id)); + assert!(!state.registry.is_connected(agent_id).await); } #[tokio::test] @@ -292,7 +293,7 @@ async fn delete_agent_rejects_agent_tokens() { // per SPEC.md § "Route authorization". let (state, auth) = fresh_state(); let agent_id = Uuid::new_v4(); - let _rx = register_fake(&state, agent_id); + let (_rx, _lease) = register_fake(&state, agent_id).await; let app = router(state.clone()); // Own token → 403. @@ -315,7 +316,7 @@ async fn delete_agent_rejects_agent_tokens() { // And the registry entry is still alive — neither attempt // took effect. - assert!(state.registry.is_connected(agent_id)); + assert!(state.registry.is_connected(agent_id).await); } // ══════════════════════════════════════════════════════════════ @@ -326,7 +327,7 @@ async fn delete_agent_rejects_agent_tokens() { async fn send_message_live_delivery_via_admin_token() { let (state, auth) = fresh_state(); let agent_id = Uuid::new_v4(); - let mut rx = register_fake(&state, agent_id); + let (mut rx, _lease) = register_fake(&state, agent_id).await; let admin_jwt = mint_admin(&auth); let app = router(state); @@ -365,7 +366,7 @@ async fn send_message_queues_when_agent_disconnected() { let (status, body) = oneshot(app, req).await; assert_eq!(status, StatusCode::ACCEPTED); assert_eq!(body["delivery"], "queued"); - assert_eq!(state.message_queue.len(agent_id), 1); + assert_eq!(state.message_queue.len(agent_id).await, 1); } #[tokio::test] @@ -409,7 +410,7 @@ async fn send_message_returns_503_on_queue_overflow() { async fn send_message_agent_token_scoped_to_own_id() { let (state, auth) = fresh_state(); let agent_id = Uuid::new_v4(); - let _rx = register_fake(&state, agent_id); + let (_rx, _lease) = register_fake(&state, agent_id).await; let app = router(state); // Own token works. @@ -467,7 +468,7 @@ async fn cancel_when_disconnected_returns_202() { async fn cancel_agent_token_scoped_to_own_id() { let (state, auth) = fresh_state(); let agent_id = Uuid::new_v4(); - let _rx = register_fake(&state, agent_id); + let (_rx, _lease) = register_fake(&state, agent_id).await; let app = router(state); // Own token → 202. @@ -499,7 +500,7 @@ async fn cancel_agent_token_scoped_to_own_id() { async fn cancel_when_connected_sends_session_cancel() { let (state, auth) = fresh_state(); let agent_id = Uuid::new_v4(); - let mut rx = register_fake(&state, agent_id); + let (mut rx, _lease) = register_fake(&state, agent_id).await; let admin_jwt = mint_admin(&auth); let app = router(state); let (status, _) = oneshot( diff --git a/server/tests/tunnel_e2e.rs b/server/tests/tunnel_e2e.rs index 08cd054..b09ac3f 100644 --- a/server/tests/tunnel_e2e.rs +++ b/server/tests/tunnel_e2e.rs @@ -256,7 +256,7 @@ async fn reconnect_drains_buffered_session_messages() { assert_eq!(parsed["delivery"], "queued"); } } - assert_eq!(state.message_queue.len(agent_id), 2); + assert_eq!(state.message_queue.len(agent_id).await, 2); // Now start the server and open the tunnel. The write loop // drains the queue before serving live traffic. @@ -272,7 +272,7 @@ async fn reconnect_drains_buffered_session_messages() { assert!(second.contains(r#""content":"second""#)); // And the queue is empty now. - assert_eq!(state.message_queue.len(agent_id), 0); + assert_eq!(state.message_queue.len(agent_id).await, 0); ws.close(None).await.ok(); } @@ -301,7 +301,7 @@ async fn turn_end_fans_out_to_sse_subscribers() { .await .unwrap(); - let fanned = timeout(Duration::from_secs(2), rx.recv()) + let fanned = timeout(Duration::from_secs(2), rx.next()) .await .expect("broker fan-out timeout") .expect("broker fan-out closed"); @@ -410,7 +410,7 @@ async fn multimodal_turn_end_fans_out_to_sse_with_blocks_intact() { .await .unwrap(); - let fanned = timeout(Duration::from_secs(2), rx.recv()) + let fanned = timeout(Duration::from_secs(2), rx.next()) .await .expect("fan-out timeout") .expect("fan-out closed"); diff --git a/tests/smoke-ha.sh b/tests/smoke-ha.sh new file mode 100755 index 0000000..0710acc --- /dev/null +++ b/tests/smoke-ha.sh @@ -0,0 +1,234 @@ +#!/usr/bin/env bash +# +# tests/smoke-ha.sh — two-instance HA smoke test for overacp-server. +# +# Starts a dockerized Valkey instance, then two overacp-server binaries +# on different ports sharing the same Valkey. Proves: +# +# 1. Cross-instance push: POST to server B, agent on server A receives. +# 2. Cross-instance SSE: SSE subscriber on B sees frames from A. +# 3. Failover: kill server A, reconnect agent to B, messages route. +# +# Dependencies on $PATH: +# - docker, curl, jq, uuidgen, openssl +# - python3 with PyJWT (`pip install PyJWT`) +# - websocat (`cargo install websocat`) +# +# Environment: +# - TARGET_DIR (optional, default ./target/debug) +# +# Exit 0 on full success, non-zero on any failure. + +set -euo pipefail + +TARGET_DIR=${TARGET_DIR:-./target/debug} +VALKEY_PORT=6399 +ALPHA_PORT=8181 +BETA_PORT=8182 +LOG_ALPHA=/tmp/overacp-ha-alpha.log +LOG_BETA=/tmp/overacp-ha-beta.log + +# ── plumbing ───────────────────────────────────────────────────── + +skip() { echo "SKIP: $1"; exit 0; } +need() { command -v "$1" >/dev/null || skip "missing dep: $1"; } +for dep in docker curl jq uuidgen openssl python3 websocat; do need "$dep"; done +python3 -c 'import jwt' 2>/dev/null || skip "missing dep: python3 PyJWT" + +# The binary must be built with --features redis. When compiled +# without the feature, OVERACP_REDIS_URL is silently ignored and +# the server runs in-memory mode — which defeats the entire test. +# We detect this by checking if the binary contains the redis +# enablement log message. +if ! grep -q "enabling redis backend" "$TARGET_DIR/overacp-server" 2>/dev/null; then + skip "binary not built with --features redis (run: cargo build -p overacp-server --features redis)" +fi + +hr() { printf '\n── %s ' "$1"; printf '%.0s─' $(seq 1 $((60 - ${#1}))); echo; } +fail() { echo "FAIL: $1"; exit 1; } + +ALPHA_PID="" BETA_PID="" VALKEY_CID="" +cleanup() { + [ -n "$ALPHA_PID" ] && kill "$ALPHA_PID" 2>/dev/null || true + [ -n "$BETA_PID" ] && kill "$BETA_PID" 2>/dev/null || true + [ -n "$VALKEY_CID" ] && docker rm -f "$VALKEY_CID" >/dev/null 2>&1 || true + # Kill any leftover background jobs (websocat, curl) + jobs -p 2>/dev/null | xargs -r kill 2>/dev/null || true + wait 2>/dev/null || true +} +trap cleanup EXIT INT TERM + +# ── 0. start Valkey ───────────────────────────────────────────── + +hr "0. start Valkey" +VALKEY_CID=$(docker run -d --rm -p "$VALKEY_PORT:6379" valkey/valkey:8) +echo "valkey container=$VALKEY_CID port=$VALKEY_PORT" + +# Wait for Valkey to be ready. +for _ in $(seq 1 30); do + docker exec "$VALKEY_CID" valkey-cli ping 2>/dev/null | grep -q PONG && break + sleep 0.2 +done +docker exec "$VALKEY_CID" valkey-cli ping | grep -q PONG || fail "valkey didn't come up" +echo "valkey=PONG" + +# ── 1. start two overacp-server instances ──────────────────────── + +hr "1. start alpha + beta" +export OVERACP_JWT_SIGNING_KEY="$(openssl rand -hex 32)" +REDIS_URL="redis://localhost:$VALKEY_PORT" + +OVERACP_REDIS_URL="$REDIS_URL" \ +OVERACP_INSTANCE_ID=alpha \ +PORT="$ALPHA_PORT" \ + "$TARGET_DIR/overacp-server" > "$LOG_ALPHA" 2>&1 & +ALPHA_PID=$! + +OVERACP_REDIS_URL="$REDIS_URL" \ +OVERACP_INSTANCE_ID=beta \ +PORT="$BETA_PORT" \ + "$TARGET_DIR/overacp-server" > "$LOG_BETA" 2>&1 & +BETA_PID=$! + +echo "alpha pid=$ALPHA_PID port=$ALPHA_PORT" +echo "beta pid=$BETA_PID port=$BETA_PORT" + +for port in $ALPHA_PORT $BETA_PORT; do + for _ in $(seq 1 40); do + curl -fs "http://localhost:$port/healthz" >/dev/null 2>&1 && break + sleep 0.2 + done + curl -fs "http://localhost:$port/healthz" >/dev/null \ + || fail "server on port $port didn't come up" + echo " port=$port healthz=ok" +done + +# ── 2. mint JWTs ──────────────────────────────────────────────── + +hr "2. mint JWTs" +ADMIN_JWT=$(python3 - <<'PY' +import os, jwt, time, uuid +print(jwt.encode({ + "sub": str(uuid.uuid4()), + "role": "admin", + "exp": int(time.time()) + 3600, + "iss": "overacp", +}, os.environ["OVERACP_JWT_SIGNING_KEY"], algorithm="HS256")) +PY +) + +AGENT_ID=$(uuidgen | tr '[:upper:]' '[:lower:]') +MINT_RESP=$(curl -fsS -X POST "http://localhost:$ALPHA_PORT/tokens" \ + -H "Authorization: Bearer $ADMIN_JWT" \ + -H "Content-Type: application/json" \ + -d "{\"agent_id\": \"$AGENT_ID\", \"ttl_secs\": 300}") +AGENT_JWT=$(echo "$MINT_RESP" | jq -r .token) +echo "agent_id=$AGENT_ID" + +# ── 3. connect agent to ALPHA ────────────────────────────────── + +hr "3. connect agent to ALPHA" +# Agent echoes any session/message back as stream/textDelta, then +# logs the raw frame to a file for assertion. +AGENT_OUT=/tmp/overacp-ha-agent.out +: > "$AGENT_OUT" + +( + websocat -tE \ + -H="Authorization: Bearer $AGENT_JWT" \ + "ws://localhost:$ALPHA_PORT/tunnel/$AGENT_ID" 2>/dev/null \ + | while IFS= read -r line; do + echo "$line" >> "$AGENT_OUT" + # Echo stream/textDelta for any session/message received. + if echo "$line" | jq -e '.method == "session/message"' >/dev/null 2>&1; then + echo '{"jsonrpc":"2.0","method":"stream/textDelta","params":{"delta":"echo"}}' + fi + done +) & +AGENT_WS_PID=$! +sleep 0.5 + +# Verify ALPHA sees the agent. +curl -fsS -H "Authorization: Bearer $ADMIN_JWT" \ + "http://localhost:$ALPHA_PORT/agents" | jq -e ".agents | length == 1" >/dev/null \ + || fail "agent not registered on alpha" +echo "agent registered on alpha" + +# ── 4. CROSS-INSTANCE PUSH ───────────────────────────────────── + +hr "4. cross-instance push: POST to BETA, agent on ALPHA" +curl -fsS -X POST "http://localhost:$BETA_PORT/agents/$AGENT_ID/messages" \ + -H "Authorization: Bearer $ADMIN_JWT" \ + -H "Content-Type: application/json" \ + -d '{"role":"user","content":{"text":"cross-instance hello"}}' | jq -c . + +sleep 1 +grep -q "cross-instance hello" "$AGENT_OUT" \ + || fail "cross-instance message did not reach agent on alpha" +echo "cross-instance push: OK" + +# ── 5. CROSS-INSTANCE SSE ────────────────────────────────────── + +hr "5. cross-instance SSE: subscribe via BETA" +SSE_OUT=/tmp/overacp-ha-sse.out +: > "$SSE_OUT" + +timeout 4 curl -sNf \ + -H "Authorization: Bearer $ADMIN_JWT" \ + "http://localhost:$BETA_PORT/agents/$AGENT_ID/stream" > "$SSE_OUT" 2>/dev/null & +SSE_PID=$! +sleep 0.5 + +# Trigger a stream/textDelta by pushing a message via ALPHA. +curl -fsS -X POST "http://localhost:$ALPHA_PORT/agents/$AGENT_ID/messages" \ + -H "Authorization: Bearer $ADMIN_JWT" \ + -H "Content-Type: application/json" \ + -d '{"role":"user","content":{"text":"trigger echo"}}' >/dev/null + +sleep 1 +wait $SSE_PID 2>/dev/null || true + +grep -q "textDelta" "$SSE_OUT" \ + || fail "SSE on BETA did not receive ALPHA's stream fan-out" +echo "cross-instance SSE: OK" + +# ── 6. FAILOVER ──────────────────────────────────────────────── + +hr "6. failover: kill ALPHA, reconnect to BETA" +kill "$ALPHA_PID" && wait "$ALPHA_PID" 2>/dev/null || true +ALPHA_PID="" +kill "$AGENT_WS_PID" 2>/dev/null || true +echo "alpha killed" + +# Wait for lease to expire or be released. +sleep 2 + +# Reconnect agent to BETA. +AGENT_OUT2=/tmp/overacp-ha-agent2.out +: > "$AGENT_OUT2" + +( + websocat -tE \ + -H="Authorization: Bearer $AGENT_JWT" \ + "ws://localhost:$BETA_PORT/tunnel/$AGENT_ID" 2>/dev/null \ + | while IFS= read -r line; do + echo "$line" >> "$AGENT_OUT2" + done +) & +sleep 0.5 + +# Push via BETA (which is now the only instance). +curl -fsS -X POST "http://localhost:$BETA_PORT/agents/$AGENT_ID/messages" \ + -H "Authorization: Bearer $ADMIN_JWT" \ + -H "Content-Type: application/json" \ + -d '{"role":"user","content":{"text":"failed over"}}' | jq -c . + +sleep 1 +grep -q "failed over" "$AGENT_OUT2" \ + || fail "failover to BETA did not route new message" +echo "failover: OK" + +# ── done ──────────────────────────────────────────────────────── + +hr "RESULT" +echo "OK: cross-instance push, SSE fan-out, and failover all work."