agent-model-contract 0.2.0

//! Transport-neutral domain contract for Model Infra.

mod usage;
pub use usage::{TokenUsage, TokenUsageCoverage};

use std::collections::{BTreeMap, BTreeSet};
use std::fmt;

use agent_runtime_identity_contract::{
    CredentialCapability, RuntimeIdentityId, TenantProjectScope,
};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;

pub const SERVICE_NAME: &str = "agent-model-infra";
pub const AUDIENCE: &str = "agent-model-infra";
pub const DEFAULT_PORT: u16 = 5130;

macro_rules! identifier {
    ($name:ident, $label:literal, $max:expr) => {
        #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
        #[serde(transparent)]
        pub struct $name(String);

        impl $name {
            pub fn new(value: impl Into<String>) -> Result<Self, String> {
                let value = value.into();
                validate_identifier($label, &value, $max)?;
                Ok(Self(value))
            }

            pub fn as_str(&self) -> &str {
                &self.0
            }
        }

        impl fmt::Display for $name {
            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                self.0.fmt(formatter)
            }
        }
    };
}

identifier!(ModelProfileId, "model_profile_id", 128);
identifier!(ModelConfigId, "model_config_id", 128);
identifier!(InvocationId, "invocation_id", 255);

#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct SecretRef(String);

impl SecretRef {
    pub fn new(value: impl Into<String>) -> Result<Self, String> {
        let value = value.into();
        if !value.starts_with("secret://") || value.len() > 512 || value.len() <= 9 {
            return Err("credential_ref must be a secret:// reference of at most 512 bytes".into());
        }
        Ok(Self(value))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl fmt::Debug for SecretRef {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("SecretRef([REDACTED])")
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelApiFamily {
    OpenAi,
    Claude,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelProfileStatus {
    Draft,
    Active,
    Disabled,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProviderKind {
    OpenAiCompatible,
    Anthropic,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Modality {
    Text,
    Image,
    Audio,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelCapabilities {
    pub input_modalities: BTreeSet<Modality>,
    pub output_modalities: BTreeSet<Modality>,
    pub tools: bool,
    pub structured_output: bool,
    pub reasoning: bool,
    pub logprobs: bool,
    pub token_count: bool,
    pub realtime: bool,
}

impl Default for ModelCapabilities {
    fn default() -> Self {
        Self {
            input_modalities: BTreeSet::from([Modality::Text]),
            output_modalities: BTreeSet::from([Modality::Text]),
            tools: true,
            structured_output: false,
            reasoning: false,
            logprobs: false,
            token_count: true,
            realtime: false,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ProviderEndpoint {
    pub id: String,
    pub base_url: String,
    pub priority: u16,
    pub weight: u16,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub region: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelPolicy {
    pub connect_timeout_ms: u64,
    pub first_event_timeout_ms: u64,
    pub total_timeout_ms: u64,
    pub stream_idle_timeout_ms: u64,
    pub max_output_tokens: u32,
    pub max_request_bytes: usize,
    pub max_stream_event_bytes: usize,
    pub max_concurrency: usize,
}

impl Default for ModelPolicy {
    fn default() -> Self {
        Self {
            connect_timeout_ms: 3_000,
            first_event_timeout_ms: 30_000,
            total_timeout_ms: 300_000,
            stream_idle_timeout_ms: 60_000,
            max_output_tokens: 8_192,
            max_request_bytes: 2 * 1024 * 1024,
            max_stream_event_bytes: 256 * 1024,
            max_concurrency: 64,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelProfile {
    pub id: ModelProfileId,
    pub scope: TenantProjectScope,
    pub display_name: String,
    pub public_model_id: String,
    pub api_families: BTreeSet<ModelApiFamily>,
    pub capabilities: ModelCapabilities,
    pub provider: ProviderKind,
    pub provider_model: String,
    pub endpoints: Vec<ProviderEndpoint>,
    pub credential_ref: SecretRef,
    pub credential_configured: bool,
    pub policy: ModelPolicy,
    pub status: ModelProfileStatus,
    pub revision: u64,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CreateModelProfileRequest {
    pub id: ModelProfileId,
    pub display_name: String,
    pub public_model_id: String,
    pub api_families: BTreeSet<ModelApiFamily>,
    #[serde(default)]
    pub capabilities: ModelCapabilities,
    pub provider: ProviderKind,
    pub provider_model: String,
    pub endpoints: Vec<ProviderEndpoint>,
    pub credential_ref: SecretRef,
    #[serde(default)]
    pub policy: ModelPolicy,
    pub idempotency_key: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdateModelProfileRequest {
    pub expected_revision: u64,
    pub display_name: String,
    pub public_model_id: String,
    pub api_families: BTreeSet<ModelApiFamily>,
    pub capabilities: ModelCapabilities,
    pub provider: ProviderKind,
    pub provider_model: String,
    pub endpoints: Vec<ProviderEndpoint>,
    pub credential_ref: SecretRef,
    pub policy: ModelPolicy,
    pub idempotency_key: String,
}

/// Provider connection shared by platform and Agent-owned model configuration.
/// Credentials are deliberately excluded and are accepted only on write requests.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ModelConnectionConfig {
    pub provider: ProviderKind,
    pub base_url: String,
    pub api_models: Vec<String>,
    pub default_model: String,
    pub api_families: BTreeSet<ModelApiFamily>,
    #[serde(default)]
    pub capabilities: ModelCapabilities,
    #[serde(default)]
    pub policy: ModelPolicy,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AgentModelConfig {
    pub id: ModelConfigId,
    pub scope: TenantProjectScope,
    pub agent_id: agent_registry_contract::AgentId,
    pub display_name: String,
    pub connection: ModelConnectionConfig,
    pub api_key_configured: bool,
    pub enabled: bool,
    pub revision: u64,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct BuiltinModelConfig {
    pub id: ModelConfigId,
    pub display_name: String,
    pub connection: ModelConnectionConfig,
    pub api_key_configured: bool,
    pub enabled: bool,
    pub is_default: bool,
    pub revision: u64,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct WriteOnlyApiKey(String);

impl WriteOnlyApiKey {
    pub fn new(value: impl Into<String>) -> Result<Self, String> {
        let value = value.into();
        if value.trim().is_empty() || value.len() > 16 * 1024 {
            return Err("api_key must be 1..=16384 bytes".into());
        }
        Ok(Self(value))
    }

    pub fn expose(&self) -> &str {
        &self.0
    }
}

impl fmt::Debug for WriteOnlyApiKey {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("WriteOnlyApiKey([REDACTED])")
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpsertAgentModelConfigRequest {
    pub display_name: String,
    pub connection: ModelConnectionConfig,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub api_key: Option<WriteOnlyApiKey>,
    #[serde(default = "default_true")]
    pub enabled: bool,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expected_revision: Option<u64>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpsertBuiltinModelConfigRequest {
    pub display_name: String,
    pub connection: ModelConnectionConfig,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub api_key: Option<WriteOnlyApiKey>,
    #[serde(default = "default_true")]
    pub enabled: bool,
    #[serde(default)]
    pub is_default: bool,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expected_revision: Option<u64>,
}

fn default_true() -> bool {
    true
}

pub const PLATFORM_ADMIN_ROLE: &str = "platform_admin";

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct PlatformAdminUser {
    pub username: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub user_id: Option<String>,
    pub role: String,
    pub active: bool,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CreatePlatformAdminUserRequest {
    pub username: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdatePlatformAdminUserRequest {
    pub active: bool,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PlatformAdminUserPage {
    pub items: Vec<PlatformAdminUser>,
}

pub fn parse_admin_username(value: &str) -> Result<String, String> {
    let value = value.trim();
    validate_identifier("username", value, 255)?;
    Ok(value.to_owned())
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ProfileCommandRequest {
    pub expected_revision: u64,
    pub idempotency_key: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ListModelProfilesRequest {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub status: Option<ModelProfileStatus>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub after: Option<ModelProfileId>,
    #[serde(default = "default_list_limit")]
    pub limit: usize,
}

impl Default for ListModelProfilesRequest {
    fn default() -> Self {
        Self {
            status: None,
            after: None,
            limit: default_list_limit(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ListModelProfilesResponse {
    pub items: Vec<ModelProfile>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub next: Option<ModelProfileId>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelActor {
    pub scope: TenantProjectScope,
    pub runtime_identity_id: RuntimeIdentityId,
    pub capabilities: Vec<CredentialCapability>,
    pub credential_expires_at_seconds: Option<u64>,
    pub relations: agent_registry_contract::ResourceRelations,
}

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

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CanonicalContent {
    Text {
        text: String,
    },
    Image {
        source: Value,
    },
    ToolUse {
        id: String,
        name: String,
        input: Value,
    },
    ToolResult {
        tool_use_id: String,
        content: Value,
    },
    Audio {
        data: String,
        format: String,
    },
    Thinking {
        text: String,
    },
}

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

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolDefinition {
    pub name: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    pub input_schema: Value,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CanonicalModelRequest {
    pub invocation_id: InvocationId,
    pub public_model_id: String,
    pub source_api: ModelApiFamily,
    pub messages: Vec<CanonicalMessage>,
    #[serde(default)]
    pub tools: Vec<ToolDefinition>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub response_format: Option<Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_output_tokens: Option<u32>,
    /// OpenAI-compatible reasoning switch forwarded only by providers that support it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub thinking: Option<Value>,
    /// Provider reasoning budget hint such as `low`, `high`, or `max`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reasoning_effort: Option<String>,
    #[serde(default)]
    pub stop_sequences: Vec<String>,
    #[serde(default)]
    pub metadata: BTreeMap<String, String>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CanonicalModelResponse {
    pub invocation_id: InvocationId,
    pub provider_request_id: Option<String>,
    pub output: Vec<CanonicalContent>,
    pub finish_reason: String,
    pub usage: TokenUsage,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CanonicalModelEvent {
    Started {
        provider_request_id: Option<String>,
    },
    TextDelta {
        delta: String,
    },
    ToolUseStart {
        index: usize,
        id: String,
        name: String,
    },
    ToolArgumentsDelta {
        index: usize,
        delta: String,
    },
    Usage {
        usage: TokenUsage,
    },
    Completed {
        finish_reason: String,
    },
    Failed {
        code: String,
        message: String,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InvocationChannel {
    Http,
    Sse,
    Websocket,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InvocationState {
    Accepted,
    Connecting,
    Streaming,
    Completed,
    Failed,
    Cancelled,
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelCredentialSource {
    Agent,
    Builtin,
    #[default]
    LegacyProfile,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelInvocation {
    pub id: InvocationId,
    pub scope: TenantProjectScope,
    pub runtime_identity_id: RuntimeIdentityId,
    #[serde(default)]
    pub relations: agent_registry_contract::ResourceRelations,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub agent_id: Option<agent_registry_contract::AgentId>,
    #[serde(default)]
    pub requested_model_id: String,
    #[serde(default)]
    pub resolved_model_id: String,
    #[serde(default)]
    pub credential_source: ModelCredentialSource,
    pub profile_id: ModelProfileId,
    pub profile_revision: u64,
    pub api_family: ModelApiFamily,
    pub channel: InvocationChannel,
    pub state: InvocationState,
    pub request_digest: String,
    pub provider_request_id: Option<String>,
    pub usage: Option<TokenUsage>,
    pub finish_reason: Option<String>,
    pub error_code: Option<String>,
    pub created_at: DateTime<Utc>,
    pub completed_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelUsageSummary {
    pub invocation_count: u64,
    pub completed_count: u64,
    pub failed_count: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub builtin_invocation_count: u64,
    pub agent_invocation_count: u64,
    #[serde(default)]
    pub usage_coverage: TokenUsageCoverage,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelUsageLedgerEntry {
    pub invocation_id: InvocationId,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub agent_id: Option<agent_registry_contract::AgentId>,
    pub requested_model_id: String,
    pub resolved_model_id: String,
    pub credential_source: ModelCredentialSource,
    pub state: InvocationState,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub usage: Option<TokenUsage>,
    pub created_at: DateTime<Utc>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub completed_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelErrorKind {
    Invalid,
    Unauthenticated,
    Forbidden,
    NotFound,
    Conflict,
    RateLimited,
    Timeout,
    Overloaded,
    Upstream,
    Internal,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelError {
    pub kind: ModelErrorKind,
    pub code: String,
    pub message: String,
    pub param: Option<String>,
    pub retryable: bool,
    pub outcome_unknown: bool,
}

impl ModelError {
    pub fn new(kind: ModelErrorKind, code: &str, message: impl Into<String>) -> Self {
        Self {
            kind,
            code: code.into(),
            message: message.into(),
            param: None,
            retryable: false,
            outcome_unknown: false,
        }
    }

    pub fn invalid(code: &str, message: impl Into<String>, param: Option<&str>) -> Self {
        let mut error = Self::new(ModelErrorKind::Invalid, code, message);
        error.param = param.map(str::to_owned);
        error
    }
}

impl fmt::Display for ModelError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.message.fmt(formatter)
    }
}

impl std::error::Error for ModelError {}

pub type ModelResult<T> = Result<T, ModelError>;

pub fn capability(profile_id: &ModelProfileId, action: &str) -> String {
    format!("model-profile:{profile_id}:{action}")
}

pub fn validate_identifier(label: &str, value: &str, max: usize) -> Result<(), String> {
    if value.is_empty() || value.len() > max {
        return Err(format!("{label} must contain 1..={max} bytes"));
    }
    if !value.bytes().all(|byte| {
        byte.is_ascii_alphanumeric() || matches!(byte, b':' | b'.' | b'-' | b'_' | b'@')
    }) {
        return Err(format!("{label} contains unsupported characters"));
    }
    Ok(())
}

fn default_list_limit() -> usize {
    50
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn secret_debug_never_reveals_reference() {
        let secret = SecretRef::new("secret://model-infra/prod").unwrap();
        assert_eq!(format!("{secret:?}"), "SecretRef([REDACTED])");
    }

    #[test]
    fn profile_body_cannot_override_scope() {
        let result = serde_json::from_value::<CreateModelProfileRequest>(serde_json::json!({
            "id":"echo", "tenantId":"attacker"
        }));
        assert!(result.is_err());
    }

    #[test]
    fn admin_username_accepts_sky_gateway_usernames() {
        assert_eq!(parse_admin_username(" jinc.cen ").unwrap(), "jinc.cen");
        assert!(parse_admin_username("").is_err());
        assert!(parse_admin_username("bad name").is_err());
    }
}