bamboo_memory/memory_store/

recall.rs

use std::cmp::Ordering;
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 super::{parse_rfc3339, DurableMemoryStatus, MemoryScope, MemoryStore, TemporalGranularity};

#[derive(Debug, Clone, PartialEq)]
pub struct MemoryRecallCandidate {
    pub id: String,
    pub title: String,
    pub score: f64,
    pub scope: MemoryScope,
    pub project_key: Option<String>,
    pub status: DurableMemoryStatus,
    pub updated_at: String,
    pub summary: String,
    /// Optional temporal granularity, used as a stable tie-breaker in
    /// [`sort_recall_candidates`]: among equally-relevant candidates, coarser
    /// (more cache-stable) memories sort first. `None` is treated as most stable.
    pub granularity: Option<TemporalGranularity>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoryRecallOptions {
    pub shortlist_limit: usize,
    pub include_global_fallback: bool,
    pub max_candidates_per_scope: usize,
}

impl Default for MemoryRecallOptions {
    fn default() -> Self {
        Self {
            shortlist_limit: 3,
            include_global_fallback: true,
            max_candidates_per_scope: 20,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryRecallStrategy {
    Lexical,
    Reranked,
    RerankFallback,
}

impl MemoryRecallStrategy {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Lexical => "lexical",
            Self::Reranked => "reranked",
            Self::RerankFallback => "rerank_fallback",
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct MemoryRecallSelection {
    pub candidates: Vec<MemoryRecallCandidate>,
    pub strategy: MemoryRecallStrategy,
}

#[derive(Clone)]
pub struct MemoryRecallRerankContext {
    pub llm: Arc<dyn LLMProvider>,
    pub model: String,
    pub session_id: Option<String>,
}

#[derive(Debug, Deserialize)]
struct MemoryRecallRerankEnvelope {
    #[serde(default)]
    ids: Vec<String>,
}

pub async fn shortlist_relevant_memories(
    store: &MemoryStore,
    project_key: Option<&str>,
    query: &str,
    options: &MemoryRecallOptions,
) -> io::Result<Vec<MemoryRecallCandidate>> {
    let limit = options.shortlist_limit.max(1);
    let mut candidates =
        lexical_shortlist_relevant_memories(store, project_key, query, options).await?;
    candidates.truncate(limit);
    Ok(candidates)
}

pub 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 mut shortlist =
        lexical_shortlist_relevant_memories(store, project_key, query, 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 {
        // A VALID but empty selection means the reranker judged NONE of the lexical
        // shortlist relevant to the query. Respect that — return no memories rather
        // than resurrecting the noisy shortlist (which otherwise surfaces unrelated
        // memories the reranker explicitly rejected). This is NOT a failure, so it
        // doesn't warn or fall back — only a genuine rerank error (Err) does.
        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,
            })
        }
    }
}

async fn lexical_shortlist_relevant_memories(
    store: &MemoryStore,
    project_key: Option<&str>,
    query: &str,
    options: &MemoryRecallOptions,
) -> io::Result<Vec<MemoryRecallCandidate>> {
    let query = query.trim();
    if query.is_empty() {
        return Ok(Vec::new());
    }

    let limit = options.shortlist_limit.max(1);
    let per_scope_limit = options.max_candidates_per_scope.max(limit);

    if let Some(project_key) = project_key.map(str::trim).filter(|value| !value.is_empty()) {
        let mut project_hits =
            shortlist_scope(store, MemoryScope::Project, Some(project_key), query).await?;
        project_hits.truncate(per_scope_limit);
        if !project_hits.is_empty() {
            return Ok(project_hits);
        }
    }

    if options.include_global_fallback {
        let mut global_hits = shortlist_scope(store, MemoryScope::Global, None, query).await?;
        global_hits.truncate(per_scope_limit);
        return Ok(global_hits);
    }

    Ok(Vec::new())
}

async fn shortlist_scope(
    store: &MemoryStore,
    scope: MemoryScope,
    project_key: Option<&str>,
    query: &str,
) -> io::Result<Vec<MemoryRecallCandidate>> {
    let Some(index) = store.read_lexical_index(scope, project_key).await? else {
        return Ok(Vec::new());
    };

    let query_tokens = super::lexical_bm25::tokenize(query);
    if query_tokens.is_empty() {
        return Ok(Vec::new());
    }

    // BM25(F) over the loaded lexical index: corpus stats (df, avgdl) are computed
    // once, then each doc is scored in O(query terms). CJK-aware tokenization makes
    // the bilingual library searchable — the previous tokenizer dropped all Chinese.
    let corpus = super::lexical_bm25::Bm25Corpus::build(&index.items);
    let mut candidates = index
        .items
        .iter()
        .enumerate()
        .filter_map(|(i, item)| corpus.score(i, &query_tokens).map(|score| (item, score)))
        .map(|(item, score)| MemoryRecallCandidate {
            id: item.id.clone(),
            title: item.title.clone(),
            score,
            scope: item.scope,
            project_key: item.project_key.clone(),
            status: item.status,
            updated_at: item.updated_at.clone(),
            summary: item.summary.clone(),
            granularity: item.granularity,
        })
        .collect::<Vec<_>>();

    sort_recall_candidates(&mut candidates);
    Ok(candidates)
}

fn sort_recall_candidates(candidates: &mut [MemoryRecallCandidate]) {
    candidates.sort_by(|left, right| {
        right
            .score
            .partial_cmp(&left.score)
            .unwrap_or(Ordering::Equal)
            // Among equally-relevant candidates, prefer coarser (more prefix-cache
            // friendly) granularity first so the recalled block is stable across
            // calls and does not churn the LLM prompt prefix (issue #61). Lower
            // cache_stability_rank = coarser = earlier.
            .then_with(|| {
                TemporalGranularity::cache_stability_rank(left.granularity).cmp(
                    &TemporalGranularity::cache_stability_rank(right.granularity),
                )
            })
            .then_with(|| {
                let left_dt = parse_rfc3339(&left.updated_at)
                    .unwrap_or(chrono::DateTime::<chrono::Utc>::MIN_UTC);
                let right_dt = parse_rfc3339(&right.updated_at)
                    .unwrap_or(chrono::DateTime::<chrono::Utc>::MIN_UTC);
                right_dt.cmp(&left_dt)
            })
            .then_with(|| left.title.cmp(&right.title))
    });
}

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,
        responses: None,
        request_purpose: Some("memory_rerank".to_string()),
        cache: None,
    };

    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 content = tokio::time::timeout(std::time::Duration::from_secs(30), async {
        let mut content = String::new();
        while let Some(chunk_result) = stream.next().await {
            match chunk_result {
                Ok(LLMChunk::Token(text)) => content.push_str(&text),
                Ok(LLMChunk::Done) => break,
                Ok(_) => {}
                Err(error) => {
                    if !content.trim().is_empty() {
                        break;
                    }
                    return Err(format!("rerank stream failed: {error}"));
                }
            }
        }
        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 stripped = strip_markdown_fence(raw);
    let fragment = extract_json_fragment(&stripped).unwrap_or(stripped.trim());
    let ids = serde_json::from_str::<MemoryRecallRerankEnvelope>(fragment)
        .map(|value| value.ids)
        .or_else(|_| serde_json::from_str::<Vec<String>>(fragment))
        .ok()?;

    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());
    }

    // The response PARSED — return the (possibly empty) selection. `None` is
    // reserved for genuinely unparseable output (handled above via `.ok()?`), so
    // an empty `{"ids":[]}` — the reranker judging NONE relevant — is a valid
    // result, not a parse failure. (Empty is acted on by the caller.)
    Some(out)
}

fn strip_markdown_fence(raw: &str) -> String {
    let trimmed = raw.trim();
    for fence in ["````", "```"] {
        if let Some(after_fence) = trimmed.strip_prefix(fence) {
            let Some(first_newline) = after_fence.find('\n') else {
                continue;
            };
            let body = &after_fence[first_newline + 1..];
            if let Some(end_idx) = body.rfind(fence) {
                return body[..end_idx].trim().to_string();
            }
        }
    }
    trimmed.to_string()
}

fn extract_json_fragment(raw: &str) -> Option<&str> {
    let trimmed = raw.trim();
    if trimmed.is_empty() {
        return None;
    }

    if let (Some(start), Some(end)) = (trimmed.find('{'), trimmed.rfind('}')) {
        if start <= end {
            return Some(trimmed[start..=end].trim());
        }
    }

    if let (Some(start), Some(end)) = (trimmed.find('['), trimmed.rfind(']')) {
        if start <= end {
            return Some(trimmed[start..=end].trim());
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::memory_store::DurableMemoryType;
    use async_trait::async_trait;
    use bamboo_domain::ReasoningEffort;
    use bamboo_llm::provider::LLMRequestOptions;
    use bamboo_llm::{LLMChunk, LLMError, LLMProvider, LLMStream};
    use futures::stream;
    use std::sync::Mutex;
    use tempfile::tempdir;

    #[derive(Clone)]
    struct StaticResponseProvider {
        response: String,
        requested_models: Arc<Mutex<Vec<String>>>,
    }

    impl StaticResponseProvider {
        fn new(response: impl Into<String>) -> Self {
            Self {
                response: response.into(),
                requested_models: Arc::new(Mutex::new(Vec::new())),
            }
        }
    }

    #[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> {
            self.requested_models
                .lock()
                .expect("lock poisoned")
                .push(model.to_string());
            Ok(Box::pin(stream::iter(vec![
                Ok(LLMChunk::Token(self.response.clone())),
                Ok(LLMChunk::Done),
            ])))
        }
    }

    fn candidate(
        id: &str,
        score: f64,
        granularity: Option<TemporalGranularity>,
    ) -> MemoryRecallCandidate {
        MemoryRecallCandidate {
            id: id.to_string(),
            title: id.to_string(),
            score,
            scope: MemoryScope::Project,
            project_key: Some("proj-1".to_string()),
            status: DurableMemoryStatus::Active,
            // Same timestamp for all so granularity is the deciding tie-breaker.
            updated_at: "2026-04-09T00:00:00Z".to_string(),
            summary: "summary".to_string(),
            granularity,
        }
    }

    #[test]
    fn equal_score_candidates_sort_coarse_granularity_first_for_cache_stability() {
        // All same score + same timestamp → granularity decides ordering. Coarser
        // (year) is more prefix-cache friendly and must sort ahead of finer (day).
        let mut candidates = vec![
            candidate("day", 5.0, Some(TemporalGranularity::Day)),
            candidate("year", 5.0, Some(TemporalGranularity::Year)),
            candidate("none", 5.0, None),
            candidate("month", 5.0, Some(TemporalGranularity::Month)),
        ];
        sort_recall_candidates(&mut candidates);
        let order: Vec<&str> = candidates.iter().map(|c| c.id.as_str()).collect();
        // None is treated as most stable (rank 0), then year, month, day.
        assert_eq!(order, vec!["none", "year", "month", "day"]);
    }

    #[test]
    fn higher_score_still_wins_over_cache_stable_granularity() {
        // Granularity is only a tie-breaker: a more relevant (higher score) day-level
        // memory must still outrank a less relevant year-level one.
        let mut candidates = vec![
            candidate("year-low", 1.0, Some(TemporalGranularity::Year)),
            candidate("day-high", 9.0, Some(TemporalGranularity::Day)),
        ];
        sort_recall_candidates(&mut candidates);
        assert_eq!(candidates[0].id, "day-high");
    }

    #[test]
    fn parse_reranked_ids_accepts_fenced_json_and_filters_unknown_ids() {
        let candidates = vec![
            MemoryRecallCandidate {
                id: "mem-a".to_string(),
                title: "A".to_string(),
                score: 10.0,
                scope: MemoryScope::Project,
                project_key: Some("proj-1".to_string()),
                status: DurableMemoryStatus::Active,
                updated_at: "2026-04-09T00:00:00Z".to_string(),
                summary: "summary a".to_string(),
                granularity: None,
            },
            MemoryRecallCandidate {
                id: "mem-b".to_string(),
                title: "B".to_string(),
                score: 9.0,
                scope: MemoryScope::Project,
                project_key: Some("proj-1".to_string()),
                status: DurableMemoryStatus::Active,
                updated_at: "2026-04-09T00:00:00Z".to_string(),
                summary: "summary b".to_string(),
                granularity: None,
            },
        ];

        let parsed = parse_reranked_ids(
            "```json\n{\"ids\":[\"mem-b\",\"unknown\",\"mem-a\",\"mem-b\"]}\n```",
            &candidates,
        )
        .expect("reranked ids should parse");

        assert_eq!(parsed, vec!["mem-b".to_string(), "mem-a".to_string()]);
    }

    #[test]
    fn reorder_candidates_by_ids_appends_remaining_lexical_candidates() {
        let lexical = vec![
            MemoryRecallCandidate {
                id: "mem-a".to_string(),
                title: "A".to_string(),
                score: 10.0,
                scope: MemoryScope::Project,
                project_key: Some("proj-1".to_string()),
                status: DurableMemoryStatus::Active,
                updated_at: "2026-04-09T00:00:00Z".to_string(),
                summary: "summary a".to_string(),
                granularity: None,
            },
            MemoryRecallCandidate {
                id: "mem-b".to_string(),
                title: "B".to_string(),
                score: 9.0,
                scope: MemoryScope::Project,
                project_key: Some("proj-1".to_string()),
                status: DurableMemoryStatus::Active,
                updated_at: "2026-04-09T00:00:00Z".to_string(),
                summary: "summary b".to_string(),
                granularity: None,
            },
            MemoryRecallCandidate {
                id: "mem-c".to_string(),
                title: "C".to_string(),
                score: 8.0,
                scope: MemoryScope::Project,
                project_key: Some("proj-1".to_string()),
                status: DurableMemoryStatus::Active,
                updated_at: "2026-04-09T00:00:00Z".to_string(),
                summary: "summary c".to_string(),
                granularity: None,
            },
        ];

        let reordered =
            reorder_candidates_by_ids(&lexical, &["mem-c".to_string(), "mem-a".to_string()], 3);

        assert_eq!(reordered[0].id, "mem-c");
        assert_eq!(reordered[1].id, "mem-a");
        assert_eq!(reordered[2].id, "mem-b");
    }

    #[tokio::test]
    async fn project_scope_shortlist_excludes_global_when_project_hits_exist() {
        let dir = tempdir().unwrap();
        let store = MemoryStore::new(dir.path());

        store
            .write_memory(
                MemoryScope::Project,
                Some("proj-1"),
                DurableMemoryType::Project,
                "Release freeze decision",
                "Project-specific release freeze note.",
                &["release".to_string()],
                Some("session-1"),
                "main-model",
                false,
                None,
            )
            .await
            .unwrap();
        store
            .write_memory(
                MemoryScope::Global,
                None,
                DurableMemoryType::Reference,
                "Global release guidance",
                "Global note that should not be used when project hits exist.",
                &["release".to_string()],
                Some("session-1"),
                "main-model",
                false,
                None,
            )
            .await
            .unwrap();

        let candidates = shortlist_relevant_memories(
            &store,
            Some("proj-1"),
            "release freeze",
            &MemoryRecallOptions::default(),
        )
        .await
        .unwrap();

        assert!(!candidates.is_empty());
        assert!(candidates
            .iter()
            .all(|candidate| candidate.scope == MemoryScope::Project));
    }

    #[tokio::test]
    async fn global_fallback_triggers_only_when_project_hits_are_absent() {
        let dir = tempdir().unwrap();
        let store = MemoryStore::new(dir.path());

        store
            .write_memory(
                MemoryScope::Global,
                None,
                DurableMemoryType::Reference,
                "Global release guidance",
                "Fallback note for release work.",
                &["release".to_string()],
                Some("session-1"),
                "main-model",
                false,
                None,
            )
            .await
            .unwrap();

        let candidates = shortlist_relevant_memories(
            &store,
            Some("proj-missing"),
            "release guidance",
            &MemoryRecallOptions::default(),
        )
        .await
        .unwrap();

        assert!(!candidates.is_empty());
        assert!(candidates
            .iter()
            .all(|candidate| candidate.scope == MemoryScope::Global));
    }

    #[tokio::test]
    async fn model_rerank_reorders_lexical_shortlist_when_enabled() {
        let dir = tempdir().unwrap();
        let store = MemoryStore::new(dir.path());

        let lexical_first = 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
            .unwrap();
        let reranked_first = 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 and should be preferred for mobile freeze requests.",
                &["mobile".to_string(), "launch".to_string()],
                Some("session-1"),
                "main-model",
                false,
                None,
            )
            .await
            .unwrap();

        let provider = StaticResponseProvider::new(format!(
            "{{\"ids\":[\"{}\",\"{}\"]}}",
            reranked_first.frontmatter.id, lexical_first.frontmatter.id
        ));
        let requested_models = provider.requested_models.clone();
        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(&MemoryRecallRerankContext {
                llm: Arc::new(provider),
                model: "rerank-fast-model".to_string(),
                session_id: Some("session-1".to_string()),
            }),
        )
        .await
        .unwrap();

        assert_eq!(selection.strategy, MemoryRecallStrategy::Reranked);
        assert_eq!(selection.candidates.len(), 2);
        assert_eq!(selection.candidates[0].id, reranked_first.frontmatter.id);
        assert_eq!(selection.candidates[1].id, lexical_first.frontmatter.id);
        assert_eq!(
            requested_models.lock().expect("lock poisoned").as_slice(),
            ["rerank-fast-model"]
        );
    }

    #[tokio::test]
    async fn invalid_model_rerank_response_falls_back_to_lexical_order() {
        let dir = tempdir().unwrap();
        let store = MemoryStore::new(dir.path());

        let lexical_first = 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
            .unwrap();
        let lexical_second = 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
            .unwrap();

        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(&MemoryRecallRerankContext {
                llm: Arc::new(StaticResponseProvider::new("not valid json")),
                model: "rerank-fast-model".to_string(),
                session_id: Some("session-1".to_string()),
            }),
        )
        .await
        .unwrap();

        assert_eq!(selection.strategy, MemoryRecallStrategy::RerankFallback);
        assert_eq!(selection.candidates.len(), 2);
        // The fallback contract is "return the lexical shortlist in its own order",
        // NOT any hard-coded doc order — verify it matches the pure lexical shortlist
        // (which BM25 orders by relevance) and preserves both durable memories.
        let lexical = shortlist_relevant_memories(
            &store,
            Some("proj-1"),
            "release freeze for mobile",
            &MemoryRecallOptions {
                shortlist_limit: 2,
                include_global_fallback: false,
                max_candidates_per_scope: 12,
            },
        )
        .await
        .unwrap();
        let fallback_ids: Vec<&str> = selection.candidates.iter().map(|c| c.id.as_str()).collect();
        let lexical_ids: Vec<&str> = lexical.iter().map(|c| c.id.as_str()).collect();
        assert_eq!(
            fallback_ids, lexical_ids,
            "fallback must preserve lexical order"
        );
        assert!(fallback_ids.contains(&lexical_first.frontmatter.id.as_str()));
        assert!(fallback_ids.contains(&lexical_second.frontmatter.id.as_str()));
    }

    #[tokio::test]
    async fn empty_rerank_selection_returns_no_memories_not_lexical() {
        let dir = tempdir().unwrap();
        let store = MemoryStore::new(dir.path());

        // Two lexically-matched candidates the reranker will reject.
        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
            .unwrap();
        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
            .unwrap();

        // The reranker judges NONE of the shortlist relevant → `{"ids":[]}`.
        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(&MemoryRecallRerankContext {
                llm: Arc::new(StaticResponseProvider::new("{\"ids\":[]}")),
                model: "rerank-fast-model".to_string(),
                session_id: Some("session-1".to_string()),
            }),
        )
        .await
        .unwrap();

        // Respect the reranker: no memories surfaced, and it is NOT treated as a
        // failure/fallback (the pre-fix bug misreported `{"ids":[]}` as a parse
        // error → lexical fallback → surfaced the rejected memories).
        assert_eq!(selection.strategy, MemoryRecallStrategy::Reranked);
        assert!(
            selection.candidates.is_empty(),
            "an empty rerank selection must surface no memories, got {}",
            selection.candidates.len()
        );
    }

    /// Provider that captures `max_output_tokens` and `reasoning_effort` from `chat_stream_with_options`.
    #[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(|o| o.reasoning_effort));
            self.chat_stream(messages, tools, max_output_tokens, model)
                .await
        }
    }

    #[tokio::test]
    async fn rerank_sufficient_max_tokens_for_high_reasoning() {
        let provider = Arc::new(RequestOptionsCaptureProvider::default());
        let candidates = vec![MemoryRecallCandidate {
            id: "mem-1".to_string(),
            score: 0.9,
            title: "Test memory".to_string(),
            scope: MemoryScope::Project,
            project_key: Some("proj-1".to_string()),
            status: DurableMemoryStatus::Active,
            updated_at: "2026-05-08T00:00:00Z".to_string(),
            summary: "A test durable memory entry".to_string(),
            granularity: None,
        }];
        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", &candidates, 5, &context).await;

        let captured_reasoning = provider
            .captured_reasoning
            .lock()
            .expect("lock should not be poisoned");
        let captured_max_tokens = provider
            .captured_max_tokens
            .lock()
            .expect("lock should not be poisoned");
        assert_eq!(
            captured_reasoning.as_slice(),
            [Some(ReasoningEffort::High)],
            "rerank should request High reasoning"
        );
        let max_tokens = captured_max_tokens[0].expect("max_output_tokens should be set");
        assert!(
            max_tokens > 4096,
            "max_output_tokens ({}) must exceed thinking budget (4096) to avoid truncation",
            max_tokens
        );
    }
}