aion-integration-acp 0.31.0

//! Two seams that did not exist before: a SECOND prompt turn on a live session, and the MCP
//! servers the harness hands the agent when a session opens.
//!
//! Both are driven the same way the rest of this module is — a real [`JsonRpcConnection`] against a
//! scripted fake agent over an in-memory duplex, with every asserted frame written as literal ACP
//! JSON rather than built from the schema's Rust types.
//!
//! [`JsonRpcConnection`]: aion_integrations::jsonrpc::JsonRpcConnection

use super::*;

use crate::mcp::McpServerSpec;

/// How long a bounded wait for the reader pump to notice end-of-stream may take.
///
/// The pump notices EOF on its next read, which is immediate; the bound exists so a regression
/// fails the test instead of hanging it.
const PUMP_END_BOUND: Duration = Duration::from_secs(5);

/// The two servers the positive control configures: one of each transport.
///
/// Two rather than one, and one of each kind, because a single stdio entry would have passed on a
/// harness that hard-coded a one-element list, and a single HTTP entry would not have exercised the
/// mandatory transport.
fn two_servers() -> Vec<McpServerSpec> {
    vec![
        McpServerSpec::Stdio {
            name: "meridian".to_owned(),
            command: PathBuf::from("/opt/meridian/mcp"),
            args: vec!["--stdio".to_owned()],
            env: vec![("MERIDIAN_HOME".to_owned(), "/var/meridian".to_owned())],
        },
        McpServerSpec::Http {
            name: "aion".to_owned(),
            url: "http://127.0.0.1:8080/mcp".to_owned(),
            headers: vec![("authorization".to_owned(), "Bearer minted".to_owned())],
        },
    ]
}

/// The exact wire form [`two_servers`] must produce, written by hand.
fn two_servers_on_the_wire() -> Value {
    json!([
        {
            "name": "meridian",
            "command": "/opt/meridian/mcp",
            "args": ["--stdio"],
            "env": [{ "name": "MERIDIAN_HOME", "value": "/var/meridian" }]
        },
        {
            "type": "http",
            "name": "aion",
            "url": "http://127.0.0.1:8080/mcp",
            "headers": [{ "name": "authorization", "value": "Bearer minted" }]
        }
    ])
}

// --- MCP servers on the wire ---

/// The positive control for the `mcpServers` pin: a harness configured with two servers sends
/// exactly those two on `session/new`.
///
/// Every other test in this module asserts the DEFAULT — an empty list. Without this cell the pin
/// could not tell "the adapter sends what it was configured with" from "the adapter always sends
/// `[]`", which is what it used to do.
#[tokio::test]
async fn a_configured_mcp_server_list_reaches_session_new_verbatim() -> TestResult {
    let (client_io, agent_io) = duplex(DUPLEX_CAPACITY);
    let (client_read, client_write) = split(client_io);
    let (agent_read, agent_write) = split(agent_io);
    let connection = Arc::new(JsonRpcConnection::new(client_read, client_write));
    let agent = Arc::new(JsonRpcConnection::new(agent_read, agent_write));

    let expected = two_servers_on_the_wire();
    let script = {
        let agent = Arc::clone(&agent);
        let expected = expected.clone();
        tokio::spawn(async move {
            agent_opens_expecting_mcp(
                agent.as_ref(),
                protocol::PROTOCOL_VERSION.as_u16(),
                &expected,
                true,
            )
            .await
        })
    };
    let opened = within(
        "the handshake carrying configured MCP servers",
        harness_with_mcp_servers(PermissionPolicy::Deny, two_servers()).connect(
            connection.as_ref(),
            PathBuf::from(WORKSPACE),
            "do the thing".to_owned(),
            None,
        ),
    )
    .await?
    .map_err(|error| fail("the handshake must succeed", &error))?;
    within("the agent script", script)
        .await?
        .map_err(|error| fail("the agent script task must not panic", &error))?
        .map_err(|error| fail("the agent script must complete", &error))?;
    assert_eq!(opened.session_id.to_string(), SESSION);
    drop(agent);
    Ok(())
}

/// The same list reaches `session/load`, so a resumed conversation keeps the tools a fresh one
/// gets.
///
/// An agent that reloaded a session with a smaller tool set would silently lose capabilities
/// mid-conversation, which is exactly the shape of defect the resume arm exists to prevent.
#[tokio::test]
async fn a_configured_mcp_server_list_reaches_session_load_verbatim() -> TestResult {
    let (client_io, agent_io) = duplex(DUPLEX_CAPACITY);
    let (client_read, client_write) = split(client_io);
    let (agent_read, agent_write) = split(agent_io);
    let connection = Arc::new(JsonRpcConnection::new(client_read, client_write));
    let agent = Arc::new(JsonRpcConnection::new(agent_read, agent_write));

    let expected = two_servers_on_the_wire();
    let script = {
        let agent = Arc::clone(&agent);
        let expected = expected.clone();
        tokio::spawn(async move {
            let initialize = expect_request(agent.as_ref(), protocol::METHOD_INITIALIZE).await?;
            agent
                .send_response(&JsonRpcResponse::success(
                    initialize.id,
                    json!({
                        "protocolVersion": protocol::PROTOCOL_VERSION.as_u16(),
                        "agentCapabilities": {
                            "loadSession": true,
                            "mcpCapabilities": { "http": true }
                        },
                        "authMethods": [],
                        "agentInfo": { "name": "fake-acp-agent", "version": "0.0.0" }
                    }),
                ))
                .await?;
            let load = expect_request(agent.as_ref(), protocol::METHOD_SESSION_LOAD).await?;
            let params = load
                .params
                .as_ref()
                .ok_or_else(|| HarnessError::protocol("session/load carried no params"))?;
            if params.get("mcpServers") != Some(&expected) {
                return Err(HarnessError::protocol(format!(
                    "session/load must send exactly the configured `mcpServers` ({expected}), got \
                     {:?}",
                    params.get("mcpServers")
                )));
            }
            agent
                .send_response(&JsonRpcResponse::success(load.id, json!({})))
                .await?;
            expect_request(agent.as_ref(), protocol::METHOD_SESSION_PROMPT).await?;
            Ok::<(), HarnessError>(())
        })
    };
    within(
        "the resume handshake carrying configured MCP servers",
        harness_with_mcp_servers(PermissionPolicy::Deny, two_servers()).connect(
            connection.as_ref(),
            PathBuf::from(WORKSPACE),
            "carry on".to_owned(),
            Some(agent_client_protocol_schema::v1::SessionId::from(
                "sess-prior".to_owned(),
            )),
        ),
    )
    .await?
    .map_err(|error| fail("the resume handshake must succeed", &error))?;
    within("the agent script", script)
        .await?
        .map_err(|error| fail("the agent script task must not panic", &error))?
        .map_err(|error| fail("the agent script must complete", &error))?;
    drop(agent);
    Ok(())
}

/// An HTTP MCP server against an agent that did not advertise `mcpCapabilities.http` is refused
/// BEFORE the session is opened, as a terminal configuration fault.
#[tokio::test]
async fn an_http_mcp_server_is_refused_when_the_agent_did_not_advertise_it() -> TestResult {
    let (client_io, agent_io) = duplex(DUPLEX_CAPACITY);
    let (client_read, client_write) = split(client_io);
    let (agent_read, agent_write) = split(agent_io);
    let connection = Arc::new(JsonRpcConnection::new(client_read, client_write));
    let agent = Arc::new(JsonRpcConnection::new(agent_read, agent_write));

    let script = {
        let agent = Arc::clone(&agent);
        tokio::spawn(async move {
            let initialize = expect_request(agent.as_ref(), protocol::METHOD_INITIALIZE).await?;
            agent
                .send_response(&JsonRpcResponse::success(
                    initialize.id,
                    json!({
                        "protocolVersion": protocol::PROTOCOL_VERSION.as_u16(),
                        "agentCapabilities": {
                            "loadSession": false,
                            "mcpCapabilities": { "http": false }
                        },
                        "authMethods": [],
                        "agentInfo": { "name": "fake-acp-agent", "version": "0.0.0" }
                    }),
                ))
                .await?;
            Ok::<(), HarnessError>(())
        })
    };
    let refused = within(
        "the handshake against an agent with no HTTP MCP support",
        harness_with_mcp_servers(PermissionPolicy::Deny, two_servers()).connect(
            connection.as_ref(),
            PathBuf::from(WORKSPACE),
            "do the thing".to_owned(),
            None,
        ),
    )
    .await?;
    within("the agent script", script)
        .await?
        .map_err(|error| fail("the agent script task must not panic", &error))?
        .map_err(|error| fail("the agent script must complete", &error))?;
    match refused {
        Ok(_opened) => Err(
            "an HTTP MCP server against an agent that cannot reach one must be refused".to_owned(),
        ),
        Err(error) => {
            assert!(
                matches!(error, HarnessError::Configuration { .. }),
                "the refusal must be terminal configuration, got {error:?}"
            );
            assert!(
                error.to_string().contains("aion"),
                "the refusal must name the server it is about: {error}"
            );
            drop(agent);
            Ok(())
        }
    }
}

// --- a second turn on a live session ---

/// Ends the turn the harness opened, draining its transcript first.
///
/// Returns the events that turn carried, so a caller can assert what did — and did not — reach it.
async fn finish_first_turn(
    session: &mut Session,
    agent: &Wire,
    prompt_id: JsonRpcId,
) -> TestResult<Vec<ActivityEvent>> {
    let mut first = session
        .take_first_turn()
        .ok_or_else(|| "the opening turn must be takeable exactly once".to_owned())?;
    agent
        .send_response(&prompt_response(prompt_id, "end_turn"))
        .await
        .map_err(|error| fail("the first terminal response must be writable", &error))?;
    let stream = first.events();
    let events = within(
        "draining the first turn's transcript",
        stream.collect::<Vec<_>>(),
    )
    .await?;
    let outcome = within("the first turn's result", first.wait())
        .await?
        .map_err(|error| fail("the first turn must complete", &error))?;
    assert_eq!(
        outcome.final_message, "",
        "the first turn streamed no message, so its final message is empty"
    );
    Ok(events)
}

/// The seam: a second `session/prompt` on the SAME live session, after the first turn's terminal
/// result, with its own transcript and its own result.
///
/// This is what makes an operator-paced conversation possible without `session/load` and without a
/// second agent process: the agent keeps everything it learned in turn one.
#[tokio::test]
async fn a_second_turn_runs_on_the_live_session() -> TestResult {
    let Wired {
        mut session,
        agent,
        prompt_id,
    } = wire_up(PermissionPolicy::Deny).await?;
    let first_events = finish_first_turn(&mut session, agent.as_ref(), prompt_id).await?;

    let mut second = within(
        "submitting the second turn",
        session.prompt("again".to_owned()),
    )
    .await?
    .map_err(|error| fail("a second turn on a live session must be accepted", &error))?;
    let request = within(
        "the second session/prompt reaching the agent",
        expect_request(agent.as_ref(), protocol::METHOD_SESSION_PROMPT),
    )
    .await?
    .map_err(|error| fail("the agent must receive the second prompt", &error))?;
    let params = request
        .params
        .as_ref()
        .ok_or_else(|| "the second prompt must carry params".to_owned())?;
    assert_eq!(
        params["sessionId"],
        json!(SESSION),
        "the second turn runs in the SAME session — that is the whole point"
    );
    assert_eq!(
        params["prompt"][0]["text"],
        json!("again"),
        "the second turn carries its own text as a `text` content block"
    );

    agent
        .send_notification(&session_update(&json!({
            "sessionUpdate": "agent_message_chunk",
            "content": { "type": "text", "text": "second answer" }
        })))
        .await
        .map_err(|error| fail("the second turn's chunk must be writable", &error))?;
    agent
        .send_response(&JsonRpcResponse::success(
            request.id,
            json!({ "stopReason": "end_turn" }),
        ))
        .await
        .map_err(|error| fail("the second terminal response must be writable", &error))?;

    let stream = second.events();
    let second_events = within(
        "draining the second turn's transcript",
        stream.collect::<Vec<_>>(),
    )
    .await?;
    let outcome = within("the second turn's result", second.wait())
        .await?
        .map_err(|error| fail("the second turn must complete", &error))?;

    assert_eq!(
        outcome.final_message, "second answer",
        "the second turn's accumulator is its own — it starts empty and carries only this turn"
    );
    assert!(
        second_events.iter().any(|event| matches!(
            &event.kind,
            ActivityEventKind::Delta { text_fragment, .. } if text_fragment == "second answer"
        )),
        "the second turn's chunk reaches the second turn's transcript"
    );
    assert!(
        !first_events.iter().any(|event| matches!(
            &event.kind,
            ActivityEventKind::Delta { text_fragment, .. } if text_fragment == "second answer"
        )),
        "the first turn's transcript closed at its own answer and carries nothing from the second"
    );
    drop(session);
    drop(agent);
    Ok(())
}

/// One turn at a time: a prompt submitted while another is open is refused, not queued.
///
/// ACP's streamed updates carry no turn id, so two open prompts would make every chunk ambiguous
/// about which turn it belongs to. The refusal is the typed `Occupied`, so a caller can tell it
/// from a dead session.
#[tokio::test]
async fn a_second_prompt_while_a_turn_is_open_is_refused_as_occupied() -> TestResult {
    let Wired {
        mut session,
        agent,
        prompt_id,
    } = wire_up(PermissionPolicy::Deny).await?;

    let refused = within(
        "submitting a turn while the opening one is still open",
        session.prompt("too soon".to_owned()),
    )
    .await?;
    match refused {
        Ok(_handle) => {
            return Err("a second concurrent prompt must be refused".to_owned());
        }
        Err(error) => assert!(
            matches!(error, HarnessError::Occupied { .. }),
            "the refusal must be `Occupied`, got {error:?}"
        ),
    }

    // The refusal did not consume the open turn: it still completes normally, which is what proves
    // the guard refused rather than damaged the session.
    finish_first_turn(&mut session, agent.as_ref(), prompt_id).await?;
    let accepted = within(
        "submitting a turn once the first has finished",
        session.prompt("now".to_owned()),
    )
    .await?;
    assert!(
        accepted.is_ok(),
        "the session is free again once the open turn has reported its result"
    );
    drop(session);
    drop(agent);
    Ok(())
}

/// A turn submitted after the agent has gone is refused IMMEDIATELY with a typed error, never left
/// waiting for an answer that cannot come.
///
/// The agent's write half is dropped, which is exactly what a killed subprocess does to its stdout;
/// the pump reads end-of-stream and closes the session to new turns. The wait for that is bounded,
/// so a regression that left `prompt` hanging fails here instead of wedging the suite.
#[tokio::test]
async fn a_turn_after_the_agent_has_gone_is_refused_rather_than_hung() -> TestResult {
    let Wired {
        mut session,
        agent,
        prompt_id,
    } = wire_up(PermissionPolicy::Deny).await?;
    finish_first_turn(&mut session, agent.as_ref(), prompt_id).await?;

    // The agent's end of the duplex goes away: its stdout closes, and the pump reads EOF.
    drop(agent);
    tokio::time::timeout(PUMP_END_BOUND, async {
        while session.is_live() {
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .map_err(|_elapsed| {
        format!("the pump did not notice end-of-stream within {PUMP_END_BOUND:?}")
    })?;

    let refused = within(
        "submitting a turn on a session whose agent has gone",
        session.prompt("anyone there".to_owned()),
    )
    .await?;
    match refused {
        Ok(_handle) => Err("a turn on a dead session must be refused".to_owned()),
        Err(error) => {
            assert!(
                matches!(error, HarnessError::StaleTarget { .. }),
                "the refusal must be `StaleTarget`, got {error:?}"
            );
            drop(session);
            Ok(())
        }
    }
}

// --- harness commands ---

/// A command turn reaches the agent in the form the ACP schema DESCRIBES.
///
/// ACP has no command transport. `AvailableCommandInput`'s only shape is
/// `unstructured`, documented as "all text that was typed after the command name is provided as
/// input" — so a command is delivered as ordinary prompt text: `/name` and whatever follows it,
/// inside a `text` content block like any other prompt.
///
/// The agent here ADVERTISES `compact` first, exactly as a real one does, and the frame reaches the
/// transcript rather than being dropped — so the surface above can decide what to offer from what
/// the agent actually said. Then the composed command line is asserted on the wire, verbatim.
///
/// The composition itself (`{name, input}` → `/name input`) is
/// `AssistantCommandInvocation::prompt_line` in `aion-core`, pinned there. This cell is the other
/// half: that the composed line is what an agent receives, with nothing added and nothing escaped.
#[tokio::test]
async fn an_advertised_command_reaches_the_agent_as_the_specs_own_prompt_line() -> TestResult {
    let Wired {
        mut session,
        agent,
        prompt_id,
    } = wire_up(PermissionPolicy::Deny).await?;

    // The agent advertises what it serves, as a real one does mid-turn.
    agent
        .send_notification(&session_update(&json!({
            "sessionUpdate": "available_commands_update",
            "availableCommands": [
                { "name": "compact", "description": "compact the conversation",
                  "input": { "hint": "what to keep" } }
            ]
        })))
        .await
        .map_err(|error| fail("the advertisement must be writable", &error))?;
    let first_events = finish_first_turn(&mut session, agent.as_ref(), prompt_id).await?;
    assert!(
        first_events.iter().any(|event| matches!(
            &event.kind,
            ActivityEventKind::Raw { source, value }
                if source.ends_with("available_commands_update")
                    && value["availableCommands"][0]["name"] == json!("compact")
        )),
        "the advertisement must reach the transcript, or nothing above can know what to offer: \
         {first_events:?}"
    );

    // The composed line, exactly as the server sends it for `{ name: "compact",
    // input: "keep the plan" }`.
    let mut turn = within(
        "submitting the command turn",
        session.prompt("/compact keep the plan".to_owned()),
    )
    .await?
    .map_err(|error| {
        fail(
            "a command turn is an ordinary prompt and must be accepted",
            &error,
        )
    })?;
    let request = within(
        "the command turn reaching the agent",
        expect_request(agent.as_ref(), protocol::METHOD_SESSION_PROMPT),
    )
    .await?
    .map_err(|error| fail("the agent must receive the command turn", &error))?;
    let params = request
        .params
        .as_ref()
        .ok_or_else(|| "the command turn must carry params".to_owned())?;
    assert_eq!(
        params["sessionId"],
        json!(SESSION),
        "a command runs in the conversation it was advertised in"
    );
    assert_eq!(
        params["prompt"][0]["type"],
        json!("text"),
        "ACP delivers a command as an ordinary text content block, not a shape of its own"
    );
    assert_eq!(
        params["prompt"][0]["text"],
        json!("/compact keep the plan"),
        "the command line reaches the agent verbatim: the name, a space, and what was typed after \
         it — which is exactly what the schema's `unstructured` input describes"
    );

    agent
        .send_response(&JsonRpcResponse::success(
            request.id,
            json!({ "stopReason": "end_turn" }),
        ))
        .await
        .map_err(|error| {
            fail(
                "the command turn's terminal response must be writable",
                &error,
            )
        })?;
    let stream = turn.events();
    drop(
        within(
            "draining the command turn's transcript",
            stream.collect::<Vec<_>>(),
        )
        .await?,
    );
    let outcome = within("the command turn's result", turn.wait())
        .await?
        .map_err(|error| fail("the command turn must complete", &error))?;
    assert_eq!(outcome.response.error, None);
    drop(session);
    drop(agent);
    Ok(())
}