llm_pipeline/
exec_ctx.rs

1//! Execution context shared across payload invocations.
2//!
3//! [`ExecCtx`] carries the HTTP client, LLM backend, endpoint, template variables,
4//! cancellation handle, and optional event handler. It is designed to be
5//! constructed once and shared across all payloads in a chain or graph.
6
7#[cfg(feature = "openai")]
8use crate::backend::OpenAiBackend;
9use crate::backend::{Backend, BackoffConfig, OllamaBackend};
10use crate::events::EventHandler;
11use crate::limits::PipelineLimits;
12#[allow(deprecated)]
13use crate::trace::TraceId;
14use reqwest::Client;
15use stack_ids::TraceCtx;
16use std::collections::HashMap;
17use std::sync::{
18    atomic::{AtomicBool, Ordering},
19    Arc,
20};
21use std::time::Duration;
22
23/// Shared execution context for payload invocations.
24///
25/// Carries everything a payload needs from the runtime environment
26/// without coupling to any specific orchestrator (Pipeline, LangGraph, etc.).
27///
28/// # Example
29///
30/// ```
31/// use llm_pipeline::ExecCtx;
32///
33/// let ctx = ExecCtx::builder("http://localhost:11434")
34///     .var("domain", "science")
35///     .var("audience", "researchers")
36///     .build();
37/// ```
38#[allow(deprecated)]
39pub struct ExecCtx {
40    /// HTTP client (cheap to clone -- uses `Arc` internally).
41    pub client: Client,
42    /// Base URL for the LLM provider (e.g. `http://localhost:11434`).
43    pub base_url: String,
44    /// LLM backend. Default: [`OllamaBackend`].
45    pub backend: Arc<dyn Backend>,
46    /// Transport retry configuration. Default: [`BackoffConfig::none()`].
47    pub backoff: BackoffConfig,
48    /// Template variables substituted into prompt `{key}` placeholders.
49    pub vars: HashMap<String, String>,
50    /// Optional cancellation flag; payloads should check before starting.
51    pub cancellation: Option<Arc<AtomicBool>>,
52    /// Optional event handler for streaming tokens and lifecycle events.
53    pub event_handler: Option<Arc<dyn EventHandler>>,
54    /// Phase status: compatibility / migration-only
55    ///
56    /// Legacy trace ID for correlating this context's operations across crates.
57    /// Auto-generated if not provided. Use [`trace_ctx`](Self::trace_ctx) for
58    /// the canonical trace form on the normal path.
59    ///
60    /// **Removal condition**: removed when all callers migrate to `TraceCtx`.
61    #[deprecated(
62        note = "Use trace_ctx instead. Will be removed when all callers migrate to TraceCtx."
63    )]
64    pub trace_id: TraceId,
65    /// Canonical trace context from `stack_ids`.
66    ///
67    /// This is the normal-path trace form. It supports parent span tracking,
68    /// bounded baggage, and W3C traceparent serialization. Automatically
69    /// generated from `trace_id` if not explicitly set.
70    pub trace_ctx: TraceCtx,
71    /// Resource limits for pipeline operations.
72    pub limits: PipelineLimits,
73}
74
75#[allow(deprecated)]
76impl ExecCtx {
77    /// Create a new builder.
78    pub fn builder(base_url: impl Into<String>) -> ExecCtxBuilder {
79        ExecCtxBuilder {
80            client: None,
81            base_url: base_url.into(),
82            backend: None,
83            backoff: None,
84            vars: HashMap::new(),
85            cancellation: None,
86            event_handler: None,
87            timeout: None,
88            trace_id: None,
89            trace_ctx: None,
90            limits: None,
91        }
92    }
93
94    /// Check whether cancellation has been requested.
95    pub fn is_cancelled(&self) -> bool {
96        self.cancellation
97            .as_ref()
98            .is_some_and(|c| c.load(Ordering::Relaxed))
99    }
100
101    /// Return an error if cancellation has been requested.
102    pub fn check_cancelled(&self) -> crate::error::Result<()> {
103        if self.is_cancelled() {
104            return Err(crate::PipelineError::Cancelled);
105        }
106        Ok(())
107    }
108
109    /// Get a reference to the cancellation AtomicBool, if set.
110    pub fn cancel_flag(&self) -> Option<&AtomicBool> {
111        self.cancellation.as_deref()
112    }
113}
114
115#[allow(deprecated)]
116impl std::fmt::Debug for ExecCtx {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        f.debug_struct("ExecCtx")
119            .field("base_url", &self.base_url)
120            .field("backend", &self.backend.name())
121            .field("backoff", &self.backoff)
122            .field("vars_count", &self.vars.len())
123            .field("has_cancellation", &self.cancellation.is_some())
124            .field("has_event_handler", &self.event_handler.is_some())
125            .field("trace_id", &self.trace_id)
126            .field("trace_ctx", &self.trace_ctx)
127            .field("limits", &self.limits)
128            .finish()
129    }
130}
131
132/// Builder for [`ExecCtx`].
133pub struct ExecCtxBuilder {
134    client: Option<Client>,
135    base_url: String,
136    backend: Option<Arc<dyn Backend>>,
137    backoff: Option<BackoffConfig>,
138    vars: HashMap<String, String>,
139    cancellation: Option<Arc<AtomicBool>>,
140    event_handler: Option<Arc<dyn EventHandler>>,
141    timeout: Option<Duration>,
142    trace_id: Option<TraceId>,
143    trace_ctx: Option<TraceCtx>,
144    limits: Option<PipelineLimits>,
145}
146
147#[allow(deprecated)]
148impl ExecCtxBuilder {
149    /// Set the HTTP client. If not set, a default client is created.
150    pub fn client(mut self, client: Client) -> Self {
151        self.client = Some(client);
152        self
153    }
154
155    /// Set the LLM backend. Default: [`OllamaBackend`].
156    pub fn backend(mut self, backend: Arc<dyn Backend>) -> Self {
157        self.backend = Some(backend);
158        self
159    }
160
161    /// Use the OpenAI-compatible backend without authentication.
162    ///
163    /// Sets the backend to [`OpenAiBackend`] with no API key. If the provider
164    /// requires authentication, use [`openai_with_key`](Self::openai_with_key) instead.
165    #[cfg(feature = "openai")]
166    pub fn openai(mut self) -> Self {
167        self.backend = Some(Arc::new(OpenAiBackend::new()));
168        self
169    }
170
171    /// Use the OpenAI-compatible backend with API key authentication.
172    ///
173    /// Sets the backend to [`OpenAiBackend`] with the given API key sent as
174    /// `Authorization: Bearer {key}`.
175    #[cfg(feature = "openai")]
176    pub fn openai_with_key(mut self, api_key: impl Into<String>) -> Self {
177        self.backend = Some(Arc::new(OpenAiBackend::new().with_api_key(api_key)));
178        self
179    }
180
181    /// Set the transport retry configuration. Default: [`BackoffConfig::none()`].
182    pub fn backoff(mut self, config: BackoffConfig) -> Self {
183        self.backoff = Some(config);
184        self
185    }
186
187    /// Set all template variables at once.
188    pub fn vars(mut self, vars: HashMap<String, String>) -> Self {
189        self.vars = vars;
190        self
191    }
192
193    /// Insert a single template variable.
194    pub fn var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
195        self.vars.insert(key.into(), value.into());
196        self
197    }
198
199    /// Set the cancellation flag.
200    pub fn cancellation(mut self, cancel: Option<Arc<AtomicBool>>) -> Self {
201        self.cancellation = cancel;
202        self
203    }
204
205    /// Set the event handler.
206    pub fn event_handler(mut self, handler: Arc<dyn EventHandler>) -> Self {
207        self.event_handler = Some(handler);
208        self
209    }
210
211    /// Set the client-level connection timeout. Default: 5 minutes (safety net).
212    ///
213    /// This is a coarse safety net on the `reqwest::Client`. Per-request timeouts
214    /// are applied individually via [`PipelineLimits::request_timeout`] or
215    /// [`LlmCall::with_timeout`](crate::llm_call::LlmCall::with_timeout).
216    ///
217    /// If a custom `Client` is provided via `.client()`, this setting is ignored
218    /// (the custom client's own timeout applies).
219    pub fn timeout(mut self, timeout: Duration) -> Self {
220        self.timeout = Some(timeout);
221        self
222    }
223
224    /// Phase status: compatibility / migration-only
225    ///
226    /// Set the legacy trace ID for correlating operations. If not set, a random
227    /// UUID v4 is generated automatically.
228    ///
229    /// Prefer [`with_trace_ctx`](Self::with_trace_ctx) for new code.
230    ///
231    /// **Removal condition**: removed when all callers migrate to `TraceCtx`.
232    #[deprecated(
233        since = "0.6.0",
234        note = "Use with_trace_ctx() instead. This method will be removed in v1.0."
235    )]
236    pub fn with_trace_id(mut self, trace_id: TraceId) -> Self {
237        self.trace_id = Some(trace_id);
238        self
239    }
240
241    /// Set the canonical trace context from `stack_ids::TraceCtx`.
242    ///
243    /// This is the normal-path trace form. When set, the legacy `trace_id` field
244    /// is derived from `trace_ctx.trace_id` for backward compatibility.
245    pub fn with_trace_ctx(mut self, trace_ctx: TraceCtx) -> Self {
246        self.trace_ctx = Some(trace_ctx);
247        self
248    }
249
250    /// Set resource limits. If not set, [`PipelineLimits::default()`] is used.
251    pub fn with_limits(mut self, limits: PipelineLimits) -> Self {
252        self.limits = Some(limits);
253        self
254    }
255
256    /// Build the execution context.
257    ///
258    /// **Preferred**: use [`with_trace_ctx()`](Self::with_trace_ctx) to set trace identity.
259    /// The legacy `with_trace_id()` method is deprecated.
260    ///
261    /// Resolution order for trace identity:
262    /// 1. **(Canonical)** If `trace_ctx` was set explicitly via `with_trace_ctx()`, use it. Derive legacy `trace_id` from it at the compatibility boundary.
263    /// 2. **(Legacy/compat)** If only `trace_id` was set via deprecated `with_trace_id()`, derive `trace_ctx` from it.
264    /// 3. **(Default)** If neither was set, generate a fresh `TraceCtx` and derive legacy `trace_id`.
265    ///
266    /// The legacy `trace_id` is always derived — never independently generated — when
267    /// `trace_ctx` is present. This ensures a single source of truth for trace identity.
268    pub fn build(self) -> ExecCtx {
269        let limits = self.limits.unwrap_or_default();
270        // Use a high safety-net timeout on the Client itself (5 minutes).
271        // Actual per-request timeouts are applied at the RequestBuilder level
272        // by each backend, driven by LlmCall.timeout or PipelineLimits.request_timeout.
273        let client_timeout = self.timeout.unwrap_or(Duration::from_secs(300));
274        let client = self.client.unwrap_or_else(|| {
275            Client::builder()
276                .timeout(client_timeout)
277                .build()
278                .expect("Failed to build HTTP client")
279        });
280
281        let (trace_id, trace_ctx) = match (self.trace_ctx, self.trace_id) {
282            // Canonical path: TraceCtx was explicitly set
283            (Some(ctx), _) => {
284                let legacy = TraceId::from_trace_ctx(&ctx);
285                (legacy, ctx)
286            }
287            // Compat path: only legacy TraceId was set
288            (None, Some(id)) => {
289                let ctx = id.to_trace_ctx();
290                (id, ctx)
291            }
292            // Default: generate fresh TraceCtx
293            (None, None) => {
294                let ctx = TraceCtx::generate();
295                let legacy = TraceId::from_trace_ctx(&ctx);
296                (legacy, ctx)
297            }
298        };
299
300        ExecCtx {
301            client,
302            base_url: normalize_base_url(&self.base_url),
303            backend: self.backend.unwrap_or_else(|| Arc::new(OllamaBackend)),
304            backoff: self.backoff.unwrap_or_else(BackoffConfig::none),
305            vars: self.vars,
306            cancellation: self.cancellation,
307            event_handler: self.event_handler,
308            trace_id,
309            trace_ctx,
310            limits,
311        }
312    }
313}
314
315/// Strip known provider path suffixes from a base URL.
316/// This prevents double-pathing when backends append their own paths.
317/// e.g., "https://api.openai.com/v1" -> "https://api.openai.com"
318/// e.g., "http://localhost:11434/api" -> "http://localhost:11434"
319fn normalize_base_url(url: &str) -> String {
320    let trimmed = url.trim_end_matches('/');
321    // Strip known suffixes (order matters — longest first)
322    for suffix in &[
323        "/v1/chat/completions",
324        "/v1/chat",
325        "/v1",
326        "/api/generate",
327        "/api/chat",
328        "/api",
329    ] {
330        if let Some(stripped) = trimmed.strip_suffix(suffix) {
331            return stripped.to_string();
332        }
333    }
334    trimmed.to_string()
335}
336
337#[allow(deprecated)]
338#[cfg(test)]
339mod tests {
340    use super::*;
341
342    #[test]
343    fn test_normalize_base_url_strips_v1() {
344        assert_eq!(
345            normalize_base_url("https://api.openai.com/v1"),
346            "https://api.openai.com"
347        );
348        assert_eq!(
349            normalize_base_url("https://api.openai.com/v1/"),
350            "https://api.openai.com"
351        );
352    }
353
354    #[test]
355    fn test_normalize_base_url_strips_api() {
356        assert_eq!(
357            normalize_base_url("http://localhost:11434/api"),
358            "http://localhost:11434"
359        );
360        assert_eq!(
361            normalize_base_url("http://localhost:11434/api/"),
362            "http://localhost:11434"
363        );
364    }
365
366    #[test]
367    fn test_normalize_base_url_preserves_clean() {
368        assert_eq!(
369            normalize_base_url("http://localhost:11434"),
370            "http://localhost:11434"
371        );
372        assert_eq!(
373            normalize_base_url("https://api.openai.com"),
374            "https://api.openai.com"
375        );
376    }
377
378    #[test]
379    fn test_normalize_base_url_strips_full_path() {
380        assert_eq!(
381            normalize_base_url("https://api.openai.com/v1/chat/completions"),
382            "https://api.openai.com"
383        );
384    }
385
386    #[test]
387    fn test_normalize_base_url_trailing_slash() {
388        assert_eq!(
389            normalize_base_url("http://localhost:11434/"),
390            "http://localhost:11434"
391        );
392    }
393
394    #[test]
395    fn test_default_timeout_applied() {
396        // Verify the builder accepts the timeout method and compiles.
397        let _ctx = ExecCtx::builder("http://localhost:11434")
398            .timeout(Duration::from_secs(120))
399            .build();
400        // Smoke test: builds without panic
401    }
402
403    #[test]
404    fn test_trace_ctx_generated_by_default() {
405        let ctx = ExecCtx::builder("http://localhost:11434").build();
406        // Both trace forms should be populated
407        assert!(!ctx.trace_id.as_str().is_empty());
408        assert!(!ctx.trace_ctx.trace_id.is_empty());
409    }
410
411    #[test]
412    fn test_trace_ctx_explicit_sets_legacy() {
413        let trace = TraceCtx::from_trace_id("0af7651916cd43dd8448eb211c80319c");
414        let ctx = ExecCtx::builder("http://localhost:11434")
415            .with_trace_ctx(trace.clone())
416            .build();
417        assert_eq!(ctx.trace_ctx.trace_id, "0af7651916cd43dd8448eb211c80319c");
418        assert_eq!(ctx.trace_id.as_str(), "0af7651916cd43dd8448eb211c80319c");
419    }
420
421    #[test]
422    fn test_legacy_trace_id_derives_trace_ctx() {
423        let id = TraceId::from_string("my-legacy-trace");
424        let ctx = ExecCtx::builder("http://localhost:11434")
425            .with_trace_id(id)
426            .build();
427        assert_eq!(ctx.trace_id.as_str(), "my-legacy-trace");
428        assert_eq!(ctx.trace_ctx.trace_id, "my-legacy-trace");
429    }
430
431    #[test]
432    fn test_trace_ctx_takes_priority_over_trace_id() {
433        let trace = TraceCtx::from_trace_id("canonical-trace");
434        let id = TraceId::from_string("legacy-trace");
435        let ctx = ExecCtx::builder("http://localhost:11434")
436            .with_trace_id(id)
437            .with_trace_ctx(trace)
438            .build();
439        // TraceCtx wins
440        assert_eq!(ctx.trace_ctx.trace_id, "canonical-trace");
441        assert_eq!(ctx.trace_id.as_str(), "canonical-trace");
442    }
443}
llm_pipeline/exec_ctx.rs

llm_pipeline/
exec_ctx.rs
