Skip to main content

vtcode_llm/copilot/
server_client.rs

1use std::path::Path;
2use std::str;
3
4use anyhow::{Context, Result, anyhow};
5use serde_json::{Value, json};
6use tokio::io::{AsyncWrite, AsyncWriteExt, BufReader};
7use tokio::process::{Child, ChildStderr, ChildStdout};
8use tokio::time::timeout;
9use vtcode_commons::sanitizer::{PROVIDER_DIAGNOSTIC_MAX_BYTES, sanitize_provider_diagnostic};
10use vtcode_config::auth::CopilotAuthConfig;
11
12use super::command::{resolve_copilot_command, spawn_copilot_server_process};
13use super::transport::{read_bounded_line, trim_line_ending};
14use super::types::CopilotDiscoveredModel;
15
16const MAX_COPILOT_HEADER_BYTES: usize = 8 * 1024;
17
18pub async fn list_available_models(
19    config: &CopilotAuthConfig,
20    workspace_root: &Path,
21) -> Result<Vec<CopilotDiscoveredModel>> {
22    let resolved = resolve_copilot_command(config)?;
23    let mut child = spawn_copilot_server_process(&resolved, workspace_root)?;
24    let stdin = child
25        .stdin
26        .take()
27        .ok_or_else(|| anyhow!("copilot cli server stdin unavailable"))?;
28    let stdout = child
29        .stdout
30        .take()
31        .ok_or_else(|| anyhow!("copilot cli server stdout unavailable"))?;
32    let stderr = child
33        .stderr
34        .take()
35        .ok_or_else(|| anyhow!("copilot cli server stderr unavailable"))?;
36
37    spawn_server_stderr(stderr);
38
39    let result = timeout(resolved.startup_timeout, async move {
40        let mut writer = stdin;
41        let mut reader = BufReader::new(stdout);
42
43        send_request(&mut writer, 1, "ping", Some(json!({ "message": "vtcode model discovery" })))
44            .await
45            .context("copilot cli ping")?;
46        let ping = read_response(&mut reader, 1).await.context("copilot cli ping")?;
47        let protocol_version = ping.get("protocolVersion").and_then(Value::as_i64).unwrap_or(0);
48        if protocol_version <= 0 {
49            return Err(anyhow!("copilot cli server did not report a protocol version"));
50        }
51
52        send_request(&mut writer, 2, "models.list", None)
53            .await
54            .context("copilot cli models.list")?;
55        let payload = read_response(&mut reader, 2).await.context("copilot cli models.list")?;
56        let models = payload
57            .get("models")
58            .and_then(Value::as_array)
59            .ok_or_else(|| anyhow!("copilot cli models.list response missing models"))?;
60
61        let mut discovered = Vec::new();
62        for model in models {
63            let id = model
64                .get("id")
65                .and_then(Value::as_str)
66                .map(str::trim)
67                .filter(|value| !value.is_empty());
68            let name = model
69                .get("name")
70                .and_then(Value::as_str)
71                .map(str::trim)
72                .filter(|value| !value.is_empty());
73            let policy_enabled = model
74                .get("policy")
75                .and_then(Value::as_object)
76                .and_then(|policy| policy.get("state"))
77                .and_then(Value::as_str)
78                .map(|state| state.eq_ignore_ascii_case("enabled"))
79                .unwrap_or(true);
80            if !policy_enabled {
81                continue;
82            }
83            let Some(id) = id else {
84                continue;
85            };
86            discovered.push(CopilotDiscoveredModel {
87                id: id.to_string(),
88                name: name.unwrap_or(id).to_string(),
89            });
90        }
91
92        discovered.sort_by(|left, right| left.id.cmp(&right.id));
93        discovered.dedup_by(|left, right| left.id.eq_ignore_ascii_case(&right.id));
94        Ok::<Vec<CopilotDiscoveredModel>, anyhow::Error>(discovered)
95    })
96    .await;
97
98    terminate_child(&mut child).await;
99    let result = result.context("copilot cli model discovery timeout")??;
100    Ok(result)
101}
102
103async fn terminate_child(child: &mut Child) {
104    if let Err(error) = child.start_kill() {
105        tracing::debug!(target: "copilot.server", error = %error, "copilot cli child already stopped");
106    }
107    if let Err(error) = child.wait().await {
108        tracing::debug!(target: "copilot.server", error = %error, "failed to reap copilot cli child");
109    }
110}
111
112fn spawn_server_stderr(stderr: ChildStderr) {
113    tokio::spawn(async move {
114        let mut reader = BufReader::new(stderr);
115        let mut line = Vec::with_capacity(PROVIDER_DIAGNOSTIC_MAX_BYTES);
116        loop {
117            match read_bounded_line(&mut reader, &mut line, PROVIDER_DIAGNOSTIC_MAX_BYTES).await {
118                Ok(Some(_truncated)) => {
119                    let trimmed = trim_line_ending(&line);
120                    if !trimmed.iter().all(u8::is_ascii_whitespace) {
121                        let safe_line = sanitize_provider_diagnostic(trimmed);
122                        tracing::debug!(target: "copilot.server.stderr", "{}", safe_line);
123                    }
124                }
125                Ok(None) => break,
126                Err(error) => {
127                    tracing::warn!(target: "copilot.server.stderr", error = %error, "stderr reader failed");
128                    break;
129                }
130            }
131        }
132    });
133}
134
135async fn send_request<W>(writer: &mut W, id: i64, method: &str, params: Option<Value>) -> Result<()>
136where
137    W: AsyncWrite + Unpin,
138{
139    let message = if let Some(params) = params {
140        json!({
141            "jsonrpc": "2.0",
142            "id": id,
143            "method": method,
144            "params": params,
145        })
146    } else {
147        json!({
148            "jsonrpc": "2.0",
149            "id": id,
150            "method": method,
151        })
152    };
153    let payload = serde_json::to_vec(&message).context("copilot cli json serialization failed")?;
154    writer
155        .write_all(format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes())
156        .await
157        .context("copilot cli write header failed")?;
158    writer.write_all(&payload).await.context("copilot cli write payload failed")?;
159    writer.flush().await.context("copilot cli flush failed")?;
160    Ok(())
161}
162
163async fn read_response(reader: &mut BufReader<ChildStdout>, expected_id: i64) -> Result<Value> {
164    loop {
165        let message = read_message(reader).await?;
166        let Some(object) = message.as_object() else {
167            continue;
168        };
169
170        if object.get("method").is_some() {
171            continue;
172        }
173
174        if let Some(error) = object.get("error") {
175            let code = error.get("code").and_then(Value::as_i64).unwrap_or_default();
176            let detail = error.get("message").and_then(Value::as_str).unwrap_or("unknown error");
177            return Err(anyhow!("copilot cli rpc error {code}: {detail}"));
178        }
179
180        if object.get("id").and_then(Value::as_i64) != Some(expected_id) {
181            continue;
182        }
183
184        return object
185            .get("result")
186            .cloned()
187            .ok_or_else(|| anyhow!("copilot cli rpc response missing result"));
188    }
189}
190
191/// Upper bound on a single Copilot CLI JSON-RPC payload.
192///
193/// `Content-Length` arrives from the spawned Copilot CLI child process and is
194/// passed straight to `read_exact_uninit`, which allocates that many bytes up
195/// front. A malformed or hostile stream advertising an enormous length would
196/// otherwise drive the process into an out-of-memory abort. Copilot responses
197/// (model lists, completions) are well under this cap; reject anything larger.
198const MAX_COPILOT_PAYLOAD_BYTES: usize = 64 * 1024 * 1024;
199
200async fn read_message(reader: &mut BufReader<ChildStdout>) -> Result<Value> {
201    let mut content_length = None;
202    let mut line = Vec::with_capacity(MAX_COPILOT_HEADER_BYTES);
203    loop {
204        let Some(truncated) = read_bounded_line(reader, &mut line, MAX_COPILOT_HEADER_BYTES)
205            .await
206            .context("copilot cli header read failed")?
207        else {
208            return Err(anyhow!("copilot cli server closed the stdio stream"));
209        };
210        if truncated {
211            return Err(anyhow!("copilot cli header exceeds {MAX_COPILOT_HEADER_BYTES} byte limit"));
212        }
213
214        let trimmed_bytes = trim_line_ending(&line);
215        let trimmed = str::from_utf8(trimmed_bytes).context("copilot cli header is not valid UTF-8")?;
216        if trimmed.is_empty() {
217            break;
218        }
219
220        if let Some(value) = trimmed.strip_prefix("Content-Length:") {
221            let parsed = value
222                .trim()
223                .parse::<usize>()
224                .context("invalid copilot cli content length header")?;
225            if parsed > MAX_COPILOT_PAYLOAD_BYTES {
226                return Err(anyhow!(
227                    "copilot cli content length {parsed} exceeds {MAX_COPILOT_PAYLOAD_BYTES} byte limit"
228                ));
229            }
230            content_length = Some(parsed);
231        }
232    }
233
234    let content_length = content_length.ok_or_else(|| anyhow!("copilot cli response missing Content-Length"))?;
235    let payload = vtcode_commons::async_utils::read_exact_uninit(reader, content_length)
236        .await
237        .context("copilot cli payload read failed")?;
238    serde_json::from_slice(&payload).context("copilot cli json decode failed")
239}