supercode_harness/tools/

context_budget.rs

//! BP-3 (catalog row "Context-budget tools", cx§1's `token_budget`
//! feature): `get_context_remaining` and `new_context`, the MODEL's doors
//! onto two mechanisms the agent already owns.
//!
//! Neither tool implements accounting or clearing of its own. BP-4 built
//! both, for the operator's `/context` and `/handoff`:
//!
//! * `crate::agent::Agent::context_usage` — the live context-window
//!   accounting, over the very estimates the context guard enforces, so
//!   what the model is told and what refuses an oversized turn can never
//!   disagree. [`GetContextRemainingTool`] reports exactly that struct.
//! * `crate::agent::Agent::new_context` — the in-session fresh window
//!   (system prompt + a handoff marker naming the objective + the curated
//!   recent tail, with the set-aside turns kept in the transcript sidecar
//!   whenever one is attached). [`NewContextTool`] asks for exactly that.
//!
//! Both travel through ONE shared [`ContextBudget`], held as an `Arc` on
//! [`crate::tools::ToolContext`]: the agent publishes the accounting when
//! the tool asks for it, and the tool parks its fresh-window request there
//! for the agent to apply the moment the tool round ends — so the very next
//! model request is the fresh window.
//!
//! **Why the park-and-apply split.** A `Tool::execute` sees its arguments
//! and the ambient context, never the agent's transcript. Keeping the
//! rewrite in the one place that owns `history` is the same reason
//! `tool_search`/`expand_reduction` are agent intrinsics — and it is what
//! keeps the model's `new_context` and the operator's `/handoff` running
//! the same code, not two drifting copies of one idea.

use std::sync::Mutex;

use async_trait::async_trait;
use serde::Deserialize;
use serde_json::{json, Value};

use crate::error::{Error, Result};
use crate::tools::{Tool, ToolContext};

/// Registered name of the remaining-budget tool.
pub const GET_CONTEXT_REMAINING: &str = "get_context_remaining";

/// Registered name of the fresh-window tool.
pub const NEW_CONTEXT: &str = "new_context";

/// One parked fresh-window request, as the model stated it. The fields are
/// exactly `Agent::new_context`'s parameters — this type carries a request
/// across the tool/agent boundary, it does not add semantics of its own.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NewContextRequest {
    /// The objective the fresh window opens with.
    pub objective: String,
    /// How many trailing messages to keep. `None` uses the same
    /// token-budget-derived count `core.compaction.keep_recent_tokens`
    /// governs, so the keep-set is curated by the budget the rest of the
    /// compaction machinery uses rather than a number the model guessed.
    pub keep_recent: Option<usize>,
}

/// The shared context accounting the two tools read and write.
#[derive(Debug, Default)]
pub struct ContextBudget {
    /// The last `Agent::context_usage()` the agent published, serialized.
    usage: Mutex<Option<Value>>,
    pending: Mutex<Option<NewContextRequest>>,
}

impl ContextBudget {
    /// A budget with nothing published yet.
    pub fn new() -> Self {
        Self::default()
    }

    /// Publish the agent's `ContextUsage` (already serialized, so this
    /// module needs no dependency on the agent's own types).
    pub fn publish(&self, usage: Value) {
        if let Ok(mut slot) = self.usage.lock() {
            *slot = Some(usage);
        }
    }

    /// The last published accounting, if the agent has published one.
    pub fn snapshot(&self) -> Option<Value> {
        self.usage.lock().ok().and_then(|slot| slot.clone())
    }

    /// Park a fresh-window request for the agent to apply.
    pub fn request_new_context(&self, request: NewContextRequest) {
        if let Ok(mut pending) = self.pending.lock() {
            *pending = Some(request);
        }
    }

    /// Take the parked request, if any (the agent calls this once per tool
    /// round).
    pub fn take_new_context(&self) -> Option<NewContextRequest> {
        self.pending.lock().ok().and_then(|mut p| p.take())
    }
}

/// `get_context_remaining` — how much of the context window is left.
#[derive(Debug, Default)]
pub struct GetContextRemainingTool;

#[async_trait]
impl Tool for GetContextRemainingTool {
    fn name(&self) -> &str {
        GET_CONTEXT_REMAINING
    }
    fn description(&self) -> &str {
        "Report how much of the model's context window this conversation is using and how many \
         tokens remain (messages, tool schemas, the reply reserve, and whether the next request \
         would still fit). Use it before starting something long, or to decide whether to call \
         new_context."
    }
    fn parameters(&self) -> Value {
        json!({"type": "object", "properties": {}, "additionalProperties": false})
    }
    fn structured_output(&self) -> bool {
        true
    }
    async fn execute(&self, _args: Value, ctx: &ToolContext) -> Result<String> {
        let Some(usage) = ctx.context_budget.snapshot() else {
            return Err(Error::tool(
                self.name(),
                "no context accounting is available for this session (this tool reports the \
                 running agent's own figures, and none have been published)",
            ));
        };
        Ok(usage.to_string())
    }
}

#[derive(Debug, Deserialize)]
struct NewContextArgs {
    objective: String,
    #[serde(default)]
    keep_recent: Option<usize>,
}

/// `new_context` — continue in a fresh window seeded with an objective and
/// the curated recent tail.
#[derive(Debug, Default)]
pub struct NewContextTool;

#[async_trait]
impl Tool for NewContextTool {
    fn name(&self) -> &str {
        NEW_CONTEXT
    }
    fn description(&self) -> &str {
        "Continue this session in a fresh context window: state the objective the new window \
         opens with. The system prompt, a handoff marker carrying that objective, and the most \
         recent messages are kept; everything earlier is set aside (and stays in the session's \
         transcript). Takes effect immediately after this tool round."
    }
    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "objective": {
                    "type": "string",
                    "description": "What the fresh window is for — the one paragraph the new \
                                    context opens with."
                },
                "keep_recent": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "How many of the most recent messages to keep. Omit to use \
                                    the session's own keep-recent token budget."
                }
            },
            "required": ["objective"],
            "additionalProperties": false
        })
    }
    async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<String> {
        let a: NewContextArgs =
            serde_json::from_value(args).map_err(|e| Error::InvalidArguments {
                tool: self.name().to_string(),
                message: e.to_string(),
            })?;
        let objective = a.objective.trim().to_string();
        if objective.is_empty() {
            return Err(Error::InvalidArguments {
                tool: self.name().to_string(),
                message: "state an objective for the fresh context window".to_string(),
            });
        }
        ctx.context_budget.request_new_context(NewContextRequest {
            objective: objective.clone(),
            keep_recent: a.keep_recent,
        });
        Ok(format!(
            "A fresh context window is queued and takes effect before your next turn. \
             Objective: {objective}. The system prompt, a handoff marker with that objective, \
             and the recent tail are kept; earlier turns are set aside."
        ))
    }
}

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

    #[tokio::test]
    async fn remaining_reports_the_agents_published_accounting_verbatim() {
        let ctx = ToolContext::new(std::env::temp_dir());
        ctx.context_budget.publish(json!({
            "model": "test/model",
            "context_limit": 200_000,
            "projected_tokens": 30_000,
            "remaining_tokens": 166_000,
            "used_pct": 15,
            "fits": true,
        }));
        let out = GetContextRemainingTool
            .execute(json!({}), &ctx)
            .await
            .unwrap();
        let v: Value = serde_json::from_str(&out).unwrap();
        assert_eq!(v["remaining_tokens"], 166_000);
        assert_eq!(v["used_pct"], 15);
        assert_eq!(v["fits"], true);
    }

    #[tokio::test]
    async fn remaining_refuses_before_anything_is_published() {
        let ctx = ToolContext::new(std::env::temp_dir());
        let err = GetContextRemainingTool
            .execute(json!({}), &ctx)
            .await
            .expect_err("nothing published yet");
        assert!(err.to_string().contains("no context accounting"), "{err}");
    }

    #[tokio::test]
    async fn new_context_parks_a_request_for_the_agent() {
        let ctx = ToolContext::new(std::env::temp_dir());
        let out = NewContextTool
            .execute(
                json!({"objective": "finish the parser", "keep_recent": 2}),
                &ctx,
            )
            .await
            .unwrap();
        assert!(out.contains("finish the parser"), "{out}");
        let parked = ctx.context_budget.take_new_context().expect("parked");
        assert_eq!(
            parked,
            NewContextRequest {
                objective: "finish the parser".into(),
                keep_recent: Some(2),
            }
        );
        assert!(
            ctx.context_budget.take_new_context().is_none(),
            "taken once"
        );
    }

    #[tokio::test]
    async fn new_context_defers_the_keep_set_to_the_session_budget_by_default() {
        let ctx = ToolContext::new(std::env::temp_dir());
        NewContextTool
            .execute(json!({"objective": "ship it"}), &ctx)
            .await
            .unwrap();
        assert_eq!(
            ctx.context_budget.take_new_context().unwrap().keep_recent,
            None
        );
    }

    #[tokio::test]
    async fn new_context_needs_an_objective() {
        let ctx = ToolContext::new(std::env::temp_dir());
        let err = NewContextTool
            .execute(json!({"objective": "   "}), &ctx)
            .await
            .expect_err("blank objective must be refused");
        assert!(err.to_string().contains("state an objective"), "{err}");
        assert!(ctx.context_budget.take_new_context().is_none());
    }
}