gemini_adk_rs/session/
vertex_ai.rs

1//! Vertex AI session service — managed session storage via Vertex AI REST API.
2//!
3//! Provides session persistence using the Vertex AI session management
4//! endpoint. Sessions are stored and managed by Google Cloud, with
5//! optional TTL-based expiration.
6
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde::Deserialize;
11use serde_json::Value;
12
13use super::{Session, SessionError, SessionId, SessionService};
14use crate::events::{Event, EventActions};
15
16// ──────────────────────────────────────────────────────────────────────────────
17// Configuration
18// ──────────────────────────────────────────────────────────────────────────────
19
20/// Configuration for the Vertex AI session service.
21#[derive(Debug, Clone)]
22pub struct VertexAiSessionConfig {
23    /// Google Cloud project ID.
24    pub project: String,
25    /// Google Cloud region (e.g., `us-central1`).
26    pub location: String,
27    /// Optional time-to-live for sessions, in seconds.
28    /// If set, sessions expire after this duration of inactivity.
29    pub ttl_seconds: Option<u64>,
30}
31
32impl VertexAiSessionConfig {
33    /// Create a new Vertex AI session config.
34    pub fn new(project: impl Into<String>, location: impl Into<String>) -> Self {
35        Self {
36            project: project.into(),
37            location: location.into(),
38            ttl_seconds: None,
39        }
40    }
41
42    /// Set the session TTL in seconds.
43    pub fn ttl_seconds(mut self, ttl: u64) -> Self {
44        self.ttl_seconds = Some(ttl);
45        self
46    }
47
48    /// Construct the base URL for the Vertex AI session endpoint.
49    ///
50    /// Format: `https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}/reasoningEngines`
51    fn base_url(&self) -> String {
52        format!(
53            "https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}",
54            project = self.project,
55            location = self.location,
56        )
57    }
58
59    /// Construct the sessions endpoint URL for a specific reasoning engine.
60    fn sessions_url(&self, engine_id: &str) -> String {
61        format!(
62            "{}/reasoningEngines/{}/sessions",
63            self.base_url(),
64            percent_encode(engine_id),
65        )
66    }
67
68    /// Construct the URL for a specific session.
69    ///
70    /// `session_id` reaches this from the caller, so it is encoded rather than
71    /// interpolated: an id containing `/` addresses a different resource, and
72    /// one containing `?` or `#` turns the rest of the path into a query or
73    /// fragment. Encoding keeps it one path segment whatever it holds.
74    fn session_url(&self, engine_id: &str, session_id: &str) -> String {
75        format!(
76            "{}/{}",
77            self.sessions_url(engine_id),
78            percent_encode(session_id)
79        )
80    }
81
82    /// Construct the events endpoint URL for a specific session.
83    fn events_url(&self, engine_id: &str, session_id: &str) -> String {
84        format!("{}/events", self.session_url(engine_id, session_id))
85    }
86}
87
88/// Percent-encode one URL path segment or query value.
89///
90/// Escapes everything outside RFC 3986's unreserved set, so the result can only
91/// ever be the single component it was meant to be — it cannot introduce a path
92/// separator, open a query string, or append another query parameter.
93fn percent_encode(value: &str) -> String {
94    let mut out = String::with_capacity(value.len());
95    for byte in value.bytes() {
96        match byte {
97            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => {
98                out.push(byte as char);
99            }
100            _ => out.push_str(&format!("%{byte:02X}")),
101        }
102    }
103    out
104}
105
106// ──────────────────────────────────────────────────────────────────────────────
107// Auth provider
108// ──────────────────────────────────────────────────────────────────────────────
109
110/// How to supply a bearer token for Vertex AI requests.
111enum TokenProvider {
112    /// No token configured — requests will fail with a clear message.
113    None,
114    /// A static, pre-fetched bearer token string.
115    Static(String),
116    /// A dynamic refresher: called before every request.
117    Refresher(Arc<dyn Fn() -> String + Send + Sync>),
118}
119
120impl TokenProvider {
121    /// Retrieve the current token, or return an error if none is configured.
122    fn get(&self) -> Result<String, SessionError> {
123        match self {
124            TokenProvider::None => Err(SessionError::Storage(
125                "missing auth token: call .with_token() or .with_token_refresher()".into(),
126            )),
127            TokenProvider::Static(t) => Ok(t.clone()),
128            TokenProvider::Refresher(f) => Ok(f()),
129        }
130    }
131}
132
133// ──────────────────────────────────────────────────────────────────────────────
134// DTO types — mirror Vertex AI JSON shapes
135// ──────────────────────────────────────────────────────────────────────────────
136
137/// Vertex AI session resource as returned by the REST API.
138#[derive(Debug, Deserialize)]
139#[serde(rename_all = "camelCase")]
140struct VertexSession {
141    /// Full resource name, e.g.
142    /// `projects/p/locations/l/reasoningEngines/e/sessions/{id}`
143    name: String,
144    /// The user ID associated with this session.
145    #[serde(default)]
146    user_id: String,
147    /// Arbitrary session state stored server-side.
148    #[serde(default)]
149    session_state: Option<Value>,
150    /// RFC 3339 creation timestamp.
151    #[serde(default)]
152    create_time: Option<String>,
153    /// RFC 3339 last-update timestamp.
154    #[serde(default)]
155    update_time: Option<String>,
156}
157
158/// Response envelope for `listSessions`.
159#[derive(Debug, Deserialize)]
160#[serde(rename_all = "camelCase")]
161struct ListSessionsResponse {
162    #[serde(default)]
163    sessions: Vec<VertexSession>,
164}
165
166/// Vertex AI event resource as returned by the REST API.
167#[derive(Debug, Deserialize)]
168#[serde(rename_all = "camelCase")]
169struct VertexEvent {
170    /// Who authored the event.
171    #[serde(default)]
172    author: String,
173    /// Invocation identifier.
174    #[serde(default)]
175    invocation_id: String,
176    /// Freeform content (may hold a `parts` array or a plain string).
177    #[serde(default)]
178    content: Option<Value>,
179    /// Actions metadata stored alongside the event.
180    #[serde(default)]
181    actions: Option<Value>,
182    /// Event identifier (last path segment of `name`).
183    #[serde(default)]
184    name: Option<String>,
185    /// Unix timestamp as a string (seconds since epoch).
186    #[serde(default)]
187    timestamp: Option<String>,
188}
189
190/// Response envelope for `listEvents`.
191#[derive(Debug, Deserialize)]
192#[serde(rename_all = "camelCase")]
193struct ListEventsResponse {
194    #[serde(default)]
195    session_events: Vec<VertexEvent>,
196}
197
198// ──────────────────────────────────────────────────────────────────────────────
199// Pure helper functions (URL/body builders + response mappers)
200// ──────────────────────────────────────────────────────────────────────────────
201
202/// Build the JSON body for a `createSession` request.
203pub(crate) fn build_create_body(user_id: &str, ttl_seconds: Option<u64>) -> Value {
204    let mut body = serde_json::json!({ "userId": user_id });
205    if let Some(ttl) = ttl_seconds {
206        body["ttl"] = Value::String(format!("{ttl}s"));
207    }
208    body
209}
210
211/// Build the JSON body for an `appendEvent` request.
212pub(crate) fn build_event_body(event: &Event) -> Value {
213    let content_val = match &event.content {
214        Some(text) => serde_json::json!({ "parts": [{ "text": text }] }),
215        None => serde_json::json!({}),
216    };
217    serde_json::json!({
218        "author":       event.author,
219        "invocationId": event.invocation_id,
220        "content":      content_val,
221        "actions":      serde_json::to_value(&event.actions)
222                            .unwrap_or(serde_json::json!({})),
223    })
224}
225
226/// Extract the short session ID from a full Vertex AI resource `name`.
227///
228/// The name looks like
229/// `projects/p/locations/l/reasoningEngines/e/sessions/SESSION_ID`.
230/// We return the last path segment.
231fn session_id_from_name(name: &str) -> &str {
232    name.rsplit('/').next().unwrap_or(name)
233}
234
235/// Map a `VertexSession` DTO into our domain `Session`.
236///
237/// `app_name` is not stored by Vertex — the caller supplies it from context.
238fn map_vertex_session(vs: VertexSession, app_name: &str) -> Session {
239    let id_str = session_id_from_name(&vs.name);
240    let state = vs
241        .session_state
242        .and_then(|v| v.as_object().cloned())
243        .map(|m| m.into_iter().collect())
244        .unwrap_or_default();
245
246    let now = vs.create_time.clone().unwrap_or_else(|| "0Z".to_string());
247
248    Session {
249        id: SessionId::from_string(id_str),
250        app_name: app_name.to_string(),
251        user_id: vs.user_id,
252        state,
253        created_at: vs.create_time.unwrap_or_else(|| now.clone()),
254        updated_at: vs.update_time.unwrap_or(now),
255        events: Vec::new(),
256    }
257}
258
259/// Map a `VertexEvent` DTO into our domain `Event`.
260fn map_vertex_event(ve: VertexEvent) -> Event {
261    // Extract plain text from the Vertex content shape:
262    // { "parts": [{ "text": "..." }] }
263    let content = ve.content.as_ref().and_then(|c| {
264        c.get("parts")
265            .and_then(|p| p.as_array())
266            .and_then(|arr| arr.first())
267            .and_then(|part| part.get("text"))
268            .and_then(|t| t.as_str())
269            .map(String::from)
270    });
271
272    let actions: EventActions = ve
273        .actions
274        .and_then(|v| serde_json::from_value(v).ok())
275        .unwrap_or_default();
276
277    let timestamp: u64 = ve
278        .timestamp
279        .as_deref()
280        .and_then(|s| s.parse().ok())
281        .unwrap_or(0);
282
283    // Derive a stable event ID from the resource name if available.
284    let id = ve
285        .name
286        .as_deref()
287        .map(|n| session_id_from_name(n).to_string())
288        .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
289
290    Event {
291        id,
292        invocation_id: ve.invocation_id,
293        author: ve.author,
294        content,
295        actions,
296        timestamp,
297    }
298}
299
300/// Pure helper: given a status code and an already-consumed body string,
301/// decide whether this is a 404 (mapped to `Ok(None)`), a non-2xx error
302/// (mapped to `Err`), or a success (caller must parse the body).
303///
304/// Returns:
305/// - `Ok(true)`  → status was 2xx; caller may parse the body.
306/// - `Ok(false)` → status was 404; caller should return `Ok(None)`.
307/// - `Err(_)`    → status was non-2xx / non-404 error.
308#[cfg_attr(not(test), allow(dead_code))]
309pub(crate) fn classify_status(status: u16, body: &str) -> Result<bool, SessionError> {
310    if status == 404 {
311        return Ok(false);
312    }
313    if (200..300).contains(&status) {
314        return Ok(true);
315    }
316    Err(SessionError::Storage(format!(
317        "Vertex AI request failed [{status}]: {body}"
318    )))
319}
320
321/// Interpret a `reqwest::Response` that should carry a JSON body.
322///
323/// * 2xx  → deserialise as `T`
324/// * 404  → `Ok(None)`
325/// * else → `Err(SessionError::Storage(...))`
326///
327/// Returns `Ok(Some(T))` on success, `Ok(None)` on 404.
328async fn parse_json_response<T: for<'de> Deserialize<'de>>(
329    resp: reqwest::Response,
330) -> Result<Option<T>, SessionError> {
331    let status = resp.status().as_u16();
332    if status == 404 {
333        return Ok(None);
334    }
335    if (200..300).contains(&status) {
336        let parsed: T = resp
337            .json()
338            .await
339            .map_err(|e| SessionError::Storage(format!("failed to parse response: {e}")))?;
340        return Ok(Some(parsed));
341    }
342    // Collect body for a useful error message.
343    let body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
344    Err(SessionError::Storage(format!(
345        "Vertex AI request failed [{status}]: {body}"
346    )))
347}
348
349/// Like [`parse_json_response`] but maps the 404 into
350/// `Err(SessionError::NotFound(id))`.
351async fn parse_json_response_required<T: for<'de> Deserialize<'de>>(
352    resp: reqwest::Response,
353    id: &SessionId,
354) -> Result<T, SessionError> {
355    match parse_json_response::<T>(resp).await? {
356        Some(v) => Ok(v),
357        None => Err(SessionError::NotFound(id.clone())),
358    }
359}
360
361/// Check that a response was successful (for void operations like DELETE).
362async fn check_success(resp: reqwest::Response) -> Result<(), SessionError> {
363    let status = resp.status().as_u16();
364    if (200..300).contains(&status) {
365        return Ok(());
366    }
367    let body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
368    Err(SessionError::Storage(format!(
369        "Vertex AI request failed [{status}]: {body}"
370    )))
371}
372
373// ──────────────────────────────────────────────────────────────────────────────
374// Service struct
375// ──────────────────────────────────────────────────────────────────────────────
376
377/// Session service backed by the Vertex AI managed session endpoint.
378///
379/// Uses the Vertex AI REST API for session CRUD and event storage.
380/// Requires a valid Google Cloud project with the AI Platform API enabled.
381///
382/// Sessions are stored server-side by Google Cloud, providing managed
383/// persistence without requiring a separate database.
384///
385/// # Quick start
386///
387/// ```rust,no_run
388/// # use gemini_adk_rs::session::{VertexAiSessionConfig, VertexAiSessionService};
389/// let svc = VertexAiSessionService::new(
390///     VertexAiSessionConfig::new("my-project", "us-central1")
391///         .ttl_seconds(3600),
392/// )
393/// .with_token("ya29.my-access-token")
394/// .reasoning_engine("my-engine-id");
395/// ```
396pub struct VertexAiSessionService {
397    config: VertexAiSessionConfig,
398    client: reqwest::Client,
399    token_provider: TokenProvider,
400    engine_id: String,
401}
402
403impl VertexAiSessionService {
404    /// Create a new Vertex AI session service.
405    ///
406    /// No auth token is configured yet — call [`with_token`](Self::with_token)
407    /// or [`with_token_refresher`](Self::with_token_refresher) before issuing
408    /// requests, otherwise they will return
409    /// `SessionError::Storage("missing auth token")`.
410    pub fn new(config: VertexAiSessionConfig) -> Self {
411        Self {
412            config,
413            client: reqwest::Client::new(),
414            token_provider: TokenProvider::None,
415            engine_id: "default".to_string(),
416        }
417    }
418
419    /// Set a static bearer token for all requests.
420    pub fn with_token(mut self, token: impl Into<String>) -> Self {
421        self.token_provider = TokenProvider::Static(token.into());
422        self
423    }
424
425    /// Set a dynamic token refresher closure.
426    ///
427    /// The closure is invoked before every HTTP request, allowing the caller
428    /// to supply a freshly-refreshed token each time.
429    pub fn with_token_refresher(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self {
430        self.token_provider = TokenProvider::Refresher(Arc::new(f));
431        self
432    }
433
434    /// Override the reasoning engine ID (defaults to `"default"`).
435    pub fn reasoning_engine(mut self, id: impl Into<String>) -> Self {
436        self.engine_id = id.into();
437        self
438    }
439
440    // ── Accessors (keep existing tests passing) ────────────────────────────
441
442    /// Returns the configured project ID.
443    pub fn project(&self) -> &str {
444        &self.config.project
445    }
446
447    /// Returns the configured location.
448    pub fn location(&self) -> &str {
449        &self.config.location
450    }
451
452    /// Returns the configured TTL in seconds, if any.
453    pub fn ttl_seconds(&self) -> Option<u64> {
454        self.config.ttl_seconds
455    }
456
457    // ── Internal helpers ───────────────────────────────────────────────────
458
459    /// Build an authorised `RequestBuilder` for a GET request.
460    fn get(&self, url: &str) -> Result<reqwest::RequestBuilder, SessionError> {
461        let token = self.token_provider.get()?;
462        Ok(self
463            .client
464            .get(url)
465            .header("Authorization", format!("Bearer {token}")))
466    }
467
468    /// Build an authorised `RequestBuilder` for a POST request with a JSON body.
469    fn post(&self, url: &str, body: Value) -> Result<reqwest::RequestBuilder, SessionError> {
470        let token = self.token_provider.get()?;
471        Ok(self
472            .client
473            .post(url)
474            .header("Authorization", format!("Bearer {token}"))
475            .header("Content-Type", "application/json")
476            .json(&body))
477    }
478
479    /// Build an authorised `RequestBuilder` for a DELETE request.
480    fn delete(&self, url: &str) -> Result<reqwest::RequestBuilder, SessionError> {
481        let token = self.token_provider.get()?;
482        Ok(self
483            .client
484            .delete(url)
485            .header("Authorization", format!("Bearer {token}")))
486    }
487}
488
489// ──────────────────────────────────────────────────────────────────────────────
490// SessionService implementation
491// ──────────────────────────────────────────────────────────────────────────────
492
493#[async_trait]
494impl SessionService for VertexAiSessionService {
495    async fn create_session(&self, app_name: &str, user_id: &str) -> Result<Session, SessionError> {
496        // Physically create the session under the configured reasoning engine
497        // (the same engine the other CRUD methods use), so it stays reachable.
498        // `app_name` is only a logical label carried on the returned session.
499        let url = self.config.sessions_url(&self.engine_id);
500        let body = build_create_body(user_id, self.config.ttl_seconds);
501
502        let resp = self
503            .post(&url, body)?
504            .send()
505            .await
506            .map_err(|e| SessionError::Storage(format!("HTTP request failed: {e}")))?;
507
508        // A successful create returns the new VertexSession.
509        // We use a placeholder ID only to satisfy the required helper signature;
510        // the real ID comes from the response body.
511        let placeholder_id = SessionId::new();
512        let vs: VertexSession = parse_json_response_required(resp, &placeholder_id).await?;
513        Ok(map_vertex_session(vs, app_name))
514    }
515
516    async fn get_session(&self, id: &SessionId) -> Result<Option<Session>, SessionError> {
517        let url = self.config.session_url(&self.engine_id, id.as_str());
518
519        let resp = self
520            .get(&url)?
521            .send()
522            .await
523            .map_err(|e| SessionError::Storage(format!("HTTP request failed: {e}")))?;
524
525        // Treat 404 as Ok(None) per the spec.
526        let opt: Option<VertexSession> = parse_json_response(resp).await?;
527        Ok(opt.map(|vs| {
528            // We don't know the app_name from the response alone — extract it
529            // from the resource name if possible, else fall back to engine_id.
530            let app_name =
531                extract_engine_id_from_name(&vs.name).unwrap_or_else(|| self.engine_id.clone());
532            map_vertex_session(vs, &app_name)
533        }))
534    }
535
536    async fn list_sessions(
537        &self,
538        app_name: &str,
539        user_id: &str,
540    ) -> Result<Vec<Session>, SessionError> {
541        // List under the configured engine (where sessions are stored), but
542        // keep `app_name` as the logical label on the returned sessions.
543        let base_url = self.config.sessions_url(&self.engine_id);
544        // `user_id` is encoded, not interpolated: a raw `&` would append a
545        // second query parameter to the request and a raw `#` would truncate
546        // the filter, in both cases listing sessions this call did not ask for.
547        let url = format!("{base_url}?filter=userId={}", percent_encode(user_id));
548
549        let resp = self
550            .get(&url)?
551            .send()
552            .await
553            .map_err(|e| SessionError::Storage(format!("HTTP request failed: {e}")))?;
554
555        // A 404 here means the engine has no sessions — treat as empty list.
556        let opt: Option<ListSessionsResponse> = parse_json_response(resp).await?;
557        let sessions = opt
558            .map(|r| {
559                r.sessions
560                    .into_iter()
561                    .map(|vs| map_vertex_session(vs, app_name))
562                    .collect()
563            })
564            .unwrap_or_default();
565        Ok(sessions)
566    }
567
568    async fn delete_session(&self, id: &SessionId) -> Result<(), SessionError> {
569        let url = self.config.session_url(&self.engine_id, id.as_str());
570
571        let resp = self
572            .delete(&url)?
573            .send()
574            .await
575            .map_err(|e| SessionError::Storage(format!("HTTP request failed: {e}")))?;
576
577        check_success(resp).await
578    }
579
580    async fn append_event(&self, id: &SessionId, event: Event) -> Result<(), SessionError> {
581        let url = self.config.events_url(&self.engine_id, id.as_str());
582        let body = build_event_body(&event);
583
584        let resp = self
585            .post(&url, body)?
586            .send()
587            .await
588            .map_err(|e| SessionError::Storage(format!("HTTP request failed: {e}")))?;
589
590        check_success(resp).await
591    }
592
593    async fn get_events(&self, id: &SessionId) -> Result<Vec<Event>, SessionError> {
594        let url = self.config.events_url(&self.engine_id, id.as_str());
595
596        let resp = self
597            .get(&url)?
598            .send()
599            .await
600            .map_err(|e| SessionError::Storage(format!("HTTP request failed: {e}")))?;
601
602        // 404 → empty list (session with no events or missing session).
603        let opt: Option<ListEventsResponse> = parse_json_response(resp).await?;
604        let events = opt
605            .map(|r| r.session_events.into_iter().map(map_vertex_event).collect())
606            .unwrap_or_default();
607        Ok(events)
608    }
609}
610
611// ──────────────────────────────────────────────────────────────────────────────
612// Small extraction helper (not pub — internal only)
613// ──────────────────────────────────────────────────────────────────────────────
614
615/// Try to pull the engine / reasoning-engine ID out of a Vertex AI resource
616/// name like `projects/p/locations/l/reasoningEngines/ENGINE/sessions/S`.
617fn extract_engine_id_from_name(name: &str) -> Option<String> {
618    let mut parts = name.split('/');
619    while let Some(segment) = parts.next() {
620        if segment == "reasoningEngines" {
621            return parts.next().map(String::from);
622        }
623    }
624    None
625}
626
627// ──────────────────────────────────────────────────────────────────────────────
628// Tests
629// ──────────────────────────────────────────────────────────────────────────────
630
631#[cfg(test)]
632mod tests {
633    use super::*;
634
635    // ── Existing tests (must remain green) ────────────────────────────────
636
637    #[test]
638    fn config_new() {
639        let config = VertexAiSessionConfig::new("my-project", "us-central1");
640        assert_eq!(config.project, "my-project");
641        assert_eq!(config.location, "us-central1");
642        assert!(config.ttl_seconds.is_none());
643    }
644
645    #[test]
646    fn config_with_ttl() {
647        let config = VertexAiSessionConfig::new("proj", "us-east1").ttl_seconds(3600);
648        assert_eq!(config.ttl_seconds, Some(3600));
649    }
650
651    #[test]
652    fn url_construction() {
653        let config = VertexAiSessionConfig::new("my-project", "us-central1");
654        assert_eq!(
655            config.base_url(),
656            "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/my-project/locations/us-central1"
657        );
658        assert!(
659            config
660                .sessions_url("engine-1")
661                .contains("reasoningEngines/engine-1/sessions")
662        );
663        assert!(
664            config
665                .session_url("engine-1", "sess-1")
666                .contains("sessions/sess-1")
667        );
668        assert!(
669            config
670                .events_url("engine-1", "sess-1")
671                .contains("sessions/sess-1/events")
672        );
673    }
674
675    #[test]
676    fn a_session_id_cannot_walk_out_of_its_collection() {
677        let config = VertexAiSessionConfig::new("my-project", "us-central1");
678
679        // Without encoding this reads as `.../sessions/../../otherEngine`,
680        // which addresses a resource under a different reasoning engine.
681        let url = config.session_url("engine-1", "../../otherEngine");
682        assert!(
683            url.ends_with("/sessions/..%2F..%2FotherEngine"),
684            "traversal was not encoded: {url}"
685        );
686
687        // And this one would end the path and start a query.
688        let url = config.session_url("engine-1", "sess-1?alt=media");
689        assert!(!url.contains('?'), "query separator survived: {url}");
690        assert!(url.ends_with("/sessions/sess-1%3Falt%3Dmedia"), "{url}");
691
692        // Ordinary ids are untouched — the unreserved set passes through.
693        assert!(
694            config
695                .session_url("engine-1", "sess-1_A.b~2")
696                .ends_with("/sessions/sess-1_A.b~2")
697        );
698    }
699
700    #[test]
701    fn a_user_id_cannot_append_its_own_query_parameter() {
702        // `&pageSize=1000` unencoded becomes a second parameter on the request
703        // rather than part of the userId being filtered on.
704        let encoded = percent_encode("alice&pageSize=1000");
705        assert_eq!(encoded, "alice%26pageSize%3D1000");
706        assert!(!encoded.contains('&'));
707
708        // Non-ASCII is escaped per byte, not dropped or passed through raw.
709        assert_eq!(percent_encode("josé"), "jos%C3%A9");
710    }
711
712    #[test]
713    fn service_accessors() {
714        let svc = VertexAiSessionService::new(
715            VertexAiSessionConfig::new("proj", "us-west1").ttl_seconds(7200),
716        );
717        assert_eq!(svc.project(), "proj");
718        assert_eq!(svc.location(), "us-west1");
719        assert_eq!(svc.ttl_seconds(), Some(7200));
720    }
721
722    // ── New: builder methods ───────────────────────────────────────────────
723
724    #[test]
725    fn with_token_sets_provider() {
726        let svc =
727            VertexAiSessionService::new(VertexAiSessionConfig::new("p", "l")).with_token("tok123");
728        // Verify the token provider returns the expected token.
729        let token = svc.token_provider.get().expect("should have token");
730        assert_eq!(token, "tok123");
731    }
732
733    #[test]
734    fn with_token_refresher_calls_closure() {
735        let svc = VertexAiSessionService::new(VertexAiSessionConfig::new("p", "l"))
736            .with_token_refresher(|| "dynamic-token".to_string());
737        let token = svc.token_provider.get().expect("should have token");
738        assert_eq!(token, "dynamic-token");
739    }
740
741    #[test]
742    fn missing_token_returns_storage_error() {
743        let svc = VertexAiSessionService::new(VertexAiSessionConfig::new("p", "l"));
744        let err = svc.token_provider.get().unwrap_err();
745        assert!(matches!(err, SessionError::Storage(_)));
746        let msg = err.to_string();
747        assert!(
748            msg.contains("missing auth token"),
749            "unexpected message: {msg}"
750        );
751    }
752
753    #[test]
754    fn reasoning_engine_overrides_default() {
755        let svc = VertexAiSessionService::new(VertexAiSessionConfig::new("p", "l"))
756            .reasoning_engine("my-engine");
757        assert_eq!(svc.engine_id, "my-engine");
758    }
759
760    #[test]
761    fn default_engine_id_is_default() {
762        let svc = VertexAiSessionService::new(VertexAiSessionConfig::new("p", "l"));
763        assert_eq!(svc.engine_id, "default");
764    }
765
766    // ── New: body builders ─────────────────────────────────────────────────
767
768    #[test]
769    fn build_create_body_without_ttl() {
770        let body = build_create_body("alice", None);
771        assert_eq!(body["userId"], "alice");
772        assert!(body.get("ttl").is_none());
773    }
774
775    #[test]
776    fn build_create_body_with_ttl() {
777        let body = build_create_body("bob", Some(3600));
778        assert_eq!(body["userId"], "bob");
779        assert_eq!(body["ttl"], "3600s");
780    }
781
782    #[test]
783    fn build_event_body_with_content() {
784        let event = Event::new("user", Some("Hello Vertex!".to_string()));
785        let body = build_event_body(&event);
786        assert_eq!(body["author"], "user");
787        assert_eq!(body["content"]["parts"][0]["text"], "Hello Vertex!");
788    }
789
790    #[test]
791    fn build_event_body_without_content() {
792        let event = Event::new("agent", None);
793        let body = build_event_body(&event);
794        assert_eq!(body["author"], "agent");
795        // content should be an empty object
796        assert!(body["content"].is_object());
797    }
798
799    #[test]
800    fn build_event_body_includes_invocation_id() {
801        let event = Event::new("model", None).with_invocation("inv-xyz");
802        let body = build_event_body(&event);
803        assert_eq!(body["invocationId"], "inv-xyz");
804    }
805
806    // ── New: response mapper helpers ───────────────────────────────────────
807
808    #[test]
809    fn session_id_from_name_extracts_last_segment() {
810        assert_eq!(
811            session_id_from_name(
812                "projects/p/locations/l/reasoningEngines/e/sessions/my-session-id"
813            ),
814            "my-session-id"
815        );
816        assert_eq!(session_id_from_name("just-an-id"), "just-an-id");
817    }
818
819    #[test]
820    fn extract_engine_id_from_name_works() {
821        let name = "projects/my-proj/locations/us-central1/reasoningEngines/eng-42/sessions/sess-1";
822        assert_eq!(
823            extract_engine_id_from_name(name),
824            Some("eng-42".to_string())
825        );
826    }
827
828    #[test]
829    fn extract_engine_id_returns_none_without_segment() {
830        assert_eq!(extract_engine_id_from_name("projects/p/locations/l"), None);
831    }
832
833    #[test]
834    fn map_vertex_session_maps_fields() {
835        let vs = VertexSession {
836            name: "projects/p/locations/l/reasoningEngines/e/sessions/sess-abc".to_string(),
837            user_id: "alice".to_string(),
838            session_state: Some(serde_json::json!({"key": "value"})),
839            create_time: Some("2024-01-01T00:00:00Z".to_string()),
840            update_time: Some("2024-01-02T00:00:00Z".to_string()),
841        };
842        let session = map_vertex_session(vs, "my-app");
843        assert_eq!(session.id.as_str(), "sess-abc");
844        assert_eq!(session.app_name, "my-app");
845        assert_eq!(session.user_id, "alice");
846        assert_eq!(session.state["key"], "value");
847        assert_eq!(session.created_at, "2024-01-01T00:00:00Z");
848        assert_eq!(session.updated_at, "2024-01-02T00:00:00Z");
849    }
850
851    #[test]
852    fn map_vertex_session_handles_missing_optionals() {
853        let vs = VertexSession {
854            name: "projects/p/locations/l/reasoningEngines/e/sessions/sess-xyz".to_string(),
855            user_id: String::new(),
856            session_state: None,
857            create_time: None,
858            update_time: None,
859        };
860        let session = map_vertex_session(vs, "app");
861        assert_eq!(session.id.as_str(), "sess-xyz");
862        assert!(session.state.is_empty());
863    }
864
865    #[test]
866    fn map_vertex_event_maps_text_content() {
867        let ve = VertexEvent {
868            author: "user".to_string(),
869            invocation_id: "inv-1".to_string(),
870            content: Some(serde_json::json!({ "parts": [{ "text": "Hi!" }] })),
871            actions: None,
872            name: Some(
873                "projects/p/locations/l/reasoningEngines/e/sessions/s/events/ev-1".to_string(),
874            ),
875            timestamp: Some("1700000000".to_string()),
876        };
877        let event = map_vertex_event(ve);
878        assert_eq!(event.author, "user");
879        assert_eq!(event.invocation_id, "inv-1");
880        assert_eq!(event.content, Some("Hi!".to_string()));
881        assert_eq!(event.id, "ev-1");
882        assert_eq!(event.timestamp, 1_700_000_000);
883    }
884
885    #[test]
886    fn map_vertex_event_handles_missing_content() {
887        let ve = VertexEvent {
888            author: "model".to_string(),
889            invocation_id: String::new(),
890            content: None,
891            actions: None,
892            name: None,
893            timestamp: None,
894        };
895        let event = map_vertex_event(ve);
896        assert_eq!(event.content, None);
897        assert_eq!(event.timestamp, 0);
898        // ID should be a generated UUID (non-empty)
899        assert!(!event.id.is_empty());
900    }
901
902    // ── New: 404 → None mapping tested via the pure classify_status helper ──
903    //
904    // We test the status-classification logic directly using the pure
905    // `classify_status` function rather than constructing synthetic
906    // `reqwest::Response` objects (which would require the `http` crate as a
907    // direct dependency). The async `parse_json_response` / `check_success`
908    // wrappers delegate to the same logic path, so these tests give full
909    // coverage without a live server.
910
911    #[test]
912    fn classify_status_404_is_none_signal() {
913        // classify_status returns Ok(false) to signal "treat as None"
914        let result = classify_status(404, "not found");
915        assert!(result.is_ok(), "expected Ok, got {result:?}");
916        assert!(!result.unwrap(), "expected false for 404");
917    }
918
919    #[test]
920    fn classify_status_200_is_success_signal() {
921        let result = classify_status(200, "");
922        assert!(result.unwrap(), "expected true for 200");
923    }
924
925    #[test]
926    fn classify_status_201_is_success_signal() {
927        let result = classify_status(201, "");
928        assert!(result.unwrap(), "expected true for 201");
929    }
930
931    #[test]
932    fn classify_status_299_is_success_signal() {
933        let result = classify_status(299, "");
934        assert!(result.unwrap(), "expected true for 299");
935    }
936
937    #[test]
938    fn classify_status_500_is_storage_error() {
939        let result = classify_status(500, "internal server error");
940        assert!(matches!(result, Err(SessionError::Storage(_))));
941        let msg = result.unwrap_err().to_string();
942        assert!(msg.contains("500"), "expected status in error: {msg}");
943        assert!(
944            msg.contains("internal server error"),
945            "expected body in error: {msg}"
946        );
947    }
948
949    #[test]
950    fn classify_status_403_is_storage_error() {
951        let result = classify_status(403, "forbidden");
952        assert!(matches!(result, Err(SessionError::Storage(_))));
953        let msg = result.unwrap_err().to_string();
954        assert!(msg.contains("403"), "expected status in error: {msg}");
955    }
956
957    #[test]
958    fn classify_status_400_is_storage_error() {
959        let result = classify_status(400, "bad request");
960        assert!(matches!(result, Err(SessionError::Storage(_))));
961    }
962
963    #[test]
964    fn classify_status_300_is_storage_error() {
965        // Redirects are not transparent in our usage — treat as error.
966        let result = classify_status(301, "moved permanently");
967        assert!(matches!(result, Err(SessionError::Storage(_))));
968    }
969}
gemini_adk_rs/session/vertex_ai.rs

gemini_adk_rs/session/
vertex_ai.rs
