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

ninox_core/lifecycle/
poller.rs
