1use std::sync::atomic::Ordering;
8use std::sync::Arc;
9use std::time::Duration;
10
11use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
12use tokio::sync::{broadcast, mpsc};
13use tokio_util::sync::CancellationToken;
14
15use super::DaemonState;
16use crate::session::handle::CMD_CHAN_CAP;
17use crate::session::transport::{TransportError, ATTACH_TIMEOUT, ATTACH_TIMEOUT_PARKED};
18use crate::session::wire::*;
19use crate::session::*;
20use crate::ipc::{ServerReadHalf, ServerStream, ServerWriteHalf};
21use crate::{SessionEvent, StreamEvent};
22
23const HELLO_TIMEOUT: Duration = Duration::from_secs(2);
24
25type Reader = BufReader<ServerReadHalf>;
26
27fn client_may_send(cmd: &SessionCommand, own: ClientId) -> Result<(), &'static str> {
36 use SessionCommand as C;
37 match cmd {
38 C::Submit { .. }
39 | C::Steer { .. }
40 | C::SubagentSteer { .. }
41 | C::Cancel
42 | C::Answer { .. }
43 | C::Set { .. }
44 | C::Compact { .. }
45 | C::NewSession
46 | C::Save
47 | C::Query { .. }
48 | C::EngineCommand { .. }
49 | C::SubmitPrepared { .. }
50 | C::PluginCommand { .. }
51 | C::Resume { .. }
52 | C::KeepWarm { .. }
53 | C::TakeInput => Ok(()),
55 C::Checkpoint { reason: crate::session::CheckpointReason::Reload } => {
58 Err("checkpoint(reload): host-only (use `synaps daemon reload`)")
59 }
60 C::Checkpoint { .. } => Ok(()),
61 C::Detach { client } if *client == own => Ok(()),
62 C::Detach { .. } => Err("detach: not your client id"),
63 C::Attach { .. } => Err("attach: already attached (one session per connection)"),
64 C::End { reason: EndReason::ClientQuit } => Ok(()),
67 C::End { .. } => Err("end: only ClientQuit may come from a client (host reasons are the daemon's)"),
68 C::Resync { .. } => Err("resync: not a client command"),
69 C::Park => Err("park: not a client command"),
70 C::AdoptWorkers { .. } => Err("adopt_workers: not a client command"),
71 C::RestoreToolActivations { .. } => Err("restore_tool_activations: not a client command"),
72 C::ReloadAborted => Err("reload_aborted: not a client command"),
73 C::ArmPark => Err("arm_park: not a client command"),
74 C::HostEvent(_) => Err("host_event: not a client command"),
75 C::DriverStart { .. } => Ok(()),
76 }
77}
78
79fn refused_while_draining(cmd: &SessionCommand) -> bool {
86 use SessionCommand as C;
87 cmd.refused_while_reloading()
88 && !matches!(
89 cmd,
90 C::Steer { .. } | C::SubagentSteer { .. } | C::Set { .. } | C::EngineCommand { .. } | C::KeepWarm { .. }
91 )
92}
93
94#[allow(clippy::large_enum_variant)]
96enum Read {
97 Frame(ClientFrame),
98 Eof,
99 Oversize,
100 Bad(String),
101}
102
103const FRAME_BUF_KEEP: usize = 64 * 1024;
105const FRAME_BUF_SHRINK_ABOVE: usize = 1024 * 1024;
107
108async fn read_frame(reader: &mut Reader, line: &mut String) -> std::io::Result<Read> {
117 let read = read_frame_into(reader, line).await;
118 release_large_frame_buf(line);
119 read
120}
121
122fn release_large_frame_buf(line: &mut String) {
123 if line.capacity() > FRAME_BUF_SHRINK_ABOVE {
124 *line = String::with_capacity(FRAME_BUF_KEEP);
125 }
126}
127
128async fn read_frame_into(reader: &mut Reader, line: &mut String) -> std::io::Result<Read> {
129 loop {
130 line.clear();
131 let n = tokio::io::AsyncReadExt::take(&mut *reader, MAX_FRAME_BYTES as u64 + 1)
132 .read_line(line)
133 .await?;
134 if n == 0 {
135 return Ok(Read::Eof);
136 }
137 if n > MAX_FRAME_BYTES {
138 return Ok(Read::Oversize);
139 }
140 let t = line.trim_end();
141 if t.is_empty() {
142 continue;
143 }
144 return Ok(match decode_line::<ClientFrame>(t) {
145 Ok(f) => Read::Frame(f),
146 Err(e) => Read::Bad(e),
147 });
148 }
149}
150
151fn spawn_writer(mut w: ServerWriteHalf) -> (mpsc::Sender<DaemonFrame>, tokio::task::JoinHandle<()>) {
153 let (tx, mut rx) = mpsc::channel::<DaemonFrame>(CMD_CHAN_CAP);
154 let task = tokio::spawn(async move {
155 while let Some(frame) = rx.recv().await {
156 let bye = matches!(frame, DaemonFrame::Bye { .. } | DaemonFrame::Refused { .. });
157 match encode_line(&frame) {
158 Ok(line) => {
159 if w.write_all(line.as_bytes()).await.is_err() {
160 break;
161 }
162 }
163 Err(e) => {
164 tracing::warn!("daemon: encode failed: {e}");
167 let err = DaemonFrame::Error { session_id: None, message: format!("daemon could not encode a frame: {e}") };
168 if let Ok(line) = encode_line(&err) {
169 let _ = w.write_all(line.as_bytes()).await;
170 }
171 let _ = w.shutdown().await;
172 break;
173 }
174 }
175 if bye {
176 let _ = w.shutdown().await;
177 break;
178 }
179 }
180 });
181 (tx, task)
182}
183
184pub async fn serve(state: Arc<DaemonState>, stream: ServerStream, shutdown: CancellationToken) {
185 let (r, w) = crate::ipc::split_server(stream);
186 let mut reader = BufReader::new(r);
187 let (tx, writer) = spawn_writer(w);
188 let mut line = String::new();
189
190 let hello = match tokio::time::timeout(HELLO_TIMEOUT, read_frame(&mut reader, &mut line)).await {
192 Ok(Ok(Read::Frame(ClientFrame::Hello(h)))) => h,
193 Ok(Ok(Read::Frame(_))) | Ok(Ok(Read::Bad(_))) => {
194 let _ = tx
195 .send(DaemonFrame::Refused { reason: RefuseReason::Protocol, message: "first frame must be hello".into() })
196 .await;
197 let _ = writer.await;
198 return;
199 }
200 Ok(Ok(Read::Oversize)) => {
201 let _ = tx.send(DaemonFrame::Error { session_id: None, message: frame_limit_msg() }).await;
202 drop(tx);
203 let _ = writer.await;
204 return;
205 }
206 _ => {
207 drop(tx);
208 let _ = writer.await;
209 return;
210 }
211 };
212 if hello.protocol_version < PROTOCOL_MIN || hello.protocol_version > PROTOCOL_MAX {
213 tracing::warn!(client = hello.protocol_version, daemon = PROTOCOL_VERSION, "daemon: refusing protocol version");
214 let _ = tx
215 .send(DaemonFrame::Refused {
216 reason: RefuseReason::Version { daemon_version: PROTOCOL_VERSION, min: PROTOCOL_MIN, max: PROTOCOL_MAX },
217 message: format!(
218 "protocol version {} not supported (daemon speaks {}; daemon binary {})",
219 hello.protocol_version,
220 PROTOCOL_VERSION,
221 binary_version()
222 ),
223 })
224 .await;
225 let _ = writer.await;
226 return;
227 }
228 if let Some(Err(message)) = hello.client.name.as_deref().map(crate::session::validate_client_name) {
230 let _ = tx.send(DaemonFrame::Refused { reason: RefuseReason::Protocol, message }).await;
231 let _ = writer.await;
232 return;
233 }
234 let daemon_version = binary_version();
235 if hello.client_version != daemon_version {
236 tracing::info!(client = %hello.client_version, daemon = %daemon_version, "daemon: client binary version differs (same protocol; allowed)");
237 }
238 let welcome = Welcome {
239 protocol_version: PROTOCOL_VERSION,
240 daemon_version,
241 pid: std::process::id(),
242 profile: state.profile.clone(),
243 sessions: state.session_metas(),
244 progressive_tool_disclosure: state.host.parts().progressive_tool_disclosure,
245 generation: state.generation,
246 daemon_build: crate::session::wire::build_id().to_string(),
247 };
248 if tx.send(DaemonFrame::Welcome(welcome)).await.is_err() {
249 return;
250 }
251 tracing::debug!(kind = ?hello.client.kind, "daemon: client connected");
252
253 let attach = loop {
255 let announce = state.announce();
256 let frame = tokio::select! {
257 _ = shutdown.cancelled() => { let _ = tx.send(DaemonFrame::Bye { reason: None }).await; let _ = writer.await; return; }
258 _ = announce.cancelled() => {
259 let _ = tx.send(DaemonFrame::Bye { reason: Some(reload_bye(&state)) }).await;
260 let _ = writer.await;
261 return;
262 }
263 r = read_frame(&mut reader, &mut line) => r,
264 };
265 match frame {
266 Ok(Read::Frame(ClientFrame::Ping)) => {
267 state.heal_daemon_json();
269 let _ = tx
270 .send(DaemonFrame::Pong { pid: std::process::id(), uptime_s: state.uptime_s(), sessions: state.live_sessions().len() })
271 .await;
272 }
273 Ok(Read::Frame(ClientFrame::Sessions)) => {
274 let _ = tx.send(DaemonFrame::SessionList { sessions: state.session_metas() }).await;
275 }
276 Ok(Read::Frame(ClientFrame::Shutdown { force })) => {
277 tracing::info!(force, "daemon: shutdown requested over the socket");
278 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
279 let _ = writer.await;
280 state.request_shutdown(force);
281 return;
282 }
283 Ok(Read::Frame(ClientFrame::Attach(a))) => {
284 if matches!(a, Attach::Create { .. }) && state.reloading.load(Ordering::SeqCst) {
285 let _ = tx
286 .send(DaemonFrame::Refused { reason: RefuseReason::Busy, message: "daemon reloading; retry in a moment".into() })
287 .await;
288 continue;
289 }
290 break a;
291 }
292 Ok(Read::Frame(ClientFrame::Cmd { session_id, .. })) => {
293 let _ = tx.send(DaemonFrame::Error { session_id: Some(session_id), message: "not attached".into() }).await;
294 }
295 Ok(Read::Frame(ref f @ ClientFrame::Reload { .. })) => {
298 let req = super::reload::ReloadRequest::from_frame(f).expect("reload frame");
299 match super::reload::prepare(&state, &state.paths, req).await {
300 Err(super::reload::ReloadError::Refused(why)) => {
301 let _ = tx
302 .send(DaemonFrame::Refused { reason: RefuseReason::ReloadRefused { why: why.clone() }, message: why })
303 .await;
304 }
305 Err(super::reload::ReloadError::ExecFailed(e)) => {
306 let _ = tx.send(DaemonFrame::Error { session_id: None, message: format!("reload: {e}") }).await;
307 }
308 Ok(prepared) => {
309 let _ = tx.send(DaemonFrame::Bye { reason: Some(reload_bye(&state)) }).await;
311 drop(tx);
312 let _ = tokio::time::timeout(
313 Duration::from_secs(crate::session::budgets::RELOAD_BYE_FLUSH_SECS),
314 writer,
315 )
316 .await;
317 let e = super::reload::exec(&state, prepared).await;
318 tracing::error!(error = %e, "daemon: reload did not exec");
319 return;
320 }
321 }
322 }
323 Ok(Read::Frame(ClientFrame::Close { session_id, force, attached })) => {
326 tracing::info!(session = %session_id, force, attached, "daemon: close requested over the socket");
327 let reply = super::close::close_session(&state, &session_id, force, attached).await;
328 let _ = tx.send(reply).await;
329 }
330 Ok(Read::Frame(ClientFrame::Pin { session_id, on })) => {
333 tracing::info!(session = %session_id, on, "daemon: pin requested over the socket");
334 let reply = super::pin::pin_session(&state, &session_id, on).await;
335 let _ = tx.send(reply).await;
336 }
337 Ok(Read::Frame(ClientFrame::Rename { session_id, name, expect })) => {
340 tracing::info!(session = %session_id, clear = name.is_none(), "daemon: rename requested over the socket");
341 let reply = super::rename::rename_session(&state, &session_id, name.as_deref(), expect).await;
342 let _ = tx.send(reply).await;
343 }
344 Ok(Read::Frame(ClientFrame::ReloadExtension { extension })) => {
347 tracing::info!(extension = %extension, "daemon: extension reload requested over the socket");
348 let reply = super::extension_reload::reload_extension(&state, &extension).await;
349 let _ = tx.send(reply).await;
350 }
351 Ok(Read::Frame(ClientFrame::Purge)) => {
353 agent_core::core::memstat::purge_arenas();
354 let _ = tx
355 .send(DaemonFrame::Pong { pid: std::process::id(), uptime_s: state.uptime_s(), sessions: state.live_sessions().len() })
356 .await;
357 }
358 Ok(Read::Frame(ClientFrame::AuthState(query))) => {
360 let _ = tx.send(auth_state_reply(&state, &query)).await;
361 }
362 Ok(Read::Frame(ClientFrame::Hello(_))) => {
363 let _ = tx.send(DaemonFrame::Error { session_id: None, message: "duplicate hello".into() }).await;
364 }
365 Ok(Read::Frame(ClientFrame::Bye)) | Ok(Read::Eof) | Err(_) => {
366 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
367 let _ = writer.await;
368 return;
369 }
370 Ok(Read::Oversize) => {
371 let _ = tx.send(DaemonFrame::Error { session_id: None, message: frame_limit_msg() }).await;
372 drop(tx);
373 let _ = writer.await;
374 return;
375 }
376 Ok(Read::Bad(e)) => {
377 let _ = tx.send(DaemonFrame::Error { session_id: None, message: format!("malformed frame: {e}") }).await;
378 }
379 }
380 };
381
382 let mut attached_to_live = false;
384 let (handle, mode) = match attach {
385 Attach::Existing { session_id, mode } => match state.attach(&session_id) {
386 Some(h) if h.meta().locked_by.is_some() => (state.retry_locked_placeholder(h).await, mode),
387 Some(h) => (h, mode),
388 None => {
389 let _ = tx
390 .send(DaemonFrame::Error { session_id: Some(session_id), message: "unknown session".into() })
391 .await;
392 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
393 let _ = writer.await;
394 return;
395 }
396 },
397 Attach::Create { mut config, mode } => {
398 config.memory_mode = None;
402 if config.cwd.is_none() {
403 config.cwd = Some(hello.cwd.clone());
404 }
405 if config.env.is_none() {
406 config.env = hello.env.clone();
407 }
408 if config.env_stripped.is_empty() {
409 config.env_stripped = hello.env_stripped.clone();
410 }
411 if let Some(cwd) = &config.cwd {
412 if !cwd.is_absolute() || !cwd.is_dir() {
413 let _ = tx
414 .send(DaemonFrame::Error {
415 session_id: None,
416 message: format!("cwd must be an absolute existing directory: {}", cwd.display()),
417 })
418 .await;
419 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
420 let _ = writer.await;
421 return;
422 }
423 }
424 let before: Vec<SessionId> = if config.continue_session.is_some() {
429 state.live_sessions().into_iter().map(|h| h.id).collect()
430 } else {
431 Vec::new()
432 };
433 match state.create(config).await {
434 Ok(h) if before.contains(&h.id) => {
435 tracing::info!(session = %h.id, "daemon: continue target is live — attaching, not creating");
436 attached_to_live = true;
437 (h, mode)
438 }
439 Ok(h) => (h, mode),
440 Err(e) => {
441 let _ = tx.send(DaemonFrame::Error { session_id: None, message: format!("create session: {e}") }).await;
442 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
443 let _ = writer.await;
444 return;
445 }
446 }
447 }
448 };
449
450 let mut rx = handle.subscribe();
451 if let Err(e) = handle.send(SessionCommand::Attach { client: hello.client.clone(), mode }).await {
452 let _ = tx.send(DaemonFrame::Error { session_id: Some(handle.id.clone()), message: format!("attach: {e}") }).await;
453 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
454 let _ = writer.await;
455 return;
456 }
457 let attach_budget = if handle.lifecycle() == SessionLifecycle::Parked { ATTACH_TIMEOUT_PARKED } else { ATTACH_TIMEOUT };
458 let attached = tokio::time::timeout(attach_budget, async {
459 loop {
460 match rx.recv().await {
461 Ok(env) => match env.event {
462 SessionEventWire::Attached { client, snapshot } => return Some(Ok((client, snapshot))),
463 SessionEventWire::AttachRefused { message } => return Some(Err(message)),
464 _ => {}
465 },
466 Err(broadcast::error::RecvError::Lagged(_)) => continue,
467 Err(broadcast::error::RecvError::Closed) => return None,
468 }
469 }
470 })
471 .await;
472 let (client, snapshot) = match attached {
473 Ok(Some(Ok(x))) => x,
474 Ok(Some(Err(message))) => {
475 let _ = tx.send(DaemonFrame::Error { session_id: Some(handle.id.clone()), message }).await;
476 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
477 let _ = writer.await;
478 return;
479 }
480 _ => {
481 let _ = tx.send(DaemonFrame::Error { session_id: Some(handle.id.clone()), message: "attach timed out".into() }).await;
482 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
483 let _ = writer.await;
484 return;
485 }
486 };
487 let legacy = legacy_shared(hello.client.capabilities, snapshot.client_input);
488 let snapshot = replay_for(hello.client.capabilities, snapshot);
489 let snapshot = if legacy { as_owner(client, snapshot) } else { snapshot };
490 if tx.send(DaemonFrame::Attached(AttachedWire::new(client, snapshot))).await.is_err() {
491 let _ = handle.send(SessionCommand::Detach { client }).await;
492 return;
493 }
494 tracing::debug!(session = %handle.id, client = client.0, "daemon: client attached");
495 if attached_to_live {
496 let env = Envelope {
497 session_id: handle.id.clone(),
498 seq: u64::MAX,
499 ts: chrono::Utc::now(),
500 event: SessionEventWire::SystemNotice(format!(
501 "session {} is already running in this daemon — attached to it instead of creating a second copy",
502 handle.id
503 )),
504 };
505 let _ = tx.send(DaemonFrame::Event(env.into())).await;
506 }
507
508 let sid = handle.id.clone();
510 let fwd_tx = tx.clone();
511 let fwd_handle = handle.clone();
512 let fwd_announce = state.announce();
513 let history = hello.client.history;
516 let caps = hello.client.capabilities;
517 let skip = move |env: &Envelope| -> bool {
518 !forwarded(history, caps, env) || (legacy && matches!(env.event, SessionEventWire::InputOwnerChanged { .. }))
519 };
520 let mut forward = tokio::spawn(async move {
522 loop {
523 let next = tokio::select! {
524 biased;
525 r = rx.recv() => r,
526 _ = fwd_announce.cancelled() => {
530 while let Ok(env) = rx.try_recv() {
531 if skip(&env) {
532 continue;
533 }
534 if fwd_tx.send(DaemonFrame::Event(env.into())).await.is_err() {
535 break;
536 }
537 }
538 break;
539 }
540 };
541 match next {
542 Ok(env) => {
543 if skip(&env) {
545 continue;
546 }
547 let ended = matches!(env.event, SessionEventWire::Ended { .. });
548 if fwd_tx.send(DaemonFrame::Event(env.into())).await.is_err() {
549 break;
550 }
551 if ended {
552 break;
553 }
554 }
555 Err(broadcast::error::RecvError::Lagged(n)) => {
556 tracing::warn!(session = %fwd_handle.id, dropped = n, "daemon: client lagged");
557 let env = Envelope {
558 session_id: fwd_handle.id.clone(),
559 seq: u64::MAX,
560 ts: chrono::Utc::now(),
561 event: SessionEventWire::SystemNotice(format!("event stream lagged; {n} dropped")),
562 };
563 if fwd_tx.send(DaemonFrame::Event(env.into())).await.is_err() {
564 break;
565 }
566 }
567 Err(broadcast::error::RecvError::Closed) => break,
568 }
569 }
570 });
571
572 let mut ended = false;
573 let mut stopping = false;
574 let mut reloading = false;
575 let announce = state.announce();
576 loop {
577 let frame = tokio::select! {
578 _ = shutdown.cancelled() => { stopping = true; break; }
579 _ = announce.cancelled() => { reloading = true; break; }
580 _ = handle.closed() => { ended = true; break; }
581 r = read_frame(&mut reader, &mut line) => r,
582 };
583 match frame {
584 Ok(Read::Frame(ClientFrame::Cmd { session_id, cmd })) => {
585 if session_id != sid {
586 let _ = tx.send(DaemonFrame::Error { session_id: Some(session_id), message: "not attached to that session".into() }).await;
587 continue;
588 }
589 if let Err(why) = client_may_send(&cmd, client) {
590 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: format!("refused: {why}") }).await;
591 continue;
592 }
593 if state.reloading.load(Ordering::SeqCst) && refused_while_draining(&cmd) {
594 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: crate::session::RELOADING_REFUSAL.into() }).await;
595 continue;
596 }
597 if matches!(&cmd, SessionCommand::Detach { client: c } if *c == client) {
602 forward.abort();
603 }
604 match handle.send_from(client, cmd).await {
605 Ok(()) => {}
606 Err(TransportError::Backpressure) => {
607 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: "backpressure: command queue full".into() }).await;
608 }
609 Err(TransportError::Closed) => {
610 ended = true;
611 break;
612 }
613 Err(e) => {
614 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: e.to_string() }).await;
615 }
616 }
617 }
618 Ok(Read::Frame(ClientFrame::Ping)) => {
619 state.heal_daemon_json();
621 let _ = tx
622 .send(DaemonFrame::Pong { pid: std::process::id(), uptime_s: state.uptime_s(), sessions: state.live_sessions().len() })
623 .await;
624 }
625 Ok(Read::Frame(ClientFrame::Sessions)) => {
626 let _ = tx.send(DaemonFrame::SessionList { sessions: state.session_metas() }).await;
627 }
628 Ok(Read::Frame(ClientFrame::Shutdown { force })) => {
629 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
630 state.request_shutdown(force);
631 break;
632 }
633 Ok(Read::Frame(ClientFrame::Attach(_))) | Ok(Read::Frame(ClientFrame::Hello(_))) => {
634 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: "already attached".into() }).await;
635 }
636 Ok(Read::Frame(ClientFrame::Reload { .. }))
639 | Ok(Read::Frame(ClientFrame::Purge))
640 | Ok(Read::Frame(ClientFrame::Close { .. }))
641 | Ok(Read::Frame(ClientFrame::Pin { .. }))
642 | Ok(Read::Frame(ClientFrame::Rename { .. }))
643 | Ok(Read::Frame(ClientFrame::ReloadExtension { .. }))
644 | Ok(Read::Frame(ClientFrame::AuthState(_))) => {
645 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: "control frames are not accepted on an attached connection".into() }).await;
646 }
647 Ok(Read::Frame(ClientFrame::Bye)) | Ok(Read::Eof) | Err(_) => break,
648 Ok(Read::Oversize) => {
649 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: frame_limit_msg() }).await;
650 break;
651 }
652 Ok(Read::Bad(e)) => {
653 let _ = tx.send(DaemonFrame::Error { session_id: Some(sid.clone()), message: format!("malformed frame: {e}") }).await;
654 }
655 }
656 }
657
658 if reloading {
661 if tokio::time::timeout(Duration::from_secs(1), &mut forward).await.is_err() {
663 forward.abort();
664 let _ = forward.await;
665 }
666 let (generation, retry_after_ms) = match reload_bye(&state) {
667 ByeReason::Reloading { generation, retry_after_ms } => (generation, retry_after_ms),
668 _ => unreachable!(),
669 };
670 let env = crate::session::Envelope {
671 session_id: sid.clone(),
672 seq: u64::MAX,
673 ts: chrono::Utc::now(),
674 event: SessionEventWire::Reloading { generation, retry_after_ms },
675 };
676 let _ = tx.send(DaemonFrame::Event(env.into())).await;
677 let _ = tx.send(DaemonFrame::Bye { reason: Some(ByeReason::Reloading { generation, retry_after_ms }) }).await;
678 let _ = tokio::time::timeout(Duration::from_secs(1), writer).await;
679 return;
680 }
681
682 if stopping {
684 match tokio::time::timeout(super::lifecycle::SESSION_END_BUDGET + Duration::from_secs(1), &mut forward).await {
686 Ok(_) => {}
687 Err(_) => {
688 forward.abort();
689 let _ = forward.await;
690 }
691 }
692 } else if ended {
693 if tokio::time::timeout(Duration::from_secs(1), &mut forward).await.is_err() {
698 forward.abort();
699 let _ = forward.await;
700 }
701 } else {
702 forward.abort();
703 let _ = forward.await;
704 }
705 if !ended && handle.is_alive() {
706 let _ = handle.send(SessionCommand::Detach { client }).await;
707 }
708 if !handle.is_alive() {
709 state.remove(&sid);
710 }
711 let _ = tx.send(DaemonFrame::Bye { reason: None }).await;
712 drop(tx);
713 let _ = tokio::time::timeout(Duration::from_millis(500), writer).await;
714 tracing::debug!(session = %sid, client = client.0, "daemon: client detached");
715}
716
717fn forwarded(history: HistoryMode, caps: ClientCapabilities, env: &Envelope) -> bool {
723 !(matches!(env.event, SessionEventWire::Attached { .. })
724 || (history == HistoryMode::Digest
725 && matches!(
726 env.event,
727 SessionEventWire::Stream(StreamEvent::Session(SessionEvent::MessageHistory(_)))
728 ))
729 || !client_decodes(caps, &env.event))
730}
731
732fn replay_for(caps: ClientCapabilities, mut snapshot: AttachSnapshot) -> AttachSnapshot {
736 snapshot.replay.retain(|env| client_decodes(caps, &env.event));
737 snapshot
738}
739
740fn as_owner(client: ClientId, mut snapshot: AttachSnapshot) -> AttachSnapshot {
744 let observes = snapshot.roster.iter().any(|c| c.id == client && c.mode == AttachMode::Observe);
745 if !observes {
746 snapshot.input_owner = Some(client);
747 }
748 snapshot.replay.retain(|env| !matches!(env.event, SessionEventWire::InputOwnerChanged { .. }));
749 snapshot
750}
751
752fn reload_bye(state: &DaemonState) -> ByeReason {
753 ByeReason::Reloading {
754 generation: state.reload_generation.load(Ordering::SeqCst),
755 retry_after_ms: super::reload::RETRY_AFTER_MS,
756 }
757}
758
759fn auth_state_reply(state: &DaemonState, query: &agent_core::auth::AuthStateQuery) -> DaemonFrame {
764 match agent_core::auth::local_auth_state(&state.host.parts().token_cache, query, true) {
765 Ok(report) => DaemonFrame::AuthState { report },
766 Err(e) => DaemonFrame::Error { session_id: None, message: format!("auth state: {e}") },
767 }
768}
769
770#[cfg(test)]
771mod tests {
772 use super::*;
773 use crate::session::CheckpointReason;
774
775 #[test]
779 fn a_reload_checkpoint_is_host_only() {
780 let own = ClientId(1);
781 let reload = SessionCommand::Checkpoint { reason: CheckpointReason::Reload };
782 assert!(client_may_send(&reload, own).is_err());
783 let host = SessionCommand::Checkpoint { reason: CheckpointReason::HostRequest };
784 assert!(client_may_send(&host, own).is_ok());
785 }
786
787 #[test]
790 fn take_input_is_a_client_command() {
791 assert!(client_may_send(&SessionCommand::TakeInput, ClientId(2)).is_ok());
792 }
793
794 fn stream_env(ev: StreamEvent) -> Envelope {
795 Envelope {
796 session_id: SessionId::from("s"),
797 seq: 1,
798 ts: chrono::Utc::now(),
799 event: SessionEventWire::Stream(ev),
800 }
801 }
802
803 #[test]
807 fn queue_signals_are_only_forwarded_to_clients_that_opted_in() {
808 use crate::LlmEvent;
809 let queued = stream_env(StreamEvent::Llm(LlmEvent::ToolQueued { tool_id: "t2".into() }));
810 let started = stream_env(StreamEvent::Llm(LlmEvent::ToolStarted { tool_id: "t2".into() }));
811 let result = stream_env(StreamEvent::Llm(LlmEvent::ToolResult {
812 tool_id: "t1".into(),
813 result: "ok".into(),
814 }));
815 let old = ClientCapabilities::default();
816 let new = ClientCapabilities::current();
817 for history in [HistoryMode::Full, HistoryMode::Digest] {
818 assert!(!forwarded(history, old, &queued));
819 assert!(!forwarded(history, old, &started));
820 assert!(forwarded(history, old, &result));
821 assert!(forwarded(history, new, &queued));
822 assert!(forwarded(history, new, &started));
823 assert!(forwarded(history, new, &result));
824 }
825
826 let session = crate::session::handle::echo::meta_for(&SessionId::from("s"));
827 let snapshot = AttachSnapshot {
828 meta: session,
829 view: crate::session::handle::echo::view_for(),
830 conversation: Default::default(),
831 streaming: true,
832 replay: vec![queued.clone(), result.clone(), started.clone()],
833 pending_prompts: vec![],
834 clients: vec![],
835 roster: Vec::new(),
836 client_input: ClientInput::wire_default(),
837 input_owner: None,
838 display_tail: None,
839 };
840 let kinds = |s: &AttachSnapshot| -> Vec<String> {
841 s.replay
842 .iter()
843 .map(|e| match &e.event {
844 SessionEventWire::Stream(StreamEvent::Llm(LlmEvent::ToolQueued { .. })) => "queued",
845 SessionEventWire::Stream(StreamEvent::Llm(LlmEvent::ToolStarted { .. })) => "started",
846 _ => "other",
847 })
848 .map(str::to_string)
849 .collect()
850 };
851 assert_eq!(kinds(&replay_for(old, snapshot.clone())), vec!["other"]);
852 assert_eq!(kinds(&replay_for(new, snapshot)), vec!["queued", "other", "started"]);
853 }
854
855 #[cfg(unix)]
860 #[tokio::test]
861 async fn a_large_frame_releases_its_buffer_after_the_read() {
862 let (ours, theirs) = tokio::net::UnixStream::pair().unwrap();
863 let (read_half, _keep) = ours.into_split();
864 let mut reader = BufReader::new(read_half);
865 let mut line = String::new();
866
867 let pad = " ".repeat(2 * 1024 * 1024);
870 let input = format!("{pad}{{\"type\":\"ping\"}}\n{pad}not json\n{{\"type\":\"ping\"}}\n");
871 let writer = tokio::spawn(async move {
872 let mut theirs = theirs;
873 theirs.write_all(input.as_bytes()).await.unwrap();
874 theirs
875 });
876
877 let first = read_frame(&mut reader, &mut line).await.unwrap();
878 assert!(matches!(first, Read::Frame(ClientFrame::Ping)));
879 assert!(line.capacity() <= FRAME_BUF_SHRINK_ABOVE, "kept {} bytes", line.capacity());
880
881 let second = read_frame(&mut reader, &mut line).await.unwrap();
882 assert!(matches!(second, Read::Bad(_)), "a malformed big frame is reported");
883 assert!(line.capacity() <= FRAME_BUF_SHRINK_ABOVE, "kept {} bytes", line.capacity());
884
885 let third = read_frame(&mut reader, &mut line).await.unwrap();
886 assert!(matches!(third, Read::Frame(ClientFrame::Ping)));
887 drop(writer.await.unwrap());
888 assert!(matches!(read_frame(&mut reader, &mut line).await.unwrap(), Read::Eof));
889 }
890
891 #[test]
893 fn a_small_frame_buffer_is_reused() {
894 let mut line = String::with_capacity(FRAME_BUF_KEEP);
895 line.push_str("{\"type\":\"ping\"}");
896 let before = line.as_ptr();
897 release_large_frame_buf(&mut line);
898 assert_eq!(line.as_ptr(), before, "a small buffer is not reallocated");
899
900 let mut big = String::with_capacity(FRAME_BUF_SHRINK_ABOVE + 1);
901 release_large_frame_buf(&mut big);
902 assert!((FRAME_BUF_KEEP..=FRAME_BUF_SHRINK_ABOVE).contains(&big.capacity()));
903 }
904
905 #[test]
908 fn the_drain_gate_keeps_steering_the_running_turn() {
909 let submit = SessionCommand::Submit { text: "x".into(), attachments: vec![] };
910 assert!(refused_while_draining(&submit));
911 assert!(refused_while_draining(&SessionCommand::NewSession));
912 assert!(refused_while_draining(&SessionCommand::Resume { id: 1, query: "q".into() }));
913 assert!(!refused_while_draining(&SessionCommand::Steer { text: "x".into() }));
914 assert!(!refused_while_draining(&SessionCommand::Cancel));
915 }
916}