1use std::{collections::HashMap, env, ffi::OsString, process::Command};
2
3use anyhow::{bail, Context, Result};
4use chrono::{DateTime, Utc};
5use serde::{Deserialize, Serialize};
6
7use crate::registry::Assignment;
8
9const OMP_EXTENSION_OWNER: &str = "omp";
10
11#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
12pub struct DiscoveredAssignment {
13 #[serde(flatten)]
14 pub assignment: Assignment,
15 #[serde(skip_serializing_if = "Option::is_none")]
16 pub runtime: Option<RuntimeProjection>,
17}
18
19#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
20pub struct RuntimeProjection {
21 pub provider: &'static str,
22 pub observed_at: DateTime<Utc>,
23 pub locations: Vec<HerdrLocation>,
24}
25
26#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
27pub struct HerdrLocation {
28 pub agent_status: String,
29 pub pane_id: String,
30 pub tab_id: String,
31 pub workspace_id: String,
32 #[serde(skip_serializing_if = "Option::is_none")]
33 pub tab_label: Option<String>,
34 #[serde(skip_serializing_if = "Option::is_none")]
35 pub workspace_label: Option<String>,
36 #[serde(skip_serializing_if = "Option::is_none")]
37 pub cwd: Option<String>,
38 #[serde(skip_serializing_if = "Option::is_none")]
39 pub foreground_cwd: Option<String>,
40 #[serde(skip_serializing_if = "Option::is_none")]
41 pub worktree: Option<HerdrWorktree>,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
45pub struct HerdrWorktree {
46 pub repo_key: String,
47 pub repo_name: String,
48 pub repo_root: String,
49 pub checkout_path: String,
50 pub is_linked_worktree: bool,
51}
52
53#[derive(Debug, Deserialize)]
54struct ApiResponse {
55 result: ApiResult,
56}
57
58#[derive(Debug, Deserialize)]
59struct ApiResult {
60 snapshot: Snapshot,
61}
62
63#[derive(Debug, Deserialize)]
64struct Snapshot {
65 agents: Vec<AgentInfo>,
66 tabs: Vec<TabInfo>,
67 workspaces: Vec<WorkspaceInfo>,
68}
69
70#[derive(Debug, Deserialize)]
71struct AgentInfo {
72 agent: Option<String>,
73 agent_session: Option<AgentSession>,
74 agent_status: String,
75 cwd: Option<String>,
76 foreground_cwd: Option<String>,
77 pane_id: String,
78 tab_id: String,
79 workspace_id: String,
80}
81
82#[derive(Debug, Deserialize)]
83struct AgentSession {
84 agent: String,
85 kind: String,
86 source: String,
87 value: String,
88}
89
90#[derive(Debug, Deserialize)]
91struct TabInfo {
92 tab_id: String,
93 label: String,
94}
95
96#[derive(Debug, Deserialize)]
97struct WorkspaceInfo {
98 workspace_id: String,
99 label: String,
100 worktree: Option<HerdrWorktree>,
101}
102
103pub fn augment_discovery(assignments: Vec<Assignment>) -> Vec<DiscoveredAssignment> {
104 if !herdr_environment()
105 || !assignments
106 .iter()
107 .any(|assignment| omp_session_file(assignment).is_some())
108 {
109 return base_records(assignments);
110 }
111
112 match load_snapshot() {
113 Ok(snapshot) => join_snapshot(assignments, snapshot, Utc::now()),
114 Err(error) => {
115 eprintln!("agent-id: unable to enrich discover from Herdr: {error:#}");
116 base_records(assignments)
117 }
118 }
119}
120
121pub fn base_records(assignments: Vec<Assignment>) -> Vec<DiscoveredAssignment> {
122 assignments
123 .into_iter()
124 .map(|assignment| DiscoveredAssignment {
125 assignment,
126 runtime: None,
127 })
128 .collect()
129}
130
131fn herdr_environment() -> bool {
132 env::var_os("HERDR_ENV").as_deref() == Some(std::ffi::OsStr::new("1"))
133 && env::var_os("HERDR_SOCKET_PATH").is_some_and(|value| !value.is_empty())
134}
135
136fn herdr_binary() -> OsString {
137 env::var_os("HERDR_BIN_PATH")
138 .filter(|value| !value.is_empty())
139 .unwrap_or_else(|| OsString::from("herdr"))
140}
141
142fn load_snapshot() -> Result<Snapshot> {
143 let output = Command::new(herdr_binary())
144 .args(["api", "snapshot"])
145 .output()
146 .context("run `herdr api snapshot`")?;
147 if !output.status.success() {
148 let detail = String::from_utf8_lossy(&output.stderr);
149 bail!("`herdr api snapshot` failed: {}", detail.trim());
150 }
151 let response: ApiResponse =
152 serde_json::from_slice(&output.stdout).context("parse Herdr session snapshot")?;
153 Ok(response.result.snapshot)
154}
155
156fn join_snapshot(
157 assignments: Vec<Assignment>,
158 snapshot: Snapshot,
159 observed_at: DateTime<Utc>,
160) -> Vec<DiscoveredAssignment> {
161 let mut by_session_id = HashMap::new();
162 let mut by_session_file = HashMap::new();
163 for (index, assignment) in assignments.iter().enumerate() {
164 by_session_id.insert(assignment.session_id.as_str(), index);
165 let Some(session_file) = omp_session_file(assignment) else {
166 continue;
167 };
168 by_session_file
169 .entry(session_file)
170 .and_modify(|index: &mut Option<usize>| *index = None)
171 .or_insert(Some(index));
172 }
173
174 let tabs: HashMap<_, _> = snapshot
175 .tabs
176 .into_iter()
177 .map(|tab| (tab.tab_id, tab.label))
178 .collect();
179 let workspaces: HashMap<_, _> = snapshot
180 .workspaces
181 .into_iter()
182 .map(|workspace| {
183 (
184 workspace.workspace_id,
185 (workspace.label, workspace.worktree),
186 )
187 })
188 .collect();
189 let mut locations = vec![Vec::new(); assignments.len()];
190
191 for agent in snapshot.agents {
192 let Some(session) = agent.agent_session.as_ref() else {
193 continue;
194 };
195 if agent.agent.as_deref() != Some("omp")
196 || session.agent != "omp"
197 || session.source != "herdr:omp"
198 {
199 continue;
200 }
201 let index = match session.kind.as_str() {
202 "id" => by_session_id.get(session.value.as_str()).copied(),
203 "path" => by_session_file
204 .get(session.value.as_str())
205 .copied()
206 .flatten(),
207 _ => None,
208 };
209 let Some(index) = index else {
210 continue;
211 };
212 let (workspace_label, worktree) = workspaces
213 .get(&agent.workspace_id)
214 .map(|(label, worktree)| (Some(label.clone()), worktree.clone()))
215 .unwrap_or((None, None));
216 locations[index].push(HerdrLocation {
217 agent_status: agent.agent_status,
218 pane_id: agent.pane_id,
219 tab_label: tabs.get(&agent.tab_id).cloned(),
220 tab_id: agent.tab_id,
221 workspace_label,
222 workspace_id: agent.workspace_id,
223 cwd: agent.cwd,
224 foreground_cwd: agent.foreground_cwd,
225 worktree,
226 });
227 }
228
229 assignments
230 .into_iter()
231 .zip(locations)
232 .map(|(assignment, mut locations)| {
233 locations.sort_by(|left, right| {
234 (&left.workspace_id, &left.tab_id, &left.pane_id).cmp(&(
235 &right.workspace_id,
236 &right.tab_id,
237 &right.pane_id,
238 ))
239 });
240 let runtime = (!locations.is_empty()).then_some(RuntimeProjection {
241 provider: "herdr",
242 observed_at,
243 locations,
244 });
245 DiscoveredAssignment {
246 assignment,
247 runtime,
248 }
249 })
250 .collect()
251}
252
253fn omp_session_file(assignment: &Assignment) -> Option<&str> {
254 assignment
255 .extensions
256 .get(OMP_EXTENSION_OWNER)?
257 .data
258 .get("session_file")?
259 .as_str()
260}
261
262#[cfg(test)]
263mod tests {
264 use std::collections::BTreeMap;
265
266 use chrono::TimeZone;
267 use serde_json::json;
268
269 use super::*;
270 use crate::registry::ExtensionMetadata;
271
272 fn assignment(session_id: &str, session_file: Option<&str>) -> Assignment {
273 let now = Utc.timestamp_opt(0, 0).single().unwrap();
274 let mut extensions = BTreeMap::new();
275 if let Some(session_file) = session_file {
276 extensions.insert(
277 "omp".to_string(),
278 ExtensionMetadata {
279 data: json!({ "session_file": session_file }),
280 updated_at: now,
281 },
282 );
283 }
284 Assignment {
285 version: 1,
286 session_id: session_id.to_string(),
287 name: format!("{session_id} Agent of Test"),
288 slug: format!("{session_id}-agent-test"),
289 first_name: session_id.to_string(),
290 family_name: "Agent".to_string(),
291 realm: "Test".to_string(),
292 summary: None,
293 state: None,
294 cwd: None,
295 extensions,
296 created_at: now,
297 updated_at: now,
298 }
299 }
300
301 fn agent(kind: &str, value: &str, pane_id: &str) -> AgentInfo {
302 AgentInfo {
303 agent: Some("omp".to_string()),
304 agent_session: Some(AgentSession {
305 agent: "omp".to_string(),
306 kind: kind.to_string(),
307 source: "herdr:omp".to_string(),
308 value: value.to_string(),
309 }),
310 agent_status: "working".to_string(),
311 cwd: Some("/work".to_string()),
312 foreground_cwd: Some("/work".to_string()),
313 pane_id: pane_id.to_string(),
314 tab_id: "w1:t1".to_string(),
315 workspace_id: "w1".to_string(),
316 }
317 }
318
319 #[test]
320 fn joins_path_and_id_references_without_guessing() {
321 let now = Utc.timestamp_opt(1, 0).single().unwrap();
322 let assignments = vec![
323 assignment("path-session", Some("/tmp/path-session.jsonl")),
324 assignment("id-session", None),
325 ];
326 let snapshot = Snapshot {
327 agents: vec![
328 agent("path", "/tmp/path-session.jsonl", "w1:p1"),
329 agent("id", "id-session", "w1:p2"),
330 agent("path", "/tmp/unmatched.jsonl", "w1:p3"),
331 ],
332 tabs: vec![TabInfo {
333 tab_id: "w1:t1".to_string(),
334 label: "agents".to_string(),
335 }],
336 workspaces: vec![WorkspaceInfo {
337 workspace_id: "w1".to_string(),
338 label: "project".to_string(),
339 worktree: None,
340 }],
341 };
342
343 let records = join_snapshot(assignments, snapshot, now);
344
345 assert_eq!(
346 records[0].runtime.as_ref().unwrap().locations[0].pane_id,
347 "w1:p1"
348 );
349 assert_eq!(
350 records[1].runtime.as_ref().unwrap().locations[0].pane_id,
351 "w1:p2"
352 );
353 assert_eq!(records[0].runtime.as_ref().unwrap().observed_at, now);
354 }
355}