ninox_core/lifecycle/
poller.rs

1use crate::{
2    config::AppConfig,
3    events::{Engine, Event},
4    github::{split_repo, CheckRun},
5    hooks,
6    lifecycle::{
7        enrichment::EnrichmentCache,
8        probe::is_pid_alive,
9        usage,
10    },
11    types::{
12        CIStatus, Comment, Notification, NotificationKind, PrId, SessionStatus, PR,
13    },
14};
15use std::{collections::HashMap, sync::Arc, time::Duration};
16use tokio_util::sync::CancellationToken;
17
18/// Unix epoch milliseconds "now" — used to stamp `Notification::created_at`.
19fn now_millis() -> i64 {
20    std::time::SystemTime::now()
21        .duration_since(std::time::UNIX_EPOCH)
22        .map(|d| d.as_millis() as i64)
23        .unwrap_or(0)
24}
25
26pub struct Poller {
27    engine:           Arc<Engine>,
28    enrichment_cache: Arc<std::sync::Mutex<EnrichmentCache>>,
29}
30
31impl Poller {
32    pub fn new(engine: Arc<Engine>) -> Self {
33        Self {
34            engine,
35            enrichment_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
36        }
37    }
38
39    pub async fn start(self, token: CancellationToken) {
40        let mut pid_interval    = tokio::time::interval(Duration::from_secs(5));
41        let mut usage_interval  = tokio::time::interval(Duration::from_secs(10));
42        let mut github_interval = tokio::time::interval(Duration::from_secs(30));
43        // Prevent a missed tick from causing back-to-back polls.
44        github_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
45        usage_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
46        loop {
47            tokio::select! {
48                _ = token.cancelled()      => break,
49                _ = pid_interval.tick()    => self.poll_pids().await,
50                _ = usage_interval.tick()  => self.poll_usage().await,
51                _ = github_interval.tick() => self.poll_github().await,
52            }
53        }
54    }
55
56    // ── PID liveness ────────────────────────────────────────────────────────
57
58    async fn poll_pids(&self) {
59        let Ok(sessions) = self.engine.store.list_sessions() else { return };
60        for mut session in sessions {
61            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated) {
62                continue;
63            }
64            if let Some(pid) = session.pid {
65                if !is_pid_alive(pid) {
66                    session.status = SessionStatus::Terminated;
67                    let _ = self.engine.store.upsert_session(&session);
68                    self.engine.emit(Event::SessionUpdated(session));
69                    continue;
70                }
71            }
72
73            // Poll metadata files for PR number on working sessions that have none yet.
74            if matches!(session.status, SessionStatus::Working | SessionStatus::Spawning)
75                && session.pr_number.is_none()
76            {
77                let sessions_dir = AppConfig::sessions_dir();
78                if let Ok(meta) = hooks::read_session_metadata(&sessions_dir, &session.id) {
79                    if let Some(pr_num) = meta.pr_number {
80                        session.pr_number = Some(pr_num);
81                        session.status    = SessionStatus::PrOpen;
82                        let _ = self.engine.store.upsert_session(&session);
83                        self.engine.emit(Event::SessionUpdated(session.clone()));
84                        tracing::info!(
85                            "session {} PR #{pr_num} detected via metadata hook",
86                            session.id
87                        );
88                    }
89                }
90            }
91        }
92    }
93
94    // ── Cost / context-window usage ─────────────────────────────────────────
95
96    /// Ingest cost/token usage for every active session by reading `claude`'s
97    /// own transcript for the session's workspace directory (see
98    /// `lifecycle::usage`). Sessions without a workspace, or whose transcript
99    /// has no usage yet (agent hasn't taken a turn), are left untouched.
100    /// Only writes + emits when something actually changed, so this doesn't
101    /// spam the store/UI every tick for idle sessions.
102    async fn poll_usage(&self) {
103        let Ok(sessions) = self.engine.store.list_sessions() else { return };
104        for mut session in sessions {
105            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated) {
106                continue;
107            }
108            let Some(workspace) = session.workspace_path.clone() else { continue };
109            let Some(snapshot) = usage::ingest_usage_for_workspace(&workspace) else { continue };
110
111            let cost_changed = (session.cost_usd - snapshot.cost_usd).abs() > 1e-9;
112            let context_changed = session.context_tokens != Some(snapshot.context_tokens);
113            if !cost_changed && !context_changed {
114                continue;
115            }
116
117            session.cost_usd = snapshot.cost_usd;
118            session.context_tokens = Some(snapshot.context_tokens);
119            if session.model.is_none() {
120                session.model = snapshot.model;
121            }
122            let _ = self.engine.store.upsert_session(&session);
123            self.engine.emit(Event::SessionUpdated(session));
124        }
125    }
126
127    // ── GitHub enrichment ────────────────────────────────────────────────────
128
129    async fn poll_github(&self) {
130        let Some(gh) = &self.engine.github else { return };
131        let Ok(sessions) = self.engine.store.list_sessions() else { return };
132
133        for session in sessions {
134            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated) {
135                continue;
136            }
137            let Some(pr_number) = session.pr_number else { continue };
138            let Some((owner, repo)) = split_repo(&session.repo) else { continue };
139
140            // -- PR state --
141            let pr_status = match gh.get_pr_status(&owner, &repo, pr_number).await {
142                Ok(s)  => s,
143                Err(e) => { tracing::warn!("github pr status: {e}"); continue }
144            };
145
146            let pr_id: PrId = pr_number as i64;
147
148            // -- Merge detection — handle before CI (no point polling CI on merged PR) --
149            if pr_status.merged && !matches!(session.status, SessionStatus::Done) {
150                self.engine.emit(Event::Notification(Notification {
151                    id:         format!("merged-{}", session.id),
152                    kind:       NotificationKind::WorkerDone,
153                    title:      format!("PR merged — {}", session.name),
154                    body:       format!("#{} merged successfully", pr_number),
155                    session_id: Some(session.id.clone()),
156                    created_at: now_millis(),
157                }));
158                if let Err(e) = self.engine.cleanup_session(&session.id).await {
159                    tracing::warn!("cleanup_session {}: {e}", session.id);
160                }
161                // Remove enrichment state for this session — it's done
162                {
163                    let mut cache = self.enrichment_cache.lock().unwrap();
164                    cache.remove(&session.id);
165                }
166                continue; // skip further enrichment for this session
167            }
168
169            // Upsert PR record — only when not merged (merged sessions stay Done after cleanup)
170            {
171                let pr = PR {
172                    id:         pr_id,
173                    number:     pr_number,
174                    title:      pr_status.title.clone(),
175                    url:        format!("https://github.com/{owner}/{repo}/pull/{pr_number}"),
176                    body:       String::new(),
177                    session_id: session.id.clone(),
178                };
179                let _ = self.engine.store.upsert_pr(&pr);
180                self.engine.emit(Event::PrOpened { session_id: session.id.clone(), pr });
181            }
182
183            // -- CI checks --
184            let checks = match gh.get_ci_checks(&owner, &repo, &pr_status.head_sha).await {
185                Ok(c)  => c,
186                Err(e) => { tracing::warn!("github ci checks: {e}"); vec![] }
187            };
188            let ci = summarize_checks(pr_id, &checks);
189            let _ = self.engine.store.upsert_ci_status(&ci);
190            self.engine.emit(Event::CiUpdated { pr_id, status: ci.clone() });
191
192            // -- Detect CI transition and update session status --
193            let (newly_failing, ci_reaction_already_sent) = {
194                let mut cache = self.enrichment_cache.lock().unwrap();
195                let state = cache.entry(session.id.clone()).or_default();
196
197                let newly_failing = state.prev_failing.is_none_or(|p| p == 0)
198                    && ci.failing > 0;
199                state.prev_failing = Some(ci.failing);
200
201                let already_sent = state.ci_reaction_sent;
202                if newly_failing && !already_sent {
203                    state.ci_reaction_sent = true;
204                }
205                if ci.failing == 0 {
206                    state.ci_reaction_sent = false;
207                }
208                (newly_failing, already_sent)
209            };
210
211            if newly_failing && !ci_reaction_already_sent {
212                self.engine.emit(Event::Notification(Notification {
213                    id:         format!("ci-{}", session.id),
214                    kind:       NotificationKind::CiFailure,
215                    title:      format!("CI failing — {}", session.name),
216                    body:       format!("{}/{} checks failing", ci.failing, ci.total),
217                    session_id: Some(session.id.clone()),
218                    created_at: now_millis(),
219                }));
220                // Send reaction to the agent in the tmux session
221                let failing_names: Vec<String> = checks.iter()
222                    .filter(|c| c.conclusion.as_deref() == Some("failure")
223                             || c.conclusion.as_deref() == Some("timed_out"))
224                    .map(|c| c.name.clone())
225                    .collect();
226                let msg = crate::lifecycle::reactions::format_ci_reaction(
227                    &session, &ci, &failing_names
228                );
229                if let Err(e) = self.engine.send_to_session(&session.id, &msg).await {
230                    tracing::warn!("send ci reaction to {}: {e}", session.id);
231                }
232            }
233
234            // -- Review threads (throttled via seen_comment_ids) --
235            let threads = match gh.get_review_threads(&owner, &repo, pr_number).await {
236                Ok(t)  => t,
237                Err(e) => { tracing::warn!("github review threads: {e}"); vec![] }
238            };
239
240            let has_changes_requested = threads.iter().any(|t| t.state == "CHANGES_REQUESTED");
241
242            let (has_new, review_reaction_already_sent, new_comments) = {
243                let mut cache = self.enrichment_cache.lock().unwrap();
244                let state = cache.entry(session.id.clone()).or_default();
245                let mut has_new = false;
246                let mut new_comments: Vec<Comment> = Vec::new();
247
248                for thread in &threads {
249                    if thread.state == "CHANGES_REQUESTED"
250                        && !state.seen_comment_ids.contains(&thread.id)
251                    {
252                        state.seen_comment_ids.insert(thread.id);
253                        has_new = true;
254                        let comment = Comment {
255                            id:         thread.id,
256                            pr_id,
257                            author:     thread.author.clone(),
258                            body:       thread.body.clone(),
259                            path:       thread.path.clone(),
260                            line:       thread.line,
261                            created_at: 0,
262                        };
263                        let _ = self.engine.store.upsert_comment(&comment);
264                        self.engine.emit(Event::ReviewComment { pr_id, comment: comment.clone() });
265                        new_comments.push(comment);
266                    }
267                }
268
269                let already_sent = state.review_reaction_sent;
270                if has_new && !already_sent {
271                    state.review_reaction_sent = true;
272                }
273                // Reset when all CHANGES_REQUESTED are resolved
274                if !has_changes_requested {
275                    state.review_reaction_sent = false;
276                }
277                (has_new, already_sent, new_comments)
278            };
279
280            // Update session status in DB (after review threads so has_changes_requested is known)
281            let new_status = derive_session_status(&session.status, &pr_status, &ci, has_changes_requested);
282            let mut updated = session.clone();
283            updated.status = new_status;
284            if updated.status != session.status {
285                let _ = self.engine.store.upsert_session(&updated);
286                self.engine.emit(Event::SessionUpdated(updated.clone()));
287            }
288
289            if has_new && !review_reaction_already_sent {
290                self.engine.emit(Event::Notification(Notification {
291                    id:         format!("review-{}", session.id),
292                    kind:       NotificationKind::PrNeedsAttention,
293                    title:      format!("Review comments — {}", session.name),
294                    body:       "Changes requested on your PR".to_string(),
295                    session_id: Some(session.id.clone()),
296                    created_at: now_millis(),
297                }));
298                if !new_comments.is_empty() {
299                    let msg = crate::lifecycle::reactions::format_review_reaction(
300                        &session, &new_comments
301                    );
302                    if let Err(e) = self.engine.send_to_session(&session.id, &msg).await {
303                        tracing::warn!("send review reaction to {}: {e}", session.id);
304                    }
305                }
306            }
307        }
308    }
309}
310
311// ── Helpers ──────────────────────────────────────────────────────────────────
312
313fn summarize_checks(pr_id: PrId, checks: &[CheckRun]) -> CIStatus {
314    let total   = checks.len() as u32;
315    let failing = checks.iter().filter(|c| {
316        c.conclusion.as_deref() == Some("failure")
317            || c.conclusion.as_deref() == Some("timed_out")
318    }).count() as u32;
319    let passing = checks.iter().filter(|c| {
320        c.conclusion.as_deref() == Some("success")
321    }).count() as u32;
322    let pending = total - failing - passing;
323    CIStatus { pr_id, total, failing, passing, pending }
324}
325
326fn derive_session_status(
327    current:               &SessionStatus,
328    pr_status:             &crate::github::PrStatus,
329    ci:                    &CIStatus,
330    has_changes_requested: bool,
331) -> SessionStatus {
332    // Terminal states are never overwritten.
333    if matches!(current, SessionStatus::Done | SessionStatus::Terminated) {
334        return current.clone();
335    }
336    if pr_status.merged {
337        return SessionStatus::Done;
338    }
339    if ci.failing > 0 {
340        return SessionStatus::CiFailed;
341    }
342    if has_changes_requested {
343        return SessionStatus::ReviewPending;
344    }
345    if pr_status.mergeable == Some(true) && ci.failing == 0 && ci.pending == 0 {
346        return SessionStatus::Mergeable;
347    }
348    SessionStatus::PrOpen
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354    use crate::types::SessionStatus;
355
356    #[test]
357    fn summarize_checks_counts_failures() {
358        let checks = vec![
359            CheckRun { name: "lint".into(), status: "completed".into(), conclusion: Some("success".into()) },
360            CheckRun { name: "test".into(), status: "completed".into(), conclusion: Some("failure".into()) },
361            CheckRun { name: "build".into(), status: "in_progress".into(), conclusion: None },
362        ];
363        let ci = summarize_checks(1, &checks);
364        assert_eq!(ci.total,   3);
365        assert_eq!(ci.passing, 1);
366        assert_eq!(ci.failing, 1);
367        assert_eq!(ci.pending, 1);
368    }
369
370    #[test]
371    fn derive_status_merged_becomes_done() {
372        let pr = crate::github::PrStatus {
373            merged: true, state: "closed".into(), mergeable: None,
374            title: "t".into(), number: 1, head_sha: String::new(),
375        };
376        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
377        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
378        assert!(matches!(s, SessionStatus::Done));
379    }
380
381    #[test]
382    fn derive_status_ci_failure_overrides_open() {
383        let pr = crate::github::PrStatus {
384            merged: false, state: "open".into(), mergeable: Some(true),
385            title: "t".into(), number: 1, head_sha: String::new(),
386        };
387        let ci = CIStatus { pr_id: 1, total: 3, failing: 1, passing: 2, pending: 0 };
388        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
389        assert!(matches!(s, SessionStatus::CiFailed));
390    }
391
392    #[test]
393    fn derive_status_all_green_becomes_mergeable() {
394        let pr = crate::github::PrStatus {
395            merged: false, state: "open".into(), mergeable: Some(true),
396            title: "t".into(), number: 1, head_sha: String::new(),
397        };
398        let ci = CIStatus { pr_id: 1, total: 3, failing: 0, passing: 3, pending: 0 };
399        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
400        assert!(matches!(s, SessionStatus::Mergeable));
401    }
402
403    #[test]
404    fn derive_status_preserves_done() {
405        let pr = crate::github::PrStatus {
406            merged: false, state: "open".into(), mergeable: Some(true),
407            title: "t".into(), number: 1, head_sha: String::new(),
408        };
409        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
410        let s  = derive_session_status(&SessionStatus::Done, &pr, &ci, false);
411        assert!(matches!(s, SessionStatus::Done));
412    }
413
414    #[test]
415    fn derive_status_preserves_terminated() {
416        let pr = crate::github::PrStatus {
417            merged: true, state: "closed".into(), mergeable: None,   // merged=true!
418            title: "t".into(), number: 1, head_sha: String::new(),
419        };
420        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
421        let s  = derive_session_status(&SessionStatus::Terminated, &pr, &ci, false);
422        assert!(matches!(s, SessionStatus::Terminated));  // must not become Done
423    }
424
425    #[test]
426    fn derive_status_changes_requested_becomes_review_pending() {
427        let pr = crate::github::PrStatus {
428            merged: false, state: "open".into(), mergeable: Some(true),
429            title: "t".into(), number: 1, head_sha: String::new(),
430        };
431        let ci = CIStatus { pr_id: 1, total: 3, failing: 0, passing: 3, pending: 0 };
432        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, true);
433        assert!(matches!(s, SessionStatus::ReviewPending));
434    }
435
436    fn test_session(id: &str, workspace: &str) -> crate::types::Session {
437        crate::types::Session {
438            id: id.into(), orchestrator_id: None, name: id.into(),
439            repo: String::new(), status: SessionStatus::Working,
440            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
441            pr_number: None, pr_id: None,
442            workspace_path: Some(workspace.into()), pid: None,
443            model: None, context_tokens: None, catalogue_path: None,
444        }
445    }
446
447    /// End-to-end (within-process) proof that the poller closes the gap
448    /// documented in `lifecycle::usage`: given a workspace whose `claude`
449    /// transcript directory has usage recorded, `poll_usage` writes the
450    /// derived cost/context/model back into the store and emits
451    /// `SessionUpdated` — the exact path the UI's $0.0000 / missing-tokens
452    /// symptom traces back to when this ingestion doesn't happen.
453    // The `ENV_TEST_GUARD` mutex is intentionally held across the `.await`
454    // points below — it serializes access to the process-global
455    // `NINOX_CLAUDE_PROJECTS_DIR` env var against other tests (in this file
456    // and in `lifecycle::usage`) for this single-threaded `#[tokio::test]`,
457    // and must stay held for the env var's entire lifetime, not just around
458    // the sync portions.
459    #[allow(clippy::await_holding_lock)]
460    #[tokio::test]
461    async fn poll_usage_ingests_transcript_into_store_and_emits_update() {
462        use crate::{lifecycle::usage::{claude_project_slug, ENV_TEST_GUARD}, store::Store};
463        use std::io::Write;
464
465        let _guard = ENV_TEST_GUARD.lock().unwrap_or_else(|e| e.into_inner());
466        let projects_dir = tempfile::tempdir().unwrap();
467        let workspace = "/tmp/poller-usage-probe-workspace";
468        let project_dir = projects_dir.path().join(claude_project_slug(workspace));
469        std::fs::create_dir_all(&project_dir).unwrap();
470        let mut f = std::fs::File::create(project_dir.join("s.jsonl")).unwrap();
471        writeln!(
472            f,
473            r#"{{"type":"assistant","timestamp":"2026-07-05T13:00:00.000Z","message":{{"model":"claude-fable-5","usage":{{"input_tokens":2,"output_tokens":300,"cache_creation_input_tokens":500,"cache_read_input_tokens":45000}}}}}}"#
474        ).unwrap();
475        drop(f);
476
477        let prior = std::env::var("NINOX_CLAUDE_PROJECTS_DIR").ok();
478        std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", projects_dir.path());
479
480        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
481        store.upsert_session(&test_session("s1", workspace)).unwrap();
482        let engine = Engine::new(store.clone());
483        let mut rx = engine.subscribe();
484        let poller = Poller::new(engine);
485
486        poller.poll_usage().await;
487
488        match prior {
489            Some(v) => std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", v),
490            None    => std::env::remove_var("NINOX_CLAUDE_PROJECTS_DIR"),
491        }
492
493        let updated = store.get_session("s1").unwrap().unwrap();
494        assert!(updated.cost_usd > 0.0, "cost_usd should be ingested, not 0.0000");
495        assert_eq!(updated.context_tokens, Some(2 + 500 + 45000));
496        assert_eq!(updated.model.as_deref(), Some("claude-fable-5"));
497
498        let evt = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv())
499            .await
500            .expect("SessionUpdated should be emitted")
501            .unwrap();
502        assert!(matches!(evt, Event::SessionUpdated(s) if s.id == "s1" && s.cost_usd > 0.0));
503    }
504
505    #[tokio::test]
506    async fn poll_usage_leaves_sessions_without_workspace_or_usage_untouched() {
507        use crate::store::Store;
508
509        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
510        let mut no_ws = test_session("no-ws", "/does/not/matter");
511        no_ws.workspace_path = None;
512        store.upsert_session(&no_ws).unwrap();
513        let engine = Engine::new(store.clone());
514        let poller = Poller::new(engine);
515
516        poller.poll_usage().await;
517
518        let unchanged = store.get_session("no-ws").unwrap().unwrap();
519        assert_eq!(unchanged.cost_usd, 0.0);
520        assert_eq!(unchanged.context_tokens, None);
521    }
522}
ninox_core/lifecycle/poller.rs

ninox_core/lifecycle/
poller.rs
