1use std::collections::{BTreeMap, BTreeSet};
11use std::path::{Path, PathBuf};
12use std::pin::Pin;
13use std::sync::atomic::Ordering;
14
15use anyhow::Result;
16use futures::Stream;
17use serde::{Deserialize, Serialize};
18use serde_json::{json, Value};
19
20use agentos_protocol::generated::v1::{
21 AcpCloseSessionRequest, AcpCreateSessionRequest, AcpGetSessionStateRequest, AcpRequest,
22 AcpResponse, AcpResumeSessionRequest, AcpRuntimeKind, AcpSessionCreatedResponse,
23 AcpSessionRequest, AcpSessionStateResponse,
24};
25use agentos_protocol::ACP_EXTENSION_NAMESPACE;
26use secure_exec_client::wire;
27
28use crate::agent_os::{AgentOs, SessionEntry};
29use crate::config::{AgentOsConfig, MountConfig, ToolKit};
30use crate::error::ClientError;
31use crate::json_rpc::{JsonRpcError, JsonRpcId, JsonRpcNotification, JsonRpcResponse};
32use crate::stream::Subscription;
33use crate::{CLOSED_SESSION_ID_RETENTION_LIMIT, PERMISSION_TIMEOUT_MS};
34
35const LEGACY_PERMISSION_METHOD: &str = "request/permission";
37
38const RESUME_ADAPTER_ENTRYPOINT_ENV: &str = "AGENT_OS_RESUME_ADAPTER_ENTRYPOINT";
44
45pub(crate) const ACP_PERMISSION_METHOD: &str = "session/request_permission";
48
49const SESSION_PENDING_REQUEST_LIMIT: usize = 1024;
51
52pub(crate) struct PermissionRouteRequest {
53 pub(crate) session_id: String,
54 pub(crate) permission_id: String,
55 pub(crate) params: Value,
56}
57
58pub(crate) struct PermissionRouteResult {
59 pub(crate) reply: Option<String>,
60}
61
62struct SessionCreatedResponse {
63 session_id: String,
64 modes: Option<Value>,
65 config_options: Vec<Value>,
66 agent_capabilities: Option<Value>,
67 agent_info: Option<Value>,
68}
69
70pub(crate) struct SessionStateResponse {
71 modes: Option<Value>,
72 config_options: Vec<Value>,
73 agent_capabilities: Option<Value>,
74 agent_info: Option<Value>,
75}
76
77const PROMPT_TEXT_CAPTURE_LIMIT_BYTES: usize = 16 * 1024 * 1024;
79
80const PROMPT_DELIVERED_CHUNK_LIMIT: usize = 262_144;
82
83pub type SessionEventStream = Pin<Box<dyn Stream<Item = JsonRpcNotification> + Send>>;
84pub type SessionEventSubscription = (SessionEventStream, Subscription);
85pub type PermissionRequestStream = Pin<Box<dyn Stream<Item = PermissionRequest> + Send>>;
86pub type PermissionRequestSubscription = (PermissionRequestStream, Subscription);
87
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
94pub struct SessionInfo {
95 #[serde(rename = "sessionId")]
96 pub session_id: String,
97 #[serde(rename = "agentType")]
98 pub agent_type: String,
99}
100
101#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
103pub struct AgentRegistryEntry {
104 pub id: String,
105 #[serde(rename = "acpAdapter")]
106 pub acp_adapter: String,
107 #[serde(rename = "agentPackage")]
108 pub agent_package: String,
109 pub installed: bool,
110}
111
112const BUILTIN_AGENT_IDS: [&str; 4] = ["pi", "pi-cli", "opencode", "claude"];
114
115struct AgentConfigDef {
118 acp_adapter: &'static str,
119 agent_package: &'static str,
120 default_env: &'static [(&'static str, &'static str)],
121}
122
123fn agent_config(agent_type: &str) -> Option<AgentConfigDef> {
125 Some(match agent_type {
126 "pi" => AgentConfigDef {
127 acp_adapter: "@agentos-software/pi",
128 agent_package: "@mariozechner/pi-coding-agent",
129 default_env: &[],
130 },
131 "pi-cli" => AgentConfigDef {
132 acp_adapter: "pi-acp",
133 agent_package: "@mariozechner/pi-coding-agent",
134 default_env: &[],
135 },
136 "opencode" => AgentConfigDef {
137 acp_adapter: "@agentos-software/opencode",
138 agent_package: "@agentos-software/opencode",
139 default_env: &[
140 ("OPENCODE_DISABLE_CONFIG_DEP_INSTALL", "1"),
141 ("OPENCODE_DISABLE_EMBEDDED_WEB_UI", "1"),
142 ],
143 },
144 "claude" => AgentConfigDef {
145 acp_adapter: "@agentos-software/claude-code",
146 agent_package: "@anthropic-ai/claude-agent-sdk",
147 default_env: &[
148 ("CLAUDE_AGENT_SDK_CLIENT_APP", "@rivet-dev/agentos"),
149 ("CLAUDE_CODE_SIMPLE", "1"),
150 ("CLAUDE_CODE_FORCE_AGENT_OS_RIPGREP", "1"),
151 ("CLAUDE_CODE_DEFER_GROWTHBOOK_INIT", "1"),
152 ("CLAUDE_CODE_DISABLE_CWD_PERSIST", "1"),
153 ("CLAUDE_CODE_DISABLE_DEV_NULL_REDIRECT", "1"),
154 ("CLAUDE_CODE_NODE_SHELL_WRAPPER", "1"),
155 ("CLAUDE_CODE_DISABLE_STREAM_JSON_HOOK_EVENTS", "1"),
156 ("CLAUDE_CODE_SHELL", "/bin/sh"),
157 ("CLAUDE_CODE_SKIP_INITIAL_MESSAGES", "1"),
158 ("CLAUDE_CODE_SKIP_SANDBOX_INIT", "1"),
159 ("CLAUDE_CODE_SIMPLE_SHELL_EXEC", "1"),
160 ("CLAUDE_CODE_SWAP_STDIO", "0"),
161 ("CLAUDE_CODE_USE_PIPE_OUTPUT", "1"),
162 ("DISABLE_TELEMETRY", "1"),
163 ("SHELL", "/bin/sh"),
164 ("USE_BUILTIN_RIPGREP", "0"),
165 ],
166 },
167 _ => return None,
168 })
169}
170
171fn resolve_package_bin(
178 config: &AgentOsConfig,
179 package_name: &str,
180 bin_name: Option<&str>,
181) -> std::result::Result<String, ClientError> {
182 let mut candidates = Vec::new();
183 let module_access_cwd = config
184 .module_access_cwd
185 .clone()
186 .unwrap_or_else(|| ".".to_string());
187 candidates.push(
188 Path::new(&module_access_cwd)
189 .join("node_modules")
190 .join(package_name)
191 .join("package.json"),
192 );
193 candidates.extend(node_modules_mount_package_json_paths(config, package_name));
194
195 let contents = candidates
196 .iter()
197 .find_map(|path| std::fs::read_to_string(path).ok())
198 .ok_or_else(|| {
199 let looked = candidates
200 .iter()
201 .map(|path| path.display().to_string())
202 .collect::<Vec<_>>()
203 .join(", ");
204 ClientError::Sidecar(format!(
205 "cannot resolve package {package_name}: no package.json found (looked in {looked})"
206 ))
207 })?;
208 let pkg: Value = serde_json::from_str(&contents).map_err(|error| {
209 ClientError::Sidecar(format!("invalid package.json for {package_name}: {error}"))
210 })?;
211 let bin_entry: Option<String> = match &pkg["bin"] {
212 Value::String(bin) => Some(bin.clone()),
213 Value::Object(map) => bin_name
214 .and_then(|name| map.get(name))
215 .or_else(|| map.get(package_name))
216 .or_else(|| map.values().next())
217 .and_then(|value| value.as_str())
218 .map(|bin| bin.to_string()),
219 _ => None,
220 };
221 let bin_entry = bin_entry.ok_or_else(|| {
222 ClientError::Sidecar(format!("No bin entry found in {package_name}/package.json"))
223 })?;
224 Ok(format!("/root/node_modules/{package_name}/{bin_entry}"))
225}
226
227fn node_modules_mount_package_json_paths(
228 config: &AgentOsConfig,
229 package_name: &str,
230) -> Vec<PathBuf> {
231 config
232 .mounts
233 .iter()
234 .filter_map(|mount| {
235 let MountConfig::Native {
236 path,
237 plugin,
238 read_only: _,
239 } = mount
240 else {
241 return None;
242 };
243 if path != "/root/node_modules" || plugin.id != "host_dir" {
244 return None;
245 }
246 let host_path = plugin
247 .config
248 .as_ref()
249 .and_then(|config| config.get("hostPath"))
250 .and_then(Value::as_str)?;
251 Some(Path::new(host_path).join(package_name).join("package.json"))
252 })
253 .collect()
254}
255
256#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
258#[serde(tag = "type", rename_all = "lowercase")]
259pub enum McpServerConfig {
260 Local {
261 command: String,
262 #[serde(default, skip_serializing_if = "Vec::is_empty")]
263 args: Vec<String>,
264 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
265 env: BTreeMap<String, String>,
266 },
267 Remote {
268 url: String,
269 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
270 headers: BTreeMap<String, String>,
271 },
272}
273
274#[derive(Debug, Clone, PartialEq, Eq, Default)]
276pub struct CreateSessionOptions {
277 pub cwd: Option<String>,
279 pub env: BTreeMap<String, String>,
280 pub mcp_servers: Vec<McpServerConfig>,
282 pub skip_os_instructions: bool,
284 pub additional_instructions: Option<String>,
285}
286
287#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
289pub struct SessionId {
290 #[serde(rename = "sessionId")]
291 pub session_id: String,
292}
293
294#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
299pub struct ResumeSessionResult {
300 #[serde(rename = "sessionId")]
301 pub session_id: String,
302 pub mode: String,
303}
304
305#[derive(Debug, Clone, PartialEq, Eq, Default)]
309pub struct ResumeSessionOptions {
310 pub transcript_path: Option<String>,
313 pub cwd: Option<String>,
315 pub env: BTreeMap<String, String>,
316}
317
318#[derive(Debug, Clone, PartialEq)]
320pub struct PromptResult {
321 pub response: JsonRpcResponse,
322 pub text: String,
323}
324
325#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
327pub struct SessionMode {
328 pub id: String,
329 #[serde(default, skip_serializing_if = "Option::is_none")]
330 pub name: Option<String>,
331 #[serde(default, skip_serializing_if = "Option::is_none")]
332 pub label: Option<String>,
333 #[serde(default, skip_serializing_if = "Option::is_none")]
334 pub description: Option<String>,
335 #[serde(flatten)]
337 pub extra: BTreeMap<String, Value>,
338}
339
340#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
346pub struct SessionModeState {
347 #[serde(default, rename = "currentModeId")]
348 pub current_mode_id: String,
349 #[serde(default, rename = "availableModes")]
350 pub available_modes: Vec<SessionMode>,
351}
352
353#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
355pub struct ConfigAllowedValue {
356 pub id: String,
357 #[serde(default, skip_serializing_if = "Option::is_none")]
358 pub label: Option<String>,
359}
360
361#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
367pub struct SessionConfigOption {
368 #[serde(default)]
369 pub id: String,
370 #[serde(default, skip_serializing_if = "Option::is_none")]
371 pub category: Option<String>,
372 #[serde(default, skip_serializing_if = "Option::is_none")]
373 pub label: Option<String>,
374 #[serde(default, skip_serializing_if = "Option::is_none")]
375 pub description: Option<String>,
376 #[serde(
377 default,
378 rename = "currentValue",
379 skip_serializing_if = "Option::is_none"
380 )]
381 pub current_value: Option<String>,
382 #[serde(
383 default,
384 rename = "allowedValues",
385 skip_serializing_if = "Option::is_none"
386 )]
387 pub allowed_values: Option<Vec<ConfigAllowedValue>>,
388 #[serde(default, rename = "readOnly", skip_serializing_if = "Option::is_none")]
389 pub read_only: Option<bool>,
390}
391
392#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
394pub struct PromptCapabilities {
395 #[serde(default, skip_serializing_if = "Option::is_none")]
396 pub audio: Option<bool>,
397 #[serde(
398 default,
399 rename = "embeddedContext",
400 skip_serializing_if = "Option::is_none"
401 )]
402 pub embedded_context: Option<bool>,
403 #[serde(default, skip_serializing_if = "Option::is_none")]
404 pub image: Option<bool>,
405 #[serde(flatten)]
406 pub extra: BTreeMap<String, Value>,
407}
408
409#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
411pub struct AgentCapabilities {
412 #[serde(default, skip_serializing_if = "Option::is_none")]
413 pub permissions: Option<bool>,
414 #[serde(default, rename = "plan_mode", skip_serializing_if = "Option::is_none")]
415 pub plan_mode: Option<bool>,
416 #[serde(default, skip_serializing_if = "Option::is_none")]
417 pub questions: Option<bool>,
418 #[serde(
419 default,
420 rename = "tool_calls",
421 skip_serializing_if = "Option::is_none"
422 )]
423 pub tool_calls: Option<bool>,
424 #[serde(
425 default,
426 rename = "text_messages",
427 skip_serializing_if = "Option::is_none"
428 )]
429 pub text_messages: Option<bool>,
430 #[serde(default, skip_serializing_if = "Option::is_none")]
431 pub images: Option<bool>,
432 #[serde(
433 default,
434 rename = "file_attachments",
435 skip_serializing_if = "Option::is_none"
436 )]
437 pub file_attachments: Option<bool>,
438 #[serde(
439 default,
440 rename = "session_lifecycle",
441 skip_serializing_if = "Option::is_none"
442 )]
443 pub session_lifecycle: Option<bool>,
444 #[serde(
445 default,
446 rename = "error_events",
447 skip_serializing_if = "Option::is_none"
448 )]
449 pub error_events: Option<bool>,
450 #[serde(default, skip_serializing_if = "Option::is_none")]
451 pub reasoning: Option<bool>,
452 #[serde(default, skip_serializing_if = "Option::is_none")]
453 pub status: Option<bool>,
454 #[serde(
455 default,
456 rename = "streaming_deltas",
457 skip_serializing_if = "Option::is_none"
458 )]
459 pub streaming_deltas: Option<bool>,
460 #[serde(default, rename = "mcp_tools", skip_serializing_if = "Option::is_none")]
461 pub mcp_tools: Option<bool>,
462 #[serde(
463 default,
464 rename = "promptCapabilities",
465 skip_serializing_if = "Option::is_none"
466 )]
467 pub prompt_capabilities: Option<PromptCapabilities>,
468 #[serde(flatten)]
469 pub extra: BTreeMap<String, Value>,
470}
471
472#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
474pub struct AgentInfo {
475 pub name: String,
476 #[serde(default, skip_serializing_if = "Option::is_none")]
477 pub title: Option<String>,
478 #[serde(default, skip_serializing_if = "Option::is_none")]
479 pub version: Option<String>,
480 #[serde(flatten)]
481 pub extra: BTreeMap<String, Value>,
482}
483
484#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
486pub struct SessionInitData {
487 #[serde(default, skip_serializing_if = "Option::is_none")]
488 pub modes: Option<SessionModeState>,
489 #[serde(
490 default,
491 rename = "configOptions",
492 skip_serializing_if = "Option::is_none"
493 )]
494 pub config_options: Option<Vec<SessionConfigOption>>,
495 #[serde(default, skip_serializing_if = "Option::is_none")]
496 pub capabilities: Option<AgentCapabilities>,
497 #[serde(default, rename = "agentInfo", skip_serializing_if = "Option::is_none")]
498 pub agent_info: Option<AgentInfo>,
499}
500
501#[derive(Clone)]
508pub struct PermissionResponder {
509 inner:
510 std::sync::Arc<parking_lot::Mutex<Option<tokio::sync::oneshot::Sender<PermissionReply>>>>,
511}
512
513impl PermissionResponder {
514 pub fn new() -> (Self, tokio::sync::oneshot::Receiver<PermissionReply>) {
516 let (tx, rx) = tokio::sync::oneshot::channel();
517 (
518 Self {
519 inner: std::sync::Arc::new(parking_lot::Mutex::new(Some(tx))),
520 },
521 rx,
522 )
523 }
524
525 pub fn respond(&self, reply: PermissionReply) {
527 if let Some(tx) = self.inner.lock().take() {
528 let _ = tx.send(reply);
529 }
530 }
531}
532
533#[derive(Clone)]
540pub struct PermissionRequest {
541 pub permission_id: String,
542 pub description: Option<String>,
543 pub params: Value,
544 pub responder: PermissionResponder,
545}
546
547impl std::fmt::Debug for PermissionRequest {
548 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
549 f.debug_struct("PermissionRequest")
550 .field("permission_id", &self.permission_id)
551 .field("description", &self.description)
552 .field("params", &self.params)
553 .finish_non_exhaustive()
554 }
555}
556
557#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
559#[serde(rename_all = "lowercase")]
560pub enum PermissionReply {
561 Once,
562 Always,
563 Reject,
564}
565
566fn permission_reply_wire(reply: PermissionReply) -> &'static str {
569 match reply {
570 PermissionReply::Once => "once",
571 PermissionReply::Always => "always",
572 PermissionReply::Reject => "reject",
573 }
574}
575
576fn agent_capabilities_is_empty(caps: &AgentCapabilities) -> bool {
584 caps.permissions.is_none()
585 && caps.plan_mode.is_none()
586 && caps.questions.is_none()
587 && caps.tool_calls.is_none()
588 && caps.text_messages.is_none()
589 && caps.images.is_none()
590 && caps.file_attachments.is_none()
591 && caps.session_lifecycle.is_none()
592 && caps.error_events.is_none()
593 && caps.reasoning.is_none()
594 && caps.status.is_none()
595 && caps.streaming_deltas.is_none()
596 && caps.mcp_tools.is_none()
597 && caps.prompt_capabilities.is_none()
598 && caps.extra.is_empty()
599}
600
601fn should_dispatch_to_session_event_handlers(notification: &JsonRpcNotification) -> bool {
604 notification.method == "session/update"
605}
606
607pub(crate) fn record_live_session_event(entry: &SessionEntry, notification: JsonRpcNotification) {
608 apply_session_update(entry, ¬ification);
609 if should_dispatch_to_session_event_handlers(¬ification) {
610 let _ = entry.event_tx.send(notification);
611 }
612}
613
614fn apply_session_update(entry: &SessionEntry, notification: &JsonRpcNotification) {
615 if notification.method != "session/update" {
616 return;
617 }
618 let Some(params) = notification.params.as_ref().and_then(Value::as_object) else {
619 return;
620 };
621 let update = params
622 .get("update")
623 .and_then(Value::as_object)
624 .unwrap_or(params);
625 match update.get("sessionUpdate").and_then(Value::as_str) {
626 Some("current_mode_update") => {
627 let Some(mode_id) = update.get("currentModeId").and_then(Value::as_str) else {
628 return;
629 };
630 let mut modes = entry.modes.lock();
631 if let Some(modes) = modes.as_mut() {
632 modes.current_mode_id = mode_id.to_string();
633 }
634 }
635 Some("config_option_update") | Some("config_options_update") => {
636 let Some(options) = update.get("configOptions").and_then(Value::as_array) else {
637 return;
638 };
639 let parsed = options
640 .iter()
641 .filter_map(|value| serde_json::from_value(value.clone()).ok())
642 .collect();
643 *entry.config_options.lock() = parsed;
644 apply_synthetic_config_overrides(entry);
645 }
646 Some("agent_message_chunk") | None | Some(_) => {}
647 }
648}
649
650fn accumulate_agent_message_chunk(
651 notification: &JsonRpcNotification,
652 delivered_chunks: &mut usize,
653 agent_text: &mut String,
654) -> std::result::Result<(), ClientError> {
655 let params = notification.params.clone().unwrap_or(Value::Null);
656 let update = params.get("update").cloned().unwrap_or(Value::Null);
657 if update.get("sessionUpdate").and_then(Value::as_str) != Some("agent_message_chunk") {
658 return Ok(());
659 }
660 if let Some(chunk) = update
661 .get("content")
662 .and_then(|content| content.get("text"))
663 .and_then(Value::as_str)
664 {
665 if *delivered_chunks >= PROMPT_DELIVERED_CHUNK_LIMIT {
666 return Err(prompt_chunk_limit_error());
667 }
668 let next_len = agent_text
669 .len()
670 .checked_add(chunk.len())
671 .ok_or_else(|| prompt_text_limit_error(usize::MAX))?;
672 if next_len > PROMPT_TEXT_CAPTURE_LIMIT_BYTES {
673 return Err(prompt_text_limit_error(next_len));
674 }
675 agent_text.push_str(chunk);
676 *delivered_chunks += 1;
677 }
678 Ok(())
679}
680
681fn pending_session_request_count(entry: &SessionEntry) -> usize {
682 let mut count = 0;
683 entry.pending_prompt_resolvers.scan(|_, _| {
684 count += 1;
685 });
686 count
687}
688
689fn prompt_text_limit_error(size: usize) -> ClientError {
690 ClientError::Sidecar(format!(
691 "prompt text capture is {size} bytes, limit is {PROMPT_TEXT_CAPTURE_LIMIT_BYTES}"
692 ))
693}
694
695fn prompt_chunk_limit_error() -> ClientError {
696 ClientError::Sidecar(format!(
697 "prompt chunk tracking limit exceeded: at most {PROMPT_DELIVERED_CHUNK_LIMIT} chunks can be captured per prompt"
698 ))
699}
700
701struct PendingSessionRequestGuard<'a> {
702 os: &'a AgentOs,
703 session_id: &'a str,
704 resolver_id: i64,
705 active: bool,
706}
707
708impl<'a> PendingSessionRequestGuard<'a> {
709 fn new(os: &'a AgentOs, session_id: &'a str, resolver_id: i64) -> Self {
710 Self {
711 os,
712 session_id,
713 resolver_id,
714 active: true,
715 }
716 }
717
718 fn cleanup(&mut self) {
719 if self.active {
720 self.os
721 .cleanup_pending_resolver(self.session_id, self.resolver_id);
722 self.active = false;
723 }
724 }
725}
726
727impl Drop for PendingSessionRequestGuard<'_> {
728 fn drop(&mut self) {
729 self.cleanup();
730 }
731}
732
733fn apply_synthetic_config_overrides(entry: &SessionEntry) {
736 let overrides = entry.config_overrides.lock().clone();
737 if overrides.is_empty() {
738 return;
739 }
740 let mut options = entry.config_options.lock();
741 for option in options.iter_mut() {
742 let override_value = overrides
745 .get(&option.id)
746 .filter(|_| !option.id.starts_with(PENDING_METHOD_PREFIX))
747 .cloned()
748 .or_else(|| {
749 option
750 .category
751 .as_ref()
752 .and_then(|category| overrides.get(category).cloned())
753 });
754 if let Some(value) = override_value {
755 option.current_value = Some(value);
756 }
757 }
758}
759
760const PENDING_METHOD_PREFIX: &str = "__pending_method::";
763
764fn sync_session_state(entry: &SessionEntry, state: &SessionStateResponse) {
767 *entry.modes.lock() = state
768 .modes
769 .as_ref()
770 .filter(|value| value.is_object())
771 .and_then(|value| serde_json::from_value(value.clone()).ok());
772
773 *entry.config_options.lock() = state
774 .config_options
775 .iter()
776 .filter_map(|value| serde_json::from_value(value.clone()).ok())
777 .collect();
778
779 apply_synthetic_config_overrides(entry);
780
781 *entry.capabilities.lock() = state
782 .agent_capabilities
783 .as_ref()
784 .filter(|value| value.is_object())
785 .and_then(|value| serde_json::from_value(value.clone()).ok());
786
787 *entry.agent_info.lock() = state
788 .agent_info
789 .as_ref()
790 .filter(|value| value.is_object())
791 .and_then(|value| serde_json::from_value(value.clone()).ok());
792}
793
794fn unsupported_config_response(agent_type: &str, category: &str) -> JsonRpcResponse {
797 let message = if agent_type == "opencode" && category == "model" {
798 "OpenCode reports available models, but model switching must be configured before createSession() because ACP session/set_config_option is not implemented.".to_string()
799 } else {
800 format!("The {category} config option is read-only for {agent_type} sessions.")
801 };
802 JsonRpcResponse {
803 jsonrpc: "2.0".to_string(),
804 id: Some(JsonRpcId::Null),
805 result: None,
806 error: Some(JsonRpcError {
807 code: -32601,
808 message,
809 data: None,
810 }),
811 }
812}
813
814fn session_closed_response(session_id: &str) -> JsonRpcResponse {
816 JsonRpcResponse {
817 jsonrpc: "2.0".to_string(),
818 id: Some(JsonRpcId::Null),
819 result: None,
820 error: Some(JsonRpcError {
821 code: -32000,
822 message: format!("Session closed: {session_id}"),
823 data: None,
824 }),
825 }
826}
827
828fn session_created_from_acp(
829 response: AcpSessionCreatedResponse,
830) -> std::result::Result<SessionCreatedResponse, ClientError> {
831 Ok(SessionCreatedResponse {
832 session_id: response.session_id,
833 modes: parse_optional_json(response.modes, "modes")?,
834 config_options: parse_json_vec(response.config_options, "configOptions")?,
835 agent_capabilities: parse_optional_json(response.agent_capabilities, "agentCapabilities")?,
836 agent_info: parse_optional_json(response.agent_info, "agentInfo")?,
837 })
838}
839
840fn session_state_from_acp(
841 response: AcpSessionStateResponse,
842) -> std::result::Result<SessionStateResponse, ClientError> {
843 Ok(SessionStateResponse {
844 modes: parse_optional_json(response.modes, "modes")?,
845 config_options: parse_json_vec(response.config_options, "configOptions")?,
846 agent_capabilities: parse_optional_json(response.agent_capabilities, "agentCapabilities")?,
847 agent_info: parse_optional_json(response.agent_info, "agentInfo")?,
848 })
849}
850
851fn parse_optional_json(
852 value: Option<String>,
853 label: &str,
854) -> std::result::Result<Option<Value>, ClientError> {
855 value
856 .map(|value| {
857 serde_json::from_str(&value).map_err(|error| {
858 ClientError::Sidecar(format!("malformed ACP {label} JSON: {error}"))
859 })
860 })
861 .transpose()
862}
863
864fn parse_json_vec(
865 values: Vec<String>,
866 label: &str,
867) -> std::result::Result<Vec<Value>, ClientError> {
868 values
869 .into_iter()
870 .map(|value| {
871 serde_json::from_str(&value).map_err(|error| {
872 ClientError::Sidecar(format!("malformed ACP {label} JSON: {error}"))
873 })
874 })
875 .collect()
876}
877
878fn unexpected_acp_response(operation: &str, response: AcpResponse) -> ClientError {
879 ClientError::Sidecar(format!("unexpected response to {operation}: {response:?}"))
880}
881
882fn combine_instructions(additional: Option<&str>, tool_reference: &str) -> Option<String> {
883 let mut parts = Vec::new();
884 if let Some(additional) = additional.map(str::trim).filter(|value| !value.is_empty()) {
885 parts.push(additional.to_string());
886 }
887 let tool_reference = tool_reference.trim();
888 if !tool_reference.is_empty() {
889 parts.push(tool_reference.to_string());
890 }
891 if parts.is_empty() {
892 None
893 } else {
894 Some(parts.join("\n\n"))
895 }
896}
897
898fn build_host_tool_reference(tool_kits: &[ToolKit]) -> String {
899 if tool_kits.is_empty() {
900 return String::new();
901 }
902
903 let mut lines = vec![
904 String::from("## Available Host Tools"),
905 String::new(),
906 String::from("Run `agentos list-tools` to see all available tools."),
907 String::new(),
908 ];
909
910 for kit in tool_kits {
911 lines.push(format!("### {}", kit.name));
912 lines.push(String::new());
913 lines.push(kit.description.clone());
914 lines.push(String::new());
915 for tool in &kit.tools {
916 let signature = build_tool_flag_signature(&tool.input_schema);
917 let suffix = if signature.is_empty() {
918 String::new()
919 } else {
920 format!(" {signature}")
921 };
922 lines.push(format!(
923 "- `agentos-{} {}{}` — {}",
924 kit.name, tool.name, suffix, tool.description
925 ));
926 }
927 lines.push(String::new());
928 lines.push(format!(
929 "Run `agentos-{} <tool> --help` for details.",
930 kit.name
931 ));
932 lines.push(String::new());
933 }
934
935 lines.join("\n")
936}
937
938fn build_tool_flag_signature(schema: &Value) -> String {
939 describe_tool_flags(schema)
940 .into_iter()
941 .map(|flag| {
942 if flag.required {
943 format!("{} <{}>", flag.name, flag.value_type)
944 } else {
945 format!("[{} <{}>]", flag.name, flag.value_type)
946 }
947 })
948 .collect::<Vec<_>>()
949 .join(" ")
950}
951
952struct ToolFlagDescription {
953 name: String,
954 value_type: String,
955 required: bool,
956}
957
958fn describe_tool_flags(schema: &Value) -> Vec<ToolFlagDescription> {
959 let properties = schema
960 .get("properties")
961 .and_then(Value::as_object)
962 .cloned()
963 .unwrap_or_default();
964 let required = schema
965 .get("required")
966 .and_then(Value::as_array)
967 .map(|items| {
968 items
969 .iter()
970 .filter_map(Value::as_str)
971 .map(str::to_owned)
972 .collect::<BTreeSet<_>>()
973 })
974 .unwrap_or_default();
975
976 properties
977 .into_iter()
978 .map(|(field_name, field_schema)| ToolFlagDescription {
979 name: format!("--{}", camel_to_kebab(&field_name)),
980 value_type: describe_tool_flag_type(&field_schema),
981 required: required.contains(&field_name),
982 })
983 .collect()
984}
985
986fn describe_tool_flag_type(schema: &Value) -> String {
987 match json_schema_type(schema) {
988 Some("array") => {
989 let item_type = schema
990 .get("items")
991 .and_then(json_schema_type)
992 .unwrap_or("string");
993 format!("{item_type}[]")
994 }
995 Some("string") => schema
996 .get("enum")
997 .and_then(Value::as_array)
998 .map(|values| values.iter().filter_map(Value::as_str).collect::<Vec<_>>())
999 .filter(|values| !values.is_empty())
1000 .map(|values| values.join("|"))
1001 .unwrap_or_else(|| String::from("string")),
1002 Some(other) => other.to_string(),
1003 None => String::from("string"),
1004 }
1005}
1006
1007fn json_schema_type(schema: &Value) -> Option<&str> {
1008 schema.get("type").and_then(Value::as_str)
1009}
1010
1011fn camel_to_kebab(value: &str) -> String {
1012 let mut output = String::new();
1013 for (index, ch) in value.chars().enumerate() {
1014 if ch.is_ascii_uppercase() && index > 0 {
1015 output.push('-');
1016 }
1017 output.push(ch.to_ascii_lowercase());
1018 }
1019 output
1020}
1021
1022impl AgentOs {
1027 fn session_ownership(&self) -> wire::OwnershipScope {
1029 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
1030 connection_id: self.connection_id().to_string(),
1031 session_id: self.wire_session_id().to_string(),
1032 vm_id: self.vm_id().to_string(),
1033 })
1034 }
1035
1036 fn require_session<R>(
1038 &self,
1039 session_id: &str,
1040 f: impl FnOnce(&SessionEntry) -> R,
1041 ) -> std::result::Result<R, ClientError> {
1042 self.inner()
1043 .sessions
1044 .read(session_id, |_, entry| f(entry))
1045 .ok_or_else(|| ClientError::SessionNotFound(session_id.to_string()))
1046 }
1047
1048 async fn hydrate_session_state(
1051 &self,
1052 session_id: &str,
1053 ) -> std::result::Result<(), ClientError> {
1054 self.require_session(session_id, |_| ())?;
1055 let response = self
1056 .send_acp_request(AcpRequest::AcpGetSessionStateRequest(
1057 AcpGetSessionStateRequest {
1058 session_id: session_id.to_string(),
1059 },
1060 ))
1061 .await?;
1062 let AcpResponse::AcpSessionStateResponse(state) = response else {
1063 return Err(unexpected_acp_response(
1064 "AcpGetSessionStateRequest",
1065 response,
1066 ));
1067 };
1068 let state = session_state_from_acp(state)?;
1069
1070 self.require_session(session_id, |entry| sync_session_state(entry, &state))?;
1071 Ok(())
1072 }
1073
1074 pub(crate) async fn send_session_request(
1078 &self,
1079 session_id: &str,
1080 method: &str,
1081 params: Option<Value>,
1082 ) -> std::result::Result<JsonRpcResponse, ClientError> {
1083 let request_params = params;
1084
1085 let resolver_id = self.inner().request_counter.fetch_add(1, Ordering::SeqCst);
1090 let (resolve_tx, resolve_rx) = tokio::sync::oneshot::channel::<JsonRpcResponse>();
1091 self.require_session(session_id, |entry| {
1092 let _guard = entry.pending_session_request_lock.lock();
1093 if pending_session_request_count(entry) >= SESSION_PENDING_REQUEST_LIMIT {
1094 return Err(ClientError::Sidecar(format!(
1095 "session pending request limit exceeded: at most {SESSION_PENDING_REQUEST_LIMIT} requests can be in flight per session"
1096 )));
1097 }
1098 let _ = entry
1099 .pending_prompt_resolvers
1100 .insert(resolver_id, resolve_tx);
1101 entry
1103 .config_overrides
1104 .lock()
1105 .entry(format!("{PENDING_METHOD_PREFIX}{resolver_id}"))
1106 .or_insert_with(|| method.to_string());
1107 Ok(())
1108 })??;
1109 let mut pending_request_guard =
1110 PendingSessionRequestGuard::new(self, session_id, resolver_id);
1111
1112 let rpc = self.send_acp_request(AcpRequest::AcpSessionRequest(AcpSessionRequest {
1113 session_id: session_id.to_string(),
1114 method: method.to_string(),
1115 params: request_params
1116 .clone()
1117 .map(|params| serde_json::to_string(¶ms))
1118 .transpose()
1119 .map_err(|error| {
1120 ClientError::Sidecar(format!("failed to encode session params: {error}"))
1121 })?,
1122 }));
1123 tokio::pin!(rpc);
1124
1125 let response = tokio::select! {
1126 biased;
1127 resolved = resolve_rx => {
1128 pending_request_guard.cleanup();
1132 match resolved {
1133 Ok(response) => return Ok(response),
1134 Err(_) => return Ok(session_closed_response(session_id)),
1135 }
1136 }
1137 result = &mut rpc => {
1138 pending_request_guard.cleanup();
1139 result?
1140 }
1141 };
1142
1143 let response = match response {
1144 AcpResponse::AcpSessionRpcResponse(rpc) => {
1145 serde_json::from_str::<JsonRpcResponse>(&rpc.response).map_err(|err| {
1146 ClientError::Sidecar(format!("malformed session rpc response: {err}"))
1147 })?
1148 }
1149 other => return Err(unexpected_acp_response("AcpSessionRequest", other)),
1150 };
1151
1152 let _ = self.hydrate_session_state(session_id).await;
1154
1155 if response.error.is_none() {
1156 self.apply_post_send_cache_updates(session_id, method, request_params.as_ref())?;
1157 }
1158
1159 Ok(response)
1160 }
1161
1162 fn cleanup_pending_resolver(&self, session_id: &str, resolver_id: i64) {
1164 let _ = self.require_session(session_id, |entry| {
1165 let _ = entry.pending_prompt_resolvers.remove(&resolver_id);
1166 entry
1167 .config_overrides
1168 .lock()
1169 .remove(&format!("{PENDING_METHOD_PREFIX}{resolver_id}"));
1170 });
1171 }
1172
1173 fn apply_post_send_cache_updates(
1175 &self,
1176 session_id: &str,
1177 method: &str,
1178 params: Option<&Value>,
1179 ) -> std::result::Result<(), ClientError> {
1180 self.require_session(session_id, |entry| {
1181 if method == "session/set_mode" {
1182 if let Some(mode_id) = params.and_then(|p| p.get("modeId")).and_then(Value::as_str)
1183 {
1184 let mut modes = entry.modes.lock();
1185 if let Some(modes) = modes.as_mut() {
1186 modes.current_mode_id = mode_id.to_string();
1187 }
1188 }
1189 }
1190 if method == "session/set_config_option" {
1191 let config_id = params
1192 .and_then(|p| p.get("configId"))
1193 .and_then(Value::as_str);
1194 let value = params.and_then(|p| p.get("value")).and_then(Value::as_str);
1195 if let (Some(config_id), Some(value)) = (config_id, value) {
1196 let mut options = entry.config_options.lock();
1197 for option in options.iter_mut() {
1198 if option.id == config_id {
1199 option.current_value = Some(value.to_string());
1200 }
1201 }
1202 }
1203 }
1204 })
1205 }
1206
1207 async fn set_session_config_by_category(
1210 &self,
1211 session_id: &str,
1212 category: &str,
1213 value: &str,
1214 ) -> std::result::Result<JsonRpcResponse, ClientError> {
1215 let (read_only, config_id, agent_type) = self.require_session(session_id, |entry| {
1216 let options = entry.config_options.lock();
1217 let option = options
1218 .iter()
1219 .find(|option| option.category.as_deref() == Some(category));
1220 (
1221 option.and_then(|option| option.read_only).unwrap_or(false),
1222 option.map(|option| option.id.clone()),
1223 entry.agent_type.clone(),
1224 )
1225 })?;
1226
1227 if read_only {
1228 return Ok(unsupported_config_response(&agent_type, category));
1229 }
1230
1231 let config_id = config_id.unwrap_or_else(|| category.to_string());
1232 let response = self
1233 .send_session_request(
1234 session_id,
1235 "session/set_config_option",
1236 Some(json!({ "configId": config_id, "value": value })),
1237 )
1238 .await?;
1239
1240 Ok(response)
1241 }
1242
1243 pub fn list_sessions(&self) -> Vec<SessionInfo> {
1245 let mut sessions = Vec::new();
1246 self.inner().sessions.scan(|session_id, entry| {
1247 sessions.push(SessionInfo {
1248 session_id: session_id.clone(),
1249 agent_type: entry.agent_type.clone(),
1250 });
1251 });
1252 sessions
1253 }
1254
1255 pub fn list_agents(&self) -> Vec<AgentRegistryEntry> {
1263 BUILTIN_AGENT_IDS
1264 .iter()
1265 .filter_map(|id| {
1266 let config = agent_config(id)?;
1267 let installed =
1268 resolve_package_bin(self.config(), config.acp_adapter, None).is_ok();
1269 Some(AgentRegistryEntry {
1270 id: (*id).to_string(),
1271 acp_adapter: config.acp_adapter.to_string(),
1272 agent_package: config.agent_package.to_string(),
1273 installed,
1274 })
1275 })
1276 .collect()
1277 }
1278
1279 pub async fn create_session(
1285 &self,
1286 agent_type: &str,
1287 options: CreateSessionOptions,
1288 ) -> Result<SessionId> {
1289 let config = agent_config(agent_type)
1290 .ok_or_else(|| ClientError::Sidecar(format!("Unknown agent type: {agent_type}")))?;
1291
1292 let adapter_entrypoint = resolve_package_bin(self.config(), config.acp_adapter, None)?;
1295
1296 let mut env: BTreeMap<String, String> = config
1298 .default_env
1299 .iter()
1300 .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
1301 .collect();
1302 for (key, value) in &options.env {
1303 env.insert(key.clone(), value.clone());
1304 }
1305 if (agent_type == "pi" || agent_type == "pi-cli") && !env.contains_key("PI_ACP_PI_COMMAND")
1306 {
1307 if let Ok(pi_command) =
1308 resolve_package_bin(self.config(), config.agent_package, Some("pi"))
1309 {
1310 env.insert("PI_ACP_PI_COMMAND".to_string(), pi_command);
1311 }
1312 }
1313
1314 let cwd = options
1315 .cwd
1316 .clone()
1317 .unwrap_or_else(|| "/workspace".to_string());
1318 let mcp_servers: Vec<Value> = options
1319 .mcp_servers
1320 .iter()
1321 .filter_map(|server| serde_json::to_value(server).ok())
1322 .collect();
1323 let client_capabilities = json!({
1324 "fs": { "readTextFile": true, "writeTextFile": true },
1325 "terminal": true,
1326 });
1327 let tool_reference = build_host_tool_reference(&self.config().tool_kits);
1328 let additional_instructions =
1329 combine_instructions(options.additional_instructions.as_deref(), &tool_reference);
1330
1331 let response = self
1332 .send_acp_request(AcpRequest::AcpCreateSessionRequest(
1333 AcpCreateSessionRequest {
1334 agent_type: agent_type.to_string(),
1335 runtime: AcpRuntimeKind::JavaScript,
1336 adapter_entrypoint,
1337 args: Vec::new(),
1338 env: env.into_iter().collect(),
1339 cwd,
1340 mcp_servers: serde_json::to_string(&mcp_servers).map_err(|error| {
1341 ClientError::Sidecar(format!("failed to encode MCP servers: {error}"))
1342 })?,
1343 protocol_version: crate::ACP_PROTOCOL_VERSION as i32,
1344 client_capabilities: serde_json::to_string(&client_capabilities).map_err(
1345 |error| {
1346 ClientError::Sidecar(format!(
1347 "failed to encode client capabilities: {error}"
1348 ))
1349 },
1350 )?,
1351 additional_instructions,
1352 skip_os_instructions: options.skip_os_instructions,
1353 },
1354 ))
1355 .await?;
1356 let AcpResponse::AcpSessionCreatedResponse(created) = response else {
1357 return Err(unexpected_acp_response("AcpCreateSessionRequest", response).into());
1358 };
1359 let created = session_created_from_acp(created)?;
1360
1361 let state = SessionStateResponse {
1364 modes: created.modes,
1365 config_options: created.config_options,
1366 agent_capabilities: created.agent_capabilities,
1367 agent_info: created.agent_info,
1368 };
1369 self.register_session(&created.session_id, agent_type, &state)
1370 .await?;
1371
1372 Ok(SessionId {
1373 session_id: created.session_id,
1374 })
1375 }
1376
1377 pub(crate) async fn register_session(
1381 &self,
1382 session_id: &str,
1383 agent_type: &str,
1384 state: &SessionStateResponse,
1385 ) -> std::result::Result<(), ClientError> {
1386 {
1387 let mut closed = self.inner().closed_session_ids.lock();
1388 closed.retain(|id| id != session_id);
1389 }
1390
1391 let (event_tx, _) = tokio::sync::broadcast::channel(1024);
1392 let (permission_tx, _) = tokio::sync::broadcast::channel(64);
1393 let entry = SessionEntry {
1394 agent_type: agent_type.to_string(),
1395 modes: parking_lot::Mutex::new(None),
1396 config_options: parking_lot::Mutex::new(Vec::new()),
1397 capabilities: parking_lot::Mutex::new(None),
1398 agent_info: parking_lot::Mutex::new(None),
1399 config_overrides: parking_lot::Mutex::new(BTreeMap::new()),
1400 event_tx,
1401 permission_tx,
1402 pending_permission_replies: scc::HashMap::new(),
1403 pending_session_request_lock: parking_lot::Mutex::new(()),
1404 pending_prompt_resolvers: scc::HashMap::new(),
1405 };
1406 sync_session_state(&entry, state);
1407 let _ = self.inner().sessions.insert(session_id.to_string(), entry);
1408
1409 match self.hydrate_session_state(session_id).await {
1410 Ok(()) => Ok(()),
1411 Err(error) => {
1412 let _ = self.inner().sessions.remove(session_id);
1413 Err(error)
1414 }
1415 }
1416 }
1417
1418 pub async fn resume_session(
1432 &self,
1433 session_id: &str,
1434 agent_type: &str,
1435 options: ResumeSessionOptions,
1436 ) -> Result<ResumeSessionResult> {
1437 let config = agent_config(agent_type)
1438 .ok_or_else(|| ClientError::Sidecar(format!("Unknown agent type: {agent_type}")))?;
1439 let adapter_entrypoint = resolve_package_bin(self.config(), config.acp_adapter, None)?;
1440
1441 let mut env: BTreeMap<String, String> = config
1445 .default_env
1446 .iter()
1447 .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
1448 .collect();
1449 for (key, value) in &options.env {
1450 env.insert(key.clone(), value.clone());
1451 }
1452 if (agent_type == "pi" || agent_type == "pi-cli") && !env.contains_key("PI_ACP_PI_COMMAND")
1453 {
1454 if let Ok(pi_command) =
1455 resolve_package_bin(self.config(), config.agent_package, Some("pi"))
1456 {
1457 env.insert("PI_ACP_PI_COMMAND".to_string(), pi_command);
1458 }
1459 }
1460 env.insert(
1461 RESUME_ADAPTER_ENTRYPOINT_ENV.to_string(),
1462 adapter_entrypoint,
1463 );
1464
1465 let cwd = options
1466 .cwd
1467 .clone()
1468 .unwrap_or_else(|| "/workspace".to_string());
1469
1470 let response = self
1471 .send_acp_request(AcpRequest::AcpResumeSessionRequest(
1472 AcpResumeSessionRequest {
1473 session_id: session_id.to_string(),
1474 agent_type: agent_type.to_string(),
1475 transcript_path: options.transcript_path.clone(),
1476 cwd,
1477 env: env.into_iter().collect(),
1478 },
1479 ))
1480 .await?;
1481 let AcpResponse::AcpSessionResumedResponse(resumed) = response else {
1482 return Err(unexpected_acp_response("AcpResumeSessionRequest", response).into());
1483 };
1484
1485 let empty_state = SessionStateResponse {
1487 modes: None,
1488 config_options: Vec::new(),
1489 agent_capabilities: None,
1490 agent_info: None,
1491 };
1492 self.register_session(&resumed.session_id, agent_type, &empty_state)
1493 .await?;
1494
1495 Ok(ResumeSessionResult {
1496 session_id: resumed.session_id,
1497 mode: resumed.mode,
1498 })
1499 }
1500
1501 pub async fn destroy_session(&self, session_id: &str) -> Result<()> {
1503 self.require_session(session_id, |_| ())?;
1504 let _ = self.cancel_session(session_id).await;
1505 self.close_session_internal(session_id).await?;
1506 Ok(())
1507 }
1508
1509 pub async fn prompt(&self, session_id: &str, text: &str) -> Result<PromptResult> {
1513 let mut rx = self.require_session(session_id, |entry| entry.event_tx.subscribe())?;
1514
1515 let mut agent_text = String::new();
1516 let mut delivered_chunks = 0;
1517 let mut prompt_text_error: Option<ClientError> = None;
1518
1519 let request = self.send_session_request(
1520 session_id,
1521 "session/prompt",
1522 Some(json!({ "prompt": [{ "type": "text", "text": text }] })),
1523 );
1524 tokio::pin!(request);
1525
1526 let response = loop {
1529 tokio::select! {
1530 biased;
1531 result = &mut request => break result,
1532 event = rx.recv() => {
1533 match event {
1534 Ok(event) => accumulate_agent_message_chunk(
1535 &event,
1536 &mut delivered_chunks,
1537 &mut agent_text,
1538 )
1539 .unwrap_or_else(|error| {
1540 prompt_text_error.get_or_insert(error);
1541 }),
1542 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
1543 Err(tokio::sync::broadcast::error::RecvError::Closed) => {
1544 break (&mut request).await;
1546 }
1547 }
1548 }
1549 }
1550 };
1551
1552 loop {
1554 match rx.try_recv() {
1555 Ok(event) => {
1556 accumulate_agent_message_chunk(&event, &mut delivered_chunks, &mut agent_text)
1557 .unwrap_or_else(|error| {
1558 prompt_text_error.get_or_insert(error);
1559 })
1560 }
1561 Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
1562 Err(tokio::sync::broadcast::error::TryRecvError::Empty)
1563 | Err(tokio::sync::broadcast::error::TryRecvError::Closed) => break,
1564 }
1565 }
1566 drop(rx);
1567
1568 let response = response?;
1569 if let Some(error) = prompt_text_error {
1570 return Err(error.into());
1571 }
1572
1573 Ok(PromptResult {
1574 response,
1575 text: agent_text,
1576 })
1577 }
1578
1579 pub async fn cancel_session(&self, session_id: &str) -> Result<JsonRpcResponse> {
1583 self.require_session(session_id, |_| ())?;
1584 let cancelled_pending_prompt = self.cancel_pending_prompt_requests(session_id)?;
1585 if cancelled_pending_prompt {
1586 let this = self.clone();
1589 let session_id_owned = session_id.to_string();
1590 tokio::spawn(async move {
1591 let _ = this
1592 .send_session_request(&session_id_owned, "session/cancel", None)
1593 .await;
1594 });
1595 return Ok(JsonRpcResponse {
1596 jsonrpc: "2.0".to_string(),
1597 id: Some(JsonRpcId::Null),
1598 result: Some(json!({
1599 "cancelled": true,
1600 "requested": true,
1601 "via": "prompt-fallback",
1602 })),
1603 error: None,
1604 });
1605 }
1606 Ok(self
1607 .send_session_request(session_id, "session/cancel", None)
1608 .await?)
1609 }
1610
1611 fn cancel_pending_prompt_requests(
1614 &self,
1615 session_id: &str,
1616 ) -> std::result::Result<bool, ClientError> {
1617 self.require_session(session_id, |entry| {
1618 let mut prompt_resolver_ids = Vec::new();
1619 {
1620 let overrides = entry.config_overrides.lock();
1621 for (key, method) in overrides.iter() {
1622 if let Some(id) = key.strip_prefix(PENDING_METHOD_PREFIX) {
1623 if method == "session/prompt" {
1624 if let Ok(id) = id.parse::<i64>() {
1625 prompt_resolver_ids.push(id);
1626 }
1627 }
1628 }
1629 }
1630 }
1631 let mut cancelled = false;
1632 for id in prompt_resolver_ids {
1633 if let Some((_, resolver)) = entry.pending_prompt_resolvers.remove(&id) {
1634 let _ = resolver.send(JsonRpcResponse {
1637 jsonrpc: "2.0".to_string(),
1638 id: Some(JsonRpcId::Null),
1639 result: Some(json!({ "stopReason": "cancelled" })),
1640 error: None,
1641 });
1642 cancelled = true;
1643 }
1644 entry
1645 .config_overrides
1646 .lock()
1647 .remove(&format!("{PENDING_METHOD_PREFIX}{id}"));
1648 }
1649 cancelled
1650 })
1651 }
1652
1653 fn abort_pending_session_requests(&self, session_id: &str) {
1656 let _ = self.require_session(session_id, |entry| {
1657 let mut ids = Vec::new();
1658 entry.pending_prompt_resolvers.scan(|id, _| ids.push(*id));
1659 for id in ids {
1660 if let Some((_, resolver)) = entry.pending_prompt_resolvers.remove(&id) {
1661 let _ = resolver.send(session_closed_response(session_id));
1664 }
1665 entry
1666 .config_overrides
1667 .lock()
1668 .remove(&format!("{PENDING_METHOD_PREFIX}{id}"));
1669 }
1670 });
1671 }
1672
1673 fn reject_pending_permission_replies(&self, session_id: &str) {
1677 let _ = self.require_session(session_id, |entry| {
1678 let mut ids = Vec::new();
1679 entry
1680 .pending_permission_replies
1681 .scan(|id, _| ids.push(id.clone()));
1682 for id in ids {
1683 let _ = entry.pending_permission_replies.remove(&id);
1684 }
1685 });
1686 }
1687
1688 pub fn close_session(&self, session_id: &str) -> std::result::Result<(), ClientError> {
1693 let known = self.inner().sessions.contains(session_id)
1694 || self.inner().closing_session_ids.contains(session_id)
1695 || self
1696 .inner()
1697 .closed_session_ids
1698 .lock()
1699 .iter()
1700 .any(|id| id == session_id);
1701 if !known {
1702 return Err(ClientError::SessionNotFound(session_id.to_string()));
1703 }
1704
1705 let _ = self
1709 .inner()
1710 .closing_session_ids
1711 .insert(session_id.to_string());
1712
1713 let this = self.clone();
1714 let session_id_owned = session_id.to_string();
1715 tokio::spawn(async move {
1716 let _ = this.close_session_internal(&session_id_owned).await;
1717 let _ = this.inner().closing_session_ids.remove(&session_id_owned);
1718 });
1719 Ok(())
1720 }
1721
1722 pub(crate) async fn close_session_internal(
1726 &self,
1727 session_id: &str,
1728 ) -> std::result::Result<(), ClientError> {
1729 if self
1730 .inner()
1731 .closed_session_ids
1732 .lock()
1733 .iter()
1734 .any(|id| id == session_id)
1735 {
1736 return Ok(());
1737 }
1738
1739 self.abort_pending_session_requests(session_id);
1740 self.reject_pending_permission_replies(session_id);
1741
1742 if !self.inner().sessions.contains(session_id) {
1744 return Err(ClientError::SessionNotFound(session_id.to_string()));
1745 }
1746 let _ = self.inner().sessions.remove(session_id);
1747 {
1748 let mut closed = self.inner().closed_session_ids.lock();
1749 closed.push_back(session_id.to_string());
1750 while closed.len() > CLOSED_SESSION_ID_RETENTION_LIMIT {
1751 closed.pop_front();
1752 }
1753 }
1754
1755 let response = self
1766 .send_acp_request(AcpRequest::AcpCloseSessionRequest(AcpCloseSessionRequest {
1767 session_id: session_id.to_string(),
1768 }))
1769 .await?;
1770 match response {
1771 AcpResponse::AcpSessionClosedResponse(_) => Ok(()),
1772 other => Err(unexpected_acp_response("AcpCloseSessionRequest", other)),
1773 }
1774 }
1775
1776 async fn send_acp_request(
1777 &self,
1778 request: AcpRequest,
1779 ) -> std::result::Result<AcpResponse, ClientError> {
1780 let payload = serde_bare::to_vec(&request).map_err(|error| {
1781 ClientError::Sidecar(format!("failed to encode ACP request: {error}"))
1782 })?;
1783 let response = self
1784 .transport()
1785 .request_wire(
1786 self.session_ownership(),
1787 wire::RequestPayload::ExtEnvelope(wire::ExtEnvelope {
1788 namespace: ACP_EXTENSION_NAMESPACE.to_string(),
1789 payload,
1790 }),
1791 )
1792 .await?;
1793 let envelope = match response {
1794 wire::ResponsePayload::ExtEnvelope(envelope) => envelope,
1795 wire::ResponsePayload::RejectedResponse(rejected) => {
1796 return Err(ClientError::Kernel {
1797 code: rejected.code,
1798 message: rejected.message,
1799 });
1800 }
1801 other => {
1802 return Err(ClientError::Sidecar(format!(
1803 "unexpected ACP Ext response: {other:?}"
1804 )));
1805 }
1806 };
1807 if envelope.namespace != ACP_EXTENSION_NAMESPACE {
1808 return Err(ClientError::Sidecar(format!(
1809 "unexpected ACP Ext namespace: {}",
1810 envelope.namespace
1811 )));
1812 }
1813 let response: AcpResponse = serde_bare::from_slice(&envelope.payload).map_err(|error| {
1814 ClientError::Sidecar(format!("failed to decode ACP response: {error}"))
1815 })?;
1816 match response {
1817 AcpResponse::AcpErrorResponse(error) => Err(ClientError::Kernel {
1818 code: error.code,
1819 message: error.message,
1820 }),
1821 response => Ok(response),
1822 }
1823 }
1824
1825 pub async fn respond_permission(
1829 &self,
1830 session_id: &str,
1831 permission_id: &str,
1832 reply: PermissionReply,
1833 ) -> Result<JsonRpcResponse> {
1834 let pending = self.require_session(session_id, |entry| {
1835 entry
1836 .pending_permission_replies
1837 .remove(permission_id)
1838 .map(|(_, responder)| responder)
1839 })?;
1840
1841 if let Some(responder) = pending {
1842 let _ = responder.send(reply);
1843 return Ok(JsonRpcResponse {
1844 jsonrpc: "2.0".to_string(),
1845 id: Some(JsonRpcId::Null),
1846 result: Some(json!({
1847 "permissionId": permission_id,
1848 "reply": reply,
1849 "via": "sidecar-request",
1850 })),
1851 error: None,
1852 });
1853 }
1854
1855 Ok(self
1856 .send_session_request(
1857 session_id,
1858 LEGACY_PERMISSION_METHOD,
1859 Some(json!({ "permissionId": permission_id, "reply": reply })),
1860 )
1861 .await?)
1862 }
1863
1864 pub async fn set_session_mode(
1866 &self,
1867 session_id: &str,
1868 mode_id: &str,
1869 ) -> Result<JsonRpcResponse> {
1870 Ok(self
1871 .send_session_request(
1872 session_id,
1873 "session/set_mode",
1874 Some(json!({ "modeId": mode_id })),
1875 )
1876 .await?)
1877 }
1878
1879 pub fn get_session_modes(&self, session_id: &str) -> Option<SessionModeState> {
1881 self.require_session(session_id, |entry| entry.modes.lock().clone())
1882 .ok()
1883 .flatten()
1884 }
1885
1886 pub async fn set_session_model(
1889 &self,
1890 session_id: &str,
1891 model: &str,
1892 ) -> Result<JsonRpcResponse> {
1893 Ok(self
1894 .set_session_config_by_category(session_id, "model", model)
1895 .await?)
1896 }
1897
1898 pub async fn set_session_thought_level(
1900 &self,
1901 session_id: &str,
1902 level: &str,
1903 ) -> Result<JsonRpcResponse> {
1904 Ok(self
1905 .set_session_config_by_category(session_id, "thought_level", level)
1906 .await?)
1907 }
1908
1909 pub fn get_session_config_options(&self, session_id: &str) -> Vec<SessionConfigOption> {
1911 self.require_session(session_id, |entry| entry.config_options.lock().clone())
1912 .unwrap_or_default()
1913 }
1914
1915 pub fn get_session_capabilities(&self, session_id: &str) -> Option<AgentCapabilities> {
1918 self.require_session(session_id, |entry| entry.capabilities.lock().clone())
1919 .ok()
1920 .flatten()
1921 .filter(|caps| !agent_capabilities_is_empty(caps))
1922 }
1923
1924 pub fn get_session_agent_info(&self, session_id: &str) -> Option<AgentInfo> {
1926 self.require_session(session_id, |entry| entry.agent_info.lock().clone())
1927 .ok()
1928 .flatten()
1929 }
1930
1931 pub async fn raw_session_send(
1934 &self,
1935 session_id: &str,
1936 method: &str,
1937 params: Option<Value>,
1938 ) -> Result<JsonRpcResponse> {
1939 Ok(self
1940 .send_session_request(session_id, method, params)
1941 .await?)
1942 }
1943
1944 pub async fn raw_send(
1946 &self,
1947 session_id: &str,
1948 method: &str,
1949 params: Option<Value>,
1950 ) -> Result<JsonRpcResponse> {
1951 self.raw_session_send(session_id, method, params).await
1952 }
1953
1954 pub fn on_session_event(
1957 &self,
1958 session_id: &str,
1959 ) -> std::result::Result<SessionEventSubscription, ClientError> {
1960 let rx = self.require_session(session_id, |entry| entry.event_tx.subscribe())?;
1961 let stream = futures::stream::unfold(rx, move |mut rx| async move {
1962 loop {
1963 match rx.recv().await {
1964 Ok(notification) => return Some((notification, rx)),
1965 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
1966 Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
1967 }
1968 }
1969 });
1970 Ok((Box::pin(stream), Subscription::noop()))
1971 }
1972
1973 pub fn on_permission_request(
1980 &self,
1981 session_id: &str,
1982 ) -> std::result::Result<PermissionRequestSubscription, ClientError> {
1983 let rx = self.require_session(session_id, |entry| entry.permission_tx.subscribe())?;
1984
1985 let stream = futures::stream::unfold(rx, move |mut rx| async move {
1989 loop {
1990 match rx.recv().await {
1991 Ok(request) => return Some((request, rx)),
1992 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
1993 Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
1994 }
1995 }
1996 });
1997
1998 Ok((Box::pin(stream), Subscription::noop()))
1999 }
2000
2001 pub(crate) async fn deliver_sidecar_permission_request(
2010 &self,
2011 request: PermissionRouteRequest,
2012 ) -> PermissionRouteResult {
2013 let PermissionRouteRequest {
2014 session_id,
2015 permission_id,
2016 params,
2017 } = request;
2018
2019 let (slot_tx, slot_rx) = tokio::sync::oneshot::channel::<PermissionReply>();
2020 let (responder, responder_rx) = PermissionResponder::new();
2021 let description = params
2022 .get("description")
2023 .and_then(Value::as_str)
2024 .map(str::to_string);
2025 let delivered = PermissionRequest {
2026 permission_id: permission_id.clone(),
2027 description,
2028 params,
2029 responder,
2030 };
2031
2032 let registered = self.require_session(&session_id, |entry| {
2035 if entry.permission_tx.receiver_count() == 0 {
2036 return false;
2037 }
2038 let _ = entry
2039 .pending_permission_replies
2040 .insert(permission_id.clone(), slot_tx);
2041 let _ = entry.permission_tx.send(delivered);
2042 true
2043 });
2044 match registered {
2045 Ok(true) => {}
2046 Ok(false) => {
2047 return PermissionRouteResult {
2048 reply: Some(permission_reply_wire(PermissionReply::Reject).to_string()),
2049 };
2050 }
2051 Err(_) => {
2052 return PermissionRouteResult { reply: None };
2053 }
2054 }
2055
2056 let this = self.clone();
2058 let bridge_session_id = session_id.clone();
2059 let bridge_permission_id = permission_id.clone();
2060 tokio::spawn(async move {
2061 if let Ok(reply) = responder_rx.await {
2062 let _ = this
2063 .respond_permission(&bridge_session_id, &bridge_permission_id, reply)
2064 .await;
2065 }
2066 });
2067
2068 let timeout = tokio::time::sleep(std::time::Duration::from_millis(PERMISSION_TIMEOUT_MS));
2069 tokio::pin!(timeout);
2070 tokio::select! {
2071 reply = slot_rx => match reply {
2072 Ok(reply) => PermissionRouteResult {
2073 reply: Some(permission_reply_wire(reply).to_string()),
2074 },
2075 Err(_) => PermissionRouteResult {
2077 reply: Some(permission_reply_wire(PermissionReply::Reject).to_string()),
2078 },
2079 },
2080 _ = &mut timeout => {
2081 let _ = self.require_session(&session_id, |entry| {
2082 let _ = entry.pending_permission_replies.remove(&permission_id);
2083 });
2084 PermissionRouteResult {
2085 reply: None,
2086 }
2087 }
2088 }
2089 }
2090}
2091
2092#[cfg(test)]
2095mod prompt_accumulation_tests {
2096 use super::*;
2097
2098 fn notification(update: Value) -> JsonRpcNotification {
2099 JsonRpcNotification {
2100 jsonrpc: "2.0".to_string(),
2101 method: "session/update".to_string(),
2102 params: Some(json!({ "update": update })),
2103 }
2104 }
2105
2106 #[test]
2107 fn non_chunk_events_do_not_affect_prompt_text() {
2108 let chunk = notification(json!({
2109 "sessionUpdate": "agent_message_chunk",
2110 "content": { "text": "hello" },
2111 }));
2112 let non_chunk = notification(json!({
2113 "sessionUpdate": "current_mode_update",
2114 "currentModeId": "default",
2115 }));
2116
2117 let mut delivered_chunks = 0;
2118 let mut text = String::new();
2119 accumulate_agent_message_chunk(&non_chunk, &mut delivered_chunks, &mut text)
2120 .expect("non-chunk");
2121 accumulate_agent_message_chunk(&chunk, &mut delivered_chunks, &mut text).expect("chunk");
2122
2123 assert_eq!(text, "hello");
2124 }
2125
2126 #[test]
2127 fn prompt_text_capture_limit_rejects_overflowing_chunk() {
2128 let chunk = notification(json!({
2129 "sessionUpdate": "agent_message_chunk",
2130 "content": { "text": "abcd" },
2131 }));
2132 let mut delivered_chunks = 0;
2133 let mut text = "x".repeat(PROMPT_TEXT_CAPTURE_LIMIT_BYTES - 3);
2134 let error = accumulate_agent_message_chunk(&chunk, &mut delivered_chunks, &mut text)
2135 .expect_err("chunk should exceed prompt text cap");
2136 assert!(
2137 error.to_string().contains("prompt text capture is"),
2138 "unexpected error: {error}"
2139 );
2140 assert_eq!(text.len(), PROMPT_TEXT_CAPTURE_LIMIT_BYTES - 3);
2141 }
2142
2143 #[test]
2144 fn prompt_chunk_limit_rejects_more_tracked_chunks() {
2145 let chunk = notification(json!({
2146 "sessionUpdate": "agent_message_chunk",
2147 "content": { "text": "x" },
2148 }));
2149 let mut delivered_chunks = PROMPT_DELIVERED_CHUNK_LIMIT;
2150 let mut text = String::new();
2151 let error = accumulate_agent_message_chunk(&chunk, &mut delivered_chunks, &mut text)
2152 .expect_err("chunk should exceed chunk tracking cap");
2153 assert!(
2154 error
2155 .to_string()
2156 .contains("prompt chunk tracking limit exceeded"),
2157 "unexpected error: {error}"
2158 );
2159 assert!(text.is_empty());
2160 }
2161
2162 #[test]
2163 fn pending_session_request_count_tracks_registered_resolvers() {
2164 let (event_tx, _) = tokio::sync::broadcast::channel(1);
2165 let (permission_tx, _) = tokio::sync::broadcast::channel(1);
2166 let entry = SessionEntry {
2167 agent_type: "pi".to_string(),
2168 modes: parking_lot::Mutex::new(None),
2169 config_options: parking_lot::Mutex::new(Vec::new()),
2170 capabilities: parking_lot::Mutex::new(None),
2171 agent_info: parking_lot::Mutex::new(None),
2172 config_overrides: parking_lot::Mutex::new(BTreeMap::new()),
2173 event_tx,
2174 permission_tx,
2175 pending_permission_replies: scc::HashMap::new(),
2176 pending_session_request_lock: parking_lot::Mutex::new(()),
2177 pending_prompt_resolvers: scc::HashMap::new(),
2178 };
2179 let (first_tx, _first_rx) = tokio::sync::oneshot::channel();
2180 let (second_tx, _second_rx) = tokio::sync::oneshot::channel();
2181 let _ = entry.pending_prompt_resolvers.insert(1, first_tx);
2182 let _ = entry.pending_prompt_resolvers.insert(2, second_tx);
2183
2184 assert_eq!(pending_session_request_count(&entry), 2);
2185 }
2186}