Skip to main content

swarm_engine_core/context/
store.rs

1//! Context System - スコープ付きコンテキスト管理
2//!
3//! # 設計思想
4//!
5//! Manager/Worker が「見えるべき」情報を制御するためのスコープシステム。
6//! データは ContextStore に正規化して格納し、ContextView で可視範囲を定義。
7//! ContextResolver が View に応じて ResolvedContext を生成。
8//!
9//! ```text
10//! ContextStore (正規化データ)
11//!       │
12//!       ├── ContextView::Global    → Manager用(全体が見える)
13//!       ├── ContextView::Local     → Worker用(自分+Neighbor)
14//!       └── ContextView::Custom    → カスタムフィルタ
15//!       │
16//!       ▼
17//! ContextResolver.resolve(store, view) → ResolvedContext
18//!       │
19//!       ▼
20//! LLM層でプロンプト生成
21//! ```
22
23use 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// ============================================================================
34// ContextStore - 正規化データストア
35// ============================================================================
36
37/// 正規化されたコンテキストデータストア
38///
39/// 全ての情報を一箇所に集約し、View によって可視範囲を制御する。
40#[derive(Debug, Clone)]
41pub struct ContextStore {
42    /// Global情報(tick, progress等)
43    pub global: GlobalContext,
44    /// 全Worker状態
45    pub workers: HashMap<WorkerId, WorkerContext>,
46    /// 全Manager情報
47    pub managers: HashMap<ManagerId, ManagerContext>,
48    /// Escalation一覧
49    pub escalations: Vec<(WorkerId, Escalation)>,
50    /// 利用可能Actions
51    pub actions: Option<ActionsConfig>,
52    /// 拡張メタデータ
53    pub metadata: HashMap<String, Value>,
54}
55
56impl ContextStore {
57    /// 新しい ContextStore を作成
58    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    /// Worker を追加
70    pub fn with_worker(mut self, ctx: WorkerContext) -> Self {
71        self.workers.insert(ctx.id, ctx);
72        self
73    }
74
75    /// Manager を追加
76    pub fn with_manager(mut self, ctx: ManagerContext) -> Self {
77        self.managers.insert(ctx.id, ctx);
78        self
79    }
80
81    /// Escalation を追加
82    pub fn with_escalation(mut self, worker_id: WorkerId, escalation: Escalation) -> Self {
83        self.escalations.push((worker_id, escalation));
84        self
85    }
86
87    /// Actions を設定
88    pub fn with_actions(mut self, actions: ActionsConfig) -> Self {
89        self.actions = Some(actions);
90        self
91    }
92
93    /// メタデータを追加
94    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    /// メタデータを取得
100    pub fn get(&self, key: &str) -> Option<&Value> {
101        self.metadata.get(key)
102    }
103
104    /// メタデータを文字列として取得
105    pub fn get_str(&self, key: &str) -> Option<&str> {
106        self.metadata.get(key).and_then(|v| v.as_str())
107    }
108}
109
110// ============================================================================
111// GlobalContext - 全体情報
112// ============================================================================
113
114/// Global情報(全リクエスト共通)
115#[derive(Debug, Clone, Default)]
116pub struct GlobalContext {
117    /// 現在の tick
118    pub tick: u64,
119    /// 最大 tick
120    pub max_ticks: u64,
121    /// 進捗 (0.0 - 1.0)
122    pub progress: f64,
123    /// 成功率 (0.0 - 1.0)
124    pub success_rate: f64,
125    /// タスク説明
126    pub task_description: Option<String>,
127    /// ヒント
128    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// ============================================================================
166// WorkerContext - Worker固有情報
167// ============================================================================
168
169/// Worker固有のコンテキスト
170#[derive(Debug, Clone)]
171pub struct WorkerContext {
172    /// Worker ID
173    pub id: WorkerId,
174    /// 連続失敗数
175    pub consecutive_failures: u32,
176    /// 最新アクション名
177    pub last_action: Option<String>,
178    /// 最新アクションの成功/失敗
179    pub last_success: Option<bool>,
180    /// 履歴の長さ
181    pub history_len: usize,
182    /// Escalation 中かどうか
183    pub has_escalation: bool,
184    /// このWorkerに適用可能なAction候補
185    pub candidates: Vec<String>,
186    /// Worker固有メタデータ
187    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// ============================================================================
232// ManagerContext - Manager固有情報
233// ============================================================================
234
235/// Manager固有のコンテキスト
236#[derive(Debug, Clone)]
237pub struct ManagerContext {
238    /// Manager ID
239    pub id: ManagerId,
240    /// Manager名
241    pub name: String,
242    /// 最後に処理した tick
243    pub last_tick: u64,
244    /// Manager固有メタデータ
245    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// ============================================================================
270// ContextView - スコープ定義
271// ============================================================================
272
273/// コンテキストの可視範囲を定義
274#[derive(Debug, Clone)]
275pub enum ContextView {
276    /// Manager用: 全体が見える
277    Global { manager_id: ManagerId },
278    /// Worker用: 自分 + Neighbor
279    Local {
280        worker_id: WorkerId,
281        /// 可視範囲の Neighbor Worker IDs
282        neighbor_ids: Vec<WorkerId>,
283    },
284    /// カスタムフィルタ(将来拡張)
285    Custom {
286        /// フィルタ名(デバッグ用)
287        name: String,
288        /// 可視 Worker IDs
289        visible_worker_ids: Vec<WorkerId>,
290        /// 可視 Manager IDs
291        visible_manager_ids: Vec<ManagerId>,
292    },
293}
294
295impl ContextView {
296    /// Manager用の Global View を作成
297    pub fn global(manager_id: ManagerId) -> Self {
298        Self::Global { manager_id }
299    }
300
301    /// Worker用の Local View を作成(Neighborなし)
302    pub fn local(worker_id: WorkerId) -> Self {
303        Self::Local {
304            worker_id,
305            neighbor_ids: Vec::new(),
306        }
307    }
308
309    /// Worker用の Local View を作成(Neighbor指定)
310    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    /// カスタム View を作成
318    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// ============================================================================
332// ActionCandidate - プロンプト構築用 Action 情報
333// ============================================================================
334
335/// プロンプト構築用の Action パラメータ情報
336#[derive(Debug, Clone)]
337pub struct ActionParam {
338    /// パラメータ名
339    pub name: String,
340    /// 説明
341    pub description: String,
342    /// 必須かどうか
343    pub required: bool,
344}
345
346/// プロンプト構築用の Action 情報
347///
348/// ActionsConfig から必要な情報だけを抽出したモデル。
349/// LLM 層でプロンプト生成時に使用する。
350#[derive(Debug, Clone)]
351pub struct ActionCandidate {
352    /// Action 名
353    pub name: String,
354    /// 説明
355    pub description: String,
356    /// パラメータ
357    pub params: Vec<ActionParam>,
358    /// 出力例(JSON 形式)
359    pub example: Option<String>,
360}
361
362impl ActionCandidate {
363    /// ActionsConfig から ActionCandidate のリストを生成
364    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// ============================================================================
386// ResolvedContext - 解決済みコンテキスト
387// ============================================================================
388
389/// 解決済みコンテキスト(LLM層に渡る)
390///
391/// ContextResolver が ContextStore + ContextView から生成。
392/// プロンプト生成に必要な情報のみを含む。
393#[derive(Debug, Clone)]
394pub struct ResolvedContext {
395    /// Global情報
396    pub global: GlobalContext,
397    /// 可視範囲の Worker 情報
398    pub visible_workers: Vec<WorkerContext>,
399    /// 可視範囲の Escalation
400    pub escalations: Vec<(WorkerId, Escalation)>,
401    /// 利用可能な Action 候補(プロンプト構築用)
402    pub candidates: Vec<ActionCandidate>,
403    /// 追加メタデータ
404    pub metadata: HashMap<String, Value>,
405    /// このコンテキストの対象
406    pub target: ContextTarget,
407    /// 自分の last_output(Scope::Minimal 用)
408    ///
409    /// WorkerScope::Minimal の場合、visible_workers は空にし、
410    /// 代わりにこのフィールドに自分の前回結果のみを格納する。
411    pub self_last_output: Option<String>,
412    /// Manager からの指示(前回 Guidance から抽出)
413    ///
414    /// 前回の Manager 判断を次回の Worker Prompt に埋め込むために使用。
415    /// Manager.prepare() で前回の Guidance から生成して設定する。
416    pub manager_instruction: Option<ManagerInstruction>,
417}
418
419/// コンテキストの対象
420#[derive(Debug, Clone)]
421pub enum ContextTarget {
422    Manager(ManagerId),
423    Worker(WorkerId),
424}
425
426impl ResolvedContext {
427    /// 新しい ResolvedContext を作成
428    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    /// Worker情報を追加
442    pub fn with_workers(mut self, workers: Vec<WorkerContext>) -> Self {
443        self.visible_workers = workers;
444        self
445    }
446
447    /// Escalationを追加
448    pub fn with_escalations(mut self, escalations: Vec<(WorkerId, Escalation)>) -> Self {
449        self.escalations = escalations;
450        self
451    }
452
453    /// 候補を設定
454    pub fn with_candidates(mut self, candidates: Vec<ActionCandidate>) -> Self {
455        self.candidates = candidates;
456        self
457    }
458
459    /// ActionsConfig から候補を設定
460    pub fn with_actions_config(mut self, config: &ActionsConfig) -> Self {
461        self.candidates = ActionCandidate::from_config(config);
462        self
463    }
464
465    /// メタデータを追加
466    pub fn with_metadata(mut self, metadata: HashMap<String, Value>) -> Self {
467        self.metadata = metadata;
468        self
469    }
470
471    /// 自分の last_output を設定(Scope::Minimal 用)
472    pub fn with_self_last_output(mut self, output: Option<String>) -> Self {
473        self.self_last_output = output;
474        self
475    }
476
477    /// Manager からの指示を設定
478    pub fn with_manager_instruction(mut self, instruction: ManagerInstruction) -> Self {
479        self.manager_instruction = Some(instruction);
480        self
481    }
482
483    /// Escalation があるか
484    pub fn has_escalations(&self) -> bool {
485        !self.escalations.is_empty()
486    }
487
488    /// 対象がManagerか
489    pub fn is_manager(&self) -> bool {
490        matches!(self.target, ContextTarget::Manager(_))
491    }
492
493    /// 対象がWorkerか
494    pub fn is_worker(&self) -> bool {
495        matches!(self.target, ContextTarget::Worker(_))
496    }
497}
498
499// ============================================================================
500// Tests
501// ============================================================================
502
503#[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}