context_harness/agent_model/
mod.rs

1//! Provider-neutral model calls, independent of the agent execution loop.
2//! Recorded calls persist metadata only; conversation snapshots belong to the
3//! runtime. No credentials or raw provider errors enter the event log.
4
5mod catalog;
6pub mod fake;
7pub mod ollama;
8pub mod openai;
9
10pub use catalog::{ModelProviderCatalog, ModelProviderFactory, ModelProviderImplementation};
11
12use anyhow::{ensure, Context};
13use async_trait::async_trait;
14use serde::{Deserialize, Serialize};
15use serde_json::{json, Value};
16use sha2::{Digest, Sha256};
17use std::sync::Arc;
18use std::{
19    collections::{BTreeMap, HashSet},
20    time::Duration,
21};
22use uuid::Uuid;
23
24use crate::agent_resource::ModelDefinition;
25use crate::agent_store::{AgentRunStore, BudgetDecision, LimitReason};
26
27pub type ModelResult<T> = Result<T, ModelError>;
28
29#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
30pub struct ToolCall {
31    pub id: String,
32    pub name: String,
33    pub arguments: Value,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct ModelTool {
38    pub name: String,
39    pub description: String,
40    pub parameters: Value,
41}
42
43/// Opaque adapter-owned continuation, serialized with the conversation. It may
44/// contain sensitive context; never include it in metadata events or diagnostics.
45#[derive(Clone, Serialize, Deserialize, PartialEq)]
46pub struct ProviderContinuation {
47    provider: String,
48    items: Vec<Value>,
49}
50impl std::fmt::Debug for ProviderContinuation {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.debug_struct("ProviderContinuation")
53            .field("provider", &self.provider)
54            .field("item_count", &self.items.len())
55            .finish()
56    }
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
60#[serde(tag = "role", rename_all = "snake_case")]
61pub enum ModelMessage {
62    System {
63        content: String,
64    },
65    User {
66        content: String,
67    },
68    Assistant {
69        text: String,
70        tool_calls: Vec<ToolCall>,
71        continuation: Option<ProviderContinuation>,
72    },
73    Tool {
74        call_id: String,
75        content: String,
76    },
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize)]
80pub struct OutputSchema {
81    pub name: String,
82    pub schema: Value,
83}
84
85#[derive(Debug, Clone, Default, Serialize, Deserialize)]
86pub struct ModelRequest {
87    pub messages: Vec<ModelMessage>,
88    pub tools: Vec<ModelTool>,
89    pub output_schema: Option<OutputSchema>,
90    pub max_output_tokens: Option<u32>,
91}
92
93impl ModelRequest {
94    pub fn validate(&self) -> ModelResult<()> {
95        let invalid = || ModelError::new(ModelErrorKind::InvalidRequest);
96        if self.messages.is_empty() || self.max_output_tokens == Some(0) {
97            return Err(invalid());
98        }
99        let mut names = HashSet::new();
100        for tool in &self.tools {
101            if tool.name.trim().is_empty()
102                || !names.insert(&tool.name)
103                || !tool.parameters.is_object()
104            {
105                return Err(invalid());
106            }
107        }
108        if self
109            .output_schema
110            .as_ref()
111            .is_some_and(|s| s.name.is_empty() || !s.schema.is_object())
112        {
113            return Err(invalid());
114        }
115        let mut pending = HashSet::new();
116        let mut seen = HashSet::new();
117        for message in &self.messages {
118            match message {
119                ModelMessage::Tool { call_id, .. } => {
120                    if !pending.remove(call_id.as_str()) {
121                        return Err(invalid());
122                    }
123                }
124                ModelMessage::Assistant { tool_calls, .. } => {
125                    if !pending.is_empty() {
126                        return Err(invalid());
127                    }
128                    for call in tool_calls {
129                        if call.id.is_empty()
130                            || call.name.is_empty()
131                            || !call.arguments.is_object()
132                            || !seen.insert(call.id.as_str())
133                        {
134                            return Err(invalid());
135                        }
136                        pending.insert(call.id.as_str());
137                    }
138                }
139                _ if !pending.is_empty() => return Err(invalid()),
140                _ => {}
141            }
142        }
143        if !pending.is_empty() {
144            return Err(invalid());
145        }
146        Ok(())
147    }
148}
149
150#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
151#[serde(rename_all = "snake_case")]
152pub enum FinishReason {
153    Completed,
154    ToolCalls,
155    Length,
156    Refusal,
157    ContentFilter,
158}
159
160#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
161pub struct Usage {
162    pub input_tokens: u64,
163    pub output_tokens: u64,
164    pub total_tokens: u64,
165}
166
167#[derive(Debug, Clone, Serialize, Deserialize)]
168pub struct ModelResponse {
169    pub text: String,
170    pub tool_calls: Vec<ToolCall>,
171    pub usage: Option<Usage>,
172    pub finish_reason: FinishReason,
173    pub structured_output: Option<Value>,
174    pub continuation: Option<ProviderContinuation>,
175}
176
177pub struct RecordedModelResponse {
178    pub response: ModelResponse,
179    pub limit: Option<LimitReason>,
180    /// Number of provider requests consumed by this conversational turn.
181    pub attempts: u32,
182}
183
184const MODEL_MAX_ATTEMPTS: u32 = 3;
185
186fn model_retry_delay(failed_attempt: u32) -> Duration {
187    match failed_attempt {
188        1 => Duration::from_millis(250),
189        _ => Duration::from_millis(750),
190    }
191}
192
193fn model_error_is_retryable(kind: ModelErrorKind) -> bool {
194    matches!(
195        kind,
196        ModelErrorKind::InvalidResponse
197            | ModelErrorKind::RateLimited
198            | ModelErrorKind::Unavailable
199            | ModelErrorKind::Timeout
200            | ModelErrorKind::Transport
201    )
202}
203impl ModelResponse {
204    pub fn text(text: impl Into<String>) -> Self {
205        Self {
206            text: text.into(),
207            tool_calls: vec![],
208            usage: None,
209            finish_reason: FinishReason::Completed,
210            structured_output: None,
211            continuation: None,
212        }
213    }
214
215    pub fn message(&self) -> ModelMessage {
216        ModelMessage::Assistant {
217            text: self.text.clone(),
218            tool_calls: self.tool_calls.clone(),
219            continuation: self.continuation.clone(),
220        }
221    }
222
223    fn validate(&self, request: &ModelRequest) -> ModelResult<()> {
224        let invalid = || ModelError::new(ModelErrorKind::InvalidResponse);
225        if (self.finish_reason == FinishReason::ToolCalls) != !self.tool_calls.is_empty() {
226            return Err(invalid());
227        }
228        let mut ids = HashSet::new();
229        let previous_ids: HashSet<_> = request
230            .messages
231            .iter()
232            .flat_map(|message| match message {
233                ModelMessage::Assistant { tool_calls, .. } => tool_calls
234                    .iter()
235                    .map(|call| call.id.as_str())
236                    .collect::<Vec<_>>(),
237                _ => vec![],
238            })
239            .collect();
240        for call in &self.tool_calls {
241            if call.id.is_empty()
242                || !ids.insert(&call.id)
243                || previous_ids.contains(call.id.as_str())
244                || !call.arguments.is_object()
245                || !request.tools.iter().any(|t| t.name == call.name)
246            {
247                return Err(invalid());
248            }
249        }
250        if request.output_schema.is_some() && self.finish_reason == FinishReason::Completed {
251            let parsed: Value = serde_json::from_str(&self.text).map_err(|_| invalid())?;
252            if self.structured_output.as_ref() != Some(&parsed) {
253                return Err(invalid());
254            }
255        }
256        Ok(())
257    }
258}
259
260/// Safe diagnostic categories: never attach a raw HTTP body, URL or credential.
261#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
262#[serde(rename_all = "snake_case")]
263pub enum ModelErrorKind {
264    InvalidRequest,
265    InvalidResponse,
266    MissingCredentials,
267    Authentication,
268    RateLimited,
269    Unavailable,
270    ModelUnavailable,
271    UnsupportedCapability,
272    Timeout,
273    Transport,
274    ProviderFailure,
275    ScriptExhausted,
276}
277
278fn wire_name(name: &str) -> String {
279    format!("{:x}", Sha256::digest(name.as_bytes()))
280}
281#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
282pub struct ModelError {
283    pub kind: ModelErrorKind,
284    pub http_status: Option<u16>,
285}
286impl ModelError {
287    pub fn new(kind: ModelErrorKind) -> Self {
288        Self {
289            kind,
290            http_status: None,
291        }
292    }
293}
294impl std::fmt::Display for ModelError {
295    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296        write!(f, "model call failed: {:?}", self.kind)?;
297        if let Some(status) = self.http_status {
298            write!(f, " (HTTP {status})")?;
299        }
300        Ok(())
301    }
302}
303impl std::error::Error for ModelError {}
304
305/// One configured provider/model. The model name and credentials are bound at
306/// construction, so requests cannot override their registry alias's model.
307#[async_trait]
308pub trait ModelProvider: Send + Sync {
309    async fn generate(&self, request: &ModelRequest) -> ModelResult<ModelResponse>;
310}
311
312struct RegisteredModel {
313    provider_name: String,
314    model_name: String,
315    implementation: ModelProviderImplementation,
316    provider: Arc<dyn ModelProvider>,
317}
318#[derive(Default)]
319pub struct ModelRegistry {
320    models: BTreeMap<String, RegisteredModel>,
321}
322impl ModelRegistry {
323    /// Configured fake models return a clearly labelled synthetic final answer.
324    /// Tests can instead register a FakeModel with a scripted sequence.
325    pub fn from_config(models: &BTreeMap<String, ModelDefinition>) -> anyhow::Result<Self> {
326        Self::from_config_with_catalog(models, &ModelProviderCatalog::with_builtins()?)
327    }
328
329    /// Construct configured models from factories explicitly trusted by the host.
330    pub fn from_config_with_catalog(
331        models: &BTreeMap<String, ModelDefinition>,
332        catalog: &ModelProviderCatalog,
333    ) -> anyhow::Result<Self> {
334        catalog.validate_config(models)?;
335        let mut registry = Self::default();
336        for (alias, definition) in models {
337            let (provider, implementation) = catalog.build(alias, definition)?;
338            registry.register_with_implementation(
339                alias,
340                &definition.provider,
341                &definition.model,
342                implementation,
343                provider,
344            )?;
345        }
346        Ok(registry)
347    }
348
349    pub fn register(
350        &mut self,
351        alias: &str,
352        provider_name: &str,
353        model_name: &str,
354        provider: Arc<dyn ModelProvider>,
355    ) -> anyhow::Result<()> {
356        self.register_with_implementation(
357            alias,
358            provider_name,
359            model_name,
360            ModelProviderImplementation::new(provider_name, "unversioned"),
361            provider,
362        )
363    }
364
365    fn register_with_implementation(
366        &mut self,
367        alias: &str,
368        provider_name: &str,
369        model_name: &str,
370        implementation: ModelProviderImplementation,
371        provider: Arc<dyn ModelProvider>,
372    ) -> anyhow::Result<()> {
373        ensure!(
374            !alias.trim().is_empty()
375                && !provider_name.trim().is_empty()
376                && !model_name.trim().is_empty(),
377            "model identity is required"
378        );
379        ensure!(
380            !self.models.contains_key(alias),
381            "duplicate model alias '{alias}'"
382        );
383        self.models.insert(
384            alias.into(),
385            RegisteredModel {
386                provider_name: provider_name.into(),
387                model_name: model_name.into(),
388                implementation,
389                provider,
390            },
391        );
392        Ok(())
393    }
394
395    /// Stable, non-secret binding used to validate persisted runtime checkpoints.
396    pub fn identity(&self, alias: &str) -> anyhow::Result<(&str, &str)> {
397        let model = self.models.get(alias).context("unknown model alias")?;
398        Ok((&model.provider_name, &model.model_name))
399    }
400
401    /// Stable adapter identity available for future checkpoint compatibility.
402    pub fn implementation_identity(
403        &self,
404        alias: &str,
405    ) -> anyhow::Result<&ModelProviderImplementation> {
406        Ok(&self
407            .models
408            .get(alias)
409            .context("unknown model alias")?
410            .implementation)
411    }
412
413    pub async fn generate(
414        &self,
415        alias: &str,
416        request: &ModelRequest,
417    ) -> anyhow::Result<ModelResponse> {
418        let model = self
419            .models
420            .get(alias)
421            .with_context(|| format!("unknown model alias '{alias}'"))?;
422        request.validate()?;
423        let response = model.provider.generate(request).await?;
424        response.validate(request)?;
425        Ok(response)
426    }
427
428    /// Persist lifecycle metadata around one model call. Run completion remains
429    /// the caller's responsibility. Cancellation leaves an unmatched requested
430    /// event; safe reconciliation belongs to the future resume layer.
431    pub async fn generate_recorded(
432        &self,
433        store: &AgentRunStore,
434        run_id: &str,
435        alias: &str,
436        request: &ModelRequest,
437    ) -> anyhow::Result<ModelResponse> {
438        let recorded = self
439            .generate_recorded_with_budget(store, run_id, alias, request)
440            .await?;
441        ensure!(
442            recorded.limit.is_none(),
443            "model response exceeded run budget"
444        );
445        Ok(recorded.response)
446    }
447
448    pub async fn generate_recorded_with_budget(
449        &self,
450        store: &AgentRunStore,
451        run_id: &str,
452        alias: &str,
453        request: &ModelRequest,
454    ) -> anyhow::Result<RecordedModelResponse> {
455        let model = self
456            .models
457            .get(alias)
458            .with_context(|| format!("unknown model alias '{alias}'"))?;
459        request.validate()?;
460        let run = store
461            .get_run(run_id)
462            .await?
463            .context("run not found in workspace")?;
464        ensure!(run.model == alias, "model alias does not match run binding");
465        let call_id = Uuid::new_v4().to_string();
466        let decision = store.record_model_requested(run_id, &json!({
467            "call_id": call_id, "attempt": 1, "max_attempts": MODEL_MAX_ATTEMPTS,
468            "model_alias": alias, "provider": model.provider_name,
469            "model": model.model_name, "message_count": request.messages.len(),
470            "tool_count": request.tools.len(), "structured_output": request.output_schema.is_some()
471        })).await?;
472        if let BudgetDecision::Exceeded(reason) = decision {
473            return Ok(RecordedModelResponse {
474                response: ModelResponse::text(String::new()),
475                limit: Some(reason),
476                attempts: 0,
477            });
478        }
479        let mut attempts = 1;
480        let mut call_id = call_id;
481        loop {
482            let result = model.provider.generate(request).await.and_then(|response| {
483                response.validate(request)?;
484                Ok(response)
485            });
486            match result {
487                Ok(response) => {
488                    let usage = response
489                        .usage
490                        .as_ref()
491                        .map(|usage| (usage.input_tokens, usage.output_tokens, usage.total_tokens));
492                    let decision = store
493                    .record_model_responded(
494                        run_id,
495                        &json!({
496                            "call_id": call_id, "attempt": attempts,
497                            "max_attempts": MODEL_MAX_ATTEMPTS,
498                            "finish_reason": response.finish_reason,
499                            "usage": response.usage, "tool_call_count": response.tool_calls.len()
500                        }),
501                        usage,
502                    )
503                    .await?;
504                    return Ok(RecordedModelResponse {
505                        response,
506                        limit: match decision {
507                            BudgetDecision::Allowed => None,
508                            BudgetDecision::Exceeded(reason) => Some(reason),
509                        },
510                        attempts,
511                    });
512                }
513                Err(error) => {
514                    let retryable = model_error_is_retryable(error.kind);
515                    let will_retry = attempts < MODEL_MAX_ATTEMPTS && retryable;
516                    store
517                        .append_event(
518                            run_id,
519                            "model.failed",
520                            &json!({
521                                "call_id": call_id,
522                                "attempt": attempts,
523                                "max_attempts": MODEL_MAX_ATTEMPTS,
524                                "error": "model request failed",
525                                "error_kind": error.kind,
526                                "http_status": error.http_status,
527                                "retryable": retryable,
528                                "will_retry": will_retry,
529                            }),
530                        )
531                        .await?;
532                    if !will_retry {
533                        return Err(error.into());
534                    }
535                    let delay = model_retry_delay(attempts);
536                    store
537                        .append_event(
538                            run_id,
539                            "model.retrying",
540                            &json!({
541                                "call_id": call_id,
542                                "failed_attempt": attempts,
543                                "next_attempt": attempts + 1,
544                                "max_attempts": MODEL_MAX_ATTEMPTS,
545                                "error_kind": error.kind,
546                                "http_status": error.http_status,
547                                "delay_ms": delay.as_millis(),
548                            }),
549                        )
550                        .await?;
551                    tokio::time::sleep(delay).await;
552                    attempts += 1;
553                    let retry_of = call_id;
554                    call_id = Uuid::new_v4().to_string();
555                    let decision = store.record_model_requested(run_id, &json!({
556                        "call_id": call_id, "attempt": attempts,
557                        "max_attempts": MODEL_MAX_ATTEMPTS,
558                        "retry_of": retry_of,
559                        "model_alias": alias, "provider": model.provider_name,
560                        "model": model.model_name, "message_count": request.messages.len(),
561                        "tool_count": request.tools.len(), "structured_output": request.output_schema.is_some()
562                    })).await?;
563                    if let BudgetDecision::Exceeded(reason) = decision {
564                        return Ok(RecordedModelResponse {
565                            response: ModelResponse::text(String::new()),
566                            limit: Some(reason),
567                            attempts: attempts - 1,
568                        });
569                    }
570                }
571            }
572        }
573    }
574}
context_harness/agent_model/mod.rs

context_harness/agent_model/
mod.rs
