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};
pub const GET_CONTEXT_REMAINING: &str = "get_context_remaining";
pub const NEW_CONTEXT: &str = "new_context";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NewContextRequest {
pub objective: String,
pub keep_recent: Option<usize>,
}
#[derive(Debug, Default)]
pub struct ContextBudget {
usage: Mutex<Option<Value>>,
pending: Mutex<Option<NewContextRequest>>,
}
impl ContextBudget {
pub fn new() -> Self {
Self::default()
}
pub fn publish(&self, usage: Value) {
if let Ok(mut slot) = self.usage.lock() {
*slot = Some(usage);
}
}
pub fn snapshot(&self) -> Option<Value> {
self.usage.lock().ok().and_then(|slot| slot.clone())
}
pub fn request_new_context(&self, request: NewContextRequest) {
if let Ok(mut pending) = self.pending.lock() {
*pending = Some(request);
}
}
pub fn take_new_context(&self) -> Option<NewContextRequest> {
self.pending.lock().ok().and_then(|mut p| p.take())
}
}
#[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>,
}
#[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."
))
}
}