boundgraph-core 0.1.0

//! The Context (`Γ`) — the exact admitted token window, and what was refused.
//!
//! Law II says a step receives only the context it declared as necessary.
//! That law is only *auditable* if you can also see what the admission
//! controller turned away. Recording admissions alone makes the law
//! unfalsifiable: "the agent saw these five things" tells you nothing about
//! the sixth thing it needed and never got.
//!
//! So [`Context`] carries both halves. When someone asks "why did the agent
//! not know X?", the answer is a lookup: X is in [`Context::evicted`] with a
//! [`EvictionReason`] attached.

use serde::{Deserialize, Serialize};

use crate::digest::Digest;

/// The token ceiling an admission ran against — `s_max` in the cost model.
///
/// This is the term that makes total cost linear rather than quadratic: a step
/// admits at most `s_max` tokens regardless of how much history exists.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct TokenBudget(pub u32);

impl TokenBudget {
    /// The ceiling as a plain count.
    pub const fn tokens(self) -> u32 {
        self.0
    }
}

/// Where an item of context came from.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContextSource {
    /// The flow's system or task prompt.
    Prompt,
    /// A condensate from an earlier step.
    Condensate,
    /// Retrieved material, e.g. from a vector store.
    Retrieval,
    /// Verified evidence from an earlier step.
    Evidence,
    /// Anything else, named by the admitter.
    Other(String),
}

/// One item that made it into the window.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdmittedItem {
    /// Where it came from.
    pub source: ContextSource,
    /// Content address of the item, so it can be retrieved and verified.
    pub digest: Digest,
    /// What it cost against the budget.
    pub tokens: u32,
    /// Why the admitter let it in.
    pub reason: String,
}

/// Why an item was refused.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EvictionReason {
    /// Admitting it would have exceeded `s_max`.
    BudgetExhausted {
        /// Tokens the item needed.
        needed: u32,
        /// Tokens still available when it was considered.
        remaining: u32,
    },
    /// The step did not declare this source as necessary.
    NotDeclared,
    /// Superseded by a fresher item covering the same ground.
    Superseded {
        /// The item that replaced it.
        by: Digest,
    },
    /// Withheld by policy, e.g. a secret or out-of-scope material.
    PolicyDenied {
        /// The policy that refused it.
        policy: String,
    },
    /// Anything else, named by the admitter.
    Other(String),
}

/// One item that was considered and refused.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct EvictionRecord {
    /// Where it would have come from.
    pub source: ContextSource,
    /// Content address, so the auditor can fetch what was withheld.
    pub digest: Digest,
    /// What it would have cost.
    pub tokens: u32,
    /// Why it did not make it in.
    pub reason: EvictionReason,
}

/// The context was over budget — an admitter bug, not a user error.
#[derive(Debug, thiserror::Error)]
#[error(
    "context exceeds its token budget: {admitted} tokens admitted against a ceiling of {budget}"
)]
pub struct ContextOverBudget {
    /// Total tokens the admitted items came to.
    pub admitted: u32,
    /// The ceiling that was supposed to hold.
    pub budget: u32,
}

/// `Γ` — the exact admitted token window. Immutable, fully attributed.
///
/// # Example
///
/// ```
/// use boundgraph_core::{AdmittedItem, Context, ContextSource, Digest, EvictionReason,
///                 EvictionRecord, TokenBudget};
///
/// let ctx = Context::new(
///     vec![AdmittedItem {
///         source: ContextSource::Prompt,
///         digest: Digest::of(b"task"),
///         tokens: 120,
///         reason: "declared by the step".into(),
///     }],
///     vec![EvictionRecord {
///         source: ContextSource::Retrieval,
///         digest: Digest::of(b"a long document"),
///         tokens: 9_000,
///         reason: EvictionReason::BudgetExhausted { needed: 9_000, remaining: 880 },
///     }],
///     TokenBudget(1_000),
/// )?;
///
/// assert_eq!(ctx.admitted_tokens(), 120);
/// // The law is auditable: we can say exactly what was withheld, and why.
/// assert_eq!(ctx.evicted.len(), 1);
/// # Ok::<(), boundgraph_core::context::ContextOverBudget>(())
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Context {
    /// What the reasoner actually saw.
    pub items: Vec<AdmittedItem>,
    /// What it did not, and why.
    pub evicted: Vec<EvictionRecord>,
    /// The ceiling this admission ran against.
    pub budget: TokenBudget,
}

impl Context {
    /// Build a context, enforcing the budget invariant.
    ///
    /// # Errors
    ///
    /// Returns [`ContextOverBudget`] if the admitted items exceed `budget`.
    /// This is deliberately a hard error rather than a silent truncation: an
    /// admitter that overruns has broken Law II, and quietly dropping the tail
    /// would hide that while also losing the eviction record that explains it.
    pub fn new(
        items: Vec<AdmittedItem>,
        evicted: Vec<EvictionRecord>,
        budget: TokenBudget,
    ) -> Result<Self, ContextOverBudget> {
        let admitted: u32 = items.iter().map(|i| i.tokens).sum();
        if admitted > budget.tokens() {
            return Err(ContextOverBudget {
                admitted,
                budget: budget.tokens(),
            });
        }
        Ok(Self {
            items,
            evicted,
            budget,
        })
    }

    /// An empty window under a given ceiling.
    pub fn empty(budget: TokenBudget) -> Self {
        Self {
            items: Vec::new(),
            evicted: Vec::new(),
            budget,
        }
    }

    /// Tokens actually admitted. Never exceeds [`Context::budget`].
    pub fn admitted_tokens(&self) -> u32 {
        self.items.iter().map(|i| i.tokens).sum()
    }

    /// Tokens left unused under the ceiling.
    pub fn headroom(&self) -> u32 {
        self.budget.tokens().saturating_sub(self.admitted_tokens())
    }

    /// Tokens that were considered and refused.
    ///
    /// A large number here against small [`headroom`](Self::headroom) means the
    /// budget was the binding constraint; a large number against *ample*
    /// headroom means something was withheld for another reason, which is
    /// usually the more interesting case when debugging.
    pub fn evicted_tokens(&self) -> u32 {
        self.evicted.iter().map(|e| e.tokens).sum()
    }

    /// Whether this item was considered and refused, and if so why.
    ///
    /// This is the "why did the agent not know X?" query.
    pub fn why_evicted(&self, digest: &Digest) -> Option<&EvictionReason> {
        self.evicted
            .iter()
            .find(|e| &e.digest == digest)
            .map(|e| &e.reason)
    }

    /// Whether this item made it into the window.
    pub fn admitted(&self, digest: &Digest) -> bool {
        self.items.iter().any(|i| &i.digest == digest)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::canonical::Canonical;

    fn item(tokens: u32, tag: &[u8]) -> AdmittedItem {
        AdmittedItem {
            source: ContextSource::Prompt,
            digest: Digest::of(tag),
            tokens,
            reason: "declared".into(),
        }
    }

    fn eviction(tokens: u32, tag: &[u8], reason: EvictionReason) -> EvictionRecord {
        EvictionRecord {
            source: ContextSource::Retrieval,
            digest: Digest::of(tag),
            tokens,
            reason,
        }
    }

    #[test]
    fn carries_items_evicted_and_budget() {
        let ctx = Context::new(
            vec![item(10, b"a")],
            vec![eviction(50, b"b", EvictionReason::NotDeclared)],
            TokenBudget(100),
        )
        .unwrap();
        assert_eq!(ctx.items.len(), 1);
        assert_eq!(ctx.evicted.len(), 1);
        assert_eq!(ctx.budget, TokenBudget(100));
    }

    #[test]
    fn admitted_items_record_source_digest_tokens_and_reason() {
        let ctx = Context::new(vec![item(10, b"a")], vec![], TokenBudget(100)).unwrap();
        let i = &ctx.items[0];
        assert_eq!(i.source, ContextSource::Prompt);
        assert_eq!(i.digest, Digest::of(b"a"));
        assert_eq!(i.tokens, 10);
        assert_eq!(i.reason, "declared");
    }

    #[test]
    fn exceeding_the_budget_is_a_typed_error_not_a_truncation() {
        let err = Context::new(vec![item(101, b"a")], vec![], TokenBudget(100)).unwrap_err();
        assert_eq!(err.admitted, 101);
        assert_eq!(err.budget, 100);
    }

    #[test]
    fn exactly_at_the_budget_is_allowed() {
        assert!(Context::new(vec![item(100, b"a")], vec![], TokenBudget(100)).is_ok());
    }

    #[test]
    fn answers_why_the_agent_did_not_know_something() {
        let missing = Digest::of(b"the missing document");
        let ctx = Context::new(
            vec![item(10, b"a")],
            vec![eviction(
                9_000,
                b"the missing document",
                EvictionReason::BudgetExhausted {
                    needed: 9_000,
                    remaining: 90,
                },
            )],
            TokenBudget(100),
        )
        .unwrap();

        assert!(!ctx.admitted(&missing));
        assert!(matches!(
            ctx.why_evicted(&missing),
            Some(EvictionReason::BudgetExhausted {
                needed: 9_000,
                remaining: 90
            })
        ));
        // An item that was never considered is distinguishable from one refused.
        assert_eq!(ctx.why_evicted(&Digest::of(b"never seen")), None);
    }

    #[test]
    fn reports_headroom_and_evicted_volume() {
        let ctx = Context::new(
            vec![item(30, b"a")],
            vec![eviction(70, b"b", EvictionReason::NotDeclared)],
            TokenBudget(100),
        )
        .unwrap();
        assert_eq!(ctx.admitted_tokens(), 30);
        assert_eq!(ctx.headroom(), 70);
        assert_eq!(ctx.evicted_tokens(), 70);
    }

    #[test]
    fn empty_context_is_within_any_budget() {
        let ctx = Context::empty(TokenBudget(0));
        assert_eq!(ctx.admitted_tokens(), 0);
        assert_eq!(ctx.headroom(), 0);
    }

    #[test]
    fn round_trips_and_has_a_stable_digest() {
        let ctx = Context::new(
            vec![item(10, b"a")],
            vec![eviction(5, b"b", EvictionReason::NotDeclared)],
            TokenBudget(100),
        )
        .unwrap();
        let bytes = crate::to_canonical_cbor(&ctx).unwrap();
        let back: Context = crate::from_canonical_cbor(&bytes).unwrap();
        assert_eq!(ctx, back);
        assert_eq!(ctx.digest().unwrap(), back.digest().unwrap());
    }
}