theway-core 0.1.21

theway core — stateful agent runtime + harness (Agent loop, skills, prompt templates, sessions, compaction) on top of theway-llm-provider.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

//! Token + cost accumulator. Subscribes to assistant `MessageEnd` events and aggregates the
//! `Usage` already attached to each assistant message by the provider. No provider-specific
//! logic — the providers populate `Usage::cost` per-message via the catalog's pricing table,
//! so the tracker is a sum-only fold.
//!
//! Issue #7 part 1: in-memory totals, snapshot accessor, the `/cost` slash command on the
//! coding-agent side. Budget caps + fallback model land in a follow-up.

use std::sync::Arc;

use parking_lot::Mutex;
use theway_llm_provider::{Message as PiMessage, Usage};

use crate::types::{AgentMessage, LoopEvent};

/// Snapshot of the running totals. Cheap to clone — plain `Copy`-able fields.
#[derive(Clone, Debug, Default)]
pub struct CostSnapshot {
    pub tokens: Usage,
    pub turn_count: u64,
}

impl CostSnapshot {
    /// Total USD (input + output + cache). Convenience for the `/cost` summary line.
    pub fn total_cost(&self) -> f64 {
        self.tokens.cost.total
    }
}

/// Thread-safe accumulator. Cloning the tracker shares state via `Arc`.
#[derive(Clone, Debug)]
pub struct CostTracker {
    inner: Arc<Mutex<CostSnapshot>>,
}

impl Default for CostTracker {
    fn default() -> Self {
        Self::new()
    }
}

impl CostTracker {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(CostSnapshot::default())),
        }
    }

    pub fn snapshot(&self) -> CostSnapshot {
        self.inner.lock().clone()
    }

    /// Reset counters. Used by `/cost reset` and on session-switch.
    pub fn reset(&self) {
        *self.inner.lock() = CostSnapshot::default();
    }

    /// Apply a single assistant usage record. The provider has already computed costs against
    /// the catalog's pricing table — we just sum.
    pub fn record(&self, usage: &Usage) {
        let mut g = self.inner.lock();
        g.tokens.input = g.tokens.input.saturating_add(usage.input);
        g.tokens.output = g.tokens.output.saturating_add(usage.output);
        g.tokens.cache_read = g.tokens.cache_read.saturating_add(usage.cache_read);
        g.tokens.cache_write = g.tokens.cache_write.saturating_add(usage.cache_write);
        g.tokens.total_tokens = g.tokens.total_tokens.saturating_add(usage.total_tokens);
        let c = &mut g.tokens.cost;
        c.input += usage.cost.input;
        c.output += usage.cost.output;
        c.cache_read += usage.cost.cache_read;
        c.cache_write += usage.cost.cache_write;
        c.total += usage.cost.total;
        g.turn_count += 1;
    }

    /// Build a synchronous [`crate::agent::LoopSyncCallback`] that records every assistant
    /// `MessageEnd`. The callback clones the tracker (cheap — `Arc` bump) so the harness
    /// can keep its handle. Registered via [`crate::Agent::subscribe_sync`].
    pub fn as_callback(&self) -> crate::agent::LoopSyncCallback {
        let tracker = self.clone();
        Arc::new(move |event| {
            if let LoopEvent::MessageEnd {
                message: AgentMessage::Llm(PiMessage::Assistant(a)),
            } = event
            {
                tracker.record(&a.usage);
            }
        })
    }
}

/// Render a one-line summary for the REPL status bar / banner. Format keeps the long
/// breakdown for `/cost` itself.
pub fn one_line_summary(snap: &CostSnapshot) -> String {
    format!(
        "tokens: in={} out={} cached={} total={} | cost ${:.4}",
        snap.tokens.input,
        snap.tokens.output,
        snap.tokens.cache_read + snap.tokens.cache_write,
        snap.tokens.total_tokens,
        snap.total_cost()
    )
}

/// Render the full breakdown — used by the `/cost` slash command.
pub fn full_breakdown(snap: &CostSnapshot) -> String {
    let c = &snap.tokens.cost;
    format!(
        "  turns:        {turns}\n\
         \n\
         Tokens:\n\
         \n  input         {input}\n  output        {output}\n  cache read    {cache_read}\n  cache write   {cache_write}\n  total         {total}\n\n\
         Cost (USD):\n\
         \n  input         ${ci:.4}\n  output        ${co:.4}\n  cache read    ${cr:.4}\n  cache write   ${cw:.4}\n  total         ${ct:.4}\n",
        turns = snap.turn_count,
        input = snap.tokens.input,
        output = snap.tokens.output,
        cache_read = snap.tokens.cache_read,
        cache_write = snap.tokens.cache_write,
        total = snap.tokens.total_tokens,
        ci = c.input,
        co = c.output,
        cr = c.cache_read,
        cw = c.cache_write,
        ct = c.total,
    )
}

#[cfg(test)]
mod cost_bridge_tests {
    tests_bridge_macro::tests_bridge!("agent/cost");
}