use super::config::Harness;
use super::exchange::Exchange;
use super::handshake::Session;
use super::host::{Running, observe};
use super::pump::{Publish, Pump};
use std::sync::Arc;
use tokio::sync::mpsc;
impl Running {
pub async fn session(
&mut self,
config: &Harness,
mut session: Session,
errors: &mut Vec<String>,
) -> anyhow::Result<i32> {
record(Arc::clone(&self.exchange), &session).await?;
self.listeners = Some(self.exchange.listen(config.prompt_queue).await?);
let (publications, receiver) = mpsc::channel(config.stream_queue);
let replayed = std::mem::take(&mut session.updates);
let replay_sent = replayed.into_iter().try_for_each(|item| {
publications
.try_send(Publish::Stream(item))
.map_err(|_| anyhow::anyhow!("handshake stream publication queue closed"))
});
let mut publisher = publish(Arc::clone(&self.exchange), receiver);
let mut publisher_consumed = false;
let pipe = self
.pipe
.as_mut()
.ok_or_else(|| anyhow::anyhow!("ACP pipe is absent"))?;
let listeners = self
.listeners
.as_mut()
.ok_or_else(|| anyhow::anyhow!("agent listeners are absent"))?;
let mut pump = Pump::new(
session,
pipe.output.clone(),
Arc::clone(&self.exchange),
config,
publications,
);
let finished: anyhow::Result<i32> = async {
replay_sent?;
let mut stdout_closed = false;
loop {
tokio::select! {
result = observe(&mut self.exited) => break result,
_ = self.term.recv() => break Ok(143),
_ = self.interrupt.recv() => break Ok(130),
result = &mut self.input => { result?; break Ok(0); },
frame = pipe.input.recv(), if !stdout_closed => match frame {
Some(Ok(Some(frame))) => pump.frame(frame)?,
Some(Ok(None)) => stdout_closed = true,
Some(Err(error)) => break Err(error),
None => break Err(anyhow::anyhow!("ACP reader stopped without an ending")),
},
event = listeners.events.recv(), if !stdout_closed => pump.delivery(event.ok_or_else(|| anyhow::anyhow!("host subscriptions stopped"))??)?,
result = pump.drafts.join_next(), if !pump.drafts.is_empty() => {
let (key, held) = result.ok_or_else(|| anyhow::anyhow!("draft worker disappeared"))??;
match held {
Ok(answer) => pump.drafted(&key, &answer)?,
Err(error) => pump.draft_failed(&key, &error)?,
}
},
result = &mut publisher => { publisher_consumed = true; result??; break Err(anyhow::anyhow!("host publisher stopped unexpectedly")); },
result = &mut pipe.writer => { self.writer_consumed = true; result??; break Err(anyhow::anyhow!("ACP writer stopped unexpectedly")); },
}
}
}.await;
if let Err(error) = pump.cancel() {
errors.push(error.to_string());
}
drain(&mut pump, errors).await;
drop(pump);
if !publisher_consumed {
match publisher.await {
Ok(Ok(())) => {}
Ok(Err(error)) => errors.push(error.to_string()),
Err(_) => errors.push("host publisher panicked".to_owned()),
}
}
finished
}
}
async fn record(exchange: Arc<Exchange>, session: &Session) -> anyhow::Result<()> {
let session_id = session.id.clone();
let version = session.version.number();
let resumable = session.resumable;
let answer = tokio::task::spawn_blocking(move || exchange.request("agents.session", &serde_json::json!({"agent":exchange.name,"session":session_id,"protocol_version":version,"load_session":resumable}))).await??;
anyhow::ensure!(
answer["recorded"] == true,
"agents.session did not confirm its durable record"
);
Ok(())
}
async fn drain(pump: &mut Pump, errors: &mut Vec<String>) {
pump.drafts.abort_all();
while let Some(result) = pump.drafts.join_next().await {
match result {
Ok((_, Ok(_))) => {}
Ok((_, Err(error))) => errors.push(error.to_string()),
Err(error) if error.is_cancelled() => {}
Err(_) => errors.push("draft worker panicked".to_owned()),
}
}
}
fn publish(
exchange: Arc<Exchange>,
mut receiver: mpsc::Receiver<Publish>,
) -> tokio::task::JoinHandle<anyhow::Result<()>> {
tokio::task::spawn_blocking(move || {
while let Some(publication) = receiver.blocking_recv() {
match publication {
Publish::Stream(item) => exchange.publish(&item)?,
Publish::Refuse(reason) => {
exchange.refuse_prompt(&reason)?;
exchange
.publish(&serde_json::json!({"type":"host_refusal","reason":reason}))?;
}
}
}
Ok(())
})
}