563 lines
18 KiB
Rust
563 lines
18 KiB
Rust
use std::{
|
|
collections::VecDeque,
|
|
fs,
|
|
sync::{
|
|
Arc,
|
|
atomic::{AtomicU64, Ordering},
|
|
},
|
|
time::{Duration, SystemTime, UNIX_EPOCH},
|
|
};
|
|
|
|
use async_trait::async_trait;
|
|
use mentra::runtime::{ProviderRetry, RunOptions, SqliteRuntimeStore};
|
|
use mentra::{
|
|
AgentConfig, BuiltinProvider, ContentBlock, Message, Role, Runtime,
|
|
agent::{AgentEvent, AgentStatus, RoundContext, RoundDecision, RoundStrategy},
|
|
error::RuntimeError,
|
|
provider::{
|
|
ContentBlockDelta, ContentBlockStart, ModelInfo, Provider, ProviderDescriptor,
|
|
ProviderError, ProviderEvent, ProviderEventStream, ProviderId, Request,
|
|
},
|
|
};
|
|
use reqwest::StatusCode;
|
|
use tokio::sync::{Mutex, broadcast, mpsc};
|
|
|
|
static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(1);
|
|
|
|
enum StreamScript {
|
|
Buffered(Vec<Result<ProviderEvent, ProviderError>>),
|
|
Fail(ProviderError),
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct ScriptedProvider {
|
|
kind: ProviderId,
|
|
models: Vec<ModelInfo>,
|
|
scripts: std::sync::Arc<Mutex<VecDeque<StreamScript>>>,
|
|
requests: std::sync::Arc<Mutex<Vec<Request<'static>>>>,
|
|
}
|
|
|
|
impl ScriptedProvider {
|
|
fn new(
|
|
kind: impl Into<ProviderId>,
|
|
models: Vec<ModelInfo>,
|
|
scripts: Vec<StreamScript>,
|
|
) -> Self {
|
|
Self {
|
|
kind: kind.into(),
|
|
models,
|
|
scripts: std::sync::Arc::new(Mutex::new(VecDeque::from(scripts))),
|
|
requests: std::sync::Arc::new(Mutex::new(Vec::new())),
|
|
}
|
|
}
|
|
|
|
async fn recorded_requests(&self) -> Vec<Request<'static>> {
|
|
self.requests.lock().await.clone()
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for ScriptedProvider {
|
|
fn descriptor(&self) -> ProviderDescriptor {
|
|
ProviderDescriptor::new(self.kind.clone())
|
|
}
|
|
|
|
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
|
Ok(self.models.clone())
|
|
}
|
|
|
|
async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
|
|
self.requests.lock().await.push(request.into_owned());
|
|
match self.scripts.lock().await.pop_front() {
|
|
Some(StreamScript::Buffered(items)) => {
|
|
let (tx, rx) = mpsc::unbounded_channel();
|
|
for item in items {
|
|
tx.send(item)
|
|
.expect("test stream receiver dropped unexpectedly");
|
|
}
|
|
Ok(rx)
|
|
}
|
|
Some(StreamScript::Fail(error)) => Err(error),
|
|
None => panic!("no scripted stream available"),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_streamed_text_turn_emits_events_and_commits_history() {
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
vec![text_stream(&model.id, "Hello")],
|
|
);
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime
|
|
.spawn_with_config(
|
|
"agent",
|
|
model,
|
|
AgentConfig {
|
|
system: Some("system prompt".to_string()),
|
|
..AgentConfig::default()
|
|
},
|
|
)
|
|
.unwrap();
|
|
let mut events = agent.subscribe_events();
|
|
|
|
let message = agent.send(vec![ContentBlock::text("hi")]).await.unwrap();
|
|
|
|
assert_eq!(message, Message::assistant(ContentBlock::text("Hello")));
|
|
assert_eq!(agent.name(), "agent");
|
|
assert_eq!(agent.model(), "model");
|
|
assert_eq!(agent.history().len(), 2);
|
|
assert_eq!(agent.config().system.as_deref(), Some("system prompt"));
|
|
assert_eq!(
|
|
agent.last_message(),
|
|
Some(&Message::assistant(ContentBlock::text("Hello")))
|
|
);
|
|
|
|
let events = collect_events(&mut events);
|
|
assert!(events.contains(&AgentEvent::RunStarted));
|
|
assert!(events.contains(&AgentEvent::TextDelta {
|
|
delta: "Hello".to_string(),
|
|
full_text: "Hello".to_string(),
|
|
}));
|
|
assert!(matches!(events.last(), Some(AgentEvent::RunFinished)));
|
|
|
|
let snapshot = agent.watch_snapshot();
|
|
assert_eq!(snapshot.borrow().status, AgentStatus::Finished);
|
|
assert_eq!(snapshot.borrow().history_len, 2);
|
|
assert!(snapshot.borrow().current_text.is_empty());
|
|
assert!(snapshot.borrow().pending_tool_uses.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_failure_rolls_history_back_and_emits_run_failed() {
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
vec![
|
|
text_stream(&model.id, "ok"),
|
|
erroring_stream(
|
|
vec![ProviderEvent::MessageStarted {
|
|
id: "msg-2".to_string(),
|
|
model: model.id.clone(),
|
|
role: Role::Assistant,
|
|
}],
|
|
ProviderError::MalformedStream("boom".to_string()),
|
|
),
|
|
],
|
|
);
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).unwrap();
|
|
agent.send(vec![ContentBlock::text("first")]).await.unwrap();
|
|
let baseline = agent.history().to_vec();
|
|
let mut events = agent.subscribe_events();
|
|
|
|
let result = agent.send(vec![ContentBlock::text("second")]).await;
|
|
assert!(result.is_err());
|
|
assert_eq!(agent.history(), baseline.as_slice());
|
|
|
|
let events = collect_events(&mut events);
|
|
assert!(matches!(events.last(), Some(AgentEvent::RunFailed { .. })));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_retries_transient_provider_error_before_streaming() {
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
vec![
|
|
failed_request(ProviderError::Http {
|
|
status: StatusCode::SERVICE_UNAVAILABLE,
|
|
body: "offline".to_string(),
|
|
retry_after: None,
|
|
}),
|
|
text_stream(&model.id, "recovered"),
|
|
],
|
|
);
|
|
let provider_handle = provider.clone();
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).expect("spawn agent");
|
|
|
|
let message = agent
|
|
.send(vec![ContentBlock::text("hello")])
|
|
.await
|
|
.expect("send should retry");
|
|
|
|
assert_eq!(message.text(), "recovered");
|
|
assert_eq!(provider_handle.recorded_requests().await.len(), 2);
|
|
assert_eq!(
|
|
agent.last_message(),
|
|
Some(&Message::assistant(ContentBlock::text("recovered")))
|
|
);
|
|
}
|
|
|
|
/// Every delay a run announced before retrying, in the order it waited them.
|
|
///
|
|
/// `RetryAttempt` reports the delay the runner is about to take, so this is the
|
|
/// schedule as it was actually applied rather than as it was configured.
|
|
fn announced_retry_delays(events: &[AgentEvent]) -> Vec<u64> {
|
|
events
|
|
.iter()
|
|
.filter_map(|event| match event {
|
|
AgentEvent::RetryAttempt { next_delay_ms, .. } => Some(*next_delay_ms),
|
|
_ => None,
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Fails `count` times with `error`, then answers `text`.
|
|
fn failing_then_text(count: usize, error: fn() -> ProviderError, text: &str) -> Vec<StreamScript> {
|
|
let mut scripts: Vec<StreamScript> = (0..count).map(|_| failed_request(error())).collect();
|
|
scripts.push(text_stream("model", text));
|
|
scripts
|
|
}
|
|
|
|
fn offline() -> ProviderError {
|
|
ProviderError::Http {
|
|
status: StatusCode::SERVICE_UNAVAILABLE,
|
|
body: "offline".to_string(),
|
|
retry_after: None,
|
|
}
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
async fn a_default_run_waits_the_schedule_mentra_has_always_waited() {
|
|
// Nobody's timing changes until they ask: 500 ms, doubling.
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
failing_then_text(3, offline, "recovered"),
|
|
);
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).expect("spawn agent");
|
|
let mut events = agent.subscribe_events();
|
|
|
|
let message = agent
|
|
.send(vec![ContentBlock::text("hello")])
|
|
.await
|
|
.expect("the run retries and then succeeds");
|
|
|
|
assert_eq!(message.text(), "recovered");
|
|
assert_eq!(
|
|
announced_retry_delays(&collect_events(&mut events)),
|
|
vec![500, 1_000, 2_000]
|
|
);
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
async fn a_host_schedule_replaces_the_default_delays() {
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
failing_then_text(3, offline, "recovered"),
|
|
);
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).expect("spawn agent");
|
|
let mut events = agent.subscribe_events();
|
|
|
|
let message = agent
|
|
.run(
|
|
vec![ContentBlock::text("hello")],
|
|
RunOptions::default().with_provider_retry(ProviderRetry {
|
|
base_delay: Duration::from_secs(2),
|
|
max_delay: Duration::from_secs(6),
|
|
..ProviderRetry::default()
|
|
}),
|
|
)
|
|
.await
|
|
.expect("the run retries and then succeeds");
|
|
|
|
assert_eq!(message.text(), "recovered");
|
|
assert_eq!(
|
|
announced_retry_delays(&collect_events(&mut events)),
|
|
vec![2_000, 4_000, 6_000],
|
|
"the host's base doubles to the host's ceiling"
|
|
);
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
async fn a_rate_limit_that_names_its_window_is_waited_out() {
|
|
// The failure this exists for: a gateway answering 429 with a window far
|
|
// longer than a blip-shaped backoff would ever reach.
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
vec![
|
|
failed_request(ProviderError::Http {
|
|
status: StatusCode::TOO_MANY_REQUESTS,
|
|
body: "rate limit exceeded".to_string(),
|
|
retry_after: Some(Duration::from_secs(45)),
|
|
}),
|
|
text_stream("model", "recovered"),
|
|
],
|
|
);
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).expect("spawn agent");
|
|
let mut events = agent.subscribe_events();
|
|
|
|
let message = agent
|
|
.send(vec![ContentBlock::text("hello")])
|
|
.await
|
|
.expect("the run waits out the window and succeeds");
|
|
|
|
assert_eq!(message.text(), "recovered");
|
|
assert_eq!(
|
|
announced_retry_delays(&collect_events(&mut events)),
|
|
vec![45_000],
|
|
"the provider's window, not the schedule's 500 ms"
|
|
);
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
async fn a_server_cannot_park_a_run_for_an_hour() {
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
vec![
|
|
failed_request(ProviderError::Http {
|
|
status: StatusCode::SERVICE_UNAVAILABLE,
|
|
body: "come back later".to_string(),
|
|
retry_after: Some(Duration::from_secs(3_600)),
|
|
}),
|
|
text_stream("model", "recovered"),
|
|
],
|
|
);
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).expect("spawn agent");
|
|
let mut events = agent.subscribe_events();
|
|
|
|
agent
|
|
.send(vec![ContentBlock::text("hello")])
|
|
.await
|
|
.expect("the run succeeds after the clamped wait");
|
|
|
|
assert_eq!(
|
|
announced_retry_delays(&collect_events(&mut events)),
|
|
vec![60_000],
|
|
"clamped to the default one-minute ceiling"
|
|
);
|
|
}
|
|
|
|
/// Counters a [`RoundStrategy`] observed at the most recent round boundary,
|
|
/// captured by [`CountingStrategy`].
|
|
#[derive(Clone, Copy, Default)]
|
|
struct RoundCounters {
|
|
rounds_completed: usize,
|
|
model_requests: usize,
|
|
transport_retries: usize,
|
|
}
|
|
|
|
/// A [`RoundStrategy`] that records [`RoundContext`]'s counters at each boundary
|
|
/// it observes, always proceeding.
|
|
struct CountingStrategy {
|
|
last: Mutex<Option<RoundCounters>>,
|
|
}
|
|
|
|
impl CountingStrategy {
|
|
fn new() -> Arc<Self> {
|
|
Arc::new(Self {
|
|
last: Mutex::new(None),
|
|
})
|
|
}
|
|
|
|
async fn last_counters(&self) -> RoundCounters {
|
|
self.last.lock().await.expect("strategy observed a round")
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl RoundStrategy for CountingStrategy {
|
|
async fn on_round(&self, ctx: RoundContext<'_>) -> RoundDecision {
|
|
*self.last.lock().await = Some(RoundCounters {
|
|
rounds_completed: ctx.rounds_completed(),
|
|
model_requests: ctx.model_requests(),
|
|
transport_retries: ctx.transport_retries(),
|
|
});
|
|
RoundDecision::proceed()
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn retry_and_round_counters_are_reported_distinctly() {
|
|
// A connection-open retry, then a success: `model_requests` (today's
|
|
// request-including-retries counter) must stay at 2, `rounds_completed` (the
|
|
// logical-round counter) must land at 1 — the round that needed a retry still
|
|
// counts as exactly one completed round — and the new `transport_retries`
|
|
// counter must isolate the one retry, distinct from both.
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
vec![
|
|
failed_request(ProviderError::Http {
|
|
status: StatusCode::SERVICE_UNAVAILABLE,
|
|
body: "offline".to_string(),
|
|
retry_after: None,
|
|
}),
|
|
text_stream(&model.id, "recovered"),
|
|
],
|
|
);
|
|
let provider_handle = provider.clone();
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).expect("spawn agent");
|
|
|
|
let strategy = CountingStrategy::new();
|
|
let message = agent
|
|
.run(
|
|
vec![ContentBlock::text("hello")],
|
|
RunOptions::default().with_round_strategy(strategy.clone()),
|
|
)
|
|
.await
|
|
.expect("run should retry then succeed");
|
|
|
|
assert_eq!(message.text(), "recovered");
|
|
assert_eq!(provider_handle.recorded_requests().await.len(), 2);
|
|
|
|
let counters = strategy.last_counters().await;
|
|
assert_eq!(counters.rounds_completed, 1, "one logical round completed");
|
|
assert_eq!(
|
|
counters.model_requests, 2,
|
|
"model_requests keeps today's semantics: it counts the retry and the success"
|
|
);
|
|
assert_eq!(
|
|
counters.transport_retries, 1,
|
|
"exactly one transient retry, isolated from rounds_completed and reported distinctly"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn resume_replays_last_failed_turn() {
|
|
let model = model_info("model", BuiltinProvider::Anthropic);
|
|
let provider = ScriptedProvider::new(
|
|
BuiltinProvider::Anthropic,
|
|
vec![model.clone()],
|
|
vec![
|
|
erroring_stream(
|
|
vec![ProviderEvent::MessageStarted {
|
|
id: "msg-1".to_string(),
|
|
model: model.id.clone(),
|
|
role: Role::Assistant,
|
|
}],
|
|
ProviderError::MalformedStream("boom".to_string()),
|
|
),
|
|
text_stream(&model.id, "done"),
|
|
],
|
|
);
|
|
|
|
let runtime = test_runtime(provider);
|
|
let mut agent = runtime.spawn("agent", model).expect("spawn agent");
|
|
|
|
let error = agent
|
|
.send(vec![ContentBlock::text("retry me")])
|
|
.await
|
|
.expect_err("first send should fail");
|
|
assert!(matches!(error, RuntimeError::FailedToStreamResponse(_)));
|
|
assert!(agent.history().is_empty());
|
|
|
|
let resumed = agent
|
|
.resume()
|
|
.await
|
|
.expect("resume should replay user turn");
|
|
|
|
assert_eq!(resumed.text(), "done");
|
|
assert_eq!(agent.history().len(), 2);
|
|
assert_eq!(
|
|
agent.history()[0],
|
|
Message::user(ContentBlock::text("retry me"))
|
|
);
|
|
assert_eq!(
|
|
agent.history()[1],
|
|
Message::assistant(ContentBlock::text("done"))
|
|
);
|
|
|
|
let error = agent
|
|
.resume()
|
|
.await
|
|
.expect_err("successful run clears resume state");
|
|
assert!(matches!(error, RuntimeError::NoResumableTurn));
|
|
}
|
|
|
|
fn collect_events(receiver: &mut broadcast::Receiver<AgentEvent>) -> Vec<AgentEvent> {
|
|
let mut events = Vec::new();
|
|
while let Ok(event) = receiver.try_recv() {
|
|
events.push(event);
|
|
}
|
|
events
|
|
}
|
|
|
|
fn test_runtime(provider: ScriptedProvider) -> Runtime {
|
|
Runtime::empty_builder()
|
|
.with_provider_instance(provider)
|
|
.with_store(temp_store("agent-runtime"))
|
|
.build()
|
|
.expect("build runtime")
|
|
}
|
|
|
|
fn temp_store(label: &str) -> SqliteRuntimeStore {
|
|
let unique = NEXT_TEMP_ID.fetch_add(1, Ordering::Relaxed);
|
|
let timestamp = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.expect("system time")
|
|
.as_nanos();
|
|
let path = std::env::temp_dir().join(format!(
|
|
"mentra-agent-runtime-{label}-{timestamp}-{unique}.sqlite"
|
|
));
|
|
if let Some(parent) = path.parent() {
|
|
fs::create_dir_all(parent).expect("create temp dir");
|
|
}
|
|
SqliteRuntimeStore::new(path)
|
|
}
|
|
|
|
fn model_info(id: &str, provider: impl Into<ProviderId>) -> ModelInfo {
|
|
ModelInfo::new(id, provider)
|
|
}
|
|
|
|
fn buffered_stream(events: Vec<ProviderEvent>) -> StreamScript {
|
|
StreamScript::Buffered(events.into_iter().map(Ok).collect())
|
|
}
|
|
|
|
fn erroring_stream(mut events: Vec<ProviderEvent>, error: ProviderError) -> StreamScript {
|
|
let mut items = events.drain(..).map(Ok).collect::<Vec<_>>();
|
|
items.push(Err(error));
|
|
StreamScript::Buffered(items)
|
|
}
|
|
|
|
fn failed_request(error: ProviderError) -> StreamScript {
|
|
StreamScript::Fail(error)
|
|
}
|
|
|
|
fn text_stream(model: &str, text: &str) -> StreamScript {
|
|
buffered_stream(vec![
|
|
ProviderEvent::MessageStarted {
|
|
id: "msg-text".to_string(),
|
|
model: model.to_string(),
|
|
role: Role::Assistant,
|
|
},
|
|
ProviderEvent::ContentBlockStarted {
|
|
index: 0,
|
|
kind: ContentBlockStart::Text,
|
|
},
|
|
ProviderEvent::ContentBlockDelta {
|
|
index: 0,
|
|
delta: ContentBlockDelta::Text(text.to_string()),
|
|
},
|
|
ProviderEvent::ContentBlockStopped { index: 0 },
|
|
ProviderEvent::MessageStopped,
|
|
])
|
|
}
|