370 lines
12 KiB
Rust
370 lines
12 KiB
Rust
//! Integration tests for the agent loop, driven by the mock provider.
|
||
|
||
use focus_core::event::VecSink;
|
||
use focus_core::model::*;
|
||
use focus_core::tool::{Tool, ToolEffects, ToolRegistry, ToolResult, ToolUpdateSink};
|
||
use focus_core::{Agent, AgentConfig};
|
||
use focus_json::JsonValue;
|
||
use mock::{make_text_response, make_tool_call_response, MockProvider};
|
||
|
||
/// A simple echo tool that returns its arguments as text.
|
||
#[derive(Debug)]
|
||
struct EchoTool;
|
||
|
||
impl Tool for EchoTool {
|
||
fn name(&self) -> &str {
|
||
"echo"
|
||
}
|
||
fn description(&self) -> &str {
|
||
"echoes the arguments back"
|
||
}
|
||
fn parameters(&self) -> JsonValue {
|
||
let mut o = JsonValue::obj();
|
||
let mut props = JsonValue::obj();
|
||
props.insert("text", JsonValue::Str("the text".into())).ok();
|
||
o.insert("type", "object".into()).ok();
|
||
o.insert("properties", props).ok();
|
||
o
|
||
}
|
||
fn effects(&self) -> ToolEffects {
|
||
ToolEffects::READ
|
||
}
|
||
fn execute(
|
||
&self,
|
||
_id: &str,
|
||
args: &JsonValue,
|
||
_on_update: Option<&ToolUpdateSink<'_>>,
|
||
) -> Result<ToolResult, focus_core::CoreError> {
|
||
let text = args.get_str("text").unwrap_or("(none)").to_string();
|
||
Ok(ToolResult::text(format!("echo: {}", text)))
|
||
}
|
||
}
|
||
|
||
fn agent_with(provider: MockProvider) -> Agent {
|
||
let config = AgentConfig {
|
||
model: "mock-model".into(),
|
||
..Default::default()
|
||
};
|
||
let tools = ToolRegistry::with(Box::new(EchoTool));
|
||
Agent::new(config, Box::new(provider), tools)
|
||
}
|
||
|
||
#[test]
|
||
fn single_text_response_completes() {
|
||
let provider = MockProvider::single_text("hello there", StopReason::Stop);
|
||
let mut agent = agent_with(provider);
|
||
let mut sink = VecSink::new();
|
||
agent.prompt("hi", &mut sink).expect("run failed");
|
||
|
||
assert_eq!(sink.count("agent_start"), 1);
|
||
assert_eq!(sink.count("agent_end"), 1);
|
||
assert_eq!(sink.count("message_end"), 2); // user + assistant
|
||
|
||
// Transcript: user + assistant
|
||
let msgs = agent.messages();
|
||
assert_eq!(msgs.len(), 2);
|
||
assert!(matches!(msgs[0], Message::User(_)));
|
||
let assistant = match &msgs[1] {
|
||
Message::Assistant(a) => a,
|
||
_ => panic!("expected assistant message"),
|
||
};
|
||
let text = assistant
|
||
.content
|
||
.iter()
|
||
.find_map(|c| c.as_text())
|
||
.expect("assistant has text");
|
||
assert_eq!(text.text, "hello there");
|
||
assert_eq!(assistant.stop_reason, StopReason::Stop);
|
||
}
|
||
|
||
#[test]
|
||
fn tool_call_then_text_completes_two_turns() {
|
||
// Turn 1: model calls "echo" tool. Turn 2: model gives final text.
|
||
let mut args = JsonValue::obj();
|
||
args.insert("text", "world".into()).ok();
|
||
let provider = MockProvider::new(vec![
|
||
make_tool_call_response(None, "echo", "call_1", args, StopReason::ToolUse),
|
||
make_text_response("all done", StopReason::Stop),
|
||
]);
|
||
let mut agent = agent_with(provider);
|
||
let mut sink = VecSink::new();
|
||
agent.prompt("run echo", &mut sink).expect("run failed");
|
||
|
||
// Two turns.
|
||
assert_eq!(sink.count("turn_end"), 2);
|
||
assert_eq!(sink.count("tool_execution_end"), 1);
|
||
assert_eq!(sink.count("agent_end"), 1);
|
||
|
||
// Transcript: user, assistant(toolcall), toolResult, assistant(text)
|
||
let msgs = agent.messages();
|
||
assert_eq!(msgs.len(), 4);
|
||
assert!(matches!(msgs[0], Message::User(_)));
|
||
assert!(matches!(
|
||
&msgs[1],
|
||
Message::Assistant(a) if a.content.iter().any(|c| c.as_tool_call().is_some())
|
||
));
|
||
assert!(matches!(
|
||
&msgs[2],
|
||
Message::ToolResult(tr) if tr.tool_name == "echo" && !tr.is_error
|
||
));
|
||
let text = match &msgs[3] {
|
||
Message::Assistant(a) => a
|
||
.content
|
||
.iter()
|
||
.find_map(|c| c.as_text())
|
||
.map(|t| t.text.clone()),
|
||
_ => None,
|
||
};
|
||
assert_eq!(text.as_deref(), Some("all done"));
|
||
|
||
// The tool result content should contain the echo.
|
||
if let Message::ToolResult(tr) = &msgs[2] {
|
||
let result_text = tr
|
||
.content
|
||
.iter()
|
||
.find_map(|c| c.as_text())
|
||
.map(|t| t.text.clone());
|
||
assert_eq!(result_text.as_deref(), Some("echo: world"));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn unknown_tool_produces_error_result() {
|
||
let provider = MockProvider::new(vec![
|
||
make_tool_call_response(
|
||
None,
|
||
"nonexistent",
|
||
"call_1",
|
||
JsonValue::obj(),
|
||
StopReason::ToolUse,
|
||
),
|
||
make_text_response("ok", StopReason::Stop),
|
||
]);
|
||
let mut agent = agent_with(provider);
|
||
let mut sink = VecSink::new();
|
||
agent.prompt("go", &mut sink).expect("run failed");
|
||
|
||
// The tool result should be an error.
|
||
let msgs = agent.messages();
|
||
assert!(matches!(msgs[2], Message::ToolResult(_)));
|
||
if let Message::ToolResult(tr) = &msgs[2] {
|
||
assert!(tr.is_error, "expected error for unknown tool");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn error_stop_reason_ends_run_immediately() {
|
||
let provider = MockProvider::new(vec![make_text_response("oops", StopReason::Error)]);
|
||
let mut agent = agent_with(provider);
|
||
let mut sink = VecSink::new();
|
||
agent.prompt("hi", &mut sink).expect("run failed");
|
||
|
||
assert_eq!(sink.count("turn_end"), 1);
|
||
assert_eq!(sink.count("agent_end"), 1);
|
||
let assistant = match agent.messages().last().unwrap() {
|
||
Message::Assistant(a) => a,
|
||
_ => panic!(),
|
||
};
|
||
assert_eq!(assistant.stop_reason, StopReason::Error);
|
||
}
|
||
|
||
#[test]
|
||
fn continue_requires_non_assistant_last_message() {
|
||
let provider = MockProvider::single_text("hi", StopReason::Stop);
|
||
let mut agent = agent_with(provider);
|
||
let mut sink = VecSink::new();
|
||
agent.prompt("hello", &mut sink).expect("run failed");
|
||
// Last message is now assistant — continue must fail.
|
||
let result = agent.r#continue(&mut sink);
|
||
assert!(result.is_err());
|
||
}
|
||
|
||
/// 回放脚本化事件的 mock provider。目前只有本测试文件使用，因此内联在此；
|
||
/// 若将来多个集成测试文件都需要，再抽成独立的测试工具 crate。
|
||
/// A mock provider that replays scripted events. Only this test file uses it
|
||
/// today, so it lives inline here; extract to a shared test-utils crate if
|
||
/// multiple integration test files need it later.
|
||
mod mock {
|
||
use focus_core::model::*;
|
||
use focus_core::provider::{
|
||
ProviderRequest, StreamEvent, StreamIterator, StreamProvider, StreamResult,
|
||
};
|
||
use focus_core::CoreError;
|
||
use std::sync::{Arc, Mutex};
|
||
|
||
/// 构造一个回放「响应列表」的 mock provider——每个响应是一个回合的事件列表。
|
||
/// Build a mock provider that replays a list of *responses*, where each
|
||
/// response is a list of events for one assistant turn.
|
||
#[derive(Debug, Clone)]
|
||
pub struct MockProvider {
|
||
/// 响应队列；每次回合消费一个。
|
||
/// Queue of responses; each response is consumed by one turn.
|
||
responses: Arc<Mutex<Vec<Vec<StreamEvent>>>>,
|
||
}
|
||
|
||
impl MockProvider {
|
||
/// 用一组响应构造 mock provider。
|
||
/// Build a mock provider from a list of responses.
|
||
pub fn new(responses: Vec<Vec<StreamEvent>>) -> Self {
|
||
Self {
|
||
responses: Arc::new(Mutex::new(responses)),
|
||
}
|
||
}
|
||
|
||
/// 一个总是返回单条文本响应（带停止原因）的 provider。
|
||
/// A provider that always returns a single text response with stop reason.
|
||
pub fn single_text(text: &str, stop: StopReason) -> Self {
|
||
Self::new(vec![make_text_response(text, stop)])
|
||
}
|
||
}
|
||
|
||
impl StreamProvider for MockProvider {
|
||
fn stream(&self, _request: &ProviderRequest) -> StreamResult {
|
||
let mut queue = self.responses.lock().expect("mock poisoned");
|
||
let events = queue
|
||
.first()
|
||
.cloned()
|
||
.ok_or_else(|| CoreError::Provider("mock provider exhausted".into()))?;
|
||
if !matches!(
|
||
events.last(),
|
||
Some(StreamEvent::Done { .. }) | Some(StreamEvent::Error { .. })
|
||
) {
|
||
return Err(CoreError::Provider(
|
||
"mock response must end with Done or Error".into(),
|
||
));
|
||
}
|
||
queue.remove(0);
|
||
Ok(Box::new(MockStream { events }))
|
||
}
|
||
}
|
||
|
||
/// 一个回放固定事件序列的流。
|
||
/// A stream that replays a fixed event sequence.
|
||
struct MockStream {
|
||
events: Vec<StreamEvent>,
|
||
}
|
||
|
||
impl StreamIterator for MockStream {
|
||
fn next_event(&mut self) -> Option<StreamEvent> {
|
||
if self.events.is_empty() {
|
||
None
|
||
} else {
|
||
Some(self.events.remove(0))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 构造一个简单的文本响应事件序列。
|
||
/// Build a simple text-response event sequence.
|
||
pub fn make_text_response(text: &str, stop: StopReason) -> Vec<StreamEvent> {
|
||
let model = "mock-model".to_string();
|
||
let timestamp = now_ms();
|
||
let partial_empty = AssistantMessage {
|
||
content: vec![],
|
||
model: model.clone(),
|
||
usage: Usage::default(),
|
||
stop_reason: stop,
|
||
error_message: None,
|
||
timestamp,
|
||
};
|
||
let partial = partial_empty.clone();
|
||
let mut text_start = partial.clone();
|
||
text_start.content.push(ContentBlock::Text(TextContent {
|
||
text: String::new(),
|
||
signature: None,
|
||
}));
|
||
let mut text_delta = text_start.clone();
|
||
if let Some(ContentBlock::Text(t)) = text_delta.content.get_mut(0) {
|
||
t.text = text.to_string();
|
||
}
|
||
let text_end = text_delta.clone();
|
||
let mut final_msg = text_end.clone();
|
||
final_msg.usage = Usage {
|
||
input_tokens: 10,
|
||
output_tokens: 5,
|
||
..Default::default()
|
||
};
|
||
final_msg.stop_reason = stop;
|
||
vec![
|
||
StreamEvent::Start {
|
||
partial: partial_empty,
|
||
},
|
||
StreamEvent::TextStart {
|
||
content_index: 0,
|
||
partial: text_start,
|
||
},
|
||
StreamEvent::TextDelta {
|
||
content_index: 0,
|
||
delta: text.into(),
|
||
partial: text_delta,
|
||
},
|
||
StreamEvent::TextEnd {
|
||
content_index: 0,
|
||
partial: text_end,
|
||
},
|
||
StreamEvent::Done { message: final_msg },
|
||
]
|
||
}
|
||
|
||
/// 构造一个工具调用响应：模型请求调用一个工具。
|
||
/// Build a tool-call response: the model asks to call a tool.
|
||
pub fn make_tool_call_response(
|
||
thinking_text: Option<&str>,
|
||
tool_name: &str,
|
||
tool_id: &str,
|
||
arguments: focus_json::JsonValue,
|
||
stop: StopReason,
|
||
) -> Vec<StreamEvent> {
|
||
let model = "mock-model".to_string();
|
||
let timestamp = now_ms();
|
||
let mut content: Vec<ContentBlock> = Vec::new();
|
||
if let Some(t) = thinking_text {
|
||
content.push(ContentBlock::text(t));
|
||
}
|
||
content.push(ContentBlock::ToolCall(ToolCall {
|
||
id: tool_id.to_string(),
|
||
name: tool_name.to_string(),
|
||
arguments,
|
||
}));
|
||
let mut msg = AssistantMessage {
|
||
content: content.clone(),
|
||
model: model.clone(),
|
||
usage: Usage::default(),
|
||
stop_reason: stop,
|
||
error_message: None,
|
||
timestamp,
|
||
};
|
||
msg.usage = Usage {
|
||
input_tokens: 10,
|
||
output_tokens: 5,
|
||
..Default::default()
|
||
};
|
||
vec![
|
||
StreamEvent::Start {
|
||
partial: AssistantMessage {
|
||
content: vec![],
|
||
model: model.clone(),
|
||
usage: Usage::default(),
|
||
stop_reason: stop,
|
||
error_message: None,
|
||
timestamp,
|
||
},
|
||
},
|
||
StreamEvent::ToolCallStart {
|
||
content_index: 0,
|
||
partial: msg.clone(),
|
||
},
|
||
StreamEvent::ToolCallEnd {
|
||
content_index: 0,
|
||
tool_call: if let Some(ContentBlock::ToolCall(tc)) = content.last() {
|
||
tc.clone()
|
||
} else {
|
||
unreachable!()
|
||
},
|
||
partial: msg.clone(),
|
||
},
|
||
StreamEvent::Done { message: msg },
|
||
]
|
||
}
|
||
}
|