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 <noreply@anthropic.com>
This commit is contained in:
Meta Vens
parent
76e4e4f269
commit
b26ee5e676
3 changed files with 1447 additions and 0 deletions
|
|
@ -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::*;
|
||||
|
|
|
|||
649
crates/model-connector/src/serve.rs
Normal file
649
crates/model-connector/src/serve.rs
Normal file
|
|
@ -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<String>,
|
||||
pub max_concurrent_requests: usize,
|
||||
pub max_payload_bytes: usize,
|
||||
pub max_output_tokens: u32,
|
||||
}
|
||||
|
||||
pub fn required_connection_capacity(callers: &[CallerConfig]) -> Result<usize, ServeError> {
|
||||
Ok(caller_request_capacity(callers)?.max(8))
|
||||
}
|
||||
|
||||
pub fn caller_request_capacity(callers: &[CallerConfig]) -> Result<usize, ServeError> {
|
||||
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<Vec<CallerAdmission>, 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<CallerAdmission>,
|
||||
replay_store: &Path,
|
||||
backend: Arc<dyn ProviderBackend>,
|
||||
) -> 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<TransportService>,
|
||||
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<TransportService>,
|
||||
backend: &dyn ProviderBackend,
|
||||
) -> Result<Vec<u8>, 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<u64, ServeError> {
|
||||
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<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl ConnectionPermit {
|
||||
fn acquire(active: Arc<AtomicUsize>, limit: usize) -> Option<Self> {
|
||||
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<Self> {
|
||||
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<CallerAdmission>,
|
||||
path: &Path,
|
||||
backend: Arc<dyn ProviderBackend>,
|
||||
) -> Option<Result<(), ServeError>> {
|
||||
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<CountingBackend>,
|
||||
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<CallerConfig> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
788
crates/model-connector/src/service.rs
Normal file
788
crates/model-connector/src/service.rs
Normal file
|
|
@ -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<String>,
|
||||
max_concurrent_requests: usize,
|
||||
max_payload_bytes: usize,
|
||||
max_output_tokens: u32,
|
||||
}
|
||||
|
||||
impl CallerAdmission {
|
||||
pub fn new(
|
||||
caller_runtime_id: impl Into<String>,
|
||||
connection_key: impl Into<String>,
|
||||
connection_key_epoch: u32,
|
||||
allowed_models: impl IntoIterator<Item = String>,
|
||||
max_concurrent_requests: usize,
|
||||
max_payload_bytes: usize,
|
||||
max_output_tokens: u32,
|
||||
) -> Result<Self, ServiceError> {
|
||||
let caller_runtime_id = caller_runtime_id.into();
|
||||
require_id(&caller_runtime_id, "caller_runtime_id")?;
|
||||
let allowed_models = allowed_models.into_iter().collect::<HashSet<_>>();
|
||||
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<ExecutionClaim>),
|
||||
Reply(TransportEnvelope),
|
||||
}
|
||||
|
||||
pub struct TransportService {
|
||||
callers: HashMap<String, CallerAdmission>,
|
||||
live: HashSet<RequestIdentity>,
|
||||
replay_store: CultCache,
|
||||
max_expiry_skew_ms: u64,
|
||||
}
|
||||
|
||||
impl TransportService {
|
||||
pub fn open(
|
||||
replay_store_path: &Path,
|
||||
admissions: impl IntoIterator<Item = CallerAdmission>,
|
||||
max_expiry_skew_ms: u64,
|
||||
) -> Result<Self, ServiceError> {
|
||||
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::<ReplayRecord>()
|
||||
.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::<ReplayRecord>()
|
||||
.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<TransportAdmission, ServiceError> {
|
||||
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::<ReplayRecord>(&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<TransportEnvelope, ServiceError> {
|
||||
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::<ReplayRecord>(&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<TransportAdmission, ServiceError> {
|
||||
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::<ReplayRecord>().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::<ReplayRecord>(&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
|
||||
);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue