use std::sync::Arc;
use crate::server::tool_trait::ToolOutput;
use lean_ctx::core::context_admission::capture::AdmissionCapture;
use lean_ctx::core::context_admission::receipt_store::{self, frame, observed, trace};
use super::{ErrorV1, agent_shell_outcome, output_withheld};
pub(super) fn protect_output(tool: &str, mut output: ToolOutput) -> Result<ToolOutput, ErrorV1> {
use crate::server::policy_guard;
use rmcp::model::{CallToolResult, ContentBlock};
if lean_ctx::core::policy::runtime::is_active() {
if output.mode.as_deref() == Some("error") {
return Err(output_withheld());
}
output.text =
policy_guard::protect_result(tool, &output.text).map_err(|_| output_withheld())?;
if let Some(blocks) = output.content_blocks.as_mut() {
for block in blocks {
let Some(text) = block.as_text() else {
return Err(output_withheld());
};
*block = ContentBlock::text(
policy_guard::protect_result(tool, &text.text)
.map_err(|_| output_withheld())?,
);
}
}
}
let blocks = output.content_blocks.as_ref().map_or(0, Vec::len);
let mut content = output.content_blocks.clone().unwrap_or_default();
content.push(ContentBlock::text(&output.text));
let released = policy_guard::release_result(tool, CallToolResult::success(content));
if released.is_error == Some(true) || released.content.len() != blocks + 1 {
return Err(output_withheld());
}
let mut metadata = CallToolResult::success(Vec::new());
metadata.structured_content = Some(serde_json::json!({
"mode": output.mode,
"shell": agent_shell_outcome(output.shell_outcome.as_ref()),
}));
if policy_guard::release_result(tool, metadata).is_error == Some(true) {
return Err(output_withheld());
}
let mut delivered = released.content;
let Some(text) = delivered
.pop()
.and_then(|block| block.as_text().map(|text| text.text.clone()))
else {
return Err(output_withheld());
};
output.text = text;
if output.content_blocks.is_some() {
output.content_blocks = Some(delivered);
}
output.saved_tokens = output
.original_tokens
.saturating_sub(lean_ctx::core::tokens::count_tokens(&output.text));
Ok(output)
}
pub(super) struct FrameCall(frame::FrameCall);
impl FrameCall {
pub(super) fn open(tool: &str) -> Self {
Self(frame::FrameCall::open(tool))
}
pub(super) fn capture(&self) -> Arc<AdmissionCapture> {
self.0.capture()
}
pub(super) fn seal(
self,
project_root: &str,
frame: &[u8],
is_error: bool,
) -> Result<Option<observed::Prepared>, frame::ReceiptMissing> {
self.0
.seal(project_root, frame, is_error, trace::Delivery::tool_session)
}
}
pub(super) fn returned(prepared: &observed::Prepared, written: &[u8]) {
prepared.returned_to_transport(&receipt_store::output_digest(written));
}
pub(super) fn not_persisted_message() -> String {
receipt_store::not_persisted_notice()
}