use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use serde_json::{Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream, duplex};
use tokio::sync::mpsc;
use tokio::time::timeout;
use super::*;
use crate::generated::api_types::{HostGitHubEnvironmentOptions, HostLocalServerOptions};
fn local_options() -> AhpHostOptions {
AhpHostOptions::new().with_local_server(HostLocalServerOptions::default())
}
const TIMEOUT: Duration = Duration::from_secs(5);
#[path = "factory_tests.rs"]
mod factory_tests;
struct Peer {
read: DuplexStream,
write: DuplexStream,
}
impl Peer {
async fn request(&mut self) -> Value {
timeout(TIMEOUT, async {
let mut header = Vec::new();
while !header.ends_with(b"\r\n\r\n") {
header.push(self.read.read_u8().await.unwrap());
}
let length: usize = std::str::from_utf8(&header)
.unwrap()
.trim()
.strip_prefix("Content-Length: ")
.unwrap()
.parse()
.unwrap();
let mut body = vec![0; length];
self.read.read_exact(&mut body).await.unwrap();
serde_json::from_slice(&body).unwrap()
})
.await
.unwrap()
}
async fn send(&mut self, value: Value) {
let body = serde_json::to_vec(&value).unwrap();
let header = format!("Content-Length: {}\r\n\r\n", body.len());
self.write.write_all(header.as_bytes()).await.unwrap();
self.write.write_all(&body).await.unwrap();
self.write.flush().await.unwrap();
}
async fn respond(&mut self, request: &Value, result: Value) {
self.send(json!({"jsonrpc": "2.0", "id": request["id"], "result": result}))
.await;
}
async fn started(&mut self, request: &Value, token: Option<&str>) {
let mut result = json!({
"hostId": request["params"]["hostId"],
"url": "http://127.0.0.1:4321"
});
if let Some(token) = token {
result["token"] = json!(token);
}
self.respond(request, result).await;
}
async fn exited(&mut self, host_id: &str) {
self.send(json!({
"jsonrpc": "2.0", "method": "host.exited",
"params": {"hostId": host_id, "reason": "exited", "exitCode": 17}
}))
.await;
}
}
fn fixture() -> (Client, Peer) {
let (client_write, read) = duplex(1 << 20);
let (write, client_read) = duplex(1 << 20);
(
Client::from_streams(client_read, client_write, PathBuf::from(".")).unwrap(),
Peer { read, write },
)
}
fn start(
client: &Client,
options: AhpHostOptions,
) -> tokio::task::JoinHandle<Result<AhpHost, Error>> {
let client = client.clone();
tokio::spawn(async move { client.start_ahp_host(options).await })
}
fn callback_options() -> (AhpHostOptions, mpsc::UnboundedReceiver<AhpHostExit>) {
let (tx, rx) = mpsc::unbounded_channel();
(
local_options().with_on_exit(move |exit| {
let _ = tx.send(exit);
}),
rx,
)
}
#[tokio::test]
async fn requires_an_explicit_transport() {
let (client, _peer) = fixture();
let error = client
.start_ahp_host(AhpHostOptions::default())
.await
.unwrap_err();
assert!(matches!(error.kind(), ErrorKind::InvalidConfig));
assert!(
error
.to_string()
.contains("requires localServer or githubEnvironment")
);
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
}
#[tokio::test]
async fn github_and_combined_transports_preserve_optional_results() {
for local in [false, true] {
let (client, mut peer) = fixture();
let mut options =
AhpHostOptions::new().with_github_environment(HostGitHubEnvironmentOptions {
name: "SDK host".into(),
compute_id: "compute".into(),
});
if local {
options = options.with_local_server(Default::default());
}
let pending = start(&client, options);
let request = peer.request().await;
let mut expected = json!({
"hostId": request["params"]["hostId"],
"githubEnvironment": {"name": "SDK host", "computeId": "compute"}
});
let mut result = json!({
"hostId": request["params"]["hostId"],
"environmentId": "environment"
});
if local {
expected["localServer"] = json!({});
result["url"] = json!("ws://127.0.0.1:4321");
}
assert_eq!(request["params"], expected);
peer.respond(&request, result).await;
let host = pending.await.unwrap().unwrap();
assert_eq!(host.environment_id.as_deref(), Some("environment"));
assert_eq!(host.url.as_deref(), local.then_some("ws://127.0.0.1:4321"));
assert_eq!(host.token, None);
assert_eq!(host.pid, None);
}
}
#[tokio::test]
async fn forwards_only_generated_options_and_returns_runtime_fields() {
let (client, mut peer) = fixture();
let (options, _exits) = callback_options();
let pending = start(
&client,
options.with_local_server(HostLocalServerOptions {
hostname: Some("::1".into()),
port: Some(0),
token: Some("explicit-token".into()),
require_connection_token: Some(true),
}),
);
let request = peer.request().await;
assert_eq!(request["method"], "host.start");
let host_id = request["params"]["hostId"].as_str().unwrap();
assert_eq!(uuid::Uuid::parse_str(host_id).unwrap().get_version_num(), 4);
assert_eq!(
request["params"],
json!({
"hostId": host_id, "localServer": {
"hostname": "::1", "port": 0,
"token": "explicit-token", "requireConnectionToken": true
}
})
);
peer.started(&request, Some("runtime-token")).await;
let host = pending.await.unwrap().unwrap();
assert_eq!(host.host_id, host_id);
assert_eq!(host.pid, None);
assert_eq!(host.url.as_deref(), Some("http://127.0.0.1:4321"));
assert_eq!(host.environment_id, None);
assert_eq!(host.token.as_deref(), Some("runtime-token"));
}
#[tokio::test]
async fn preserves_optional_legacy_separate_host_pid() {
let (client, mut peer) = fixture();
let pending = start(&client, local_options());
let request = peer.request().await;
peer.respond(
&request,
json!({
"hostId": request["params"]["hostId"],
"url": "ws://127.0.0.1:4321",
"pid": 1234
}),
)
.await;
assert_eq!(pending.await.unwrap().unwrap().pid, Some(1234));
}
#[tokio::test]
async fn listener_task_exit_has_no_process_exit_code() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
peer.send(json!({
"jsonrpc": "2.0", "method": "host.exited",
"params": {"hostId": host.host_id, "reason": "exited"}
}))
.await;
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_eq!(exit.reason, AhpHostExitReason::Exited);
assert_eq!(exit.exit_code, None);
}
#[tokio::test]
async fn explicit_local_defaults_are_forwarded_and_token_can_be_absent() {
let (client, mut peer) = fixture();
let pending = start(&client, local_options());
let request = peer.request().await;
assert_eq!(request["params"].as_object().unwrap().len(), 2);
assert_eq!(request["params"]["localServer"], json!({}));
peer.started(&request, None).await;
assert!(pending.await.unwrap().unwrap().token.is_none());
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
}
#[tokio::test]
async fn publication_forwards_existing_identity_to_the_owning_listener() {
let (client, mut peer) = fixture();
let pending = start(&client, local_options());
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
let host_id = host.host_id.clone();
let publication = tokio::spawn(async move { host.publish_session("resident").await });
let request = peer.request().await;
assert_eq!(request["method"], "host.publishSession");
assert_eq!(
request["params"],
json!({"hostId": host_id, "sessionId": "resident"})
);
peer.respond(
&request,
json!({"sessionId": "resident", "sessionUri": "copilot:/resident"}),
)
.await;
let result = publication.await.unwrap().unwrap();
assert_eq!(result.session_id, "resident");
assert_eq!(result.session_uri, "copilot:/resident");
}
#[tokio::test]
async fn listener_validation_is_left_to_the_runtime() {
let (client, mut peer) = fixture();
let pending = start(
&client,
AhpHostOptions::default().with_local_server(HostLocalServerOptions {
hostname: Some("".into()),
port: Some(-1),
token: Some("".into()),
require_connection_token: Some(false),
}),
);
let request = peer.request().await;
assert_eq!(request["params"]["localServer"]["port"], -1);
assert_eq!(request["params"]["localServer"]["hostname"], "");
assert_eq!(request["params"]["localServer"]["token"], "");
assert_eq!(
request["params"]["localServer"]["requireConnectionToken"],
false
);
peer.send(json!({
"jsonrpc": "2.0", "id": request["id"],
"error": {"code": -32602, "message": "invalid listener"}
}))
.await;
let error = pending.await.unwrap().unwrap_err();
assert!(matches!(error.kind(), ErrorKind::Rpc { code: -32602 }));
assert!(error.to_string().contains("invalid listener"));
}
#[tokio::test]
async fn correlates_early_exit_and_delivers_at_most_once() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
let host_id = request["params"]["hostId"].as_str().unwrap();
peer.exited("unrelated-host").await;
peer.exited(host_id).await;
peer.exited(host_id).await;
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_eq!(exit.host_id, host_id);
assert_eq!(exit.reason, AhpHostExitReason::Exited);
assert_eq!(exit.exit_code, Some(17));
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
peer.started(&request, None).await;
assert_eq!(pending.await.unwrap().unwrap().host_id, host_id);
}
#[tokio::test]
async fn concurrent_hosts_have_independent_callbacks_and_ids() {
let (client, mut peer) = fixture();
let (first_options, mut first_exits) = callback_options();
let first = start(&client, first_options);
let first_request = peer.request().await;
let (second_options, mut second_exits) = callback_options();
let second = start(&client, second_options);
let second_request = peer.request().await;
let first_id = first_request["params"]["hostId"].as_str().unwrap();
let second_id = second_request["params"]["hostId"].as_str().unwrap();
assert_ne!(first_id, second_id);
peer.started(&first_request, None).await;
peer.started(&second_request, None).await;
first.await.unwrap().unwrap();
second.await.unwrap().unwrap();
peer.exited(second_id).await;
assert_eq!(
timeout(TIMEOUT, second_exits.recv())
.await
.unwrap()
.unwrap()
.host_id,
second_id
);
assert!(first_exits.try_recv().is_err());
peer.exited(first_id).await;
assert_eq!(
timeout(TIMEOUT, first_exits.recv())
.await
.unwrap()
.unwrap()
.host_id,
first_id
);
}
#[tokio::test]
async fn failed_start_releases_callback_without_synthetic_exit() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
assert_eq!(client.inner.ahp_host_callbacks.lock().len(), 1);
peer.send(json!({
"jsonrpc": "2.0", "id": request["id"],
"error": {"code": -32000, "message": "could not start"}
}))
.await;
assert!(pending.await.unwrap().is_err());
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
peer.exited(request["params"]["hostId"].as_str().unwrap())
.await;
}
#[tokio::test]
async fn cancelled_start_disposes_after_late_success() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
pending.abort();
assert!(pending.await.unwrap_err().is_cancelled());
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
assert!(
timeout(Duration::from_millis(50), peer.read.read_u8())
.await
.is_err()
);
peer.started(&request, None).await;
let disposal = peer.request().await;
assert_eq!(disposal["method"], "host.dispose");
assert_eq!(disposal["params"]["hostId"], request["params"]["hostId"]);
peer.respond(&disposal, json!({})).await;
}
#[tokio::test]
async fn cancelled_start_does_not_dispose_after_late_failure() {
let (client, mut peer) = fixture();
let pending = start(&client, local_options());
let request = peer.request().await;
pending.abort();
assert!(pending.await.unwrap_err().is_cancelled());
peer.send(json!({
"jsonrpc": "2.0", "id": request["id"],
"error": {"code": -32000, "message": "could not start"}
}))
.await;
assert!(
timeout(Duration::from_millis(50), peer.read.read_u8())
.await
.is_err()
);
}
#[tokio::test]
async fn start_future_dropped_on_plain_thread_still_disposes() {
let (client, mut peer) = fixture();
let mut pending = Box::pin(client.start_ahp_host(local_options()));
assert!(futures_util::poll!(&mut pending).is_pending());
let request = peer.request().await;
std::thread::scope(|scope| scope.spawn(move || drop(pending)).join().unwrap());
peer.started(&request, None).await;
let disposal = peer.request().await;
assert_eq!(disposal["method"], "host.dispose");
assert_eq!(disposal["params"]["hostId"], request["params"]["hostId"]);
peer.respond(&disposal, json!({})).await;
}
#[tokio::test]
async fn cancelled_start_does_not_keep_owner_alive_while_waiting_for_response() {
let (client, mut peer) = fixture();
let owner = Arc::downgrade(&client.inner);
let pending = start(&client, local_options());
peer.request().await;
pending.abort();
assert!(pending.await.unwrap_err().is_cancelled());
drop(client);
assert!(owner.upgrade().is_none());
assert!(
timeout(TIMEOUT, peer.read.read_u8())
.await
.unwrap()
.is_err()
);
}
#[tokio::test]
async fn dispose_forwards_concurrent_repeated_calls_and_runtime_errors() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
let first = tokio::spawn({
let host = host.clone();
async move { host.dispose().await }
});
let second = tokio::spawn({
let host = host.clone();
async move { host.dispose().await }
});
let first_request = peer.request().await;
let second_request = peer.request().await;
assert_ne!(first_request["id"], second_request["id"]);
for request in [&first_request, &second_request] {
assert_eq!(request["method"], "host.dispose");
assert_eq!(request["params"], json!({"hostId": host.host_id}));
}
peer.respond(&second_request, json!({})).await;
peer.respond(&first_request, json!({})).await;
first.await.unwrap().unwrap();
second.await.unwrap().unwrap();
assert_eq!(client.inner.ahp_host_callbacks.lock().len(), 1);
assert!(
exits.try_recv().is_err(),
"dispose must not synthesize exit"
);
let third = tokio::spawn({
let host = host.clone();
async move { host.dispose().await }
});
let request = peer.request().await;
assert_eq!(request["method"], "host.dispose");
peer.send(json!({
"jsonrpc": "2.0", "id": request["id"],
"error": {"code": -32001, "message": "runtime disposal error"}
}))
.await;
assert!(matches!(
third.await.unwrap().unwrap_err().kind(),
ErrorKind::Rpc { code: -32001 }
));
peer.exited(&host.host_id).await;
timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
let fourth = tokio::spawn(async move { host.dispose().await });
let request = peer.request().await;
assert_eq!(request["method"], "host.dispose");
peer.respond(&request, json!({})).await;
fourth.await.unwrap().unwrap();
}
#[tokio::test]
async fn disconnect_notifies_once_without_claiming_reaping() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
drop(peer);
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_owner_disconnected(&exit, &host.host_id);
client.force_stop();
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
assert!(host.dispose().await.is_err());
}
#[tokio::test]
async fn force_stop_notifies_once_without_dispose_loop() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
client.force_stop();
client.force_stop();
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_owner_disconnected(&exit, &host.host_id);
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
assert_eq!(
timeout(TIMEOUT, peer.read.read_u8())
.await
.unwrap()
.unwrap_err()
.kind(),
std::io::ErrorKind::UnexpectedEof
);
}
fn assert_owner_disconnected(exit: &AhpHostExit, host_id: &str) {
assert_eq!(exit.host_id, host_id);
assert_eq!(exit.reason, AhpHostExitReason::OwnerDisconnected);
assert_eq!(exit.exit_code, None);
assert!(
exit.error
.as_deref()
.unwrap()
.contains("runtime cleanup cannot be acknowledged")
);
}
#[tokio::test]
async fn queued_real_exit_precedes_disconnect_and_is_not_duplicated() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
let (options, mut remaining_exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let remaining = pending.await.unwrap().unwrap();
for _ in 0..2 {
peer.exited(&host.host_id).await;
}
peer.write.shutdown().await.unwrap();
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_eq!(exit.host_id, host.host_id);
assert_eq!(exit.reason, AhpHostExitReason::Exited);
assert_eq!(exit.exit_code, Some(17));
assert_eq!(exit.error, None);
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
let exit = timeout(TIMEOUT, remaining_exits.recv())
.await
.unwrap()
.unwrap();
assert_owner_disconnected(&exit, &remaining.host_id);
assert!(
timeout(TIMEOUT, remaining_exits.recv())
.await
.unwrap()
.is_none()
);
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
}
#[tokio::test]
async fn wire_exit_immediately_before_eof_is_not_replaced_by_disconnect() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
peer.exited(&host.host_id).await;
drop(peer);
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_eq!(exit.reason, AhpHostExitReason::Exited);
assert_eq!(exit.exit_code, Some(17));
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
}
#[tokio::test]
async fn disconnect_notifies_all_callbacks_even_when_they_panic() {
let (client, mut peer) = fixture();
let (tx, mut exits) = mpsc::unbounded_channel();
let mut host_ids = Vec::new();
for _ in 0..3 {
let tx = tx.clone();
let pending = start(
&client,
local_options().with_on_exit(move |exit| {
tx.send(exit).unwrap();
panic!("test disconnect callback panic");
}),
);
let request = peer.request().await;
peer.started(&request, None).await;
host_ids.push(pending.await.unwrap().unwrap().host_id);
}
drop(tx);
client.force_stop();
for _ in 0..3 {
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
let index = host_ids.iter().position(|id| id == &exit.host_id).unwrap();
assert_owner_disconnected(&exit, &host_ids.remove(index));
}
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
}
#[tokio::test]
async fn callback_panic_does_not_break_other_callbacks_or_rpc() {
let (client, mut peer) = fixture();
let calls = Arc::new(AtomicUsize::new(0));
let pending = start(
&client,
local_options().with_on_exit({
let calls = calls.clone();
move |_| {
calls.fetch_add(1, Ordering::SeqCst);
panic!("test callback panic");
}
}),
);
let request = peer.request().await;
peer.started(&request, None).await;
let first = pending.await.unwrap().unwrap();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let second = pending.await.unwrap().unwrap();
peer.exited(&first.host_id).await;
peer.exited(&first.host_id).await;
peer.exited(&second.host_id).await;
timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
let dispose = tokio::spawn(async move { second.dispose().await });
let request = peer.request().await;
peer.respond(&request, json!({})).await;
dispose.await.unwrap().unwrap();
}
#[tokio::test]
async fn handle_and_dispatcher_do_not_keep_client_alive() {
let (client, mut peer) = fixture();
let client_weak = Arc::downgrade(&client.inner);
let callbacks_weak = Arc::downgrade(&client.inner.ahp_host_callbacks);
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
let host = pending.await.unwrap().unwrap();
drop(client);
assert!(client_weak.upgrade().is_none());
let exit = timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert_owner_disconnected(&exit, &host.host_id);
assert!(timeout(TIMEOUT, exits.recv()).await.unwrap().is_none());
assert!(callbacks_weak.upgrade().is_none());
assert!(host.dispose().await.is_err());
}
#[tokio::test]
async fn dropping_handle_does_not_dispose_or_unregister_callback() {
let (client, mut peer) = fixture();
let (options, mut exits) = callback_options();
let pending = start(&client, options);
let request = peer.request().await;
peer.started(&request, None).await;
drop(pending.await.unwrap().unwrap());
assert_eq!(client.inner.ahp_host_callbacks.lock().len(), 1);
assert!(
timeout(Duration::from_millis(50), peer.read.read_u8())
.await
.is_err()
);
peer.exited(request["params"]["hostId"].as_str().unwrap())
.await;
timeout(TIMEOUT, exits.recv()).await.unwrap().unwrap();
assert!(client.inner.ahp_host_callbacks.lock().is_empty());
}