kcode-k1-codex-runtime-observations 0.1.0

//! Provider observation routing for the serialized runtime actor.

use kcode_k1_codex_conversations::{Error, ErrorKind};
use kcode_k1_codex_events::{
    Inbound, ScopedKind, TokenUsage, TokenUsageScope, TokenUsageUpdated, decode,
    decode_token_usage_updated,
};
use kcode_k1_codex_runtime_session::Session;
use serde_json::Value;
use std::collections::HashMap;
use tokio::sync::mpsc;

type Sink = mpsc::UnboundedSender<RuntimeObservation>;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UsageObservation {
    pub model: String,
    pub usage: TokenUsage,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RuntimeObservation {
    Usage(UsageObservation),
    TurnCompleted(TokenUsageScope),
}

#[derive(Default)]
pub struct ObservationRoutes {
    sinks: HashMap<String, Sink>,
}

impl ObservationRoutes {
    pub fn subscribe(
        &mut self,
        key: String,
    ) -> Result<mpsc::UnboundedReceiver<RuntimeObservation>, Error> {
        if self.sinks.get(&key).is_some_and(|sink| !sink.is_closed()) {
            return Err(Error::new(
                ErrorKind::Busy,
                "observation subscription is already live",
            ));
        }
        let (sink, receiver) = mpsc::unbounded_channel();
        self.sinks.insert(key, sink);
        Ok(receiver)
    }

    pub fn observe(&mut self, session: &Session, message: &Value) {
        if let TokenUsageUpdated::Valid(usage) = decode_token_usage_updated(message.clone()) {
            let thread = usage.scope.thread_id.clone();
            if let Some((key, model)) = owner_model(session, &thread) {
                self.send(
                    &key,
                    RuntimeObservation::Usage(UsageObservation { model, usage }),
                );
            }
        }
        if let Inbound::Scoped(event) = decode(message.clone())
            && matches!(event.kind, ScopedKind::TurnCompleted(_))
            && let Some((key, _)) = owner_model(session, &event.scope.thread_id)
        {
            self.send(
                &key,
                RuntimeObservation::TurnCompleted(TokenUsageScope {
                    thread_id: event.scope.thread_id,
                    turn_id: event.scope.turn_id,
                }),
            );
        }
    }

    fn send(&mut self, key: &str, observation: RuntimeObservation) {
        if self
            .sinks
            .get(key)
            .is_some_and(|sink| sink.send(observation).is_err())
        {
            self.sinks.remove(key);
        }
    }
}

fn owner_model(session: &Session, thread: &str) -> Option<(String, String)> {
    let key = session.owner(thread)?;
    Some((key.to_owned(), session.bound_profile(key)?.model.clone()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use kcode_k1_codex_conversations::Config;
    use serde_json::json;
    use std::sync::Arc;

    fn session() -> Session {
        let mut session = Session::default();
        let (events, _) = mpsc::unbounded_channel();
        session
            .begin_start(
                "key",
                Arc::new(Config {
                    executable: "codex".into(),
                    working_directory: "/work".into(),
                    model: "model".into(),
                    reasoning_effort: None,
                    base_instructions: String::new(),
                    tools: Vec::new(),
                }),
                events,
            )
            .unwrap();
        session.set_thread("key", "thread").unwrap();
        session
    }

    fn counts(value: i64) -> Value {
        json!({
            "inputTokens": value,
            "cachedInputTokens": value,
            "outputTokens": value,
            "reasoningOutputTokens": value,
            "totalTokens": value
        })
    }

    fn usage(thread: &str) -> Value {
        json!({
            "method": "thread/tokenUsage/updated",
            "params": {
                "threadId": thread,
                "turnId": "turn",
                "tokenUsage": {
                    "total": counts(10),
                    "last": counts(1),
                    "modelContextWindow": 128
                }
            }
        })
    }

    #[test]
    fn subscription_rejects_live_sink_and_replaces_closed_sink() {
        let mut routes = ObservationRoutes::default();
        let first = routes.subscribe("key".into()).unwrap();
        assert_eq!(
            routes.subscribe("key".into()).unwrap_err().kind,
            ErrorKind::Busy
        );
        drop(first);
        assert!(routes.subscribe("key".into()).is_ok());
    }

    #[test]
    fn routes_valid_usage_with_bound_model_and_preserves_duplicates() {
        let mut routes = ObservationRoutes::default();
        let mut receiver = routes.subscribe("key".into()).unwrap();
        let session = session();
        routes.observe(&session, &usage("thread"));
        routes.observe(&session, &usage("thread"));
        let first = receiver.try_recv().unwrap();
        let second = receiver.try_recv().unwrap();
        assert_eq!(first, second);
        assert!(
            matches!(first, RuntimeObservation::Usage(UsageObservation { ref model, .. }) if model == "model")
        );
    }

    #[test]
    fn routes_valid_scoped_completion() {
        let mut routes = ObservationRoutes::default();
        let mut receiver = routes.subscribe("key".into()).unwrap();
        let session = session();
        routes.observe(
            &session,
            &json!({
                "method": "turn/completed",
                "params": {"threadId": "thread", "turn": {"id": "turn", "status": "completed"}}
            }),
        );
        assert!(
            matches!(receiver.try_recv(), Ok(RuntimeObservation::TurnCompleted(TokenUsageScope { ref thread_id, .. })) if thread_id == "thread")
        );
    }

    #[test]
    fn ignores_unknown_threads_and_malformed_messages() {
        let mut routes = ObservationRoutes::default();
        let mut receiver = routes.subscribe("key".into()).unwrap();
        let session = session();
        routes.observe(&session, &usage("other"));
        routes.observe(
            &session,
            &json!({"method": "thread/tokenUsage/updated", "params": {}}),
        );
        assert!(receiver.try_recv().is_err());
    }

    #[test]
    fn closed_sink_is_removed_after_delivery_attempt() {
        let mut routes = ObservationRoutes::default();
        let receiver = routes.subscribe("key".into()).unwrap();
        drop(receiver);
        routes.observe(&session(), &usage("thread"));
        assert!(routes.sinks.is_empty());
    }
}