agent-runtime-extension-api 0.1.1

//! Public extension contract used by first-party and third-party Agent definitions.

use std::{collections::BTreeMap, sync::Arc, time::Instant};

use async_trait::async_trait;
use runtime_types::{DefinitionId, ExecutionId, ToolCallId};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use tokio_util::sync::CancellationToken;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ToolDescriptor {
    pub name: String,
    pub description: String,
    pub input_schema: Value,
    #[serde(default)]
    pub required_capabilities: Vec<String>,
    #[serde(default)]
    pub side_effect: SideEffect,
}

#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum SideEffect {
    #[default]
    ReadOnly,
    WorkspaceWrite,
    Process,
    External,
}

#[derive(Debug, Error, Clone)]
pub enum ExtensionError {
    #[error("invalid extension input: {0}")]
    Invalid(String),
    #[error("extension permission denied: {0}")]
    Denied(String),
    #[error("extension operation failed: {0}")]
    Failed(String),
    #[error("extension operation canceled")]
    Canceled,
    #[error("extension operation timed out")]
    Timeout,
}

#[derive(Debug, Clone)]
pub struct ExtensionContext {
    pub execution_id: ExecutionId,
    pub deadline: Instant,
    pub cancellation: CancellationToken,
}

#[derive(Clone)]
pub struct ToolContext {
    pub extension: ExtensionContext,
    pub call_id: ToolCallId,
    pub workspace: Arc<dyn WorkspaceFacade>,
    pub interaction: Arc<dyn InteractionFacade>,
    pub subagent: Arc<dyn SubagentFacade>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct ToolOutcome {
    pub content: String,
    pub failed: bool,
    pub completion: Option<String>,
    pub metadata: BTreeMap<String, Value>,
    pub context_invalidation: ContextInvalidation,
}

impl ToolOutcome {
    pub fn success(content: impl Into<String>) -> Self {
        Self {
            content: content.into(),
            failed: false,
            completion: None,
            metadata: BTreeMap::new(),
            context_invalidation: ContextInvalidation::NONE,
        }
    }
    pub fn failure(content: impl Into<String>) -> Self {
        Self {
            content: content.into(),
            failed: true,
            completion: None,
            metadata: BTreeMap::new(),
            context_invalidation: ContextInvalidation::NONE,
        }
    }

    pub fn complete(answer: impl Into<String>) -> Self {
        let answer = answer.into();
        Self {
            content: answer.clone(),
            failed: false,
            completion: Some(answer),
            metadata: BTreeMap::new(),
            context_invalidation: ContextInvalidation::NONE,
        }
    }

    pub fn with_context_invalidation(mut self, invalidation: ContextInvalidation) -> Self {
        self.context_invalidation = invalidation;
        self
    }
}

#[async_trait]
pub trait Tool: Send + Sync {
    fn descriptor(&self) -> ToolDescriptor;
    async fn invoke(
        &self,
        context: ToolContext,
        arguments: Value,
    ) -> Result<ToolOutcome, ExtensionError>;
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WorkspaceEntry {
    pub path: String,
    pub is_dir: bool,
    pub size: Option<u64>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CommandOutcome {
    pub exit_code: i32,
    pub stdout: String,
    pub stderr: String,
    pub truncated: bool,
}

#[async_trait]
pub trait WorkspaceFacade: Send + Sync {
    async fn describe(&self) -> Result<String, ExtensionError>;
    async fn read_text(&self, path: &str) -> Result<String, ExtensionError>;
    async fn write_text(&self, path: &str, content: &str) -> Result<(), ExtensionError>;
    async fn list(&self, path: &str) -> Result<Vec<WorkspaceEntry>, ExtensionError>;
    async fn search(
        &self,
        path: &str,
        query: &str,
        limit: usize,
    ) -> Result<Vec<String>, ExtensionError>;
    async fn execute(
        &self,
        command: &str,
        cwd: Option<&str>,
    ) -> Result<CommandOutcome, ExtensionError>;
}

#[async_trait]
pub trait InteractionFacade: Send + Sync {
    async fn request_text(&self, prompt: &str) -> Result<String, ExtensionError>;
}

#[async_trait]
pub trait SubagentFacade: Send + Sync {
    async fn execute(&self, prompt: &str, max_model_turns: usize)
    -> Result<String, ExtensionError>;
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextItem {
    pub source: String,
    pub content: String,
    pub priority: i32,
    pub required: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ContextSlot {
    System,
    Instructions,
    ToolCatalog,
    Skills,
    Memory,
    Plan,
    RuntimeState,
    History,
    Observations,
    Checkpoint,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContextSensitivity {
    Public,
    Internal,
    Sensitive,
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ContextInvalidation(u16);

impl ContextInvalidation {
    pub const NONE: Self = Self(0);
    pub const SYSTEM: Self = Self(1 << 0);
    pub const TOOL_CATALOG: Self = Self(1 << 1);
    pub const SKILLS: Self = Self(1 << 2);
    pub const MEMORY: Self = Self(1 << 3);
    pub const PLAN: Self = Self(1 << 4);
    pub const RUNTIME_STATE: Self = Self(1 << 5);

    pub const fn union(self, other: Self) -> Self {
        Self(self.0 | other.0)
    }

    pub const fn contains(self, other: Self) -> bool {
        self.0 & other.0 == other.0
    }

    pub const fn is_empty(self) -> bool {
        self.0 == 0
    }
}

impl std::ops::BitOr for ContextInvalidation {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        self.union(rhs)
    }
}

impl std::ops::BitOrAssign for ContextInvalidation {
    fn bitor_assign(&mut self, rhs: Self) {
        *self = self.union(rhs);
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextResource {
    pub id: String,
    pub slot: ContextSlot,
    pub source: String,
    pub content: String,
    pub required: bool,
    pub sensitivity: ContextSensitivity,
    pub version: String,
    /// Stable content identity used by the Kernel to avoid treating an
    /// unchanged resource as a refresh.
    pub digest: String,
    /// Agent-side estimate for diagnostics. The Kernel always measures the
    /// final serialized request independently and remains authoritative.
    pub estimated_tokens: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextResourceKey {
    pub slot: ContextSlot,
    pub id: String,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ContextPatch {
    pub resources: Vec<ContextResource>,
    pub removals: Vec<ContextResourceKey>,
    pub invalidations: ContextInvalidation,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ContextEvent {
    ExecutionStart,
    BeforeModel {
        turn: usize,
    },
    AfterTool {
        turn: usize,
        tool_name: String,
        failed: bool,
        invalidation: ContextInvalidation,
    },
    BeforeCompaction {
        turn: usize,
    },
    AfterCompaction {
        turn: usize,
    },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextFrame {
    pub execution_id: ExecutionId,
    pub turn: usize,
    pub user_prompt: String,
    pub visible_tools: Vec<String>,
    pub pending_invalidations: ContextInvalidation,
    pub resources: Vec<ContextResource>,
    pub budget: Option<ContextBudgetFrame>,
    pub latest_tool: Option<ContextToolFrame>,
    pub transcript_summary: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextBudgetFrame {
    pub max_input_tokens: u64,
    pub message_tokens: u64,
    pub tool_schema_tokens: u64,
    pub protocol_overhead_tokens: u64,
    pub reserved_output_tokens: u64,
    pub total_tokens: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextToolFrame {
    pub name: String,
    pub failed: bool,
}

#[async_trait]
pub trait ContextHook: Send + Sync {
    fn name(&self) -> &str;

    /// Stable authoring fingerprint for the Hook implementation/configuration.
    ///
    /// The Definition builder incorporates this value into its immutable
    /// digest. Change it whenever Hook behavior or embedded static resources
    /// change. Per-execution dynamic resource content still carries its own
    /// `ContextResource.digest` at runtime.
    fn definition_fingerprint(&self) -> String;

    async fn apply(
        &self,
        context: &ExtensionContext,
        event: &ContextEvent,
        frame: &ContextFrame,
    ) -> Result<ContextPatch, ExtensionError>;
}

#[derive(Debug, Clone)]
pub struct ToolInvocation {
    pub execution_id: ExecutionId,
    pub call_id: ToolCallId,
    pub descriptor: ToolDescriptor,
    pub arguments: Value,
}

#[async_trait]
pub trait Guard: Send + Sync {
    async fn authorize(&self, invocation: &ToolInvocation) -> Result<(), ExtensionError>;
}

#[derive(Debug, Clone)]
pub enum ObserverEvent {
    ExecutionStarted {
        execution_id: ExecutionId,
    },
    ModelCompleted {
        execution_id: ExecutionId,
        turn: usize,
    },
    ToolCompleted {
        execution_id: ExecutionId,
        call_id: ToolCallId,
        name: String,
        failed: bool,
    },
    ExecutionFinished {
        execution_id: ExecutionId,
        succeeded: bool,
    },
}

#[async_trait]
pub trait Observer: Send + Sync {
    async fn observe(&self, event: ObserverEvent) -> Result<(), ExtensionError>;
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct DefinitionManifest {
    pub id: DefinitionId,
    pub version: String,
    pub runtime_type: String,
    pub schema_version: u32,
    pub digest: String,
}

pub struct AgentDefinition {
    pub manifest: DefinitionManifest,
    pub system_prompt: String,
    pub tools: Vec<Arc<dyn Tool>>,
    pub context_hooks: Vec<Arc<dyn ContextHook>>,
    pub guards: Vec<Arc<dyn Guard>>,
    pub observers: Vec<Arc<dyn Observer>>,
    pub completion_mode: CompletionMode,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CompletionMode {
    ModelOrTool,
    RequiredTool(String),
}

impl AgentDefinition {
    pub fn tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.tools
            .iter()
            .find(|tool| tool.descriptor().name == name)
            .cloned()
    }
}