M3: All providers + auth #4
@@ -254,6 +254,8 @@ impl EngineHandle {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
max_steps: DEFAULT_MAX_STEPS,
|
max_steps: DEFAULT_MAX_STEPS,
|
||||||
instructions: self.inner.config.instructions.clone(),
|
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
|
// Reserve the run slot *before* spawning. If we inserted after spawning, a run that
|
||||||
|
|||||||
@@ -20,6 +20,33 @@ pub struct RunConfig {
|
|||||||
pub temperature: Option<f32>,
|
pub temperature: Option<f32>,
|
||||||
pub max_steps: u32,
|
pub max_steps: u32,
|
||||||
pub instructions: Vec<String>,
|
pub instructions: Vec<String>,
|
||||||
|
/// Pricing for `model`, from models.dev metadata. `None` leaves cost at 0.
|
||||||
|
pub cost: Option<crate::types::ModelCost>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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
|
/// opencode's exit condition: keep stepping while the last assistant turn asked for more
|
||||||
@@ -221,17 +248,24 @@ pub async fn run_session(
|
|||||||
&ctx,
|
&ctx,
|
||||||
run_config.model.clone(),
|
run_config.model.clone(),
|
||||||
&run_config.agent_name,
|
&run_config.agent_name,
|
||||||
|
run_config.cost,
|
||||||
&mut doomloop,
|
&mut doomloop,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
match step {
|
match step {
|
||||||
Ok(outcome) if outcome.aborted => return RunOutcome::Aborted,
|
Ok(outcome) if outcome.aborted => {
|
||||||
Ok(outcome) => match outcome.result {
|
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::Continue => continue,
|
||||||
StepResult::Stop => return RunOutcome::Stopped,
|
StepResult::Stop => return RunOutcome::Stopped,
|
||||||
StepResult::Compact => return RunOutcome::Stopped, // stub until M6
|
StepResult::Compact => return RunOutcome::Stopped, // stub until M6
|
||||||
},
|
}
|
||||||
|
}
|
||||||
Err(step_err) if matches!(step_err.source, ProviderError::Cancelled) => {
|
Err(step_err) if matches!(step_err.source, ProviderError::Cancelled) => {
|
||||||
return RunOutcome::Aborted;
|
return RunOutcome::Aborted;
|
||||||
}
|
}
|
||||||
@@ -422,11 +456,28 @@ mod tests {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
max_steps: 10,
|
max_steps: 10,
|
||||||
instructions: Vec::new(),
|
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;
|
let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await;
|
||||||
assert!(matches!(outcome, RunOutcome::Stopped));
|
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();
|
let messages = store.messages(session_id.clone()).await.unwrap();
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
messages.len(),
|
messages.len(),
|
||||||
@@ -517,6 +568,7 @@ mod tests {
|
|||||||
temperature: None,
|
temperature: None,
|
||||||
max_steps: 10,
|
max_steps: 10,
|
||||||
instructions: Vec::new(),
|
instructions: Vec::new(),
|
||||||
|
cost: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await;
|
let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await;
|
||||||
|
|||||||
@@ -26,6 +26,8 @@ pub struct StepOutcome {
|
|||||||
pub result: StepResult,
|
pub result: StepResult,
|
||||||
pub message_id: Option<MessageId>,
|
pub message_id: Option<MessageId>,
|
||||||
pub usage: TokenUsage,
|
pub usage: TokenUsage,
|
||||||
|
/// Dollar cost of this step's usage (0.0 when no pricing is available).
|
||||||
|
pub cost: f64,
|
||||||
pub aborted: bool,
|
pub aborted: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -513,6 +515,7 @@ impl<'a> Run<'a> {
|
|||||||
&mut self,
|
&mut self,
|
||||||
reason: FinishReason,
|
reason: FinishReason,
|
||||||
usage: TokenUsage,
|
usage: TokenUsage,
|
||||||
|
cost: f64,
|
||||||
) -> Result<StepResult, ProviderError> {
|
) -> Result<StepResult, ProviderError> {
|
||||||
let part = Part {
|
let part = Part {
|
||||||
id: PartId::new(),
|
id: PartId::new(),
|
||||||
@@ -521,7 +524,7 @@ impl<'a> Run<'a> {
|
|||||||
idx: self.next_idx,
|
idx: self.next_idx,
|
||||||
body: PartBody::StepFinish {
|
body: PartBody::StepFinish {
|
||||||
usage,
|
usage,
|
||||||
cost: 0.0,
|
cost,
|
||||||
reason: reason.clone(),
|
reason: reason.clone(),
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
@@ -567,10 +570,12 @@ pub async fn process_step(
|
|||||||
ctx: &StepContext,
|
ctx: &StepContext,
|
||||||
model: crate::types::ModelRef,
|
model: crate::types::ModelRef,
|
||||||
agent: &str,
|
agent: &str,
|
||||||
|
cost: Option<crate::types::ModelCost>,
|
||||||
doomloop: &mut DoomLoopGuard,
|
doomloop: &mut DoomLoopGuard,
|
||||||
) -> Result<StepOutcome, StepError> {
|
) -> Result<StepOutcome, StepError> {
|
||||||
let mut run = Run::new(ctx);
|
let mut run = Run::new(ctx);
|
||||||
let mut usage = TokenUsage::default();
|
let mut usage = TokenUsage::default();
|
||||||
|
let mut step_cost = 0.0;
|
||||||
let mut result = StepResult::Stop;
|
let mut result = StepResult::Stop;
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
@@ -582,6 +587,7 @@ pub async fn process_step(
|
|||||||
result: StepResult::Stop,
|
result: StepResult::Stop,
|
||||||
message_id: run.assistant.as_ref().map(|m| m.id.clone()),
|
message_id: run.assistant.as_ref().map(|m| m.id.clone()),
|
||||||
usage,
|
usage,
|
||||||
|
cost: step_cost,
|
||||||
aborted: true,
|
aborted: true,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
@@ -630,7 +636,8 @@ pub async fn process_step(
|
|||||||
usage: finish_usage,
|
usage: finish_usage,
|
||||||
} => {
|
} => {
|
||||||
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) => {
|
Ok(r) => {
|
||||||
result = r;
|
result = r;
|
||||||
Ok(())
|
Ok(())
|
||||||
@@ -652,6 +659,7 @@ pub async fn process_step(
|
|||||||
result,
|
result,
|
||||||
message_id: run.assistant.as_ref().map(|m| m.id.clone()),
|
message_id: run.assistant.as_ref().map(|m| m.id.clone()),
|
||||||
usage,
|
usage,
|
||||||
|
cost: step_cost,
|
||||||
aborted: false,
|
aborted: false,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ pub enum StoreCmd {
|
|||||||
UpsertSession(Session, Reply<()>),
|
UpsertSession(Session, Reply<()>),
|
||||||
UpsertMessage(Message, Reply<()>),
|
UpsertMessage(Message, Reply<()>),
|
||||||
UpsertPart(Part, Reply<()>),
|
UpsertPart(Part, Reply<()>),
|
||||||
|
Session(SessionId, Reply<Option<Session>>),
|
||||||
Sessions(Reply<Vec<Session>>),
|
Sessions(Reply<Vec<Session>>),
|
||||||
Messages(SessionId, Reply<Vec<Message>>),
|
Messages(SessionId, Reply<Vec<Message>>),
|
||||||
Parts(MessageId, Reply<Vec<Part>>),
|
Parts(MessageId, Reply<Vec<Part>>),
|
||||||
@@ -69,6 +70,15 @@ fn upsert_part(conn: &Connection, part: &Part) -> Result<(), StoreError> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn get_session(conn: &Connection, id: &SessionId) -> Result<Option<Session>, 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<Vec<Session>, StoreError> {
|
fn list_sessions(conn: &Connection) -> Result<Vec<Session>, StoreError> {
|
||||||
let mut stmt = conn.prepare("SELECT data FROM session ORDER BY id")?;
|
let mut stmt = conn.prepare("SELECT data FROM session ORDER BY id")?;
|
||||||
let rows = stmt
|
let rows = stmt
|
||||||
@@ -116,6 +126,9 @@ pub fn run(conn: Connection, mut rx: tokio::sync::mpsc::Receiver<StoreCmd>) {
|
|||||||
StoreCmd::UpsertPart(part, reply) => {
|
StoreCmd::UpsertPart(part, reply) => {
|
||||||
let _ = reply.send(upsert_part(&conn, &part));
|
let _ = reply.send(upsert_part(&conn, &part));
|
||||||
}
|
}
|
||||||
|
StoreCmd::Session(id, reply) => {
|
||||||
|
let _ = reply.send(get_session(&conn, &id));
|
||||||
|
}
|
||||||
StoreCmd::Sessions(reply) => {
|
StoreCmd::Sessions(reply) => {
|
||||||
let _ = reply.send(list_sessions(&conn));
|
let _ = reply.send(list_sessions(&conn));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -77,6 +77,11 @@ impl Store {
|
|||||||
self.call(|reply| StoreCmd::UpsertPart(part, reply)).await
|
self.call(|reply| StoreCmd::UpsertPart(part, reply)).await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn session(&self, session_id: SessionId) -> Result<Option<Session>, StoreError> {
|
||||||
|
self.call(|reply| StoreCmd::Session(session_id, reply))
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn sessions(&self) -> Result<Vec<Session>, StoreError> {
|
pub async fn sessions(&self) -> Result<Vec<Session>, StoreError> {
|
||||||
self.call(StoreCmd::Sessions).await
|
self.call(StoreCmd::Sessions).await
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -6,6 +6,6 @@ pub mod session;
|
|||||||
|
|
||||||
pub use ids::{MessageId, PartId, SessionId};
|
pub use ids::{MessageId, PartId, SessionId};
|
||||||
pub use message::{Message, MessageError, Role};
|
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 part::{Part, PartBody, ToolState};
|
||||||
pub use session::Session;
|
pub use session::Session;
|
||||||
|
|||||||
@@ -34,11 +34,71 @@ impl TokenUsage {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Full cost/context-limit metadata is populated by harness-providers (models.dev) in M1/M3;
|
/// Per-model pricing in USD per **one million** tokens. Populated from models.dev metadata
|
||||||
// this placeholder only carries what harness-core needs to key on.
|
/// 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)]
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||||
pub struct ModelInfo {
|
pub struct ModelInfo {
|
||||||
pub model: ModelRef,
|
pub model: ModelRef,
|
||||||
pub context_limit: u64,
|
pub context_limit: u64,
|
||||||
pub output_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,
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user