atd_runtime/
context.rs

1//! Per-call context passed to every `Tool::call` invocation.
2
3use std::path::PathBuf;
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use crate::capability::CapabilitySet;
8use crate::cursor::CursorIssuer;
9use crate::secrets::SecretBundle;
10use crate::tier::ToolTier;
11use crate::tracker::ReadTracker;
12
13#[non_exhaustive]
14pub struct CallContext {
15    /// Working directory for relative-path tools (Read / Bash / Glob / ...).
16    pub cwd: PathBuf,
17    /// Advisory truncation budget. Tools should respect this and return
18    /// truncation markers when producing larger output. In SP-12 this is
19    /// derived from the tool's tier via `TierPolicy::max_output`.
20    pub max_output_bytes: usize,
21    /// Unique id for tracing/logging; not emitted on the wire.
22    pub call_id: ulid::Ulid,
23    /// Absolute deadline. Tools that wrap long operations in tokio::time::timeout
24    /// should pass `remaining_time()` as the budget.
25    pub deadline: Option<Instant>,
26    /// Shared-per-connection read tracker. `None` in isolated unit tests;
27    /// server always attaches one via `Arc::clone` in per-connection state.
28    pub read_tracker: Option<Arc<ReadTracker>>,
29    /// Connection-scoped capability allow-list. Populated from the `Hello`
30    /// handshake; shared across all calls on a single connection. Tools that
31    /// gate side effects on caller authority may read this, though dispatch
32    /// already enforces `required_capabilities` before invocation.
33    pub capabilities: Arc<CapabilitySet>,
34    /// Tier the current call resolved to. Informational for tools; dispatch
35    /// uses the tier to pick the deadline / max_output budget above.
36    pub tier: ToolTier,
37    /// Optional per-call agent identity. Populated from the `Hello.client_id`
38    /// handshake; `None` for unit tests and in-process dispatch with no
39    /// client-id source. Audit events copy this field into `CallEvent.caller_id`.
40    pub caller_id: Option<String>,
41    /// Optional per-call secret bundle resolved by `ServerConfig::token_broker`
42    /// (SP-token-broker-phase1). `None` means either (a) no broker is
43    /// configured on this server, or (b) the broker returned `Ok(None)` for
44    /// this caller — both cases pass through to the tool's existing fallback
45    /// (env vars, saved file, etc.). Tools that need a secret access via
46    /// `ctx.secrets()`. Audit `CallEvent.secrets_resolved` is `true` iff
47    /// this is `Some(_)`.
48    pub secrets: Option<Arc<SecretBundle>>,
49    /// SP-pagination-v1 — cursor issuer for tools that override
50    /// `Tool::call_paginated`. `None` for non-paginating dispatch paths
51    /// (test fixtures, listeners that haven't wired pagination). Tools
52    /// use `ctx.cursor_issuer().issue(payload)` to mint a `next_cursor`.
53    pub cursor_issuer: Option<Arc<CursorIssuer>>,
54}
55
56impl CallContext {
57    /// Construct a `CallContext` from its required fields. This is the
58    /// stable constructor across crate boundaries — the type is
59    /// `#[non_exhaustive]`, so adding new fields in the future means
60    /// callers migrate through `CallContext::new(...)` (or a future
61    /// `CallContextBuilder`), not via struct-literal breakage.
62    #[allow(clippy::too_many_arguments)]
63    pub fn new(
64        cwd: PathBuf,
65        max_output_bytes: usize,
66        call_id: ulid::Ulid,
67        deadline: Option<Instant>,
68        read_tracker: Option<Arc<ReadTracker>>,
69        capabilities: Arc<CapabilitySet>,
70        tier: ToolTier,
71        caller_id: Option<String>,
72        secrets: Option<Arc<SecretBundle>>,
73    ) -> Self {
74        Self {
75            cwd,
76            max_output_bytes,
77            call_id,
78            deadline,
79            read_tracker,
80            capabilities,
81            tier,
82            caller_id,
83            secrets,
84            cursor_issuer: None,
85        }
86    }
87
88    /// Builder-style attach of the cursor issuer. Call after `new` from
89    /// dispatch code that wants to enable pagination for the call.
90    pub fn with_cursor_issuer(mut self, issuer: Arc<CursorIssuer>) -> Self {
91        self.cursor_issuer = Some(issuer);
92        self
93    }
94
95    /// Convenience accessor for paginating tools. Returns the issuer if
96    /// one was attached by dispatch; tools that override `call_paginated`
97    /// but find this `None` must fall back to returning `next_cursor: None`
98    /// (or error if cursors are mandatory).
99    pub fn cursor_issuer(&self) -> Option<&CursorIssuer> {
100        self.cursor_issuer.as_deref()
101    }
102
103    pub fn remaining_time(&self) -> Option<Duration> {
104        self.deadline
105            .map(|d| d.saturating_duration_since(Instant::now()))
106    }
107
108    /// Returns the resolved secret bundle, if any. `None` when no broker
109    /// is configured on the server, or when the broker returned `None`
110    /// for this caller.
111    pub fn secrets(&self) -> Option<&SecretBundle> {
112        self.secrets.as_deref()
113    }
114}
115
116#[cfg(any(test, feature = "testing"))]
117impl CallContext {
118    /// Construct a sensible default for unit tests. cwd = current dir,
119    /// 1 MiB output budget, fresh call_id, no deadline, no tracker,
120    /// empty capability set, Warm tier.
121    pub fn for_test() -> Self {
122        Self {
123            cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
124            max_output_bytes: 1_048_576,
125            call_id: ulid::Ulid::new(),
126            deadline: None,
127            read_tracker: None,
128            capabilities: Arc::new(CapabilitySet::empty()),
129            tier: ToolTier::Warm,
130            caller_id: None,
131            secrets: None,
132            cursor_issuer: None,
133        }
134    }
135
136    /// Test-only: build a CallContext with a fresh tracker attached.
137    /// Returns both so tests can also record/check on the same tracker.
138    pub fn for_test_with_tracker() -> (Self, Arc<ReadTracker>) {
139        let tracker = Arc::new(ReadTracker::new());
140        let ctx = Self {
141            cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
142            max_output_bytes: 1_048_576,
143            call_id: ulid::Ulid::new(),
144            deadline: None,
145            read_tracker: Some(tracker.clone()),
146            capabilities: Arc::new(CapabilitySet::empty()),
147            tier: ToolTier::Warm,
148            caller_id: None,
149            secrets: None,
150            cursor_issuer: None,
151        };
152        (ctx, tracker)
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    #[test]
161    fn for_test_has_sensible_defaults() {
162        let ctx = CallContext::for_test();
163        assert!(ctx.cwd.exists(), "cwd should be a real directory");
164        assert_eq!(ctx.max_output_bytes, 1_048_576);
165        assert!(ctx.deadline.is_none());
166        assert!(ctx.read_tracker.is_none());
167    }
168
169    #[test]
170    fn for_test_with_tracker_shares_arc() {
171        let (ctx, tracker) = CallContext::for_test_with_tracker();
172        assert!(ctx.read_tracker.is_some());
173        let ctx_tracker = ctx.read_tracker.as_ref().unwrap();
174        assert!(Arc::ptr_eq(ctx_tracker, &tracker));
175    }
176
177    #[test]
178    fn remaining_time_is_none_when_no_deadline() {
179        let ctx = CallContext::for_test();
180        assert!(ctx.remaining_time().is_none());
181    }
182
183    #[test]
184    fn remaining_time_counts_down_from_deadline() {
185        let ctx = CallContext {
186            cwd: PathBuf::from("."),
187            max_output_bytes: 1024,
188            call_id: ulid::Ulid::new(),
189            deadline: Some(Instant::now() + Duration::from_secs(5)),
190            read_tracker: None,
191            capabilities: Arc::new(CapabilitySet::empty()),
192            tier: ToolTier::Warm,
193            caller_id: None,
194            secrets: None,
195            cursor_issuer: None,
196        };
197        let r = ctx.remaining_time().unwrap();
198        assert!(r <= Duration::from_secs(5));
199        assert!(r > Duration::from_secs(4));
200    }
201
202    #[test]
203    fn remaining_time_saturates_to_zero_after_deadline() {
204        let ctx = CallContext {
205            cwd: PathBuf::from("."),
206            max_output_bytes: 1024,
207            call_id: ulid::Ulid::new(),
208            deadline: Some(Instant::now() - Duration::from_secs(10)),
209            read_tracker: None,
210            capabilities: Arc::new(CapabilitySet::empty()),
211            tier: ToolTier::Warm,
212            caller_id: None,
213            secrets: None,
214            cursor_issuer: None,
215        };
216        assert_eq!(ctx.remaining_time().unwrap(), Duration::ZERO);
217    }
218
219    #[test]
220    fn for_test_has_empty_capabilities_and_warm_tier() {
221        let ctx = CallContext::for_test();
222        assert!(ctx.capabilities.granted().is_empty());
223        assert_eq!(ctx.tier, ToolTier::Warm);
224    }
225}
atd_runtime/context.rs

atd_runtime/
context.rs
