Skip to main content

heartbit_core/
execution_context.rs

1//! Per-request execution context threaded through tool dispatch.
2//!
3//! Every `Tool::execute` call receives an `&ExecutionContext`. The context
4//! carries tenant/user identity, the workspace root, and resolvers for
5//! per-tenant secrets and audit sinks. It is constructed at the request
6//! boundary (CLI command, Restate workflow activity, daemon dispatch) and
7//! threaded through the agent runner unchanged.
8
9use std::future::Future;
10use std::path::PathBuf;
11use std::pin::Pin;
12use std::sync::Arc;
13
14use crate::error::Error;
15
16/// Per-request context carried into every tool invocation.
17#[derive(Clone, Default)]
18pub struct ExecutionContext {
19    /// Tenant identifier (multi-tenant deployments). `None` outside of multi-tenant flows.
20    pub tenant_id: Option<String>,
21    /// User identifier on whose behalf the agent runs. `None` outside of authenticated flows.
22    pub user_id: Option<String>,
23    /// Workspace root for filesystem-aware tools. `None` when no workspace is configured.
24    pub workspace: Option<PathBuf>,
25    /// Resolver for per-tenant secrets (API keys, OAuth tokens). `None` when no resolver is configured.
26    pub credentials: Option<Arc<dyn CredentialResolver>>,
27    /// Sink for tool-level audit records. `None` when no audit sink is configured.
28    pub audit_sink: Option<Arc<dyn AuditSink>>,
29    /// Snapshot of the conversation at tool-dispatch time. `None` outside an
30    /// agent run. Lets introspection tools (e.g. the advisor) see the full
31    /// history the way Claude Code forwards it to its advisor.
32    pub transcript: Option<Arc<Vec<crate::llm::types::Message>>>,
33}
34
35impl std::fmt::Debug for ExecutionContext {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("ExecutionContext")
38            .field("tenant_id", &self.tenant_id)
39            .field("user_id", &self.user_id)
40            .field("workspace", &self.workspace)
41            .field(
42                "credentials",
43                &self.credentials.as_ref().map(|_| "<resolver>"),
44            )
45            .field("audit_sink", &self.audit_sink.as_ref().map(|_| "<sink>"))
46            .finish()
47    }
48}
49
50/// Resolves a named secret (API key, token) for the current tenant.
51pub trait CredentialResolver: Send + Sync {
52    /// Resolve a secret by logical name (e.g. `"X_API_KEY"`).
53    fn resolve(
54        &self,
55        name: &str,
56    ) -> Pin<Box<dyn Future<Output = Result<Secret, Error>> + Send + '_>>;
57}
58
59/// Receives per-tool audit records emitted by tools that opt in.
60pub trait AuditSink: Send + Sync {
61    /// Record a structured audit entry. Implementations must not block.
62    fn record(
63        &self,
64        record: serde_json::Value,
65    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;
66}
67
68/// A secret value with redacted `Debug`/`Display` formatting.
69#[derive(Clone)]
70pub struct Secret(String);
71
72impl Secret {
73    /// Wrap a secret string. Use `expose()` to read the inner value.
74    pub fn new(value: impl Into<String>) -> Self {
75        Self(value.into())
76    }
77
78    /// Read the inner secret string. Caller is responsible for not logging the result.
79    pub fn expose(&self) -> &str {
80        &self.0
81    }
82}
83
84impl std::fmt::Debug for Secret {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        write!(f, "Secret(<redacted>)")
87    }
88}
89
90impl std::fmt::Display for Secret {
91    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92        write!(f, "<redacted>")
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn execution_context_default_has_no_identity() {
102        let ctx = ExecutionContext::default();
103        assert!(ctx.tenant_id.is_none());
104        assert!(ctx.user_id.is_none());
105        assert!(ctx.workspace.is_none());
106        assert!(ctx.credentials.is_none());
107        assert!(ctx.audit_sink.is_none());
108    }
109
110    #[test]
111    fn execution_context_clone_preserves_fields() {
112        let ctx = ExecutionContext {
113            tenant_id: Some("tenant-1".into()),
114            user_id: Some("user-2".into()),
115            workspace: Some(PathBuf::from("/tmp/ws")),
116            credentials: None,
117            audit_sink: None,
118            transcript: None,
119        };
120        let cloned = ctx.clone();
121        assert_eq!(cloned.tenant_id.as_deref(), Some("tenant-1"));
122        assert_eq!(cloned.user_id.as_deref(), Some("user-2"));
123        assert_eq!(cloned.workspace, Some(PathBuf::from("/tmp/ws")));
124    }
125
126    #[test]
127    fn secret_debug_redacts() {
128        let s = Secret::new("super-secret-token");
129        let debug = format!("{:?}", s);
130        assert!(!debug.contains("super-secret-token"));
131        assert!(debug.contains("<redacted>"));
132    }
133
134    #[test]
135    fn secret_display_redacts() {
136        let s = Secret::new("super-secret-token");
137        let display = format!("{}", s);
138        assert!(!display.contains("super-secret-token"));
139        assert!(display.contains("<redacted>"));
140    }
141
142    #[test]
143    fn secret_expose_returns_inner() {
144        let s = Secret::new("super-secret-token");
145        assert_eq!(s.expose(), "super-secret-token");
146    }
147
148    #[test]
149    fn execution_context_debug_does_not_leak_resolver_internals() {
150        struct DummyResolver;
151        impl CredentialResolver for DummyResolver {
152            fn resolve(
153                &self,
154                _name: &str,
155            ) -> Pin<Box<dyn Future<Output = Result<Secret, Error>> + Send + '_>> {
156                Box::pin(async { Ok(Secret::new("x")) })
157            }
158        }
159
160        let ctx = ExecutionContext {
161            credentials: Some(Arc::new(DummyResolver)),
162            ..ExecutionContext::default()
163        };
164        let debug = format!("{:?}", ctx);
165        assert!(debug.contains("<resolver>"));
166        assert!(!debug.contains("DummyResolver"));
167    }
168}