From 1045bf86388be53ecedf95921acddca62b287560 Mon Sep 17 00:00:00 2001 From: Assaf Sapir Date: Mon, 2 Mar 2026 16:02:26 +0200 Subject: [PATCH 1/2] feat: add exit tool for agent-initiated session exit Add an ExitTool that the LLM can invoke when the user wants to quit (e.g. says 'bye', 'goodbye', 'I'm done'). The tool uses a shared AtomicBool flag checked by the REPL after each engine run. - src/tools/exit.rs: ExitTool with triggered()/reset() API - src/tools/mod.rs: register exit module - src/main.rs: register ExitTool, check flag after engine.run() - tests/tools_test.rs: integration tests for registry interaction - 6 unit tests + 2 integration tests added - AGENTS.md and README.md updated --- AGENTS.md | 2 +- README.md | 2 +- src/main.rs | 10 +++++ src/tools/exit.rs | 102 ++++++++++++++++++++++++++++++++++++++++++++ src/tools/mod.rs | 1 + tests/tools_test.rs | 41 +++++++++++++++++- 6 files changed, 155 insertions(+), 3 deletions(-) create mode 100644 src/tools/exit.rs diff --git a/AGENTS.md b/AGENTS.md index f80114b..9f26253 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,7 +66,7 @@ src/ │ ├── gemini.rs # Google Gemini generateContent API │ ├── human.rs # Human-in-the-loop thinker │ └── mock.rs # MockThinker for tests -├── tools/ # Tool trait + ToolRegistry + ShellTool +├── tools/ # Tool trait + ToolRegistry + ShellTool + ExitTool ├── memory/ # Memory trait + SqliteMemory (task + session memory) └── spinner.rs # Thinking spinner during LLM calls ``` diff --git a/README.md b/README.md index 41ef440..6a4d2ad 100644 --- a/README.md +++ b/README.md @@ -142,7 +142,7 @@ Everything is a trait. Everything is swappable. - **`ProviderConfig`** — defines a provider's identity, auth flow, and thinker construction - **`Engine`** — the outermost boundary (`fn run(task) -> answer`) - **`Thinker`** — the brain (Anthropic, Gemini, human, mock — picked via `--provider`) -- **`Tool`** — something the agent can do (shell commands, more coming) +- **`Tool`** — something the agent can do (shell commands, exit, more coming) - **`Command`** — built-in REPL commands (`/help`, `/model`, `/new`, etc.) - **`Memory`** — what the agent remembers (task iterations + session history, SQLite-backed) - **`Config`** — persistent key-value settings (model preference, etc.) diff --git a/src/main.rs b/src/main.rs index 7331fd4..8c544ee 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ use golem::engine::react::{ReactConfig, ReactEngine}; use golem::memory::sqlite::SqliteMemory; use golem::provider::{LoginProvider, Provider, build_provider, handle_login, handle_logout}; use golem::tools::ToolRegistry; +use golem::tools::exit::ExitTool; use golem::tools::shell::{ShellConfig, ShellMode, ShellTool}; #[derive(Parser)] @@ -147,6 +148,10 @@ async fn main() -> anyhow::Result<()> { let tools = Arc::new(ToolRegistry::new()); tools.register(Arc::new(ShellTool::new(shell_config))).await; + let exit_tool = Arc::new(ExitTool::new()); + tools + .register(Arc::clone(&exit_tool) as Arc) + .await; // Collect tool names for /tools command let tool_names: Vec = tools @@ -257,6 +262,11 @@ async fn main() -> anyhow::Result<()> { println!("\n\ninterrupted"); } } + + // Check if the agent invoked the exit tool + if exit_tool.triggered() { + break; + } } print_session_summary(engine.session_usage()); diff --git a/src/tools/exit.rs b/src/tools/exit.rs new file mode 100644 index 0000000..8e5ee81 --- /dev/null +++ b/src/tools/exit.rs @@ -0,0 +1,102 @@ +use anyhow::Result; +use async_trait::async_trait; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use super::Tool; + +/// A tool the agent can invoke to exit the REPL. +/// +/// When the user says something like "bye" or "quit", the LLM can call this +/// tool to signal a graceful shutdown. The caller checks [`ExitTool::triggered`] +/// after each engine run to decide whether to break the loop. +pub struct ExitTool { + flag: Arc, +} + +impl ExitTool { + pub fn new() -> Self { + Self { + flag: Arc::new(AtomicBool::new(false)), + } + } + + /// Returns `true` if the tool was invoked (the agent wants to exit). + pub fn triggered(&self) -> bool { + self.flag.load(Ordering::Relaxed) + } + + /// Reset the flag (useful if the caller decides not to exit after all). + pub fn reset(&self) { + self.flag.store(false, Ordering::Relaxed); + } +} + +impl Default for ExitTool { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Tool for ExitTool { + fn name(&self) -> &str { + "exit" + } + + fn description(&self) -> &str { + "Exit the session. Call this when the user wants to quit, say goodbye, or end the conversation. No arguments required." + } + + async fn execute(&self, _args: &HashMap) -> Result { + self.flag.store(true, Ordering::Relaxed); + Ok("Goodbye! Exiting session.".to_string()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn not_triggered_initially() { + let tool = ExitTool::new(); + assert!(!tool.triggered()); + } + + #[tokio::test] + async fn triggered_after_execute() { + let tool = ExitTool::new(); + let result = tool.execute(&HashMap::new()).await; + assert!(result.is_ok()); + assert!(tool.triggered()); + } + + #[tokio::test] + async fn reset_clears_flag() { + let tool = ExitTool::new(); + tool.execute(&HashMap::new()).await.unwrap(); + assert!(tool.triggered()); + tool.reset(); + assert!(!tool.triggered()); + } + + #[test] + fn name_is_exit() { + let tool = ExitTool::new(); + assert_eq!(Tool::name(&tool), "exit"); + } + + #[test] + fn description_mentions_quit() { + let tool = ExitTool::new(); + assert!(Tool::description(&tool).contains("quit")); + } + + #[test] + fn default_impl_works() { + let tool = ExitTool::default(); + assert!(!tool.triggered()); + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 9d15d57..593eb4f 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -1,3 +1,4 @@ +pub mod exit; pub mod shell; use anyhow::Result; diff --git a/tests/tools_test.rs b/tests/tools_test.rs index be61b7f..33205db 100644 --- a/tests/tools_test.rs +++ b/tests/tools_test.rs @@ -1,8 +1,9 @@ use std::collections::HashMap; use std::sync::Arc; +use golem::tools::exit::ExitTool; use golem::tools::shell::{ShellConfig, ShellMode, ShellTool}; -use golem::tools::{Outcome, ToolRegistry}; +use golem::tools::{Outcome, Tool, ToolRegistry}; /// Helper: build a shell tool with no confirmation, read-write mode, cwd as work dir. fn test_shell() -> ShellTool { @@ -201,3 +202,41 @@ async fn registry_unregister_removes_tool() { assert_eq!(registry.descriptions().await.len(), 0); } + +// --- ExitTool integration tests --- + +#[tokio::test] +async fn exit_tool_executes_via_registry() { + let registry = ToolRegistry::new(); + let exit_tool = Arc::new(ExitTool::new()); + registry + .register(Arc::clone(&exit_tool) as Arc) + .await; + + let result = registry.execute("exit", &HashMap::new()).await; + assert!(matches!(result.outcome, Outcome::Success(ref s) if s.contains("Goodbye"))); + assert!(exit_tool.triggered()); +} + +#[tokio::test] +async fn exit_tool_coexists_with_shell() { + let registry = ToolRegistry::new(); + registry.register(Arc::new(test_shell())).await; + let exit_tool = Arc::new(ExitTool::new()); + registry + .register(Arc::clone(&exit_tool) as Arc) + .await; + + assert_eq!(registry.descriptions().await.len(), 2); + + // Shell still works + let args = HashMap::from([("command".to_string(), "echo hi".to_string())]); + let result = registry.execute("shell", &args).await; + assert!(matches!(result.outcome, Outcome::Success(_))); + assert!(!exit_tool.triggered()); + + // Exit works + let result = registry.execute("exit", &HashMap::new()).await; + assert!(matches!(result.outcome, Outcome::Success(_))); + assert!(exit_tool.triggered()); +} From 45420313614c14c4d4efdc57f755ddc796624bec Mon Sep 17 00:00:00 2001 From: Assaf Sapir Date: Wed, 4 Mar 2026 14:04:52 +0200 Subject: [PATCH 2/2] fix: address PR review comments on ExitTool - Remove unnecessary inner Arc; store AtomicBool directly since ExitTool is already shared via Arc - Replace Ordering::Relaxed with Release (stores) and Acquire (loads) for correct cross-thread synchronization semantics --- src/tools/exit.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/tools/exit.rs b/src/tools/exit.rs index 8e5ee81..52251c7 100644 --- a/src/tools/exit.rs +++ b/src/tools/exit.rs @@ -1,7 +1,6 @@ use anyhow::Result; use async_trait::async_trait; use std::collections::HashMap; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use super::Tool; @@ -11,25 +10,28 @@ use super::Tool; /// When the user says something like "bye" or "quit", the LLM can call this /// tool to signal a graceful shutdown. The caller checks [`ExitTool::triggered`] /// after each engine run to decide whether to break the loop. +/// +/// Share via `Arc` — the `AtomicBool` is stored inline; no inner +/// `Arc` is needed. pub struct ExitTool { - flag: Arc, + flag: AtomicBool, } impl ExitTool { pub fn new() -> Self { Self { - flag: Arc::new(AtomicBool::new(false)), + flag: AtomicBool::new(false), } } /// Returns `true` if the tool was invoked (the agent wants to exit). pub fn triggered(&self) -> bool { - self.flag.load(Ordering::Relaxed) + self.flag.load(Ordering::Acquire) } /// Reset the flag (useful if the caller decides not to exit after all). pub fn reset(&self) { - self.flag.store(false, Ordering::Relaxed); + self.flag.store(false, Ordering::Release); } } @@ -50,7 +52,7 @@ impl Tool for ExitTool { } async fn execute(&self, _args: &HashMap) -> Result { - self.flag.store(true, Ordering::Relaxed); + self.flag.store(true, Ordering::Release); Ok("Goodbye! Exiting session.".to_string()) } }