1#![doc = include_str!("../README.md")]
2#![warn(missing_docs)]
3#![deny(rustdoc::broken_intra_doc_links)]
4#![cfg_attr(test, allow(clippy::unwrap_used))]
5
6pub mod canvas;
8mod canvas_dispatch;
9#[cfg(feature = "bundled-cli")]
11pub(crate) mod embeddedcli;
12mod errors;
13pub use errors::*;
14pub mod copilot_request_handler;
18#[doc(hidden)]
21pub mod github_telemetry;
22pub mod handler;
24pub mod hooks;
26mod jsonrpc;
27pub mod permission;
29pub mod provider_token;
31mod provider_token_dispatch;
32pub(crate) mod resolve;
34mod router;
35pub mod session;
37pub mod session_fs;
39mod session_fs_dispatch;
40pub mod subscription;
42pub mod tool;
44pub mod trace_context;
46pub mod transforms;
48pub mod types;
50mod wire;
51
52pub mod session_events;
54
55pub mod rpc;
58
59pub(crate) mod generated;
64
65pub mod mode;
68
69use std::ffi::OsString;
70use std::path::{Path, PathBuf};
71use std::process::Stdio;
72use std::sync::{Arc, OnceLock};
73use std::time::{Duration, Instant};
74
75use async_trait::async_trait;
76pub use indexmap::IndexMap;
80pub(crate) use jsonrpc::{
83 JsonRpcClient, JsonRpcError, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, error_codes,
84};
85pub use mode::{BUILTIN_TOOLS_ISOLATED, ClientMode, ToolSet};
86pub use provider_token::{BearerTokenError, BearerTokenProvider, ProviderTokenArgs};
87
88#[cfg(feature = "test-support")]
90pub mod test_support {
91 pub use crate::jsonrpc::{
92 JsonRpcClient, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse,
93 error_codes,
94 };
95}
96use serde::{Deserialize, Serialize};
97use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, BufReader};
98use tokio::net::TcpStream;
99use tokio::process::{Child, Command};
100use tokio::sync::{broadcast, mpsc, oneshot};
101use tracing::{Instrument, debug, error, info, warn};
102pub use types::*;
103
104mod sdk_protocol_version;
105pub use sdk_protocol_version::{SDK_PROTOCOL_VERSION, get_sdk_protocol_version};
106pub use subscription::{EventSubscription, LifecycleSubscription};
107
108const MIN_PROTOCOL_VERSION: u32 = 3;
110const RUNTIME_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
111
112#[derive(Debug, Default)]
114#[non_exhaustive]
115pub enum Transport {
116 #[default]
118 Stdio,
119 Tcp {
121 port: u16,
123 connection_token: Option<String>,
127 },
128 External {
130 host: String,
132 port: u16,
134 connection_token: Option<String>,
137 },
138}
139
140#[derive(Debug, Clone, Default)]
142pub enum CliProgram {
143 #[default]
146 Resolve,
147 Path(PathBuf),
149}
150
151impl From<PathBuf> for CliProgram {
152 fn from(path: PathBuf) -> Self {
153 Self::Path(path)
154 }
155}
156
157pub const HAS_BUNDLED_CLI: bool = cfg!(has_bundled_cli);
164
165pub fn install_bundled_cli() -> Option<PathBuf> {
189 #[cfg(feature = "bundled-cli")]
190 {
191 embeddedcli::path()
192 }
193 #[cfg(not(feature = "bundled-cli"))]
194 {
195 None
196 }
197}
198
199#[non_exhaustive]
209pub struct ClientOptions {
210 pub program: CliProgram,
212 pub prefix_args: Vec<OsString>,
214 pub working_directory: PathBuf,
216 pub env: Vec<(OsString, OsString)>,
218 pub env_remove: Vec<OsString>,
220 pub extra_args: Vec<String>,
222 pub transport: Transport,
224 pub github_token: Option<String>,
229 pub use_logged_in_user: Option<bool>,
233 pub log_level: Option<LogLevel>,
237 pub session_idle_timeout_seconds: Option<u64>,
243 pub on_list_models: Option<Arc<dyn ListModelsHandler>>,
251 pub session_fs: Option<SessionFsConfig>,
259 pub request_handler: Option<Arc<dyn crate::copilot_request_handler::CopilotRequestHandler>>,
268 #[doc(hidden)]
276 pub on_github_telemetry: Option<crate::github_telemetry::GitHubTelemetryCallback>,
277 pub on_get_trace_context: Option<Arc<dyn TraceContextProvider>>,
287 pub telemetry: Option<TelemetryConfig>,
291 pub base_directory: Option<PathBuf>,
296 pub enable_remote_sessions: bool,
302 pub bundled_cli_extract_dir: Option<PathBuf>,
321 pub mode: ClientMode,
325}
326
327impl std::fmt::Debug for ClientOptions {
328 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
329 f.debug_struct("ClientOptions")
330 .field("program", &self.program)
331 .field("prefix_args", &self.prefix_args)
332 .field("working_directory", &self.working_directory)
333 .field("env", &self.env)
334 .field("env_remove", &self.env_remove)
335 .field("extra_args", &self.extra_args)
336 .field("transport", &self.transport)
337 .field(
338 "github_token",
339 &self.github_token.as_ref().map(|_| "<redacted>"),
340 )
341 .field("use_logged_in_user", &self.use_logged_in_user)
342 .field("log_level", &self.log_level)
343 .field(
344 "session_idle_timeout_seconds",
345 &self.session_idle_timeout_seconds,
346 )
347 .field(
348 "on_list_models",
349 &self.on_list_models.as_ref().map(|_| "<set>"),
350 )
351 .field("session_fs", &self.session_fs)
352 .field(
353 "request_handler",
354 &self.request_handler.as_ref().map(|_| "<set>"),
355 )
356 .field(
357 "on_github_telemetry",
358 &self.on_github_telemetry.as_ref().map(|_| "<set>"),
359 )
360 .field(
361 "on_get_trace_context",
362 &self.on_get_trace_context.as_ref().map(|_| "<set>"),
363 )
364 .field("telemetry", &self.telemetry)
365 .field("base_directory", &self.base_directory)
366 .field("enable_remote_sessions", &self.enable_remote_sessions)
367 .field("bundled_cli_extract_dir", &self.bundled_cli_extract_dir)
368 .finish()
369 }
370}
371
372#[async_trait]
381pub trait ListModelsHandler: Send + Sync + 'static {
382 async fn list_models(&self) -> Result<Vec<Model>>;
384}
385
386#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
388#[serde(rename_all = "lowercase")]
389pub enum LogLevel {
390 None,
392 Error,
394 Warning,
396 Info,
398 Debug,
400 All,
402}
403
404impl LogLevel {
405 pub fn as_str(self) -> &'static str {
407 match self {
408 Self::None => "none",
409 Self::Error => "error",
410 Self::Warning => "warning",
411 Self::Info => "info",
412 Self::Debug => "debug",
413 Self::All => "all",
414 }
415 }
416}
417
418impl std::fmt::Display for LogLevel {
419 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
420 f.write_str(self.as_str())
421 }
422}
423
424#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
429#[serde(rename_all = "kebab-case")]
430#[non_exhaustive]
431pub enum OtelExporterType {
432 OtlpHttp,
435 File,
438}
439
440impl OtelExporterType {
441 pub fn as_str(self) -> &'static str {
443 match self {
444 Self::OtlpHttp => "otlp-http",
445 Self::File => "file",
446 }
447 }
448}
449
450#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
456#[non_exhaustive]
457pub enum OtlpHttpProtocol {
458 #[serde(rename = "http/json")]
460 HttpJson,
461 #[serde(rename = "http/protobuf")]
463 HttpProtobuf,
464}
465
466impl OtlpHttpProtocol {
467 pub fn as_str(self) -> &'static str {
469 match self {
470 Self::HttpJson => "http/json",
471 Self::HttpProtobuf => "http/protobuf",
472 }
473 }
474}
475
476#[derive(Debug, Clone, Default)]
511#[non_exhaustive]
512pub struct TelemetryConfig {
513 pub otlp_endpoint: Option<String>,
515 pub otlp_protocol: Option<OtlpHttpProtocol>,
517 pub file_path: Option<PathBuf>,
519 pub exporter_type: Option<OtelExporterType>,
522 pub source_name: Option<String>,
526 pub capture_content: Option<bool>,
530}
531
532impl TelemetryConfig {
533 pub fn new() -> Self {
536 Self::default()
537 }
538
539 pub fn with_otlp_endpoint(mut self, endpoint: impl Into<String>) -> Self {
541 self.otlp_endpoint = Some(endpoint.into());
542 self
543 }
544
545 pub fn with_otlp_protocol(mut self, protocol: OtlpHttpProtocol) -> Self {
547 self.otlp_protocol = Some(protocol);
548 self
549 }
550
551 pub fn with_file_path(mut self, path: impl Into<PathBuf>) -> Self {
553 self.file_path = Some(path.into());
554 self
555 }
556
557 pub fn with_exporter_type(mut self, exporter_type: OtelExporterType) -> Self {
559 self.exporter_type = Some(exporter_type);
560 self
561 }
562
563 pub fn with_source_name(mut self, source_name: impl Into<String>) -> Self {
567 self.source_name = Some(source_name.into());
568 self
569 }
570
571 pub fn with_capture_content(mut self, capture: bool) -> Self {
575 self.capture_content = Some(capture);
576 self
577 }
578
579 pub fn is_empty(&self) -> bool {
582 self.otlp_endpoint.is_none()
583 && self.otlp_protocol.is_none()
584 && self.file_path.is_none()
585 && self.exporter_type.is_none()
586 && self.source_name.is_none()
587 && self.capture_content.is_none()
588 }
589}
590
591impl Default for ClientOptions {
592 fn default() -> Self {
593 Self {
594 program: CliProgram::Resolve,
595 prefix_args: Vec::new(),
596 working_directory: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
597 env: Vec::new(),
598 env_remove: Vec::new(),
599 extra_args: Vec::new(),
600 transport: Transport::default(),
601 github_token: None,
602 use_logged_in_user: None,
603 log_level: None,
604 session_idle_timeout_seconds: None,
605 on_list_models: None,
606 session_fs: None,
607 request_handler: None,
608 on_github_telemetry: None,
609 on_get_trace_context: None,
610 telemetry: None,
611 base_directory: None,
612 enable_remote_sessions: false,
613 bundled_cli_extract_dir: None,
614 mode: ClientMode::default(),
615 }
616 }
617}
618
619impl ClientOptions {
620 pub fn new() -> Self {
636 Self::default()
637 }
638
639 pub fn with_program(mut self, program: impl Into<CliProgram>) -> Self {
641 self.program = program.into();
642 self
643 }
644
645 pub fn with_prefix_args<I, S>(mut self, args: I) -> Self
647 where
648 I: IntoIterator<Item = S>,
649 S: Into<OsString>,
650 {
651 self.prefix_args = args.into_iter().map(Into::into).collect();
652 self
653 }
654
655 pub fn with_cwd(mut self, cwd: impl Into<PathBuf>) -> Self {
657 self.working_directory = cwd.into();
658 self
659 }
660
661 pub fn with_env<I, K, V>(mut self, env: I) -> Self
663 where
664 I: IntoIterator<Item = (K, V)>,
665 K: Into<OsString>,
666 V: Into<OsString>,
667 {
668 self.env = env.into_iter().map(|(k, v)| (k.into(), v.into())).collect();
669 self
670 }
671
672 pub fn with_env_remove<I, S>(mut self, names: I) -> Self
674 where
675 I: IntoIterator<Item = S>,
676 S: Into<OsString>,
677 {
678 self.env_remove = names.into_iter().map(Into::into).collect();
679 self
680 }
681
682 pub fn with_extra_args<I, S>(mut self, args: I) -> Self
684 where
685 I: IntoIterator<Item = S>,
686 S: Into<String>,
687 {
688 self.extra_args = args.into_iter().map(Into::into).collect();
689 self
690 }
691
692 pub fn with_transport(mut self, transport: Transport) -> Self {
694 self.transport = transport;
695 self
696 }
697
698 pub fn with_github_token(mut self, token: impl Into<String>) -> Self {
701 self.github_token = Some(token.into());
702 self
703 }
704
705 pub fn with_use_logged_in_user(mut self, use_logged_in: bool) -> Self {
708 self.use_logged_in_user = Some(use_logged_in);
709 self
710 }
711
712 pub fn with_log_level(mut self, level: LogLevel) -> Self {
714 self.log_level = Some(level);
715 self
716 }
717
718 pub fn with_session_idle_timeout_seconds(mut self, seconds: u64) -> Self {
721 self.session_idle_timeout_seconds = Some(seconds);
722 self
723 }
724
725 pub fn with_list_models_handler<H>(mut self, handler: H) -> Self
728 where
729 H: ListModelsHandler + 'static,
730 {
731 self.on_list_models = Some(Arc::new(handler));
732 self
733 }
734
735 pub fn with_session_fs(mut self, config: SessionFsConfig) -> Self {
737 self.session_fs = Some(config);
738 self
739 }
740
741 pub fn with_request_handler<H>(mut self, handler: H) -> Self
746 where
747 H: crate::copilot_request_handler::CopilotRequestHandler,
748 {
749 self.request_handler = Some(Arc::new(handler));
750 self
751 }
752
753 #[doc(hidden)]
759 pub fn with_on_github_telemetry<F>(mut self, callback: F) -> Self
760 where
761 F: Fn(crate::github_telemetry::GitHubTelemetryNotification) + Send + Sync + 'static,
762 {
763 self.on_github_telemetry = Some(Arc::new(callback));
764 self
765 }
766
767 pub fn with_trace_context_provider<P>(mut self, provider: P) -> Self
771 where
772 P: TraceContextProvider + 'static,
773 {
774 self.on_get_trace_context = Some(Arc::new(provider));
775 self
776 }
777
778 pub fn with_telemetry(mut self, config: TelemetryConfig) -> Self {
780 self.telemetry = Some(config);
781 self
782 }
783
784 pub fn with_base_directory(mut self, dir: impl Into<PathBuf>) -> Self {
787 self.base_directory = Some(dir.into());
788 self
789 }
790
791 pub fn with_enable_remote_sessions(mut self, enabled: bool) -> Self {
794 self.enable_remote_sessions = enabled;
795 self
796 }
797
798 pub fn with_bundled_cli_extract_dir(mut self, dir: impl Into<PathBuf>) -> Self {
808 self.bundled_cli_extract_dir = Some(dir.into());
809 self
810 }
811
812 pub fn with_mode(mut self, mode: ClientMode) -> Self {
817 self.mode = mode;
818 self
819 }
820}
821
822fn validate_session_fs_config(cfg: &SessionFsConfig) -> Result<()> {
824 if cfg.initial_cwd.trim().is_empty() {
825 return Err(Error::with_message(
826 ErrorKind::Session(SessionErrorKind::InvalidSessionFsConfig),
827 "invalid SessionFsConfig: initial_cwd must not be empty",
828 ));
829 }
830 if cfg.session_state_path.trim().is_empty() {
831 return Err(Error::with_message(
832 ErrorKind::Session(SessionErrorKind::InvalidSessionFsConfig),
833 "invalid SessionFsConfig: session_state_path must not be empty",
834 ));
835 }
836 Ok(())
837}
838
839fn generate_connection_token() -> String {
846 let mut bytes = [0u8; 16];
847 getrandom::getrandom(&mut bytes)
848 .expect("OS CSPRNG (getrandom) is unavailable; cannot generate connection token");
849 let mut hex = String::with_capacity(32);
850 for byte in bytes {
851 use std::fmt::Write;
852 let _ = write!(hex, "{byte:02x}");
853 }
854 hex
855}
856
857#[derive(Clone)]
862pub struct Client {
863 inner: Arc<ClientInner>,
864}
865
866impl std::fmt::Debug for Client {
867 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
868 f.debug_struct("Client")
869 .field("working_directory", &self.inner.cwd)
870 .field("pid", &self.pid())
871 .finish()
872 }
873}
874
875struct ClientInner {
876 child: parking_lot::Mutex<Option<Child>>,
877 rpc: JsonRpcClient,
878 cwd: PathBuf,
879 request_rx: parking_lot::Mutex<Option<mpsc::UnboundedReceiver<JsonRpcRequest>>>,
880 notification_tx: broadcast::Sender<JsonRpcNotification>,
881 router: router::SessionRouter,
882 negotiated_protocol_version: OnceLock<u32>,
883 state: parking_lot::Mutex<ConnectionState>,
884 lifecycle_tx: broadcast::Sender<SessionLifecycleEvent>,
885 on_list_models: Option<Arc<dyn ListModelsHandler>>,
886 models_cache: parking_lot::Mutex<Arc<tokio::sync::OnceCell<Vec<Model>>>>,
887 session_fs_configured: bool,
888 session_fs_sqlite_declared: bool,
889 llm_inference: OnceLock<Arc<copilot_request_handler::CopilotRequestDispatcher>>,
892 on_github_telemetry: Option<crate::github_telemetry::GitHubTelemetryCallback>,
897 on_get_trace_context: Option<Arc<dyn TraceContextProvider>>,
898 effective_connection_token: Option<String>,
903 pub(crate) mode: ClientMode,
906}
907
908impl Client {
909 pub async fn start(options: ClientOptions) -> Result<Self> {
922 let start_time = Instant::now();
923 if options.mode == ClientMode::Empty
924 && options.base_directory.is_none()
925 && options.session_fs.is_none()
926 {
927 return Err(Error::with_message(
928 ErrorKind::InvalidConfig,
929 "ClientMode::Empty requires either `base_directory` or \
930 `session_fs` to be set (no implicit ~/.copilot fallback).",
931 ));
932 }
933 if let Some(cfg) = &options.session_fs {
934 validate_session_fs_config(cfg)?;
935 }
936 if matches!(options.transport, Transport::External { .. }) {
939 if options.github_token.is_some() {
940 return Err(Error::with_message(
941 ErrorKind::InvalidConfig,
942 "invalid client configuration: github_token cannot be used with \
943 Transport::External (external server manages its own auth)",
944 ));
945 }
946 if options.use_logged_in_user == Some(true) {
947 return Err(Error::with_message(
948 ErrorKind::InvalidConfig,
949 "invalid client configuration: use_logged_in_user cannot be used with \
950 Transport::External (external server manages its own auth)",
951 ));
952 }
953 }
954 match &options.transport {
958 Transport::Tcp {
959 connection_token: Some(t),
960 ..
961 }
962 | Transport::External {
963 connection_token: Some(t),
964 ..
965 } if t.is_empty() => {
966 return Err(Error::with_message(
967 ErrorKind::InvalidConfig,
968 "invalid client configuration: connection_token must be a non-empty string",
969 ));
970 }
971 _ => {}
972 }
973 let mut options = options;
978 let effective_connection_token: Option<String> = match &mut options.transport {
979 Transport::Stdio => None,
980 Transport::Tcp {
981 connection_token, ..
982 } => Some(
983 connection_token
984 .get_or_insert_with(generate_connection_token)
985 .clone(),
986 ),
987 Transport::External {
988 connection_token, ..
989 } => connection_token.clone(),
990 };
991 let session_fs_config = options.session_fs.clone();
992 let request_handler = options.request_handler.clone();
993 let session_fs_sqlite_declared = session_fs_config
994 .as_ref()
995 .and_then(|c| c.capabilities.as_ref())
996 .is_some_and(|caps| caps.sqlite);
997 let program = match &options.program {
998 CliProgram::Path(path) => {
999 info!(path = %path.display(), "using explicit copilot CLI path");
1000 path.clone()
1001 }
1002 CliProgram::Resolve => {
1003 let resolved = resolve::copilot_binary_with_extract_dir(
1004 options.bundled_cli_extract_dir.as_deref(),
1005 )?;
1006 info!(path = %resolved.display(), "resolved copilot CLI");
1007 #[cfg(windows)]
1008 {
1009 if let Some(ext) = resolved.extension().and_then(|e| e.to_str()).filter(|ext| {
1010 ext.eq_ignore_ascii_case("cmd") || ext.eq_ignore_ascii_case("bat")
1011 }) {
1012 warn!(
1013 path = %resolved.display(),
1014 ext = %ext,
1015 "resolved copilot CLI is a .cmd/.bat wrapper; \
1016 this may cause console window flashes on Windows"
1017 );
1018 }
1019 }
1020 resolved
1021 }
1022 };
1023
1024 let client = match options.transport {
1025 Transport::External {
1026 ref host,
1027 port,
1028 connection_token: _,
1029 } => {
1030 info!(host = %host, port = %port, "connecting to external CLI server");
1031 let connect_start = Instant::now();
1032 let stream = TcpStream::connect((host.as_str(), port)).await?;
1033 debug!(
1034 elapsed_ms = connect_start.elapsed().as_millis(),
1035 host = %host,
1036 port,
1037 "Client::start TCP connect complete"
1038 );
1039 let (reader, writer) = tokio::io::split(stream);
1040 Self::from_transport(
1041 reader,
1042 writer,
1043 None,
1044 options.working_directory,
1045 options.on_list_models,
1046 session_fs_config.is_some(),
1047 session_fs_sqlite_declared,
1048 options.on_get_trace_context,
1049 options.on_github_telemetry,
1050 effective_connection_token.clone(),
1051 options.mode,
1052 )?
1053 }
1054 Transport::Tcp {
1055 port,
1056 connection_token: _,
1057 } => {
1058 let (mut child, actual_port) = Self::spawn_tcp(&program, &options, port).await?;
1059 let connect_start = Instant::now();
1060 let stream = TcpStream::connect(("127.0.0.1", actual_port)).await?;
1061 debug!(
1062 elapsed_ms = connect_start.elapsed().as_millis(),
1063 port = actual_port,
1064 "Client::start TCP connect complete"
1065 );
1066 let (reader, writer) = tokio::io::split(stream);
1067 Self::drain_stderr(&mut child);
1068 Self::from_transport(
1069 reader,
1070 writer,
1071 Some(child),
1072 options.working_directory,
1073 options.on_list_models,
1074 session_fs_config.is_some(),
1075 session_fs_sqlite_declared,
1076 options.on_get_trace_context,
1077 options.on_github_telemetry,
1078 effective_connection_token.clone(),
1079 options.mode,
1080 )?
1081 }
1082 Transport::Stdio => {
1083 let mut child = Self::spawn_stdio(&program, &options)?;
1084 let stdin = child.stdin.take().expect("stdin is piped");
1085 let stdout = child.stdout.take().expect("stdout is piped");
1086 Self::drain_stderr(&mut child);
1087 Self::from_transport(
1088 stdout,
1089 stdin,
1090 Some(child),
1091 options.working_directory,
1092 options.on_list_models,
1093 session_fs_config.is_some(),
1094 session_fs_sqlite_declared,
1095 options.on_get_trace_context,
1096 options.on_github_telemetry,
1097 effective_connection_token.clone(),
1098 options.mode,
1099 )?
1100 }
1101 };
1102
1103 debug!(
1104 elapsed_ms = start_time.elapsed().as_millis(),
1105 "Client::start transport setup complete"
1106 );
1107 client.verify_protocol_version().await?;
1108 debug!(
1109 elapsed_ms = start_time.elapsed().as_millis(),
1110 "Client::start protocol verification complete"
1111 );
1112 if let Some(cfg) = session_fs_config {
1113 let session_fs_start = Instant::now();
1114 let capabilities = cfg.capabilities.as_ref().map(|c| {
1115 crate::generated::api_types::SessionFsSetProviderCapabilities {
1116 sqlite: Some(c.sqlite),
1117 }
1118 });
1119 let request = crate::generated::api_types::SessionFsSetProviderRequest {
1120 capabilities,
1121 conventions: cfg.conventions.into_wire(),
1122 initial_cwd: cfg.initial_cwd,
1123 session_state_path: cfg.session_state_path,
1124 };
1125 client.rpc().session_fs().set_provider(request).await?;
1126 debug!(
1127 elapsed_ms = session_fs_start.elapsed().as_millis(),
1128 "Client::start session filesystem setup complete"
1129 );
1130 }
1131 if let Some(handler) = request_handler {
1132 let llm_inference_start = Instant::now();
1133 let dispatcher = Arc::new(copilot_request_handler::CopilotRequestDispatcher::new(
1134 handler,
1135 ));
1136 dispatcher.set_client(Arc::downgrade(&client.inner));
1137 let _ = client.inner.llm_inference.set(dispatcher.clone());
1138 client.inner.router.ensure_started(
1141 &client.inner.notification_tx,
1142 &client.inner.request_rx,
1143 Some(dispatcher.clone()),
1144 client.inner.on_github_telemetry.clone(),
1145 );
1146 client.rpc().llm_inference().set_provider().await?;
1147 debug!(
1148 elapsed_ms = llm_inference_start.elapsed().as_millis(),
1149 "Client::start Copilot request handler registration complete"
1150 );
1151 }
1152 debug!(
1153 elapsed_ms = start_time.elapsed().as_millis(),
1154 "Client::start complete"
1155 );
1156 Ok(client)
1157 }
1158
1159 pub fn from_streams(
1163 reader: impl AsyncRead + Unpin + Send + 'static,
1164 writer: impl AsyncWrite + Unpin + Send + 'static,
1165 cwd: PathBuf,
1166 ) -> Result<Self> {
1167 Self::from_transport(
1168 reader,
1169 writer,
1170 None,
1171 cwd,
1172 None,
1173 false,
1174 false,
1175 None,
1176 None,
1177 None,
1178 ClientMode::default(),
1179 )
1180 }
1181
1182 #[cfg(any(test, feature = "test-support"))]
1190 pub fn from_streams_with_trace_provider(
1191 reader: impl AsyncRead + Unpin + Send + 'static,
1192 writer: impl AsyncWrite + Unpin + Send + 'static,
1193 cwd: PathBuf,
1194 provider: Arc<dyn TraceContextProvider>,
1195 ) -> Result<Self> {
1196 Self::from_transport(
1197 reader,
1198 writer,
1199 None,
1200 cwd,
1201 None,
1202 false,
1203 false,
1204 Some(provider),
1205 None,
1206 None,
1207 ClientMode::default(),
1208 )
1209 }
1210
1211 #[cfg(any(test, feature = "test-support"))]
1215 pub fn from_streams_with_connection_token(
1216 reader: impl AsyncRead + Unpin + Send + 'static,
1217 writer: impl AsyncWrite + Unpin + Send + 'static,
1218 cwd: PathBuf,
1219 token: Option<String>,
1220 ) -> Result<Self> {
1221 Self::from_transport(
1222 reader,
1223 writer,
1224 None,
1225 cwd,
1226 None,
1227 false,
1228 false,
1229 None,
1230 None,
1231 token,
1232 ClientMode::default(),
1233 )
1234 }
1235
1236 #[doc(hidden)]
1239 #[cfg(any(test, feature = "test-support"))]
1240 pub fn from_streams_with_github_telemetry(
1241 reader: impl AsyncRead + Unpin + Send + 'static,
1242 writer: impl AsyncWrite + Unpin + Send + 'static,
1243 cwd: PathBuf,
1244 on_github_telemetry: crate::github_telemetry::GitHubTelemetryCallback,
1245 ) -> Result<Self> {
1246 Self::from_transport(
1247 reader,
1248 writer,
1249 None,
1250 cwd,
1251 None,
1252 false,
1253 false,
1254 None,
1255 Some(on_github_telemetry),
1256 None,
1257 ClientMode::default(),
1258 )
1259 }
1260
1261 #[cfg(any(test, feature = "test-support"))]
1267 pub fn generate_connection_token_for_test() -> String {
1268 generate_connection_token()
1269 }
1270
1271 #[allow(clippy::too_many_arguments)]
1272 fn from_transport(
1273 reader: impl AsyncRead + Unpin + Send + 'static,
1274 writer: impl AsyncWrite + Unpin + Send + 'static,
1275 child: Option<Child>,
1276 cwd: PathBuf,
1277 on_list_models: Option<Arc<dyn ListModelsHandler>>,
1278 session_fs_configured: bool,
1279 session_fs_sqlite_declared: bool,
1280 on_get_trace_context: Option<Arc<dyn TraceContextProvider>>,
1281 on_github_telemetry: Option<crate::github_telemetry::GitHubTelemetryCallback>,
1282 effective_connection_token: Option<String>,
1283 mode: ClientMode,
1284 ) -> Result<Self> {
1285 let setup_start = Instant::now();
1286 let (request_tx, request_rx) = mpsc::unbounded_channel::<JsonRpcRequest>();
1287 let (notification_broadcast_tx, _) = broadcast::channel::<JsonRpcNotification>(1024);
1288 let rpc = JsonRpcClient::new(
1289 writer,
1290 reader,
1291 notification_broadcast_tx.clone(),
1292 request_tx,
1293 );
1294
1295 let pid = child.as_ref().and_then(|c| c.id());
1296 info!(pid = ?pid, "copilot CLI client ready");
1297
1298 let client = Self {
1299 inner: Arc::new(ClientInner {
1300 child: parking_lot::Mutex::new(child),
1301 rpc,
1302 cwd,
1303 request_rx: parking_lot::Mutex::new(Some(request_rx)),
1304 notification_tx: notification_broadcast_tx,
1305 router: router::SessionRouter::new(),
1306 negotiated_protocol_version: OnceLock::new(),
1307 state: parking_lot::Mutex::new(ConnectionState::Connected),
1308 lifecycle_tx: broadcast::channel(256).0,
1309 on_list_models,
1310 models_cache: parking_lot::Mutex::new(Arc::new(tokio::sync::OnceCell::new())),
1311 session_fs_configured,
1312 session_fs_sqlite_declared,
1313 llm_inference: OnceLock::new(),
1314 on_github_telemetry,
1315 on_get_trace_context,
1316 effective_connection_token,
1317 mode,
1318 }),
1319 };
1320 client.spawn_lifecycle_dispatcher();
1321 debug!(
1322 elapsed_ms = setup_start.elapsed().as_millis(),
1323 pid = ?pid,
1324 "Client::from_transport setup complete"
1325 );
1326 Ok(client)
1327 }
1328
1329 fn spawn_lifecycle_dispatcher(&self) {
1333 let inner = Arc::clone(&self.inner);
1334 let mut notif_rx = inner.notification_tx.subscribe();
1335 tokio::spawn(async move {
1336 loop {
1337 match notif_rx.recv().await {
1338 Ok(notification) => {
1339 if notification.method != "session.lifecycle" {
1340 continue;
1341 }
1342 let Some(params) = notification.params.as_ref() else {
1343 continue;
1344 };
1345 let event: SessionLifecycleEvent =
1346 match serde_json::from_value(params.clone()) {
1347 Ok(e) => e,
1348 Err(e) => {
1349 warn!(
1350 error = %e,
1351 "failed to deserialize session.lifecycle notification"
1352 );
1353 continue;
1354 }
1355 };
1356 let _ = inner.lifecycle_tx.send(event);
1359 }
1360 Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
1361 warn!(missed = n, "lifecycle dispatcher lagged");
1362 }
1363 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
1364 }
1365 }
1366 });
1367 }
1368
1369 fn build_command(program: &Path, options: &ClientOptions) -> Command {
1370 let mut command = Command::new(program);
1371 for arg in &options.prefix_args {
1372 command.arg(arg);
1373 }
1374 if let Some(token) = &options.github_token {
1377 command.env("COPILOT_SDK_AUTH_TOKEN", token);
1378 }
1379 if let Some(telemetry) = &options.telemetry {
1382 command.env("COPILOT_OTEL_ENABLED", "true");
1383 if let Some(endpoint) = &telemetry.otlp_endpoint {
1384 command.env("OTEL_EXPORTER_OTLP_ENDPOINT", endpoint);
1385 }
1386 if let Some(protocol) = telemetry.otlp_protocol {
1387 command.env("OTEL_EXPORTER_OTLP_PROTOCOL", protocol.as_str());
1388 }
1389 if let Some(path) = &telemetry.file_path {
1390 command.env("COPILOT_OTEL_FILE_EXPORTER_PATH", path);
1391 }
1392 if let Some(exporter) = telemetry.exporter_type {
1393 command.env("COPILOT_OTEL_EXPORTER_TYPE", exporter.as_str());
1394 }
1395 if let Some(source) = &telemetry.source_name {
1396 command.env("COPILOT_OTEL_SOURCE_NAME", source);
1397 }
1398 if let Some(capture) = telemetry.capture_content {
1399 command.env(
1400 "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",
1401 if capture { "true" } else { "false" },
1402 );
1403 }
1404 }
1405 if let Some(dir) = &options.base_directory {
1406 command.env("COPILOT_HOME", dir);
1407 }
1408 if options.mode == ClientMode::Empty {
1411 command.env("COPILOT_DISABLE_KEYTAR", "1");
1412 }
1413 if let Transport::Tcp {
1414 connection_token: Some(token),
1415 ..
1416 } = &options.transport
1417 {
1418 command.env("COPILOT_CONNECTION_TOKEN", token);
1419 }
1420 for (key, value) in &options.env {
1421 command.env(key, value);
1422 }
1423 for key in &options.env_remove {
1424 command.env_remove(key);
1425 }
1426 command
1427 .current_dir(&options.working_directory)
1428 .stdout(Stdio::piped())
1429 .stderr(Stdio::piped());
1430
1431 #[cfg(windows)]
1432 {
1433 use std::os::windows::process::CommandExt;
1434 const CREATE_NO_WINDOW: u32 = 0x08000000;
1435 command.as_std_mut().creation_flags(CREATE_NO_WINDOW);
1436 }
1437
1438 command
1439 }
1440
1441 fn auth_args(options: &ClientOptions) -> Vec<&'static str> {
1449 let mut args: Vec<&'static str> = Vec::new();
1450 if options.github_token.is_some() {
1451 args.push("--auth-token-env");
1452 args.push("COPILOT_SDK_AUTH_TOKEN");
1453 }
1454 let use_logged_in = options
1455 .use_logged_in_user
1456 .unwrap_or(options.github_token.is_none());
1457 if !use_logged_in {
1458 args.push("--no-auto-login");
1459 }
1460 args
1461 }
1462
1463 fn session_idle_timeout_args(options: &ClientOptions) -> Vec<String> {
1467 match options.session_idle_timeout_seconds {
1468 Some(secs) if secs > 0 => {
1469 vec!["--session-idle-timeout".to_string(), secs.to_string()]
1470 }
1471 _ => Vec::new(),
1472 }
1473 }
1474
1475 fn remote_args(options: &ClientOptions) -> Vec<String> {
1476 if options.enable_remote_sessions {
1477 vec!["--remote".to_string()]
1478 } else {
1479 Vec::new()
1480 }
1481 }
1482
1483 fn log_level_args(options: &ClientOptions) -> Vec<&'static str> {
1484 match options.log_level {
1485 Some(level) => vec!["--log-level", level.as_str()],
1486 None => Vec::new(),
1487 }
1488 }
1489
1490 fn spawn_stdio(program: &Path, options: &ClientOptions) -> Result<Child> {
1491 info!(cwd = ?options.working_directory, program = %program.display(), "spawning copilot CLI (stdio)");
1492 let mut command = Self::build_command(program, options);
1493 command
1494 .args(["--server", "--stdio", "--no-auto-update"])
1495 .args(Self::log_level_args(options))
1496 .args(Self::auth_args(options))
1497 .args(Self::session_idle_timeout_args(options))
1498 .args(Self::remote_args(options))
1499 .args(&options.extra_args)
1500 .stdin(Stdio::piped());
1501 let spawn_start = Instant::now();
1502 let child = command.spawn()?;
1503 debug!(
1504 elapsed_ms = spawn_start.elapsed().as_millis(),
1505 "Client::spawn_stdio subprocess spawned"
1506 );
1507 Ok(child)
1508 }
1509
1510 async fn spawn_tcp(program: &Path, options: &ClientOptions, port: u16) -> Result<(Child, u16)> {
1511 info!(cwd = ?options.working_directory, program = %program.display(), port = %port, "spawning copilot CLI (tcp)");
1512 let mut command = Self::build_command(program, options);
1513 command
1514 .args(["--server", "--port", &port.to_string(), "--no-auto-update"])
1515 .args(Self::log_level_args(options))
1516 .args(Self::auth_args(options))
1517 .args(Self::session_idle_timeout_args(options))
1518 .args(Self::remote_args(options))
1519 .args(&options.extra_args)
1520 .stdin(Stdio::null());
1521 let spawn_start = Instant::now();
1522 let mut child = command.spawn()?;
1523 debug!(
1524 elapsed_ms = spawn_start.elapsed().as_millis(),
1525 "Client::spawn_tcp subprocess spawned"
1526 );
1527 let stdout = child.stdout.take().expect("stdout is piped");
1528
1529 let (port_tx, port_rx) = oneshot::channel::<u16>();
1530 let span = tracing::error_span!("copilot_cli_port_scan");
1531 tokio::spawn(
1532 async move {
1533 let port_re = regex::Regex::new(r"listening on port (\d+)").expect("valid regex");
1535 let mut lines = BufReader::new(stdout).lines();
1536 let mut port_tx = Some(port_tx);
1537 while let Ok(Some(line)) = lines.next_line().await {
1538 debug!(line = %line, "CLI stdout");
1539 if let Some(tx) = port_tx.take() {
1540 if let Some(caps) = port_re.captures(&line)
1541 && let Some(p) =
1542 caps.get(1).and_then(|m| m.as_str().parse::<u16>().ok())
1543 {
1544 let _ = tx.send(p);
1545 continue;
1546 }
1547 port_tx = Some(tx);
1549 }
1550 }
1551 }
1552 .instrument(span),
1553 );
1554
1555 let port_wait_start = Instant::now();
1556 let actual_port = tokio::time::timeout(std::time::Duration::from_secs(10), port_rx)
1557 .await
1558 .map_err(|_| Error::from(ErrorKind::Protocol(ProtocolErrorKind::CliStartupTimeout)))?
1559 .map_err(|_| Error::from(ErrorKind::Protocol(ProtocolErrorKind::CliStartupFailed)))?;
1560
1561 debug!(
1562 elapsed_ms = port_wait_start.elapsed().as_millis(),
1563 port = actual_port,
1564 "Client::spawn_tcp TCP port wait complete"
1565 );
1566 info!(port = %actual_port, "CLI server listening");
1567 Ok((child, actual_port))
1568 }
1569
1570 fn drain_stderr(child: &mut Child) {
1571 if let Some(stderr) = child.stderr.take() {
1572 let span = tracing::error_span!("copilot_cli");
1573 tokio::spawn(
1574 async move {
1575 let mut reader = BufReader::new(stderr).lines();
1576 while let Ok(Some(line)) = reader.next_line().await {
1577 warn!(line = %line, "CLI stderr");
1578 }
1579 }
1580 .instrument(span),
1581 );
1582 }
1583 }
1584
1585 pub fn cwd(&self) -> &PathBuf {
1587 &self.inner.cwd
1588 }
1589
1590 pub fn mode(&self) -> ClientMode {
1592 self.inner.mode
1593 }
1594
1595 pub fn rpc(&self) -> crate::generated::rpc::ClientRpc<'_> {
1606 crate::generated::rpc::ClientRpc { client: self }
1607 }
1608
1609 #[allow(dead_code, reason = "convenience for future internal use")]
1611 pub(crate) async fn send_request(
1612 &self,
1613 method: &str,
1614 params: Option<serde_json::Value>,
1615 ) -> Result<JsonRpcResponse> {
1616 self.inner.rpc.send_request(method, params).await
1617 }
1618
1619 pub async fn call(
1639 &self,
1640 method: &str,
1641 params: Option<serde_json::Value>,
1642 ) -> Result<serde_json::Value> {
1643 self.call_with_inline_callback(method, params, None).await
1644 }
1645
1646 pub(crate) async fn call_with_inline_callback(
1661 &self,
1662 method: &str,
1663 params: Option<serde_json::Value>,
1664 inline_callback: Option<crate::jsonrpc::InlineResponseCallback>,
1665 ) -> Result<serde_json::Value> {
1666 let session_id: Option<SessionId> = params
1667 .as_ref()
1668 .and_then(|p| p.get("sessionId"))
1669 .and_then(|v| v.as_str())
1670 .map(SessionId::from);
1671 let response = self
1672 .inner
1673 .rpc
1674 .send_request_with_inline_callback(method, params, inline_callback)
1675 .await?;
1676 if let Some(err) = response.error {
1677 if err.message.contains("Session not found") {
1678 return Err(ErrorKind::Session(SessionErrorKind::NotFound(
1679 session_id.unwrap_or_else(|| "unknown".into()),
1680 ))
1681 .into());
1682 }
1683 return Err(Error::with_message(
1684 ErrorKind::Rpc { code: err.code },
1685 err.message,
1686 ));
1687 }
1688 Ok(response.result.unwrap_or(serde_json::Value::Null))
1689 }
1690
1691 pub(crate) async fn send_response(&self, response: &JsonRpcResponse) -> Result<()> {
1693 self.inner.rpc.write(response).await
1694 }
1695
1696 pub(crate) fn from_inner(inner: Arc<ClientInner>) -> Self {
1698 Self { inner }
1699 }
1700
1701 #[expect(dead_code, reason = "reserved for future pub(crate) use")]
1705 pub(crate) fn take_request_rx(&self) -> Option<mpsc::UnboundedReceiver<JsonRpcRequest>> {
1706 self.inner.request_rx.lock().take()
1707 }
1708
1709 pub(crate) fn register_session(
1717 &self,
1718 session_id: &SessionId,
1719 ) -> crate::router::SessionChannels {
1720 self.inner.router.ensure_started(
1721 &self.inner.notification_tx,
1722 &self.inner.request_rx,
1723 self.inner.llm_inference.get().cloned(),
1724 self.inner.on_github_telemetry.clone(),
1725 );
1726 self.inner.router.register(session_id)
1727 }
1728
1729 pub(crate) fn unregister_session(&self, session_id: &SessionId) {
1731 self.inner.router.unregister(session_id);
1732 }
1733
1734 pub fn protocol_version(&self) -> Option<u32> {
1741 self.inner.negotiated_protocol_version.get().copied()
1742 }
1743
1744 pub async fn verify_protocol_version(&self) -> Result<()> {
1768 let handshake_start = Instant::now();
1769 let mut used_fallback_ping = false;
1770 let server_version = match self.connect_handshake().await {
1774 Ok(v) => v,
1775 Err(ref e) if e.rpc_code() == Some(error_codes::METHOD_NOT_FOUND) => {
1776 used_fallback_ping = true;
1777 self.ping(None).await?.protocol_version
1778 }
1779 Err(e) => return Err(e),
1780 };
1781
1782 match server_version {
1783 None => {
1784 warn!("CLI server did not report protocolVersion; skipping version check");
1785 }
1786 Some(v) if !(MIN_PROTOCOL_VERSION..=SDK_PROTOCOL_VERSION).contains(&v) => {
1787 return Err(ErrorKind::Protocol(ProtocolErrorKind::VersionMismatch {
1788 server: v,
1789 min: MIN_PROTOCOL_VERSION,
1790 max: SDK_PROTOCOL_VERSION,
1791 })
1792 .into());
1793 }
1794 Some(v) => {
1795 if let Some(&existing) = self.inner.negotiated_protocol_version.get() {
1796 if existing != v {
1797 return Err(ErrorKind::Protocol(ProtocolErrorKind::VersionChanged {
1798 previous: existing,
1799 current: v,
1800 })
1801 .into());
1802 }
1803 } else {
1804 let _ = self.inner.negotiated_protocol_version.set(v);
1805 }
1806 }
1807 }
1808
1809 debug!(
1810 elapsed_ms = handshake_start.elapsed().as_millis(),
1811 protocol_version = ?server_version,
1812 used_fallback_ping,
1813 "Client::verify_protocol_version protocol handshake complete"
1814 );
1815 Ok(())
1816 }
1817
1818 async fn connect_handshake(&self) -> Result<Option<u32>> {
1825 let params = crate::generated::api_types::ConnectRequest {
1826 token: self.inner.effective_connection_token.clone(),
1827 enable_git_hub_telemetry_forwarding: self
1828 .inner
1829 .on_github_telemetry
1830 .is_some()
1831 .then_some(true),
1832 };
1833 let value = self
1834 .call(
1835 crate::generated::api_types::rpc_methods::CONNECT,
1836 Some(serde_json::to_value(params)?),
1837 )
1838 .await?;
1839 let result: crate::generated::api_types::ConnectResult = serde_json::from_value(value)?;
1840 Ok(Some(u32::try_from(result.protocol_version).map_err(
1841 |_| ProtocolErrorKind::InvalidProtocolVersion {
1842 server: result.protocol_version,
1843 },
1844 )?))
1845 }
1846
1847 pub async fn ping(&self, message: Option<&str>) -> Result<crate::types::PingResponse> {
1855 let params = match message {
1856 Some(m) => serde_json::json!({ "message": m }),
1857 None => serde_json::json!({}),
1858 };
1859 let value = self
1860 .call(generated::api_types::rpc_methods::PING, Some(params))
1861 .await?;
1862 Ok(serde_json::from_value(value)?)
1863 }
1864
1865 pub async fn list_sessions(
1868 &self,
1869 filter: Option<SessionListFilter>,
1870 ) -> Result<Vec<SessionMetadata>> {
1871 let params = match filter {
1872 Some(f) => serde_json::json!({ "filter": f }),
1873 None => serde_json::json!({}),
1874 };
1875 let result = self.call("session.list", Some(params)).await?;
1876 let response: ListSessionsResponse = serde_json::from_value(result)?;
1877 Ok(response.sessions)
1878 }
1879
1880 pub async fn get_session_metadata(
1898 &self,
1899 session_id: &SessionId,
1900 ) -> Result<Option<SessionMetadata>> {
1901 let result = self
1902 .call(
1903 "session.getMetadata",
1904 Some(serde_json::json!({ "sessionId": session_id })),
1905 )
1906 .await?;
1907 let response: GetSessionMetadataResponse = serde_json::from_value(result)?;
1908 Ok(response.session)
1909 }
1910
1911 pub async fn delete_session(&self, session_id: &SessionId) -> Result<()> {
1913 self.call(
1914 "session.delete",
1915 Some(serde_json::json!({ "sessionId": session_id })),
1916 )
1917 .await?;
1918 Ok(())
1919 }
1920
1921 pub async fn get_last_session_id(&self) -> Result<Option<SessionId>> {
1937 let result = self
1938 .call("session.getLastId", Some(serde_json::json!({})))
1939 .await?;
1940 let response: GetLastSessionIdResponse = serde_json::from_value(result)?;
1941 Ok(response.session_id)
1942 }
1943
1944 pub async fn get_foreground_session_id(&self) -> Result<Option<SessionId>> {
1949 let result = self
1950 .call("session.getForeground", Some(serde_json::json!({})))
1951 .await?;
1952 let response: GetForegroundSessionResponse = serde_json::from_value(result)?;
1953 Ok(response.session_id)
1954 }
1955
1956 pub async fn set_foreground_session_id(&self, session_id: &SessionId) -> Result<()> {
1961 self.call(
1962 "session.setForeground",
1963 Some(serde_json::json!({ "sessionId": session_id })),
1964 )
1965 .await?;
1966 Ok(())
1967 }
1968
1969 pub async fn get_status(&self) -> Result<GetStatusResponse> {
1971 let result = self.call("status.get", Some(serde_json::json!({}))).await?;
1972 Ok(serde_json::from_value(result)?)
1973 }
1974
1975 pub async fn get_auth_status(&self) -> Result<GetAuthStatusResponse> {
1977 let result = self
1978 .call("auth.getStatus", Some(serde_json::json!({})))
1979 .await?;
1980 Ok(serde_json::from_value(result)?)
1981 }
1982
1983 pub async fn list_models(&self) -> Result<Vec<Model>> {
1988 let cache = self.inner.models_cache.lock().clone();
1989 let models = cache
1990 .get_or_try_init(|| async {
1991 if let Some(handler) = &self.inner.on_list_models {
1992 handler.list_models().await
1993 } else {
1994 Ok(self.rpc().models().list().await?.models)
1995 }
1996 })
1997 .await?;
1998 Ok(models.clone())
1999 }
2000
2001 pub(crate) async fn resolve_trace_context(&self) -> TraceContext {
2004 if let Some(provider) = &self.inner.on_get_trace_context {
2005 provider.get_trace_context().await
2006 } else {
2007 TraceContext::default()
2008 }
2009 }
2010
2011 pub fn pid(&self) -> Option<u32> {
2013 self.inner.child.lock().as_ref().and_then(|c| c.id())
2014 }
2015
2016 pub async fn stop(&self) -> std::result::Result<(), StopErrors> {
2043 let pid = self.pid();
2044 info!(pid = ?pid, "stopping CLI process");
2045 let mut errors: Vec<Error> = Vec::new();
2046
2047 for session_id in self.inner.router.session_ids() {
2050 match self
2051 .call(
2052 "session.destroy",
2053 Some(serde_json::json!({ "sessionId": session_id })),
2054 )
2055 .await
2056 {
2057 Ok(_) => {}
2058 Err(e) => {
2059 warn!(
2060 session_id = %session_id,
2061 error = %e,
2062 "session.destroy failed during Client::stop",
2063 );
2064 errors.push(e);
2065 }
2066 }
2067 self.inner.router.unregister(&session_id);
2068 }
2069
2070 let should_shutdown_runtime = self.inner.child.lock().is_some();
2071 if should_shutdown_runtime {
2072 let runtime_shutdown_start = Instant::now();
2073 match tokio::time::timeout(RUNTIME_SHUTDOWN_TIMEOUT, self.rpc().runtime().shutdown())
2074 .await
2075 {
2076 Ok(Ok(())) => {
2077 debug!(
2078 elapsed_ms = runtime_shutdown_start.elapsed().as_millis(),
2079 "Client::stop runtime shutdown complete"
2080 );
2081 }
2082 Ok(Err(e)) => {
2083 warn!(
2084 elapsed_ms = runtime_shutdown_start.elapsed().as_millis(),
2085 error = %e,
2086 "runtime.shutdown failed during Client::stop",
2087 );
2088 errors.push(e);
2089 }
2090 Err(_) => {
2091 let e = std::io::Error::new(
2092 std::io::ErrorKind::TimedOut,
2093 "runtime.shutdown timed out during Client::stop",
2094 );
2095 warn!(
2096 elapsed_ms = runtime_shutdown_start.elapsed().as_millis(),
2097 timeout = ?RUNTIME_SHUTDOWN_TIMEOUT,
2098 error = %e,
2099 "runtime.shutdown timed out during Client::stop",
2100 );
2101 errors.push(e.into());
2102 }
2103 }
2104 }
2105
2106 let child = self.inner.child.lock().take();
2107 *self.inner.state.lock() = ConnectionState::Disconnected;
2108 *self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new());
2109 if let Some(mut child) = child {
2110 match child.try_wait() {
2111 Ok(Some(_status)) => {}
2112 Ok(None) => {
2113 if let Err(e) = child.kill().await {
2120 errors.push(e.into());
2121 }
2122 }
2123 Err(e) => errors.push(e.into()),
2124 }
2125 }
2126
2127 info!(pid = ?pid, errors = errors.len(), "CLI process stopped");
2128 if errors.is_empty() {
2129 Ok(())
2130 } else {
2131 Err(StopErrors(errors))
2132 }
2133 }
2134
2135 pub fn force_stop(&self) {
2165 let pid = self.pid();
2166 info!(pid = ?pid, "force-stopping CLI process");
2167 if let Some(mut child) = self.inner.child.lock().take()
2168 && let Err(e) = child.start_kill()
2169 {
2170 error!(pid = ?pid, error = %e, "failed to send kill signal");
2171 }
2172 self.inner.rpc.force_close();
2173 self.inner.router.clear();
2176 *self.inner.state.lock() = ConnectionState::Disconnected;
2177 *self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new());
2178 }
2179
2180 pub fn subscribe_lifecycle(&self) -> LifecycleSubscription {
2215 LifecycleSubscription::new(self.inner.lifecycle_tx.subscribe())
2216 }
2217}
2218
2219impl Drop for ClientInner {
2220 fn drop(&mut self) {
2221 if let Some(ref mut child) = *self.child.lock() {
2222 let pid = child.id();
2223 if let Err(e) = child.start_kill() {
2224 error!(pid = ?pid, error = %e, "failed to kill CLI process on drop");
2225 } else {
2226 info!(pid = ?pid, "kill signal sent for CLI process on drop");
2227 }
2228 }
2229 }
2230}
2231
2232#[cfg(test)]
2233mod tests {
2234 use super::*;
2235
2236 #[test]
2237 fn is_transport_failure_matches_request_cancelled() {
2238 let err = Error::from(ErrorKind::Protocol(ProtocolErrorKind::RequestCancelled));
2239 assert!(err.is_transport_failure());
2240 }
2241
2242 #[test]
2243 fn is_transport_failure_matches_io_error() {
2244 let err = Error::from(std::io::Error::new(std::io::ErrorKind::BrokenPipe, "gone"));
2245 assert!(err.is_transport_failure());
2246 }
2247
2248 #[test]
2249 fn is_transport_failure_rejects_rpc_error() {
2250 let err = Error::with_message(ErrorKind::Rpc { code: -1 }, "bad");
2251 assert!(!err.is_transport_failure());
2252 }
2253
2254 #[test]
2255 fn is_transport_failure_rejects_session_error() {
2256 let err = Error::from(ErrorKind::Session(SessionErrorKind::NotFound("s1".into())));
2257 assert!(!err.is_transport_failure());
2258 }
2259
2260 #[test]
2261 fn client_options_builder_composes() {
2262 let opts = ClientOptions::new()
2263 .with_program(CliProgram::Path(PathBuf::from("/usr/local/bin/copilot")))
2264 .with_prefix_args(["node"])
2265 .with_cwd(PathBuf::from("/tmp"))
2266 .with_env([("KEY", "value")])
2267 .with_env_remove(["UNWANTED"])
2268 .with_extra_args(["--quiet"])
2269 .with_github_token("ghp_test")
2270 .with_use_logged_in_user(false)
2271 .with_log_level(LogLevel::Debug)
2272 .with_session_idle_timeout_seconds(120)
2273 .with_enable_remote_sessions(true);
2274 assert!(matches!(opts.program, CliProgram::Path(_)));
2275 assert_eq!(opts.prefix_args, vec![std::ffi::OsString::from("node")]);
2276 assert_eq!(opts.working_directory, PathBuf::from("/tmp"));
2277 assert_eq!(
2278 opts.env,
2279 vec![(
2280 std::ffi::OsString::from("KEY"),
2281 std::ffi::OsString::from("value")
2282 )]
2283 );
2284 assert_eq!(opts.env_remove, vec![std::ffi::OsString::from("UNWANTED")]);
2285 assert_eq!(opts.extra_args, vec!["--quiet".to_string()]);
2286 assert_eq!(opts.github_token.as_deref(), Some("ghp_test"));
2287 assert_eq!(opts.use_logged_in_user, Some(false));
2288 assert!(matches!(opts.log_level, Some(LogLevel::Debug)));
2289 assert_eq!(opts.session_idle_timeout_seconds, Some(120));
2290 assert!(opts.enable_remote_sessions);
2291 }
2292
2293 #[test]
2294 fn is_transport_failure_rejects_other_protocol_errors() {
2295 let err = Error::from(ErrorKind::Protocol(ProtocolErrorKind::CliStartupTimeout));
2296 assert!(!err.is_transport_failure());
2297 }
2298
2299 #[test]
2300 fn build_command_lets_env_remove_strip_injected_token() {
2301 let opts = ClientOptions {
2302 github_token: Some("secret".to_string()),
2303 env_remove: vec![std::ffi::OsString::from("COPILOT_SDK_AUTH_TOKEN")],
2304 ..Default::default()
2305 };
2306 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2307 let action = cmd
2309 .as_std()
2310 .get_envs()
2311 .find(|(k, _)| *k == std::ffi::OsStr::new("COPILOT_SDK_AUTH_TOKEN"))
2312 .map(|(_, v)| v);
2313 assert_eq!(
2314 action,
2315 Some(None),
2316 "env_remove should win over github_token"
2317 );
2318 }
2319
2320 #[test]
2321 fn build_command_lets_env_override_injected_token() {
2322 let opts = ClientOptions {
2323 github_token: Some("from-options".to_string()),
2324 env: vec![(
2325 std::ffi::OsString::from("COPILOT_SDK_AUTH_TOKEN"),
2326 std::ffi::OsString::from("from-env"),
2327 )],
2328 ..Default::default()
2329 };
2330 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2331 let value = cmd
2332 .as_std()
2333 .get_envs()
2334 .find(|(k, _)| *k == std::ffi::OsStr::new("COPILOT_SDK_AUTH_TOKEN"))
2335 .and_then(|(_, v)| v);
2336 assert_eq!(value, Some(std::ffi::OsStr::new("from-env")));
2337 }
2338
2339 #[test]
2340 fn build_command_injects_github_token_by_default() {
2341 let opts = ClientOptions {
2342 github_token: Some("just-the-token".to_string()),
2343 ..Default::default()
2344 };
2345 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2346 let value = cmd
2347 .as_std()
2348 .get_envs()
2349 .find(|(k, _)| *k == std::ffi::OsStr::new("COPILOT_SDK_AUTH_TOKEN"))
2350 .and_then(|(_, v)| v);
2351 assert_eq!(value, Some(std::ffi::OsStr::new("just-the-token")));
2352 }
2353
2354 fn env_value<'a>(cmd: &'a tokio::process::Command, key: &str) -> Option<&'a std::ffi::OsStr> {
2355 cmd.as_std()
2356 .get_envs()
2357 .find(|(k, _)| *k == std::ffi::OsStr::new(key))
2358 .and_then(|(_, v)| v)
2359 }
2360
2361 #[test]
2362 fn telemetry_config_builder_composes() {
2363 let cfg = TelemetryConfig::new()
2364 .with_otlp_endpoint("http://collector:4318")
2365 .with_otlp_protocol(OtlpHttpProtocol::HttpProtobuf)
2366 .with_file_path(PathBuf::from("/var/log/copilot.jsonl"))
2367 .with_exporter_type(OtelExporterType::OtlpHttp)
2368 .with_source_name("my-app")
2369 .with_capture_content(true);
2370
2371 assert_eq!(cfg.otlp_endpoint.as_deref(), Some("http://collector:4318"));
2372 assert_eq!(cfg.otlp_protocol, Some(OtlpHttpProtocol::HttpProtobuf));
2373 assert_eq!(
2374 cfg.file_path.as_deref(),
2375 Some(Path::new("/var/log/copilot.jsonl")),
2376 );
2377 assert_eq!(cfg.exporter_type, Some(OtelExporterType::OtlpHttp));
2378 assert_eq!(cfg.source_name.as_deref(), Some("my-app"));
2379 assert_eq!(cfg.capture_content, Some(true));
2380 assert!(!cfg.is_empty());
2381 assert!(TelemetryConfig::new().is_empty());
2382 }
2383
2384 #[test]
2385 fn otlp_http_protocol_serde_matches_env_value() {
2386 for (protocol, wire) in [
2387 (OtlpHttpProtocol::HttpJson, "http/json"),
2388 (OtlpHttpProtocol::HttpProtobuf, "http/protobuf"),
2389 ] {
2390 assert_eq!(protocol.as_str(), wire);
2391
2392 let serialized = serde_json::to_string(&protocol).unwrap();
2393 assert_eq!(serialized, format!("\"{wire}\""));
2394
2395 let deserialized: OtlpHttpProtocol = serde_json::from_str(&serialized).unwrap();
2396 assert_eq!(deserialized, protocol);
2397 }
2398 }
2399
2400 #[test]
2401 fn build_command_sets_otel_env_when_telemetry_enabled() {
2402 let opts = ClientOptions {
2403 telemetry: Some(TelemetryConfig {
2404 otlp_endpoint: Some("http://collector:4318".to_string()),
2405 otlp_protocol: Some(OtlpHttpProtocol::HttpProtobuf),
2406 file_path: Some(PathBuf::from("/var/log/copilot.jsonl")),
2407 exporter_type: Some(OtelExporterType::OtlpHttp),
2408 source_name: Some("my-app".to_string()),
2409 capture_content: Some(true),
2410 }),
2411 ..Default::default()
2412 };
2413 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2414 assert_eq!(
2415 env_value(&cmd, "COPILOT_OTEL_ENABLED"),
2416 Some(std::ffi::OsStr::new("true")),
2417 );
2418 assert_eq!(
2419 env_value(&cmd, "OTEL_EXPORTER_OTLP_ENDPOINT"),
2420 Some(std::ffi::OsStr::new("http://collector:4318")),
2421 );
2422 assert_eq!(
2423 env_value(&cmd, "OTEL_EXPORTER_OTLP_PROTOCOL"),
2424 Some(std::ffi::OsStr::new("http/protobuf")),
2425 );
2426 assert_eq!(
2427 env_value(&cmd, "COPILOT_OTEL_FILE_EXPORTER_PATH"),
2428 Some(std::ffi::OsStr::new("/var/log/copilot.jsonl")),
2429 );
2430 assert_eq!(
2431 env_value(&cmd, "COPILOT_OTEL_EXPORTER_TYPE"),
2432 Some(std::ffi::OsStr::new("otlp-http")),
2433 );
2434 assert_eq!(
2435 env_value(&cmd, "COPILOT_OTEL_SOURCE_NAME"),
2436 Some(std::ffi::OsStr::new("my-app")),
2437 );
2438 assert_eq!(
2439 env_value(&cmd, "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"),
2440 Some(std::ffi::OsStr::new("true")),
2441 );
2442 }
2443
2444 #[test]
2445 fn build_command_omits_otel_env_when_telemetry_none() {
2446 let opts = ClientOptions::default();
2447 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2448 for key in [
2449 "COPILOT_OTEL_ENABLED",
2450 "OTEL_EXPORTER_OTLP_ENDPOINT",
2451 "OTEL_EXPORTER_OTLP_PROTOCOL",
2452 "COPILOT_OTEL_FILE_EXPORTER_PATH",
2453 "COPILOT_OTEL_EXPORTER_TYPE",
2454 "COPILOT_OTEL_SOURCE_NAME",
2455 "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",
2456 ] {
2457 assert!(
2458 env_value(&cmd, key).is_none(),
2459 "expected {key} to be unset when telemetry is None",
2460 );
2461 }
2462 }
2463
2464 #[test]
2465 fn build_command_omits_unset_telemetry_fields() {
2466 let opts = ClientOptions {
2467 telemetry: Some(TelemetryConfig {
2468 otlp_endpoint: Some("http://collector:4318".to_string()),
2469 ..Default::default()
2470 }),
2471 ..Default::default()
2472 };
2473 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2474 assert_eq!(
2476 env_value(&cmd, "COPILOT_OTEL_ENABLED"),
2477 Some(std::ffi::OsStr::new("true")),
2478 );
2479 assert_eq!(
2480 env_value(&cmd, "OTEL_EXPORTER_OTLP_ENDPOINT"),
2481 Some(std::ffi::OsStr::new("http://collector:4318")),
2482 );
2483 for key in [
2485 "OTEL_EXPORTER_OTLP_PROTOCOL",
2486 "COPILOT_OTEL_FILE_EXPORTER_PATH",
2487 "COPILOT_OTEL_EXPORTER_TYPE",
2488 "COPILOT_OTEL_SOURCE_NAME",
2489 "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",
2490 ] {
2491 assert!(env_value(&cmd, key).is_none(), "{key} should be unset");
2492 }
2493 }
2494
2495 #[test]
2496 fn build_command_lets_user_env_override_telemetry() {
2497 let opts = ClientOptions {
2498 telemetry: Some(TelemetryConfig {
2499 otlp_endpoint: Some("http://from-config:4318".to_string()),
2500 ..Default::default()
2501 }),
2502 env: vec![(
2503 std::ffi::OsString::from("OTEL_EXPORTER_OTLP_ENDPOINT"),
2504 std::ffi::OsString::from("http://from-user-env:4318"),
2505 )],
2506 ..Default::default()
2507 };
2508 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2509 assert_eq!(
2510 env_value(&cmd, "OTEL_EXPORTER_OTLP_ENDPOINT"),
2511 Some(std::ffi::OsStr::new("http://from-user-env:4318")),
2512 "user-supplied options.env should override telemetry config",
2513 );
2514 }
2515
2516 #[test]
2517 fn build_command_sets_copilot_home_env_when_configured() {
2518 let opts = ClientOptions::new().with_base_directory(PathBuf::from("/custom/copilot"));
2519 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2520 assert_eq!(
2521 env_value(&cmd, "COPILOT_HOME"),
2522 Some(std::ffi::OsStr::new("/custom/copilot")),
2523 );
2524
2525 let opts = ClientOptions::default();
2526 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2527 assert!(env_value(&cmd, "COPILOT_HOME").is_none());
2528 }
2529
2530 #[test]
2531 fn build_command_sets_connection_token_env_when_configured() {
2532 let opts = ClientOptions::new().with_transport(Transport::Tcp {
2533 port: 0,
2534 connection_token: Some("secret-token".to_string()),
2535 });
2536 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2537 assert_eq!(
2538 env_value(&cmd, "COPILOT_CONNECTION_TOKEN"),
2539 Some(std::ffi::OsStr::new("secret-token")),
2540 );
2541
2542 let opts = ClientOptions::default();
2543 let cmd = Client::build_command(Path::new("/bin/echo"), &opts);
2544 assert!(env_value(&cmd, "COPILOT_CONNECTION_TOKEN").is_none());
2545 }
2546
2547 #[tokio::test]
2548 async fn start_rejects_empty_connection_token() {
2549 let opts = ClientOptions::new()
2550 .with_transport(Transport::Tcp {
2551 port: 0,
2552 connection_token: Some(String::new()),
2553 })
2554 .with_program(CliProgram::Path(PathBuf::from("/bin/echo")));
2555 let err = Client::start(opts).await.unwrap_err();
2556 assert!(
2557 matches!(err.kind(), ErrorKind::InvalidConfig),
2558 "got {err:?}"
2559 );
2560 }
2561
2562 #[tokio::test]
2563 async fn start_rejects_empty_external_connection_token() {
2564 let opts = ClientOptions::new()
2565 .with_transport(Transport::External {
2566 host: "127.0.0.1".to_string(),
2567 port: 1,
2568 connection_token: Some(String::new()),
2569 })
2570 .with_program(CliProgram::Path(PathBuf::from("/bin/echo")));
2571 let err = Client::start(opts).await.unwrap_err();
2572 assert!(
2573 matches!(err.kind(), ErrorKind::InvalidConfig),
2574 "got {err:?}"
2575 );
2576 }
2577
2578 #[test]
2579 fn telemetry_config_capture_content_serializes_as_lowercase_bool() {
2580 let opts_true = ClientOptions {
2581 telemetry: Some(TelemetryConfig {
2582 capture_content: Some(true),
2583 ..Default::default()
2584 }),
2585 ..Default::default()
2586 };
2587 let opts_false = ClientOptions {
2588 telemetry: Some(TelemetryConfig {
2589 capture_content: Some(false),
2590 ..Default::default()
2591 }),
2592 ..Default::default()
2593 };
2594 let cmd_true = Client::build_command(Path::new("/bin/echo"), &opts_true);
2595 let cmd_false = Client::build_command(Path::new("/bin/echo"), &opts_false);
2596 assert_eq!(
2597 env_value(
2598 &cmd_true,
2599 "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
2600 ),
2601 Some(std::ffi::OsStr::new("true")),
2602 );
2603 assert_eq!(
2604 env_value(
2605 &cmd_false,
2606 "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
2607 ),
2608 Some(std::ffi::OsStr::new("false")),
2609 );
2610 }
2611
2612 #[test]
2613 fn session_idle_timeout_args_are_omitted_by_default() {
2614 let opts = ClientOptions::default();
2615 assert!(Client::session_idle_timeout_args(&opts).is_empty());
2616 }
2617
2618 #[test]
2619 fn session_idle_timeout_args_omitted_for_zero() {
2620 let opts = ClientOptions {
2621 session_idle_timeout_seconds: Some(0),
2622 ..Default::default()
2623 };
2624 assert!(Client::session_idle_timeout_args(&opts).is_empty());
2625 }
2626
2627 #[test]
2628 fn session_idle_timeout_args_emit_flag_for_positive_value() {
2629 let opts = ClientOptions {
2630 session_idle_timeout_seconds: Some(300),
2631 ..Default::default()
2632 };
2633 assert_eq!(
2634 Client::session_idle_timeout_args(&opts),
2635 vec!["--session-idle-timeout".to_string(), "300".to_string()]
2636 );
2637 }
2638
2639 #[test]
2640 fn remote_args_omitted_by_default() {
2641 let opts = ClientOptions::default();
2642 assert!(Client::remote_args(&opts).is_empty());
2643 }
2644
2645 #[test]
2646 fn remote_args_emit_flag_when_enabled() {
2647 let opts = ClientOptions {
2648 enable_remote_sessions: true,
2649 ..Default::default()
2650 };
2651 assert_eq!(Client::remote_args(&opts), vec!["--remote".to_string()]);
2652 }
2653
2654 #[test]
2655 fn log_level_args_omitted_when_unset() {
2656 let opts = ClientOptions::default();
2657 assert!(opts.log_level.is_none());
2658 assert!(
2659 Client::log_level_args(&opts).is_empty(),
2660 "with no caller-supplied log_level the SDK must not pass --log-level"
2661 );
2662 }
2663
2664 #[test]
2665 fn log_level_args_emit_flag_when_set() {
2666 let opts = ClientOptions::default().with_log_level(LogLevel::Debug);
2667 assert_eq!(Client::log_level_args(&opts), vec!["--log-level", "debug"]);
2668 }
2669
2670 #[test]
2671 fn log_level_str_round_trips() {
2672 for level in [
2673 LogLevel::None,
2674 LogLevel::Error,
2675 LogLevel::Warning,
2676 LogLevel::Info,
2677 LogLevel::Debug,
2678 LogLevel::All,
2679 ] {
2680 let s = level.as_str();
2681 let json = serde_json::to_string(&level).unwrap();
2682 assert_eq!(json, format!("\"{s}\""));
2683 let parsed: LogLevel = serde_json::from_str(&json).unwrap();
2684 assert_eq!(parsed, level);
2685 }
2686 }
2687
2688 #[test]
2689 fn client_options_debug_redacts_handler() {
2690 struct StubHandler;
2691 #[async_trait]
2692 impl ListModelsHandler for StubHandler {
2693 async fn list_models(&self) -> Result<Vec<Model>> {
2694 Ok(vec![])
2695 }
2696 }
2697 let opts = ClientOptions {
2698 on_list_models: Some(Arc::new(StubHandler)),
2699 github_token: Some("secret-token".into()),
2700 ..Default::default()
2701 };
2702 let debug = format!("{opts:?}");
2703 assert!(debug.contains("on_list_models: Some(\"<set>\")"));
2704 assert!(debug.contains("github_token: Some(\"<redacted>\")"));
2705 assert!(!debug.contains("secret-token"));
2706 }
2707
2708 #[tokio::test]
2709 async fn list_models_uses_on_list_models_handler_when_set() {
2710 use std::sync::atomic::{AtomicUsize, Ordering};
2711
2712 struct CountingHandler {
2713 calls: Arc<AtomicUsize>,
2714 models: Vec<Model>,
2715 }
2716 #[async_trait]
2717 impl ListModelsHandler for CountingHandler {
2718 async fn list_models(&self) -> Result<Vec<Model>> {
2719 self.calls.fetch_add(1, Ordering::SeqCst);
2720 Ok(self.models.clone())
2721 }
2722 }
2723
2724 let calls = Arc::new(AtomicUsize::new(0));
2725 let model = Model {
2726 id: "byok-gpt-4".into(),
2727 name: "BYOK GPT-4".into(),
2728 ..Default::default()
2729 };
2730 let handler: Arc<dyn ListModelsHandler> = Arc::new(CountingHandler {
2731 calls: Arc::clone(&calls),
2732 models: vec![model.clone()],
2733 });
2734
2735 let client = client_with_list_models_handler(handler);
2736
2737 let result = client.list_models().await.unwrap();
2738 assert_eq!(result.len(), 1);
2739 assert_eq!(result[0].id, "byok-gpt-4");
2740 assert_eq!(calls.load(Ordering::SeqCst), 1);
2741 }
2742
2743 #[tokio::test]
2744 async fn list_models_serializes_concurrent_cache_misses() {
2745 use std::sync::atomic::{AtomicUsize, Ordering};
2746
2747 struct SlowCountingHandler {
2748 calls: Arc<AtomicUsize>,
2749 models: Vec<Model>,
2750 }
2751 #[async_trait]
2752 impl ListModelsHandler for SlowCountingHandler {
2753 async fn list_models(&self) -> Result<Vec<Model>> {
2754 self.calls.fetch_add(1, Ordering::SeqCst);
2755 tokio::time::sleep(std::time::Duration::from_millis(25)).await;
2756 Ok(self.models.clone())
2757 }
2758 }
2759
2760 let calls = Arc::new(AtomicUsize::new(0));
2761 let model = Model {
2762 id: "single-flight-model".into(),
2763 name: "Single Flight Model".into(),
2764 ..Default::default()
2765 };
2766 let handler: Arc<dyn ListModelsHandler> = Arc::new(SlowCountingHandler {
2767 calls: Arc::clone(&calls),
2768 models: vec![model],
2769 });
2770 let client = client_with_list_models_handler(handler);
2771
2772 let (first, second) = tokio::join!(client.list_models(), client.list_models());
2773 assert_eq!(first.unwrap()[0].id, "single-flight-model");
2774 assert_eq!(second.unwrap()[0].id, "single-flight-model");
2775 assert_eq!(calls.load(Ordering::SeqCst), 1);
2776 }
2777
2778 #[tokio::test]
2779 async fn cancelled_resume_session_unregisters_pending_session() {
2780 let (client_write, _server_read) = tokio::io::duplex(8192);
2781 let (_server_write, client_read) = tokio::io::duplex(8192);
2782 let client = Client::from_streams(client_read, client_write, std::env::temp_dir()).unwrap();
2783 let session_id = SessionId::new("resume-cancel-test");
2784 let handle = tokio::spawn({
2785 let client = client.clone();
2786 async move {
2787 client
2788 .resume_session(ResumeSessionConfig::new(session_id))
2789 .await
2790 }
2791 });
2792
2793 wait_for_pending_session_registration(&client).await;
2794 handle.abort();
2795 let _ = handle.await;
2796
2797 assert!(client.inner.router.session_ids().is_empty());
2798 client.force_stop();
2799 }
2800
2801 fn client_with_list_models_handler(handler: Arc<dyn ListModelsHandler>) -> Client {
2802 Client {
2803 inner: Arc::new(ClientInner {
2804 child: parking_lot::Mutex::new(None),
2805 rpc: {
2806 let (req_tx, _req_rx) = mpsc::unbounded_channel();
2807 let (notif_tx, _notif_rx) = broadcast::channel(16);
2808 let (read_pipe, _write_pipe) = tokio::io::duplex(64);
2809 let (_unused_read, write_pipe) = tokio::io::duplex(64);
2810 JsonRpcClient::new(write_pipe, read_pipe, notif_tx, req_tx)
2811 },
2812 cwd: PathBuf::from("."),
2813 request_rx: parking_lot::Mutex::new(None),
2814 notification_tx: broadcast::channel(16).0,
2815 router: router::SessionRouter::new(),
2816 negotiated_protocol_version: OnceLock::new(),
2817 state: parking_lot::Mutex::new(ConnectionState::Connected),
2818 lifecycle_tx: broadcast::channel(16).0,
2819 on_list_models: Some(handler),
2820 models_cache: parking_lot::Mutex::new(Arc::new(tokio::sync::OnceCell::new())),
2821 session_fs_configured: false,
2822 session_fs_sqlite_declared: false,
2823 llm_inference: OnceLock::new(),
2824 on_github_telemetry: None,
2825 on_get_trace_context: None,
2826 effective_connection_token: None,
2827 mode: ClientMode::default(),
2828 }),
2829 }
2830 }
2831
2832 async fn wait_for_pending_session_registration(client: &Client) {
2833 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
2834 while client.inner.router.session_ids().is_empty() {
2835 assert!(
2836 tokio::time::Instant::now() < deadline,
2837 "session was not registered"
2838 );
2839 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
2840 }
2841 }
2842}