agent-runtime-ports 0.1.1

//! Provider-neutral capabilities required by the Runtime Kernel.
//!
//! This crate contains no HTTP client, database driver, Axum type, Agent Infra
//! DTO, or provider implementation.

use std::{
    collections::BTreeMap,
    time::{Duration, Instant},
};

use async_trait::async_trait;
use runtime_types::{
    CallerScope, ConversationId, DelegationLeaseRef, ExecutionId, OperationId, RequestAuthority,
    RuntimeInstanceId, WorkspaceId,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use tokio_util::sync::CancellationToken;

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

impl OperationContext {
    pub fn remaining(&self) -> Result<Duration, PortFailure> {
        if self.cancellation.is_cancelled() {
            return Err(PortFailure::canceled());
        }
        self.deadline
            .checked_duration_since(Instant::now())
            .filter(|remaining| !remaining.is_zero())
            .ok_or_else(PortFailure::timeout)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommitDisposition {
    NotCommitted,
    Committed,
    Unknown,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PortFailureKind {
    Invalid,
    NotFound,
    Forbidden,
    Conflict,
    RateLimited,
    Timeout,
    Canceled,
    Unavailable,
    Protocol,
    Internal,
}

#[derive(Debug, Error, Clone)]
#[error("{code}: {message}")]
pub struct PortFailure {
    pub kind: PortFailureKind,
    pub code: String,
    pub message: String,
    pub retryable: bool,
    pub commit: CommitDisposition,
}

impl PortFailure {
    pub fn new(kind: PortFailureKind, code: impl Into<String>, message: impl Into<String>) -> Self {
        Self {
            kind,
            code: code.into(),
            message: message.into(),
            retryable: false,
            commit: CommitDisposition::NotCommitted,
        }
    }
    pub fn timeout() -> Self {
        Self::new(
            PortFailureKind::Timeout,
            "DEADLINE_EXCEEDED",
            "operation deadline exceeded",
        )
    }
    pub fn canceled() -> Self {
        Self::new(PortFailureKind::Canceled, "CANCELED", "operation canceled")
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ModelContent {
    Text {
        text: String,
    },
    ToolUse {
        id: String,
        name: String,
        arguments: Value,
    },
    ToolResult {
        tool_use_id: String,
        name: String,
        content: String,
        failed: bool,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelRole {
    System,
    User,
    Assistant,
    Tool,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelMessage {
    pub role: ModelRole,
    pub content: Vec<ModelContent>,
}

impl ModelMessage {
    pub fn text(role: ModelRole, text: impl Into<String>) -> Self {
        Self {
            role,
            content: vec![ModelContent::Text { text: text.into() }],
        }
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelToolDefinition {
    pub name: String,
    pub description: String,
    pub input_schema: Value,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelRequest {
    /// Stable across transport attempts and shared with Context invocation
    /// audit. This is correlation identity, not permission to retry a model
    /// side effect whose commit disposition is unknown.
    pub invocation_id: String,
    pub model: String,
    pub messages: Vec<ModelMessage>,
    pub tools: Vec<ModelToolDefinition>,
    pub max_output_tokens: Option<u32>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelFinish {
    Stop,
    ToolCalls,
    Length,
    ContentFilter,
    Other,
}

impl ModelFinish {
    pub fn is_complete(self) -> bool {
        matches!(self, Self::Stop | Self::ToolCalls)
    }
}

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TokenUsage {
    /// `None` means the provider did not report this value. It must never be
    /// collapsed to zero because zero is a real, materially different usage.
    pub input_tokens: Option<u64>,
    pub output_tokens: Option<u64>,
}

impl TokenUsage {
    /// Additive identity used only when beginning a known aggregate.
    pub const fn zero() -> Self {
        Self {
            input_tokens: Some(0),
            output_tokens: Some(0),
        }
    }

    /// Aggregate without inventing missing provider measurements. Once any
    /// turn is unknown, the corresponding execution total stays unknown.
    pub fn add_assign(&mut self, value: &Self) {
        self.input_tokens = add_known(self.input_tokens, value.input_tokens);
        self.output_tokens = add_known(self.output_tokens, value.output_tokens);
    }
}

fn add_known(left: Option<u64>, right: Option<u64>) -> Option<u64> {
    left.zip(right)
        .map(|(left, right)| left.saturating_add(right))
}

#[cfg(test)]
mod token_usage_tests {
    use super::TokenUsage;

    #[test]
    fn aggregate_preserves_unknown_and_saturates_known_values() {
        let mut usage = TokenUsage::zero();
        usage.add_assign(&TokenUsage {
            input_tokens: Some(u64::MAX),
            output_tokens: Some(3),
        });
        usage.add_assign(&TokenUsage {
            input_tokens: Some(1),
            output_tokens: None,
        });
        assert_eq!(usage.input_tokens, Some(u64::MAX));
        assert_eq!(usage.output_tokens, None);
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelResponse {
    pub output: Vec<ModelContent>,
    pub finish: ModelFinish,
    pub usage: TokenUsage,
}

#[async_trait]
pub trait ModelSession: Send + Sync {
    async fn invoke(
        &self,
        operation: &OperationContext,
        request: ModelRequest,
    ) -> Result<ModelResponse, PortFailure>;
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConversationMessage {
    pub role: ModelRole,
    pub text: String,
}

#[async_trait]
pub trait ConversationSession: Send + Sync {
    async fn load_recent(
        &self,
        operation: &OperationContext,
        limit: usize,
    ) -> Result<Vec<ConversationMessage>, PortFailure>;
    async fn append(
        &self,
        operation: &OperationContext,
        messages: &[ConversationMessage],
    ) -> Result<(), PortFailure>;
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContextAuditStage {
    Original,
    Assembled,
    Compressed,
}

#[derive(Debug, Clone)]
pub struct ContextAuditSnapshot {
    pub id: String,
    pub stage: ContextAuditStage,
    pub purpose: String,
    pub messages: Vec<ModelMessage>,
    pub tools: Vec<ModelToolDefinition>,
    pub protocol_overhead_tokens: u64,
    pub reserved_output_tokens: u64,
    pub estimated_tokens: u64,
    /// Digest of the exact messages, Tool projection, and output reservation
    /// used for the corresponding model request.
    pub request_digest: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContextDerivationKind {
    Assemble,
    Compress,
    Trim,
    Reorder,
}

#[derive(Debug, Clone)]
pub struct ContextAuditDerivation {
    pub id: String,
    pub source_snapshot_id: String,
    pub target_snapshot_id: String,
    pub kind: ContextDerivationKind,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub input_items: usize,
    pub output_items: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelInvocationAuditState {
    Prepared,
    Succeeded,
    Failed,
}

#[derive(Debug, Clone)]
pub struct ModelInvocationAudit {
    pub id: String,
    pub request_snapshot_id: String,
    pub model: String,
    pub state: ModelInvocationAuditState,
    pub usage: TokenUsage,
    pub error_code: Option<String>,
}

#[async_trait]
pub trait ContextAuditSession: Send + Sync {
    async fn snapshot(
        &self,
        operation: &OperationContext,
        snapshot: ContextAuditSnapshot,
    ) -> Result<(), PortFailure>;
    async fn derivation(
        &self,
        operation: &OperationContext,
        derivation: ContextAuditDerivation,
    ) -> Result<(), PortFailure>;
    async fn model_invocation(
        &self,
        operation: &OperationContext,
        invocation: ModelInvocationAudit,
    ) -> Result<(), PortFailure>;
}

#[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 CommandRequest {
    pub command: String,
    pub cwd: Option<String>,
}

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

#[async_trait]
pub trait WorkspaceSession: Send + Sync {
    async fn describe(&self, operation: &OperationContext) -> Result<String, PortFailure>;
    async fn read(&self, operation: &OperationContext, path: &str) -> Result<Vec<u8>, PortFailure>;
    async fn write(
        &self,
        operation: &OperationContext,
        path: &str,
        content: &[u8],
    ) -> Result<(), PortFailure>;
    async fn list(
        &self,
        operation: &OperationContext,
        path: &str,
    ) -> Result<Vec<WorkspaceEntry>, PortFailure>;
    async fn search(
        &self,
        operation: &OperationContext,
        path: &str,
        query: &str,
        limit: usize,
    ) -> Result<Vec<String>, PortFailure>;
    async fn execute(
        &self,
        operation: &OperationContext,
        request: CommandRequest,
    ) -> Result<CommandOutput, PortFailure>;
}

#[derive(Debug, Clone)]
pub struct ResolvedExecutionContext {
    pub runtime_instance_id: RuntimeInstanceId,
    pub agent_id: String,
    pub runtime_type: String,
    pub model: String,
    pub workspace_id: WorkspaceId,
    pub package_id: String,
    pub package_version: String,
    pub package_digest: Option<String>,
    pub metadata: BTreeMap<String, String>,
}

#[async_trait]
pub trait RuntimeInstanceResolver: Send + Sync {
    async fn resolve(
        &self,
        operation: &OperationContext,
        caller: &CallerScope,
        id: &RuntimeInstanceId,
    ) -> Result<ResolvedExecutionContext, PortFailure>;
}

#[async_trait]
pub trait TraceSession: Send + Sync {
    async fn append(
        &self,
        operation: &OperationContext,
        event_type: &str,
        payload: Value,
    ) -> Result<(), PortFailure>;
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InteractionRequest {
    pub request_id: String,
    pub prompt: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InteractionResponse {
    pub text: String,
}

#[async_trait]
pub trait InteractionSession: Send + Sync {
    /// Register a pending request before it becomes externally visible.
    async fn prepare(
        &self,
        operation: &OperationContext,
        request: InteractionRequest,
    ) -> Result<(), PortFailure>;

    /// Wait for a response to a request registered by [`Self::prepare`].
    async fn wait(
        &self,
        operation: &OperationContext,
        request_id: &str,
    ) -> Result<InteractionResponse, PortFailure>;

    async fn request(
        &self,
        operation: &OperationContext,
        request: InteractionRequest,
    ) -> Result<InteractionResponse, PortFailure> {
        let request_id = request.request_id.clone();
        self.prepare(operation, request).await?;
        self.wait(operation, &request_id).await
    }
}

#[derive(Debug, Clone)]
pub struct SubagentRequest {
    pub request_id: String,
    pub prompt: String,
    pub max_model_turns: usize,
}

#[derive(Debug, Clone)]
pub struct SubagentOutcome {
    pub answer: String,
}

#[async_trait]
pub trait SubagentSession: Send + Sync {
    async fn execute(
        &self,
        operation: &OperationContext,
        request: SubagentRequest,
    ) -> Result<SubagentOutcome, PortFailure>;
}

#[derive(Debug, Clone)]
pub struct ResourceSnapshot {
    pub instructions: Vec<String>,
    pub skills: BTreeMap<String, String>,
    pub digest: String,
}

#[async_trait]
pub trait ResourceSession: Send + Sync {
    async fn load(&self, operation: &OperationContext) -> Result<ResourceSnapshot, PortFailure>;
}

#[derive(Debug, Clone)]
pub struct SessionScope {
    pub execution_id: runtime_types::ExecutionId,
    pub conversation_id: Option<ConversationId>,
    pub resolved: ResolvedExecutionContext,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExecutionDelegationLease {
    pub lease_ref: DelegationLeaseRef,
    pub expires_at_seconds: u64,
    pub revision: u64,
}

pub struct ExecutionSessions {
    pub model: std::sync::Arc<dyn ModelSession>,
    pub conversation: std::sync::Arc<dyn ConversationSession>,
    pub context_audit: std::sync::Arc<dyn ContextAuditSession>,
    pub workspace: std::sync::Arc<dyn WorkspaceSession>,
    pub trace: std::sync::Arc<dyn TraceSession>,
    pub interaction: std::sync::Arc<dyn InteractionSession>,
    pub subagent: std::sync::Arc<dyn SubagentSession>,
    pub resources: std::sync::Arc<dyn ResourceSession>,
}

#[async_trait]
pub trait ExecutionSessionFactory: Send + Sync {
    async fn establish_delegation(
        &self,
        operation: &OperationContext,
        authority: &RequestAuthority,
        scope: &SessionScope,
    ) -> Result<ExecutionDelegationLease, PortFailure>;

    async fn create(
        &self,
        operation: &OperationContext,
        caller: &CallerScope,
        delegation: &ExecutionDelegationLease,
        scope: &SessionScope,
    ) -> Result<ExecutionSessions, PortFailure>;

    async fn renew_delegation(
        &self,
        operation: &OperationContext,
        caller: &CallerScope,
        delegation: &ExecutionDelegationLease,
    ) -> Result<ExecutionDelegationLease, PortFailure>;

    async fn revoke_delegation(
        &self,
        operation: &OperationContext,
        caller: &CallerScope,
        delegation: &ExecutionDelegationLease,
    ) -> Result<(), PortFailure>;

    /// Deliver input to a currently waiting execution. Implementations must
    /// deduplicate `operation_id` and reject input when no request is pending.
    async fn submit_input(
        &self,
        operation: &OperationContext,
        caller: &CallerScope,
        execution_id: &runtime_types::ExecutionId,
        operation_id: &OperationId,
        request_id: &str,
        response: InteractionResponse,
    ) -> Result<(), PortFailure>;
}