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/// Last-seen `(cost_usd, context_used_pct, context_total_tokens)` snapshot per session.
19/// Used by `poll_context_updates` to detect external changes.
20type ContextSnapshot = (f64, Option<f64>, Option<u64>);
21
22/// Unix epoch milliseconds "now" — used to stamp `Notification::created_at`.
23fn now_millis() -> i64 {
24    std::time::SystemTime::now()
25        .duration_since(std::time::UNIX_EPOCH)
26        .map(|d| d.as_millis() as i64)
27        .unwrap_or(0)
28}
29
30pub struct Poller {
31    engine:           Arc<Engine>,
32    enrichment_cache: Arc<std::sync::Mutex<EnrichmentCache>>,
33    /// Last-seen `(cost_usd, context_used_pct, context_total_tokens)` per
34    /// session, used solely to detect changes written externally by the
35    /// `ninox statusline` subcommand — see `poll_context_updates`.
36    context_cache:    Arc<std::sync::Mutex<HashMap<String, ContextSnapshot>>>,
37}
38
39impl Poller {
40    pub fn new(engine: Arc<Engine>) -> Self {
41        Self {
42            engine,
43            enrichment_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
44            context_cache:    Arc::new(std::sync::Mutex::new(HashMap::new())),
45        }
46    }
47
48    pub async fn start(self, token: CancellationToken) {
49        let mut pid_interval    = tokio::time::interval(Duration::from_secs(5));
50        let mut usage_interval  = tokio::time::interval(Duration::from_secs(10));
51        let mut github_interval = tokio::time::interval(Duration::from_secs(30));
52        // Prevent a missed tick from causing back-to-back polls.
53        github_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
54        usage_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
55        loop {
56            tokio::select! {
57                _ = token.cancelled()      => break,
58                _ = pid_interval.tick()    => {
59                    self.poll_pids().await;
60                    self.poll_context_updates().await;
61                }
62                _ = usage_interval.tick()  => self.poll_usage().await,
63                _ = github_interval.tick() => self.poll_github().await,
64            }
65        }
66    }
67
68    // ── PID liveness ────────────────────────────────────────────────────────
69
70    async fn poll_pids(&self) {
71        // Metadata first: a dying worker's last acts (PR create, work
72        // request) are processed before the reap below marks it Terminated.
73        self.sync_sessions_metadata(&AppConfig::sessions_dir()).await;
74
75        let Ok(sessions) = self.engine.store.list_sessions() else { return };
76        for mut session in sessions {
77            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated | SessionStatus::Interrupted) {
78                continue;
79            }
80            if let Some(pid) = session.pid {
81                if !is_pid_alive(pid) {
82                    session.status = SessionStatus::Terminated;
83                    let _ = self.engine.store.upsert_session(&session);
84                    self.engine.emit(Event::SessionUpdated(session));
85                }
86            }
87        }
88    }
89
90    // ── Session metadata (wrapper hooks + `ninox request-work`) ────────────
91
92    /// One pass over every non-terminal session's metadata file: adopt the
93    /// first reported PR as the session's canonical one, record + notify any
94    /// PR opened beyond it, and deliver pending work requests to the
95    /// orchestrator. The dir is a parameter so tests can drive this against
96    /// a tempdir instead of `AppConfig::sessions_dir()`.
97    async fn sync_sessions_metadata(&self, sessions_dir: &std::path::Path) {
98        let Ok(sessions) = self.engine.store.list_sessions() else { return };
99        for mut session in sessions {
100            // Work requests are about *new* work, not this session — deliver
101            // them even when the requesting worker has already finished or
102            // died (a worker's last act is often "request follow-up, exit").
103            self.deliver_work_requests(&session, sessions_dir).await;
104
105            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated | SessionStatus::Interrupted) {
106                continue;
107            }
108            let Ok(meta) = hooks::read_session_metadata(sessions_dir, &session.id) else {
109                continue;
110            };
111
112            // -- First reported PR becomes the session's tracked PR --
113            if session.pr_number.is_none() {
114                if let Some(first) = meta.pr_reports.first() {
115                    session.pr_number = Some(first.number);
116                    session.status    = SessionStatus::PrOpen;
117                    let _ = self.engine.store.upsert_session(&session);
118                    self.engine.emit(Event::SessionUpdated(session.clone()));
119                    tracing::info!(
120                        "session {} PR #{} detected via metadata hook",
121                        session.id, first.number
122                    );
123                }
124            }
125
126            // -- Every reported PR beyond the tracked one --
127            let Some(tracked) = session.pr_number else { continue };
128
129            // Ledger rows first, every tick: a store error at notification
130            // time must only defer the row to the next tick, never lose it.
131            // Only write when the row id (bare PR number — collides across
132            // repos) is free: never steal another session's row.
133            for report in meta.pr_reports.iter().filter(|r| r.number != tracked) {
134                if let Ok(None) = self.engine.store.get_pr(report.number as i64) {
135                    let url = report.url.clone().unwrap_or_else(|| {
136                        format!("https://github.com/{}/pull/{}", session.repo, report.number)
137                    });
138                    let _ = self.engine.store.upsert_pr(&PR {
139                        id:         report.number as i64,
140                        number:     report.number,
141                        title:      String::new(),
142                        url,
143                        body:       String::new(),
144                        session_id: session.id.clone(),
145                    });
146                }
147            }
148
149            // Notifications second, deduped via the poller-owned side file.
150            let notified = hooks::read_notified_extra_prs(sessions_dir, &session.id);
151            let mut fresh: Vec<(u64, Option<String>)> = Vec::new();
152            for report in meta.pr_reports.iter()
153                .filter(|r| r.number != tracked && !notified.contains(&r.number))
154            {
155                self.engine.emit(Event::Notification(Notification {
156                    id:         format!("extra-pr-{}-{}", session.id, report.number),
157                    kind:       NotificationKind::ExtraPr,
158                    title:      format!("Extra PR — {}", session.name),
159                    body:       format!("#{} opened beyond tracked #{tracked}", report.number),
160                    session_id: Some(session.id.clone()),
161                    created_at: now_millis(),
162                }));
163                fresh.push((report.number, report.url.clone()));
164            }
165            if fresh.is_empty() {
166                continue;
167            }
168            let numbers: Vec<u64> = fresh.iter().map(|(n, _)| *n).collect();
169            if let Err(e) = hooks::mark_extra_prs_notified(sessions_dir, &session.id, &numbers) {
170                tracing::warn!("mark extra PRs notified for {}: {e}", session.id);
171            }
172            if let Some(orch) = session.orchestrator_id.clone() {
173                let msg = crate::lifecycle::reactions::format_extra_pr_reaction(
174                    &session, tracked, &fresh,
175                );
176                if let Err(e) = self.engine.send_to_session(&orch, &msg).await {
177                    tracing::warn!("send extra-PR reaction to orchestrator {orch}: {e}");
178                }
179            }
180        }
181    }
182
183    /// Forward every pending `ninox request-work` entry for this session to
184    /// the UI and the orchestrator, then move it out of the pending set.
185    async fn deliver_work_requests(&self, session: &crate::types::Session, sessions_dir: &std::path::Path) {
186        let pending = match hooks::read_pending_work_requests(sessions_dir, &session.id) {
187            Ok(p) if !p.is_empty() => p,
188            _ => return,
189        };
190        for request in &pending {
191            self.engine.emit(Event::Notification(Notification {
192                id:         format!("work-request-{}", request.id),
193                kind:       NotificationKind::WorkRequested,
194                title:      format!("Work requested — {}", session.name),
195                body:       request.description.clone(),
196                session_id: Some(session.id.clone()),
197                created_at: now_millis(),
198            }));
199            if let Some(orch) = session.orchestrator_id.clone() {
200                let msg = crate::lifecycle::reactions::format_work_request_reaction(
201                    session, &request.description,
202                );
203                if let Err(e) = self.engine.send_to_session(&orch, &msg).await {
204                    tracing::warn!("send work request to orchestrator {orch}: {e}");
205                }
206            }
207        }
208        // Marked delivered even when the tmux nudge failed — the UI
209        // notification is already out, and retrying every tick would spam
210        // both channels.
211        let ids: Vec<String> = pending.iter().map(|r| r.id.clone()).collect();
212        if let Err(e) = hooks::mark_work_requests_delivered(sessions_dir, &session.id, &ids) {
213            tracing::warn!("mark work requests delivered for {}: {e}", session.id);
214        }
215    }
216
217    // ── Cost / context-window usage ─────────────────────────────────────────
218
219    /// Ingest cost/token usage for every active session by reading `claude`'s
220    /// own transcript for the session's workspace directory (see
221    /// `lifecycle::usage`). Sessions without a workspace, or whose transcript
222    /// has no usage yet (agent hasn't taken a turn), are left untouched.
223    /// Only writes + emits when something actually changed, so this doesn't
224    /// spam the store/UI every tick for idle sessions.
225    async fn poll_usage(&self) {
226        let Ok(sessions) = self.engine.store.list_sessions() else { return };
227        for mut session in sessions {
228            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated | SessionStatus::Interrupted) {
229                continue;
230            }
231            let Some(workspace) = session.workspace_path.clone() else { continue };
232            let Some(snapshot) = usage::ingest_usage_for_workspace(&workspace) else { continue };
233
234            let cost_changed = (session.cost_usd - snapshot.cost_usd).abs() > 1e-9;
235            let context_changed = session.context_tokens != Some(snapshot.context_tokens);
236            if !cost_changed && !context_changed {
237                continue;
238            }
239
240            session.cost_usd = snapshot.cost_usd;
241            session.context_tokens = Some(snapshot.context_tokens);
242            if session.model.is_none() {
243                session.model = snapshot.model;
244            }
245            let _ = self.engine.store.upsert_session(&session);
246            self.engine.emit(Event::SessionUpdated(session));
247        }
248    }
249
250    // ── Statusline-sourced cost/context updates (external writer) ──────────
251
252    /// The `ninox statusline` subcommand (invoked by Claude Code's own
253    /// `statusLine` hook — see `lifecycle::statusline`) writes cost/context
254    /// fields directly into the store from a separate short-lived process.
255    /// Unlike every other poll method, this data doesn't arrive via a
256    /// read-modify-write cycle this poller drives, so there's nothing to
257    /// diff against except a cache of the last-seen values. Detects
258    /// external changes and re-broadcasts them as `SessionUpdated` so the
259    /// GUI picks them up.
260    async fn poll_context_updates(&self) {
261        let Ok(sessions) = self.engine.store.list_sessions() else { return };
262        let mut changed = Vec::new();
263        {
264            let mut cache = self.context_cache.lock().unwrap();
265            for session in sessions {
266                let key = (session.cost_usd, session.context_used_pct, session.context_total_tokens);
267                // `None` means this session has never been cached — seed it
268                // silently rather than treating "no prior state" as a change
269                // (that would spam an event for every session on startup).
270                if let Some(prev) = cache.insert(session.id.clone(), key) {
271                    if prev != key {
272                        changed.push(session);
273                    }
274                }
275            }
276        }
277        for session in changed {
278            self.engine.emit(Event::SessionUpdated(session));
279        }
280    }
281
282    // ── GitHub enrichment ────────────────────────────────────────────────────
283
284    async fn poll_github(&self) {
285        let Some(gh) = &self.engine.github else { return };
286        let Ok(sessions) = self.engine.store.list_sessions() else { return };
287
288        for session in sessions {
289            if matches!(session.status, SessionStatus::Done | SessionStatus::Terminated | SessionStatus::Interrupted) {
290                continue;
291            }
292            let Some(pr_number) = session.pr_number else { continue };
293            let Some((owner, repo)) = split_repo(&session.repo) else { continue };
294
295            // -- PR state --
296            let pr_status = match gh.get_pr_status(&owner, &repo, pr_number).await {
297                Ok(s)  => s,
298                Err(e) => { tracing::warn!("github pr status: {e}"); continue }
299            };
300
301            let pr_id: PrId = pr_number as i64;
302
303            // -- Merge detection — handle before CI (no point polling CI on merged PR) --
304            if pr_status.merged && !matches!(session.status, SessionStatus::Done) {
305                self.engine.emit(Event::Notification(Notification {
306                    id:         format!("merged-{}", session.id),
307                    kind:       NotificationKind::WorkerDone,
308                    title:      format!("PR merged — {}", session.name),
309                    body:       format!("#{} merged successfully", pr_number),
310                    session_id: Some(session.id.clone()),
311                    created_at: now_millis(),
312                }));
313                if let Err(e) = self.engine.cleanup_session(&session.id).await {
314                    tracing::warn!("cleanup_session {}: {e}", session.id);
315                }
316                // Remove enrichment state for this session — it's done
317                {
318                    let mut cache = self.enrichment_cache.lock().unwrap();
319                    cache.remove(&session.id);
320                }
321                continue; // skip further enrichment for this session
322            }
323
324            // Upsert PR record — only when not merged (merged sessions stay Done after cleanup)
325            {
326                let pr = PR {
327                    id:         pr_id,
328                    number:     pr_number,
329                    title:      pr_status.title.clone(),
330                    url:        format!("https://github.com/{owner}/{repo}/pull/{pr_number}"),
331                    body:       String::new(),
332                    session_id: session.id.clone(),
333                };
334                let _ = self.engine.store.upsert_pr(&pr);
335                self.engine.emit(Event::PrOpened { session_id: session.id.clone(), pr });
336            }
337
338            // -- CI checks --
339            let checks = match gh.get_ci_checks(&owner, &repo, &pr_status.head_sha).await {
340                Ok(c)  => c,
341                Err(e) => { tracing::warn!("github ci checks: {e}"); vec![] }
342            };
343            let ci = summarize_checks(pr_id, &checks);
344            let _ = self.engine.store.upsert_ci_status(&ci);
345            self.engine.emit(Event::CiUpdated { pr_id, status: ci.clone() });
346
347            // -- Detect CI transition and update session status --
348            let (newly_failing, ci_reaction_already_sent) = {
349                let mut cache = self.enrichment_cache.lock().unwrap();
350                let state = cache.entry(session.id.clone()).or_default();
351
352                let newly_failing = state.prev_failing.is_none_or(|p| p == 0)
353                    && ci.failing > 0;
354                state.prev_failing = Some(ci.failing);
355
356                let already_sent = state.ci_reaction_sent;
357                if newly_failing && !already_sent {
358                    state.ci_reaction_sent = true;
359                }
360                if ci.failing == 0 {
361                    state.ci_reaction_sent = false;
362                }
363                (newly_failing, already_sent)
364            };
365
366            if newly_failing && !ci_reaction_already_sent {
367                self.engine.emit(Event::Notification(Notification {
368                    id:         format!("ci-{}", session.id),
369                    kind:       NotificationKind::CiFailure,
370                    title:      format!("CI failing — {}", session.name),
371                    body:       format!("{}/{} checks failing", ci.failing, ci.total),
372                    session_id: Some(session.id.clone()),
373                    created_at: now_millis(),
374                }));
375                // Send reaction to the agent in the tmux session
376                let failing_names: Vec<String> = checks.iter()
377                    .filter(|c| c.conclusion.as_deref() == Some("failure")
378                             || c.conclusion.as_deref() == Some("timed_out"))
379                    .map(|c| c.name.clone())
380                    .collect();
381                let msg = crate::lifecycle::reactions::format_ci_reaction(
382                    &session, &ci, &failing_names
383                );
384                if let Err(e) = self.engine.send_to_session(&session.id, &msg).await {
385                    tracing::warn!("send ci reaction to {}: {e}", session.id);
386                }
387            }
388
389            // -- Review threads (throttled via seen_comment_ids) --
390            let threads = match gh.get_review_threads(&owner, &repo, pr_number).await {
391                Ok(t)  => t,
392                Err(e) => { tracing::warn!("github review threads: {e}"); vec![] }
393            };
394
395            let has_changes_requested = threads.iter().any(|t| t.state == "CHANGES_REQUESTED");
396
397            let (has_new, review_reaction_already_sent, new_comments) = {
398                let mut cache = self.enrichment_cache.lock().unwrap();
399                let state = cache.entry(session.id.clone()).or_default();
400                let mut has_new = false;
401                let mut new_comments: Vec<Comment> = Vec::new();
402
403                for thread in &threads {
404                    if thread.state == "CHANGES_REQUESTED"
405                        && !state.seen_comment_ids.contains(&thread.id)
406                    {
407                        state.seen_comment_ids.insert(thread.id);
408                        has_new = true;
409                        let comment = Comment {
410                            id:         thread.id,
411                            pr_id,
412                            author:     thread.author.clone(),
413                            body:       thread.body.clone(),
414                            path:       thread.path.clone(),
415                            line:       thread.line,
416                            created_at: 0,
417                        };
418                        let _ = self.engine.store.upsert_comment(&comment);
419                        self.engine.emit(Event::ReviewComment { pr_id, comment: comment.clone() });
420                        new_comments.push(comment);
421                    }
422                }
423
424                let already_sent = state.review_reaction_sent;
425                if has_new && !already_sent {
426                    state.review_reaction_sent = true;
427                }
428                // Reset when all CHANGES_REQUESTED are resolved
429                if !has_changes_requested {
430                    state.review_reaction_sent = false;
431                }
432                (has_new, already_sent, new_comments)
433            };
434
435            // Update session status in DB (after review threads so has_changes_requested is known)
436            let new_status = derive_session_status(&session.status, &pr_status, &ci, has_changes_requested);
437            let mut updated = session.clone();
438            updated.status = new_status;
439            if updated.status != session.status {
440                let _ = self.engine.store.upsert_session(&updated);
441                self.engine.emit(Event::SessionUpdated(updated.clone()));
442            }
443
444            if has_new && !review_reaction_already_sent {
445                self.engine.emit(Event::Notification(Notification {
446                    id:         format!("review-{}", session.id),
447                    kind:       NotificationKind::PrNeedsAttention,
448                    title:      format!("Review comments — {}", session.name),
449                    body:       "Changes requested on your PR".to_string(),
450                    session_id: Some(session.id.clone()),
451                    created_at: now_millis(),
452                }));
453                if !new_comments.is_empty() {
454                    let msg = crate::lifecycle::reactions::format_review_reaction(
455                        &session, &new_comments
456                    );
457                    if let Err(e) = self.engine.send_to_session(&session.id, &msg).await {
458                        tracing::warn!("send review reaction to {}: {e}", session.id);
459                    }
460                }
461            }
462        }
463    }
464}
465
466// ── Helpers ──────────────────────────────────────────────────────────────────
467
468fn summarize_checks(pr_id: PrId, checks: &[CheckRun]) -> CIStatus {
469    let total   = checks.len() as u32;
470    let failing = checks.iter().filter(|c| {
471        c.conclusion.as_deref() == Some("failure")
472            || c.conclusion.as_deref() == Some("timed_out")
473    }).count() as u32;
474    let passing = checks.iter().filter(|c| {
475        c.conclusion.as_deref() == Some("success")
476    }).count() as u32;
477    let pending = total - failing - passing;
478    CIStatus { pr_id, total, failing, passing, pending }
479}
480
481fn derive_session_status(
482    current:               &SessionStatus,
483    pr_status:             &crate::github::PrStatus,
484    ci:                    &CIStatus,
485    has_changes_requested: bool,
486) -> SessionStatus {
487    // Terminal states are never overwritten.
488    if matches!(current, SessionStatus::Done | SessionStatus::Terminated | SessionStatus::Interrupted) {
489        return current.clone();
490    }
491    if pr_status.merged {
492        return SessionStatus::Done;
493    }
494    if ci.failing > 0 {
495        return SessionStatus::CiFailed;
496    }
497    if has_changes_requested {
498        return SessionStatus::ReviewPending;
499    }
500    if pr_status.mergeable == Some(true) && ci.failing == 0 && ci.pending == 0 {
501        return SessionStatus::Mergeable;
502    }
503    SessionStatus::PrOpen
504}
505
506#[cfg(test)]
507mod tests {
508    use super::*;
509    use crate::types::SessionStatus;
510
511    #[test]
512    fn summarize_checks_counts_failures() {
513        let checks = vec![
514            CheckRun { name: "lint".into(), status: "completed".into(), conclusion: Some("success".into()) },
515            CheckRun { name: "test".into(), status: "completed".into(), conclusion: Some("failure".into()) },
516            CheckRun { name: "build".into(), status: "in_progress".into(), conclusion: None },
517        ];
518        let ci = summarize_checks(1, &checks);
519        assert_eq!(ci.total,   3);
520        assert_eq!(ci.passing, 1);
521        assert_eq!(ci.failing, 1);
522        assert_eq!(ci.pending, 1);
523    }
524
525    #[test]
526    fn derive_status_merged_becomes_done() {
527        let pr = crate::github::PrStatus {
528            merged: true, state: "closed".into(), mergeable: None,
529            title: "t".into(), number: 1, head_sha: String::new(),
530        };
531        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
532        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
533        assert!(matches!(s, SessionStatus::Done));
534    }
535
536    #[test]
537    fn derive_status_ci_failure_overrides_open() {
538        let pr = crate::github::PrStatus {
539            merged: false, state: "open".into(), mergeable: Some(true),
540            title: "t".into(), number: 1, head_sha: String::new(),
541        };
542        let ci = CIStatus { pr_id: 1, total: 3, failing: 1, passing: 2, pending: 0 };
543        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
544        assert!(matches!(s, SessionStatus::CiFailed));
545    }
546
547    #[test]
548    fn derive_status_all_green_becomes_mergeable() {
549        let pr = crate::github::PrStatus {
550            merged: false, state: "open".into(), mergeable: Some(true),
551            title: "t".into(), number: 1, head_sha: String::new(),
552        };
553        let ci = CIStatus { pr_id: 1, total: 3, failing: 0, passing: 3, pending: 0 };
554        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, false);
555        assert!(matches!(s, SessionStatus::Mergeable));
556    }
557
558    #[test]
559    fn derive_status_preserves_done() {
560        let pr = crate::github::PrStatus {
561            merged: false, state: "open".into(), mergeable: Some(true),
562            title: "t".into(), number: 1, head_sha: String::new(),
563        };
564        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
565        let s  = derive_session_status(&SessionStatus::Done, &pr, &ci, false);
566        assert!(matches!(s, SessionStatus::Done));
567    }
568
569    #[test]
570    fn derive_status_preserves_terminated() {
571        let pr = crate::github::PrStatus {
572            merged: true, state: "closed".into(), mergeable: None,   // merged=true!
573            title: "t".into(), number: 1, head_sha: String::new(),
574        };
575        let ci = CIStatus { pr_id: 1, total: 0, failing: 0, passing: 0, pending: 0 };
576        let s  = derive_session_status(&SessionStatus::Terminated, &pr, &ci, false);
577        assert!(matches!(s, SessionStatus::Terminated));  // must not become Done
578    }
579
580    #[test]
581    fn derive_status_changes_requested_becomes_review_pending() {
582        let pr = crate::github::PrStatus {
583            merged: false, state: "open".into(), mergeable: Some(true),
584            title: "t".into(), number: 1, head_sha: String::new(),
585        };
586        let ci = CIStatus { pr_id: 1, total: 3, failing: 0, passing: 3, pending: 0 };
587        let s  = derive_session_status(&SessionStatus::PrOpen, &pr, &ci, true);
588        assert!(matches!(s, SessionStatus::ReviewPending));
589    }
590
591    fn test_session(id: &str, workspace: &str) -> crate::types::Session {
592        crate::types::Session {
593            id: id.into(), orchestrator_id: None, name: id.into(),
594            repo: String::new(), status: SessionStatus::Working,
595            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
596            pr_number: None, pr_id: None,
597            workspace_path: Some(workspace.into()), pid: None,
598            model: None, context_tokens: None, catalogue_path: None,
599            context_used_pct: None, context_total_tokens: None, context_window_size: None,
600            claude_session_id: None,
601        }
602    }
603
604    #[tokio::test]
605    async fn poll_pids_leaves_interrupted_sessions_alone() {
606        use crate::store::Store;
607
608        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
609        let mut s = test_session("interrupted-1", "/ws");
610        s.status = SessionStatus::Interrupted;
611        s.pid = Some(999_999); // a pid that is almost certainly dead
612        store.upsert_session(&s).unwrap();
613        let engine = Engine::new(store.clone());
614        let poller = Poller::new(engine);
615
616        poller.poll_pids().await;
617
618        let after = store.get_session("interrupted-1").unwrap().unwrap();
619        assert!(
620            matches!(after.status, SessionStatus::Interrupted),
621            "poll_pids must not re-terminate an Interrupted session just because its stale pid is dead",
622        );
623    }
624
625    /// End-to-end (within-process) proof that the poller closes the gap
626    /// documented in `lifecycle::usage`: given a workspace whose `claude`
627    /// transcript directory has usage recorded, `poll_usage` writes the
628    /// derived cost/context/model back into the store and emits
629    /// `SessionUpdated` — the exact path the UI's $0.0000 / missing-tokens
630    /// symptom traces back to when this ingestion doesn't happen.
631    // The `ENV_TEST_GUARD` mutex is intentionally held across the `.await`
632    // points below — it serializes access to the process-global
633    // `NINOX_CLAUDE_PROJECTS_DIR` env var against other tests (in this file
634    // and in `lifecycle::usage`) for this single-threaded `#[tokio::test]`,
635    // and must stay held for the env var's entire lifetime, not just around
636    // the sync portions.
637    #[allow(clippy::await_holding_lock)]
638    #[tokio::test]
639    async fn poll_usage_ingests_transcript_into_store_and_emits_update() {
640        use crate::{lifecycle::usage::{claude_project_slug, ENV_TEST_GUARD}, store::Store};
641        use std::io::Write;
642
643        let _guard = ENV_TEST_GUARD.lock().unwrap_or_else(|e| e.into_inner());
644        let projects_dir = tempfile::tempdir().unwrap();
645        let workspace = "/tmp/poller-usage-probe-workspace";
646        let project_dir = projects_dir.path().join(claude_project_slug(workspace));
647        std::fs::create_dir_all(&project_dir).unwrap();
648        let mut f = std::fs::File::create(project_dir.join("s.jsonl")).unwrap();
649        writeln!(
650            f,
651            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}}}}}}"#
652        ).unwrap();
653        drop(f);
654
655        let prior = std::env::var("NINOX_CLAUDE_PROJECTS_DIR").ok();
656        std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", projects_dir.path());
657
658        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
659        store.upsert_session(&test_session("s1", workspace)).unwrap();
660        let engine = Engine::new(store.clone());
661        let mut rx = engine.subscribe();
662        let poller = Poller::new(engine);
663
664        poller.poll_usage().await;
665
666        match prior {
667            Some(v) => std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", v),
668            None    => std::env::remove_var("NINOX_CLAUDE_PROJECTS_DIR"),
669        }
670
671        let updated = store.get_session("s1").unwrap().unwrap();
672        assert!(updated.cost_usd > 0.0, "cost_usd should be ingested, not 0.0000");
673        assert_eq!(updated.context_tokens, Some(2 + 500 + 45000));
674        assert_eq!(updated.model.as_deref(), Some("claude-fable-5"));
675
676        let evt = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv())
677            .await
678            .expect("SessionUpdated should be emitted")
679            .unwrap();
680        assert!(matches!(evt, Event::SessionUpdated(s) if s.id == "s1" && s.cost_usd > 0.0));
681    }
682
683    /// The `ninox statusline` subcommand (a separate short-lived process)
684    /// writes cost/context fields directly into the store — outside any
685    /// read-modify-write cycle this poller drives. This proves the diff
686    /// cache detects that external write and re-broadcasts it, and that an
687    /// untouched session generates no spurious event.
688    #[tokio::test]
689    async fn poll_context_updates_emits_only_for_changed_sessions() {
690        use crate::store::Store;
691
692        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
693        let mut s1 = test_session("s1", "/ws1");
694        let s2 = test_session("s2", "/ws2");
695        store.upsert_session(&s1).unwrap();
696        store.upsert_session(&s2).unwrap();
697        let engine = Engine::new(store.clone());
698        let mut rx = engine.subscribe();
699        let poller = Poller::new(engine);
700
701        // First tick establishes the baseline — nothing to diff against yet,
702        // so it must not emit for sessions that already exist with no prior
703        // cached state.
704        poller.poll_context_updates().await;
705        let baseline_events = drain_events(&mut rx);
706        assert!(baseline_events.is_empty(), "no prior cached state means no change to report");
707
708        // Simulate the statusline hook writing directly into the store for s1 only.
709        s1.context_used_pct = Some(42.0);
710        s1.cost_usd = 3.5;
711        store.upsert_session(&s1).unwrap();
712
713        poller.poll_context_updates().await;
714        let events = drain_events(&mut rx);
715        assert_eq!(events.len(), 1, "only the changed session should emit");
716        assert!(matches!(
717            &events[0],
718            Event::SessionUpdated(s) if s.id == "s1" && s.context_used_pct == Some(42.0) && s.cost_usd == 3.5
719        ));
720
721        // A third tick with no further changes emits nothing.
722        poller.poll_context_updates().await;
723        assert!(drain_events(&mut rx).is_empty());
724    }
725
726    /// Drain every event currently buffered on the receiver.
727    fn drain_events(rx: &mut tokio::sync::broadcast::Receiver<Event>) -> Vec<Event> {
728        let mut events = Vec::new();
729        while let Ok(e) = rx.try_recv() {
730            events.push(e);
731        }
732        events
733    }
734
735    /// A worker that opened three PRs: the first becomes the session's
736    /// tracked PR, every later one is recorded in the store and raised as an
737    /// ExtraPr notification — and only once, however often the poller ticks.
738    #[tokio::test]
739    async fn metadata_sync_adopts_first_pr_and_flags_every_extra_once() {
740        use crate::store::Store;
741
742        let sessions_dir = tempfile::tempdir().unwrap();
743        let meta = serde_json::json!({
744            "agentReportedPrNumber": "44",
745            "agentReportedPrUrl": "https://github.com/org/repo/pull/44",
746            "agentReportedPrs": [
747                {"number": "42", "url": "https://github.com/org/repo/pull/42"},
748                {"number": "43", "url": "https://github.com/org/repo/pull/43"},
749                {"number": "44", "url": "https://github.com/org/repo/pull/44"},
750            ],
751        });
752        std::fs::write(
753            sessions_dir.path().join("s1.json"),
754            serde_json::to_string(&meta).unwrap(),
755        ).unwrap();
756
757        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
758        store.upsert_session(&test_session("s1", "/ws")).unwrap();
759        let engine = Engine::new(store.clone());
760        let mut rx = engine.subscribe();
761        let poller = Poller::new(engine);
762
763        poller.sync_sessions_metadata(sessions_dir.path()).await;
764
765        let session = store.get_session("s1").unwrap().unwrap();
766        assert_eq!(session.pr_number, Some(42), "first PR is the canonical one");
767        assert!(matches!(session.status, SessionStatus::PrOpen));
768        assert!(store.get_pr(43).unwrap().is_some(), "extra PR #43 recorded");
769        assert!(store.get_pr(44).unwrap().is_some(), "extra PR #44 recorded");
770        assert_eq!(
771            store.get_pr(43).unwrap().unwrap().url,
772            "https://github.com/org/repo/pull/43",
773        );
774
775        let events = drain_events(&mut rx);
776        let extra_notifs: Vec<_> = events.iter().filter(|e| matches!(
777            e, Event::Notification(n) if n.kind == crate::types::NotificationKind::ExtraPr
778        )).collect();
779        assert_eq!(extra_notifs.len(), 2, "one ExtraPr notification per extra PR");
780
781        // Second tick: nothing new — no duplicate notifications.
782        poller.sync_sessions_metadata(sessions_dir.path()).await;
783        let events = drain_events(&mut rx);
784        assert!(
785            !events.iter().any(|e| matches!(e, Event::Notification(_))),
786            "extra PRs must not be re-notified on every tick",
787        );
788    }
789
790    /// A single reported PR (the normal case) adopts it with no extra-PR
791    /// noise — the pre-existing first-PR-detection behavior.
792    #[tokio::test]
793    async fn metadata_sync_single_pr_has_no_extra_notifications() {
794        use crate::store::Store;
795
796        let sessions_dir = tempfile::tempdir().unwrap();
797        let meta = serde_json::json!({
798            "agentReportedPrNumber": "5",
799            "agentReportedPrUrl": "https://github.com/org/repo/pull/5",
800        });
801        std::fs::write(
802            sessions_dir.path().join("s1.json"),
803            serde_json::to_string(&meta).unwrap(),
804        ).unwrap();
805
806        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
807        store.upsert_session(&test_session("s1", "/ws")).unwrap();
808        let engine = Engine::new(store.clone());
809        let mut rx = engine.subscribe();
810        let poller = Poller::new(engine);
811
812        poller.sync_sessions_metadata(sessions_dir.path()).await;
813
814        let session = store.get_session("s1").unwrap().unwrap();
815        assert_eq!(session.pr_number, Some(5));
816        assert!(matches!(session.status, SessionStatus::PrOpen));
817        let events = drain_events(&mut rx);
818        assert!(!events.iter().any(|e| matches!(e, Event::Notification(_))));
819    }
820
821    /// Work requests recorded by `ninox request-work` surface exactly one
822    /// WorkRequested notification each, then are marked delivered.
823    #[tokio::test]
824    async fn metadata_sync_delivers_work_requests_exactly_once() {
825        use crate::store::Store;
826
827        let sessions_dir = tempfile::tempdir().unwrap();
828        hooks::append_work_request(sessions_dir.path(), "s1", "Migrate the config loader").unwrap();
829
830        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
831        store.upsert_session(&test_session("s1", "/ws")).unwrap();
832        let engine = Engine::new(store.clone());
833        let mut rx = engine.subscribe();
834        let poller = Poller::new(engine);
835
836        poller.sync_sessions_metadata(sessions_dir.path()).await;
837
838        let events = drain_events(&mut rx);
839        let notif = events.iter().find_map(|e| match e {
840            Event::Notification(n) if n.kind == crate::types::NotificationKind::WorkRequested => Some(n),
841            _ => None,
842        }).expect("WorkRequested notification emitted");
843        assert!(notif.body.contains("Migrate the config loader"));
844        assert_eq!(notif.session_id.as_deref(), Some("s1"));
845
846        assert!(
847            hooks::read_pending_work_requests(sessions_dir.path(), "s1").unwrap().is_empty(),
848            "delivered requests must leave the pending set",
849        );
850
851        poller.sync_sessions_metadata(sessions_dir.path()).await;
852        let events = drain_events(&mut rx);
853        assert!(
854            !events.iter().any(|e| matches!(e, Event::Notification(_))),
855            "delivered work requests must not fire again",
856        );
857    }
858
859    /// A worker can request work and exit before the next tick — the request
860    /// must still reach the orchestrator, not die with the session.
861    #[tokio::test]
862    async fn metadata_sync_delivers_work_requests_from_terminated_sessions() {
863        use crate::store::Store;
864
865        let sessions_dir = tempfile::tempdir().unwrap();
866        hooks::append_work_request(sessions_dir.path(), "s1", "Follow-up refactor").unwrap();
867
868        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
869        let mut session = test_session("s1", "/ws");
870        session.status = SessionStatus::Terminated;
871        store.upsert_session(&session).unwrap();
872        let engine = Engine::new(store.clone());
873        let mut rx = engine.subscribe();
874        let poller = Poller::new(engine);
875
876        poller.sync_sessions_metadata(sessions_dir.path()).await;
877
878        let events = drain_events(&mut rx);
879        assert!(
880            events.iter().any(|e| matches!(
881                e, Event::Notification(n) if n.kind == crate::types::NotificationKind::WorkRequested
882            )),
883            "work requests outlive their session",
884        );
885    }
886
887    /// The ledger row for an extra PR is best-effort at notification time (a
888    /// busy store must not kill the alert) — but it must self-heal on later
889    /// ticks rather than be lost forever, and healing must not re-notify.
890    #[tokio::test]
891    async fn extra_pr_ledger_row_backfills_after_notification_without_renotifying() {
892        use crate::store::Store;
893
894        let sessions_dir = tempfile::tempdir().unwrap();
895        let meta = serde_json::json!({
896            "agentReportedPrs": [
897                {"number": "7", "url": "https://github.com/org/repo/pull/7"},
898                {"number": "9", "url": "https://github.com/org/repo/pull/9"},
899            ],
900        });
901        std::fs::write(
902            sessions_dir.path().join("s1.json"),
903            serde_json::to_string(&meta).unwrap(),
904        ).unwrap();
905        // Simulate "notified previously, but the row write failed that tick".
906        hooks::mark_extra_prs_notified(sessions_dir.path(), "s1", &[9]).unwrap();
907
908        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
909        store.upsert_session(&test_session("s1", "/ws")).unwrap();
910        let engine = Engine::new(store.clone());
911        let mut rx = engine.subscribe();
912        let poller = Poller::new(engine);
913
914        poller.sync_sessions_metadata(sessions_dir.path()).await;
915
916        assert!(
917            store.get_pr(9).unwrap().is_some(),
918            "already-notified extra PR must still get its ledger row backfilled",
919        );
920        let events = drain_events(&mut rx);
921        assert!(
922            !events.iter().any(|e| matches!(
923                e, Event::Notification(n) if n.kind == crate::types::NotificationKind::ExtraPr
924            )),
925            "backfilling the row must not re-notify",
926        );
927    }
928
929    /// Extra-PR dedup must not be fooled by an unrelated session in another
930    /// repo already owning the `prs` row for that number (prs.id is the bare
931    /// PR number, which collides across repos) — and must not steal that row.
932    #[tokio::test]
933    async fn extra_pr_detection_survives_cross_repo_pr_number_collision() {
934        use crate::store::Store;
935
936        let sessions_dir = tempfile::tempdir().unwrap();
937        let meta = serde_json::json!({
938            "agentReportedPrs": [
939                {"number": "7", "url": "https://github.com/org/repo-a/pull/7"},
940                {"number": "9", "url": "https://github.com/org/repo-a/pull/9"},
941            ],
942        });
943        std::fs::write(
944            sessions_dir.path().join("s1.json"),
945            serde_json::to_string(&meta).unwrap(),
946        ).unwrap();
947
948        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
949        store.upsert_session(&test_session("s1", "/ws")).unwrap();
950        // Another repo's session already tracks its own PR #9.
951        let other = PR {
952            id: 9, number: 9, title: "other repo's PR".into(),
953            url: "https://github.com/org/repo-b/pull/9".into(),
954            body: String::new(), session_id: "other".into(),
955        };
956        store.upsert_pr(&other).unwrap();
957
958        let engine = Engine::new(store.clone());
959        let mut rx = engine.subscribe();
960        let poller = Poller::new(engine);
961
962        poller.sync_sessions_metadata(sessions_dir.path()).await;
963
964        let events = drain_events(&mut rx);
965        assert!(
966            events.iter().any(|e| matches!(
967                e, Event::Notification(n) if n.kind == crate::types::NotificationKind::ExtraPr
968            )),
969            "the collision must not suppress the extra-PR alert",
970        );
971        let row = store.get_pr(9).unwrap().unwrap();
972        assert_eq!(row.session_id, "other", "the other repo's row must not be stolen");
973        assert_eq!(row.url, "https://github.com/org/repo-b/pull/9");
974
975        poller.sync_sessions_metadata(sessions_dir.path()).await;
976        let events = drain_events(&mut rx);
977        assert!(
978            !events.iter().any(|e| matches!(e, Event::Notification(_))),
979            "dedup must hold across ticks even without a prs row of our own",
980        );
981    }
982
983    #[tokio::test]
984    async fn poll_usage_leaves_sessions_without_workspace_or_usage_untouched() {
985        use crate::store::Store;
986
987        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
988        let mut no_ws = test_session("no-ws", "/does/not/matter");
989        no_ws.workspace_path = None;
990        store.upsert_session(&no_ws).unwrap();
991        let engine = Engine::new(store.clone());
992        let poller = Poller::new(engine);
993
994        poller.poll_usage().await;
995
996        let unchanged = store.get_session("no-ws").unwrap().unwrap();
997        assert_eq!(unchanged.cost_usd, 0.0);
998        assert_eq!(unchanged.context_tokens, None);
999    }
1000}
ninox_core/lifecycle/poller.rs

ninox_core/lifecycle/
poller.rs
