agy-bridge 0.11.0

//! Core types for the agent SDK bridge.
//!
//! This module defines the data structures that model an agent's execution
//! trajectory: individual [`Step`](crate::types::Step)s,
//! [`ToolCallInfo`](crate::types::ToolCallInfo) requests,
//! [`ToolResult`](crate::types::ToolResult) responses, and
//! [`UsageMetadata`](crate::types::UsageMetadata) for token accounting. All types derive
//! `Serialize`/`Deserialize` for JSON interchange with the Python SDK.

use std::{fmt, str::FromStr};

use serde::{Deserialize, Serialize};
use typed_builder::TypedBuilder;

// =============================================================================
// Step / ToolCall / ToolResult types (§1.6)
// =============================================================================

/// Define an SDK enum with `SCREAMING_SNAKE_CASE` serde rename and auto-generated
/// `Display` and `FromStr` impls.
///
/// Each variant maps to a wire-format string. Unrecognized strings parse as `Err`
/// via `FromStr` — they never panic.
///
/// # Syntax
///
/// ```text
/// define_sdk_enum! {
///     /// Doc comment for the enum.
///     EnumName {
///         Variant1 => "WIRE_STRING_1",
///         Variant2 => "WIRE_STRING_2",
///         #[default]
///         Unknown => "UNKNOWN",
///     }
/// }
/// ```
macro_rules! define_sdk_enum {
    (
        $(#[$meta:meta])*
        $name:ident {
            $(
                $(#[$vmeta:meta])*
                $variant:ident => $wire:literal
            ),+ $(,)?
        }
    ) => {
        $(#[$meta])*
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
        #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
        pub enum $name {
            $(
                $(#[$vmeta])*
                $variant,
            )+
        }

        impl fmt::Display for $name {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                let s = match self {
                    $( Self::$variant => $wire, )+
                };
                f.write_str(s)
            }
        }

        impl FromStr for $name {
            type Err = String;

            fn from_str(s: &str) -> Result<Self, Self::Err> {
                match s {
                    $( $wire => Ok(Self::$variant), )+
                    other => Err(format!(concat!("Unrecognized ", stringify!($name), ": {:?}"), other)),
                }
            }
        }
    };
}

/// Like [`define_sdk_enum!`] but for enums where the serde wire format uses
/// per-variant `#[serde(rename = "...")]` instead of `rename_all`.
///
/// This is needed for [`StepTarget`] whose SDK strings have a `TARGET_` prefix
/// that doesn't match the `SCREAMING_SNAKE_CASE` of the enum name.
macro_rules! define_sdk_enum_custom_serde {
    (
        $(#[$meta:meta])*
        $name:ident {
            $(
                $(#[$vmeta:meta])*
                $variant:ident => $wire:literal
            ),+ $(,)?
        }
    ) => {
        $(#[$meta])*
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
        pub enum $name {
            $(
                $(#[$vmeta])*
                #[serde(rename = $wire)]
                $variant,
            )+
        }

        impl fmt::Display for $name {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                let s = match self {
                    $( Self::$variant => $wire, )+
                };
                f.write_str(s)
            }
        }

        impl FromStr for $name {
            type Err = String;

            fn from_str(s: &str) -> Result<Self, Self::Err> {
                match s {
                    $( $wire => Ok(Self::$variant), )+
                    other => Err(format!(concat!("Unrecognized ", stringify!($name), ": {:?}"), other)),
                }
            }
        }
    };
}

define_sdk_enum! {
    /// The high-level type of a step in the agent trajectory.
    StepType {
        /// A textual response from the model.
        TextResponse => "TEXT_RESPONSE",
        /// A tool invocation requested by the model.
        ToolCall => "TOOL_CALL",
        /// A system-generated message (e.g. context injection).
        SystemMessage => "SYSTEM_MESSAGE",
        /// A context-window compaction event.
        Compaction => "COMPACTION",
        /// The agent has signaled task completion.
        Finish => "FINISH",
        /// Internal model thinking / reasoning step.
        Thinking => "THINKING",
        /// Unrecognized step type (forward-compatibility fallback).
        #[default]
        Unknown => "UNKNOWN",
    }
}

define_sdk_enum! {
    /// The source that generated a step.
    StepSource {
        /// Generated by the system runtime.
        System => "SYSTEM",
        /// Provided by the user.
        User => "USER",
        /// Generated by the model.
        Model => "MODEL",
        /// Unrecognized source (forward-compatibility fallback).
        #[default]
        Unknown => "UNKNOWN",
    }
}

define_sdk_enum! {
    /// The execution status of a step.
    StepStatus {
        /// Step is currently executing.
        Active => "ACTIVE",
        /// Step completed successfully.
        Done => "DONE",
        /// Step is blocked waiting for user input.
        WaitingForUser => "WAITING_FOR_USER",
        /// Step failed with an error.
        Error => "ERROR",
        /// Step was canceled before completion.
        Canceled => "CANCELED",
        /// Unrecognized status (forward-compatibility fallback).
        #[default]
        Unknown => "UNKNOWN",
    }
}

define_sdk_enum_custom_serde! {
    /// Target of a step interaction, mirroring the Python SDK's `StepTarget`.
    ///
    /// The Python SDK uses `TARGET_` prefixed strings (e.g. `TARGET_USER`).
    /// Uses per-variant `#[serde(rename)]` because the SDK's wire format has a
    /// `TARGET_` prefix that doesn't follow `SCREAMING_SNAKE_CASE` of the enum name.
    StepTarget {
        /// Step is directed at the model.
        Model => "TARGET_MODEL",
        /// Step is directed at the user.
        User => "TARGET_USER",
        /// Step is directed at the environment (tool execution).
        Environment => "TARGET_ENVIRONMENT",
        /// Target is unspecified.
        Unspecified => "TARGET_UNSPECIFIED",
        /// Unknown target (fallback).
        #[default]
        Unknown => "UNKNOWN",
    }
}

define_sdk_enum! {
    /// Reason why an agent trajectory stopped or halted.
    StopReason {
        /// Reason unspecified.
        Unspecified => "UNSPECIFIED",
        /// Maximum number of model calls exceeded.
        MaxModelCallsExceeded => "MAX_MODEL_CALLS_EXCEEDED",
        /// Maximum number of tool calls exceeded.
        MaxToolCallsExceeded => "MAX_TOOL_CALLS_EXCEEDED",
        /// Maximum prompt/input tokens exceeded.
        MaxInputTokensExceeded => "MAX_INPUT_TOKENS_EXCEEDED",
        /// Maximum output tokens exceeded.
        MaxOutputTokensExceeded => "MAX_OUTPUT_TOKENS_EXCEEDED",
        /// Maximum total tokens exceeded.
        MaxTotalTokensExceeded => "MAX_TOTAL_TOKENS_EXCEEDED",
        /// Remote API quota exhausted.
        QuotaExhausted => "QUOTA_EXHAUSTED",
        /// Unrecognized stop reason (fallback).
        #[default]
        Unknown => "UNKNOWN",
    }
}

define_sdk_enum! {
    /// Token modality.
    Modality {
        /// Unspecified modality.
        Unspecified => "MODALITY_UNSPECIFIED",
        /// Text tokens.
        Text => "TEXT",
        /// Image tokens.
        Image => "IMAGE",
        /// Video tokens.
        Video => "VIDEO",
        /// Audio tokens.
        Audio => "AUDIO",
        /// Document tokens.
        Document => "DOCUMENT",
        /// Unrecognized modality.
        #[default]
        Unknown => "UNKNOWN",
    }
}

/// Token count for a specific modality.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct ModalityTokenCount {
    /// The modality of tokens (e.g. TEXT, IMAGE).
    pub modality: Modality,
    /// The number of tokens.
    pub token_count: u64,
}

/// A tool call from the model, mirroring the Python SDK's `ToolCall`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCallInfo {
    /// Tool name — either a `BuiltinTools` string or a custom tool name.
    pub name: String,
    /// Arguments as a JSON value (typically an object/dict).
    #[serde(default)]
    pub args: serde_json::Value,
    /// Optional unique identifier for the call.
    #[serde(default)]
    pub id: Option<String>,
    /// Optional normalized filesystem path for file-related tools.
    #[serde(default)]
    pub canonical_path: Option<String>,
}

/// Result of a single tool execution, mirroring the Python SDK's `ToolResult`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolResult {
    /// The name of the tool that was executed.
    pub name: String,
    /// Optional identifier correlating this result with a `ToolCallInfo.id`.
    #[serde(default)]
    pub id: Option<String>,
    /// The tool's return value (any JSON-serializable value).
    #[serde(default)]
    pub result: serde_json::Value,
    /// An error message if execution failed, or `None` on success.
    #[serde(default)]
    pub error: Option<String>,
}

/// Token usage metadata from the model API, mirroring the SDK's `UsageMetadata`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct UsageMetadata {
    /// Number of tokens in the prompt.
    #[serde(default)]
    pub prompt_token_count: Option<u64>,
    /// Number of tokens from cached content (subset of prompt tokens).
    #[serde(default)]
    pub cached_content_token_count: Option<u64>,
    /// Number of tokens in the generated candidates (excluding thinking).
    #[serde(default)]
    pub candidates_token_count: Option<u64>,
    /// Number of tokens used for thinking/reasoning.
    #[serde(default)]
    pub thoughts_token_count: Option<u64>,
    /// Sum of prompt + candidates + thinking tokens.
    #[serde(default)]
    pub total_token_count: Option<u64>,
    /// Detailed prompt tokens by modality.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub prompt_tokens_details: Vec<ModalityTokenCount>,
    /// Detailed cache tokens by modality.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub cache_tokens_details: Vec<ModalityTokenCount>,
    /// Detailed candidate output tokens by modality.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub candidates_tokens_details: Vec<ModalityTokenCount>,
    /// Detailed tool use prompt tokens by modality.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_use_prompt_tokens_details: Vec<ModalityTokenCount>,
}

/// The role of a message author in the conversation.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[derive(Default)]
pub enum MessageRole {
    /// A user-authored message.
    #[default]
    User,
    /// A model-generated message.
    Model,
    /// A system-level message.
    System,
    /// An unrecognized role — preserves the original string for forward
    /// compatibility with new SDK roles.
    #[serde(untagged)]
    Unknown(String),
}

impl std::fmt::Display for MessageRole {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::User => f.write_str("user"),
            Self::Model => f.write_str("model"),
            Self::System => f.write_str("system"),
            Self::Unknown(s) => f.write_str(s),
        }
    }
}

/// A single message in the conversation history, mirroring the Python SDK's
/// `ConversationMessage`.
///
/// Each message has a [`MessageRole`] and textual `content`.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ConversationMessage {
    /// The role of the message author.
    #[serde(default)]
    pub role: MessageRole,
    /// The textual content of the message.
    #[serde(default)]
    pub content: String,
}

/// A single step in the agent trajectory, mirroring the SDK's `Step`.
///
/// # Construction
///
/// `Step` is `#[non_exhaustive]`, so outside this crate it can only be built
/// with the [`TypedBuilder`]. Every field defaults (matching [`Default`]), so
/// callers set only the fields they care about:
///
/// ```
/// use agy_bridge::Step;
///
/// let step = Step::builder()
///     .id("traj:0")
///     .content("Running command...")
///     .build();
/// assert_eq!(step.id, "traj:0");
/// // Unset fields fall back to their defaults.
/// assert_eq!(step.step_index, 0);
/// assert!(step.tool_calls.is_empty());
/// ```
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default, TypedBuilder)]
#[builder(field_defaults(default))]
pub struct Step {
    /// Unique string identifier for the step.
    #[serde(default)]
    #[builder(setter(into))]
    pub id: String,
    /// Integer index of the step in the trajectory.
    #[serde(default)]
    pub step_index: u32,
    /// Identifier of the trajectory this step belongs to.
    ///
    /// The SDK's connection layer assigns one trajectory id per agent
    /// trajectory. Compare with [`cascade_id`](Self::cascade_id) to tell a
    /// primary-agent step (`trajectory_id == cascade_id`) apart from a subagent
    /// step (`trajectory_id != cascade_id`); see [`Step::is_subagent_step`].
    /// Empty when the harness did not report one.
    #[serde(default)]
    #[builder(setter(into))]
    pub trajectory_id: String,
    /// Identifier of the top-level cascade — i.e. the primary agent trajectory.
    ///
    /// Every step in a turn, primary and subagent alike, shares the same
    /// `cascade_id`, so it identifies the parent trajectory. A step originates
    /// from a subagent when its [`trajectory_id`](Self::trajectory_id) differs
    /// from this value. Empty when the harness did not report one.
    #[serde(default)]
    #[builder(setter(into))]
    pub cascade_id: String,
    /// Identifier of the direct parent trajectory for nested subagents.
    #[serde(default)]
    #[builder(setter(into))]
    pub parent_trajectory_id: String,
    /// Nesting depth of this step (0 for primary agent, 1 for direct subagent, etc.).
    #[serde(default)]
    pub depth: u32,
    /// The high-level type of the step.
    #[serde(default, rename = "type")]
    pub step_type: StepType,
    /// The source that generated the step.
    #[serde(default)]
    pub source: StepSource,
    /// The target of the step interaction.
    #[serde(default)]
    pub target: StepTarget,
    /// The status of the step.
    #[serde(default)]
    pub status: StepStatus,
    /// The text content/output of the step.
    #[serde(default)]
    #[builder(setter(into))]
    pub content: String,
    /// Incremental text content added since the last update for this step.
    #[serde(default)]
    #[builder(setter(into))]
    pub content_delta: String,
    /// Full model reasoning/thinking text for planner responses.
    #[serde(default)]
    #[builder(setter(into))]
    pub thinking: String,
    /// Incremental thinking text added since the last update for this step.
    #[serde(default)]
    #[builder(setter(into))]
    pub thinking_delta: String,
    /// List of tool calls associated with the step.
    #[serde(default)]
    #[builder(setter(transform = |v: impl IntoIterator<Item = impl Into<ToolCallInfo>>| v.into_iter().map(Into::into).collect()))]
    pub tool_calls: Vec<ToolCallInfo>,
    /// Short error message if the step failed.
    #[serde(default)]
    #[builder(setter(into))]
    pub error: String,
    /// HTTP status code from the harness error, if any (e.g. 400, 429, 503).
    ///
    /// The SDK populates this from the harness's `error.http_code` field.
    /// Used by error detection in `forward_step_to_writer` for logging.
    #[serde(default)]
    pub http_code: u16,
    /// Whether this step is a completed model response directed at the user.
    ///
    /// Multiple steps per turn may have this flag set; consumers wanting only
    /// the last response should iterate fully.
    #[serde(default)]
    #[builder(setter(strip_option))]
    pub is_complete_response: Option<bool>,
    /// Structured output payload extracted from the FINISH step.
    ///
    /// This is `serde_json::Value` because it contains user-defined schema data
    /// whose shape is not known at compile time.
    #[serde(default)]
    #[builder(setter(strip_option))]
    pub structured_output: Option<serde_json::Value>,
    /// Token usage for this step's model invocation.
    #[serde(default)]
    #[builder(setter(strip_option))]
    pub usage_metadata: Option<UsageMetadata>,
}

impl Step {
    /// Whether this step originates from a subagent trajectory rather than the
    /// primary agent's.
    ///
    /// Mirrors the SDK's own parent/subagent discrimination
    /// (`cascade_id AND trajectory_id AND trajectory_id != cascade_id`): a step
    /// is a subagent step only when it carries a known parent
    /// [`cascade_id`](Self::cascade_id) *and* a
    /// [`trajectory_id`](Self::trajectory_id) that differs from it. Steps
    /// missing either id — e.g. from mocks or older harnesses — are treated as
    /// primary.
    #[must_use]
    pub fn is_subagent_step(&self) -> bool {
        !self.cascade_id.is_empty()
            && !self.trajectory_id.is_empty()
            && self.trajectory_id != self.cascade_id
    }
}

#[cfg(feature = "python")]
macro_rules! impl_from_py_object {
    ($($t:ty),+) => {
        $(
            impl<'a, 'py> pyo3::FromPyObject<'a, 'py> for $t {
                type Error = pyo3::PyErr;

                fn extract(ob: pyo3::Borrowed<'a, 'py, pyo3::PyAny>) -> pyo3::PyResult<Self> {
                    crate::runtime::py_scripts::warm_up_lazy_imports(ob.py());
                    pythonize::depythonize(&*ob).map_err(|e| {
                        pyo3::exceptions::PyValueError::new_err(format!(
                            "Failed to deserialize {} from Python dict: {}",
                            stringify!($t),
                            e
                        ))
                    })
                }
            }
        )+
    };
}

#[cfg(feature = "python")]
impl_from_py_object!(
    StepType,
    StepSource,
    StepStatus,
    StepTarget,
    StopReason,
    Modality,
    ModalityTokenCount,
    ToolCallInfo,
    ToolResult,
    UsageMetadata,
    MessageRole,
    ConversationMessage,
    Step
);

#[cfg(test)]
#[cfg(test)]
#[path = "types_tests.rs"]
mod tests;