1#[cfg(unix)]
2use std::os::unix::fs::{OpenOptionsExt as _, PermissionsExt};
3use std::{
4 fs,
5 io::{ErrorKind, Write},
6 path::{Path, PathBuf},
7};
8
9use atomic_write_file::AtomicWriteFile;
10#[cfg(unix)]
11use atomic_write_file::unix::OpenOptionsExt as _;
12use serde::{Deserialize, Serialize};
13
14use super::yaml::state_dir_path;
15use crate::{
16 adapter::process_identity::LocalProcessIdentity,
17 constants::MUSTER_AGENT_SESSION_STATE_FILE_ENV,
18 domain::{
19 agent_session::{
20 AgentProcessId, AgentProcessStartToken, AgentSession, AgentSessionId,
21 AgentSessionState, NativeSessionId,
22 },
23 config::ConfigError,
24 port::AgentSessionStore,
25 process::AgentTool,
26 },
27};
28
29const AGENT_SESSIONS_FILE: &str = "agent-sessions.yml";
31const SESSION_FILE_VERSION: u8 = 1;
33const MAX_SESSION_STATE_SYMLINKS: usize = 40;
35#[cfg(unix)]
37const PRIVATE_SESSION_FILE_MODE: u32 = 0o600;
38#[cfg(unix)]
40const FILE_PERMISSION_MASK: u32 = 0o777;
41#[cfg(unix)]
43const OWNER_PERMISSION_MASK: u32 = 0o700;
44
45#[derive(Serialize, Deserialize)]
47struct SessionFile {
48 version: u8,
49 sessions: Vec<AgentSession>,
50}
51
52impl SessionFile {
53 fn empty() -> Self {
55 Self {
56 version: SESSION_FILE_VERSION,
57 sessions: Vec::new(),
58 }
59 }
60}
61
62#[derive(Default)]
64pub struct YamlAgentSessionStore;
65
66impl YamlAgentSessionStore {
67 fn path() -> Option<PathBuf> {
69 std::env::var_os(MUSTER_AGENT_SESSION_STATE_FILE_ENV)
70 .map(PathBuf::from)
71 .filter(|path| path.is_absolute())
72 .or_else(|| state_dir_path(AGENT_SESSIONS_FILE))
73 }
74
75 fn load_from(path: &Path) -> Result<SessionFile, ConfigError> {
80 if !path.exists() {
81 return Ok(SessionFile::empty());
82 }
83 let raw = fs::read_to_string(path).map_err(|source| ConfigError::Read {
84 path: path.to_path_buf(),
85 source,
86 })?;
87 let file: SessionFile = serde_yaml_ng::from_str(&raw)?;
88 if file.version != SESSION_FILE_VERSION {
89 return Err(ConfigError::UnsupportedAgentSessionVersion(file.version));
90 }
91 Ok(file)
92 }
93
94 fn update(
99 path: &Path,
100 mutate: impl FnOnce(&mut Vec<AgentSession>) -> Result<(), ConfigError>,
101 ) -> Result<(), ConfigError> {
102 let path = Self::write_destination(path)?;
103 let _guard = Self::lock(&path)?;
104 let mut file = Self::load_from(&path)?;
105 mutate(&mut file.sessions)?;
106 Self::write(&path, &file)
107 }
108
109 fn write_destination(path: &Path) -> Result<PathBuf, ConfigError> {
115 let mut destination = path.to_path_buf();
116 for depth in 0..=MAX_SESSION_STATE_SYMLINKS {
117 match fs::symlink_metadata(&destination) {
118 Ok(metadata) if metadata.file_type().is_symlink() => {
119 if depth == MAX_SESSION_STATE_SYMLINKS {
120 return Self::symlink_depth_error(destination);
121 }
122 let target =
123 fs::read_link(&destination).map_err(|source| ConfigError::Read {
124 path: destination.clone(),
125 source,
126 })?;
127 destination = if target.is_absolute() {
128 target
129 } else {
130 destination
131 .parent()
132 .map_or(target.clone(), |parent| parent.join(target))
133 };
134 },
135 Ok(_) => return Ok(destination),
136 Err(error) if error.kind() == ErrorKind::NotFound => return Ok(destination),
137 Err(source) => {
138 return Err(ConfigError::Read {
139 path: destination,
140 source,
141 });
142 },
143 }
144 }
145 Self::symlink_depth_error(destination)
146 }
147
148 fn symlink_depth_error(path: PathBuf) -> Result<PathBuf, ConfigError> {
150 Err(ConfigError::Read {
151 path,
152 source: std::io::Error::other("agent session-state symlink depth exceeded"),
153 })
154 }
155
156 fn write(path: &Path, value: &SessionFile) -> Result<(), ConfigError> {
163 if let Some(parent) = path.parent()
164 && !parent.as_os_str().is_empty()
165 {
166 fs::create_dir_all(parent).map_err(|source| ConfigError::Write {
167 path: parent.to_path_buf(),
168 source,
169 })?;
170 }
171 let raw = serde_yaml_ng::to_string(value)?;
172 let mut options = AtomicWriteFile::options();
173 #[cfg(unix)]
174 {
175 options.preserve_mode(false);
176 options.mode(PRIVATE_SESSION_FILE_MODE);
177 }
178 let mut file = options.open(path).map_err(|source| ConfigError::Write {
179 path: path.to_path_buf(),
180 source,
181 })?;
182 #[cfg(unix)]
183 file.set_permissions(fs::Permissions::from_mode(Self::secure_file_mode(path)?))
184 .map_err(|source| ConfigError::Write {
185 path: path.to_path_buf(),
186 source,
187 })?;
188 file.write_all(raw.as_bytes())
189 .map_err(|source| ConfigError::Write {
190 path: path.to_path_buf(),
191 source,
192 })?;
193 file.commit().map_err(|source| ConfigError::Write {
194 path: path.to_path_buf(),
195 source,
196 })
197 }
198
199 #[cfg(unix)]
206 fn secure_file_mode(path: &Path) -> Result<u32, ConfigError> {
207 match fs::metadata(path) {
208 Ok(metadata) => {
209 let mode = metadata.permissions().mode() & FILE_PERMISSION_MASK;
210 Ok(if mode & !OWNER_PERMISSION_MASK == 0 {
211 mode
212 } else {
213 PRIVATE_SESSION_FILE_MODE
214 })
215 },
216 Err(error) if error.kind() == ErrorKind::NotFound => Ok(PRIVATE_SESSION_FILE_MODE),
217 Err(source) => Err(ConfigError::Read {
218 path: path.to_path_buf(),
219 source,
220 }),
221 }
222 }
223
224 fn lock(path: &Path) -> Result<fs::File, ConfigError> {
230 let lock_path = Self::lock_path(path);
231 if let Some(parent) = lock_path.parent() {
232 fs::create_dir_all(parent).map_err(|source| ConfigError::Write {
233 path: parent.to_path_buf(),
234 source,
235 })?;
236 }
237 let file = fs::OpenOptions::new()
238 .create(true)
239 .truncate(false)
240 .read(true)
241 .write(true)
242 .open(&lock_path)
243 .map_err(|source| ConfigError::Write {
244 path: lock_path.clone(),
245 source,
246 })?;
247 file.lock().map_err(|source| ConfigError::Write {
248 path: lock_path,
249 source,
250 })?;
251 Ok(file)
252 }
253
254 fn lock_path(path: &Path) -> PathBuf {
256 let mut name = path.file_name().unwrap_or_default().to_os_string();
257 name.push(".lock");
258 path.with_file_name(name)
259 }
260
261 fn replace(sessions: &mut Vec<AgentSession>, session: AgentSession) {
263 sessions.retain(|candidate| candidate.id() != session.id());
264 sessions.push(session);
265 }
266
267 fn claim_owner(
274 session: &mut AgentSession,
275 id: &AgentSessionId,
276 process_id: AgentProcessId,
277 process_start_token: Option<AgentProcessStartToken>,
278 wrapper_process_id: Option<AgentProcessId>,
279 ) -> Result<(), ConfigError> {
280 if let (Some(owner), Some(token)) = (
281 session.owner_process_id(),
282 session.owner_process_start_token(),
283 ) && LocalProcessIdentity::start_token(*owner) == Some(*token)
284 && *owner != process_id
285 {
286 return Err(ConfigError::AgentSessionAlreadyOwned {
287 id: id.clone(),
288 owner: *owner,
289 });
290 }
291 *session = session.clone().with_launch_processes(
292 process_id,
293 process_start_token,
294 wrapper_process_id,
295 );
296 Ok(())
297 }
298
299 fn assign_native_id(
306 sessions: &mut [AgentSession],
307 id: &AgentSessionId,
308 provider: AgentTool,
309 process_id: AgentProcessId,
310 parent_process_id: Option<AgentProcessId>,
311 native_id: NativeSessionId,
312 ) -> Result<(), ConfigError> {
313 Self::assign_native_id_with_start_token(
314 sessions,
315 id,
316 provider,
317 process_id,
318 parent_process_id,
319 LocalProcessIdentity::start_token(process_id),
320 native_id,
321 )
322 }
323
324 fn assign_native_id_with_start_token(
331 sessions: &mut [AgentSession],
332 id: &AgentSessionId,
333 provider: AgentTool,
334 process_id: AgentProcessId,
335 parent_process_id: Option<AgentProcessId>,
336 process_start_token: Option<AgentProcessStartToken>,
337 native_id: NativeSessionId,
338 ) -> Result<(), ConfigError> {
339 let session = sessions
340 .iter_mut()
341 .find(|session| session.id() == id)
342 .ok_or_else(|| ConfigError::AgentSessionNotFound(id.clone()))?;
343 if *session.tool() != provider {
344 return Err(ConfigError::AgentSessionProviderMismatch {
345 id: id.clone(),
346 expected: *session.tool(),
347 reported: provider,
348 });
349 }
350 let owns_process = session.owner_process_id().as_ref() == Some(&process_id)
351 && session
352 .owner_process_start_token()
353 .as_ref()
354 .is_none_or(|expected| process_start_token.as_ref() == Some(expected));
355 let is_wrapper_handoff = session.wrapper_process_id().is_some()
356 && session.wrapper_process_id().as_ref() == parent_process_id.as_ref();
357 if !owns_process && !is_wrapper_handoff {
358 return Err(ConfigError::AgentSessionProcessMismatch {
359 id: id.clone(),
360 expected: *session.owner_process_id(),
361 reported: process_id,
362 });
363 }
364 *session = session.clone().with_native_id(native_id);
365 Ok(())
366 }
367}
368
369impl AgentSessionStore for YamlAgentSessionStore {
370 fn sessions(&self) -> Result<Vec<AgentSession>, ConfigError> {
371 let path = Self::path().ok_or(ConfigError::NoConfigDir)?;
372 let path = Self::write_destination(&path)?;
373 Ok(Self::load_from(&path)?.sessions)
374 }
375
376 fn state_file_path(&self) -> Result<Option<PathBuf>, ConfigError> {
377 Ok(Self::path())
378 }
379
380 fn upsert(&self, session: &AgentSession) -> Result<(), ConfigError> {
381 let path = Self::path().ok_or(ConfigError::NoConfigDir)?;
382 Self::update(&path, |sessions| {
383 Self::replace(sessions, session.clone());
384 Ok(())
385 })
386 }
387
388 fn set_state(&self, id: &AgentSessionId, state: AgentSessionState) -> Result<(), ConfigError> {
389 let path = Self::path().ok_or(ConfigError::NoConfigDir)?;
390 Self::update(&path, |sessions| {
391 let index = sessions
392 .iter()
393 .position(|session| session.id() == id)
394 .ok_or_else(|| ConfigError::AgentSessionNotFound(id.clone()))?;
395 let session = sessions.remove(index).with_state(state);
396 sessions.push(session);
397 Ok(())
398 })
399 }
400
401 fn set_owner_process_id(
402 &self,
403 id: &AgentSessionId,
404 process_id: AgentProcessId,
405 process_start_token: Option<AgentProcessStartToken>,
406 wrapper_process_id: Option<AgentProcessId>,
407 ) -> Result<(), ConfigError> {
408 let path = Self::path().ok_or(ConfigError::NoConfigDir)?;
409 Self::update(&path, |sessions| {
410 let session = sessions
411 .iter_mut()
412 .find(|session| session.id() == id)
413 .ok_or_else(|| ConfigError::AgentSessionNotFound(id.clone()))?;
414 Self::claim_owner(
415 session,
416 id,
417 process_id,
418 process_start_token,
419 wrapper_process_id,
420 )
421 })
422 }
423
424 fn capture_native_id(
425 &self,
426 id: &AgentSessionId,
427 provider: AgentTool,
428 process_id: AgentProcessId,
429 parent_process_id: Option<AgentProcessId>,
430 native_id: NativeSessionId,
431 ) -> Result<(), ConfigError> {
432 let path = Self::path().ok_or(ConfigError::NoConfigDir)?;
433 Self::update(&path, |sessions| {
434 Self::assign_native_id(
435 sessions,
436 id,
437 provider,
438 process_id,
439 parent_process_id,
440 native_id,
441 )
442 })
443 }
444}
445
446#[cfg(test)]
447mod tests {
448 use super::*;
449 use crate::domain::{
450 process::AgentTool,
451 value::{CommandLine, ProcessName},
452 };
453
454 fn session() -> AgentSession {
456 AgentSession::builder()
457 .id(AgentSessionId::generate().unwrap())
458 .name(ProcessName::try_new("Ada").unwrap())
459 .tool(AgentTool::Codex)
460 .project(PathBuf::from("/repo/muster.yml"))
461 .launch_command(CommandLine::try_new("codex").unwrap())
462 .state(AgentSessionState::Open)
463 .build()
464 }
465
466 #[test]
468 fn updates_session_state_and_native_identity() {
469 let dir =
470 std::env::temp_dir().join(format!("muster-agent-sessions-{}", uuid::Uuid::new_v4()));
471 let path = dir.join(AGENT_SESSIONS_FILE);
472 let original = session().with_owner_process_id(AgentProcessId::try_new(1).unwrap());
473
474 YamlAgentSessionStore::update(&path, |sessions| {
475 YamlAgentSessionStore::replace(sessions, original.clone());
476 Ok(())
477 })
478 .unwrap();
479 YamlAgentSessionStore::update(&path, |sessions| {
480 let item = sessions.first_mut().unwrap();
481 *item = item
482 .clone()
483 .with_native_id(NativeSessionId::try_new("native").unwrap())
484 .with_state(AgentSessionState::Closed);
485 Ok(())
486 })
487 .unwrap();
488
489 let loaded = YamlAgentSessionStore::load_from(&path).unwrap();
490 assert_eq!(loaded.sessions[0].id(), original.id());
491 assert_eq!(
492 loaded.sessions[0].native_id().as_ref().map(AsRef::as_ref),
493 Some("native")
494 );
495 assert_eq!(*loaded.sessions[0].state(), AgentSessionState::Closed);
496 fs::remove_dir_all(dir).unwrap();
497 }
498
499 #[cfg(unix)]
502 #[test]
503 fn writes_session_state_through_a_symlink() {
504 use std::os::unix::fs::symlink;
505
506 let dir =
507 std::env::temp_dir().join(format!("muster-agent-sessions-{}", uuid::Uuid::new_v4()));
508 fs::create_dir_all(&dir).unwrap();
509 let target = dir.join("shared.yml");
510 let link = dir.join(AGENT_SESSIONS_FILE);
511 symlink(&target, &link).unwrap();
512 let record = session();
513
514 YamlAgentSessionStore::update(&link, |sessions| {
515 YamlAgentSessionStore::replace(sessions, record.clone());
516 Ok(())
517 })
518 .unwrap();
519
520 assert!(
521 fs::symlink_metadata(&link)
522 .unwrap()
523 .file_type()
524 .is_symlink()
525 );
526 assert_eq!(
527 YamlAgentSessionStore::load_from(&target)
528 .unwrap()
529 .sessions
530 .len(),
531 1
532 );
533 fs::remove_dir_all(dir).unwrap();
534 }
535
536 #[test]
539 fn updates_identity_only_for_the_owning_provider() {
540 let original = session();
541 let id = original.id().clone();
542 let first = NativeSessionId::try_new("first-native").unwrap();
543 let second = NativeSessionId::try_new("second-native").unwrap();
544 let descendant = NativeSessionId::try_new("descendant-native").unwrap();
545 let mut sessions = vec![original];
546
547 let owner = AgentProcessId::try_new(1).unwrap();
548 sessions[0] = sessions[0].clone().with_owner_process_id(owner);
549 YamlAgentSessionStore::assign_native_id(
550 &mut sessions,
551 &id,
552 AgentTool::Codex,
553 owner,
554 None,
555 first,
556 )
557 .unwrap();
558 YamlAgentSessionStore::assign_native_id(
559 &mut sessions,
560 &id,
561 AgentTool::Codex,
562 owner,
563 None,
564 second.clone(),
565 )
566 .unwrap();
567 let result = YamlAgentSessionStore::assign_native_id(
568 &mut sessions,
569 &id,
570 AgentTool::Claude,
571 owner,
572 None,
573 descendant,
574 );
575
576 assert!(matches!(
577 result,
578 Err(ConfigError::AgentSessionProviderMismatch {
579 id: conflict,
580 expected: AgentTool::Codex,
581 reported: AgentTool::Claude,
582 }) if conflict == id
583 ));
584 assert_eq!(sessions[0].native_id().as_ref(), Some(&second));
585 }
586
587 #[test]
589 fn rejects_same_provider_identity_from_another_process() {
590 let original = session().with_owner_process_id(AgentProcessId::try_new(1).unwrap());
591 let id = original.id().clone();
592 let mut sessions = vec![original];
593 let owner = AgentProcessId::try_new(1).unwrap();
594 let descendant = AgentProcessId::try_new(2).unwrap();
595 let first = NativeSessionId::try_new("parent-native").unwrap();
596 let child = NativeSessionId::try_new("child-native").unwrap();
597
598 YamlAgentSessionStore::assign_native_id(
599 &mut sessions,
600 &id,
601 AgentTool::Codex,
602 owner,
603 None,
604 first.clone(),
605 )
606 .unwrap();
607 let result = YamlAgentSessionStore::assign_native_id(
608 &mut sessions,
609 &id,
610 AgentTool::Codex,
611 descendant,
612 None,
613 child,
614 );
615
616 assert!(matches!(
617 result,
618 Err(ConfigError::AgentSessionProcessMismatch {
619 id: conflict,
620 expected: Some(expected),
621 reported,
622 }) if conflict == id && expected == owner && reported == descendant
623 ));
624 assert_eq!(sessions[0].native_id().as_ref(), Some(&first));
625 }
626
627 #[test]
630 fn rejects_a_reused_owner_pid_with_a_different_start_token() {
631 let owner = AgentProcessId::try_new(1).unwrap();
632 let stored_token = AgentProcessStartToken::try_new(10).unwrap();
633 let reused_token = AgentProcessStartToken::try_new(11).unwrap();
634 let original = session().with_launch_processes(owner, Some(stored_token), None);
635 let id = original.id().clone();
636 let mut sessions = vec![original];
637
638 let result = YamlAgentSessionStore::assign_native_id_with_start_token(
639 &mut sessions,
640 &id,
641 AgentTool::Codex,
642 owner,
643 None,
644 Some(reused_token),
645 NativeSessionId::try_new("reused-native").unwrap(),
646 );
647
648 assert!(matches!(
649 result,
650 Err(ConfigError::AgentSessionProcessMismatch {
651 id: conflict,
652 expected: Some(expected),
653 reported,
654 }) if conflict == id && expected == owner && reported == owner
655 ));
656 assert!(sessions[0].native_id().is_none());
657 }
658
659 #[test]
662 fn rejects_a_second_claim_while_the_existing_owner_is_live() {
663 let owner = AgentProcessId::try_new(std::process::id()).unwrap();
664 let token = LocalProcessIdentity::start_token(owner).unwrap();
665 let mut record = session().with_launch_processes(owner, Some(token), None);
666 let id = record.id().clone();
667 let claimant = AgentProcessId::try_new(owner.into_inner().saturating_add(1)).unwrap();
668
669 let result = YamlAgentSessionStore::claim_owner(&mut record, &id, claimant, None, None);
670
671 assert!(matches!(
672 result,
673 Err(ConfigError::AgentSessionAlreadyOwned {
674 id: conflict,
675 owner: live_owner,
676 }) if conflict == id && live_owner == owner
677 ));
678 assert_eq!(record.owner_process_id(), &Some(owner));
679 assert_eq!(record.owner_process_start_token(), &Some(token));
680 }
681
682 #[test]
685 fn accepts_identity_from_a_direct_provider_child() {
686 let owner = AgentProcessId::try_new(1).unwrap();
687 let provider = AgentProcessId::try_new(2).unwrap();
688 let original = session().with_launch_processes(owner, None, Some(owner));
689 let id = original.id().clone();
690 let native = NativeSessionId::try_new("pipeline-native").unwrap();
691 let mut sessions = vec![original];
692
693 YamlAgentSessionStore::assign_native_id(
694 &mut sessions,
695 &id,
696 AgentTool::Codex,
697 provider,
698 Some(owner),
699 native.clone(),
700 )
701 .unwrap();
702
703 assert_eq!(sessions[0].native_id().as_ref(), Some(&native));
704 assert_eq!(sessions[0].owner_process_id(), &Some(owner));
705
706 let switched = NativeSessionId::try_new("resumed-native").unwrap();
707 YamlAgentSessionStore::assign_native_id(
708 &mut sessions,
709 &id,
710 AgentTool::Codex,
711 provider,
712 Some(owner),
713 switched.clone(),
714 )
715 .unwrap();
716 assert_eq!(sessions[0].native_id().as_ref(), Some(&switched));
717
718 let nested = AgentProcessId::try_new(3).unwrap();
719 let result = YamlAgentSessionStore::assign_native_id(
720 &mut sessions,
721 &id,
722 AgentTool::Codex,
723 nested,
724 Some(provider),
725 NativeSessionId::try_new("nested-native").unwrap(),
726 );
727 assert!(matches!(
728 result,
729 Err(ConfigError::AgentSessionProcessMismatch { .. })
730 ));
731 assert_eq!(sessions[0].native_id().as_ref(), Some(&switched));
732 }
733
734 #[cfg(unix)]
737 #[test]
738 fn writes_session_state_with_owner_only_permissions() {
739 use std::os::unix::fs::PermissionsExt;
740
741 const OWNER_READ_ONLY_MODE: u32 = 0o400;
742 const LEGACY_SHARED_MODE: u32 = 0o644;
743 let dir = std::env::temp_dir().join(format!("muster-agent-mode-{}", uuid::Uuid::new_v4()));
744 let path = dir.join(AGENT_SESSIONS_FILE);
745 let original = session();
746 let write = || {
747 YamlAgentSessionStore::update(&path, |sessions| {
748 YamlAgentSessionStore::replace(sessions, original.clone());
749 Ok(())
750 })
751 .unwrap();
752 };
753
754 write();
755 assert_eq!(
756 fs::metadata(&path).unwrap().permissions().mode() & FILE_PERMISSION_MASK,
757 PRIVATE_SESSION_FILE_MODE
758 );
759
760 fs::set_permissions(&path, fs::Permissions::from_mode(OWNER_READ_ONLY_MODE)).unwrap();
761 write();
762 assert_eq!(
763 fs::metadata(&path).unwrap().permissions().mode() & FILE_PERMISSION_MASK,
764 OWNER_READ_ONLY_MODE
765 );
766
767 fs::set_permissions(&path, fs::Permissions::from_mode(LEGACY_SHARED_MODE)).unwrap();
768 write();
769 assert_eq!(
770 fs::metadata(&path).unwrap().permissions().mode() & FILE_PERMISSION_MASK,
771 PRIVATE_SESSION_FILE_MODE
772 );
773 fs::remove_dir_all(dir).unwrap();
774 }
775}