199 lines
7.8 KiB
Rust
199 lines
7.8 KiB
Rust
//! 全链路集成：agent 跑对话 → harness 压缩规划 → provider 生成摘要 →
|
||
//! 摘要写回 → agent 继续（对应未来 TUI 的自动压缩流程）。
|
||
//! Full-chain integration: agent runs a conversation → harness plans
|
||
//! compaction → the provider summarizes → the summary is written back → the
|
||
//! agent continues (mirroring the TUI's future auto-compaction flow).
|
||
|
||
mod common;
|
||
|
||
use focus_core::event::VecSink;
|
||
use focus_core::model::*;
|
||
use focus_core::provider::{ProviderRequest, StreamEvent};
|
||
use focus_core::tool::{ToolEffects, ToolRegistry, ToolResult};
|
||
use focus_core::{Agent, AgentConfig, CoreError};
|
||
use focus_harness::compaction::{
|
||
apply_summary, plan_compaction, DEFAULT_KEEP_RATIO, DEFAULT_THRESHOLD_RATIO,
|
||
};
|
||
use focus_harness::prompt::{default_context, ContextUsage, SystemPromptTemplate};
|
||
use focus_json::JsonValue;
|
||
use focus_providers::anthropic::AnthropicProvider;
|
||
use focus_providers::config::{resolve_context_window, ProviderConfig};
|
||
use std::sync::Arc;
|
||
|
||
/// 简单无副作用工具。
|
||
/// A trivial side-effect-free tool.
|
||
#[derive(Debug)]
|
||
struct NoopTool;
|
||
impl focus_core::Tool for NoopTool {
|
||
fn name(&self) -> &str {
|
||
"noop"
|
||
}
|
||
fn description(&self) -> &str {
|
||
"does nothing"
|
||
}
|
||
fn parameters(&self) -> JsonValue {
|
||
let mut o = JsonValue::obj();
|
||
o.insert("type", "object".into()).ok();
|
||
o
|
||
}
|
||
fn effects(&self) -> ToolEffects {
|
||
ToolEffects::NONE
|
||
}
|
||
fn execute(
|
||
&self,
|
||
_id: &str,
|
||
_args: &JsonValue,
|
||
_u: Option<&focus_core::tool::ToolUpdateSink>,
|
||
) -> Result<ToolResult, CoreError> {
|
||
Ok(ToolResult::text("ok"))
|
||
}
|
||
}
|
||
|
||
/// 构造一条「只输出文本」的 Anthropic SSE 回复。
|
||
/// Build an Anthropic SSE reply that only outputs text.
|
||
fn text_sse(text: &str) -> String {
|
||
format!(
|
||
"event: message_start\ndata: {{\"type\":\"message_start\",\"message\":{{\"usage\":{{\"input_tokens\":10}}}}}}\n\n\
|
||
event: content_block_start\ndata: {{\"type\":\"content_block_start\",\"index\":0,\"content_block\":{{\"type\":\"text\",\"text\":\"\"}}}}\n\n\
|
||
event: content_block_delta\ndata: {{\"type\":\"content_block_delta\",\"index\":0,\"delta\":{{\"type\":\"text_delta\",\"text\":\"{}\"}}}}\n\n\
|
||
event: content_block_stop\ndata: {{\"type\":\"content_block_stop\",\"index\":0}}\n\n\
|
||
event: message_delta\ndata: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}},\"usage\":{{\"output_tokens\":5}}}}\n\n\
|
||
event: message_stop\ndata: {{\"type\":\"message_stop\"}}\n\n",
|
||
text
|
||
)
|
||
}
|
||
|
||
fn agent_with(mock: common::MockTransport) -> Agent {
|
||
let provider =
|
||
AnthropicProvider::with_transport(ProviderConfig::new("sk-test"), Arc::new(mock));
|
||
let config = AgentConfig {
|
||
model: "claude-sonnet-4".into(),
|
||
system_prompt: "You are helpful.".into(),
|
||
max_tokens: Some(256),
|
||
..Default::default()
|
||
};
|
||
Agent::new(
|
||
config,
|
||
Box::new(provider),
|
||
ToolRegistry::with(Box::new(NoopTool)),
|
||
)
|
||
}
|
||
|
||
#[test]
|
||
fn full_conversation_compaction_and_continue() {
|
||
// ---- 阶段 1：跑几轮对话（mock 每次都回同样的话）。----
|
||
// ---- Phase 1: run a few turns (the mock replies the same text). ----
|
||
let mut mock = common::MockTransport::new();
|
||
for _ in 0..6 {
|
||
mock.push_body(text_sse("some answer"));
|
||
}
|
||
let mut agent = agent_with(mock);
|
||
let mut sink = VecSink::new();
|
||
for i in 0..6 {
|
||
agent
|
||
.prompt(format!("question {}", i), &mut sink)
|
||
.expect("run failed");
|
||
}
|
||
let messages = agent.messages().to_vec();
|
||
assert_eq!(messages.len(), 12); // 6 user + 6 assistant / 6 user + 6 assistant
|
||
|
||
// ---- 阶段 2：上下文窗口未知时，按三层解析取窗口。----
|
||
// ---- Phase 2: unknown window → resolve via the three-layer lookup. ----
|
||
let window = resolve_context_window(None, "claude-sonnet-4");
|
||
assert_eq!(window, 200_000);
|
||
|
||
// ---- 阶段 3：估算并规划压缩。----
|
||
// ---- Phase 3: estimate and plan compaction. ----
|
||
let estimated = focus_harness::estimate_messages(&messages);
|
||
assert!(estimated > 0);
|
||
// 小窗口强制触发压缩。
|
||
// Force a trigger with a tiny window.
|
||
let plan = plan_compaction(&messages, 64, DEFAULT_THRESHOLD_RATIO, DEFAULT_KEEP_RATIO)
|
||
.expect("compaction plan");
|
||
assert!(!plan.summarize.is_empty());
|
||
assert!(!plan.keep.is_empty());
|
||
|
||
// ---- 阶段 4：用 provider「生成」摘要（同样走 mock 流）。----
|
||
// ---- Phase 4: "generate" the summary via the provider (mock stream). ----
|
||
let mut mock2 = common::MockTransport::new();
|
||
mock2.push_body(text_sse(
|
||
"<summary>earlier talk</summary><key-facts>- f1</key-facts>",
|
||
));
|
||
let provider =
|
||
AnthropicProvider::with_transport(ProviderConfig::new("sk-test"), Arc::new(mock2));
|
||
let summary_request = ProviderRequest {
|
||
model: "claude-sonnet-4".into(),
|
||
system_prompt: plan.summary_instruction.clone(),
|
||
messages: plan.summarize.clone(),
|
||
tools: JsonValue::arr(),
|
||
max_tokens: Some(512),
|
||
temperature: None,
|
||
};
|
||
let events = common::collect(&provider, &summary_request);
|
||
let summary = match events.last().unwrap() {
|
||
StreamEvent::Done { message } => message
|
||
.content
|
||
.iter()
|
||
.find_map(|c| c.as_text())
|
||
.map(|t| t.text.clone())
|
||
.unwrap_or_default(),
|
||
other => panic!("expected done, got {:?}", other),
|
||
};
|
||
assert!(summary.contains("<summary>"), "got: {}", summary);
|
||
|
||
// ---- 阶段 5：写回摘要并注入动态系统提示（方案 E）。----
|
||
// ---- Phase 5: write the summary back and inject a dynamic system
|
||
// prompt (scheme E). ----
|
||
let compacted = apply_summary(&summary, &plan);
|
||
agent.replace_messages(compacted);
|
||
|
||
let mut ctx = default_context();
|
||
ctx.context_usage = Some(ContextUsage {
|
||
estimated_tokens: focus_harness::estimate_messages(agent.messages()),
|
||
context_window: window,
|
||
window_known: true,
|
||
});
|
||
let tpl = SystemPromptTemplate::default();
|
||
agent.set_system_prompt(tpl.render(&ctx));
|
||
assert!(agent.config().system_prompt.contains("Context usage"));
|
||
|
||
// ---- 阶段 6：压缩后继续对话。----
|
||
// ---- Phase 6: continue the conversation after compaction. ----
|
||
let mut mock3 = common::MockTransport::new();
|
||
mock3.push_body(text_sse("post-compaction answer"));
|
||
// 重新注入 provider（mock 是一次性的）。
|
||
// Re-inject a provider (the mock is single-use).
|
||
let provider =
|
||
AnthropicProvider::with_transport(ProviderConfig::new("sk-test"), Arc::new(mock3));
|
||
agent.replace_messages(agent.messages().to_vec());
|
||
// 通过内部字段替换 provider：直接重建 agent 以保留消息。
|
||
// Rebuild the agent keeping the transcript, since the provider is not
|
||
// swappable on an existing agent.
|
||
let mut agent2 = Agent::new(
|
||
AgentConfig {
|
||
model: "claude-sonnet-4".into(),
|
||
system_prompt: agent.config().system_prompt.clone(),
|
||
max_tokens: Some(256),
|
||
..Default::default()
|
||
},
|
||
Box::new(provider),
|
||
ToolRegistry::with(Box::new(NoopTool)),
|
||
);
|
||
agent2.replace_messages(agent.messages().to_vec());
|
||
|
||
let mut sink2 = VecSink::new();
|
||
agent2.prompt("what now?", &mut sink2).expect("run failed");
|
||
let msgs = agent2.messages();
|
||
// 摘要消息 + 保留消息 + 新 user + 新 assistant。
|
||
// Summary message + kept messages + new user + new assistant.
|
||
assert_eq!(
|
||
msgs.len(),
|
||
plan.keep.len() + 3,
|
||
"transcript: summary + keep + user + assistant"
|
||
);
|
||
let first = &msgs[0];
|
||
assert!(matches!(first, Message::User(_)));
|
||
let last = msgs.last().unwrap();
|
||
assert!(matches!(last, Message::Assistant(_)));
|
||
}
|