1use std::collections::{HashMap, HashSet};
4
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7
8use crate::issue::Issue;
9use crate::session::RetryEntry;
10
11#[derive(Debug, Serialize, Deserialize)]
13pub struct OrchestratorState {
14 pub poll_interval_ms: u64,
15 pub max_concurrent_agents: u32,
16 pub running: HashMap<String, RunningEntry>,
17 pub claimed: HashSet<String>,
18 pub retry_attempts: HashMap<String, RetryEntry>,
19 pub completed: HashSet<String>,
21 pub codex_totals: CodexTotals,
22 pub codex_rate_limits: Option<serde_json::Value>,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct RunningEntry {
28 pub identifier: String,
29 pub issue: Issue,
30 pub session_id: Option<String>,
31 pub codex_app_server_pid: Option<String>,
32 pub last_codex_message: Option<String>,
33 pub last_codex_event: Option<String>,
34 pub last_codex_timestamp: Option<DateTime<Utc>>,
35 pub codex_input_tokens: u64,
36 pub codex_output_tokens: u64,
37 pub codex_total_tokens: u64,
38 pub last_reported_input_tokens: u64,
39 pub last_reported_output_tokens: u64,
40 pub last_reported_total_tokens: u64,
41 pub retry_attempt: Option<u32>,
42 pub started_at: DateTime<Utc>,
43 pub turn_count: u32,
44}
45
46#[derive(Debug, Default, Clone, Serialize, Deserialize)]
48pub struct CodexTotals {
49 pub input_tokens: u64,
50 pub output_tokens: u64,
51 pub total_tokens: u64,
52 pub seconds_running: f64,
53}
54
55impl OrchestratorState {
56 pub fn new(poll_interval_ms: u64, max_concurrent_agents: u32) -> Self {
57 Self {
58 poll_interval_ms,
59 max_concurrent_agents,
60 running: HashMap::new(),
61 claimed: HashSet::new(),
62 retry_attempts: HashMap::new(),
63 completed: HashSet::new(),
64 codex_totals: CodexTotals::default(),
65 codex_rate_limits: None,
66 }
67 }
68
69 pub fn available_slots(&self) -> u32 {
71 self.max_concurrent_agents
72 .saturating_sub(self.running.len() as u32)
73 }
74
75 pub fn is_claimed(&self, issue_id: &str) -> bool {
77 self.claimed.contains(issue_id)
78 }
79}
80
81#[cfg(test)]
82mod tests {
83 use super::*;
84
85 #[test]
86 fn new_state_has_full_slots() {
87 let state = OrchestratorState::new(30000, 10);
88 assert_eq!(state.available_slots(), 10);
89 }
90
91 #[test]
92 fn available_slots_decrements() {
93 let mut state = OrchestratorState::new(30000, 2);
94 state.running.insert(
95 "issue-1".into(),
96 RunningEntry {
97 identifier: "T-1".into(),
98 issue: crate::issue::Issue {
99 id: "issue-1".into(),
100 identifier: "T-1".into(),
101 title: "Test".into(),
102 description: None,
103 priority: None,
104 state: "Todo".into(),
105 branch_name: None,
106 url: None,
107 labels: vec![],
108 blocked_by: vec![],
109 created_at: None,
110 updated_at: None,
111 },
112 session_id: None,
113 codex_app_server_pid: None,
114 last_codex_message: None,
115 last_codex_event: None,
116 last_codex_timestamp: None,
117 codex_input_tokens: 0,
118 codex_output_tokens: 0,
119 codex_total_tokens: 0,
120 last_reported_input_tokens: 0,
121 last_reported_output_tokens: 0,
122 last_reported_total_tokens: 0,
123 retry_attempt: None,
124 started_at: chrono::Utc::now(),
125 turn_count: 0,
126 },
127 );
128 assert_eq!(state.available_slots(), 1);
129 }
130}