use std::collections::BTreeMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use serde_json::json;
use supercode::server::{run_http, RpcEngine};
use supercode::{
AcpFrontendCheckpoint, AcpFrontendConnectOptions, AcpFrontendRuntime, Agent, ApprovalPolicy,
ChatMessage, ChatRequest, Config, FrontendApprovalDecision, FrontendRequest, FrontendResponse,
FrontendRuntime, FrontendRuntimeMetadata, FunctionCall, HttpFrontendRuntime, Provider, Role,
RuntimeLaunch, ToolCall, Usage,
};
fn bin() -> &'static str {
env!("CARGO_BIN_EXE_supercode")
}
struct ApprovalStreamingProvider {
calls: AtomicUsize,
streaming: Arc<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
}
#[async_trait]
impl Provider for ApprovalStreamingProvider {
async fn complete(
&self,
request: &ChatRequest,
on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
) -> supercode::Result<(ChatMessage, Usage)> {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
return Ok((
ChatMessage {
role: Role::Assistant,
content: None,
content_parts: None,
tool_calls: Some(vec![ToolCall {
id: "shared-approval-call".into(),
kind: "function".into(),
function: FunctionCall {
name: "bash".into(),
arguments: json!({"command":"printf tool-approved"}).to_string(),
},
}]),
tool_call_id: None,
name: None,
metadata: Default::default(),
},
Usage::default(),
));
}
let tool_result = request
.messages
.last()
.and_then(|message| message.content.as_deref())
.unwrap_or_default();
assert!(tool_result.contains("tool-approved"), "{tool_result}");
assert!(!tool_result.contains("was not approved"), "{tool_result}");
on_delta("stream-before-restart ");
self.streaming.notify_one();
self.release.notified().await;
on_delta("stream-after-restart");
Ok((
ChatMessage::assistant("stream-before-restart stream-after-restart"),
Usage::default(),
))
}
}
fn new_runtime(
cwd: &std::path::Path,
name: &str,
) -> (
Arc<RpcEngine>,
Arc<tokio::sync::Notify>,
Arc<tokio::sync::Notify>,
) {
let streaming = Arc::new(tokio::sync::Notify::new());
let release = Arc::new(tokio::sync::Notify::new());
let mut config = Config::builder()
.cwd(cwd.to_path_buf())
.approval(ApprovalPolicy::OnRequest)
.build();
config.permissions_enabled = true;
let agent = Agent::with_provider(
config,
Box::new(ApprovalStreamingProvider {
calls: AtomicUsize::new(0),
streaming: streaming.clone(),
release: release.clone(),
}),
);
(
RpcEngine::new_named_with_frontend_requests(
agent,
name,
FrontendRuntimeMetadata::default(),
None,
),
streaming,
release,
)
}
async fn approve_next_request(
attachment: &mut supercode::FrontendAttachment,
runtime: &dyn FrontendRuntime,
) {
loop {
let event = attachment.next_event().await.unwrap();
if event.kind == "request" {
let request: FrontendRequest =
serde_json::from_value(event.payload["request"].clone()).unwrap();
runtime
.respond(FrontendResponse::Approval {
request_id: request.id,
decision: FrontendApprovalDecision::Allow,
})
.await
.unwrap();
return;
}
}
}
async fn drive(runtime: Arc<dyn FrontendRuntime>, release: Arc<tokio::sync::Notify>) {
let mut attachment = runtime.attach(100).await.unwrap();
let submit_runtime = runtime.clone();
let submit =
tokio::spawn(async move { submit_runtime.submit("run approved tool".into()).await });
approve_next_request(&mut attachment, runtime.as_ref()).await;
release.notify_one();
assert_eq!(
submit.await.unwrap().unwrap(),
"stream-before-restart stream-after-restart"
);
}
async fn spawn_acp_bridge(
url: &str,
token: &str,
checkpoint: Option<AcpFrontendCheckpoint>,
) -> Result<Arc<AcpFrontendRuntime>> {
let mut env = BTreeMap::new();
env.insert("SUPERCODE_SERVER_TOKEN".into(), token.into());
let runtime = AcpFrontendRuntime::connect(AcpFrontendConnectOptions {
launch: RuntimeLaunch {
program: bin().into(),
arguments: vec!["acp".into(), "--connect".into(), url.into()],
env,
},
cwd: None,
session_id: checkpoint.as_ref().map(|value| value.session_id.clone()),
after_sequence: checkpoint.as_ref().map(|value| value.acknowledged_sequence),
})
.await?;
if let Some(checkpoint) = checkpoint {
runtime.restore_checkpoint(checkpoint)?;
}
Ok(runtime)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 6)]
async fn headless_acp_restart_matches_local_and_http_transcript_byte_for_byte() -> Result<()> {
let cwd = std::env::temp_dir().join(format!("supercode-sup48-proof-{}", std::process::id()));
std::fs::create_dir_all(&cwd)?;
let (local_engine, _, local_release) = new_runtime(&cwd, "proof-runtime");
drive(local_engine.clone(), local_release).await;
let local_bytes = serde_json::to_vec(&local_engine.history(100).await)?;
let (http_engine, _, http_release) = new_runtime(&cwd, "proof-runtime");
let http_token: Arc<str> = "sup48-http-token".into();
let http_address = run_http(http_engine.clone(), "127.0.0.1:0", http_token.clone()).await?;
let http =
HttpFrontendRuntime::connect(format!("http://{http_address}"), http_token.to_string())
.await?;
drive(http, http_release).await;
let http_bytes = serde_json::to_vec(&http_engine.history(100).await)?;
let (acp_engine, streaming, release) = new_runtime(&cwd, "proof-runtime");
let acp_token: Arc<str> = "sup48-acp-token".into();
let acp_address = run_http(acp_engine.clone(), "127.0.0.1:0", acp_token.clone()).await?;
let acp_url = format!("http://{acp_address}");
let first = spawn_acp_bridge(&acp_url, acp_token.as_ref(), None).await?;
let mut first_attachment = first.attach(100).await?;
let submit_runtime = first.clone();
let submit =
tokio::spawn(async move { submit_runtime.submit("run approved tool".into()).await });
approve_next_request(&mut first_attachment, first.as_ref()).await;
tokio::time::timeout(Duration::from_secs(3), streaming.notified())
.await
.expect("ACP turn never began streaming");
loop {
if first_attachment.next_event().await?.kind == "text_delta" {
break;
}
}
let checkpoint = first.checkpoint();
first.detach().await?;
submit.abort();
drop(first_attachment);
drop(first);
assert!(!acp_engine.is_shutting_down());
let second = spawn_acp_bridge(&acp_url, acp_token.as_ref(), Some(checkpoint.clone())).await?;
let mut second_attachment = second.attach(100).await?;
release.notify_one();
let mut resumed = Vec::new();
loop {
let event = second_attachment.next_event().await?;
resumed.push(event.sequence);
if event.kind == "turn_succeeded" {
break;
}
}
assert_eq!(
resumed.first(),
Some(&(checkpoint.acknowledged_sequence + 1))
);
assert!(
resumed.windows(2).all(|pair| pair[1] > pair[0]),
"resumed ACP sequence was not strictly increasing: {resumed:?}"
);
assert!(!acp_engine.is_shutting_down());
let acp_bytes = serde_json::to_vec(&acp_engine.history(100).await)?;
second.detach().await?;
drop(second_attachment);
drop(second);
assert_eq!(http_bytes, local_bytes);
assert_eq!(acp_bytes, local_bytes);
local_engine.shutdown().await;
http_engine.shutdown().await;
acp_engine.shutdown().await;
std::fs::remove_dir_all(cwd).ok();
Ok(())
}