1use std::env;
27use std::ffi::OsString;
28use std::fs;
29use std::path::{Path, PathBuf};
30use std::sync::Arc;
31use std::time::Instant;
32
33use async_trait::async_trait;
34use base64::engine::general_purpose::STANDARD as BASE64;
35use base64::Engine as _;
36use github_copilot_sdk::handler::{
37 ApproveAllHandler, DenyAllHandler, PermissionHandler, PermissionResult,
38};
39use github_copilot_sdk::session_events::{
40 AssistantMessageData, AssistantMessageDeltaData, AssistantUsageData, SessionErrorData,
41 SessionEventType, ToolExecutionCompleteData, ToolExecutionStartData,
42};
43use github_copilot_sdk::subscription::RecvErrorKind;
44use github_copilot_sdk::{
45 Attachment, CliProgram, Client, ClientMode, ClientOptions, Error as SdkError,
46 ErrorKind as SdkErrorKind, EventSubscription, IndexMap, LogLevel, McpHttpServerConfig,
47 McpServerConfig, McpStdioServerConfig, MessageOptions, PermissionRequestData,
48 PermissionRequestKind, ProtocolErrorKind, ProviderConfig, RequestId, SessionConfig,
49 SessionErrorKind, SessionEvent, SessionId, SystemMessageConfig, Transport,
50};
51use serde_json::Value;
52use tokio::sync::{mpsc, Mutex};
53use tokio::time;
54use tokio_stream::wrappers::UnboundedReceiverStream;
55use tokio_stream::StreamExt;
56use tracing::{debug, info, warn};
57
58use crate::copilot_common::{
59 render_turn, runtime_env_keys, system_prompt, HeadlessEventStream, HeadlessStreamEvent,
60 HeadlessToolResponse, HeadlessTurnProvider, ObservedToolCall, PermissionPolicy,
61};
62use crate::copilot_models::catalog_ids;
63use crate::copilot_sdk_config::{CopilotSdkConfig, CopilotSdkProvider};
64use crate::types::{
65 ChatMessage, ChatRequest, ChatResponse, ChatStream, ErrorKind, LlmCapabilities, LlmProvider,
66 McpTransport, RunnerError, StreamChunk, TokenUsage,
67};
68
69const SERVICE: &str = "copilot-sdk";
71
72type EventSink<'a> = Option<&'a mpsc::UnboundedSender<Result<HeadlessStreamEvent, RunnerError>>>;
74
75#[derive(Debug)]
79pub struct CopilotSdkRunner {
80 inner: Arc<Inner>,
83 models: Vec<String>,
84}
85
86#[derive(Debug)]
87struct Inner {
88 config: CopilotSdkConfig,
89 client: Mutex<Option<Arc<Client>>>,
93}
94
95impl CopilotSdkRunner {
96 #[must_use]
98 pub fn from_env() -> Self {
99 Self::with_config(CopilotSdkConfig::from_env())
100 }
101
102 #[must_use]
104 pub fn with_config(config: CopilotSdkConfig) -> Self {
105 Self {
106 inner: Arc::new(Inner {
107 config,
108 client: Mutex::new(None),
109 }),
110 models: catalog_ids(),
111 }
112 }
113
114 #[must_use]
116 pub fn config(&self) -> &CopilotSdkConfig {
117 &self.inner.config
118 }
119}
120
121const SDK_SET_ENV_KEYS: &[&str] = &["COPILOT_HOME", "COPILOT_DISABLE_KEYTAR"];
128
129fn client_options(
137 config: &CopilotSdkConfig,
138 host_env: impl Iterator<Item = OsString>,
139) -> ClientOptions {
140 let mut options = ClientOptions::new()
141 .with_transport(Transport::Stdio)
142 .with_mode(ClientMode::Empty)
143 .with_base_directory(config.base_directory.clone())
144 .with_log_level(LogLevel::Warning);
145 if let Some(path) = &config.runtime_path {
146 options = options.with_program(CliProgram::Path(path.clone()));
147 }
148 if let Some(token) = &config.github_token {
158 options = options.with_env([("COPILOT_GITHUB_TOKEN", token.as_str())]);
159 }
160
161 let allowed = runtime_env_keys();
162 let set_here: Vec<OsString> = options.env.iter().map(|(name, _)| name.clone()).collect();
163 let inherited: Vec<OsString> = host_env
164 .filter(|name| {
165 let kept = name.to_str().is_some_and(|name| {
166 allowed.iter().any(|key| key == name) || SDK_SET_ENV_KEYS.contains(&name)
167 });
168 !kept && !set_here.contains(name)
169 })
170 .collect();
171 options.with_env_remove(inherited)
172}
173
174impl Inner {
175 async fn client(&self) -> Result<Arc<Client>, RunnerError> {
177 let mut slot = self.client.lock().await;
178 if let Some(client) = slot.as_ref() {
179 return Ok(Arc::clone(client));
180 }
181
182 let options = client_options(&self.config, env::vars_os().map(|(name, _)| name));
183
184 let started = Instant::now();
185 let client = time::timeout(self.config.session_timeout, Client::start(options))
186 .await
187 .map_err(|_| {
188 RunnerError::timeout(format!(
189 "{SERVICE}: runtime did not finish its handshake within {:?}",
190 self.config.session_timeout
191 ))
192 })?
193 .map_err(|e| map_sdk_error(&e))?;
194 info!(
195 elapsed_ms = started.elapsed().as_millis(),
196 runtime = ?self.config.runtime_path,
197 "{SERVICE}: runtime started"
198 );
199
200 let client = Arc::new(client);
201 *slot = Some(Arc::clone(&client));
202 Ok(client)
203 }
204
205 async fn discard(&self, failed: &Arc<Client>) {
208 let mut slot = self.client.lock().await;
209 if slot.as_ref().is_some_and(|c| Arc::ptr_eq(c, failed)) {
210 *slot = None;
211 }
212 drop(slot);
213 if let Err(e) = failed.stop().await {
214 debug!(error = %e, "{SERVICE}: stopping a failed client");
215 }
216 }
217
218 async fn validated_model(
230 &self,
231 client: &Client,
232 requested: Option<&str>,
233 ) -> Result<String, RunnerError> {
234 let model = requested.map_or_else(|| self.config.model.clone(), str::to_owned);
235 if self.config.provider.is_some() {
236 return Ok(model);
237 }
238 let catalogue = client.list_models().await.map_err(|e| map_sdk_error(&e))?;
239 if catalogue.iter().any(|m| m.id == model) {
240 Ok(model)
241 } else {
242 Err(RunnerError::model_unavailable(model))
243 }
244 }
245
246 async fn run_turn(
248 &self,
249 request: &ChatRequest,
250 sink: EventSink<'_>,
251 ) -> Result<HeadlessToolResponse, RunnerError> {
252 let started = Instant::now();
253 let client = self.client().await?;
254 let model = self
255 .validated_model(&client, request.model.as_deref())
256 .await?;
257
258 let scratch = tempfile::Builder::new()
262 .prefix("embacle-copilot-sdk-")
263 .tempdir()
264 .map_err(|e| RunnerError::internal(format!("{SERVICE}: scratch dir: {e}")))?;
265
266 let rendered = render_turn(request, self.config.max_history_turns, false);
267 let attachments = write_image_attachments(scratch.path(), rendered.last_user)?;
268 let session_config = self.session_config(request, &model, scratch.path());
269
270 let session = match time::timeout(
271 self.config.session_timeout,
272 client.create_session(session_config),
273 )
274 .await
275 {
276 Ok(Ok(session)) => session,
277 Ok(Err(e)) => {
278 if e.is_transport_failure() {
279 self.discard(&client).await;
280 }
281 return Err(map_sdk_error(&e));
282 }
283 Err(_) => {
284 return Err(RunnerError::timeout(format!(
285 "{SERVICE}: session.create exceeded {:?}",
286 self.config.session_timeout
287 )));
288 }
289 };
290 let mut events = session.subscribe();
291
292 let mut message = MessageOptions::from(rendered.text);
293 if !attachments.is_empty() {
294 message.attachments = Some(attachments);
295 }
296 if let Err(e) = session.send(message).await {
297 if e.is_transport_failure() {
298 self.discard(&client).await;
299 }
300 let _ = session.disconnect().await;
301 return Err(map_sdk_error(&e));
302 }
303
304 let outcome =
305 time::timeout(self.config.prompt_timeout, collect_turn(&mut events, sink)).await;
306 let result = match outcome {
307 Ok(Ok(acc)) => Ok(acc.finish(model)),
308 Ok(Err(e)) => Err(naming_model(e, &model)),
309 Err(_) => {
310 if let Err(e) = session.abort().await {
313 debug!(error = %e, "{SERVICE}: abort after prompt timeout");
314 }
315 Err(RunnerError::timeout(format!(
316 "{SERVICE}: turn exceeded {:?}",
317 self.config.prompt_timeout
318 )))
319 }
320 };
321
322 if let Err(e) = session.disconnect().await {
323 debug!(error = %e, "{SERVICE}: session.disconnect");
324 }
325 drop(scratch);
326
327 match &result {
328 Ok(response) => info!(
329 model = %response.model,
330 tool_calls = response.tool_calls.len(),
331 elapsed_ms = started.elapsed().as_millis(),
332 "{SERVICE}: turn complete"
333 ),
334 Err(e) => {
335 warn!(error = %e, elapsed_ms = started.elapsed().as_millis(), "{SERVICE}: turn failed");
336 }
337 }
338 result
339 }
340
341 fn session_config(&self, request: &ChatRequest, model: &str, scratch: &Path) -> SessionConfig {
342 let mut config = SessionConfig::default()
343 .with_model(model)
344 .with_streaming(true)
345 .with_working_directory(scratch)
346 .with_client_name(concat!("embacle/", env!("CARGO_PKG_VERSION")));
347
348 if let Some(provider) = &self.config.provider {
349 config = config.with_provider(provider_config(provider));
350 }
351
352 if let Some(system) = system_prompt(request) {
353 config = config.with_system_message(
354 SystemMessageConfig::new()
355 .with_mode("replace")
356 .with_content(system),
357 );
358 }
359
360 let mut available: Vec<String> = Vec::new();
363 let mut registered: Vec<String> = Vec::new();
364 if !request.mcp_servers.is_empty() {
365 let mut servers = IndexMap::new();
366 for server in &request.mcp_servers {
367 servers.insert(server.name.clone(), mcp_server(&server.transport));
368 registered.push(server.name.clone());
369 }
370 config = config.with_mcp_servers(servers);
371 available.push("mcp:*".to_owned());
372 }
373 config
374 .with_available_tools(available)
375 .with_permission_handler(self.permission_handler(registered))
376 }
377
378 fn permission_handler(&self, servers: Vec<String>) -> Arc<dyn PermissionHandler> {
385 match self.config.permission_policy {
386 PermissionPolicy::AutoApprove => Arc::new(ApproveAllHandler),
387 PermissionPolicy::DenyAll => Arc::new(RegisteredMcpHandler { servers }),
388 }
389 }
390}
391
392#[derive(Debug)]
400struct RegisteredMcpHandler {
401 servers: Vec<String>,
402}
403
404impl RegisteredMcpHandler {
405 fn permits(&self, data: &PermissionRequestData) -> bool {
406 if !matches!(data.kind, Some(PermissionRequestKind::Mcp))
407 || data.managed_approval_required == Some(true)
408 {
409 return false;
410 }
411 let request = data.extra.get("permissionRequest").unwrap_or(&data.extra);
412 request
413 .get("serverName")
414 .and_then(Value::as_str)
415 .is_some_and(|server| self.servers.iter().any(|known| known == server))
416 }
417}
418
419#[async_trait]
420impl PermissionHandler for RegisteredMcpHandler {
421 async fn handle(
422 &self,
423 session_id: SessionId,
424 request_id: RequestId,
425 data: PermissionRequestData,
426 ) -> PermissionResult {
427 if self.permits(&data) {
428 PermissionResult::approve_once()
429 } else {
430 DenyAllHandler.handle(session_id, request_id, data).await
431 }
432 }
433}
434
435fn provider_config(provider: &CopilotSdkProvider) -> ProviderConfig {
438 let mut config = ProviderConfig::new(provider.base_url.clone());
439 if let Some(provider_type) = &provider.provider_type {
440 config = config.with_provider_type(provider_type.clone());
441 }
442 if let Some(api_key) = &provider.api_key {
443 config = config.with_api_key(api_key.clone());
444 }
445 if let Some(max) = provider.max_prompt_tokens {
446 config = config.with_max_prompt_tokens(max);
447 }
448 config
449}
450
451fn mcp_server(transport: &McpTransport) -> McpServerConfig {
453 match transport {
454 McpTransport::Http { url, headers } => McpServerConfig::Http(McpHttpServerConfig {
455 url: url.clone(),
456 headers: headers
457 .iter()
458 .map(|h| (h.name.clone(), h.value.clone()))
459 .collect(),
460 ..McpHttpServerConfig::default()
461 }),
462 McpTransport::Sse { url, headers } => McpServerConfig::Sse(McpHttpServerConfig {
463 url: url.clone(),
464 headers: headers
465 .iter()
466 .map(|h| (h.name.clone(), h.value.clone()))
467 .collect(),
468 ..McpHttpServerConfig::default()
469 }),
470 McpTransport::Stdio { command, args, env } => {
471 McpServerConfig::Stdio(McpStdioServerConfig {
472 command: command.clone(),
473 args: args.clone(),
474 env: env
475 .iter()
476 .map(|h| (h.name.clone(), h.value.clone()))
477 .collect(),
478 ..McpStdioServerConfig::default()
479 })
480 }
481 }
482}
483
484fn write_image_attachments(
487 dir: &Path,
488 last_user: Option<&ChatMessage>,
489) -> Result<Vec<Attachment>, RunnerError> {
490 let Some(images) = last_user.and_then(|m| m.images.as_ref()) else {
491 return Ok(Vec::new());
492 };
493 let mut attachments = Vec::with_capacity(images.len());
494 for (index, image) in images.iter().enumerate() {
495 let bytes = BASE64.decode(&image.data).map_err(|e| {
496 RunnerError::internal(format!("{SERVICE}: image {index} is not valid base64: {e}"))
497 })?;
498 let extension = match image.mime_type.as_str() {
499 "image/png" => "png",
500 "image/jpeg" => "jpg",
501 "image/webp" => "webp",
502 "image/gif" => "gif",
503 _ => "bin",
504 };
505 let path: PathBuf = dir.join(format!("image-{index}.{extension}"));
506 fs::write(&path, bytes)
507 .map_err(|e| RunnerError::internal(format!("{SERVICE}: writing image {index}: {e}")))?;
508 attachments.push(Attachment::File {
509 path,
510 display_name: None,
511 line_range: None,
512 });
513 }
514 Ok(attachments)
515}
516
517#[derive(Debug, Default)]
519struct TurnAccumulator {
520 content: String,
521 served_model: Option<String>,
522 prompt_tokens: u32,
523 completion_tokens: u32,
524 cached_read_tokens: u32,
525 cached_write_tokens: u32,
526 reasoning_tokens: u32,
527 saw_usage: bool,
528 tool_calls: Vec<ObservedToolCall>,
529 finish_reason: Option<String>,
530}
531
532impl TurnAccumulator {
533 fn absorb(&mut self, event: &SessionEvent, sink: EventSink<'_>) -> Result<bool, RunnerError> {
536 match event.parsed_type() {
537 SessionEventType::AssistantMessageDelta => {
538 if event.agent_id.is_none() {
541 if let Some(delta) = event.typed_data::<AssistantMessageDeltaData>() {
542 emit(sink, HeadlessStreamEvent::TextDelta(delta.delta_content));
543 }
544 }
545 }
546 SessionEventType::AssistantMessage => {
547 if let Some(message) = event.typed_data::<AssistantMessageData>() {
548 if !message.content.is_empty() {
551 self.content = message.content;
552 }
553 if let Some(model) = message.model {
554 self.served_model = Some(model);
555 }
556 }
557 }
558 SessionEventType::AssistantUsage => {
559 if let Some(usage) = event.typed_data::<AssistantUsageData>() {
560 self.add_usage(&usage);
561 }
562 }
563 SessionEventType::ToolExecutionStart => {
564 if let Some(start) = event.typed_data::<ToolExecutionStartData>() {
565 let observed = ObservedToolCall {
566 id: start.tool_call_id,
567 title: start.tool_name.clone(),
568 status: "InProgress".to_owned(),
569 name: Some(start.tool_name),
570 arguments: start.arguments,
571 result: None,
572 };
573 self.tool_calls.push(observed.clone());
574 emit(sink, HeadlessStreamEvent::ToolCall(observed));
575 }
576 }
577 SessionEventType::ToolExecutionComplete => {
578 if let Some(done) = event.typed_data::<ToolExecutionCompleteData>() {
579 if let Some(existing) = self
580 .tool_calls
581 .iter_mut()
582 .find(|t| t.id == done.tool_call_id)
583 {
584 existing.status =
585 String::from(if done.success { "Completed" } else { "Failed" });
586 existing.result = done
587 .result
588 .map(|r| r.content)
589 .or_else(|| done.error.map(|e| e.message));
590 emit(sink, HeadlessStreamEvent::ToolCall(existing.clone()));
591 }
592 }
593 }
594 SessionEventType::SessionError => {
595 return Err(map_session_error(event));
596 }
597 SessionEventType::SessionIdle => return Ok(true),
598 _ => {}
599 }
600 Ok(false)
601 }
602
603 fn add_usage(&mut self, usage: &AssistantUsageData) {
604 self.saw_usage = true;
605 self.prompt_tokens = self.prompt_tokens.saturating_add(count(usage.input_tokens));
606 self.completion_tokens = self
607 .completion_tokens
608 .saturating_add(count(usage.output_tokens));
609 self.cached_read_tokens = self
610 .cached_read_tokens
611 .saturating_add(count(usage.cache_read_tokens));
612 self.cached_write_tokens = self
613 .cached_write_tokens
614 .saturating_add(count(usage.cache_write_tokens));
615 self.reasoning_tokens = self
616 .reasoning_tokens
617 .saturating_add(count(usage.reasoning_tokens));
618 self.served_model = Some(usage.model.clone());
619 if usage.finish_reason.is_some() {
620 self.finish_reason.clone_from(&usage.finish_reason);
621 }
622 }
623
624 fn finish(self, requested_model: String) -> HeadlessToolResponse {
625 let usage = self.saw_usage.then(|| TokenUsage {
626 prompt_tokens: self.prompt_tokens,
627 completion_tokens: self.completion_tokens,
628 total_tokens: self.prompt_tokens.saturating_add(self.completion_tokens),
629 cached_read_tokens: Some(self.cached_read_tokens),
630 cached_write_tokens: Some(self.cached_write_tokens),
631 reasoning_tokens: Some(self.reasoning_tokens),
632 });
633 HeadlessToolResponse {
634 content: self.content,
635 model: self.served_model.unwrap_or(requested_model),
636 tool_calls: self.tool_calls,
637 usage,
638 finish_reason: self.finish_reason,
639 }
640 }
641}
642
643fn count(value: Option<i64>) -> u32 {
645 value.and_then(|v| u32::try_from(v).ok()).unwrap_or(0)
646}
647
648fn emit(sink: EventSink<'_>, event: HeadlessStreamEvent) {
649 if let Some(tx) = sink {
650 let _ = tx.send(Ok(event));
653 }
654}
655
656async fn collect_turn(
658 events: &mut EventSubscription,
659 sink: EventSink<'_>,
660) -> Result<TurnAccumulator, RunnerError> {
661 let mut acc = TurnAccumulator::default();
662 loop {
663 let event = match events.recv().await {
664 Ok(event) => event,
665 Err(e) => match e.kind() {
666 RecvErrorKind::Lagged(lag) => {
667 warn!(%lag, "{SERVICE}: event subscription lagged; continuing from the next event");
668 continue;
669 }
670 RecvErrorKind::Closed => {
671 return Err(RunnerError::external_service(
672 SERVICE,
673 "event stream closed before the turn went idle",
674 ));
675 }
676 other => {
679 return Err(RunnerError::external_service(
680 SERVICE,
681 format!("event subscription failed: {other:?}"),
682 ));
683 }
684 },
685 };
686 if acc.absorb(&event, sink)? {
687 return Ok(acc);
688 }
689 }
690}
691
692fn map_session_error(event: &SessionEvent) -> RunnerError {
703 let Some(data) = event.typed_data::<SessionErrorData>() else {
704 return RunnerError::external_service(SERVICE, format!("session.error: {}", event.data));
705 };
706 let mut message = data.message;
707 if let Some(code) = &data.error_code {
708 message.push_str(" (");
709 message.push_str(code);
710 message.push(')');
711 }
712 if data.status_code == Some(404) {
713 return RunnerError {
714 kind: ErrorKind::ModelUnavailable,
715 message: format!("{SERVICE}: {message}"),
716 };
717 }
718 match data.error_type.as_str() {
719 "authentication" | "authorization" => RunnerError::auth_failure(message),
720 "quota" | "rate_limit" => RunnerError::rate_limit(SERVICE, message),
721 "context_limit" => RunnerError::context_length(message),
722 _ => RunnerError::external_service(SERVICE, message),
723 }
724}
725
726fn naming_model(mut error: RunnerError, model: &str) -> RunnerError {
730 if error.kind == ErrorKind::ModelUnavailable && !error.message.contains(model) {
731 error.message = format!("{} (model {model:?})", error.message);
732 }
733 error
734}
735
736fn map_sdk_error(error: &SdkError) -> RunnerError {
738 let message = error.to_string();
739 match error.kind() {
740 SdkErrorKind::BinaryNotFound { .. } => RunnerError::binary_not_found("copilot-runtime"),
741 SdkErrorKind::InvalidConfig => RunnerError::config(message),
742 SdkErrorKind::Session(SessionErrorKind::Timeout(_)) => RunnerError::timeout(message),
743 SdkErrorKind::Protocol(ProtocolErrorKind::VersionMismatch { .. }) => {
744 RunnerError::config(format!("{SERVICE}: {message}"))
745 }
746 _ => RunnerError::external_service(SERVICE, message),
747 }
748}
749
750#[async_trait]
751impl HeadlessTurnProvider for CopilotSdkRunner {
752 async fn converse(&self, request: &ChatRequest) -> Result<HeadlessToolResponse, RunnerError> {
753 self.inner.run_turn(request, None).await
754 }
755
756 async fn converse_stream(
757 &self,
758 request: &ChatRequest,
759 ) -> Result<HeadlessEventStream, RunnerError> {
760 let (tx, rx) = mpsc::unbounded_channel();
761 let inner = Arc::clone(&self.inner);
762 let request = request.clone();
763 tokio::spawn(async move {
767 let outcome = inner.run_turn(&request, Some(&tx)).await;
768 let _ = match outcome {
769 Ok(response) => tx.send(Ok(HeadlessStreamEvent::Done(response))),
770 Err(e) => tx.send(Err(e)),
771 };
772 });
773 Ok(Box::pin(UnboundedReceiverStream::new(rx)))
774 }
775}
776
777#[async_trait]
778impl LlmProvider for CopilotSdkRunner {
779 fn name(&self) -> &'static str {
780 "copilot_sdk"
781 }
782
783 fn display_name(&self) -> &str {
784 "GitHub Copilot (SDK)"
785 }
786
787 fn capabilities(&self) -> LlmCapabilities {
788 let base =
789 LlmCapabilities::STREAMING | LlmCapabilities::SYSTEM_MESSAGES | LlmCapabilities::VISION;
790 if self.inner.config.mcp_tool_calling {
795 base | LlmCapabilities::SDK_TOOL_CALLING
796 } else {
797 base
798 }
799 }
800
801 fn default_model(&self) -> &str {
802 &self.inner.config.model
803 }
804
805 fn available_models(&self) -> &[String] {
806 &self.models
807 }
808
809 async fn complete(&self, request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
810 let turn = self.converse(request).await?;
811 Ok(ChatResponse {
812 content: turn.content,
813 model: turn.model,
814 usage: turn.usage,
815 finish_reason: turn.finish_reason,
816 warnings: None,
817 tool_calls: None,
818 })
819 }
820
821 async fn complete_stream(&self, request: &ChatRequest) -> Result<ChatStream, RunnerError> {
822 let events = self.converse_stream(request).await?;
823 let chunks = events.map(|event| {
824 event.map(|event| match event {
825 HeadlessStreamEvent::TextDelta(delta) => StreamChunk {
826 delta,
827 is_final: false,
828 finish_reason: None,
829 },
830 HeadlessStreamEvent::ToolCall(_) => StreamChunk {
831 delta: String::new(),
832 is_final: false,
833 finish_reason: None,
834 },
835 HeadlessStreamEvent::Done(response) => StreamChunk {
836 delta: String::new(),
837 is_final: true,
838 finish_reason: response.finish_reason,
839 },
840 })
841 });
842 Ok(Box::pin(chunks))
843 }
844
845 async fn health_check(&self) -> Result<bool, RunnerError> {
846 let client = self.inner.client().await?;
847 client.ping(None).await.map_err(|e| map_sdk_error(&e))?;
848 Ok(true)
849 }
850}
851
852#[cfg(test)]
853mod tests {
854 use super::*;
855 use crate::types::McpHeader;
856 use serde_json::json;
857
858 fn approves_once(result: &PermissionResult) -> bool {
861 format!("{result:?}").contains("ApproveOnce")
862 }
863
864 fn event(event_type: &str, data: serde_json::Value) -> SessionEvent {
865 let mut envelope = json!({
866 "id": "evt-1",
867 "timestamp": "2026-09-18T00:00:00Z",
868 "type": event_type,
869 });
870 envelope["data"] = data;
871 serde_json::from_value(envelope).unwrap_or_else(|e| unreachable!("test event: {e}"))
872 }
873
874 #[test]
879 fn the_runtime_child_is_stripped_of_every_host_variable_outside_the_allowlist() {
880 let config = CopilotSdkConfig {
881 github_token: Some("tok-sdk-1".to_owned()),
882 ..CopilotSdkConfig::default()
883 };
884 let host_env = [
885 "PATH",
886 "HOME",
887 "DATABASE_URL",
888 "PIERRE_MASTER_ENCRYPTION_KEY",
889 "COPILOT_GITHUB_TOKEN",
890 "COPILOT_HOME",
891 "COPILOT_PROVIDER_BASE_URL",
892 "HTTPS_PROXY",
893 ]
894 .map(OsString::from);
895
896 let options = client_options(&config, host_env.into_iter());
897
898 let removed: Vec<&str> = options
899 .env_remove
900 .iter()
901 .filter_map(|name| name.to_str())
902 .collect();
903 assert_eq!(removed, ["DATABASE_URL", "PIERRE_MASTER_ENCRYPTION_KEY"]);
904 assert_eq!(
905 options.env,
906 [(
907 OsString::from("COPILOT_GITHUB_TOKEN"),
908 OsString::from("tok-sdk-1")
909 )]
910 );
911 }
912
913 #[test]
914 fn the_published_list_carries_the_5_5_sonnet_and_haiku() {
915 let runner = CopilotSdkRunner::with_config(CopilotSdkConfig::default());
918 let published = runner.available_models();
919 for model in ["claude-sonnet-5.5", "claude-haiku-5.5"] {
920 assert!(
921 published.iter().any(|id| id == model),
922 "{model} missing from {published:?}"
923 );
924 }
925 }
926
927 #[test]
928 fn usage_accumulates_across_calls_and_names_the_served_model() {
929 let mut acc = TurnAccumulator::default();
930 let first = event(
931 "assistant.usage",
932 json!({"model": "claude-haiku-4.5", "inputTokens": 21_678, "outputTokens": 154, "cacheReadTokens": 0, "cacheWriteTokens": 21_678}),
933 );
934 let second = event(
935 "assistant.usage",
936 json!({"model": "claude-haiku-4.5", "inputTokens": 21_860, "outputTokens": 9, "cacheReadTokens": 21_678, "cacheWriteTokens": 176, "reasoningTokens": 113, "finishReason": "stop"}),
937 );
938 assert_eq!(acc.absorb(&first, None).ok(), Some(false));
939 assert_eq!(acc.absorb(&second, None).ok(), Some(false));
940 let response = acc.finish("claude-sonnet-5".to_owned());
941 assert_eq!(response.model, "claude-haiku-4.5");
942 let usage = response
943 .usage
944 .unwrap_or_else(|| unreachable!("usage recorded"));
945 assert_eq!(usage.prompt_tokens, 43_538);
946 assert_eq!(usage.completion_tokens, 163);
947 assert_eq!(usage.total_tokens, 43_701);
948 assert_eq!(usage.cached_read_tokens, Some(21_678));
949 assert_eq!(usage.cached_write_tokens, Some(21_854));
950 assert_eq!(usage.reasoning_tokens, Some(113));
951 assert_eq!(response.finish_reason.as_deref(), Some("stop"));
952 }
953
954 #[test]
955 fn tool_events_carry_name_arguments_and_result() {
956 let mut acc = TurnAccumulator::default();
957 let (tx, mut rx) = mpsc::unbounded_channel();
958 let start = event(
959 "tool.execution_start",
960 json!({"toolCallId": "call-1", "toolName": "get_activities", "arguments": {"limit": 3}}),
961 );
962 let done = event(
963 "tool.execution_complete",
964 json!({"toolCallId": "call-1", "success": true, "result": {"content": "three rides"}}),
965 );
966 assert_eq!(acc.absorb(&start, Some(&tx)).ok(), Some(false));
967 assert_eq!(acc.absorb(&done, Some(&tx)).ok(), Some(false));
968
969 let Ok(Ok(HeadlessStreamEvent::ToolCall(started))) = rx.try_recv() else {
970 unreachable!("first event is the tool start");
971 };
972 assert_eq!(started.status, "InProgress");
973 assert_eq!(started.name.as_deref(), Some("get_activities"));
974 assert_eq!(started.arguments, Some(json!({"limit": 3})));
975
976 let Ok(Ok(HeadlessStreamEvent::ToolCall(finished))) = rx.try_recv() else {
977 unreachable!("second event is the tool completion");
978 };
979 assert_eq!(finished.status, "Completed");
980 assert_eq!(finished.result.as_deref(), Some("three rides"));
981
982 let response = acc.finish("m".to_owned());
983 assert_eq!(response.tool_calls.len(), 1);
984 assert_eq!(
985 response.tool_calls[0].result.as_deref(),
986 Some("three rides")
987 );
988 }
989
990 #[test]
991 fn the_last_non_empty_assistant_message_is_the_answer() {
992 let mut acc = TurnAccumulator::default();
993 let before_tool = event(
994 "assistant.message",
995 json!({"content": "", "messageId": "m1", "model": "claude-haiku-4.5"}),
996 );
997 let answer = event(
998 "assistant.message",
999 json!({"content": "PELICAN-7741", "messageId": "m2"}),
1000 );
1001 let idle = event("session.idle", json!({}));
1002 assert_eq!(acc.absorb(&before_tool, None).ok(), Some(false));
1003 assert_eq!(acc.absorb(&answer, None).ok(), Some(false));
1004 assert_eq!(acc.absorb(&idle, None).ok(), Some(true));
1005 assert_eq!(acc.finish("m".to_owned()).content, "PELICAN-7741");
1006 }
1007
1008 #[test]
1009 fn session_error_types_map_onto_the_host_error_kinds() {
1010 let cases = [
1011 ("authentication", "auth"),
1012 ("authorization", "auth"),
1013 ("quota", "rate_limit"),
1014 ("rate_limit", "rate_limit"),
1015 ("context_limit", "context_length"),
1016 ("model_call", "external"),
1017 ];
1018 for (error_type, expected) in cases {
1019 let err = map_session_error(&event(
1020 "session.error",
1021 json!({"errorType": error_type, "message": "boom", "errorCode": "user_weekly_rate_limited"}),
1022 ));
1023 let text = err.to_string();
1024 assert!(text.contains("boom"), "{error_type}: message kept: {text}");
1025 assert!(
1026 text.contains("user_weekly_rate_limited"),
1027 "{error_type}: error code kept: {text}"
1028 );
1029 let kind = format!("{:?}", err.kind).to_lowercase();
1030 let matched = match expected {
1031 "auth" => kind.contains("auth"),
1032 "rate_limit" => kind.contains("ratelimit"),
1033 "context_length" => kind.contains("contextlength"),
1034 _ => kind.contains("externalservice"),
1035 };
1036 assert!(matched, "{error_type} -> {kind}");
1037 }
1038 }
1039
1040 #[test]
1041 fn a_provider_404_is_a_model_nobody_serves() {
1042 let err = map_session_error(&event(
1043 "session.error",
1044 json!({
1045 "errorType": "query",
1046 "statusCode": 404,
1047 "message": "Model 'no-such-model' not found on provider at http://localhost:11434/v1 (HTTP 404).",
1048 }),
1049 ));
1050 assert_eq!(err.kind, ErrorKind::ModelUnavailable, "{err}");
1051 assert!(err.to_string().contains("no-such-model"), "{err}");
1052 assert!(err.to_string().contains("localhost:11434"), "{err}");
1053 }
1054
1055 #[test]
1056 fn an_unavailable_model_is_named_once() {
1057 let generic = RunnerError {
1058 kind: ErrorKind::ModelUnavailable,
1059 message: "Resource not found on provider (HTTP 404).".to_owned(),
1060 };
1061 let named = naming_model(generic, "qwen2.5:3b");
1062 assert!(named.message.ends_with("(model \"qwen2.5:3b\")"), "{named}");
1063 assert_eq!(
1064 naming_model(named.clone(), "qwen2.5:3b").message,
1065 named.message,
1066 "a message that already names the model is left alone"
1067 );
1068
1069 let other = naming_model(RunnerError::timeout("slow"), "qwen2.5:3b");
1070 assert!(!other.message.contains("qwen"), "only this kind: {other}");
1071 }
1072
1073 #[test]
1074 fn a_turn_ends_on_session_error() {
1075 let mut acc = TurnAccumulator::default();
1076 let err = event(
1077 "session.error",
1078 json!({"errorType": "quota", "message": "premium requests exhausted"}),
1079 );
1080 assert!(acc.absorb(&err, None).is_err());
1081 }
1082
1083 #[test]
1084 fn mcp_servers_keep_their_headers() {
1085 let transport = McpTransport::Http {
1086 url: "http://127.0.0.1:1/mcp".to_owned(),
1087 headers: vec![McpHeader {
1088 name: "Authorization".to_owned(),
1089 value: "Bearer x".to_owned(),
1090 }],
1091 };
1092 let McpServerConfig::Http(http) = mcp_server(&transport) else {
1093 unreachable!("http stays http");
1094 };
1095 assert_eq!(http.url, "http://127.0.0.1:1/mcp");
1096 assert_eq!(
1097 http.headers.get("Authorization").map(String::as_str),
1098 Some("Bearer x")
1099 );
1100 }
1101
1102 #[test]
1103 fn missing_or_negative_counts_are_zero() {
1104 assert_eq!(count(None), 0);
1105 assert_eq!(count(Some(-5)), 0);
1106 assert_eq!(count(Some(42)), 42);
1107 }
1108
1109 fn mcp_permission(server: &str) -> PermissionRequestData {
1110 let mut data: PermissionRequestData =
1111 serde_json::from_value(json!({"kind": "mcp"})).unwrap_or_else(|e| unreachable!("{e}"));
1112 data.extra = json!({
1113 "requestId": "r-1",
1114 "permissionRequest": {"kind": "mcp", "serverName": server, "toolName": "get_secret_number"}
1115 });
1116 data
1117 }
1118
1119 #[tokio::test]
1120 async fn registered_mcp_servers_are_approved_and_nothing_else_is() {
1121 let handler = RegisteredMcpHandler {
1122 servers: vec!["dravr".to_owned()],
1123 };
1124 assert!(handler.permits(&mcp_permission("dravr")));
1125 assert!(!handler.permits(&mcp_permission("github-mcp-server")));
1126
1127 let mut shell: PermissionRequestData = serde_json::from_value(json!({"kind": "shell"}))
1128 .unwrap_or_else(|e| unreachable!("{e}"));
1129 shell.extra = json!({"permissionRequest": {"kind": "shell", "serverName": "dravr"}});
1130 assert!(
1131 !handler.permits(&shell),
1132 "a shell prompt is not an MCP call"
1133 );
1134
1135 let mut managed = mcp_permission("dravr");
1136 managed.managed_approval_required = Some(true);
1137 assert!(!handler.permits(&managed), "managed policy wants a human");
1138
1139 let approved = handler
1140 .handle(
1141 SessionId::from("s"),
1142 RequestId::new("1"),
1143 mcp_permission("dravr"),
1144 )
1145 .await;
1146 assert!(approves_once(&approved), "{approved:?}");
1147 let refused = handler
1148 .handle(
1149 SessionId::from("s"),
1150 RequestId::new("2"),
1151 mcp_permission("other"),
1152 )
1153 .await;
1154 assert!(!approves_once(&refused), "{refused:?}");
1155 }
1156}