1use std::sync::Arc;
19
20use rx4::provider::{
21 Message, Provider as Rx4Provider, ProviderError as Rx4ProviderError, Role, StreamEvent,
22};
23
24use crate::agent::hooks::{run_post_hooks, run_pre_hooks, HookDecision, ToolHook};
25use crate::agent::stream::{emit, AgentStreamEvent, AgentStreamTx};
26use crate::cost::{ContextSnapshot, CostTracker, TokenUsage};
27use crate::plugin::{HookManager, LifecycleEvent, PluginRegistry};
28use crate::providers::{ChatMessage, ChatRequest, Provider as UnthinkclawProvider};
29use crate::tools::{Tool as UnthinkclawTool, ToolResult as UnthinkclawToolResult, ToolSpec};
30
31#[derive(Clone, Default)]
39pub struct ToolHookContext {
40 hooks: Vec<Arc<dyn ToolHook>>,
41 plugins: Option<Arc<tokio::sync::RwLock<PluginRegistry>>>,
42 hook_manager: Option<Arc<HookManager>>,
43 stream: Option<AgentStreamTx>,
44}
45
46impl ToolHookContext {
47 pub fn new(
48 hooks: Vec<Arc<dyn ToolHook>>,
49 plugins: Option<Arc<tokio::sync::RwLock<PluginRegistry>>>,
50 ) -> Self {
51 Self {
52 hooks,
53 plugins,
54 hook_manager: None,
55 stream: None,
56 }
57 }
58
59 pub fn with_hook_manager(mut self, hook_manager: Arc<HookManager>) -> Self {
62 self.hook_manager = Some(hook_manager);
63 self
64 }
65
66 pub fn with_stream(mut self, stream: Option<AgentStreamTx>) -> Self {
69 self.stream = stream;
70 self
71 }
72
73 async fn emit_lifecycle(&self, event: LifecycleEvent) {
74 if let Some(manager) = &self.hook_manager {
75 manager.emit(&event).await;
76 }
77 }
78
79 pub async fn check_pre_tool(&self, name: &str, arguments: &str) -> HookDecision {
81 if let Some(plugins) = &self.plugins {
82 let registry = plugins.read().await;
83 if let HookDecision::Block(reason) = registry.check_pre_tool(name, arguments).await {
84 return HookDecision::Block(format!("Blocked by plugin: {reason}"));
85 }
86 }
87 match run_pre_hooks(&self.hooks, name, arguments).await {
88 HookDecision::Block(reason) => {
89 HookDecision::Block(format!("Blocked by policy: {reason}"))
90 }
91 HookDecision::Allow => HookDecision::Allow,
92 }
93 }
94
95 pub async fn notify_post_tool(
97 &self,
98 name: &str,
99 arguments: &str,
100 result: &UnthinkclawToolResult,
101 ) {
102 run_post_hooks(&self.hooks, name, arguments, result).await;
103 self.emit_lifecycle(LifecycleEvent::AfterToolCall(
104 name.to_string(),
105 arguments.to_string(),
106 result.clone(),
107 ))
108 .await;
109 if let Some(plugins) = &self.plugins {
110 let registry = plugins.read().await;
111 registry.notify_post_tool(name, arguments, result).await;
112 }
113 }
114}
115
116pub async fn execute_tool_with_hooks(
122 ctx: &ToolHookContext,
123 name: &str,
124 arguments: &str,
125 tool: Option<&Arc<dyn UnthinkclawTool>>,
126) -> UnthinkclawToolResult {
127 ctx.emit_lifecycle(LifecycleEvent::BeforeToolCall(
128 name.to_string(),
129 arguments.to_string(),
130 ))
131 .await;
132 emit(
133 &ctx.stream,
134 AgentStreamEvent::ToolStart {
135 name: name.to_string(),
136 hint: crate::agent::loop_runner::extract_tool_hint(name, arguments),
137 },
138 );
139
140 let started = std::time::Instant::now();
141 let result = match ctx.check_pre_tool(name, arguments).await {
142 HookDecision::Block(reason) => {
143 tracing::info!("blocked '{}': {}", name, reason);
144 UnthinkclawToolResult::error(reason)
145 }
146 HookDecision::Allow => match tool {
147 Some(tool) => match tool.execute(arguments).await {
148 Ok(result) => result,
149 Err(e) => UnthinkclawToolResult::error(crate::redaction::redact_text(&format!(
150 "Tool error: {e}"
151 ))),
152 },
153 None => UnthinkclawToolResult::error(format!("Unknown tool: {name}")),
154 },
155 };
156
157 ctx.notify_post_tool(name, arguments, &result).await;
158
159 emit(
160 &ctx.stream,
161 AgentStreamEvent::ToolEnd {
162 name: name.to_string(),
163 ok: !result.is_error,
164 elapsed_secs: started.elapsed().as_secs(),
165 },
166 );
167
168 result
169}
170
171pub fn chat_message_to_rx4(msg: &ChatMessage) -> Message {
175 let role = match msg.role.as_str() {
176 "system" => Role::System,
177 "user" => Role::User,
178 "assistant" | "assistant_tool_use" => Role::Assistant,
179 "tool_result" => Role::Tool,
180 _ => Role::User,
181 };
182 Message {
183 role,
184 content: msg.content.clone(),
185 tool_call_id: msg.tool_use_id.clone(),
186 tool_calls: Vec::new(),
187 }
188}
189
190pub fn rx4_message_to_chat(msg: &Message) -> ChatMessage {
192 let role = match msg.role {
193 Role::System => "system",
194 Role::User => "user",
195 Role::Assistant => "assistant",
196 Role::Tool => "tool_result",
197 };
198 ChatMessage {
199 role: role.to_string(),
200 content: msg.content.clone(),
201 tool_use_id: msg.tool_call_id.clone(),
202 }
203}
204
205pub struct RotaryProviderAdapter {
215 inner: Arc<dyn UnthinkclawProvider>,
216 id: String,
217 name: String,
218 cost_tracker: Option<Arc<CostTracker>>,
219}
220
221impl RotaryProviderAdapter {
222 pub fn new(
223 provider: Arc<dyn UnthinkclawProvider>,
224 cost_tracker: Option<Arc<CostTracker>>,
225 ) -> Self {
226 let id = provider.name().to_string();
227 let name = format!("apollo-{}", provider.name());
228 Self {
229 inner: provider,
230 id,
231 name,
232 cost_tracker,
233 }
234 }
235}
236
237#[async_trait::async_trait]
238impl Rx4Provider for RotaryProviderAdapter {
239 fn id(&self) -> &str {
240 &self.id
241 }
242
243 fn name(&self) -> &str {
244 &self.name
245 }
246
247 async fn stream(
248 &self,
249 messages: &[Message],
250 system: &Option<String>,
251 model: &str,
252 tools: &[serde_json::Value],
253 _reasoning_effort: Option<&str>,
254 ) -> Result<rx4::provider::StreamResult, Rx4ProviderError> {
255 let mut chat_messages: Vec<ChatMessage> = Vec::new();
257
258 if let Some(sys) = system {
260 chat_messages.push(ChatMessage::system(sys));
261 }
262
263 for msg in messages {
264 chat_messages.push(rx4_message_to_chat(msg));
265 }
266
267 let tool_specs: Vec<ToolSpec> = tools
269 .iter()
270 .filter_map(|t| {
271 let name = t.get("name")?.as_str()?.to_string();
272 let description = t
273 .get("description")
274 .and_then(|d| d.as_str())
275 .unwrap_or("")
276 .to_string();
277 let parameters = t
278 .get("parameters")
279 .cloned()
280 .unwrap_or(serde_json::Value::Null);
281 Some(ToolSpec {
282 name,
283 description,
284 parameters,
285 })
286 })
287 .collect();
288
289 let tool_refs: &[ToolSpec] = if tool_specs.is_empty() {
290 &[]
291 } else {
292 &tool_specs
295 };
296
297 let request = ChatRequest {
298 messages: &chat_messages,
299 tools: if tool_refs.is_empty() {
300 None
301 } else {
302 Some(tool_refs)
303 },
304 model,
305 temperature: 0.7,
306 max_tokens: Some(8192),
307 };
308
309 if let Some(tracker) = &self.cost_tracker {
310 let system_chars = system
311 .as_ref()
312 .map(|value| value.chars().count())
313 .unwrap_or(0);
314 let history_chars = messages
315 .iter()
316 .map(|message| message.content.chars().count())
317 .sum::<usize>();
318 let tool_chars = tools
319 .iter()
320 .map(|tool| {
321 serde_json::to_string(tool)
322 .unwrap_or_default()
323 .chars()
324 .count()
325 })
326 .sum::<usize>();
327 tracker
328 .record_context(ContextSnapshot {
329 system_chars,
330 history_chars,
331 tool_chars,
332 estimated_input_tokens: (system_chars + history_chars + tool_chars).div_ceil(4),
333 })
334 .await;
335 }
336
337 let response = self
338 .inner
339 .chat(&request)
340 .await
341 .map_err(|e| Rx4ProviderError::Api(e.to_string()))?;
342
343 if let (Some(tracker), Some(usage)) = (&self.cost_tracker, response.usage.as_ref()) {
344 let _ = tracker
345 .record(
346 model,
347 TokenUsage {
348 input_tokens: usage.input_tokens as usize,
349 output_tokens: usage.output_tokens as usize,
350 total_tokens: usage.input_tokens as usize + usage.output_tokens as usize,
351 },
352 )
353 .await;
354 }
355
356 let text = response.text.unwrap_or_default();
358 let tool_calls = response.tool_calls;
359
360 let events: Vec<Result<StreamEvent, Rx4ProviderError>> = {
362 let mut evs = Vec::new();
363 if !text.is_empty() {
364 evs.push(Ok(StreamEvent::Delta(text)));
365 }
366 for tc in tool_calls {
367 evs.push(Ok(StreamEvent::ToolCall(rx4::ToolCall {
368 id: tc.id,
369 name: tc.name,
370 arguments: tc.arguments,
371 })));
372 }
373 evs.push(Ok(StreamEvent::Done));
374 evs
375 };
376
377 use futures_util::stream;
379 Ok(Box::new(Box::pin(stream::iter(events))))
380 }
381}
382
383pub fn register_apollo_tools(
395 registry: &mut rx4::ToolRegistry,
396 tools: &[Arc<dyn UnthinkclawTool>],
397 hook_ctx: &ToolHookContext,
398) {
399 use rx4::guardrails::classify_tool;
400 use rx4::{ToolDefinition, ToolEffect, ToolExecuteBox};
401
402 for tool in tools {
403 let spec = tool.spec();
404 let name = spec.name.clone();
405 let description = spec.description.clone();
406 let parameters_json = serde_json::to_string(&spec.parameters).unwrap_or_default();
407
408 let tool_clone = Arc::clone(tool);
409 let hook_ctx = hook_ctx.clone();
410 let tool_name = name.clone();
411 let execute: ToolExecuteBox = Box::new(move |_ctx, args| {
412 let tool = Arc::clone(&tool_clone);
413 let hook_ctx = hook_ctx.clone();
414 let tool_name = tool_name.clone();
415 Box::pin(async move {
416 let result =
417 execute_tool_with_hooks(&hook_ctx, &tool_name, &args, Some(&tool)).await;
418
419 rx4::ToolResult {
420 id: String::new(),
421 content: result.output,
422 is_error: result.is_error,
423 error_kind: None,
424 }
425 })
426 });
427
428 let effect = match classify_tool(&name) {
429 rx4::guardrails::ToolClass::Idempotent => ToolEffect::Read,
430 rx4::guardrails::ToolClass::Mutating => ToolEffect::Write,
431 };
432
433 registry.register(
434 ToolDefinition::new_boxed(name, description, parameters_json, execute)
435 .with_effect(effect),
436 );
437 }
438}
439
440pub struct RotaryBridgeConfig {
444 pub provider: Arc<dyn UnthinkclawProvider>,
445 pub tools: Vec<Arc<dyn UnthinkclawTool>>,
446 pub system_prompt: String,
447 pub model: String,
448 pub workspace: std::path::PathBuf,
449 pub max_tool_iterations: usize,
450 pub auto_compact_after: usize,
453 pub cost_tracker: Option<Arc<CostTracker>>,
455 pub hook_ctx: ToolHookContext,
458}
459
460fn model_registry_for(provider: &dyn UnthinkclawProvider, model: &str) -> rx4::ModelRegistry {
461 let mut registry = rx4::ModelRegistry::new();
462 let capabilities = provider.capabilities();
463 let mut info = rx4::ModelInfo::new(
464 provider.name(),
465 model,
466 capabilities.max_context.max(128_000) as usize,
467 8_192,
468 );
469 info.supports_tools = capabilities.native_tools;
470 info.supports_vision = capabilities.vision;
471 registry.register(info);
472 registry
473}
474
475pub struct RotaryAgentBridge {
487 agent: rx4::Agent,
488 hook_ctx: ToolHookContext,
489 messages: Vec<Message>,
491}
492
493impl RotaryAgentBridge {
494 pub fn new(config: RotaryBridgeConfig) -> Self {
496 Self::new_with_model_registry(config, rx4::ModelRegistry::new())
497 }
498
499 pub fn new_with_model_registry(
502 config: RotaryBridgeConfig,
503 model_registry: rx4::ModelRegistry,
504 ) -> Self {
505 let rx4_provider = Arc::new(RotaryProviderAdapter::new(
506 Arc::clone(&config.provider),
507 config.cost_tracker,
508 ));
509
510 let mut agent = rx4::Agent::new();
511 let model_registry = if model_registry.is_empty() {
512 model_registry_for(config.provider.as_ref(), &config.model)
513 } else {
514 model_registry
515 };
516 agent.set_model_registry(model_registry);
517 agent.set_model(&config.model);
518 agent.set_system_prompt(&config.system_prompt);
519 agent.set_provider(rx4_provider);
520 agent.set_workspace_root(&config.workspace);
521 agent.max_tool_iterations = config.max_tool_iterations;
522 agent.auto_compact_after = config.auto_compact_after;
526
527 agent.set_policy(rx4::Policy::full_access());
542
543 let mut tool_registry = rx4::ToolRegistry::new();
545 register_apollo_tools(&mut tool_registry, &config.tools, &config.hook_ctx);
546 agent.tools = Arc::new(tool_registry);
547
548 Self {
549 agent,
550 hook_ctx: config.hook_ctx,
551 messages: Vec::new(),
552 }
553 }
554
555 pub fn agent(&self) -> &rx4::Agent {
557 &self.agent
558 }
559
560 pub fn agent_mut(&mut self) -> &mut rx4::Agent {
562 &mut self.agent
563 }
564
565 pub fn clear_messages(&mut self) {
567 self.messages.clear();
568 self.agent.clear_messages();
569 }
570
571 pub fn message_count(&self) -> usize {
573 self.messages.len()
574 }
575
576 pub fn set_model(&mut self, model: &str) {
578 self.agent.set_model(model);
579 }
580
581 pub fn set_system_prompt(&mut self, prompt: &str) {
583 self.agent.set_system_prompt(prompt);
584 }
585
586 pub fn set_workspace_root(&mut self, path: &std::path::Path) {
588 self.agent.set_workspace_root(path);
589 }
590
591 pub fn set_scope(&mut self, scope: rx4::Scope) {
593 self.agent.set_scope(scope);
594 }
595
596 pub fn subscribe(&mut self, callback: impl Fn(&rx4::Event) + Send + Sync + 'static) {
598 self.agent.subscribe(callback);
599 }
600
601 pub async fn run_prompt(&mut self, prompt: &str) -> anyhow::Result<String> {
612 let last_response = Arc::new(parking_lot::RwLock::new(String::new()));
614 let last_response_clone = Arc::clone(&last_response);
615
616 self.agent.subscribe(move |event| {
617 if let rx4::Event::MessageEnd {
618 content,
619 role: Role::Assistant,
620 } = event
621 {
622 *last_response_clone.write() = content.clone();
623 }
624 });
625
626 self.agent.prompt(prompt).await?;
627
628 let response = last_response.read().clone();
629 Ok(response)
630 }
631
632 pub async fn run_prompt_with_history(
638 &mut self,
639 prompt: &str,
640 history: &[ChatMessage],
641 ) -> anyhow::Result<String> {
642 self.agent.clear_messages();
644 for msg in history {
645 let rx4_msg = chat_message_to_rx4(msg);
646 self.agent.messages.write().push(rx4_msg);
649 }
650
651 self.run_prompt(prompt).await
652 }
653
654 pub fn register_tools(&mut self, tools: &[Arc<dyn UnthinkclawTool>]) {
656 if let Some(registry) = Arc::get_mut(&mut self.agent.tools) {
657 register_apollo_tools(registry, tools, &self.hook_ctx);
658 } else {
659 tracing::warn!("cannot register rx4 tools while the registry is shared");
660 }
661 }
662
663 pub fn list_tools(&self) -> Vec<String> {
665 self.agent
666 .tools
667 .definitions()
668 .iter()
669 .filter_map(|d| {
670 d.get("name")
671 .and_then(|n| n.as_str())
672 .map(|s| s.to_string())
673 })
674 .collect()
675 }
676
677 pub fn compact(&mut self, reason: &str) {
679 self.agent.compact(reason);
680 }
681
682 pub fn messages_handle(&self) -> Arc<parking_lot::RwLock<Vec<Message>>> {
690 self.agent.messages_handle()
691 }
692
693 pub fn enable_skill_engine(&mut self, workspace: &std::path::Path) {
700 let mut engine = build_rx4_skill_engine(workspace);
701 if let Err(error) = engine.load() {
702 tracing::warn!("rx4 skill engine load failed, leaving it unset: {error}");
703 return;
704 }
705 self.agent.set_skill_engine(engine);
706 }
707
708 pub fn enable_graph_memory(&mut self, workspace: &std::path::Path, auto_dream: bool) {
714 self.agent
715 .set_graph_memory(rx4::GraphMemory::from_workspace(workspace));
716 self.agent.enable_auto_dream(auto_dream);
717 }
718}
719
720pub fn build_rx4_skill_engine(workspace: &std::path::Path) -> rx4::SkillEngine {
742 let home = dirs::home_dir().unwrap_or_default();
743
744 let managed_dir = workspace.join(".apollo/skills");
746
747 let mut engine = rx4::SkillEngine::new(managed_dir);
748
749 let openclaw_skills = home.join(".npm-global/lib/node_modules/openclaw/skills");
751 engine.add_extra_dir(openclaw_skills);
752
753 let shared_skills = home.join(".openclaw/workspace/skills");
754 engine.add_extra_dir(shared_skills);
755
756 engine
757}
758
759#[cfg(test)]
760mod tests {
761 use super::*;
762
763 #[test]
764 fn test_chat_message_to_rx4_system() {
765 let msg = ChatMessage::system("hello");
766 let rx4_msg = chat_message_to_rx4(&msg);
767 assert_eq!(rx4_msg.role, Role::System);
768 assert_eq!(rx4_msg.content, "hello");
769 }
770
771 #[test]
772 fn test_chat_message_to_rx4_user() {
773 let msg = ChatMessage::user("test");
774 let rx4_msg = chat_message_to_rx4(&msg);
775 assert_eq!(rx4_msg.role, Role::User);
776 assert_eq!(rx4_msg.content, "test");
777 }
778
779 #[test]
780 fn test_chat_message_to_rx4_tool_result() {
781 let msg = ChatMessage::tool_result("tc_123", "result text");
782 let rx4_msg = chat_message_to_rx4(&msg);
783 assert_eq!(rx4_msg.role, Role::Tool);
784 assert_eq!(rx4_msg.content, "result text");
785 assert_eq!(rx4_msg.tool_call_id.as_deref(), Some("tc_123"));
786 }
787
788 #[test]
789 fn test_rx4_message_to_chat() {
790 let msg = Message::assistant("hello back");
791 let chat_msg = rx4_message_to_chat(&msg);
792 assert_eq!(chat_msg.role, "assistant");
793 assert_eq!(chat_msg.content, "hello back");
794 }
795
796 #[test]
797 fn test_roundtrip_translation() {
798 let original = ChatMessage::user("roundtrip test");
799 let rx4_msg = chat_message_to_rx4(&original);
800 let back = rx4_message_to_chat(&rx4_msg);
801 assert_eq!(back.role, "user");
802 assert_eq!(back.content, "roundtrip test");
803 }
804
805 #[test]
806 fn test_build_rx4_skill_engine() {
807 let tmp = tempfile::tempdir().unwrap();
809 let engine = build_rx4_skill_engine(tmp.path());
810 assert!(
811 engine.skills_dir().exists()
812 || engine.skills_dir() == tmp.path().join(".apollo/skills")
813 );
814 }
815
816 struct RecordingTool {
817 ran: Arc<std::sync::atomic::AtomicBool>,
818 }
819
820 #[async_trait::async_trait]
821 impl UnthinkclawTool for RecordingTool {
822 fn name(&self) -> &str {
823 "exec"
824 }
825
826 fn spec(&self) -> ToolSpec {
827 ToolSpec {
828 name: "exec".to_string(),
829 description: "test tool".to_string(),
830 parameters: serde_json::json!({"type": "object"}),
831 }
832 }
833
834 async fn execute(&self, _arguments: &str) -> anyhow::Result<UnthinkclawToolResult> {
835 self.ran.store(true, std::sync::atomic::Ordering::SeqCst);
836 Ok(UnthinkclawToolResult::success("ran"))
837 }
838 }
839
840 async fn run_exec_through_rx4(hook_ctx: ToolHookContext) -> (rx4::ToolResult, bool) {
841 let ran = Arc::new(std::sync::atomic::AtomicBool::new(false));
842 let tool: Arc<dyn UnthinkclawTool> = Arc::new(RecordingTool {
843 ran: Arc::clone(&ran),
844 });
845 let mut registry = rx4::ToolRegistry::new();
846 register_apollo_tools(&mut registry, &[tool], &hook_ctx);
847
848 let ctx = Arc::new(rx4::ToolContext::new("."));
849 let result = registry
850 .execute("exec", &ctx, r#"{"command":"rm -rf /"}"#)
851 .await
852 .expect("tool registered");
853 (result, ran.load(std::sync::atomic::Ordering::SeqCst))
854 }
855
856 #[tokio::test]
857 async fn rx4_bridge_enforces_blocking_hooks() {
858 let hook: Arc<dyn ToolHook> = Arc::new(crate::agent::hooks::PermissionHook::new(
859 vec!["exec".to_string()],
860 vec![],
861 ));
862 let (result, ran) = run_exec_through_rx4(ToolHookContext::new(vec![hook], None)).await;
863 assert!(result.is_error, "blocked tool must report an error");
864 assert!(
865 result.content.contains("Blocked by policy"),
866 "unexpected content: {}",
867 result.content
868 );
869 assert!(!ran, "a blocked tool must not execute under rx4");
870 }
871
872 #[tokio::test]
873 async fn rx4_bridge_allows_unblocked_tools() {
874 let (result, ran) = run_exec_through_rx4(ToolHookContext::default()).await;
875 assert!(!result.is_error);
876 assert_eq!(result.content, "ran");
877 assert!(ran);
878 }
879
880 #[tokio::test]
881 async fn rx4_bridge_enforces_plugin_pre_tool_block() {
882 let mut registry = PluginRegistry::new();
883 registry.register_pre_tool_hook(Arc::new(BlockingPluginHook));
884 let ctx = ToolHookContext::new(
885 Vec::new(),
886 Some(Arc::new(tokio::sync::RwLock::new(registry))),
887 );
888 let (result, ran) = run_exec_through_rx4(ctx).await;
889 assert!(result.is_error);
890 assert!(
891 result.content.contains("Blocked by plugin"),
892 "unexpected content: {}",
893 result.content
894 );
895 assert!(!ran);
896 }
897
898 struct BlockingPluginHook;
899
900 #[async_trait::async_trait]
901 impl crate::plugin::PreToolHook for BlockingPluginHook {
902 fn name(&self) -> &str {
903 "blocking-test-hook"
904 }
905
906 async fn before_tool_call(&self, _name: &str, _arguments: &str) -> HookDecision {
907 HookDecision::Block("plugin says no".to_string())
908 }
909 }
910
911 struct RecordingLifecycleHook {
913 seen: Arc<std::sync::Mutex<Vec<String>>>,
914 }
915
916 #[async_trait::async_trait]
917 impl crate::plugin::LifecycleHook for RecordingLifecycleHook {
918 fn name(&self) -> &str {
919 "recording-lifecycle"
920 }
921
922 async fn on_event(&self, event: &LifecycleEvent) -> anyhow::Result<()> {
923 let label = match event {
924 LifecycleEvent::BeforeToolCall(name, _) => format!("before:{name}"),
925 LifecycleEvent::AfterToolCall(name, _, _) => format!("after:{name}"),
926 other => format!("other:{other:?}"),
927 };
928 self.seen.lock().unwrap().push(label);
929 Ok(())
930 }
931 }
932
933 fn stream_labels(
934 rx: &mut tokio::sync::mpsc::UnboundedReceiver<AgentStreamEvent>,
935 ) -> Vec<String> {
936 let mut labels = Vec::new();
937 while let Ok(event) = rx.try_recv() {
938 labels.push(match event {
939 AgentStreamEvent::ToolStart { name, .. } => format!("tool_start:{name}"),
940 AgentStreamEvent::ToolEnd { name, ok, .. } => format!("tool_end:{name}:{ok}"),
941 other => format!("other:{other:?}"),
942 });
943 }
944 labels
945 }
946
947 fn recording_context() -> (
949 ToolHookContext,
950 Arc<std::sync::Mutex<Vec<String>>>,
951 tokio::sync::mpsc::UnboundedReceiver<AgentStreamEvent>,
952 ) {
953 let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
954 let mut manager = HookManager::new();
955 manager.register_lifecycle(Arc::new(RecordingLifecycleHook {
956 seen: Arc::clone(&seen),
957 }));
958 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
959 let ctx = ToolHookContext::default()
960 .with_hook_manager(Arc::new(manager))
961 .with_stream(Some(tx));
962 (ctx, seen, rx)
963 }
964
965 #[tokio::test]
971 async fn both_engines_emit_the_same_hooks_and_events() {
972 let args = r#"{"command":"ls"}"#;
973
974 let (ctx, rx4_seen, mut rx4_stream) = recording_context();
976 let ran = Arc::new(std::sync::atomic::AtomicBool::new(false));
977 let tool: Arc<dyn UnthinkclawTool> = Arc::new(RecordingTool {
978 ran: Arc::clone(&ran),
979 });
980 let mut registry = rx4::ToolRegistry::new();
981 register_apollo_tools(&mut registry, &[Arc::clone(&tool)], &ctx);
982 let tool_ctx = Arc::new(rx4::ToolContext::new("."));
983 registry
984 .execute("exec", &tool_ctx, args)
985 .await
986 .expect("tool registered");
987 let rx4_events = rx4_seen.lock().unwrap().clone();
988 let rx4_stream_events = stream_labels(&mut rx4_stream);
989
990 let (ctx, legacy_seen, mut legacy_stream) = recording_context();
992 execute_tool_with_hooks(&ctx, "exec", args, Some(&tool)).await;
993 let legacy_events = legacy_seen.lock().unwrap().clone();
994 let legacy_stream_events = stream_labels(&mut legacy_stream);
995
996 assert_eq!(
997 rx4_events, legacy_events,
998 "the paths disagree on lifecycle hooks"
999 );
1000 assert_eq!(
1001 rx4_stream_events, legacy_stream_events,
1002 "the paths disagree on stream events"
1003 );
1004 assert_eq!(legacy_events, vec!["before:exec", "after:exec"]);
1005 assert_eq!(
1006 legacy_stream_events,
1007 vec!["tool_start:exec", "tool_end:exec:true"]
1008 );
1009 }
1010
1011 #[tokio::test]
1012 async fn a_blocked_tool_still_reports_start_and_end() {
1013 let hook: Arc<dyn ToolHook> = Arc::new(crate::agent::hooks::PermissionHook::new(
1014 vec!["exec".to_string()],
1015 vec![],
1016 ));
1017 let (ctx, seen, mut stream) = recording_context();
1018 let ctx = ToolHookContext::new(vec![hook], None)
1019 .with_hook_manager(Arc::new({
1020 let mut manager = HookManager::new();
1021 manager.register_lifecycle(Arc::new(RecordingLifecycleHook {
1022 seen: Arc::clone(&seen),
1023 }));
1024 manager
1025 }))
1026 .with_stream(ctx.stream.clone());
1027 let ran = Arc::new(std::sync::atomic::AtomicBool::new(false));
1028 let tool: Arc<dyn UnthinkclawTool> = Arc::new(RecordingTool {
1029 ran: Arc::clone(&ran),
1030 });
1031 let result = execute_tool_with_hooks(&ctx, "exec", "{}", Some(&tool)).await;
1032 assert!(result.is_error);
1033 assert!(!ran.load(std::sync::atomic::Ordering::SeqCst));
1034 assert_eq!(
1035 stream_labels(&mut stream),
1036 vec!["tool_start:exec", "tool_end:exec:false"],
1037 "a blocked call must still open and close its progress line"
1038 );
1039 assert_eq!(
1040 seen.lock().unwrap().clone(),
1041 vec!["before:exec", "after:exec"]
1042 );
1043 }
1044}