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