use anyhow::{anyhow, Result};
use chrono::Utc;
use crate::warehouse::agent_model_runs::ModelCaller;
use crate::warehouse::dep_graph::WorkspaceGraph;
use super::event::{Event, ItemKind, PlanStatus};
use super::ids::{IdeaId, NodeId, PlanId};
use super::store::Store;
#[cfg(feature = "vector")]
pub use crate::vector::store::Embedder as ProposerEmbedder;
#[cfg(not(feature = "vector"))]
pub trait ProposerEmbedder {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ComponentProposal {
pub picks: Vec<String>,
pub confident: bool,
pub all: Vec<String>,
pub source: ProposalSource,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProposalSource {
Llm,
Embedder,
PickList,
}
impl ProposalSource {
pub fn as_str(&self) -> &'static str {
match self {
ProposalSource::Llm => "llm",
ProposalSource::Embedder => "embedder",
ProposalSource::PickList => "pick_list",
}
}
}
impl ComponentProposal {
pub fn to_json(&self) -> serde_json::Value {
serde_json::json!({
"picks": self.picks,
"confident": self.confident,
"source": self.source.as_str(),
"all": self.all,
})
}
}
pub fn component_list(graph: &WorkspaceGraph) -> Vec<String> {
graph.component_names()
}
pub fn propose_components(
input: &str,
graph: &WorkspaceGraph,
llm: Option<&dyn ModelCaller>,
embedder: Option<&dyn ProposerEmbedder>,
) -> ComponentProposal {
let all = component_list(graph);
if let Some(caller) = llm {
if let Some(picks) = llm_pick(caller, input, &all) {
if !picks.is_empty() {
return ComponentProposal {
picks,
confident: true,
all,
source: ProposalSource::Llm,
};
}
}
}
#[cfg(feature = "vector")]
if let Some(emb) = embedder {
if let Some(pick) = embed_pick(emb, input, &all) {
return ComponentProposal {
picks: vec![pick],
confident: true,
all,
source: ProposalSource::Embedder,
};
}
}
#[cfg(not(feature = "vector"))]
let _ = &embedder;
ComponentProposal { picks: Vec::new(), confident: false, all, source: ProposalSource::PickList }
}
fn llm_pick(caller: &dyn ModelCaller, input: &str, components: &[String]) -> Option<Vec<String>> {
let prompt = build_classify_prompt(input, components);
let answer = caller.call("planner", "classify", &prompt).ok()?;
Some(extract_components(&answer.output, components))
}
pub fn build_classify_prompt(input: &str, components: &[String]) -> String {
let mut p = String::new();
p.push_str(
"You map a change request to the software component(s) it touches.\n\
Reply with ONLY the matching component name(s) from this exact list, \
comma-separated, nothing else.\n\nComponents:\n",
);
for c in components {
p.push_str("- ");
p.push_str(c);
p.push('\n');
}
p.push_str("\nChange request:\n");
p.push_str(input);
p.push_str("\n\nComponent(s):");
p
}
pub fn extract_components(answer: &str, components: &[String]) -> Vec<String> {
let lc: std::collections::BTreeMap<String, &String> =
components.iter().map(|c| (c.to_ascii_lowercase(), c)).collect();
let mut picks: Vec<String> = Vec::new();
let cleaned: String = answer
.chars()
.map(|c| if c.is_alphanumeric() || c == '-' || c == '_' { c } else { ' ' })
.collect();
for tok in cleaned.split_whitespace() {
if let Some(name) = lc.get(&tok.to_ascii_lowercase()) {
if !picks.contains(*name) {
picks.push((*name).clone());
}
}
}
picks
}
#[cfg(feature = "vector")]
fn embed_pick(
emb: &dyn crate::vector::store::Embedder,
input: &str,
components: &[String],
) -> Option<String> {
if components.is_empty() {
return None;
}
let mut texts = Vec::with_capacity(components.len() + 1);
texts.push(input.to_string());
texts.extend(components.iter().cloned());
let vecs = emb.embed(&texts).ok()?;
if vecs.len() != texts.len() {
return None;
}
let q = &vecs[0];
let mut best: Option<(usize, f32)> = None;
for (i, comp_vec) in vecs[1..].iter().enumerate() {
let s = cosine(q, comp_vec);
if best.map(|(_, bs)| s > bs).unwrap_or(true) {
best = Some((i, s));
}
}
best.map(|(i, _)| components[i].clone())
}
#[cfg(feature = "vector")]
fn cosine(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let mut dot = 0.0f32;
let mut na = 0.0f32;
let mut nb = 0.0f32;
for (x, y) in a.iter().zip(b.iter()) {
dot += x * y;
na += x * x;
nb += y * y;
}
if na == 0.0 || nb == 0.0 {
0.0
} else {
dot / (na.sqrt() * nb.sqrt())
}
}
pub fn affected_components(graph: &WorkspaceGraph, touched: &[String]) -> Vec<String> {
graph.affected_by_change(touched)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DerivedPlan {
pub touched: Vec<String>,
pub affected: Vec<String>,
pub node_kind: String,
}
impl DerivedPlan {
pub fn to_json(&self) -> serde_json::Value {
serde_json::json!({
"touched": self.touched,
"affected": self.affected,
"node_kind": self.node_kind,
})
}
}
pub fn derive_plan(graph: &WorkspaceGraph, touched: &[String], kind: ItemKind) -> DerivedPlan {
DerivedPlan {
touched: touched.to_vec(),
affected: affected_components(graph, touched),
node_kind: kind.node_kind().to_string(),
}
}
pub fn plan_from_component(
store: &mut Store,
idea_id: &IdeaId,
touched: &[String],
graph: &WorkspaceGraph,
) -> Result<PlanId> {
if touched.is_empty() {
return Err(anyhow!("plan_from_component: no touched component(s) given"));
}
let unknown: Vec<&String> = touched.iter().filter(|c| !graph.has_component(c)).collect();
if !unknown.is_empty() {
let known = component_list(graph).join(", ");
return Err(anyhow!(
"plan_from_component: unknown component(s) {unknown:?}; known: {known}"
));
}
let kind = store
.funnel
.ideas
.get(idea_id)
.map(|i| i.item_kind)
.unwrap_or(ItemKind::Idea);
let node_kind = kind.node_kind().to_string();
let affected = affected_components(graph, touched);
if affected.is_empty() {
return Err(anyhow!(
"plan_from_component: component(s) {touched:?} are not in the dep graph"
));
}
let base = Utc::now();
let mut tick: i64 = 0;
let mut at = || {
let t = base + chrono::Duration::microseconds(tick);
tick += 1;
t
};
let plan_id = PlanId::seq(store.funnel.next_plan);
let summary = format!(
"auto-plan ({}) for {} → {} affected component(s)",
kind.as_str(),
touched.join(", "),
affected.len(),
);
store.record(Event::PlanCreated {
id: plan_id.clone(),
idea_id: idea_id.clone(),
summary,
planner: "dep-graph".into(),
ts: at(),
})?;
store.record(Event::PlanStatusChanged {
plan_id: plan_id.clone(),
status: PlanStatus::Active,
why: Some("auto-planned from component dep graph".into()),
ts: at(),
})?;
let mut prev: Option<NodeId> = None;
for comp in &affected {
let node_id = NodeId::seq(store.funnel.next_node);
let mut params = serde_json::Map::new();
params.insert("title".into(), serde_json::Value::String(format!("{node_kind} {comp}")));
params.insert("component".into(), serde_json::Value::String(comp.clone()));
store.record(Event::NodeAdded {
plan_id: plan_id.clone(),
node_id: node_id.clone(),
kind: node_kind.clone(),
params,
targets: vec![comp.clone()],
prompt_excerpt: None,
ts: at(),
})?;
if let Some(from) = prev {
store.record(Event::EdgeAdded {
plan_id: plan_id.clone(),
from_node: from,
to_node: node_id.clone(),
ts: at(),
})?;
}
prev = Some(node_id);
}
store.funnel.promote_ready();
Ok(plan_id)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::warehouse::agent_model_runs::{MockCaller, ModelAnswer};
use crate::warehouse::dep_graph::{topo_order_from_edges, CrossRepoEdge, RepoFacts, WorkspaceGraph};
#[cfg(feature = "vector")]
use crate::vector::store::{Embedder, ModelProfile};
use std::collections::{BTreeMap, BTreeSet};
use std::path::PathBuf;
fn facts(name: &str, produces: &[&str], consumes: &[&str]) -> RepoFacts {
RepoFacts {
name: name.to_string(),
root: PathBuf::from("/dev/null"),
produces: produces.iter().map(|s| s.to_string()).collect(),
consumes: consumes.iter().map(|s| s.to_string()).collect(),
}
}
fn edge(from: &str, to: &str, via: &[&str]) -> CrossRepoEdge {
CrossRepoEdge {
from: from.to_string(),
to: to.to_string(),
via: via.iter().map(|s| s.to_string()).collect::<BTreeSet<_>>(),
}
}
fn diamond() -> WorkspaceGraph {
let facts = [
facts("app", &["app_c"], &["a_c", "b_c"]),
facts("liba", &["a_c"], &["util_c"]),
facts("libb", &["b_c"], &["util_c"]),
facts("util", &["util_c"], &[]),
];
let mut fmap = BTreeMap::new();
for f in facts {
fmap.insert(f.name.clone(), f);
}
WorkspaceGraph::from_query_parts(
fmap,
vec![
edge("app", "liba", &["a_c"]),
edge("app", "libb", &["b_c"]),
edge("liba", "util", &["util_c"]),
edge("libb", "util", &["util_c"]),
],
)
}
#[cfg(feature = "vector")]
struct KeywordEmbedder;
#[cfg(feature = "vector")]
impl Embedder for KeywordEmbedder {
fn profile(&self) -> ModelProfile {
ModelProfile {
model_name: "kw".into(),
weights_sha: "w".into(),
tokenizer_sha: "t".into(),
pooling: "mean".into(),
normalize: true,
dim: 4,
dtype: "f32".into(),
}
}
fn embed(&self, texts: &[String]) -> anyhow::Result<Vec<Vec<f32>>> {
let axis = |t: &str| -> usize {
let t = t.to_ascii_lowercase();
if t.contains("util") {
0
} else if t.contains("liba") {
1
} else if t.contains("libb") {
2
} else {
3
}
};
Ok(texts
.iter()
.map(|t| {
let mut v = vec![0.0f32; 4];
v[axis(t)] = 1.0;
v
})
.collect())
}
}
#[test]
fn component_list_is_sorted_repo_names() {
let g = diamond();
assert_eq!(component_list(&g), vec!["app", "liba", "libb", "util"]);
}
#[test]
fn affected_set_is_blast_radius_in_build_order() {
let g = diamond();
let affected = affected_components(&g, &["util".to_string()]);
assert_eq!(
affected.iter().cloned().collect::<BTreeSet<_>>(),
["app", "liba", "libb", "util"].iter().map(|s| s.to_string()).collect()
);
let pos = |n: &str| affected.iter().position(|x| x == n).unwrap();
assert!(pos("util") < pos("liba"));
assert!(pos("util") < pos("libb"));
assert!(pos("liba") < pos("app"));
assert!(pos("libb") < pos("app"));
assert_eq!(affected_components(&g, &["app".to_string()]), vec!["app".to_string()]);
}
#[test]
fn llm_leg_picks_the_named_component() {
let g = diamond();
let caller = MockCaller::new().with_cell(
"planner",
"classify",
ModelAnswer::basic("The component is util.", 1.0, 1, 1, 1.0, 1.0),
);
let p = propose_components("speed up the shared util layer", &g, Some(&caller), None);
assert!(p.confident);
assert_eq!(p.source, ProposalSource::Llm);
assert_eq!(p.picks, vec!["util".to_string()]);
assert_eq!(p.all, vec!["app", "liba", "libb", "util"]);
}
#[test]
fn llm_hallucination_falls_through_to_pick_list() {
let g = diamond();
let caller = MockCaller::new().with_cell(
"planner",
"classify",
ModelAnswer::basic("ghost-component", 1.0, 1, 1, 1.0, 1.0),
);
let p = propose_components("touch the ghost", &g, Some(&caller), None);
assert!(!p.confident);
assert_eq!(p.source, ProposalSource::PickList);
assert!(p.picks.is_empty());
assert_eq!(p.all, vec!["app", "liba", "libb", "util"]);
}
#[cfg(feature = "vector")]
#[test]
fn embedder_leg_used_when_llm_absent() {
let g = diamond();
let p = propose_components("rework liba internals", &g, None, Some(&KeywordEmbedder));
assert!(p.confident);
assert_eq!(p.source, ProposalSource::Embedder);
assert_eq!(p.picks, vec!["liba".to_string()]);
}
#[test]
fn pick_list_when_no_backend() {
let g = diamond();
let p = propose_components("anything", &g, None, None);
assert!(!p.confident);
assert_eq!(p.source, ProposalSource::PickList);
assert!(p.picks.is_empty());
assert_eq!(p.all, vec!["app", "liba", "libb", "util"]);
}
#[test]
fn plan_from_component_generates_topo_plan_nodes() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path().join("wh");
let mut store = Store::open(&root).unwrap();
let idea = IdeaId::seq(store.funnel.next_idea);
store
.record(Event::IdeaSubmitted {
id: idea.clone(),
source: "test".into(),
text: "add tests for the shared util layer".into(),
refs: vec![],
item_kind: ItemKind::Test,
ts: Utc::now(),
})
.unwrap();
let g = diamond();
let caller = MockCaller::new().with_cell(
"planner",
"classify",
ModelAnswer::basic("util", 1.0, 1, 1, 1.0, 1.0),
);
let prop = propose_components(&store.funnel.ideas[&idea].text, &g, Some(&caller), None);
assert_eq!(prop.picks, vec!["util".to_string()]);
let plan_id = plan_from_component(&mut store, &idea, &prop.picks, &g).unwrap();
let plan = store.funnel.plans.get(&plan_id).expect("plan created");
assert_eq!(plan.idea_id, idea);
assert_eq!(plan.status, PlanStatus::Active);
assert_eq!(plan.nodes.len(), 4, "one node per affected component");
let mut targeted: BTreeSet<String> = BTreeSet::new();
for n in plan.nodes.values() {
assert_eq!(n.kind, "test:write", "test idea → test:write node kind");
assert_eq!(n.targets.len(), 1);
targeted.insert(n.targets[0].clone());
}
assert_eq!(
targeted,
["app", "liba", "libb", "util"].iter().map(|s| s.to_string()).collect()
);
assert_eq!(plan.edges.len(), 3, "nodes chained prev→cur");
let ready: Vec<String> = plan
.nodes
.values()
.filter(|n| n.status == super::super::NodeStatus::Ready)
.flat_map(|n| n.targets.clone())
.collect();
assert_eq!(ready, vec!["util".to_string()], "only the deps-first node is Ready");
}
#[test]
fn plan_from_component_rejects_empty_and_unknown() {
let dir = tempfile::tempdir().unwrap();
let mut store = Store::open(dir.path().join("wh")).unwrap();
let idea = IdeaId::seq(0);
store
.record(Event::IdeaSubmitted {
id: idea.clone(),
source: "t".into(),
text: "x".into(),
refs: vec![],
item_kind: ItemKind::Idea,
ts: Utc::now(),
})
.unwrap();
let g = diamond();
assert!(plan_from_component(&mut store, &idea, &[], &g).is_err());
assert!(plan_from_component(&mut store, &idea, &["ghost".to_string()], &g).is_err());
}
#[test]
fn derive_plan_uses_item_kind_node_kind() {
let g = diamond();
assert_eq!(
derive_plan(&g, &["liba".to_string()], ItemKind::Error).node_kind,
"code:fix"
);
assert_eq!(
derive_plan(&g, &["liba".to_string()], ItemKind::Idea).node_kind,
"code:write"
);
let d = derive_plan(&g, &["util".to_string()], ItemKind::Test);
assert_eq!(d.node_kind, "test:write");
assert_eq!(
d.affected.iter().cloned().collect::<BTreeSet<_>>(),
["app", "liba", "libb", "util"].iter().map(|s| s.to_string()).collect()
);
}
#[test]
fn topo_order_is_stable_for_subset() {
let edges = vec![edge("liba", "util", &["c"]), edge("app", "liba", &["c"])];
let repos = vec!["app".to_string(), "liba".to_string(), "util".to_string()];
let order = topo_order_from_edges(&repos, &edges);
let pos = |n: &str| order.iter().position(|x| x == n).unwrap();
assert!(pos("util") < pos("liba"));
assert!(pos("liba") < pos("app"));
}
}