use super::tree::TreeIndex;
use anyhow::{Context as _, Result};
use serde::{Deserialize, Serialize};
pub const MAX_TOOL_ROUNDS: usize = 12;
pub const SYSTEM_PROMPT: &str = "\
You are a spec retrieval agent for an SDD (spec-driven development) project.
Given a task and optional file paths, find which specs are relevant.
NAVIGATION PROTOCOL (follow this order):
1. Call list_specs() to see all available spec documents and their purposes.
2. For specs whose purpose seems relevant to the task, call get_document_structure(spec_id)\
to see their requirement titles (cheap, no full text).
3. For requirements that look relevant, call get_spec_content(spec_id, req_ids) to read\
the full MUST/SHALL statements.
4. Finally, output ONLY a JSON object (no other text, no code fences):
{\"direct\": [{\"id\": \"<spec_id>\", \"reason\": \"<one sentence why this MUST be read>\"}],\
\"related\": [{\"id\": \"<spec_id>\", \"reason\": \"<one sentence>\"}]}
CLASSIFICATION RULES:
- \"direct\" = specs whose behavior contract (any MUST/SHALL statement, command behavior,\
exit code, output format, validation rule, or CLI surface) is affected by this change.
- \"related\" = specs that provide useful context but whose contract won't change.
- If the task changes behavior, the governing spec MUST be in \"direct\" — do NOT leave\
\"direct\" empty just because you are unsure. An empty result means the task touches NO\
spec behavior at all (e.g. a typo or formatting fix in a non-spec file, a pure dependency\
version bump, or a file wholly outside every spec's scope).
- Decide \"direct\" vs \"related\" vs omit based on the requirement text you read, not on the\
spec's purpose line alone.
- Be precise: prefer fewer, well-reasoned entries over many guesses.\
";
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Msg {
System(String),
User(String),
Assistant {
content: Option<String>,
tool_calls: Vec<ToolCall>,
},
Tool {
tool_call_id: String,
content: String,
},
}
#[derive(Clone, Debug)]
pub struct ToolSchema {
pub name: &'static str,
pub description: &'static str,
pub parameters: serde_json::Value,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: String,
}
#[derive(Clone, Debug, Default)]
pub struct ChatTurn {
pub content: Option<String>,
pub tool_calls: Vec<ToolCall>,
}
#[allow(async_fn_in_trait)]
pub trait ChatInvoker {
async fn chat_turn(&self, messages: &[Msg], tools: &[ToolSchema]) -> Result<ChatTurn>;
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct TierEntry {
pub id: String,
pub reason: String,
}
#[derive(Clone, Debug, Default)]
pub struct RetrievalOutput {
pub direct: Vec<TierEntry>,
pub related: Vec<TierEntry>,
pub tool_calls: usize,
pub truncated: bool,
}
impl RetrievalOutput {
pub fn truncated(tool_calls: usize) -> Self {
Self {
direct: Vec::new(),
related: Vec::new(),
tool_calls,
truncated: true,
}
}
}
pub fn build_tool_schemas() -> Vec<ToolSchema> {
vec![
ToolSchema {
name: "list_specs",
description: "List all spec documents with metadata. Call this first to see what specs exist.",
parameters: serde_json::json!({ "type": "object", "properties": {} }),
},
ToolSchema {
name: "get_document_structure",
description: "Get the tree structure of one spec (titles + req_ids only, no full text, to save tokens).",
parameters: serde_json::json!({
"type": "object",
"properties": { "spec_id": { "type": "string" } },
"required": ["spec_id"],
}),
},
ToolSchema {
name: "get_spec_content",
description: "Get the full statement text of specific requirements in a spec.",
parameters: serde_json::json!({
"type": "object",
"properties": {
"spec_id": { "type": "string" },
"req_ids": { "type": "array", "items": { "type": "string" } },
},
"required": ["spec_id"],
}),
},
]
}
pub fn dispatch_tool(name: &str, arguments: &str, tree: &TreeIndex) -> String {
let args: serde_json::Value =
serde_json::from_str(arguments).unwrap_or_else(|_| serde_json::json!({}));
match name {
"list_specs" => list_specs(tree),
"get_document_structure" => {
let spec_id = args["spec_id"].as_str().unwrap_or("");
get_document_structure(tree, spec_id)
}
"get_spec_content" => {
let spec_id = args["spec_id"].as_str().unwrap_or("");
let req_ids: Vec<String> = args["req_ids"]
.as_array()
.map(|a| {
a.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
get_spec_content(tree, spec_id, &req_ids)
}
other => serde_json::json!({ "error": format!("unknown tool: {other}") }).to_string(),
}
}
fn list_specs(tree: &TreeIndex) -> String {
let arr: Vec<serde_json::Value> = tree
.docs
.iter()
.map(|d| {
serde_json::json!({
"spec_id": d.spec_id,
"purpose": d.purpose,
"req_count": d.reqs.len(),
})
})
.collect();
serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_string())
}
fn get_document_structure(tree: &TreeIndex, spec_id: &str) -> String {
match tree.docs.iter().find(|d| d.spec_id == spec_id) {
Some(d) => {
let reqs: Vec<serde_json::Value> = d
.reqs
.iter()
.map(|r| serde_json::json!({ "req_id": r.req_id, "title": r.title }))
.collect();
serde_json::json!({
"spec_id": d.spec_id,
"purpose": d.purpose,
"reqs": reqs,
})
.to_string()
}
None => {
serde_json::json!({ "error": format!("spec_id {spec_id:?} not found") }).to_string()
}
}
}
fn get_spec_content(tree: &TreeIndex, spec_id: &str, req_ids: &[String]) -> String {
match tree.docs.iter().find(|d| d.spec_id == spec_id) {
Some(d) => {
let want: std::collections::HashSet<&str> =
req_ids.iter().map(|s| s.as_str()).collect();
let result: Vec<serde_json::Value> = d
.reqs
.iter()
.filter(|r| want.contains(r.req_id.as_str()))
.map(|r| serde_json::json!({ "req_id": r.req_id, "statement": r.statement }))
.collect();
serde_json::to_string(&result).unwrap_or_else(|_| "[]".to_string())
}
None => {
serde_json::json!({ "error": format!("spec_id {spec_id:?} not found") }).to_string()
}
}
}
pub async fn retrieve<I: ChatInvoker>(
invoker: &I,
tree: &TreeIndex,
task: &str,
paths: &[String],
) -> Result<RetrievalOutput> {
let tools = build_tool_schemas();
let user = if paths.is_empty() {
format!("Task: {task}")
} else {
format!("Task: {task}\nFile paths involved: {}", paths.join(", "))
};
let mut messages: Vec<Msg> = vec![Msg::System(SYSTEM_PROMPT.to_string()), Msg::User(user)];
let mut tool_calls_total = 0usize;
let debug = std::env::var("LLMAN_SDD_INDEX_DEBUG").is_ok();
if debug {
eprintln!("[pageindex] task={task:?} paths={paths:?}");
}
for _round in 0..MAX_TOOL_ROUNDS {
let turn = invoker.chat_turn(&messages, &tools).await?;
if debug {
eprintln!(
"[pageindex] turn: content={:?} tool_calls={{{}}}",
turn.content.as_deref().unwrap_or(""),
turn.tool_calls
.iter()
.map(|tc| format!("{}({})", tc.name, tc.arguments))
.collect::<Vec<_>>()
.join(", ")
);
}
if turn.tool_calls.is_empty() {
let mut out = parse_final_answer(turn.content.as_deref().unwrap_or(""))?;
out.tool_calls = tool_calls_total;
if debug {
eprintln!(
"[pageindex] final: direct={:?} related={:?}",
out.direct, out.related
);
}
return Ok(out);
}
messages.push(Msg::Assistant {
content: turn.content.clone(),
tool_calls: turn.tool_calls.clone(),
});
for tc in &turn.tool_calls {
tool_calls_total += 1;
let result = dispatch_tool(&tc.name, &tc.arguments, tree);
messages.push(Msg::Tool {
tool_call_id: tc.id.clone(),
content: result,
});
}
}
let empty_tools: Vec<ToolSchema> = Vec::new();
let salvage_prompt = Msg::User(
"You have used all available tool calls. Based ONLY on what you have read so far, \
output the final direct/related JSON now. Do not request any more tools."
.to_string(),
);
messages.push(salvage_prompt);
let turn = invoker.chat_turn(&messages, &empty_tools).await?;
if debug {
eprintln!(
"[pageindex] salvage turn (no tools): content={:?}",
turn.content.as_deref().unwrap_or("")
);
}
match parse_final_answer(turn.content.as_deref().unwrap_or("")) {
Ok(mut out) => {
out.tool_calls = tool_calls_total;
out.truncated = true;
if debug {
eprintln!(
"[pageindex] final (salvaged, truncated): direct={:?} related={:?}",
out.direct, out.related
);
}
return Ok(out);
}
Err(e) => {
if debug {
eprintln!("[pageindex] salvage parse failed: {e}");
}
}
}
Ok(RetrievalOutput::truncated(tool_calls_total))
}
pub fn parse_final_answer(content: &str) -> Result<RetrievalOutput> {
let json_str = extract_json_object(content)
.ok_or_else(|| anyhow::anyhow!("model did not return a JSON object; got: {content}"))?;
#[derive(Deserialize)]
struct Answer {
#[serde(default)]
direct: Vec<TierEntry>,
#[serde(default)]
related: Vec<TierEntry>,
}
let ans: Answer =
serde_json::from_str(&json_str).context("failed to parse direct/related JSON")?;
Ok(RetrievalOutput {
direct: ans.direct,
related: ans.related,
tool_calls: 0,
truncated: false,
})
}
fn extract_json_object(content: &str) -> Option<String> {
let start = content.find('{')?;
let end = content.rfind('}')?;
if end < start {
return None;
}
Some(content[start..=end].to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sdd::spec::ir::{MainSpecDoc, RequirementEntry};
use std::sync::atomic::{AtomicUsize, Ordering};
fn build_test_tree() -> TreeIndex {
let mk_req = |rid: &str, title: &str, stmt: &str| RequirementEntry {
req_id: rid.to_string(),
title: title.to_string(),
statement: stmt.to_string(),
};
let docs = crate::sdd::context::tree::build_docs(&[
(
"sdd-workflow".to_string(),
MainSpecDoc {
kind: "llman.sdd.spec".into(),
name: "sdd-workflow".into(),
purpose: "Define the SDD workflow.".into(),
valid_scope: vec![],
requirements: vec![
mk_req("r1", "Init", "`llman sdd init` MUST create dirs."),
mk_req(
"r12",
"Validate",
"`llman sdd validate` MUST check format and exit non-zero on error.",
),
],
scenarios: vec![],
feature_refs: None,
},
),
(
"cli".to_string(),
MainSpecDoc {
kind: "llman.sdd.spec".into(),
name: "cli".into(),
purpose: "CLI surface.".into(),
valid_scope: vec![],
requirements: vec![mk_req("r1", "Commands", "MUST expose subcommands.")],
scenarios: vec![],
feature_refs: None,
},
),
]);
TreeIndex::new(docs, "hash".into(), "ts".into(), "model".into())
}
#[test]
fn test_list_specs_tool() {
let tree = build_test_tree();
let out: serde_json::Value =
serde_json::from_str(&dispatch_tool("list_specs", "{}", &tree)).unwrap();
let arr = out.as_array().unwrap();
assert_eq!(arr.len(), 2);
assert!(arr.iter().any(|v| v["spec_id"] == "sdd-workflow"));
}
#[test]
fn test_get_document_structure_strips_statement() {
let tree = build_test_tree();
let out: serde_json::Value = serde_json::from_str(&dispatch_tool(
"get_document_structure",
r#"{"spec_id":"sdd-workflow"}"#,
&tree,
))
.unwrap();
assert_eq!(out["spec_id"], "sdd-workflow");
let reqs = out["reqs"].as_array().unwrap();
assert_eq!(reqs.len(), 2);
assert!(reqs[0]["title"].as_str().is_some());
assert!(reqs[0].get("statement").is_none());
}
#[test]
fn test_get_spec_content_returns_statements() {
let tree = build_test_tree();
let out: serde_json::Value = serde_json::from_str(&dispatch_tool(
"get_spec_content",
r#"{"spec_id":"sdd-workflow","req_ids":["r12"]}"#,
&tree,
))
.unwrap();
let arr = out.as_array().unwrap();
assert_eq!(arr.len(), 1);
assert_eq!(arr[0]["req_id"], "r12");
assert!(arr[0]["statement"].as_str().unwrap().contains("MUST"));
}
#[test]
fn test_unknown_tool_errors() {
let tree = build_test_tree();
let out: serde_json::Value =
serde_json::from_str(&dispatch_tool("bogus", "{}", &tree)).unwrap();
assert!(out["error"].as_str().unwrap().contains("unknown tool"));
}
#[test]
fn test_parse_final_answer_plain_and_fenced() {
let plain = r#"{"direct":[{"id":"sdd-workflow","reason":"x"}],"related":[]}"#;
let out = parse_final_answer(plain).unwrap();
assert_eq!(out.direct.len(), 1);
assert_eq!(out.direct[0].id, "sdd-workflow");
let fenced = "Here you go:\n```json\n{\"direct\":[],\"related\":[]}\n```\n";
let out2 = parse_final_answer(fenced).unwrap();
assert!(out2.direct.is_empty());
}
struct ScriptedInvoker {
turns: Vec<ChatTurn>,
idx: AtomicUsize,
}
impl ScriptedInvoker {
fn new(turns: Vec<ChatTurn>) -> Self {
Self {
turns,
idx: AtomicUsize::new(0),
}
}
}
impl ChatInvoker for ScriptedInvoker {
async fn chat_turn(&self, _messages: &[Msg], _tools: &[ToolSchema]) -> Result<ChatTurn> {
let i = self.idx.fetch_add(1, Ordering::SeqCst);
let last = self.turns.len().saturating_sub(1);
Ok(self.turns[i.min(last)].clone())
}
}
fn block_on<F: std::future::Future>(fut: F) -> F::Output {
tokio::runtime::Builder::new_current_thread()
.build()
.unwrap()
.block_on(fut)
}
fn tool_call(id: &str, name: &str, args: &str) -> ChatTurn {
ChatTurn {
content: None,
tool_calls: vec![ToolCall {
id: id.to_string(),
name: name.to_string(),
arguments: args.to_string(),
}],
}
}
#[test]
fn test_agentic_loop_navigates_then_classifies() {
let tree = build_test_tree();
let invoker = ScriptedInvoker::new(vec![
tool_call("a", "list_specs", "{}"),
tool_call("b", "get_document_structure", r#"{"spec_id":"sdd-workflow"}"#),
tool_call("c", "get_spec_content", r#"{"spec_id":"sdd-workflow","req_ids":["r12"]}"#),
ChatTurn {
content: Some(
r#"{"direct":[{"id":"sdd-workflow","reason":"validate exit code lives here"}],"related":[{"id":"cli","reason":"cli surface"}]}"#.to_string(),
),
tool_calls: vec![],
},
]);
let out = block_on(retrieve(&invoker, &tree, "fix validate exit code", &[])).unwrap();
assert!(!out.truncated);
assert_eq!(out.tool_calls, 3);
assert_eq!(out.direct.len(), 1);
assert_eq!(out.direct[0].id, "sdd-workflow");
assert_eq!(out.related.len(), 1);
assert_eq!(out.related[0].id, "cli");
}
#[test]
fn test_agentic_loop_truncates_after_max_rounds() {
let tree = build_test_tree();
let invoker = ScriptedInvoker::new(vec![tool_call("x", "list_specs", "{}")]);
let out = block_on(retrieve(&invoker, &tree, "task", &[])).unwrap();
assert!(out.truncated, "loop must stop at MAX_TOOL_ROUNDS");
assert_eq!(out.tool_calls, MAX_TOOL_ROUNDS);
assert!(out.direct.is_empty(), "salvage turn had no JSON → empty");
}
#[test]
fn test_agentic_loop_salvages_partial_answer_on_truncation() {
let tree = build_test_tree();
let mut turns: Vec<ChatTurn> = (0..MAX_TOOL_ROUNDS)
.map(|_| tool_call("x", "list_specs", "{}"))
.collect();
turns.push(ChatTurn {
content: Some(
r#"{"direct":[{"id":"sdd-workflow","reason":"salvaged"}],"related":[]}"#
.to_string(),
),
tool_calls: vec![],
});
let invoker = ScriptedInvoker::new(turns);
let out = block_on(retrieve(&invoker, &tree, "task", &[])).unwrap();
assert!(out.truncated, "hit the round limit");
assert_eq!(out.tool_calls, MAX_TOOL_ROUNDS);
assert_eq!(out.direct.len(), 1, "salvage answer preserved");
assert_eq!(out.direct[0].id, "sdd-workflow");
}
}