use super::{
config::Seat,
handshake::{Session, Version},
queue::{Prompt, Prompts},
wire::{Frame, Output, RpcError},
};
use serde_json::{Value, json};
use std::collections::{BTreeMap, BTreeSet};
use tokio::sync::mpsc;
#[path = "pump_config.rs"]
mod config;
#[path = "pump_permission.rs"]
mod permission;
#[path = "pump_sessions.rs"]
mod sessions;
#[path = "pump_turn.rs"]
mod turn;
struct Turn {
call: u64,
request: Value,
acknowledged: bool,
idle: Option<Value>,
cancelled: bool,
}
pub(super) struct Pump {
pub session: Session,
seat: Seat,
agent: String,
person: String,
output: Output,
publications: mpsc::Sender<Value>,
prompts: Prompts,
next_id: u64,
turn: Option<Turn>,
pub closing: bool,
pub approvals: super::approval_watch::ApprovalWatch,
known_drafts: std::collections::VecDeque<String>,
commands: BTreeSet<String>,
configuring: Option<(u64, Value)>,
opening: Option<sessions::Opening>,
pending: BTreeMap<String, crate::assistant::permission::Pending>,
answers: BTreeMap<String, (String, Value)>,
issuer: crate::assistant::permission::Issuer,
pub drafts: tokio::task::JoinSet<(String, anyhow::Result<String>)>,
}
impl Pump {
pub async fn new(
session: Session,
seat: Seat,
agent: String,
person: String,
output: Output,
publications: mpsc::Sender<Value>,
issuer: crate::assistant::permission::Issuer,
) -> anyhow::Result<Self> {
let approvals =
super::approval_watch::ApprovalWatch::start(issuer.clone(), seat.harness.stream_queue);
let mut pump = Self {
session,
agent,
person,
output,
publications,
prompts: Prompts::new(seat.harness.prompt_queue),
seat,
next_id: 3,
turn: None,
closing: false,
approvals,
known_drafts: std::collections::VecDeque::new(),
commands: BTreeSet::new(),
configuring: None,
opening: None,
pending: BTreeMap::new(),
answers: BTreeMap::new(),
issuer,
drafts: tokio::task::JoinSet::new(),
};
pump.bootstrap()?;
for (_, params) in std::mem::take(&mut pump.session.updates) {
pump.notification(params).await?;
}
Ok(pump)
}
async fn emit(&self, value: Value) -> anyhow::Result<()> {
self.publications
.send(value)
.await
.map_err(|_| anyhow::anyhow!("stream_queue_closed"))
}
async fn say(&self, item: Value) -> anyhow::Result<()> {
self.say_in(&self.session.id.clone(), item).await
}
async fn say_in(&self, conversation: &str, mut item: Value) -> anyhow::Result<()> {
item["conversation"] = json!(conversation);
item["seq"] = json!(u64::MAX);
item["durable"] = json!(false);
super::stream::fit(&mut item, self.seat.harness.content_bytes)?;
self.emit(json!({"item":item})).await
}
async fn ack(&self, id: &Value) -> anyhow::Result<()> {
self.emit(json!({"ack":{"id":id}})).await
}
async fn refuse(&self, id: &Value, code: &str) -> anyhow::Result<()> {
self.emit(json!({"refused":{"id":id,"code":code,"message":code.replace('_'," ")}}))
.await
}
fn call_id(&mut self) -> anyhow::Result<u64> {
let id = self.next_id;
self.next_id = id
.checked_add(1)
.ok_or_else(|| anyhow::anyhow!("acp_request_id_exhausted"))?;
Ok(id)
}
pub async fn operation(&mut self, value: Value) -> anyhow::Result<()> {
let op = value
.get("op")
.and_then(Value::as_object)
.ok_or_else(|| anyhow::anyhow!("door_op_invalid"))?;
let id = op
.get("id")
.filter(|id| id.is_string() || id.as_i64().is_some() || id.as_u64().is_some())
.ok_or_else(|| anyhow::anyhow!("door_op_id_invalid"))?;
if op
.get("person")
.and_then(Value::as_str)
.is_none_or(str::is_empty)
|| (op.contains_key("decided") && op["person"] == self.agent)
|| (!op.contains_key("decided") && op["person"] != self.person)
{
return self.refuse(id, "person_mismatch").await;
}
let names = ["text", "command", "config", "session", "cancel", "decided"];
let selected: Vec<_> = names.iter().filter(|key| op.contains_key(**key)).collect();
if selected.len() != 1
|| op.keys().any(|key| {
!names.contains(&key.as_str())
&& !["id", "person", "context"].contains(&key.as_str())
})
{
return self.refuse(id, "invalid_op").await;
}
let body = Value::Object(op.clone());
match *selected[0] {
"text" | "command" => self.prompt_op(id, &body).await,
"config" => self.configure(id, &body["config"]).await,
"session" => self.session_control(id, &body["session"]).await,
"cancel" if body["cancel"] == true => {
self.cancel()?;
self.ack(id).await
}
"decided" => {
self.decision(
id,
&body["decided"],
body["person"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("decision_person_missing"))?,
)
.await
}
_ => self.refuse(id, "invalid_op").await,
}
}
async fn prompt_op(&mut self, id: &Value, body: &Value) -> anyhow::Result<()> {
if self
.opening
.as_ref()
.is_some_and(sessions::Opening::changes_session)
{
return self.refuse(id, "session_busy").await;
}
let text = match self.prompt_text(body) {
Ok(text) => text,
Err(code) => return self.refuse(id, code).await,
};
if self.turn.is_some() {
if let Err(code) = self.prompts.push(Prompt {
id: id.clone(),
text: text.clone(),
}) {
return self.refuse(id, code).await;
}
self.say(json!({"kind":"request","request_id":id,"text":text}))
.await?;
self.say(json!({"kind":"prompt_queued","request_id":id,"count":self.prompts.len()}))
.await?;
} else {
self.start(
Prompt {
id: id.clone(),
text,
},
true,
)
.await?;
}
self.ack(id).await
}
fn prompt_text(&self, body: &Value) -> Result<String, &'static str> {
let mut text = if let Some(command) = body.get("command") {
let name = command["name"].as_str().ok_or("invalid_command")?;
if !self.commands.contains(name) {
return Err("command_not_offered");
}
let input = command["input"].as_str().ok_or("invalid_command")?;
if input.is_empty() {
format!("/{name}")
} else {
format!("/{name} {input}")
}
} else {
body["text"]
.as_str()
.filter(|text| !text.trim().is_empty())
.ok_or("invalid_text")?
.to_owned()
};
if let Some(context) = body.get("context") {
let object = context.as_object().ok_or("invalid_context")?;
if object.keys().any(|key| {
!["store", "branch", "version", "screen", "selection"].contains(&key.as_str())
}) {
return Err("invalid_context");
}
text = format!(
"Studio context: {}\n\n{text}",
serde_json::to_string(context).map_err(|_| "invalid_context")?
);
}
if text.len() > self.seat.harness.content_bytes {
return Err("content_bytes_exceeded");
}
Ok(text)
}
async fn start(&mut self, prompt: Prompt, publish_request: bool) -> anyhow::Result<()> {
let call = self.call_id()?;
self.output.request(
call,
"session/prompt",
json!({"sessionId":self.session.id,"prompt":[{"type":"text","text":prompt.text}]}),
)?;
if publish_request {
self.say(json!({"kind":"request","request_id":prompt.id,"text":prompt.text}))
.await?;
}
self.say(json!({"kind":"prompt_accepted","request_id":prompt.id}))
.await?;
self.turn = Some(Turn {
call,
request: prompt.id,
acknowledged: false,
idle: None,
cancelled: false,
});
Ok(())
}
pub async fn frame(&mut self, frame: Frame) -> anyhow::Result<()> {
match frame {
Frame::Notification { method, params } if method == "session/update" => {
self.notification(params).await
}
Frame::Notification { method, .. }
if method.starts_with('_')
|| method
== agent_client_protocol_schema::v1::PROTOCOL_LEVEL_METHOD_NAMES
.cancel_request =>
{
Ok(())
}
Frame::Notification { .. } => anyhow::bail!("acp_notification_not_advertised"),
Frame::Response { id, result }
if self
.opening
.as_ref()
.is_some_and(|opening| id == json!(opening.call)) =>
{
self.opened(result).await
}
Frame::Response { id, result }
if self
.configuring
.as_ref()
.is_some_and(|(call, _)| id == json!(call)) =>
{
self.configured(result).await
}
Frame::Response { id, result } => self.turn_response(&id, result).await,
Frame::Request { id, method, params } if method == "session/request_permission" => {
self.permission(id, ¶ms)
}
Frame::Request { id, .. } => self.output.refuse(
id,
-32601,
"no filesystem, terminal or credential callbacks",
),
}
}
async fn update(&mut self, params: Value, replayed: bool) -> anyhow::Result<()> {
let idle = params["update"]["sessionUpdate"] == "state_update"
&& params["update"]["state"] == "idle";
let stop = params["update"]["stopReason"].clone();
let advertised = params["update"]["availableCommands"]
.as_array()
.map(|commands| {
commands
.iter()
.filter_map(|command| command["name"].as_str().map(str::to_owned))
.collect::<BTreeSet<_>>()
});
let config = params["update"].get("configOptions").cloned();
if let Some(mut item) = super::stream::project(
self.session.version,
params,
&self.session.id,
self.seat.harness.content_bytes,
)? {
if item["kind"] == "available_commands" {
self.commands = advertised.unwrap_or_default();
}
if item["kind"] == "config_options"
&& let Some(config) = config
{
self.session.config = config;
}
if replayed {
item["replayed"] = json!(true);
}
self.say(item).await?;
}
if idle
&& !replayed
&& self.session.version == Version::Two
&& let Some(turn) = self.turn.as_mut()
{
turn.idle = Some(stop);
if turn.acknowledged {
self.end_turn().await?;
}
}
Ok(())
}
pub fn cancel(&mut self) -> anyhow::Result<()> {
self.cancel_pending()?;
if let Some(turn) = self.turn.as_mut() {
turn.cancelled = true;
self.output
.notification("session/cancel", json!({"sessionId":self.session.id}))?;
}
Ok(())
}
async fn finished(&mut self) -> anyhow::Result<()> {
self.turn = None;
self.cancel_pending()?;
if !self.closing
&& let Some(prompt) = self.prompts.pop()
{
self.start(prompt, false).await?;
}
Ok(())
}
}
impl Pump {
pub async fn closing_update(&mut self, params: Value) -> anyhow::Result<()> {
self.notification(params).await
}
pub async fn end(&mut self, status: i32) -> anyhow::Result<()> {
let mut errors = Vec::new();
if let Err(error) = self.approvals.close().await {
errors.push(error.to_string());
}
self.drafts.abort_all();
while let Some(result) = self.drafts.join_next().await {
if let Err(error) = result
&& !error.is_cancelled()
{
errors.push("permission_hold_worker_panicked".to_owned());
}
}
let status = if status == 0 && !errors.is_empty() {
1
} else {
status
};
if self.session.id.is_empty() {
anyhow::ensure!(errors.is_empty(), "{}", errors.join("; "));
return Ok(());
}
let published=self.say(json!({"kind":"ended","exit_status":status,"state":if self.session.resumable {"dormant"}else{"ended"}})).await;
if let Err(error) = published {
errors.push(error.to_string());
}
anyhow::ensure!(errors.is_empty(), "{}", errors.join("; "));
Ok(())
}
}