objectiveai_cli/command/agents/logs/token_usage/
subscribe.rs

1//! `agents logs token-usage subscribe` — wait for an agent's stored
2//! `total_tokens` snapshot to change, OR for its instance lock to
3//! drop. One-shot: emits exactly one item then ends.
4//!
5//! Flow:
6//! 1. Read the current stored value (`baseline`).
7//! 2. Fast path: if `--previous` is set AND a value is stored AND it
8//!    differs from `previous`, return that value immediately.
9//! 3. Otherwise subscribe: race the token-usage change (postgres
10//!    LISTEN, filtered + value-compared against `baseline`) against the
11//!    instance-lock release. On change → `Item`; on lock drop → a final
12//!    change check, else `"agents_inactive"`.
13
14use std::pin::Pin;
15
16use futures::Stream;
17use objectiveai_sdk::cli::command::agents::logs::token_usage::subscribe::{
18    AgentsInactiveTag, Request, ResponseItem, TokenUsage,
19};
20
21use crate::context::Context;
22use crate::error::Error;
23
24type ItemStream = Pin<Box<dyn Stream<Item = Result<ResponseItem, Error>> + Send>>;
25
26fn once(item: Result<ResponseItem, Error>) -> ItemStream {
27    Box::pin(futures::stream::once(async move { item }))
28}
29
30pub async fn execute(ctx: &Context, request: Request) -> Result<ItemStream, Error> {
31    let db = ctx.db_client().await?.clone();
32    let state_dir = ctx.filesystem.state_dir();
33    let aih = request.agent_instance_hierarchy;
34    let previous = request.previous;
35
36    let baseline = crate::db::logs::get_agent_token_usage(&db, &aih).await?;
37
38    // Fast path: `--previous` given and the stored value already
39    // differs — return it without subscribing.
40    if let (Some(prev), Some(current)) = (previous, baseline) {
41        if current != prev {
42            return Ok(once(Ok(ResponseItem::Item(TokenUsage {
43                agent_instance_hierarchy: aih,
44                total_tokens: current,
45            }))));
46        }
47    }
48
49    // Subscribe: race token-usage change vs instance-lock release. The
50    // command is one-shot, so we await the race here and emit a single
51    // item (rather than a generator that yields once).
52    let (lock_dir, lock_key) =
53        crate::command::agents::locks::agent_instance_lock(&state_dir, &aih);
54
55    let item: Result<ResponseItem, Error> = tokio::select! {
56        result = crate::db::logs::wait_for_token_usage_change(&db, &aih, baseline) => {
57            result
58                .map(|total_tokens| ResponseItem::Item(TokenUsage {
59                    agent_instance_hierarchy: aih.clone(),
60                    total_tokens,
61                }))
62                .map_err(Error::from)
63        }
64        result = objectiveai_sdk::lockfile::wait_released(&lock_dir, &lock_key) => {
65            match result {
66                Ok(()) => {
67                    // A change may have landed as the lock dropped —
68                    // report it rather than a bare inactive signal.
69                    match crate::db::logs::get_agent_token_usage(&db, &aih).await {
70                        Ok(current) if current != baseline => match current {
71                            Some(total_tokens) => Ok(ResponseItem::Item(TokenUsage {
72                                agent_instance_hierarchy: aih.clone(),
73                                total_tokens,
74                            })),
75                            None => Ok(ResponseItem::AgentsInactive(
76                                AgentsInactiveTag::AgentsInactive,
77                            )),
78                        },
79                        Ok(_) => Ok(ResponseItem::AgentsInactive(
80                            AgentsInactiveTag::AgentsInactive,
81                        )),
82                        Err(e) => Err(Error::from(e)),
83                    }
84                }
85                Err(e) => Err(Error::Lockfile { key: lock_key.clone(), source: e }),
86            }
87        }
88    };
89
90    Ok(once(item))
91}
92
93pub mod request_schema {
94    use objectiveai_sdk::cli::command::agents::logs::token_usage::subscribe as sdk;
95    use objectiveai_sdk::cli::command::agents::logs::token_usage::subscribe::request_schema::{
96        Request, Response,
97    };
98
99    use crate::context::Context;
100    use crate::error::Error;
101
102    pub async fn execute(_ctx: &Context, _request: Request) -> Result<Response, Error> {
103        Ok(objectiveai_sdk::cli::command::ResponseSchema(schemars::schema_for!(sdk::Request)))
104    }
105}
106
107pub mod response_schema {
108    use objectiveai_sdk::cli::command::agents::logs::token_usage::subscribe as sdk;
109    use objectiveai_sdk::cli::command::agents::logs::token_usage::subscribe::response_schema::{
110        Request, Response,
111    };
112
113    use crate::context::Context;
114    use crate::error::Error;
115
116    pub async fn execute(_ctx: &Context, _request: Request) -> Result<Response, Error> {
117        Ok(objectiveai_sdk::cli::command::ResponseSchema(schemars::schema_for!(sdk::ResponseItem)))
118    }
119}
objectiveai_cli/command/agents/logs/token_usage/subscribe.rs

objectiveai_cli/command/agents/logs/token_usage/
subscribe.rs
