255 lines
9.6 KiB
Rust
255 lines
9.6 KiB
Rust
//! 后台任务：自动压缩 + 运行 agent（在独立线程中执行同步的 agent 循环）。
|
||
//! Background jobs: auto-compaction + running the agent (the synchronous agent
|
||
//! loop executes on a dedicated thread).
|
||
|
||
use crate::config::{ProviderKind, TuiConfig};
|
||
use focus_core::event::{AgentEvent, EventSink};
|
||
use focus_core::model::{ContentBlock, Message};
|
||
use focus_core::provider::{ProviderRequest, StreamEvent, StreamProvider};
|
||
use focus_core::tool::ToolRegistry;
|
||
use focus_core::Agent;
|
||
use focus_harness::compaction::{
|
||
apply_summary, plan_compaction, DEFAULT_KEEP_RATIO, DEFAULT_THRESHOLD_RATIO,
|
||
};
|
||
use focus_harness::estimate_messages;
|
||
use focus_providers::anthropic::AnthropicProvider;
|
||
use focus_providers::config::{resolve_context_window, ProviderConfig};
|
||
use focus_providers::openai::OpenAiProvider;
|
||
use focus_tools::{EditTool, ReadTool, ShellTool, WriteTool};
|
||
use std::path::PathBuf;
|
||
use std::sync::mpsc::Sender;
|
||
use std::thread::JoinHandle;
|
||
use std::time::Duration;
|
||
|
||
/// 压缩触发阈值（估算占用比例）。
|
||
/// Compaction trigger threshold (estimated occupancy ratio).
|
||
pub const COMPACT_THRESHOLD: f64 = DEFAULT_THRESHOLD_RATIO;
|
||
|
||
/// 后台任务发给 UI 的事件。
|
||
/// Events the background job sends to the UI.
|
||
#[derive(Debug)]
|
||
pub enum JobEvent {
|
||
/// agent 事件（转发）。
|
||
/// An agent event (forwarded).
|
||
Agent(AgentEvent),
|
||
/// 状态说明（如"正在压缩…"）。
|
||
/// A status note (e.g. "compacting…").
|
||
Note(String),
|
||
/// 任务结束。`messages` 为**最终完整 transcript**；`to_persist` 为需要新写入
|
||
/// 会话的消息（含压缩产生的摘要消息）；`appended_count` 为本轮 agent 实际追加
|
||
/// 的消息数（不含摘要消息，用于思考耗时映射）。
|
||
/// Job finished. `messages` is the **final full transcript**;
|
||
/// `to_persist` are the messages that must be newly written to the session
|
||
/// (including a compaction summary); `appended_count` is how many messages
|
||
/// the agent run actually appended (excluding the summary; used to map
|
||
/// thinking durations).
|
||
Done {
|
||
messages: Vec<Message>,
|
||
to_persist: Vec<Message>,
|
||
appended_count: usize,
|
||
error: Option<String>,
|
||
},
|
||
}
|
||
|
||
/// 一次后台任务（prompt / continue / 仅压缩）。
|
||
/// One background job (prompt / continue / compact-only).
|
||
pub struct RunJob {
|
||
/// 任务开始时的 transcript。
|
||
/// The transcript at job start.
|
||
pub transcript: Vec<Message>,
|
||
/// 用户消息（`None` 表示仅压缩，不运行 agent）。
|
||
/// The user message (`None` means compact-only, no agent run).
|
||
pub user_text: Option<String>,
|
||
/// TUI 配置（provider / 模型 / 端点）。
|
||
/// TUI config (provider / model / endpoints).
|
||
pub config: TuiConfig,
|
||
/// 项目根目录（工具的相对路径基准 + shell cwd）。
|
||
/// Project root (tools' relative-path base + shell cwd).
|
||
pub cwd: PathBuf,
|
||
/// 渲染好的系统提示（含 Usage/上下文占用）。
|
||
/// The rendered system prompt (with usage / context occupancy).
|
||
pub system_prompt: String,
|
||
/// 注入的 provider（测试用；`None` 时按配置构造）。
|
||
/// Injected provider (for tests; `None` builds one from the config).
|
||
pub provider: Option<Box<dyn StreamProvider>>,
|
||
}
|
||
|
||
/// 启动一个后台任务。
|
||
/// Spawn a background job.
|
||
pub fn spawn(job: RunJob, tx: Sender<JobEvent>) -> JoinHandle<()> {
|
||
std::thread::Builder::new()
|
||
.name("focus-tui-job".into())
|
||
.spawn(move || run_job(job, tx))
|
||
.expect("spawn tui job thread")
|
||
}
|
||
|
||
/// 执行任务。
|
||
/// Execute the job.
|
||
fn run_job(job: RunJob, tx: Sender<JobEvent>) {
|
||
let _ = tx.send(JobEvent::Note("starting…".into()));
|
||
|
||
// 0. provider：测试注入或按配置构造；摘要与 agent 共用。
|
||
// 0. Provider: injected (tests) or built from config; shared between the
|
||
// summarizer and the agent.
|
||
let provider = match job.provider {
|
||
Some(b) => b,
|
||
None => build_provider(&job.config),
|
||
};
|
||
|
||
// 1. 自动压缩（无用户确认，超过阈值即在回合间执行）。
|
||
// 1. Auto-compaction (no user confirmation; runs between turns when over
|
||
// the threshold).
|
||
let mut transcript = job.transcript.clone();
|
||
let mut compacted = false;
|
||
let window = resolve_context_window(job.config.context_window, &job.config.model);
|
||
let estimated = estimate_messages(&transcript);
|
||
let over = !transcript.is_empty() && (estimated as f64) >= (window as f64) * COMPACT_THRESHOLD;
|
||
if over {
|
||
let _ = tx.send(JobEvent::Note(format!(
|
||
"compacting… ({} tokens used of {})",
|
||
crate::format::format_tokens(estimated),
|
||
crate::format::format_tokens(window)
|
||
)));
|
||
if let Some(plan) =
|
||
plan_compaction(&transcript, window, COMPACT_THRESHOLD, DEFAULT_KEEP_RATIO)
|
||
{
|
||
match summarize(provider.as_ref(), &job.config, &plan) {
|
||
Ok(summary) => {
|
||
transcript = apply_summary(&summary, &plan);
|
||
compacted = true;
|
||
}
|
||
Err(e) => {
|
||
let _ = tx.send(JobEvent::Note(format!("compaction failed: {}", e)));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 2. 构建 agent 并运行（compaction 之后）。
|
||
// 2. Build the agent and run (after compaction).
|
||
let tools = build_tools(&job.cwd);
|
||
let config = focus_core::AgentConfig {
|
||
model: job.config.model.clone(),
|
||
system_prompt: job.system_prompt,
|
||
max_tokens: Some(4096),
|
||
..Default::default()
|
||
};
|
||
let mut agent = Agent::new(config, provider, tools);
|
||
agent.replace_messages(transcript.clone());
|
||
|
||
let mut sink = ChannelSink(tx.clone());
|
||
let result = match &job.user_text {
|
||
Some(text) => agent.prompt(text.clone(), &mut sink),
|
||
None => Ok(()), // 仅压缩任务，不运行 agent / compact-only job
|
||
};
|
||
|
||
let final_transcript = agent.messages().to_vec();
|
||
// 需要新写入会话的消息：压缩摘要（如有）+ 本轮追加的消息。
|
||
// Messages to newly persist: the compaction summary (if any) + the
|
||
// messages appended by this run.
|
||
let appended = final_transcript[transcript.len().min(final_transcript.len())..].to_vec();
|
||
let mut to_persist = Vec::new();
|
||
if compacted {
|
||
if let Some(first) = final_transcript.first() {
|
||
to_persist.push(first.clone());
|
||
}
|
||
}
|
||
to_persist.extend(appended.iter().cloned());
|
||
|
||
match result {
|
||
Ok(()) => {
|
||
let _ = tx.send(JobEvent::Done {
|
||
messages: final_transcript,
|
||
to_persist,
|
||
appended_count: appended.len(),
|
||
error: None,
|
||
});
|
||
}
|
||
Err(e) => {
|
||
let _ = tx.send(JobEvent::Done {
|
||
messages: final_transcript,
|
||
to_persist,
|
||
appended_count: appended.len(),
|
||
error: Some(e.to_string()),
|
||
});
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 用 provider 生成摘要（收集流事件中的 Done 文本）。
|
||
/// Generate a summary via the provider (collect the Done text from the stream).
|
||
fn summarize(
|
||
provider: &dyn StreamProvider,
|
||
config: &TuiConfig,
|
||
plan: &focus_harness::compaction::CompactionPlan,
|
||
) -> Result<String, String> {
|
||
let request = ProviderRequest {
|
||
model: config.model.clone(),
|
||
system_prompt: plan.summary_instruction.clone(),
|
||
messages: plan.summarize.clone(),
|
||
tools: focus_json::JsonValue::arr(),
|
||
max_tokens: Some(2000),
|
||
temperature: None,
|
||
};
|
||
let mut iter = provider.stream(&request).map_err(|e| e.to_string())?;
|
||
let mut text = String::new();
|
||
while let Some(ev) = iter.next_event() {
|
||
match ev {
|
||
StreamEvent::Done { message } => {
|
||
for block in &message.content {
|
||
if let ContentBlock::Text(t) = block {
|
||
text.push_str(&t.text);
|
||
}
|
||
}
|
||
break;
|
||
}
|
||
StreamEvent::Error { error } => {
|
||
return Err(error
|
||
.error_message
|
||
.unwrap_or_else(|| "summary stream errored".into()));
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
if text.trim().is_empty() {
|
||
return Err("summary stream produced no text".into());
|
||
}
|
||
Ok(text)
|
||
}
|
||
|
||
/// 依据配置构造 provider。
|
||
/// Build a provider from the config.
|
||
pub fn build_provider(config: &TuiConfig) -> Box<dyn StreamProvider> {
|
||
let pc = ProviderConfig {
|
||
api_key: config.api_key.clone(),
|
||
base_url: config.base_url.clone(),
|
||
context_window: config.context_window,
|
||
timeout: Duration::from_secs(120),
|
||
};
|
||
match config.provider {
|
||
ProviderKind::Anthropic => Box::new(AnthropicProvider::new(pc)),
|
||
ProviderKind::OpenAI => Box::new(OpenAiProvider::new(pc, config.openai_protocol)),
|
||
}
|
||
}
|
||
|
||
/// 依据项目根目录构造工具注册表。
|
||
/// Build the tool registry for a project root.
|
||
pub fn build_tools(cwd: &std::path::Path) -> ToolRegistry {
|
||
let mut registry = ToolRegistry::new();
|
||
registry.register(Box::new(ReadTool::new(cwd)));
|
||
registry.register(Box::new(WriteTool::new(cwd)));
|
||
registry.register(Box::new(EditTool::new(cwd)));
|
||
registry.register(Box::new(ShellTool::new(cwd)));
|
||
registry
|
||
}
|
||
|
||
/// 把 agent 事件转发到 UI channel 的 sink。
|
||
/// A sink that forwards agent events to the UI channel.
|
||
struct ChannelSink(Sender<JobEvent>);
|
||
|
||
impl EventSink for ChannelSink {
|
||
fn emit(&mut self, event: AgentEvent) {
|
||
let _ = self.0.send(JobEvent::Agent(event));
|
||
}
|
||
}
|