zag_agent/process_registration.rs
1//! Helper for registering an agent process in zag's `ProcessStore` and
2//! producing the env vars (`ZAG_PROCESS_ID`, `ZAG_SESSION_ID`, etc.) that
3//! `zag ps kill self` / `zig self terminate` use to resolve the running
4//! agent from inside the agent's own subshell.
5//!
6//! Two callers need this exact sequence:
7//!
8//! - `zag-cli/src/commands/agent_action.rs` — the `zag agent` CLI path.
9//! - `zag-agent::AgentBuilder` — used by library consumers (e.g. `zig run`
10//! interactive steps in the `zig` workflow tool) that drive an agent
11//! programmatically without going through the `zag` CLI.
12//!
13//! Before this module existed, only `agent_action.rs` did the registration
14//! (and via `unsafe std::env::set_var`, polluting the parent process's env).
15//! Library callers got no registration and no env vars, so `zig self
16//! terminate` failed with `Cannot resolve "self": ZAG_PROCESS_ID is not set`
17//! from inside any zig-run interactive step.
18//!
19//! This module is now the single source of truth. Callers either:
20//!
21//! - Call [`register`] directly and wire the returned [`ProcessRegistration`]
22//! to their `Agent` / `AgentBuilder` via `set_env_vars` + `set_on_spawn_hook`
23//! (or [`AgentBuilder::env`] / [`AgentBuilder::on_spawn`] / the convenience
24//! [`AgentBuilder::register_process`]), then call
25//! [`ProcessRegistration::update_status`] when the agent exits.
26//! - Or pass a [`RegisterOptionsOwned`] to [`AgentBuilder::register_process`]
27//! and let the builder handle registration + finalisation automatically.
28
29use crate::agent::OnSpawnHook;
30use crate::process_store::{ProcessEntry, ProcessStore};
31use std::sync::Arc;
32
33/// Borrowed options for [`register`]. Field semantics mirror
34/// [`ProcessEntry`] one-for-one.
35pub struct RegisterOptions<'a> {
36 /// Provider name (e.g. `"claude"`, `"codex"`).
37 pub provider: &'a str,
38 /// Effective model string (e.g. `"sonnet"`).
39 pub model: &'a str,
40 /// Subcommand label stored in the entry: `"run"`, `"exec"`, `"plan"`, etc.
41 /// Used by `zag ps` for display.
42 pub command: &'a str,
43 /// First ~100 chars of the prompt, if any.
44 pub prompt_preview: Option<&'a str>,
45 /// Session id this process is associated with (zag's internal session_id,
46 /// not the provider-native one). Becomes `ZAG_SESSION_ID` for the child.
47 pub session_id: Option<&'a str>,
48 /// Optional human-friendly session name. Becomes `ZAG_SESSION_NAME`.
49 pub session_name: Option<&'a str>,
50 /// Project root passed through to the entry and `ZAG_ROOT`.
51 pub root: Option<&'a str>,
52}
53
54/// Owned variant of [`RegisterOptions`] for storage on a builder before the
55/// terminal method runs.
56#[derive(Debug, Clone, Default)]
57pub struct RegisterOptionsOwned {
58 pub provider: String,
59 pub model: String,
60 pub command: String,
61 pub prompt_preview: Option<String>,
62 pub session_id: Option<String>,
63 pub session_name: Option<String>,
64 pub root: Option<String>,
65}
66
67impl RegisterOptionsOwned {
68 pub(crate) fn as_borrowed(&self) -> RegisterOptions<'_> {
69 RegisterOptions {
70 provider: &self.provider,
71 model: &self.model,
72 command: &self.command,
73 prompt_preview: self.prompt_preview.as_deref(),
74 session_id: self.session_id.as_deref(),
75 session_name: self.session_name.as_deref(),
76 root: self.root.as_deref(),
77 }
78 }
79}
80
81/// A live process registration. Holds the generated `proc_id` and the env
82/// vars to inject into the agent subprocess. Callers wire `env_vars()` and
83/// `on_spawn_hook()` to the agent and call `update_status()` once the agent
84/// exits.
85pub struct ProcessRegistration {
86 proc_id: String,
87 env_vars: Vec<(String, String)>,
88}
89
90impl ProcessRegistration {
91 /// The UUID assigned to this registration. Matches the `id` field of the
92 /// `ProcessEntry` in zag's process store and the `ZAG_PROCESS_ID` env
93 /// var injected into the child.
94 pub fn proc_id(&self) -> &str {
95 &self.proc_id
96 }
97
98 /// The env vars to inject into the agent subprocess so it can resolve
99 /// `"self"` for `zag ps kill self` / `zig self terminate`. Use with
100 /// [`crate::agent::Agent::set_env_vars`] (CLI path) or repeated
101 /// [`crate::builder::AgentBuilder::env`] calls (library path).
102 pub fn env_vars(&self) -> &[(String, String)] {
103 &self.env_vars
104 }
105
106 /// `on_spawn` hook that retargets the registry entry's `pid` from zag's
107 /// own pid (registered up-front so `zag ps` is populated immediately) to
108 /// the actual agent subprocess pid. Wire via
109 /// [`crate::agent::Agent::set_on_spawn_hook`] or
110 /// [`crate::builder::AgentBuilder::on_spawn`].
111 ///
112 /// Without this, `zag ps kill self` would SIGTERM the parent zag/zig
113 /// orchestrator instead of the agent child, taking the workflow down
114 /// with it.
115 pub fn on_spawn_hook(&self) -> OnSpawnHook {
116 let proc_id = self.proc_id.clone();
117 Arc::new(move |pid: u32| {
118 if let Ok(mut pstore) = ProcessStore::load() {
119 pstore.update_pid(&proc_id, pid);
120 let _ = pstore.save();
121 }
122 })
123 }
124
125 /// Mark the registration's entry as finished. Pass `"exited"` for clean
126 /// completion or `"killed"` when the agent crashed / was signalled.
127 pub fn update_status(&self, status: &str, exit_code: Option<i32>) {
128 if let Ok(mut pstore) = ProcessStore::load() {
129 pstore.update_status(&self.proc_id, status, exit_code);
130 let _ = pstore.save();
131 }
132 }
133}
134
135/// Build the env-var list for a registration. Pure: no I/O, no env reads.
136/// Exposed so unit tests can assert the var set without touching the real
137/// `~/.zag/processes.json` or mutating process-global env state.
138pub(crate) fn build_env_vars(proc_id: &str, opts: &RegisterOptions<'_>) -> Vec<(String, String)> {
139 let mut env_vars = Vec::with_capacity(6);
140 if let Some(sid) = opts.session_id {
141 env_vars.push(("ZAG_SESSION_ID".to_string(), sid.to_string()));
142 }
143 env_vars.push(("ZAG_PROCESS_ID".to_string(), proc_id.to_string()));
144 env_vars.push(("ZAG_PROVIDER".to_string(), opts.provider.to_string()));
145 env_vars.push(("ZAG_MODEL".to_string(), opts.model.to_string()));
146 if let Some(r) = opts.root {
147 env_vars.push(("ZAG_ROOT".to_string(), r.to_string()));
148 }
149 if let Some(name) = opts.session_name {
150 env_vars.push(("ZAG_SESSION_NAME".to_string(), name.to_string()));
151 }
152 env_vars
153}
154
155/// Build a `ProcessEntry` for a fresh registration. Pure: takes parent
156/// linkage explicitly so tests don't need to mutate `ZAG_PROCESS_ID` /
157/// `ZAG_SESSION_ID` to exercise it.
158pub(crate) fn build_entry(
159 proc_id: &str,
160 opts: &RegisterOptions<'_>,
161 parent_process_id: Option<String>,
162 parent_session_id: Option<String>,
163) -> ProcessEntry {
164 ProcessEntry {
165 id: proc_id.to_string(),
166 pid: std::process::id(),
167 session_id: opts.session_id.map(String::from),
168 provider: opts.provider.to_string(),
169 model: opts.model.to_string(),
170 command: opts.command.to_string(),
171 prompt: opts.prompt_preview.map(String::from),
172 started_at: chrono::Utc::now().to_rfc3339(),
173 status: "running".to_string(),
174 exit_code: None,
175 exited_at: None,
176 root: opts.root.map(String::from),
177 parent_process_id,
178 parent_session_id,
179 }
180}
181
182/// Register a new process entry in zag's `ProcessStore` and return a
183/// [`ProcessRegistration`] holding the proc_id and the env vars to inject
184/// into the agent subprocess.
185///
186/// Reads `ZAG_PROCESS_ID` / `ZAG_SESSION_ID` from the current process env to
187/// record parent-process linkage for nested invocations.
188///
189/// The entry is registered with `pid = std::process::id()` (the caller's
190/// pid). Wire [`ProcessRegistration::on_spawn_hook`] to the agent so the pid
191/// is retargeted to the agent subprocess once it spawns.
192pub fn register(opts: RegisterOptions<'_>) -> ProcessRegistration {
193 let proc_id = uuid::Uuid::new_v4().to_string();
194 let parent_process_id = std::env::var("ZAG_PROCESS_ID").ok();
195 let parent_session_id = std::env::var("ZAG_SESSION_ID").ok();
196
197 let entry = build_entry(&proc_id, &opts, parent_process_id, parent_session_id);
198 let env_vars = build_env_vars(&proc_id, &opts);
199
200 if let Ok(mut pstore) = ProcessStore::load() {
201 pstore.add(entry);
202 let _ = pstore.save();
203 }
204
205 ProcessRegistration { proc_id, env_vars }
206}
207
208#[cfg(test)]
209#[path = "process_registration_tests.rs"]
210mod tests;