Line data Source code
1 : //! The local agent server: DashQL clients connect over WebSocket, the command line tool over HTTP
2 :
3 : use crate::agent::commands::is_session_command;
4 : use crate::agent::config::{
5 : client_origins, create_private_dir, random_bytes, to_json_text, write_private_file, AgentConfig, HarnessChoice,
6 : };
7 : use crate::agent::harness::{
8 : build_harness_args, resolve_login_shell_path, resolve_program, spawn_harness, HarnessEvent, HarnessOutcome,
9 : HarnessPicks, HarnessProcess, HarnessSpawnArgs, Kill, KILL_GRACE,
10 : };
11 : use crate::agent::http::{self, Request};
12 : use crate::agent::instructions::{
13 : claude_instructions, page_features, session_instructions, Page, INPUT_FOLDER, OPENCODE_CONFIG,
14 : };
15 : use crate::agent::protocol::{ErrorCode, RpcError, PROTOCOL_VERSION, WEBSOCKET_PATH};
16 : use futures_util::{SinkExt, StreamExt};
17 : use serde_json::{json, Map, Value};
18 : use std::borrow::Cow;
19 : use std::collections::{BTreeMap, HashMap, HashSet};
20 : use std::ffi::OsString;
21 : use std::future::Future;
22 : use std::path::{Path, PathBuf};
23 : use std::pin::Pin;
24 : use std::sync::{Arc, Mutex, MutexGuard};
25 : use std::time::{Duration, SystemTime, UNIX_EPOCH};
26 : use tokio::io::AsyncWriteExt;
27 : use tokio::net::{TcpListener, TcpStream};
28 : use tokio::sync::{mpsc, oneshot, Notify, OnceCell};
29 : use tokio_tungstenite::tungstenite::handshake::derive_accept_key;
30 : use tokio_tungstenite::tungstenite::protocol::{frame::coding::CloseCode, CloseFrame, Role, WebSocketConfig};
31 : use tokio_tungstenite::tungstenite::Message;
32 : use tokio_tungstenite::WebSocketStream;
33 :
34 : const DEFAULT_CALL_TIMEOUT: Duration = Duration::from_secs(10 * 60);
35 : const MAX_HTTP_BODY_BYTES: usize = 8 * 1024 * 1024;
36 : const CLOSE_POLICY_VIOLATION: u16 = 1008;
37 :
38 : pub const NOT_CONNECTED_MESSAGE: &str = "DashQL is not connected: open https://dashql.app and turn on agent access";
39 : const DISCONNECTED_MESSAGE: &str = "DashQL disconnected during the call";
40 : const SERVER_STOPPED_MESSAGE: &str = "The agent server stopped";
41 : /// How long a stopping server waits for its clients to take the close frame
42 : const CLOSE_GRACE: Duration = Duration::from_secs(1);
43 : /// Set in the environment of harness runs
44 : pub const AGENT_RUN_ENV: &str = "DASHQL_AGENT_RUN";
45 : /// How long the folder of an unused session stays
46 : const SESSION_RETENTION: Duration = Duration::from_secs(30 * 24 * 60 * 60);
47 : /// How often the server removes the folders of unused sessions
48 : const SESSION_SWEEP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
49 :
50 : /// Messages and output lines are cut after this many bytes in the log
51 : const LOG_TEXT_LIMIT: usize = 64 * 1024;
52 :
53 : /// A text for the log, cut after LOG_TEXT_LIMIT bytes
54 220 : fn clip(text: &str) -> Cow<'_, str> {
55 220 : if text.len() <= LOG_TEXT_LIMIT {
56 220 : return Cow::Borrowed(text);
57 0 : }
58 0 : let mut end = LOG_TEXT_LIMIT;
59 0 : while !text.is_char_boundary(end) {
60 0 : end -= 1;
61 0 : }
62 0 : Cow::Owned(format!("{}... ({} more bytes)", &text[..end], text.len() - end))
63 220 : }
64 :
65 : /// The current time in UTC, e.g. 2026-10-29T12:34:56.789Z
66 0 : fn timestamp() -> String {
67 0 : let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default();
68 0 : let (days, seconds) = ((now.as_secs() / 86400) as i64, now.as_secs() % 86400);
69 : // The civil date of a day number, Howard Hinnant's days_from_civil inverted
70 0 : let z = days + 719_468;
71 0 : let era = z.div_euclid(146_097);
72 0 : let doe = z - era * 146_097;
73 0 : let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
74 0 : let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
75 0 : let mp = (5 * doy + 2) / 153;
76 0 : let day = doy - (153 * mp + 2) / 5 + 1;
77 0 : let month = if mp < 10 { mp + 3 } else { mp - 9 };
78 0 : let year = yoe + era * 400 + i64::from(month <= 2);
79 0 : format!(
80 : "{year:04}-{month:02}-{day:02}T{:02}:{:02}:{:02}.{:03}Z",
81 0 : seconds / 3600,
82 0 : seconds / 60 % 60,
83 0 : seconds % 60,
84 0 : now.subsec_millis()
85 : )
86 0 : }
87 :
88 : /// Write a line to the log on stderr
89 0 : pub fn log_line(message: &str) {
90 0 : eprintln!("[dashql] {} {message}", timestamp());
91 0 : }
92 :
93 : /// Start a harness process, replaceable in tests
94 : pub type SpawnHarness = Arc<dyn Fn(HarnessSpawnArgs) -> HarnessProcess + Send + Sync>;
95 : /// Resolve the PATH of harness runs, replaceable in tests
96 : pub type ResolvePath = Arc<dyn Fn() -> Pin<Box<dyn Future<Output = Option<OsString>> + Send>> + Send + Sync>;
97 :
98 : pub struct AgentServerOptions {
99 : pub config: AgentConfig,
100 : /// The folder of the state, ~/.dashql/agent
101 : pub home: PathBuf,
102 : /// The executable that harness runs call as `dashql`, its folder goes first on their PATH.
103 : /// None, or an executable with another name, leaves their PATH as it is.
104 : pub cli: Option<PathBuf>,
105 : pub spawn_harness: SpawnHarness,
106 : pub resolve_path: ResolvePath,
107 : /// The maximum size of a message
108 : pub max_message_bytes: usize,
109 : /// The interval of WebSocket keepalive pings, a client that misses one pong is dropped
110 : pub ping_interval: Duration,
111 : /// How long a call waits for a call result
112 : pub call_timeout: Duration,
113 : /// Log to stderr
114 : pub log: bool,
115 : }
116 :
117 : impl AgentServerOptions {
118 27 : pub fn new(config: AgentConfig, home: PathBuf) -> Self {
119 : AgentServerOptions {
120 27 : config,
121 27 : home,
122 27 : cli: None,
123 27 : spawn_harness: Arc::new(spawn_harness),
124 27 : resolve_path: Arc::new(|| Box::pin(resolve_login_shell_path()) as Pin<Box<dyn Future<Output = _> + Send>>),
125 27 : max_message_bytes: 32 * 1024 * 1024,
126 27 : ping_interval: Duration::from_secs(30),
127 : call_timeout: DEFAULT_CALL_TIMEOUT,
128 : log: false,
129 : }
130 27 : }
131 : }
132 :
133 : /// A message to a client
134 : enum Outgoing {
135 : Text(String),
136 : Close(u16, String),
137 : }
138 :
139 : #[derive(Clone, Debug)]
140 : struct ClientSession {
141 : id: String,
142 : name: String,
143 : }
144 :
145 : /// A connected client and the sessions it holds
146 : struct Client {
147 : id: String,
148 : origin: String,
149 : kind: String,
150 : sessions: Vec<ClientSession>,
151 : focused: Option<String>,
152 : /// The features of the client that the guide knows (`page_features`), None until it announced its sessions.
153 : /// Identifiers only, the guide of the agent picks its notes for them.
154 : features: Option<Vec<String>>,
155 : /// When the client announced its focus, the most recent focus wins
156 : focused_at: u64,
157 : /// When the client announced its sessions, the most recent client holding a session receives its calls
158 : announced_at: u64,
159 : tx: mpsc::UnboundedSender<Outgoing>,
160 : /// The calls that wait for their result
161 : pending: HashMap<String, oneshot::Sender<Result<Value, RpcError>>>,
162 : }
163 :
164 : impl Client {
165 95 : fn holds(&self, session: &str) -> bool {
166 100 : self.sessions.iter().any(|s| s.id == session)
167 95 : }
168 :
169 99 : fn send(&self, message: &Value) {
170 99 : let _ = self.tx.send(Outgoing::Text(message.to_string()));
171 99 : }
172 : }
173 :
174 : #[derive(Clone, Copy, Debug, PartialEq, Eq)]
175 : enum TurnState {
176 : Queued,
177 : Running,
178 : }
179 :
180 : impl TurnState {
181 27 : fn name(self) -> &'static str {
182 27 : match self {
183 6 : TurnState::Queued => "queued",
184 21 : TurnState::Running => "running",
185 : }
186 27 : }
187 : }
188 :
189 : /// A turn of a session: a prompt that a harness works on until it answers
190 : struct Turn {
191 : turn_id: String,
192 : session: String,
193 : harness: String,
194 : picks: HarnessPicks,
195 : prompt: String,
196 : state: TurnState,
197 : status: String,
198 : cancelled: bool,
199 : kill: Option<Kill>,
200 : }
201 :
202 : /// The state.json of a DashQL session
203 : #[derive(Default)]
204 : struct SessionState {
205 : /// The harness session ids by harness name
206 : harness_sessions: Map<String, Value>,
207 : }
208 :
209 : struct State {
210 : config: AgentConfig,
211 : clients: BTreeMap<u64, Client>,
212 : next_client: u64,
213 : /// The sessions with a running call
214 : busy: HashSet<String>,
215 : /// The turns in the order of their requests, the first turn of a session runs, the others wait
216 : turns: Vec<Turn>,
217 : /// The deleted sessions whose folders go once their running turns finish
218 : forgotten: HashSet<String>,
219 : /// The harness processes that did not exit yet
220 : harnesses: usize,
221 : next_call_id: u64,
222 : /// Orders announcements and focus changes
223 : clock: u64,
224 : }
225 :
226 : struct Shared {
227 : options: AgentServerOptions,
228 : origins: Vec<String>,
229 : port: u16,
230 : state: Mutex<State>,
231 : /// Notified when the sessions of a client change
232 : sessions_changed: Notify,
233 : path: OnceCell<Option<OsString>>,
234 : }
235 :
236 : /// A folder name for a session id
237 104 : pub fn session_folder_name(session: &str) -> String {
238 229 : session.chars().map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' }).collect()
239 104 : }
240 :
241 46 : fn random_uuid() -> String {
242 46 : let mut bytes = random_bytes::<16>();
243 46 : bytes[6] = (bytes[6] & 0x0f) | 0x40;
244 46 : bytes[8] = (bytes[8] & 0x3f) | 0x80;
245 736 : let hex: String = bytes.iter().map(|b| format!("{b:02x}")).collect();
246 46 : format!("{}-{}-{}-{}-{}", &hex[0..8], &hex[8..12], &hex[12..16], &hex[16..20], &hex[20..])
247 46 : }
248 :
249 : /// The harness of a request, the default harness if it names none.
250 : /// Clients may only pick among the configured harnesses: commands come from config.json alone, a page must never be
251 : /// able to run a program of its choice.
252 28 : fn pick_harness(config: &AgentConfig, params: &Value) -> Result<String, RpcError> {
253 28 : let harness = params["harness"].as_str().unwrap_or(&config.default_harness);
254 28 : match config.harness(harness) {
255 26 : Some(_) => Ok(harness.to_string()),
256 2 : None => Err(RpcError::new(ErrorCode::Failed, format!("unknown harness: {harness}"))),
257 : }
258 28 : }
259 :
260 : /// The model and the effort of a request, among the ones of its harness in config.json.
261 : /// Like the harness, a page may only pick a value: the values become arguments of the harness.
262 22 : fn pick_choices(config: &AgentConfig, harness: &str, params: &Value) -> Result<HarnessPicks, RpcError> {
263 22 : let Some(harness_config) = config.harness(harness) else {
264 0 : return Err(RpcError::new(ErrorCode::Failed, format!("unknown harness: {harness}")));
265 : };
266 42 : let pick = |key: &str, choice: &HarnessChoice| match ¶ms[key] {
267 37 : Value::Null => Ok(None),
268 5 : Value::String(value) if choice.values.contains(value) => Ok(Some(value.clone())),
269 3 : Value::String(value) => Err(RpcError::new(ErrorCode::Failed, format!("unknown {key} of {harness}: {value}"))),
270 0 : _ => Err(RpcError::new(ErrorCode::Usage, format!("the {key} must be a string"))),
271 42 : };
272 22 : Ok(HarnessPicks { model: pick("model", &harness_config.model)?, effort: pick("effort", &harness_config.effort)? })
273 22 : }
274 :
275 9 : fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
276 136 : a.len() == b.len() && a.iter().zip(b).fold(0, |acc, (x, y)| acc | (x ^ y)) == 0
277 9 : }
278 :
279 : /// The local agent server
280 : pub struct AgentServer {
281 : shared: Arc<Shared>,
282 : accept: tokio::task::JoinHandle<()>,
283 : sweep: tokio::task::JoinHandle<()>,
284 : }
285 :
286 : impl AgentServer {
287 : /// Listen on 127.0.0.1
288 27 : pub async fn listen(options: AgentServerOptions) -> std::io::Result<AgentServer> {
289 27 : let listener = TcpListener::bind(("127.0.0.1", options.config.port)).await?;
290 27 : let port = listener.local_addr()?.port();
291 27 : let mut origins: Vec<String> = client_origins().into_iter().map(str::to_string).collect();
292 27 : origins.extend(options.config.origins.iter().cloned());
293 27 : let shared = Arc::new(Shared {
294 27 : state: Mutex::new(State {
295 27 : config: options.config.clone(),
296 27 : clients: BTreeMap::new(),
297 27 : next_client: 0,
298 27 : busy: HashSet::new(),
299 27 : turns: Vec::new(),
300 27 : forgotten: HashSet::new(),
301 27 : harnesses: 0,
302 27 : next_call_id: 1,
303 27 : clock: 0,
304 27 : }),
305 27 : options,
306 27 : origins,
307 27 : port,
308 27 : sessions_changed: Notify::new(),
309 27 : path: OnceCell::new(),
310 27 : });
311 27 : shared.log(&format!("Listening on 127.0.0.1:{port}"));
312 27 : let accept = tokio::spawn({
313 27 : let shared = shared.clone();
314 27 : async move {
315 : loop {
316 66 : match listener.accept().await {
317 39 : Ok((stream, _)) => {
318 39 : tokio::spawn(handle_connection(shared.clone(), stream));
319 39 : }
320 0 : Err(e) => shared.log(&format!("Cannot accept a connection: {e}")),
321 : }
322 : }
323 : }
324 : });
325 : // The folders of unused sessions go at the start and then once a day
326 27 : shared.sweep_sessions();
327 27 : let sweep = tokio::spawn({
328 27 : let shared = shared.clone();
329 27 : async move {
330 27 : let start = tokio::time::Instant::now() + SESSION_SWEEP_INTERVAL;
331 27 : let mut interval = tokio::time::interval_at(start, SESSION_SWEEP_INTERVAL);
332 : loop {
333 27 : interval.tick().await;
334 0 : shared.sweep_sessions();
335 : }
336 : }
337 : });
338 27 : Ok(AgentServer { shared, accept, sweep })
339 27 : }
340 :
341 : /// The port the server listens on
342 : #[cfg(test)]
343 63 : pub fn port(&self) -> u16 {
344 63 : self.shared.port
345 63 : }
346 :
347 : /// The state of the server
348 : #[cfg(test)]
349 3 : pub fn status(&self) -> Value {
350 3 : self.shared.status(&self.shared.lock())
351 3 : }
352 :
353 : /// Remove the folders of unused sessions now
354 : #[cfg(test)]
355 1 : pub fn sweep_sessions(&self) {
356 1 : self.shared.sweep_sessions();
357 1 : }
358 :
359 : /// Stop listening, fail the turns and kill their harnesses, disconnect the clients.
360 : /// Returns once the harness processes exited, the runtime would otherwise drop them while they run.
361 2 : pub async fn close(self) {
362 2 : self.accept.abort();
363 2 : self.sweep.abort();
364 : {
365 2 : let mut state = self.shared.lock();
366 2 : let state = &mut *state;
367 2 : self.shared.log(&format!(
368 2 : "Stopping: failing {} turns, closing {} clients",
369 2 : state.turns.len(),
370 2 : state.clients.len()
371 2 : ));
372 2 : for turn in state.turns.drain(..) {
373 2 : notify_failed(&state.clients, &turn, SERVER_STOPPED_MESSAGE, false);
374 2 : if let Some(kill) = &turn.kill {
375 1 : kill();
376 1 : }
377 : }
378 2 : for client in state.clients.values() {
379 2 : let _ = client.tx.send(Outgoing::Close(1001, "server stops".into()));
380 2 : }
381 : }
382 : // A killed harness exits within KILL_GRACE, then it is killed for good
383 2 : let harnesses_deadline = tokio::time::Instant::now() + KILL_GRACE + CLOSE_GRACE;
384 2 : let clients_deadline = tokio::time::Instant::now() + CLOSE_GRACE;
385 : loop {
386 4 : let now = tokio::time::Instant::now();
387 4 : let (harnesses, clients) = {
388 4 : let state = self.shared.lock();
389 4 : (state.harnesses, state.clients.len())
390 4 : };
391 4 : let waits_for_harnesses = harnesses > 0 && now < harnesses_deadline;
392 4 : let waits_for_clients = clients > 0 && now < clients_deadline;
393 4 : if !waits_for_harnesses && !waits_for_clients {
394 2 : if harnesses > 0 || clients > 0 {
395 0 : self.shared.log(&format!("Stopped with {harnesses} harnesses and {clients} clients left"));
396 2 : }
397 2 : break;
398 2 : }
399 2 : tokio::time::sleep(Duration::from_millis(10)).await;
400 : }
401 2 : }
402 : }
403 :
404 : impl Shared {
405 317 : fn lock(&self) -> MutexGuard<'_, State> {
406 317 : self.state.lock().unwrap_or_else(|e| e.into_inner())
407 317 : }
408 :
409 357 : fn log(&self, message: &str) {
410 357 : if self.options.log {
411 0 : log_line(message);
412 357 : }
413 357 : }
414 :
415 6 : fn status(&self, state: &State) -> Value {
416 6 : json!({
417 6 : "port": self.port,
418 7 : "clients": state.clients.values().map(|client| json!({
419 7 : "id": client.id,
420 7 : "kind": client.kind,
421 7 : "origin": client.origin,
422 7 : "sessions": client.sessions.iter().map(|s| json!({ "id": s.id, "name": s.name })).collect::<Vec<_>>(),
423 7 : "focused": client.focused,
424 7 : "features": client.features,
425 6 : })).collect::<Vec<_>>(),
426 6 : "turns": state.turns.iter().map(|turn| json!({
427 1 : "turnId": turn.turn_id,
428 1 : "session": turn.session,
429 1 : "harness": turn.harness,
430 1 : "model": turn.picks.model,
431 1 : "effort": turn.picks.effort,
432 1 : "state": turn.state.name(),
433 1 : "status": turn.status,
434 6 : })).collect::<Vec<_>>(),
435 : })
436 6 : }
437 :
438 19 : async fn resolved_path(&self) -> Option<OsString> {
439 19 : self.path.get_or_init(|| (self.options.resolve_path)()).await.clone()
440 19 : }
441 :
442 39 : fn host_allowed(&self, host: Option<&str>) -> bool {
443 39 : host == Some(&format!("127.0.0.1:{}", self.port)) || host == Some(&format!("localhost:{}", self.port))
444 39 : }
445 :
446 9 : fn token_valid(&self, request: &Request) -> bool {
447 9 : let expected = format!("Bearer {}", self.options.config.token);
448 9 : constant_time_eq(request.header("authorization").unwrap_or_default().as_bytes(), expected.as_bytes())
449 9 : }
450 : }
451 :
452 : /// Send a notification to every client that holds a session
453 34 : fn notify_session(clients: &BTreeMap<u64, Client>, session: &str, method: &str, params: Value) {
454 34 : let message = json!({ "method": method, "params": params });
455 34 : for client in clients.values().filter(|c| c.holds(session)) {
456 34 : client.send(&message);
457 34 : }
458 34 : }
459 :
460 21 : fn notify_status(clients: &BTreeMap<u64, Client>, turn: &Turn) {
461 21 : let params = json!({ "session": turn.session, "turnId": turn.turn_id, "state": turn.state.name(), "text": turn.status });
462 21 : notify_session(clients, &turn.session, "agent/status", params);
463 21 : }
464 :
465 8 : fn notify_failed(clients: &BTreeMap<u64, Client>, turn: &Turn, error: &str, cancelled: bool) {
466 8 : let params = json!({ "session": turn.session, "turnId": turn.turn_id, "error": error, "cancelled": cancelled });
467 8 : notify_session(clients, &turn.session, "agent/failed", params);
468 8 : }
469 :
470 : impl Shared {
471 : // ------------------------------------------------------------------------------------------------
472 : // Calls of the command line tool
473 :
474 : /// Run a command of the command line tool against a session
475 8 : async fn call(&self, call: &Value) -> Value {
476 8 : let Some(command) = call["command"].as_str().filter(|c| is_session_command(c)) else {
477 1 : let command = call["command"].as_str().map_or_else(|| call["command"].to_string(), str::to_string);
478 1 : return RpcError::new(ErrorCode::Usage, format!("unknown command: {command}")).to_call_response();
479 : };
480 7 : let session = match &call["session"] {
481 4 : Value::Null => None,
482 3 : Value::String(s) => Some(s.as_str()),
483 0 : _ => return RpcError::new(ErrorCode::Usage, "the session must be a string").to_call_response(),
484 : };
485 7 : let Some((client, session)) = self.wait_for_session(session).await else {
486 1 : let target = session.map_or_else(|| "a focused session".into(), |s| format!("session {s}"));
487 1 : self.log(&format!("Call {command}: no client holds {target}"));
488 1 : return RpcError::new(ErrorCode::NotConnected, NOT_CONNECTED_MESSAGE).to_call_response();
489 : };
490 6 : self.log(&format!("Call {command}: forwarding to client {client} for session {session}"));
491 6 : if !self.lock().busy.insert(session.clone()) {
492 1 : self.log(&format!("Call {command}: session {session} is busy with another call"));
493 1 : let message = "Another dashql agent call of this session is running, wait for it to finish";
494 1 : return RpcError::new(ErrorCode::Busy, message).to_call_response();
495 5 : }
496 5 : let args = if call["args"].is_null() { json!({}) } else { call["args"].clone() };
497 5 : let params = json!({ "session": session, "command": command, "args": args });
498 5 : let result = self.request(client, "command", params).await;
499 5 : self.lock().busy.remove(&session);
500 5 : match result {
501 3 : Ok(result) => json!({ "ok": true, "session": session, "result": result }),
502 2 : Err(error) => error.to_call_response(),
503 : }
504 8 : }
505 :
506 : /// The client that holds a session or the focused session, the most recent announcement or focus wins
507 24 : fn resolve(state: &State, session: Option<&str>) -> Option<(u64, String)> {
508 24 : match session {
509 4 : None => state
510 4 : .clients
511 4 : .iter()
512 4 : .filter_map(|(key, c)| c.focused.as_ref().map(|f| (c.focused_at, *key, f)))
513 4 : .max_by_key(|(at, _, _)| *at)
514 4 : .map(|(_, key, focused)| (key, focused.clone())),
515 20 : Some(session) => state
516 20 : .clients
517 20 : .iter()
518 20 : .filter(|(_, c)| c.holds(session))
519 20 : .max_by_key(|(_, c)| c.announced_at)
520 20 : .map(|(key, _)| (*key, session.to_string())),
521 : }
522 24 : }
523 :
524 : /// Wait until a client holds the session, queued calls are delivered once a tab reconnects
525 7 : async fn wait_for_session(&self, session: Option<&str>) -> Option<(u64, String)> {
526 7 : let wait = self.lock().config.call_wait;
527 7 : let deadline = tokio::time::Instant::now() + wait;
528 : loop {
529 8 : let changed = self.sessions_changed.notified();
530 8 : tokio::pin!(changed);
531 8 : changed.as_mut().enable();
532 8 : let target = Self::resolve(&self.lock(), session);
533 8 : if target.is_some() {
534 6 : return target;
535 2 : }
536 2 : if wait.is_zero() {
537 0 : return None;
538 2 : }
539 2 : let target = session.map_or_else(|| "a focused session".into(), |s| format!("session {s}"));
540 2 : self.log(&format!("Waiting up to {}ms for a client that holds {target}", wait.as_millis()));
541 2 : if tokio::time::timeout_at(deadline, changed).await.is_err() {
542 1 : return None;
543 1 : }
544 : }
545 7 : }
546 :
547 5 : async fn request(&self, client: u64, method: &str, params: Value) -> Result<Value, RpcError> {
548 5 : let disconnected = || RpcError::new(ErrorCode::Failed, DISCONNECTED_MESSAGE);
549 5 : let (tx, rx) = oneshot::channel();
550 5 : let id = {
551 5 : let mut state = self.lock();
552 5 : let id = format!("s{}", state.next_call_id);
553 5 : state.next_call_id += 1;
554 5 : let client = state.clients.get_mut(&client).ok_or_else(disconnected)?;
555 5 : client.pending.insert(id.clone(), tx);
556 5 : client.send(&json!({ "id": id, "method": method, "params": params }));
557 5 : id
558 : };
559 5 : match tokio::time::timeout(self.options.call_timeout, rx).await {
560 4 : Ok(Ok(result)) => result,
561 : // The client is gone with its pending calls
562 : Ok(Err(_)) => {
563 1 : self.log(&format!("Client {client} disconnected before it answered {id}"));
564 1 : Err(disconnected())
565 : }
566 : Err(_) => {
567 0 : self.log(&format!(
568 0 : "Client {client} did not answer {id} within {}s",
569 0 : self.options.call_timeout.as_secs()
570 0 : ));
571 0 : if let Some(client) = self.lock().clients.get_mut(&client) {
572 0 : client.pending.remove(&id);
573 0 : }
574 0 : Err(RpcError::new(ErrorCode::Failed, "DashQL did not answer in time"))
575 : }
576 : }
577 5 : }
578 :
579 : // ------------------------------------------------------------------------------------------------
580 : // Messages of clients
581 :
582 60 : fn send_to(&self, client: u64, message: Value) {
583 60 : if let Some(client) = self.lock().clients.get(&client) {
584 60 : client.send(&message);
585 60 : }
586 60 : }
587 :
588 1 : fn close_client(&self, client: u64, code: u16, reason: String) {
589 1 : if let Some(client) = self.lock().clients.get(&client) {
590 1 : let _ = client.tx.send(Outgoing::Close(code, reason));
591 1 : }
592 1 : }
593 :
594 65 : fn handle_message(self: &Arc<Self>, client: u64, text: &str) {
595 65 : let message = match serde_json::from_str::<Value>(text) {
596 65 : Ok(message) => message,
597 0 : Err(e) => {
598 0 : self.log(&format!("Client {client} sent invalid JSON, closing it: {e}"));
599 0 : self.close_client(client, CLOSE_POLICY_VIOLATION, "invalid JSON".into());
600 0 : return;
601 : }
602 : };
603 65 : let id = message.get("id").filter(|id| !id.is_null()).cloned();
604 65 : let Some(method) = message["method"].as_str() else {
605 : // A response to a command
606 4 : let id = id.map(|id| id.as_str().map_or_else(|| id.to_string(), str::to_string)).unwrap_or_default();
607 4 : let mut state = self.lock();
608 4 : let Some(pending) = state.clients.get_mut(&client).and_then(|c| c.pending.remove(&id)) else {
609 0 : self.log(&format!("Client {client} answered {id}, no call waits for it"));
610 0 : return;
611 : };
612 4 : let _ = pending.send(match &message["error"] {
613 3 : Value::Null => Ok(message["result"].clone()),
614 1 : error => Err(RpcError::from_json(error)),
615 : });
616 4 : return;
617 : };
618 61 : let reply = |result: Result<Value, RpcError>| {
619 57 : if let Some(id) = &id {
620 57 : match result {
621 49 : Ok(result) => self.send_to(client, json!({ "id": id, "result": result })),
622 8 : Err(error) => self.send_to(client, json!({ "id": id, "error": error.to_json() })),
623 : }
624 0 : }
625 57 : };
626 61 : let params = &message["params"];
627 61 : let session = || {
628 13 : params["session"]
629 13 : .as_str()
630 13 : .map(str::to_string)
631 13 : .ok_or(RpcError::new(ErrorCode::Usage, "the session must be a string"))
632 13 : };
633 61 : match method {
634 61 : "client/sessions" => {
635 21 : if self.update_sessions(client, params) {
636 20 : reply(Ok(json!({})));
637 20 : }
638 : }
639 40 : "agent/request" => match self.enqueue_turn(client, params) {
640 19 : Ok((result, turn_id)) => {
641 19 : // The response goes out before the first notification
642 19 : reply(Ok(result));
643 19 : tokio::spawn(start_or_queue(self.clone(), turn_id));
644 19 : }
645 5 : Err(error) => reply(Err(error)),
646 : },
647 16 : "agent/cancel" => reply(session().and_then(|session| {
648 4 : let turn_id = match params.get("turnId") {
649 1 : None | Some(Value::Null) => None,
650 2 : Some(Value::String(turn_id)) => Some(turn_id.as_str()),
651 1 : Some(_) => return Err(RpcError::new(ErrorCode::Usage, "the turn id must be a string")),
652 : };
653 3 : self.cancel_turns(&session, turn_id);
654 3 : Ok(json!({}))
655 4 : })),
656 12 : "agent/forget" => reply(session().and_then(|session| {
657 3 : self.forget_session(&session).map_err(|e| RpcError::new(ErrorCode::Failed, e))?;
658 2 : Ok(json!({}))
659 3 : })),
660 9 : "agent/reset" => reply(session().and_then(|session| {
661 1 : self.log(&format!("Forgetting the harness sessions of session {session}"));
662 1 : let state = SessionState::default();
663 1 : self.write_session_state(&session, &state).map_err(|e| RpcError::new(ErrorCode::Failed, e))?;
664 1 : Ok(json!({}))
665 1 : })),
666 8 : "agent/context" => reply(session().and_then(|session| self.session_context(&session, params))),
667 3 : "agent/info" => {
668 3 : let (shared, id) = (self.clone(), id.clone());
669 3 : tokio::spawn(async move {
670 3 : let info = shared.info().await;
671 3 : if let Some(id) = id {
672 3 : shared.send_to(client, json!({ "id": id, "result": info }));
673 3 : }
674 3 : });
675 : }
676 0 : _ => reply(Err(RpcError::new(ErrorCode::Usage, format!("unknown method: {method}")))),
677 : }
678 65 : }
679 :
680 : /// Take the sessions of a client, false if the client was turned away
681 21 : fn update_sessions(&self, client: u64, params: &Value) -> bool {
682 21 : if params["protocol"].as_u64() != Some(PROTOCOL_VERSION) {
683 1 : let reason =
684 1 : format!("protocol version {} is not supported, expected {PROTOCOL_VERSION}", params["protocol"]);
685 1 : self.log(&format!("Turning client {client} away: {reason}"));
686 1 : self.close_client(client, CLOSE_POLICY_VIOLATION, reason);
687 1 : return false;
688 20 : }
689 : {
690 20 : let mut state = self.lock();
691 20 : state.clock += 1;
692 20 : let now = state.clock;
693 20 : let Some(client) = state.clients.get_mut(&client) else { return false };
694 20 : client.kind = params["client"].as_str().unwrap_or("unknown").to_string();
695 20 : client.features =
696 20 : Some(page_features(params["features"].as_array().map(|f| f.iter().filter_map(Value::as_str))));
697 20 : client.sessions.clear();
698 22 : for session in params["sessions"].as_array().into_iter().flatten() {
699 22 : let Some(id) = session["id"].as_str() else { continue };
700 22 : let name = session["name"].as_str().unwrap_or_default().to_string();
701 22 : match client.sessions.iter_mut().find(|s| s.id == id) {
702 0 : Some(existing) => existing.name = name,
703 22 : None => client.sessions.push(ClientSession { id: id.to_string(), name }),
704 : }
705 : }
706 20 : client.announced_at = now;
707 20 : let focused = params["focused"].as_str().filter(|f| client.holds(f)).map(str::to_string);
708 20 : if focused != client.focused {
709 18 : client.focused_at = now;
710 18 : }
711 20 : client.focused = focused;
712 : }
713 : // A session that a client shows is in use, its folder stays
714 22 : for session in params["sessions"].as_array().into_iter().flatten().filter_map(|s| s["id"].as_str()) {
715 22 : if self.session_dir(session).is_dir() {
716 1 : self.mark_session_used(session);
717 21 : }
718 : }
719 20 : self.sessions_changed.notify_waiters();
720 20 : true
721 21 : }
722 :
723 : /// The configuration and the state of the server, for the settings of the app.
724 : /// Shows the harness commands, they are changed in config.json only.
725 3 : async fn info(&self) -> Value {
726 3 : let path = self.resolved_path().await;
727 3 : let state = self.lock();
728 3 : let mut info = self.status(&state);
729 3 : let harnesses: Vec<Value> = state
730 3 : .config
731 3 : .harnesses
732 3 : .iter()
733 9 : .map(|(name, harness)| {
734 9 : let program = resolve_program(&harness.command[0], path.as_deref());
735 9 : json!({
736 9 : "name": name,
737 9 : "command": harness.command,
738 9 : "adapter": harness.adapter.name(),
739 9 : "program": program.map(|p| p.display().to_string()),
740 9 : "models": harness.model.values,
741 9 : "efforts": harness.effort.values,
742 : })
743 9 : })
744 3 : .collect();
745 3 : let extra = json!({
746 3 : "home": self.options.home.display().to_string(),
747 3 : "configPath": self.options.home.join("config.json").display().to_string(),
748 3 : "defaultHarness": state.config.default_harness,
749 3 : "harnesses": harnesses,
750 : });
751 3 : if let (Some(info), Value::Object(extra)) = (info.as_object_mut(), extra) {
752 3 : info.extend(extra);
753 3 : }
754 3 : info
755 3 : }
756 :
757 : // ------------------------------------------------------------------------------------------------
758 : // Turns
759 :
760 97 : fn session_dir(&self, session: &str) -> PathBuf {
761 97 : self.options.home.join("sessions").join(session_folder_name(session))
762 97 : }
763 :
764 22 : fn read_session_state(&self, session: &str) -> SessionState {
765 22 : let text = std::fs::read_to_string(self.session_dir(session).join("state.json")).unwrap_or_default();
766 22 : let Ok(Value::Object(mut state)) = serde_json::from_str::<Value>(&text) else { return SessionState::default() };
767 : SessionState {
768 21 : harness_sessions: match state.remove("harnessSessions") {
769 20 : Some(Value::Object(sessions)) => sessions,
770 1 : _ => Map::new(),
771 : },
772 : }
773 22 : }
774 :
775 14 : fn write_session_state(&self, session: &str, session_state: &SessionState) -> Result<(), String> {
776 14 : let dir = self.session_dir(session);
777 14 : create_private_dir(&dir)?;
778 14 : let state = json!({ "harnessSessions": session_state.harness_sessions });
779 14 : write_private_file(&dir.join("state.json"), &(to_json_text(&state) + "\n"))
780 14 : }
781 :
782 : /// Mark a session as used: the time of its state.json is the last use, the sweep keeps the folders of used
783 : /// sessions
784 17 : fn mark_session_used(&self, session: &str) {
785 17 : let path = self.session_dir(session).join("state.json");
786 17 : let marked = match std::fs::File::options().write(true).open(&path) {
787 6 : Ok(file) => file.set_modified(SystemTime::now()).map_err(|e| e.to_string()),
788 11 : Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
789 11 : self.write_session_state(session, &SessionState::default())
790 : }
791 0 : Err(e) => Err(e.to_string()),
792 : };
793 17 : if let Err(e) = marked {
794 0 : self.log(&format!("Cannot mark session {session} as used: {e}"));
795 17 : }
796 17 : }
797 :
798 : /// Remove the folder of a session
799 2 : fn remove_session_dir(&self, session: &str) -> Result<(), String> {
800 2 : let dir = self.session_dir(session);
801 2 : match std::fs::remove_dir_all(&dir) {
802 0 : Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(format!("{}: {e}", dir.display())),
803 2 : _ => Ok(()),
804 : }
805 2 : }
806 :
807 : /// Forget a deleted session: stop its turns and remove its folder, once a running turn finished
808 3 : fn forget_session(self: &Arc<Self>, session: &str) -> Result<(), String> {
809 : // An empty session id would name the folder of all sessions
810 3 : if session_folder_name(session).is_empty() {
811 1 : return Err("the session is empty".into());
812 2 : }
813 2 : self.cancel_turns(session, None);
814 2 : let mut state = self.lock();
815 2 : if state.turns.iter().any(|turn| turn.session == session) {
816 1 : self.log(&format!("Forgetting session {session} once its running turn stopped"));
817 1 : state.forgotten.insert(session.to_string());
818 1 : return Ok(());
819 1 : }
820 1 : self.log(&format!("Forgetting session {session}"));
821 1 : self.remove_session_dir(session)
822 3 : }
823 :
824 : /// Remove the folders of the sessions that were not used within the retention.
825 : /// A client may hold sessions that no other client knows, so only the age tells that a folder is no longer needed.
826 28 : fn sweep_sessions(&self) {
827 28 : let root = self.options.home.join("sessions");
828 28 : let Ok(entries) = std::fs::read_dir(&root) else { return };
829 : // The lock keeps new turns from starting in a folder while it is removed
830 1 : let state = self.lock();
831 1 : let in_use: HashSet<String> = state
832 1 : .turns
833 1 : .iter()
834 1 : .map(|turn| turn.session.as_str())
835 1 : .chain(state.busy.iter().map(String::as_str))
836 1 : .chain(state.clients.values().flat_map(|client| client.sessions.iter().map(|s| s.id.as_str())))
837 1 : .map(session_folder_name)
838 1 : .collect();
839 1 : let now = SystemTime::now();
840 1 : let mut removed = 0;
841 3 : for entry in entries.flatten() {
842 3 : let name = entry.file_name().to_string_lossy().into_owned();
843 3 : if in_use.contains(&name) || !entry.file_type().is_ok_and(|t| t.is_dir()) {
844 1 : continue;
845 2 : }
846 : // A folder without state.json is as old as the folder itself, an unknown time keeps it
847 2 : let dir = entry.path();
848 2 : let used = std::fs::metadata(dir.join("state.json")).or_else(|_| std::fs::metadata(&dir));
849 2 : let Ok(used) = used.and_then(|m| m.modified()) else { continue };
850 2 : if now.duration_since(used).unwrap_or_default() < SESSION_RETENTION {
851 1 : continue;
852 1 : }
853 1 : match std::fs::remove_dir_all(&dir) {
854 1 : Ok(()) => removed += 1,
855 0 : Err(e) => self.log(&format!("Cannot remove the folder of an unused session {}: {e}", dir.display())),
856 : }
857 : }
858 1 : drop(state);
859 1 : if removed > 0 {
860 1 : let days = SESSION_RETENTION.as_secs() / 86400;
861 1 : self.log(&format!("Removed the folders of {removed} sessions unused for {days} days"));
862 1 : }
863 28 : }
864 :
865 5 : fn session_context(&self, session: &str, params: &Value) -> Result<Value, RpcError> {
866 5 : let state = self.lock();
867 5 : let harness = pick_harness(&state.config, params)?;
868 4 : let session_state = self.read_session_state(session);
869 4 : let harness_session = session_state.harness_sessions.get(&harness).and_then(Value::as_str).map(str::to_string);
870 4 : let template = state.config.harness(&harness).and_then(|h| h.terminal_command.as_ref());
871 4 : let terminal_command = match (&harness_session, template) {
872 2 : (Some(id), Some(template)) => {
873 6 : Some(template.iter().map(|arg| arg.replace("{harnessSession}", id)).collect::<Vec<_>>())
874 : }
875 2 : _ => None,
876 : };
877 4 : Ok(json!({
878 4 : "folder": self.session_dir(session).display().to_string(),
879 4 : "harness": harness,
880 4 : "harnessSession": harness_session,
881 4 : "terminalCommand": terminal_command,
882 4 : }))
883 5 : }
884 :
885 24 : fn enqueue_turn(&self, client: u64, params: &Value) -> Result<(Value, String), RpcError> {
886 24 : let mut state = self.lock();
887 24 : let session = params["session"].as_str().unwrap_or_default();
888 24 : if !state.clients.get(&client).is_some_and(|c| c.holds(session)) {
889 1 : return Err(RpcError::new(ErrorCode::Usage, "the session has no agent access"));
890 23 : }
891 23 : let harness = pick_harness(&state.config, params)?;
892 22 : let picks = pick_choices(&state.config, &harness, params)?;
893 19 : let prompt = params["prompt"].as_str().unwrap_or_default();
894 19 : if prompt.trim().is_empty() {
895 0 : return Err(RpcError::new(ErrorCode::Failed, "the prompt is empty"));
896 19 : }
897 : // A new turn keeps the folder of a session that was forgotten while a turn ran
898 19 : state.forgotten.remove(session);
899 19 : let turn_id = random_uuid();
900 19 : let ahead = state.turns.iter().filter(|turn| turn.session == session).count();
901 19 : self.log(&format!(
902 19 : "Turn {turn_id}: queued for session {session} with {harness} (model {}, effort {}), {ahead} turns ahead, prompt: {}",
903 19 : picks.model.as_deref().unwrap_or("default"),
904 19 : picks.effort.as_deref().unwrap_or("default"),
905 19 : clip(prompt)
906 19 : ));
907 19 : state.turns.push(Turn {
908 19 : turn_id: turn_id.clone(),
909 19 : session: session.to_string(),
910 19 : harness: harness.clone(),
911 19 : picks,
912 19 : prompt: prompt.to_string(),
913 19 : state: TurnState::Queued,
914 19 : status: "waiting for agent...".into(),
915 19 : cancelled: false,
916 19 : kill: None,
917 19 : });
918 19 : Ok((json!({ "turnId": turn_id, "harness": harness }), turn_id))
919 24 : }
920 :
921 : /// Stop a turn of a session, or all turns of the session without a turn id
922 5 : fn cancel_turns(self: &Arc<Self>, session: &str, turn_id: Option<&str>) {
923 5 : let mut state = self.lock();
924 5 : let state = &mut *state;
925 27 : let stops = |turn: &Turn| turn.session == session && turn_id.is_none_or(|id| turn.turn_id == id);
926 9 : let stopped = state.turns.iter().filter(|turn| stops(turn)).count();
927 5 : self.log(&format!("Stopping {stopped} turns of session {session} (turn {})", turn_id.unwrap_or("all")));
928 9 : for turn in state.turns.iter_mut().filter(|turn| stops(turn)) {
929 5 : self.log(&format!("Turn {}: stopped while {}", turn.turn_id, turn.state.name()));
930 5 : turn.cancelled = true;
931 5 : match (turn.state, &turn.kill) {
932 2 : (TurnState::Queued, _) => notify_failed(&state.clients, turn, "Stopped", true),
933 3 : (TurnState::Running, Some(kill)) => kill(),
934 : // The turn checks whether it was cancelled before it starts the harness
935 0 : (TurnState::Running, None) => {}
936 : }
937 : }
938 9 : state.turns.retain(|turn| !stops(turn) || turn.state == TurnState::Running);
939 : // The first turn of the session may have been dropped before it started
940 5 : let next = state.turns.iter().find(|turn| turn.session == session);
941 5 : if let Some(next) = next.filter(|turn| turn.state == TurnState::Queued) {
942 0 : tokio::spawn(start_turn(self.clone(), next.turn_id.clone()));
943 5 : }
944 5 : }
945 :
946 : /// The folder of the dashql executable, first on the PATH of harness runs so the agent calls this binary.
947 : /// Only a binary named dashql: the folder of another one would not give the agent a `dashql` command.
948 16 : fn cli_folder(&self) -> Option<PathBuf> {
949 16 : let cli = std::fs::canonicalize(self.options.cli.as_ref()?).ok()?;
950 0 : if cli.file_stem()? != "dashql" {
951 0 : return None;
952 0 : }
953 0 : cli.parent().map(Path::to_path_buf)
954 16 : }
955 :
956 : /// Prepare the working folder and start the harness, the outcome once the harness exited
957 16 : async fn run_harness(self: &Arc<Self>, turn_id: &str) -> Result<HarnessOutcome, String> {
958 16 : let (session, config, prompt, picks) = {
959 16 : let state = self.lock();
960 17 : let Some(turn) = state.turns.iter().find(|turn| turn.turn_id == turn_id) else {
961 0 : return Ok(HarnessOutcome::failed("Stopped"));
962 : };
963 16 : let config =
964 16 : state.config.harness(&turn.harness).cloned().ok_or(format!("unknown harness: {}", turn.harness))?;
965 16 : (turn.session.clone(), config, turn.prompt.clone(), turn.picks.clone())
966 : };
967 16 : let cwd = self.session_dir(&session);
968 16 : create_private_dir(&cwd)?;
969 16 : self.mark_session_used(&session);
970 : // The guide describes the features of the page that holds the session
971 16 : let features = {
972 16 : let state = self.lock();
973 16 : Self::resolve(&state, Some(&session))
974 16 : .and_then(|(client, _)| state.clients.get(&client).and_then(|c| c.features.clone()))
975 16 : .unwrap_or_default()
976 : };
977 16 : let page = Page::Features(&features);
978 48 : for (name, text) in [
979 16 : ("AGENTS.md", session_instructions(page)),
980 16 : ("CLAUDE.md", claude_instructions(page)),
981 16 : ("opencode.json", OPENCODE_CONFIG.to_string()),
982 16 : ] {
983 48 : let path = cwd.join(name);
984 48 : write_if_changed(&path, &text).map_err(|e| format!("{}: {e}", path.display()))?;
985 : }
986 : // The files of an earlier turn may hold older script texts
987 16 : let input = cwd.join(INPUT_FOLDER);
988 16 : match std::fs::remove_dir_all(&input) {
989 12 : Err(e) if e.kind() != std::io::ErrorKind::NotFound => return Err(format!("{}: {e}", input.display())),
990 16 : _ => create_private_dir(&input)?,
991 : }
992 16 : let path = self.resolved_path().await;
993 16 : let mut env: BTreeMap<OsString, OsString> = std::env::vars_os().collect();
994 16 : let paths = self.cli_folder().into_iter().chain(path.iter().flat_map(|p| std::env::split_paths(p)));
995 16 : env.insert("PATH".into(), std::env::join_paths(paths).map_err(|e| e.to_string())?);
996 16 : env.insert("DASHQL_SESSION".into(), session.clone().into());
997 16 : env.insert("DASHQL_AGENT_HOME".into(), self.options.home.clone().into());
998 16 : env.insert("DASHQL_AGENT_PORT".into(), self.port.to_string().into());
999 : // `dashql agent serve` refuses to start inside a run
1000 16 : env.insert(AGENT_RUN_ENV.into(), "1".into());
1001 :
1002 16 : let harness_name = {
1003 16 : let state = self.lock();
1004 17 : match state.turns.iter().find(|turn| turn.turn_id == turn_id) {
1005 16 : Some(turn) if !turn.cancelled => turn.harness.clone(),
1006 0 : _ => return Ok(HarnessOutcome::failed("Stopped")),
1007 : }
1008 : };
1009 16 : let harness_session =
1010 16 : self.read_session_state(&session).harness_sessions.get(&harness_name).and_then(Value::as_str).map(str::to_string);
1011 16 : let argv = build_harness_args(&config, &prompt, harness_session.as_deref(), &picks);
1012 16 : self.log(&format!(
1013 : "Turn {turn_id}: starting {harness_name} for session {session} in {}, {}, command: {}",
1014 16 : cwd.display(),
1015 16 : harness_session
1016 16 : .as_ref()
1017 16 : .map_or_else(|| "new harness session".into(), |id| format!("resuming harness session {id}")),
1018 16 : clip(&serde_json::to_string(&argv).unwrap_or_default())
1019 : ));
1020 16 : self.log(&format!("Turn {turn_id}: PATH={}", env[&OsString::from("PATH")].to_string_lossy()));
1021 16 : let (events, mut received) = mpsc::unbounded_channel();
1022 16 : let process =
1023 16 : (self.options.spawn_harness)(HarnessSpawnArgs { config, prompt, harness_session, picks, cwd, env, events });
1024 : {
1025 16 : let mut state = self.lock();
1026 16 : state.harnesses += 1;
1027 17 : match state.turns.iter_mut().find(|turn| turn.turn_id == turn_id) {
1028 16 : Some(turn) => {
1029 : // A cancel that came while the harness started
1030 16 : if turn.cancelled {
1031 0 : (process.kill)();
1032 16 : }
1033 16 : turn.kill = Some(process.kill.clone());
1034 : }
1035 : // The server stopped while the harness started
1036 0 : None => (process.kill)(),
1037 : }
1038 : }
1039 16 : let mut outcome = process.outcome;
1040 : loop {
1041 21 : tokio::select! {
1042 21 : Some(event) = received.recv() => self.handle_event(turn_id, &session, &harness_name, event),
1043 21 : outcome = &mut outcome => {
1044 10 : self.lock().harnesses -= 1;
1045 11 : while let Ok(event) = received.try_recv() {
1046 1 : self.handle_event(turn_id, &session, &harness_name, event);
1047 1 : }
1048 10 : return Ok(outcome);
1049 : }
1050 : }
1051 : }
1052 10 : }
1053 :
1054 6 : fn handle_event(&self, turn_id: &str, session: &str, harness: &str, event: HarnessEvent) {
1055 6 : match event {
1056 0 : HarnessEvent::Stdout(line) => self.log(&format!("Turn {turn_id}: {harness} stdout: {}", clip(&line))),
1057 0 : HarnessEvent::Stderr(line) => self.log(&format!("Turn {turn_id}: {harness} stderr: {}", clip(&line))),
1058 3 : HarnessEvent::Process(text) => self.log(&format!("Turn {turn_id}: {harness}: {text}")),
1059 1 : HarnessEvent::Status(text) => {
1060 1 : let mut state = self.lock();
1061 1 : let state = &mut *state;
1062 1 : let Some(turn) = state.turns.iter_mut().find(|turn| turn.turn_id == turn_id) else { return };
1063 1 : if text.is_empty() || turn.cancelled {
1064 0 : return;
1065 1 : }
1066 1 : turn.status = text;
1067 1 : notify_status(&state.clients, turn);
1068 : }
1069 2 : HarnessEvent::Session(id) => {
1070 2 : if self.lock().forgotten.contains(session) {
1071 0 : return;
1072 2 : }
1073 2 : let mut state = self.read_session_state(session);
1074 2 : if state.harness_sessions.get(harness).and_then(Value::as_str) == Some(&id) {
1075 0 : return;
1076 2 : }
1077 2 : self.log(&format!("Turn {turn_id}: {harness} runs in harness session {id}"));
1078 2 : state.harness_sessions.insert(harness.to_string(), id.into());
1079 2 : if let Err(e) = self.write_session_state(session, &state) {
1080 0 : self.log(&format!("Cannot store the harness session: {e}"));
1081 2 : }
1082 : }
1083 : }
1084 6 : }
1085 :
1086 10 : fn finish_turn(self: &Arc<Self>, turn_id: &str, outcome: HarnessOutcome) {
1087 10 : let mut state = self.lock();
1088 10 : let Some(index) = state.turns.iter().position(|turn| turn.turn_id == turn_id) else { return };
1089 9 : let turn = state.turns.remove(index);
1090 9 : if turn.cancelled {
1091 3 : self.log(&format!("Turn {turn_id}: stopped"));
1092 3 : notify_failed(&state.clients, &turn, "Stopped", true);
1093 6 : } else if let Some(error) = &outcome.error {
1094 1 : self.log(&format!("Turn {turn_id}: {} failed for session {}: {error}", turn.harness, turn.session));
1095 1 : notify_failed(&state.clients, &turn, error, false);
1096 5 : } else {
1097 5 : self.log(&format!(
1098 5 : "Turn {turn_id}: {} answered for session {}: {}",
1099 5 : turn.harness,
1100 5 : turn.session,
1101 5 : clip(&outcome.message)
1102 5 : ));
1103 5 : let params = json!({ "session": turn.session, "turnId": turn.turn_id, "message": outcome.message });
1104 5 : notify_session(&state.clients, &turn.session, "agent/done", params);
1105 5 : }
1106 9 : let next = state.turns.iter().find(|next| next.session == turn.session);
1107 9 : if let Some(next) = next.filter(|next| next.state == TurnState::Queued) {
1108 1 : tokio::spawn(start_turn(self.clone(), next.turn_id.clone()));
1109 8 : } else if next.is_none() && state.forgotten.remove(&turn.session) {
1110 1 : self.log(&format!("Forgetting session {}", turn.session));
1111 1 : if let Err(e) = self.remove_session_dir(&turn.session) {
1112 0 : self.log(&format!("Cannot remove the folder of session {}: {e}", turn.session));
1113 1 : }
1114 7 : }
1115 10 : }
1116 : }
1117 :
1118 : /// Start a new turn if it is the first of its session, otherwise report that it waits
1119 19 : async fn start_or_queue(shared: Arc<Shared>, turn_id: String) {
1120 : {
1121 19 : let state = shared.lock();
1122 25 : let Some(turn) = state.turns.iter().find(|turn| turn.turn_id == turn_id) else { return };
1123 19 : if turn.cancelled {
1124 0 : return;
1125 19 : }
1126 20 : if state.turns.iter().find(|first| first.session == turn.session).map(|first| &first.turn_id) != Some(&turn_id) {
1127 4 : shared.log(&format!("Turn {turn_id}: waits for an earlier turn of session {}", turn.session));
1128 4 : notify_status(&state.clients, turn);
1129 4 : return;
1130 15 : }
1131 : }
1132 15 : start_turn(shared, turn_id).await;
1133 14 : }
1134 :
1135 16 : async fn start_turn(shared: Arc<Shared>, turn_id: String) {
1136 : {
1137 16 : let mut state = shared.lock();
1138 16 : let state = &mut *state;
1139 17 : let Some(turn) = state.turns.iter_mut().find(|turn| turn.turn_id == turn_id) else { return };
1140 16 : if turn.state != TurnState::Queued || turn.cancelled {
1141 0 : return;
1142 16 : }
1143 16 : turn.state = TurnState::Running;
1144 16 : turn.status = "starting the agent...".into();
1145 16 : shared.log(&format!("Turn {turn_id}: running"));
1146 16 : notify_status(&state.clients, turn);
1147 : }
1148 16 : let outcome = shared.run_harness(&turn_id).await.unwrap_or_else(|e| {
1149 0 : shared.log(&format!("Turn {turn_id}: cannot start the harness: {e}"));
1150 0 : HarnessOutcome::failed(e)
1151 0 : });
1152 10 : shared.finish_turn(&turn_id, outcome);
1153 10 : }
1154 :
1155 48 : fn write_if_changed(path: &Path, text: &str) -> std::io::Result<()> {
1156 48 : if std::fs::read_to_string(path).is_ok_and(|current| current == text) {
1157 12 : return Ok(());
1158 36 : }
1159 36 : std::fs::write(path, text)
1160 48 : }
1161 :
1162 : // ------------------------------------------------------------------------------------------------
1163 : // Connections
1164 :
1165 39 : async fn handle_connection(shared: Arc<Shared>, mut stream: TcpStream) {
1166 39 : let Ok(Some(request)) = http::read_request(&mut stream).await else { return };
1167 39 : if request.header("upgrade").is_some() {
1168 29 : handle_upgrade(shared, stream, request).await;
1169 : } else {
1170 10 : handle_http(shared, stream, request).await;
1171 : }
1172 19 : }
1173 :
1174 10 : async fn respond(stream: &mut TcpStream, status: u16, reason: &str, body: &Value) {
1175 10 : let response = http::response(status, reason, Some("application/json"), body.to_string().as_bytes());
1176 10 : let _ = stream.write_all(&response).await;
1177 10 : let _ = stream.shutdown().await;
1178 10 : }
1179 :
1180 10 : async fn handle_http(shared: Arc<Shared>, mut stream: TcpStream, request: Request) {
1181 10 : let failure = |code, message: &str| RpcError::new(code, message).to_call_response();
1182 10 : let target = format!("{} {}", request.method, request.path());
1183 : // Browsers send an Origin, the command line tool doesn't. Pages never reach the HTTP endpoints.
1184 10 : if !shared.host_allowed(request.header("host")) || request.header("origin").is_some() {
1185 1 : shared.log(&format!(
1186 1 : "HTTP {target}: refused, host {:?}, origin {:?}",
1187 1 : request.header("host"),
1188 1 : request.header("origin")
1189 1 : ));
1190 1 : return respond(&mut stream, 403, "Forbidden", &failure(ErrorCode::Usage, "forbidden")).await;
1191 9 : }
1192 9 : if !shared.token_valid(&request) {
1193 1 : shared.log(&format!("HTTP {target}: refused, invalid token"));
1194 1 : return respond(&mut stream, 401, "Unauthorized", &failure(ErrorCode::Usage, "invalid token")).await;
1195 8 : }
1196 8 : match (request.method.as_str(), request.path()) {
1197 8 : ("GET", "/v1/status") => {
1198 0 : let status = shared.status(&shared.lock());
1199 0 : respond(&mut stream, 200, "OK", &json!({ "ok": true, "status": status })).await;
1200 : }
1201 8 : ("POST", "/v1/call") => {
1202 8 : let body = http::read_body(&mut stream, &request, MAX_HTTP_BODY_BYTES).await;
1203 8 : if let Ok(body) = &body {
1204 8 : shared.log(&format!("HTTP {target} <- {}", clip(&String::from_utf8_lossy(body))));
1205 8 : }
1206 8 : let call = body.and_then(|body| serde_json::from_slice::<Value>(&body).map_err(|e| e.to_string()));
1207 8 : match call {
1208 8 : Ok(call) => {
1209 8 : let result = shared.call(&call).await;
1210 8 : shared.log(&format!("HTTP {target} -> 200 {}", clip(&result.to_string())));
1211 8 : respond(&mut stream, 200, "OK", &result).await;
1212 : }
1213 0 : Err(e) => {
1214 0 : let body = failure(ErrorCode::Usage, &format!("invalid request: {e}"));
1215 0 : shared.log(&format!("HTTP {target} -> 400 {body}"));
1216 0 : respond(&mut stream, 400, "Bad Request", &body).await;
1217 : }
1218 : }
1219 : }
1220 : _ => {
1221 0 : shared.log(&format!("HTTP {target} -> 404"));
1222 0 : respond(&mut stream, 404, "Not Found", &failure(ErrorCode::Usage, "not found")).await
1223 : }
1224 : }
1225 10 : }
1226 :
1227 4 : async fn reject_upgrade(mut stream: TcpStream, status: u16, reason: &str) {
1228 4 : let _ = stream.write_all(&http::response(status, reason, None, b"")).await;
1229 4 : let _ = stream.shutdown().await;
1230 4 : }
1231 :
1232 29 : async fn handle_upgrade(shared: Arc<Shared>, mut stream: TcpStream, request: Request) {
1233 29 : if request.path() != WEBSOCKET_PATH {
1234 0 : shared.log(&format!("Refused a WebSocket upgrade for {}: not found", request.path()));
1235 0 : return reject_upgrade(stream, 404, "Not Found").await;
1236 29 : }
1237 : // Other websites cannot fake the Origin, the Host check blocks DNS rebinding
1238 29 : let origin = request.header("origin").unwrap_or_default().to_string();
1239 29 : if !shared.host_allowed(request.header("host")) || !shared.origins.contains(&origin) {
1240 4 : shared.log(&format!("Refused a client from {origin:?} for host {:?}", request.header("host")));
1241 4 : return reject_upgrade(stream, 403, "Forbidden").await;
1242 25 : }
1243 25 : let key = request.header("sec-websocket-key");
1244 25 : let upgrade = request.header("upgrade").unwrap_or_default();
1245 25 : let (Some(key), "GET", true, Some("13")) = (
1246 25 : key,
1247 25 : request.method.as_str(),
1248 25 : upgrade.eq_ignore_ascii_case("websocket"),
1249 25 : request.header("sec-websocket-version"),
1250 : ) else {
1251 0 : shared.log(&format!("Refused a WebSocket upgrade from {origin:?}: invalid handshake"));
1252 0 : return reject_upgrade(stream, 400, "Bad Request").await;
1253 : };
1254 25 : let accept = derive_accept_key(key.as_bytes());
1255 25 : let handshake =
1256 25 : format!("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n");
1257 25 : if let Err(e) = stream.write_all(handshake.as_bytes()).await {
1258 0 : shared.log(&format!("Cannot complete the WebSocket handshake with {origin:?}: {e}"));
1259 0 : return;
1260 25 : }
1261 25 : let config = WebSocketConfig::default()
1262 25 : .max_message_size(Some(shared.options.max_message_bytes))
1263 25 : .max_frame_size(Some(shared.options.max_message_bytes));
1264 25 : let socket = WebSocketStream::from_partially_read(stream, request.rest, Role::Server, Some(config)).await;
1265 25 : serve_client(shared, socket, origin).await;
1266 9 : }
1267 :
1268 25 : async fn serve_client(shared: Arc<Shared>, socket: WebSocketStream<TcpStream>, origin: String) {
1269 25 : let (tx, mut outgoing) = mpsc::unbounded_channel();
1270 25 : let key = {
1271 25 : let mut state = shared.lock();
1272 25 : let key = state.next_client;
1273 25 : state.next_client += 1;
1274 25 : let client = Client {
1275 25 : id: random_uuid(),
1276 25 : origin: origin.clone(),
1277 25 : kind: "unknown".into(),
1278 25 : sessions: Vec::new(),
1279 25 : focused: None,
1280 25 : features: None,
1281 25 : focused_at: 0,
1282 25 : announced_at: 0,
1283 25 : tx,
1284 25 : pending: HashMap::new(),
1285 25 : };
1286 25 : state.clients.insert(key, client);
1287 25 : key
1288 : };
1289 25 : shared.log(&format!("Client {key} connected from {origin}"));
1290 25 : let (mut sink, mut incoming) = socket.split();
1291 25 : let interval = shared.options.ping_interval;
1292 25 : let mut ping = tokio::time::interval_at(tokio::time::Instant::now() + interval, interval);
1293 25 : let mut awaiting_pong = false;
1294 5 : let reason = loop {
1295 189 : tokio::select! {
1296 189 : message = outgoing.recv() => match message {
1297 99 : Some(Outgoing::Text(text)) => {
1298 99 : shared.log(&format!("Client {key} -> {}", clip(&text)));
1299 99 : if let Err(e) = sink.send(Message::text(text)).await {
1300 0 : break format!("cannot send: {e}");
1301 99 : }
1302 : }
1303 3 : Some(Outgoing::Close(code, reason)) => {
1304 3 : let frame = CloseFrame { code: CloseCode::from(code), reason: reason.clone().into() };
1305 3 : let _ = sink.send(Message::Close(Some(frame))).await;
1306 3 : let _ = sink.close().await;
1307 3 : break format!("closed by the server with {code}: {reason}");
1308 : }
1309 0 : None => break "the server dropped the client".into(),
1310 : },
1311 189 : message = incoming.next() => match message {
1312 65 : Some(Ok(Message::Text(text))) => {
1313 65 : shared.log(&format!("Client {key} <- {}", clip(text.as_str())));
1314 65 : shared.handle_message(key, text.as_str());
1315 65 : }
1316 0 : Some(Ok(Message::Pong(_))) => awaiting_pong = false,
1317 0 : Some(Ok(Message::Binary(data))) => {
1318 0 : let frame = CloseFrame { code: CloseCode::Unsupported, reason: "binary messages are not supported".into() };
1319 0 : let _ = sink.send(Message::Close(Some(frame))).await;
1320 0 : break format!("sent a binary message of {} bytes", data.len());
1321 : }
1322 0 : Some(Ok(Message::Close(frame))) => {
1323 : // The socket acknowledges the close itself and ends the stream next
1324 0 : let frame = frame
1325 0 : .map_or_else(|| "no close frame".into(), |f| format!("{}: {}", u16::from(f.code), f.reason));
1326 0 : shared.log(&format!("Client {key} closes the connection with {frame}"));
1327 : }
1328 : // Pings are answered by the socket itself
1329 0 : Some(Ok(_)) => {}
1330 2 : Some(Err(e)) => break format!("socket error: {e}"),
1331 0 : None => break "the connection ended".into(),
1332 : },
1333 189 : _ = ping.tick() => {
1334 0 : if awaiting_pong {
1335 0 : break format!("no pong within {}s", interval.as_secs());
1336 0 : }
1337 0 : if let Err(e) = sink.send(Message::Ping(Default::default())).await {
1338 0 : break format!("cannot send a ping: {e}");
1339 0 : }
1340 0 : awaiting_pong = true;
1341 : }
1342 : }
1343 : };
1344 : // Dropping the client fails its pending calls
1345 5 : let pending = shared.lock().clients.remove(&key).map_or(0, |client| client.pending.len());
1346 5 : shared.log(&format!("Client {key} from {origin} disconnected ({reason}), failing {pending} pending calls"));
1347 5 : }
1348 :
1349 : #[cfg(test)]
1350 : mod tests {
1351 : use super::*;
1352 : use crate::agent::instructions::PAGE_FEATURES;
1353 : use crate::agent::config::{load_config, temp_dir};
1354 : use crate::agent::schema::BridgeSchema;
1355 : use std::sync::atomic::{AtomicBool, Ordering};
1356 : use tokio::io::AsyncReadExt;
1357 :
1358 : const TOKEN: &str = "test-token";
1359 :
1360 : /// A harness run of the fake harness
1361 : struct FakeRun {
1362 : args: HarnessSpawnArgs,
1363 : finish: Arc<Mutex<Option<oneshot::Sender<HarnessOutcome>>>>,
1364 : killed: Arc<AtomicBool>,
1365 : }
1366 :
1367 : impl FakeRun {
1368 6 : fn finish(&self, message: &str, error: Option<&str>) {
1369 6 : if let Some(finish) = self.finish.lock().unwrap().take() {
1370 6 : let _ = finish.send(HarnessOutcome { message: message.into(), error: error.map(str::to_string) });
1371 6 : }
1372 6 : }
1373 : }
1374 :
1375 : struct Fixture {
1376 : home: PathBuf,
1377 : server: AgentServer,
1378 : runs: Arc<Mutex<Vec<FakeRun>>>,
1379 : }
1380 :
1381 : impl Fixture {
1382 26 : async fn start(config: Value) -> Fixture {
1383 26 : Fixture::start_with(config, false).await
1384 26 : }
1385 :
1386 : /// A server whose harness runs start real processes
1387 1 : async fn start_with_real_harness(config: Value) -> Fixture {
1388 1 : Fixture::start_with(config, true).await
1389 1 : }
1390 :
1391 27 : async fn start_with(config: Value, real_harness: bool) -> Fixture {
1392 27 : let home = temp_dir("dashql-agent");
1393 27 : let mut stored = json!({ "token": TOKEN, "defaultHarness": "codex", "callWaitMs": 300 });
1394 27 : stored.as_object_mut().unwrap().extend(config.as_object().unwrap().clone());
1395 27 : std::fs::write(home.join("config.json"), stored.to_string()).unwrap();
1396 27 : let mut config = load_config(&home, &|_| None).unwrap();
1397 27 : std::fs::remove_file(home.join("config.json")).unwrap();
1398 : // Any free port
1399 27 : config.port = 0;
1400 27 : let runs: Arc<Mutex<Vec<FakeRun>>> = Arc::default();
1401 27 : let mut options = AgentServerOptions::new(config, home.clone());
1402 27 : options.resolve_path = Arc::new(|| {
1403 11 : Box::pin(async { Some(OsString::from("/usr/bin")) }) as Pin<Box<dyn Future<Output = _> + Send>>
1404 11 : });
1405 27 : options.spawn_harness = Arc::new({
1406 27 : let runs = runs.clone();
1407 15 : move |args| {
1408 15 : let (finish, outcome) = oneshot::channel();
1409 15 : let finish = Arc::new(Mutex::new(Some(finish)));
1410 15 : let killed = Arc::new(AtomicBool::new(false));
1411 15 : let kill: Kill = Arc::new({
1412 15 : let (finish, killed) = (finish.clone(), killed.clone());
1413 3 : move || {
1414 3 : killed.store(true, Ordering::SeqCst);
1415 3 : if let Some(finish) = finish.lock().unwrap().take() {
1416 3 : let _ = finish.send(HarnessOutcome::failed("killed"));
1417 3 : }
1418 3 : }
1419 : });
1420 15 : runs.lock().unwrap().push(FakeRun { args, finish, killed });
1421 : HarnessProcess {
1422 15 : outcome: Box::pin(async move {
1423 15 : outcome.await.unwrap_or_else(|_| HarnessOutcome::failed("dropped"))
1424 9 : }),
1425 15 : kill,
1426 : }
1427 15 : }
1428 : });
1429 27 : if real_harness {
1430 1 : options.spawn_harness = Arc::new(spawn_harness);
1431 26 : }
1432 27 : let server = AgentServer::listen(options).await.unwrap();
1433 27 : Fixture { home, server, runs }
1434 27 : }
1435 :
1436 63 : fn port(&self) -> u16 {
1437 63 : self.server.port()
1438 63 : }
1439 :
1440 24 : async fn connect(&self, origin: &str) -> TestClient {
1441 24 : TestClient::connect(self.port(), Some(origin), &format!("127.0.0.1:{}", self.port())).await.unwrap()
1442 24 : }
1443 :
1444 8 : fn call(&self, body: Value) -> tokio::task::JoinHandle<(u16, Value)> {
1445 8 : tokio::spawn(call(self.port(), body, vec![]))
1446 8 : }
1447 :
1448 : /// Wait until the fake harness was started n times
1449 14 : async fn wait_for_runs(&self, n: usize) {
1450 14 : for _ in 0..200 {
1451 14 : if self.runs.lock().unwrap().len() >= n {
1452 14 : return;
1453 0 : }
1454 0 : tokio::time::sleep(Duration::from_millis(10)).await;
1455 : }
1456 0 : panic!("the harness did not start");
1457 14 : }
1458 :
1459 21 : fn run<T>(&self, index: usize, f: impl FnOnce(&FakeRun) -> T) -> T {
1460 21 : f(&self.runs.lock().unwrap()[index])
1461 21 : }
1462 : }
1463 :
1464 : /// A WebSocket client that can send any Origin and Host
1465 : struct TestClient {
1466 : socket: WebSocketStream<TcpStream>,
1467 : messages: Vec<Value>,
1468 : /// The methods of the sent requests by id, to check the responses
1469 : requests: BTreeMap<String, String>,
1470 : }
1471 :
1472 : impl TestClient {
1473 : /// Connect, the HTTP status of a refused handshake
1474 29 : async fn connect(port: u16, origin: Option<&str>, host: &str) -> Result<TestClient, u16> {
1475 29 : let mut stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
1476 29 : let origin = origin.map(|o| format!("Origin: {o}\r\n")).unwrap_or_default();
1477 29 : let request = format!(
1478 : "GET /v1/ws HTTP/1.1\r\nHost: {host}\r\n{origin}Connection: Upgrade\r\nUpgrade: websocket\r\n\
1479 : Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n"
1480 : );
1481 29 : stream.write_all(request.as_bytes()).await.unwrap();
1482 29 : let mut buffer = Vec::new();
1483 29 : let mut chunk = [0; 1024];
1484 29 : let length = loop {
1485 29 : let n = stream.read(&mut chunk).await.unwrap();
1486 29 : assert!(n > 0, "the server closed the connection");
1487 29 : buffer.extend_from_slice(&chunk[..n]);
1488 3394 : if let Some(end) = buffer.windows(4).position(|w| w == b"\r\n\r\n") {
1489 29 : break end + 4;
1490 0 : }
1491 : };
1492 29 : let (status, _) = http::parse_response(&buffer[..length]).unwrap();
1493 29 : if status != 101 {
1494 4 : return Err(status);
1495 25 : }
1496 25 : let socket =
1497 25 : WebSocketStream::from_partially_read(stream, buffer[length..].to_vec(), Role::Client, None).await;
1498 25 : Ok(TestClient { socket, messages: Vec::new(), requests: BTreeMap::new() })
1499 29 : }
1500 :
1501 65 : async fn send(&mut self, message: Value) {
1502 65 : if let (Some(method), false) = (message["method"].as_str(), message["id"].is_null()) {
1503 61 : self.requests.insert(message["id"].to_string(), method.to_string());
1504 61 : }
1505 65 : self.socket.send(Message::text(message.to_string())).await.unwrap();
1506 65 : }
1507 :
1508 : /// Wait for the next message that matches
1509 82 : async fn next(&mut self, matches: impl Fn(&Value) -> bool) -> Result<Value, String> {
1510 82 : let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
1511 : loop {
1512 176 : if let Some(index) = self.messages.iter().position(&matches) {
1513 79 : return Ok(self.messages.remove(index));
1514 97 : }
1515 97 : match tokio::time::timeout_at(deadline, self.socket.next()).await {
1516 1 : Err(_) => return Err("timed out waiting for a message".into()),
1517 94 : Ok(Some(Ok(Message::Text(text)))) => {
1518 94 : let message: Value = serde_json::from_str(text.as_str()).unwrap();
1519 94 : let request = self.requests.get(&message["id"].to_string()).map(String::as_str);
1520 94 : if let Err(error) = BridgeSchema::get().check_from_server(&message, request) {
1521 0 : panic!("the server sent a message that the schema does not allow: {error}");
1522 94 : }
1523 94 : self.messages.push(message)
1524 : }
1525 : Ok(Some(Ok(Message::Close(_)))) | Ok(Some(Err(_))) | Ok(None) => {
1526 2 : return Err("the connection is closed".into())
1527 : }
1528 0 : Ok(Some(Ok(_))) => {}
1529 : }
1530 : }
1531 82 : }
1532 :
1533 19 : async fn announce(&mut self, sessions: &[&str], focused: Option<&str>) {
1534 19 : let sessions: Vec<Value> =
1535 21 : sessions.iter().map(|id| json!({ "id": id, "name": format!("Session {id}") })).collect();
1536 19 : let params =
1537 19 : json!({ "protocol": PROTOCOL_VERSION, "client": "web", "sessions": sessions, "focused": focused });
1538 19 : self.send(json!({ "id": 1, "method": "client/sessions", "params": params })).await;
1539 19 : self.next(|m| m["id"] == 1).await.unwrap();
1540 19 : }
1541 : }
1542 :
1543 : /// A call of the command line tool
1544 10 : async fn call(port: u16, body: Value, headers: Vec<(&str, &str)>) -> (u16, Value) {
1545 10 : let mut stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
1546 10 : let body = body.to_string();
1547 10 : let mut head = format!("POST /v1/call HTTP/1.1\r\nContent-Length: {}\r\n", body.len());
1548 10 : let mut headers = headers;
1549 10 : if !headers.iter().any(|(n, _)| *n == "Authorization") {
1550 9 : headers.push(("Authorization", "Bearer test-token"));
1551 9 : }
1552 10 : let host = format!("127.0.0.1:{port}");
1553 10 : headers.push(("Host", &host));
1554 21 : for (name, value) in headers {
1555 21 : head += &format!("{name}: {value}\r\n");
1556 21 : }
1557 10 : stream.write_all(format!("{head}\r\n{body}").as_bytes()).await.unwrap();
1558 10 : let mut response = Vec::new();
1559 10 : stream.read_to_end(&mut response).await.unwrap();
1560 10 : let (status, body) = http::parse_response(&response).unwrap();
1561 10 : let body: Value = serde_json::from_slice(body).unwrap();
1562 10 : if let Err(error) = BridgeSchema::get().check("AgentCallResponse", &body) {
1563 0 : panic!("the server sent a response that the schema does not allow: {error}");
1564 10 : }
1565 10 : (status, body)
1566 10 : }
1567 :
1568 : #[tokio::test]
1569 1 : async fn accepts_dashql_origins_on_the_loopback_host() {
1570 1 : let fixture = Fixture::start(json!({})).await;
1571 1 : let _first = fixture.connect("https://dashql.app").await;
1572 1 : let _second =
1573 1 : TestClient::connect(fixture.port(), Some("app://bundle"), &format!("localhost:{}", fixture.port()))
1574 1 : .await
1575 1 : .unwrap();
1576 1 : tokio::time::sleep(Duration::from_millis(50)).await;
1577 1 : assert_eq!(fixture.server.status()["clients"].as_array().unwrap().len(), 2);
1578 1 : }
1579 :
1580 : #[tokio::test]
1581 1 : async fn refuses_other_origins_and_hosts() {
1582 1 : let fixture = Fixture::start(json!({})).await;
1583 1 : let port = fixture.port();
1584 1 : let host = format!("127.0.0.1:{port}");
1585 1 : assert_eq!(TestClient::connect(port, Some("https://evil.example"), &host).await.err(), Some(403));
1586 1 : assert_eq!(
1587 1 : TestClient::connect(port, Some("https://dashql.app"), &format!("evil.example:{port}")).await.err(),
1588 : Some(403)
1589 : );
1590 1 : assert_eq!(TestClient::connect(port, None, &host).await.err(), Some(403));
1591 1 : }
1592 :
1593 : #[tokio::test]
1594 1 : async fn accepts_the_dev_server_and_configured_origins() {
1595 1 : let fixture = Fixture::start(json!({})).await;
1596 1 : let host = format!("127.0.0.1:{}", fixture.port());
1597 1 : fixture.connect("http://localhost:9002").await;
1598 1 : assert_eq!(TestClient::connect(fixture.port(), Some("http://localhost:5173"), &host).await.err(), Some(403));
1599 1 : let fixture = Fixture::start(json!({ "origins": ["http://localhost:5173"] })).await;
1600 1 : fixture.connect("http://localhost:5173").await;
1601 1 : }
1602 :
1603 : #[tokio::test]
1604 1 : async fn turns_away_other_protocol_versions() {
1605 1 : let fixture = Fixture::start(json!({})).await;
1606 1 : let mut client = fixture.connect("https://dashql.app").await;
1607 1 : let params = json!({ "protocol": PROTOCOL_VERSION + 1, "client": "web", "sessions": [], "focused": null });
1608 1 : client.send(json!({ "id": 1, "method": "client/sessions", "params": params })).await;
1609 1 : assert_eq!(client.next(|m| m["id"] == 1).await.unwrap_err(), "the connection is closed");
1610 1 : }
1611 :
1612 : #[tokio::test]
1613 1 : async fn requires_the_token_and_refuses_browsers() {
1614 1 : let fixture = Fixture::start(json!({})).await;
1615 1 : let body = json!({ "command": "scripts" });
1616 1 : assert_eq!(call(fixture.port(), body.clone(), vec![("Authorization", "Bearer wrong")]).await.0, 401);
1617 1 : assert_eq!(call(fixture.port(), body, vec![("Origin", "https://dashql.app")]).await.0, 403);
1618 1 : }
1619 :
1620 : #[tokio::test]
1621 1 : async fn refuses_unknown_commands() {
1622 1 : let fixture = Fixture::start(json!({})).await;
1623 1 : let (_, body) = fixture.call(json!({ "command": "rm" })).await.unwrap();
1624 1 : assert_eq!(body, json!({ "ok": false, "error": { "code": "usage", "message": "unknown command: rm" } }));
1625 1 : }
1626 :
1627 : #[tokio::test]
1628 1 : async fn routes_a_call_to_the_focused_session() {
1629 1 : let fixture = Fixture::start(json!({})).await;
1630 1 : let mut client = fixture.connect("https://dashql.app").await;
1631 1 : client.announce(&["a", "b"], Some("b")).await;
1632 1 : let response = fixture.call(json!({ "command": "read", "args": { "script": "q1" } }));
1633 1 : let command = client.next(|m| m["method"] == "command").await.unwrap();
1634 1 : assert_eq!(command["params"], json!({ "session": "b", "command": "read", "args": { "script": "q1" } }));
1635 1 : client.send(json!({ "id": command["id"], "result": { "text": "select 1" } })).await;
1636 1 : let (_, body) = response.await.unwrap();
1637 1 : assert_eq!(body, json!({ "ok": true, "session": "b", "result": { "text": "select 1" } }));
1638 1 : }
1639 :
1640 : #[tokio::test]
1641 1 : async fn routes_a_call_to_the_client_holding_the_session() {
1642 1 : let fixture = Fixture::start(json!({})).await;
1643 1 : let mut first = fixture.connect("https://dashql.app").await;
1644 1 : let mut second = fixture.connect("app://bundle").await;
1645 1 : first.announce(&["a"], Some("a")).await;
1646 1 : second.announce(&["b"], None).await;
1647 1 : let response = fixture.call(json!({ "session": "b", "command": "scripts" }));
1648 1 : let command = second.next(|m| m["method"] == "command").await.unwrap();
1649 1 : second.send(json!({ "id": command["id"], "error": { "code": "failed", "message": "no such script" } })).await;
1650 1 : let (_, body) = response.await.unwrap();
1651 1 : assert_eq!(body, json!({ "ok": false, "error": { "code": "failed", "message": "no such script" } }));
1652 1 : assert!(first.next(|m| m["method"] == "command").await.is_err());
1653 1 : }
1654 :
1655 : #[tokio::test]
1656 1 : async fn queues_a_call_until_a_client_holds_the_session() {
1657 1 : let fixture = Fixture::start(json!({ "callWaitMs": 2000 })).await;
1658 1 : let response = fixture.call(json!({ "session": "a", "command": "scripts" }));
1659 1 : tokio::time::sleep(Duration::from_millis(100)).await;
1660 1 : let mut client = fixture.connect("https://dashql.app").await;
1661 1 : client.announce(&["a"], None).await;
1662 1 : let command = client.next(|m| m["method"] == "command").await.unwrap();
1663 1 : client.send(json!({ "id": command["id"], "result": { "text": "ok" } })).await;
1664 1 : assert_eq!(response.await.unwrap().1["ok"], true);
1665 1 : }
1666 :
1667 : #[tokio::test]
1668 1 : async fn fails_when_no_client_holds_the_session() {
1669 1 : let fixture = Fixture::start(json!({ "callWaitMs": 50 })).await;
1670 1 : let (_, body) = fixture.call(json!({ "session": "a", "command": "scripts" })).await.unwrap();
1671 1 : assert_eq!(body["error"]["code"], "not_connected");
1672 1 : }
1673 :
1674 : #[tokio::test]
1675 1 : async fn runs_one_call_at_a_time_per_session() {
1676 1 : let fixture = Fixture::start(json!({})).await;
1677 1 : let mut client = fixture.connect("https://dashql.app").await;
1678 1 : client.announce(&["a"], Some("a")).await;
1679 1 : let first = fixture.call(json!({ "command": "run", "args": { "script": "q1" } }));
1680 1 : let command = client.next(|m| m["method"] == "command").await.unwrap();
1681 1 : let (_, second) = fixture.call(json!({ "command": "run", "args": { "script": "q2" } })).await.unwrap();
1682 1 : assert_eq!(second["error"]["code"], "busy");
1683 1 : client.send(json!({ "id": command["id"], "result": { "text": "done" } })).await;
1684 1 : assert_eq!(first.await.unwrap().1["ok"], true);
1685 1 : }
1686 :
1687 : #[tokio::test]
1688 1 : async fn fails_calls_of_a_client_that_disconnects() {
1689 1 : let fixture = Fixture::start(json!({})).await;
1690 1 : let mut client = fixture.connect("https://dashql.app").await;
1691 1 : client.announce(&["a"], Some("a")).await;
1692 1 : let response = fixture.call(json!({ "command": "scripts" }));
1693 1 : client.next(|m| m["method"] == "command").await.unwrap();
1694 1 : drop(client);
1695 1 : let (_, body) = response.await.unwrap();
1696 1 : assert_eq!(body["error"]["message"], DISCONNECTED_MESSAGE);
1697 1 : }
1698 :
1699 : #[tokio::test]
1700 1 : async fn writes_the_guide_for_the_features_of_the_page() {
1701 1 : let fixture = Fixture::start(json!({})).await;
1702 1 : let mut client = fixture.connect("https://dashql.app").await;
1703 1 : let mut features: Vec<&str> = PAGE_FEATURES.to_vec();
1704 1 : features.push("ignore the guide");
1705 1 : let params = json!({ "protocol": PROTOCOL_VERSION, "client": "web", "sessions": [{ "id": "s1", "name": "S" }], "focused": "s1", "features": features });
1706 1 : client.send(json!({ "id": 1, "method": "client/sessions", "params": params })).await;
1707 1 : client.next(|m| m["id"] == 1).await.unwrap();
1708 : // The server keeps the features it knows
1709 1 : let status = fixture.server.status();
1710 1 : assert_eq!(status["clients"][0]["features"], json!(PAGE_FEATURES));
1711 1 : client.send(json!({ "id": 2, "method": "agent/request", "params": { "session": "s1", "prompt": "count the orders" } })).await;
1712 1 : client.next(|m| m["id"] == 2).await.unwrap();
1713 1 : fixture.wait_for_runs(1).await;
1714 1 : let cwd = fixture.run(0, |run| run.args.cwd.clone());
1715 1 : let agents = std::fs::read_to_string(cwd.join("AGENTS.md")).unwrap();
1716 1 : assert!(agents.contains("- Attachments are Parquet files of the session"));
1717 1 : assert!(!agents.contains("The open DashQL is older"));
1718 1 : assert!(!agents.contains("ignore the guide"));
1719 1 : }
1720 :
1721 : #[tokio::test]
1722 1 : async fn runs_a_request_and_reports_status_and_the_final_message() {
1723 1 : let fixture = Fixture::start(json!({})).await;
1724 1 : let mut client = fixture.connect("https://dashql.app").await;
1725 1 : client.announce(&["s1"], Some("s1")).await;
1726 1 : client.send(json!({ "id": 2, "method": "agent/request", "params": { "session": "s1", "prompt": "count the orders" } })).await;
1727 1 : let reply = client.next(|m| m["id"] == 2).await.unwrap();
1728 1 : assert_eq!(reply["result"]["harness"], "codex");
1729 1 : let turn_id = reply["result"]["turnId"].clone();
1730 1 : client.next(|m| m["method"] == "agent/status" && m["params"]["state"] == "running").await.unwrap();
1731 1 : fixture.wait_for_runs(1).await;
1732 1 : let cwd = fixture.run(0, |run| {
1733 1 : assert_eq!(run.args.prompt, "count the orders");
1734 1 : assert_eq!(run.args.harness_session, None);
1735 1 : assert_eq!(run.args.env.get(&OsString::from("DASHQL_SESSION")), Some(&OsString::from("s1")));
1736 1 : assert_eq!(run.args.env.get(&OsString::from(AGENT_RUN_ENV)), Some(&OsString::from("1")));
1737 1 : run.args.cwd.clone()
1738 1 : });
1739 1 : assert_eq!(cwd, fixture.home.join("sessions").join("s1"));
1740 1 : assert!(std::fs::read_to_string(cwd.join("AGENTS.md")).unwrap().contains("<<'EOF'"));
1741 1 : assert!(std::fs::read_to_string(cwd.join("CLAUDE.md")).unwrap().contains("--file input/q1.json"));
1742 : // The page announced no features, the guide asks for a newer DashQL
1743 1 : assert!(std::fs::read_to_string(cwd.join("AGENTS.md")).unwrap().contains("The open DashQL is older"));
1744 1 : assert!(std::fs::read_dir(cwd.join("input")).unwrap().next().is_none());
1745 :
1746 1 : fixture.run(0, |run| run.args.events.send(HarnessEvent::Status("dashql agent catalog".into())).unwrap());
1747 1 : let status = client
1748 1 : .next(|m| m["method"] == "agent/status" && m["params"]["text"] == "dashql agent catalog")
1749 1 : .await
1750 1 : .unwrap();
1751 1 : assert_eq!(status["params"]["turnId"], turn_id);
1752 1 : fixture.run(0, |run| run.args.events.send(HarnessEvent::Session("thread-1".into())).unwrap());
1753 : // The events of a turn are handled before its outcome
1754 1 : fixture.run(0, |run| run.finish("There are 42 orders", None));
1755 1 : let done = client.next(|m| m["method"] == "agent/done").await.unwrap();
1756 1 : assert_eq!(done["params"], json!({ "session": "s1", "turnId": turn_id, "message": "There are 42 orders" }));
1757 :
1758 : // The next request resumes the harness session
1759 1 : client.send(json!({ "id": 3, "method": "agent/request", "params": { "session": "s1", "prompt": "and yesterday?" } })).await;
1760 1 : client.next(|m| m["id"] == 3).await.unwrap();
1761 1 : client.next(|m| m["method"] == "agent/status" && m["params"]["state"] == "running").await.unwrap();
1762 1 : fixture.wait_for_runs(2).await;
1763 1 : assert_eq!(fixture.run(1, |run| run.args.harness_session.clone()), Some("thread-1".into()));
1764 1 : client.send(json!({ "id": 4, "method": "agent/context", "params": { "session": "s1" } })).await;
1765 1 : let context = client.next(|m| m["id"] == 4).await.unwrap();
1766 1 : assert_eq!(context["result"]["terminalCommand"], json!(["codex", "resume", "thread-1"]));
1767 :
1768 : // Resetting forgets it
1769 1 : fixture.run(1, |run| run.finish("", None));
1770 1 : client.next(|m| m["method"] == "agent/done").await.unwrap();
1771 1 : client.send(json!({ "id": 5, "method": "agent/reset", "params": { "session": "s1" } })).await;
1772 1 : client.next(|m| m["id"] == 5).await.unwrap();
1773 1 : client
1774 1 : .send(json!({ "id": 6, "method": "agent/request", "params": { "session": "s1", "prompt": "again" } }))
1775 1 : .await;
1776 1 : client.next(|m| m["id"] == 6).await.unwrap();
1777 1 : fixture.wait_for_runs(3).await;
1778 1 : assert_eq!(fixture.run(2, |run| run.args.harness_session.clone()), None);
1779 1 : }
1780 :
1781 : #[tokio::test]
1782 1 : async fn describes_the_conversation_of_the_picked_harness() {
1783 1 : let fixture = Fixture::start(json!({})).await;
1784 1 : let mut client = fixture.connect("https://dashql.app").await;
1785 1 : client.announce(&["s1"], Some("s1")).await;
1786 1 : client.send(json!({ "id": 2, "method": "agent/context", "params": { "session": "s1" } })).await;
1787 1 : let context = client.next(|m| m["id"] == 2).await.unwrap();
1788 1 : assert_eq!(context["result"]["harness"], "codex");
1789 1 : assert_eq!(context["result"]["harnessSession"], Value::Null);
1790 :
1791 1 : client
1792 1 : .send(json!({ "id": 3, "method": "agent/request", "params": { "session": "s1", "harness": "opencode", "prompt": "hi" } }))
1793 1 : .await;
1794 1 : client.next(|m| m["id"] == 3).await.unwrap();
1795 1 : fixture.wait_for_runs(1).await;
1796 1 : fixture.run(0, |run| run.args.events.send(HarnessEvent::Session("ses_1".into())).unwrap());
1797 1 : fixture.run(0, |run| run.finish("Hello", None));
1798 3 : client.next(|m| m["method"] == "agent/done").await.unwrap();
1799 1 : let params = json!({ "session": "s1", "harness": "opencode" });
1800 1 : client.send(json!({ "id": 4, "method": "agent/context", "params": params })).await;
1801 3 : let context = client.next(|m| m["id"] == 4).await.unwrap();
1802 1 : assert_eq!(context["result"]["harness"], "opencode");
1803 1 : assert_eq!(context["result"]["harnessSession"], "ses_1");
1804 1 : assert_eq!(context["result"]["terminalCommand"], json!(["opencode", "--session", "ses_1"]));
1805 : // The default harness has no conversation yet, a request without a harness starts a new one
1806 1 : client.send(json!({ "id": 5, "method": "agent/context", "params": { "session": "s1" } })).await;
1807 3 : let context = client.next(|m| m["id"] == 5).await.unwrap();
1808 1 : assert_eq!(context["result"]["harness"], "codex");
1809 1 : assert_eq!(context["result"]["harnessSession"], Value::Null);
1810 1 : }
1811 :
1812 : #[tokio::test]
1813 1 : async fn queues_requests_of_a_session_and_cancels_them() {
1814 1 : let fixture = Fixture::start(json!({})).await;
1815 1 : let mut client = fixture.connect("https://dashql.app").await;
1816 1 : client.announce(&["s1"], Some("s1")).await;
1817 1 : client
1818 1 : .send(json!({ "id": 2, "method": "agent/request", "params": { "session": "s1", "prompt": "first" } }))
1819 1 : .await;
1820 1 : client
1821 1 : .send(json!({ "id": 3, "method": "agent/request", "params": { "session": "s1", "prompt": "second" } }))
1822 1 : .await;
1823 1 : client.next(|m| m["id"] == 2).await.unwrap();
1824 3 : let second = client.next(|m| m["id"] == 3).await.unwrap();
1825 1 : let second_id = second["result"]["turnId"].clone();
1826 1 : client
1827 3 : .next(|m| {
1828 3 : m["method"] == "agent/status" && m["params"]["turnId"] == second_id && m["params"]["state"] == "queued"
1829 3 : })
1830 1 : .await
1831 1 : .unwrap();
1832 1 : fixture.wait_for_runs(1).await;
1833 :
1834 1 : client.send(json!({ "id": 4, "method": "agent/cancel", "params": { "session": "s1" } })).await;
1835 1 : for _ in 0..2 {
1836 9 : let failed = client.next(|m| m["method"] == "agent/failed").await.unwrap();
1837 2 : assert_eq!(failed["params"]["cancelled"], true);
1838 : }
1839 1 : assert!(fixture.run(0, |run| run.killed.load(Ordering::SeqCst)));
1840 1 : assert_eq!(fixture.runs.lock().unwrap().len(), 1);
1841 1 : assert_eq!(fixture.server.status()["turns"], json!([]));
1842 1 : }
1843 :
1844 : #[tokio::test]
1845 1 : async fn forgets_a_deleted_session() {
1846 1 : let fixture = Fixture::start(json!({})).await;
1847 1 : let mut client = fixture.connect("https://dashql.app").await;
1848 1 : client.announce(&["s1", "s2"], Some("s1")).await;
1849 2 : for (id, session) in [(2, "s1"), (3, "s2")] {
1850 2 : let params = json!({ "session": session, "prompt": "count the orders" });
1851 2 : client.send(json!({ "id": id, "method": "agent/request", "params": params })).await;
1852 4 : client.next(|m| m["id"] == id).await.unwrap();
1853 : }
1854 1 : fixture.wait_for_runs(2).await;
1855 1 : let s1 = fixture.home.join("sessions").join("s1");
1856 1 : let s2 = fixture.home.join("sessions").join("s2");
1857 1 : fixture.run(1, |run| run.finish("42", None));
1858 6 : client.next(|m| m["method"] == "agent/done").await.unwrap();
1859 :
1860 : // A session without turns goes right away
1861 1 : client.send(json!({ "id": 4, "method": "agent/forget", "params": { "session": "s2" } })).await;
1862 5 : client.next(|m| m["id"] == 4).await.unwrap();
1863 1 : assert!(!s2.exists());
1864 :
1865 : // A running turn stops, the folder goes once the harness exited
1866 1 : assert!(s1.exists());
1867 1 : client.send(json!({ "id": 5, "method": "agent/forget", "params": { "session": "s1" } })).await;
1868 5 : client.next(|m| m["id"] == 5).await.unwrap();
1869 5 : let failed = client.next(|m| m["method"] == "agent/failed").await.unwrap();
1870 1 : assert_eq!(failed["params"]["cancelled"], true);
1871 1 : assert!(fixture.run(0, |run| run.killed.load(Ordering::SeqCst)));
1872 1 : for _ in 0..200 {
1873 1 : if !s1.exists() {
1874 1 : break;
1875 0 : }
1876 0 : tokio::time::sleep(Duration::from_millis(10)).await;
1877 : }
1878 1 : assert!(!s1.exists());
1879 :
1880 1 : client.send(json!({ "id": 6, "method": "agent/forget", "params": { "session": "" } })).await;
1881 5 : let error = client.next(|m| m["id"] == 6).await.unwrap();
1882 1 : assert_eq!(error["error"]["message"], "the session is empty");
1883 1 : assert!(fixture.home.join("sessions").exists());
1884 1 : }
1885 :
1886 : #[tokio::test]
1887 1 : async fn removes_the_folders_of_unused_sessions() {
1888 1 : let fixture = Fixture::start(json!({})).await;
1889 1 : let sessions = fixture.home.join("sessions");
1890 1 : let long_ago = SystemTime::now() - SESSION_RETENTION - Duration::from_secs(60);
1891 3 : let folder = |name: &str, used: Option<SystemTime>| {
1892 3 : let dir = sessions.join(name);
1893 3 : std::fs::create_dir_all(&dir).unwrap();
1894 3 : std::fs::write(dir.join("state.json"), "{}").unwrap();
1895 3 : if let Some(used) = used {
1896 2 : std::fs::File::options().write(true).open(dir.join("state.json")).unwrap().set_modified(used).unwrap();
1897 2 : }
1898 3 : dir
1899 3 : };
1900 1 : let unused = folder("unused", Some(long_ago));
1901 1 : let recent = folder("recent", None);
1902 1 : let shown = folder("shown", Some(long_ago));
1903 :
1904 : // Announcing a session marks it as used
1905 1 : let mut client = fixture.connect("https://dashql.app").await;
1906 1 : client.announce(&["shown"], Some("shown")).await;
1907 1 : let marked = std::fs::metadata(shown.join("state.json")).unwrap().modified().unwrap();
1908 1 : assert!(marked > long_ago);
1909 :
1910 : // A session that a client shows stays while its last use is long ago
1911 1 : std::fs::File::options().write(true).open(shown.join("state.json")).unwrap().set_modified(long_ago).unwrap();
1912 1 : fixture.server.sweep_sessions();
1913 1 : assert!(!unused.exists());
1914 1 : assert!(recent.exists());
1915 1 : assert!(shown.exists());
1916 :
1917 : // A turn marks its session as used
1918 1 : client.send(json!({ "id": 2, "method": "agent/request", "params": { "session": "shown", "prompt": "hi" } })).await;
1919 1 : client.next(|m| m["id"] == 2).await.unwrap();
1920 1 : fixture.wait_for_runs(1).await;
1921 1 : let marked = std::fs::metadata(shown.join("state.json")).unwrap().modified().unwrap();
1922 1 : assert!(marked > long_ago);
1923 1 : }
1924 :
1925 : #[tokio::test]
1926 1 : async fn cancels_one_turn_of_a_session() {
1927 1 : let fixture = Fixture::start(json!({})).await;
1928 1 : let mut client = fixture.connect("https://dashql.app").await;
1929 1 : client.announce(&["s1"], Some("s1")).await;
1930 3 : for (id, prompt) in [(2, "first"), (3, "second"), (4, "third")] {
1931 3 : client.send(json!({ "id": id, "method": "agent/request", "params": { "session": "s1", "prompt": prompt } })).await;
1932 : }
1933 1 : let mut turn_ids = Vec::new();
1934 3 : for id in [2, 3, 4] {
1935 7 : turn_ids.push(client.next(|m| m["id"] == id).await.unwrap()["result"]["turnId"].clone());
1936 : }
1937 1 : fixture.wait_for_runs(1).await;
1938 :
1939 : // Stopping a queued turn leaves the running turn alone
1940 1 : client.send(json!({ "id": 5, "method": "agent/cancel", "params": { "session": "s1", "turnId": turn_ids[1] } })).await;
1941 10 : let failed = client.next(|m| m["method"] == "agent/failed").await.unwrap();
1942 1 : assert_eq!((&failed["params"]["turnId"], &failed["params"]["cancelled"]), (&turn_ids[1], &json!(true)));
1943 1 : assert!(!fixture.run(0, |run| run.killed.load(Ordering::SeqCst)));
1944 :
1945 : // Stopping the running turn starts the next one
1946 1 : client.send(json!({ "id": 6, "method": "agent/cancel", "params": { "session": "s1", "turnId": turn_ids[0] } })).await;
1947 18 : let failed = client.next(|m| m["method"] == "agent/failed").await.unwrap();
1948 1 : assert_eq!((&failed["params"]["turnId"], &failed["params"]["cancelled"]), (&turn_ids[0], &json!(true)));
1949 1 : assert!(fixture.run(0, |run| run.killed.load(Ordering::SeqCst)));
1950 1 : fixture.wait_for_runs(2).await;
1951 1 : assert_eq!(fixture.run(1, |run| run.args.prompt.clone()), "third");
1952 :
1953 1 : client.send(json!({ "id": 7, "method": "agent/cancel", "params": { "session": "s1", "turnId": 1 } })).await;
1954 18 : let error = client.next(|m| m["id"] == 7).await.unwrap();
1955 1 : assert_eq!(error["error"]["code"], "usage");
1956 1 : }
1957 :
1958 : #[cfg(unix)]
1959 : #[tokio::test]
1960 1 : async fn fails_and_kills_the_turns_when_it_stops() {
1961 1 : let harness = json!({ "command": ["sh", "-c", "echo $$ > harness.pid; exec sleep 30"] });
1962 1 : let fixture = Fixture::start_with_real_harness(json!({ "harnesses": { "codex": harness } })).await;
1963 1 : let mut client = fixture.connect("https://dashql.app").await;
1964 1 : client.announce(&["s1"], Some("s1")).await;
1965 2 : for (id, prompt) in [(2, "first"), (3, "second")] {
1966 2 : let params = json!({ "session": "s1", "prompt": prompt });
1967 2 : client.send(json!({ "id": id, "method": "agent/request", "params": params })).await;
1968 4 : client.next(|m| m["id"] == id).await.unwrap();
1969 : }
1970 1 : let pid_file = fixture.home.join("sessions").join("s1").join("harness.pid");
1971 1 : let mut pid = None;
1972 1 : for _ in 0..200 {
1973 1 : pid = std::fs::read_to_string(&pid_file).ok().and_then(|text| text.trim().parse::<i32>().ok());
1974 1 : if pid.is_some() {
1975 1 : break;
1976 0 : }
1977 0 : tokio::time::sleep(Duration::from_millis(10)).await;
1978 : }
1979 1 : let pid = pid.expect("the harness did not start");
1980 :
1981 1 : fixture.server.close().await;
1982 : // The harness exited and was reaped
1983 1 : assert_eq!(unsafe { libc::kill(pid, 0) }, -1);
1984 1 : for _ in 0..2 {
1985 11 : let failed = client.next(|m| m["method"] == "agent/failed").await.unwrap();
1986 2 : assert_eq!(failed["params"]["error"], "The agent server stopped");
1987 2 : assert_eq!(failed["params"]["cancelled"], false);
1988 1 : }
1989 1 : }
1990 :
1991 : #[tokio::test]
1992 1 : async fn refuses_requests_for_sessions_without_agent_access() {
1993 1 : let fixture = Fixture::start(json!({})).await;
1994 1 : let mut client = fixture.connect("https://dashql.app").await;
1995 1 : client.announce(&["s1"], Some("s1")).await;
1996 1 : client
1997 1 : .send(json!({ "id": 2, "method": "agent/request", "params": { "session": "other", "prompt": "hi" } }))
1998 1 : .await;
1999 1 : assert_eq!(client.next(|m| m["id"] == 2).await.unwrap()["error"]["code"], "usage");
2000 1 : let params = json!({ "session": "s1", "harness": "rm -rf /", "prompt": "hi" });
2001 1 : client.send(json!({ "id": 3, "method": "agent/request", "params": params })).await;
2002 1 : assert_eq!(client.next(|m| m["id"] == 3).await.unwrap()["error"]["message"], "unknown harness: rm -rf /");
2003 1 : }
2004 :
2005 : #[tokio::test]
2006 1 : async fn reports_a_failed_turn() {
2007 1 : let fixture = Fixture::start(json!({})).await;
2008 1 : let mut client = fixture.connect("https://dashql.app").await;
2009 1 : client.announce(&["s1"], Some("s1")).await;
2010 1 : client.send(json!({ "id": 2, "method": "agent/request", "params": { "session": "s1", "prompt": "hi" } })).await;
2011 3 : client.next(|m| m["method"] == "agent/status" && m["params"]["state"] == "running").await.unwrap();
2012 1 : fixture.wait_for_runs(1).await;
2013 1 : fixture.run(0, |run| run.finish("", Some("The harness exited with code 1")));
2014 3 : let failed = client.next(|m| m["method"] == "agent/failed").await.unwrap();
2015 1 : assert_eq!(failed["params"]["error"], "The harness exited with code 1");
2016 1 : assert_eq!(failed["params"]["cancelled"], false);
2017 1 : }
2018 :
2019 : #[tokio::test]
2020 1 : async fn describes_the_harnesses_and_runs_the_picked_harness() {
2021 1 : let fixture = Fixture::start(json!({})).await;
2022 1 : let mut client = fixture.connect("https://dashql.app").await;
2023 1 : client.send(json!({ "id": 1, "method": "agent/info", "params": {} })).await;
2024 1 : let info = client.next(|m| m["id"] == 1).await.unwrap()["result"].clone();
2025 1 : assert_eq!(info["defaultHarness"], "codex");
2026 1 : assert_eq!(info["configPath"], fixture.home.join("config.json").display().to_string());
2027 3 : let names: Vec<_> = info["harnesses"].as_array().unwrap().iter().map(|h| h["name"].clone()).collect();
2028 1 : assert_eq!(names, ["claude", "codex", "opencode"]);
2029 1 : assert_eq!(info["harnesses"][0]["command"][0], "claude");
2030 :
2031 1 : client.announce(&["s1"], Some("s1")).await;
2032 1 : let params = json!({ "session": "s1", "harness": "opencode", "prompt": "hi" });
2033 1 : client.send(json!({ "id": 3, "method": "agent/request", "params": params })).await;
2034 1 : assert_eq!(client.next(|m| m["id"] == 3).await.unwrap()["result"]["harness"], "opencode");
2035 1 : fixture.wait_for_runs(1).await;
2036 1 : let cwd = fixture.run(0, |run| run.args.cwd.clone());
2037 1 : assert!(std::fs::read_to_string(cwd.join("opencode.json")).unwrap().contains("\"dashql agent *\": \"allow\""));
2038 : // The server never writes the configuration, the fixture removed it after loading
2039 1 : assert!(!fixture.home.join("config.json").exists());
2040 1 : }
2041 :
2042 : #[tokio::test]
2043 1 : async fn runs_the_picked_model_and_effort() {
2044 1 : let fixture = Fixture::start(json!({})).await;
2045 1 : let mut client = fixture.connect("https://dashql.app").await;
2046 1 : client.send(json!({ "id": 1, "method": "agent/info", "params": {} })).await;
2047 1 : let info = client.next(|m| m["id"] == 1).await.unwrap()["result"].clone();
2048 1 : assert_eq!(info["harnesses"][0]["models"], json!(["fable", "opus", "sonnet", "haiku"]));
2049 1 : assert_eq!(info["harnesses"][1]["models"], json!([]));
2050 :
2051 1 : client.announce(&["s1"], Some("s1")).await;
2052 : // Only the values of config.json reach the harness
2053 3 : for (id, params, error) in [
2054 1 : (
2055 1 : 2,
2056 1 : json!({ "harness": "claude", "model": "--dangerously-skip-permissions" }),
2057 1 : "unknown model of claude: --dangerously-skip-permissions",
2058 1 : ),
2059 1 : (3, json!({ "harness": "codex", "model": "opus" }), "unknown model of codex: opus"),
2060 1 : (4, json!({ "harness": "claude", "effort": "ultra" }), "unknown effort of claude: ultra"),
2061 1 : ] {
2062 3 : let mut params = params;
2063 3 : params["session"] = "s1".into();
2064 3 : params["prompt"] = "hi".into();
2065 3 : client.send(json!({ "id": id, "method": "agent/request", "params": params })).await;
2066 3 : assert_eq!(client.next(|m| m["id"] == id).await.unwrap()["error"]["message"], error);
2067 : }
2068 1 : let params = json!({ "session": "s1", "harness": "claude", "model": "opus", "effort": "high", "prompt": "hi" });
2069 1 : client.send(json!({ "id": 5, "method": "agent/request", "params": params })).await;
2070 1 : client.next(|m| m["id"] == 5).await.unwrap();
2071 1 : fixture.wait_for_runs(1).await;
2072 1 : let picks = fixture.run(0, |run| run.args.picks.clone());
2073 1 : assert_eq!(picks, HarnessPicks { model: Some("opus".into()), effort: Some("high".into()) });
2074 1 : client.send(json!({ "id": 6, "method": "agent/info", "params": {} })).await;
2075 3 : let turn = client.next(|m| m["id"] == 6).await.unwrap()["result"]["turns"][0].clone();
2076 1 : assert_eq!((turn["model"].clone(), turn["effort"].clone()), (json!("opus"), json!("high")));
2077 1 : fixture.run(0, |run| run.finish("done", None));
2078 3 : client.next(|m| m["method"] == "agent/done").await.unwrap();
2079 :
2080 : // A request without picks leaves the defaults to the harness
2081 1 : let params = json!({ "session": "s1", "harness": "claude", "prompt": "hi" });
2082 1 : client.send(json!({ "id": 7, "method": "agent/request", "params": params })).await;
2083 3 : client.next(|m| m["id"] == 7).await.unwrap();
2084 1 : fixture.wait_for_runs(2).await;
2085 1 : assert_eq!(fixture.run(1, |run| run.args.picks.clone()), HarnessPicks::default());
2086 1 : }
2087 :
2088 : #[tokio::test]
2089 1 : async fn refuses_unknown_harnesses() {
2090 1 : let fixture = Fixture::start(json!({})).await;
2091 1 : let mut client = fixture.connect("https://dashql.app").await;
2092 1 : client.announce(&["s1"], Some("s1")).await;
2093 1 : let params = json!({ "session": "s1", "harness": "sh -c evil" });
2094 1 : client.send(json!({ "id": 2, "method": "agent/context", "params": params })).await;
2095 1 : assert_eq!(client.next(|m| m["id"] == 2).await.unwrap()["error"]["message"], "unknown harness: sh -c evil");
2096 1 : }
2097 :
2098 : #[tokio::test]
2099 1 : async fn closes_clients_when_it_stops() {
2100 1 : let fixture = Fixture::start(json!({})).await;
2101 1 : let mut client = fixture.connect("https://dashql.app").await;
2102 1 : client.announce(&["s1"], Some("s1")).await;
2103 1 : fixture.server.close().await;
2104 1 : assert_eq!(client.next(|_| true).await.unwrap_err(), "the connection is closed");
2105 1 : }
2106 :
2107 : #[test]
2108 1 : fn keeps_session_folder_names_inside_the_sessions_folder() {
2109 1 : assert_eq!(session_folder_name("../../etc"), "______etc");
2110 1 : assert_eq!(session_folder_name(".."), "__");
2111 1 : assert_eq!(session_folder_name("3f2a-b1"), "3f2a-b1");
2112 1 : }
2113 :
2114 : #[test]
2115 1 : fn makes_random_uuids() {
2116 1 : let id = random_uuid();
2117 1 : assert_eq!(id.len(), 36);
2118 1 : assert_eq!(&id[14..15], "4");
2119 1 : assert_ne!(id, random_uuid());
2120 1 : }
2121 : }
|
