From aab7675670ef20df10de8844f2477780fb08cb80 Mon Sep 17 00:00:00 2001 From: dynamder Date: Thu, 20 Aug 2026 00:37:58 +0800 Subject: [PATCH 1/5] feat(core): serde-gated messages, stable session id, usage, parallel tools, reasoning levels, shell marker, cancellable loop - serde: gate Serialize/Deserialize derives on chat message types behind a new \serde\ feature, with round-trip tests - session: FuneraSession::with_id; AgentRuntime carries a stable session_id reused by send/send_stream (fire keeps a fresh one-shot id) - usage: TokenUsage + TokenEvent::Usage; requests set stream_options.include_usage; DeepSeek/OpenAI providers report usage - parallel tools: max_concurrent_tools workers over a shared tool bus - reasoning: ReasoningLevel (Off..Max) with hot-reload; DeepSeek mirrors the official harness (thinking + reasoning_effort), OpenAI maps to reasoning_effort (xhigh/max clamped to high) - Tool::is_shell_tool marker (default false); ShellPolicy applies by marker instead of hard-coded tool names; ShellTool opts in - cancellation: per-call CancellationToken through ReActLoopConfig, ToolExecCommand, and every blocking await (biased select); the loop emits MiddlewareEvent::cancelled on cancel - fix CI-found issues: biased cancellation selects, clippy clone_on_copy --- funera_core/Cargo.toml | 3 + funera_core/src/chat/message.rs | 57 ++++- funera_core/src/chat/session.rs | 36 ++- funera_core/src/env.rs | 46 ++++ funera_core/src/env_actor.rs | 30 ++- funera_core/src/event_bus/env_state_bus.rs | 2 + funera_core/src/event_bus/token_bus.rs | 48 ++++ funera_core/src/event_bus/tool_bus.rs | 8 +- funera_core/src/middleware.rs | 9 +- funera_core/src/provider.rs | 38 +++- funera_core/src/provider/deepseek.rs | 151 ++++++++++++- funera_core/src/provider/openai.rs | 62 +++++- funera_core/src/re_act.rs | 242 ++++++++++++++++++--- funera_core/src/re_act/tool.rs | 25 +++ funera_core/src/re_act/tool_executor.rs | 99 ++++++--- funera_core/src/security/policy.rs | 76 ++++++- funera_core/src/security/registry.rs | 55 ++++- 17 files changed, 886 insertions(+), 101 deletions(-) diff --git a/funera_core/Cargo.toml b/funera_core/Cargo.toml index 21ea1b1..1d678c7 100644 --- a/funera_core/Cargo.toml +++ b/funera_core/Cargo.toml @@ -55,3 +55,6 @@ test-utils = [] middleware = [] tool = [] skill = [] +# Enables serde Serialize/Deserialize on the chat message types +# (FuneraMessage, MsgVariant, Role, TextMessage, ToolRequestMessage, ToolResponseMessage). +serde = [] diff --git a/funera_core/src/chat/message.rs b/funera_core/src/chat/message.rs index 286afee..4dd50d6 100644 --- a/funera_core/src/chat/message.rs +++ b/funera_core/src/chat/message.rs @@ -1,6 +1,7 @@ use std::{fmt, sync::Arc}; use chrono::{DateTime, Utc}; +#[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; use serde_json::{Value as JsonValue, json}; use uuid::Uuid; @@ -12,7 +13,8 @@ pub trait Message { fn to_prompt_content(&self) -> String; } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub enum Role { User, Assistant, @@ -30,7 +32,8 @@ impl fmt::Display for Role { } } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub enum MsgVariant { Text(TextMessage), #[cfg(feature = "tool")] @@ -39,7 +42,8 @@ pub enum MsgVariant { ToolResponse(ToolResponseMessage), } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct FuneraMessage { id: Uuid, role: Role, @@ -47,14 +51,16 @@ pub struct FuneraMessage { msg_variant: MsgVariant, } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct TextMessage { pub text: Arc, pub reasoning_content: Option>, } #[cfg(feature = "tool")] -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct ToolRequestMessage { pub tool_call_id: Arc, pub tool_type: ToolType, @@ -64,7 +70,8 @@ pub struct ToolRequestMessage { } #[cfg(feature = "tool")] -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct ToolResponseMessage { pub tool_call_id: Arc, pub result: Arc, @@ -181,3 +188,41 @@ impl FuneraMessage { } } } + +#[cfg(all(test, feature = "serde"))] +mod serde_tests { + use super::*; + use serde_json::json; + + #[test] + fn funera_message_roundtrips_text() { + let msg = FuneraMessage::new( + Role::User, + MsgVariant::Text(TextMessage { + text: Arc::from("hello"), + reasoning_content: Some(Arc::from("thinking...")), + }), + ); + let s = serde_json::to_string(&msg).unwrap(); + let back: FuneraMessage = serde_json::from_str(&s).unwrap(); + assert_eq!(msg, back); + } + + #[cfg(feature = "tool")] + #[test] + fn funera_message_roundtrips_tool() { + let msg = FuneraMessage::new( + Role::Assistant, + MsgVariant::ToolRequest(ToolRequestMessage { + tool_call_id: Arc::from("call_1"), + tool_type: ToolType::Function, + function_name: Arc::from("read"), + function_args: json!({"path": "a.txt"}), + reasoning_content: None, + }), + ); + let s = serde_json::to_string(&msg).unwrap(); + let back: FuneraMessage = serde_json::from_str(&s).unwrap(); + assert_eq!(msg, back); + } +} diff --git a/funera_core/src/chat/session.rs b/funera_core/src/chat/session.rs index 21f7e0b..f9b30a2 100644 --- a/funera_core/src/chat/session.rs +++ b/funera_core/src/chat/session.rs @@ -76,10 +76,17 @@ pub struct FuneraSession { impl FuneraSession { pub fn new(session_tx: mpsc::UnboundedSender) -> Self { - Self { - id: Uuid::new_v4(), - session_tx, - } + Self::with_id(Uuid::new_v4(), session_tx) + } + + /// Create a session wrapper bound to a specific, stable session id. + /// + /// `Agent::send` / `Agent::send_stream` pass the runtime's stable session + /// id here so every turn of a conversation shares one identity; `fire` / + /// `fire_stream` keep the fresh-id `new` constructor (one-shot fork + /// semantics). + pub fn with_id(id: Uuid, session_tx: mpsc::UnboundedSender) -> Self { + Self { id, session_tx } } pub fn id(&self) -> Uuid { @@ -141,3 +148,24 @@ impl FuneraSession { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn new_generates_fresh_ids() { + let (tx, _rx) = mpsc::unbounded_channel(); + let a = FuneraSession::new(tx.clone()); + let b = FuneraSession::new(tx); + assert_ne!(a.id(), b.id()); + } + + #[test] + fn with_id_uses_provided_id() { + let (tx, _rx) = mpsc::unbounded_channel(); + let id = Uuid::new_v4(); + let session = FuneraSession::with_id(id, tx); + assert_eq!(session.id(), id); + } +} diff --git a/funera_core/src/env.rs b/funera_core/src/env.rs index 3518982..2dbb571 100644 --- a/funera_core/src/env.rs +++ b/funera_core/src/env.rs @@ -3,6 +3,8 @@ use std::sync::Arc; use async_openai::config::OpenAIConfig; use parking_lot::Mutex; +use crate::provider::ReasoningLevel; + #[cfg(feature = "skill")] use crate::re_act::skills::{Skill, SkillRegistry}; #[cfg(feature = "tool")] @@ -29,10 +31,12 @@ pub struct FuneraEnv { pub(crate) skill_registry: Arc>, llm_client: async_openai::Client, model: String, + reasoning_level: ReasoningLevel, #[cfg(feature = "tool")] tool_tx: watch::Sender, client_tx: watch::Sender>, model_tx: watch::Sender, + reasoning_tx: watch::Sender, #[cfg(feature = "skill")] skill_tx: watch::Sender, #[cfg(feature = "sandbox")] @@ -50,6 +54,7 @@ impl FuneraEnv { let model = model.into(); let (client_tx, client_rx) = watch::channel(llm_client.clone()); let (model_tx, model_rx) = watch::channel(model.clone()); + let (reasoning_tx, reasoning_rx) = watch::channel(ReasoningLevel::default()); #[cfg(feature = "tool")] let tool_registry = Arc::new(RwLock::new(ToolRegistry::new())); @@ -69,10 +74,12 @@ impl FuneraEnv { skill_registry, llm_client, model, + reasoning_level: ReasoningLevel::default(), #[cfg(feature = "tool")] tool_tx, client_tx, model_tx, + reasoning_tx, #[cfg(feature = "skill")] skill_tx, #[cfg(feature = "sandbox")] @@ -84,6 +91,7 @@ impl FuneraEnv { tool_rx, client_rx, model_rx, + reasoning_rx, #[cfg(feature = "skill")] skill_rx, }, @@ -173,6 +181,23 @@ impl FuneraEnv { let _ = self.model_tx.send(model); } + /// Set the current reasoning level at build time (pushes the watch + /// channel so the value is live immediately). + pub fn with_reasoning_level(mut self, level: ReasoningLevel) -> Self { + self.reasoning_level = level; + let _ = self.reasoning_tx.send(level); + self + } + + pub(crate) fn set_reasoning_level(&mut self, level: ReasoningLevel) { + self.reasoning_level = level; + let _ = self.reasoning_tx.send(level); + } + + pub(crate) fn reasoning_level(&self) -> ReasoningLevel { + self.reasoning_level + } + #[cfg(feature = "skill")] pub(crate) async fn add_skill(&mut self, skill: Skill) { let mut registry = self.skill_registry.write().await; @@ -265,6 +290,7 @@ pub struct FuneraEnvWatcher { tool_rx: watch::Receiver, client_rx: watch::Receiver>, model_rx: watch::Receiver, + reasoning_rx: watch::Receiver, #[cfg(feature = "skill")] skill_rx: watch::Receiver, } @@ -283,6 +309,10 @@ impl FuneraEnvWatcher { self.model_rx.borrow_and_update().clone() } + pub fn watch_reasoning_level(&mut self) -> ReasoningLevel { + *self.reasoning_rx.borrow_and_update() + } + #[cfg(feature = "skill")] pub fn watch_skill(&mut self) -> String { self.skill_rx.borrow_and_update().clone() @@ -301,6 +331,10 @@ impl FuneraEnvWatcher { self.model_rx.has_changed().unwrap_or(false) } + pub fn has_reasoning_level_changed(&self) -> bool { + self.reasoning_rx.has_changed().unwrap_or(false) + } + #[cfg(feature = "skill")] pub fn has_skill_changed(&self) -> bool { self.skill_rx.has_changed().unwrap_or(false) @@ -323,6 +357,10 @@ impl FuneraEnvWatcher { self.model_rx.changed().await } + pub async fn reasoning_level_changed(&mut self) -> Result<(), RecvError> { + self.reasoning_rx.changed().await + } + #[cfg(feature = "skill")] pub async fn skill_changed(&mut self) -> Result<(), RecvError> { self.skill_rx.changed().await @@ -397,6 +435,14 @@ mod tests { assert_eq!(watcher.watch_model(), "m2"); } + #[test] + fn set_reasoning_level_updates_watcher() { + let (mut env, mut watcher) = test_env(); + assert_eq!(watcher.watch_reasoning_level(), ReasoningLevel::Medium); + env.set_reasoning_level(ReasoningLevel::High); + assert_eq!(watcher.watch_reasoning_level(), ReasoningLevel::High); + } + #[test] fn has_model_changed_reflects_changes() { let (mut env, mut watcher) = test_env(); diff --git a/funera_core/src/env_actor.rs b/funera_core/src/env_actor.rs index 7ef8c57..594f068 100644 --- a/funera_core/src/env_actor.rs +++ b/funera_core/src/env_actor.rs @@ -21,6 +21,7 @@ use crate::security::sandbox::SandboxPolicy; use crate::env::{FuneraEnv, FuneraEnvWatcher}; use crate::event_bus::env_state_bus::EnvStateEvent; +use crate::provider::ReasoningLevel; // ═══════════════════════════════════════════════════════════════ // Config structs — bundle params to keep fn arg count ≤ 7 @@ -36,6 +37,10 @@ pub struct EnvToolConfig { pub tool_bus: ToolBus, #[cfg(feature = "tool")] pub exec_rx: mpsc::Receiver, + /// How many concurrent tool workers to spawn (multiple tool calls in one + /// turn execute in parallel). + #[cfg(feature = "tool")] + pub max_concurrent_tools: usize, } /// Bundle of security resources for the EnvActor. @@ -68,6 +73,7 @@ pub enum EnvCmd { // ── Mutation (fire-and-forget) ──────────────────────────── SetModel(String), SetClient(async_openai::Client), + SetReasoningLevel(ReasoningLevel), #[cfg(feature = "tool")] AddTool(Arc), #[cfg(feature = "tool")] @@ -108,6 +114,9 @@ pub enum EnvCmd { GetModel { respond: oneshot::Sender, }, + GetReasoningLevel { + respond: oneshot::Sender, + }, #[cfg(feature = "tool")] GetToolNames { respond: oneshot::Sender>, @@ -159,9 +168,17 @@ pub fn spawn_env_actor( if let Some(tc) = tool { let reg = env.tool_registry.clone(); let tb = tc.tool_bus.clone(); - tokio::spawn(async move { - ToolExecutor::new(reg, tc.exec_rx).run().await; - }); + // Fan the single tool-bus receiver out to `max_concurrent_tools` + // workers: receive is serialized (brief lock), execution runs in + // parallel across workers, and results are matched back by call_id. + let shared_rx = Arc::new(tokio::sync::Mutex::new(tc.exec_rx)); + for _ in 0..tc.max_concurrent_tools.max(1) { + let reg = Arc::clone(®); + let rx = Arc::clone(&shared_rx); + tokio::spawn(async move { + ToolExecutor::with_shared_receiver(reg, rx).run().await; + }); + } Some(tb) } else { None @@ -210,6 +227,10 @@ pub fn spawn_env_actor( EnvCmd::SetClient(client) => { env.set_client(client); } + EnvCmd::SetReasoningLevel(level) => { + env.set_reasoning_level(level); + let _ = state_tx.send(EnvStateEvent::ReasoningLevelChanged(level)); + } #[cfg(feature = "tool")] EnvCmd::AddTool(tool) => { let name = tool.name().to_string(); @@ -280,6 +301,9 @@ pub fn spawn_env_actor( EnvCmd::GetModel { respond } => { let _ = respond.send(env.model().to_string()); } + EnvCmd::GetReasoningLevel { respond } => { + let _ = respond.send(env.reasoning_level()); + } #[cfg(feature = "tool")] EnvCmd::GetToolNames { respond } => { let names = if let Ok(guard) = env.tool_registry.try_read() { diff --git a/funera_core/src/event_bus/env_state_bus.rs b/funera_core/src/event_bus/env_state_bus.rs index 4b06249..aa0d4d4 100644 --- a/funera_core/src/event_bus/env_state_bus.rs +++ b/funera_core/src/event_bus/env_state_bus.rs @@ -4,12 +4,14 @@ use crate::event_bus::{ react_bus::{ReactBus, ReactEvent}, token_bus::TokenEvent, }; +use crate::provider::ReasoningLevel; #[derive(Debug, Clone)] pub enum EnvStateEvent { SessionStart, SessionClosed, LlmChanged(String), + ReasoningLevelChanged(ReasoningLevel), #[cfg(feature = "tool")] ToolAdded(String), #[cfg(feature = "tool")] diff --git a/funera_core/src/event_bus/token_bus.rs b/funera_core/src/event_bus/token_bus.rs index 591d547..8adeace 100644 --- a/funera_core/src/event_bus/token_bus.rs +++ b/funera_core/src/event_bus/token_bus.rs @@ -1,10 +1,56 @@ use async_openai::error::OpenAIError; use async_openai::types::stream::StreamResponse; use futures::StreamExt; +use serde::{Deserialize, Serialize}; use tokio::sync::broadcast; use crate::provider::StreamChunkExt; +/// Token usage statistics for one LLM turn. +/// +/// Providers report usage on the final streamed chunk (when the request asks +/// for it via `stream_options.include_usage`). The framework only tracks +/// tokens — cost computation is left to callers. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct TokenUsage { + /// Tokens in the request prompt. + pub prompt_tokens: u32, + /// Tokens in the generated completion. + pub completion_tokens: u32, + /// `prompt_tokens + completion_tokens`. + pub total_tokens: u32, + /// Prompt tokens served from the provider's cache. + /// + /// Serialized as `prompt_cache_hit_tokens` (DeepSeek's field name), or + /// mapped from OpenAI's `prompt_tokens_details.cached_tokens`. + #[serde(default, rename = "prompt_cache_hit_tokens")] + pub cache_read_tokens: u32, + /// Prompt tokens written to the provider's cache. + /// + /// Serialized as `prompt_cache_miss_tokens` (DeepSeek's field name); + /// OpenAI does not report cache writes, so this stays `0` there. + #[serde(default, rename = "prompt_cache_miss_tokens")] + pub cache_write_tokens: u32, +} + +impl TokenUsage { + /// Map OpenAI's `CompletionUsage` (present on the final streamed chunk) + /// into the framework's unified usage record. + pub fn from_openai(usage: &async_openai::types::chat::CompletionUsage) -> Self { + Self { + prompt_tokens: usage.prompt_tokens, + completion_tokens: usage.completion_tokens, + total_tokens: usage.total_tokens, + cache_read_tokens: usage + .prompt_tokens_details + .as_ref() + .and_then(|d| d.cached_tokens) + .unwrap_or(0), + cache_write_tokens: 0, + } + } +} + #[derive(Debug, Clone)] pub enum TokenEvent { Text(String), @@ -16,6 +62,8 @@ pub enum TokenEvent { }, Finish(async_openai::types::chat::FinishReason), Reasoning(String), + /// Per-turn token usage, emitted once (usually on the final chunk). + Usage(TokenUsage), } pub struct TokenBus { diff --git a/funera_core/src/event_bus/tool_bus.rs b/funera_core/src/event_bus/tool_bus.rs index f0c7de2..af18897 100644 --- a/funera_core/src/event_bus/tool_bus.rs +++ b/funera_core/src/event_bus/tool_bus.rs @@ -1,5 +1,6 @@ use serde_json::Value as JsonValue; use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; use crate::event_bus::react_bus::ReactBus; use crate::re_act::tool::ToolCallError; @@ -10,6 +11,9 @@ pub struct ToolExecCommand { pub args: JsonValue, pub resp_tx: oneshot::Sender>, pub react_bus: Option, + /// Cancel token of the agent call that issued this command: when it fires, + /// the executing worker abandons the in-flight tool execution. + pub cancel: CancellationToken, } #[derive(Clone)] @@ -19,7 +23,7 @@ pub struct ToolBus { impl ToolBus { pub fn new() -> (Self, mpsc::Receiver) { - let (exec_tx, exec_rx) = mpsc::channel(10); + let (exec_tx, exec_rx) = mpsc::channel(64); (Self { exec_tx }, exec_rx) } @@ -29,6 +33,7 @@ impl ToolBus { name: String, args: JsonValue, react_bus: Option, + cancel: CancellationToken, ) -> Result { let (resp_tx, resp_rx) = oneshot::channel(); self.exec_tx @@ -38,6 +43,7 @@ impl ToolBus { args, resp_tx, react_bus, + cancel, }) .await .map_err(|_| ToolCallError::ToolExecutionError(anyhow::anyhow!("tool bus closed")))?; diff --git a/funera_core/src/middleware.rs b/funera_core/src/middleware.rs index f8a17ba..ddb911b 100644 --- a/funera_core/src/middleware.rs +++ b/funera_core/src/middleware.rs @@ -83,6 +83,7 @@ use serde_json::Value as JsonValue; use tokio::sync::mpsc; use crate::chat::message::{MsgVariant, Role}; +use crate::event_bus::token_bus::TokenUsage; // ═══════════════════════════════════════════════════════════════ // Inspector — 只读观察，后台并行，不等待 @@ -622,12 +623,16 @@ pub trait MiddlewareEvent: Clone + Send + 'static { /// Factory：turn 开始。 fn turn_start() -> Self; - /// Factory：turn 结束，携带 finish_reason。 - fn turn_end(finish_reason: Option) -> Self; + /// Factory：turn 结束，携带 finish_reason 与本轮 token 用量。 + fn turn_end(finish_reason: Option, usage: Option) -> Self; /// Factory：会话结束。 fn done() -> Self; + /// Factory：会话被外部取消（middleware 可借此清理自己持有的外部服务； + /// 该事件流经管道但不会写入 session 历史）。 + fn cancelled() -> Self; + /// 转换为 session 历史消息。 /// /// 返回 `Some((role, variant))` 用于构造 `FuneraMessage`。 diff --git a/funera_core/src/provider.rs b/funera_core/src/provider.rs index a1ad049..05a6e2e 100644 --- a/funera_core/src/provider.rs +++ b/funera_core/src/provider.rs @@ -5,10 +5,11 @@ use async_openai::{ types::{chat::CreateChatCompletionStreamResponse, stream::StreamResponse}, }; use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; use serde_json::Value as JsonValue; use std::future::Future; -use crate::event_bus::token_bus::TokenEvent; +use crate::event_bus::token_bus::{TokenEvent, TokenUsage}; #[cfg(feature = "deepseek")] pub mod deepseek; @@ -24,6 +25,29 @@ pub trait StreamChunkExt: DeserializeOwned + Send + 'static { fn extract_events(&self) -> Vec; } +/// Provider-neutral reasoning intensity for the agent's LLM calls. +/// +/// The framework-level enum is model-agnostic; each provider interprets it +/// as it sees fit (levels mirror the official DeepSeek harness): +/// +/// | Level | DeepSeek | OpenAI | +/// |-------|----------|--------| +/// | `Off` | `thinking: {"type":"disabled"}` | `reasoning_effort` omitted | +/// | `Minimal`/`Low`/`Medium`/`XHigh` | `thinking` enabled, deployment default intensity | `reasoning_effort: minimal/low/medium/high` (`XHigh` → `high`) | +/// | `High` | `thinking` enabled + `reasoning_effort: "high"` | `reasoning_effort: "high"` | +/// | `Max` | `thinking` enabled + `reasoning_effort: "max"` | `reasoning_effort: "high"` (clamped) | +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] +pub enum ReasoningLevel { + Off, + Minimal, + Low, + #[default] + Medium, + High, + XHigh, + Max, +} + /// Abstraction over an LLM backend. /// /// Implementations handle provider-specific request construction and stream @@ -37,12 +61,14 @@ pub trait ChatProvider: Send + Sync + 'static { /// Build the JSON request body sent to the LLM API. /// /// Merges conversation messages, active skill content (as a system - /// message), and tool definitions into the provider's expected format. + /// message), tool definitions, and the current reasoning level into the + /// provider's expected format. fn build_request_json( model: &str, messages: &[JsonValue], skill_content: &str, tools_json: &JsonValue, + reasoning_level: ReasoningLevel, ) -> JsonValue; /// Create a streaming completion request. @@ -60,6 +86,9 @@ pub trait ChatProvider: Send + Sync + 'static { impl StreamChunkExt for CreateChatCompletionStreamResponse { fn extract_events(&self) -> Vec { let mut events = Vec::new(); + if let Some(usage) = &self.usage { + events.push(TokenEvent::Usage(TokenUsage::from_openai(usage))); + } for choice in &self.choices { if let Some(finish_reason) = choice.finish_reason { events.push(TokenEvent::Finish(finish_reason)); @@ -130,6 +159,11 @@ pub fn build_standard_request_json( .into_iter() .collect(), ); + // Ask the provider to include per-request token usage on the final chunk. + req.as_object_mut().unwrap().insert( + "stream_options".into(), + serde_json::json!({ "include_usage": true }), + ); if let Some(arr) = tools_json.as_array() && !arr.is_empty() { diff --git a/funera_core/src/provider/deepseek.rs b/funera_core/src/provider/deepseek.rs index 87a2d3c..d6dc718 100644 --- a/funera_core/src/provider/deepseek.rs +++ b/funera_core/src/provider/deepseek.rs @@ -10,8 +10,8 @@ use async_openai::{ use serde::Deserialize; use serde_json::Value as JsonValue; -use crate::event_bus::token_bus::TokenEvent; -use crate::provider::{ChatProvider, StreamChunkExt, build_standard_request_json}; +use crate::event_bus::token_bus::{TokenEvent, TokenUsage}; +use crate::provider::{ChatProvider, ReasoningLevel, StreamChunkExt, build_standard_request_json}; #[derive(Debug, Deserialize)] pub struct Delta { @@ -39,6 +39,8 @@ pub struct StreamChunk { pub model: String, #[serde(default)] pub system_fingerprint: Option, + #[serde(default)] + pub usage: Option, pub object: String, } @@ -89,6 +91,9 @@ impl StreamChunkExt for StreamChunk { (None, None) => {} } } + if let Some(usage) = &self.usage { + events.push(TokenEvent::Usage(usage.clone())); + } events } } @@ -103,6 +108,7 @@ impl ChatProvider for DeepSeekProvider { messages: &[JsonValue], skill_content: &str, tools_json: &JsonValue, + reasoning_level: ReasoningLevel, ) -> JsonValue { let mut json = build_standard_request_json(model, messages, skill_content, tools_json); @@ -168,9 +174,41 @@ impl ChatProvider for DeepSeekProvider { *msgs = merged; } - json.as_object_mut() - .unwrap() - .insert("thinking".into(), serde_json::json!({"type": "enabled"})); + // Reasoning: mirror the official DeepSeek harness (`llm-deepseek`). + // `off` disables thinking entirely; `high`/`max` add `reasoning_effort`; + // the remaining levels keep thinking enabled at the deployment's + // default intensity. + match reasoning_level { + ReasoningLevel::Off => { + json.as_object_mut() + .unwrap() + .insert("thinking".into(), serde_json::json!({"type": "disabled"})); + } + ReasoningLevel::High => { + json.as_object_mut() + .unwrap() + .insert("thinking".into(), serde_json::json!({"type": "enabled"})); + json.as_object_mut() + .unwrap() + .insert("reasoning_effort".into(), serde_json::json!("high")); + } + ReasoningLevel::Max => { + json.as_object_mut() + .unwrap() + .insert("thinking".into(), serde_json::json!({"type": "enabled"})); + json.as_object_mut() + .unwrap() + .insert("reasoning_effort".into(), serde_json::json!("max")); + } + ReasoningLevel::Minimal + | ReasoningLevel::Low + | ReasoningLevel::Medium + | ReasoningLevel::XHigh => { + json.as_object_mut() + .unwrap() + .insert("thinking".into(), serde_json::json!({"type": "enabled"})); + } + } json } @@ -201,7 +239,13 @@ mod tests { json!({"role": "tool", "tool_call_id": "call_2", "content": "18°C"}), json!({"role": "tool", "tool_call_id": "call_3", "content": "25°C"}), ]; - let result = DeepSeekProvider::build_request_json("test-model", &msgs, "", &json!([])); + let result = DeepSeekProvider::build_request_json( + "test-model", + &msgs, + "", + &json!([]), + ReasoningLevel::Medium, + ); let msgs_out = result["messages"].as_array().unwrap(); assert_eq!( msgs_out.len(), @@ -231,7 +275,13 @@ mod tests { json!({"role": "tool", "tool_call_id": "call_2", "content": "20°C"}), json!({"role": "assistant", "content": "Done."}), ]; - let result = DeepSeekProvider::build_request_json("test-model", &msgs, "", &json!([])); + let result = DeepSeekProvider::build_request_json( + "test-model", + &msgs, + "", + &json!([]), + ReasoningLevel::Medium, + ); let msgs_out = result["messages"].as_array().unwrap(); assert_eq!( msgs_out[1]["reasoning_content"].as_str().unwrap(), @@ -248,7 +298,13 @@ mod tests { json!({"role": "assistant", "content": "I'll get that."}), json!({"role": "assistant", "content": null, "tool_calls": [{"id": "c2", "function": {"name": "t2", "arguments": "{}"}}]}), ]; - let result = DeepSeekProvider::build_request_json("test-model", &msgs, "", &json!([])); + let result = DeepSeekProvider::build_request_json( + "test-model", + &msgs, + "", + &json!([]), + ReasoningLevel::Medium, + ); let msgs_out = result["messages"].as_array().unwrap(); assert_eq!(msgs_out.len(), 4, "non-consecutive should NOT merge"); let calls0 = msgs_out[1]["tool_calls"].as_array().unwrap(); @@ -263,7 +319,13 @@ mod tests { json!({"role": "user", "content": "hi"}), json!({"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "t1", "arguments": "{}"}}]}), ]; - let result = DeepSeekProvider::build_request_json("test-model", &msgs, "", &json!([])); + let result = DeepSeekProvider::build_request_json( + "test-model", + &msgs, + "", + &json!([]), + ReasoningLevel::Medium, + ); let msgs_out = result["messages"].as_array().unwrap(); assert!(msgs_out[1]["content"].is_null()); assert_eq!(msgs_out[1]["tool_calls"].as_array().unwrap().len(), 1); @@ -272,7 +334,76 @@ mod tests { #[test] fn build_request_includes_thinking() { let msgs = vec![json!({"role": "user", "content": "hi"})]; - let result = DeepSeekProvider::build_request_json("test-model", &msgs, "", &json!([])); + let result = DeepSeekProvider::build_request_json( + "test-model", + &msgs, + "", + &json!([]), + ReasoningLevel::Medium, + ); assert_eq!(result["thinking"]["type"], "enabled"); } + + #[test] + fn reasoning_level_mapping() { + let msgs = vec![json!({"role": "user", "content": "hi"})]; + let tools = json!([]); + + // Off → thinking disabled, no reasoning_effort. + let req = DeepSeekProvider::build_request_json("m", &msgs, "", &tools, ReasoningLevel::Off); + assert_eq!(req["thinking"]["type"], "disabled"); + assert!(req.get("reasoning_effort").is_none()); + + // High → thinking enabled + reasoning_effort "high". + let req = + DeepSeekProvider::build_request_json("m", &msgs, "", &tools, ReasoningLevel::High); + assert_eq!(req["thinking"]["type"], "enabled"); + assert_eq!(req["reasoning_effort"], "high"); + + // Max → thinking enabled + reasoning_effort "max". + let req = DeepSeekProvider::build_request_json("m", &msgs, "", &tools, ReasoningLevel::Max); + assert_eq!(req["thinking"]["type"], "enabled"); + assert_eq!(req["reasoning_effort"], "max"); + + // Medium → thinking enabled at the deployment's default intensity. + let req = + DeepSeekProvider::build_request_json("m", &msgs, "", &tools, ReasoningLevel::Medium); + assert_eq!(req["thinking"]["type"], "enabled"); + assert!(req.get("reasoning_effort").is_none()); + } + + #[test] + fn extract_events_includes_usage() { + let chunk = StreamChunk { + id: "chunk_1".into(), + choices: vec![], + created: 0, + model: "deepseek-chat".into(), + system_fingerprint: None, + usage: Some(TokenUsage { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + cache_read_tokens: 4, + cache_write_tokens: 6, + }), + object: "chat.completion.chunk".into(), + }; + let events = chunk.extract_events(); + assert!(matches!( + events.as_slice(), + [TokenEvent::Usage(u)] if u.total_tokens == 15 + )); + } + + #[test] + fn usage_deserializes_from_deepseek_shape() { + let raw = r#"{"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "prompt_cache_hit_tokens": 4, "prompt_cache_miss_tokens": 6}"#; + let u: TokenUsage = serde_json::from_str(raw).unwrap(); + assert_eq!(u.prompt_tokens, 10); + assert_eq!(u.completion_tokens, 5); + assert_eq!(u.total_tokens, 15); + assert_eq!(u.cache_read_tokens, 4); + assert_eq!(u.cache_write_tokens, 6); + } } diff --git a/funera_core/src/provider/openai.rs b/funera_core/src/provider/openai.rs index a4f4f2d..509cf7a 100644 --- a/funera_core/src/provider/openai.rs +++ b/funera_core/src/provider/openai.rs @@ -6,7 +6,7 @@ use async_openai::{ }; use serde_json::Value as JsonValue; -use crate::provider::{ChatProvider, build_standard_request_json}; +use crate::provider::{ChatProvider, ReasoningLevel, build_standard_request_json}; pub struct OpenAIProvider; @@ -18,8 +18,39 @@ impl ChatProvider for OpenAIProvider { messages: &[JsonValue], skill_content: &str, tools_json: &JsonValue, + reasoning_level: ReasoningLevel, ) -> JsonValue { - build_standard_request_json(model, messages, skill_content, tools_json) + let mut json = build_standard_request_json(model, messages, skill_content, tools_json); + match reasoning_level { + ReasoningLevel::Off => {} + ReasoningLevel::Minimal => { + json.as_object_mut() + .unwrap() + .insert("reasoning_effort".into(), serde_json::json!("minimal")); + } + ReasoningLevel::Low => { + json.as_object_mut() + .unwrap() + .insert("reasoning_effort".into(), serde_json::json!("low")); + } + ReasoningLevel::Medium => { + json.as_object_mut() + .unwrap() + .insert("reasoning_effort".into(), serde_json::json!("medium")); + } + ReasoningLevel::High => { + json.as_object_mut() + .unwrap() + .insert("reasoning_effort".into(), serde_json::json!("high")); + } + // OpenAI exposes only up to `high`; clamp the stronger levels. + ReasoningLevel::XHigh | ReasoningLevel::Max => { + json.as_object_mut() + .unwrap() + .insert("reasoning_effort".into(), serde_json::json!("high")); + } + } + json } async fn create_stream( @@ -32,3 +63,30 @@ impl ChatProvider for OpenAIProvider { .await } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn reasoning_effort_mapping() { + let msgs = vec![json!({"role": "user", "content": "hi"})]; + let tools = json!([]); + for (level, expected) in [ + (ReasoningLevel::Off, None), + (ReasoningLevel::Minimal, Some("minimal")), + (ReasoningLevel::Low, Some("low")), + (ReasoningLevel::Medium, Some("medium")), + (ReasoningLevel::High, Some("high")), + (ReasoningLevel::XHigh, Some("high")), + (ReasoningLevel::Max, Some("high")), + ] { + let req = OpenAIProvider::build_request_json("m", &msgs, "", &tools, level); + match expected { + Some(v) => assert_eq!(req["reasoning_effort"], json!(v), "level {level:?}"), + None => assert!(req.get("reasoning_effort").is_none(), "level {level:?}"), + } + } + } +} diff --git a/funera_core/src/re_act.rs b/funera_core/src/re_act.rs index cbdec31..00e6a1d 100644 --- a/funera_core/src/re_act.rs +++ b/funera_core/src/re_act.rs @@ -16,11 +16,13 @@ use crate::event_bus::env_state_bus::{EnvStateEvent, TurnHighWayHandle}; use crate::event_bus::react_bus::{ReactBus, ReactEvent}; #[cfg(feature = "tool")] use crate::event_bus::react_bus::{ToolCallErrorInfo, ToolCallRequest, ToolCallResponse}; -use crate::event_bus::token_bus::{TokenBus, TokenEvent}; +use crate::event_bus::token_bus::{TokenBus, TokenEvent, TokenUsage}; #[cfg(feature = "tool")] use crate::event_bus::tool_bus::ToolBus; use crate::middleware::{ErrorsEnabled, EventSenderFn, MiddlewareChain, MiddlewareEvent}; use crate::provider::ChatProvider; +#[cfg(test)] +use crate::provider::ReasoningLevel; #[cfg(feature = "skill")] pub mod skills; @@ -49,6 +51,10 @@ pub struct ReActLoopConfig { pub turn_highway_handle: TurnHighWayHandle, /// Optional sender to the session actor — enables message history persistence. pub session_tx: Option>, + /// Per-call cancellation token; when fired the loop stops receiving + /// output, abandons in-flight tool executions, and notifies middleware + /// via [`MiddlewareEvent::cancelled`](crate::middleware::MiddlewareEvent::cancelled). + pub cancel: CancellationToken, } impl ReActLoopConfig { @@ -71,9 +77,16 @@ impl ReActLoopConfig { env_state_tx, turn_highway_handle, session_tx: None, + cancel: CancellationToken::new(), } } + /// Bind this loop to a per-call cancellation token. + pub fn with_cancel(mut self, cancel: CancellationToken) -> Self { + self.cancel = cancel; + self + } + #[cfg(feature = "tool")] pub fn with_tool_bus(mut self, tool_bus: ToolBus) -> Self { // Placeholder: tool_bus is already set via the field; this allows @@ -92,6 +105,7 @@ pub struct ReActLoop { #[allow(dead_code)] env_state_tx: broadcast::Sender, turn_highway_handle: TurnHighWayHandle, + cancel: CancellationToken, _phantom: PhantomData

, } @@ -114,6 +128,7 @@ impl ReActLoop

{ }, env_state_tx, turn_highway_handle, + cancel: CancellationToken::new(), _phantom: PhantomData, } } @@ -133,6 +148,7 @@ impl ReActLoop

{ tool_bus: config.tool_bus, env_state_tx: config.env_state_tx, turn_highway_handle: config.turn_highway_handle, + cancel: config.cancel, _phantom: PhantomData, } } @@ -148,7 +164,7 @@ impl ReActLoop

{ middleware: Option>>, event_sender: Option>, ) -> ReActLoopHandle { - let token = CancellationToken::new(); + let token = self.cancel.clone(); let token_clone = token.clone(); let task = tokio::spawn(async move { @@ -156,12 +172,22 @@ impl ReActLoop

{ let mut env_watcher = self.env_watcher; while iteration < self.max_iteration { + if token.is_cancelled() { + filter_and_store( + vec![E::cancelled()], + &middleware, + &event_sender, + &self.session_tx, + ); + break; + } let client = env_watcher.watch_client(); #[cfg(feature = "tool")] let tools_json = env_watcher.watch_tool(); #[cfg(not(feature = "tool"))] let tools_json = JsonValue::Array(Vec::new()); let model = env_watcher.watch_model(); + let reasoning_level = env_watcher.watch_reasoning_level(); #[cfg(feature = "skill")] let skill_content = env_watcher.watch_skill(); #[cfg(not(feature = "skill"))] @@ -175,18 +201,53 @@ impl ReActLoop

{ let (token_tx, react_bus) = self.turn_highway_handle.prepare_turn().await; - let request_json = - P::build_request_json(&model, &history_json, &skill_content, &tools_json); + let request_json = P::build_request_json( + &model, + &history_json, + &skill_content, + &tools_json, + reasoning_level, + ); react_bus.send(ReactEvent::TurnStart).ok(); emit_event(&event_sender, E::turn_start()); - let stream = P::create_stream(&client, request_json).await?; + // Cancellation can arrive while the initial provider request + // is in flight: don't wait for it — exit cooperatively + // (`biased` so cancellation wins over a fast response). + let stream = tokio::select! { + biased; + _ = token.cancelled() => { + filter_and_store( + vec![E::cancelled()], + &middleware, + &event_sender, + &self.session_tx, + ); + break; + } + s = P::create_stream(&client, request_json) => s?, + }; let mut token_bus = TokenBus::::with_sender(token_tx, stream); - let (assistant_content, reasoning_content, tool_call_accums, turn_finish_reason) = - process_token_stream(&mut token_bus, &react_bus).await?; + let ( + assistant_content, + reasoning_content, + tool_call_accums, + turn_finish_reason, + turn_usage, + ) = process_token_stream(&mut token_bus, &react_bus, &token).await?; + + if token.is_cancelled() { + filter_and_store( + vec![E::cancelled()], + &middleware, + &event_sender, + &self.session_tx, + ); + break; + } let finish_reason_str = turn_finish_reason.as_ref().map(|r| format!("{:?}", r)); @@ -222,9 +283,20 @@ impl ReActLoop

{ &react_bus, #[cfg(feature = "tool")] &self.tool_bus, + &token, ) .await?; + if token.is_cancelled() { + filter_and_store( + vec![E::cancelled()], + &middleware, + &event_sender, + &self.session_tx, + ); + break; + } + // `tool_results` is only consumed when the `tool` feature is on. #[cfg(not(feature = "tool"))] let _ = tool_results; @@ -242,7 +314,7 @@ impl ReActLoop

{ filter_and_store(result_events, &middleware, &event_sender, &self.session_tx); } - emit_event(&event_sender, E::turn_end(finish_reason_str)); + emit_event(&event_sender, E::turn_end(finish_reason_str, turn_usage)); react_bus.send(ReactEvent::TurnEnd).ok(); if !should_continue { @@ -306,19 +378,32 @@ fn filter_and_store( async fn process_token_stream( token_bus: &mut TokenBus, react_bus: &ReactBus, + cancel: &CancellationToken, ) -> Result<( String, String, HashMap, Option, + Option, )> { let mut assistant_content = String::new(); let mut reasoning_content = String::new(); let mut tool_call_accums: HashMap = HashMap::new(); let mut turn_finish_reason: Option = None; - - while let Some(result) = token_bus.recv().await { - let events = result?; + let mut turn_usage: Option = None; + + loop { + // Cancel → stop receiving output immediately. `biased` makes + // cancellation win over a simultaneously-ready stream chunk. + let events = tokio::select! { + biased; + _ = cancel.cancelled() => break, + r = token_bus.recv() => match r { + None => break, + Some(Err(e)) => return Err(e.into()), + Some(Ok(evs)) => evs, + }, + }; for event in events { match event { TokenEvent::Text(t) => { @@ -363,6 +448,9 @@ async fn process_token_stream( TokenEvent::Finish(reason) => { turn_finish_reason = Some(reason); } + TokenEvent::Usage(usage) => { + turn_usage = Some(usage); + } } } } @@ -372,6 +460,7 @@ async fn process_token_stream( reasoning_content, tool_call_accums, turn_finish_reason, + turn_usage, )) } @@ -381,6 +470,7 @@ async fn handle_turn_finish( tool_call_accums: &HashMap, react_bus: &ReactBus, tool_bus: &ToolBus, + cancel: &CancellationToken, ) -> Result<(bool, Vec)> { match finish_reason { None | Some(FinishReason::Stop) => Ok((false, Vec::new())), @@ -411,6 +501,7 @@ async fn handle_turn_finish( acc.name.clone(), args, Some(react_bus.clone()), + cancel.clone(), ) })) .await; @@ -461,6 +552,7 @@ async fn handle_turn_finish( finish_reason: Option<&FinishReason>, tool_call_accums: &HashMap, _react_bus: &ReactBus, + _cancel: &CancellationToken, ) -> Result<(bool, Vec)> { match finish_reason { Some(FinishReason::ToolCalls) | Some(FinishReason::Length) @@ -521,12 +613,15 @@ mod tests { fn turn_start() -> Self { TestEvent("turn_start".into()) } - fn turn_end(_finish_reason: Option) -> Self { + fn turn_end(_finish_reason: Option, _usage: Option) -> Self { TestEvent("turn_end".into()) } fn done() -> Self { TestEvent("done".into()) } + fn cancelled() -> Self { + TestEvent("cancelled".into()) + } fn into_session_message(self) -> Option<(Role, MsgVariant)> { Some(( Role::Assistant, @@ -547,9 +642,10 @@ mod tests { let mut token_bus = TokenBus::with_sender(tx, stream); let react_bus = ReactBus::new(); - let (content, reasoning, accums, reason) = process_token_stream(&mut token_bus, &react_bus) - .await - .unwrap(); + let (content, reasoning, accums, reason, _usage) = + process_token_stream(&mut token_bus, &react_bus, &CancellationToken::new()) + .await + .unwrap(); assert_eq!(content, ""); assert_eq!(reasoning, ""); assert!(accums.is_empty()); @@ -563,9 +659,10 @@ mod tests { let mut token_bus = TokenBus::with_sender(tx, stream); let react_bus = ReactBus::new(); - let (content, reasoning, accums, reason) = process_token_stream(&mut token_bus, &react_bus) - .await - .unwrap(); + let (content, reasoning, accums, reason, _usage) = + process_token_stream(&mut token_bus, &react_bus, &CancellationToken::new()) + .await + .unwrap(); assert_eq!(content, "Hello World"); assert_eq!(reasoning, ""); assert!(accums.is_empty()); @@ -579,9 +676,10 @@ mod tests { let mut token_bus = TokenBus::with_sender(tx, stream); let react_bus = ReactBus::new(); - let (content, reasoning, accums, reason) = process_token_stream(&mut token_bus, &react_bus) - .await - .unwrap(); + let (content, reasoning, accums, reason, _usage) = + process_token_stream(&mut token_bus, &react_bus, &CancellationToken::new()) + .await + .unwrap(); assert_eq!(content, ""); assert_eq!(reasoning, ""); assert_eq!(accums.len(), 1); @@ -595,10 +693,31 @@ mod tests { let mut token_bus = TokenBus::with_sender(tx, stream); let react_bus = ReactBus::new(); - let result = process_token_stream(&mut token_bus, &react_bus).await; + let result = + process_token_stream(&mut token_bus, &react_bus, &CancellationToken::new()).await; assert!(result.is_err()); } + #[tokio::test] + async fn process_stream_cancelled_stops_immediately() { + let (tx, _) = tokio::sync::broadcast::channel(50); + let stream = test_helpers::mock_text_stream(vec!["Hello", "World"]); + let mut token_bus = TokenBus::with_sender(tx, stream); + let react_bus = ReactBus::new(); + let cancel = CancellationToken::new(); + cancel.cancel(); + + let (content, reasoning, accums, reason, usage) = + process_token_stream(&mut token_bus, &react_bus, &cancel) + .await + .unwrap(); + assert_eq!(content, ""); + assert_eq!(reasoning, ""); + assert!(accums.is_empty()); + assert!(reason.is_none()); + assert!(usage.is_none()); + } + #[tokio::test] async fn process_stream_text_and_tool() { let (tx, _) = tokio::sync::broadcast::channel(50); @@ -611,9 +730,10 @@ mod tests { let mut token_bus = TokenBus::with_sender(tx, stream); let react_bus = ReactBus::new(); - let (content, reasoning, accums, reason) = process_token_stream(&mut token_bus, &react_bus) - .await - .unwrap(); + let (content, reasoning, accums, reason, _usage) = + process_token_stream(&mut token_bus, &react_bus, &CancellationToken::new()) + .await + .unwrap(); assert_eq!(content, "Thinking..."); assert_eq!(reasoning, ""); assert_eq!(accums.len(), 1); @@ -633,6 +753,7 @@ mod tests { &HashMap::new(), &react_bus, &tool_bus, + &CancellationToken::new(), ) .await .unwrap(); @@ -647,10 +768,15 @@ mod tests { let react_bus = ReactBus::new(); let (tool_bus, _rx) = crate::event_bus::tool_bus::ToolBus::new(); - let (should_continue, results) = - handle_turn_finish(None, &HashMap::new(), &react_bus, &tool_bus) - .await - .unwrap(); + let (should_continue, results) = handle_turn_finish( + None, + &HashMap::new(), + &react_bus, + &tool_bus, + &CancellationToken::new(), + ) + .await + .unwrap(); assert!(!should_continue); assert!(results.is_empty()); @@ -667,6 +793,7 @@ mod tests { &HashMap::new(), &react_bus, &tool_bus, + &CancellationToken::new(), ) .await .unwrap(); @@ -747,6 +874,7 @@ mod tests { messages: &[JsonValue], _skill_content: &str, _tools_json: &JsonValue, + _reasoning_level: ReasoningLevel, ) -> JsonValue { serde_json::json!({ "model": "mock", @@ -821,4 +949,60 @@ mod tests { let inner = result.unwrap(); assert!(inner.is_ok(), "loop should return Ok(())"); } + + #[cfg(feature = "tool")] + #[tokio::test] + async fn react_loop_cancelled_notifies_middleware() { + let (env_state_tx, _env_state_rx) = tokio::sync::broadcast::channel(20); + let (tool_bus, _exec_rx) = crate::event_bus::tool_bus::ToolBus::new(); + + // Per-call token, cancelled before the loop starts. + let cancel = CancellationToken::new(); + cancel.cancel(); + + // Highway stub (mirrors the completing test). + let (loop_tx, mut rx_from_loop) = tokio::sync::mpsc::channel(5); + let (tx_to_loop, loop_rx) = tokio::sync::mpsc::channel(5); + tokio::spawn(async move { + let _req = rx_from_loop.recv().await; + let (token_tx, _) = tokio::sync::broadcast::channel(50); + let react_bus = ReactBus::new(); + let _ = tx_to_loop + .send(TurnHighWayEvent::TurnPrepareResponse { + token_tx, + react_bus, + }) + .await; + }); + let handle = TurnHighWayHandle { + turn_high_way_tx: loop_tx, + turn_high_way_rx: loop_rx, + }; + + let (_env, env_watcher) = + crate::env::FuneraEnv::new(async_openai::Client::new(), "mock-model"); + + let config = ReActLoopConfig::new(32, 5, env_watcher, env_state_tx, handle) + .with_tool_bus(tool_bus) + .with_cancel(cancel); + + // Capture the events the loop emits (middleware/event-sender channel). + let (event_tx, mut event_rx) = tokio::sync::mpsc::unbounded_channel(); + let sender: EventSenderFn = Box::new(move |e| { + let _ = event_tx.send(e); + }); + + let loop_handle = + ReActLoop::::from_config(config).run::(None, Some(sender)); + let result = loop_handle.task.await; + assert!(result.is_ok()); + assert!(result.unwrap().is_ok()); + + // Cancellation must be delivered as an event to the middleware layer. + let got = event_rx.try_recv(); + assert!( + matches!(&got, Ok(TestEvent(s)) if s.as_str() == "cancelled"), + "got {got:?}" + ); + } } diff --git a/funera_core/src/re_act/tool.rs b/funera_core/src/re_act/tool.rs index 7e8fbdf..40aca7c 100644 --- a/funera_core/src/re_act/tool.rs +++ b/funera_core/src/re_act/tool.rs @@ -26,6 +26,21 @@ impl Display for ToolType { /// Implement this trait to define custom tools. The framework will expose /// the tool's [`schema`](Tool::schema) to the LLM and invoke /// [`execute`](Tool::execute) when the LLM requests it. +/// +/// ## Tools are CLI-like and self-contained +/// +/// A tool behaves like a single command-line program: it receives JSON +/// arguments and returns a string result. It must **not** depend on another +/// tool's *registration* — no access to the tool registry, no assumption +/// that a peer tool is (or ever was) registered, and no cross-tool +/// bookkeeping. Invoking external programs as subprocesses (the way the +/// built-in `shell` tool runs `git`, `cargo`, …) is fine — that is exactly +/// what makes a tool CLI-like. Composing *agent* capabilities happens at the +/// LLM level through multiple tool calls, not inside a tool. +/// +/// Optional markers refine how the framework treats a tool: returning `true` +/// from [`is_shell_tool`](Tool::is_shell_tool) opts the tool into shell +/// command policy scrutiny regardless of its registered name. #[async_trait] pub trait Tool: Send + Sync { /// Unique name for this tool (e.g. `"read"`, `"shell"`). @@ -48,6 +63,16 @@ pub trait Tool: Send + Sync { /// /// This is sent to the LLM so it can generate well-formed invocations. fn schema(&self) -> JsonValue; + + /// Whether this tool executes shell commands (e.g. `shell`, `bash`, `cmd`). + /// + /// Defaults to `false`. Override to `true` for tools whose + /// [`execute`](Tool::execute) runs shell commands, so the security + /// layer's `ShellPolicy` scrutiny applies — based on this explicit + /// opt-in, not on the tool's registered name. + fn is_shell_tool(&self) -> bool { + false + } } /// Errors that can occur during tool execution. diff --git a/funera_core/src/re_act/tool_executor.rs b/funera_core/src/re_act/tool_executor.rs index 8ad4b88..9190452 100644 --- a/funera_core/src/re_act/tool_executor.rs +++ b/funera_core/src/re_act/tool_executor.rs @@ -2,20 +2,37 @@ use std::sync::Arc; -use tokio::sync::{RwLock, mpsc}; +use tokio::sync::{Mutex, RwLock, mpsc}; +#[cfg(test)] +use tokio_util::sync::CancellationToken; use crate::event_bus::tool_bus::ToolExecCommand; -use crate::re_act::tool::ToolRegistry; +use crate::re_act::tool::{ToolCallError, ToolRegistry}; pub struct ToolExecutor { tool_registry: Arc>, - exec_rx: mpsc::Receiver, + exec_rx: Arc>>, } impl ToolExecutor { pub fn new( tool_registry: Arc>, exec_rx: mpsc::Receiver, + ) -> Self { + Self { + tool_registry, + exec_rx: Arc::new(Mutex::new(exec_rx)), + } + } + + /// Build an executor over a receiver shared by multiple workers. + /// + /// Used by the env actor to fan one tool bus out to + /// `max_concurrent_tools` workers: receive is serialized (brief lock), + /// execution runs in parallel across workers. + pub fn with_shared_receiver( + tool_registry: Arc>, + exec_rx: Arc>>, ) -> Self { Self { tool_registry, @@ -23,42 +40,57 @@ impl ToolExecutor { } } - pub async fn run(mut self) { - while let Some(cmd) = self.exec_rx.recv().await { - let result = { - // Without the security layer, clone the tool `Arc` and execute - // it outside the registry lock so slow tools do not block - // dynamic tool add/remove/availability changes. - #[cfg(not(feature = "security"))] - { - use crate::re_act::tool::ToolCallError; - let tool = { - let registry = self.tool_registry.read().await; - registry.get_tool_arc(&cmd.name) - }; - match tool { - Some(tool) => tool.execute(cmd.args).await, - None => Err(ToolCallError::ToolNotFound(cmd.name.clone())), - } - } + pub async fn run(self) { + loop { + let cmd = { self.exec_rx.lock().await.recv().await }; + let Some(cmd) = cmd else { break }; - // Clone the guarded registry so policy/boundary/approval/audit - // run against a snapshot while the tool itself executes outside - // the registry lock. Pending approvals and the react bus are - // `Arc`-shared, so cloned registries still cooperate. - #[cfg(feature = "security")] - { - let registry = { - let guard = self.tool_registry.read().await; - guard.clone() - }; - registry.set_react_bus(cmd.react_bus.clone()); - registry.call_tool(&cmd.name, cmd.args).await + // Cooperative cancellation: if the issuing agent call was + // cancelled, abandon this in-flight tool execution (the worker + // itself survives — it is shared across calls). `biased` makes + // cancellation win over a simultaneously-completed tool. + let result = tokio::select! { + biased; + _ = cmd.cancel.cancelled() => { + Err(ToolCallError::ToolUnavailable("cancelled".into())) } + r = self.execute_command(&cmd) => r, }; let _ = cmd.resp_tx.send(result); } } + + /// Run one tool command (policy/registry/execution) without cancellation. + async fn execute_command(&self, cmd: &ToolExecCommand) -> Result { + // Without the security layer, clone the tool `Arc` and execute + // it outside the registry lock so slow tools do not block + // dynamic tool add/remove/availability changes. + #[cfg(not(feature = "security"))] + { + let tool = { + let registry = self.tool_registry.read().await; + registry.get_tool_arc(&cmd.name) + }; + match tool { + Some(tool) => tool.execute(cmd.args.clone()).await, + None => Err(ToolCallError::ToolNotFound(cmd.name.clone())), + } + } + + // Clone the guarded registry so policy/boundary/approval/audit + // run against a snapshot while the tool itself executes outside + // the registry lock. Pending approvals and the react bus are + // `Arc`-shared, so cloned registries still cooperate. + #[cfg(feature = "security")] + { + let registry = { + let guard = self.tool_registry.read().await; + guard.clone() + }; + registry.set_react_bus(cmd.react_bus.clone()); + registry.call_tool(&cmd.name, cmd.args.clone()).await + } + } } #[cfg(test)] @@ -101,6 +133,7 @@ mod tests { args: json!({}), resp_tx, react_bus: None, + cancel: CancellationToken::new(), }) .await .expect("executor must consume commands"); diff --git a/funera_core/src/security/policy.rs b/funera_core/src/security/policy.rs index 4df467b..c22a0d2 100644 --- a/funera_core/src/security/policy.rs +++ b/funera_core/src/security/policy.rs @@ -29,7 +29,7 @@ pub struct ToolPolicy { /// Maximum allowed timeout in seconds. pub max_timeout_secs: f64, - /// Shell command policy (applies to shell/bash/sh/cmd/powershell tools). + /// Shell command policy (applies to tools that opt in via `Tool::is_shell_tool`). pub shell_policy: Option, /// Allowed working directories for shell tools. @@ -100,18 +100,18 @@ impl ToolPolicy { /// Check a shell command against the configured shell policy. /// - /// Only applies if the tool is recognized as a shell tool and a + /// Only applies if the tool opts in via `Tool::is_shell_tool` and a /// [`ShellPolicy`] is configured. pub fn check_shell_command( &self, - tool_name: &str, + tool: &dyn crate::re_act::tool::Tool, args: &JsonValue, ) -> Result<(), PolicyError> { let policy = match self.shell_policy { Some(ref p) => p, None => return Ok(()), }; - if !policy.is_relevant_tool(tool_name) { + if !tool.is_shell_tool() { return Ok(()); } let command = match args.get("command").and_then(|v| v.as_str()) { @@ -149,6 +149,10 @@ impl ToolPolicy { /// /// Controls which commands are allowed or denied, with built-in detection of /// dangerous patterns (e.g., `rm -rf`, `diskpart`, `reg add`, etc.). +/// +/// Applies only to tools that opt in via `Tool::is_shell_tool` (the built-in +/// `shell` tool does; third-party shell tools should do the same) — never by +/// tool name. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ShellPolicy { /// Commands matching these prefixes are allowed. @@ -189,10 +193,6 @@ impl ShellPolicy { } } - fn is_relevant_tool(&self, tool_name: &str) -> bool { - matches!(tool_name, "shell" | "bash" | "sh" | "cmd" | "powershell") - } - /// Evaluate a shell command against all allow/deny rules. pub fn check_command(&self, command: &str) -> Result<(), PolicyError> { if !self.allow_commands.is_empty() { @@ -431,6 +431,66 @@ mod tests { assert!(policy.check_command("echo hello").is_ok()); } + #[test] + fn shell_policy_applies_only_to_opted_in_tools() { + use crate::re_act::tool::ToolCallError; + + struct NotShell; + #[async_trait::async_trait] + impl crate::re_act::tool::Tool for NotShell { + fn name(&self) -> &str { + "shell" // name alone must NOT trigger shell policy + } + fn description(&self) -> &str { + "not a shell tool" + } + fn schema(&self) -> JsonValue { + JsonValue::Null + } + async fn execute(&self, _args: JsonValue) -> Result { + Ok(String::new()) + } + } + + struct MyShell; + #[async_trait::async_trait] + impl crate::re_act::tool::Tool for MyShell { + fn name(&self) -> &str { + "my-shell" // unrelated name, but opts in explicitly + } + fn description(&self) -> &str { + "runs shell commands" + } + fn schema(&self) -> JsonValue { + JsonValue::Null + } + async fn execute(&self, _args: JsonValue) -> Result { + Ok(String::new()) + } + fn is_shell_tool(&self) -> bool { + true + } + } + + let policy = ToolPolicy { + shell_policy: Some(ShellPolicy::strict()), + ..Default::default() + }; + + // Not opted in → the dangerous command passes policy scrutiny. + assert!( + policy + .check_shell_command(&NotShell, &serde_json::json!({"command": "rm -rf /"})) + .is_ok() + ); + // Opted in → the same command is rejected. + assert!( + policy + .check_shell_command(&MyShell, &serde_json::json!({"command": "rm -rf /"})) + .is_err() + ); + } + #[cfg(feature = "regex")] #[test] fn shell_policy_regex_deny() { diff --git a/funera_core/src/security/registry.rs b/funera_core/src/security/registry.rs index a59934c..cd2044d 100644 --- a/funera_core/src/security/registry.rs +++ b/funera_core/src/security/registry.rs @@ -180,7 +180,11 @@ impl GuardedToolRegistry { if let Some(workdir) = args.get("workdir").and_then(|v| v.as_str()) { self.policy.check_workdir(workdir)?; } - self.policy.check_shell_command(name, &args)?; + // Shell scrutiny is opt-in via `Tool::is_shell_tool` — never by name. + if let Some(entry) = self.inner.get_tool(name) { + self.policy + .check_shell_command(entry.tool.as_ref(), &args)?; + } #[cfg(feature = "sandbox")] self.audit_sandbox(name); Ok(()) @@ -744,4 +748,53 @@ mod tests { "shell-like command with sandbox+path_guard must not be rejected; got: {result:?}" ); } + + /// A tool that opts in via `is_shell_tool` gets shell policy applied + /// even when its name is unrelated to "shell". + #[tokio::test] + async fn opted_in_shell_tool_gets_shell_policy() { + use crate::security::policy::ShellPolicy; + + struct MyShell; + #[async_trait] + impl Tool for MyShell { + fn name(&self) -> &str { + "my-shell" + } + fn description(&self) -> &str { + "runs shell commands" + } + fn schema(&self) -> serde_json::Value { + serde_json::json!({}) + } + async fn execute(&self, _args: serde_json::Value) -> Result { + Ok("done".into()) + } + fn is_shell_tool(&self) -> bool { + true + } + } + + let policy = ToolPolicy { + shell_policy: Some(ShellPolicy::with_allowed(vec!["ls".into()])), + ..Default::default() + }; + let mut registry = GuardedToolRegistry::new_from_policy(policy); + registry.add_tool(Arc::new(MyShell)); + + // Allowed command passes through the full guarded path. + let ok = registry + .call_tool("my-shell", json!({"command": "ls -la"})) + .await; + assert!(ok.is_ok(), "allowed command should pass: {ok:?}"); + + // A dangerous command is rejected by shell policy. + let bad = registry + .call_tool("my-shell", json!({"command": "rm -rf /"})) + .await; + assert!( + bad.is_err(), + "dangerous command should be rejected: {bad:?}" + ); + } } From eb7d9b15e23b67d9c3e2b8af2400b25dbeca99eb Mon Sep 17 00:00:00 2001 From: dynamder Date: Thu, 20 Aug 2026 00:38:06 +0800 Subject: [PATCH 2/5] feat(orchestrate): cancellable handles (FireHandle), usage/reasoning/parallel wiring, cancel example - Agent::fire returns FireHandle; all four entry points carry a per-call CancellationToken exposed as cancel(); dropping any handle auto-cancels - AgentEvent::Cancelled flows through the middleware chain and subscribers (never into session history); awaiting a cancelled handle yields OrchestrateError::Cancelled; recv() returns None once cancelled - ChatResponse.usage carries the last turn's TokenUsage - builder knobs: max_concurrent_tools, reasoning_level; set_reasoning_level hot-reload; AgentRuntime::session_id - new example: cancel (explicit cancel, drop-to-cancel, timeout, middleware notification) - fire() call sites updated across examples, integration tests, and doctests --- Cargo.lock | 1 + funera-orchestrate/Cargo.toml | 1 + funera-orchestrate/examples/builtin_tools.rs | 1 + funera-orchestrate/examples/cancel.rs | 123 +++++++++++ funera-orchestrate/examples/middleware.rs | 1 + funera-orchestrate/examples/minimal.rs | 1 + funera-orchestrate/examples/sandbox.rs | 10 +- funera-orchestrate/examples/security.rs | 1 + funera-orchestrate/examples/skills.rs | 1 + funera-orchestrate/src/agent.rs | 133 ++++-------- funera-orchestrate/src/event.rs | 26 ++- funera-orchestrate/src/lib.rs | 14 +- funera-orchestrate/src/response.rs | 4 + funera-orchestrate/src/runtime.rs | 119 ++++++++++- funera-orchestrate/src/send_handle.rs | 205 ++++++++++++++++--- funera-orchestrate/tests/integration.rs | 16 +- 16 files changed, 521 insertions(+), 136 deletions(-) create mode 100644 funera-orchestrate/examples/cancel.rs diff --git a/Cargo.lock b/Cargo.lock index 4af3386..f125f81 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -975,6 +975,7 @@ dependencies = [ "serde_json", "thiserror", "tokio", + "tokio-util", "tracing", "uuid", ] diff --git a/funera-orchestrate/Cargo.toml b/funera-orchestrate/Cargo.toml index 6a5b327..f8b8ca5 100644 --- a/funera-orchestrate/Cargo.toml +++ b/funera-orchestrate/Cargo.toml @@ -22,6 +22,7 @@ thiserror = "2.0.18" anyhow = "1.0.103" futures = "0.3.32" tracing = "0.1" +tokio-util = "0.7.18" [features] default = ["deepseek", "tool"] diff --git a/funera-orchestrate/examples/builtin_tools.rs b/funera-orchestrate/examples/builtin_tools.rs index 8e4d7f2..0e3cb1f 100644 --- a/funera-orchestrate/examples/builtin_tools.rs +++ b/funera-orchestrate/examples/builtin_tools.rs @@ -26,6 +26,7 @@ async fn main() -> Result<(), Box> { // Ask the agent to read Cargo.toml (it will use the Read tool) let resp = agent .fire("Read Cargo.toml and tell me the dependencies.", &runtime) + .await? .await?; println!("{}", resp.content); diff --git a/funera-orchestrate/examples/cancel.rs b/funera-orchestrate/examples/cancel.rs new file mode 100644 index 0000000..f08e6ad --- /dev/null +++ b/funera-orchestrate/examples/cancel.rs @@ -0,0 +1,123 @@ +//! Cancellation — interrupt an agent call mid-flight. +//! +//! ```bash +//! cargo run -p funera-orchestrate --example cancel +//! # with the middleware demo: +//! cargo run -p funera-orchestrate --example cancel --features middleware +//! ``` +//! +//! Every funera entry point returns a cancellable handle: +//! +//! - [`Agent::fire`] / [`Agent::fire_stream`] → one-shot call +//! - [`Agent::send`] / [`Agent::send_stream`] → multi-turn call +//! +//! [`cancel()`](FireStreamHandle::cancel) immediately stops receiving output, +//! abandons in-flight tool executions, and emits [`AgentEvent::Cancelled`] to +//! subscribers — and through the middleware chain when the `middleware` +//! feature is enabled — so middleware holding external services can clean up. +//! **Dropping the handle cancels automatically**: no detached background work +//! is left running. +//! +//! Requires an API key (set `OPENAI_API_KEY`, or call `.api_key()`). + +use std::time::Duration; + +use funera_orchestrate::{Agent, AgentEvent, AgentRuntime, DeepSeekProvider}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + #[cfg_attr(not(feature = "middleware"), allow(unused_mut))] + let mut builder = AgentRuntime::::builder() + .api_key(std::env::var("OPENAI_API_KEY")?) + .model(std::env::var("OPENAI_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".into())); + + // A tiny middleware inspector that logs AgentEvent::Cancelled, so the + // "cancellation reaches middleware" contract is visible. Requires the + // `middleware` feature. + #[cfg(feature = "middleware")] + { + use funera_orchestrate::middleware::{InspectorError, InspectorMiddleware}; + use funera_orchestrate::middleware_bundle::MiddlewareBundle; + + struct CancelLogger; + impl InspectorMiddleware for CancelLogger { + fn name(&self) -> &str { + "cancel_logger" + } + fn inspect(&self, event: &AgentEvent) -> Result<(), InspectorError> { + if matches!(event, AgentEvent::Cancelled) { + eprintln!( + "[middleware] saw AgentEvent::Cancelled — clean up external services here" + ); + } + Ok(()) + } + } + + let bundle = MiddlewareBundle::from_chain( + funera_orchestrate::middleware::MiddlewareChain::::new() + .with_inspector(CancelLogger), + ); + builder = builder.with_middleware_bundle(bundle); + } + + let runtime = builder.build()?; + let agent = Agent::builder() + .system_prompt("You are a helpful assistant.") + .build(); + + // ── 1. Explicit cancel() on a streaming call ───────────────── + let mut stream = agent + .fire_stream("Write a very long essay about Rust.", &runtime) + .await?; + + // Subscribe before cancelling so we can observe AgentEvent::Cancelled. + let mut events = agent.subscribe_events(); + tokio::time::sleep(Duration::from_millis(200)).await; + stream.cancel(); + + // recv() returns None once cancelled. + assert!( + stream.recv().await.is_none(), + "recv should end after cancel" + ); + + // The cancellation event reaches subscribers (and middleware, if enabled). + loop { + match events.recv().await { + Ok(AgentEvent::Cancelled) => { + println!("[1] explicit cancel: subscriber received AgentEvent::Cancelled"); + break; + } + Ok(_) => {} + Err(_) => break, + } + } + + // Awaiting a cancelled handle yields OrchestrateError::Cancelled. + match stream.await { + Err(funera_orchestrate::OrchestrateError::Cancelled) => { + println!("[1] explicit cancel: wait() reported OrchestrateError::Cancelled"); + } + other => println!("[1] unexpected result: {other:?}"), + } + + // ── 2. Dropping the handle auto-cancels ───────────────────── + let handle = agent.fire("Explain the borrow checker.", &runtime).await?; + // Drop without awaiting: the background call is cancelled, nothing leaks. + drop(handle); + println!("[2] dropped handle: background call auto-cancelled"); + + // ── 3. Timeout via tokio::select! ─────────────────────────── + let handle = agent.fire("Write a 10-page novel.", &runtime).await?; + let resp = tokio::select! { + r = handle => r?, + _ = tokio::time::sleep(Duration::from_secs(3)) => { + println!("[3] timeout: the fire future was dropped, call cancelled"); + return Ok(()); + } + }; + println!("[3] completed before timeout: {} chars", resp.content.len()); + + Ok(()) +} diff --git a/funera-orchestrate/examples/middleware.rs b/funera-orchestrate/examples/middleware.rs index 0059ca5..0096a26 100644 --- a/funera-orchestrate/examples/middleware.rs +++ b/funera-orchestrate/examples/middleware.rs @@ -63,6 +63,7 @@ impl InspectorMiddleware for EventLogger { } => { eprintln!("[log] approval required: {tool_name} — {reason}") } + AgentEvent::Cancelled => eprintln!("[log] call cancelled"), } Ok(()) } diff --git a/funera-orchestrate/examples/minimal.rs b/funera-orchestrate/examples/minimal.rs index ad7baf0..b8438db 100644 --- a/funera-orchestrate/examples/minimal.rs +++ b/funera-orchestrate/examples/minimal.rs @@ -28,6 +28,7 @@ async fn main() -> Result<(), Box> { // fire() shares the runtime (&) — no session state is mutated let resp = agent .fire("Tell me about Rust programming.", &runtime) + .await? .await?; println!("=== Response ==="); diff --git a/funera-orchestrate/examples/sandbox.rs b/funera-orchestrate/examples/sandbox.rs index e1c251a..7daeba5 100644 --- a/funera-orchestrate/examples/sandbox.rs +++ b/funera-orchestrate/examples/sandbox.rs @@ -152,7 +152,7 @@ async fn main() { }) .build(); - match agent + let handle = match agent .fire( "Explore the current directory, then attempt operations \ that violate the sandbox restrictions to verify the \ @@ -161,6 +161,14 @@ async fn main() { ) .await { + Ok(h) => h, + Err(e) => { + eprintln!("error: agent request failed: {e}"); + std::process::exit(1); + } + }; + + match handle.await { Ok(resp) => { println!("=== Agent Response ===\n{}", resp.content); eprintln!( diff --git a/funera-orchestrate/examples/security.rs b/funera-orchestrate/examples/security.rs index 6648f5e..26c8721 100644 --- a/funera-orchestrate/examples/security.rs +++ b/funera-orchestrate/examples/security.rs @@ -101,6 +101,7 @@ async fn main() -> Result<(), Box> { // `approve_tool_call` because both go through the shared EnvActor. let resp = agent .fire("List all .rs files in the current directory", &runtime) + .await? .await?; println!("Agent: {}", resp.content); diff --git a/funera-orchestrate/examples/skills.rs b/funera-orchestrate/examples/skills.rs index eb14d7f..da10daf 100644 --- a/funera-orchestrate/examples/skills.rs +++ b/funera-orchestrate/examples/skills.rs @@ -72,6 +72,7 @@ async fn main() -> Result<(), Box> { // response should be noticeably short. let resp = agent .fire("What is the Rust programming language?", &runtime) + .await? .await?; println!("=== Response ==="); diff --git a/funera-orchestrate/src/agent.rs b/funera-orchestrate/src/agent.rs index 26e3f47..fb5b250 100644 --- a/funera-orchestrate/src/agent.rs +++ b/funera-orchestrate/src/agent.rs @@ -10,13 +10,13 @@ use funera_core::middleware::EventSenderFn; use funera_core::middleware::{ErrorsEnabled, MiddlewareChain}; use funera_core::provider::ChatProvider; use funera_core::re_act::ReActLoopConfig; +use tokio_util::sync::CancellationToken; use crate::dispatcher::{CallbackDispatcher, CallbackRegistry}; use crate::error::OrchestrateError; use crate::event::{AgentEvent, RawAgentEvent}; -use crate::response::{ChatResponse, ToolCallInfo}; use crate::runtime::{AgentRuntime, Idle}; -use crate::send_handle::{FireStreamHandle, SendHandle, SendStreamHandle}; +use crate::send_handle::{FireHandle, FireStreamHandle, SendHandle, SendStreamHandle}; // --------------------------------------------------------------------------- // AgentBuilder @@ -238,13 +238,17 @@ impl Agent { /// One-shot query. Creates a temporary session, runs a single ReAct loop, /// and discards the session. + /// + /// Returns a [`FireHandle`]: `await` it for the final [`ChatResponse`], + /// call `cancel()` to interrupt, or simply drop it — dropping cancels all + /// background work spawned for the call. pub async fn fire( &self, msg: impl Into, runtime: &AgentRuntime, - ) -> Result { + ) -> Result { let text = msg.into(); - let mut event_rx = self.subscribe_events(); + let event_rx = self.subscribe_events(); let (env_state_bus, turn_highway_handle) = EnvStateBus::new(); let env_state_tx = env_state_bus.env_state_tx.clone(); @@ -256,10 +260,9 @@ impl Agent { self.event_tx.clone(), self.raw_event_tx.clone(), ); - let _ = env_state_tx.send(EnvStateEvent::SessionStart); - // Temporary actor for one-shot — dropped after react_loop completes + // Temporary actor for one-shot — dropped after react_loop completes. let session_tx = funera_core::chat::session::spawn_session_actor(); let session = FuneraSession::new(session_tx); if let Some(ref sys) = self.system_prompt { @@ -271,7 +274,6 @@ impl Agent { }), )); } - let init_msg = FuneraMessage::new( Role::User, MsgVariant::Text(TextMessage { @@ -293,21 +295,25 @@ impl Agent { { config = config.with_tool_bus(react.tool_bus); } + let cancel = CancellationToken::new(); + config = config.with_cancel(cancel.clone()); let event_sender = build_event_sender(self.callbacks.clone(), self.event_tx.clone()); - let result = session - .react_loop::( - init_msg, - config, - env_state_tx.clone(), - middleware_opt(runtime), - Some(event_sender), - ) - .await; + let mw = middleware_opt(runtime); + let env_tx = env_state_tx.clone(); + let handle = tokio::spawn(async move { + session + .react_loop::(init_msg, config, env_tx, mw, Some(event_sender)) + .await + }); - let _ = env_state_tx.send(EnvStateEvent::SessionClosed); - aggregate_response(&mut event_rx, result).await + Ok(FireHandle { + handle: Some(handle), + cancel, + event_rx: Some(event_rx), + env_state_tx, + }) } /// Streaming variant of [`fire`](Self::fire). @@ -372,6 +378,8 @@ impl Agent { { config = config.with_tool_bus(react.tool_bus); } + let cancel = CancellationToken::new(); + config = config.with_cancel(cancel.clone()); let event_sender = build_event_sender(self.callbacks.clone(), self.event_tx.clone()); // Spawn react_loop as background task @@ -384,10 +392,11 @@ impl Agent { }); Ok(FireStreamHandle { - handle, - event_rx, - stream_rx, + handle: Some(handle), + event_rx: Some(event_rx), + stream_rx: Some(stream_rx), env_state_tx, + cancel, }) } @@ -417,7 +426,7 @@ impl Agent { ); let _ = env_state_tx.send(EnvStateEvent::SessionStart); - let session = FuneraSession::new(runtime.session_tx()); + let session = FuneraSession::with_id(runtime.session_id(), runtime.session_tx()); if let Some(ref sys) = self.system_prompt { let msgs = session.session_context().await; if msgs.is_empty() { @@ -450,6 +459,8 @@ impl Agent { { config = config.with_tool_bus(react.tool_bus); } + let cancel = CancellationToken::new(); + config = config.with_cancel(cancel.clone()); let event_sender = build_event_sender(self.callbacks.clone(), self.event_tx.clone()); let env_tx = env_state_tx.clone(); @@ -461,10 +472,11 @@ impl Agent { }); Ok(SendHandle { - runtime: runtime.into_acquired(), - handle, - event_rx, + runtime: Some(runtime.into_acquired()), + handle: Some(handle), + event_rx: Some(event_rx), env_state_tx, + cancel, }) } @@ -496,7 +508,7 @@ impl Agent { ); let _ = env_state_tx.send(EnvStateEvent::SessionStart); - let session = FuneraSession::new(runtime.session_tx()); + let session = FuneraSession::with_id(runtime.session_id(), runtime.session_tx()); if let Some(ref sys) = self.system_prompt { let msgs = session.session_context().await; if msgs.is_empty() { @@ -529,6 +541,8 @@ impl Agent { { config = config.with_tool_bus(react.tool_bus); } + let cancel = CancellationToken::new(); + config = config.with_cancel(cancel.clone()); let event_sender = build_event_sender(self.callbacks.clone(), self.event_tx.clone()); let env_tx = env_state_tx.clone(); @@ -540,11 +554,12 @@ impl Agent { }); Ok(SendStreamHandle { - runtime: runtime.into_acquired(), - handle, - event_rx, - stream_rx, + runtime: Some(runtime.into_acquired()), + handle: Some(handle), + event_rx: Some(event_rx), + stream_rx: Some(stream_rx), env_state_tx, + cancel, }) } } @@ -602,64 +617,6 @@ async fn relay_broadcast_to_mpsc( } } -/// Aggregate middleware-filtered events from the event stream into a ChatResponse. -async fn aggregate_response( - event_rx: &mut broadcast::Receiver, - react_result: Result<(), anyhow::Error>, -) -> Result { - react_result.map_err(OrchestrateError::Session)?; - - let mut content = String::new(); - let mut tool_calls = Vec::new(); - let mut iterations = 0usize; - let mut finish_reason: Option = None; - - // Track pending tool call requests to match with results - let mut pending_requests: Vec<(Arc, String, serde_json::Value)> = Vec::new(); - - loop { - match event_rx.recv().await { - Ok(AgentEvent::Text(t)) => { - content = t; - } - Ok(AgentEvent::ToolCallRequest { - call_id, - name, - args, - .. - }) => { - pending_requests.push((call_id, name, args)); - } - Ok(AgentEvent::ToolCallResult { - call_id, - name: _, - result, - }) => { - if let Some(pos) = pending_requests - .iter() - .position(|(id, _, _)| *id == call_id) - { - let (_, name, args) = pending_requests.remove(pos); - tool_calls.push(ToolCallInfo { name, args, result }); - } - } - Ok(AgentEvent::TurnStart) => iterations += 1, - Ok(AgentEvent::TurnEnd { finish_reason: fr }) => finish_reason = fr, - Ok(AgentEvent::Done) => break, - Err(broadcast::error::RecvError::Closed) => break, - Err(broadcast::error::RecvError::Lagged(_)) => continue, - _ => {} - } - } - - Ok(ChatResponse { - content, - tool_calls, - iterations, - finish_reason, - }) -} - #[cfg(test)] mod tests { use super::*; diff --git a/funera-orchestrate/src/event.rs b/funera-orchestrate/src/event.rs index 5107a30..b66e23b 100644 --- a/funera-orchestrate/src/event.rs +++ b/funera-orchestrate/src/event.rs @@ -7,7 +7,7 @@ use funera_core::chat::message::{MsgVariant, Role, TextMessage}; use funera_core::chat::message::{ToolRequestMessage, ToolResponseMessage}; use funera_core::event_bus::env_state_bus::EnvStateEvent; use funera_core::event_bus::react_bus::ReactEvent; -use funera_core::event_bus::token_bus::TokenEvent; +use funera_core::event_bus::token_bus::{TokenEvent, TokenUsage}; use funera_core::middleware::MiddlewareEvent; #[cfg(feature = "tool")] use funera_core::re_act::tool::ToolType; @@ -30,6 +30,7 @@ pub enum AgentEvent { TurnStart, TurnEnd { finish_reason: Option, + usage: Option, }, Error(String), ToolApprovalRequired { @@ -37,6 +38,9 @@ pub enum AgentEvent { tool_name: String, reason: String, }, + /// 本次调用被 `cancel()` 打断。会流经 middleware 链和事件订阅者， + /// 但不会写入 session 历史。 + Cancelled, Done, } @@ -68,14 +72,21 @@ impl MiddlewareEvent for AgentEvent { AgentEvent::TurnStart } - fn turn_end(finish_reason: Option) -> Self { - AgentEvent::TurnEnd { finish_reason } + fn turn_end(finish_reason: Option, usage: Option) -> Self { + AgentEvent::TurnEnd { + finish_reason, + usage, + } } fn done() -> Self { AgentEvent::Done } + fn cancelled() -> Self { + AgentEvent::Cancelled + } + fn into_session_message(self) -> Option<(Role, MsgVariant)> { match self { AgentEvent::Text(text) => Some(( @@ -232,8 +243,15 @@ mod tests { fn agent_event_turn_end() { let e = AgentEvent::TurnEnd { finish_reason: Some("Stop".into()), + usage: None, }; - assert!(matches!(e, AgentEvent::TurnEnd { finish_reason: Some(ref r) } if r == "Stop")); + assert!(matches!( + e, + AgentEvent::TurnEnd { + finish_reason: Some(ref r), + .. + } if r == "Stop" + )); } #[test] diff --git a/funera-orchestrate/src/lib.rs b/funera-orchestrate/src/lib.rs index f5beb5c..e50af16 100644 --- a/funera-orchestrate/src/lib.rs +++ b/funera-orchestrate/src/lib.rs @@ -37,7 +37,7 @@ //! .system_prompt("You are a helpful assistant.") //! .build(); //! -//! let resp = agent.fire("Hello!", &runtime).await?; +//! let resp = agent.fire("Hello!", &runtime).await?.await?; //! println!("{}", resp.content); //! Ok(()) //! } @@ -122,7 +122,7 @@ //! let agent = Agent::builder().build(); //! //! let (fast, _) = agent.send("Hello", fast).await?.await?; // fast model -//! agent.fire("What is Rust?", &powerful).await?; // powerful model (temp) +//! agent.fire("What is Rust?", &powerful).await?.await?; // powerful model (temp) //! let (_fast, _) = agent.send("Tell me more", fast).await?.await?; // back to fast //! # Ok(()) //! # } @@ -135,9 +135,10 @@ //! while the agent is running — works with `fire()`, `send()`, and `send_stream()`. //! //! ```rust,no_run +//! # #[cfg(all(feature = "tool", feature = "security"))] +//! # async fn example() -> Result<(), Box> { //! # use funera_orchestrate::{Agent, AgentRuntime, ApprovalHandle, DeepSeekProvider}; //! # use std::time::Duration; -//! # async fn example() -> Result<(), Box> { //! let (approval_tx, mut approval_rx) = tokio::sync::mpsc::unbounded_channel(); //! //! let runtime = AgentRuntime::::builder() @@ -169,7 +170,7 @@ //! //! - [`runtime`] — [`AgentRuntimeBuilder`] and [`AgentRuntime`] //! - [`agent`] — [`AgentBuilder`] and [`Agent`] -//! - [`send_handle`] — [`SendHandle`], [`SendStreamHandle`], [`FireStreamHandle`], [`ApprovalHandle`] +//! - [`send_handle`] — [`FireHandle`], [`SendHandle`], [`SendStreamHandle`], [`FireStreamHandle`], [`ApprovalHandle`] //! - [`dispatcher`] — Event bus subscription and callback dispatch //! - [`event`] — [`AgentEvent`] enum //! - [`response`] — [`ChatResponse`] and [`ToolCallInfo`] @@ -198,7 +199,7 @@ pub use response::{ChatResponse, ToolCallInfo}; pub use runtime::{Acquired, AgentRuntime, AgentRuntimeBuilder, Idle}; #[cfg(all(feature = "tool", feature = "security"))] pub use send_handle::ApprovalHandle; -pub use send_handle::{FireStreamHandle, SendHandle, SendStreamHandle}; +pub use send_handle::{FireHandle, FireStreamHandle, SendHandle, SendStreamHandle}; // Re-export security policy types for convenience. #[cfg(feature = "security")] @@ -211,7 +212,8 @@ pub use funera_core::event_bus::env_state_bus::EnvStateEvent; pub use funera_core::event_bus::react_bus::{ ReactEvent, ToolCallErrorInfo, ToolCallRequest, ToolCallResponse, }; -pub use funera_core::event_bus::token_bus::TokenEvent; +pub use funera_core::event_bus::token_bus::{TokenEvent, TokenUsage}; +pub use funera_core::provider::ReasoningLevel; /// Middleware 相关的类型和 trait。 /// diff --git a/funera-orchestrate/src/response.rs b/funera-orchestrate/src/response.rs index 88bd340..36a4156 100644 --- a/funera-orchestrate/src/response.rs +++ b/funera-orchestrate/src/response.rs @@ -1,3 +1,4 @@ +use funera_core::event_bus::token_bus::TokenUsage; use serde_json::Value as JsonValue; #[derive(Debug, Clone)] @@ -6,6 +7,8 @@ pub struct ChatResponse { pub tool_calls: Vec, pub iterations: usize, pub finish_reason: Option, + /// Token usage for the last LLM turn, when the provider reports it. + pub usage: Option, } #[derive(Debug, Clone)] @@ -26,6 +29,7 @@ mod tests { tool_calls: vec![], iterations: 1, finish_reason: Some("stop".into()), + usage: None, }; assert_eq!(resp.content, "Hello"); assert_eq!(resp.iterations, 1); diff --git a/funera-orchestrate/src/runtime.rs b/funera-orchestrate/src/runtime.rs index c30016f..e476ca9 100644 --- a/funera-orchestrate/src/runtime.rs +++ b/funera-orchestrate/src/runtime.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use async_openai::config::OpenAIConfig; use tokio::sync::{broadcast, mpsc}; +use uuid::Uuid; #[cfg(test)] use funera_core::chat::session::FuneraSession; @@ -19,9 +20,9 @@ use funera_core::env_actor::{EnvCmd, ReActConfig, spawn_env_actor}; use funera_core::event_bus::env_state_bus::EnvStateEvent; #[cfg(feature = "tool")] use funera_core::event_bus::tool_bus::ToolBus; -use funera_core::provider::ChatProvider; #[cfg(feature = "deepseek")] use funera_core::provider::deepseek::DeepSeekProvider; +use funera_core::provider::{ChatProvider, ReasoningLevel}; #[cfg(feature = "skill")] use funera_core::re_act::skills::{Skill, SkillRegistry}; #[cfg(feature = "tool")] @@ -65,8 +66,10 @@ pub struct AgentRuntimeBuilder { base_url: Option, client: Option>, model: Option, + reasoning_level: Option, max_iterations: usize, channel_buffer: usize, + max_concurrent_tools: usize, #[cfg(feature = "tool")] tools: Vec>, #[cfg(feature = "skill")] @@ -102,8 +105,10 @@ impl AgentRuntimeBuilder { base_url: None, client: None, model: None, + reasoning_level: None, max_iterations: 10, channel_buffer: 32, + max_concurrent_tools: 4, #[cfg(feature = "tool")] tools: Vec::new(), #[cfg(feature = "skill")] @@ -153,6 +158,12 @@ impl AgentRuntimeBuilder { self } + /// Initial reasoning level (default [`ReasoningLevel::Medium`]). + pub fn reasoning_level(mut self, level: ReasoningLevel) -> Self { + self.reasoning_level = Some(level); + self + } + /// Directly provide an OpenAI client (overrides api_key + base_url). pub fn client(mut self, client: async_openai::Client) -> Self { self.client = Some(client); @@ -171,6 +182,16 @@ impl AgentRuntimeBuilder { self } + /// Maximum number of tools that may execute concurrently (default 4). + /// + /// The `ToolExecutor` runs this many workers over the shared tool bus, so + /// multiple tool calls in one LLM turn execute in parallel. Results are + /// matched back by `call_id`, so execution order does not matter. + pub fn max_concurrent_tools(mut self, n: usize) -> Self { + self.max_concurrent_tools = n; + self + } + /// Load a skill from a SKILL.md file. #[cfg(feature = "skill")] pub fn with_skill_file(mut self, path: impl Into) -> Self { @@ -451,6 +472,11 @@ impl AgentRuntimeBuilder { env }; + let env = match self.reasoning_level { + Some(level) => env.with_reasoning_level(level), + None => env, + }; + #[cfg(feature = "tool")] let env = env.with_tool_registry(registry); #[cfg(feature = "skill")] @@ -492,7 +518,11 @@ impl AgentRuntimeBuilder { let tool_cfg = { #[cfg(feature = "tool")] { - Some(EnvToolConfig { tool_bus, exec_rx }) + Some(EnvToolConfig { + tool_bus, + exec_rx, + max_concurrent_tools: self.max_concurrent_tools, + }) } #[cfg(not(feature = "tool"))] { @@ -523,10 +553,15 @@ impl AgentRuntimeBuilder { let env_cmd_tx = spawn_env_actor(env, env_watcher, max_iters, chan_buf, tool_cfg, sec_cfg); let session_tx = spawn_session_actor(); + // Stable identity for this runtime's conversation session: `send` / + // `send_stream` reuse it across calls, while `fire` uses a fresh + // one-shot id per call (fork semantics). + let session_id = Uuid::new_v4(); Ok(AgentRuntime::

{ env_cmd_tx, session_tx, + session_id, #[cfg(feature = "middleware")] middleware_chain, _state: PhantomData, @@ -579,6 +614,7 @@ pub struct Acquired; pub struct AgentRuntime { pub(crate) env_cmd_tx: mpsc::UnboundedSender, pub(crate) session_tx: mpsc::UnboundedSender, + session_id: Uuid, #[cfg(feature = "middleware")] pub(crate) middleware_chain: Arc>, _state: PhantomData, @@ -603,6 +639,15 @@ impl AgentRuntime { self.session_tx.clone() } + /// The stable identifier of this runtime's conversation session. + /// + /// `send` / `send_stream` keep this id across calls so every turn of a + /// conversation shares one identity; `fire` / `fire_stream` use a fresh + /// one-shot id per call (fork semantics). + pub fn session_id(&self) -> Uuid { + self.session_id + } + /// Subscribe to runtime-level environment state events. pub async fn subscribe_env_state(&self) -> broadcast::Receiver { let (respond, rx) = tokio::sync::oneshot::channel(); @@ -653,6 +698,22 @@ impl AgentRuntime { let _ = self.env_cmd_tx.send(EnvCmd::SetModel(model.into())); } + /// Change the reasoning level at runtime. + /// + /// Pushes to the internal watch channel (picked up by the ReAct loop on + /// the next iteration) and broadcasts + /// [`EnvStateEvent::ReasoningLevelChanged`]. + pub fn set_reasoning_level(&self, level: ReasoningLevel) { + let _ = self.env_cmd_tx.send(EnvCmd::SetReasoningLevel(level)); + } + + /// Query the current reasoning level from the env actor. + pub async fn reasoning_level(&self) -> ReasoningLevel { + let (respond, rx) = tokio::sync::oneshot::channel(); + let _ = self.env_cmd_tx.send(EnvCmd::GetReasoningLevel { respond }); + rx.await.unwrap_or_default() + } + /// Change the LLM client at runtime (endpoint, key, etc.). pub fn set_client(&self, client: async_openai::Client) { let _ = self.env_cmd_tx.send(EnvCmd::SetClient(client)); @@ -778,6 +839,7 @@ impl AgentRuntime { AgentRuntime:: { env_cmd_tx: self.env_cmd_tx, session_tx: self.session_tx, + session_id: self.session_id, #[cfg(feature = "middleware")] middleware_chain: self.middleware_chain, _state: PhantomData, @@ -800,6 +862,7 @@ impl AgentRuntime { AgentRuntime:: { env_cmd_tx: self.env_cmd_tx, session_tx: self.session_tx, + session_id: self.session_id, #[cfg(feature = "middleware")] middleware_chain: self.middleware_chain, _state: PhantomData, @@ -944,12 +1007,26 @@ mod tests { assert_eq!(b.channel_buffer, 64); } + #[test] + fn builder_max_concurrent_tools_default_and_set() { + assert_eq!(AgentRuntimeBuilder::new().max_concurrent_tools, 4); + let b = AgentRuntimeBuilder::new().max_concurrent_tools(8); + assert_eq!(b.max_concurrent_tools, 8); + } + #[test] fn builder_set_model() { let b = AgentRuntimeBuilder::new().model("test-model"); assert_eq!(b.model, Some("test-model".into())); } + #[test] + fn builder_reasoning_level() { + assert!(AgentRuntimeBuilder::new().reasoning_level.is_none()); + let b = AgentRuntimeBuilder::new().reasoning_level(ReasoningLevel::Max); + assert_eq!(b.reasoning_level, Some(ReasoningLevel::Max)); + } + #[test] fn builder_set_api_key() { let b = AgentRuntimeBuilder::new().api_key("sk-test"); @@ -1036,6 +1113,22 @@ mod tests { assert!(tx.send(SessionCmd::Clear).is_ok()); } + #[tokio::test] + async fn session_id_stable_across_wrappers() { + let rt = AgentRuntimeBuilder::new() + .api_key("sk-test") + .model("x") + .build() + .unwrap(); + let id = rt.session_id(); + // Every send-style wrapper over the same runtime shares its session id. + let s1 = FuneraSession::with_id(rt.session_id(), rt.session_tx()); + let s2 = FuneraSession::with_id(rt.session_id(), rt.session_tx()); + assert_eq!(s1.id(), id); + assert_eq!(s2.id(), id); + assert_eq!(rt.session_id(), id); + } + #[tokio::test] async fn session_context_works_immediately() { let rt = AgentRuntimeBuilder::new() @@ -1087,6 +1180,28 @@ mod tests { )); } + #[tokio::test] + async fn reasoning_level_hot_reloads() { + let rt = AgentRuntimeBuilder::new() + .api_key("sk-test") + .model("x") + .build() + .unwrap(); + // Default is Medium. + assert_eq!(rt.reasoning_level().await, ReasoningLevel::Medium); + + let mut rx = rt.subscribe_env_state().await; + rt.set_reasoning_level(ReasoningLevel::High); + let got = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv()).await; + assert!(matches!( + got, + Ok(Ok(EnvStateEvent::ReasoningLevelChanged( + ReasoningLevel::High + ))) + )); + assert_eq!(rt.reasoning_level().await, ReasoningLevel::High); + } + // ── sandbox integration tests ─────────────────────────────────── #[cfg(feature = "sandbox")] diff --git a/funera-orchestrate/src/send_handle.rs b/funera-orchestrate/src/send_handle.rs index b5d14c2..afbe593 100644 --- a/funera-orchestrate/src/send_handle.rs +++ b/funera-orchestrate/src/send_handle.rs @@ -8,7 +8,9 @@ use tokio::task::JoinHandle; use funera_core::chat::session::SessionCmd; use funera_core::event_bus::env_state_bus::EnvStateEvent; +use funera_core::event_bus::token_bus::TokenUsage; use funera_core::provider::ChatProvider; +use tokio_util::sync::CancellationToken; use crate::error::OrchestrateError; use crate::event::AgentEvent; @@ -39,10 +41,11 @@ use crate::runtime::{Acquired, AgentRuntime, Idle}; /// # } /// ``` pub struct SendHandle { - pub(crate) runtime: AgentRuntime, - pub(crate) handle: JoinHandle>, - pub(crate) event_rx: broadcast::Receiver, + pub(crate) runtime: Option>, + pub(crate) handle: Option>>, + pub(crate) event_rx: Option>, pub(crate) env_state_tx: broadcast::Sender, + pub(crate) cancel: CancellationToken, } impl SendHandle

{ @@ -51,6 +54,8 @@ impl SendHandle

{ let (respond, rx) = tokio::sync::oneshot::channel(); let _ = self .runtime + .as_ref() + .expect("runtime present") .session_tx .send(SessionCmd::FetchContext { respond }); rx.await.unwrap_or_default() @@ -60,21 +65,32 @@ impl SendHandle

{ /// during react_loop execution. #[cfg(all(feature = "tool", feature = "security"))] pub fn approval_handle(&self) -> ApprovalHandle { - ApprovalHandle::new(self.runtime.env_cmd_tx.clone()) + ApprovalHandle::new( + self.runtime + .as_ref() + .expect("runtime present") + .env_cmd_tx + .clone(), + ) } - async fn wait(self) -> Result<(AgentRuntime, ChatResponse), OrchestrateError> { - // Wait for react_loop to complete - self.handle - .await - .map_err(|e| OrchestrateError::Session(e.into()))??; + async fn wait(mut self) -> Result<(AgentRuntime, ChatResponse), OrchestrateError> { + // Wait for react_loop to complete (an aborted task surfaces as Cancelled). + let handle = self.handle.take().expect("handle taken once"); + handle.await.map_err(|_| OrchestrateError::Cancelled)??; let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed); + if self.cancel.is_cancelled() { + return Err(OrchestrateError::Cancelled); + } + // Aggregate filtered events into ChatResponse - let resp = aggregate_from_broadcast(self.event_rx).await?; + let event_rx = self.event_rx.take().expect("event rx taken once"); + let resp = aggregate_from_broadcast(event_rx).await?; - Ok((self.runtime.into_idle(), resp)) + let runtime = self.runtime.take().expect("runtime taken once"); + Ok((runtime.into_idle(), resp)) } } @@ -87,6 +103,12 @@ impl IntoFuture for SendHandle

{ } } +impl Drop for SendHandle

{ + fn drop(&mut self) { + drop_cancel(&self.cancel); + } +} + // ═══════════════════════════════════════════════════════════════ // SendStreamHandle — stateful send + streaming // ═══════════════════════════════════════════════════════════════ @@ -95,17 +117,23 @@ impl IntoFuture for SendHandle

{ /// /// Like [`SendHandle`] but also provides `recv()` for per-event streaming. pub struct SendStreamHandle { - pub(crate) runtime: AgentRuntime, - pub(crate) handle: JoinHandle>, - pub(crate) event_rx: broadcast::Receiver, - pub(crate) stream_rx: mpsc::Receiver, + pub(crate) runtime: Option>, + pub(crate) handle: Option>>, + pub(crate) event_rx: Option>, + pub(crate) stream_rx: Option>, pub(crate) env_state_tx: broadcast::Sender, + pub(crate) cancel: CancellationToken, } impl SendStreamHandle

{ /// Receive the next streaming event. + /// + /// Returns `None` once the call is cancelled (or the stream is exhausted). pub async fn recv(&mut self) -> Option { - self.stream_rx.recv().await + if self.cancel.is_cancelled() { + return None; + } + self.stream_rx.as_mut()?.recv().await } /// Query session context while streaming is in progress. @@ -113,6 +141,8 @@ impl SendStreamHandle

{ let (respond, rx) = tokio::sync::oneshot::channel(); let _ = self .runtime + .as_ref() + .expect("runtime present") .session_tx .send(SessionCmd::FetchContext { respond }); rx.await.unwrap_or_default() @@ -122,16 +152,26 @@ impl SendStreamHandle

{ /// during react_loop execution. #[cfg(all(feature = "tool", feature = "security"))] pub fn approval_handle(&self) -> ApprovalHandle { - ApprovalHandle::new(self.runtime.env_cmd_tx.clone()) + ApprovalHandle::new( + self.runtime + .as_ref() + .expect("runtime present") + .env_cmd_tx + .clone(), + ) } - async fn wait(self) -> Result<(AgentRuntime, ChatResponse), OrchestrateError> { - self.handle - .await - .map_err(|e| OrchestrateError::Session(e.into()))??; + async fn wait(mut self) -> Result<(AgentRuntime, ChatResponse), OrchestrateError> { + let handle = self.handle.take().expect("handle taken once"); + handle.await.map_err(|_| OrchestrateError::Cancelled)??; let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed); - let resp = aggregate_from_broadcast(self.event_rx).await?; - Ok((self.runtime.into_idle(), resp)) + if self.cancel.is_cancelled() { + return Err(OrchestrateError::Cancelled); + } + let event_rx = self.event_rx.take().expect("event rx taken once"); + let resp = aggregate_from_broadcast(event_rx).await?; + let runtime = self.runtime.take().expect("runtime taken once"); + Ok((runtime.into_idle(), resp)) } } @@ -144,6 +184,63 @@ impl IntoFuture for SendStreamHandle

{ } } +impl Drop for SendStreamHandle

{ + fn drop(&mut self) { + drop_cancel(&self.cancel); + } +} + +// ═══════════════════════════════════════════════════════════════ +// FireHandle — one-shot result handle (no per-event stream) +// ═══════════════════════════════════════════════════════════════ + +/// Handle returned by [`Agent::fire`](crate::Agent::fire). +/// +/// One-shot (temporary session): `cancel()` interrupts the call and stops all +/// background work spawned for it, `await`/`wait()` yields the final +/// [`ChatResponse`]. Dropping the handle cancels automatically. +pub struct FireHandle { + pub(crate) handle: Option>>, + pub(crate) cancel: CancellationToken, + pub(crate) event_rx: Option>, + pub(crate) env_state_tx: broadcast::Sender, +} + +impl FireHandle { + /// Immediately interrupt this call: stop receiving output, abandon + /// in-flight tool executions, and notify middleware via + /// [`AgentEvent::Cancelled`]. + pub fn cancel(&self) { + drop_cancel(&self.cancel); + } + + async fn wait(mut self) -> Result { + let handle = self.handle.take().expect("handle taken once"); + handle.await.map_err(|_| OrchestrateError::Cancelled)??; + let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed); + if self.cancel.is_cancelled() { + return Err(OrchestrateError::Cancelled); + } + let event_rx = self.event_rx.take().expect("event rx taken once"); + aggregate_from_broadcast(event_rx).await + } +} + +impl IntoFuture for FireHandle { + type Output = Result; + type IntoFuture = Pin + Send>>; + + fn into_future(self) -> Self::IntoFuture { + Box::pin(async move { self.wait().await }) + } +} + +impl Drop for FireHandle { + fn drop(&mut self) { + drop_cancel(&self.cancel); + } +} + // ═══════════════════════════════════════════════════════════════ // FireStreamHandle — stateless stream (no runtime binding) // ═══════════════════════════════════════════════════════════════ @@ -154,24 +251,40 @@ impl IntoFuture for SendStreamHandle

{ /// temporary. Provides `recv()` for streaming and `wait()` (`IntoFuture`) for /// the final [`ChatResponse`]. pub struct FireStreamHandle { - pub(crate) handle: JoinHandle>, - pub(crate) event_rx: broadcast::Receiver, - pub(crate) stream_rx: mpsc::Receiver, + pub(crate) handle: Option>>, + pub(crate) event_rx: Option>, + pub(crate) stream_rx: Option>, pub(crate) env_state_tx: broadcast::Sender, + pub(crate) cancel: CancellationToken, } impl FireStreamHandle { /// Receive the next streaming event. + /// + /// Returns `None` once the call is cancelled (or the stream is exhausted). pub async fn recv(&mut self) -> Option { - self.stream_rx.recv().await + if self.cancel.is_cancelled() { + return None; + } + self.stream_rx.as_mut()?.recv().await + } + + /// Immediately interrupt this call: stop receiving output, abandon + /// in-flight tool executions, and notify middleware via + /// [`AgentEvent::Cancelled`]. + pub fn cancel(&self) { + drop_cancel(&self.cancel); } - async fn wait(self) -> Result { - self.handle - .await - .map_err(|e| OrchestrateError::Session(e.into()))??; + async fn wait(mut self) -> Result { + let handle = self.handle.take().expect("handle taken once"); + handle.await.map_err(|_| OrchestrateError::Cancelled)??; let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed); - aggregate_from_broadcast(self.event_rx).await + if self.cancel.is_cancelled() { + return Err(OrchestrateError::Cancelled); + } + let event_rx = self.event_rx.take().expect("event rx taken once"); + aggregate_from_broadcast(event_rx).await } } @@ -184,10 +297,28 @@ impl IntoFuture for FireStreamHandle { } } +impl Drop for FireStreamHandle { + fn drop(&mut self) { + drop_cancel(&self.cancel); + } +} + // ═══════════════════════════════════════════════════════════════ // Internal helpers // ═══════════════════════════════════════════════════════════════ +/// Cancel a call's token. +/// +/// Cancellation is fully cooperative: every blocking await point of the ReAct +/// loop — stream consumption, the initial provider request, in-flight tool +/// executions — listens for the token via `select!`, so cancelling makes the +/// loop exit promptly and emit [`AgentEvent::Cancelled`] to middleware and +/// subscribers. No abort is needed, which keeps the notification reliable. +/// Used by every handle's `cancel()` and `Drop`. +fn drop_cancel(cancel: &CancellationToken) { + cancel.cancel(); +} + /// Aggregate filtered events from a broadcast receiver into a ChatResponse. async fn aggregate_from_broadcast( mut event_rx: broadcast::Receiver, @@ -196,6 +327,7 @@ async fn aggregate_from_broadcast( let mut tool_calls = Vec::new(); let mut iterations = 0usize; let mut finish_reason: Option = None; + let mut usage: Option = None; let mut pending: Vec<(Arc, String, serde_json::Value)> = Vec::new(); loop { @@ -220,7 +352,13 @@ async fn aggregate_from_broadcast( } } Ok(AgentEvent::TurnStart) => iterations += 1, - Ok(AgentEvent::TurnEnd { finish_reason: fr }) => finish_reason = fr, + Ok(AgentEvent::TurnEnd { + finish_reason: fr, + usage: u, + }) => { + finish_reason = fr; + usage = u; + } Ok(AgentEvent::Done) => break, Err(broadcast::error::RecvError::Closed) => break, Err(broadcast::error::RecvError::Lagged(_)) => continue, @@ -233,6 +371,7 @@ async fn aggregate_from_broadcast( tool_calls, iterations, finish_reason, + usage, }) } diff --git a/funera-orchestrate/tests/integration.rs b/funera-orchestrate/tests/integration.rs index 58cb3b8..6d00263 100644 --- a/funera-orchestrate/tests/integration.rs +++ b/funera-orchestrate/tests/integration.rs @@ -36,6 +36,8 @@ async fn fire_simple_response() { let resp = agent .fire("Say exactly: hello world", &runtime) .await + .unwrap() + .await .unwrap(); assert!( non_empty(&resp.content), @@ -51,7 +53,12 @@ async fn fire_with_custom_model() { .system_prompt("Reply with one word.") .build(); - let resp = agent.fire("Say: Rust", &runtime).await.unwrap(); + let resp = agent + .fire("Say: Rust", &runtime) + .await + .unwrap() + .await + .unwrap(); assert!(non_empty(&resp.content)); } @@ -66,7 +73,12 @@ async fn fire_respects_max_iterations() { .unwrap(); let agent = Agent::builder().system_prompt("You are concise.").build(); - let resp = agent.fire("What is Rust?", &runtime).await.unwrap(); + let resp = agent + .fire("What is Rust?", &runtime) + .await + .unwrap() + .await + .unwrap(); assert!(non_empty(&resp.content)); } From 4829f274d03c3594fb1e1b6a42c0c53e7e0d4d3c Mon Sep 17 00:00:00 2001 From: dynamder Date: Thu, 20 Aug 2026 00:38:30 +0800 Subject: [PATCH 3/5] docs: proofread docs against code, rebuild mdBook, fix tool_policy command - README: cancellation feature bullet + cancel example row; fire() double-await in quick starts - docs/architecture.md: Cancellation section (handles, cooperative exit, AgentEvent::Cancelled never written to session) - docs/examples/README.md: cancel example listed - docs/getting-started.md: fix tool_policy run command (missing --features security); fire() double-await - docs/concepts/effects.md: tool registration is an internal pairing, not an extension API (external code must not couple to the tool registry) - CHANGELOG: serde/session-id/usage/parallel/reasoning/cancellation entries - book/: rebuilt with mdbook (includes the Cancellation section) --- CHANGELOG.md | 29 +++++++++++++++++++++++++++++ README.md | 30 +++++++++++++++++++++++------- docs/architecture.md | 17 +++++++++++++++++ docs/concepts/effects.md | 31 ++++++++++++++++++++++++++----- docs/examples/README.md | 2 +- docs/getting-started.md | 4 ++-- src/lib.rs | 2 +- 7 files changed, 99 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c5dd18..e7368e2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,35 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ### Added +- Cancellation — every entry point (`fire` → `FireHandle`, `fire_stream`, + `send`, `send_stream`) carries a per-call `CancellationToken`; `cancel()` + or dropping the handle immediately stops receiving output, abandons + in-flight tool executions (workers survive), and notifies middleware via + `AgentEvent::Cancelled` (never written to session history). `fire` now + returns a `FireHandle` — `await` it for the `ChatResponse`. +- `serde` feature on `funera-core` — gated `Serialize`/`Deserialize` derives + for the chat message types (`FuneraMessage`, `MsgVariant`, `Role`, + `TextMessage`, `ToolRequestMessage`, `ToolResponseMessage`). +- `AgentRuntime::session_id` — stable conversation id reused by `send` / + `send_stream`; `fire` / `fire_stream` keep a fresh one-shot id (fork + semantics). +- Token usage tracking — `TokenUsage` + `TokenEvent::Usage`; requests ask for + `stream_options.include_usage`; `ChatResponse.usage` carries the last turn's + usage. Cost computation is left to callers. +- Parallel tool execution — the tool bus fans out to `max_concurrent_tools` + (default 4) workers via `AgentRuntimeBuilder::max_concurrent_tools`, so + multiple tool calls in one turn run concurrently. +- Reasoning levels — `ReasoningLevel` + (`Off`/`Minimal`/`Low`/`Medium`/`High`/`XHigh`/`Max`, default `Medium`) with + hot-reload via `AgentRuntime::set_reasoning_level`; providers interpret it + (DeepSeek mirrors the official harness: `thinking` + `reasoning_effort` for + `high`/`max`; OpenAI maps to `reasoning_effort`, clamping `xhigh`/`max` to + `high`). +- `Tool::is_shell_tool` marker (default `false`) — tools that execute shell + commands opt in explicitly, and `ShellPolicy` scrutiny is now driven by the + marker instead of a hard-coded tool-name list + (`shell`/`bash`/`sh`/`cmd`/`powershell`). `ToolPolicy::check_shell_command` + now takes the tool itself. - `FuneraEnv::effect` / `dispose` — reversible effects with LIFO teardown: every registration can be paired with its inverse, and disposal runs all inverses in reverse registration order (idempotent and panic-isolated). `EnvActor` calls diff --git a/README.md b/README.md index 620f59a..3b1cb62 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ async fn main() -> Result<(), Box> { .system_prompt("You are a helpful assistant.") .build(); - let resp = agent.fire("Hello!", &runtime).await?; + let resp = agent.fire("Hello!", &runtime).await?.await?.await?; println!("{}", resp.content); Ok(()) } @@ -179,7 +179,8 @@ sequenceDiagram - **Skill system** — load prompt templates from YAML-frontmatter Markdown files - **Middleware pipeline** — intercept agent events with inspectors (read-only, parallel) and mutators (pass/modify/block, sequential) - **Security layer** — tool/shell policies, path allowlisting, audit logging, secure API key storage -- **Type-state session** — compile-time enforcement of session ownership (`Idle` / `Acquired`) +- **Type-state session** — compile-time enforcement of session ownership (Idle / Acquired) +- **Cancellation** — every call returns a cancellable handle ( ire → FireHandle, ire_stream / send / send_stream); cancel() or dropping the handle immediately stops receiving output, abandons in-flight tool executions, and notifies middleware via AgentEvent::Cancelled (never written to session history). ## Examples @@ -192,6 +193,7 @@ sequenceDiagram | `middleware` | Inspector/Mutator middleware pipeline | | `reversible_effects` | LIFO teardown of registered effects (no LLM) | | `tool_policy` / `security` / `sandbox` | Security policies, audit, and sandboxing | +| `cancel` | Cancellation: explicit `cancel()`, drop-to-cancel, timeout via `select!` | ## Reversible effects @@ -217,9 +219,11 @@ env.effect(|| { disposer is caught and logged; the rest still run). - `EnvActor` calls `dispose()` automatically once the runtime is dropped, so anything registered against the env is reverted — no memory or service leaks. -- For tool registrations, the safe inverse is `remove_tool_if_same`: it removes a tool only if - the registered entry is the *same* `Arc`, so a stale teardown never deletes a replacement - tool that reuses the same name. +- Tool registrations are paired with their exact inverse internally + (`remove_tool_if_same` removes a tool only if the registered entry is the + *same* `Arc`, so a stale teardown never deletes a replacement that reuses + the same name); external extensions register tools through the runtime API + and do not reach into the tool registry themselves. Run the demo: @@ -277,7 +281,7 @@ async fn main() -> Result<(), Box> { .system_prompt("You are a helpful assistant.") .build(); - let resp = agent.fire("Hello!", &runtime).await?; + let resp = agent.fire("Hello!", &runtime).await?.await?.await?; println!("{}", resp.content); Ok(()) } @@ -387,6 +391,18 @@ let runtime = AgentRuntime::::builder() .build()?; ``` +If your tool executes shell commands, opt in explicitly via `is_shell_tool` +so the security layer's `ShellPolicy` applies — based on the marker, not on +the tool's registered name: + +```rust +impl Tool for Shell { + fn name(&self) -> &str { "shell" } + // ... + fn is_shell_tool(&self) -> bool { true } +} +``` + ### Security configuration Requires the `security` feature (and optionally `funera-builtin-tools`, `sandbox`): @@ -410,7 +426,7 @@ async fn main() -> Result<(), Box> { .build()?; let agent = Agent::builder().build(); - let resp = agent.fire("List the git log.", &runtime).await?; + let resp = agent.fire("List the git log.", &runtime).await?.await?.await?; println!("{}", resp.content); Ok(()) } diff --git a/docs/architecture.md b/docs/architecture.md index 995484d..83393a1 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -35,3 +35,20 @@ env is torn down: This guarantees that anything registered against the env — tools, listeners, connections — is reverted on teardown, preventing memory and service leaks. See [Reversible Effects](concepts/effects.md) for the pattern. + +## Cancellation + +Every agent call returns a cancellable handle — `FireHandle` (one-shot `fire`), +`FireStreamHandle`, `SendHandle`, `SendStreamHandle`. `cancel()` — or simply +dropping the handle — cancels the call's `CancellationToken`. The `ReActLoop` +exits cooperatively at its next blocking point (stream consumption, the initial +provider request, in-flight tool execution) and: + +- stops receiving output immediately; +- abandons in-flight tool executions (the shared `ToolExecutor` workers survive); +- emits `AgentEvent::Cancelled` through the middleware chain and to event + subscribers — **never** into session history — so middleware holding external + services can clean up. + +Awaiting a cancelled handle yields `OrchestrateError::Cancelled`; `recv()` on a +streaming handle returns `None` once cancelled. \ No newline at end of file diff --git a/docs/concepts/effects.md b/docs/concepts/effects.md index 4a1140d..3d7b749 100644 --- a/docs/concepts/effects.md +++ b/docs/concepts/effects.md @@ -28,13 +28,33 @@ env.effect(|| { sender is gone — i.e. when the runtime is dropped. You can also call `env.dispose()` manually to tear down an env you created directly. -## Leak-safe tool registration +## Tool registration is an internal pairing — not an extension API -The safe inverse of `add_tool` is `remove_tool_if_same`: it removes the tool only if the -registered entry is the *same* `Arc` value, so a stale teardown can never delete a replacement -tool that reuses the same name: +Inside the framework, a tool registration is paired with its exact inverse so +teardown is leak-safe: the safe inverse of `add_tool` is `remove_tool_if_same`, +which removes a tool only if the registered entry is the *same* `Arc` value. A +stale teardown can therefore never delete a replacement tool that reuses the +same name. `EnvActor` performs this pairing for the registrations it owns, and +`dispose()` runs the paired inverses in LIFO order. + +This registry-level pairing is an **internal implementation detail**, not a +public extension pattern: + +- `FuneraEnv::tool_registry` and its mutation methods (`add_tool`, + `remove_tool_if_same`) are crate-private; external code cannot reach them. +- External extensions should not couple their effects to tool-registry state. + If you need tools, register them through the public runtime API — + `with_tool_instance` / `add_tool` / `remove_tool` on `AgentRuntime` — and + let the framework own the registration lifecycle (including disposal). +- Reserve `FuneraEnv::effect` for effects you create yourself (connections, + listeners, resources), whose inverses you control. + +For reference, the internal pattern — how the framework's own env pairs a +registration with its inverse (mirrored in the crate's tests): ```rust,ignore +// internal: `tool_registry` and `add_tool` are crate-private — this is how +// the framework pairs its own registrations, not a public extension API. let tool: Arc = Arc::new(MyTool); env.add_tool(Arc::clone(&tool)).await; @@ -50,7 +70,8 @@ env.effect(move || { }); ``` -When the env is disposed, exactly this tool is removed — never a later replacement. +When the env is disposed, exactly this tool is removed — never a later +replacement. ## Run the example diff --git a/docs/examples/README.md b/docs/examples/README.md index e4c10b0..b21f9b8 100644 --- a/docs/examples/README.md +++ b/docs/examples/README.md @@ -17,4 +17,4 @@ LLM examples require an API key: - `minimal`, `multi_turn`, `streaming`, `custom_tool`, `middleware`, `skills`, `builtin_tools`, `session_reset`, `raw_events`, `multi_runtime`, `sandbox`, - `security`, `streaming_with_tools` + `security`, `streaming_with_tools`, `cancel` diff --git a/docs/getting-started.md b/docs/getting-started.md index 4310e3b..39a7bf2 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -23,7 +23,7 @@ async fn main() -> Result<(), Box> { .system_prompt("You are a helpful assistant.") .build(); - let resp = agent.fire("Hello!", &runtime).await?; + let resp = agent.fire("Hello!", &runtime).await?.await?.await?; println!("{}", resp.content); Ok(()) } @@ -39,5 +39,5 @@ Examples that do not require an LLM API key: ```bash cargo run -p funera-orchestrate --example reversible_effects -cargo run -p funera-orchestrate --example tool_policy +cargo run -p funera-orchestrate --example tool_policy --features security ``` diff --git a/src/lib.rs b/src/lib.rs index 1dc1f49..1bd8b0c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,7 +17,7 @@ //! .system_prompt("You are a helpful assistant.") //! .build(); //! -//! let resp = agent.fire("Hello!", &runtime).await?; +//! let resp = agent.fire("Hello!", &runtime).await?.await?.await?; //! println!("{}", resp.content); //! ``` //! From 339a5ea5a015f145aa56580ba3ae0c05f1b9aeea Mon Sep 17 00:00:00 2001 From: dynamder Date: Thu, 20 Aug 2026 00:39:54 +0800 Subject: [PATCH 4/5] feat(builtin-tools): ShellTool opts in via Tool::is_shell_tool ShellPolicy scrutiny now applies by the marker instead of a hard-coded tool-name list. --- funera_builtin_tools/src/shell.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/funera_builtin_tools/src/shell.rs b/funera_builtin_tools/src/shell.rs index 405cbe3..895140c 100644 --- a/funera_builtin_tools/src/shell.rs +++ b/funera_builtin_tools/src/shell.rs @@ -69,6 +69,10 @@ impl Tool for ShellTool { "Execute shell commands. Use with caution." } + fn is_shell_tool(&self) -> bool { + true + } + fn schema(&self) -> JsonValue { json!({ "type": "function", From f5f41ba4d9d7709d1f599079014b5cec399c5902 Mon Sep 17 00:00:00 2001 From: dynamder Date: Sun, 23 Aug 2026 13:57:37 +0800 Subject: [PATCH 5/5] test(core,builtin-tools): wire set_tool_availability, close mutation-test gaps - core: FuneraEnv::set_tool_availability now actually toggles tool availability in the registry (RawToolRegistry / GuardedToolRegistry) and rebroadcasts the watch snapshot only when the availability changed - builtin-tools: add 19 tests covering end-anchor range replaces, multi-line prepend, exact anchor output, the LINE#HASH prefix guard, tool metadata, shell output formatting and the pinned anchor hash; verified via cargo-mutants (funera-builtin-tools: 42 missed -> 0 missed) - chore: add .cargo/mutants.toml excluding four provably-equivalent mutants (edit.rs range bounds, shell exit-code default) with rationale --- .cargo/mutants.toml | 24 ++++ funera_builtin_tools/src/edit.rs | 197 +++++++++++++++++++++++++++ funera_builtin_tools/src/hashline.rs | 8 ++ funera_builtin_tools/src/read.rs | 42 ++++++ funera_builtin_tools/src/shell.rs | 63 +++++++++ funera_builtin_tools/src/write.rs | 10 ++ funera_core/src/env.rs | 49 ++++++- funera_core/src/re_act/tool.rs | 50 +++++++ funera_core/src/security/registry.rs | 24 ++++ 9 files changed, 463 insertions(+), 4 deletions(-) create mode 100644 .cargo/mutants.toml diff --git a/.cargo/mutants.toml b/.cargo/mutants.toml new file mode 100644 index 0000000..4552100 --- /dev/null +++ b/.cargo/mutants.toml @@ -0,0 +1,24 @@ +# cargo-mutants configuration. +# +# Options docs: https://github.com/sourcefrog/cargo-mutants + +exclude_re = [ + # Equivalent mutants in EditTool::apply_ops (funera_builtin_tools/src/edit.rs:163). + # Verified anchors always satisfy `idx <= lines.len() - 1`, so + # `idx >= lines.len() - 1` and its `+ 1` / `/ 1` alternatives compute the + # same insert position in every reachable state — no black-box test can + # distinguish them. + "replace - with \\+ in EditTool::apply_ops", + "replace - with / in EditTool::apply_ops", + # Equivalent in EditTool::apply_ops (edit.rs:157): when the range bounds + # check is reached, both operands are statically false — the end anchor was + # verified within the file (end_idx < len) and the start guard at line 135 + # guarantees end_idx >= idx. `false || false` and `false && false` are + # indistinguishable. + "replace \\|\\| with && in EditTool::apply_ops", + # Equivalent on Windows: `ExitStatus::code()` is always `Some` there, so the + # `unwrap_or(-1)` default in ShellTool::format_output is unreachable and the + # `- 1` replacement cannot be observed. (On Unix it can be exercised by a + # signal-death test; remove this pattern to re-enable it on such CI.) + "delete - in format_output", +] diff --git a/funera_builtin_tools/src/edit.rs b/funera_builtin_tools/src/edit.rs index 3594652..213f2bf 100644 --- a/funera_builtin_tools/src/edit.rs +++ b/funera_builtin_tools/src/edit.rs @@ -512,6 +512,59 @@ mod tests { assert_eq!(result, "a\nx\ny\nc"); } + // --- apply_ops: end anchors (range replace) --- + + #[test] + fn replace_range_with_end_after_start() { + let lines: Vec = vec!["a".into(), "b".into(), "c".into()]; + let h1 = compute_line_anchor(&lines, 1); + let h2 = compute_line_anchor(&lines, 2); + let edit = Edit { + operation: serde_json::from_value(json!("replace")).unwrap(), + pos: Some(format!("1#{}", h1)), + end: Some(format!("2#{}", h2)), + lines: vec!["x".into(), "y".into()], + old_text: None, + new_text: None, + }; + let result = EditTool::apply_ops("a\nb\nc", &[edit]).unwrap(); + assert_eq!(result, "x\ny\nc"); + } + + #[test] + fn replace_range_with_end_equal_start() { + let lines: Vec = vec!["a".into(), "b".into(), "c".into()]; + let h2 = compute_line_anchor(&lines, 2); + let edit = Edit { + operation: serde_json::from_value(json!("replace")).unwrap(), + pos: Some(format!("2#{}", h2)), + end: Some(format!("2#{}", h2)), + lines: vec!["x".into()], + old_text: None, + new_text: None, + }; + let result = EditTool::apply_ops("a\nb\nc", &[edit]).unwrap(); + assert_eq!(result, "a\nx\nc"); + } + + #[test] + fn replace_range_with_end_before_start_rejected() { + let lines: Vec = vec!["a".into(), "b".into(), "c".into()]; + let h1 = compute_line_anchor(&lines, 1); + let h2 = compute_line_anchor(&lines, 2); + let edit = Edit { + operation: serde_json::from_value(json!("replace")).unwrap(), + pos: Some(format!("2#{}", h2)), + end: Some(format!("1#{}", h1)), + lines: vec!["x".into()], + old_text: None, + new_text: None, + }; + let result = EditTool::apply_ops("a\nb\nc", &[edit]); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("end line before start line")); + } + // --- apply_ops: append --- #[test] @@ -561,6 +614,22 @@ mod tests { assert_eq!(result, "a\nx\nb"); } + #[test] + fn prepend_multiple_lines() { + let lines: Vec = vec!["a".into(), "b".into(), "c".into()]; + let h2 = compute_line_anchor(&lines, 2); + let edit = Edit { + operation: serde_json::from_value(json!("prepend")).unwrap(), + pos: Some(format!("2#{}", h2)), + end: None, + lines: vec!["x".into(), "y".into()], + old_text: None, + new_text: None, + }; + let result = EditTool::apply_ops("a\nb\nc", &[edit]).unwrap(); + assert_eq!(result, "a\nx\ny\nb\nc"); + } + // --- apply_ops: replace_text --- #[test] @@ -609,6 +678,42 @@ mod tests { assert!(result.unwrap_err().contains("STALE_ANCHOR")); } + // --- build_result_anchors: exact output --- + + #[test] + fn build_result_anchors_exact() { + let new_content = "a\nb\nc"; + let lines: Vec = vec!["a".into(), "b".into(), "c".into()]; + let h1 = compute_line_anchor(&lines, 1); + let edit = Edit { + operation: serde_json::from_value(json!("replace")).unwrap(), + pos: Some(format!("1#{}", h1)), + end: None, + lines: vec!["x".into()], + old_text: None, + new_text: None, + }; + + // An edit anchored at line 1 renders the affected context (lines 1-2) + // with exact hashline anchors — pins the wire format of the result. + let expected = format!( + "--- Anchors A-B ---\nline#hash:content\n{}\n{}\n", + hashline::format_line_trimmed(1, &hashline::compute_anchor("", "a", "b"), "a"), + hashline::format_line_trimmed(2, &hashline::compute_anchor("a", "b", "c"), "b"), + ); + let out = EditTool::build_result_anchors(new_content, &[edit]); + assert_eq!(out, expected); + } + + // --- description --- + + #[test] + fn description_mentions_anchors() { + let d = EditTool.description(); + assert!(d.len() > 20, "description too short: {d}"); + assert!(d.contains("anchor"), "description should mention anchors: {d}"); + } + // --- Full execute flow --- #[tokio::test] @@ -630,6 +735,15 @@ mod tests { })) .await; assert!(result.is_ok(), "replace failed: {:?}", result.err()); + let result_str = result.unwrap(); + assert!( + result_str.contains("--- Anchors A-B ---"), + "result should include the anchors section: {result_str}" + ); + assert!( + result_str.contains("line#hash:content"), + "result should include the anchors header: {result_str}" + ); let new_content = tokio::fs::read_to_string(&path).await.unwrap(); assert_eq!(new_content, "line1\nmodified!\nline3\n"); cleanup("replace").await; @@ -678,6 +792,89 @@ mod tests { } } + // ── LINE#HASH prefix guard in edit lines ───────────────────── + + /// A 2-part "N#HASH" prefix (uppercase hash) must be rejected. + #[tokio::test] + async fn edit_execute_rejects_line_hash_prefix() { + let path = write_test_file("hashprefix", "hashprefix.txt", "hello\n").await; + let content = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec = content.lines().map(|l| l.to_string()).collect(); + let hash = compute_line_anchor(&lines, 1); + + let tool = EditTool; + let result = tool + .execute(json!({ + "filePath": path.to_string_lossy(), + "edits": [{ + "op": "append", + "pos": format!("1#{}", hash), + "lines": [" 11#KT world"] + }] + })) + .await; + assert!( + matches!(result, Err(ToolCallError::ParameterMismatch(_))), + "a 2-part LINE#HASH prefix must be rejected: {result:?}" + ); + cleanup("hashprefix").await; + } + + /// A hash of exactly two uppercase chars must also be rejected. + #[tokio::test] + async fn edit_execute_rejects_two_char_hash_prefix() { + let path = write_test_file("hash2", "hash2.txt", "hello\n").await; + let content = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec = content.lines().map(|l| l.to_string()).collect(); + let hash = compute_line_anchor(&lines, 1); + + let tool = EditTool; + let result = tool + .execute(json!({ + "filePath": path.to_string_lossy(), + "edits": [{ + "op": "append", + "pos": format!("1#{}", hash), + "lines": [" 1#AB"] + }] + })) + .await; + assert!( + matches!(result, Err(ToolCallError::ParameterMismatch(_))), + "a two-char uppercase hash prefix must be rejected: {result:?}" + ); + cleanup("hash2").await; + } + + /// Mixed-case 2-char hashes are NOT anchor prefixes: the edit must pass + /// the guard and apply successfully. + #[tokio::test] + async fn edit_execute_accepts_mixed_case_hash_prefix() { + let path = write_test_file("hashmix", "hashmix.txt", "hello\n").await; + let content = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec = content.lines().map(|l| l.to_string()).collect(); + let hash = compute_line_anchor(&lines, 1); + + let tool = EditTool; + let result = tool + .execute(json!({ + "filePath": path.to_string_lossy(), + "edits": [{ + "op": "append", + "pos": format!("1#{}", hash), + "lines": ["1#Ab"] + }] + })) + .await; + assert!( + result.is_ok(), + "mixed-case hash is not an anchor prefix and must pass: {result:?}" + ); + let new_content = tokio::fs::read_to_string(&path).await.unwrap(); + assert_eq!(new_content, "hello\n1#Ab\n"); + cleanup("hashmix").await; + } + #[tokio::test] async fn edit_execute_noop_loop() { let path = write_test_file("noop", "noop.txt", "hello\n").await; diff --git a/funera_builtin_tools/src/hashline.rs b/funera_builtin_tools/src/hashline.rs index cf0e899..8ed4764 100644 --- a/funera_builtin_tools/src/hashline.rs +++ b/funera_builtin_tools/src/hashline.rs @@ -63,6 +63,14 @@ mod tests { assert_eq!(a, b); } + #[test] + fn compute_hash_pins_known_input() { + // Pins the exact anchor byte for a known context so the low-byte + // mask (`& 0xFF`) cannot silently become a different bitwise op + // without breaking the read/edit anchor wire format. + assert_eq!(compute_hash("a", "b", "c"), 0x9E); + } + #[test] fn hash_different_context_different_result() { let h1 = compute_hash("", "line", ""); diff --git a/funera_builtin_tools/src/read.rs b/funera_builtin_tools/src/read.rs index d853c63..ab9f15e 100644 --- a/funera_builtin_tools/src/read.rs +++ b/funera_builtin_tools/src/read.rs @@ -164,6 +164,13 @@ mod tests { } } + #[test] + fn description_mentions_hash_anchors() { + let d = ReadTool.description(); + assert!(d.len() > 20, "description too short: {d}"); + assert!(d.contains("LINE#HASH"), "description should mention anchors: {d}"); + } + #[tokio::test] async fn read_nonexistent_file() { let tool = ReadTool; @@ -230,6 +237,41 @@ mod tests { cleanup("limit").await; } + /// `limit: 0` must not truncate: only a positive limit applies. + #[tokio::test] + async fn read_limit_zero_returns_whole_file() { + let path = setup_file("limitzero", "limitzero.txt", "1\n2\n3\n").await; + let tool = ReadTool; + let result = tool + .execute(json!({"filePath": path.to_string_lossy(), "limit": 0})) + .await; + assert!(result.is_ok()); + let output = result.unwrap(); + assert!(output.contains("1")); + assert!(output.contains("3"), "limit=0 must return the whole file: {output}"); + cleanup("limitzero").await; + } + + /// Pins the exact anchor output so the read/write anchor wire format + /// cannot silently change (prev/curr/next hashing, line numbering). + #[tokio::test] + async fn read_with_anchors_exact() { + let path = setup_file("anchors", "anchors.txt", "a\nb\nc\n").await; + let tool = ReadTool; + let result = tool + .execute(json!({"filePath": path.to_string_lossy()})) + .await; + assert!(result.is_ok()); + let expected = format!( + "{}\n{}\n{}", + hashline::format_line_trimmed(1, &hashline::compute_anchor("", "a", "b"), "a"), + hashline::format_line_trimmed(2, &hashline::compute_anchor("a", "b", "c"), "b"), + hashline::format_line_trimmed(3, &hashline::compute_anchor("b", "c", ""), "c"), + ); + assert_eq!(result.unwrap(), expected); + cleanup("anchors").await; + } + #[tokio::test] async fn read_empty_file() { let path = setup_file("empty", "empty.txt", "").await; diff --git a/funera_builtin_tools/src/shell.rs b/funera_builtin_tools/src/shell.rs index 895140c..a1c252a 100644 --- a/funera_builtin_tools/src/shell.rs +++ b/funera_builtin_tools/src/shell.rs @@ -273,6 +273,69 @@ mod tests { assert!(result.unwrap().contains("42")); } + // ── metadata / config tests ───────────────────────────────────── + + #[test] + fn shell_description_mentions_caution() { + let tool = ShellTool::new(); + let d = tool.description(); + assert!(d.len() > 20, "description too short: {d}"); + assert!(d.contains("caution"), "description should warn: {d}"); + } + + #[test] + fn shell_is_shell_tool_true() { + assert!( + ShellTool::new().is_shell_tool(), + "shell tool must opt into shell policy scrutiny" + ); + } + + #[cfg(feature = "sandbox")] + #[test] + fn with_sandbox_stores_policy() { + let policy = funera_core::security::sandbox::SandboxPolicy::default(); + let tool = ShellTool::with_sandbox(policy); + assert!( + tool.sandbox_policy.is_some(), + "with_sandbox must store the policy" + ); + } + + // ── output formatting ─────────────────────────────────────────── + + #[tokio::test] + async fn shell_stdout_only_has_no_stderr_section() { + let tool = ShellTool::new(); + let result = tool.execute(json!({"command": "echo hello"})).await; + assert!(result.is_ok()); + let output = result.unwrap(); + assert!(output.contains("stdout:")); + assert!( + !output.contains("stderr:"), + "stdout-only output must not contain a stderr section: {output:?}" + ); + } + + #[tokio::test] + async fn shell_stderr_only_no_leading_newline() { + let tool = ShellTool::new(); + let cmd = if cfg!(target_os = "windows") { + "cmd /c echo oops 1>&2" + } else { + "sh -c 'echo oops >&2'" + }; + let result = tool.execute(json!({"command": cmd})).await; + assert!(result.is_ok()); + let output = result.unwrap(); + assert!(output.contains("stderr:")); + assert!(output.contains("oops")); + assert!( + !output.starts_with('\n'), + "stderr-only output must not start with a newline: {output:?}" + ); + } + #[tokio::test] async fn shell_timeout_param() { let tool = ShellTool::new(); diff --git a/funera_builtin_tools/src/write.rs b/funera_builtin_tools/src/write.rs index 22bbbfb..ffe8d4c 100644 --- a/funera_builtin_tools/src/write.rs +++ b/funera_builtin_tools/src/write.rs @@ -97,6 +97,16 @@ mod tests { let _ = tokio::fs::remove_dir_all(test_dir(label)).await; } + #[test] + fn description_mentions_parent_directories() { + let d = WriteTool.description(); + assert!(d.len() > 20, "description too short: {d}"); + assert!( + d.contains("parent directories"), + "description should mention parent dirs: {d}" + ); + } + #[tokio::test] async fn write_missing_file_path() { let tool = WriteTool; diff --git a/funera_core/src/env.rs b/funera_core/src/env.rs index 2dbb571..ca22be6 100644 --- a/funera_core/src/env.rs +++ b/funera_core/src/env.rs @@ -165,9 +165,11 @@ impl FuneraEnv { } #[cfg(feature = "tool")] - pub(crate) async fn set_tool_availability(&mut self, _name: &str, _available: bool) { - let registry = self.tool_registry.read().await; - let _ = self.tool_tx.send(registry.available_tools_json()); + pub(crate) async fn set_tool_availability(&mut self, name: &str, available: bool) { + let mut registry = self.tool_registry.write().await; + if registry.set_tool_availability(name, available) { + let _ = self.tool_tx.send(registry.available_tools_json()); + } } pub(crate) fn set_client(&mut self, client: async_openai::Client) { @@ -598,13 +600,52 @@ mod tests { } #[tokio::test] - async fn set_tool_availability_rebroadcasts_snapshot() { + async fn set_tool_availability_toggles_tool() { let (mut env, mut watcher) = FuneraEnv::new(async_openai::Client::new(), "m"); env.add_tool(Arc::new(MockTool)).await; let _ = watcher.watch_tool(); assert!(!watcher.has_tool_changed()); + + // Disabling drops the tool from the snapshot and blocks execution. env.set_tool_availability("mock", false).await; assert!(watcher.has_tool_changed()); + assert!( + watcher + .watch_tool() + .as_array() + .is_some_and(|a| a.is_empty()) + ); + assert!(matches!( + env.tool_registry + .read() + .await + .call_tool("mock", json!({})) + .await, + Err(ToolCallError::ToolUnavailable(_)) + )); + + // Re-enabling restores the tool in the snapshot and execution. + let _ = watcher.watch_tool(); + env.set_tool_availability("mock", true).await; + assert!( + watcher + .watch_tool() + .as_array() + .is_some_and(|a| a.len() == 1) + ); + assert!( + env.tool_registry + .read() + .await + .call_tool("mock", json!({})) + .await + .is_ok() + ); + + // Unknown tool is a silent no-op: no spurious notification. + let _ = watcher.watch_tool(); + env.set_tool_availability("ghost", false).await; + assert!(!watcher.has_tool_changed()); } #[test] diff --git a/funera_core/src/re_act/tool.rs b/funera_core/src/re_act/tool.rs index 40aca7c..c4a7a0e 100644 --- a/funera_core/src/re_act/tool.rs +++ b/funera_core/src/re_act/tool.rs @@ -181,6 +181,21 @@ impl RawToolRegistry { self.tools.remove(name); } + /// Set whether a registered tool is available for execution. + /// + /// Returns `true` if the tool exists and its availability actually + /// changed (so callers can skip rebroadcasting when nothing changed), + /// `false` for an unknown name or a no-op set. + pub fn set_tool_availability(&mut self, name: &str, available: bool) -> bool { + match self.tools.get_mut(name) { + Some(entry) if entry.available != available => { + entry.available = available; + true + } + _ => false, + } + } + /// Remove the tool only if the registered entry is the same `Arc` value. /// This prevents a stale disposer from deleting a replacement tool that /// reuses the same name (e.g. HMR replacement). @@ -314,4 +329,39 @@ mod tests { assert!(reg.remove_tool_if_same("mock", &original)); assert!(!reg.tool_exists("mock")); } + + #[test] + fn set_tool_availability_toggles_entry() { + let mut reg = RawToolRegistry::new(); + // Unknown tool: no-op. + assert!(!reg.set_tool_availability("mock", false)); + + reg.add_tool(Arc::new(MockTool)); + assert!(reg.set_tool_availability("mock", false)); + assert!( + reg.get_tool("mock").is_some_and(|e| !e.is_available()), + "entry must be marked unavailable" + ); + assert!( + reg.get_tool_arc("mock").is_none(), + "unavailable tool must not be runnable" + ); + assert!( + reg.available_tools_json().as_array().is_some_and(|a| a.is_empty()), + "unavailable tool must drop out of the snapshot" + ); + + // Setting the same value again is a no-op. + assert!(!reg.set_tool_availability("mock", false)); + + assert!(reg.set_tool_availability("mock", true)); + assert!( + reg.get_tool_arc("mock").is_some(), + "re-enabled tool must be runnable again" + ); + assert!( + reg.available_tools_json().as_array().is_some_and(|a| a.len() == 1), + "re-enabled tool must return to the snapshot" + ); + } } diff --git a/funera_core/src/security/registry.rs b/funera_core/src/security/registry.rs index cd2044d..dd8690b 100644 --- a/funera_core/src/security/registry.rs +++ b/funera_core/src/security/registry.rs @@ -152,6 +152,10 @@ impl GuardedToolRegistry { self.inner.remove_tool_if_same(name, tool) } + pub fn set_tool_availability(&mut self, name: &str, available: bool) -> bool { + self.inner.set_tool_availability(name, available) + } + pub fn tool_exists(&self, name: &str) -> bool { self.inner.tool_exists(name) } @@ -431,6 +435,26 @@ mod tests { assert!(!registry.tool_exists("ok_tool")); } + #[tokio::test] + async fn set_tool_availability_delegates_to_inner() { + let mut registry = GuardedToolRegistry::new(); + assert!(!registry.set_tool_availability("ok_tool", false)); + + registry.add_tool(Arc::new(OkTool)); + assert!(registry.set_tool_availability("ok_tool", false)); + assert!( + registry.get_tool_arc("ok_tool").is_none(), + "unavailable tool must not be runnable" + ); + let result = registry.call_tool("ok_tool", json!({})).await; + assert!(matches!(result, Err(ToolCallError::ToolUnavailable(_)))); + + assert!(!registry.set_tool_availability("ok_tool", false)); + assert!(registry.set_tool_availability("ok_tool", true)); + let result = registry.call_tool("ok_tool", json!({})).await; + assert_eq!(result.unwrap(), "done"); + } + #[tokio::test] async fn denied_tool_blocked() { let mut policy = ToolPolicy::default();