diff --git a/crates/harness-app/src/lib.rs b/crates/harness-app/src/lib.rs index d81b365..d9fd86c 100644 --- a/crates/harness-app/src/lib.rs +++ b/crates/harness-app/src/lib.rs @@ -254,6 +254,8 @@ impl EngineHandle { temperature: None, max_steps: DEFAULT_MAX_STEPS, instructions: self.inner.config.instructions.clone(), + // Populated from models.dev metadata in the follow-up increment. + cost: None, }; // 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, }