From b26ee5e676df5c4986e0a8cae77221550ef63586 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Wed, 28 Oct 2026 04:40:15 +0100 Subject: [PATCH] Lift the replay service and serve loop behind a ProviderBackend trait service.rs is CodexConnector lib.rs 954-1313 verbatim under the rename rule. serve.rs is the accept loop of daemon.rs over ProviderBackend (ready, execute), ServeLimits, CallerConfig and load_admissions; no config document, no provider, no Idunn route observation. serve refuses a non-loopback listener and drops non-loopback peers. ready() is consulted once, before replay state is opened and before the first accept, as base consulted readiness at startup; execute is called on the request path without consulting it (base only checked readiness for health). New tests: end-to-end replay with a counting backend, readiness timing, limits, peer and listener refusal, connection permits, and behavioural kills for survivors in service policy, payload bound, refusal mapping, restart record validation and completion mismatch. Co-Authored-By: Claude Opus 5.5 --- crates/model-connector/src/lib.rs | 10 + crates/model-connector/src/serve.rs | 649 +++++++++++++++++++++ crates/model-connector/src/service.rs | 788 ++++++++++++++++++++++++++ 3 files changed, 1447 insertions(+) create mode 100644 crates/model-connector/src/serve.rs create mode 100644 crates/model-connector/src/service.rs diff --git a/crates/model-connector/src/lib.rs b/crates/model-connector/src/lib.rs index f288476..ba1642b 100644 --- a/crates/model-connector/src/lib.rs +++ b/crates/model-connector/src/lib.rs @@ -5,6 +5,16 @@ mod envelope; #[cfg(test)] mod fixtures; +#[cfg(feature = "daemon")] +mod serve; +#[cfg(feature = "daemon")] +mod service; + pub use client::*; pub use contract::*; pub use envelope::*; + +#[cfg(feature = "daemon")] +pub use serve::*; +#[cfg(feature = "daemon")] +pub use service::*; diff --git a/crates/model-connector/src/serve.rs b/crates/model-connector/src/serve.rs new file mode 100644 index 0000000..1b613f5 --- /dev/null +++ b/crates/model-connector/src/serve.rs @@ -0,0 +1,649 @@ +use std::fs; +use std::net::{IpAddr, TcpListener, TcpStream}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::thread; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use zeroize::Zeroizing; + +use crate::contract::{TransportInvocation, TransportResult}; +use crate::envelope::{ + ServiceError, TransportEnvelope, TransportFrameError, read_transport_frame, + write_transport_frame, +}; +use crate::service::{CallerAdmission, TransportAdmission, TransportService}; + +/// The only provider seam of the connector kit. A connector binary implements +/// it over one model provider; the kit names no provider. +pub trait ProviderBackend: Send + Sync { + /// Whether the provider can take work. `serve` consults it once, before it + /// opens the replay store or accepts a connection; the request path does + /// not consult it. + fn ready(&self) -> bool; + /// Executes one admitted invocation. The service has already bound the + /// invocation to its caller and recorded its replay claim. + fn execute(&self, invocation: &TransportInvocation) -> TransportResult; +} + +/// Bounds of one serve loop. The connector binary's own document owns where +/// they come from. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ServeLimits { + pub max_frame_bytes: usize, + pub max_connections: usize, + pub socket_timeout_ms: u64, + pub max_expiry_skew_ms: u64, +} + +impl ServeLimits { + fn validate(&self) -> Result<(), ServeError> { + if self.max_frame_bytes == 0 + || self.max_frame_bytes > u32::MAX as usize + || self.max_connections == 0 + || self.socket_timeout_ms == 0 + || self.max_expiry_skew_ms == 0 + { + return Err(ServeError::InvalidConfig("serve limits")); + } + Ok(()) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct CallerConfig { + pub caller_runtime_id: String, + pub connection_key_file: PathBuf, + pub connection_key_epoch: u32, + pub allowed_models: Vec, + pub max_concurrent_requests: usize, + pub max_payload_bytes: usize, + pub max_output_tokens: u32, +} + +pub fn required_connection_capacity(callers: &[CallerConfig]) -> Result { + Ok(caller_request_capacity(callers)?.max(8)) +} + +pub fn caller_request_capacity(callers: &[CallerConfig]) -> Result { + callers + .iter() + .try_fold(0_usize, |total, caller| { + total.checked_add(caller.max_concurrent_requests) + }) + .ok_or(ServeError::InvalidConfig("connection capacity")) +} + +pub fn load_admissions(callers: &[CallerConfig]) -> Result, ServeError> { + let mut admissions = Vec::with_capacity(callers.len()); + for caller in callers { + let raw = Zeroizing::new( + fs::read(&caller.connection_key_file).map_err(ServeError::ConnectionKey)?, + ); + let secret = std::str::from_utf8(raw.as_slice()) + .map_err(|_| ServeError::InvalidConfig("caller connection key encoding"))? + .trim_end_matches(['\r', '\n']); + admissions.push( + CallerAdmission::new( + caller.caller_runtime_id.clone(), + secret.to_string(), + caller.connection_key_epoch, + caller.allowed_models.clone(), + caller.max_concurrent_requests, + caller.max_payload_bytes, + caller.max_output_tokens, + ) + .map_err(ServeError::Service)?, + ); + } + Ok(admissions) +} + +/// Serves admitted callers on a loopback listener until the listener fails. +/// The caller binds; a listener that is not bound to loopback is refused, and +/// peers that are not loopback are dropped. +pub fn serve( + listener: TcpListener, + limits: ServeLimits, + admissions: Vec, + replay_store: &Path, + backend: Arc, +) -> Result<(), ServeError> { + limits.validate()?; + if admissions.is_empty() { + return Err(ServeError::InvalidConfig("callers")); + } + if !listener + .local_addr() + .map_err(ServeError::Listen)? + .ip() + .is_loopback() + { + return Err(ServeError::InvalidConfig("listener address")); + } + if !backend.ready() { + return Err(ServeError::BackendUnready); + } + let service = Arc::new(Mutex::new(TransportService::open( + replay_store, + admissions, + limits.max_expiry_skew_ms, + )?)); + let active_connections = Arc::new(AtomicUsize::new(0)); + let timeout = Duration::from_millis(limits.socket_timeout_ms); + + loop { + let (stream, peer) = listener.accept().map_err(ServeError::Listen)?; + if !peer_admitted(peer.ip()) { + continue; + } + let Some(permit) = + ConnectionPermit::acquire(active_connections.clone(), limits.max_connections) + else { + continue; + }; + let service = service.clone(); + let backend = backend.clone(); + let max_frame_bytes = limits.max_frame_bytes; + thread::Builder::new() + .name("model-connector-request".to_string()) + .spawn(move || { + let _permit = permit; + let _ = serve_connection( + stream, + &service, + backend.as_ref(), + max_frame_bytes, + timeout, + ); + }) + .map_err(ServeError::Thread)?; + } +} + +fn peer_admitted(peer: IpAddr) -> bool { + peer.is_loopback() +} + +fn serve_connection( + mut stream: TcpStream, + service: &Mutex, + backend: &dyn ProviderBackend, + max_frame_bytes: usize, + timeout: Duration, +) -> Result<(), ServeError> { + stream + .set_read_timeout(Some(timeout)) + .and_then(|()| stream.set_write_timeout(Some(timeout))) + .map_err(ServeError::Connection)?; + let request = read_transport_frame(&mut stream, max_frame_bytes).map_err(frame_error)?; + let response = execute_inference_frame(&request, service, backend)?; + write_transport_frame(&mut stream, &response, max_frame_bytes).map_err(frame_error) +} + +fn execute_inference_frame( + request: &[u8], + service: &Mutex, + backend: &dyn ProviderBackend, +) -> Result, ServeError> { + let envelope: TransportEnvelope = + rmp_serde::from_slice(request).map_err(|_| ServeError::FrameEncoding)?; + let admission = service + .lock() + .map_err(|_| ServeError::ServicePoisoned)? + .begin(&envelope, unix_ms()?)?; + let response = match admission { + TransportAdmission::Reply(response) => response, + TransportAdmission::Execute(claim) => { + let result = backend.execute(claim.invocation()); + service + .lock() + .map_err(|_| ServeError::ServicePoisoned)? + .complete(*claim, result)? + } + }; + rmp_serde::to_vec(&response).map_err(|_| ServeError::FrameEncoding) +} + +fn frame_error(error: TransportFrameError) -> ServeError { + match error { + TransportFrameError::Connection(error) => ServeError::Connection(error), + TransportFrameError::Size => ServeError::FrameSize, + } +} + +fn unix_ms() -> Result { + let millis = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|_| ServeError::Clock)? + .as_millis(); + u64::try_from(millis).map_err(|_| ServeError::Clock) +} + +struct ConnectionPermit { + active: Arc, +} + +impl ConnectionPermit { + fn acquire(active: Arc, limit: usize) -> Option { + let mut observed = active.load(Ordering::Acquire); + loop { + if observed >= limit { + return None; + } + match active.compare_exchange_weak( + observed, + observed + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Some(Self { active }), + Err(current) => observed = current, + } + } + } +} + +impl Drop for ConnectionPermit { + fn drop(&mut self) { + self.active.fetch_sub(1, Ordering::AcqRel); + } +} + +#[derive(Debug, Error)] +pub enum ServeError { + #[error("invalid serve configuration field {0}")] + InvalidConfig(&'static str), + #[error("connector connection key is unavailable")] + ConnectionKey(#[source] std::io::Error), + #[error(transparent)] + Service(#[from] ServiceError), + #[error("provider backend is not ready")] + BackendUnready, + #[error("connector service lock was poisoned")] + ServicePoisoned, + #[error("connector listener failed")] + Listen(#[source] std::io::Error), + #[error("connector request thread failed")] + Thread(#[source] std::io::Error), + #[error("connector connection failed")] + Connection(#[source] std::io::Error), + #[error("connector frame exceeded its bound")] + FrameSize, + #[error("connector frame was not valid MessagePack")] + FrameEncoding, + #[error("system clock is before the Unix epoch")] + Clock, +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + use std::net::Ipv4Addr; + + use super::*; + use crate::client::ConnectorClient; + use crate::contract::{ + Refusal, TransportDisposition, TransportEvent, TransportEventPayload, TransportResult, + }; + use crate::fixtures::{TestStore, caller, invocation_for, receipt, request}; + + const CALLER: &str = "epiphany-yggdrasil"; + const KEY: &str = "epiphany-distinct-test-key"; + + struct CountingBackend { + ready: bool, + ready_calls: AtomicUsize, + executions: AtomicUsize, + } + + impl CountingBackend { + fn new(ready: bool) -> Arc { + Arc::new(Self { + ready, + ready_calls: AtomicUsize::new(0), + executions: AtomicUsize::new(0), + }) + } + } + + impl ProviderBackend for CountingBackend { + fn ready(&self) -> bool { + self.ready_calls.fetch_add(1, Ordering::SeqCst); + self.ready + } + + fn execute(&self, invocation: &TransportInvocation) -> TransportResult { + self.executions.fetch_add(1, Ordering::SeqCst); + TransportResult::transported( + invocation, + vec![TransportEvent { + sequence: 0, + payload: TransportEventPayload::TextDelta { + text: "done".to_string(), + }, + }], + receipt(invocation), + ) + } + } + + fn limits() -> ServeLimits { + ServeLimits { + max_frame_bytes: 64 * 1024, + max_connections: 8, + socket_timeout_ms: 5_000, + max_expiry_skew_ms: 300_000, + } + } + + /// The test's own clock, so a mutated `unix_ms` cannot move the fixture with it. + fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + } + + /// What `serve` returned, or `None` when it is still serving after two seconds. + fn outcome( + listener: TcpListener, + limits: ServeLimits, + admissions: Vec, + path: &Path, + backend: Arc, + ) -> Option> { + let (sender, receiver) = std::sync::mpsc::channel(); + let path = path.to_path_buf(); + thread::spawn(move || { + let _ = sender.send(serve(listener, limits, admissions, &path, backend)); + }); + receiver.recv_timeout(Duration::from_secs(2)).ok() + } + + fn fresh(request_id: &str) -> TransportInvocation { + let mut invocation = invocation_for(CALLER, request_id); + invocation.expires_at_unix_ms = now_ms() + 60_000; + invocation + } + + fn start( + backend: Arc, + store: &TestStore, + ) -> (ConnectorClient, std::net::SocketAddr) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let endpoint = listener.local_addr().unwrap(); + let path = store.path.clone(); + thread::spawn(move || { + serve( + listener, + limits(), + vec![caller(CALLER, KEY, 4)], + &path, + backend, + ) + }); + let client = + ConnectorClient::new(endpoint, KEY, 64 * 1024, Some(Duration::from_secs(5))).unwrap(); + (client, endpoint) + } + + fn callers(capacities: &[usize]) -> Vec { + capacities + .iter() + .enumerate() + .map(|(index, capacity)| CallerConfig { + caller_runtime_id: format!("caller-{index}"), + connection_key_file: PathBuf::from(format!("caller-{index}.key")), + connection_key_epoch: 1, + allowed_models: vec!["gpt-5.6-luna".to_string()], + max_concurrent_requests: *capacity, + max_payload_bytes: 1_048_576, + max_output_tokens: 16_384, + }) + .collect() + } + + #[test] + fn serve_end_to_end_replays_without_a_second_execution() { + let store = TestStore::new(); + let backend = CountingBackend::new(true); + let (client, _) = start(backend.clone(), &store); + let invocation = fresh("request-1"); + + let first = client.execute(&invocation).unwrap(); + let second = client.execute(&invocation).unwrap(); + + assert!(matches!( + first.disposition, + TransportDisposition::Transported { .. } + )); + assert_eq!(first, second); + assert_eq!(backend.executions.load(Ordering::SeqCst), 1); + + let mut conflicting_request = request(); + conflicting_request.request_id = "request-1".to_string(); + conflicting_request.instructions = "A different instruction.".to_string(); + let mut conflicting = TransportInvocation::new( + CALLER, + invocation.expires_at_unix_ms, + [7; 32], + conflicting_request, + ) + .unwrap(); + conflicting.expires_at_unix_ms = invocation.expires_at_unix_ms; + let refused = client.execute(&conflicting).unwrap(); + assert_eq!( + refused.disposition, + TransportDisposition::Refused(Refusal::ReplayConflict) + ); + assert_eq!(backend.executions.load(Ordering::SeqCst), 1); + } + + #[test] + fn serve_replay_is_byte_identical_on_the_wire() { + let store = TestStore::new(); + let backend = CountingBackend::new(true); + let (_client, endpoint) = start(backend.clone(), &store); + let security = crate::fixtures::security(KEY); + let request = rmp_serde::to_vec( + &crate::envelope::encrypt_invocation(&fresh("request-wire"), &security).unwrap(), + ) + .unwrap(); + let mut replies = Vec::new(); + for _ in 0..2 { + let mut stream = TcpStream::connect(endpoint).unwrap(); + write_transport_frame(&mut stream, &request, 64 * 1024).unwrap(); + replies.push(read_transport_frame(&mut stream, 64 * 1024).unwrap()); + } + assert_eq!(replies[0], replies[1]); + assert_eq!(backend.executions.load(Ordering::SeqCst), 1); + } + + #[test] + fn serve_refuses_to_start_on_an_unready_backend() { + let store = TestStore::new(); + let backend = CountingBackend::new(false); + let result = outcome( + TcpListener::bind("127.0.0.1:0").unwrap(), + limits(), + vec![caller(CALLER, KEY, 4)], + &store.path, + backend.clone(), + ); + assert!(matches!(result, Some(Err(ServeError::BackendUnready)))); + assert_eq!(backend.executions.load(Ordering::SeqCst), 0); + assert!( + !store.path.exists(), + "an unready backend must not open replay state" + ); + } + + #[test] + fn serve_consults_readiness_at_start_and_never_on_the_request_path() { + let store = TestStore::new(); + let backend = CountingBackend::new(true); + let (client, _) = start(backend.clone(), &store); + client.execute(&fresh("request-a")).unwrap(); + client.execute(&fresh("request-b")).unwrap(); + assert_eq!(backend.executions.load(Ordering::SeqCst), 2); + assert_eq!(backend.ready_calls.load(Ordering::SeqCst), 1); + } + + #[test] + fn serve_refuses_a_listener_that_is_not_loopback() { + let store = TestStore::new(); + let backend = CountingBackend::new(true); + let result = outcome( + TcpListener::bind((Ipv4Addr::UNSPECIFIED, 0)).unwrap(), + limits(), + vec![caller(CALLER, KEY, 4)], + &store.path, + backend.clone(), + ); + assert!(matches!( + result, + Some(Err(ServeError::InvalidConfig("listener address"))) + )); + assert_eq!(backend.ready_calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn serve_refuses_unbounded_limits_and_admits_the_widest_frame() { + let store = TestStore::new(); + let mutations: [fn(&mut ServeLimits); 5] = [ + |limits| limits.max_frame_bytes = 0, + |limits| limits.max_frame_bytes = u32::MAX as usize + 1, + |limits| limits.max_connections = 0, + |limits| limits.socket_timeout_ms = 0, + |limits| limits.max_expiry_skew_ms = 0, + ]; + for mutate in mutations { + let mut bounded = limits(); + mutate(&mut bounded); + let result = outcome( + TcpListener::bind("127.0.0.1:0").unwrap(), + bounded, + vec![caller(CALLER, KEY, 4)], + &store.path, + CountingBackend::new(true), + ); + assert!(matches!( + result, + Some(Err(ServeError::InvalidConfig("serve limits"))) + )); + } + let widest = ServeLimits { + max_frame_bytes: u32::MAX as usize, + ..limits() + }; + let result = outcome( + TcpListener::bind("127.0.0.1:0").unwrap(), + widest, + vec![caller(CALLER, KEY, 4)], + &store.path, + CountingBackend::new(true), + ); + assert!(result.is_none(), "the widest frame the wire can carry must serve"); + } + + #[test] + fn serve_refuses_an_empty_caller_set() { + let store = TestStore::new(); + let result = outcome( + TcpListener::bind("127.0.0.1:0").unwrap(), + limits(), + Vec::new(), + &store.path, + CountingBackend::new(true), + ); + assert!(matches!( + result, + Some(Err(ServeError::InvalidConfig("callers"))) + )); + } + + #[test] + fn only_loopback_peers_are_admitted() { + assert!(peer_admitted(IpAddr::from([127, 0, 0, 1]))); + assert!(peer_admitted(IpAddr::from(std::net::Ipv6Addr::LOCALHOST))); + assert!(!peer_admitted(IpAddr::from([192, 168, 1, 7]))); + assert!(!peer_admitted(IpAddr::from([172, 19, 0, 1]))); + } + + #[test] + fn connection_permits_are_bounded_and_returned_on_drop() { + let active = Arc::new(AtomicUsize::new(0)); + let first = ConnectionPermit::acquire(active.clone(), 2).unwrap(); + let second = ConnectionPermit::acquire(active.clone(), 2).unwrap(); + assert!(ConnectionPermit::acquire(active.clone(), 2).is_none()); + drop(first); + let third = ConnectionPermit::acquire(active.clone(), 2).unwrap(); + assert!(ConnectionPermit::acquire(active.clone(), 2).is_none()); + drop((second, third)); + assert_eq!(active.load(Ordering::SeqCst), 0); + } + + #[test] + fn caller_admission_reserves_every_independently_admitted_request_slot() { + let both = callers(&[8, 8]); + assert_eq!(caller_request_capacity(&both).unwrap(), 16); + assert_eq!(required_connection_capacity(&both).unwrap(), 16); + assert_eq!(required_connection_capacity(&callers(&[3])).unwrap(), 8); + assert!(matches!( + caller_request_capacity(&callers(&[usize::MAX, 1])), + Err(ServeError::InvalidConfig("connection capacity")) + )); + } + + #[test] + fn provider_capacity_is_caller_execution_capacity_not_listener_capacity() { + let callers = callers(&[4]); + assert_eq!(caller_request_capacity(&callers).unwrap(), 4); + assert_eq!(required_connection_capacity(&callers).unwrap(), 8); + } + + #[test] + fn direct_pipe_frame_is_big_endian_and_refuses_oversize_before_allocation() { + let mut encoded = Vec::new(); + write_transport_frame(&mut encoded, b"typed", 16).unwrap(); + assert_eq!(&encoded[..4], &5_u32.to_be_bytes()); + assert_eq!( + read_transport_frame(&mut Cursor::new(encoded), 16).unwrap(), + b"typed" + ); + + let oversized = 17_u32.to_be_bytes().to_vec(); + assert!(matches!( + read_transport_frame(&mut Cursor::new(oversized), 16), + Err(TransportFrameError::Size) + )); + } + + #[test] + fn load_admissions_reads_the_key_file_without_its_trailing_newline() { + let store = TestStore::new(); + let key_file = store.root.join("caller.key"); + fs::write(&key_file, format!("{KEY}\r\n")).unwrap(); + let mut config = callers(&[4]); + config[0].caller_runtime_id = CALLER.to_string(); + config[0].connection_key_file = key_file; + config[0].allowed_models = vec!["gpt-5.4".to_string()]; + config[0].max_payload_bytes = 64 * 1024; + let admissions = load_admissions(&config).unwrap(); + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let endpoint = listener.local_addr().unwrap(); + let backend = CountingBackend::new(true); + let path = store.path.clone(); + let served = backend.clone(); + thread::spawn(move || serve(listener, limits(), admissions, &path, served)); + let client = + ConnectorClient::new(endpoint, KEY, 64 * 1024, Some(Duration::from_secs(5))).unwrap(); + client.execute(&fresh("request-key")).unwrap(); + assert_eq!(backend.executions.load(Ordering::SeqCst), 1); + } +} diff --git a/crates/model-connector/src/service.rs b/crates/model-connector/src/service.rs new file mode 100644 index 0000000..1413537 --- /dev/null +++ b/crates/model-connector/src/service.rs @@ -0,0 +1,788 @@ +use std::collections::{HashMap, HashSet}; +use std::path::Path; + +use cultcache_rs::{CultCache, DatabaseEntry, OwnedRedbMessagePackBackingStore}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use zeroize::Zeroizing; + +use crate::contract::{ + ContractError, Refusal, TransportDisposition, TransportInvocation, TransportKey, + TransportResult, canonical_invocation_bytes, require_id, +}; +use crate::envelope::{ + EnvelopeKind, ServiceError, TransportEnvelope, decrypt_invocation, encrypt_result, + validate_envelope, +}; + +pub struct CallerAdmission { + caller_runtime_id: String, + connection_key_epoch: u32, + security: TransportKey, + allowed_models: HashSet, + max_concurrent_requests: usize, + max_payload_bytes: usize, + max_output_tokens: u32, +} + +impl CallerAdmission { + pub fn new( + caller_runtime_id: impl Into, + connection_key: impl Into, + connection_key_epoch: u32, + allowed_models: impl IntoIterator, + max_concurrent_requests: usize, + max_payload_bytes: usize, + max_output_tokens: u32, + ) -> Result { + let caller_runtime_id = caller_runtime_id.into(); + require_id(&caller_runtime_id, "caller_runtime_id")?; + let allowed_models = allowed_models.into_iter().collect::>(); + if connection_key_epoch == 0 + || allowed_models.is_empty() + || allowed_models + .iter() + .any(|model| require_id(model, "allowed_model").is_err()) + || max_concurrent_requests == 0 + || max_payload_bytes == 0 + || max_output_tokens == 0 + { + return Err(ServiceError::InvalidAdmission); + } + let connection_key = Zeroizing::new(connection_key.into()); + let security = TransportKey::from_connection_secret(connection_key.as_str())?; + Ok(Self { + caller_runtime_id, + connection_key_epoch, + security, + allowed_models, + max_concurrent_requests, + max_payload_bytes, + max_output_tokens, + }) + } + + pub fn caller_runtime_id(&self) -> &str { + &self.caller_runtime_id + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct RequestIdentity { + caller_runtime_id: String, + request_id: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +enum ReplayState { + Active, + Completed { response: TransportEnvelope }, +} + +#[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] +#[cultcache(type = "gamecult.codex.replay_record.v1")] +struct ReplayRecord { + #[cultcache(key = 0)] + caller_runtime_id: String, + #[cultcache(key = 1)] + request_id: String, + #[cultcache(key = 2)] + connection_key_epoch: u32, + #[cultcache(key = 3)] + invocation_sha256: [u8; 32], + #[cultcache(key = 4)] + expires_at_unix_ms: u64, + #[cultcache(key = 5)] + state: ReplayState, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecutionClaim { + invocation: TransportInvocation, + invocation_sha256: [u8; 32], +} + +impl ExecutionClaim { + pub fn invocation(&self) -> &TransportInvocation { + &self.invocation + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TransportAdmission { + Execute(Box), + Reply(TransportEnvelope), +} + +pub struct TransportService { + callers: HashMap, + live: HashSet, + replay_store: CultCache, + max_expiry_skew_ms: u64, +} + +impl TransportService { + pub fn open( + replay_store_path: &Path, + admissions: impl IntoIterator, + max_expiry_skew_ms: u64, + ) -> Result { + if max_expiry_skew_ms == 0 { + return Err(ServiceError::InvalidAdmission); + } + let mut callers = HashMap::new(); + let mut keys = HashSet::new(); + for admission in admissions { + if !keys.insert(admission.security.0) { + return Err(ServiceError::SharedCallerKey); + } + let caller_runtime_id = admission.caller_runtime_id.clone(); + if callers.insert(caller_runtime_id, admission).is_some() { + return Err(ServiceError::DuplicateCaller); + } + } + if callers.is_empty() { + return Err(ServiceError::InvalidAdmission); + } + + let mut replay_store = CultCache::new(); + replay_store + .register_entry_type::() + .map_err(replay_store_error)?; + replay_store.add_generic_backing_store( + OwnedRedbMessagePackBackingStore::new(replay_store_path).map_err(replay_store_error)?, + ); + replay_store + .pull_all_backing_stores() + .map_err(replay_store_error)?; + + for (key, record) in replay_store + .get_all_with_keys::() + .map_err(replay_store_error)? + { + let identity = RequestIdentity { + caller_runtime_id: record.caller_runtime_id.clone(), + request_id: record.request_id.clone(), + }; + if key != replay_key(&identity) + || record.invocation_sha256 == [0; 32] + || record.expires_at_unix_ms == 0 + { + return Err(ServiceError::InvalidReplayRecord); + } + let Some(caller) = callers.get(&record.caller_runtime_id) else { + continue; + }; + if record.connection_key_epoch != caller.connection_key_epoch { + return Err(ServiceError::ReplayCallerKeyMismatch); + } + } + Ok(Self { + callers, + live: HashSet::new(), + replay_store, + max_expiry_skew_ms, + }) + } + + pub fn begin( + &mut self, + envelope: &TransportEnvelope, + now_unix_ms: u64, + ) -> Result { + validate_envelope(envelope, EnvelopeKind::Invocation)?; + let caller = self + .callers + .get(&envelope.caller_runtime_id) + .ok_or(ServiceError::CallerNotAdmitted)?; + if envelope.ciphertext.len() > caller.max_payload_bytes { + return Err(ServiceError::PayloadTooLarge); + } + let security = caller.security.clone(); + let max_concurrent_requests = caller.max_concurrent_requests; + let invocation = decrypt_invocation(envelope, &security)?; + if let Err(error) = invocation.validate(now_unix_ms, self.max_expiry_skew_ms) { + return self.reply_refusal(&invocation, contract_refusal(&error), &security); + } + if !caller.allowed_models.contains(&invocation.request.model) + || invocation + .request + .max_output_tokens + .is_some_and(|limit| limit > caller.max_output_tokens) + { + return self.reply_refusal(&invocation, Refusal::Policy, &security); + } + + let identity = RequestIdentity { + caller_runtime_id: invocation.caller_runtime_id.clone(), + request_id: invocation.request_id().to_string(), + }; + let invocation_sha256: [u8; 32] = + Sha256::digest(canonical_invocation_bytes(&invocation)?).into(); + if let Some(record) = self + .replay_store + .get::(&replay_key(&identity)) + .map_err(replay_store_error)? + { + if record.invocation_sha256 != invocation_sha256 { + return self.reply_refusal(&invocation, Refusal::ReplayConflict, &security); + } + return match record.state { + ReplayState::Completed { response } => { + Ok(TransportAdmission::Reply(response)) + } + ReplayState::Active => { + let refusal = if self.live.contains(&identity) { + Refusal::InFlight + } else { + Refusal::Indeterminate + }; + self.reply_refusal(&invocation, refusal, &security) + } + }; + } + let active_for_caller = self + .live + .iter() + .filter(|key| key.caller_runtime_id == invocation.caller_runtime_id) + .count(); + if active_for_caller >= max_concurrent_requests { + return self.reply_refusal(&invocation, Refusal::Capacity, &security); + } + + let replay_record = ReplayRecord { + caller_runtime_id: identity.caller_runtime_id.clone(), + request_id: identity.request_id.clone(), + connection_key_epoch: caller.connection_key_epoch, + invocation_sha256, + expires_at_unix_ms: invocation.expires_at_unix_ms, + state: ReplayState::Active, + }; + self.replay_store + .put(replay_key(&identity), &replay_record) + .map_err(replay_store_error)?; + self.live.insert(identity); + Ok(TransportAdmission::Execute(Box::new( + ExecutionClaim { + invocation, + invocation_sha256, + }, + ))) + } + + pub fn complete( + &mut self, + claim: ExecutionClaim, + result: TransportResult, + ) -> Result { + if !matches!( + &result.disposition, + TransportDisposition::Transported { .. } + ) { + return Err(ServiceError::InvalidCompletion); + } + result.validate_against(&claim.invocation)?; + let identity = RequestIdentity { + caller_runtime_id: claim.invocation.caller_runtime_id.clone(), + request_id: claim.invocation.request_id().to_string(), + }; + if !self.live.contains(&identity) { + return Err(ServiceError::NoActiveRequest); + } + let active = self + .replay_store + .get::(&replay_key(&identity)) + .map_err(replay_store_error)? + .ok_or(ServiceError::NoActiveRequest)?; + if active.invocation_sha256 != claim.invocation_sha256 + || !matches!(active.state, ReplayState::Active) + { + return Err(ServiceError::ActiveRequestMismatch); + } + let caller = self + .callers + .get(&identity.caller_runtime_id) + .ok_or(ServiceError::CallerNotAdmitted)?; + let response = encrypt_result(&result, &caller.security)?; + let replay_record = ReplayRecord { + caller_runtime_id: identity.caller_runtime_id.clone(), + request_id: identity.request_id.clone(), + connection_key_epoch: caller.connection_key_epoch, + invocation_sha256: claim.invocation_sha256, + expires_at_unix_ms: claim.invocation.expires_at_unix_ms, + state: ReplayState::Completed { + response: response.clone(), + }, + }; + self.replay_store + .put(replay_key(&identity), &replay_record) + .map_err(replay_store_error)?; + self.live.remove(&identity); + Ok(response) + } + + fn reply_refusal( + &self, + invocation: &TransportInvocation, + refusal: Refusal, + security: &TransportKey, + ) -> Result { + Ok(TransportAdmission::Reply(encrypt_result( + &TransportResult::refused(invocation, refusal), + security, + )?)) + } +} + +fn replay_key(identity: &RequestIdentity) -> String { + let mut hasher = Sha256::new(); + hasher.update((identity.caller_runtime_id.len() as u64).to_be_bytes()); + hasher.update(identity.caller_runtime_id.as_bytes()); + hasher.update((identity.request_id.len() as u64).to_be_bytes()); + hasher.update(identity.request_id.as_bytes()); + const HEX: &[u8; 16] = b"0123456789abcdef"; + let digest = hasher.finalize(); + let mut key = String::with_capacity(64); + for byte in digest { + key.push(HEX[(byte >> 4) as usize] as char); + key.push(HEX[(byte & 0x0f) as usize] as char); + } + key +} + +fn replay_store_error(error: impl std::fmt::Display) -> ServiceError { + ServiceError::ReplayStore(error.to_string()) +} + +fn contract_refusal(error: &ContractError) -> Refusal { + match error { + ContractError::Expiry => Refusal::Expired, + ContractError::Identity => Refusal::IdentitySubstitution, + ContractError::ProviderDigest => Refusal::ProviderDigestSubstitution, + _ => Refusal::Malformed, + } +} + + +#[cfg(all(test, feature = "daemon"))] +mod tests { + use super::*; + use crate::fixtures::*; + use crate::{contract::*, envelope::*}; + + #[test] + fn service_refuses_shared_keys_and_isolates_caller_request_identities() { + let store = TestStore::new(); + assert!(matches!( + TransportService::open( + &store.path, + [ + caller("epiphany-yggdrasil", "shared-key", 1), + caller("ghostlight-yggdrasil", "shared-key", 1), + ], + 5_000, + ), + Err(ServiceError::SharedCallerKey) + )); + + let epiphany = invocation_for("epiphany-yggdrasil", "shared-request-id"); + let ghostlight = invocation_for("ghostlight-yggdrasil", "shared-request-id"); + let mut service = open_service(&store, 1); + assert!(matches!( + service + .begin( + &encrypt_invocation(&epiphany, &security("epiphany-distinct-test-key")) + .unwrap(), + 1_000, + ) + .unwrap(), + TransportAdmission::Execute(_) + )); + assert!(matches!( + service + .begin( + &encrypt_invocation(&ghostlight, &security("ghostlight-distinct-test-key")) + .unwrap(), + 1_000, + ) + .unwrap(), + TransportAdmission::Execute(_) + )); + } + + #[test] + fn service_returns_exact_completed_replay_and_refuses_conflicting_replay() { + let store = TestStore::new(); + let invocation = invocation(); + let security = security("epiphany-distinct-test-key"); + let mut service = open_service(&store, 1); + let claim = match service + .begin(&encrypt_invocation(&invocation, &security).unwrap(), 1_000) + .unwrap() + { + TransportAdmission::Execute(claim) => claim, + TransportAdmission::Reply(_) => panic!("first invocation must execute"), + }; + let response = service + .complete( + *claim, + TransportResult::transported( + &invocation, + vec![TransportEvent { + sequence: 0, + payload: TransportEventPayload::ToolCall { + call_id: "call-1".to_string(), + name: "read_source".to_string(), + arguments: "{}".to_string(), + }, + }], + receipt(&invocation), + ), + ) + .unwrap(); + drop(service); + let mut service = open_service(&store, 1); + + let replay = match service + .begin(&encrypt_invocation(&invocation, &security).unwrap(), 1_000) + .unwrap() + { + TransportAdmission::Reply(response) => response, + TransportAdmission::Execute(_) => panic!("completed invocation re-executed"), + }; + assert_eq!(replay, response); + let decoded = decrypt_result(&replay, &security, &invocation).unwrap(); + assert!(matches!( + decoded.disposition, + TransportDisposition::Transported { .. } + )); + + let mut changed_request = request(); + changed_request.instructions.push_str(" Different cargo."); + let conflicting = + TransportInvocation::new("epiphany-yggdrasil", 2_000, [7; 32], changed_request) + .unwrap(); + let refusal = match service + .begin(&encrypt_invocation(&conflicting, &security).unwrap(), 1_000) + .unwrap() + { + TransportAdmission::Reply(response) => response, + TransportAdmission::Execute(_) => panic!("conflicting replay executed"), + }; + assert!(matches!( + decrypt_result(&refusal, &security, &conflicting) + .unwrap() + .disposition, + TransportDisposition::Refused(Refusal::ReplayConflict) + )); + } + + #[test] + fn service_refuses_duplicate_inflight_and_per_caller_capacity() { + let store = TestStore::new(); + let first = invocation_for("epiphany-yggdrasil", "request-1"); + let second = invocation_for("epiphany-yggdrasil", "request-2"); + let security = security("epiphany-distinct-test-key"); + let mut service = open_service(&store, 1); + let first_envelope = encrypt_invocation(&first, &security).unwrap(); + assert!(matches!( + service.begin(&first_envelope, 1_000).unwrap(), + TransportAdmission::Execute(_) + )); + + for (invocation, expected) in [ + (&first, Refusal::InFlight), + (&second, Refusal::Capacity), + ] { + let response = match service + .begin(&encrypt_invocation(invocation, &security).unwrap(), 1_000) + .unwrap() + { + TransportAdmission::Reply(response) => response, + TransportAdmission::Execute(_) => panic!("refused invocation executed"), + }; + assert_eq!( + decrypt_result(&response, &security, invocation) + .unwrap() + .disposition, + TransportDisposition::Refused(expected) + ); + } + } + + #[test] + fn restart_refuses_ambiguous_claim_without_consuming_live_capacity() { + let store = TestStore::new(); + let first = invocation_for("epiphany-yggdrasil", "request-1"); + let second = invocation_for("epiphany-yggdrasil", "request-2"); + let security = security("epiphany-distinct-test-key"); + let mut service = open_service(&store, 1); + assert!(matches!( + service + .begin(&encrypt_invocation(&first, &security).unwrap(), 1_000) + .unwrap(), + TransportAdmission::Execute(_) + )); + drop(service); + + let mut restarted = open_service(&store, 1); + let ambiguous = match restarted + .begin(&encrypt_invocation(&first, &security).unwrap(), 1_000) + .unwrap() + { + TransportAdmission::Reply(response) => response, + TransportAdmission::Execute(_) => panic!("ambiguous claim re-executed"), + }; + assert_eq!( + decrypt_result(&ambiguous, &security, &first) + .unwrap() + .disposition, + TransportDisposition::Refused(Refusal::Indeterminate) + ); + assert!(matches!( + restarted + .begin(&encrypt_invocation(&second, &security).unwrap(), 1_000) + .unwrap(), + TransportAdmission::Execute(_) + )); + drop(restarted); + assert!(matches!( + TransportService::open( + &store.path, + [CallerAdmission::new( + "epiphany-yggdrasil", + "rotated-connection-key", + 2, + ["gpt-5.4".to_string()], + 1, + 64 * 1024, + 32_768, + ) + .unwrap()], + 5_000, + ), + Err(ServiceError::ReplayCallerKeyMismatch) + )); + } + + const CALLER: &str = "epiphany-yggdrasil"; + const KEY: &str = "epiphany-distinct-test-key"; + + fn run( + service: &mut TransportService, + request_id: &str, + edit: impl FnOnce(&mut ProviderRequest), + ) -> (TransportAdmission, TransportInvocation) { + let mut request = invocation_for(CALLER, request_id).request; + edit(&mut request); + let invocation = TransportInvocation::new(CALLER, 2_000, [7; 32], request).unwrap(); + let envelope = encrypt_invocation(&invocation, &security(KEY)).unwrap(); + (service.begin(&envelope, 1_000).unwrap(), invocation) + } + + fn refusal_of(admission: TransportAdmission, invocation: &TransportInvocation) -> Refusal { + let TransportAdmission::Reply(envelope) = admission else { + panic!("expected a refusal reply"); + }; + match decrypt_result(&envelope, &security(KEY), invocation) + .unwrap() + .disposition + { + TransportDisposition::Refused(refusal) => refusal, + TransportDisposition::Transported { .. } => panic!("expected a refusal"), + } + } + + fn raw_envelope(invocation: &TransportInvocation) -> TransportEnvelope { + encrypt_message( + &invocation.caller_runtime_id, + invocation.request_id(), + EnvelopeKind::Invocation, + &rmp_serde::to_vec(invocation).unwrap(), + &security(KEY), + ) + .unwrap() + } + + #[test] + fn caller_admission_reports_its_runtime_id() { + assert_eq!( + caller("ghostlight-yggdrasil", "some-key", 1).caller_runtime_id(), + "ghostlight-yggdrasil" + ); + } + + #[test] + fn policy_refuses_unlisted_models_and_output_tokens_over_the_caller_cap() { + let store = TestStore::new(); + let mut service = open_service(&store, 8); + + let (admission, invocation) = run(&mut service, "unlisted-model", |request| { + request.model = "other-model".to_string(); + }); + assert_eq!(refusal_of(admission, &invocation), Refusal::Policy); + + let (admission, invocation) = run(&mut service, "over-cap", |request| { + request.max_output_tokens = Some(32_769); + }); + assert_eq!(refusal_of(admission, &invocation), Refusal::Policy); + + for (id, tokens) in [ + ("at-cap", Some(32_768)), + ("under-cap", Some(100)), + ("unset", None), + ] { + let (admission, _) = run(&mut service, id, |request| { + request.max_output_tokens = tokens; + }); + assert!(matches!(admission, TransportAdmission::Execute(_)), "{id}"); + } + } + + #[test] + fn payload_bound_is_inclusive_of_the_encrypted_size() { + let invocation = invocation_for(CALLER, "sized"); + let envelope = encrypt_invocation(&invocation, &security(KEY)).unwrap(); + let size = envelope.ciphertext.len(); + let admit = |max_payload_bytes: usize| { + let store = TestStore::new(); + let admission = CallerAdmission::new( + CALLER, + KEY, + 1, + ["gpt-5.4".to_string()], + 4, + max_payload_bytes, + 32_768, + ) + .unwrap(); + TransportService::open(&store.path, [admission], 5_000) + .unwrap() + .begin(&envelope, 1_000) + }; + assert!(matches!(admit(size), Ok(TransportAdmission::Execute(_)))); + assert_eq!(admit(size - 1).unwrap_err(), ServiceError::PayloadTooLarge); + } + + #[test] + fn contract_failures_are_refused_with_the_matching_refusal() { + let store = TestStore::new(); + let mut service = open_service(&store, 8); + + let mut expired = invocation_for(CALLER, "expired"); + expired.expires_at_unix_ms = 10; + let reply = service.begin(&raw_envelope(&expired), 1_000).unwrap(); + assert_eq!(refusal_of(reply, &expired), Refusal::Expired); + + let mut substituted = invocation_for(CALLER, "substituted"); + substituted.provider_request_sha256 = [9; 32]; + let reply = service.begin(&raw_envelope(&substituted), 1_000).unwrap(); + assert_eq!( + refusal_of(reply, &substituted), + Refusal::ProviderDigestSubstitution + ); + + let mut malformed = invocation_for(CALLER, "malformed"); + malformed.native_request_sha256 = [0; 32]; + let reply = service.begin(&raw_envelope(&malformed), 1_000).unwrap(); + assert_eq!(refusal_of(reply, &malformed), Refusal::Malformed); + } + + fn seed(path: &Path, key: &str, edit: impl FnOnce(&mut ReplayRecord)) { + let mut record = ReplayRecord { + caller_runtime_id: CALLER.to_string(), + request_id: "seeded".to_string(), + connection_key_epoch: 1, + invocation_sha256: [1; 32], + expires_at_unix_ms: 5_000, + state: ReplayState::Active, + }; + edit(&mut record); + let mut cache = CultCache::new(); + cache.register_entry_type::().unwrap(); + cache.add_generic_backing_store(OwnedRedbMessagePackBackingStore::new(path).unwrap()); + cache.pull_all_backing_stores().unwrap(); + cache.put(key.to_string(), &record).unwrap(); + } + + fn seeded_key() -> String { + replay_key(&RequestIdentity { + caller_runtime_id: CALLER.to_string(), + request_id: "seeded".to_string(), + }) + } + + #[test] + fn restart_refuses_each_malformed_replay_record() { + let control = TestStore::new(); + seed(&control.path, &seeded_key(), |_| {}); + assert!(TransportService::open(&control.path, [caller(CALLER, KEY, 1)], 5_000).is_ok()); + + let cases: [(&str, fn(&mut ReplayRecord)); 3] = [ + ("wrong-key", |_| {}), + ("zero-digest", |record| record.invocation_sha256 = [0; 32]), + ("zero-expiry", |record| record.expires_at_unix_ms = 0), + ]; + for (name, edit) in cases { + let store = TestStore::new(); + let key = if name == "wrong-key" { + "not-the-replay-key".to_string() + } else { + seeded_key() + }; + seed(&store.path, &key, edit); + assert!( + matches!( + TransportService::open(&store.path, [caller(CALLER, KEY, 1)], 5_000), + Err(ServiceError::InvalidReplayRecord) + ), + "{name}" + ); + } + } + + #[test] + fn completion_refuses_a_claim_that_is_not_the_active_request() { + let store = TestStore::new(); + let mut service = open_service(&store, 4); + let (admission, invocation) = run(&mut service, "request-1", |_| {}); + let TransportAdmission::Execute(claim) = admission else { + panic!("first invocation must execute"); + }; + let result = TransportResult::transported( + claim.invocation(), + vec![TransportEvent { + sequence: 0, + payload: TransportEventPayload::TextDelta { + text: "done".to_string(), + }, + }], + receipt(claim.invocation()), + ); + + let mut forged = (*claim).clone(); + forged.invocation_sha256 = [1; 32]; + assert_eq!( + service.complete(forged, result.clone()).unwrap_err(), + ServiceError::ActiveRequestMismatch + ); + + let key = replay_key(&RequestIdentity { + caller_runtime_id: CALLER.to_string(), + request_id: invocation.request_id().to_string(), + }); + let mut record = service + .replay_store + .get::(&key) + .unwrap() + .unwrap(); + record.state = ReplayState::Completed { + response: encrypt_result(&result, &security(KEY)).unwrap(), + }; + service.replay_store.put(key, &record).unwrap(); + assert_eq!( + service.complete(*claim, result).unwrap_err(), + ServiceError::ActiveRequestMismatch + ); + } +}