abort.rs

use std::sync::Arc;

use async_trait::async_trait;
use github_copilot_sdk::handler::ApproveAllHandler;
use github_copilot_sdk::session_events::{AssistantMessageDeltaData, SessionEventType};
use github_copilot_sdk::tool::ToolHandler;
use github_copilot_sdk::{Error, SessionConfig, Tool, ToolInvocation, ToolResult};
use serde_json::json;
use tokio::sync::{Mutex, mpsc, oneshot};

use super::support::{
    DEFAULT_TEST_TOKEN, assistant_message_content, recv_with_timeout, wait_for_event,
};

#[tokio::test]
async fn should_abort_during_active_streaming() {
    super::support::with_dedicated_e2e_context(
        "abort",
        "should_abort_during_active_streaming",
        |ctx| {
            Box::pin(async move {
                ctx.set_default_copilot_user();
                let client = ctx.start_client().await;
                let session = client
                    .create_session(ctx.approve_all_session_config().with_streaming(true))
                    .await
                    .expect("create session");
                let events = session.subscribe();

                session
                .send(
                    "Write a very long essay about the history of computing, covering every decade \
                         from the 1940s to the 2020s in great detail.",
                )
                .await
                .expect("send long streaming turn");

                let delta = wait_for_event(events, "assistant.message_delta", |event| {
                    event.parsed_type() == SessionEventType::AssistantMessageDelta
                })
                .await;
                assert!(
                    !delta
                        .typed_data::<AssistantMessageDeltaData>()
                        .expect("assistant.message_delta data")
                        .delta_content
                        .is_empty()
                );

                session.abort().await.expect("abort session");

                // Session should be usable after abort. Wait for the specific recovery
                // message rather than racing against a late idle from the aborted turn.
                let recovery_events = session.subscribe();
                session
                    .send("Say 'abort_recovery_ok'.")
                    .await
                    .expect("send recovery");
                let recovery = wait_for_event(
                    recovery_events,
                    "assistant.message containing abort_recovery_ok",
                    |event| {
                        event.parsed_type() == SessionEventType::AssistantMessage
                            && assistant_message_content(event)
                                .to_lowercase()
                                .contains("abort_recovery_ok")
                    },
                )
                .await;
                assert!(
                    assistant_message_content(&recovery)
                        .to_lowercase()
                        .contains("abort_recovery_ok")
                );

                session.disconnect().await.expect("disconnect session");
                client.stop().await.expect("stop client");
            })
        },
    )
    .await;
}

#[tokio::test]
async fn should_abort_during_active_tool_execution() {
    super::support::with_dedicated_e2e_context(
        "abort",
        "should_abort_during_active_tool_execution",
        |ctx| {
            Box::pin(async move {
                ctx.set_default_copilot_user();
                let client = ctx.start_client().await;
                let (started_tx, mut started_rx) = mpsc::unbounded_channel();
                let (release_tx, release_rx) = oneshot::channel();
                let slow_tool = Arc::new(SlowAnalysisTool {
                    started_tx,
                    release_rx: Mutex::new(Some(release_rx)),
                });
                let session = client
                    .create_session(
                        SessionConfig::default()
                            .with_github_token(DEFAULT_TEST_TOKEN)
                            .with_permission_handler(Arc::new(ApproveAllHandler))
                            .with_tools(vec![
                                Tool::new("slow_analysis")
                                    .with_description(
                                        "A slow analysis tool that blocks until released",
                                    )
                                    .with_parameters(json!({
                                        "type": "object",
                                        "properties": {
                                            "value": {
                                                "type": "string",
                                                "description": "Value to analyze"
                                            }
                                        },
                                        "required": ["value"]
                                    }))
                                    .with_handler(slow_tool),
                            ]),
                    )
                    .await
                    .expect("create session");
                let events = session.subscribe();

                session
                    .send("Use slow_analysis with value 'test_abort'. Wait for the result.")
                    .await
                    .expect("send tool turn");

                let tool_value = recv_with_timeout(&mut started_rx, "slow tool start").await;
                assert_eq!(tool_value, "test_abort");

                session.abort().await.expect("abort session");
                release_tx
                    .send("RELEASED_AFTER_ABORT".to_string())
                    .expect("release slow tool");
                wait_for_event(events, "session.idle after abort", |event| {
                    event.parsed_type() == SessionEventType::SessionIdle
                })
                .await;

                let recovery = session
                    .send_and_wait("Say 'tool_abort_recovery_ok'.")
                    .await
                    .expect("send recovery")
                    .expect("assistant message");
                assert!(
                    assistant_message_content(&recovery)
                        .to_lowercase()
                        .contains("tool_abort_recovery_ok")
                );

                session.disconnect().await.expect("disconnect session");
                client.stop().await.expect("stop client");
            })
        },
    )
    .await;
}

struct SlowAnalysisTool {
    started_tx: mpsc::UnboundedSender<String>,
    release_rx: Mutex<Option<oneshot::Receiver<String>>>,
}

#[async_trait]
impl ToolHandler for SlowAnalysisTool {
    async fn call(&self, invocation: ToolInvocation) -> Result<ToolResult, Error> {
        let value = invocation
            .arguments
            .get("value")
            .and_then(serde_json::Value::as_str)
            .unwrap_or_default()
            .to_string();
        let _ = self.started_tx.send(value);
        let release_rx = self
            .release_rx
            .lock()
            .await
            .take()
            .expect("slow tool called once");
        let released = release_rx.await.unwrap_or_else(|_| "released".to_string());
        Ok(ToolResult::Text(released))
    }
}