Skip to main content

agent_code_lib/query/
source.rs

1//! Query source tagging.
2//!
3//! Every API call is tagged with its originating source so costs
4//! and errors can be attributed correctly.
5
6/// Identifies where an API call originated from.
7#[derive(Debug, Clone, PartialEq, Eq, Hash)]
8pub enum QuerySource {
9    /// Interactive REPL main thread.
10    ReplMainThread,
11    /// Subagent execution.
12    SubAgent { agent_id: String },
13    /// History compaction.
14    Compact,
15    /// Session memory summarization.
16    SessionMemory,
17    /// Hook execution.
18    Hook { event: String },
19    /// Background task.
20    BackgroundTask { task_id: String },
21    /// One-shot CLI invocation.
22    OneShot,
23}
24
25impl QuerySource {
26    /// Whether this source should retry on transient errors.
27    pub fn should_retry_on_overload(&self) -> bool {
28        match self {
29            Self::ReplMainThread | Self::OneShot => true,
30            Self::SubAgent { .. } | Self::BackgroundTask { .. } => false,
31            Self::Compact | Self::SessionMemory | Self::Hook { .. } => true,
32        }
33    }
34
35    /// String label for telemetry.
36    pub fn label(&self) -> String {
37        match self {
38            Self::ReplMainThread => "repl_main".to_string(),
39            Self::SubAgent { agent_id } => format!("subagent_{agent_id}"),
40            Self::Compact => "compact".to_string(),
41            Self::SessionMemory => "session_memory".to_string(),
42            Self::Hook { event } => format!("hook_{event}"),
43            Self::BackgroundTask { task_id } => format!("bg_task_{task_id}"),
44            Self::OneShot => "oneshot".to_string(),
45        }
46    }
47}