objectiveai_sdk/cli/command/agents/logs/token_usage/
mod.rs

1//! `agents logs token-usage` — per-AIH token-usage snapshot subtree.
2//! Leaves:
3//!
4//! - `get` — read the current stored `total_tokens` (no waiting).
5//! - `subscribe` — wait for the stored `total_tokens` to change (or the
6//!   instance lock to drop).
7
8use crate::cli::command::CommandRequest;
9
10pub mod get;
11pub mod subscribe;
12
13#[derive(clap::Subcommand)]
14pub enum Command {
15    /// Read an agent's current token-usage snapshot.
16    Get(get::Command),
17    /// Wait for an agent's token-usage snapshot to change.
18    Subscribe(subscribe::Command),
19}
20
21#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
22#[serde(untagged)]
23#[schemars(rename = "cli.command.agents.logs.token_usage.Request")]
24pub enum Request {
25    #[schemars(title = "Get")]
26    Get(get::Request),
27    #[schemars(title = "GetRequestSchema")]
28    GetRequestSchema(get::request_schema::Request),
29    #[schemars(title = "GetResponseSchema")]
30    GetResponseSchema(get::response_schema::Request),
31    #[schemars(title = "Subscribe")]
32    Subscribe(subscribe::Request),
33    #[schemars(title = "SubscribeRequestSchema")]
34    SubscribeRequestSchema(subscribe::request_schema::Request),
35    #[schemars(title = "SubscribeResponseSchema")]
36    SubscribeResponseSchema(subscribe::response_schema::Request),
37}
38
39// Exempt from json-schema coverage: tier aggregate (see the root
40// `ResponseItem` in command.rs - TS7056).
41#[objectiveai_sdk_macros::json_schema_ignore]
42#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
43#[schemars(rename = "cli.command.agents.logs.token_usage.ResponseItem")]
44#[serde(untagged)]
45pub enum ResponseItem {
46    #[schemars(title = "Get")]
47    Get(get::Response),
48    #[schemars(title = "GetRequestSchema")]
49    GetRequestSchema(get::request_schema::Response),
50    #[schemars(title = "GetResponseSchema")]
51    GetResponseSchema(get::response_schema::Response),
52    #[schemars(title = "Subscribe")]
53    Subscribe(subscribe::ResponseItem),
54    #[schemars(title = "SubscribeRequestSchema")]
55    SubscribeRequestSchema(subscribe::request_schema::Response),
56    #[schemars(title = "SubscribeResponseSchema")]
57    SubscribeResponseSchema(subscribe::response_schema::Response),
58}
59
60#[cfg(feature = "mcp")]
61impl crate::cli::command::CommandResponse for ResponseItem {
62    fn into_mcp(self) -> crate::cli::command::McpResponseItem {
63        match self {
64            ResponseItem::Get(v) => v.into_mcp(),
65            ResponseItem::GetRequestSchema(v) => v.into_mcp(),
66            ResponseItem::GetResponseSchema(v) => v.into_mcp(),
67            ResponseItem::Subscribe(v) => v.into_mcp(),
68            ResponseItem::SubscribeRequestSchema(v) => v.into_mcp(),
69            ResponseItem::SubscribeResponseSchema(v) => v.into_mcp(),
70        }
71    }
72}
73
74impl TryFrom<Command> for Request {
75    type Error = crate::cli::command::FromArgsError;
76    fn try_from(command: Command) -> Result<Self, Self::Error> {
77        match command {
78            Command::Get(cmd) => match cmd.schema {
79                None => Ok(Request::Get(get::Request::try_from(cmd.args)?)),
80                Some(get::Schema::RequestSchema(args)) => Ok(
81                    Request::GetRequestSchema(get::request_schema::Request::try_from(args)?),
82                ),
83                Some(get::Schema::ResponseSchema(args)) => Ok(
84                    Request::GetResponseSchema(get::response_schema::Request::try_from(args)?),
85                ),
86            },
87            Command::Subscribe(cmd) => match cmd.schema {
88                None => Ok(Request::Subscribe(subscribe::Request::try_from(cmd.args)?)),
89                Some(subscribe::Schema::RequestSchema(args)) => Ok(
90                    Request::SubscribeRequestSchema(subscribe::request_schema::Request::try_from(args)?),
91                ),
92                Some(subscribe::Schema::ResponseSchema(args)) => Ok(
93                    Request::SubscribeResponseSchema(subscribe::response_schema::Request::try_from(args)?),
94                ),
95            },
96        }
97    }
98}
99
100impl CommandRequest for Request {
101    fn request_base(&self) -> &crate::cli::command::RequestBase {
102        match self {
103            Request::Get(inner) => inner.request_base(),
104            Request::GetRequestSchema(inner) => inner.request_base(),
105            Request::GetResponseSchema(inner) => inner.request_base(),
106            Request::Subscribe(inner) => inner.request_base(),
107            Request::SubscribeRequestSchema(inner) => inner.request_base(),
108            Request::SubscribeResponseSchema(inner) => inner.request_base(),
109        }
110    }
111
112    fn request_base_mut(&mut self) -> Option<&mut crate::cli::command::RequestBase> {
113        match self {
114            Request::Get(inner) => inner.request_base_mut(),
115            Request::GetRequestSchema(inner) => inner.request_base_mut(),
116            Request::GetResponseSchema(inner) => inner.request_base_mut(),
117            Request::Subscribe(inner) => inner.request_base_mut(),
118            Request::SubscribeRequestSchema(inner) => inner.request_base_mut(),
119            Request::SubscribeResponseSchema(inner) => inner.request_base_mut(),
120        }
121    }
122}
123
124#[cfg(feature = "cli-executor")]
125pub async fn execute<E: crate::cli::command::CommandExecutor>(
126    executor: &E,
127    request: Request,
128    agent_arguments: Option<&crate::cli::command::AgentArguments>,
129) -> Result<
130    std::pin::Pin<Box<dyn futures::Stream<Item = Result<ResponseItem, E::Error>> + Send>>,
131    E::Error,
132> {
133    use futures::StreamExt;
134    let stream: std::pin::Pin<
135        Box<dyn futures::Stream<Item = Result<ResponseItem, E::Error>> + Send>,
136    > = match request {
137        Request::Get(req) => {
138            let value = get::execute(executor, req, agent_arguments).await?;
139            Box::pin(crate::cli::command::StreamOnce::new(Ok(ResponseItem::Get(value))))
140        }
141        Request::GetRequestSchema(req) => {
142            let value = get::request_schema::execute(executor, req, agent_arguments).await?;
143            Box::pin(crate::cli::command::StreamOnce::new(Ok(
144                ResponseItem::GetRequestSchema(value),
145            )))
146        }
147        Request::GetResponseSchema(req) => {
148            let value = get::response_schema::execute(executor, req, agent_arguments).await?;
149            Box::pin(crate::cli::command::StreamOnce::new(Ok(
150                ResponseItem::GetResponseSchema(value),
151            )))
152        }
153        Request::Subscribe(req) => {
154            let inner = subscribe::execute(executor, req, agent_arguments).await?;
155            Box::pin(inner.map(|r| r.map(ResponseItem::Subscribe)))
156        }
157        Request::SubscribeRequestSchema(req) => {
158            let value = subscribe::request_schema::execute(executor, req, agent_arguments).await?;
159            Box::pin(crate::cli::command::StreamOnce::new(Ok(
160                ResponseItem::SubscribeRequestSchema(value),
161            )))
162        }
163        Request::SubscribeResponseSchema(req) => {
164            let value = subscribe::response_schema::execute(executor, req, agent_arguments).await?;
165            Box::pin(crate::cli::command::StreamOnce::new(Ok(
166                ResponseItem::SubscribeResponseSchema(value),
167            )))
168        }
169    };
170    Ok(stream)
171}
172
173#[cfg(feature = "cli-executor")]
174pub async fn execute_transform<E: crate::cli::command::CommandExecutor>(
175    executor: &E,
176    request: Request,
177    transform: crate::cli::command::Transform,
178    agent_arguments: Option<&crate::cli::command::AgentArguments>,
179) -> Result<
180    std::pin::Pin<Box<dyn futures::Stream<Item = Result<serde_json::Value, E::Error>> + Send>>,
181    E::Error,
182> {
183    let stream: std::pin::Pin<
184        Box<dyn futures::Stream<Item = Result<serde_json::Value, E::Error>> + Send>,
185    > = match request {
186        Request::Get(req) => {
187            let value = get::execute_transform(executor, req, transform, agent_arguments).await?;
188            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
189        }
190        Request::GetRequestSchema(req) => {
191            let value =
192                get::request_schema::execute_transform(executor, req, transform, agent_arguments).await?;
193            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
194        }
195        Request::GetResponseSchema(req) => {
196            let value =
197                get::response_schema::execute_transform(executor, req, transform, agent_arguments).await?;
198            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
199        }
200        Request::Subscribe(req) => {
201            let inner = subscribe::execute_transform(executor, req, transform, agent_arguments).await?;
202            Box::pin(inner)
203        }
204        Request::SubscribeRequestSchema(req) => {
205            let value =
206                subscribe::request_schema::execute_transform(executor, req, transform, agent_arguments).await?;
207            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
208        }
209        Request::SubscribeResponseSchema(req) => {
210            let value =
211                subscribe::response_schema::execute_transform(executor, req, transform, agent_arguments).await?;
212            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
213        }
214    };
215    Ok(stream)
216}
objectiveai_sdk/cli/command/agents/logs/token_usage/mod.rs

objectiveai_sdk/cli/command/agents/logs/token_usage/
mod.rs
