use std::collections::HashSet;
use std::io;
use std::sync::Arc;
use futures::StreamExt;
use serde::Deserialize;
use bamboo_agent_core::Message;
use bamboo_domain::ReasoningEffort;
use bamboo_llm::{LLMChunk, LLMProvider, LLMRequestOptions};
use bamboo_memory::memory_store::{
shortlist_relevant_memories, MemoryRecallCandidate, MemoryRecallOptions, MemoryStore,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum MemoryRecallStrategy {
Lexical,
Reranked,
RerankFallback,
}
impl MemoryRecallStrategy {
pub(super) const fn as_str(self) -> &'static str {
match self {
Self::Lexical => "lexical",
Self::Reranked => "reranked",
Self::RerankFallback => "rerank_fallback",
}
}
}
pub(super) struct MemoryRecallSelection {
pub(super) candidates: Vec<MemoryRecallCandidate>,
pub(super) strategy: MemoryRecallStrategy,
}
#[derive(Clone)]
pub(super) struct MemoryRecallRerankContext {
pub(super) llm: Arc<dyn LLMProvider>,
pub(super) model: String,
pub(super) session_id: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct MemoryRecallRerankEnvelope {
ids: Vec<String>,
}
pub(super) async fn select_relevant_memories(
store: &MemoryStore,
project_key: Option<&str>,
query: &str,
options: &MemoryRecallOptions,
rerank_context: Option<&MemoryRecallRerankContext>,
) -> io::Result<MemoryRecallSelection> {
let query = query.trim();
if query.is_empty() {
return Ok(MemoryRecallSelection {
candidates: Vec::new(),
strategy: MemoryRecallStrategy::Lexical,
});
}
let limit = options.shortlist_limit.max(1);
let candidate_limit = options.max_candidates_per_scope.max(limit);
let candidate_options = MemoryRecallOptions {
shortlist_limit: candidate_limit,
include_global_fallback: options.include_global_fallback,
max_candidates_per_scope: candidate_limit,
};
let mut shortlist =
shortlist_relevant_memories(store, project_key, query, &candidate_options).await?;
if shortlist.is_empty() {
return Ok(MemoryRecallSelection {
candidates: shortlist,
strategy: MemoryRecallStrategy::Lexical,
});
}
let Some(rerank_context) = rerank_context else {
shortlist.truncate(limit);
return Ok(MemoryRecallSelection {
candidates: shortlist,
strategy: MemoryRecallStrategy::Lexical,
});
};
if shortlist.len() <= 1 {
shortlist.truncate(limit);
return Ok(MemoryRecallSelection {
candidates: shortlist,
strategy: MemoryRecallStrategy::Lexical,
});
}
match rerank_candidate_ids(query, &shortlist, limit, rerank_context).await {
Ok(ids) if ids.is_empty() => Ok(MemoryRecallSelection {
candidates: Vec::new(),
strategy: MemoryRecallStrategy::Reranked,
}),
Ok(ids) => Ok(MemoryRecallSelection {
candidates: reorder_candidates_by_ids(&shortlist, &ids, limit),
strategy: MemoryRecallStrategy::Reranked,
}),
Err(error) => {
tracing::warn!(
"Relevant memory rerank failed for model '{}': {}. Falling back to lexical shortlist.",
rerank_context.model,
error
);
shortlist.truncate(limit);
Ok(MemoryRecallSelection {
candidates: shortlist,
strategy: MemoryRecallStrategy::RerankFallback,
})
}
}
}
fn build_rerank_prompt(query: &str, candidates: &[MemoryRecallCandidate], limit: usize) -> String {
let mut prompt = String::from("# Bamboo Relevant Memory Recall Rerank\n\n");
prompt.push_str(
"Select the durable memory candidates that are most relevant to the user query.\n",
);
prompt.push_str("Return JSON only in the form {\"ids\":[\"candidate-id\", ...]}.\n");
prompt
.push_str("Do not include commentary, markdown fences, explanations, or unknown ids.\n\n");
prompt.push_str("## User query\n");
prompt.push_str(query.trim());
prompt.push_str("\n\n## Candidate memories\n");
for (index, candidate) in candidates.iter().enumerate() {
prompt.push_str(&format!(
"{}. id={}\n title: {}\n scope: {}\n status: {}\n updated_at: {}\n lexical_score: {:.2}\n summary: {}\n",
index + 1,
candidate.id,
candidate.title,
candidate.scope.as_str(),
candidate.status.as_str(),
candidate.updated_at,
candidate.score,
candidate.summary.replace('\n', " "),
));
}
prompt.push_str(&format!(
"\n## Selection rules\n- Return at most {limit} ids.\n- Use only ids from the candidate list above.\n- Prefer candidates that best answer the user query or encode active preferences/constraints relevant to it.\n- Prefer active memories over stale ones when relevance is otherwise similar.\n- Keep the ids ordered best-to-worst.\n"
));
prompt
}
async fn rerank_candidate_ids(
query: &str,
candidates: &[MemoryRecallCandidate],
limit: usize,
context: &MemoryRecallRerankContext,
) -> Result<Vec<String>, String> {
let model = context.model.trim();
if model.is_empty() {
return Err("rerank model is empty".to_string());
}
let messages = vec![
Message::system(
"You rerank Bamboo durable-memory recall candidates. Return strict JSON only in the form {\"ids\":[...]} using only candidate ids from the prompt.",
),
Message::user(build_rerank_prompt(query, candidates, limit)),
];
let options = LLMRequestOptions {
session_id: context.session_id.clone(),
reasoning_effort: Some(ReasoningEffort::High),
parallel_tool_calls: None,
required_tool: None,
responses: None,
request_purpose: Some("memory_rerank".to_string()),
cache: None,
};
let content = tokio::time::timeout(std::time::Duration::from_secs(30), async {
let mut stream = context
.llm
.chat_stream_with_options(&messages, &[], Some(8192), model, Some(&options))
.await
.map_err(|error| format!("rerank provider call failed: {error}"))?;
let mut content = String::new();
let mut terminal_done = false;
while let Some(chunk_result) = stream.next().await {
match chunk_result {
Ok(LLMChunk::Token(text)) => content.push_str(&text),
Ok(LLMChunk::Done) => {
terminal_done = true;
break;
}
Ok(_) => {}
Err(error) => return Err(format!("rerank stream failed: {error}")),
}
}
if !terminal_done {
return Err("rerank stream ended without terminal completion".to_string());
}
Ok(content)
})
.await
.unwrap_or_else(|_| Err("rerank timed out after 30s".to_string()))?;
parse_reranked_ids(&content, candidates)
.ok_or_else(|| format!("failed to parse rerank response: {}", content.trim()))
}
fn reorder_candidates_by_ids(
lexical_candidates: &[MemoryRecallCandidate],
preferred_ids: &[String],
limit: usize,
) -> Vec<MemoryRecallCandidate> {
if lexical_candidates.is_empty() || limit == 0 {
return Vec::new();
}
let allowed = lexical_candidates
.iter()
.map(|candidate| candidate.id.as_str())
.collect::<HashSet<_>>();
let mut seen = HashSet::new();
let mut ordered = Vec::new();
for id in preferred_ids {
let trimmed = id.trim();
if trimmed.is_empty() || !allowed.contains(trimmed) || !seen.insert(trimmed.to_string()) {
continue;
}
if let Some(candidate) = lexical_candidates
.iter()
.find(|candidate| candidate.id == trimmed)
.cloned()
{
ordered.push(candidate);
if ordered.len() >= limit {
return ordered;
}
}
}
for candidate in lexical_candidates {
if seen.insert(candidate.id.clone()) {
ordered.push(candidate.clone());
if ordered.len() >= limit {
break;
}
}
}
ordered
}
fn parse_reranked_ids(raw: &str, candidates: &[MemoryRecallCandidate]) -> Option<Vec<String>> {
let ids = serde_json::from_str::<MemoryRecallRerankEnvelope>(raw.trim())
.ok()?
.ids;
let explicit_empty_selection = ids.is_empty();
let allowed = candidates
.iter()
.map(|candidate| candidate.id.as_str())
.collect::<HashSet<_>>();
let mut seen = HashSet::new();
let mut out = Vec::new();
for id in ids {
let trimmed = id.trim();
if trimmed.is_empty() || !allowed.contains(trimmed) || !seen.insert(trimmed.to_string()) {
continue;
}
out.push(trimmed.to_string());
}
if out.is_empty() && !explicit_empty_selection {
return None;
}
Some(out)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
use async_trait::async_trait;
use bamboo_llm::{LLMError, LLMStream};
use bamboo_memory::memory_store::{
DurableMemoryStatus, DurableMemoryType, MemoryScope, TemporalGranularity,
};
use futures::stream;
#[derive(Clone)]
struct StaticResponseProvider {
response: String,
}
#[async_trait]
impl LLMProvider for StaticResponseProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token(self.response.clone())),
Ok(LLMChunk::Done),
])))
}
}
#[derive(Clone, Copy)]
enum FailingProvider {
Call,
PendingCall,
PendingStream,
PartialThenError,
EofWithoutDone,
}
#[async_trait]
impl LLMProvider for FailingProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
match self {
Self::Call => Err(LLMError::Api("rerank unavailable".to_string())),
Self::PendingCall => std::future::pending::<Result<LLMStream, LLMError>>().await,
Self::PendingStream => Ok(Box::pin(stream::pending())),
Self::PartialThenError => Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("{\"ids\":[]}".to_string())),
Err(LLMError::Stream("connection reset".to_string())),
]))),
Self::EofWithoutDone => Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Token(
"{\"ids\":[]}".to_string(),
))]))),
}
}
}
fn candidate(id: &str, score: f64) -> MemoryRecallCandidate {
MemoryRecallCandidate {
id: id.to_string(),
title: id.to_string(),
score,
scope: MemoryScope::Project,
project_key: Some("proj-1".to_string()),
status: DurableMemoryStatus::Active,
updated_at: "2026-04-09T00:00:00Z".to_string(),
summary: format!("summary for {id}"),
granularity: Some(TemporalGranularity::Month),
}
}
#[test]
fn parse_reranked_ids_filters_unknown_and_duplicate_ids() {
let candidates = vec![candidate("mem-a", 10.0), candidate("mem-b", 9.0)];
let parsed = parse_reranked_ids(
" \n{\"ids\":[\"mem-b\",\"unknown\",\"mem-a\",\"mem-b\"]}\t ",
&candidates,
)
.expect("reranked ids should parse");
assert_eq!(parsed, vec!["mem-b".to_string(), "mem-a".to_string()]);
}
#[test]
fn parse_reranked_ids_requires_an_explicit_well_typed_ids_field() {
let candidates = vec![candidate("mem-a", 10.0)];
assert!(parse_reranked_ids("{}", &candidates).is_none());
assert!(parse_reranked_ids("{\"other\":[]}", &candidates).is_none());
assert!(parse_reranked_ids("{\"ids\":\"mem-a\"}", &candidates).is_none());
assert!(
parse_reranked_ids("{\"ids\":[],\"error\":\"rate limited\"}", &candidates).is_none()
);
assert!(parse_reranked_ids("```json\n{\"ids\":[\"mem-a\"]}\n```", &candidates).is_none());
assert!(parse_reranked_ids("result: {\"ids\":[\"mem-a\"]}", &candidates).is_none());
assert!(parse_reranked_ids("{\"ids\":[\"mem-a\"]} done", &candidates).is_none());
}
#[test]
fn parse_reranked_ids_accepts_only_explicit_empty_object_selection() {
let candidates = vec![candidate("mem-a", 10.0)];
assert_eq!(
parse_reranked_ids("{\"ids\":[]}", &candidates),
Some(Vec::new())
);
assert!(parse_reranked_ids("[]", &candidates).is_none());
assert!(parse_reranked_ids("[\"mem-a\"]", &candidates).is_none());
assert!(parse_reranked_ids("{\"ids\":[\"unknown\",\" \"]}", &candidates).is_none());
assert!(parse_reranked_ids("[\"unknown\",\"\"]", &candidates).is_none());
}
#[test]
fn reorder_candidates_by_ids_appends_remaining_lexical_candidates() {
let lexical = vec![
candidate("mem-a", 10.0),
candidate("mem-b", 9.0),
candidate("mem-c", 8.0),
];
let reordered =
reorder_candidates_by_ids(&lexical, &["mem-c".to_string(), "mem-a".to_string()], 3);
assert_eq!(
reordered
.iter()
.map(|candidate| candidate.id.as_str())
.collect::<Vec<_>>(),
vec!["mem-c", "mem-a", "mem-b"]
);
}
async fn recall_store() -> (tempfile::TempDir, MemoryStore) {
let dir = tempfile::tempdir().expect("temp dir");
let store = MemoryStore::new(dir.path());
store
.write_memory(
MemoryScope::Project,
Some("proj-1"),
DurableMemoryType::Project,
"Release freeze checklist",
"Generic release freeze checklist for shipping work.",
&["release".to_string(), "freeze".to_string()],
Some("session-1"),
"main-model",
false,
None,
)
.await
.expect("write first memory");
store
.write_memory(
MemoryScope::Project,
Some("proj-1"),
DurableMemoryType::Project,
"Mobile launch blocker",
"This durable note captures the release freeze decision for the mobile app.",
&["mobile".to_string(), "launch".to_string()],
Some("session-1"),
"main-model",
false,
None,
)
.await
.expect("write second memory");
(dir, store)
}
fn rerank_context(response: &str) -> MemoryRecallRerankContext {
MemoryRecallRerankContext {
llm: Arc::new(StaticResponseProvider {
response: response.to_string(),
}),
model: "rerank-fast-model".to_string(),
session_id: Some("session-1".to_string()),
}
}
async fn assert_lexical_fallback(provider: Arc<dyn LLMProvider>) {
let (_dir, store) = recall_store().await;
let options = MemoryRecallOptions {
shortlist_limit: 2,
include_global_fallback: false,
max_candidates_per_scope: 12,
};
let expected = shortlist_relevant_memories(
&store,
Some("proj-1"),
"release freeze for mobile",
&options,
)
.await
.expect("deterministic shortlist");
let selection = select_relevant_memories(
&store,
Some("proj-1"),
"release freeze for mobile",
&options,
Some(&MemoryRecallRerankContext {
llm: provider,
model: "rerank-fast-model".to_string(),
session_id: Some("session-1".to_string()),
}),
)
.await
.expect("fallback selection");
assert_eq!(selection.strategy, MemoryRecallStrategy::RerankFallback);
assert_eq!(selection.candidates, expected);
}
#[tokio::test]
async fn invalid_or_empty_after_filter_response_falls_back_to_deterministic_shortlist() {
let (_dir, store) = recall_store().await;
let options = MemoryRecallOptions {
shortlist_limit: 2,
include_global_fallback: false,
max_candidates_per_scope: 12,
};
let expected = shortlist_relevant_memories(
&store,
Some("proj-1"),
"release freeze for mobile",
&options,
)
.await
.expect("deterministic shortlist");
for response in [
"not valid json",
"[]",
"[\"mem-a\"]",
"{}",
"{\"other\":[]}",
"{\"ids\":[],\"error\":\"rate limited\"}",
"{\"ids\":[\"unknown\",\" \"]}",
] {
let selection = select_relevant_memories(
&store,
Some("proj-1"),
"release freeze for mobile",
&options,
Some(&rerank_context(response)),
)
.await
.expect("fallback selection");
assert_eq!(selection.strategy, MemoryRecallStrategy::RerankFallback);
assert_eq!(selection.candidates, expected);
}
let known_id = expected
.first()
.expect("lexical shortlist should contain a known candidate")
.id
.clone();
for response in [
format!("```json\n{{\"ids\":[\"{known_id}\"]}}\n```"),
format!("result: {{\"ids\":[\"{known_id}\"]}}"),
format!("{{\"ids\":[\"{known_id}\"]}} done"),
] {
let selection = select_relevant_memories(
&store,
Some("proj-1"),
"release freeze for mobile",
&options,
Some(&rerank_context(&response)),
)
.await
.expect("wrapped known id should fall back to lexical selection");
assert_eq!(selection.strategy, MemoryRecallStrategy::RerankFallback);
assert_eq!(selection.candidates, expected);
}
}
#[tokio::test]
async fn valid_empty_model_selection_surfaces_no_memories() {
let (_dir, store) = recall_store().await;
let selection = select_relevant_memories(
&store,
Some("proj-1"),
"release freeze for mobile",
&MemoryRecallOptions {
shortlist_limit: 2,
include_global_fallback: false,
max_candidates_per_scope: 12,
},
Some(&rerank_context("{\"ids\":[]}")),
)
.await
.expect("reranked selection");
assert_eq!(selection.strategy, MemoryRecallStrategy::Reranked);
assert!(selection.candidates.is_empty());
}
#[tokio::test]
async fn provider_failure_falls_back_to_deterministic_shortlist() {
assert_lexical_fallback(Arc::new(FailingProvider::Call)).await;
}
#[tokio::test(start_paused = true)]
async fn rerank_timeout_falls_back_to_deterministic_shortlist() {
assert_lexical_fallback(Arc::new(FailingProvider::PendingStream)).await;
}
#[tokio::test(start_paused = true)]
async fn provider_connect_timeout_falls_back_to_deterministic_shortlist() {
assert_lexical_fallback(Arc::new(FailingProvider::PendingCall)).await;
}
#[tokio::test]
async fn partial_tokens_followed_by_stream_error_fall_back_to_deterministic_shortlist() {
assert_lexical_fallback(Arc::new(FailingProvider::PartialThenError)).await;
}
#[tokio::test]
async fn eof_without_done_falls_back_to_deterministic_shortlist() {
assert_lexical_fallback(Arc::new(FailingProvider::EofWithoutDone)).await;
}
#[derive(Default)]
struct PromptCaptureProvider {
candidate_ids: Mutex<Vec<String>>,
}
#[async_trait]
impl LLMProvider for PromptCaptureProvider {
async fn chat_stream(
&self,
messages: &[Message],
_tools: &[bamboo_agent_core::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
let prompt = messages
.iter()
.rev()
.find(|message| matches!(message.role, bamboo_agent_core::Role::User))
.map(|message| message.content.as_str())
.unwrap_or_default();
let ids = prompt
.lines()
.filter_map(|line| {
let (position, id) = line.split_once(". id=")?;
position.trim().parse::<usize>().ok()?;
Some(id.trim().to_string())
})
.collect::<Vec<_>>();
*self
.candidate_ids
.lock()
.expect("lock should not be poisoned") = ids.clone();
let response = serde_json::json!({ "ids": ids }).to_string();
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token(response)),
Ok(LLMChunk::Done),
])))
}
}
#[tokio::test]
async fn rerank_sees_candidate_pool_but_final_selection_respects_shortlist_limit() {
let dir = tempfile::tempdir().expect("temp dir");
let store = MemoryStore::new(dir.path());
for index in 0..12 {
store
.write_memory(
MemoryScope::Project,
Some("proj-1"),
DurableMemoryType::Project,
&format!("Release freeze component {index}"),
&format!(
"Release freeze evidence for independent component {index} with unique-marker-{index}."
),
&[format!("component-{index}")],
Some("session-1"),
"main-model",
false,
None,
)
.await
.expect("write matching memory");
}
let provider = Arc::new(PromptCaptureProvider::default());
let selection = select_relevant_memories(
&store,
Some("proj-1"),
"release freeze",
&MemoryRecallOptions {
shortlist_limit: 3,
include_global_fallback: false,
max_candidates_per_scope: 12,
},
Some(&MemoryRecallRerankContext {
llm: provider.clone(),
model: "rerank-fast-model".to_string(),
session_id: Some("session-1".to_string()),
}),
)
.await
.expect("reranked selection");
assert_eq!(selection.strategy, MemoryRecallStrategy::Reranked);
assert_eq!(
provider
.candidate_ids
.lock()
.expect("lock should not be poisoned")
.len(),
12,
"the model should see the configured rerank candidate pool"
);
assert_eq!(selection.candidates.len(), 3);
}
#[derive(Default)]
struct RequestOptionsCaptureProvider {
captured_max_tokens: Mutex<Vec<Option<u32>>>,
captured_reasoning: Mutex<Vec<Option<ReasoningEffort>>>,
}
#[async_trait]
impl LLMProvider for RequestOptionsCaptureProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("{\"ids\":[]}".to_string())),
Ok(LLMChunk::Done),
])))
}
async fn chat_stream_with_options(
&self,
messages: &[Message],
tools: &[bamboo_agent_core::ToolSchema],
max_output_tokens: Option<u32>,
model: &str,
options: Option<&LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
self.captured_max_tokens
.lock()
.expect("lock should not be poisoned")
.push(max_output_tokens);
self.captured_reasoning
.lock()
.expect("lock should not be poisoned")
.push(options.and_then(|options| options.reasoning_effort));
self.chat_stream(messages, tools, max_output_tokens, model)
.await
}
}
#[tokio::test]
async fn rerank_preserves_high_reasoning_token_budget() {
let provider = Arc::new(RequestOptionsCaptureProvider::default());
let context = MemoryRecallRerankContext {
llm: provider.clone(),
model: "deepseek-v4-pro".to_string(),
session_id: Some("test-session".to_string()),
};
let _ = rerank_candidate_ids("test query", &[candidate("mem-1", 0.9)], 5, &context).await;
assert_eq!(
provider
.captured_reasoning
.lock()
.expect("lock should not be poisoned")
.as_slice(),
[Some(ReasoningEffort::High)]
);
let max_tokens = provider.captured_max_tokens.lock().expect("lock")[0]
.expect("max_output_tokens should be set");
assert!(max_tokens > 4096);
}
}