502 lines
17 KiB
Rust
502 lines
17 KiB
Rust
//! Turn state, one stdin owner, typed events, and durable checkpoints.
|
||
use crate::{ChatMessage, Session};
|
||
use anyhow::{anyhow, Result};
|
||
use serde::{Deserialize, Serialize};
|
||
use serde_json::{json, Value};
|
||
use std::{
|
||
collections::VecDeque,
|
||
future::Future,
|
||
io::BufRead,
|
||
path::PathBuf,
|
||
sync::{Arc, Mutex},
|
||
time::Duration,
|
||
};
|
||
use tokio::sync::{oneshot, watch, Notify};
|
||
|
||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||
#[serde(rename_all = "snake_case")]
|
||
pub enum StepStatus {
|
||
Pending,
|
||
InProgress,
|
||
Completed,
|
||
}
|
||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||
pub struct PlanStep {
|
||
pub step: String,
|
||
pub status: StepStatus,
|
||
}
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
#[serde(tag = "type", rename_all = "snake_case")]
|
||
pub enum AgentEvent {
|
||
Status {
|
||
message: String,
|
||
},
|
||
Progress {
|
||
message: String,
|
||
},
|
||
PlanUpdated {
|
||
plan: Vec<PlanStep>,
|
||
explanation: Option<String>,
|
||
},
|
||
ToolStarted {
|
||
id: String,
|
||
name: String,
|
||
},
|
||
ToolFinished {
|
||
id: String,
|
||
name: String,
|
||
success: bool,
|
||
},
|
||
Waiting {
|
||
stage: String,
|
||
elapsed_secs: u64,
|
||
},
|
||
Question {
|
||
question: Value,
|
||
},
|
||
Steering {
|
||
message: String,
|
||
},
|
||
TurnEnded {
|
||
verdict: String,
|
||
message: String,
|
||
},
|
||
}
|
||
|
||
#[derive(Default)]
|
||
struct InputState {
|
||
lines: VecDeque<String>,
|
||
steering: VecDeque<String>,
|
||
question: Option<oneshot::Sender<Option<String>>>,
|
||
eof: bool,
|
||
running: bool,
|
||
}
|
||
pub struct InputBroker {
|
||
state: Mutex<InputState>,
|
||
notify: Notify,
|
||
cancel: watch::Sender<bool>,
|
||
persistent: bool,
|
||
}
|
||
impl std::fmt::Debug for InputBroker {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
f.write_str("InputBroker")
|
||
}
|
||
}
|
||
impl InputBroker {
|
||
pub fn new() -> Arc<Self> {
|
||
Self::with_persistence(false)
|
||
}
|
||
pub(crate) fn persistent() -> Arc<Self> {
|
||
Self::with_persistence(true)
|
||
}
|
||
fn with_persistence(persistent: bool) -> Arc<Self> {
|
||
Arc::new(Self {
|
||
persistent,
|
||
state: Mutex::new(InputState::default()),
|
||
notify: Notify::new(),
|
||
cancel: watch::channel(false).0,
|
||
})
|
||
}
|
||
pub fn stdin() -> Arc<Self> {
|
||
let broker = Self::new();
|
||
let reader = broker.clone();
|
||
// This is the only stdin reader for run/agent, including human tools.
|
||
std::thread::spawn(move || {
|
||
for line in std::io::stdin().lock().lines() {
|
||
match line {
|
||
Ok(line) => reader.feed(line),
|
||
Err(_) => break,
|
||
}
|
||
}
|
||
reader.close();
|
||
});
|
||
broker
|
||
}
|
||
pub(crate) fn can_ask(&self) -> bool {
|
||
!self.state.lock().unwrap().eof
|
||
}
|
||
pub(crate) fn has_question(&self) -> bool {
|
||
self.state.lock().unwrap().question.is_some()
|
||
}
|
||
pub fn feed(&self, line: String) {
|
||
let mut state = self.state.lock().unwrap();
|
||
if state.running && matches!(line.trim(), "/stop" | "/exit" | "/quit") {
|
||
self.cancel.send_replace(true);
|
||
return;
|
||
}
|
||
if let Some(answer) = state.question.take() {
|
||
let _ = answer.send(Some(line));
|
||
} else if state.running {
|
||
state.steering.push_back(line);
|
||
} else {
|
||
state.lines.push_back(line);
|
||
}
|
||
drop(state);
|
||
self.notify.notify_one();
|
||
}
|
||
pub fn close(&self) {
|
||
let mut state = self.state.lock().unwrap();
|
||
state.eof = true;
|
||
if let Some(answer) = state.question.take() {
|
||
let _ = answer.send(None);
|
||
}
|
||
drop(state);
|
||
self.notify.notify_one();
|
||
}
|
||
pub fn begin(&self) {
|
||
self.cancel.send_replace(false);
|
||
self.state.lock().unwrap().running = true;
|
||
}
|
||
pub fn end(&self) {
|
||
let mut s = self.state.lock().unwrap();
|
||
s.running = false;
|
||
s.question.take();
|
||
let rest = s.steering.drain(..).collect::<Vec<_>>();
|
||
s.lines.extend(rest);
|
||
}
|
||
pub fn interrupt(&self) {
|
||
if self.state.lock().unwrap().running {
|
||
self.cancel.send_replace(true);
|
||
} else {
|
||
self.feed("/exit".into());
|
||
}
|
||
}
|
||
pub fn cancelled(&self) -> bool {
|
||
*self.cancel.borrow()
|
||
}
|
||
pub async fn cancellation(&self) {
|
||
let mut rx = self.cancel.subscribe();
|
||
loop {
|
||
if *rx.borrow_and_update() {
|
||
return;
|
||
}
|
||
if rx.changed().await.is_err() {
|
||
return;
|
||
}
|
||
}
|
||
}
|
||
pub fn drain(&self) -> Vec<String> {
|
||
self.state
|
||
.lock()
|
||
.unwrap()
|
||
.steering
|
||
.drain(..)
|
||
.filter(|s| !s.trim().is_empty())
|
||
.collect()
|
||
}
|
||
pub async fn next(&self) -> Option<String> {
|
||
loop {
|
||
let notified = self.notify.notified();
|
||
{
|
||
let mut s = self.state.lock().unwrap();
|
||
if let Some(line) = s.lines.pop_front() {
|
||
return Some(line);
|
||
}
|
||
if s.eof {
|
||
return None;
|
||
}
|
||
}
|
||
notified.await;
|
||
}
|
||
}
|
||
pub async fn ask(&self, seconds: u64, on_ready: impl FnOnce()) -> Result<String> {
|
||
let (tx, rx) = oneshot::channel();
|
||
{
|
||
let mut s = self.state.lock().unwrap();
|
||
if s.eof {
|
||
drop(s);
|
||
on_ready();
|
||
return Err(anyhow!("stdin closed; answer in a resumed session"));
|
||
}
|
||
s.question = Some(tx);
|
||
}
|
||
on_ready();
|
||
let result = tokio::select! {
|
||
biased;
|
||
_ = self.cancellation() => Err(anyhow!("cancelled")),
|
||
answer = async { if self.persistent { Ok(rx.await) } else { tokio::time::timeout(Duration::from_secs(seconds.clamp(1, 3600)), rx).await } } => match answer {
|
||
Ok(Ok(Some(line))) => Ok(line),
|
||
Ok(_) => Err(anyhow!("stdin closed")),
|
||
Err(_) => Err(anyhow!("user input timed out")),
|
||
}
|
||
};
|
||
self.state.lock().unwrap().question.take();
|
||
result
|
||
}
|
||
}
|
||
|
||
type CheckpointSink = Arc<dyn Fn(&Session) -> Result<()> + Send + Sync>;
|
||
type EventSink = Arc<dyn Fn(&AgentEvent) + Send + Sync>;
|
||
|
||
#[derive(Default)]
|
||
pub struct Runtime {
|
||
pub input: Option<Arc<InputBroker>>,
|
||
pub plan: Mutex<Vec<PlanStep>>,
|
||
pub pending_question: Mutex<Option<Value>>,
|
||
pub active_command: Mutex<Option<Value>>,
|
||
pub browser_url: Mutex<Option<String>>,
|
||
pub(crate) browser_profile: Mutex<Option<PathBuf>>,
|
||
checkpoint: Mutex<Option<Session>>,
|
||
pub events: Mutex<Vec<AgentEvent>>,
|
||
event_sink: Mutex<Option<EventSink>>,
|
||
checkpoint_sink: Mutex<Option<CheckpointSink>>,
|
||
}
|
||
impl std::fmt::Debug for Runtime {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
f.write_str("Runtime")
|
||
}
|
||
}
|
||
impl Runtime {
|
||
pub fn with_input(input: Arc<InputBroker>) -> Arc<Self> {
|
||
Arc::new(Self {
|
||
input: Some(input),
|
||
..Default::default()
|
||
})
|
||
}
|
||
pub fn for_session(session: &Session, input: Arc<InputBroker>) -> Arc<Self> {
|
||
Arc::new(Self {
|
||
input: Some(input),
|
||
plan: Mutex::new(session.plan.clone()),
|
||
pending_question: Mutex::new(session.pending_question.clone()),
|
||
active_command: Mutex::new(session.active_command.clone()),
|
||
browser_url: Mutex::new(session.last_browser_url.clone()),
|
||
browser_profile: Mutex::new(None),
|
||
checkpoint: Mutex::new(Some(session.clone())),
|
||
events: Mutex::new(vec![]),
|
||
event_sink: Mutex::new(None),
|
||
checkpoint_sink: Mutex::new(None),
|
||
})
|
||
}
|
||
pub(crate) fn set_checkpoint_handler(
|
||
&self,
|
||
handler: impl Fn(&Session) -> Result<()> + Send + Sync + 'static,
|
||
) {
|
||
*self.checkpoint_sink.lock().unwrap() = Some(Arc::new(handler));
|
||
}
|
||
fn save_checkpoint(&self, session: &Session) -> Result<()> {
|
||
if let Some(sink) = self.checkpoint_sink.lock().unwrap().clone() {
|
||
sink(session)
|
||
} else {
|
||
crate::save_session(session).map(|_| ())
|
||
}
|
||
}
|
||
pub fn set_event_handler(&self, handler: impl Fn(&AgentEvent) + Send + Sync + 'static) {
|
||
*self.event_sink.lock().unwrap() = Some(Arc::new(handler));
|
||
}
|
||
pub fn emit(&self, event: AgentEvent) {
|
||
let sink = self.event_sink.lock().unwrap().clone();
|
||
if let Some(sink) = sink {
|
||
sink(&event);
|
||
}
|
||
self.events.lock().unwrap().push(event);
|
||
}
|
||
pub async fn wait<T>(&self, stage: &str, future: impl Future<Output = Result<T>>) -> Result<T> {
|
||
tokio::pin!(future);
|
||
let start = std::time::Instant::now();
|
||
let mut tick = tokio::time::interval(Duration::from_secs(20));
|
||
tick.tick().await;
|
||
loop {
|
||
tokio::select! {
|
||
biased;
|
||
_ = async { match &self.input { Some(i) => i.cancellation().await, None => std::future::pending().await } } => return Err(anyhow!("cancelled")),
|
||
result = &mut future => return result,
|
||
_ = tick.tick() => self.emit(AgentEvent::Waiting { stage: stage.into(), elapsed_secs: start.elapsed().as_secs() }),
|
||
}
|
||
}
|
||
}
|
||
pub fn cancelled(&self) -> bool {
|
||
self.input.as_ref().is_some_and(|i| i.cancelled())
|
||
}
|
||
pub fn steering(&self) -> Vec<String> {
|
||
self.input.as_ref().map(|i| i.drain()).unwrap_or_default()
|
||
}
|
||
pub fn update_plan(&self, args: &Value) -> Result<Value> {
|
||
let next: Vec<PlanStep> = serde_json::from_value(args["plan"].clone())?;
|
||
if next.is_empty()
|
||
|| next.len() > 12
|
||
|| next.iter().any(|s| s.step.trim().is_empty())
|
||
|| next
|
||
.iter()
|
||
.filter(|s| s.status == StepStatus::InProgress)
|
||
.count()
|
||
> 1
|
||
{
|
||
return Err(anyhow!(
|
||
"plan requires 1–12 nonempty steps, at most one in_progress"
|
||
));
|
||
}
|
||
let mut current = self.plan.lock().unwrap();
|
||
let explanation = args["explanation"]
|
||
.as_str()
|
||
.filter(|s| !s.trim().is_empty())
|
||
.map(str::to_string);
|
||
if !current.is_empty()
|
||
&& current.iter().map(|s| &s.step).collect::<Vec<_>>()
|
||
!= next.iter().map(|s| &s.step).collect::<Vec<_>>()
|
||
&& explanation.is_none()
|
||
{
|
||
return Err(anyhow!("explain why the plan steps changed"));
|
||
}
|
||
*current = next.clone();
|
||
drop(current);
|
||
self.emit(AgentEvent::PlanUpdated {
|
||
plan: next,
|
||
explanation,
|
||
});
|
||
Ok(json!({"updated":true}))
|
||
}
|
||
pub fn unfinished(&self) -> bool {
|
||
self.plan
|
||
.lock()
|
||
.unwrap()
|
||
.iter()
|
||
.any(|s| s.status != StepStatus::Completed)
|
||
}
|
||
pub fn checkpoint(&self, messages: &[ChatMessage], pending_tool: Option<&str>) -> Result<()> {
|
||
let mut guard = self.checkpoint.lock().unwrap();
|
||
if let Some(s) = guard.as_mut() {
|
||
s.messages = messages.to_vec();
|
||
s.plan = self.plan.lock().unwrap().clone();
|
||
s.pending_question = self.pending_question.lock().unwrap().clone();
|
||
s.pending_tool = pending_tool.map(str::to_string);
|
||
s.active_command = self.active_command.lock().unwrap().clone();
|
||
s.last_browser_url = self
|
||
.browser_url
|
||
.lock()
|
||
.unwrap()
|
||
.clone()
|
||
.or(s.last_browser_url.clone());
|
||
s.touch();
|
||
self.save_checkpoint(s)?;
|
||
}
|
||
Ok(())
|
||
}
|
||
pub fn sync_session(&self, session: &mut Session) {
|
||
session.plan = self.plan.lock().unwrap().clone();
|
||
session.pending_question = self.pending_question.lock().unwrap().clone();
|
||
session.pending_tool = None;
|
||
session.active_command = self.active_command.lock().unwrap().clone();
|
||
}
|
||
pub fn profile_dir(&self) -> Option<PathBuf> {
|
||
if let Some(path) = self.browser_profile.lock().unwrap().clone() {
|
||
return Some(path);
|
||
}
|
||
self.checkpoint.lock().unwrap().as_ref().and_then(|s| {
|
||
(if self.checkpoint_sink.lock().unwrap().is_some() {
|
||
Ok(crate::team::data_dir().join("profiles"))
|
||
} else {
|
||
crate::sessions_dir()
|
||
})
|
||
.ok()
|
||
.map(|d| d.join(format!("{}.browser", s.id)))
|
||
})
|
||
}
|
||
pub fn save_output(&self, output: &str) -> Result<Option<String>> {
|
||
if output.len() <= 16000 {
|
||
return Ok(None);
|
||
}
|
||
let guard = self.checkpoint.lock().unwrap();
|
||
let Some(s) = guard.as_ref() else {
|
||
return Ok(None);
|
||
};
|
||
// Artifacts belong to the workspace so read_file can access them.
|
||
let dir = s.cwd.join(".grokboy-output").join(&s.id);
|
||
std::fs::create_dir_all(&dir)?;
|
||
let path = dir.join(format!("{}.txt", uuid::Uuid::new_v4()));
|
||
std::fs::write(&path, output)?;
|
||
Ok(Some(json!({"preview":output.chars().take(4000).collect::<String>(),"output_file":path,"bytes":output.len(),"truncated":true,"hint":"read_file with offset/limit to inspect full output"}).to_string()))
|
||
}
|
||
pub async fn question(&self, args: &Value) -> Result<Value> {
|
||
let input = self
|
||
.input
|
||
.as_ref()
|
||
.ok_or_else(|| anyhow!("no interactive input channel"))?;
|
||
*self.pending_question.lock().unwrap() = Some(args.clone());
|
||
// Persist pending question against latest pre-tool checkpoint.
|
||
{
|
||
let mut guard = self.checkpoint.lock().unwrap();
|
||
if let Some(s) = guard.as_mut() {
|
||
s.pending_question = Some(args.clone());
|
||
s.last_browser_url = self
|
||
.browser_url
|
||
.lock()
|
||
.unwrap()
|
||
.clone()
|
||
.or(s.last_browser_url.clone());
|
||
self.save_checkpoint(s)?;
|
||
}
|
||
}
|
||
let line = input
|
||
.ask(args["timeout_secs"].as_u64().unwrap_or(300), || {
|
||
self.emit(AgentEvent::Question {
|
||
question: args.clone(),
|
||
})
|
||
})
|
||
.await?;
|
||
*self.pending_question.lock().unwrap() = None;
|
||
let answer = line
|
||
.trim()
|
||
.parse::<usize>()
|
||
.ok()
|
||
.and_then(|n| n.checked_sub(1))
|
||
.or_else(|| {
|
||
let text = line.trim().to_ascii_lowercase();
|
||
if text.len() == 1 {
|
||
text.as_bytes()[0]
|
||
.checked_sub(b'a')
|
||
.filter(|n| *n < 26)
|
||
.map(usize::from)
|
||
} else {
|
||
None
|
||
}
|
||
})
|
||
.and_then(|n| args["options"].get(n))
|
||
.and_then(Value::as_str)
|
||
.unwrap_or(line.trim())
|
||
.to_string();
|
||
Ok(json!({"answer":answer}))
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
#[tokio::test]
|
||
async fn queued_tasks_steering_and_answers_do_not_steal_each_other() {
|
||
let input = InputBroker::new();
|
||
input.feed("task one".into());
|
||
input.feed("task two".into());
|
||
assert_eq!(input.next().await.as_deref(), Some("task one"));
|
||
input.begin();
|
||
input.feed("keep original goal, add constraint".into());
|
||
assert_eq!(input.drain(), vec!["keep original goal, add constraint"]);
|
||
let answer = input.ask(2, || input.feed("answer".into())).await.unwrap();
|
||
assert_eq!(answer, "answer");
|
||
assert!(input.drain().is_empty());
|
||
input.end();
|
||
assert_eq!(input.next().await.as_deref(), Some("task two"));
|
||
}
|
||
#[tokio::test]
|
||
async fn cancellation_interrupts_a_pending_question_without_losing_next_input() {
|
||
let input = InputBroker::new();
|
||
input.begin();
|
||
assert!(input.ask(30, || input.interrupt()).await.is_err());
|
||
input.end();
|
||
input.feed("resume".into());
|
||
assert_eq!(input.next().await.as_deref(), Some("resume"));
|
||
}
|
||
#[test]
|
||
fn plan_changes_require_explanation_and_one_active_step() {
|
||
let runtime = Runtime::default();
|
||
runtime
|
||
.update_plan(&json!({"plan":[{"step":"inspect","status":"in_progress"}]}))
|
||
.unwrap();
|
||
assert!(runtime
|
||
.update_plan(&json!({"plan":[{"step":"other","status":"in_progress"}]}))
|
||
.is_err());
|
||
assert!(runtime.update_plan(&json!({"explanation":"new discovery","plan":[{"step":"other","status":"in_progress"},{"step":"third","status":"in_progress"}]})).is_err());
|
||
runtime.update_plan(&json!({"explanation":"new discovery","plan":[{"step":"other","status":"completed"}]})).unwrap();
|
||
assert!(!runtime.unfinished());
|
||
}
|
||
}
|