Skip to content

Commit 8794594

Browse files
stephentoubCopilot
andauthored
Expand Rust E2E coverage (#1250)
* Expand Rust E2E coverage Add a replay-backed Rust E2E suite matching the .NET coverage, update Rust SDK session lifecycle support for session filesystem and multi-client scenarios, and add reliability fixes for model caching and cancellation-safe pending session registration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address Rust E2E review feedback Move struct-only E2E placeholders into unit tests, exercise invalid external-auth client options, avoid logging session IDs from test assertions, and harden failing Rust E2Es on CI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix Rust CI after review updates Allow inert use_logged_in_user(false) on external transports so shared E2E client options continue to work, while still rejecting true external auth requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Stabilize Rust selection attachment E2E Use a workspace-relative selection file path so the replayed prompt is stable across platforms instead of containing platform-specific temp directory relatives. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Make Rust E2E CI fail fast Add a per-test E2E timeout and run Rust CI tests serially with uncaptured output so stuck replay-backed tests expose the active test instead of hanging silently. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix Rust E2E suite runtime Replace the global Rust E2E lock with bounded replay-test concurrency so the suite no longer serializes every replay-backed case. Wait for disconnected-client tool removal before sending the follow-up multi-client prompt to keep the concurrent run deterministic. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Use bounded Rust E2E concurrency in CI Set libtest to four worker threads so GitHub-hosted runners actually exercise the harness concurrency limit instead of defaulting to the runner CPU count. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix Rust session tests for generated IDs Update session_test fake servers to echo the session id requested by the SDK instead of returning arbitrary ids. This keeps the tests aligned with the SDK's session-id validation and prevents cargo test from failing before the E2E suite runs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address Rust E2E review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix Rust multi-client clippy Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 4901dff commit 8794594

164 files changed

Lines changed: 17958 additions & 783 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/workflows/rust-sdk-tests.yml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,10 +94,12 @@ jobs:
9494
run: pwsh.exe -Command "Write-Host 'PowerShell ready'"
9595

9696
- name: cargo test
97+
timeout-minutes: 90
9798
env:
99+
RUST_E2E_CONCURRENCY: 4
98100
COPILOT_HMAC_KEY: ${{ secrets.COPILOT_DEVELOPER_CLI_INTEGRATION_HMAC_KEY }}
99101
COPILOT_CLI_PATH: ${{ steps.setup-copilot.outputs.cli-path }}
100-
run: cargo test --features test-support
102+
run: cargo test --features test-support -- --test-threads=4 --nocapture
101103

102104
# Validates the `embedded-cli` build path on all three supported
103105
# platforms. This is the only place `build.rs` actually runs (the

rust/Cargo.lock

Lines changed: 10 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

rust/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ parking_lot = "0.12"
5151
regex = "1"
5252
sha2 = { version = "0.10", optional = true }
5353
getrandom = "0.2"
54+
uuid = { version = "1", default-features = false, features = ["v4"] }
5455
zstd = { version = "0.13", optional = true }
5556

5657
[dev-dependencies]

rust/src/jsonrpc.rs

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,12 @@ use std::sync::Arc;
33
use std::sync::atomic::{AtomicU64, Ordering};
44
use std::time::Instant;
55

6-
use parking_lot::RwLock;
6+
use parking_lot::{Mutex, RwLock};
77
use serde::{Deserialize, Serialize};
88
use serde_json::Value;
99
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
1010
use tokio::sync::{broadcast, mpsc, oneshot};
11+
use tokio::task::JoinHandle;
1112
use tracing::{Instrument, debug, error, warn};
1213

1314
use crate::{Error, ProtocolError};
@@ -184,6 +185,8 @@ pub struct JsonRpcClient {
184185
pending_requests: Arc<RwLock<HashMap<u64, oneshot::Sender<JsonRpcResponse>>>>,
185186
notification_tx: broadcast::Sender<JsonRpcNotification>,
186187
request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
188+
read_task: Mutex<Option<JoinHandle<()>>>,
189+
write_task: Mutex<Option<JoinHandle<()>>>,
187190
}
188191

189192
impl JsonRpcClient {
@@ -202,22 +205,24 @@ impl JsonRpcClient {
202205
let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteCommand>();
203206

204207
let writer_span = tracing::error_span!("jsonrpc_write_loop");
205-
tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span));
208+
let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span));
206209

207210
let client = Self {
208211
request_id: AtomicU64::new(1),
209212
write_tx,
210213
pending_requests: Arc::new(RwLock::new(HashMap::new())),
211214
notification_tx,
212215
request_tx,
216+
read_task: Mutex::new(None),
217+
write_task: Mutex::new(Some(write_task)),
213218
};
214219

215220
let pending_requests = client.pending_requests.clone();
216221
let notification_tx_clone = client.notification_tx.clone();
217222
let request_tx_clone = client.request_tx.clone();
218223
let reader_span = tracing::error_span!("jsonrpc_read_loop");
219224

220-
tokio::spawn(
225+
let read_task = tokio::spawn(
221226
async move {
222227
Self::read_loop(
223228
reader,
@@ -229,10 +234,21 @@ impl JsonRpcClient {
229234
}
230235
.instrument(reader_span),
231236
);
237+
*client.read_task.lock() = Some(read_task);
232238

233239
client
234240
}
235241

242+
pub(crate) fn force_close(&self) {
243+
if let Some(task) = self.read_task.lock().take() {
244+
task.abort();
245+
}
246+
if let Some(task) = self.write_task.lock().take() {
247+
task.abort();
248+
}
249+
self.pending_requests.write().clear();
250+
}
251+
236252
/// Writer-actor task. Owns the `AsyncWrite`, drains the command queue,
237253
/// and writes each frame atomically (header + body + flush) before
238254
/// signaling the ack.

rust/src/lib.rs

Lines changed: 166 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,15 @@ pub enum SessionError {
267267
/// non-empty.
268268
#[error("invalid SessionFsConfig: {0}")]
269269
InvalidSessionFsConfig(String),
270+
271+
/// The CLI returned a different session ID than the one the SDK registered.
272+
#[error("CLI returned session ID {returned} after SDK registered {requested}")]
273+
SessionIdMismatch {
274+
/// Session ID registered by the SDK before the RPC was sent.
275+
requested: SessionId,
276+
/// Session ID returned by the CLI.
277+
returned: SessionId,
278+
},
270279
}
271280

272281
/// How the SDK communicates with the CLI server.
@@ -873,6 +882,7 @@ struct ClientInner {
873882
state: parking_lot::Mutex<ConnectionState>,
874883
lifecycle_tx: broadcast::Sender<SessionLifecycleEvent>,
875884
on_list_models: Option<Arc<dyn ListModelsHandler>>,
885+
models_cache: parking_lot::Mutex<Arc<tokio::sync::OnceCell<Vec<Model>>>>,
876886
session_fs_configured: bool,
877887
on_get_trace_context: Option<Arc<dyn TraceContextProvider>>,
878888
/// Token sent in the `connect` handshake. Auto-generated when the
@@ -900,6 +910,24 @@ impl Client {
900910
if let Some(cfg) = &options.session_fs {
901911
validate_session_fs_config(cfg)?;
902912
}
913+
// Auth options only make sense when the SDK spawns the CLI; with an
914+
// external server, the server manages its own auth.
915+
if matches!(options.transport, Transport::External { .. }) {
916+
if options.github_token.is_some() {
917+
return Err(Error::InvalidConfig(
918+
"github_token cannot be used with Transport::External \
919+
(external server manages its own auth)"
920+
.to_string(),
921+
));
922+
}
923+
if options.use_logged_in_user == Some(true) {
924+
return Err(Error::InvalidConfig(
925+
"use_logged_in_user cannot be used with Transport::External \
926+
(external server manages its own auth)"
927+
.to_string(),
928+
));
929+
}
930+
}
903931
// Validate token + transport combination. Stdio cannot use a
904932
// connection token; auto-generate a UUID when the SDK spawns
905933
// its own CLI in TCP mode and no explicit token was set.
@@ -1138,6 +1166,7 @@ impl Client {
11381166
state: parking_lot::Mutex::new(ConnectionState::Connected),
11391167
lifecycle_tx: broadcast::channel(256).0,
11401168
on_list_models,
1169+
models_cache: parking_lot::Mutex::new(Arc::new(tokio::sync::OnceCell::new())),
11411170
session_fs_configured,
11421171
on_get_trace_context,
11431172
effective_connection_token,
@@ -1752,10 +1781,17 @@ impl Client {
17521781
/// When [`ClientOptions::on_list_models`] is set, returns the handler's
17531782
/// result without making a `models.list` RPC. Otherwise queries the CLI.
17541783
pub async fn list_models(&self) -> Result<Vec<Model>, Error> {
1755-
if let Some(handler) = &self.inner.on_list_models {
1756-
return handler.list_models().await;
1757-
}
1758-
Ok(self.rpc().models().list().await?.models)
1784+
let cache = self.inner.models_cache.lock().clone();
1785+
let models = cache
1786+
.get_or_try_init(|| async {
1787+
if let Some(handler) = &self.inner.on_list_models {
1788+
handler.list_models().await
1789+
} else {
1790+
Ok(self.rpc().models().list().await?.models)
1791+
}
1792+
})
1793+
.await?;
1794+
Ok(models.clone())
17591795
}
17601796

17611797
/// Invoke [`ClientOptions::on_get_trace_context`] when configured,
@@ -1828,6 +1864,7 @@ impl Client {
18281864

18291865
let child = self.inner.child.lock().take();
18301866
*self.inner.state.lock() = ConnectionState::Disconnected;
1867+
*self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new());
18311868
if let Some(mut child) = child
18321869
&& let Err(e) = child.kill().await
18331870
{
@@ -1879,10 +1916,12 @@ impl Client {
18791916
{
18801917
error!(pid = ?pid, error = %e, "failed to send kill signal");
18811918
}
1919+
self.inner.rpc.force_close();
18821920
// Drop all session channels so any awaiters see a closed channel
18831921
// instead of waiting for responses that will never arrive.
18841922
self.inner.router.clear();
18851923
*self.inner.state.lock() = ConnectionState::Disconnected;
1924+
*self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new());
18861925
}
18871926

18881927
/// Subscribe to lifecycle events.
@@ -2405,43 +2444,137 @@ mod tests {
24052444
policy: None,
24062445
supported_reasoning_efforts: Vec::new(),
24072446
};
2408-
let handler = Arc::new(CountingHandler {
2447+
let handler: Arc<dyn ListModelsHandler> = Arc::new(CountingHandler {
24092448
calls: Arc::clone(&calls),
24102449
models: vec![model.clone()],
24112450
});
24122451

2413-
// We can't call list_models() through Client::start without a CLI, but we
2414-
// can exercise the override path by directly constructing a Client whose
2415-
// inner has the handler set. This is the same dispatch path as the real
2416-
// call; from_streams's None default is replaced via inner construction.
2417-
let inner = ClientInner {
2418-
child: parking_lot::Mutex::new(None),
2419-
rpc: {
2420-
let (req_tx, _req_rx) = mpsc::unbounded_channel();
2421-
let (notif_tx, _notif_rx) = broadcast::channel(16);
2422-
let (read_pipe, _write_pipe) = tokio::io::duplex(64);
2423-
let (_unused_read, write_pipe) = tokio::io::duplex(64);
2424-
JsonRpcClient::new(write_pipe, read_pipe, notif_tx, req_tx)
2425-
},
2426-
cwd: PathBuf::from("."),
2427-
request_rx: parking_lot::Mutex::new(None),
2428-
notification_tx: broadcast::channel(16).0,
2429-
router: router::SessionRouter::new(),
2430-
negotiated_protocol_version: OnceLock::new(),
2431-
state: parking_lot::Mutex::new(ConnectionState::Connected),
2432-
lifecycle_tx: broadcast::channel(16).0,
2433-
on_list_models: Some(handler),
2434-
session_fs_configured: false,
2435-
on_get_trace_context: None,
2436-
effective_connection_token: None,
2437-
};
2438-
let client = Client {
2439-
inner: Arc::new(inner),
2440-
};
2452+
let client = client_with_list_models_handler(handler);
24412453

24422454
let result = client.list_models().await.unwrap();
24432455
assert_eq!(result.len(), 1);
24442456
assert_eq!(result[0].id, "byok-gpt-4");
24452457
assert_eq!(calls.load(Ordering::SeqCst), 1);
24462458
}
2459+
2460+
#[tokio::test]
2461+
async fn list_models_serializes_concurrent_cache_misses() {
2462+
use std::sync::atomic::{AtomicUsize, Ordering};
2463+
2464+
struct SlowCountingHandler {
2465+
calls: Arc<AtomicUsize>,
2466+
models: Vec<Model>,
2467+
}
2468+
#[async_trait]
2469+
impl ListModelsHandler for SlowCountingHandler {
2470+
async fn list_models(&self) -> Result<Vec<Model>, Error> {
2471+
self.calls.fetch_add(1, Ordering::SeqCst);
2472+
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
2473+
Ok(self.models.clone())
2474+
}
2475+
}
2476+
2477+
let calls = Arc::new(AtomicUsize::new(0));
2478+
let model = Model {
2479+
billing: None,
2480+
capabilities: ModelCapabilities {
2481+
limits: None,
2482+
supports: None,
2483+
},
2484+
default_reasoning_effort: None,
2485+
id: "single-flight-model".into(),
2486+
name: "Single Flight Model".into(),
2487+
policy: None,
2488+
supported_reasoning_efforts: Vec::new(),
2489+
};
2490+
let handler: Arc<dyn ListModelsHandler> = Arc::new(SlowCountingHandler {
2491+
calls: Arc::clone(&calls),
2492+
models: vec![model],
2493+
});
2494+
let client = client_with_list_models_handler(handler);
2495+
2496+
let (first, second) = tokio::join!(client.list_models(), client.list_models());
2497+
assert_eq!(first.unwrap()[0].id, "single-flight-model");
2498+
assert_eq!(second.unwrap()[0].id, "single-flight-model");
2499+
assert_eq!(calls.load(Ordering::SeqCst), 1);
2500+
}
2501+
2502+
#[tokio::test]
2503+
async fn cancelled_create_session_unregisters_pending_session() {
2504+
let (client_write, _server_read) = tokio::io::duplex(8192);
2505+
let (_server_write, client_read) = tokio::io::duplex(8192);
2506+
let client = Client::from_streams(client_read, client_write, std::env::temp_dir()).unwrap();
2507+
let handle = tokio::spawn({
2508+
let client = client.clone();
2509+
async move { client.create_session(SessionConfig::default()).await }
2510+
});
2511+
2512+
wait_for_pending_session_registration(&client).await;
2513+
handle.abort();
2514+
let _ = handle.await;
2515+
2516+
assert!(client.inner.router.session_ids().is_empty());
2517+
client.force_stop();
2518+
}
2519+
2520+
#[tokio::test]
2521+
async fn cancelled_resume_session_unregisters_pending_session() {
2522+
let (client_write, _server_read) = tokio::io::duplex(8192);
2523+
let (_server_write, client_read) = tokio::io::duplex(8192);
2524+
let client = Client::from_streams(client_read, client_write, std::env::temp_dir()).unwrap();
2525+
let session_id = SessionId::new("resume-cancel-test");
2526+
let handle = tokio::spawn({
2527+
let client = client.clone();
2528+
async move {
2529+
client
2530+
.resume_session(ResumeSessionConfig::new(session_id))
2531+
.await
2532+
}
2533+
});
2534+
2535+
wait_for_pending_session_registration(&client).await;
2536+
handle.abort();
2537+
let _ = handle.await;
2538+
2539+
assert!(client.inner.router.session_ids().is_empty());
2540+
client.force_stop();
2541+
}
2542+
2543+
fn client_with_list_models_handler(handler: Arc<dyn ListModelsHandler>) -> Client {
2544+
Client {
2545+
inner: Arc::new(ClientInner {
2546+
child: parking_lot::Mutex::new(None),
2547+
rpc: {
2548+
let (req_tx, _req_rx) = mpsc::unbounded_channel();
2549+
let (notif_tx, _notif_rx) = broadcast::channel(16);
2550+
let (read_pipe, _write_pipe) = tokio::io::duplex(64);
2551+
let (_unused_read, write_pipe) = tokio::io::duplex(64);
2552+
JsonRpcClient::new(write_pipe, read_pipe, notif_tx, req_tx)
2553+
},
2554+
cwd: PathBuf::from("."),
2555+
request_rx: parking_lot::Mutex::new(None),
2556+
notification_tx: broadcast::channel(16).0,
2557+
router: router::SessionRouter::new(),
2558+
negotiated_protocol_version: OnceLock::new(),
2559+
state: parking_lot::Mutex::new(ConnectionState::Connected),
2560+
lifecycle_tx: broadcast::channel(16).0,
2561+
on_list_models: Some(handler),
2562+
models_cache: parking_lot::Mutex::new(Arc::new(tokio::sync::OnceCell::new())),
2563+
session_fs_configured: false,
2564+
on_get_trace_context: None,
2565+
effective_connection_token: None,
2566+
}),
2567+
}
2568+
}
2569+
2570+
async fn wait_for_pending_session_registration(client: &Client) {
2571+
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
2572+
while client.inner.router.session_ids().is_empty() {
2573+
assert!(
2574+
tokio::time::Instant::now() < deadline,
2575+
"session was not registered"
2576+
);
2577+
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
2578+
}
2579+
}
24472580
}

0 commit comments

Comments
 (0)