use std::collections::{BTreeMap, BTreeSet};
use std::pin::Pin;
use std::sync::atomic::Ordering;
use anyhow::Result;
use futures::Stream;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use agentos_protocol::generated::v1::{
AcpCloseSessionRequest, AcpCreateSessionRequest, AcpGetSessionStateRequest,
AcpListAgentsRequest, AcpRequest, AcpResponse, AcpResumeSessionRequest, AcpRuntimeKind,
AcpSessionCreatedResponse, AcpSessionRequest, AcpSessionStateResponse,
};
use agentos_protocol::ACP_EXTENSION_NAMESPACE;
use agentos_sidecar_client::wire;
use crate::agent_os::{AgentOs, SessionEntry};
use crate::config::ToolKit;
use crate::error::ClientError;
use crate::json_rpc::{JsonRpcError, JsonRpcId, JsonRpcNotification, JsonRpcResponse};
use crate::stream::Subscription;
use crate::{CLOSED_SESSION_ID_RETENTION_LIMIT, PERMISSION_TIMEOUT_MS};
const LEGACY_PERMISSION_METHOD: &str = "request/permission";
pub(crate) const ACP_PERMISSION_METHOD: &str = "session/request_permission";
const SESSION_PENDING_REQUEST_LIMIT: usize = 1024;
pub(crate) struct PermissionRouteRequest {
pub(crate) session_id: String,
pub(crate) permission_id: String,
pub(crate) params: Value,
}
pub(crate) struct PermissionRouteResult {
pub(crate) reply: Option<String>,
}
struct SessionCreatedResponse {
session_id: String,
modes: Option<Value>,
config_options: Vec<Value>,
agent_capabilities: Option<Value>,
agent_info: Option<Value>,
}
pub(crate) struct SessionStateResponse {
modes: Option<Value>,
config_options: Vec<Value>,
agent_capabilities: Option<Value>,
agent_info: Option<Value>,
}
const PROMPT_TEXT_CAPTURE_LIMIT_BYTES: usize = 16 * 1024 * 1024;
const PROMPT_DELIVERED_CHUNK_LIMIT: usize = 262_144;
pub type SessionEventStream = Pin<Box<dyn Stream<Item = JsonRpcNotification> + Send>>;
pub type SessionEventSubscription = (SessionEventStream, Subscription);
pub type PermissionRequestStream = Pin<Box<dyn Stream<Item = PermissionRequest> + Send>>;
pub type PermissionRequestSubscription = (PermissionRequestStream, Subscription);
pub type AgentExitStream = Pin<Box<dyn Stream<Item = AgentExitEvent> + Send>>;
pub type AgentExitSubscription = (AgentExitStream, Subscription);
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentExitEvent {
#[serde(rename = "sessionId")]
pub session_id: String,
#[serde(rename = "agentType")]
pub agent_type: String,
#[serde(rename = "processId")]
pub process_id: String,
#[serde(rename = "exitCode")]
pub exit_code: Option<i32>,
pub restart: String,
#[serde(rename = "restartCount")]
pub restart_count: u32,
#[serde(rename = "maxRestarts")]
pub max_restarts: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionInfo {
#[serde(rename = "sessionId")]
pub session_id: String,
#[serde(rename = "agentType")]
pub agent_type: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentRegistryEntry {
pub id: String,
pub installed: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum McpServerConfig {
Local {
command: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
args: Vec<String>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
env: BTreeMap<String, String>,
},
Remote {
url: String,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
headers: BTreeMap<String, String>,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct CreateSessionOptions {
pub cwd: Option<String>,
pub env: BTreeMap<String, String>,
pub mcp_servers: Vec<McpServerConfig>,
pub skip_os_instructions: bool,
pub additional_instructions: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionId {
#[serde(rename = "sessionId")]
pub session_id: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResumeSessionResult {
#[serde(rename = "sessionId")]
pub session_id: String,
pub mode: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ResumeSessionOptions {
pub transcript_path: Option<String>,
pub cwd: Option<String>,
pub env: BTreeMap<String, String>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PromptResult {
pub response: JsonRpcResponse,
pub text: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SessionMode {
pub id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub label: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(flatten)]
pub extra: BTreeMap<String, Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct SessionModeState {
#[serde(default, rename = "currentModeId")]
pub current_mode_id: String,
#[serde(default, rename = "availableModes")]
pub available_modes: Vec<SessionMode>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ConfigAllowedValue {
pub id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub label: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionConfigOption {
#[serde(default)]
pub id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub category: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub label: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(
default,
rename = "currentValue",
skip_serializing_if = "Option::is_none"
)]
pub current_value: Option<String>,
#[serde(
default,
rename = "allowedValues",
skip_serializing_if = "Option::is_none"
)]
pub allowed_values: Option<Vec<ConfigAllowedValue>>,
#[serde(default, rename = "readOnly", skip_serializing_if = "Option::is_none")]
pub read_only: Option<bool>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct PromptCapabilities {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub audio: Option<bool>,
#[serde(
default,
rename = "embeddedContext",
skip_serializing_if = "Option::is_none"
)]
pub embedded_context: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub image: Option<bool>,
#[serde(flatten)]
pub extra: BTreeMap<String, Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct AgentCapabilities {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub permissions: Option<bool>,
#[serde(default, rename = "plan_mode", skip_serializing_if = "Option::is_none")]
pub plan_mode: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub questions: Option<bool>,
#[serde(
default,
rename = "tool_calls",
skip_serializing_if = "Option::is_none"
)]
pub tool_calls: Option<bool>,
#[serde(
default,
rename = "text_messages",
skip_serializing_if = "Option::is_none"
)]
pub text_messages: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub images: Option<bool>,
#[serde(
default,
rename = "file_attachments",
skip_serializing_if = "Option::is_none"
)]
pub file_attachments: Option<bool>,
#[serde(
default,
rename = "session_lifecycle",
skip_serializing_if = "Option::is_none"
)]
pub session_lifecycle: Option<bool>,
#[serde(
default,
rename = "error_events",
skip_serializing_if = "Option::is_none"
)]
pub error_events: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub status: Option<bool>,
#[serde(
default,
rename = "streaming_deltas",
skip_serializing_if = "Option::is_none"
)]
pub streaming_deltas: Option<bool>,
#[serde(default, rename = "mcp_tools", skip_serializing_if = "Option::is_none")]
pub mcp_tools: Option<bool>,
#[serde(
default,
rename = "promptCapabilities",
skip_serializing_if = "Option::is_none"
)]
pub prompt_capabilities: Option<PromptCapabilities>,
#[serde(flatten)]
pub extra: BTreeMap<String, Value>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AgentInfo {
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub version: Option<String>,
#[serde(flatten)]
pub extra: BTreeMap<String, Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct SessionInitData {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub modes: Option<SessionModeState>,
#[serde(
default,
rename = "configOptions",
skip_serializing_if = "Option::is_none"
)]
pub config_options: Option<Vec<SessionConfigOption>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub capabilities: Option<AgentCapabilities>,
#[serde(default, rename = "agentInfo", skip_serializing_if = "Option::is_none")]
pub agent_info: Option<AgentInfo>,
}
#[derive(Clone)]
pub struct PermissionResponder {
inner:
std::sync::Arc<parking_lot::Mutex<Option<tokio::sync::oneshot::Sender<PermissionReply>>>>,
}
impl PermissionResponder {
pub fn new() -> (Self, tokio::sync::oneshot::Receiver<PermissionReply>) {
let (tx, rx) = tokio::sync::oneshot::channel();
(
Self {
inner: std::sync::Arc::new(parking_lot::Mutex::new(Some(tx))),
},
rx,
)
}
pub fn respond(&self, reply: PermissionReply) {
if let Some(tx) = self.inner.lock().take() {
let _ = tx.send(reply);
}
}
}
#[derive(Clone)]
pub struct PermissionRequest {
pub permission_id: String,
pub description: Option<String>,
pub params: Value,
pub responder: PermissionResponder,
}
impl std::fmt::Debug for PermissionRequest {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PermissionRequest")
.field("permission_id", &self.permission_id)
.field("description", &self.description)
.field("params", &self.params)
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PermissionReply {
Once,
Always,
Reject,
}
fn permission_reply_wire(reply: PermissionReply) -> &'static str {
match reply {
PermissionReply::Once => "once",
PermissionReply::Always => "always",
PermissionReply::Reject => "reject",
}
}
fn agent_capabilities_is_empty(caps: &AgentCapabilities) -> bool {
caps.permissions.is_none()
&& caps.plan_mode.is_none()
&& caps.questions.is_none()
&& caps.tool_calls.is_none()
&& caps.text_messages.is_none()
&& caps.images.is_none()
&& caps.file_attachments.is_none()
&& caps.session_lifecycle.is_none()
&& caps.error_events.is_none()
&& caps.reasoning.is_none()
&& caps.status.is_none()
&& caps.streaming_deltas.is_none()
&& caps.mcp_tools.is_none()
&& caps.prompt_capabilities.is_none()
&& caps.extra.is_empty()
}
fn should_dispatch_to_session_event_handlers(notification: &JsonRpcNotification) -> bool {
notification.method == "session/update"
}
pub(crate) fn record_live_session_event(entry: &SessionEntry, notification: JsonRpcNotification) {
apply_session_update(entry, ¬ification);
if should_dispatch_to_session_event_handlers(¬ification) {
let _ = entry.event_tx.send(notification);
}
}
fn apply_session_update(entry: &SessionEntry, notification: &JsonRpcNotification) {
if notification.method != "session/update" {
return;
}
let Some(params) = notification.params.as_ref().and_then(Value::as_object) else {
return;
};
let update = params
.get("update")
.and_then(Value::as_object)
.unwrap_or(params);
match update.get("sessionUpdate").and_then(Value::as_str) {
Some("current_mode_update") => {
let Some(mode_id) = update.get("currentModeId").and_then(Value::as_str) else {
return;
};
let mut modes = entry.modes.lock();
if let Some(modes) = modes.as_mut() {
modes.current_mode_id = mode_id.to_string();
}
}
Some("config_option_update") | Some("config_options_update") => {
let Some(options) = update.get("configOptions").and_then(Value::as_array) else {
return;
};
let parsed = options
.iter()
.filter_map(|value| serde_json::from_value(value.clone()).ok())
.collect();
*entry.config_options.lock() = parsed;
apply_synthetic_config_overrides(entry);
}
Some("agent_message_chunk") | None | Some(_) => {}
}
}
fn accumulate_agent_message_chunk(
notification: &JsonRpcNotification,
delivered_chunks: &mut usize,
agent_text: &mut String,
) -> std::result::Result<(), ClientError> {
let params = notification.params.clone().unwrap_or(Value::Null);
let update = params.get("update").cloned().unwrap_or(Value::Null);
if update.get("sessionUpdate").and_then(Value::as_str) != Some("agent_message_chunk") {
return Ok(());
}
if let Some(chunk) = update
.get("content")
.and_then(|content| content.get("text"))
.and_then(Value::as_str)
{
if *delivered_chunks >= PROMPT_DELIVERED_CHUNK_LIMIT {
return Err(prompt_chunk_limit_error());
}
let next_len = agent_text
.len()
.checked_add(chunk.len())
.ok_or_else(|| prompt_text_limit_error(usize::MAX))?;
if next_len > PROMPT_TEXT_CAPTURE_LIMIT_BYTES {
return Err(prompt_text_limit_error(next_len));
}
agent_text.push_str(chunk);
*delivered_chunks += 1;
}
Ok(())
}
fn pending_session_request_count(entry: &SessionEntry) -> usize {
let mut count = 0;
entry.pending_prompt_resolvers.scan(|_, _| {
count += 1;
});
count
}
fn prompt_text_limit_error(size: usize) -> ClientError {
ClientError::Sidecar(format!(
"prompt text capture is {size} bytes, limit is {PROMPT_TEXT_CAPTURE_LIMIT_BYTES}"
))
}
fn prompt_chunk_limit_error() -> ClientError {
ClientError::Sidecar(format!(
"prompt chunk tracking limit exceeded: at most {PROMPT_DELIVERED_CHUNK_LIMIT} chunks can be captured per prompt"
))
}
struct PendingSessionRequestGuard<'a> {
os: &'a AgentOs,
session_id: &'a str,
resolver_id: i64,
active: bool,
}
impl<'a> PendingSessionRequestGuard<'a> {
fn new(os: &'a AgentOs, session_id: &'a str, resolver_id: i64) -> Self {
Self {
os,
session_id,
resolver_id,
active: true,
}
}
fn cleanup(&mut self) {
if self.active {
self.os
.cleanup_pending_resolver(self.session_id, self.resolver_id);
self.active = false;
}
}
}
impl Drop for PendingSessionRequestGuard<'_> {
fn drop(&mut self) {
self.cleanup();
}
}
fn apply_synthetic_config_overrides(entry: &SessionEntry) {
let overrides = entry.config_overrides.lock().clone();
if overrides.is_empty() {
return;
}
let mut options = entry.config_options.lock();
for option in options.iter_mut() {
let override_value = overrides
.get(&option.id)
.filter(|_| !option.id.starts_with(PENDING_METHOD_PREFIX))
.cloned()
.or_else(|| {
option
.category
.as_ref()
.and_then(|category| overrides.get(category).cloned())
});
if let Some(value) = override_value {
option.current_value = Some(value);
}
}
}
const PENDING_METHOD_PREFIX: &str = "__pending_method::";
fn sync_session_state(entry: &SessionEntry, state: &SessionStateResponse) {
*entry.modes.lock() = state
.modes
.as_ref()
.filter(|value| value.is_object())
.and_then(|value| serde_json::from_value(value.clone()).ok());
*entry.config_options.lock() = state
.config_options
.iter()
.filter_map(|value| serde_json::from_value(value.clone()).ok())
.collect();
apply_synthetic_config_overrides(entry);
*entry.capabilities.lock() = state
.agent_capabilities
.as_ref()
.filter(|value| value.is_object())
.and_then(|value| serde_json::from_value(value.clone()).ok());
*entry.agent_info.lock() = state
.agent_info
.as_ref()
.filter(|value| value.is_object())
.and_then(|value| serde_json::from_value(value.clone()).ok());
}
fn unsupported_config_response(agent_type: &str, category: &str) -> JsonRpcResponse {
let message = if agent_type == "opencode" && category == "model" {
"OpenCode reports available models, but model switching must be configured before createSession() because ACP session/set_config_option is not implemented.".to_string()
} else {
format!("The {category} config option is read-only for {agent_type} sessions.")
};
JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id: Some(JsonRpcId::Null),
result: None,
error: Some(JsonRpcError {
code: -32601,
message,
data: None,
}),
}
}
fn session_closed_response(session_id: &str) -> JsonRpcResponse {
JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id: Some(JsonRpcId::Null),
result: None,
error: Some(JsonRpcError {
code: -32000,
message: format!("Session closed: {session_id}"),
data: None,
}),
}
}
fn session_created_from_acp(
response: AcpSessionCreatedResponse,
) -> std::result::Result<SessionCreatedResponse, ClientError> {
Ok(SessionCreatedResponse {
session_id: response.session_id,
modes: parse_optional_json(response.modes, "modes")?,
config_options: parse_json_vec(response.config_options, "configOptions")?,
agent_capabilities: parse_optional_json(response.agent_capabilities, "agentCapabilities")?,
agent_info: parse_optional_json(response.agent_info, "agentInfo")?,
})
}
fn session_state_from_acp(
response: AcpSessionStateResponse,
) -> std::result::Result<SessionStateResponse, ClientError> {
Ok(SessionStateResponse {
modes: parse_optional_json(response.modes, "modes")?,
config_options: parse_json_vec(response.config_options, "configOptions")?,
agent_capabilities: parse_optional_json(response.agent_capabilities, "agentCapabilities")?,
agent_info: parse_optional_json(response.agent_info, "agentInfo")?,
})
}
fn parse_optional_json(
value: Option<String>,
label: &str,
) -> std::result::Result<Option<Value>, ClientError> {
value
.map(|value| {
serde_json::from_str(&value).map_err(|error| {
ClientError::Sidecar(format!("malformed ACP {label} JSON: {error}"))
})
})
.transpose()
}
fn parse_json_vec(
values: Vec<String>,
label: &str,
) -> std::result::Result<Vec<Value>, ClientError> {
values
.into_iter()
.map(|value| {
serde_json::from_str(&value).map_err(|error| {
ClientError::Sidecar(format!("malformed ACP {label} JSON: {error}"))
})
})
.collect()
}
fn unexpected_acp_response(operation: &str, response: AcpResponse) -> ClientError {
ClientError::Sidecar(format!("unexpected response to {operation}: {response:?}"))
}
fn combine_instructions(additional: Option<&str>, tool_reference: &str) -> Option<String> {
let mut parts = Vec::new();
if let Some(additional) = additional.map(str::trim).filter(|value| !value.is_empty()) {
parts.push(additional.to_string());
}
let tool_reference = tool_reference.trim();
if !tool_reference.is_empty() {
parts.push(tool_reference.to_string());
}
if parts.is_empty() {
None
} else {
Some(parts.join("\n\n"))
}
}
fn build_host_tool_reference(tool_kits: &[ToolKit]) -> String {
if tool_kits.is_empty() {
return String::new();
}
let mut lines = vec![
String::from("## Available Host Tools"),
String::new(),
String::from("Run `agentos list-tools` to see all available tools."),
String::new(),
];
for kit in tool_kits {
lines.push(format!("### {}", kit.name));
lines.push(String::new());
lines.push(kit.description.clone());
lines.push(String::new());
for tool in &kit.tools {
let signature = build_tool_flag_signature(&tool.input_schema);
let suffix = if signature.is_empty() {
String::new()
} else {
format!(" {signature}")
};
lines.push(format!(
"- `agentos-{} {}{}` — {}",
kit.name, tool.name, suffix, tool.description
));
}
lines.push(String::new());
lines.push(format!(
"Run `agentos-{} <tool> --help` for details.",
kit.name
));
lines.push(String::new());
}
lines.join("\n")
}
fn build_tool_flag_signature(schema: &Value) -> String {
describe_tool_flags(schema)
.into_iter()
.map(|flag| {
if flag.required {
format!("{} <{}>", flag.name, flag.value_type)
} else {
format!("[{} <{}>]", flag.name, flag.value_type)
}
})
.collect::<Vec<_>>()
.join(" ")
}
struct ToolFlagDescription {
name: String,
value_type: String,
required: bool,
}
fn describe_tool_flags(schema: &Value) -> Vec<ToolFlagDescription> {
let properties = schema
.get("properties")
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let required = schema
.get("required")
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(Value::as_str)
.map(str::to_owned)
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
properties
.into_iter()
.map(|(field_name, field_schema)| ToolFlagDescription {
name: format!("--{}", camel_to_kebab(&field_name)),
value_type: describe_tool_flag_type(&field_schema),
required: required.contains(&field_name),
})
.collect()
}
fn describe_tool_flag_type(schema: &Value) -> String {
match json_schema_type(schema) {
Some("array") => {
let item_type = schema
.get("items")
.and_then(json_schema_type)
.unwrap_or("string");
format!("{item_type}[]")
}
Some("string") => schema
.get("enum")
.and_then(Value::as_array)
.map(|values| values.iter().filter_map(Value::as_str).collect::<Vec<_>>())
.filter(|values| !values.is_empty())
.map(|values| values.join("|"))
.unwrap_or_else(|| String::from("string")),
Some(other) => other.to_string(),
None => String::from("string"),
}
}
fn json_schema_type(schema: &Value) -> Option<&str> {
schema.get("type").and_then(Value::as_str)
}
fn camel_to_kebab(value: &str) -> String {
let mut output = String::new();
for (index, ch) in value.chars().enumerate() {
if ch.is_ascii_uppercase() && index > 0 {
output.push('-');
}
output.push(ch.to_ascii_lowercase());
}
output
}
impl AgentOs {
fn session_ownership(&self) -> wire::OwnershipScope {
wire::OwnershipScope::VmOwnership(wire::VmOwnership {
connection_id: self.connection_id().to_string(),
session_id: self.wire_session_id().to_string(),
vm_id: self.vm_id().to_string(),
})
}
fn require_session<R>(
&self,
session_id: &str,
f: impl FnOnce(&SessionEntry) -> R,
) -> std::result::Result<R, ClientError> {
self.inner()
.sessions
.read(session_id, |_, entry| f(entry))
.ok_or_else(|| ClientError::SessionNotFound(session_id.to_string()))
}
async fn hydrate_session_state(
&self,
session_id: &str,
) -> std::result::Result<(), ClientError> {
self.require_session(session_id, |_| ())?;
let response = self
.send_acp_request(AcpRequest::AcpGetSessionStateRequest(
AcpGetSessionStateRequest {
session_id: session_id.to_string(),
},
))
.await?;
let AcpResponse::AcpSessionStateResponse(state) = response else {
return Err(unexpected_acp_response(
"AcpGetSessionStateRequest",
response,
));
};
let state = session_state_from_acp(state)?;
self.require_session(session_id, |entry| sync_session_state(entry, &state))?;
Ok(())
}
pub(crate) async fn send_session_request(
&self,
session_id: &str,
method: &str,
params: Option<Value>,
) -> std::result::Result<JsonRpcResponse, ClientError> {
let request_params = params;
let resolver_id = self.inner().request_counter.fetch_add(1, Ordering::SeqCst);
let (resolve_tx, resolve_rx) = tokio::sync::oneshot::channel::<JsonRpcResponse>();
self.require_session(session_id, |entry| {
let _guard = entry.pending_session_request_lock.lock();
if pending_session_request_count(entry) >= SESSION_PENDING_REQUEST_LIMIT {
return Err(ClientError::Sidecar(format!(
"session pending request limit exceeded: at most {SESSION_PENDING_REQUEST_LIMIT} requests can be in flight per session"
)));
}
let _ = entry
.pending_prompt_resolvers
.insert(resolver_id, resolve_tx);
entry
.config_overrides
.lock()
.entry(format!("{PENDING_METHOD_PREFIX}{resolver_id}"))
.or_insert_with(|| method.to_string());
Ok(())
})??;
let mut pending_request_guard =
PendingSessionRequestGuard::new(self, session_id, resolver_id);
let rpc = self.send_acp_request(AcpRequest::AcpSessionRequest(AcpSessionRequest {
session_id: session_id.to_string(),
method: method.to_string(),
params: request_params
.clone()
.map(|params| serde_json::to_string(¶ms))
.transpose()
.map_err(|error| {
ClientError::Sidecar(format!("failed to encode session params: {error}"))
})?,
}));
tokio::pin!(rpc);
let response = tokio::select! {
biased;
resolved = resolve_rx => {
pending_request_guard.cleanup();
match resolved {
Ok(response) => return Ok(response),
Err(_) => return Ok(session_closed_response(session_id)),
}
}
result = &mut rpc => {
pending_request_guard.cleanup();
result?
}
};
let response = match response {
AcpResponse::AcpSessionRpcResponse(rpc) => {
serde_json::from_str::<JsonRpcResponse>(&rpc.response).map_err(|err| {
ClientError::Sidecar(format!("malformed session rpc response: {err}"))
})?
}
other => return Err(unexpected_acp_response("AcpSessionRequest", other)),
};
let _ = self.hydrate_session_state(session_id).await;
if response.error.is_none() {
self.apply_post_send_cache_updates(session_id, method, request_params.as_ref())?;
}
Ok(response)
}
fn cleanup_pending_resolver(&self, session_id: &str, resolver_id: i64) {
let _ = self.require_session(session_id, |entry| {
let _ = entry.pending_prompt_resolvers.remove(&resolver_id);
entry
.config_overrides
.lock()
.remove(&format!("{PENDING_METHOD_PREFIX}{resolver_id}"));
});
}
fn apply_post_send_cache_updates(
&self,
session_id: &str,
method: &str,
params: Option<&Value>,
) -> std::result::Result<(), ClientError> {
self.require_session(session_id, |entry| {
if method == "session/set_mode" {
if let Some(mode_id) = params.and_then(|p| p.get("modeId")).and_then(Value::as_str)
{
let mut modes = entry.modes.lock();
if let Some(modes) = modes.as_mut() {
modes.current_mode_id = mode_id.to_string();
}
}
}
if method == "session/set_config_option" {
let config_id = params
.and_then(|p| p.get("configId"))
.and_then(Value::as_str);
let value = params.and_then(|p| p.get("value")).and_then(Value::as_str);
if let (Some(config_id), Some(value)) = (config_id, value) {
let mut options = entry.config_options.lock();
for option in options.iter_mut() {
if option.id == config_id {
option.current_value = Some(value.to_string());
}
}
}
}
})
}
async fn set_session_config_by_category(
&self,
session_id: &str,
category: &str,
value: &str,
) -> std::result::Result<JsonRpcResponse, ClientError> {
let (read_only, config_id, agent_type) = self.require_session(session_id, |entry| {
let options = entry.config_options.lock();
let option = options
.iter()
.find(|option| option.category.as_deref() == Some(category));
(
option.and_then(|option| option.read_only).unwrap_or(false),
option.map(|option| option.id.clone()),
entry.agent_type.clone(),
)
})?;
if read_only {
return Ok(unsupported_config_response(&agent_type, category));
}
let config_id = config_id.unwrap_or_else(|| category.to_string());
let response = self
.send_session_request(
session_id,
"session/set_config_option",
Some(json!({ "configId": config_id, "value": value })),
)
.await?;
Ok(response)
}
pub fn list_sessions(&self) -> Vec<SessionInfo> {
let mut sessions = Vec::new();
self.inner().sessions.scan(|session_id, entry| {
sessions.push(SessionInfo {
session_id: session_id.clone(),
agent_type: entry.agent_type.clone(),
});
});
sessions
}
pub async fn list_agents(&self) -> Result<Vec<AgentRegistryEntry>> {
let response = self
.send_acp_request(AcpRequest::AcpListAgentsRequest(AcpListAgentsRequest {
reserved: false,
}))
.await?;
let AcpResponse::AcpListAgentsResponse(listed) = response else {
return Err(unexpected_acp_response("AcpListAgentsRequest", response).into());
};
Ok(listed
.agents
.into_iter()
.map(|agent| AgentRegistryEntry {
id: agent.id,
installed: agent.installed,
})
.collect())
}
pub async fn create_session(
&self,
agent_type: &str,
options: CreateSessionOptions,
) -> Result<SessionId> {
let env: BTreeMap<String, String> = options.env.clone();
let cwd = options
.cwd
.clone()
.unwrap_or_else(|| "/workspace".to_string());
let mcp_servers: Vec<Value> = options
.mcp_servers
.iter()
.filter_map(|server| serde_json::to_value(server).ok())
.collect();
let client_capabilities = json!({
"fs": { "readTextFile": true, "writeTextFile": true },
"terminal": true,
});
let tool_reference = build_host_tool_reference(&self.config().tool_kits);
let additional_instructions =
combine_instructions(options.additional_instructions.as_deref(), &tool_reference);
let response = self
.send_acp_request(AcpRequest::AcpCreateSessionRequest(
AcpCreateSessionRequest {
agent_type: agent_type.to_string(),
runtime: AcpRuntimeKind::JavaScript,
args: Vec::new(),
env: env.into_iter().collect(),
cwd,
mcp_servers: serde_json::to_string(&mcp_servers).map_err(|error| {
ClientError::Sidecar(format!("failed to encode MCP servers: {error}"))
})?,
protocol_version: crate::ACP_PROTOCOL_VERSION as i32,
client_capabilities: serde_json::to_string(&client_capabilities).map_err(
|error| {
ClientError::Sidecar(format!(
"failed to encode client capabilities: {error}"
))
},
)?,
additional_instructions,
skip_os_instructions: options.skip_os_instructions,
},
))
.await?;
let AcpResponse::AcpSessionCreatedResponse(created) = response else {
return Err(unexpected_acp_response("AcpCreateSessionRequest", response).into());
};
let created = session_created_from_acp(created)?;
let state = SessionStateResponse {
modes: created.modes,
config_options: created.config_options,
agent_capabilities: created.agent_capabilities,
agent_info: created.agent_info,
};
self.register_session(&created.session_id, agent_type, &state)
.await?;
Ok(SessionId {
session_id: created.session_id,
})
}
pub(crate) async fn register_session(
&self,
session_id: &str,
agent_type: &str,
state: &SessionStateResponse,
) -> std::result::Result<(), ClientError> {
{
let mut closed = self.inner().closed_session_ids.lock();
closed.retain(|id| id != session_id);
}
let (event_tx, _) = tokio::sync::broadcast::channel(1024);
let (permission_tx, _) = tokio::sync::broadcast::channel(64);
let (agent_exit_tx, _) = tokio::sync::broadcast::channel(16);
let entry = SessionEntry {
agent_type: agent_type.to_string(),
modes: parking_lot::Mutex::new(None),
config_options: parking_lot::Mutex::new(Vec::new()),
capabilities: parking_lot::Mutex::new(None),
agent_info: parking_lot::Mutex::new(None),
config_overrides: parking_lot::Mutex::new(BTreeMap::new()),
event_tx,
permission_tx,
agent_exit_tx,
pending_permission_replies: scc::HashMap::new(),
pending_session_request_lock: parking_lot::Mutex::new(()),
pending_prompt_resolvers: scc::HashMap::new(),
};
sync_session_state(&entry, state);
let _ = self.inner().sessions.insert(session_id.to_string(), entry);
match self.hydrate_session_state(session_id).await {
Ok(()) => Ok(()),
Err(error) => {
let _ = self.inner().sessions.remove(session_id);
Err(error)
}
}
}
pub async fn resume_session(
&self,
session_id: &str,
agent_type: &str,
options: ResumeSessionOptions,
) -> Result<ResumeSessionResult> {
let env: BTreeMap<String, String> = options.env.clone();
let cwd = options
.cwd
.clone()
.unwrap_or_else(|| "/workspace".to_string());
let response = self
.send_acp_request(AcpRequest::AcpResumeSessionRequest(
AcpResumeSessionRequest {
session_id: session_id.to_string(),
agent_type: agent_type.to_string(),
transcript_path: options.transcript_path.clone(),
cwd,
env: env.into_iter().collect(),
},
))
.await?;
let AcpResponse::AcpSessionResumedResponse(resumed) = response else {
return Err(unexpected_acp_response("AcpResumeSessionRequest", response).into());
};
let empty_state = SessionStateResponse {
modes: None,
config_options: Vec::new(),
agent_capabilities: None,
agent_info: None,
};
self.register_session(&resumed.session_id, agent_type, &empty_state)
.await?;
Ok(ResumeSessionResult {
session_id: resumed.session_id,
mode: resumed.mode,
})
}
pub async fn destroy_session(&self, session_id: &str) -> Result<()> {
self.require_session(session_id, |_| ())?;
let _ = self.cancel_session(session_id).await;
self.close_session_internal(session_id).await?;
Ok(())
}
pub async fn prompt(&self, session_id: &str, text: &str) -> Result<PromptResult> {
let mut rx = self.require_session(session_id, |entry| entry.event_tx.subscribe())?;
let mut agent_text = String::new();
let mut delivered_chunks = 0;
let mut prompt_text_error: Option<ClientError> = None;
let request = self.send_session_request(
session_id,
"session/prompt",
Some(json!({ "prompt": [{ "type": "text", "text": text }] })),
);
tokio::pin!(request);
let response = loop {
tokio::select! {
biased;
result = &mut request => break result,
event = rx.recv() => {
match event {
Ok(event) => accumulate_agent_message_chunk(
&event,
&mut delivered_chunks,
&mut agent_text,
)
.unwrap_or_else(|error| {
prompt_text_error.get_or_insert(error);
}),
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
break (&mut request).await;
}
}
}
}
};
loop {
match rx.try_recv() {
Ok(event) => {
accumulate_agent_message_chunk(&event, &mut delivered_chunks, &mut agent_text)
.unwrap_or_else(|error| {
prompt_text_error.get_or_insert(error);
})
}
Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::TryRecvError::Empty)
| Err(tokio::sync::broadcast::error::TryRecvError::Closed) => break,
}
}
drop(rx);
let response = response?;
if let Some(error) = prompt_text_error {
return Err(error.into());
}
Ok(PromptResult {
response,
text: agent_text,
})
}
pub async fn cancel_session(&self, session_id: &str) -> Result<JsonRpcResponse> {
self.require_session(session_id, |_| ())?;
let cancelled_pending_prompt = self.cancel_pending_prompt_requests(session_id)?;
if cancelled_pending_prompt {
let this = self.clone();
let session_id_owned = session_id.to_string();
tokio::spawn(async move {
let _ = this
.send_session_request(&session_id_owned, "session/cancel", None)
.await;
});
return Ok(JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id: Some(JsonRpcId::Null),
result: Some(json!({
"cancelled": true,
"requested": true,
"via": "prompt-fallback",
})),
error: None,
});
}
Ok(self
.send_session_request(session_id, "session/cancel", None)
.await?)
}
fn cancel_pending_prompt_requests(
&self,
session_id: &str,
) -> std::result::Result<bool, ClientError> {
self.require_session(session_id, |entry| {
let mut prompt_resolver_ids = Vec::new();
{
let overrides = entry.config_overrides.lock();
for (key, method) in overrides.iter() {
if let Some(id) = key.strip_prefix(PENDING_METHOD_PREFIX) {
if method == "session/prompt" {
if let Ok(id) = id.parse::<i64>() {
prompt_resolver_ids.push(id);
}
}
}
}
}
let mut cancelled = false;
for id in prompt_resolver_ids {
if let Some((_, resolver)) = entry.pending_prompt_resolvers.remove(&id) {
let _ = resolver.send(JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id: Some(JsonRpcId::Null),
result: Some(json!({ "stopReason": "cancelled" })),
error: None,
});
cancelled = true;
}
entry
.config_overrides
.lock()
.remove(&format!("{PENDING_METHOD_PREFIX}{id}"));
}
cancelled
})
}
fn abort_pending_session_requests(&self, session_id: &str) {
let _ = self.require_session(session_id, |entry| {
let mut ids = Vec::new();
entry.pending_prompt_resolvers.scan(|id, _| ids.push(*id));
for id in ids {
if let Some((_, resolver)) = entry.pending_prompt_resolvers.remove(&id) {
let _ = resolver.send(session_closed_response(session_id));
}
entry
.config_overrides
.lock()
.remove(&format!("{PENDING_METHOD_PREFIX}{id}"));
}
});
}
fn reject_pending_permission_replies(&self, session_id: &str) {
let _ = self.require_session(session_id, |entry| {
let mut ids = Vec::new();
entry
.pending_permission_replies
.scan(|id, _| ids.push(id.clone()));
for id in ids {
let _ = entry.pending_permission_replies.remove(&id);
}
});
}
pub fn close_session(&self, session_id: &str) -> std::result::Result<(), ClientError> {
let known = self.inner().sessions.contains(session_id)
|| self.inner().closing_session_ids.contains(session_id)
|| self
.inner()
.closed_session_ids
.lock()
.iter()
.any(|id| id == session_id);
if !known {
return Err(ClientError::SessionNotFound(session_id.to_string()));
}
let _ = self
.inner()
.closing_session_ids
.insert(session_id.to_string());
let this = self.clone();
let session_id_owned = session_id.to_string();
tokio::spawn(async move {
let _ = this.close_session_internal(&session_id_owned).await;
let _ = this.inner().closing_session_ids.remove(&session_id_owned);
});
Ok(())
}
pub(crate) async fn close_session_internal(
&self,
session_id: &str,
) -> std::result::Result<(), ClientError> {
if self
.inner()
.closed_session_ids
.lock()
.iter()
.any(|id| id == session_id)
{
return Ok(());
}
self.abort_pending_session_requests(session_id);
self.reject_pending_permission_replies(session_id);
if !self.inner().sessions.contains(session_id) {
return Err(ClientError::SessionNotFound(session_id.to_string()));
}
let _ = self.inner().sessions.remove(session_id);
{
let mut closed = self.inner().closed_session_ids.lock();
closed.push_back(session_id.to_string());
while closed.len() > CLOSED_SESSION_ID_RETENTION_LIMIT {
closed.pop_front();
}
}
let response = self
.send_acp_request(AcpRequest::AcpCloseSessionRequest(AcpCloseSessionRequest {
session_id: session_id.to_string(),
}))
.await?;
match response {
AcpResponse::AcpSessionClosedResponse(_) => Ok(()),
other => Err(unexpected_acp_response("AcpCloseSessionRequest", other)),
}
}
async fn send_acp_request(
&self,
request: AcpRequest,
) -> std::result::Result<AcpResponse, ClientError> {
let payload = serde_bare::to_vec(&request).map_err(|error| {
ClientError::Sidecar(format!("failed to encode ACP request: {error}"))
})?;
let response = self
.transport()
.request_wire(
self.session_ownership(),
wire::RequestPayload::ExtEnvelope(wire::ExtEnvelope {
namespace: ACP_EXTENSION_NAMESPACE.to_string(),
payload,
}),
)
.await?;
let envelope = match response {
wire::ResponsePayload::ExtEnvelope(envelope) => envelope,
wire::ResponsePayload::RejectedResponse(rejected) => {
return Err(ClientError::Kernel {
code: rejected.code,
message: rejected.message,
});
}
other => {
return Err(ClientError::Sidecar(format!(
"unexpected ACP Ext response: {other:?}"
)));
}
};
if envelope.namespace != ACP_EXTENSION_NAMESPACE {
return Err(ClientError::Sidecar(format!(
"unexpected ACP Ext namespace: {}",
envelope.namespace
)));
}
let response: AcpResponse = serde_bare::from_slice(&envelope.payload).map_err(|error| {
ClientError::Sidecar(format!("failed to decode ACP response: {error}"))
})?;
match response {
AcpResponse::AcpErrorResponse(error) => Err(ClientError::Kernel {
code: error.code,
message: error.message,
}),
response => Ok(response),
}
}
pub async fn respond_permission(
&self,
session_id: &str,
permission_id: &str,
reply: PermissionReply,
) -> Result<JsonRpcResponse> {
let pending = self.require_session(session_id, |entry| {
entry
.pending_permission_replies
.remove(permission_id)
.map(|(_, responder)| responder)
})?;
if let Some(responder) = pending {
let _ = responder.send(reply);
return Ok(JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id: Some(JsonRpcId::Null),
result: Some(json!({
"permissionId": permission_id,
"reply": reply,
"via": "sidecar-request",
})),
error: None,
});
}
Ok(self
.send_session_request(
session_id,
LEGACY_PERMISSION_METHOD,
Some(json!({ "permissionId": permission_id, "reply": reply })),
)
.await?)
}
pub async fn set_session_mode(
&self,
session_id: &str,
mode_id: &str,
) -> Result<JsonRpcResponse> {
Ok(self
.send_session_request(
session_id,
"session/set_mode",
Some(json!({ "modeId": mode_id })),
)
.await?)
}
pub fn get_session_modes(&self, session_id: &str) -> Option<SessionModeState> {
self.require_session(session_id, |entry| entry.modes.lock().clone())
.ok()
.flatten()
}
pub async fn set_session_model(
&self,
session_id: &str,
model: &str,
) -> Result<JsonRpcResponse> {
Ok(self
.set_session_config_by_category(session_id, "model", model)
.await?)
}
pub async fn set_session_thought_level(
&self,
session_id: &str,
level: &str,
) -> Result<JsonRpcResponse> {
Ok(self
.set_session_config_by_category(session_id, "thought_level", level)
.await?)
}
pub fn get_session_config_options(&self, session_id: &str) -> Vec<SessionConfigOption> {
self.require_session(session_id, |entry| entry.config_options.lock().clone())
.unwrap_or_default()
}
pub fn get_session_capabilities(&self, session_id: &str) -> Option<AgentCapabilities> {
self.require_session(session_id, |entry| entry.capabilities.lock().clone())
.ok()
.flatten()
.filter(|caps| !agent_capabilities_is_empty(caps))
}
pub fn get_session_agent_info(&self, session_id: &str) -> Option<AgentInfo> {
self.require_session(session_id, |entry| entry.agent_info.lock().clone())
.ok()
.flatten()
}
pub async fn raw_session_send(
&self,
session_id: &str,
method: &str,
params: Option<Value>,
) -> Result<JsonRpcResponse> {
Ok(self
.send_session_request(session_id, method, params)
.await?)
}
pub async fn raw_send(
&self,
session_id: &str,
method: &str,
params: Option<Value>,
) -> Result<JsonRpcResponse> {
self.raw_session_send(session_id, method, params).await
}
pub fn on_session_event(
&self,
session_id: &str,
) -> std::result::Result<SessionEventSubscription, ClientError> {
let rx = self.require_session(session_id, |entry| entry.event_tx.subscribe())?;
let stream = futures::stream::unfold(rx, move |mut rx| async move {
loop {
match rx.recv().await {
Ok(notification) => return Some((notification, rx)),
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
}
}
});
Ok((Box::pin(stream), Subscription::noop()))
}
pub fn on_permission_request(
&self,
session_id: &str,
) -> std::result::Result<PermissionRequestSubscription, ClientError> {
let rx = self.require_session(session_id, |entry| entry.permission_tx.subscribe())?;
let stream = futures::stream::unfold(rx, move |mut rx| async move {
loop {
match rx.recv().await {
Ok(request) => return Some((request, rx)),
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
}
}
});
Ok((Box::pin(stream), Subscription::noop()))
}
pub fn on_agent_exit(
&self,
session_id: &str,
) -> std::result::Result<AgentExitSubscription, ClientError> {
let rx = self.require_session(session_id, |entry| entry.agent_exit_tx.subscribe())?;
let stream = futures::stream::unfold(rx, move |mut rx| async move {
loop {
match rx.recv().await {
Ok(event) => return Some((event, rx)),
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
}
}
});
Ok((Box::pin(stream), Subscription::noop()))
}
pub(crate) async fn deliver_sidecar_permission_request(
&self,
request: PermissionRouteRequest,
) -> PermissionRouteResult {
let PermissionRouteRequest {
session_id,
permission_id,
params,
} = request;
let (slot_tx, slot_rx) = tokio::sync::oneshot::channel::<PermissionReply>();
let (responder, responder_rx) = PermissionResponder::new();
let description = params
.get("description")
.and_then(Value::as_str)
.map(str::to_string);
let delivered = PermissionRequest {
permission_id: permission_id.clone(),
description,
params,
responder,
};
let registered = self.require_session(&session_id, |entry| {
if entry.permission_tx.receiver_count() == 0 {
return false;
}
let _ = entry
.pending_permission_replies
.insert(permission_id.clone(), slot_tx);
let _ = entry.permission_tx.send(delivered);
true
});
match registered {
Ok(true) => {}
Ok(false) => {
return PermissionRouteResult {
reply: Some(permission_reply_wire(PermissionReply::Reject).to_string()),
};
}
Err(_) => {
return PermissionRouteResult { reply: None };
}
}
let this = self.clone();
let bridge_session_id = session_id.clone();
let bridge_permission_id = permission_id.clone();
tokio::spawn(async move {
if let Ok(reply) = responder_rx.await {
let _ = this
.respond_permission(&bridge_session_id, &bridge_permission_id, reply)
.await;
}
});
let timeout = tokio::time::sleep(std::time::Duration::from_millis(PERMISSION_TIMEOUT_MS));
tokio::pin!(timeout);
tokio::select! {
reply = slot_rx => match reply {
Ok(reply) => PermissionRouteResult {
reply: Some(permission_reply_wire(reply).to_string()),
},
Err(_) => PermissionRouteResult {
reply: Some(permission_reply_wire(PermissionReply::Reject).to_string()),
},
},
_ = &mut timeout => {
let _ = self.require_session(&session_id, |entry| {
let _ = entry.pending_permission_replies.remove(&permission_id);
});
PermissionRouteResult {
reply: None,
}
}
}
}
}
#[cfg(test)]
mod prompt_accumulation_tests {
use super::*;
fn notification(update: Value) -> JsonRpcNotification {
JsonRpcNotification {
jsonrpc: "2.0".to_string(),
method: "session/update".to_string(),
params: Some(json!({ "update": update })),
}
}
#[test]
fn non_chunk_events_do_not_affect_prompt_text() {
let chunk = notification(json!({
"sessionUpdate": "agent_message_chunk",
"content": { "text": "hello" },
}));
let non_chunk = notification(json!({
"sessionUpdate": "current_mode_update",
"currentModeId": "default",
}));
let mut delivered_chunks = 0;
let mut text = String::new();
accumulate_agent_message_chunk(&non_chunk, &mut delivered_chunks, &mut text)
.expect("non-chunk");
accumulate_agent_message_chunk(&chunk, &mut delivered_chunks, &mut text).expect("chunk");
assert_eq!(text, "hello");
}
#[test]
fn prompt_text_capture_limit_rejects_overflowing_chunk() {
let chunk = notification(json!({
"sessionUpdate": "agent_message_chunk",
"content": { "text": "abcd" },
}));
let mut delivered_chunks = 0;
let mut text = "x".repeat(PROMPT_TEXT_CAPTURE_LIMIT_BYTES - 3);
let error = accumulate_agent_message_chunk(&chunk, &mut delivered_chunks, &mut text)
.expect_err("chunk should exceed prompt text cap");
assert!(
error.to_string().contains("prompt text capture is"),
"unexpected error: {error}"
);
assert_eq!(text.len(), PROMPT_TEXT_CAPTURE_LIMIT_BYTES - 3);
}
#[test]
fn prompt_chunk_limit_rejects_more_tracked_chunks() {
let chunk = notification(json!({
"sessionUpdate": "agent_message_chunk",
"content": { "text": "x" },
}));
let mut delivered_chunks = PROMPT_DELIVERED_CHUNK_LIMIT;
let mut text = String::new();
let error = accumulate_agent_message_chunk(&chunk, &mut delivered_chunks, &mut text)
.expect_err("chunk should exceed chunk tracking cap");
assert!(
error
.to_string()
.contains("prompt chunk tracking limit exceeded"),
"unexpected error: {error}"
);
assert!(text.is_empty());
}
#[test]
fn pending_session_request_count_tracks_registered_resolvers() {
let (event_tx, _) = tokio::sync::broadcast::channel(1);
let (permission_tx, _) = tokio::sync::broadcast::channel(1);
let (agent_exit_tx, _) = tokio::sync::broadcast::channel(1);
let entry = SessionEntry {
agent_type: "pi".to_string(),
modes: parking_lot::Mutex::new(None),
config_options: parking_lot::Mutex::new(Vec::new()),
capabilities: parking_lot::Mutex::new(None),
agent_info: parking_lot::Mutex::new(None),
config_overrides: parking_lot::Mutex::new(BTreeMap::new()),
event_tx,
permission_tx,
agent_exit_tx,
pending_permission_replies: scc::HashMap::new(),
pending_session_request_lock: parking_lot::Mutex::new(()),
pending_prompt_resolvers: scc::HashMap::new(),
};
let (first_tx, _first_rx) = tokio::sync::oneshot::channel();
let (second_tx, _second_rx) = tokio::sync::oneshot::channel();
let _ = entry.pending_prompt_resolvers.insert(1, first_tx);
let _ = entry.pending_prompt_resolvers.insert(2, second_tx);
assert_eq!(pending_session_request_count(&entry), 2);
}
}