swarm_engine_core/context/
summary.rs

1//! TaskContext - Analyzer が生成するタスク状況
2//!
3//! # 設計
4//!
5//! Analyzer が SwarmState を分析して TaskContext を生成。
6//! Manager は TaskContext を見て Request を生成する。
7//!
8//! ```text
9//! SwarmState → [Analyzer] → TaskContext → [Manager] → BatchDecisionRequest
10//! ```
11//!
12//! # 拡張性
13//!
14//! ベース情報（tick, workers, success_rate 等）に加え、
15//! `metadata: HashMap<String, Value>` で任意の追加情報を格納可能。
16//! 軽量LLMで分析した結果等を入れることを想定。
17
18use std::collections::{HashMap, HashSet};
19use std::sync::Arc;
20
21use serde_json::Value;
22
23use crate::actions::ActionsConfig;
24use crate::agent::Guidance;
25use crate::state::Escalation;
26use crate::types::WorkerId;
27
28// ============================================================================
29// WorkerSummary - Worker 状態の要約
30// ============================================================================
31
32/// Worker 状態の要約
33#[derive(Debug, Clone)]
34pub struct WorkerSummary {
35    /// Worker ID
36    pub id: WorkerId,
37    /// 連続失敗数
38    pub consecutive_failures: u32,
39    /// 最新アクション名
40    pub last_action: Option<String>,
41    /// 最新アクションの成功/失敗
42    pub last_success: Option<bool>,
43    /// 最新アクションの出力（Environment からの結果）
44    pub last_output: Option<String>,
45    /// 履歴の長さ
46    pub history_len: usize,
47    /// Escalation 中かどうか
48    pub has_escalation: bool,
49}
50
51impl WorkerSummary {
52    pub fn new(id: WorkerId) -> Self {
53        Self {
54            id,
55            consecutive_failures: 0,
56            last_action: None,
57            last_success: None,
58            last_output: None,
59            history_len: 0,
60            has_escalation: false,
61        }
62    }
63
64    /// 連続失敗数を設定
65    pub fn with_failures(mut self, count: u32) -> Self {
66        self.consecutive_failures = count;
67        self
68    }
69
70    /// 最新アクションを設定
71    pub fn with_last_action(mut self, action: impl Into<String>, success: bool) -> Self {
72        self.last_action = Some(action.into());
73        self.last_success = Some(success);
74        self
75    }
76
77    /// 履歴の長さを設定
78    pub fn with_history_len(mut self, len: usize) -> Self {
79        self.history_len = len;
80        self
81    }
82
83    /// Escalation を設定
84    pub fn with_escalation(mut self, has_escalation: bool) -> Self {
85        self.has_escalation = has_escalation;
86        self
87    }
88}
89
90// ============================================================================
91// TaskContext - タスク状況
92// ============================================================================
93
94/// タスク状況（Analyzer が生成、Manager が消費）
95///
96/// # 構成
97///
98/// - **ベース情報**: tick, workers, success_rate, progress, escalations, available_actions
99/// - **探索情報**: v2_guidances, excluded_actions
100/// - **拡張 KV**: metadata で任意の追加情報を格納
101///
102/// # 使用例
103///
104/// ```ignore
105/// let context = TaskContext::new(tick)
106///     .with_worker(WorkerSummary::new(WorkerId(0)).with_phase(WorkerPhase::Working))
107///     .with_progress(0.5)
108///     .insert("llm_summary", "探索フェーズ中");
109/// ```
110#[derive(Debug, Clone)]
111pub struct TaskContext {
112    // === ベース情報 ===
113    /// 現在の tick
114    pub tick: u64,
115    /// 各 Worker の状態サマリ
116    pub workers: HashMap<WorkerId, WorkerSummary>,
117    /// 成功率 (0.0 - 1.0)
118    pub success_rate: f64,
119    /// 進捗 (0.0 - 1.0)
120    pub progress: f64,
121    /// Escalation 一覧（WorkerId, Escalation）
122    pub escalations: Vec<(WorkerId, Escalation)>,
123    /// 利用可能なアクション
124    pub available_actions: Option<ActionsConfig>,
125
126    // === 探索情報 ===
127    /// ExplorationSpaceV2 から生成された Guidance
128    ///
129    /// select_nodes() → Guidance 変換で直接生成。
130    /// Manager はこれをそのまま Worker に配布できる。
131    pub v2_guidances: Option<Vec<Guidance>>,
132    /// 除外すべきアクション（成功済み/クローズ済み）- プロンプトから除外
133    pub excluded_actions: Vec<String>,
134
135    // === Manager 指示情報 ===
136    /// 前回の Guidance（Worker ID -> Arc<Guidance>）
137    ///
138    /// Orchestrator が Manager.prepare() 呼び出し前に設定。
139    /// Manager は prepare() でこれを使って ResolvedContext に ManagerInstruction を埋め込む。
140    /// Arc で共有することでクローン時のディープコピーを回避。
141    pub previous_guidances: HashMap<WorkerId, Arc<Guidance>>,
142
143    /// Goal Action（Terminal Action）を達成した Worker
144    /// これらのWorkerは Manager.prepare() でリクエスト対象から除外される
145    pub done_workers: HashSet<WorkerId>,
146
147    // === 拡張用 KV ===
148    /// 追加メタデータ（軽量LLM分析結果等）
149    pub metadata: HashMap<String, Value>,
150}
151
152impl TaskContext {
153    /// 新しい TaskContext を作成
154    pub fn new(tick: u64) -> Self {
155        Self {
156            tick,
157            workers: HashMap::new(),
158            success_rate: 0.0,
159            progress: 0.0,
160            escalations: Vec::new(),
161            available_actions: None,
162            v2_guidances: None,
163            excluded_actions: Vec::new(),
164            previous_guidances: HashMap::new(),
165            done_workers: HashSet::new(),
166            metadata: HashMap::new(),
167        }
168    }
169
170    // === Builder methods ===
171
172    /// Worker サマリを追加
173    pub fn with_worker(mut self, summary: WorkerSummary) -> Self {
174        self.workers.insert(summary.id, summary);
175        self
176    }
177
178    /// 成功率を設定
179    pub fn with_success_rate(mut self, rate: f64) -> Self {
180        self.success_rate = rate;
181        self
182    }
183
184    /// 進捗を設定
185    pub fn with_progress(mut self, progress: f64) -> Self {
186        self.progress = progress;
187        self
188    }
189
190    /// Escalation を追加
191    pub fn with_escalation(mut self, worker_id: WorkerId, escalation: Escalation) -> Self {
192        self.escalations.push((worker_id, escalation));
193        self
194    }
195
196    /// 利用可能なアクションを設定
197    pub fn with_actions(mut self, actions: ActionsConfig) -> Self {
198        self.available_actions = Some(actions);
199        self
200    }
201
202    /// 前回の Guidance を設定（Arc で共有）
203    pub fn with_previous_guidances(mut self, guidances: HashMap<WorkerId, Arc<Guidance>>) -> Self {
204        self.previous_guidances = guidances;
205        self
206    }
207
208    /// 単一 Worker の前回 Guidance を追加
209    pub fn with_previous_guidance(mut self, worker_id: WorkerId, guidance: Arc<Guidance>) -> Self {
210        self.previous_guidances.insert(worker_id, guidance);
211        self
212    }
213
214    // === Metadata (KV) methods ===
215
216    /// メタデータを追加
217    pub fn insert<V: Into<Value>>(mut self, key: impl Into<String>, value: V) -> Self {
218        self.metadata.insert(key.into(), value.into());
219        self
220    }
221
222    /// メタデータを追加（mutable）
223    pub fn set<V: Into<Value>>(&mut self, key: impl Into<String>, value: V) {
224        self.metadata.insert(key.into(), value.into());
225    }
226
227    /// メタデータを取得
228    pub fn get(&self, key: &str) -> Option<&Value> {
229        self.metadata.get(key)
230    }
231
232    /// メタデータを文字列として取得
233    pub fn get_str(&self, key: &str) -> Option<&str> {
234        self.metadata.get(key).and_then(|v| v.as_str())
235    }
236
237    /// メタデータを数値として取得
238    pub fn get_f64(&self, key: &str) -> Option<f64> {
239        self.metadata.get(key).and_then(|v| v.as_f64())
240    }
241
242    /// メタデータを整数として取得
243    pub fn get_i64(&self, key: &str) -> Option<i64> {
244        self.metadata.get(key).and_then(|v| v.as_i64())
245    }
246
247    /// メタデータを真偽値として取得
248    pub fn get_bool(&self, key: &str) -> Option<bool> {
249        self.metadata.get(key).and_then(|v| v.as_bool())
250    }
251
252    // === Query methods ===
253
254    /// Escalation があるか
255    pub fn has_escalations(&self) -> bool {
256        !self.escalations.is_empty()
257    }
258
259    /// 特定の Worker が Escalation 中かどうか
260    pub fn has_escalation_for(&self, worker_id: WorkerId) -> bool {
261        self.escalations.iter().any(|(id, _)| *id == worker_id)
262    }
263
264    /// 特定 Worker の情報を取得
265    pub fn worker(&self, id: WorkerId) -> Option<&WorkerSummary> {
266        self.workers.get(&id)
267    }
268
269    /// Escalation 中の Worker 数を取得
270    pub fn escalated_worker_count(&self) -> usize {
271        self.workers.values().filter(|w| w.has_escalation).count()
272    }
273
274    /// Worker ID 一覧を取得
275    pub fn worker_ids(&self) -> Vec<WorkerId> {
276        self.workers.keys().copied().collect()
277    }
278}
279
280impl TaskContext {
281    // === Exploration methods ===
282
283    /// 探索が有効か（v2_guidances がある場合）
284    pub fn has_exploration(&self) -> bool {
285        self.v2_guidances.is_some()
286    }
287
288    /// 指定した Worker のみを含むフィルタ済み TaskContext を作成
289    ///
290    /// Manager のパーティショニングで使用。
291    /// 各 Manager は担当する Worker のみの TaskContext を受け取る。
292    pub fn filter_for_workers(&self, worker_ids: &[WorkerId]) -> TaskContext {
293        use std::collections::HashSet;
294        let worker_set: HashSet<WorkerId> = worker_ids.iter().copied().collect();
295
296        // フィルタ済み workers
297        let filtered_workers: HashMap<WorkerId, WorkerSummary> = self
298            .workers
299            .iter()
300            .filter(|(id, _)| worker_set.contains(id))
301            .map(|(id, summary)| (*id, summary.clone()))
302            .collect();
303
304        // フィルタ済み escalations
305        let filtered_escalations: Vec<(WorkerId, Escalation)> = self
306            .escalations
307            .iter()
308            .filter(|(id, _)| worker_set.contains(id))
309            .cloned()
310            .collect();
311
312        // フィルタ済み previous_guidances（Arc::clone は軽量）
313        let filtered_guidances: HashMap<WorkerId, Arc<Guidance>> = self
314            .previous_guidances
315            .iter()
316            .filter(|(id, _)| worker_set.contains(id))
317            .map(|(id, g)| (*id, Arc::clone(g)))
318            .collect();
319
320        // フィルタ済み done_workers
321        let filtered_done_workers: HashSet<WorkerId> = self
322            .done_workers
323            .iter()
324            .filter(|id| worker_set.contains(id))
325            .copied()
326            .collect();
327
328        TaskContext {
329            tick: self.tick,
330            workers: filtered_workers,
331            success_rate: self.success_rate,
332            progress: self.progress,
333            escalations: filtered_escalations,
334            available_actions: self.available_actions.clone(),
335            v2_guidances: self.v2_guidances.clone(),
336            excluded_actions: self.excluded_actions.clone(),
337            previous_guidances: filtered_guidances,
338            done_workers: filtered_done_workers,
339            metadata: self.metadata.clone(),
340        }
341    }
342}
343
344impl Default for TaskContext {
345    fn default() -> Self {
346        Self::new(0)
347    }
348}
349
350// ============================================================================
351// Tests
352// ============================================================================
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357
358    #[test]
359    fn test_task_context_new() {
360        let ctx = TaskContext::new(10);
361        assert_eq!(ctx.tick, 10);
362        assert!(ctx.workers.is_empty());
363        assert_eq!(ctx.success_rate, 0.0);
364        assert_eq!(ctx.progress, 0.0);
365    }
366
367    #[test]
368    fn test_task_context_builder() {
369        let ctx = TaskContext::new(5)
370            .with_worker(WorkerSummary::new(WorkerId(0)))
371            .with_worker(WorkerSummary::new(WorkerId(1)).with_escalation(true))
372            .with_success_rate(0.8)
373            .with_progress(0.5)
374            .insert("key1", "value1")
375            .insert("count", 42);
376
377        assert_eq!(ctx.tick, 5);
378        assert_eq!(ctx.workers.len(), 2);
379        assert_eq!(ctx.success_rate, 0.8);
380        assert_eq!(ctx.progress, 0.5);
381        assert_eq!(ctx.get_str("key1"), Some("value1"));
382        assert_eq!(ctx.get_i64("count"), Some(42));
383    }
384
385    #[test]
386    fn test_worker_summary() {
387        let summary = WorkerSummary::new(WorkerId(0))
388            .with_failures(2)
389            .with_last_action("read:/path", true)
390            .with_history_len(10)
391            .with_escalation(true);
392
393        assert_eq!(summary.id, WorkerId(0));
394        assert_eq!(summary.consecutive_failures, 2);
395        assert_eq!(summary.last_action, Some("read:/path".to_string()));
396        assert_eq!(summary.last_success, Some(true));
397        assert_eq!(summary.history_len, 10);
398        assert!(summary.has_escalation);
399    }
400
401    #[test]
402    fn test_query_methods() {
403        let ctx = TaskContext::new(0)
404            .with_worker(WorkerSummary::new(WorkerId(0)))
405            .with_worker(WorkerSummary::new(WorkerId(1)).with_escalation(true))
406            .with_worker(WorkerSummary::new(WorkerId(2)));
407
408        assert_eq!(ctx.escalated_worker_count(), 1);
409        assert_eq!(ctx.worker_ids().len(), 3);
410    }
411
412    #[test]
413    fn test_filter_for_workers() {
414        // Setup: 4 workers with different states
415        let ctx = TaskContext::new(10)
416            .with_worker(WorkerSummary::new(WorkerId(0)).with_failures(1))
417            .with_worker(WorkerSummary::new(WorkerId(1)).with_escalation(true))
418            .with_worker(WorkerSummary::new(WorkerId(2)).with_history_len(5))
419            .with_worker(WorkerSummary::new(WorkerId(3)).with_last_action("read", true))
420            .with_escalation(WorkerId(1), Escalation::consecutive_failures(3, 5))
421            .with_success_rate(0.75)
422            .with_progress(0.5)
423            .insert("meta_key", "meta_value");
424
425        // Filter for workers 0 and 2 only
426        let filtered = ctx.filter_for_workers(&[WorkerId(0), WorkerId(2)]);
427
428        // Verify filtered context
429        assert_eq!(filtered.tick, 10);
430        assert_eq!(filtered.workers.len(), 2);
431        assert!(filtered.workers.contains_key(&WorkerId(0)));
432        assert!(filtered.workers.contains_key(&WorkerId(2)));
433        assert!(!filtered.workers.contains_key(&WorkerId(1)));
434        assert!(!filtered.workers.contains_key(&WorkerId(3)));
435
436        // Worker 0 should have its failures preserved
437        assert_eq!(
438            filtered
439                .workers
440                .get(&WorkerId(0))
441                .unwrap()
442                .consecutive_failures,
443            1
444        );
445
446        // Worker 2 should have its history_len preserved
447        assert_eq!(filtered.workers.get(&WorkerId(2)).unwrap().history_len, 5);
448
449        // Escalations should be filtered (worker 1's escalation should be excluded)
450        assert!(filtered.escalations.is_empty());
451
452        // Global metrics should be preserved
453        assert_eq!(filtered.success_rate, 0.75);
454        assert_eq!(filtered.progress, 0.5);
455
456        // Metadata should be preserved
457        assert_eq!(filtered.get_str("meta_key"), Some("meta_value"));
458    }
459
460    #[test]
461    fn test_filter_for_workers_empty() {
462        let ctx = TaskContext::new(5)
463            .with_worker(WorkerSummary::new(WorkerId(0)))
464            .with_worker(WorkerSummary::new(WorkerId(1)));
465
466        let filtered = ctx.filter_for_workers(&[]);
467
468        assert_eq!(filtered.tick, 5);
469        assert!(filtered.workers.is_empty());
470    }
471}
swarm_engine_core/context/summary.rs

swarm_engine_core/context/
summary.rs
