rs-adk 0.5.0

//! Vertex AI session service — managed session storage via Vertex AI REST API.
//!
//! Provides session persistence using the Vertex AI session management
//! endpoint. Sessions are stored and managed by Google Cloud, with
//! optional TTL-based expiration.

use async_trait::async_trait;

use super::{Session, SessionError, SessionId, SessionService};
use crate::events::Event;

/// Configuration for the Vertex AI session service.
#[derive(Debug, Clone)]
pub struct VertexAiSessionConfig {
    /// Google Cloud project ID.
    pub project: String,
    /// Google Cloud region (e.g., `us-central1`).
    pub location: String,
    /// Optional time-to-live for sessions, in seconds.
    /// If set, sessions expire after this duration of inactivity.
    pub ttl_seconds: Option<u64>,
}

impl VertexAiSessionConfig {
    /// Create a new Vertex AI session config.
    pub fn new(project: impl Into<String>, location: impl Into<String>) -> Self {
        Self {
            project: project.into(),
            location: location.into(),
            ttl_seconds: None,
        }
    }

    /// Set the session TTL in seconds.
    pub fn ttl_seconds(mut self, ttl: u64) -> Self {
        self.ttl_seconds = Some(ttl);
        self
    }

    /// Construct the base URL for the Vertex AI session endpoint.
    ///
    /// Format: `https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}/reasoningEngines`
    fn base_url(&self) -> String {
        format!(
            "https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}",
            project = self.project,
            location = self.location,
        )
    }

    /// Construct the sessions endpoint URL for a specific reasoning engine.
    fn sessions_url(&self, engine_id: &str) -> String {
        format!(
            "{}/reasoningEngines/{}/sessions",
            self.base_url(),
            engine_id,
        )
    }

    /// Construct the URL for a specific session.
    fn session_url(&self, engine_id: &str, session_id: &str) -> String {
        format!("{}/{}", self.sessions_url(engine_id), session_id)
    }

    /// Construct the events endpoint URL for a specific session.
    fn events_url(&self, engine_id: &str, session_id: &str) -> String {
        format!("{}/events", self.session_url(engine_id, session_id))
    }
}

/// Session service backed by the Vertex AI managed session endpoint.
///
/// Uses the Vertex AI REST API for session CRUD and event storage.
/// Requires a valid Google Cloud project with the AI Platform API enabled.
///
/// Sessions are stored server-side by Google Cloud, providing managed
/// persistence without requiring a separate database.
pub struct VertexAiSessionService {
    config: VertexAiSessionConfig,
    // In a real implementation, this would hold:
    // - An HTTP client (e.g., `reqwest::Client`)
    // - An auth token provider (e.g., `rs_genai::VertexAIAuth`)
    // - An optional reasoning engine ID
}

impl VertexAiSessionService {
    /// Create a new Vertex AI session service.
    pub fn new(config: VertexAiSessionConfig) -> Self {
        Self { config }
    }

    /// Returns the configured project ID.
    pub fn project(&self) -> &str {
        &self.config.project
    }

    /// Returns the configured location.
    pub fn location(&self) -> &str {
        &self.config.location
    }

    /// Returns the configured TTL in seconds, if any.
    pub fn ttl_seconds(&self) -> Option<u64> {
        self.config.ttl_seconds
    }
}

#[async_trait]
impl SessionService for VertexAiSessionService {
    async fn create_session(&self, app_name: &str, user_id: &str) -> Result<Session, SessionError> {
        let _url = self.config.sessions_url(app_name);
        let _user = user_id;

        // Real implementation would:
        // POST {sessions_url}
        // Authorization: Bearer {token}
        // Content-Type: application/json
        //
        // {
        //   "userId": "{user_id}",
        //   "ttl": "{ttl_seconds}s"   // if configured
        // }
        //
        // Response: { "name": "...sessions/{id}", "userId": "...", ... }
        let _ttl_body = self
            .config
            .ttl_seconds
            .map(|t| format!("\"ttl\": \"{t}s\""));

        todo!("POST to {_url} to create Vertex AI session for user={_user}")
    }

    async fn get_session(&self, id: &SessionId) -> Result<Option<Session>, SessionError> {
        let _url = self.config.session_url("default", id.as_str());

        // Real implementation would:
        // GET {session_url}
        // Authorization: Bearer {token}
        //
        // Response: { "name": "...sessions/{id}", "userId": "...", "state": {...}, ... }
        // Returns None if 404.
        todo!("GET {_url} to fetch Vertex AI session")
    }

    async fn list_sessions(
        &self,
        app_name: &str,
        user_id: &str,
    ) -> Result<Vec<Session>, SessionError> {
        let _url = self.config.sessions_url(app_name);
        let _user = user_id;

        // Real implementation would:
        // GET {sessions_url}?filter=userId={user_id}
        // Authorization: Bearer {token}
        //
        // Response: { "sessions": [{ "name": "...", ... }, ...] }
        todo!("GET {_url} to list Vertex AI sessions for user={_user}")
    }

    async fn delete_session(&self, id: &SessionId) -> Result<(), SessionError> {
        let _url = self.config.session_url("default", id.as_str());

        // Real implementation would:
        // DELETE {session_url}
        // Authorization: Bearer {token}
        //
        // Response: {} (empty on success)
        todo!("DELETE {_url} to remove Vertex AI session")
    }

    async fn append_event(&self, id: &SessionId, event: Event) -> Result<(), SessionError> {
        let _url = self.config.events_url("default", id.as_str());
        let _event_json =
            serde_json::to_value(&event).map_err(|e| SessionError::Storage(e.to_string()))?;

        // Real implementation would:
        // POST {events_url}
        // Authorization: Bearer {token}
        // Content-Type: application/json
        //
        // {
        //   "author": "{event.author}",
        //   "invocationId": "{event.invocation_id}",
        //   "content": { "parts": [{ "text": "{event.content}" }] },
        //   "actions": {event.actions}
        // }
        todo!("POST to {_url} to append event to Vertex AI session")
    }

    async fn get_events(&self, id: &SessionId) -> Result<Vec<Event>, SessionError> {
        let _url = self.config.events_url("default", id.as_str());

        // Real implementation would:
        // GET {events_url}
        // Authorization: Bearer {token}
        //
        // Response: { "sessionEvents": [{ "author": "...", ... }, ...] }
        // Transform each Vertex AI event into our Event type.
        todo!("GET {_url} to fetch events for Vertex AI session")
    }
}

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

    #[test]
    fn config_new() {
        let config = VertexAiSessionConfig::new("my-project", "us-central1");
        assert_eq!(config.project, "my-project");
        assert_eq!(config.location, "us-central1");
        assert!(config.ttl_seconds.is_none());
    }

    #[test]
    fn config_with_ttl() {
        let config = VertexAiSessionConfig::new("proj", "us-east1").ttl_seconds(3600);
        assert_eq!(config.ttl_seconds, Some(3600));
    }

    #[test]
    fn url_construction() {
        let config = VertexAiSessionConfig::new("my-project", "us-central1");
        assert_eq!(
            config.base_url(),
            "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/my-project/locations/us-central1"
        );
        assert!(config
            .sessions_url("engine-1")
            .contains("reasoningEngines/engine-1/sessions"));
        assert!(config
            .session_url("engine-1", "sess-1")
            .contains("sessions/sess-1"));
        assert!(config
            .events_url("engine-1", "sess-1")
            .contains("sessions/sess-1/events"));
    }

    #[test]
    fn service_accessors() {
        let svc = VertexAiSessionService::new(
            VertexAiSessionConfig::new("proj", "us-west1").ttl_seconds(7200),
        );
        assert_eq!(svc.project(), "proj");
        assert_eq!(svc.location(), "us-west1");
        assert_eq!(svc.ttl_seconds(), Some(7200));
    }
}