mod session;
mod translate;
use std::path::PathBuf;
use clap::Args;
use leviath_agent_client::{
AgentCapabilities, AgentInfo, ContentBlock, InitializeParams, InitializeResult, JsonRpcMessage,
PROTOCOL_VERSION, PromptCapabilities, RequestPermissionResult, SessionCancelParams,
SessionNewParams, SessionNewResult, SessionPromptParams, SessionPromptResult, SessionUpdate,
SessionUpdateParams, StopReason, error_codes, flatten_prompt, is_permission_request,
parse_region_markers, permission_request,
};
use leviath_core::interaction::{ApprovalScope, InteractionRequest, InteractionResponse};
use leviath_core::run_meta::RunStatus;
use leviath_runtime::control_socket::{ControlClient, ControlRequest, ControlResponse};
use leviath_runtime::host::WorldEvent;
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt};
use self::mapping::{PermissionChoice, interpret_permission};
use self::session::{ResolvedBlueprint, resolve_blueprint, spawn_args};
use self::translate::{StageTail, split_chunks};
const OUTPUT_POLL: std::time::Duration = std::time::Duration::from_millis(250);
#[derive(Args, Debug, Clone, Default)]
pub struct AgentClientArgs {
#[arg(long)]
pub agent: Option<String>,
#[arg(long)]
pub yolo: bool,
#[arg(long)]
pub allow: Vec<String>,
#[arg(long)]
pub max_depth: Option<usize>,
#[arg(long)]
pub no_seed_commands: bool,
#[arg(long, value_name = "LABEL")]
pub output_format: Option<String>,
#[arg(long, value_name = "TEXT")]
pub output_instructions: Option<String>,
}
pub async fn serve_over<R, W>(
reader: R,
writer: W,
control: ControlClient,
args: AgentClientArgs,
runs_dir: PathBuf,
default_cwd: String,
) -> anyhow::Result<()>
where
R: AsyncBufRead + Send + 'static,
W: AsyncWrite + Send + 'static,
{
let mut reader: BoxReader = Box::pin(reader);
let mut server = Server {
control,
args,
runs_dir,
writer: Box::pin(writer),
caps_present: false,
session: None,
next_request_id: 0,
io_alive: true,
default_cwd,
};
server.run(&mut reader).await;
Ok(())
}
type BoxReader = std::pin::Pin<Box<dyn AsyncBufRead + Send>>;
type BoxWriter = std::pin::Pin<Box<dyn AsyncWrite + Send>>;
struct ActiveSession {
session_id: String,
blueprint: ResolvedBlueprint,
cwd: String,
run_id: Option<String>,
}
struct Server {
control: ControlClient,
args: AgentClientArgs,
runs_dir: PathBuf,
writer: BoxWriter,
caps_present: bool,
session: Option<ActiveSession>,
next_request_id: i64,
default_cwd: String,
io_alive: bool,
}
enum RunStart {
Ready(String),
MessageUndeliverable,
SpawnFailed,
}
enum InteractionOutcome {
Continue,
Park(StopReason),
}
impl Server {
async fn run(&mut self, reader: &mut BoxReader) {
while self.io_alive {
let Some(line) = read_line(reader).await else {
break; };
let trimmed = line.trim();
if trimmed.is_empty() {
continue; }
match serde_json::from_str::<JsonRpcMessage>(trimmed) {
Ok(msg) => self.dispatch(reader, msg).await,
Err(_) => {
self.write(&JsonRpcMessage::error_response(
serde_json::Value::Null,
error_codes::PARSE_ERROR,
"invalid JSON",
))
.await;
}
}
}
}
async fn dispatch(&mut self, reader: &mut BoxReader, msg: JsonRpcMessage) {
match (msg.method.as_deref(), msg.id.clone()) {
(Some("initialize"), Some(id)) => self.on_initialize(id, msg.params).await,
(Some("session/new"), Some(id)) => self.on_session_new(id, msg.params).await,
(Some("session/prompt"), Some(id)) => {
self.on_session_prompt(reader, id, msg.params).await
}
(Some(_other), Some(id)) => {
self.write(&JsonRpcMessage::error_response(
id,
error_codes::METHOD_NOT_FOUND,
"method not supported",
))
.await
}
(Some("session/cancel"), None) => self.on_cancel_notification(msg.params).await,
_ => {}
}
}
async fn on_initialize(&mut self, id: serde_json::Value, params: Option<serde_json::Value>) {
let params: InitializeParams = params
.and_then(|p| serde_json::from_value(p).ok())
.unwrap_or_default();
self.caps_present = params.client_capabilities.is_some();
let result = InitializeResult {
protocol_version: PROTOCOL_VERSION,
agent_capabilities: AgentCapabilities {
load_session: false,
prompt_capabilities: PromptCapabilities {
image: false,
audio: false,
embedded_context: true,
},
},
agent_info: AgentInfo {
name: "leviath".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
},
auth_methods: vec![],
};
self.write(&JsonRpcMessage::response(id, &result)).await;
}
async fn on_session_new(&mut self, id: serde_json::Value, params: Option<serde_json::Value>) {
let params: SessionNewParams = params
.and_then(|p| serde_json::from_value(p).ok())
.unwrap_or_default();
if !params.mcp_servers.is_empty() {
let ignored = params.mcp_servers.len();
tracing::info!(
mcp_server_count = ignored,
"session/new supplied MCP servers; ignoring in favour of the blueprint's own"
);
}
let cwd = if params.cwd.trim().is_empty() {
self.default_cwd.clone()
} else {
params.cwd
};
match resolve_blueprint(self.args.agent.as_deref(), &cwd) {
Ok(blueprint) => {
let session_id = new_session_id(&blueprint.agent_name);
self.session = Some(ActiveSession {
session_id: session_id.clone(),
blueprint,
cwd,
run_id: None,
});
self.write(&JsonRpcMessage::response(
id,
&SessionNewResult { session_id },
))
.await;
}
Err(e) => {
self.write(&JsonRpcMessage::error_response(
id,
error_codes::INVALID_PARAMS,
format!("no blueprint for this session: {e}"),
))
.await;
}
}
}
async fn on_session_prompt(
&mut self,
reader: &mut BoxReader,
id: serde_json::Value,
params: Option<serde_json::Value>,
) {
if self.session.is_none() {
self.write(&JsonRpcMessage::error_response(
id,
error_codes::INVALID_REQUEST,
"no active session; call session/new first",
))
.await;
return;
}
let params: SessionPromptParams = params
.and_then(|p| serde_json::from_value(p).ok())
.unwrap_or_default();
let text = flatten_prompt(¶ms.prompt);
if text.is_empty() {
self.write(&JsonRpcMessage::error_response(
id,
error_codes::INVALID_PARAMS,
"prompt has no usable text content",
))
.await;
return;
}
let regions = parse_region_markers(&text);
let task = regions.get("task").cloned().unwrap_or_default();
let stop_reason = self.run_turn(reader, task, regions).await;
self.write(&JsonRpcMessage::response(
id,
&SessionPromptResult { stop_reason },
))
.await;
}
async fn on_cancel_notification(&mut self, params: Option<serde_json::Value>) {
let _: SessionCancelParams = params
.and_then(|p| serde_json::from_value(p).ok())
.unwrap_or_default();
if let Some(run_id) = self.session.as_ref().and_then(|s| s.run_id.clone()) {
let _ = self
.control
.request(&ControlRequest::Cancel { run_id })
.await;
}
}
async fn run_turn(
&mut self,
reader: &mut BoxReader,
task: String,
regions: std::collections::HashMap<String, String>,
) -> StopReason {
let Ok(mut stream) = self.control.subscribe().await else {
return StopReason::Refusal;
};
let session_id = self
.session
.as_ref()
.expect("session present")
.session_id
.clone();
let run_id = match self.start_run(task, regions).await {
RunStart::Ready(run_id) => run_id,
RunStart::MessageUndeliverable => return StopReason::EndTurn,
RunStart::SpawnFailed => return StopReason::Refusal,
};
let mut tail = StageTail::new();
while self.io_alive {
tokio::select! {
biased;
event = stream.next() => {
let Some(event) = event else {
self.flush_output(&session_id, &mut tail, &run_id).await;
return StopReason::EndTurn;
};
if event.run_id() != run_id {
continue; }
self.flush_output(&session_id, &mut tail, &run_id).await;
match event {
WorldEvent::Completed { status, final_output, .. } => {
self.emit_final_output(&session_id, final_output.as_ref()).await;
return leviath_agent_client::stop_reason_for_label(&status);
}
WorldEvent::Context { total_tokens, max_tokens, .. } => {
self.emit_usage(&session_id, total_tokens, max_tokens).await;
}
WorldEvent::Interaction { request, .. } => {
match self.handle_interaction(reader, &session_id, &run_id, request).await {
InteractionOutcome::Continue => {}
InteractionOutcome::Park(reason) => return reason,
}
}
_ => {}
}
if let Some(reason) = self.run_finished(&run_id) {
return reason;
}
}
_ = tokio::time::sleep(OUTPUT_POLL) => {
self.flush_output(&session_id, &mut tail, &run_id).await;
if let Some(reason) = self.run_finished(&run_id) {
return reason;
}
}
incoming = read_line(reader) => {
match incoming {
None => {
self.flush_output(&session_id, &mut tail, &run_id).await;
return StopReason::EndTurn;
}
Some(line) => self.handle_midturn_input(&run_id, &line).await,
}
}
}
}
StopReason::EndTurn
}
fn run_finished(&self, run_id: &str) -> Option<StopReason> {
let status = read_run_status(&self.runs_dir, run_id)?;
leviath_agent_client::stop_reason_for(&status)
}
async fn start_run(
&mut self,
task: String,
regions: std::collections::HashMap<String, String>,
) -> RunStart {
let existing = self
.session
.as_ref()
.expect("session present")
.run_id
.clone();
match existing {
Some(run_id) => {
let delivered = matches!(
self.control
.request(&ControlRequest::Message {
agent_id: run_id.clone(),
content: task,
target_region: None,
})
.await,
Ok(ControlResponse::Ok { ok: true })
);
if delivered {
RunStart::Ready(run_id)
} else {
RunStart::MessageUndeliverable
}
}
None => {
let session = self.session.as_ref().expect("session present");
let spawn =
spawn_args(&session.blueprint, &task, &session.cwd, &self.args, regions);
match self.control.spawn(spawn).await {
Ok(ControlResponse::Spawned { run_id }) => {
self.session.as_mut().expect("session present").run_id =
Some(run_id.clone());
RunStart::Ready(run_id)
}
_ => RunStart::SpawnFailed,
}
}
}
}
async fn handle_interaction(
&mut self,
reader: &mut BoxReader,
session_id: &str,
run_id: &str,
request: InteractionRequest,
) -> InteractionOutcome {
if self.caps_present {
if is_permission_request(&request) {
return self
.request_permission(reader, session_id, run_id, request)
.await;
}
self.emit_chunk(session_id, &format!("{}\n", request.prompt))
.await;
return InteractionOutcome::Park(StopReason::EndTurn);
}
self.emit_chunk(session_id, &format!("{}\n", request.prompt))
.await;
InteractionOutcome::Continue
}
async fn request_permission(
&mut self,
reader: &mut BoxReader,
session_id: &str,
run_id: &str,
request: InteractionRequest,
) -> InteractionOutcome {
let request_id = serde_json::json!(self.next_id());
let params = permission_request(session_id, &request);
self.write(&JsonRpcMessage::request(
request_id.clone(),
"session/request_permission",
¶ms,
))
.await;
loop {
let Some(line) = read_line(reader).await else {
return InteractionOutcome::Park(StopReason::EndTurn);
};
let Ok(msg) = serde_json::from_str::<JsonRpcMessage>(line.trim()) else {
continue;
};
if msg.method.as_deref() == Some("session/cancel") {
let _ = self
.control
.request(&ControlRequest::Cancel {
run_id: run_id.to_string(),
})
.await;
self.answer_interaction(&request.id, false, ApprovalScope::Once)
.await;
return InteractionOutcome::Continue;
}
if msg.id.as_ref() == Some(&request_id) {
let choice = msg
.result
.and_then(|r| serde_json::from_value::<RequestPermissionResult>(r).ok())
.map(|r| interpret_permission(&r.outcome))
.unwrap_or(PermissionChoice {
approved: false,
scope: ApprovalScope::Once,
});
self.answer_interaction(&request.id, choice.approved, choice.scope)
.await;
return InteractionOutcome::Continue;
}
}
}
async fn answer_interaction(&mut self, request_id: &str, approved: bool, scope: ApprovalScope) {
let response = InteractionResponse {
request_id: request_id.to_string(),
value: None,
choice_index: None,
approved: Some(approved),
scope: Some(scope),
};
let _ = self
.control
.request(&ControlRequest::AnswerInteraction { response })
.await;
}
async fn handle_midturn_input(&mut self, run_id: &str, line: &str) {
let Ok(msg) = serde_json::from_str::<JsonRpcMessage>(line.trim()) else {
return;
};
if msg.method.as_deref() == Some("session/cancel") {
let _ = self
.control
.request(&ControlRequest::Cancel {
run_id: run_id.to_string(),
})
.await;
}
}
async fn flush_output(&mut self, session_id: &str, tail: &mut StageTail, run_id: &str) {
let text = tail.pump(&self.runs_dir, run_id);
for chunk in split_chunks(&text) {
self.emit_chunk(session_id, chunk).await;
}
}
async fn emit_chunk(&mut self, session_id: &str, text: &str) {
let params = SessionUpdateParams {
session_id: session_id.to_string(),
update: SessionUpdate::AgentMessageChunk {
content: ContentBlock::text(text),
},
};
self.write(&JsonRpcMessage::notification("session/update", ¶ms))
.await;
}
async fn emit_final_output(
&mut self,
session_id: &str,
output: Option<&leviath_core::output::FinalOutput>,
) {
let Some(output) = output else { return };
let shape = output
.format
.as_deref()
.map(|f| format!(" ({f})"))
.unwrap_or_default();
let text = format!("\n\n--- final output{shape} ---\n{}", output.content);
for chunk in split_chunks(&text) {
self.emit_chunk(session_id, chunk).await;
}
}
async fn emit_usage(&mut self, session_id: &str, used: usize, size: usize) {
let params = SessionUpdateParams {
session_id: session_id.to_string(),
update: SessionUpdate::UsageUpdate { used, size },
};
self.write(&JsonRpcMessage::notification("session/update", ¶ms))
.await;
}
async fn write(&mut self, msg: &JsonRpcMessage) {
let mut line = serde_json::to_string(msg).expect("JsonRpcMessage always serializes");
line.push('\n');
let ok = self.writer.write_all(line.as_bytes()).await.is_ok()
&& self.writer.flush().await.is_ok();
if !ok {
self.io_alive = false;
}
}
fn next_id(&mut self) -> i64 {
self.next_request_id += 1;
self.next_request_id
}
}
fn read_run_status(runs_dir: &std::path::Path, run_id: &str) -> Option<RunStatus> {
#[derive(serde::Deserialize)]
struct StatusOnly {
status: RunStatus,
}
let path = runs_dir.join(run_id).join("meta.json");
let json = std::fs::read_to_string(path).ok()?;
serde_json::from_str::<StatusOnly>(&json)
.ok()
.map(|s| s.status)
}
fn new_session_id(agent_name: &str) -> String {
crate::runstate::new_run_id(agent_name)
}
async fn read_line(reader: &mut BoxReader) -> Option<String> {
let mut line = String::new();
match reader.read_line(&mut line).await.unwrap_or(0) {
0 => None,
_ => Some(line),
}
}
mod mapping {
use leviath_agent_client::PermissionOutcome;
use leviath_agent_client::mapping::{
OPTION_ALLOW_ALWAYS, OPTION_ALLOW_ONCE, OPTION_REJECT_ONCE,
};
use leviath_core::interaction::ApprovalScope;
pub(super) struct PermissionChoice {
pub(super) approved: bool,
pub(super) scope: ApprovalScope,
}
pub(super) fn interpret_permission(outcome: &PermissionOutcome) -> PermissionChoice {
match outcome {
PermissionOutcome::Selected { option_id } if option_id == OPTION_ALLOW_ONCE => {
PermissionChoice {
approved: true,
scope: ApprovalScope::Once,
}
}
PermissionOutcome::Selected { option_id } if option_id == OPTION_ALLOW_ALWAYS => {
PermissionChoice {
approved: true,
scope: ApprovalScope::Run,
}
}
PermissionOutcome::Selected { option_id } if option_id == OPTION_REJECT_ONCE => {
PermissionChoice {
approved: false,
scope: ApprovalScope::Once,
}
}
_ => PermissionChoice {
approved: false,
scope: ApprovalScope::Once,
},
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn interpret_maps_every_offered_option() {
let allow_once = interpret_permission(&PermissionOutcome::Selected {
option_id: OPTION_ALLOW_ONCE.to_string(),
});
assert!(allow_once.approved);
assert_eq!(allow_once.scope, ApprovalScope::Once);
let allow_always = interpret_permission(&PermissionOutcome::Selected {
option_id: OPTION_ALLOW_ALWAYS.to_string(),
});
assert!(allow_always.approved);
assert_eq!(allow_always.scope, ApprovalScope::Run);
let reject = interpret_permission(&PermissionOutcome::Selected {
option_id: OPTION_REJECT_ONCE.to_string(),
});
assert!(!reject.approved);
assert_eq!(reject.scope, ApprovalScope::Once);
}
#[test]
fn interpret_denies_unknown_option_and_cancellation() {
let unknown = interpret_permission(&PermissionOutcome::Selected {
option_id: "made-up".to_string(),
});
assert!(!unknown.approved);
assert_eq!(unknown.scope, ApprovalScope::Once);
let cancelled = interpret_permission(&PermissionOutcome::Cancelled);
assert!(!cancelled.approved);
assert_eq!(cancelled.scope, ApprovalScope::Once);
}
}
}
#[cfg(test)]
mod tests;