aion-server 0.30.0

//! The session bearer: what it admits, what it refuses, and how long for.
//!
//! Every cell here drives the REAL resolution against a real store, so what is
//! measured is the admission decision the route makes rather than a restatement
//! of it. No process is spawned: admission is decided entirely from the record's
//! digest and the session's projected lifecycle, both of which are store facts.

use std::sync::Arc;

use aion_core::{AssistantSessionId, AssistantSessionState};
use aion_store::InMemoryStore;
use aion_store::assistant::AssistantSessionStore;

use crate::assistant::sessions::fixture::{
    OPERATOR, open_session, registry, registry_over, shared_context,
};
use crate::assistant::sessions::registry::AssistantSessions;
use crate::assistant::sessions::token::{self, MintedSessionToken, SESSION_ID_HEADER};

use super::caller::{AssistantMcpAuthError, resolve};

type TestResult = Result<(), String>;

/// Open a session and put a KNOWN bearer's digest on its record.
///
/// The mint is done here rather than through a spawn because a spawn would put
/// a process between the secret and the assertion — and the secret is exactly
/// what has to be held to present it.
async fn session_with_bearer(
    sessions: &AssistantSessions,
    store: &dyn AssistantSessionStore,
) -> Result<(AssistantSessionId, MintedSessionToken), String> {
    let summary = open_session(sessions).await?;
    let minted = MintedSessionToken::mint();
    let mut record = store
        .get_assistant_session(&summary.session_id)
        .await
        .map_err(|error| error.to_string())?
        .ok_or_else(|| "the created session has a record".to_owned())?;
    record.mcp_token_digest = Some(minted.digest().to_owned());
    store
        .put_assistant_session(record)
        .await
        .map_err(|error| error.to_string())?;
    Ok((summary.session_id, minted))
}

/// The headers an agent presents.
fn presented(session_id: AssistantSessionId, secret: &str) -> Vec<(String, String)> {
    vec![
        (
            "authorization".to_owned(),
            token::authorization_value(secret),
        ),
        (
            SESSION_ID_HEADER.to_owned(),
            token::session_header_value(session_id),
        ),
    ]
}

/// The baseline: a live session's own bearer is admitted, and the caller it
/// produces speaks for that session and no other.
#[tokio::test]
async fn a_live_sessions_own_bearer_is_admitted_for_that_session() -> TestResult {
    let (sessions, store) = registry();
    let (session_id, minted) = session_with_bearer(&sessions, store.as_ref()).await?;
    let caller = resolve(&sessions, &presented(session_id, minted.secret()))
        .await
        .map_err(|error| format!("a live session's own bearer must be admitted: {error}"))?;
    assert_eq!(caller.session_id(), session_id);
    Ok(())
}

/// Row (b), first pin: A's token after A ENDS is refused — even though nothing
/// expired it. The session's projected lifecycle IS the revocation, and it is a
/// store fact rather than an in-memory set.
#[tokio::test]
async fn a_bearer_stops_being_honoured_the_moment_its_session_ends() -> TestResult {
    let (sessions, store) = registry();
    let (session_id, minted) = session_with_bearer(&sessions, store.as_ref()).await?;
    // The control: it worked a moment ago, so the refusal below is the ending
    // and not a token that never verified.
    resolve(&sessions, &presented(session_id, minted.secret()))
        .await
        .map_err(|error| format!("the bearer must be good before the session ends: {error}"))?;

    sessions
        .delete(OPERATOR, session_id)
        .await
        .map_err(|error| error.to_string())?;
    let (state, _reason) = sessions
        .state_of_session(session_id)
        .await
        .map_err(|error| error.to_string())?;
    assert_eq!(state, AssistantSessionState::Ended);

    match resolve(&sessions, &presented(session_id, minted.secret())).await {
        Err(AssistantMcpAuthError::SessionEnded { session_id: named }) => {
            assert_eq!(named, session_id);
            Ok(())
        }
        other => Err(format!(
            "an ended session's bearer must be refused however unexpired it is, got {other:?}"
        )),
    }
}

/// Row (b), second pin: A's token presented against B is refused. The digest
/// that would have to match is B's record's, and it is not A's.
#[tokio::test]
async fn one_sessions_bearer_does_not_authorize_another_session() -> TestResult {
    let (sessions, store) = registry();
    let (first, first_token) = session_with_bearer(&sessions, store.as_ref()).await?;
    let (second, second_token) = session_with_bearer(&sessions, store.as_ref()).await?;
    assert_ne!(first, second);

    match resolve(&sessions, &presented(second, first_token.secret())).await {
        Err(AssistantMcpAuthError::WrongToken { session_id }) => assert_eq!(session_id, second),
        other => {
            return Err(format!(
                "session A's bearer against session B must be refused, got {other:?}"
            ));
        }
    }
    // Both directions, so the pin is not passing because one record happened to
    // carry no digest.
    match resolve(&sessions, &presented(first, second_token.secret())).await {
        Err(AssistantMcpAuthError::WrongToken { session_id }) => assert_eq!(session_id, first),
        other => {
            return Err(format!(
                "session B's bearer against session A must be refused, got {other:?}"
            ));
        }
    }
    Ok(())
}

/// Row (b), third pin: a RESTART between the mint and the call does not
/// invalidate the bearer. The digest is a store fact, so a token minted by a
/// process that has since died still verifies for as long as its session lives.
///
/// The restart is modelled as what it actually is — a NEW registry over the same
/// durable store, holding none of the previous process's memory. A resolution
/// that consulted anything in-process would fail here.
#[tokio::test]
async fn a_bearer_survives_the_restart_of_the_process_that_minted_it() -> TestResult {
    let store = Arc::new(InMemoryStore::default());
    let minting = registry_over(Arc::clone(&store) as Arc<dyn AssistantSessionStore>);
    let (session_id, minted) = session_with_bearer(&minting, store.as_ref()).await?;
    // The session is dormant: created, and nothing has said anything in it —
    // which is a continuable session, and the state a restart leaves one in.
    let (state, _reason) = minting
        .state_of_session(session_id)
        .await
        .map_err(|error| error.to_string())?;
    assert!(
        state.is_continuable(),
        "the fixture session must still be continuable for this pin to be about the restart"
    );
    drop(minting);

    // A second registry over the SAME store, holding nothing the first held.
    let after_restart = registry_over(Arc::clone(&store) as Arc<dyn AssistantSessionStore>);
    let caller = resolve(&after_restart, &presented(session_id, minted.secret()))
        .await
        .map_err(|error| {
            format!("a bearer minted before a restart must still be honoured after one: {error}")
        })?;
    assert_eq!(caller.session_id(), session_id);
    Ok(())
}

/// A HUMAN token is refused. This route accepts one credential — the bearer this
/// server minted for one session — and an operator's own token is not it,
/// whatever authority it carries elsewhere.
#[tokio::test]
async fn a_human_token_does_not_authorize_the_assistant_route() -> TestResult {
    let (sessions, store) = registry();
    let (session_id, _minted) = session_with_bearer(&sessions, store.as_ref()).await?;
    // A plausible operator bearer, presented with the session it would like to
    // read. It is not the minted secret, so it does not verify.
    match resolve(
        &sessions,
        &presented(session_id, "an-operator-bearer-token"),
    )
    .await
    {
        Err(AssistantMcpAuthError::WrongToken { .. }) => {}
        other => {
            return Err(format!(
                "a human token must not authorize a session's context, got {other:?}"
            ));
        }
    }
    // And with no session named at all — the shape a human token actually
    // arrives in, since nothing tells a person to send the session header.
    match resolve(
        &sessions,
        &[(
            "authorization".to_owned(),
            "Bearer an-operator-bearer-token".to_owned(),
        )],
    )
    .await
    {
        Err(AssistantMcpAuthError::NoSession { header }) => {
            assert_eq!(header, SESSION_ID_HEADER);
            Ok(())
        }
        other => Err(format!(
            "a call naming no session must be refused, got {other:?}"
        )),
    }
}

/// T3's isolation pin, end to end through the tool: session A's credential
/// answers with A's context and never B's.
#[tokio::test]
async fn one_sessions_credential_never_reads_another_sessions_context() -> TestResult {
    let (sessions, store) = registry();
    let (first, first_token) = session_with_bearer(&sessions, store.as_ref()).await?;
    let (second, _second_token) = session_with_bearer(&sessions, store.as_ref()).await?;
    shared_context(&sessions, first, "/studio/a.awl").await?;
    shared_context(&sessions, second, "/runs/b").await?;

    let caller = resolve(&sessions, &presented(first, first_token.secret()))
        .await
        .map_err(|error| error.to_string())?;
    let answered = super::context_tool::assistant_context(&sessions, &caller)
        .await
        .map_err(|failure| failure.to_string())?;
    assert_eq!(
        answered.structured["url"],
        serde_json::json!("/studio/a.awl")
    );
    assert_ne!(
        answered.structured["url"],
        serde_json::json!("/runs/b"),
        "session A's credential must never read session B's screen"
    );
    Ok(())
}

/// A session that was opened with no MCP server of ours mints no bearer, and a
/// call presenting anything against it is refused rather than admitted on an
/// absent digest — the arm where "no digest stored" could silently become "any
/// token matches".
#[tokio::test]
async fn a_session_with_no_minted_bearer_admits_nothing() -> TestResult {
    let (sessions, _store) = registry();
    let summary = open_session(&sessions).await?;
    match resolve(&sessions, &presented(summary.session_id, "anything")).await {
        Err(AssistantMcpAuthError::NoTokenMinted { session_id }) => {
            assert_eq!(session_id, summary.session_id);
            Ok(())
        }
        other => Err(format!(
            "a session with no stored digest must admit nothing, got {other:?}"
        )),
    }
}

/// A session id that names no record is refused as unknown, not admitted and
/// not crashed on.
#[tokio::test]
async fn a_bearer_naming_a_session_that_does_not_exist_is_refused() -> TestResult {
    let (sessions, _store) = registry();
    let absent = AssistantSessionId::new_v4();
    match resolve(&sessions, &presented(absent, "anything")).await {
        Err(AssistantMcpAuthError::UnknownSession { session_id }) => {
            assert_eq!(session_id, absent);
            Ok(())
        }
        other => Err(format!(
            "a bearer naming no session must be refused, got {other:?}"
        )),
    }
}