use alloc::sync::Arc;
use core::{
fmt,
future::Future,
hash::{BuildHasher as _, Hash as _, Hasher as _},
pin::Pin,
sync::atomic::{AtomicU64, Ordering},
task::{Context, Poll},
};
use std::{
collections::HashMap,
hash::{DefaultHasher, RandomState},
io,
sync::{Mutex, MutexGuard, PoisonError},
time::SystemTime,
};
use serde_json::{Map, Value};
use tokio::{
io::{AsyncRead, AsyncWrite},
sync::{mpsc, oneshot, watch},
task::{AbortHandle, JoinHandle},
};
use crate::{
framing::{AsyncLineReader, DEFAULT_MAX_LINE, FrameReader, FrameWriter, encode_line},
protocol::{
CommandLifecycle, ControlCancelRequest, ControlRequest, ControlRequestBody,
ControlResponse, ControlResponseBody, Inbound, InterruptRequest, SetModelRequest,
SetPermissionModeRequest, StopTaskRequest, System, UserInput,
},
};
pub type ControlReply = Result<Option<Value>, String>;
pub type ControlFuture = Pin<Box<dyn Future<Output = ControlReply> + Send>>;
pub trait ControlHandler: Send + Sync + 'static {
fn handle(&self, request: ControlRequest) -> ControlFuture;
}
impl<F> ControlHandler for F
where
F: Fn(ControlRequest) -> ControlFuture + Send + Sync + 'static,
{
fn handle(&self, request: ControlRequest) -> ControlFuture {
self(request)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum SessionState {
Idle,
InTurn,
Closed,
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
#[error("writing to the CLI: {0}")]
Write(#[source] io::Error),
#[error("reading from the CLI: {0}")]
Read(#[source] io::Error),
#[error("the CLI's input is closed")]
InputClosed,
#[error("the session is closed")]
Closed,
#[error("the CLI's output ended mid-turn")]
EndedMidTurn,
#[error("control request {request_id} failed: {message}")]
Control { request_id: String, message: String },
#[error("serializing a frame: {0}")]
Serialize(#[from] serde_json::Error),
}
#[derive(Clone)]
pub struct SessionBuilder {
max_line: usize,
start_offset: u64,
replay_until: u64,
handler: Option<Arc<dyn ControlHandler>>,
answers_control: bool,
}
impl Default for SessionBuilder {
fn default() -> Self {
Self {
max_line: DEFAULT_MAX_LINE,
start_offset: 0,
replay_until: 0,
handler: None,
answers_control: false,
}
}
}
impl fmt::Debug for SessionBuilder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SessionBuilder")
.field("max_line", &self.max_line)
.field("start_offset", &self.start_offset)
.field("replay_until", &self.replay_until)
.field("handler", &self.handler.is_some())
.field("answers_control", &self.answers_control)
.finish()
}
}
impl SessionBuilder {
#[must_use]
pub fn max_line_len(mut self, max: usize) -> Self {
self.max_line = max;
self
}
#[must_use]
pub fn start_offset(mut self, offset: u64) -> Self {
self.start_offset = offset;
self
}
#[must_use]
pub fn replay_until(mut self, offset: u64) -> Self {
self.replay_until = offset;
self
}
#[must_use]
pub fn control_handler(mut self, handler: impl ControlHandler) -> Self {
self.handler = Some(Arc::new(handler));
self
}
#[must_use]
pub fn answers_control_requests(mut self, yes: bool) -> Self {
self.answers_control = yes;
self
}
pub fn start<R, W>(self, reader: R, writer: W) -> Session
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let (write_tx, write_rx) = mpsc::unbounded_channel();
let (out_tx, out_rx) = mpsc::unbounded_channel();
let (state, _) = watch::channel(SessionState::Idle);
let shared = Arc::new(Shared {
inner: Mutex::new(Inner::default()),
state,
writer: write_tx,
errors: out_tx.downgrade(),
next_id: AtomicU64::new(1),
nonce: format!("{:08x}", random_u64() & 0xffff_ffff),
replay_until: self.replay_until,
callbacks: self.handler.is_some() || self.answers_control,
handler: self.handler,
});
let frames = FrameReader::from_lines(
AsyncLineReader::new(reader)
.max_line_len(self.max_line)
.start_offset(self.start_offset),
);
let reader = tokio::spawn(read_loop(Arc::clone(&shared), frames, out_tx));
let writer = tokio::spawn(write_loop(
Arc::clone(&shared),
FrameWriter::new(writer),
write_rx,
));
Session {
handle: SessionHandle { shared },
rx: out_rx,
reader,
writer,
}
}
}
pub struct Session {
handle: SessionHandle,
rx: mpsc::UnboundedReceiver<Item>,
reader: JoinHandle<()>,
writer: JoinHandle<()>,
}
type Item = Result<(Inbound, u64), Error>;
impl fmt::Debug for Session {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Session")
.field("state", &self.state())
.finish_non_exhaustive()
}
}
impl Session {
pub fn new<R, W>(reader: R, writer: W) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
SessionBuilder::default().start(reader, writer)
}
pub fn builder() -> SessionBuilder {
SessionBuilder::default()
}
pub fn handle(&self) -> SessionHandle {
self.handle.clone()
}
pub async fn next(&mut self) -> Option<Result<Inbound, Error>> {
Some(self.rx.recv().await?.map(|(frame, _)| frame))
}
pub async fn next_at(&mut self) -> Option<Result<(Inbound, u64), Error>> {
self.rx.recv().await
}
pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Inbound, Error>>> {
self.rx
.poll_recv(cx)
.map(|item| item.map(|r| r.map(|(frame, _)| frame)))
}
pub async fn close(mut self) -> Result<(), Error> {
self.handle.end_input();
let mut err = None;
while let Some(item) = self.rx.recv().await {
if let Err(e) = item {
err.get_or_insert(e);
}
}
let _ = (&mut self.reader).await;
let _ = (&mut self.writer).await;
err.map_or(Ok(()), Err)
}
}
impl core::ops::Deref for Session {
type Target = SessionHandle;
fn deref(&self) -> &SessionHandle {
&self.handle
}
}
impl Drop for Session {
fn drop(&mut self) {
self.reader.abort();
self.writer.abort();
let shared = &self.handle.shared;
let mut inner = shared.lock();
inner.closed = true;
inner.input_closed = true;
inner.pending.clear();
for (_, task) in inner.incoming.drain() {
if let Some(task) = task {
task.abort();
}
}
shared.settle(&mut inner);
}
}
#[derive(Clone)]
pub struct SessionHandle {
shared: Arc<Shared>,
}
impl fmt::Debug for SessionHandle {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SessionHandle")
.field("state", &self.state())
.finish_non_exhaustive()
}
}
impl SessionHandle {
pub async fn send(&self, input: impl Into<UserInput>) -> Result<(), Error> {
let mut input = input.into();
let uuid = input.uuid.get_or_insert_with(new_uuid).clone();
let line = encode_line(&input)?;
let (ack, done) = oneshot::channel();
self.shared
.queue(&self.shared.lock(), line, Some(ack), Some(uuid))?;
wait_written(done).await
}
pub async fn send_raw<T: serde::Serialize + ?Sized>(&self, frame: &T) -> Result<(), Error> {
let line = encode_line(frame)?;
let (ack, done) = oneshot::channel();
self.shared
.queue(&self.shared.lock(), line, Some(ack), None)?;
wait_written(done).await
}
pub async fn request(
&self,
body: impl Into<ControlRequestBody>,
) -> Result<Option<Value>, Error> {
let n = self.shared.next_id.fetch_add(1, Ordering::Relaxed);
let request_id = format!("req_{n}_{}", self.shared.nonce);
let frame = Inbound::ControlRequest(ControlRequest::new(request_id.clone(), body));
let line = encode_line(&frame)?;
let (ack, written) = oneshot::channel();
let (answer_tx, answer) = oneshot::channel();
let _pending = {
let mut inner = self.shared.lock();
self.shared.queue(&inner, line, Some(ack), None)?;
inner.pending.insert(request_id.clone(), answer_tx);
self.shared.settle(&mut inner);
Pending {
shared: &self.shared,
id: &request_id,
}
};
wait_written(written).await?;
let failed = |message| Error::Control {
request_id: request_id.clone(),
message,
};
match answer.await.map_err(|_| Error::Closed)? {
ControlResponseBody::Success(s) => Ok(s.response.into_option()),
ControlResponseBody::Error(e) => Err(failed(e.error)),
other => Err(failed(format!(
"unexpected response subtype {:?}",
other.tag()
))),
}
}
pub async fn request_raw(
&self,
subtype: &str,
mut fields: Map<String, Value>,
) -> Result<Option<Value>, Error> {
fields.insert("subtype".to_owned(), Value::String(subtype.to_owned()));
self.request(ControlRequestBody::Other {
subtype: subtype.to_owned(),
raw: Value::Object(fields),
})
.await
}
pub async fn interrupt(&self) -> Result<(), Error> {
self.request(InterruptRequest::default()).await.map(drop)
}
pub async fn set_model(&self, model: Option<&str>) -> Result<(), Error> {
let req = SetModelRequest {
model: model.map(str::to_owned).into(),
..SetModelRequest::default()
};
self.request(req).await.map(drop)
}
pub async fn set_permission_mode(&self, mode: &str) -> Result<(), Error> {
let req = SetPermissionModeRequest {
mode: mode.to_owned(),
..SetPermissionModeRequest::default()
};
self.request(req).await.map(drop)
}
pub async fn stop_task(&self, task_id: &str) -> Result<(), Error> {
let req = StopTaskRequest {
task_id: task_id.to_owned(),
..StopTaskRequest::default()
};
self.request(req).await.map(drop)
}
pub async fn respond(&self, response: ControlResponse) -> Result<bool, Error> {
let Some(id) = response.request_id().map(str::to_owned) else {
return Err(Error::Control {
request_id: String::new(),
message: "a control response without a request_id".to_owned(),
});
};
let line = encode_line(&Inbound::ControlResponse(response))?;
let (ack, done) = oneshot::channel();
{
let mut inner = self.shared.lock();
if !inner.incoming.contains_key(&id) {
return Ok(false);
}
self.shared.queue(&inner, line, Some(ack), None)?;
inner.incoming.remove(&id);
self.shared.settle(&mut inner);
}
wait_written(done).await.map(|()| true)
}
pub fn finish(&self) {
let mut inner = self.shared.lock();
inner.finish = true;
if !self.shared.callbacks && inner.incoming.is_empty() {
self.shared.close_input(&mut inner);
}
self.shared.settle(&mut inner);
}
pub fn end_input(&self) {
let mut inner = self.shared.lock();
self.shared.close_input(&mut inner);
self.shared.settle(&mut inner);
}
pub fn state(&self) -> SessionState {
*self.shared.state.borrow()
}
pub fn watch_state(&self) -> watch::Receiver<SessionState> {
self.shared.state.subscribe()
}
pub fn is_idle(&self) -> bool {
self.shared.lock().idle()
}
pub fn session_id(&self) -> Option<String> {
self.shared.lock().session_id.clone()
}
}
async fn wait_written(done: oneshot::Receiver<io::Result<()>>) -> Result<(), Error> {
match done.await {
Ok(Ok(())) => Ok(()),
Ok(Err(e)) => Err(Error::Write(e)),
Err(_) => Err(Error::InputClosed),
}
}
struct Pending<'a> {
shared: &'a Shared,
id: &'a str,
}
impl Drop for Pending<'_> {
fn drop(&mut self) {
let mut inner = self.shared.lock();
if inner.pending.remove(self.id).is_some() {
let cancel = Inbound::ControlCancelRequest(ControlCancelRequest {
request_id: self.id.to_owned(),
..ControlCancelRequest::default()
});
if let Ok(line) = encode_line(&cancel) {
let _ = self.shared.queue(&inner, line, None, None);
}
}
self.shared.settle(&mut inner);
}
}
fn random_u64() -> u64 {
static COUNTER: AtomicU64 = AtomicU64::new(0);
let mut h: DefaultHasher = RandomState::new().build_hasher();
COUNTER.fetch_add(1, Ordering::Relaxed).hash(&mut h);
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or_default()
.hash(&mut h);
h.finish()
}
fn new_uuid() -> String {
let (a, b) = (random_u64(), random_u64());
let a = (a & !0xf000) | 0x4000;
let b = (b & !(0b11 << 62)) | (0b10 << 62);
format!(
"{:08x}-{:04x}-{:04x}-{:04x}-{:012x}",
a >> 32,
(a >> 16) & 0xffff,
a & 0xffff,
b >> 48,
b & 0xffff_ffff_ffff
)
}
enum WriteCmd {
Line {
line: Vec<u8>,
ack: Option<oneshot::Sender<io::Result<()>>>,
turn: Option<String>,
},
Close,
}
struct Shared {
inner: Mutex<Inner>,
state: watch::Sender<SessionState>,
writer: mpsc::UnboundedSender<WriteCmd>,
errors: mpsc::WeakUnboundedSender<Item>,
next_id: AtomicU64,
nonce: String,
replay_until: u64,
callbacks: bool,
handler: Option<Arc<dyn ControlHandler>>,
}
#[derive(Default)]
struct Inner {
pending: HashMap<String, oneshot::Sender<ControlResponseBody>>,
incoming: HashMap<String, Option<AbortHandle>>,
running: bool,
in_flight: Vec<(String, bool)>,
lifecycle: bool,
finish: bool,
input_closed: bool,
closed: bool,
session_id: Option<String>,
}
impl Inner {
fn in_turn(&self) -> bool {
self.running || !self.in_flight.is_empty()
}
fn idle(&self) -> bool {
!self.in_turn() && self.pending.is_empty() && self.incoming.is_empty()
}
fn end_turn(&mut self) {
self.running = false;
if self.lifecycle {
self.in_flight.retain(|&(_, started)| !started);
} else {
self.in_flight.clear();
}
}
fn lifecycle(&mut self, c: &CommandLifecycle) {
let (Some(uuid), Some(state)) = (c.command_uuid.as_deref(), c.state.as_deref()) else {
return;
};
self.lifecycle = true;
let at = self.in_flight.iter().position(|(u, _)| u == uuid);
match (state, at) {
("queued", None) => self.in_flight.push((uuid.to_owned(), false)),
("started", Some(i)) => self.in_flight[i].1 = true,
("started", None) => self.in_flight.push((uuid.to_owned(), true)),
("completed", Some(i)) => {
self.in_flight.remove(i);
}
_ => {}
}
}
}
impl Shared {
fn lock(&self) -> MutexGuard<'_, Inner> {
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
fn queue(
&self,
inner: &Inner,
line: Vec<u8>,
ack: Option<oneshot::Sender<io::Result<()>>>,
turn: Option<String>,
) -> Result<(), Error> {
if inner.input_closed {
return Err(Error::InputClosed);
}
self.writer
.send(WriteCmd::Line { line, ack, turn })
.map_err(|_| Error::InputClosed)
}
fn report(&self, e: Error) {
if let Some(out) = self.errors.upgrade() {
let _ = out.send(Err(e));
}
}
fn close_input(&self, inner: &mut Inner) {
if !inner.input_closed {
inner.input_closed = true;
let _ = self.writer.send(WriteCmd::Close);
}
}
fn settle(&self, inner: &mut Inner) {
let state = if inner.closed {
SessionState::Closed
} else if inner.in_turn() {
SessionState::InTurn
} else {
SessionState::Idle
};
self.state.send_if_modified(|s| {
let changed = *s != state;
*s = state;
changed
});
if inner.finish && inner.idle() {
self.close_input(inner);
}
}
}
async fn write_loop<W: AsyncWrite + Unpin>(
shared: Arc<Shared>,
mut w: FrameWriter<W>,
mut rx: mpsc::UnboundedReceiver<WriteCmd>,
) {
while let Some(cmd) = rx.recv().await {
let WriteCmd::Line { line, ack, turn } = cmd else {
break;
};
if let Some(uuid) = &turn {
let mut inner = shared.lock();
inner.in_flight.push((uuid.clone(), false));
shared.settle(&mut inner);
}
let r = w.write_line(&line).await;
let failed = r.is_err();
if let (true, Some(uuid)) = (failed, &turn) {
let mut inner = shared.lock();
inner.in_flight.retain(|(u, _)| u != uuid);
shared.settle(&mut inner);
}
match (ack, r) {
(Some(ack), r) => {
let _ = ack.send(r);
}
(None, Err(e)) => shared.report(Error::Write(e)),
(None, Ok(())) => {}
}
if failed {
break;
}
}
{
let mut inner = shared.lock();
inner.input_closed = true;
shared.settle(&mut inner);
}
rx.close();
if let Err(e) = w.shutdown().await {
shared.report(Error::Write(e));
}
}
async fn read_loop<R: AsyncRead + Unpin>(
shared: Arc<Shared>,
mut frames: FrameReader<R>,
out: mpsc::UnboundedSender<Item>,
) {
let end = loop {
match frames.next_at().await {
Ok(Some((frame, at))) => {
let frame = if at <= shared.replay_until {
replay(&shared, frame)
} else {
dispatch(&shared, frame)
};
if let Some(frame) = frame {
let _ = out.send(Ok((frame, at)));
}
}
Ok(None) => break None,
Err(e) => break Some(Error::Read(e)),
}
};
let mid_turn = {
let mut inner = shared.lock();
let mid_turn = inner.in_turn();
inner.closed = true;
inner.running = false;
inner.in_flight.clear();
inner.pending.clear();
for (_, task) in inner.incoming.drain() {
if let Some(task) = task {
task.abort();
}
}
shared.close_input(&mut inner);
shared.settle(&mut inner);
mid_turn
};
if let Some(e) = end.or_else(|| mid_turn.then_some(Error::EndedMidTurn)) {
let _ = out.send(Err(e));
}
}
fn dispatch(shared: &Arc<Shared>, frame: Inbound) -> Option<Inbound> {
match frame {
Inbound::ControlResponse(resp) => {
let mut inner = shared.lock();
let waiter = resp.request_id().and_then(|id| inner.pending.remove(id));
let Some(waiter) = waiter else {
return Some(Inbound::ControlResponse(resp));
};
let _ = waiter.send(resp.response);
shared.settle(&mut inner);
None
}
Inbound::ControlRequest(req) => {
if let Some(handler) = &shared.handler {
spawn_handler(shared, handler.as_ref(), req);
return None;
}
let mut inner = shared.lock();
inner.incoming.insert(req.request_id.clone(), None);
shared.settle(&mut inner);
Some(Inbound::ControlRequest(req))
}
Inbound::ControlCancelRequest(cancel) => {
let mut inner = shared.lock();
let open = inner.incoming.remove(&cancel.request_id);
shared.settle(&mut inner);
match open {
Some(Some(task)) => {
task.abort();
None
}
_ => Some(Inbound::ControlCancelRequest(cancel)),
}
}
frame => {
track(shared, &frame);
Some(frame)
}
}
}
fn replay(shared: &Shared, frame: Inbound) -> Option<Inbound> {
track(shared, &frame);
Some(frame)
}
fn track(shared: &Shared, frame: &Inbound) {
let mut inner = shared.lock();
match frame {
f if f.ends_turn() => inner.end_turn(),
Inbound::System(System::Init(init)) => {
inner.running = true;
if let Some(id) = init.session_id.get() {
inner.session_id = Some(id.clone());
}
}
Inbound::CommandLifecycle(c) => inner.lifecycle(c),
_ => return,
}
shared.settle(&mut inner);
}
fn spawn_handler(shared: &Arc<Shared>, handler: &dyn ControlHandler, req: ControlRequest) {
let id = req.request_id.clone();
{
let mut inner = shared.lock();
inner.incoming.insert(id.clone(), None);
shared.settle(&mut inner);
}
let answer = handler.handle(req);
let guard = Answer {
shared: Arc::clone(shared),
id: id.clone(),
done: false,
};
let task = tokio::spawn(async move {
let reply = answer.await;
guard.send(reply);
});
if let Some(slot) = shared.lock().incoming.get_mut(&id) {
*slot = Some(task.abort_handle());
}
}
struct Answer {
shared: Arc<Shared>,
id: String,
done: bool,
}
impl Answer {
fn send(mut self, reply: ControlReply) {
self.write(reply);
self.done = true;
}
fn write(&self, reply: ControlReply) {
let mut inner = self.shared.lock();
if inner.incoming.remove(&self.id).is_none() {
return;
}
let resp = match reply {
Ok(payload) => ControlResponse::success(&self.id, payload),
Err(message) => ControlResponse::error(&self.id, message),
};
match encode_line(&Inbound::ControlResponse(resp)) {
Ok(line) => {
if let Err(e) = self.shared.queue(&inner, line, None, None) {
self.shared.report(e);
}
}
Err(e) => self.shared.report(e.into()),
}
self.shared.settle(&mut inner);
}
}
impl Drop for Answer {
fn drop(&mut self) {
if !self.done {
self.write(Err("the control handler ended without an answer".to_owned()));
}
}
}