aion-server 0.29.0

//! `assistant_context` — the tool that answers "what is the operator looking
//! at".
//!
//! The answer is read off the SESSION'S OWN TRANSCRIPT: the latest
//! `request`/`context_shared` record is the shared context, so both console
//! surfaces and this tool read one fact from one place, and a restart loses
//! none of it.
//!
//! # Positions are one-based here
//!
//! The wire carries zero-based editor coordinates (that is what a document model
//! holds). A model reading a selection is reading the numbers a human quotes off
//! a gutter, so the `+1` happens here — in the same direction and for the same
//! reason `sessions/prompt.rs` does it for the prose.

use aion_core::{AssistantDocumentContext, AssistantTurnContext};
use aion_mcp::tools::service::{ToolFailure, ToolOutcome};
use serde_json::{Value, json};

use crate::assistant::sessions::AssistantSessions;

use super::caller::AssistantSessionCaller;

/// Answer `assistant_context` for the session the caller speaks for.
///
/// # Errors
///
/// [`ToolFailure`] when the session's transcript cannot be read. Phrased for the
/// model that reads it: it says the context is unavailable, not that the screen
/// is empty, so an agent does not conclude the operator is looking at nothing.
pub(crate) async fn assistant_context(
    sessions: &AssistantSessions,
    caller: &AssistantSessionCaller,
) -> Result<ToolOutcome, ToolFailure> {
    let projection = sessions
        .projection(caller.session_id())
        .await
        .map_err(|error| {
            ToolFailure::new(
                format!(
                    "the operator's on-screen context could not be read for this conversation: \
                     {error}. This does NOT mean the screen is empty — it means this server could \
                     not read it, so do not conclude anything about what is open."
                ),
                json!({ "code": "context_unreadable" }),
            )
        })?;
    Ok(outcome(
        projection.latest_context.as_ref(),
        projection.document_revision,
    ))
}

/// Render the answer, shared or not.
fn outcome(context: Option<&AssistantTurnContext>, revision: u64) -> ToolOutcome {
    let Some(context) = context else {
        return ToolOutcome {
            structured: json!({ "shared": false }),
            summary: "The operator has not shared what is on their screen yet.".to_owned(),
        };
    };
    ToolOutcome {
        structured: json!({
            "shared": true,
            "url": context.url,
            "concepts": context.concepts,
            "document": context.document.as_ref().map(document),
            "revision": revision,
        }),
        summary: summary(context),
    }
}

/// The document block, with positions moved into the numbers a human reads.
fn document(document: &AssistantDocumentContext) -> Value {
    json!({
        "path": document.path,
        "text": document.text,
        "selection": document.selection.map(|selection| json!({
            "from_line": selection.from.line.saturating_add(1),
            "from_column": selection.from.column.saturating_add(1),
            "to_line": selection.to.line.saturating_add(1),
            "to_column": selection.to.column.saturating_add(1),
        })),
        "cursor": document.cursor.map(|cursor| json!({
            "line": cursor.line.saturating_add(1),
            "column": cursor.column.saturating_add(1),
        })),
    })
}

/// Where in the document the operator is, for the summary line: a selection
/// when there is one, else the caret's one-based line, else nothing selected.
fn describe_where(document: &AssistantDocumentContext) -> String {
    match (document.selection.as_ref(), document.cursor.as_ref()) {
        (Some(_), _) => " with a selection".to_owned(),
        (None, Some(cursor)) => format!(
            " with nothing selected and the cursor on line {}",
            cursor.line.saturating_add(1)
        ),
        (None, None) => " with nothing selected".to_owned(),
    }
}

/// One line naming what came back, so a model can act without parsing the
/// structured block first.
fn summary(context: &AssistantTurnContext) -> String {
    match (context.document.as_ref(), context.url.as_deref()) {
        (Some(document), _) => format!(
            "The operator is editing `{}`{}.",
            document.path,
            describe_where(document)
        ),
        (None, Some(url)) => format!("The operator is on `{url}` with no document open."),
        (None, None) => {
            "The operator shared a context that names neither a page nor a document.".to_owned()
        }
    }
}

#[cfg(test)]
mod tests {
    use aion_core::{AssistantDocumentPosition, AssistantDocumentSelection};

    use super::*;

    fn context() -> AssistantTurnContext {
        AssistantTurnContext {
            url: Some("/studio/pipeline.awl".to_owned()),
            concepts: vec!["awl.step".to_owned()],
            document: Some(AssistantDocumentContext {
                path: "pipeline.awl".to_owned(),
                text: "workflow demo {}".to_owned(),
                selection: Some(AssistantDocumentSelection {
                    from: AssistantDocumentPosition { line: 0, column: 0 },
                    to: AssistantDocumentPosition { line: 2, column: 4 },
                }),
                cursor: Some(AssistantDocumentPosition { line: 2, column: 4 }),
            }),
        }
    }

    /// The caret is answered in the same one-based numbers as a selection, and
    /// a document that carries no caret answers `null` rather than a made-up
    /// position.
    #[test]
    fn the_caret_is_answered_one_based_or_null() {
        let answered = outcome(Some(&context()), 0);
        assert_eq!(
            answered.structured["document"]["cursor"],
            json!({ "line": 3, "column": 5 })
        );
        let mut without = context();
        if let Some(document) = without.document.as_mut() {
            document.cursor = None;
            document.selection = None;
        }
        let answered = outcome(Some(&without), 0);
        assert_eq!(answered.structured["document"]["cursor"], Value::Null);
        assert!(answered.summary.ends_with("with nothing selected."));
        let mut caret_only = context();
        if let Some(document) = caret_only.document.as_mut() {
            document.selection = None;
        }
        let answered = outcome(Some(&caret_only), 0);
        assert!(
            answered
                .summary
                .ends_with("with nothing selected and the cursor on line 3.")
        );
    }

    /// Nothing shared is `shared: false` and NOTHING ELSE — never an empty
    /// document or a null path, either of which an agent would read as "the
    /// operator has an empty file open".
    #[test]
    fn an_unshared_context_says_so_rather_than_inventing_an_empty_screen() {
        let answered = outcome(None, 0);
        assert_eq!(answered.structured, json!({ "shared": false }));
        assert!(answered.structured.get("document").is_none());
        assert!(answered.summary.contains("not shared"));
    }

    /// The `+1` happens exactly once, here. A selection quoted to a model in
    /// zero-based numbers would send it to the line above the one the operator
    /// is pointing at.
    #[test]
    fn a_selection_is_answered_in_the_numbers_the_operator_reads() {
        let answered = outcome(Some(&context()), 3);
        let selection = &answered.structured["document"]["selection"];
        assert_eq!(selection["from_line"], json!(1));
        assert_eq!(selection["from_column"], json!(1));
        assert_eq!(selection["to_line"], json!(3));
        assert_eq!(selection["to_column"], json!(5));
    }

    /// The path is in the answer, which is the whole reason the tool exists:
    /// an agent that has this never has to ask which directory anything is in.
    #[test]
    fn the_document_path_is_answered_so_nothing_has_to_ask_where_the_repo_is() {
        let answered = outcome(Some(&context()), 3);
        assert_eq!(
            answered.structured["document"]["path"],
            json!("pipeline.awl")
        );
        assert_eq!(answered.structured["shared"], json!(true));
        assert!(answered.summary.contains("pipeline.awl"));
    }

    /// The revision travels with the answer, so an agent re-reading the
    /// context can tell a rebased document from the one it last edited without
    /// diffing the text.
    #[test]
    fn a_shared_context_carries_the_document_revision() {
        let answered = outcome(Some(&context()), 7);
        assert_eq!(answered.structured["revision"], json!(7));
    }

    /// A document with no selection offers the whole file, and says so rather
    /// than omitting the field.
    #[test]
    fn a_document_with_no_selection_answers_a_null_selection() {
        let mut unselected = context();
        if let Some(document) = unselected.document.as_mut() {
            document.selection = None;
        }
        let answered = outcome(Some(&unselected), 0);
        assert_eq!(answered.structured["document"]["selection"], Value::Null);
        assert!(answered.summary.contains("nothing selected"));
    }
}