Skip to main content

oharness_tools/
context.rs

1//! `ToolContext` and `Workspace` (§7.2).
2
3use oharness_core::{
4    ApprovalChannel, BudgetHandle, Cancellation, EventSink, MetadataMap, NullApprovalChannel,
5    NullBudget, NullSink,
6};
7use std::path::{Path, PathBuf};
8use std::sync::Arc;
9
10/// Threaded through every `ToolSet::execute` call. Carries cross-cutting concerns —
11/// event emission, budget, cancellation, approval, workspace scoping.
12pub struct ToolContext {
13    pub events: Arc<dyn EventSink>,
14    pub budget: Arc<dyn BudgetHandle>,
15    pub cancellation: Cancellation,
16    pub approval: Arc<dyn ApprovalChannel>,
17    pub workspace: Option<Arc<Workspace>>,
18    /// Reverse-DNS-namespaced extras — like request extensions.
19    pub extensions: MetadataMap,
20}
21
22impl ToolContext {
23    /// Null-everything context, useful for tests and M1a smoke paths.
24    pub fn null() -> Self {
25        Self {
26            events: Arc::new(NullSink),
27            budget: Arc::new(NullBudget),
28            cancellation: Cancellation::new(),
29            approval: Arc::new(NullApprovalChannel),
30            workspace: None,
31            extensions: MetadataMap::new(),
32        }
33    }
34
35    pub fn workspace_path(&self) -> Option<&Path> {
36        self.workspace.as_ref().map(|w| w.path.as_path())
37    }
38}
39
40/// A scratch directory + optional cleanup strategy. Benchmarks use this to hand
41/// per-task directories to tools.
42pub struct Workspace {
43    pub path: PathBuf,
44    /// Cleanup policy; run on `teardown()` or on `Drop` (sync cleanups only run
45    /// inline; async cleanups on Drop are best-effort).
46    cleanup: std::sync::Mutex<Option<WorkspaceCleanup>>,
47}
48
49impl Workspace {
50    pub fn new(path: PathBuf) -> Self {
51        Self {
52            path,
53            cleanup: std::sync::Mutex::new(None),
54        }
55    }
56
57    pub fn with_sync_cleanup(mut self, f: impl FnOnce() + Send + 'static) -> Self {
58        self.cleanup = std::sync::Mutex::new(Some(WorkspaceCleanup::Sync(Box::new(f))));
59        self
60    }
61
62    /// Drop the workspace, running its cleanup.
63    pub async fn teardown(self) {
64        let cleanup = {
65            let mut guard = self.cleanup.lock().unwrap();
66            guard.take()
67        };
68        if let Some(cleanup) = cleanup {
69            match cleanup {
70                WorkspaceCleanup::Sync(f) => f(),
71                WorkspaceCleanup::Async(fut) => fut.await,
72            }
73        }
74    }
75}
76
77impl std::fmt::Debug for Workspace {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        f.debug_struct("Workspace")
80            .field("path", &self.path)
81            .finish()
82    }
83}
84
85impl Drop for Workspace {
86    fn drop(&mut self) {
87        if let Ok(mut guard) = self.cleanup.lock() {
88            if let Some(cleanup) = guard.take() {
89                match cleanup {
90                    WorkspaceCleanup::Sync(f) => f(),
91                    // On drop, async cleanups are best-effort. If there's a runtime,
92                    // we spawn; otherwise we silently drop the work. Benchmarks that
93                    // require guaranteed async cleanup call `teardown().await`.
94                    WorkspaceCleanup::Async(fut) => {
95                        if let Ok(handle) = tokio::runtime::Handle::try_current() {
96                            handle.spawn(fut);
97                        }
98                    }
99                }
100            }
101        }
102    }
103}
104
105pub enum WorkspaceCleanup {
106    Sync(Box<dyn FnOnce() + Send>),
107    Async(futures::future::BoxFuture<'static, ()>),
108}