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(assignments: Vec<Assignment>) -> Vec<DiscoveredAssignment> {
106 if !herdr_environment()
107 || !assignments
108 .iter()
109 .any(|assignment| omp_session_file(assignment).is_some())
110 {
111 return base_records(assignments);
112 }
113
114 match load_snapshot() {
115 Ok(snapshot) => join_snapshot(assignments, snapshot, Utc::now()),
116 Err(error) => {
117 eprintln!("agent-id: unable to enrich discover from Herdr: {error:#}");
118 base_records(assignments)
119 }
120 }
121}
122
123pub fn base_records(assignments: Vec<Assignment>) -> Vec<DiscoveredAssignment> {
124 assignments
125 .into_iter()
126 .map(|assignment| DiscoveredAssignment {
127 assignment,
128 runtime: None,
129 })
130 .collect()
131}
132
133fn herdr_environment() -> bool {
134 env::var_os("HERDR_ENV").as_deref() == Some(std::ffi::OsStr::new("1"))
135 && env::var_os("HERDR_SOCKET_PATH").is_some_and(|value| !value.is_empty())
136}
137
138fn herdr_binary() -> OsString {
139 env::var_os("HERDR_BIN_PATH")
140 .filter(|value| !value.is_empty())
141 .unwrap_or_else(|| OsString::from("herdr"))
142}
143
144fn load_snapshot() -> Result<Snapshot> {
145 let output = Command::new(herdr_binary())
146 .args(["api", "snapshot"])
147 .output()
148 .context("run `herdr api snapshot`")?;
149 if !output.status.success() {
150 let detail = String::from_utf8_lossy(&output.stderr);
151 bail!("`herdr api snapshot` failed: {}", detail.trim());
152 }
153 let response: ApiResponse =
154 serde_json::from_slice(&output.stdout).context("parse Herdr session snapshot")?;
155 Ok(response.result.snapshot)
156}
157
158fn join_snapshot(
159 assignments: Vec<Assignment>,
160 snapshot: Snapshot,
161 observed_at: DateTime<Utc>,
162) -> Vec<DiscoveredAssignment> {
163 let mut by_session_id = HashMap::new();
164 let mut by_session_file = HashMap::new();
165 for (index, assignment) in assignments.iter().enumerate() {
166 by_session_id.insert(assignment.session_id.as_str(), index);
167 let Some(session_file) = omp_session_file(assignment) else {
168 continue;
169 };
170 by_session_file
171 .entry(session_file)
172 .and_modify(|index: &mut Option<usize>| *index = None)
173 .or_insert(Some(index));
174 }
175
176 let tabs: HashMap<_, _> = snapshot
177 .tabs
178 .into_iter()
179 .map(|tab| (tab.tab_id, tab.label))
180 .collect();
181 let workspaces: HashMap<_, _> = snapshot
182 .workspaces
183 .into_iter()
184 .map(|workspace| {
185 (
186 workspace.workspace_id,
187 (workspace.label, workspace.worktree),
188 )
189 })
190 .collect();
191 let mut locations = vec![Vec::new(); assignments.len()];
192
193 for agent in snapshot.agents {
194 let Some(session) = agent.agent_session.as_ref() else {
195 continue;
196 };
197 if agent.agent.as_deref() != Some("omp")
198 || session.agent != "omp"
199 || session.source != "herdr:omp"
200 {
201 continue;
202 }
203 let index = match session.kind.as_str() {
204 "id" => by_session_id.get(session.value.as_str()).copied(),
205 "path" => by_session_file
206 .get(session.value.as_str())
207 .copied()
208 .flatten(),
209 _ => None,
210 };
211 let Some(index) = index else {
212 continue;
213 };
214 let (workspace_label, worktree) = workspaces
215 .get(&agent.workspace_id)
216 .map(|(label, worktree)| (Some(label.clone()), worktree.clone()))
217 .unwrap_or((None, None));
218 locations[index].push(HerdrLocation {
219 agent_status: agent.agent_status,
220 pane_id: agent.pane_id,
221 tab_label: tabs.get(&agent.tab_id).cloned(),
222 tab_id: agent.tab_id,
223 workspace_label,
224 workspace_id: agent.workspace_id,
225 cwd: agent.cwd,
226 foreground_cwd: agent.foreground_cwd,
227 worktree,
228 });
229 }
230
231 assignments
232 .into_iter()
233 .zip(locations)
234 .map(|(mut assignment, mut locations)| {
235 locations.sort_by(|left, right| {
236 (&left.workspace_id, &left.tab_id, &left.pane_id).cmp(&(
237 &right.workspace_id,
238 &right.tab_id,
239 &right.pane_id,
240 ))
241 });
242 let runtime = (!locations.is_empty()).then(|| {
243 let state = locations[0].agent_status.clone();
244 assignment.state = ActivityState::from_external(&state, observed_at);
245 RuntimeProjection {
246 provider: "herdr",
247 state,
248 observed_at,
249 locations,
250 }
251 });
252 DiscoveredAssignment {
253 assignment,
254 runtime,
255 }
256 })
257 .collect()
258}
259
260fn omp_session_file(assignment: &Assignment) -> Option<&str> {
261 assignment
262 .extensions
263 .get(OMP_EXTENSION_OWNER)?
264 .data
265 .get("session_file")?
266 .as_str()
267}
268
269#[cfg(test)]
270mod tests {
271 use std::collections::BTreeMap;
272
273 use chrono::TimeZone;
274 use serde_json::json;
275
276 use super::*;
277 use crate::registry::ExtensionMetadata;
278
279 fn assignment(session_id: &str, session_file: Option<&str>) -> Assignment {
280 let now = Utc.timestamp_opt(0, 0).single().unwrap();
281 let mut extensions = BTreeMap::new();
282 if let Some(session_file) = session_file {
283 extensions.insert(
284 "omp".to_string(),
285 ExtensionMetadata {
286 data: json!({ "session_file": session_file }),
287 updated_at: now,
288 },
289 );
290 }
291 Assignment {
292 version: 1,
293 session_id: session_id.to_string(),
294 name: format!("{session_id} Agent of Test"),
295 slug: format!("{session_id}-agent-test"),
296 first_name: session_id.to_string(),
297 family_name: "Agent".to_string(),
298 realm: "Test".to_string(),
299 summary: None,
300 state: ActivityState::unknown(now),
301 cwd: None,
302 extensions,
303 created_at: now,
304 updated_at: now,
305 }
306 }
307
308 fn agent(kind: &str, value: &str, pane_id: &str) -> AgentInfo {
309 AgentInfo {
310 agent: Some("omp".to_string()),
311 agent_session: Some(AgentSession {
312 agent: "omp".to_string(),
313 kind: kind.to_string(),
314 source: "herdr:omp".to_string(),
315 value: value.to_string(),
316 }),
317 agent_status: "working".to_string(),
318 cwd: Some("/work".to_string()),
319 foreground_cwd: Some("/work".to_string()),
320 pane_id: pane_id.to_string(),
321 tab_id: "w1:t1".to_string(),
322 workspace_id: "w1".to_string(),
323 }
324 }
325
326 #[test]
327 fn joins_path_and_id_references_without_guessing() {
328 let now = Utc.timestamp_opt(1, 0).single().unwrap();
329 let assignments = vec![
330 assignment("path-session", Some("/tmp/path-session.jsonl")),
331 assignment("id-session", None),
332 ];
333 let snapshot = Snapshot {
334 agents: vec![
335 agent("path", "/tmp/path-session.jsonl", "w1:p1"),
336 agent("id", "id-session", "w1:p2"),
337 agent("path", "/tmp/unmatched.jsonl", "w1:p3"),
338 ],
339 tabs: vec![TabInfo {
340 tab_id: "w1:t1".to_string(),
341 label: "agents".to_string(),
342 }],
343 workspaces: vec![WorkspaceInfo {
344 workspace_id: "w1".to_string(),
345 label: "project".to_string(),
346 worktree: None,
347 }],
348 };
349
350 let records = join_snapshot(assignments, snapshot, now);
351
352 assert_eq!(
353 records[0].runtime.as_ref().unwrap().locations[0].pane_id,
354 "w1:p1"
355 );
356 assert_eq!(records[0].runtime.as_ref().unwrap().state, "working");
357 assert_eq!(records[0].assignment.state.value.to_string(), "working");
358 assert_eq!(records[0].assignment.state.updated_at, now);
359 assert_eq!(
360 records[1].runtime.as_ref().unwrap().locations[0].pane_id,
361 "w1:p2"
362 );
363 assert_eq!(records[0].runtime.as_ref().unwrap().observed_at, now);
364 }
365}