use std::collections::BTreeSet;
use crate::api::{RedDBError, RedDBResult};
pub use crate::runtime::agent_tools::{QueryPhase, RedToolCall as PlannerToolCall};
use super::ask_rql_planner::{validate_candidate, CandidateDisposition, ValidatedCandidate};
pub const ASK_PLANNER_SYSTEM_PROMPT: &str = include_str!("prompts/ask_planner.txt");
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AskIntent {
Factual,
Synthesis,
HowTo,
}
impl AskIntent {
pub fn as_str(&self) -> &'static str {
match self {
AskIntent::Factual => "factual",
AskIntent::Synthesis => "synthesis",
AskIntent::HowTo => "how_to",
}
}
fn parse(raw: &str) -> Option<AskIntent> {
match raw.trim().to_ascii_lowercase().as_str() {
"factual" => Some(AskIntent::Factual),
"synthesis" => Some(AskIntent::Synthesis),
"how_to" | "howto" | "how-to" => Some(AskIntent::HowTo),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ScoredCollection {
pub collection: String,
pub score: f32,
pub columns: Vec<String>,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct NarrowedSlice {
pub collections: Vec<ScoredCollection>,
}
impl NarrowedSlice {
pub fn collection_names(&self) -> Vec<&str> {
self.collections
.iter()
.map(|c| c.collection.as_str())
.collect()
}
pub fn is_empty(&self) -> bool {
self.collections.is_empty()
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct PlannerObservation {
pub call: PlannerToolCall,
pub output: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RawSuggestion {
pub rql: String,
pub rationale: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SuggestedStatement {
pub rql: String,
pub mutating: bool,
pub statement_type: &'static str,
pub rationale: String,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AskPlan {
pub intent: AskIntent,
pub query: Option<String>,
pub tool: Option<PlannerToolCall>,
pub answer: String,
pub suggestion: Vec<RawSuggestion>,
pub rationale: String,
}
impl AskPlan {
pub fn summary(&self) -> String {
let mut out = format!("intent={}", self.intent.as_str());
if let Some(query) = &self.query {
out.push_str("; query=");
out.push_str(query.trim());
}
if let Some(tool) = &self.tool {
out.push_str("; tool=");
out.push_str(tool.name());
}
out
}
}
pub fn validate_suggestions(raw: &[RawSuggestion]) -> Vec<SuggestedStatement> {
let mut out = Vec::new();
for item in raw {
if let Ok(candidate) = validate_candidate(&item.rql) {
out.push(SuggestedStatement {
mutating: !candidate.is_read_only(),
statement_type: candidate.statement_type,
rql: candidate.rql,
rationale: item.rationale.clone(),
});
}
}
out
}
pub trait PlannerModel {
fn plan(&self, prompt: &str) -> RedDBResult<String>;
}
impl<F> PlannerModel for F
where
F: Fn(&str) -> RedDBResult<String>,
{
fn plan(&self, prompt: &str) -> RedDBResult<String> {
self(prompt)
}
}
pub trait PlannerChatModel {
fn chat(
&self,
messages: &[crate::ai::AiChatMessage],
tools: &[crate::ai::AiToolDefinition],
) -> RedDBResult<crate::ai::AiChatResponse>;
}
impl<F> PlannerChatModel for F
where
F: Fn(
&[crate::ai::AiChatMessage],
&[crate::ai::AiToolDefinition],
) -> RedDBResult<crate::ai::AiChatResponse>,
{
fn chat(
&self,
messages: &[crate::ai::AiChatMessage],
tools: &[crate::ai::AiToolDefinition],
) -> RedDBResult<crate::ai::AiChatResponse> {
self(messages, tools)
}
}
pub fn ask_tool_definitions() -> Vec<crate::ai::AiToolDefinition> {
crate::runtime::agent_tools::definitions(crate::runtime::agent_tools::ToolSurface::Ask)
.into_iter()
.map(|definition| crate::ai::AiToolDefinition {
name: definition.name.to_string(),
description: definition.description.to_string(),
input_schema: definition.input_schema,
})
.collect()
}
#[derive(Debug, Clone, PartialEq)]
pub enum PlanRouting {
Execute {
candidate: ValidatedCandidate,
params: Vec<crate::json::Value>,
purpose: String,
},
UseTool { call: PlannerToolCall },
RefuseMutating {
statement_type: &'static str,
rql: String,
},
Suggest {
answer: String,
suggestion: Vec<SuggestedStatement>,
},
Unsupported { intent: AskIntent },
}
#[derive(Debug, Clone, PartialEq)]
pub struct PlannedRoute {
pub prompt: String,
pub plan: AskPlan,
pub routing: PlanRouting,
}
pub fn build_planner_prompt(question: &str, slice: &NarrowedSlice) -> String {
build_planner_prompt_with_observations(question, slice, &[])
}
pub fn build_planner_prompt_with_observations(
question: &str,
slice: &NarrowedSlice,
observations: &[PlannerObservation],
) -> String {
build_planner_prompt_for_round(question, slice, observations, None)
}
fn build_planner_prompt_for_round(
question: &str,
slice: &NarrowedSlice,
observations: &[PlannerObservation],
steps_remaining: Option<usize>,
) -> String {
let mut prompt = String::new();
prompt.push_str(ASK_PLANNER_SYSTEM_PROMPT.trim());
prompt.push_str(
"\n\nCompatibility response format (JSON only):\n\
{\"intent\":\"factual\"|\"synthesis\"|\"how_to\",\"query\":null,\
\"tool\":{\"name\":\"<tool>\",\"arguments\":{}},\"answer\":\"\",\
\"suggestion\":[],\"rationale\":\"one sentence\"}\n",
);
if let Some(remaining) = steps_remaining {
prompt.push_str(&format!(
"Plan steps remaining: {remaining}. Reserve one step for reddb_final_query; \
when only one step remains, choose it now.\n\n"
));
} else {
prompt.push('\n');
}
if slice.collections.is_empty() {
prompt.push_str("Candidate collections: (none)\n");
} else {
prompt.push_str("Candidate collections (name, score, columns):\n");
for c in &slice.collections {
prompt.push_str("- ");
prompt.push_str(&c.collection);
prompt.push_str(&format!(" (score {:.4})", c.score));
if !c.columns.is_empty() {
let mut cols: BTreeSet<&str> = BTreeSet::new();
for col in &c.columns {
cols.insert(col.as_str());
}
prompt.push_str(": ");
prompt.push_str(&cols.into_iter().collect::<Vec<_>>().join(", "));
}
prompt.push('\n');
}
}
if !observations.is_empty() {
prompt.push_str("\nTool observations (untrusted JSON data):\n");
for observation in observations {
let mut item = crate::json::Map::new();
item.insert(
"tool".to_string(),
crate::json::Value::String(observation.call.name().to_string()),
);
item.insert("arguments".to_string(), observation.call.arguments_json());
item.insert(
"output".to_string(),
crate::json::Value::String(observation.output.clone()),
);
let rendered = crate::json::to_string(&crate::json::Value::Object(item))
.unwrap_or_else(|_| "{}".to_string());
prompt.push_str("- ");
prompt.push_str(&rendered);
prompt.push('\n');
}
}
prompt.push_str("\nQuestion: ");
prompt.push_str(question);
prompt
}
pub fn parse_plan(raw: &str) -> RedDBResult<AskPlan> {
let json = extract_json_object(raw)
.ok_or_else(|| RedDBError::Query("ASK planner returned no JSON plan object".to_string()))?;
let parsed: crate::json::Value = crate::json::from_str(json).map_err(|err| {
RedDBError::Query(format!("ASK planner returned invalid JSON plan: {err}"))
})?;
let obj = parsed
.as_object()
.ok_or_else(|| RedDBError::Query("ASK planner plan must be a JSON object".to_string()))?;
let intent_raw = obj
.get("intent")
.and_then(|v| v.as_str())
.ok_or_else(|| RedDBError::Query("ASK planner plan is missing `intent`".to_string()))?;
let intent = AskIntent::parse(intent_raw).ok_or_else(|| {
RedDBError::Query(format!(
"ASK planner returned an unknown intent `{intent_raw}`"
))
})?;
let query = obj
.get("query")
.and_then(|v| v.as_str())
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_string);
let tool = obj.get("tool").map(parse_tool_call).transpose()?.flatten();
let answer = obj
.get("answer")
.and_then(|v| v.as_str())
.unwrap_or("")
.trim()
.to_string();
let suggestion = obj
.get("suggestion")
.and_then(|v| v.as_array())
.map(parse_suggestions)
.unwrap_or_default();
let rationale = obj
.get("rationale")
.and_then(|v| v.as_str())
.unwrap_or("")
.trim()
.to_string();
Ok(AskPlan {
intent,
query,
tool,
answer,
suggestion,
rationale,
})
}
fn parse_tool_call(value: &crate::json::Value) -> RedDBResult<Option<PlannerToolCall>> {
if matches!(value, crate::json::Value::Null) {
return Ok(None);
}
let object = value.as_object().ok_or_else(|| {
RedDBError::Query("ASK planner `tool` must be an object or null".to_string())
})?;
let name = object
.get("name")
.and_then(crate::json::Value::as_str)
.ok_or_else(|| RedDBError::Query("ASK planner tool is missing `name`".to_string()))?;
let arguments = object
.get("arguments")
.cloned()
.unwrap_or_else(|| crate::json::Value::Object(crate::json::Map::new()));
crate::runtime::agent_tools::parse_call(
name,
&arguments,
crate::runtime::agent_tools::ToolParseMode::AskLegacy,
)
.map(Some)
}
fn parse_native_tool_call(call: &crate::ai::AiToolCall) -> RedDBResult<PlannerToolCall> {
crate::runtime::agent_tools::parse_call(
&call.name,
&call.arguments,
crate::runtime::agent_tools::ToolParseMode::AskNative,
)
}
fn parse_suggestions(items: &[crate::json::Value]) -> Vec<RawSuggestion> {
let mut out = Vec::new();
for item in items {
if let Some(rql) = item.as_str() {
let rql = rql.trim();
if !rql.is_empty() {
out.push(RawSuggestion {
rql: rql.to_string(),
rationale: String::new(),
});
}
} else if let Some(obj) = item.as_object() {
let rql = obj
.get("rql")
.or_else(|| obj.get("statement"))
.and_then(|v| v.as_str())
.unwrap_or("")
.trim()
.to_string();
if rql.is_empty() {
continue;
}
let rationale = obj
.get("rationale")
.and_then(|v| v.as_str())
.unwrap_or("")
.trim()
.to_string();
out.push(RawSuggestion { rql, rationale });
}
}
out
}
pub fn route_plan(plan: &AskPlan) -> RedDBResult<PlanRouting> {
match plan.intent {
AskIntent::Factual => match (&plan.query, &plan.tool) {
(Some(_), Some(_)) => Err(RedDBError::Query(
"ASK planner factual plan must choose exactly one of `query` or `tool`".to_string(),
)),
(Some(rql), None) => {
route_query_candidate(rql, &[], &plan.rationale, QueryPhase::Final)
}
(
None,
Some(PlannerToolCall::Query {
rql,
params,
purpose,
phase,
}),
) => route_query_candidate(rql, params, purpose, *phase),
(None, Some(call)) => Ok(PlanRouting::UseTool { call: call.clone() }),
(None, None) => Err(RedDBError::Query(
"ASK planner classified the question as factual but produced no query or \
tool call"
.to_string(),
)),
},
AskIntent::HowTo => {
let suggestion = validate_suggestions(&plan.suggestion);
Ok(PlanRouting::Suggest {
answer: plan.answer.clone(),
suggestion,
})
}
other => Ok(PlanRouting::Unsupported { intent: other }),
}
}
fn route_query_candidate(
rql: &str,
params: &[crate::json::Value],
purpose: &str,
phase: QueryPhase,
) -> RedDBResult<PlanRouting> {
let candidate = validate_candidate(rql)?;
match candidate.disposition {
CandidateDisposition::ReadOnly if phase == QueryPhase::Explore => {
Ok(PlanRouting::UseTool {
call: PlannerToolCall::Query {
rql: candidate.rql,
params: params.to_vec(),
purpose: purpose.to_string(),
phase,
},
})
}
CandidateDisposition::ReadOnly => Ok(PlanRouting::Execute {
candidate,
params: params.to_vec(),
purpose: purpose.to_string(),
}),
CandidateDisposition::Mutating => Ok(PlanRouting::RefuseMutating {
statement_type: candidate.statement_type,
rql: candidate.rql,
}),
}
}
pub fn plan_and_route<M: PlannerModel>(
question: &str,
slice: &NarrowedSlice,
model: &M,
) -> RedDBResult<PlannedRoute> {
plan_and_route_with_observations(question, slice, &[], model)
}
pub fn plan_and_route_with_observations<M: PlannerModel>(
question: &str,
slice: &NarrowedSlice,
observations: &[PlannerObservation],
model: &M,
) -> RedDBResult<PlannedRoute> {
let prompt = build_planner_prompt_with_observations(question, slice, observations);
let raw = model.plan(&prompt)?;
let plan = parse_plan(&raw)?;
let routing = route_plan(&plan)?;
Ok(PlannedRoute {
prompt,
plan,
routing,
})
}
#[derive(Debug, Clone, PartialEq)]
pub enum AgentLoopOutcome {
Routed {
route: PlannedRoute,
observations: Vec<PlannerObservation>,
},
BudgetExhausted(BudgetExhausted),
}
pub fn run_bounded_agent_loop<M, E>(
question: &str,
slice: &NarrowedSlice,
budget: &mut PlanBudget,
model: &M,
executor: &E,
) -> RedDBResult<AgentLoopOutcome>
where
M: PlannerModel,
E: Fn(&PlannerToolCall) -> RedDBResult<String>,
{
let mut observations = Vec::new();
loop {
if budget.remaining() == 0 {
return Ok(AgentLoopOutcome::BudgetExhausted(BudgetExhausted {
attempted: PlanStep::Query,
executed_steps: budget.executed_count(),
max_steps: budget.max_steps(),
}));
}
let prompt = build_planner_prompt_for_round(
question,
slice,
&observations,
Some(budget.remaining()),
);
let raw = model.plan(&prompt)?;
let plan = parse_plan(&raw)?;
let route = PlannedRoute {
prompt,
routing: route_plan(&plan)?,
plan,
};
match &route.routing {
PlanRouting::UseTool { call } => {
if budget.remaining() <= 1 {
return Ok(AgentLoopOutcome::BudgetExhausted(BudgetExhausted {
attempted: planner_tool_step(call),
executed_steps: budget.executed_count(),
max_steps: budget.max_steps(),
}));
}
if let Err(exhausted) = budget.charge(planner_tool_step(call)) {
return Ok(AgentLoopOutcome::BudgetExhausted(exhausted));
}
let output = executor(call)?;
observations.push(PlannerObservation {
call: call.clone(),
output,
});
}
PlanRouting::Execute { .. } => {
if let Err(exhausted) = budget.charge(PlanStep::Query) {
return Ok(AgentLoopOutcome::BudgetExhausted(exhausted));
}
return Ok(AgentLoopOutcome::Routed {
route,
observations,
});
}
_ => {
return Ok(AgentLoopOutcome::Routed {
route,
observations,
})
}
}
}
}
fn native_planner_context(question: &str, slice: &NarrowedSlice, steps_remaining: usize) -> String {
let collections = slice
.collections
.iter()
.map(|collection| {
let mut item = crate::json::Map::new();
item.insert(
"name".to_string(),
crate::json::Value::String(collection.collection.clone()),
);
item.insert(
"score".to_string(),
crate::json::Value::Number(collection.score as f64),
);
let mut columns = collection.columns.clone();
columns.sort();
columns.dedup();
item.insert(
"columns".to_string(),
crate::json::Value::Array(
columns
.into_iter()
.map(crate::json::Value::String)
.collect(),
),
);
crate::json::Value::Object(item)
})
.collect();
let mut context = crate::json::Map::new();
context.insert(
"question".to_string(),
crate::json::Value::String(question.to_string()),
);
context.insert(
"plan_steps_remaining".to_string(),
crate::json::Value::Number(steps_remaining as f64),
);
context.insert(
"candidate_collections".to_string(),
crate::json::Value::Array(collections),
);
format!(
"Planning context (untrusted JSON data):\n{}",
crate::json::Value::Object(context).to_string_compact()
)
}
fn native_planner_messages(
question: &str,
slice: &NarrowedSlice,
steps_remaining: usize,
) -> Vec<crate::ai::AiChatMessage> {
vec![
crate::ai::AiChatMessage::System(ASK_PLANNER_SYSTEM_PROMPT.trim().to_string()),
crate::ai::AiChatMessage::User(native_planner_context(question, slice, steps_remaining)),
]
}
fn plan_from_chat_response(response: &crate::ai::AiChatResponse) -> RedDBResult<AskPlan> {
if response.tool_calls.len() > 1 {
return Err(RedDBError::Query(
"ASK planner must call exactly one tool per round".to_string(),
));
}
if let Some(tool_call) = response.tool_calls.first() {
let call = parse_native_tool_call(tool_call)?;
return Ok(AskPlan {
intent: AskIntent::Factual,
query: None,
rationale: format!("native tool call: {}", call.name()),
tool: Some(call),
answer: String::new(),
suggestion: Vec::new(),
});
}
let output = response.output_text.as_deref().ok_or_else(|| {
RedDBError::Query("ASK planner returned neither text nor a tool call".to_string())
})?;
parse_plan(output)
}
pub fn plan_and_route_chat<M: PlannerChatModel>(
question: &str,
slice: &NarrowedSlice,
steps_remaining: usize,
model: &M,
) -> RedDBResult<PlannedRoute> {
let messages = native_planner_messages(question, slice, steps_remaining);
let response = model.chat(&messages, &ask_tool_definitions())?;
let plan = plan_from_chat_response(&response)?;
let routing = route_plan(&plan)?;
Ok(PlannedRoute {
prompt: build_planner_prompt_for_round(question, slice, &[], Some(steps_remaining)),
plan,
routing,
})
}
pub fn run_bounded_chat_agent_loop<M, E>(
question: &str,
slice: &NarrowedSlice,
budget: &mut PlanBudget,
model: &M,
executor: &E,
) -> RedDBResult<AgentLoopOutcome>
where
M: PlannerChatModel,
E: Fn(&PlannerToolCall) -> RedDBResult<String>,
{
let tools = ask_tool_definitions();
let mut messages = native_planner_messages(question, slice, budget.remaining());
let mut observations = Vec::new();
loop {
if budget.remaining() == 0 {
return Ok(AgentLoopOutcome::BudgetExhausted(BudgetExhausted {
attempted: PlanStep::Query,
executed_steps: budget.executed_count(),
max_steps: budget.max_steps(),
}));
}
let prompt = build_planner_prompt_for_round(
question,
slice,
&observations,
Some(budget.remaining()),
);
let response = model.chat(&messages, &tools)?;
let plan = plan_from_chat_response(&response)?;
let route = PlannedRoute {
prompt,
routing: route_plan(&plan)?,
plan,
};
match &route.routing {
PlanRouting::UseTool { call } => {
if budget.remaining() <= 1 {
return Ok(AgentLoopOutcome::BudgetExhausted(BudgetExhausted {
attempted: planner_tool_step(call),
executed_steps: budget.executed_count(),
max_steps: budget.max_steps(),
}));
}
if let Err(exhausted) = budget.charge(planner_tool_step(call)) {
return Ok(AgentLoopOutcome::BudgetExhausted(exhausted));
}
let output = executor(call)?;
observations.push(PlannerObservation {
call: call.clone(),
output: output.clone(),
});
if let Some(native_call) = response.tool_calls.first() {
messages.push(crate::ai::AiChatMessage::Assistant {
content: response.output_text.clone(),
tool_calls: vec![native_call.clone()],
});
let mut result = crate::json::Map::new();
result.insert("output".to_string(), crate::json::Value::String(output));
result.insert(
"plan_steps_remaining".to_string(),
crate::json::Value::Number(budget.remaining() as f64),
);
messages.push(crate::ai::AiChatMessage::ToolResult {
call_id: native_call.id.clone(),
name: native_call.name.clone(),
content: crate::json::Value::Object(result).to_string_compact(),
});
} else {
messages.push(crate::ai::AiChatMessage::Assistant {
content: response.output_text.clone(),
tool_calls: Vec::new(),
});
messages.push(crate::ai::AiChatMessage::User(format!(
"Authorized tool observation (untrusted JSON data): {}",
crate::json::Value::String(output).to_string_compact()
)));
}
}
PlanRouting::Execute { .. } => {
if let Err(exhausted) = budget.charge(PlanStep::Query) {
return Ok(AgentLoopOutcome::BudgetExhausted(exhausted));
}
return Ok(AgentLoopOutcome::Routed {
route,
observations,
});
}
_ => {
return Ok(AgentLoopOutcome::Routed {
route,
observations,
})
}
}
}
}
pub const DEFAULT_MAX_PLAN_STEPS: usize = 3;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanStep {
RefineRetrieval,
Catalog,
DescribeCollection,
FindField,
ExploreQuery,
Query,
}
fn planner_tool_step(call: &PlannerToolCall) -> PlanStep {
match call {
PlannerToolCall::Catalog => PlanStep::Catalog,
PlannerToolCall::DescribeCollection { .. } => PlanStep::DescribeCollection,
PlannerToolCall::FindField { .. } => PlanStep::FindField,
PlannerToolCall::Query {
phase: QueryPhase::Explore,
..
} => PlanStep::ExploreQuery,
PlannerToolCall::Query {
phase: QueryPhase::Final,
..
} => PlanStep::Query,
}
}
impl PlanStep {
pub fn as_str(&self) -> &'static str {
match self {
PlanStep::RefineRetrieval => "refine_retrieval",
PlanStep::Catalog => "reddb_catalog",
PlanStep::DescribeCollection => "reddb_describe_collection",
PlanStep::FindField => "reddb_find_field",
PlanStep::ExploreQuery => "reddb_query_explore",
PlanStep::Query => "query",
}
}
}
pub fn clamp_plan_steps(requested: Option<usize>, cap: usize) -> usize {
let cap = cap.max(1);
match requested {
Some(n) => n.clamp(1, cap),
None => cap,
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BudgetExhausted {
pub attempted: PlanStep,
pub executed_steps: usize,
pub max_steps: usize,
}
#[derive(Debug, Clone)]
pub struct PlanBudget {
max_steps: usize,
executed: Vec<PlanStep>,
}
impl PlanBudget {
pub fn new(requested_steps: Option<usize>, cap: usize) -> Self {
Self {
max_steps: clamp_plan_steps(requested_steps, cap),
executed: Vec::new(),
}
}
pub fn max_steps(&self) -> usize {
self.max_steps
}
pub fn executed_count(&self) -> usize {
self.executed.len()
}
pub fn remaining(&self) -> usize {
self.max_steps.saturating_sub(self.executed.len())
}
pub fn executed_steps(&self) -> &[PlanStep] {
&self.executed
}
pub fn charge(&mut self, step: PlanStep) -> Result<(), BudgetExhausted> {
if self.executed.len() >= self.max_steps {
return Err(BudgetExhausted {
attempted: step,
executed_steps: self.executed.len(),
max_steps: self.max_steps,
});
}
self.executed.push(step);
Ok(())
}
}
pub trait RetrievalFunnel {
fn funnel(&self, expanded: bool) -> RedDBResult<NarrowedSlice>;
}
impl<F> RetrievalFunnel for F
where
F: Fn(bool) -> RedDBResult<NarrowedSlice>,
{
fn funnel(&self, expanded: bool) -> RedDBResult<NarrowedSlice> {
self(expanded)
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum GroundingOutcome {
Grounded { slice: NarrowedSlice, refined: bool },
NoMatchingSources,
}
pub fn ground_with_refine<F: RetrievalFunnel>(funnel: &F) -> RedDBResult<GroundingOutcome> {
let first = funnel.funnel(false)?;
if !first.is_empty() {
return Ok(GroundingOutcome::Grounded {
slice: first,
refined: false,
});
}
let refined = funnel.funnel(true)?;
if !refined.is_empty() {
return Ok(GroundingOutcome::Grounded {
slice: refined,
refined: true,
});
}
Ok(GroundingOutcome::NoMatchingSources)
}
fn extract_json_object(raw: &str) -> Option<&str> {
let start = raw.find('{')?;
let end = raw.rfind('}')?;
if end <= start {
return None;
}
Some(&raw[start..=end])
}
#[cfg(test)]
mod tests {
use std::cell::RefCell;
use std::collections::VecDeque;
use super::*;
fn slice() -> NarrowedSlice {
NarrowedSlice {
collections: vec![
ScoredCollection {
collection: "travelers".to_string(),
score: 0.91,
columns: vec!["passport".to_string(), "name".to_string()],
},
ScoredCollection {
collection: "trips".to_string(),
score: 0.44,
columns: vec!["passport".to_string(), "city".to_string()],
},
],
}
}
fn mock_model(plan_json: &'static str) -> impl PlannerModel {
move |_prompt: &str| Ok(plan_json.to_string())
}
struct TestChatModel {
responses: RefCell<VecDeque<crate::ai::AiChatResponse>>,
requests: RefCell<Vec<(Vec<crate::ai::AiChatMessage>, Vec<String>)>>,
}
impl TestChatModel {
fn new(responses: Vec<crate::ai::AiChatResponse>) -> Self {
Self {
responses: RefCell::new(responses.into()),
requests: RefCell::new(Vec::new()),
}
}
}
impl PlannerChatModel for TestChatModel {
fn chat(
&self,
messages: &[crate::ai::AiChatMessage],
tools: &[crate::ai::AiToolDefinition],
) -> RedDBResult<crate::ai::AiChatResponse> {
self.requests.borrow_mut().push((
messages.to_vec(),
tools.iter().map(|tool| tool.name.clone()).collect(),
));
self.responses
.borrow_mut()
.pop_front()
.ok_or_else(|| RedDBError::Internal("test model script exhausted".to_string()))
}
}
fn tool_response(
id: &str,
name: &str,
arguments: Vec<(&str, crate::json::Value)>,
) -> crate::ai::AiChatResponse {
let mut object = crate::json::Map::new();
for (key, value) in arguments {
object.insert(key.to_string(), value);
}
crate::ai::AiChatResponse {
provider: "test",
model: "test-model".to_string(),
output_text: None,
tool_calls: vec![crate::ai::AiToolCall {
id: id.to_string(),
name: name.to_string(),
arguments: crate::json::Value::Object(object),
}],
prompt_tokens: None,
completion_tokens: None,
total_tokens: None,
stop_reason: Some("tool_use".to_string()),
}
}
#[test]
fn prompt_contains_only_narrowed_slice() {
let prompt = build_planner_prompt("who owns passport FDD-1?", &slice());
assert!(prompt.contains("travelers"));
assert!(prompt.contains("trips"));
assert!(prompt.contains("passport"));
assert!(!prompt.contains("secret_ledger"));
}
#[test]
fn native_harness_keeps_tool_transcript_and_requires_final_query() {
let model = TestChatModel::new(vec![
tool_response(
"call_1",
"reddb_find_field",
vec![("field", crate::json::Value::String("passport".to_string()))],
),
tool_response(
"call_2",
"reddb_final_query",
vec![
(
"sql",
crate::json::Value::String(
"SELECT name FROM travelers WHERE passport = $1".to_string(),
),
),
(
"params",
crate::json::Value::Array(vec![crate::json::Value::String(
"FDD-1".to_string(),
)]),
),
(
"purpose",
crate::json::Value::String("answer the passport lookup".to_string()),
),
],
),
]);
let mut budget = PlanBudget::new(Some(3), 3);
let executor = |call: &PlannerToolCall| {
assert!(matches!(call, PlannerToolCall::FindField { .. }));
Ok("[{\"collection\":\"travelers\",\"name\":\"passport\"}]".to_string())
};
let outcome = run_bounded_chat_agent_loop(
"who owns passport FDD-1?",
&slice(),
&mut budget,
&model,
&executor,
)
.expect("native harness");
let AgentLoopOutcome::Routed {
route,
observations,
} = outcome
else {
panic!("expected routed outcome")
};
assert_eq!(observations.len(), 1);
assert_eq!(
budget.executed_steps(),
&[PlanStep::FindField, PlanStep::Query]
);
assert!(matches!(route.routing, PlanRouting::Execute { .. }));
let requests = model.requests.borrow();
assert_eq!(requests.len(), 2);
assert!(requests[0].1.iter().any(|name| name == "reddb_final_query"));
assert!(requests[1]
.0
.iter()
.any(|message| matches!(message, crate::ai::AiChatMessage::ToolResult { call_id, .. } if call_id == "call_1")));
}
#[test]
fn native_query_tool_cannot_smuggle_a_final_phase() {
let call = crate::ai::AiToolCall {
id: "call_1".to_string(),
name: "reddb_query".to_string(),
arguments: {
let mut arguments = crate::json::Map::new();
arguments.insert(
"sql".to_string(),
crate::json::Value::String("SELECT 1".to_string()),
);
arguments.insert("params".to_string(), crate::json::Value::Array(Vec::new()));
arguments.insert(
"purpose".to_string(),
crate::json::Value::String("inspect".to_string()),
);
arguments.insert(
"phase".to_string(),
crate::json::Value::String("final".to_string()),
);
crate::json::Value::Object(arguments)
},
};
assert!(matches!(
parse_native_tool_call(&call).expect("native call"),
PlannerToolCall::Query {
phase: QueryPhase::Explore,
..
}
));
}
#[test]
fn parse_plan_reads_intent_and_query() {
let plan = parse_plan(
"{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"lookup\"}",
)
.unwrap();
assert_eq!(plan.intent, AskIntent::Factual);
assert_eq!(
plan.query.as_deref(),
Some("SELECT * FROM travelers WHERE passport = 'FDD-1'")
);
assert_eq!(plan.rationale, "lookup");
}
#[test]
fn metadata_tool_call_routes_to_a_bounded_agent_step() {
let route = plan_and_route(
"which collections contain passport_number?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":null,\"tool\":{\"name\":\"reddb_find_field\",\"arguments\":{\"field\":\"passport_number\"}},\"rationale\":\"discover schema\"}",
),
)
.expect("metadata tool plan must route");
assert_eq!(
route.routing,
PlanRouting::UseTool {
call: PlannerToolCall::FindField {
field: "passport_number".to_string()
}
}
);
}
#[test]
fn query_tool_call_uses_the_same_read_only_validator() {
let route = plan_and_route(
"who owns passport FDD-1?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":null,\"tool\":{\"name\":\"reddb_query\",\"arguments\":{\"sql\":\"SELECT * FROM travelers WHERE passport = $1\",\"params\":[\"FDD-1\"],\"purpose\":\"find the traveler\",\"phase\":\"final\"}}}",
),
)
.expect("read-only query tool must route");
match route.routing {
PlanRouting::Execute {
params, purpose, ..
} => {
assert_eq!(
params,
vec![crate::json::Value::String("FDD-1".to_string())]
);
assert_eq!(purpose, "find the traveler");
}
other => panic!("expected final query, got {other:?}"),
}
}
#[test]
fn exploratory_query_cannot_bypass_the_mutation_refusal() {
let route = plan_and_route(
"remove FDD-1",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"tool\":{\"name\":\"reddb_query\",\"arguments\":{\"sql\":\"DELETE FROM travelers WHERE passport = $1\",\"params\":[\"FDD-1\"],\"purpose\":\"remove a row\",\"phase\":\"explore\"}}}",
),
)
.expect("mutating tool calls become structured refusals");
assert!(matches!(
route.routing,
PlanRouting::RefuseMutating {
statement_type: "delete",
..
}
));
}
#[test]
fn query_tool_rejects_unbounded_or_invalid_params() {
let not_an_array = parse_plan(
"{\"intent\":\"factual\",\"tool\":{\"name\":\"reddb_query\",\"arguments\":{\"sql\":\"SELECT 1\",\"params\":\"secret\",\"phase\":\"final\"}}}",
)
.expect_err("params must be an array");
assert!(not_an_array.to_string().contains("must be an array"));
let params = vec![
crate::json::Value::Null;
crate::runtime::agent_tools::ASK_QUERY_PARAMETER_COUNT_MAX + 1
];
let mut arguments = crate::json::Map::new();
arguments.insert(
"sql".to_string(),
crate::json::Value::String("SELECT 1".to_string()),
);
arguments.insert("params".to_string(), crate::json::Value::Array(params));
let mut tool = crate::json::Map::new();
tool.insert(
"name".to_string(),
crate::json::Value::String("reddb_query".to_string()),
);
tool.insert(
"arguments".to_string(),
crate::json::Value::Object(arguments),
);
let mut root = crate::json::Map::new();
root.insert(
"intent".to_string(),
crate::json::Value::String("factual".to_string()),
);
root.insert("tool".to_string(), crate::json::Value::Object(tool));
let raw = crate::json::to_string(&crate::json::Value::Object(root)).unwrap();
let oversized = parse_plan(&raw).expect_err("parameter count must be bounded");
assert!(oversized.to_string().contains("exceeds 64 values"));
}
#[test]
fn planner_rejects_unknown_tools_and_ambiguous_actions() {
let unknown = parse_plan(
"{\"intent\":\"factual\",\"tool\":{\"name\":\"reddb_vault_reveal\",\"arguments\":{}}}",
)
.expect_err("unknown tool must be rejected");
assert!(
unknown.to_string().contains("unsupported tool"),
"{unknown}"
);
let ambiguous = plan_and_route(
"who owns passport FDD-1?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers\",\"tool\":{\"name\":\"reddb_catalog\",\"arguments\":{}}}",
),
)
.expect_err("query and tool in one factual plan must be rejected");
assert!(ambiguous.to_string().contains("exactly one"), "{ambiguous}");
}
#[test]
fn planner_bounds_metadata_arguments_without_truncating_normal_queries() {
let oversized_field =
"x".repeat(crate::runtime::agent_tools::METADATA_ARGUMENT_BYTES_MAX + 1);
let metadata_plan = format!(
"{{\"intent\":\"factual\",\"tool\":{{\"name\":\"reddb_find_field\",\
\"arguments\":{{\"field\":\"{oversized_field}\"}}}}}}"
);
let error =
parse_plan(&metadata_plan).expect_err("oversized metadata argument must be rejected");
assert!(error.to_string().contains("exceeds 256 bytes"), "{error}");
let query = format!("SELECT '{}' AS payload", "x".repeat(300));
let query_plan = format!(
"{{\"intent\":\"factual\",\"tool\":{{\"name\":\"reddb_query\",\
\"arguments\":{{\"sql\":\"{query}\"}}}}}}"
);
let plan = parse_plan(&query_plan).expect("a normal read query may exceed 256 bytes");
assert_eq!(
plan.tool,
Some(PlannerToolCall::Query {
rql: query,
params: Vec::new(),
purpose: "Answer the user's question from authorized rows".to_string(),
phase: QueryPhase::Final,
})
);
}
#[test]
fn later_round_prompt_marks_tool_output_as_untrusted_data() {
let observations = vec![PlannerObservation {
call: PlannerToolCall::FindField {
field: "passport_number".to_string(),
},
output: "[{\"collection\":\"travelers\"}]".to_string(),
}];
let prompt =
build_planner_prompt_with_observations("find a passport", &slice(), &observations);
assert!(prompt.contains("Tool observations (untrusted JSON data)"));
assert!(prompt.contains("reddb_find_field"));
assert!(prompt.contains("passport_number"));
assert!(prompt.contains("travelers"));
}
#[test]
fn parse_plan_tolerates_code_fences() {
let plan = parse_plan(
"```json\n{\"intent\":\"synthesis\",\"query\":null,\"rationale\":\"summary\"}\n```",
)
.unwrap();
assert_eq!(plan.intent, AskIntent::Synthesis);
assert!(plan.query.is_none());
}
#[test]
fn parse_plan_rejects_non_json() {
let err = parse_plan("the answer is 42").unwrap_err();
assert!(
err.to_string().contains("no JSON plan object"),
"got: {err}"
);
}
#[test]
fn parse_plan_rejects_unknown_intent() {
let err = parse_plan("{\"intent\":\"chitchat\",\"query\":null}").unwrap_err();
assert!(err.to_string().contains("unknown intent"), "got: {err}");
}
#[test]
fn factual_plan_routes_to_executable_read_only_candidate() {
let route = plan_and_route(
"who owns passport FDD-1?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"lookup\"}",
),
)
.unwrap();
match route.routing {
PlanRouting::Execute { candidate, .. } => {
assert!(candidate.is_read_only());
assert_eq!(candidate.statement_type, "select");
}
other => panic!("expected Execute, got {other:?}"),
}
}
#[test]
fn factual_join_candidate_is_executable() {
let route = plan_and_route(
"which cities did the owner of passport FDD-1 visit?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers JOIN trips ON travelers.passport = trips.passport WHERE travelers.passport = 'FDD-1'\",\"rationale\":\"join\"}",
),
)
.unwrap();
assert!(matches!(route.routing, PlanRouting::Execute { .. }));
}
#[test]
fn factual_plan_with_mutating_candidate_is_refused_not_executed() {
let route = plan_and_route(
"delete traveler FDD-1",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":\"DELETE FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"oops\"}",
),
)
.unwrap();
match route.routing {
PlanRouting::RefuseMutating { statement_type, .. } => {
assert_eq!(statement_type, "delete")
}
other => panic!("expected RefuseMutating, got {other:?}"),
}
}
#[test]
fn factual_plan_with_malformed_candidate_is_rejected() {
let err = plan_and_route(
"who owns passport FDD-1?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":\"this is not rql\",\"rationale\":\"x\"}",
),
)
.unwrap_err();
assert!(
err.to_string().contains("invalid RQL candidate"),
"got: {err}"
);
}
#[test]
fn factual_plan_without_query_is_rejected() {
let err = plan_and_route(
"who owns passport FDD-1?",
&slice(),
&mock_model("{\"intent\":\"factual\",\"query\":null,\"rationale\":\"x\"}"),
)
.unwrap_err();
assert!(
err.to_string().contains("no query or tool call"),
"got: {err}"
);
}
#[test]
fn synthesis_intent_routes_unsupported_in_this_slice() {
let route = plan_and_route(
"summarise yesterday's incidents",
&slice(),
&mock_model("{\"intent\":\"synthesis\",\"query\":null,\"rationale\":\"rag\"}"),
)
.unwrap();
assert_eq!(
route.routing,
PlanRouting::Unsupported {
intent: AskIntent::Synthesis
}
);
}
#[test]
fn synthesis_intent_carries_the_routed_intent_for_the_rag_fallthrough() {
let route = plan_and_route(
"summarise the incidents from yesterday",
&slice(),
&mock_model("{\"intent\":\"synthesis\",\"query\":null,\"rationale\":\"rag\"}"),
)
.unwrap();
match route.routing {
PlanRouting::Unsupported { intent } => {
assert_eq!(intent, AskIntent::Synthesis);
assert_eq!(intent.as_str(), "synthesis");
}
other => panic!("expected Unsupported{{synthesis}}, got {other:?}"),
}
}
#[test]
fn calculation_question_routes_factual_not_synthesis() {
let route = plan_and_route(
"how many trips went to Lisbon?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":\"SELECT COUNT(*) FROM trips WHERE city = 'Lisbon'\",\"rationale\":\"aggregate count\"}",
),
)
.unwrap();
assert_ne!(
route.plan.intent,
AskIntent::Synthesis,
"a calculation must never classify as synthesis (the LLM must not invent numbers)"
);
match route.routing {
PlanRouting::Execute { candidate, .. } => {
assert!(candidate.is_read_only());
assert_eq!(candidate.statement_type, "select");
}
other => {
panic!("a calculation must route to an executable factual candidate, got {other:?}")
}
}
}
#[test]
fn how_to_intent_routes_to_a_validated_suggestion_envelope() {
let route = plan_and_route(
"how would I capture events from orders into a queue?",
&slice(),
&mock_model(
"{\"intent\":\"how_to\",\"answer\":\"Create a queue and backfill events into it.\",\
\"suggestion\":[\
{\"rql\":\"CREATE QUEUE events_q WORK\",\"rationale\":\"the sink queue\"},\
{\"rql\":\"EVENTS BACKFILL orders TO events_q\",\"rationale\":\"seed history\"},\
{\"rql\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"inspect\"}\
],\"rationale\":\"guide\"}",
),
)
.unwrap();
match route.routing {
PlanRouting::Suggest { answer, suggestion } => {
assert!(answer.contains("Create a queue"));
assert_eq!(suggestion.len(), 3);
assert!(suggestion[0].mutating, "CREATE QUEUE is mutating DDL");
assert!(suggestion[1].mutating, "EVENTS BACKFILL is mutating");
assert!(!suggestion[2].mutating, "the SELECT is read-only");
assert_eq!(suggestion[2].statement_type, "select");
assert_eq!(suggestion[0].rationale, "the sink queue");
}
other => panic!("expected Suggest, got {other:?}"),
}
}
#[test]
fn how_to_suggestion_drops_unparseable_statements_never_returns_them_raw() {
let route = plan_and_route(
"how do I list travelers?",
&slice(),
&mock_model(
"{\"intent\":\"how_to\",\"answer\":\"Select from the collection.\",\
\"suggestion\":[\
{\"rql\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\"},\
{\"rql\":\"this is not rql at all\"},\
\"DELETE FROM travelers WHERE passport = 'FDD-1'\"\
]}",
),
)
.unwrap();
match route.routing {
PlanRouting::Suggest { suggestion, .. } => {
assert_eq!(suggestion.len(), 2, "the unparseable statement is dropped");
assert_eq!(suggestion[0].statement_type, "select");
assert_eq!(suggestion[1].statement_type, "delete");
assert!(suggestion[1].mutating);
for s in &suggestion {
assert!(
!s.rql.contains("not rql"),
"raw unparseable text must never survive"
);
}
}
other => panic!("expected Suggest, got {other:?}"),
}
}
#[test]
fn routing_distinguishes_how_to_from_factual_at_the_closure_model_seam() {
let factual = plan_and_route(
"how would I capture events into a queue?",
&slice(),
&mock_model(
"{\"intent\":\"factual\",\"query\":\"SELECT * FROM travelers WHERE passport = 'FDD-1'\",\"rationale\":\"x\"}",
),
)
.unwrap();
assert!(matches!(factual.routing, PlanRouting::Execute { .. }));
let how_to = plan_and_route(
"how would I capture events into a queue?",
&slice(),
&mock_model(
"{\"intent\":\"how_to\",\"answer\":\"Use a queue.\",\"suggestion\":[\"CREATE QUEUE events_q WORK\"],\"rationale\":\"x\"}",
),
)
.unwrap();
assert!(matches!(how_to.routing, PlanRouting::Suggest { .. }));
}
#[test]
fn steps_clause_is_clamped_to_the_config_cap_never_exceeding_it() {
assert_eq!(clamp_plan_steps(Some(9), 3), 3);
assert_eq!(clamp_plan_steps(Some(2), 3), 2);
assert_eq!(clamp_plan_steps(Some(3), 3), 3);
assert_eq!(clamp_plan_steps(None, 3), 3);
assert_eq!(clamp_plan_steps(Some(0), 3), 1);
assert_eq!(clamp_plan_steps(None, 0), 1);
}
#[test]
fn plan_budget_charges_until_exhausted_then_refuses_more() {
let mut budget = PlanBudget::new(Some(2), 3);
assert_eq!(budget.max_steps(), 2);
assert_eq!(budget.remaining(), 2);
assert!(budget.charge(PlanStep::RefineRetrieval).is_ok());
assert_eq!(budget.remaining(), 1);
assert!(budget.charge(PlanStep::Query).is_ok());
assert_eq!(budget.remaining(), 0);
assert_eq!(
budget.executed_steps(),
&[PlanStep::RefineRetrieval, PlanStep::Query]
);
let err = budget.charge(PlanStep::Query).unwrap_err();
assert_eq!(err.max_steps, 2);
assert_eq!(err.executed_steps, 2);
assert_eq!(err.attempted, PlanStep::Query);
assert_eq!(budget.executed_count(), 2);
}
#[test]
fn plan_budget_request_above_cap_cannot_buy_extra_steps() {
let mut budget = PlanBudget::new(Some(5), 1);
assert_eq!(budget.max_steps(), 1);
assert!(budget.charge(PlanStep::Query).is_ok());
assert!(budget.charge(PlanStep::RefineRetrieval).is_err());
}
#[test]
fn metadata_calls_and_terminal_query_share_one_step_budget() {
let mut budget = PlanBudget::new(Some(3), 3);
budget
.charge(PlanStep::Catalog)
.expect("catalog step fits budget");
budget
.charge(PlanStep::FindField)
.expect("field step fits budget");
budget
.charge(PlanStep::Query)
.expect("query step fits budget");
assert_eq!(budget.remaining(), 0);
assert_eq!(
budget.executed_steps(),
&[PlanStep::Catalog, PlanStep::FindField, PlanStep::Query]
);
assert!(budget.charge(PlanStep::DescribeCollection).is_err());
}
#[test]
fn exploratory_and_final_queries_are_distinct_budgeted_actions() {
let mut budget = PlanBudget::new(Some(2), 2);
budget
.charge(PlanStep::ExploreQuery)
.expect("exploration fits budget");
budget
.charge(PlanStep::Query)
.expect("final query fits budget");
assert_eq!(
budget.executed_steps(),
&[PlanStep::ExploreQuery, PlanStep::Query]
);
}
#[test]
fn bounded_agent_loop_feeds_metadata_to_the_next_round_then_routes_query() {
let responses = std::rc::Rc::new(std::cell::RefCell::new(
std::collections::VecDeque::from([
"{\"intent\":\"factual\",\"query\":null,\"tool\":{\"name\":\"reddb_find_field\",\"arguments\":{\"field\":\"passport_number\"}}}"
.to_string(),
"{\"intent\":\"factual\",\"query\":\"SELECT name FROM travelers WHERE passport_number = 'FDD-1'\",\"tool\":null}"
.to_string(),
]),
));
let prompts = std::rc::Rc::new(std::cell::RefCell::new(Vec::new()));
let responses_inner = responses.clone();
let prompts_inner = prompts.clone();
let model = move |prompt: &str| {
prompts_inner.borrow_mut().push(prompt.to_string());
Ok(responses_inner
.borrow_mut()
.pop_front()
.expect("scripted planner response"))
};
let calls = std::rc::Rc::new(std::cell::RefCell::new(Vec::new()));
let calls_inner = calls.clone();
let executor = move |call: &PlannerToolCall| {
calls_inner.borrow_mut().push(call.clone());
Ok("{\"rows\":[{\"collection\":\"travelers\"}],\"truncated\":false}".to_string())
};
let mut budget = PlanBudget::new(Some(2), 3);
let outcome =
run_bounded_agent_loop("who owns FDD-1?", &slice(), &mut budget, &model, &executor)
.expect("bounded loop must reach the terminal query");
assert!(matches!(
outcome,
AgentLoopOutcome::Routed {
route: PlannedRoute {
routing: PlanRouting::Execute { .. },
..
},
..
}
));
assert_eq!(
budget.executed_steps(),
&[PlanStep::FindField, PlanStep::Query]
);
assert_eq!(
calls.borrow().as_slice(),
&[PlannerToolCall::FindField {
field: "passport_number".to_string()
}]
);
let prompts = prompts.borrow();
assert_eq!(prompts.len(), 2);
assert!(prompts[1].contains("Tool observations (untrusted JSON data)"));
assert!(prompts[1].contains("travelers"));
}
#[test]
fn bounded_agent_loop_executes_parameterized_exploration_before_final_query() {
let responses = std::rc::Rc::new(std::cell::RefCell::new(
std::collections::VecDeque::from([
"{\"intent\":\"factual\",\"tool\":{\"name\":\"reddb_query\",\"arguments\":{\"sql\":\"SELECT id FROM travelers WHERE passport = $1 LIMIT 5\",\"params\":[\"FDD-1\"],\"purpose\":\"resolve person id\",\"phase\":\"explore\"}}}"
.to_string(),
"{\"intent\":\"factual\",\"tool\":{\"name\":\"reddb_query\",\"arguments\":{\"sql\":\"SELECT name FROM travelers WHERE passport = $1\",\"params\":[\"FDD-1\"],\"purpose\":\"answer who owns the passport\",\"phase\":\"final\"}}}"
.to_string(),
]),
));
let responses_inner = responses.clone();
let model = move |_prompt: &str| {
Ok(responses_inner
.borrow_mut()
.pop_front()
.expect("scripted planner response"))
};
let calls = std::rc::Rc::new(std::cell::RefCell::new(Vec::new()));
let calls_inner = calls.clone();
let executor = move |call: &PlannerToolCall| {
calls_inner.borrow_mut().push(call.clone());
Ok("{\"rows\":[{\"id\":\"t1\"}],\"truncated\":false}".to_string())
};
let mut budget = PlanBudget::new(Some(2), 2);
let outcome =
run_bounded_agent_loop("who owns FDD-1?", &slice(), &mut budget, &model, &executor)
.expect("exploration must lead to a final query");
let AgentLoopOutcome::Routed {
route,
observations,
} = outcome
else {
panic!("expected routed outcome");
};
assert_eq!(observations.len(), 1);
assert_eq!(calls.borrow().len(), 1);
assert_eq!(
budget.executed_steps(),
&[PlanStep::ExploreQuery, PlanStep::Query]
);
match route.routing {
PlanRouting::Execute {
candidate,
params,
purpose,
} => {
assert_eq!(
candidate.rql,
"SELECT name FROM travelers WHERE passport = $1"
);
assert_eq!(
params,
vec![crate::json::Value::String("FDD-1".to_string())]
);
assert_eq!(purpose, "answer who owns the passport");
}
other => panic!("expected final query, got {other:?}"),
}
}
#[test]
fn bounded_agent_loop_does_not_call_the_model_after_the_last_tool_step() {
let planner_calls = std::rc::Rc::new(std::cell::Cell::new(0usize));
let planner_calls_inner = planner_calls.clone();
let model = move |_prompt: &str| {
planner_calls_inner.set(planner_calls_inner.get() + 1);
Ok("{\"intent\":\"factual\",\"query\":null,\"tool\":{\"name\":\"reddb_catalog\",\"arguments\":{}}}"
.to_string())
};
let tool_calls = std::rc::Rc::new(std::cell::Cell::new(0usize));
let tool_calls_inner = tool_calls.clone();
let executor = move |_call: &PlannerToolCall| {
tool_calls_inner.set(tool_calls_inner.get() + 1);
Ok("{\"rows\":[],\"truncated\":false}".to_string())
};
let mut budget = PlanBudget::new(Some(1), 3);
let outcome = run_bounded_agent_loop(
"what data is available?",
&slice(),
&mut budget,
&model,
&executor,
)
.expect("bounded loop must return structured exhaustion");
assert_eq!(
outcome,
AgentLoopOutcome::BudgetExhausted(BudgetExhausted {
attempted: PlanStep::Catalog,
executed_steps: 0,
max_steps: 1,
})
);
assert_eq!(planner_calls.get(), 1, "no extra planner round");
assert_eq!(tool_calls.get(), 0, "the final step stays reserved");
assert!(budget.executed_steps().is_empty());
}
fn scripted_funnel(
script: Vec<bool>,
) -> (impl RetrievalFunnel, std::rc::Rc<std::cell::Cell<usize>>) {
let calls = std::rc::Rc::new(std::cell::Cell::new(0usize));
let calls_inner = calls.clone();
let funnel = move |_expanded: bool| -> RedDBResult<NarrowedSlice> {
let n = calls_inner.get();
calls_inner.set(n + 1);
let non_empty = script.get(n).copied().unwrap_or(false);
Ok(if non_empty {
slice()
} else {
NarrowedSlice::default()
})
};
(funnel, calls)
}
#[test]
fn grounding_succeeds_on_first_pass_without_refining() {
let (funnel, calls) = scripted_funnel(vec![true]);
let outcome = ground_with_refine(&funnel).unwrap();
assert_eq!(calls.get(), 1, "no refine when the first pass grounds");
match outcome {
GroundingOutcome::Grounded { refined, .. } => assert!(!refined),
other => panic!("expected Grounded, got {other:?}"),
}
}
#[test]
fn empty_first_pass_triggers_exactly_one_refine_retrieval() {
let (funnel, calls) = scripted_funnel(vec![false, true]);
let outcome = ground_with_refine(&funnel).unwrap();
assert_eq!(calls.get(), 2, "exactly one refine re-funnel");
match outcome {
GroundingOutcome::Grounded { refined, .. } => assert!(refined),
other => panic!("expected Grounded after refine, got {other:?}"),
}
}
#[test]
fn second_grounding_failure_returns_no_matching_sources_not_an_invented_answer() {
let (funnel, calls) = scripted_funnel(vec![false, false]);
let outcome = ground_with_refine(&funnel).unwrap();
assert_eq!(calls.get(), 2, "exactly one refine retry, then give up");
assert_eq!(outcome, GroundingOutcome::NoMatchingSources);
}
#[test]
fn plan_summary_includes_intent_and_query() {
let plan = AskPlan {
intent: AskIntent::Factual,
query: Some("SELECT * FROM travelers WHERE passport = 'FDD-1'".to_string()),
tool: None,
answer: String::new(),
suggestion: Vec::new(),
rationale: "lookup".to_string(),
};
let summary = plan.summary();
assert!(summary.contains("intent=factual"));
assert!(summary.contains("SELECT * FROM travelers"));
}
}