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) {
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) {
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) {
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) {
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) {
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,
555        }
556    }
557
558    /// End-to-end (within-process) proof that the poller closes the gap
559    /// documented in `lifecycle::usage`: given a workspace whose `claude`
560    /// transcript directory has usage recorded, `poll_usage` writes the
561    /// derived cost/context/model back into the store and emits
562    /// `SessionUpdated` — the exact path the UI's $0.0000 / missing-tokens
563    /// symptom traces back to when this ingestion doesn't happen.
564    // The `ENV_TEST_GUARD` mutex is intentionally held across the `.await`
565    // points below — it serializes access to the process-global
566    // `NINOX_CLAUDE_PROJECTS_DIR` env var against other tests (in this file
567    // and in `lifecycle::usage`) for this single-threaded `#[tokio::test]`,
568    // and must stay held for the env var's entire lifetime, not just around
569    // the sync portions.
570    #[allow(clippy::await_holding_lock)]
571    #[tokio::test]
572    async fn poll_usage_ingests_transcript_into_store_and_emits_update() {
573        use crate::{lifecycle::usage::{claude_project_slug, ENV_TEST_GUARD}, store::Store};
574        use std::io::Write;
575
576        let _guard = ENV_TEST_GUARD.lock().unwrap_or_else(|e| e.into_inner());
577        let projects_dir = tempfile::tempdir().unwrap();
578        let workspace = "/tmp/poller-usage-probe-workspace";
579        let project_dir = projects_dir.path().join(claude_project_slug(workspace));
580        std::fs::create_dir_all(&project_dir).unwrap();
581        let mut f = std::fs::File::create(project_dir.join("s.jsonl")).unwrap();
582        writeln!(
583            f,
584            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}}}}}}"#
585        ).unwrap();
586        drop(f);
587
588        let prior = std::env::var("NINOX_CLAUDE_PROJECTS_DIR").ok();
589        std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", projects_dir.path());
590
591        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
592        store.upsert_session(&test_session("s1", workspace)).unwrap();
593        let engine = Engine::new(store.clone());
594        let mut rx = engine.subscribe();
595        let poller = Poller::new(engine);
596
597        poller.poll_usage().await;
598
599        match prior {
600            Some(v) => std::env::set_var("NINOX_CLAUDE_PROJECTS_DIR", v),
601            None    => std::env::remove_var("NINOX_CLAUDE_PROJECTS_DIR"),
602        }
603
604        let updated = store.get_session("s1").unwrap().unwrap();
605        assert!(updated.cost_usd > 0.0, "cost_usd should be ingested, not 0.0000");
606        assert_eq!(updated.context_tokens, Some(2 + 500 + 45000));
607        assert_eq!(updated.model.as_deref(), Some("claude-fable-5"));
608
609        let evt = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv())
610            .await
611            .expect("SessionUpdated should be emitted")
612            .unwrap();
613        assert!(matches!(evt, Event::SessionUpdated(s) if s.id == "s1" && s.cost_usd > 0.0));
614    }
615
616    /// Drain every event currently buffered on the receiver.
617    fn drain_events(rx: &mut tokio::sync::broadcast::Receiver<Event>) -> Vec<Event> {
618        let mut events = Vec::new();
619        while let Ok(e) = rx.try_recv() {
620            events.push(e);
621        }
622        events
623    }
624
625    /// A worker that opened three PRs: the first becomes the session's
626    /// tracked PR, every later one is recorded in the store and raised as an
627    /// ExtraPr notification — and only once, however often the poller ticks.
628    #[tokio::test]
629    async fn metadata_sync_adopts_first_pr_and_flags_every_extra_once() {
630        use crate::store::Store;
631
632        let sessions_dir = tempfile::tempdir().unwrap();
633        let meta = serde_json::json!({
634            "agentReportedPrNumber": "44",
635            "agentReportedPrUrl": "https://github.com/org/repo/pull/44",
636            "agentReportedPrs": [
637                {"number": "42", "url": "https://github.com/org/repo/pull/42"},
638                {"number": "43", "url": "https://github.com/org/repo/pull/43"},
639                {"number": "44", "url": "https://github.com/org/repo/pull/44"},
640            ],
641        });
642        std::fs::write(
643            sessions_dir.path().join("s1.json"),
644            serde_json::to_string(&meta).unwrap(),
645        ).unwrap();
646
647        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
648        store.upsert_session(&test_session("s1", "/ws")).unwrap();
649        let engine = Engine::new(store.clone());
650        let mut rx = engine.subscribe();
651        let poller = Poller::new(engine);
652
653        poller.sync_sessions_metadata(sessions_dir.path()).await;
654
655        let session = store.get_session("s1").unwrap().unwrap();
656        assert_eq!(session.pr_number, Some(42), "first PR is the canonical one");
657        assert!(matches!(session.status, SessionStatus::PrOpen));
658        assert!(store.get_pr(43).unwrap().is_some(), "extra PR #43 recorded");
659        assert!(store.get_pr(44).unwrap().is_some(), "extra PR #44 recorded");
660        assert_eq!(
661            store.get_pr(43).unwrap().unwrap().url,
662            "https://github.com/org/repo/pull/43",
663        );
664
665        let events = drain_events(&mut rx);
666        let extra_notifs: Vec<_> = events.iter().filter(|e| matches!(
667            e, Event::Notification(n) if n.kind == crate::types::NotificationKind::ExtraPr
668        )).collect();
669        assert_eq!(extra_notifs.len(), 2, "one ExtraPr notification per extra PR");
670
671        // Second tick: nothing new — no duplicate notifications.
672        poller.sync_sessions_metadata(sessions_dir.path()).await;
673        let events = drain_events(&mut rx);
674        assert!(
675            !events.iter().any(|e| matches!(e, Event::Notification(_))),
676            "extra PRs must not be re-notified on every tick",
677        );
678    }
679
680    /// A single reported PR (the normal case) adopts it with no extra-PR
681    /// noise — the pre-existing first-PR-detection behavior.
682    #[tokio::test]
683    async fn metadata_sync_single_pr_has_no_extra_notifications() {
684        use crate::store::Store;
685
686        let sessions_dir = tempfile::tempdir().unwrap();
687        let meta = serde_json::json!({
688            "agentReportedPrNumber": "5",
689            "agentReportedPrUrl": "https://github.com/org/repo/pull/5",
690        });
691        std::fs::write(
692            sessions_dir.path().join("s1.json"),
693            serde_json::to_string(&meta).unwrap(),
694        ).unwrap();
695
696        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
697        store.upsert_session(&test_session("s1", "/ws")).unwrap();
698        let engine = Engine::new(store.clone());
699        let mut rx = engine.subscribe();
700        let poller = Poller::new(engine);
701
702        poller.sync_sessions_metadata(sessions_dir.path()).await;
703
704        let session = store.get_session("s1").unwrap().unwrap();
705        assert_eq!(session.pr_number, Some(5));
706        assert!(matches!(session.status, SessionStatus::PrOpen));
707        let events = drain_events(&mut rx);
708        assert!(!events.iter().any(|e| matches!(e, Event::Notification(_))));
709    }
710
711    /// Work requests recorded by `ninox request-work` surface exactly one
712    /// WorkRequested notification each, then are marked delivered.
713    #[tokio::test]
714    async fn metadata_sync_delivers_work_requests_exactly_once() {
715        use crate::store::Store;
716
717        let sessions_dir = tempfile::tempdir().unwrap();
718        hooks::append_work_request(sessions_dir.path(), "s1", "Migrate the config loader").unwrap();
719
720        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
721        store.upsert_session(&test_session("s1", "/ws")).unwrap();
722        let engine = Engine::new(store.clone());
723        let mut rx = engine.subscribe();
724        let poller = Poller::new(engine);
725
726        poller.sync_sessions_metadata(sessions_dir.path()).await;
727
728        let events = drain_events(&mut rx);
729        let notif = events.iter().find_map(|e| match e {
730            Event::Notification(n) if n.kind == crate::types::NotificationKind::WorkRequested => Some(n),
731            _ => None,
732        }).expect("WorkRequested notification emitted");
733        assert!(notif.body.contains("Migrate the config loader"));
734        assert_eq!(notif.session_id.as_deref(), Some("s1"));
735
736        assert!(
737            hooks::read_pending_work_requests(sessions_dir.path(), "s1").unwrap().is_empty(),
738            "delivered requests must leave the pending set",
739        );
740
741        poller.sync_sessions_metadata(sessions_dir.path()).await;
742        let events = drain_events(&mut rx);
743        assert!(
744            !events.iter().any(|e| matches!(e, Event::Notification(_))),
745            "delivered work requests must not fire again",
746        );
747    }
748
749    /// A worker can request work and exit before the next tick — the request
750    /// must still reach the orchestrator, not die with the session.
751    #[tokio::test]
752    async fn metadata_sync_delivers_work_requests_from_terminated_sessions() {
753        use crate::store::Store;
754
755        let sessions_dir = tempfile::tempdir().unwrap();
756        hooks::append_work_request(sessions_dir.path(), "s1", "Follow-up refactor").unwrap();
757
758        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
759        let mut session = test_session("s1", "/ws");
760        session.status = SessionStatus::Terminated;
761        store.upsert_session(&session).unwrap();
762        let engine = Engine::new(store.clone());
763        let mut rx = engine.subscribe();
764        let poller = Poller::new(engine);
765
766        poller.sync_sessions_metadata(sessions_dir.path()).await;
767
768        let events = drain_events(&mut rx);
769        assert!(
770            events.iter().any(|e| matches!(
771                e, Event::Notification(n) if n.kind == crate::types::NotificationKind::WorkRequested
772            )),
773            "work requests outlive their session",
774        );
775    }
776
777    /// The ledger row for an extra PR is best-effort at notification time (a
778    /// busy store must not kill the alert) — but it must self-heal on later
779    /// ticks rather than be lost forever, and healing must not re-notify.
780    #[tokio::test]
781    async fn extra_pr_ledger_row_backfills_after_notification_without_renotifying() {
782        use crate::store::Store;
783
784        let sessions_dir = tempfile::tempdir().unwrap();
785        let meta = serde_json::json!({
786            "agentReportedPrs": [
787                {"number": "7", "url": "https://github.com/org/repo/pull/7"},
788                {"number": "9", "url": "https://github.com/org/repo/pull/9"},
789            ],
790        });
791        std::fs::write(
792            sessions_dir.path().join("s1.json"),
793            serde_json::to_string(&meta).unwrap(),
794        ).unwrap();
795        // Simulate "notified previously, but the row write failed that tick".
796        hooks::mark_extra_prs_notified(sessions_dir.path(), "s1", &[9]).unwrap();
797
798        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
799        store.upsert_session(&test_session("s1", "/ws")).unwrap();
800        let engine = Engine::new(store.clone());
801        let mut rx = engine.subscribe();
802        let poller = Poller::new(engine);
803
804        poller.sync_sessions_metadata(sessions_dir.path()).await;
805
806        assert!(
807            store.get_pr(9).unwrap().is_some(),
808            "already-notified extra PR must still get its ledger row backfilled",
809        );
810        let events = drain_events(&mut rx);
811        assert!(
812            !events.iter().any(|e| matches!(
813                e, Event::Notification(n) if n.kind == crate::types::NotificationKind::ExtraPr
814            )),
815            "backfilling the row must not re-notify",
816        );
817    }
818
819    /// Extra-PR dedup must not be fooled by an unrelated session in another
820    /// repo already owning the `prs` row for that number (prs.id is the bare
821    /// PR number, which collides across repos) — and must not steal that row.
822    #[tokio::test]
823    async fn extra_pr_detection_survives_cross_repo_pr_number_collision() {
824        use crate::store::Store;
825
826        let sessions_dir = tempfile::tempdir().unwrap();
827        let meta = serde_json::json!({
828            "agentReportedPrs": [
829                {"number": "7", "url": "https://github.com/org/repo-a/pull/7"},
830                {"number": "9", "url": "https://github.com/org/repo-a/pull/9"},
831            ],
832        });
833        std::fs::write(
834            sessions_dir.path().join("s1.json"),
835            serde_json::to_string(&meta).unwrap(),
836        ).unwrap();
837
838        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
839        store.upsert_session(&test_session("s1", "/ws")).unwrap();
840        // Another repo's session already tracks its own PR #9.
841        let other = PR {
842            id: 9, number: 9, title: "other repo's PR".into(),
843            url: "https://github.com/org/repo-b/pull/9".into(),
844            body: String::new(), session_id: "other".into(),
845        };
846        store.upsert_pr(&other).unwrap();
847
848        let engine = Engine::new(store.clone());
849        let mut rx = engine.subscribe();
850        let poller = Poller::new(engine);
851
852        poller.sync_sessions_metadata(sessions_dir.path()).await;
853
854        let events = drain_events(&mut rx);
855        assert!(
856            events.iter().any(|e| matches!(
857                e, Event::Notification(n) if n.kind == crate::types::NotificationKind::ExtraPr
858            )),
859            "the collision must not suppress the extra-PR alert",
860        );
861        let row = store.get_pr(9).unwrap().unwrap();
862        assert_eq!(row.session_id, "other", "the other repo's row must not be stolen");
863        assert_eq!(row.url, "https://github.com/org/repo-b/pull/9");
864
865        poller.sync_sessions_metadata(sessions_dir.path()).await;
866        let events = drain_events(&mut rx);
867        assert!(
868            !events.iter().any(|e| matches!(e, Event::Notification(_))),
869            "dedup must hold across ticks even without a prs row of our own",
870        );
871    }
872
873    #[tokio::test]
874    async fn poll_usage_leaves_sessions_without_workspace_or_usage_untouched() {
875        use crate::store::Store;
876
877        let store = std::sync::Arc::new(Store::open(tempfile::tempdir().unwrap().keep().join("t.db")).unwrap());
878        let mut no_ws = test_session("no-ws", "/does/not/matter");
879        no_ws.workspace_path = None;
880        store.upsert_session(&no_ws).unwrap();
881        let engine = Engine::new(store.clone());
882        let poller = Poller::new(engine);
883
884        poller.poll_usage().await;
885
886        let unchanged = store.get_session("no-ws").unwrap().unwrap();
887        assert_eq!(unchanged.cost_usd, 0.0);
888        assert_eq!(unchanged.context_tokens, None);
889    }
890}
ninox_core/lifecycle/poller.rs

ninox_core/lifecycle/
poller.rs
