kcode_k1_codex_runtime_observations/
lib.rs

1//! Provider observation routing for the serialized runtime actor.
2
3use kcode_k1_codex_conversations::{Error, ErrorKind};
4use kcode_k1_codex_events::{
5    Inbound, ScopedKind, TokenUsage, TokenUsageScope, TokenUsageUpdated, decode,
6    decode_token_usage_updated,
7};
8use kcode_k1_codex_runtime_session::Session;
9use serde_json::Value;
10use std::collections::HashMap;
11use tokio::sync::mpsc;
12
13type Sink = mpsc::UnboundedSender<RuntimeObservation>;
14
15#[derive(Clone, Debug, PartialEq, Eq)]
16pub struct UsageObservation {
17    pub model: String,
18    pub usage: TokenUsage,
19}
20
21#[derive(Clone, Debug, PartialEq, Eq)]
22pub enum RuntimeObservation {
23    Usage(UsageObservation),
24    TurnCompleted(TokenUsageScope),
25}
26
27#[derive(Default)]
28pub struct ObservationRoutes {
29    sinks: HashMap<String, Sink>,
30}
31
32impl ObservationRoutes {
33    pub fn subscribe(
34        &mut self,
35        key: String,
36    ) -> Result<mpsc::UnboundedReceiver<RuntimeObservation>, Error> {
37        if self.sinks.get(&key).is_some_and(|sink| !sink.is_closed()) {
38            return Err(Error::new(
39                ErrorKind::Busy,
40                "observation subscription is already live",
41            ));
42        }
43        let (sink, receiver) = mpsc::unbounded_channel();
44        self.sinks.insert(key, sink);
45        Ok(receiver)
46    }
47
48    pub fn observe(&mut self, session: &Session, message: &Value) {
49        if let TokenUsageUpdated::Valid(usage) = decode_token_usage_updated(message.clone()) {
50            let thread = usage.scope.thread_id.clone();
51            if let Some((key, model)) = owner_model(session, &thread) {
52                self.send(
53                    &key,
54                    RuntimeObservation::Usage(UsageObservation { model, usage }),
55                );
56            }
57        }
58        if let Inbound::Scoped(event) = decode(message.clone())
59            && matches!(event.kind, ScopedKind::TurnCompleted(_))
60            && let Some((key, _)) = owner_model(session, &event.scope.thread_id)
61        {
62            self.send(
63                &key,
64                RuntimeObservation::TurnCompleted(TokenUsageScope {
65                    thread_id: event.scope.thread_id,
66                    turn_id: event.scope.turn_id,
67                }),
68            );
69        }
70    }
71
72    fn send(&mut self, key: &str, observation: RuntimeObservation) {
73        if self
74            .sinks
75            .get(key)
76            .is_some_and(|sink| sink.send(observation).is_err())
77        {
78            self.sinks.remove(key);
79        }
80    }
81}
82
83fn owner_model(session: &Session, thread: &str) -> Option<(String, String)> {
84    let key = session.owner(thread)?;
85    Some((key.to_owned(), session.bound_profile(key)?.model.clone()))
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use kcode_k1_codex_conversations::Config;
92    use serde_json::json;
93    use std::sync::Arc;
94
95    fn session() -> Session {
96        let mut session = Session::default();
97        let (events, _) = mpsc::unbounded_channel();
98        session
99            .begin_start(
100                "key",
101                Arc::new(Config {
102                    executable: "codex".into(),
103                    working_directory: "/work".into(),
104                    model: "model".into(),
105                    reasoning_effort: None,
106                    base_instructions: String::new(),
107                    tools: Vec::new(),
108                }),
109                events,
110            )
111            .unwrap();
112        session.set_thread("key", "thread").unwrap();
113        session
114    }
115
116    fn counts(value: i64) -> Value {
117        json!({
118            "inputTokens": value,
119            "cachedInputTokens": value,
120            "outputTokens": value,
121            "reasoningOutputTokens": value,
122            "totalTokens": value
123        })
124    }
125
126    fn usage(thread: &str) -> Value {
127        json!({
128            "method": "thread/tokenUsage/updated",
129            "params": {
130                "threadId": thread,
131                "turnId": "turn",
132                "tokenUsage": {
133                    "total": counts(10),
134                    "last": counts(1),
135                    "modelContextWindow": 128
136                }
137            }
138        })
139    }
140
141    #[test]
142    fn subscription_rejects_live_sink_and_replaces_closed_sink() {
143        let mut routes = ObservationRoutes::default();
144        let first = routes.subscribe("key".into()).unwrap();
145        assert_eq!(
146            routes.subscribe("key".into()).unwrap_err().kind,
147            ErrorKind::Busy
148        );
149        drop(first);
150        assert!(routes.subscribe("key".into()).is_ok());
151    }
152
153    #[test]
154    fn routes_valid_usage_with_bound_model_and_preserves_duplicates() {
155        let mut routes = ObservationRoutes::default();
156        let mut receiver = routes.subscribe("key".into()).unwrap();
157        let session = session();
158        routes.observe(&session, &usage("thread"));
159        routes.observe(&session, &usage("thread"));
160        let first = receiver.try_recv().unwrap();
161        let second = receiver.try_recv().unwrap();
162        assert_eq!(first, second);
163        assert!(
164            matches!(first, RuntimeObservation::Usage(UsageObservation { ref model, .. }) if model == "model")
165        );
166    }
167
168    #[test]
169    fn routes_valid_scoped_completion() {
170        let mut routes = ObservationRoutes::default();
171        let mut receiver = routes.subscribe("key".into()).unwrap();
172        let session = session();
173        routes.observe(
174            &session,
175            &json!({
176                "method": "turn/completed",
177                "params": {"threadId": "thread", "turn": {"id": "turn", "status": "completed"}}
178            }),
179        );
180        assert!(
181            matches!(receiver.try_recv(), Ok(RuntimeObservation::TurnCompleted(TokenUsageScope { ref thread_id, .. })) if thread_id == "thread")
182        );
183    }
184
185    #[test]
186    fn ignores_unknown_threads_and_malformed_messages() {
187        let mut routes = ObservationRoutes::default();
188        let mut receiver = routes.subscribe("key".into()).unwrap();
189        let session = session();
190        routes.observe(&session, &usage("other"));
191        routes.observe(
192            &session,
193            &json!({"method": "thread/tokenUsage/updated", "params": {}}),
194        );
195        assert!(receiver.try_recv().is_err());
196    }
197
198    #[test]
199    fn closed_sink_is_removed_after_delivery_attempt() {
200        let mut routes = ObservationRoutes::default();
201        let receiver = routes.subscribe("key".into()).unwrap();
202        drop(receiver);
203        routes.observe(&session(), &usage("thread"));
204        assert!(routes.sinks.is_empty());
205    }
206}
kcode_k1_codex_runtime_observations/lib.rs

kcode_k1_codex_runtime_observations/
lib.rs
