gitdataai/libs/agent/tool/context.rs
2026-04-14 19:02:01 +08:00

134 lines
3.3 KiB
Rust

//! Execution context passed to each tool handler.
//!
//! Carries runtime information a tool handler needs: database, cache,
//! request metadata, and the tool registry. Cheap to clone via `Arc`.
use std::sync::Arc;
use db::cache::AppCache;
use db::database::AppDatabase;
use uuid::Uuid;
use super::registry::ToolRegistry;
/// Context available during tool execution. Cheap to clone via `Arc`.
#[derive(Clone)]
pub struct ToolContext {
inner: Arc<Inner>,
}
#[derive(Clone)]
struct Inner {
pub db: AppDatabase,
pub cache: AppCache,
pub room_id: Uuid,
pub sender_id: Option<Uuid>,
pub project_id: Uuid,
pub registry: ToolRegistry,
depth: u32,
max_depth: u32,
tool_call_count: usize,
max_tool_calls: usize,
}
impl ToolContext {
pub fn new(db: AppDatabase, cache: AppCache, room_id: Uuid, sender_id: Option<Uuid>) -> Self {
Self {
inner: Arc::new(Inner {
db,
cache,
room_id,
sender_id,
project_id: Uuid::nil(),
registry: ToolRegistry::new(),
depth: 0,
max_depth: 5,
tool_call_count: 0,
max_tool_calls: 128,
}),
}
}
pub fn with_project(mut self, project_id: Uuid) -> Self {
Arc::make_mut(&mut self.inner).project_id = project_id;
self
}
pub fn with_registry(mut self, registry: ToolRegistry) -> Self {
Arc::make_mut(&mut self.inner).registry = registry;
self
}
pub fn with_max_depth(mut self, max_depth: u32) -> Self {
Arc::make_mut(&mut self.inner).max_depth = max_depth;
self
}
pub fn with_max_tool_calls(mut self, max: usize) -> Self {
Arc::make_mut(&mut self.inner).max_tool_calls = max;
self
}
pub fn recursion_exceeded(&self) -> bool {
self.inner.depth >= self.inner.max_depth
}
pub fn tool_calls_exceeded(&self) -> bool {
self.inner.tool_call_count >= self.inner.max_tool_calls
}
/// Current recursion depth.
pub fn depth(&self) -> u32 {
self.inner.depth
}
/// Current tool call count.
pub fn tool_call_count(&self) -> usize {
self.inner.tool_call_count
}
/// Increments the tool call count.
pub(crate) fn increment_tool_calls(&mut self) {
Arc::make_mut(&mut self.inner).tool_call_count += 1;
}
/// Returns a child context for a recursive tool call (depth + 1).
pub(crate) fn child_context(&self) -> Self {
let mut inner = (*self.inner).clone();
inner.depth += 1;
Self {
inner: Arc::new(inner),
}
}
/// Database connection.
pub fn db(&self) -> &AppDatabase {
&self.inner.db
}
/// Redis cache.
pub fn cache(&self) -> &AppCache {
&self.inner.cache
}
/// Room where the original message was sent.
pub fn room_id(&self) -> Uuid {
self.inner.room_id
}
/// User who sent the original message.
pub fn sender_id(&self) -> Option<Uuid> {
self.inner.sender_id
}
/// Project context for the room.
pub fn project_id(&self) -> Uuid {
self.inner.project_id
}
/// Tool registry for this request.
pub fn registry(&self) -> &ToolRegistry {
&self.inner.registry
}
}