use crate::types::WorkerId;
use super::store::{
ActionCandidate, ContextStore, ContextTarget, ContextView, ResolvedContext, WorkerContext,
};
use crate::agent::{ManagerId, WorkerScope};
pub struct ContextResolver;
impl ContextResolver {
pub fn resolve(store: &ContextStore, view: &ContextView) -> ResolvedContext {
match view {
ContextView::Global { manager_id } => Self::resolve_global(store, *manager_id),
ContextView::Local {
worker_id,
neighbor_ids,
} => Self::resolve_local(store, *worker_id, neighbor_ids),
ContextView::Custom {
name: _,
visible_worker_ids,
visible_manager_ids: _,
} => Self::resolve_custom(store, visible_worker_ids),
}
}
fn resolve_global(store: &ContextStore, manager_id: ManagerId) -> ResolvedContext {
let global = store.global.clone();
let visible_workers: Vec<WorkerContext> = store.workers.values().cloned().collect();
let escalations = store.escalations.clone();
let candidates = store
.actions
.as_ref()
.map(ActionCandidate::from_config)
.unwrap_or_default();
ResolvedContext::new(global, ContextTarget::Manager(manager_id))
.with_workers(visible_workers)
.with_escalations(escalations)
.with_candidates(candidates)
.with_metadata(store.metadata.clone())
}
fn resolve_local(
store: &ContextStore,
worker_id: WorkerId,
neighbor_ids: &[WorkerId],
) -> ResolvedContext {
let global = store.global.clone();
let mut visible_worker_ids = vec![worker_id];
visible_worker_ids.extend(neighbor_ids.iter().copied());
let visible_workers: Vec<WorkerContext> = visible_worker_ids
.iter()
.filter_map(|id| store.workers.get(id).cloned())
.collect();
let escalations: Vec<_> = store
.escalations
.iter()
.filter(|(wid, _)| visible_worker_ids.contains(wid))
.cloned()
.collect();
let candidates = store
.actions
.as_ref()
.map(ActionCandidate::from_config)
.unwrap_or_default();
let mut metadata = std::collections::HashMap::new();
if let Some(task) = store.get("task") {
metadata.insert("task".to_string(), task.clone());
}
if let Some(hint) = store.get("hint") {
metadata.insert("hint".to_string(), hint.clone());
}
ResolvedContext::new(global, ContextTarget::Worker(worker_id))
.with_workers(visible_workers)
.with_escalations(escalations)
.with_candidates(candidates)
.with_metadata(metadata)
}
fn resolve_custom(store: &ContextStore, visible_worker_ids: &[WorkerId]) -> ResolvedContext {
let global = store.global.clone();
let visible_workers: Vec<WorkerContext> = visible_worker_ids
.iter()
.filter_map(|id| store.workers.get(id).cloned())
.collect();
let escalations: Vec<_> = store
.escalations
.iter()
.filter(|(wid, _)| visible_worker_ids.contains(wid))
.cloned()
.collect();
let candidates = store
.actions
.as_ref()
.map(ActionCandidate::from_config)
.unwrap_or_default();
let target = visible_worker_ids
.first()
.map(|&id| ContextTarget::Worker(id))
.unwrap_or(ContextTarget::Worker(WorkerId(0)));
ResolvedContext::new(global, target)
.with_workers(visible_workers)
.with_escalations(escalations)
.with_candidates(candidates)
.with_metadata(store.metadata.clone())
}
pub fn resolve_with_scope(
store: &ContextStore,
worker_id: WorkerId,
scope: &WorkerScope,
) -> ResolvedContext {
let global = store.global.clone();
let candidates = store
.actions
.as_ref()
.map(ActionCandidate::from_config)
.unwrap_or_default();
match scope {
WorkerScope::Idle => {
ResolvedContext::new(global, ContextTarget::Worker(worker_id))
.with_candidates(candidates)
}
WorkerScope::Minimal => {
let self_output = store
.workers
.get(&worker_id)
.and_then(|w| w.metadata.get("last_output"))
.and_then(|v| v.as_str())
.map(String::from);
ResolvedContext::new(global, ContextTarget::Worker(worker_id))
.with_self_last_output(self_output)
.with_candidates(candidates)
}
WorkerScope::SelfDetail => {
let self_ctx = store.workers.get(&worker_id).cloned();
let visible_workers = self_ctx.into_iter().collect();
ResolvedContext::new(global, ContextTarget::Worker(worker_id))
.with_workers(visible_workers)
.with_candidates(candidates)
}
WorkerScope::WithTeamSummary => {
let visible_workers: Vec<WorkerContext> = store
.workers
.values()
.map(|w| {
WorkerContext {
id: w.id,
consecutive_failures: 0,
last_action: w.last_action.clone(),
last_success: w.last_success,
history_len: 0,
has_escalation: w.has_escalation,
candidates: Vec::new(),
metadata: std::collections::HashMap::new(),
}
})
.collect();
ResolvedContext::new(global, ContextTarget::Worker(worker_id))
.with_workers(visible_workers)
.with_candidates(candidates)
}
WorkerScope::WithTeamDetail => {
let visible_workers: Vec<WorkerContext> = store.workers.values().cloned().collect();
let escalations = store.escalations.clone();
ResolvedContext::new(global, ContextTarget::Worker(worker_id))
.with_workers(visible_workers)
.with_escalations(escalations)
.with_candidates(candidates)
}
}
}
}
pub trait NeighborStrategy: Send + Sync {
fn neighbors(&self, worker_id: WorkerId, all_workers: &[WorkerId]) -> Vec<WorkerId>;
}
#[derive(Debug, Clone, Default)]
pub struct AllVisibleStrategy;
impl NeighborStrategy for AllVisibleStrategy {
fn neighbors(&self, _worker_id: WorkerId, all_workers: &[WorkerId]) -> Vec<WorkerId> {
all_workers.to_vec()
}
}
#[derive(Debug, Clone)]
pub struct AdjacentStrategy {
pub radius: usize,
}
impl AdjacentStrategy {
pub fn new(radius: usize) -> Self {
Self { radius }
}
}
impl NeighborStrategy for AdjacentStrategy {
fn neighbors(&self, worker_id: WorkerId, all_workers: &[WorkerId]) -> Vec<WorkerId> {
let idx = worker_id.0;
all_workers
.iter()
.enumerate()
.filter(|(i, _)| {
let diff = (*i as isize - idx as isize).unsigned_abs();
diff > 0 && diff <= self.radius
})
.map(|(_, &id)| id)
.collect()
}
}
#[derive(Debug, Clone, Default)]
pub struct EscalationOnlyStrategy;
impl NeighborStrategy for EscalationOnlyStrategy {
fn neighbors(&self, _worker_id: WorkerId, _all_workers: &[WorkerId]) -> Vec<WorkerId> {
Vec::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::ManagerContext;
fn create_test_store() -> ContextStore {
ContextStore::new(10)
.with_worker(
WorkerContext::new(WorkerId(0))
.with_last_action("read:/a", true)
.with_candidates(vec!["read".into(), "grep".into()]),
)
.with_worker(
WorkerContext::new(WorkerId(1))
.with_last_action("grep:pattern", false)
.with_escalation(true),
)
.with_worker(WorkerContext::new(WorkerId(2)).with_failures(3))
.with_manager(ManagerContext::new(ManagerId(0)))
.insert("task", "Find the bug")
.insert("hint", "Check error logs")
}
#[test]
fn test_resolve_global() {
let store = create_test_store();
let view = ContextView::global(ManagerId(0));
let resolved = ContextResolver::resolve(&store, &view);
assert_eq!(resolved.visible_workers.len(), 3);
assert!(resolved.is_manager());
assert_eq!(resolved.global.tick, 10);
}
#[test]
fn test_resolve_local_no_neighbors() {
let store = create_test_store();
let view = ContextView::local(WorkerId(0));
let resolved = ContextResolver::resolve(&store, &view);
assert_eq!(resolved.visible_workers.len(), 1);
assert_eq!(resolved.visible_workers[0].id, WorkerId(0));
assert!(resolved.is_worker());
}
#[test]
fn test_resolve_local_with_neighbors() {
let store = create_test_store();
let view = ContextView::local_with_neighbors(WorkerId(0), vec![WorkerId(1)]);
let resolved = ContextResolver::resolve(&store, &view);
assert_eq!(resolved.visible_workers.len(), 2);
let ids: Vec<_> = resolved.visible_workers.iter().map(|w| w.id).collect();
assert!(ids.contains(&WorkerId(0)));
assert!(ids.contains(&WorkerId(1)));
}
#[test]
fn test_resolve_local_filters_escalations() {
let store = create_test_store()
.with_escalation(
WorkerId(1),
crate::state::Escalation::consecutive_failures(3, 5),
)
.with_escalation(
WorkerId(2),
crate::state::Escalation::consecutive_failures(5, 8),
);
let view = ContextView::local_with_neighbors(WorkerId(0), vec![WorkerId(1)]);
let resolved = ContextResolver::resolve(&store, &view);
assert_eq!(resolved.escalations.len(), 1);
assert_eq!(resolved.escalations[0].0, WorkerId(1));
}
#[test]
fn test_adjacent_strategy() {
let strategy = AdjacentStrategy::new(1);
let all = vec![WorkerId(0), WorkerId(1), WorkerId(2), WorkerId(3)];
let neighbors = strategy.neighbors(WorkerId(1), &all);
assert!(neighbors.contains(&WorkerId(0)));
assert!(neighbors.contains(&WorkerId(2)));
assert!(!neighbors.contains(&WorkerId(3)));
}
}