diff --git a/Cargo.lock b/Cargo.lock index edca5ee..2ef4d1f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -679,6 +679,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-util", + "tracing", ] [[package]] @@ -723,12 +724,14 @@ dependencies = [ "async-stream", "async-trait", "bytes", + "dirs", "eventsource-stream", "futures", "harness-core", "reqwest", "serde", "serde_json", + "tempfile", "thiserror 2.0.18", "tokio", "tokio-util", diff --git a/crates/harness-app/Cargo.toml b/crates/harness-app/Cargo.toml index ceb4f67..2eb3ea2 100644 --- a/crates/harness-app/Cargo.toml +++ b/crates/harness-app/Cargo.toml @@ -14,6 +14,7 @@ tokio = { workspace = true } tokio-util = { workspace = true } dirs = { workspace = true } thiserror = { workspace = true } +tracing = { workspace = true } [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/harness-app/src/lib.rs b/crates/harness-app/src/lib.rs index 999bcb3..44073c2 100644 --- a/crates/harness-app/src/lib.rs +++ b/crates/harness-app/src/lib.rs @@ -19,7 +19,7 @@ use harness_core::tool::ToolRegistry; use harness_core::types::{ Message, MessageId, ModelRef, Part, PartBody, PartId, Session, SessionId, }; -use harness_providers::{AnthropicProvider, ProviderRegistry}; +use harness_providers::{AnthropicProvider, ModelCatalog, OpenAiProvider, ProviderRegistry}; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; @@ -83,6 +83,7 @@ struct EngineInner { permissions: Arc, tools: ToolRegistry, providers: ProviderRegistry, + catalog: ModelCatalog, cwd: PathBuf, data_dir: PathBuf, runs: Mutex>, @@ -97,7 +98,9 @@ impl EngineHandle { pub fn init(cwd: PathBuf) -> Result { let path = db_path(&cwd); let store = Store::open(&path)?; - Self::new(cwd, store) + let handle = Self::new(cwd, store)?; + handle.spawn_catalog_refresh(); + Ok(handle) } /// In-memory store — useful for tests and ephemeral sessions. @@ -122,12 +125,31 @@ impl EngineHandle { { providers.register(Arc::new(AnthropicProvider::new(key))); } + if let Some(key) = config + .providers + .get("openai") + .and_then(|p| p.api_key.clone()) + { + let provider = match config + .providers + .get("openai") + .and_then(|p| p.base_url.clone()) + { + Some(base_url) => OpenAiProvider::with_base_url(key, base_url), + None => OpenAiProvider::new(key), + }; + providers.register(Arc::new(provider)); + } let data_dir = dirs::data_dir() .unwrap_or_else(std::env::temp_dir) .join("ai-harness") .join("tool-output"); + // No network at construction: use the cached copy if fresh, else the baked snapshot. + // `init` warms the cache in the background for the next launch. + let catalog = ModelCatalog::load_cached_or_baked(&ModelCatalog::default_cache_path()); + Ok(Self { inner: Arc::new(EngineInner { config, @@ -136,6 +158,7 @@ impl EngineHandle { permissions, tools, providers, + catalog, cwd, data_dir, runs: Mutex::new(HashMap::new()), @@ -143,6 +166,15 @@ impl EngineHandle { }) } + /// Fire-and-forget refresh of the models.dev cache so the next launch has current pricing. + fn spawn_catalog_refresh(&self) { + tokio::spawn(async { + if let Err(e) = ModelCatalog::refresh_default_cache().await { + tracing::debug!(error = %e, "models.dev refresh failed; using cached/baked metadata"); + } + }); + } + pub fn bus(&self) -> EventBus { self.inner.bus.clone() } @@ -239,6 +271,7 @@ impl EngineHandle { temperature: None, max_steps: DEFAULT_MAX_STEPS, instructions: self.inner.config.instructions.clone(), + cost: Some(self.inner.catalog.cost(provider_id, model_id)), }; // Reserve the run slot *before* spawning. If we inserted after spawning, a run that diff --git a/crates/harness-core/src/engine/loop.rs b/crates/harness-core/src/engine/loop.rs index e495bba..2b8423a 100644 --- a/crates/harness-core/src/engine/loop.rs +++ b/crates/harness-core/src/engine/loop.rs @@ -20,6 +20,33 @@ pub struct RunConfig { pub temperature: Option, pub max_steps: u32, pub instructions: Vec, + /// Pricing for `model`, from models.dev metadata. `None` leaves cost at 0. + pub cost: Option, +} + +/// Adds a step's usage/cost onto the persisted session and republishes it. Cost accounting is +/// best-effort: a store error here is logged, not surfaced as a run failure. +async fn accumulate_session_usage( + ctx: &StepContext, + usage: &crate::types::TokenUsage, + cost: f64, + now: i64, +) { + match ctx.store.session(ctx.session_id.clone()).await { + Ok(Some(mut session)) => { + session.usage.add(usage); + session.cost += cost; + session.updated_at = now; + if let Err(e) = ctx.store.upsert_session(session.clone()).await { + tracing::warn!(error = %e, "failed to persist session usage"); + return; + } + ctx.bus + .publish(crate::event::AppEvent::SessionUpdated { session }); + } + Ok(None) => {} + Err(e) => tracing::warn!(error = %e, "failed to load session for usage accounting"), + } } /// opencode's exit condition: keep stepping while the last assistant turn asked for more @@ -221,17 +248,24 @@ pub async fn run_session( &ctx, run_config.model.clone(), &run_config.agent_name, + run_config.cost, &mut doomloop, ) .await; match step { - Ok(outcome) if outcome.aborted => return RunOutcome::Aborted, - Ok(outcome) => match outcome.result { - StepResult::Continue => continue, - StepResult::Stop => return RunOutcome::Stopped, - StepResult::Compact => return RunOutcome::Stopped, // stub until M6 - }, + Ok(outcome) if outcome.aborted => { + accumulate_session_usage(&ctx, &outcome.usage, outcome.cost, now_fn()).await; + return RunOutcome::Aborted; + } + Ok(outcome) => { + accumulate_session_usage(&ctx, &outcome.usage, outcome.cost, now_fn()).await; + match outcome.result { + StepResult::Continue => continue, + StepResult::Stop => return RunOutcome::Stopped, + StepResult::Compact => return RunOutcome::Stopped, // stub until M6 + } + } Err(step_err) if matches!(step_err.source, ProviderError::Cancelled) => { return RunOutcome::Aborted; } @@ -422,11 +456,28 @@ mod tests { temperature: None, max_steps: 10, instructions: Vec::new(), + // $3/1M input, $15/1M output. + cost: Some(crate::types::ModelCost { + input: 3.0, + output: 15.0, + ..Default::default() + }), }; let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await; assert!(matches!(outcome, RunOutcome::Stopped)); + // Both steps' usage and cost accumulate onto the session. + let session = store.session(session_id.clone()).await.unwrap().unwrap(); + assert_eq!(session.usage.input, 30); + assert_eq!(session.usage.output, 13); + // (10*3 + 5*15)/1e6 + (20*3 + 8*15)/1e6 = 0.000105 + 0.00018 + assert!( + (session.cost - 0.000_285).abs() < 1e-9, + "cost = {}", + session.cost + ); + let messages = store.messages(session_id.clone()).await.unwrap(); assert_eq!( messages.len(), @@ -517,6 +568,7 @@ mod tests { temperature: None, max_steps: 10, instructions: Vec::new(), + cost: None, }; let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await; diff --git a/crates/harness-core/src/engine/processor.rs b/crates/harness-core/src/engine/processor.rs index a4ff597..bdc2839 100644 --- a/crates/harness-core/src/engine/processor.rs +++ b/crates/harness-core/src/engine/processor.rs @@ -26,6 +26,8 @@ pub struct StepOutcome { pub result: StepResult, pub message_id: Option, pub usage: TokenUsage, + /// Dollar cost of this step's usage (0.0 when no pricing is available). + pub cost: f64, pub aborted: bool, } @@ -513,6 +515,7 @@ impl<'a> Run<'a> { &mut self, reason: FinishReason, usage: TokenUsage, + cost: f64, ) -> Result { let part = Part { id: PartId::new(), @@ -521,7 +524,7 @@ impl<'a> Run<'a> { idx: self.next_idx, body: PartBody::StepFinish { usage, - cost: 0.0, + cost, reason: reason.clone(), }, }; @@ -567,10 +570,12 @@ pub async fn process_step( ctx: &StepContext, model: crate::types::ModelRef, agent: &str, + cost: Option, doomloop: &mut DoomLoopGuard, ) -> Result { let mut run = Run::new(ctx); let mut usage = TokenUsage::default(); + let mut step_cost = 0.0; let mut result = StepResult::Stop; loop { @@ -582,6 +587,7 @@ pub async fn process_step( result: StepResult::Stop, message_id: run.assistant.as_ref().map(|m| m.id.clone()), usage, + cost: step_cost, aborted: true, }); } @@ -630,7 +636,8 @@ pub async fn process_step( usage: finish_usage, } => { usage = finish_usage; - match run.on_finish(reason, finish_usage).await { + step_cost = cost.map(|c| c.cost_of(&finish_usage)).unwrap_or(0.0); + match run.on_finish(reason, finish_usage, step_cost).await { Ok(r) => { result = r; Ok(()) @@ -652,6 +659,7 @@ pub async fn process_step( result, message_id: run.assistant.as_ref().map(|m| m.id.clone()), usage, + cost: step_cost, aborted: false, }) } diff --git a/crates/harness-core/src/store/actor.rs b/crates/harness-core/src/store/actor.rs index 098c550..05671de 100644 --- a/crates/harness-core/src/store/actor.rs +++ b/crates/harness-core/src/store/actor.rs @@ -13,6 +13,7 @@ pub enum StoreCmd { UpsertSession(Session, Reply<()>), UpsertMessage(Message, Reply<()>), UpsertPart(Part, Reply<()>), + Session(SessionId, Reply>), Sessions(Reply>), Messages(SessionId, Reply>), Parts(MessageId, Reply>), @@ -69,6 +70,15 @@ fn upsert_part(conn: &Connection, part: &Part) -> Result<(), StoreError> { Ok(()) } +fn get_session(conn: &Connection, id: &SessionId) -> Result, StoreError> { + let mut stmt = conn.prepare("SELECT data FROM session WHERE id = ?1")?; + let mut rows = stmt.query_map(params![id.as_ref()], |row| row.get::<_, String>(0))?; + match rows.next() { + Some(data) => Ok(Some(serde_json::from_str(&data?)?)), + None => Ok(None), + } +} + fn list_sessions(conn: &Connection) -> Result, StoreError> { let mut stmt = conn.prepare("SELECT data FROM session ORDER BY id")?; let rows = stmt @@ -116,6 +126,9 @@ pub fn run(conn: Connection, mut rx: tokio::sync::mpsc::Receiver) { StoreCmd::UpsertPart(part, reply) => { let _ = reply.send(upsert_part(&conn, &part)); } + StoreCmd::Session(id, reply) => { + let _ = reply.send(get_session(&conn, &id)); + } StoreCmd::Sessions(reply) => { let _ = reply.send(list_sessions(&conn)); } diff --git a/crates/harness-core/src/store/api.rs b/crates/harness-core/src/store/api.rs index 0a6b29f..3b00129 100644 --- a/crates/harness-core/src/store/api.rs +++ b/crates/harness-core/src/store/api.rs @@ -77,6 +77,11 @@ impl Store { self.call(|reply| StoreCmd::UpsertPart(part, reply)).await } + pub async fn session(&self, session_id: SessionId) -> Result, StoreError> { + self.call(|reply| StoreCmd::Session(session_id, reply)) + .await + } + pub async fn sessions(&self) -> Result, StoreError> { self.call(StoreCmd::Sessions).await } diff --git a/crates/harness-core/src/types/mod.rs b/crates/harness-core/src/types/mod.rs index ed9e940..24c4b6d 100644 --- a/crates/harness-core/src/types/mod.rs +++ b/crates/harness-core/src/types/mod.rs @@ -6,6 +6,6 @@ pub mod session; pub use ids::{MessageId, PartId, SessionId}; pub use message::{Message, MessageError, Role}; -pub use model::{ModelInfo, ModelRef, TokenUsage}; +pub use model::{ModelCost, ModelInfo, ModelRef, TokenUsage}; pub use part::{Part, PartBody, ToolState}; pub use session::Session; diff --git a/crates/harness-core/src/types/model.rs b/crates/harness-core/src/types/model.rs index ac2c491..675dac6 100644 --- a/crates/harness-core/src/types/model.rs +++ b/crates/harness-core/src/types/model.rs @@ -34,11 +34,71 @@ impl TokenUsage { } } -// Full cost/context-limit metadata is populated by harness-providers (models.dev) in M1/M3; -// this placeholder only carries what harness-core needs to key on. +/// Per-model pricing in USD per **one million** tokens. Populated from models.dev metadata +/// by `harness-providers`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)] +pub struct ModelCost { + pub input: f64, + pub output: f64, + pub cache_read: f64, + pub cache_write: f64, +} + +impl ModelCost { + /// Dollar cost of a usage sample. Reasoning tokens are billed within `output` by the + /// providers we support (OpenAI reports them as a subset of `output_tokens`; Anthropic + /// counts thinking in output), so they are intentionally not charged separately here. + pub fn cost_of(&self, usage: &TokenUsage) -> f64 { + let per_million = |tokens: u64, rate: f64| (tokens as f64) * rate / 1_000_000.0; + per_million(usage.input, self.input) + + per_million(usage.output, self.output) + + per_million(usage.cache_read, self.cache_read) + + per_million(usage.cache_write, self.cache_write) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cost_of_sums_components_per_million_excluding_reasoning() { + let cost = ModelCost { + input: 3.0, + output: 15.0, + cache_read: 0.30, + cache_write: 3.75, + }; + let usage = TokenUsage { + input: 1_000_000, + output: 1_000_000, + reasoning: 500_000, // billed within output — must not add extra cost + cache_read: 1_000_000, + cache_write: 1_000_000, + }; + // 3 + 15 + 0.30 + 3.75, with reasoning contributing nothing. + assert!((cost.cost_of(&usage) - 22.05).abs() < 1e-9); + } + + #[test] + fn default_cost_is_zero() { + assert_eq!(ModelCost::default().cost_of(&TokenUsage::default()), 0.0); + } +} + +/// Model metadata, keyed on `model`. Cost/limits are populated from models.dev by +/// `harness-providers`; `reasoning`/`tool_call`/`attachment` are capability flags. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ModelInfo { pub model: ModelRef, pub context_limit: u64, pub output_limit: u64, + #[serde(default)] + pub cost: ModelCost, + #[serde(default)] + pub reasoning: bool, + #[serde(default)] + pub tool_call: bool, + #[serde(default)] + pub attachment: bool, } diff --git a/crates/harness-providers/Cargo.toml b/crates/harness-providers/Cargo.toml index 9972f4c..ea6b8a7 100644 --- a/crates/harness-providers/Cargo.toml +++ b/crates/harness-providers/Cargo.toml @@ -18,9 +18,11 @@ serde = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } tracing = { workspace = true } +dirs = { workspace = true } [dev-dependencies] tokio = { workspace = true, features = ["test-util", "macros"] } +tempfile = { workspace = true } [lints] workspace = true diff --git a/crates/harness-providers/assets/models-snapshot.json b/crates/harness-providers/assets/models-snapshot.json new file mode 100644 index 0000000..932efa7 --- /dev/null +++ b/crates/harness-providers/assets/models-snapshot.json @@ -0,0 +1,67 @@ +{ + "anthropic": { + "id": "anthropic", + "name": "Anthropic", + "models": { + "claude-sonnet-4-5": { + "id": "claude-sonnet-4-5", + "reasoning": true, + "tool_call": true, + "attachment": true, + "cost": { "input": 3, "output": 15, "cache_read": 0.3, "cache_write": 3.75 }, + "limit": { "context": 200000, "output": 64000 } + }, + "claude-opus-4-1": { + "id": "claude-opus-4-1", + "reasoning": true, + "tool_call": true, + "attachment": true, + "cost": { "input": 15, "output": 75, "cache_read": 1.5, "cache_write": 18.75 }, + "limit": { "context": 200000, "output": 32000 } + }, + "claude-haiku-4-5": { + "id": "claude-haiku-4-5", + "reasoning": true, + "tool_call": true, + "attachment": true, + "cost": { "input": 1, "output": 5, "cache_read": 0.1, "cache_write": 1.25 }, + "limit": { "context": 200000, "output": 64000 } + } + } + }, + "openai": { + "id": "openai", + "name": "OpenAI", + "models": { + "gpt-4o": { + "id": "gpt-4o", + "tool_call": true, + "attachment": true, + "cost": { "input": 2.5, "output": 10, "cache_read": 1.25 }, + "limit": { "context": 128000, "output": 16384 } + }, + "gpt-4o-mini": { + "id": "gpt-4o-mini", + "tool_call": true, + "attachment": true, + "cost": { "input": 0.15, "output": 0.6, "cache_read": 0.075 }, + "limit": { "context": 128000, "output": 16384 } + }, + "gpt-5": { + "id": "gpt-5", + "reasoning": true, + "tool_call": true, + "attachment": true, + "cost": { "input": 1.25, "output": 10, "cache_read": 0.125 }, + "limit": { "context": 400000, "output": 128000 } + }, + "o3": { + "id": "o3", + "reasoning": true, + "tool_call": true, + "cost": { "input": 2, "output": 8, "cache_read": 0.5 }, + "limit": { "context": 200000, "output": 100000 } + } + } + } +} diff --git a/crates/harness-providers/src/auth.rs b/crates/harness-providers/src/auth.rs new file mode 100644 index 0000000..4340081 --- /dev/null +++ b/crates/harness-providers/src/auth.rs @@ -0,0 +1,228 @@ +//! Credential storage for providers that authenticate outside the config file. +//! +//! A single JSON file (`~/.local/share/ai-harness/auth.json`, mode `0600`) keyed by provider +//! id. OAuth providers (Copilot) store the token triple and refresh it in place; API-key +//! providers store the bare key. Config-supplied keys take precedence over this file — this +//! is only for interactively-obtained credentials. + +use std::collections::HashMap; +use std::io; +use std::path::{Path, PathBuf}; + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AuthRecord { + /// OAuth credentials. `expires` is a unix-ms timestamp; `0` means the token never expires + /// (opencode's direct-Bearer Copilot mode). + OAuth { + access: String, + refresh: String, + expires: i64, + }, + Api { + key: String, + }, +} + +impl AuthRecord { + /// True when an OAuth token is at or past `expires` (with a safety skew), given `now_ms`. + /// API keys and never-expiring OAuth tokens (`expires == 0`) are never considered expired. + pub fn is_expired(&self, now_ms: i64, skew_ms: i64) -> bool { + match self { + AuthRecord::OAuth { expires, .. } if *expires > 0 => now_ms + skew_ms >= *expires, + _ => false, + } + } +} + +/// Read/modify/write access to the auth file. Cheap to construct; every operation reloads so +/// concurrent writers (a refresh in one provider, a login in another) don't clobber each other. +#[derive(Debug, Clone)] +pub struct AuthStorage { + path: PathBuf, +} + +impl AuthStorage { + pub fn new(path: PathBuf) -> Self { + Self { path } + } + + /// `~/.local/share/ai-harness/auth.json` (falling back to the temp dir if there is no + /// data directory). + pub fn default_path() -> PathBuf { + dirs::data_dir() + .unwrap_or_else(std::env::temp_dir) + .join("ai-harness") + .join("auth.json") + } + + pub fn with_default_path() -> Self { + Self::new(Self::default_path()) + } + + /// Loads all records. A missing file yields an empty map; a corrupt file is an error. + pub fn load(&self) -> io::Result> { + match std::fs::read(&self.path) { + Ok(bytes) => serde_json::from_slice(&bytes) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)), + Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(HashMap::new()), + Err(e) => Err(e), + } + } + + pub fn get(&self, provider: &str) -> io::Result> { + Ok(self.load()?.remove(provider)) + } + + pub fn set(&self, provider: &str, record: AuthRecord) -> io::Result<()> { + let mut records = self.load()?; + records.insert(provider.to_string(), record); + self.write(&records) + } + + pub fn remove(&self, provider: &str) -> io::Result<()> { + let mut records = self.load()?; + if records.remove(provider).is_some() { + self.write(&records)?; + } + Ok(()) + } + + /// Serializes `records` and writes them with `0600` permissions via a temp-file rename so + /// a partial write can never leave a truncated auth file behind. + fn write(&self, records: &HashMap) -> io::Result<()> { + if let Some(parent) = self.path.parent() { + std::fs::create_dir_all(parent)?; + } + let json = serde_json::to_vec_pretty(records) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + + let tmp = self.path.with_extension("json.tmp"); + std::fs::write(&tmp, &json)?; + set_owner_only(&tmp)?; + std::fs::rename(&tmp, &self.path)?; + Ok(()) + } +} + +#[cfg(unix)] +fn set_owner_only(path: &Path) -> io::Result<()> { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)) +} + +#[cfg(not(unix))] +fn set_owner_only(_path: &Path) -> io::Result<()> { + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn storage() -> (tempfile::TempDir, AuthStorage) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("nested").join("auth.json"); + (dir, AuthStorage::new(path)) + } + + #[test] + fn missing_file_loads_empty() { + let (_dir, store) = storage(); + assert!(store.load().unwrap().is_empty()); + assert_eq!(store.get("anthropic").unwrap(), None); + } + + #[test] + fn set_get_roundtrip_and_overwrite() { + let (_dir, store) = storage(); + store + .set("anthropic", AuthRecord::Api { key: "k1".into() }) + .unwrap(); + assert_eq!( + store.get("anthropic").unwrap(), + Some(AuthRecord::Api { key: "k1".into() }) + ); + + store + .set("anthropic", AuthRecord::Api { key: "k2".into() }) + .unwrap(); + assert_eq!( + store.get("anthropic").unwrap(), + Some(AuthRecord::Api { key: "k2".into() }) + ); + } + + #[test] + fn multiple_providers_coexist() { + let (_dir, store) = storage(); + store + .set("openai", AuthRecord::Api { key: "sk".into() }) + .unwrap(); + store + .set( + "github-copilot", + AuthRecord::OAuth { + access: "a".into(), + refresh: "r".into(), + expires: 0, + }, + ) + .unwrap(); + let all = store.load().unwrap(); + assert_eq!(all.len(), 2); + assert!(all.contains_key("openai")); + assert!(all.contains_key("github-copilot")); + } + + #[test] + fn remove_deletes_only_named_provider() { + let (_dir, store) = storage(); + store + .set("openai", AuthRecord::Api { key: "sk".into() }) + .unwrap(); + store + .set("anthropic", AuthRecord::Api { key: "an".into() }) + .unwrap(); + store.remove("openai").unwrap(); + assert_eq!(store.get("openai").unwrap(), None); + assert!(store.get("anthropic").unwrap().is_some()); + // Removing an absent provider is a no-op, not an error. + store.remove("openai").unwrap(); + } + + #[test] + fn oauth_expiry_respects_skew_and_never_expires() { + let never = AuthRecord::OAuth { + access: "a".into(), + refresh: "r".into(), + expires: 0, + }; + assert!(!never.is_expired(i64::MAX, 0)); + + let expiring = AuthRecord::OAuth { + access: "a".into(), + refresh: "r".into(), + expires: 1_000, + }; + assert!(!expiring.is_expired(800, 120)); + assert!(expiring.is_expired(880, 120)); // within skew window + assert!(expiring.is_expired(1_000, 0)); + + assert!(!AuthRecord::Api { key: "k".into() }.is_expired(i64::MAX, 0)); + } + + #[cfg(unix)] + #[test] + fn file_is_written_owner_only() { + use std::os::unix::fs::PermissionsExt; + let (_dir, store) = storage(); + store + .set("openai", AuthRecord::Api { key: "sk".into() }) + .unwrap(); + let mode = std::fs::metadata(&store.path).unwrap().permissions().mode(); + assert_eq!(mode & 0o777, 0o600); + } +} diff --git a/crates/harness-providers/src/codec/mod.rs b/crates/harness-providers/src/codec/mod.rs index e529997..daffdf8 100644 --- a/crates/harness-providers/src/codec/mod.rs +++ b/crates/harness-providers/src/codec/mod.rs @@ -1 +1,3 @@ pub mod anthropic; +pub mod openai_chat; +pub mod openai_responses; diff --git a/crates/harness-providers/src/codec/openai_chat.rs b/crates/harness-providers/src/codec/openai_chat.rs new file mode 100644 index 0000000..a9310eb --- /dev/null +++ b/crates/harness-providers/src/codec/openai_chat.rs @@ -0,0 +1,502 @@ +//! Request builder + SSE decoder for OpenAI's `/chat/completions` streaming API. +//! +//! Also used as the fallback codec for OpenAI-compatible endpoints (including Copilot models +//! whose `supported_endpoints` list `/chat/completions`). + +use std::collections::HashMap; + +use async_stream::try_stream; +use eventsource_stream::Eventsource; +use futures::Stream; +use harness_core::llm::{ + FinishReason, LlmEvent, LlmEventStream, LlmRequest, ProviderError, ReasoningEffort, Role, + WireContent, +}; +use harness_core::types::TokenUsage; +use serde_json::{json, Value}; + +fn effort_str(effort: ReasoningEffort) -> &'static str { + match effort { + ReasoningEffort::Low => "low", + ReasoningEffort::Medium => "medium", + ReasoningEffort::High => "high", + } +} + +fn role_str(role: Role) -> &'static str { + match role { + Role::System => "system", + Role::User => "user", + Role::Assistant => "assistant", + Role::Tool => "tool", + } +} + +/// Flattens our grouped `WireMessage`s into the flat chat-completions message list. Tool +/// results become their own `role: "tool"` messages (one per result), which is what the API +/// expects regardless of how the engine grouped them. +fn build_messages(system: &[String], messages: &[harness_core::llm::WireMessage]) -> Vec { + let mut out: Vec = Vec::new(); + + if !system.is_empty() { + out.push(json!({"role": "system", "content": system.join("\n\n")})); + } + + for m in messages { + // Tool results are always emitted as standalone `tool` messages. + for c in &m.content { + if let WireContent::ToolResult { + call_id, output, .. + } = c + { + out.push(json!({ + "role": "tool", + "tool_call_id": call_id, + "content": output, + })); + } + } + + let mut text = String::new(); + let mut images: Vec = Vec::new(); + let mut tool_calls: Vec = Vec::new(); + for c in &m.content { + match c { + WireContent::Text { text: t } => { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(t); + } + WireContent::ToolCall { + call_id, + name, + input, + } => tool_calls.push(json!({ + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": input.to_string()}, + })), + WireContent::Image { mime_type, data } => images.push(json!({ + "type": "image_url", + "image_url": {"url": format!("data:{mime_type};base64,{data}")}, + })), + WireContent::ToolResult { .. } => {} // handled above + } + } + + // A message that carried nothing but tool results contributes no further entry. + if text.is_empty() && images.is_empty() && tool_calls.is_empty() { + continue; + } + + let mut msg = json!({"role": role_str(m.role)}); + if images.is_empty() { + msg["content"] = json!(text); + } else { + let mut parts = vec![json!({"type": "text", "text": text})]; + parts.extend(images); + msg["content"] = json!(parts); + } + if !tool_calls.is_empty() { + msg["tool_calls"] = json!(tool_calls); + } + out.push(msg); + } + + out +} + +pub fn build_request(req: &LlmRequest) -> Value { + let mut body = json!({ + "model": req.model, + "messages": build_messages(&req.system, &req.messages), + "stream": true, + "stream_options": {"include_usage": true}, + }); + + if !req.tools.is_empty() { + body["tools"] = Value::Array( + req.tools + .iter() + .map(|t| { + json!({ + "type": "function", + "function": { + "name": t.name, + "description": t.description, + "parameters": t.parameters, + }, + }) + }) + .collect(), + ); + } + if let Some(temp) = req.temperature { + body["temperature"] = json!(temp); + } + if let Some(max) = req.max_tokens { + body["max_completion_tokens"] = json!(max); + } + if let Some(effort) = req.reasoning.as_ref().and_then(|r| r.effort) { + body["reasoning_effort"] = json!(effort_str(effort)); + } + body +} + +fn map_finish_reason(reason: &str) -> FinishReason { + match reason { + "stop" => FinishReason::Stop, + "tool_calls" | "function_call" => FinishReason::ToolCalls, + "length" => FinishReason::Length, + "content_filter" => FinishReason::ContentFilter, + other => FinishReason::Unknown(other.to_string()), + } +} + +#[derive(Default)] +struct ToolAccum { + call_id: String, + name: String, + args: String, +} + +/// Decodes a chat-completions SSE byte stream into normalized `LlmEvent`s. Tool-call deltas +/// are accumulated by their `index` and flushed as `ToolCall`s once the stream ends. +pub fn decode(byte_stream: S) -> LlmEventStream +where + S: Stream> + Send + 'static, + E: std::error::Error + Send + Sync + 'static, +{ + let events = byte_stream.eventsource(); + Box::pin(try_stream! { + futures::pin_mut!(events); + let mut usage = TokenUsage::default(); + let mut reason = FinishReason::Stop; + let mut text_open = false; + let mut reasoning_open = false; + let mut tools: HashMap = HashMap::new(); + let mut tool_order: Vec = Vec::new(); + + while let Some(item) = futures::StreamExt::next(&mut events).await { + let event = item.map_err(|e| ProviderError::Decode(e.to_string()))?; + let data = event.data.trim(); + if data.is_empty() { + continue; + } + if data == "[DONE]" { + break; + } + let value: Value = serde_json::from_str(data) + .map_err(|e| ProviderError::Decode(format!("{e}: {data}")))?; + + if let Some(u) = value.get("usage").filter(|u| u.is_object()) { + usage.input = u["prompt_tokens"].as_u64().unwrap_or(usage.input); + usage.output = u["completion_tokens"].as_u64().unwrap_or(usage.output); + usage.reasoning = u["completion_tokens_details"]["reasoning_tokens"] + .as_u64() + .unwrap_or(usage.reasoning); + usage.cache_read = u["prompt_tokens_details"]["cached_tokens"] + .as_u64() + .unwrap_or(usage.cache_read); + } + + let choice = &value["choices"][0]; + let delta = &choice["delta"]; + + if let Some(rc) = delta["reasoning_content"].as_str().filter(|s| !s.is_empty()) { + if !reasoning_open { + reasoning_open = true; + yield LlmEvent::ReasoningStart { id: "reasoning".into() }; + } + yield LlmEvent::ReasoningDelta { id: "reasoning".into(), text: rc.to_string() }; + } + + if let Some(text) = delta["content"].as_str().filter(|s| !s.is_empty()) { + if reasoning_open { + reasoning_open = false; + yield LlmEvent::ReasoningEnd { id: "reasoning".into(), signature: None }; + } + if !text_open { + text_open = true; + yield LlmEvent::TextStart { id: "0".into() }; + } + yield LlmEvent::TextDelta { id: "0".into(), text: text.to_string() }; + } + + if let Some(calls) = delta["tool_calls"].as_array() { + for call in calls { + let index = call["index"].as_u64().unwrap_or(0); + let entry = tools.entry(index).or_insert_with(|| { + tool_order.push(index); + ToolAccum::default() + }); + if let Some(id) = call["id"].as_str().filter(|s| !s.is_empty()) { + entry.call_id = id.to_string(); + } + if let Some(name) = call["function"]["name"].as_str().filter(|s| !s.is_empty()) { + entry.name = name.to_string(); + yield LlmEvent::ToolInputStart { + call_id: entry.call_id.clone(), + name: entry.name.clone(), + }; + } + if let Some(args) = call["function"]["arguments"].as_str().filter(|s| !s.is_empty()) { + entry.args.push_str(args); + yield LlmEvent::ToolInputDelta { + call_id: entry.call_id.clone(), + json: args.to_string(), + }; + } + } + } + + if let Some(fr) = choice["finish_reason"].as_str() { + reason = map_finish_reason(fr); + } + } + + if reasoning_open { + yield LlmEvent::ReasoningEnd { id: "reasoning".into(), signature: None }; + } + if text_open { + yield LlmEvent::TextEnd { id: "0".into() }; + } + for index in tool_order { + if let Some(entry) = tools.remove(&index) { + let input: Value = if entry.args.trim().is_empty() { + json!({}) + } else { + serde_json::from_str(&entry.args).unwrap_or(Value::Null) + }; + yield LlmEvent::ToolCall { call_id: entry.call_id, name: entry.name, input }; + } + } + yield LlmEvent::Finish { reason, usage }; + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt; + use harness_core::llm::{Initiator, ReasoningOpts, ToolSchema, WireMessage}; + + fn sse_stream(raw: &'static str) -> LlmEventStream { + let chunks: Vec> = + vec![Ok(bytes::Bytes::from_static(raw.as_bytes()))]; + decode(futures::stream::iter(chunks)) + } + + fn req(model: &str) -> LlmRequest { + LlmRequest { + model: model.into(), + system: vec![], + messages: vec![], + tools: vec![], + temperature: None, + max_tokens: None, + reasoning: None, + initiator: Initiator::User, + } + } + + #[tokio::test] + async fn decodes_text_only_response() { + let raw = concat!( + "data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"},\"finish_reason\":null}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"lo\"},\"finish_reason\":null}]}\n\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5}}\n\n", + "data: [DONE]\n\n", + ); + let events: Vec = sse_stream(raw).map(|e| e.unwrap()).collect().await; + assert_eq!( + events, + vec![ + LlmEvent::TextStart { id: "0".into() }, + LlmEvent::TextDelta { + id: "0".into(), + text: "Hel".into() + }, + LlmEvent::TextDelta { + id: "0".into(), + text: "lo".into() + }, + LlmEvent::TextEnd { id: "0".into() }, + LlmEvent::Finish { + reason: FinishReason::Stop, + usage: TokenUsage { + input: 10, + output: 5, + ..Default::default() + }, + }, + ] + ); + } + + #[tokio::test] + async fn decodes_tool_call_accumulated_by_index() { + let raw = concat!( + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"read\",\"arguments\":\"\"}}]}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"file\\\"\"}}]}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\":\\\"a.txt\\\"}\"}}]}}]}\n\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":8}}\n\n", + "data: [DONE]\n\n", + ); + let events: Vec = sse_stream(raw).map(|e| e.unwrap()).collect().await; + assert_eq!( + events, + vec![ + LlmEvent::ToolInputStart { + call_id: "call_1".into(), + name: "read".into() + }, + LlmEvent::ToolInputDelta { + call_id: "call_1".into(), + json: "{\"file\"".into() + }, + LlmEvent::ToolInputDelta { + call_id: "call_1".into(), + json: ":\"a.txt\"}".into() + }, + LlmEvent::ToolCall { + call_id: "call_1".into(), + name: "read".into(), + input: json!({"file": "a.txt"}), + }, + LlmEvent::Finish { + reason: FinishReason::ToolCalls, + usage: TokenUsage { + input: 1, + output: 8, + ..Default::default() + }, + }, + ] + ); + } + + #[tokio::test] + async fn decodes_reasoning_content_before_text() { + let raw = concat!( + "data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"hmm\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n", + ); + let events: Vec = sse_stream(raw).map(|e| e.unwrap()).collect().await; + assert_eq!( + events, + vec![ + LlmEvent::ReasoningStart { + id: "reasoning".into() + }, + LlmEvent::ReasoningDelta { + id: "reasoning".into(), + text: "hmm".into() + }, + LlmEvent::ReasoningEnd { + id: "reasoning".into(), + signature: None + }, + LlmEvent::TextStart { id: "0".into() }, + LlmEvent::TextDelta { + id: "0".into(), + text: "answer".into() + }, + LlmEvent::TextEnd { id: "0".into() }, + LlmEvent::Finish { + reason: FinishReason::Stop, + usage: TokenUsage::default() + }, + ] + ); + } + + #[test] + fn build_request_wraps_tools_in_function_envelope() { + let mut r = req("gpt-4o"); + r.tools = vec![ToolSchema { + name: "read".into(), + description: "reads a file".into(), + parameters: json!({"type": "object"}), + }]; + let body = build_request(&r); + assert_eq!(body["tools"][0]["type"], "function"); + assert_eq!(body["tools"][0]["function"]["name"], "read"); + assert_eq!( + body["tools"][0]["function"]["parameters"], + json!({"type": "object"}) + ); + assert_eq!(body["stream_options"]["include_usage"], true); + } + + #[test] + fn build_request_prepends_system_message() { + let mut r = req("gpt-4o"); + r.system = vec!["env".into(), "agent".into()]; + r.messages = vec![WireMessage { + role: Role::User, + content: vec![WireContent::Text { text: "hi".into() }], + }]; + let body = build_request(&r); + let msgs = body["messages"].as_array().unwrap(); + assert_eq!(msgs[0]["role"], "system"); + assert_eq!(msgs[0]["content"], "env\n\nagent"); + assert_eq!(msgs[1]["role"], "user"); + assert_eq!(msgs[1]["content"], "hi"); + } + + #[test] + fn build_request_expands_tool_results_to_tool_messages() { + let mut r = req("gpt-4o"); + r.messages = vec![WireMessage { + role: Role::Tool, + content: vec![WireContent::ToolResult { + call_id: "call_1".into(), + output: "contents".into(), + is_error: false, + }], + }]; + let body = build_request(&r); + let msgs = body["messages"].as_array().unwrap(); + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0]["role"], "tool"); + assert_eq!(msgs[0]["tool_call_id"], "call_1"); + assert_eq!(msgs[0]["content"], "contents"); + } + + #[test] + fn build_request_emits_assistant_tool_calls() { + let mut r = req("gpt-4o"); + r.messages = vec![WireMessage { + role: Role::Assistant, + content: vec![WireContent::ToolCall { + call_id: "call_1".into(), + name: "read".into(), + input: json!({"file": "a.txt"}), + }], + }]; + let body = build_request(&r); + let msgs = body["messages"].as_array().unwrap(); + assert_eq!(msgs[0]["tool_calls"][0]["id"], "call_1"); + assert_eq!(msgs[0]["tool_calls"][0]["function"]["name"], "read"); + assert_eq!( + msgs[0]["tool_calls"][0]["function"]["arguments"], + "{\"file\":\"a.txt\"}" + ); + } + + #[test] + fn build_request_sets_reasoning_effort() { + let mut r = req("o3"); + r.reasoning = Some(ReasoningOpts { + effort: Some(ReasoningEffort::High), + budget_tokens: None, + }); + let body = build_request(&r); + assert_eq!(body["reasoning_effort"], "high"); + } +} diff --git a/crates/harness-providers/src/codec/openai_responses.rs b/crates/harness-providers/src/codec/openai_responses.rs new file mode 100644 index 0000000..fc5de23 --- /dev/null +++ b/crates/harness-providers/src/codec/openai_responses.rs @@ -0,0 +1,402 @@ +//! Request builder + SSE decoder for OpenAI's `/responses` streaming API. +//! +//! The responses API is the preferred surface for `gpt-*` / `o-*` models: it carries +//! reasoning items natively and reports reasoning-token usage separately. + +use std::collections::HashMap; + +use async_stream::try_stream; +use eventsource_stream::Eventsource; +use futures::Stream; +use harness_core::llm::{ + FinishReason, LlmEvent, LlmEventStream, LlmRequest, ProviderError, ReasoningEffort, Role, + WireContent, +}; +use harness_core::types::TokenUsage; +use serde_json::{json, Value}; + +fn effort_str(effort: ReasoningEffort) -> &'static str { + match effort { + ReasoningEffort::Low => "low", + ReasoningEffort::Medium => "medium", + ReasoningEffort::High => "high", + } +} + +/// Builds the `input` item list. Text/image content become `message` items; tool calls and +/// their results become `function_call` / `function_call_output` items (the responses API +/// keeps these as top-level items rather than nesting them inside messages). +fn build_input(messages: &[harness_core::llm::WireMessage]) -> Vec { + let mut out: Vec = Vec::new(); + for m in messages { + let (role, text_type) = match m.role { + Role::Assistant => ("assistant", "output_text"), + _ => ("user", "input_text"), + }; + let mut content_parts: Vec = Vec::new(); + for c in &m.content { + match c { + WireContent::Text { text } => { + content_parts.push(json!({"type": text_type, "text": text})); + } + WireContent::Image { mime_type, data } => { + content_parts.push(json!({ + "type": "input_image", + "image_url": format!("data:{mime_type};base64,{data}"), + })); + } + WireContent::ToolCall { + call_id, + name, + input, + } => out.push(json!({ + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": input.to_string(), + })), + WireContent::ToolResult { + call_id, output, .. + } => out.push(json!({ + "type": "function_call_output", + "call_id": call_id, + "output": output, + })), + } + } + if !content_parts.is_empty() { + out.push(json!({"type": "message", "role": role, "content": content_parts})); + } + } + out +} + +pub fn build_request(req: &LlmRequest) -> Value { + let mut body = json!({ + "model": req.model, + "input": build_input(&req.messages), + "stream": true, + "store": false, + }); + + if !req.system.is_empty() { + body["instructions"] = json!(req.system.join("\n\n")); + } + if !req.tools.is_empty() { + body["tools"] = Value::Array( + req.tools + .iter() + .map(|t| { + json!({ + "type": "function", + "name": t.name, + "description": t.description, + "parameters": t.parameters, + }) + }) + .collect(), + ); + } + if let Some(temp) = req.temperature { + body["temperature"] = json!(temp); + } + if let Some(max) = req.max_tokens { + body["max_output_tokens"] = json!(max); + } + if let Some(effort) = req.reasoning.as_ref().and_then(|r| r.effort) { + body["reasoning"] = json!({"effort": effort_str(effort), "summary": "auto"}); + } + body +} + +fn map_status(status: Option<&str>, had_tool_call: bool) -> FinishReason { + match status { + Some("completed") if had_tool_call => FinishReason::ToolCalls, + Some("completed") => FinishReason::Stop, + Some("incomplete") => FinishReason::Length, + Some(other) => FinishReason::Unknown(other.to_string()), + None if had_tool_call => FinishReason::ToolCalls, + None => FinishReason::Stop, + } +} + +/// Tracks which normalized stream an `item_id` belongs to so deltas route correctly. +enum ItemKind { + Text, + Reasoning, + FunctionCall { call_id: String, name: String }, +} + +/// Decodes a `/responses` SSE byte stream into normalized `LlmEvent`s. +pub fn decode(byte_stream: S) -> LlmEventStream +where + S: Stream> + Send + 'static, + E: std::error::Error + Send + Sync + 'static, +{ + let events = byte_stream.eventsource(); + Box::pin(try_stream! { + futures::pin_mut!(events); + let mut items: HashMap = HashMap::new(); + let mut fn_args: HashMap = HashMap::new(); + let mut usage = TokenUsage::default(); + let mut had_tool_call = false; + + while let Some(item) = futures::StreamExt::next(&mut events).await { + let event = item.map_err(|e| ProviderError::Decode(e.to_string()))?; + if event.data.trim().is_empty() { + continue; + } + let value: Value = serde_json::from_str(&event.data) + .map_err(|e| ProviderError::Decode(format!("{e}: {}", event.data)))?; + let kind = value["type"].as_str().unwrap_or_default(); + + match kind { + "response.output_item.added" => { + let item = &value["item"]; + let id = item["id"].as_str().unwrap_or_default().to_string(); + match item["type"].as_str().unwrap_or_default() { + "message" => { + items.insert(id.clone(), ItemKind::Text); + yield LlmEvent::TextStart { id }; + } + "reasoning" => { + items.insert(id.clone(), ItemKind::Reasoning); + yield LlmEvent::ReasoningStart { id }; + } + "function_call" => { + let call_id = item["call_id"].as_str().unwrap_or_default().to_string(); + let name = item["name"].as_str().unwrap_or_default().to_string(); + fn_args.insert(id.clone(), String::new()); + items.insert(id, ItemKind::FunctionCall { call_id: call_id.clone(), name: name.clone() }); + had_tool_call = true; + yield LlmEvent::ToolInputStart { call_id, name }; + } + _ => {} + } + } + "response.output_text.delta" => { + let id = value["item_id"].as_str().unwrap_or_default().to_string(); + let text = value["delta"].as_str().unwrap_or_default().to_string(); + yield LlmEvent::TextDelta { id, text }; + } + "response.reasoning_summary_text.delta" => { + let id = value["item_id"].as_str().unwrap_or_default().to_string(); + let text = value["delta"].as_str().unwrap_or_default().to_string(); + yield LlmEvent::ReasoningDelta { id, text }; + } + "response.function_call_arguments.delta" => { + let id = value["item_id"].as_str().unwrap_or_default().to_string(); + let delta = value["delta"].as_str().unwrap_or_default(); + if let (Some(buf), Some(ItemKind::FunctionCall { call_id, .. })) = + (fn_args.get_mut(&id), items.get(&id)) + { + buf.push_str(delta); + yield LlmEvent::ToolInputDelta { call_id: call_id.clone(), json: delta.to_string() }; + } + } + "response.output_item.done" => { + let id = value["item"]["id"].as_str().unwrap_or_default().to_string(); + match items.remove(&id) { + Some(ItemKind::Text) => yield LlmEvent::TextEnd { id }, + Some(ItemKind::Reasoning) => yield LlmEvent::ReasoningEnd { id, signature: None }, + Some(ItemKind::FunctionCall { call_id, name }) => { + let raw = fn_args.remove(&id).unwrap_or_default(); + let input: Value = if raw.trim().is_empty() { + json!({}) + } else { + serde_json::from_str(&raw).unwrap_or(Value::Null) + }; + yield LlmEvent::ToolCall { call_id, name, input }; + } + None => {} + } + } + "response.completed" | "response.incomplete" => { + let response = &value["response"]; + let u = &response["usage"]; + usage.input = u["input_tokens"].as_u64().unwrap_or(0); + usage.output = u["output_tokens"].as_u64().unwrap_or(0); + usage.reasoning = u["output_tokens_details"]["reasoning_tokens"].as_u64().unwrap_or(0); + usage.cache_read = u["input_tokens_details"]["cached_tokens"].as_u64().unwrap_or(0); + let status = response["status"].as_str(); + yield LlmEvent::Finish { reason: map_status(status, had_tool_call), usage }; + } + "error" | "response.failed" => { + let message = value["response"]["error"]["message"] + .as_str() + .or_else(|| value["message"].as_str()) + .unwrap_or("unknown error") + .to_string(); + Err(ProviderError::Http { status: 0, body: message })?; + } + _ => {} + } + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt; + use harness_core::llm::{Initiator, ReasoningOpts, ToolSchema, WireMessage}; + + fn sse_stream(raw: &'static str) -> LlmEventStream { + let chunks: Vec> = + vec![Ok(bytes::Bytes::from_static(raw.as_bytes()))]; + decode(futures::stream::iter(chunks)) + } + + fn req(model: &str) -> LlmRequest { + LlmRequest { + model: model.into(), + system: vec![], + messages: vec![], + tools: vec![], + temperature: None, + max_tokens: None, + reasoning: None, + initiator: Initiator::User, + } + } + + #[tokio::test] + async fn decodes_text_response() { + let raw = concat!( + "data: {\"type\":\"response.output_item.added\",\"item\":{\"id\":\"msg_1\",\"type\":\"message\"}}\n\n", + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"delta\":\"Hello\"}\n\n", + "data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"msg_1\"}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":10,\"output_tokens\":5}}}\n\n", + ); + let events: Vec = sse_stream(raw).map(|e| e.unwrap()).collect().await; + assert_eq!( + events, + vec![ + LlmEvent::TextStart { id: "msg_1".into() }, + LlmEvent::TextDelta { + id: "msg_1".into(), + text: "Hello".into() + }, + LlmEvent::TextEnd { id: "msg_1".into() }, + LlmEvent::Finish { + reason: FinishReason::Stop, + usage: TokenUsage { + input: 10, + output: 5, + ..Default::default() + }, + }, + ] + ); + } + + #[tokio::test] + async fn decodes_function_call_with_reasoning_tokens() { + let raw = concat!( + "data: {\"type\":\"response.output_item.added\",\"item\":{\"id\":\"fc_1\",\"type\":\"function_call\",\"call_id\":\"call_1\",\"name\":\"read\"}}\n\n", + "data: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"fc_1\",\"delta\":\"{\\\"file\\\":\"}\n\n", + "data: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"fc_1\",\"delta\":\"\\\"a.txt\\\"}\"}\n\n", + "data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"fc_1\"}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":3,\"output_tokens\":9,\"output_tokens_details\":{\"reasoning_tokens\":4}}}}\n\n", + ); + let events: Vec = sse_stream(raw).map(|e| e.unwrap()).collect().await; + assert_eq!( + events, + vec![ + LlmEvent::ToolInputStart { + call_id: "call_1".into(), + name: "read".into() + }, + LlmEvent::ToolInputDelta { + call_id: "call_1".into(), + json: "{\"file\":".into() + }, + LlmEvent::ToolInputDelta { + call_id: "call_1".into(), + json: "\"a.txt\"}".into() + }, + LlmEvent::ToolCall { + call_id: "call_1".into(), + name: "read".into(), + input: json!({"file": "a.txt"}), + }, + LlmEvent::Finish { + reason: FinishReason::ToolCalls, + usage: TokenUsage { + input: 3, + output: 9, + reasoning: 4, + ..Default::default() + }, + }, + ] + ); + } + + #[tokio::test] + async fn failed_response_surfaces_as_err() { + let raw = "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"boom\"}}}\n\n"; + let events: Vec> = sse_stream(raw).collect().await; + assert!( + matches!(events.last(), Some(Err(ProviderError::Http { body, .. })) if body == "boom") + ); + } + + #[test] + fn build_request_puts_system_in_instructions() { + let mut r = req("gpt-5"); + r.system = vec!["env".into(), "agent".into()]; + let body = build_request(&r); + assert_eq!(body["instructions"], "env\n\nagent"); + assert!(body.get("messages").is_none()); + } + + #[test] + fn build_request_flattens_tool_calls_and_results_to_items() { + let mut r = req("gpt-5"); + r.messages = vec![ + WireMessage { + role: Role::Assistant, + content: vec![WireContent::ToolCall { + call_id: "call_1".into(), + name: "read".into(), + input: json!({"file": "a.txt"}), + }], + }, + WireMessage { + role: Role::Tool, + content: vec![WireContent::ToolResult { + call_id: "call_1".into(), + output: "contents".into(), + is_error: false, + }], + }, + ]; + let body = build_request(&r); + let input = body["input"].as_array().unwrap(); + assert_eq!(input[0]["type"], "function_call"); + assert_eq!(input[0]["call_id"], "call_1"); + assert_eq!(input[0]["arguments"], "{\"file\":\"a.txt\"}"); + assert_eq!(input[1]["type"], "function_call_output"); + assert_eq!(input[1]["output"], "contents"); + } + + #[test] + fn build_request_uses_flat_function_tool_shape_and_reasoning() { + let mut r = req("o3"); + r.tools = vec![ToolSchema { + name: "read".into(), + description: "reads a file".into(), + parameters: json!({"type": "object"}), + }]; + r.reasoning = Some(ReasoningOpts { + effort: Some(ReasoningEffort::Medium), + budget_tokens: None, + }); + let body = build_request(&r); + assert_eq!(body["tools"][0]["type"], "function"); + assert_eq!(body["tools"][0]["name"], "read"); + assert_eq!(body["reasoning"]["effort"], "medium"); + assert_eq!(body["reasoning"]["summary"], "auto"); + } +} diff --git a/crates/harness-providers/src/copilot/device_flow.rs b/crates/harness-providers/src/copilot/device_flow.rs new file mode 100644 index 0000000..406c58e --- /dev/null +++ b/crates/harness-providers/src/copilot/device_flow.rs @@ -0,0 +1,209 @@ +//! GitHub OAuth device flow (RFC 8628) for Copilot login. +//! +//! ai-harness must register its own GitHub OAuth app and supply its client id (we do not +//! hardcode opencode's). The client id is passed in by the caller — see [`request_device_code`]. + +use std::time::Duration; + +use serde::Deserialize; +use serde_json::Value; + +const DEVICE_CODE_URL: &str = "https://github.com/login/device/code"; +const ACCESS_TOKEN_URL: &str = "https://github.com/login/oauth/access_token"; +const GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code"; +/// Minimum scope needed to call the Copilot token-exchange endpoint. +pub const SCOPE: &str = "read:user"; + +#[derive(Debug, Clone, Deserialize)] +pub struct DeviceCode { + pub device_code: String, + pub user_code: String, + pub verification_uri: String, + /// Seconds between polls; GitHub requires honoring this and any `slow_down` bumps. + #[serde(default = "default_interval")] + pub interval: u64, + #[serde(default)] + pub expires_in: u64, +} + +fn default_interval() -> u64 { + 5 +} + +/// Result of one poll of the access-token endpoint. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PollOutcome { + /// The user hasn't authorized yet; keep polling at the current interval. + Pending, + /// GitHub asked us to slow down; add 5s to the interval (RFC 8628). + SlowDown, + /// Authorization complete. + Success { access_token: String }, + /// Terminal failure (expired code, denied, unknown error) with a human-readable reason. + Failed(String), +} + +#[derive(Debug, thiserror::Error)] +pub enum DeviceFlowError { + #[error("network: {0}")] + Network(String), + #[error("unexpected response: {0}")] + Unexpected(String), +} + +/// Pure mapping of an access-token poll response body to a [`PollOutcome`], so the state +/// machine is testable without a live GitHub. +pub fn parse_poll_response(body: &Value) -> PollOutcome { + if let Some(token) = body["access_token"].as_str() { + return PollOutcome::Success { + access_token: token.to_string(), + }; + } + match body["error"].as_str() { + Some("authorization_pending") => PollOutcome::Pending, + Some("slow_down") => PollOutcome::SlowDown, + Some(other) => { + let desc = body["error_description"].as_str().unwrap_or(other); + PollOutcome::Failed(desc.to_string()) + } + None => PollOutcome::Failed("no access_token and no error in response".to_string()), + } +} + +/// Applies a `slow_down` to a poll interval per RFC 8628 (+5s). +pub fn bump_interval(interval: u64) -> u64 { + interval + 5 +} + +/// Step 1: request a device + user code for `client_id`. +pub async fn request_device_code( + client: &reqwest::Client, + client_id: &str, +) -> Result { + let resp = client + .post(DEVICE_CODE_URL) + .header("accept", "application/json") + .json(&serde_json::json!({"client_id": client_id, "scope": SCOPE})) + .send() + .await + .map_err(|e| DeviceFlowError::Network(e.to_string()))?; + let value: Value = resp + .json() + .await + .map_err(|e| DeviceFlowError::Network(e.to_string()))?; + serde_json::from_value(value.clone()) + .map_err(|_| DeviceFlowError::Unexpected(value.to_string())) +} + +/// Step 2 (single poll): exchange the device code for an access token, once. +pub async fn poll_once( + client: &reqwest::Client, + client_id: &str, + device_code: &str, +) -> Result { + let resp = client + .post(ACCESS_TOKEN_URL) + .header("accept", "application/json") + .json(&serde_json::json!({ + "client_id": client_id, + "device_code": device_code, + "grant_type": GRANT_TYPE, + })) + .send() + .await + .map_err(|e| DeviceFlowError::Network(e.to_string()))?; + let value: Value = resp + .json() + .await + .map_err(|e| DeviceFlowError::Network(e.to_string()))?; + Ok(parse_poll_response(&value)) +} + +/// Step 2 (full loop): polls until success or terminal failure, honoring `interval` and +/// `slow_down`. `sleep` is injected so tests can drive it without real time. +pub async fn poll_for_token( + client: &reqwest::Client, + client_id: &str, + device: &DeviceCode, + sleep: S, +) -> Result +where + S: Fn(Duration) -> Fut, + Fut: std::future::Future, +{ + let mut interval = device.interval; + loop { + sleep(Duration::from_secs(interval)).await; + match poll_once(client, client_id, &device.device_code).await? { + PollOutcome::Pending => {} + PollOutcome::SlowDown => interval = bump_interval(interval), + PollOutcome::Success { access_token } => return Ok(access_token), + PollOutcome::Failed(reason) => return Err(DeviceFlowError::Unexpected(reason)), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn parses_success() { + let outcome = + parse_poll_response(&json!({"access_token": "gho_abc", "token_type": "bearer"})); + assert_eq!( + outcome, + PollOutcome::Success { + access_token: "gho_abc".into() + } + ); + } + + #[test] + fn parses_pending_and_slow_down() { + assert_eq!( + parse_poll_response(&json!({"error": "authorization_pending"})), + PollOutcome::Pending + ); + assert_eq!( + parse_poll_response(&json!({"error": "slow_down", "interval": 10})), + PollOutcome::SlowDown + ); + } + + #[test] + fn parses_terminal_errors_with_description() { + match parse_poll_response( + &json!({"error": "expired_token", "error_description": "code expired"}), + ) { + PollOutcome::Failed(msg) => assert_eq!(msg, "code expired"), + other => panic!("expected Failed, got {other:?}"), + } + assert!(matches!( + parse_poll_response(&json!({"error": "access_denied"})), + PollOutcome::Failed(_) + )); + assert!(matches!( + parse_poll_response(&json!({})), + PollOutcome::Failed(_) + )); + } + + #[test] + fn slow_down_adds_five_seconds() { + assert_eq!(bump_interval(5), 10); + assert_eq!(bump_interval(10), 15); + } + + #[test] + fn device_code_defaults_interval() { + let dc: DeviceCode = serde_json::from_value(json!({ + "device_code": "d", + "user_code": "WXYZ-1234", + "verification_uri": "https://github.com/login/device" + })) + .unwrap(); + assert_eq!(dc.interval, 5); + } +} diff --git a/crates/harness-providers/src/copilot/mod.rs b/crates/harness-providers/src/copilot/mod.rs new file mode 100644 index 0000000..7416c0a --- /dev/null +++ b/crates/harness-providers/src/copilot/mod.rs @@ -0,0 +1,10 @@ +//! GitHub Copilot: OAuth device-flow login, token exchange, and a provider that multiplexes +//! the chat/responses/anthropic codecs behind `api.githubcopilot.com`. + +pub mod device_flow; +pub mod provider; +pub mod token; + +pub use device_flow::{DeviceCode, PollOutcome}; +pub use provider::{CopilotCodec, CopilotModel, CopilotProvider}; +pub use token::{CopilotToken, TokenProvider}; diff --git a/crates/harness-providers/src/copilot/provider.rs b/crates/harness-providers/src/copilot/provider.rs new file mode 100644 index 0000000..f2e754b --- /dev/null +++ b/crates/harness-providers/src/copilot/provider.rs @@ -0,0 +1,296 @@ +//! Copilot provider: multiplexes all three wire codecs behind `api.githubcopilot.com`. +//! +//! Each model's `supported_endpoints` (from `GET /models`) decides which codec/endpoint to +//! use. The token comes from [`super::token::TokenProvider`]. Network paths here are the +//! flagged live-verification risk — the routing/header logic is unit-tested; end-to-end +//! streaming must be checked against the live API. + +use std::collections::HashMap; +use std::time::{SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use harness_core::llm::{ + Initiator, LlmEventStream, LlmRequest, Provider, ProviderError, WireContent, +}; +use harness_core::types::ModelInfo; +use serde_json::Value; +use tokio_util::sync::CancellationToken; + +use crate::codec::{anthropic, openai_chat, openai_responses}; + +use super::token::TokenProvider; + +pub const DEFAULT_BASE_URL: &str = "https://api.githubcopilot.com"; +const API_VERSION: &str = "2026-06-01"; +const ANTHROPIC_BETA: &str = "interleaved-thinking-2025-05-14"; +const USER_AGENT: &str = concat!("ai-harness/", env!("CARGO_PKG_VERSION")); + +/// Which wire codec a Copilot model speaks. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CopilotCodec { + Chat, + Responses, + Anthropic, +} + +impl CopilotCodec { + fn path(self) -> &'static str { + match self { + CopilotCodec::Chat => "/chat/completions", + CopilotCodec::Responses => "/responses", + CopilotCodec::Anthropic => "/v1/messages", + } + } + + fn build_request(self, req: &LlmRequest) -> Value { + match self { + CopilotCodec::Chat => openai_chat::build_request(req), + CopilotCodec::Responses => openai_responses::build_request(req), + CopilotCodec::Anthropic => anthropic::build_request(req), + } + } + + fn decode( + self, + stream: impl futures::Stream> + Send + 'static, + ) -> LlmEventStream { + match self { + CopilotCodec::Chat => openai_chat::decode(stream), + CopilotCodec::Responses => openai_responses::decode(stream), + CopilotCodec::Anthropic => anthropic::decode(stream), + } + } +} + +/// A single model as advertised by `GET /models`. +#[derive(Debug, Clone)] +pub struct CopilotModel { + pub id: String, + pub codec: CopilotCodec, + /// Only `model_picker_enabled` models are offered in the UI picker. + pub picker_enabled: bool, +} + +/// Picks a codec from a model's `supported_endpoints`, preferring the responses API, then the +/// Anthropic messages API, then chat completions. +pub fn codec_for_endpoints(endpoints: &[String]) -> CopilotCodec { + let has = |needle: &str| endpoints.iter().any(|e| e.contains(needle)); + if has("/responses") { + CopilotCodec::Responses + } else if has("/v1/messages") || has("/messages") { + CopilotCodec::Anthropic + } else { + CopilotCodec::Chat + } +} + +/// Parses a Copilot `GET /models` response body. +pub fn parse_models(body: &Value) -> Vec { + body["data"] + .as_array() + .map(|arr| { + arr.iter() + .filter_map(|m| { + let id = m["id"].as_str()?.to_string(); + let endpoints: Vec = m["supported_endpoints"] + .as_array() + .map(|a| { + a.iter() + .filter_map(|e| e.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + Some(CopilotModel { + id, + codec: codec_for_endpoints(&endpoints), + picker_enabled: m["model_picker_enabled"].as_bool().unwrap_or(false), + }) + }) + .collect() + }) + .unwrap_or_default() +} + +fn has_images(req: &LlmRequest) -> bool { + req.messages + .iter() + .flat_map(|m| &m.content) + .any(|c| matches!(c, WireContent::Image { .. })) +} + +fn initiator_header(initiator: Initiator) -> &'static str { + match initiator { + Initiator::User => "user", + Initiator::Agent => "agent", + } +} + +pub struct CopilotProvider { + tokens: TokenProvider, + base_url: String, + client: reqwest::Client, + /// model id → codec, from `GET /models`. Unknown models default to chat completions. + codecs: HashMap, +} + +impl CopilotProvider { + pub fn new(oauth_token: impl Into, models: Vec) -> Self { + Self::with_base_url(oauth_token, DEFAULT_BASE_URL, models) + } + + pub fn with_base_url( + oauth_token: impl Into, + base_url: impl Into, + models: Vec, + ) -> Self { + let client = reqwest::Client::new(); + let codecs = models.into_iter().map(|m| (m.id, m.codec)).collect(); + Self { + tokens: TokenProvider::new(oauth_token, client.clone()), + base_url: base_url.into(), + client, + codecs, + } + } + + fn codec_for(&self, model: &str) -> CopilotCodec { + self.codecs + .get(model) + .copied() + .unwrap_or(CopilotCodec::Chat) + } + + fn classify_error(status: reqwest::StatusCode, body: String) -> ProviderError { + match status.as_u16() { + 401 | 403 => ProviderError::Auth(body), + 429 => ProviderError::RateLimited { retry_after: None }, + s if (500..600).contains(&s) => ProviderError::Overloaded, + s => ProviderError::Http { status: s, body }, + } + } +} + +#[async_trait] +impl Provider for CopilotProvider { + fn id(&self) -> &str { + "github-copilot" + } + + async fn list_models(&self) -> Result, ProviderError> { + // Cost/limit metadata is layered on by models.dev; routing metadata lives in `codecs`. + Ok(Vec::new()) + } + + async fn stream( + &self, + req: LlmRequest, + cancel: CancellationToken, + ) -> Result { + let codec = self.codec_for(&req.model); + let body = codec.build_request(&req); + let url = format!("{}{}", self.base_url, codec.path()); + + let now_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() as i64; + let token = self.tokens.token(now_secs).await; + + let mut builder = self + .client + .post(&url) + .header("authorization", format!("Bearer {}", token.token)) + .header("user-agent", USER_AGENT) + .header("x-github-api-version", API_VERSION) + .header("openai-intent", "conversation-edits") + .header("x-initiator", initiator_header(req.initiator)); + if has_images(&req) { + builder = builder.header("copilot-vision-request", "true"); + } + if codec == CopilotCodec::Anthropic { + builder = builder.header("anthropic-beta", ANTHROPIC_BETA); + } + + let send = builder.json(&body).send(); + let response = tokio::select! { + result = send => result.map_err(|e| ProviderError::Network(e.to_string()))?, + _ = cancel.cancelled() => return Err(ProviderError::Cancelled), + }; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(Self::classify_error(status, body)); + } + + Ok(codec.decode(response.bytes_stream())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn codec_routing_prefers_responses_then_anthropic_then_chat() { + assert_eq!( + codec_for_endpoints(&["/chat/completions".into(), "/responses".into()]), + CopilotCodec::Responses + ); + assert_eq!( + codec_for_endpoints(&["/v1/messages".into()]), + CopilotCodec::Anthropic + ); + assert_eq!( + codec_for_endpoints(&["/chat/completions".into()]), + CopilotCodec::Chat + ); + assert_eq!(codec_for_endpoints(&[]), CopilotCodec::Chat); + } + + #[test] + fn parses_models_with_endpoints_and_picker_flag() { + let body = json!({ + "data": [ + {"id": "gpt-4o", "supported_endpoints": ["/chat/completions"], "model_picker_enabled": true}, + {"id": "claude-sonnet-4-5", "supported_endpoints": ["/v1/messages"], "model_picker_enabled": true}, + {"id": "o3", "supported_endpoints": ["/responses"], "model_picker_enabled": false}, + {"id": "no-endpoints"} + ] + }); + let models = parse_models(&body); + assert_eq!(models.len(), 4); + assert_eq!(models[0].codec, CopilotCodec::Chat); + assert!(models[0].picker_enabled); + assert_eq!(models[1].codec, CopilotCodec::Anthropic); + assert_eq!(models[2].codec, CopilotCodec::Responses); + assert!(!models[2].picker_enabled); + // Missing supported_endpoints → default chat, picker false. + assert_eq!(models[3].codec, CopilotCodec::Chat); + assert!(!models[3].picker_enabled); + } + + #[test] + fn unknown_model_defaults_to_chat_codec() { + let provider = CopilotProvider::new("oauth", vec![]); + assert_eq!(provider.codec_for("whatever"), CopilotCodec::Chat); + assert_eq!(provider.id(), "github-copilot"); + } + + #[test] + fn known_model_uses_mapped_codec() { + let provider = CopilotProvider::new( + "oauth", + vec![CopilotModel { + id: "claude-sonnet-4-5".into(), + codec: CopilotCodec::Anthropic, + picker_enabled: true, + }], + ); + assert_eq!( + provider.codec_for("claude-sonnet-4-5"), + CopilotCodec::Anthropic + ); + } +} diff --git a/crates/harness-providers/src/copilot/token.rs b/crates/harness-providers/src/copilot/token.rs new file mode 100644 index 0000000..e4165a7 --- /dev/null +++ b/crates/harness-providers/src/copilot/token.rs @@ -0,0 +1,144 @@ +//! Copilot API token strategy (flagged-risk area — verify against the live API). +//! +//! opencode sends the GitHub OAuth token directly as the Bearer (`expires: 0`). The classic +//! Copilot API instead wants a short-lived token from `copilot_internal/v2/token`. We try the +//! exchange first and cache it until shortly before expiry; if the endpoint rejects us, we fall +//! back to the direct-Bearer behavior. Refresh is single-flight under a `tokio::Mutex`. + +use serde::Deserialize; +use tokio::sync::Mutex; + +const EXCHANGE_URL: &str = "https://api.github.com/copilot_internal/v2/token"; +/// Refresh this many seconds before the reported expiry. +const EXPIRY_SKEW_SECS: i64 = 120; +const USER_AGENT: &str = concat!("ai-harness/", env!("CARGO_PKG_VERSION")); + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CopilotToken { + pub token: String, + /// Unix seconds; `0` means "never expires" (direct-Bearer fallback). + pub expires_at: i64, +} + +impl CopilotToken { + /// True when the token should be refreshed at `now_secs` (never, for `expires_at == 0`). + pub fn needs_refresh(&self, now_secs: i64) -> bool { + self.expires_at != 0 && now_secs + EXPIRY_SKEW_SECS >= self.expires_at + } +} + +#[derive(Debug, thiserror::Error)] +pub enum TokenError { + #[error("network: {0}")] + Network(String), + /// The exchange endpoint is unavailable/unauthorized — the caller should fall back to the + /// direct-Bearer strategy. + #[error("exchange unsupported (status {0})")] + ExchangeUnsupported(u16), + #[error("unexpected response: {0}")] + Unexpected(String), +} + +#[derive(Deserialize)] +struct ExchangeResponse { + token: String, + #[serde(default)] + expires_at: i64, +} + +/// Performs the token exchange once. `ExchangeUnsupported` signals the caller to fall back. +pub async fn exchange( + client: &reqwest::Client, + oauth_token: &str, +) -> Result { + let resp = client + .get(EXCHANGE_URL) + .header("authorization", format!("token {oauth_token}")) + .header("accept", "application/json") + .header("user-agent", USER_AGENT) + .send() + .await + .map_err(|e| TokenError::Network(e.to_string()))?; + + let status = resp.status(); + if !status.is_success() { + // 401/403/404 → this deployment doesn't support the exchange; fall back to direct Bearer. + if matches!(status.as_u16(), 401 | 403 | 404) { + return Err(TokenError::ExchangeUnsupported(status.as_u16())); + } + return Err(TokenError::Network(format!("status {}", status.as_u16()))); + } + + let parsed: ExchangeResponse = resp + .json() + .await + .map_err(|e| TokenError::Unexpected(e.to_string()))?; + Ok(CopilotToken { + token: parsed.token, + expires_at: parsed.expires_at, + }) +} + +/// Caches the exchanged Copilot token and refreshes it single-flight. Falls back to using the +/// OAuth token directly (never-expiring) when the exchange endpoint rejects the request. +pub struct TokenProvider { + oauth_token: String, + client: reqwest::Client, + cached: Mutex>, +} + +impl TokenProvider { + pub fn new(oauth_token: impl Into, client: reqwest::Client) -> Self { + Self { + oauth_token: oauth_token.into(), + client, + cached: Mutex::new(None), + } + } + + /// Returns a usable Bearer token, refreshing/exchanging as needed. `now_secs` is injected + /// for testability. + pub async fn token(&self, now_secs: i64) -> CopilotToken { + let mut guard = self.cached.lock().await; + if let Some(tok) = guard.as_ref() { + if !tok.needs_refresh(now_secs) { + return tok.clone(); + } + } + let fresh = match exchange(&self.client, &self.oauth_token).await { + Ok(tok) => tok, + Err(_) => CopilotToken { + // Direct-Bearer fallback: use the OAuth token itself, never expiring. + token: self.oauth_token.clone(), + expires_at: 0, + }, + }; + *guard = Some(fresh.clone()); + fresh + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn never_expiring_token_does_not_refresh() { + let tok = CopilotToken { + token: "t".into(), + expires_at: 0, + }; + assert!(!tok.needs_refresh(i64::MAX)); + } + + #[test] + fn refresh_triggers_within_skew_window() { + let tok = CopilotToken { + token: "t".into(), + expires_at: 1_000, + }; + assert!(!tok.needs_refresh(800)); // 800 + 120 < 1000 + assert!(tok.needs_refresh(881)); // 881 + 120 >= 1000 + assert!(tok.needs_refresh(1_000)); + } +} diff --git a/crates/harness-providers/src/lib.rs b/crates/harness-providers/src/lib.rs index 63c8f2d..d7f8fb0 100644 --- a/crates/harness-providers/src/lib.rs +++ b/crates/harness-providers/src/lib.rs @@ -1,6 +1,14 @@ pub mod anthropic; +pub mod auth; pub mod codec; +pub mod copilot; +pub mod modelsdev; +pub mod openai; pub mod registry; pub use anthropic::AnthropicProvider; +pub use auth::{AuthRecord, AuthStorage}; +pub use copilot::CopilotProvider; +pub use modelsdev::ModelCatalog; +pub use openai::OpenAiProvider; pub use registry::ProviderRegistry; diff --git a/crates/harness-providers/src/modelsdev.rs b/crates/harness-providers/src/modelsdev.rs new file mode 100644 index 0000000..6d11bc2 --- /dev/null +++ b/crates/harness-providers/src/modelsdev.rs @@ -0,0 +1,296 @@ +//! models.dev metadata: per-model context/output limits, pricing, and capability flags. +//! +//! At runtime we prefer a cached copy of `https://models.dev/api.json` (refreshed every 24h); +//! if the cache is missing/stale and the network is unavailable, we fall back to a baked +//! snapshot so cost display and limits still work offline. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::time::{Duration, SystemTime}; + +use harness_core::types::{ModelCost, ModelInfo, ModelRef}; +use serde::Deserialize; + +const API_URL: &str = "https://models.dev/api.json"; +const CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60); +/// Baked fallback, refreshed manually. Keeps cost/limits working with no cache and no network. +const SNAPSHOT: &str = include_str!("../assets/models-snapshot.json"); + +/// models.dev top-level shape: `{ provider_id: { models: { model_id: {...} } } }`. Unknown +/// keys (provider metadata, per-model extras) are ignored. +#[derive(Deserialize)] +struct ApiProvider { + #[serde(default)] + models: HashMap, +} + +#[derive(Deserialize)] +struct ApiModel { + #[serde(default)] + reasoning: bool, + #[serde(default)] + tool_call: bool, + #[serde(default)] + attachment: bool, + #[serde(default)] + cost: ApiCost, + #[serde(default)] + limit: ApiLimit, +} + +#[derive(Deserialize, Default)] +struct ApiCost { + #[serde(default)] + input: f64, + #[serde(default)] + output: f64, + #[serde(default)] + cache_read: f64, + #[serde(default)] + cache_write: f64, +} + +#[derive(Deserialize, Default)] +struct ApiLimit { + #[serde(default)] + context: u64, + #[serde(default)] + output: u64, +} + +/// Parsed metadata keyed by `(provider_id, model_id)`. +#[derive(Debug, Clone, Default)] +pub struct ModelCatalog { + models: HashMap<(String, String), ModelInfo>, +} + +impl ModelCatalog { + /// Parses a models.dev `api.json` document. + pub fn from_api_json(bytes: &[u8]) -> Result { + let raw: HashMap = serde_json::from_slice(bytes)?; + let mut models = HashMap::new(); + for (provider_id, provider) in raw { + for (model_id, m) in provider.models { + let info = ModelInfo { + model: ModelRef::new(provider_id.clone(), model_id.clone()), + context_limit: m.limit.context, + output_limit: m.limit.output, + cost: ModelCost { + input: m.cost.input, + output: m.cost.output, + cache_read: m.cost.cache_read, + cache_write: m.cost.cache_write, + }, + reasoning: m.reasoning, + tool_call: m.tool_call, + attachment: m.attachment, + }; + models.insert((provider_id.clone(), model_id), info); + } + } + Ok(Self { models }) + } + + /// The baked fallback snapshot — always available, never fails. + pub fn baked() -> Self { + Self::from_api_json(SNAPSHOT.as_bytes()).expect("baked models snapshot must parse") + } + + pub fn get(&self, provider: &str, model: &str) -> Option<&ModelInfo> { + self.models.get(&(provider.to_string(), model.to_string())) + } + + /// Pricing for a model, or zero-cost if unknown. + pub fn cost(&self, provider: &str, model: &str) -> ModelCost { + self.get(provider, model) + .map(|m| m.cost) + .unwrap_or_default() + } + + pub fn len(&self) -> usize { + self.models.len() + } + + pub fn is_empty(&self) -> bool { + self.models.is_empty() + } + + /// Default cache location: `~/.cache/ai-harness/models.json`. + pub fn default_cache_path() -> PathBuf { + dirs::cache_dir() + .unwrap_or_else(std::env::temp_dir) + .join("ai-harness") + .join("models.json") + } + + /// Loads from the cache if present and younger than [`CACHE_TTL`]; otherwise returns the + /// baked snapshot. Never touches the network — call [`ModelCatalog::refresh`] for that. + pub fn load_cached_or_baked(cache_path: &Path) -> Self { + if cache_fresh(cache_path, CACHE_TTL, SystemTime::now()) { + if let Ok(bytes) = std::fs::read(cache_path) { + if let Ok(catalog) = Self::from_api_json(&bytes) { + if !catalog.is_empty() { + return catalog; + } + } + } + } + Self::baked() + } + + /// Refreshes the default cache location with a fresh HTTP client. Convenience wrapper for + /// callers (the app) that don't want to depend on `reqwest` directly. + pub async fn refresh_default_cache() -> Result { + let client = reqwest::Client::new(); + Self::refresh(&client, &Self::default_cache_path()).await + } + + /// Fetches the latest metadata and writes it to `cache_path` (best-effort). Returns the + /// freshly-parsed catalog. Intended to run in the background so the *next* launch is current. + pub async fn refresh( + client: &reqwest::Client, + cache_path: &Path, + ) -> Result { + let bytes = client + .get(API_URL) + .send() + .await + .map_err(|e| RefreshError::Network(e.to_string()))? + .error_for_status() + .map_err(|e| RefreshError::Network(e.to_string()))? + .bytes() + .await + .map_err(|e| RefreshError::Network(e.to_string()))?; + + let catalog = + Self::from_api_json(&bytes).map_err(|e| RefreshError::Parse(e.to_string()))?; + + if let Some(parent) = cache_path.parent() { + let _ = std::fs::create_dir_all(parent); + } + if let Err(e) = std::fs::write(cache_path, &bytes) { + tracing::warn!(error = %e, "failed to cache models.dev metadata"); + } + Ok(catalog) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum RefreshError { + #[error("network: {0}")] + Network(String), + #[error("parse: {0}")] + Parse(String), +} + +/// True when `path` exists and was modified within `ttl` of `now`. +fn cache_fresh(path: &Path, ttl: Duration, now: SystemTime) -> bool { + let Ok(modified) = std::fs::metadata(path).and_then(|m| m.modified()) else { + return false; + }; + match now.duration_since(modified) { + Ok(age) => age < ttl, + Err(_) => true, // modified in the future (clock skew) — treat as fresh + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const FIXTURE: &str = r#"{ + "anthropic": { + "id": "anthropic", + "name": "Anthropic", + "models": { + "claude-sonnet-4-5": { + "id": "claude-sonnet-4-5", + "reasoning": true, + "tool_call": true, + "attachment": true, + "cost": {"input": 3, "output": 15, "cache_read": 0.3, "cache_write": 3.75}, + "limit": {"context": 200000, "output": 64000} + } + } + }, + "openai": { + "id": "openai", + "models": { + "gpt-4o": { + "id": "gpt-4o", + "tool_call": true, + "cost": {"input": 2.5, "output": 10}, + "limit": {"context": 128000, "output": 16384} + } + } + } + }"#; + + #[test] + fn parses_providers_and_models() { + let catalog = ModelCatalog::from_api_json(FIXTURE.as_bytes()).unwrap(); + assert_eq!(catalog.len(), 2); + + let sonnet = catalog.get("anthropic", "claude-sonnet-4-5").unwrap(); + assert_eq!(sonnet.context_limit, 200_000); + assert_eq!(sonnet.output_limit, 64_000); + assert_eq!(sonnet.cost.input, 3.0); + assert_eq!(sonnet.cost.cache_write, 3.75); + assert!(sonnet.reasoning && sonnet.tool_call && sonnet.attachment); + } + + #[test] + fn missing_cost_fields_default_to_zero() { + let catalog = ModelCatalog::from_api_json(FIXTURE.as_bytes()).unwrap(); + let gpt = catalog.get("openai", "gpt-4o").unwrap(); + assert_eq!(gpt.cost.cache_read, 0.0); + assert_eq!(gpt.cost.cache_write, 0.0); + assert!(!gpt.reasoning); // absent → false + assert!(gpt.tool_call); + } + + #[test] + fn cost_helper_is_zero_for_unknown_model() { + let catalog = ModelCatalog::from_api_json(FIXTURE.as_bytes()).unwrap(); + assert_eq!(catalog.cost("openai", "nonexistent"), ModelCost::default()); + assert_eq!(catalog.cost("anthropic", "claude-sonnet-4-5").input, 3.0); + } + + #[test] + fn baked_snapshot_parses_and_is_nonempty() { + let catalog = ModelCatalog::baked(); + assert!(!catalog.is_empty()); + } + + #[test] + fn cache_freshness_respects_ttl() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("models.json"); + assert!(!cache_fresh(&path, CACHE_TTL, SystemTime::now())); // missing + + std::fs::write(&path, "{}").unwrap(); + let now = SystemTime::now(); + assert!(cache_fresh(&path, CACHE_TTL, now)); + // A "now" far in the future makes the file look stale. + let future = now + Duration::from_secs(48 * 60 * 60); + assert!(!cache_fresh(&path, CACHE_TTL, future)); + } + + #[test] + fn load_cached_or_baked_falls_back_when_cache_absent() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("models.json"); + let catalog = ModelCatalog::load_cached_or_baked(&path); + assert!(!catalog.is_empty()); // baked fallback + } + + #[test] + fn load_cached_or_baked_reads_fresh_cache() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("models.json"); + std::fs::write(&path, FIXTURE).unwrap(); + let catalog = ModelCatalog::load_cached_or_baked(&path); + assert_eq!(catalog.len(), 2); + assert!(catalog.get("anthropic", "claude-sonnet-4-5").is_some()); + } +} diff --git a/crates/harness-providers/src/openai.rs b/crates/harness-providers/src/openai.rs new file mode 100644 index 0000000..9961d17 --- /dev/null +++ b/crates/harness-providers/src/openai.rs @@ -0,0 +1,173 @@ +use std::time::Duration; + +use async_trait::async_trait; +use harness_core::llm::{LlmEventStream, LlmRequest, Provider, ProviderError}; +use harness_core::types::ModelInfo; +use tokio_util::sync::CancellationToken; + +use crate::codec::{openai_chat, openai_responses}; + +const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; + +/// Which OpenAI wire format to speak for a given model. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ApiFlavor { + Responses, + Chat, +} + +/// `gpt-*` and `o-*` models use the responses API; everything else (including +/// OpenAI-compatible third-party endpoints) falls back to chat completions. +fn flavor_for(model: &str) -> ApiFlavor { + let is_native = model.starts_with("gpt-") + || model.starts_with("o1") + || model.starts_with("o3") + || model.starts_with("o4") + || model == "o1" + || model == "o3"; + if is_native { + ApiFlavor::Responses + } else { + ApiFlavor::Chat + } +} + +pub struct OpenAiProvider { + api_key: String, + base_url: String, + client: reqwest::Client, +} + +impl OpenAiProvider { + pub fn new(api_key: impl Into) -> Self { + Self { + api_key: api_key.into(), + base_url: DEFAULT_BASE_URL.to_string(), + client: reqwest::Client::new(), + } + } + + pub fn with_base_url(api_key: impl Into, base_url: impl Into) -> Self { + Self { + api_key: api_key.into(), + base_url: base_url.into(), + client: reqwest::Client::new(), + } + } + + fn classify_error( + status: reqwest::StatusCode, + body: String, + retry_after: Option, + ) -> ProviderError { + match status.as_u16() { + 400 if body.contains("context_length_exceeded") + || body.contains("maximum context length") => + { + ProviderError::ContextOverflow + } + 401 | 403 => ProviderError::Auth(body), + 429 => ProviderError::RateLimited { retry_after }, + s if (500..600).contains(&s) => ProviderError::Overloaded, + s => ProviderError::Http { status: s, body }, + } + } +} + +#[async_trait] +impl Provider for OpenAiProvider { + fn id(&self) -> &str { + "openai" + } + + async fn list_models(&self) -> Result, ProviderError> { + // models.dev metadata is layered on top by the caller (see `modelsdev.rs`). + Ok(Vec::new()) + } + + async fn stream( + &self, + req: LlmRequest, + cancel: CancellationToken, + ) -> Result { + let (path, body) = match flavor_for(&req.model) { + ApiFlavor::Responses => ("/responses", openai_responses::build_request(&req)), + ApiFlavor::Chat => ("/chat/completions", openai_chat::build_request(&req)), + }; + let url = format!("{}{}", self.base_url, path); + + let send = self + .client + .post(&url) + .header("authorization", format!("Bearer {}", self.api_key)) + .json(&body) + .send(); + + let response = tokio::select! { + result = send => result.map_err(|e| ProviderError::Network(e.to_string()))?, + _ = cancel.cancelled() => return Err(ProviderError::Cancelled), + }; + + if !response.status().is_success() { + let status = response.status(); + let retry_after = response + .headers() + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + .map(Duration::from_secs); + let body = response.text().await.unwrap_or_default(); + return Err(Self::classify_error(status, body, retry_after)); + } + + Ok(match flavor_for(&req.model) { + ApiFlavor::Responses => openai_responses::decode(response.bytes_stream()), + ApiFlavor::Chat => openai_chat::decode(response.bytes_stream()), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn routes_gpt_and_o_series_to_responses() { + assert_eq!(flavor_for("gpt-4o"), ApiFlavor::Responses); + assert_eq!(flavor_for("gpt-5"), ApiFlavor::Responses); + assert_eq!(flavor_for("o1"), ApiFlavor::Responses); + assert_eq!(flavor_for("o3-mini"), ApiFlavor::Responses); + assert_eq!(flavor_for("o4-mini"), ApiFlavor::Responses); + } + + #[test] + fn routes_other_models_to_chat() { + assert_eq!(flavor_for("llama-3.1-70b"), ApiFlavor::Chat); + assert_eq!(flavor_for("deepseek-chat"), ApiFlavor::Chat); + } + + #[test] + fn id_is_openai() { + assert_eq!(OpenAiProvider::new("k").id(), "openai"); + } + + #[test] + fn classifies_context_overflow_from_400_body() { + assert!(matches!( + OpenAiProvider::classify_error( + reqwest::StatusCode::BAD_REQUEST, + "context_length_exceeded".into(), + None, + ), + ProviderError::ContextOverflow + )); + assert!(matches!( + OpenAiProvider::classify_error( + reqwest::StatusCode::BAD_REQUEST, + "some other error".into(), + None, + ), + ProviderError::Http { status: 400, .. } + )); + } +} diff --git a/crates/harness-tui/src/input.rs b/crates/harness-tui/src/input.rs index 26ac273..ee8097c 100644 --- a/crates/harness-tui/src/input.rs +++ b/crates/harness-tui/src/input.rs @@ -187,6 +187,8 @@ pub async fn apply_action(action: InputAction, state: &mut AppState, engine: &En state.session_id = Some(id); state.messages.clear(); state.scroll_offset = 0; + state.session_cost = 0.0; + state.session_tokens = 0; state.dirty = true; } Err(e) => tracing::error!(error = %e, "failed to create session"), diff --git a/crates/harness-tui/src/render.rs b/crates/harness-tui/src/render.rs index 8cd3ffd..43994da 100644 --- a/crates/harness-tui/src/render.rs +++ b/crates/harness-tui/src/render.rs @@ -171,11 +171,25 @@ fn render_status(frame: &mut Frame, state: &AppState, area: Rect) { } else { "Ctrl+S: sessions | Esc: abort | Ctrl+C: quit" }; - let text = format!("{spinner}{status} · {hints}"); + let usage = format!( + "{} · ${:.4}", + format_tokens(state.session_tokens), + state.session_cost + ); + let text = format!("{spinner}{status} · {usage} · {hints}"); let paragraph = Paragraph::new(text).style(Style::new().fg(Color::Gray)); frame.render_widget(paragraph, area); } +/// Compact token count: `1234` → `1.2k`, `2_000_000` → `2.0M`. +fn format_tokens(n: u64) -> String { + match n { + 0..=999 => format!("{n} tok"), + 1_000..=999_999 => format!("{:.1}k tok", n as f64 / 1_000.0), + _ => format!("{:.1}M tok", n as f64 / 1_000_000.0), + } +} + fn render_permission_modal( frame: &mut Frame, request: &harness_core::event::PermissionRequest, diff --git a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_empty_session.snap b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_empty_session.snap index de758c7..bc0e440 100644 --- a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_empty_session.snap +++ b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_empty_session.snap @@ -1,6 +1,6 @@ --- source: crates/harness-tui/src/render.rs -assertion_line: 277 +assertion_line: 300 expression: buffer_to_string(terminal.backend()) --- new session · orchestrator · anthropic/claude-sonnet-4-5 @@ -26,4 +26,4 @@ expression: buffer_to_string(terminal.backend()) ┌ input ───────────────────────────────────────────────────────────────────────┐ │ │ └──────────────────────────────────────────────────────────────────────────────┘ -idle · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit +idle · 0 tok · $0.0000 · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit diff --git a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_permission_modal.snap b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_permission_modal.snap index e83d418..348638e 100644 --- a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_permission_modal.snap +++ b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_permission_modal.snap @@ -1,6 +1,6 @@ --- source: crates/harness-tui/src/render.rs -assertion_line: 397 +assertion_line: 430 expression: buffer_to_string(terminal.backend()) --- new session · orchestrator · anthropic/claude-sonnet-4-5 @@ -26,4 +26,4 @@ expression: buffer_to_string(terminal.backend()) ┌ input ───────────────────────────────────────────────────────────────────────┐ │ │ └──────────────────────────────────────────────────────────────────────────────┘ -idle · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit +idle · 0 tok · $0.0000 · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit diff --git a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_session_picker_modal.snap b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_session_picker_modal.snap index 82c4456..7411661 100644 --- a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_session_picker_modal.snap +++ b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_session_picker_modal.snap @@ -1,6 +1,6 @@ --- source: crates/harness-tui/src/render.rs -assertion_line: 449 +assertion_line: 487 expression: buffer_to_string(terminal.backend()) --- new session · orchestrator · anthropic/claude-sonnet-4-5 @@ -26,4 +26,4 @@ expression: buffer_to_string(terminal.backend()) ┌ input ───────────────────────────────────────────────────────────────────────┐ │ │ └──────────────────────────────────────────────────────────────────────────────┘ -idle · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit +idle · 0 tok · $0.0000 · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit diff --git a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_completed.snap b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_completed.snap index c35c01c..f7c0a66 100644 --- a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_completed.snap +++ b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_completed.snap @@ -1,6 +1,6 @@ --- source: crates/harness-tui/src/render.rs -assertion_line: 330 +assertion_line: 357 expression: buffer_to_string(terminal.backend()) --- new session · orchestrator · anthropic/claude-sonnet-4-5 @@ -26,4 +26,4 @@ expression: buffer_to_string(terminal.backend()) ┌ input ───────────────────────────────────────────────────────────────────────┐ │ │ └──────────────────────────────────────────────────────────────────────────────┘ -idle · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit +idle · 0 tok · $0.0000 · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit diff --git a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_error.snap b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_error.snap index 24d9403..c0ec904 100644 --- a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_error.snap +++ b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_error.snap @@ -1,6 +1,6 @@ --- source: crates/harness-tui/src/render.rs -assertion_line: 376 +assertion_line: 407 expression: buffer_to_string(terminal.backend()) --- new session · orchestrator · anthropic/claude-sonnet-4-5 @@ -26,4 +26,4 @@ expression: buffer_to_string(terminal.backend()) ┌ input ───────────────────────────────────────────────────────────────────────┐ │ │ └──────────────────────────────────────────────────────────────────────────────┘ -idle · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit +idle · 0 tok · $0.0000 · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit diff --git a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_running.snap b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_running.snap index 0b90d93..f5f4ce9 100644 --- a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_running.snap +++ b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_tool_card_running.snap @@ -1,6 +1,6 @@ --- source: crates/harness-tui/src/render.rs -assertion_line: 353 +assertion_line: 382 expression: buffer_to_string(terminal.backend()) --- new session · orchestrator · anthropic/claude-sonnet-4-5 @@ -26,4 +26,4 @@ expression: buffer_to_string(terminal.backend()) ┌ input ───────────────────────────────────────────────────────────────────────┐ │ │ └──────────────────────────────────────────────────────────────────────────────┘ -idle · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit +idle · 0 tok · $0.0000 · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit diff --git a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_user_and_assistant_messages.snap b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_user_and_assistant_messages.snap index f4cb1ee..4cc4add 100644 --- a/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_user_and_assistant_messages.snap +++ b/crates/harness-tui/src/snapshots/harness__render__tests__snapshot_user_and_assistant_messages.snap @@ -1,6 +1,6 @@ --- source: crates/harness-tui/src/render.rs -assertion_line: 306 +assertion_line: 331 expression: buffer_to_string(terminal.backend()) --- new session · orchestrator · anthropic/claude-sonnet-4-5 @@ -26,4 +26,4 @@ expression: buffer_to_string(terminal.backend()) ┌ input ───────────────────────────────────────────────────────────────────────┐ │ │ └──────────────────────────────────────────────────────────────────────────────┘ -idle · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit +idle · 0 tok · $0.0000 · Ctrl+S: sessions | Esc: abort | Ctrl+C: quit diff --git a/crates/harness-tui/src/state.rs b/crates/harness-tui/src/state.rs index 1b00b18..cb507e1 100644 --- a/crates/harness-tui/src/state.rs +++ b/crates/harness-tui/src/state.rs @@ -105,6 +105,10 @@ pub struct AppState { pub ctrl_c_pressed: bool, /// Width the currently cached part lines were wrapped to; a change invalidates them. pub render_width: u16, + /// Accumulated session cost in USD, from `SessionCreated`/`SessionUpdated` events. + pub session_cost: f64, + /// Accumulated session tokens (input + output), for the status bar. + pub session_tokens: u64, } impl AppState { @@ -125,6 +129,8 @@ impl AppState { dirty: true, ctrl_c_pressed: false, render_width: 0, + session_cost: 0.0, + session_tokens: 0, } } @@ -147,6 +153,8 @@ impl AppState { self.session_title = session.title; self.agent = session.agent; self.model_ref = format!("{}/{}", session.model.provider_id, session.model.model_id); + self.session_cost = session.cost; + self.session_tokens = session.usage.input + session.usage.output; self.messages.clear(); let mut messages = messages; @@ -171,6 +179,8 @@ impl AppState { self.agent = session.agent; self.model_ref = format!("{}/{}", session.model.provider_id, session.model.model_id); + self.session_cost = session.cost; + self.session_tokens = session.usage.input + session.usage.output; self.dirty = true; } }