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