From 55b6525ac03f69dce24debf51cf289e3533cef5c Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Mon, 29 Jun 2026 23:23:32 -0700 Subject: [PATCH] =?UTF-8?q?feat(agent):=20M3+M4=20=E2=80=94=20gateway=20tr?= =?UTF-8?q?ansport,=20mock=20transport,=20and=20the=20agent=20loop?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit M3 (transport): - client::GatewayClient — the harness's only network surface. Picks the dialect by model, POSTs provider wire to the gateway with a bai_v1 bearer (adds anthropic-version for the Anthropic dialect), and frames the chunked SSE body line-by-line into StreamEvents via async-stream. No provider keys. - mock::MockTransport — replays scripted StreamEvent turns and records requests, so the loop and the binaries test with no network. - dialect::push_sse_line — incremental per-line decode the client drives off the socket (decode_sse now builds on it). M4 (loop): - agent::Agent — Pi's pi-agent-core runtime. One iteration = stream a turn, assemble assistant content blocks from the event stream (dialect-blind), append, run requested tools, feed results back, repeat until end_turn or max_steps. Tool failures become error tool_results, not aborted runs. Proven: 21 agent-core unit tests (loop dispatch, tool round-trip, unknown/ failing tool, max_steps, event streaming) + 2 real-socket integration tests driving GatewayClient against a bare TCP SSE server (success + HTTP error). Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.lock | 40 +++ crates/agent-core/Cargo.toml | 6 + crates/agent-core/src/agent.rs | 399 +++++++++++++++++++++++ crates/agent-core/src/client.rs | 90 +++++ crates/agent-core/src/dialect/mod.rs | 30 +- crates/agent-core/src/lib.rs | 6 + crates/agent-core/src/mock.rs | 93 ++++++ crates/agent-core/tests/client_socket.rs | 106 ++++++ 8 files changed, 759 insertions(+), 11 deletions(-) create mode 100644 crates/agent-core/src/agent.rs create mode 100644 crates/agent-core/src/client.rs create mode 100644 crates/agent-core/src/mock.rs create mode 100644 crates/agent-core/tests/client_socket.rs diff --git a/Cargo.lock b/Cargo.lock index 0d772f5..5e316fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -219,6 +219,28 @@ dependencies = [ "url", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -356,8 +378,10 @@ dependencies = [ name = "beyond-ai-agent-core" version = "0.1.0" dependencies = [ + "async-stream", "async-trait", "futures", + "reqwest", "serde", "serde_json", "thiserror 2.0.18", @@ -2851,6 +2875,7 @@ dependencies = [ "base64", "bytes", "futures-core", + "futures-util", "http", "http-body", "http-body-util", @@ -2870,12 +2895,14 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", ] @@ -4086,6 +4113,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" diff --git a/crates/agent-core/Cargo.toml b/crates/agent-core/Cargo.toml index 0659680..1de32e3 100644 --- a/crates/agent-core/Cargo.toml +++ b/crates/agent-core/Cargo.toml @@ -23,6 +23,12 @@ serde_json.workspace = true thiserror.workspace = true tracing.workspace = true +# The default model transport (`client::GatewayClient`) speaks HTTP to the gateway. `async-stream` +# turns the chunked SSE body into a `Stream` without hand-rolling a poll state machine. +# reqwest mirrors the gateway crate's pins (rustls, no OpenSSL) plus `stream` for the SSE body. +async-stream = "0.3" +reqwest = { version = "0.13", default-features = false, features = ["json", "rustls", "stream"] } + [dev-dependencies] # Async test runtime for the `Tool`/loop tests. The library itself is runtime-agnostic — it pulls in # no executor, only the `futures` stream traits — so the runtime lives in dev-deps. diff --git a/crates/agent-core/src/agent.rs b/crates/agent-core/src/agent.rs new file mode 100644 index 0000000..90741fa --- /dev/null +++ b/crates/agent-core/src/agent.rs @@ -0,0 +1,399 @@ +//! The agent loop — Pi's `pi-agent-core` runtime, ported. +//! +//! One iteration = one model turn: stream a completion, assemble the assistant message from the +//! event stream, append it to the session, and — if the model asked for tools — run each tool and +//! feed the results back as a new user turn. Repeat until the model ends its turn (or `max_steps`). +//! +//! The loop is dialect-blind (both wire dialects normalize to the same `StreamEvent` sequence) and +//! network-blind (it depends only on [`ModelTransport`], so tests drive it with `MockTransport`). + +use std::sync::Arc; + +use futures::StreamExt; +use serde_json::{Value, json}; + +use crate::error::{Error, Result}; +use crate::message::{ContentBlock, Message, StopReason, StreamEvent}; +use crate::session::Session; +use crate::tool::ToolRegistry; +use crate::transport::{ModelRequest, ModelTransport}; + +/// Default per-turn output token ceiling. +const DEFAULT_MAX_TOKENS: u32 = 4096; +/// Default ceiling on loop iterations before bailing — a runaway-tool-call backstop. +const DEFAULT_MAX_STEPS: u32 = 24; + +/// A configured agent: a model, a transport, a tool set, and loop bounds. Cheap to clone-construct; +/// `run` borrows it so one agent can drive many sessions. +pub struct Agent { + transport: Arc, + tools: ToolRegistry, + model: String, + system: Option, + max_tokens: u32, + max_steps: u32, +} + +impl Agent { + /// An agent over `transport` using `model`, with no tools and default bounds. + pub fn new(transport: Arc, model: impl Into) -> Self { + Self { + transport, + tools: ToolRegistry::new(), + model: model.into(), + system: None, + max_tokens: DEFAULT_MAX_TOKENS, + max_steps: DEFAULT_MAX_STEPS, + } + } + + /// Set the tools the model may call. + pub fn with_tools(mut self, tools: ToolRegistry) -> Self { + self.tools = tools; + self + } + + /// Set the system prompt. + pub fn with_system(mut self, system: impl Into) -> Self { + self.system = Some(system.into()); + self + } + + /// Set the per-turn output token ceiling. + pub fn with_max_tokens(mut self, max_tokens: u32) -> Self { + self.max_tokens = max_tokens; + self + } + + /// Set the loop-iteration ceiling. + pub fn with_max_steps(mut self, max_steps: u32) -> Self { + self.max_steps = max_steps; + self + } + + /// Drive the loop to completion against `session`, invoking `on_event` for every streamed event + /// (use it to render assistant text/tool activity live). Returns when the model ends its turn + /// without requesting tools, or errors with [`Error::MaxSteps`] if it never does. + pub async fn run(&self, session: &mut Session, mut on_event: F) -> Result<()> + where + F: FnMut(&StreamEvent), + { + loop { + if session.steps >= self.max_steps { + return Err(Error::MaxSteps(self.max_steps)); + } + + let mut req = ModelRequest::new( + self.model.clone(), + session.messages.clone(), + self.max_tokens, + ) + .with_tools(self.tools.definitions()); + if let Some(system) = &self.system { + req = req.with_system(system.clone()); + } + + let turn = self.run_turn(req, &mut on_event).await?; + session.push(Message::assistant(turn.blocks)); + session.record_usage(turn.input_tokens, turn.output_tokens); + session.steps += 1; + + // Collect the tool calls the assistant just made. + let calls: Vec<(String, String, Value)> = session + .messages + .last() + .map(|m| { + m.tool_uses() + .map(|(id, name, input)| (id.to_string(), name.to_string(), input.clone())) + .collect() + }) + .unwrap_or_default(); + + if calls.is_empty() || turn.stop_reason != StopReason::ToolUse { + return Ok(()); // model ended its turn — done. + } + + // Run each tool and feed results back as a user turn. A tool's own failure becomes an + // error `tool_result`, not an aborted run — the model can react to it next turn. + for (id, name, input) in calls { + let (content, is_error) = match self.tools.get(&name) { + Some(tool) => match tool.run(input).await { + Ok(out) => (out, false), + Err(e) => (e.to_string(), true), + }, + None => (format!("unknown tool: {name}"), true), + }; + session.push(Message::tool_result(id, content, is_error)); + } + } + } + + /// Stream and assemble a single model turn into content blocks + accounting. + async fn run_turn(&self, req: ModelRequest, on_event: &mut F) -> Result + where + F: FnMut(&StreamEvent), + { + let mut stream = self.transport.stream(req).await?; + let mut acc = Accumulator::default(); + while let Some(ev) = stream.next().await { + let ev = ev?; + on_event(&ev); + acc.apply(ev); + } + Ok(acc.finish()) + } +} + +/// The assembled result of one model turn. +struct Turn { + blocks: Vec, + stop_reason: StopReason, + input_tokens: u32, + output_tokens: u32, +} + +/// Folds a `StreamEvent` sequence into content blocks. Text accrues into the current text run; a +/// tool call accrues its streamed JSON argument fragments; `ContentBlockStop` finalizes whichever is +/// open. Works identically for both dialects because they emit the same event shape. +#[derive(Default)] +struct Accumulator { + blocks: Vec, + text: String, + tool: Option<(String, String, String)>, // (id, name, json-arg buffer) + stop_reason: StopReason, + input_tokens: u32, + output_tokens: u32, +} + +impl Accumulator { + fn apply(&mut self, ev: StreamEvent) { + match ev { + StreamEvent::MessageStart => {} + StreamEvent::TextDelta { text } => self.text.push_str(&text), + StreamEvent::ToolUseStart { id, name } => { + self.flush_text(); + self.tool = Some((id, name, String::new())); + } + StreamEvent::InputJsonDelta { partial_json } => { + if let Some((_, _, buf)) = &mut self.tool { + buf.push_str(&partial_json); + } + } + StreamEvent::ContentBlockStop => self.flush_block(), + StreamEvent::Usage { + input_tokens, + output_tokens, + } => { + self.input_tokens = input_tokens; + self.output_tokens = output_tokens; + } + StreamEvent::MessageStop { stop_reason } => self.stop_reason = stop_reason, + } + } + + fn flush_text(&mut self) { + if !self.text.is_empty() { + self.blocks.push(ContentBlock::Text { + text: std::mem::take(&mut self.text), + }); + } + } + + fn flush_block(&mut self) { + if let Some((id, name, args)) = self.tool.take() { + let input = if args.trim().is_empty() { + json!({}) + } else { + serde_json::from_str(&args).unwrap_or(Value::Null) + }; + self.blocks.push(ContentBlock::ToolUse { id, name, input }); + } else { + self.flush_text(); + } + } + + fn finish(mut self) -> Turn { + // A stream that ended without a trailing ContentBlockStop (or with leftover text) still + // contributes its text. + self.flush_block(); + Turn { + blocks: self.blocks, + stop_reason: self.stop_reason, + input_tokens: self.input_tokens, + output_tokens: self.output_tokens, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mock::{MockTransport, turn}; + use crate::tool::Tool; + use async_trait::async_trait; + use serde_json::json; + + struct EchoTool; + #[async_trait] + impl Tool for EchoTool { + fn name(&self) -> &str { + "echo" + } + fn description(&self) -> &str { + "echo the text arg" + } + fn input_schema(&self) -> Value { + json!({ "type": "object", "properties": { "text": { "type": "string" } }, "required": ["text"] }) + } + async fn run(&self, input: Value) -> std::result::Result { + input + .get("text") + .and_then(Value::as_str) + .map(str::to_string) + .ok_or_else(|| crate::error::ToolError::InvalidInput("missing text".into())) + } + } + + fn agent_with( + turns: Vec>, + tools: ToolRegistry, + ) -> (Agent, Arc) { + let mock = Arc::new(MockTransport::new(turns)); + let agent = Agent::new(mock.clone(), "claude-opus-4-8") + .with_tools(tools) + .with_max_steps(8); + (agent, mock) + } + + #[tokio::test] + async fn single_text_turn_completes() { + let (agent, mock) = agent_with(vec![turn::text("hello world")], ToolRegistry::new()); + let mut session = Session::new(); + session.user("hi"); + agent.run(&mut session, |_| {}).await.unwrap(); + + assert_eq!(mock.calls(), 1); + assert_eq!(session.steps, 1); + // user + assistant + assert_eq!(session.messages.len(), 2); + assert_eq!( + session.messages[1].content, + vec![ContentBlock::Text { + text: "hello world".into() + }] + ); + } + + #[tokio::test] + async fn tool_call_round_trips_and_continues() { + let mut tools = ToolRegistry::new(); + tools.register(Arc::new(EchoTool)); + let (agent, mock) = agent_with( + vec![ + turn::tool_call("tu_1", "echo", r#"{"text":"pong"}"#), + turn::text("done"), + ], + tools, + ); + let mut session = Session::new(); + session.user("say pong"); + agent.run(&mut session, |_| {}).await.unwrap(); + + assert_eq!(mock.calls(), 2); + assert_eq!(session.steps, 2); + // user, assistant(tool_use), user(tool_result), assistant(text) + assert_eq!(session.messages.len(), 4); + assert_eq!( + session.messages[2].content, + vec![ContentBlock::ToolResult { + tool_use_id: "tu_1".into(), + content: "pong".into(), + is_error: false + }] + ); + // The second request the loop sent must include the tool result. + let second = &mock.requests()[1]; + assert!( + second + .messages + .iter() + .any(|m| matches!(m.content.first(), Some(ContentBlock::ToolResult { .. }))) + ); + } + + #[tokio::test] + async fn unknown_tool_yields_error_result() { + let (agent, _mock) = agent_with( + vec![ + turn::tool_call("tu_1", "nonexistent", "{}"), + turn::text("ok"), + ], + ToolRegistry::new(), + ); + let mut session = Session::new(); + session.user("go"); + agent.run(&mut session, |_| {}).await.unwrap(); + + match &session.messages[2].content[0] { + ContentBlock::ToolResult { + is_error, content, .. + } => { + assert!(is_error); + assert!(content.contains("unknown tool")); + } + other => panic!("expected error tool_result, got {other:?}"), + } + } + + #[tokio::test] + async fn failing_tool_is_reported_not_fatal() { + let mut tools = ToolRegistry::new(); + tools.register(Arc::new(EchoTool)); + let (agent, _mock) = agent_with( + vec![ + turn::tool_call("tu_1", "echo", r#"{"wrong":"key"}"#), + turn::text("recovered"), + ], + tools, + ); + let mut session = Session::new(); + session.user("go"); + agent.run(&mut session, |_| {}).await.unwrap(); + match &session.messages[2].content[0] { + ContentBlock::ToolResult { is_error, .. } => assert!(is_error), + other => panic!("expected error tool_result, got {other:?}"), + } + } + + #[tokio::test] + async fn max_steps_is_enforced() { + // The model keeps asking for a tool forever. + let turns = vec![turn::tool_call("t", "echo", r#"{"text":"x"}"#); 10]; + let mut tools = ToolRegistry::new(); + tools.register(Arc::new(EchoTool)); + let mock = Arc::new(MockTransport::new(turns)); + let agent = Agent::new(mock.clone(), "claude-opus-4-8") + .with_tools(tools) + .with_max_steps(3); + let mut session = Session::new(); + session.user("loop"); + let err = agent.run(&mut session, |_| {}).await.unwrap_err(); + assert!(matches!(err, Error::MaxSteps(3))); + } + + #[tokio::test] + async fn streams_events_to_callback() { + let (agent, _mock) = agent_with(vec![turn::text("stream me")], ToolRegistry::new()); + let mut session = Session::new(); + session.user("hi"); + let mut seen = Vec::new(); + agent + .run(&mut session, |ev| seen.push(ev.clone())) + .await + .unwrap(); + assert!( + seen.iter() + .any(|e| matches!(e, StreamEvent::TextDelta { text } if text == "stream me")) + ); + } +} diff --git a/crates/agent-core/src/client.rs b/crates/agent-core/src/client.rs new file mode 100644 index 0000000..aae4925 --- /dev/null +++ b/crates/agent-core/src/client.rs @@ -0,0 +1,90 @@ +//! The default model transport: an HTTP client that speaks provider wire to the Beyond gateway. +//! +//! This is the harness's whole network surface. It never holds a provider key or picks a provider — +//! it sends `Authorization: Bearer ` to the gateway, which swaps in the pool key, routes to +//! the real provider, and meters usage. The client only picks the *dialect* (by model id), builds +//! the request body, and frames the streaming SSE response back into [`StreamEvent`]s. + +use async_trait::async_trait; +use futures::StreamExt; + +use crate::dialect::{Dialect, push_sse_line}; +use crate::error::{Error, Result}; +use crate::transport::{EventStream, ModelRequest, ModelTransport}; + +/// Anthropic's Messages API requires this header; the gateway relays it to the upstream verbatim. +const ANTHROPIC_VERSION: &str = "2023-06-01"; + +/// An HTTP client pointed at a Beyond gateway base URL, authenticated with a `bai_v1` virtual key. +pub struct GatewayClient { + http: reqwest::Client, + base_url: String, + api_key: String, +} + +impl GatewayClient { + /// Build a client for `base_url` (e.g. `http://ai.internal` or `http://127.0.0.1:8080`) using + /// `api_key` (a `bai_v1…` virtual key, or a BYO provider key the gateway forwards untouched). + pub fn new(base_url: impl Into, api_key: impl Into) -> Result { + let http = reqwest::Client::builder() + .build() + .map_err(|e| Error::Transport(e.to_string()))?; + Ok(Self { + http, + base_url: base_url.into().trim_end_matches('/').to_string(), + api_key: api_key.into(), + }) + } +} + +#[async_trait] +impl ModelTransport for GatewayClient { + async fn stream(&self, req: ModelRequest) -> Result { + let dialect = Dialect::for_model(&req.model); + let url = format!("{}{}", self.base_url, dialect.endpoint_path()); + let body = dialect.build_body(&req); + let http = self.http.clone(); + let api_key = self.api_key.clone(); + + let stream = async_stream::try_stream! { + let mut builder = http.post(&url).bearer_auth(&api_key).json(&body); + if dialect == Dialect::Anthropic { + builder = builder.header("anthropic-version", ANTHROPIC_VERSION); + } + + let resp = builder.send().await.map_err(|e| Error::Transport(e.to_string()))?; + let status = resp.status(); + if !status.is_success() { + let detail = resp.text().await.unwrap_or_default(); + Err(Error::Transport(format!("gateway returned {status}: {}", detail.trim())))?; + } else { + // Frame the chunked body line-by-line. SSE for both providers carries one JSON object + // per `data:` line, so a line splitter suffices; a partial trailing line is buffered + // across chunks until its newline arrives. + let mut decoder = dialect.decoder(); + let mut buf = String::new(); + let mut bytes = resp.bytes_stream(); + while let Some(chunk) = bytes.next().await { + let chunk = chunk.map_err(|e| Error::Transport(e.to_string()))?; + buf.push_str(&String::from_utf8_lossy(&chunk)); + while let Some(nl) = buf.find('\n') { + let line: String = buf.drain(..=nl).collect(); + for ev in push_sse_line(decoder.as_mut(), &line)? { + yield ev; + } + } + } + if !buf.is_empty() { + for ev in push_sse_line(decoder.as_mut(), &buf)? { + yield ev; + } + } + for ev in decoder.finish() { + yield ev; + } + } + }; + + Ok(Box::pin(stream)) + } +} diff --git a/crates/agent-core/src/dialect/mod.rs b/crates/agent-core/src/dialect/mod.rs index 3f308b5..d2c12b1 100644 --- a/crates/agent-core/src/dialect/mod.rs +++ b/crates/agent-core/src/dialect/mod.rs @@ -78,22 +78,30 @@ pub trait StreamDecoder: Send { pub fn decode_sse(decoder: &mut dyn StreamDecoder, raw: &str) -> Result> { let mut out = Vec::new(); for line in raw.lines() { - let line = line.trim_end_matches('\r'); - let Some(payload) = line.strip_prefix("data:") else { - continue; - }; - let payload = payload.trim(); - if payload.is_empty() || payload == "[DONE]" { - continue; - } - let v: Value = serde_json::from_str(payload) - .map_err(|e| Error::Transport(format!("malformed SSE json: {e}")))?; - out.extend(decoder.push(&v)); + out.extend(push_sse_line(decoder, line)?); } out.extend(decoder.finish()); Ok(out) } +/// Feed a single SSE line to the decoder, returning any events it produced. Non-`data:` lines +/// (comments, `event:`, blanks) and the OpenAI `[DONE]` sentinel produce nothing. This is the +/// incremental entry point the streaming HTTP client drives line-by-line off the socket; the caller +/// is responsible for calling [`StreamDecoder::finish`] once the stream closes. +pub fn push_sse_line(decoder: &mut dyn StreamDecoder, line: &str) -> Result> { + let line = line.trim_end_matches('\r'); + let Some(payload) = line.strip_prefix("data:") else { + return Ok(Vec::new()); + }; + let payload = payload.trim(); + if payload.is_empty() || payload == "[DONE]" { + return Ok(Vec::new()); + } + let v: Value = serde_json::from_str(payload) + .map_err(|e| Error::Transport(format!("malformed SSE json: {e}")))?; + Ok(decoder.push(&v)) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/agent-core/src/lib.rs b/crates/agent-core/src/lib.rs index 27694e6..ee64d5d 100644 --- a/crates/agent-core/src/lib.rs +++ b/crates/agent-core/src/lib.rs @@ -18,15 +18,21 @@ // job is to assert a precondition holds, so `.unwrap()` *is* the assertion — allow it under `test`. #![cfg_attr(test, allow(clippy::unwrap_used, clippy::expect_used, clippy::panic))] +pub mod agent; +pub mod client; pub mod dialect; pub mod error; pub mod message; +pub mod mock; pub mod session; pub mod tool; pub mod transport; +pub use agent::Agent; +pub use client::GatewayClient; pub use error::{Error, Result, ToolError}; pub use message::{ContentBlock, Message, Role, StopReason, StreamEvent, ToolDef}; +pub use mock::MockTransport; pub use session::Session; pub use tool::{Tool, ToolRegistry}; pub use transport::{ModelRequest, ModelTransport}; diff --git a/crates/agent-core/src/mock.rs b/crates/agent-core/src/mock.rs new file mode 100644 index 0000000..c4afc58 --- /dev/null +++ b/crates/agent-core/src/mock.rs @@ -0,0 +1,93 @@ +//! A scripted [`ModelTransport`] for tests. +//! +//! Replaying a fixed sequence of `StreamEvent` "turns" lets the agent loop — and the binaries that +//! drive it — be exercised end to end with no network and no model. Each call to `stream` pops the +//! next turn and records the request, so a test can assert exactly what the loop sent back. + +use std::collections::VecDeque; +use std::sync::Mutex; + +use async_trait::async_trait; + +use crate::error::{Error, Result}; +use crate::message::StreamEvent; +use crate::transport::{EventStream, ModelRequest, ModelTransport}; + +/// A transport that yields pre-scripted turns. Construct with one `Vec` per model turn. +pub struct MockTransport { + turns: Mutex>>, + requests: Mutex>, +} + +impl MockTransport { + /// Script the transport with the turns it will return, in order. + pub fn new(turns: Vec>) -> Self { + Self { + turns: Mutex::new(turns.into()), + requests: Mutex::new(Vec::new()), + } + } + + /// The requests the loop has sent so far, in order (for asserting what the loop fed back). + pub fn requests(&self) -> Vec { + self.requests.lock().map(|r| r.clone()).unwrap_or_default() + } + + /// How many turns the loop has consumed. + pub fn calls(&self) -> usize { + self.requests.lock().map(|r| r.len()).unwrap_or(0) + } +} + +#[async_trait] +impl ModelTransport for MockTransport { + async fn stream(&self, req: ModelRequest) -> Result { + if let Ok(mut reqs) = self.requests.lock() { + reqs.push(req); + } + let turn = self + .turns + .lock() + .ok() + .and_then(|mut t| t.pop_front()) + .ok_or_else(|| Error::Transport("MockTransport: no more scripted turns".into()))?; + Ok(Box::pin(futures::stream::iter(turn.into_iter().map(Ok)))) + } +} + +/// Convenience builders for the events that make up a scripted turn. +pub mod turn { + use crate::message::{StopReason, StreamEvent}; + + /// A turn that emits one text block and ends the conversation. + pub fn text(s: &str) -> Vec { + vec![ + StreamEvent::MessageStart, + StreamEvent::TextDelta { + text: s.to_string(), + }, + StreamEvent::ContentBlockStop, + StreamEvent::MessageStop { + stop_reason: StopReason::EndTurn, + }, + ] + } + + /// A turn that calls one tool with the given (whole) JSON argument string. + pub fn tool_call(id: &str, name: &str, args_json: &str) -> Vec { + vec![ + StreamEvent::MessageStart, + StreamEvent::ToolUseStart { + id: id.to_string(), + name: name.to_string(), + }, + StreamEvent::InputJsonDelta { + partial_json: args_json.to_string(), + }, + StreamEvent::ContentBlockStop, + StreamEvent::MessageStop { + stop_reason: StopReason::ToolUse, + }, + ] + } +} diff --git a/crates/agent-core/tests/client_socket.rs b/crates/agent-core/tests/client_socket.rs new file mode 100644 index 0000000..6fc5623 --- /dev/null +++ b/crates/agent-core/tests/client_socket.rs @@ -0,0 +1,106 @@ +//! `GatewayClient` over a real TCP socket. +//! +//! A bare TCP server writes a canned `text/event-stream` response (connection-close framing, like a +//! provider streaming SSE). This proves the client's chunked-body line framing + dialect decoding +//! end to end over a socket — the unit tests cover decoding from a string; this covers it from the +//! wire. No hyper, no gateway, no network. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +use std::io::{Read, Write}; +use std::net::TcpListener; +use std::thread; + +use agent_core::message::StreamEvent; +use agent_core::transport::ModelTransport; +use agent_core::{GatewayClient, Message, ModelRequest}; +use futures::StreamExt; + +/// Spawn a one-shot server that returns `body` as an SSE response, and return its base URL. +fn spawn_sse_server(body: &'static str) -> String { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr"); + thread::spawn(move || { + if let Ok((mut stream, _)) = listener.accept() { + // Drain the request (headers + small JSON body) so the client's write completes. + let mut buf = [0u8; 8192]; + let _ = stream.read(&mut buf); + let resp = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n{body}" + ); + let _ = stream.write_all(resp.as_bytes()); + let _ = stream.flush(); + } + }); + format!("http://{addr}") +} + +const ANTHROPIC_SSE: &str = "event: message_start\n\ +data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":7,\"output_tokens\":1}}}\n\ +\n\ +event: content_block_start\n\ +data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\ +\n\ +event: content_block_delta\n\ +data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi over the wire\"}}\n\ +\n\ +event: content_block_stop\n\ +data: {\"type\":\"content_block_stop\",\"index\":0}\n\ +\n\ +event: message_delta\n\ +data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\ +\n\ +event: message_stop\n\ +data: {\"type\":\"message_stop\"}\n\ +\n"; + +#[tokio::test] +async fn gateway_client_decodes_sse_over_socket() { + let base = spawn_sse_server(ANTHROPIC_SSE); + let client = GatewayClient::new(base, "bai_v1.test").expect("client"); + + let req = ModelRequest::new("claude-opus-4-8", vec![Message::user("hi")], 64); + let mut stream = client.stream(req).await.expect("stream"); + + let mut events = Vec::new(); + while let Some(ev) = stream.next().await { + events.push(ev.expect("event")); + } + + assert!(events.contains(&StreamEvent::TextDelta { + text: "hi over the wire".into() + })); + assert!(events.contains(&StreamEvent::Usage { + input_tokens: 7, + output_tokens: 5 + })); + assert!(matches!( + events.last(), + Some(StreamEvent::MessageStop { .. }) + )); +} + +#[tokio::test] +async fn gateway_client_surfaces_http_error() { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr"); + thread::spawn(move || { + if let Ok((mut stream, _)) = listener.accept() { + let mut buf = [0u8; 8192]; + let _ = stream.read(&mut buf); + let _ = stream.write_all( + b"HTTP/1.1 403 Forbidden\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n{\"error\":\"denied\"}", + ); + } + }); + let client = GatewayClient::new(format!("http://{addr}"), "bai_v1.test").expect("client"); + let mut stream = client + .stream(ModelRequest::new( + "claude-opus-4-8", + vec![Message::user("hi")], + 64, + )) + .await + .expect("stream"); + let first = stream.next().await.expect("an item"); + assert!(first.is_err(), "a 403 must surface as a stream error"); +}