use anyhow::Result;
use rusqlite::Connection;
use crate::context::{
load_session_start_candidates_with_limits, ContextLimits, LoadedBundleCandidates,
SessionStartRelevancePlan,
};
use crate::retrieval::embedding::local_only_embedding_profile_fingerprint;
use crate::retrieval_router::{
plan_context_bundle_with_limits, plan_session_start_with_limits, RetrievalPlan,
};
use super::domain::{ContextBundle, ContextItem, ContextRequest};
use super::executor::{
blocked_before_load, execute, execute_with_trace, BudgetEnforcement, ExecutorInputs,
};
pub(crate) struct SessionStartCompile {
pub bundle: ContextBundle,
pub relevance_plan: SessionStartRelevancePlan,
}
pub fn compile_session_start_bundle(
conn: &Connection,
request: &ContextRequest,
cwd: &str,
current_branch: Option<&str>,
enrichment_available: bool,
) -> Result<ContextBundle> {
let limits = ContextLimits::from_env();
let local_embedding_fingerprint = local_only_embedding_profile_fingerprint();
let compiled = plan_context_bundle_with_limits(request, &limits, &local_embedding_fingerprint)?;
Ok(bundle_for_plan(
conn,
&compiled,
&request.project.key,
cwd,
current_branch,
&limits,
enrichment_available,
))
}
#[allow(clippy::too_many_arguments)]
fn bundle_for_plan(
conn: &Connection,
compiled: &RetrievalPlan,
project: &str,
cwd: &str,
current_branch: Option<&str>,
limits: &ContextLimits,
enrichment_available: bool,
) -> ContextBundle {
match load_session_start_candidates_with_limits(conn, project, cwd, current_branch, limits) {
Ok(LoadedBundleCandidates {
candidates,
poisoning_drops,
preselection_drops,
}) => execute(
compiled,
&ExecutorInputs {
candidates,
poisoning_drops,
preselection_drops,
enrichment_available,
},
),
Err(error) => blocked_before_load(compiled, &error.to_string()),
}
}
pub(crate) fn compile_session_start_for_renderer(
request: &ContextRequest,
limits: &ContextLimits,
candidates: Vec<ContextItem>,
poisoning_drops: Vec<ContextItem>,
preselection_drops: Vec<super::executor::PreselectionDrop>,
enrichment_available: bool,
) -> Result<SessionStartCompile> {
let compiled = plan_session_start_with_limits(request, limits)?;
let trace = execute_with_trace(
&compiled,
&ExecutorInputs {
candidates,
poisoning_drops,
preselection_drops,
enrichment_available,
},
BudgetEnforcement::DeferToRenderer,
);
Ok(SessionStartCompile {
bundle: trace.bundle,
relevance_plan: trace.relevance_plan,
})
}
pub(crate) fn seal_session_start_bundle(
bundle: &mut ContextBundle,
selected_keys: &std::collections::HashSet<String>,
total_truncated_keys: &std::collections::HashSet<String>,
output_chars: usize,
) {
retain_selected_sections(bundle, selected_keys);
for entry in &mut bundle.audit.entries {
if entry.selected && !selected_keys.contains(&entry.stable_key) {
entry.selected = false;
entry.reason = if total_truncated_keys.contains(&entry.stable_key) {
"total_char_limit"
} else {
"section_budget"
}
.to_string();
}
}
bundle.audit.selected_count = bundle
.audit
.entries
.iter()
.filter(|entry| entry.selected)
.count() as u32;
bundle.audit.dropped_count = bundle.audit.candidates_considered - bundle.audit.selected_count;
bundle.audit.token_estimate = (output_chars as u32).div_ceil(4);
if !total_truncated_keys.is_empty() {
bundle.audit.truncation_reason = Some("total_char_limit".to_string());
}
}
pub(crate) fn reseal_after_emission_gate(
bundle: &mut ContextBundle,
selected_keys: &std::collections::HashSet<String>,
output_chars: usize,
drop_reason: &str,
output_truncated: bool,
) {
retain_selected_sections(bundle, selected_keys);
let mut dropped_by_gate = false;
for entry in &mut bundle.audit.entries {
if entry.selected && !selected_keys.contains(&entry.stable_key) {
entry.selected = false;
entry.reason = drop_reason.to_string();
dropped_by_gate = true;
}
}
bundle.audit.selected_count = bundle
.audit
.entries
.iter()
.filter(|entry| entry.selected)
.count() as u32;
bundle.audit.dropped_count = bundle.audit.candidates_considered - bundle.audit.selected_count;
bundle.audit.token_estimate = (output_chars as u32).div_ceil(4);
if dropped_by_gate || output_truncated {
bundle.audit.truncation_reason = Some(drop_reason.to_string());
}
}
fn retain_selected_sections(
bundle: &mut ContextBundle,
selected_keys: &std::collections::HashSet<String>,
) {
for section in [
&mut bundle.preferences,
&mut bundle.failure_lessons,
&mut bundle.current_truth,
&mut bundle.workstreams,
&mut bundle.memory_index,
&mut bundle.recent_sessions,
] {
section.retain(|item| selected_keys.contains(&item.stable_key));
}
}