atomr_agents_agent/
model_pin.rs

1//! Model/provider version pinning + drift detection (FR-3).
2//!
3//! A model change is a gate-governed promotion event: a silently
4//! upgraded model changes live trading behavior un-gated and breaks the
5//! "which model produced this trade?" audit chain. [`PinnedClient`] is an
6//! [`InferenceClient`] middleware that, at call time, resolves the
7//! concrete model + version, compares it to a declared [`ModelPin`], and:
8//!
9//! * refuses to run (typed [`ModelPinViolation`]) when `strict_pin` is on
10//!   and no pin is set;
11//! * emits a [`RunEventKind::ModelDrift`] onto the telemetry backbone
12//!   (FR-19) whenever the resolved model differs from the pin;
13//! * otherwise delegates to the inner client unchanged.
14//!
15//! The resolved pin is also what the recording path stamps into every
16//! [`InferenceRecord`](atomr_agents_state::InferenceRecord).
17
18use std::sync::Arc;
19
20use async_trait::async_trait;
21use atomr_agents_core::Result;
22use atomr_agents_observability::{ModelPinRef, RunEvent, RunEventKind, Telemetry};
23use atomr_agents_tool::Provider;
24use atomr_infer_core::batch::ExecuteBatch;
25use thiserror::Error;
26
27use crate::inference::{InferenceClient, TurnResult};
28
29/// A first-class pin on `(provider, model, version, params)`.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct ModelPin {
32    pub provider: Provider,
33    pub model_id: String,
34    /// Semantic version or content digest of the model.
35    pub model_version: String,
36    /// Hash of the sampling/params that affect output.
37    pub params_hash: String,
38}
39
40impl ModelPin {
41    pub fn new(
42        provider: Provider,
43        model_id: impl Into<String>,
44        model_version: impl Into<String>,
45        params_hash: impl Into<String>,
46    ) -> Self {
47        Self {
48            provider,
49            model_id: model_id.into(),
50            model_version: model_version.into(),
51            params_hash: params_hash.into(),
52        }
53    }
54
55    fn as_ref(&self) -> ModelPinRef {
56        ModelPinRef {
57            provider: provider_str(self.provider).into(),
58            model_id: self.model_id.clone(),
59            model_version: self.model_version.clone(),
60            params_hash: self.params_hash.clone(),
61        }
62    }
63}
64
65/// What a provider actually resolves to at call time.
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub struct ResolvedModel {
68    pub provider: Provider,
69    pub model_id: String,
70    pub model_version: String,
71    pub params_hash: String,
72}
73
74/// Kind of drift detected.
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub enum DriftKind {
77    VersionDrift,
78    ParamsDrift,
79}
80
81impl DriftKind {
82    fn as_str(self) -> &'static str {
83        match self {
84            DriftKind::VersionDrift => "version_drift",
85            DriftKind::ParamsDrift => "params_drift",
86        }
87    }
88}
89
90/// Resolves the concrete model/version a client would use right now.
91/// Hosts implement this (often by querying the provider) so a deployed
92/// pin can be re-validated on a heartbeat to catch provider-side silent
93/// upgrades.
94pub trait ModelResolver: Send + Sync + 'static {
95    fn resolve_version(&self) -> ResolvedModel;
96}
97
98/// A static resolver (tests / known-version deployments).
99pub struct StaticResolver(pub ResolvedModel);
100
101impl ModelResolver for StaticResolver {
102    fn resolve_version(&self) -> ResolvedModel {
103        self.0.clone()
104    }
105}
106
107/// Pin violation: strict mode with no pin, or (optionally) a hard
108/// refusal on drift.
109#[derive(Debug, Error)]
110pub enum ModelPinViolation {
111    #[error("strict_pin: a run requires a ModelPin but none was set")]
112    NoPinInStrictMode,
113    #[error("model pin mismatch ({kind:?}): expected {expected:?}, got {actual:?}")]
114    Mismatch {
115        kind: DriftKind,
116        expected: ModelPin,
117        actual: ResolvedModel,
118    },
119}
120
121impl From<ModelPinViolation> for atomr_agents_core::AgentError {
122    fn from(e: ModelPinViolation) -> Self {
123        atomr_agents_core::AgentError::PolicyDenied(e.to_string())
124    }
125}
126
127/// Compare a resolved model to a pin; `None` if they match.
128pub fn detect_drift(pin: &ModelPin, actual: &ResolvedModel) -> Option<DriftKind> {
129    if pin.provider != actual.provider
130        || pin.model_id != actual.model_id
131        || pin.model_version != actual.model_version
132    {
133        Some(DriftKind::VersionDrift)
134    } else if pin.params_hash != actual.params_hash {
135        Some(DriftKind::ParamsDrift)
136    } else {
137        None
138    }
139}
140
141/// `InferenceClient` middleware enforcing a [`ModelPin`].
142pub struct PinnedClient {
143    inner: Arc<dyn InferenceClient>,
144    resolver: Arc<dyn ModelResolver>,
145    pin: Option<ModelPin>,
146    strict: bool,
147    /// If true, drift is a hard refusal (not just an emitted event).
148    refuse_on_drift: bool,
149    telemetry: Option<Telemetry>,
150    run_id: Option<String>,
151}
152
153impl PinnedClient {
154    pub fn new(inner: Arc<dyn InferenceClient>, resolver: Arc<dyn ModelResolver>) -> Self {
155        Self {
156            inner,
157            resolver,
158            pin: None,
159            strict: false,
160            refuse_on_drift: false,
161            telemetry: None,
162            run_id: None,
163        }
164    }
165
166    /// Pin the model (see [`crate::Agent`] wiring as `pin_model`).
167    pub fn pin_model(mut self, pin: ModelPin) -> Self {
168        self.pin = Some(pin);
169        self
170    }
171
172    /// Require a pin to be set, else refuse to run.
173    pub fn strict(mut self, strict: bool) -> Self {
174        self.strict = strict;
175        self
176    }
177
178    /// Make drift a hard refusal in addition to an emitted event.
179    pub fn refuse_on_drift(mut self, refuse: bool) -> Self {
180        self.refuse_on_drift = refuse;
181        self
182    }
183
184    pub fn with_telemetry(mut self, telemetry: Telemetry, run_id: impl Into<String>) -> Self {
185        self.telemetry = Some(telemetry);
186        self.run_id = Some(run_id.into());
187        self
188    }
189
190    /// The pin this client carries, if any (for stamping into records).
191    pub fn pin(&self) -> Option<&ModelPin> {
192        self.pin.as_ref()
193    }
194
195    /// Re-validate the pin against a freshly resolved version (heartbeat).
196    /// Returns the drift kind if the provider has drifted.
197    pub fn revalidate(&self) -> Option<DriftKind> {
198        let actual = self.resolver.resolve_version();
199        self.pin.as_ref().and_then(|p| detect_drift(p, &actual))
200    }
201
202    // The error variant intentionally carries the full pin + resolved
203    // model for audit; size is not a hot-path concern here.
204    #[allow(clippy::result_large_err)]
205    fn enforce(&self) -> std::result::Result<(), ModelPinViolation> {
206        let Some(pin) = self.pin.as_ref() else {
207            if self.strict {
208                return Err(ModelPinViolation::NoPinInStrictMode);
209            }
210            return Ok(());
211        };
212        let actual = self.resolver.resolve_version();
213        if let Some(kind) = detect_drift(pin, &actual) {
214            if let (Some(t), Some(run)) = (&self.telemetry, &self.run_id) {
215                t.emit(
216                    RunEvent::new(RunEventKind::ModelDrift {
217                        expected: format!("{:?}", pin.as_ref()),
218                        actual: format!(
219                            "{:?}",
220                            ModelPinRef {
221                                provider: provider_str(actual.provider).into(),
222                                model_id: actual.model_id.clone(),
223                                model_version: actual.model_version.clone(),
224                                params_hash: actual.params_hash.clone(),
225                            }
226                        ),
227                        drift: kind.as_str().into(),
228                    })
229                    .with_model_pin(pin.as_ref()),
230                );
231                let _ = run; // run id retained for callers that set it
232            }
233            if self.refuse_on_drift {
234                return Err(ModelPinViolation::Mismatch {
235                    kind,
236                    expected: pin.clone(),
237                    actual,
238                });
239            }
240        }
241        Ok(())
242    }
243}
244
245#[async_trait]
246impl InferenceClient for PinnedClient {
247    fn provider(&self) -> Provider {
248        self.inner.provider()
249    }
250
251    async fn run(&self, batch: ExecuteBatch) -> Result<TurnResult> {
252        self.enforce()?;
253        self.inner.run(batch).await
254    }
255}
256
257fn provider_str(p: Provider) -> &'static str {
258    match p {
259        Provider::OpenAi => "open_ai",
260        Provider::Anthropic => "anthropic",
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use atomr_agents_observability::InMemoryTelemetrySink;
268    use atomr_infer_core::tokens::TokenUsage;
269
270    struct FixedClient(Provider);
271    #[async_trait]
272    impl InferenceClient for FixedClient {
273        fn provider(&self) -> Provider {
274            self.0
275        }
276        async fn run(&self, _batch: ExecuteBatch) -> Result<TurnResult> {
277            Ok(TurnResult {
278                text: "ok".into(),
279                usage: TokenUsage::default(),
280                finish_reason: None,
281                tool_calls: vec![],
282            })
283        }
284    }
285
286    fn batch() -> ExecuteBatch {
287        ExecuteBatch {
288            request_id: "r".into(),
289            model: "m".into(),
290            messages: vec![],
291            sampling: Default::default(),
292            stream: false,
293            estimated_tokens: 0,
294        }
295    }
296
297    fn resolved(version: &str, params: &str) -> ResolvedModel {
298        ResolvedModel {
299            provider: Provider::Anthropic,
300            model_id: "claude".into(),
301            model_version: version.into(),
302            params_hash: params.into(),
303        }
304    }
305
306    #[tokio::test]
307    async fn strict_mode_without_pin_refuses() {
308        let c = PinnedClient::new(
309            Arc::new(FixedClient(Provider::Anthropic)),
310            Arc::new(StaticResolver(resolved("1", "p"))),
311        )
312        .strict(true);
313        assert!(c.run(batch()).await.is_err());
314    }
315
316    #[tokio::test]
317    async fn matching_pin_passes_through() {
318        let c = PinnedClient::new(
319            Arc::new(FixedClient(Provider::Anthropic)),
320            Arc::new(StaticResolver(resolved("1", "p"))),
321        )
322        .pin_model(ModelPin::new(Provider::Anthropic, "claude", "1", "p"))
323        .strict(true);
324        let r = c.run(batch()).await.unwrap();
325        assert_eq!(r.text, "ok");
326        assert!(c.revalidate().is_none());
327    }
328
329    #[tokio::test]
330    async fn version_drift_emits_event_and_can_refuse() {
331        let sink = Arc::new(InMemoryTelemetrySink::new());
332        let tel = Telemetry::new().with_sink(sink.clone());
333        // pin says version "1", provider resolves "2" -> VersionDrift
334        let c = PinnedClient::new(
335            Arc::new(FixedClient(Provider::Anthropic)),
336            Arc::new(StaticResolver(resolved("2", "p"))),
337        )
338        .pin_model(ModelPin::new(Provider::Anthropic, "claude", "1", "p"))
339        .with_telemetry(tel, "run-1")
340        .refuse_on_drift(true);
341
342        let err = c.run(batch()).await.unwrap_err();
343        assert!(err.to_string().contains("mismatch"));
344        let events = sink.events();
345        assert_eq!(events.len(), 1);
346        match &events[0].kind {
347            RunEventKind::ModelDrift { drift, .. } => assert_eq!(drift, "version_drift"),
348            other => panic!("unexpected {other:?}"),
349        }
350        assert!(events[0].model_pin.is_some());
351    }
352
353    #[tokio::test]
354    async fn params_drift_detected() {
355        assert_eq!(
356            detect_drift(
357                &ModelPin::new(Provider::Anthropic, "claude", "1", "p1"),
358                &resolved("1", "p2")
359            ),
360            Some(DriftKind::ParamsDrift)
361        );
362    }
363}
atomr_agents_agent/model_pin.rs

atomr_agents_agent/
model_pin.rs
