objectiveai_mcp_proxy/

session_manager.rs

//! Session registry.
//!
//! Maps the **objectiveai response id** (`X-OBJECTIVEAI-RESPONSE-ID`) to
//! the [`Session`] holding that response's live upstream MCP
//! connections. Routing is keyed entirely by response id; the
//! `Mcp-Session-Id` the proxy mints on `initialize` is a plain random
//! UUID returned for MCP-spec compliance only (so 3rd-party clients like
//! the Claude Agent SDK get a valid session header) — the proxy never
//! routes on it.
//!
//! All per-session dispatch (list, call, read) lives on [`Session`]
//! itself; this file only packs the opened connections into a [`Session`]
//! under their response id and looks sessions back up.

use std::sync::Arc;

use dashmap::DashMap;
use indexmap::IndexMap;

use crate::reverse_channel::Upstream;
use crate::session::Session;

/// Maps an objectiveai response id to its [`Session`] state.
#[derive(Debug, Default)]
pub struct SessionManager {
    sessions: DashMap<String, Arc<Session>>,
}

impl SessionManager {
    pub fn new() -> Self {
        Self::default()
    }

    /// Register the upstreams opened for one objectiveai response under
    /// its response id. Builds the `prefix -> Upstream` routing map (see
    /// [`build_prefix_map`]) and inserts a fresh [`Session`]. An existing
    /// entry for the same response id is replaced.
    pub fn add(&self, response_id: String, connections: Vec<Upstream>) {
        let by_prefix = build_prefix_map(connections);
        self.sessions
            .insert(response_id, Arc::new(Session::new(by_prefix)));
    }

    /// Cheap clone-out of a [`Session`] by response id — never holds a
    /// DashMap guard across the await boundary.
    pub fn get(&self, response_id: &str) -> Option<Arc<Session>> {
        self.sessions.get(response_id).map(|e| e.value().clone())
    }

    /// Remove a session from the registry by response id. Returns
    /// `Some(_)` if a session was present, `None` if the id was unknown.
    ///
    /// Once every `Arc<Session>` to the removed session has dropped, the
    /// session's `IndexMap<String, Connection>` drops, every `Connection`'s
    /// `Drop` fires its upstream's wakeup signal, and each upstream's
    /// listener task wakes to re-check liveness. The listener sees
    /// `Arc::strong_count == 1` (only itself) and exits, which drops the
    /// inner state and closes the upstream HTTP session.
    pub fn remove(&self, response_id: &str) -> Option<Arc<Session>> {
        self.sessions.remove(response_id).map(|(_, session)| session)
    }
}

/// Parse an `MCP_ENCRYPTION_KEY` env-var value: a single base64-encoded
/// 32-byte key. Empty string → `None`. Malformed → `Err`.
///
/// Retained as a no-op compatibility shim: the proxy no longer encrypts
/// session ids (they're plain UUIDs now), but the API crate still parses
/// and forwards the configured key. Kept so that the env plumbing in the
/// API/CLI compiles unchanged; the parsed value is ignored by the proxy.
/// Removal is deferred to a later step of the session-id refactor.
pub fn parse_key_env(s: &str) -> Result<Option<[u8; 32]>, String> {
    use base64::Engine;
    let trimmed = s.trim();
    if trimmed.is_empty() {
        return Ok(None);
    }
    let decoded = base64::engine::general_purpose::STANDARD
        .decode(trimmed)
        .map_err(|e| format!("MCP_ENCRYPTION_KEY: not valid base64: {e}"))?;
    let key: [u8; 32] = decoded.try_into().map_err(|got: Vec<u8>| {
        format!(
            "MCP_ENCRYPTION_KEY: expected 32 bytes after base64-decode, got {}",
            got.len(),
        )
    })?;
    Ok(Some(key))
}

/// Normalize a server name or version into a routing-prefix-safe token:
/// every `_` and `.` becomes `-`. The resulting prefix is free of both the
/// split separator (`_`) and `.`, so the first `_` in a prefixed identifier
/// is always the prefix/original-name boundary.
fn normalize_prefix_token(s: &str) -> String {
    s.replace(['_', '.'], "-")
}

/// Build the `prefix -> Connection` routing map.
///
/// Connections are sorted by `url` first so the prefixes and indices are
/// stable across calls — a client that re-lists must see the exact tool
/// names it already holds.
///
/// Each connection's prefix escalates only as far as needed for global
/// uniqueness:
///   1. `normalize(server_info.name)`
///   2. on collision -> `{name}-{normalize(version)}` (all colliding members)
///   3. still colliding -> `{name}-{version}-{index}` (index = url-sorted
///      position, globally unique, so this tier always resolves)
/// Uniqueness is re-checked over the full set after each tier so a rare
/// cross-tier collision escalates too.
fn build_prefix_map(
    mut connections: Vec<Upstream>,
) -> IndexMap<String, Upstream> {
    connections.sort_by(|a, b| a.url().cmp(b.url()));
    let n = connections.len();

    let names: Vec<String> = connections
        .iter()
        .map(|c| normalize_prefix_token(c.server_name()))
        .collect();
    let versions: Vec<String> = connections
        .iter()
        .map(|c| normalize_prefix_token(c.server_version()))
        .collect();

    // tier: 1 = name, 2 = name-version, 3 = name-version-index.
    let prefix_at = |i: usize, tier: u8| -> String {
        match tier {
            1 => names[i].clone(),
            2 => format!("{}-{}", names[i], versions[i]),
            _ => format!("{}-{}-{}", names[i], versions[i], i),
        }
    };

    let mut tier: Vec<u8> = vec![1; n];
    loop {
        let current: Vec<String> = (0..n).map(|i| prefix_at(i, tier[i])).collect();
        let mut counts: std::collections::HashMap<&str, usize> =
            std::collections::HashMap::new();
        for p in &current {
            *counts.entry(p.as_str()).or_insert(0) += 1;
        }
        let mut changed = false;
        for i in 0..n {
            if counts[current[i].as_str()] > 1 && tier[i] < 3 {
                tier[i] += 1;
                changed = true;
            }
        }
        if !changed {
            break;
        }
    }

    let mut by_prefix: IndexMap<String, Upstream> = IndexMap::with_capacity(n);
    for (i, connection) in connections.into_iter().enumerate() {
        let key = prefix_at(i, tier[i]);
        // The index tier guarantees uniqueness; a residual duplicate would
        // be a logic bug rather than a real-world collision.
        debug_assert!(
            !by_prefix.contains_key(&key),
            "duplicate routing prefix after escalation: {key}",
        );
        by_prefix.insert(key, connection);
    }
    by_prefix
}

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

    #[test]
    fn parse_key_env_round_trip() {
        use base64::Engine;
        let key = [0xAAu8; 32];
        let env = base64::engine::general_purpose::STANDARD.encode(key);
        let parsed = parse_key_env(&env).expect("parse").expect("Some");
        assert_eq!(parsed, key);

        assert!(parse_key_env("").unwrap().is_none());
        assert!(parse_key_env("   ").unwrap().is_none());
        assert!(parse_key_env("not-base64!@#").is_err());
        // Wrong-length payload (16 bytes after b64 decode):
        let short =
            base64::engine::general_purpose::STANDARD.encode(&[0u8; 16][..]);
        assert!(parse_key_env(&short).is_err());
    }
}