agent_proxy_rust_storage/

lib.rs

//! Backend-agnostic storage abstraction for agent-proxy-rust.
//!
//! Defines the [`Storage`] trait and all data types shared across backends.
//! Middleware crates depend on `Box<dyn Storage>` injected at construction time,
//! and never know whether the backend is `SQLite`, `PostgreSQL`, or an in-memory mock.

#![forbid(unsafe_code)]
#![warn(missing_docs, missing_debug_implementations)]

mod error;
mod types;

use std::fmt::Debug;

use async_trait::async_trait;
pub use error::StorageError;
use secrecy::SecretString;
pub use types::{
    AvailableChannelInfo, AvailableModelInfo, Channel, CompressionSavingsReport, CostAggregate,
    CostFilter, CostGroupBy, CostRecord, Model, ModelMapping, ProtocolEntry, Provider,
    SeedEntryStatus, SeedManifest, SeedManifestEntry, SeedStatus, SubscriptionFee, SwitchLog,
    TimeRange,
};

/// Backend-agnostic storage for providers, models, channels, and cost records.
///
/// Every method except [`max_connections`](Self::max_connections) is async and
/// returns `Result<T, StorageError>`. Implementations must be `Send + Sync`
/// so the trait object can be shared across Tokio tasks behind an `Arc`.
#[async_trait]
pub trait Storage: Send + Sync + Debug {
    // ── Provider ────────────────────────────────────────────

    /// List all providers.
    async fn list_providers(&self) -> Result<Vec<Provider>, StorageError>;

    /// Get a single provider by ID.
    async fn get_provider(&self, id: &str) -> Result<Option<Provider>, StorageError>;

    // ── Model ───────────────────────────────────────────────

    /// List models, optionally filtered by provider.
    async fn list_models(&self, provider_id: Option<&str>) -> Result<Vec<Model>, StorageError>;

    /// Get a single model by ID.
    async fn get_model(&self, id: &str) -> Result<Option<Model>, StorageError>;

    // ── Channel ─────────────────────────────────────────────

    /// List all channels, optionally filtered by model ID.
    async fn list_channels(&self, model_id: Option<&str>) -> Result<Vec<Channel>, StorageError>;

    /// Get a single channel by ID.
    async fn get_channel(&self, id: &str) -> Result<Option<Channel>, StorageError>;

    /// Insert or replace a channel (upsert).
    async fn upsert_channel(&self, channel: &Channel) -> Result<(), StorageError>;

    /// Toggle a channel enabled/disabled.
    async fn set_channel_enabled(&self, id: &str, enabled: bool) -> Result<(), StorageError>;

    /// Update just the API key for a channel.
    async fn set_channel_api_key(&self, id: &str, key: &SecretString) -> Result<(), StorageError>;

    /// Update channel fields (name, enabled, priority, quota, protocols, `force_protocol`).
    #[allow(clippy::too_many_arguments)]
    async fn update_channel(
        &self,
        id: &str,
        name: Option<&str>,
        enabled: Option<bool>,
        priority: Option<u32>,
        monthly_quota: Option<u64>,
        quota_policy: Option<&str>,
        protocols: Option<&str>,
        force_protocol: Option<&str>,
    ) -> Result<Channel, StorageError>;

    /// Delete a channel and its model mappings (cascade).
    async fn delete_channel(&self, id: &str) -> Result<(), StorageError>;

    /// Mark a channel as healthy (reset failures).
    async fn mark_channel_healthy(&self, id: &str) -> Result<(), StorageError>;

    /// Record a channel failure (increments counter, may set Degraded/Cooldown).
    async fn record_channel_failure(&self, id: &str) -> Result<(), StorageError>;

    // ── Model Mapping ───────────────────────────────────────

    /// List all model mappings for a channel.
    async fn list_mappings(&self, channel_id: &str) -> Result<Vec<ModelMapping>, StorageError>;

    /// Upsert a single model mapping.
    async fn upsert_mapping(&self, mapping: &ModelMapping) -> Result<(), StorageError>;

    /// Toggle a model mapping enabled/disabled.
    async fn set_mapping_enabled(&self, id: &str, enabled: bool) -> Result<(), StorageError>;

    /// Delete a model mapping.
    async fn delete_mapping(&self, id: &str) -> Result<(), StorageError>;

    /// List all model mappings.
    async fn list_all_mappings(&self) -> Result<Vec<ModelMapping>, StorageError>;

    // ── Cost Records ────────────────────────────────────────

    /// Record a completed request.
    async fn insert_cost_record(&self, record: &CostRecord) -> Result<(), StorageError>;

    /// Query cost records with optional filters.
    async fn query_cost_records(&self, filter: CostFilter)
    -> Result<Vec<CostRecord>, StorageError>;

    /// Aggregate costs grouped by the given dimension within a time range.
    async fn aggregate_costs(
        &self,
        group_by: CostGroupBy,
        range: TimeRange,
    ) -> Result<Vec<CostAggregate>, StorageError>;

    /// Delete records older than N days, returning the count of deleted rows.
    async fn prune_cost_records(&self, older_than_days: u32) -> Result<u64, StorageError>;

    /// List distinct project paths from cost records, sorted alphabetically.
    async fn list_projects(&self) -> Result<Vec<String>, StorageError>;

    // ── Switch Log ──────────────────────────────────────────

    /// Record a channel switch event.
    async fn insert_switch_log(&self, log: &SwitchLog) -> Result<(), StorageError>;

    /// Query recent switch logs, optionally limited.
    async fn query_switch_logs(&self, limit: Option<u32>) -> Result<Vec<SwitchLog>, StorageError>;

    // ── Available Channels ───────────────────────────────────

    /// List all enabled channels with their bound models.
    /// Used by token-fleet-switch for Claude direct-connect mode.
    async fn list_available_channels(&self) -> Result<Vec<AvailableChannelInfo>, StorageError>;

    // ── Subscription Fees ───────────────────────────────────

    /// Record a monthly subscription fee.
    async fn insert_subscription_fee(&self, fee: &SubscriptionFee) -> Result<(), StorageError>;

    /// Query subscription fees optionally filtered by channel and/or month.
    async fn query_subscription_fees(
        &self,
        channel: Option<&str>,
        month: Option<&str>,
    ) -> Result<Vec<SubscriptionFee>, StorageError>;

    // ── Lifecycle ───────────────────────────────────────────

    /// Run migrations / schema init. Idempotent — safe to call on every startup.
    async fn migrate(&self) -> Result<(), StorageError>;

    /// Health check — returns `true` if the backend is reachable.
    async fn health_check(&self) -> Result<bool, StorageError>;

    /// Maximum number of concurrent connections this backend supports.
    fn max_connections(&self) -> usize;
}

/// Seed data manager — initializes and refreshes reference data.
///
/// Seed data includes providers, models, channels, and model mappings
/// with pricing. It is versioned separately from schema migrations so
/// it can be hot-updated without redeploying.
#[async_trait]
pub trait SeedManager: Send + Sync {
    /// Initialize seed data from the embedded fallback (compile-time JSON).
    ///
    /// Idempotent — safe to call on every startup. Only inserts when the
    /// local database has no seed data yet (version = 0).
    async fn seed_init(&self) -> Result<SeedStatus, StorageError>;

    /// Fetch and apply seed data from a remote URL.
    ///
    /// `url` is the base URL to the seed directory (e.g.
    /// `"https://raw.githubusercontent.com/.../refs/tags/seed-v3/crates/storage-sqlite/seed/"`).
    /// When `None`, tries the environment variable `AGENT_PROXY_SEED_URL`,
    /// then falls back to the last-used URL stored in `seed_metadata`.
    async fn seed_refresh(&self, url: Option<&str>) -> Result<SeedStatus, StorageError>;

    /// Query the current seed data status without fetching.
    async fn seed_status(&self) -> Result<SeedStatus, StorageError>;

    /// Fetch the remote manifest and compare with local status, but do NOT
    /// apply any changes. Used by the admin API `?remote=true` query parameter.
    async fn seed_check_remote(&self, url: Option<&str>) -> Result<SeedStatus, StorageError>;
}