1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use futures::StreamExt as _;
5use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
6use tokio::sync::watch;
7use tokio_util::sync::CancellationToken;
8
9use crate::compaction::{self, CompactOutcome, CompactionPolicy};
10use crate::event::{AgentEvent, LlmStreamEvent};
11use crate::hooks::{HooksSet, PlanModeHandle};
12use crate::loop_::{LoopConfig, LoopError, run_loop};
13use crate::message::{AgentMessage, StopReason, Usage};
14use crate::session::{Session, SessionEntry};
15use crate::streamfn::{LlmOpts, StreamFn};
16use crate::tool::Tool;
17
18pub const PLAN_SECTION: &str = "# Plan mode\nYou are planning, not executing. Explore and propose only — do NOT modify the\nworkspace: the editing tools are disabled and mutating shell commands are refused.\nFine-tune the plan with the user first. When the plan is final, record it as a\ntodo list (status pending), then tell the user to run `/plan off` to execute.";
21
22pub struct AgentConfig {
23 pub system: String,
24 pub tools: Vec<Tool>,
25 pub llm: LlmOpts,
26 pub stream_fn: StreamFn,
27 pub hooks: HooksSet,
28 pub session: Option<Session>, pub context: Vec<AgentMessage>,
31 pub working_dir: PathBuf,
35 pub max_turns: usize,
39 pub parallel_tools: bool,
42 pub compaction: CompactionPolicy,
45 pub plan_mode: PlanModeHandle,
48}
49
50pub struct Agent {
51 system: String,
52 base_system: String,
54 tools: Vec<Tool>,
55 llm: LlmOpts,
56 stream_fn: StreamFn,
57 hooks: HooksSet,
58 session: Option<Session>,
59 working_dir: PathBuf,
60 max_turns: usize,
61 parallel_tools: bool,
62 compaction: CompactionPolicy,
63 plan_mode: PlanModeHandle,
64 ctx: Vec<AgentMessage>,
65 steer_tx: UnboundedSender<AgentMessage>,
66 steer_rx: Option<UnboundedReceiver<AgentMessage>>,
67 follow_tx: UnboundedSender<AgentMessage>,
68 follow_rx: Option<UnboundedReceiver<AgentMessage>>,
69 cancel: CancellationToken,
70 ctx_snapshot: watch::Sender<Vec<AgentMessage>>,
74}
75
76impl Agent {
77 pub fn new(mut cfg: AgentConfig) -> Agent {
78 cfg.llm.session_id = cfg
83 .session
84 .as_ref()
85 .and_then(session_header_id)
86 .or_else(|| cfg.llm.session_id.clone())
87 .or_else(|| Some(uuid::Uuid::new_v4().to_string()));
88 let (steer_tx, steer_rx) = mpsc::unbounded_channel();
89 let (follow_tx, follow_rx) = mpsc::unbounded_channel();
90 let (ctx_snapshot, _ctx_rx) = watch::channel(cfg.context.clone());
94 let working_dir = if cfg.working_dir.as_os_str().is_empty() {
96 default_working_dir()
97 } else {
98 cfg.working_dir
99 };
100 Agent {
101 system: cfg.system.clone(),
102 base_system: cfg.system,
103 tools: cfg.tools,
104 llm: cfg.llm,
105 stream_fn: cfg.stream_fn,
106 hooks: cfg.hooks,
107 session: cfg.session,
108 working_dir,
109 max_turns: cfg.max_turns,
110 parallel_tools: cfg.parallel_tools,
111 compaction: cfg.compaction,
112 plan_mode: cfg.plan_mode,
113 ctx: cfg.context,
114 steer_tx,
115 steer_rx: Some(steer_rx),
116 follow_tx,
117 follow_rx: Some(follow_rx),
118 cancel: CancellationToken::new(),
119 ctx_snapshot,
120 }
121 }
122
123 pub fn steer(&self, m: AgentMessage) {
125 let _ = self.steer_tx.send(m);
126 }
127
128 pub fn set_plan_mode(&mut self, on: bool) {
133 self.plan_mode.set(on);
134 self.system = if on {
135 format!("{}\n\n{}", self.base_system, PLAN_SECTION)
136 } else {
137 self.base_system.clone()
138 };
139 }
140 pub fn set_model(&mut self, model: String) -> std::io::Result<()> {
146 if let Some(session) = &mut self.session {
149 session.append(SessionEntry::ModelChange {
150 id: uuid::Uuid::new_v4().to_string(),
151 model: model.clone(),
152 })?;
153 }
154 self.llm.model = model;
155 self.llm.settle_provider();
158 Ok(())
159 }
160
161 pub fn effort(&self) -> Option<&str> {
163 self.llm.effort.as_deref()
164 }
165
166 pub fn set_effort(&mut self, effort: Option<String>) -> std::io::Result<()> {
169 if let Some(session) = &mut self.session {
170 session.append(SessionEntry::EffortChange {
171 id: uuid::Uuid::new_v4().to_string(),
172 effort: effort.clone(),
173 })?;
174 }
175 self.llm.effort = effort;
176 Ok(())
177 }
178
179 pub fn steer_sender(&self) -> UnboundedSender<AgentMessage> {
186 self.steer_tx.clone()
187 }
188
189 pub fn follow_up(&self, m: AgentMessage) {
191 let _ = self.follow_tx.send(m);
192 }
193
194 pub fn follow_up_sender(&self) -> UnboundedSender<AgentMessage> {
198 self.follow_tx.clone()
199 }
200
201 pub fn cancel(&self) {
202 self.cancel.cancel();
203 }
204
205 pub fn cancel_token(&self) -> CancellationToken {
208 self.cancel.clone()
209 }
210
211 pub fn hooks(&self) -> &HooksSet {
214 &self.hooks
215 }
216
217 pub async fn run(
218 &mut self,
219 user_text: &str,
220 sink: UnboundedSender<AgentEvent>,
221 ) -> Result<StopReason, LoopError> {
222 let user = AgentMessage::user_text(user_text);
227 if let Some(session) = &mut self.session {
228 let appended = session
229 .append(SessionEntry::Message {
230 id: uuid::Uuid::new_v4().to_string(),
231 parent_id: None,
232 message: user.clone(),
233 })
234 .map_err(LoopError::Session);
235 if let Err(e) = appended {
236 self.hooks.after_run(&self.ctx, StopReason::Error).await;
242 return Err(e);
243 }
244 }
245 self.ctx.push(user);
246
247 let _ = self.ctx_snapshot.send_replace(self.ctx.clone());
250
251 let steering = self.steer_rx.take().expect("steer_rx present between runs");
252 let follow_ups = self
253 .follow_rx
254 .take()
255 .expect("follow_rx present between runs");
256 let cfg = LoopConfig {
257 system: self.system.clone(),
258 tools: self.tools.clone(),
259 llm: self.llm.clone(),
260 stream_fn: Arc::clone(&self.stream_fn),
261 hooks: self.hooks.clone(),
262 steering,
263 follow_ups,
264 cancel: self.cancel.clone(),
265 working_dir: self.working_dir.clone(),
266 max_turns: self.max_turns,
267 parallel: self.parallel_tools,
268 compaction: self.compaction,
269 ctx_snapshot: self.ctx_snapshot.clone(),
270 session: self.session.as_mut(),
271 };
272 let res = run_loop(&mut self.ctx, cfg, sink).await;
273
274 self.cancel = CancellationToken::new();
277
278 match res {
279 Ok(result) => {
280 self.steer_rx = Some(result.steering);
285 self.follow_rx = Some(result.follow_ups);
286 self.hooks.after_run(&self.ctx, result.stop_reason).await;
289 Ok(result.stop_reason)
290 }
291 Err(e) => {
292 let (steer_tx, steer_rx) = mpsc::unbounded_channel();
296 self.steer_tx = steer_tx;
297 self.steer_rx = Some(steer_rx);
298 let (follow_tx, follow_rx) = mpsc::unbounded_channel();
299 self.follow_tx = follow_tx;
300 self.follow_rx = Some(follow_rx);
301 self.hooks.after_run(&self.ctx, StopReason::Error).await;
306 Err(e)
307 }
308 }
309 }
310
311 pub fn messages(&self) -> &[AgentMessage] {
312 &self.ctx
313 }
314
315 pub fn notify(&mut self, content: String) -> std::io::Result<AgentMessage> {
321 let message = AgentMessage::user_text(content);
322 if let Some(session) = &mut self.session {
323 session.append(SessionEntry::Message {
324 id: uuid::Uuid::new_v4().to_string(),
325 parent_id: None,
326 message: message.clone(),
327 })?;
328 }
329 self.ctx.push(message.clone());
330 Ok(message)
331 }
332
333 pub async fn side_ask(&self, text: &str) -> Result<SideAnswer, LoopError> {
346 side_ask_with(&self.stream_fn, &self.ctx, &self.system, &self.llm, text).await
347 }
348
349 pub fn side_handle(&self) -> SideHandle {
354 SideHandle {
355 stream_fn: Arc::clone(&self.stream_fn),
356 system: self.system.clone(),
357 llm: self.llm.clone(),
358 ctx: self.ctx_snapshot.clone(),
359 }
360 }
361
362 pub async fn compact(&mut self, instructions: Option<&str>) -> Result<CompactOutcome, String> {
371 compaction::compact_ctx(
372 &mut self.ctx,
373 &self.compaction,
374 &self.stream_fn,
375 &self.llm,
376 instructions,
377 self.session.as_mut(),
378 )
379 .await
380 }
381
382 pub fn session_path(&self) -> Option<&Path> {
383 self.session.as_ref().and_then(|s| s.path())
384 }
385
386 pub fn session_id(&self) -> Option<String> {
389 self.session.as_ref().and_then(session_header_id)
390 }
391
392 pub(crate) fn llm(&self) -> &LlmOpts {
396 &self.llm
397 }
398}
399
400#[derive(Clone, Debug, Default, PartialEq)]
403pub struct SideAnswer {
404 pub text: String,
405 pub usage: Option<Usage>,
406}
407
408#[derive(Clone)]
414pub struct SideHandle {
415 stream_fn: StreamFn,
416 system: String,
417 llm: LlmOpts,
418 ctx: watch::Sender<Vec<AgentMessage>>,
419}
420
421impl SideHandle {
422 pub fn history(&self) -> Vec<AgentMessage> {
425 self.ctx.subscribe().borrow().to_vec()
426 }
427
428 pub fn last_assistant_text(&self) -> Option<String> {
432 let rx = self.ctx.subscribe();
433 let messages = rx.borrow();
434 crate::message::last_assistant_text(&messages[..])
435 }
436
437 pub(crate) fn llm(&self) -> &LlmOpts {
440 &self.llm
441 }
442
443 pub async fn side_ask(&self, text: &str) -> Result<SideAnswer, LoopError> {
446 side_ask_with(&self.stream_fn, &self.history(), &self.system, &self.llm, text).await
447 }
448}
449
450async fn side_ask_with(
455 stream_fn: &StreamFn,
456 ctx: &[AgentMessage],
457 system: &str,
458 llm: &LlmOpts,
459 text: &str,
460) -> Result<SideAnswer, LoopError> {
461 let mut msgs = ctx.to_vec();
462 msgs.push(AgentMessage::user_text(frame_btw(text)));
463 let mut stream = (stream_fn)(&msgs, system, &[], llm);
464 let mut out = String::new();
465 let mut usage = None;
466 while let Some(ev) = stream.next().await {
467 match ev {
468 LlmStreamEvent::TextDelta(delta) => out.push_str(&delta),
469 LlmStreamEvent::Done { usage: u, .. } => {
470 usage = u;
471 break;
472 }
473 LlmStreamEvent::Error { message, .. } => return Err(LoopError::Stream(message)),
474 _ => {}
475 }
476 }
477 let text = out.trim().to_string();
478 if text.is_empty() {
479 return Err(LoopError::Stream("side answer produced no text".to_string()));
480 }
481 Ok(SideAnswer { text, usage })
482}
483
484fn frame_btw(q: &str) -> String {
487 format!("[side question — no tools available; answer briefly and directly from the conversation above]\n{q}")
488}
489
490fn default_working_dir() -> PathBuf {
492 std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
493}
494
495fn session_header_id(session: &Session) -> Option<String> {
497 session.entries().iter().find_map(|e| match e {
498 SessionEntry::Header { id, .. } => Some(id.clone()),
499 _ => None,
500 })
501}
502
503#[cfg(test)]
504mod tests {
505 use super::*;
506 use crate::compaction::CompactionPolicy;
507 use crate::hooks::HooksSet;
508 use crate::streamfn::{LlmStream, StreamFn};
509
510 fn scripted(script: Vec<LlmStreamEvent>) -> StreamFn {
512 Arc::new(move |_ctx: &[AgentMessage], _system, _tools, _opts: &LlmOpts| {
513 Box::pin(futures::stream::iter(script.clone())) as LlmStream
514 })
515 }
516
517 fn make_agent(session: Option<Session>, stream_fn: StreamFn, context: Vec<AgentMessage>) -> Agent {
518 Agent::new(AgentConfig {
519 system: "sys".into(),
520 tools: vec![],
521 llm: LlmOpts {
522 model: "m1".into(),
523 ..LlmOpts::default()
524 },
525 stream_fn,
526 hooks: HooksSet::default(),
527 session,
528 context,
529 working_dir: PathBuf::new(),
530 max_turns: crate::loop_::DEFAULT_MAX_TURNS,
531 parallel_tools: true,
532 compaction: CompactionPolicy::default(),
533 plan_mode: PlanModeHandle::new(),
534 })
535 }
536
537 #[tokio::test]
538 async fn side_ask_leaves_ctx_and_session_untouched() {
539 let dir = tempfile::tempdir().unwrap();
540 let session = Session::create(dir.path()).unwrap();
541 let stream = scripted(vec![
542 LlmStreamEvent::TextDelta("an".into()),
543 LlmStreamEvent::TextDelta("swer".into()),
544 LlmStreamEvent::Done {
545 stop_reason: StopReason::Stop,
546 usage: None,
547 },
548 ]);
549 let agent = make_agent(Some(session), stream, vec![AgentMessage::user_text("earlier")]);
550
551 let ctx_before = agent.ctx.clone();
552 let path = agent.session.as_ref().unwrap().path().unwrap().to_path_buf();
553 let entries_before = agent.session.as_ref().unwrap().entries().to_vec();
554 let bytes_before = std::fs::read(&path).unwrap();
555
556 let answer = agent.side_ask("why?").await.unwrap();
557 assert_eq!(answer.text, "answer");
558
559 assert_eq!(agent.ctx, ctx_before, "ctx is unchanged");
560 assert_eq!(agent.messages(), &ctx_before[..]);
561 assert_eq!(
562 agent.session.as_ref().unwrap().entries(),
563 &entries_before[..],
564 "session entries unchanged"
565 );
566 assert_eq!(
567 std::fs::read(&path).unwrap(),
568 bytes_before,
569 "session file bytes unchanged"
570 );
571 }
572
573 #[tokio::test]
574 async fn side_ask_folds_deltas_and_captures_usage() {
575 let usage = Usage {
576 input_tokens: 7,
577 output_tokens: 3,
578 cache_read_tokens: Some(2),
579 cache_write_tokens: None,
580 };
581 let stream = scripted(vec![
582 LlmStreamEvent::TextDelta("he".into()),
583 LlmStreamEvent::TextDelta("llo".into()),
584 LlmStreamEvent::Done {
585 stop_reason: StopReason::Stop,
586 usage: Some(usage),
587 },
588 ]);
589 let agent = make_agent(None, stream, vec![]);
590 assert_eq!(
591 agent.side_ask("hi").await.unwrap(),
592 SideAnswer {
593 text: "hello".into(),
594 usage: Some(usage),
595 }
596 );
597 }
598
599 #[tokio::test]
600 async fn side_ask_surfaces_a_stream_error() {
601 let stream = scripted(vec![LlmStreamEvent::Error {
602 message: "boom".into(),
603 fatal: true,
604 }]);
605 let agent = make_agent(None, stream, vec![]);
606 let err = agent.side_ask("hi").await.unwrap_err();
607 assert!(err.to_string().contains("boom"), "{err}");
608 }
609
610 #[tokio::test]
611 async fn side_ask_errors_on_empty_output() {
612 let stream = scripted(vec![LlmStreamEvent::Done {
613 stop_reason: StopReason::Stop,
614 usage: None,
615 }]);
616 let agent = make_agent(None, stream, vec![]);
617 let err = agent.side_ask("hi").await.unwrap_err();
618 assert!(err.to_string().contains("no text"), "{err}");
619 }
620
621 #[test]
622 fn frame_btw_wraps_the_question() {
623 let framed = frame_btw("why?");
624 assert!(framed.starts_with("[side question"));
625 assert!(framed.ends_with("\nwhy?"));
626 }
627
628 #[test]
629 fn set_plan_mode_appends_and_restores_the_section() {
630 let mut agent = make_agent(None, scripted(vec![]), vec![]);
631 let base = agent.system.clone();
632
633 agent.set_plan_mode(true);
634 assert!(agent.plan_mode.get());
635 assert_eq!(agent.system, format!("{base}\n\n{PLAN_SECTION}"));
636 assert!(agent.system.contains("# Plan mode"));
637 assert!(
638 agent.system.contains("record it as a\ntodo list (status pending)"),
639 "the P2 plan\u{2192}todo instruction is present"
640 );
641
642 agent.set_plan_mode(false);
643 assert!(!agent.plan_mode.get());
644 assert_eq!(agent.system, base, "restored exactly");
645 }
646}