1use std::collections::HashMap;
24
25use serde_json::Value;
26
27use crate::actions::ActionsConfig;
28use crate::state::Escalation;
29use crate::types::WorkerId;
30
31use crate::agent::{ManagerId, ManagerInstruction};
32
33#[derive(Debug, Clone)]
41pub struct ContextStore {
42 pub global: GlobalContext,
44 pub workers: HashMap<WorkerId, WorkerContext>,
46 pub managers: HashMap<ManagerId, ManagerContext>,
48 pub escalations: Vec<(WorkerId, Escalation)>,
50 pub actions: Option<ActionsConfig>,
52 pub metadata: HashMap<String, Value>,
54}
55
56impl ContextStore {
57 pub fn new(tick: u64) -> Self {
59 Self {
60 global: GlobalContext::new(tick),
61 workers: HashMap::new(),
62 managers: HashMap::new(),
63 escalations: Vec::new(),
64 actions: None,
65 metadata: HashMap::new(),
66 }
67 }
68
69 pub fn with_worker(mut self, ctx: WorkerContext) -> Self {
71 self.workers.insert(ctx.id, ctx);
72 self
73 }
74
75 pub fn with_manager(mut self, ctx: ManagerContext) -> Self {
77 self.managers.insert(ctx.id, ctx);
78 self
79 }
80
81 pub fn with_escalation(mut self, worker_id: WorkerId, escalation: Escalation) -> Self {
83 self.escalations.push((worker_id, escalation));
84 self
85 }
86
87 pub fn with_actions(mut self, actions: ActionsConfig) -> Self {
89 self.actions = Some(actions);
90 self
91 }
92
93 pub fn insert<V: Into<Value>>(mut self, key: impl Into<String>, value: V) -> Self {
95 self.metadata.insert(key.into(), value.into());
96 self
97 }
98
99 pub fn get(&self, key: &str) -> Option<&Value> {
101 self.metadata.get(key)
102 }
103
104 pub fn get_str(&self, key: &str) -> Option<&str> {
106 self.metadata.get(key).and_then(|v| v.as_str())
107 }
108}
109
110#[derive(Debug, Clone, Default)]
116pub struct GlobalContext {
117 pub tick: u64,
119 pub max_ticks: u64,
121 pub progress: f64,
123 pub success_rate: f64,
125 pub task_description: Option<String>,
127 pub hint: Option<String>,
129}
130
131impl GlobalContext {
132 pub fn new(tick: u64) -> Self {
133 Self {
134 tick,
135 ..Default::default()
136 }
137 }
138
139 pub fn with_max_ticks(mut self, max: u64) -> Self {
140 self.max_ticks = max;
141 self
142 }
143
144 pub fn with_progress(mut self, progress: f64) -> Self {
145 self.progress = progress;
146 self
147 }
148
149 pub fn with_success_rate(mut self, rate: f64) -> Self {
150 self.success_rate = rate;
151 self
152 }
153
154 pub fn with_task(mut self, description: impl Into<String>) -> Self {
155 self.task_description = Some(description.into());
156 self
157 }
158
159 pub fn with_hint(mut self, hint: impl Into<String>) -> Self {
160 self.hint = Some(hint.into());
161 self
162 }
163}
164
165#[derive(Debug, Clone)]
171pub struct WorkerContext {
172 pub id: WorkerId,
174 pub consecutive_failures: u32,
176 pub last_action: Option<String>,
178 pub last_success: Option<bool>,
180 pub history_len: usize,
182 pub has_escalation: bool,
184 pub candidates: Vec<String>,
186 pub metadata: HashMap<String, Value>,
188}
189
190impl WorkerContext {
191 pub fn new(id: WorkerId) -> Self {
192 Self {
193 id,
194 consecutive_failures: 0,
195 last_action: None,
196 last_success: None,
197 history_len: 0,
198 has_escalation: false,
199 candidates: Vec::new(),
200 metadata: HashMap::new(),
201 }
202 }
203
204 pub fn with_failures(mut self, count: u32) -> Self {
205 self.consecutive_failures = count;
206 self
207 }
208
209 pub fn with_last_action(mut self, action: impl Into<String>, success: bool) -> Self {
210 self.last_action = Some(action.into());
211 self.last_success = Some(success);
212 self
213 }
214
215 pub fn with_history_len(mut self, len: usize) -> Self {
216 self.history_len = len;
217 self
218 }
219
220 pub fn with_escalation(mut self, has: bool) -> Self {
221 self.has_escalation = has;
222 self
223 }
224
225 pub fn with_candidates(mut self, candidates: Vec<String>) -> Self {
226 self.candidates = candidates;
227 self
228 }
229}
230
231#[derive(Debug, Clone)]
237pub struct ManagerContext {
238 pub id: ManagerId,
240 pub name: String,
242 pub last_tick: u64,
244 pub metadata: HashMap<String, Value>,
246}
247
248impl ManagerContext {
249 pub fn new(id: ManagerId) -> Self {
250 Self {
251 id,
252 name: format!("Manager_{}", id.0),
253 last_tick: 0,
254 metadata: HashMap::new(),
255 }
256 }
257
258 pub fn with_name(mut self, name: impl Into<String>) -> Self {
259 self.name = name.into();
260 self
261 }
262
263 pub fn with_last_tick(mut self, tick: u64) -> Self {
264 self.last_tick = tick;
265 self
266 }
267}
268
269#[derive(Debug, Clone)]
275pub enum ContextView {
276 Global { manager_id: ManagerId },
278 Local {
280 worker_id: WorkerId,
281 neighbor_ids: Vec<WorkerId>,
283 },
284 Custom {
286 name: String,
288 visible_worker_ids: Vec<WorkerId>,
290 visible_manager_ids: Vec<ManagerId>,
292 },
293}
294
295impl ContextView {
296 pub fn global(manager_id: ManagerId) -> Self {
298 Self::Global { manager_id }
299 }
300
301 pub fn local(worker_id: WorkerId) -> Self {
303 Self::Local {
304 worker_id,
305 neighbor_ids: Vec::new(),
306 }
307 }
308
309 pub fn local_with_neighbors(worker_id: WorkerId, neighbor_ids: Vec<WorkerId>) -> Self {
311 Self::Local {
312 worker_id,
313 neighbor_ids,
314 }
315 }
316
317 pub fn custom(
319 name: impl Into<String>,
320 visible_workers: Vec<WorkerId>,
321 visible_managers: Vec<ManagerId>,
322 ) -> Self {
323 Self::Custom {
324 name: name.into(),
325 visible_worker_ids: visible_workers,
326 visible_manager_ids: visible_managers,
327 }
328 }
329}
330
331#[derive(Debug, Clone)]
337pub struct ActionParam {
338 pub name: String,
340 pub description: String,
342 pub required: bool,
344}
345
346#[derive(Debug, Clone)]
351pub struct ActionCandidate {
352 pub name: String,
354 pub description: String,
356 pub params: Vec<ActionParam>,
358 pub example: Option<String>,
360}
361
362impl ActionCandidate {
363 pub fn from_config(config: &ActionsConfig) -> Vec<Self> {
365 config
366 .all_actions()
367 .map(|def| ActionCandidate {
368 name: def.name.clone(),
369 description: def.description.clone(),
370 params: def
371 .params
372 .iter()
373 .map(|p| ActionParam {
374 name: p.name.clone(),
375 description: p.description.clone(),
376 required: p.required,
377 })
378 .collect(),
379 example: def.example.clone(),
380 })
381 .collect()
382 }
383}
384
385#[derive(Debug, Clone)]
394pub struct ResolvedContext {
395 pub global: GlobalContext,
397 pub visible_workers: Vec<WorkerContext>,
399 pub escalations: Vec<(WorkerId, Escalation)>,
401 pub candidates: Vec<ActionCandidate>,
403 pub metadata: HashMap<String, Value>,
405 pub target: ContextTarget,
407 pub self_last_output: Option<String>,
412 pub manager_instruction: Option<ManagerInstruction>,
417}
418
419#[derive(Debug, Clone)]
421pub enum ContextTarget {
422 Manager(ManagerId),
423 Worker(WorkerId),
424}
425
426impl ResolvedContext {
427 pub fn new(global: GlobalContext, target: ContextTarget) -> Self {
429 Self {
430 global,
431 visible_workers: Vec::new(),
432 escalations: Vec::new(),
433 candidates: Vec::new(),
434 metadata: HashMap::new(),
435 target,
436 self_last_output: None,
437 manager_instruction: None,
438 }
439 }
440
441 pub fn with_workers(mut self, workers: Vec<WorkerContext>) -> Self {
443 self.visible_workers = workers;
444 self
445 }
446
447 pub fn with_escalations(mut self, escalations: Vec<(WorkerId, Escalation)>) -> Self {
449 self.escalations = escalations;
450 self
451 }
452
453 pub fn with_candidates(mut self, candidates: Vec<ActionCandidate>) -> Self {
455 self.candidates = candidates;
456 self
457 }
458
459 pub fn with_actions_config(mut self, config: &ActionsConfig) -> Self {
461 self.candidates = ActionCandidate::from_config(config);
462 self
463 }
464
465 pub fn with_metadata(mut self, metadata: HashMap<String, Value>) -> Self {
467 self.metadata = metadata;
468 self
469 }
470
471 pub fn with_self_last_output(mut self, output: Option<String>) -> Self {
473 self.self_last_output = output;
474 self
475 }
476
477 pub fn with_manager_instruction(mut self, instruction: ManagerInstruction) -> Self {
479 self.manager_instruction = Some(instruction);
480 self
481 }
482
483 pub fn has_escalations(&self) -> bool {
485 !self.escalations.is_empty()
486 }
487
488 pub fn is_manager(&self) -> bool {
490 matches!(self.target, ContextTarget::Manager(_))
491 }
492
493 pub fn is_worker(&self) -> bool {
495 matches!(self.target, ContextTarget::Worker(_))
496 }
497}
498
499#[cfg(test)]
504mod tests {
505 use super::*;
506
507 #[test]
508 fn test_context_store_builder() {
509 let store = ContextStore::new(10)
510 .with_worker(WorkerContext::new(WorkerId(0)))
511 .with_worker(WorkerContext::new(WorkerId(1)))
512 .with_manager(ManagerContext::new(ManagerId(0)))
513 .insert("task", "Find the bug");
514
515 assert_eq!(store.global.tick, 10);
516 assert_eq!(store.workers.len(), 2);
517 assert_eq!(store.managers.len(), 1);
518 assert_eq!(store.get_str("task"), Some("Find the bug"));
519 }
520
521 #[test]
522 fn test_context_view_creation() {
523 let global = ContextView::global(ManagerId(0));
524 assert!(matches!(global, ContextView::Global { .. }));
525
526 let local = ContextView::local(WorkerId(0));
527 assert!(matches!(local, ContextView::Local { .. }));
528
529 let local_with_neighbors =
530 ContextView::local_with_neighbors(WorkerId(0), vec![WorkerId(1), WorkerId(2)]);
531 if let ContextView::Local { neighbor_ids, .. } = local_with_neighbors {
532 assert_eq!(neighbor_ids.len(), 2);
533 }
534 }
535
536 #[test]
537 fn test_worker_context_builder() {
538 let ctx = WorkerContext::new(WorkerId(0))
539 .with_failures(2)
540 .with_last_action("read:/path", true)
541 .with_history_len(10)
542 .with_escalation(true)
543 .with_candidates(vec!["read".into(), "grep".into()]);
544
545 assert_eq!(ctx.id, WorkerId(0));
546 assert_eq!(ctx.consecutive_failures, 2);
547 assert_eq!(ctx.last_action, Some("read:/path".to_string()));
548 assert!(ctx.has_escalation);
549 assert_eq!(ctx.candidates.len(), 2);
550 }
551
552 #[test]
553 fn test_resolved_context() {
554 let global = GlobalContext::new(5)
555 .with_progress(0.5)
556 .with_task("Test task");
557
558 let candidates = vec![ActionCandidate {
559 name: "action1".to_string(),
560 description: "Test action".to_string(),
561 params: vec![],
562 example: None,
563 }];
564
565 let resolved = ResolvedContext::new(global, ContextTarget::Worker(WorkerId(0)))
566 .with_workers(vec![WorkerContext::new(WorkerId(0))])
567 .with_candidates(candidates);
568
569 assert!(resolved.is_worker());
570 assert!(!resolved.is_manager());
571 assert_eq!(resolved.visible_workers.len(), 1);
572 assert_eq!(resolved.candidates.len(), 1);
573 }
574}