mod definitions;
mod mapping;
mod validation;
pub(crate) use definitions::definitions as contract_tool_definitions;
use saya_agent::ToolError;
use saya_types::{DatabaseObjectKind, DatabaseObjectRef, ProfileIdentity};
use mapping::{
REASON_NO_CONTRACT, REASON_NO_IDENTITY, REASON_NO_MATCH, REASON_PRIVACY, REASON_STALE,
REASON_STORE, contract, contract_payload, contracts, empty_for, read_payload,
};
use validation::validate_arguments;
use super::DatabaseTools;
use crate::commands::cached_schema_availability;
use crate::contracts::args::parse_qualified;
use crate::contracts::{
RecallBounds, RecallMode, RecallRequest, RetrievalPolicy, now_unix_ms, recall,
show as show_contract,
};
impl DatabaseTools {
pub(super) async fn execute_contract_tool(
&self,
name: &str,
arguments: serde_json::Value,
) -> Result<serde_json::Value, ToolError> {
validate_arguments(name, &arguments)?;
let connection = arguments
.get("connection")
.and_then(serde_json::Value::as_str);
let entry = self.registry.resolve(connection)?;
let profile_name = resolve_name(self.registry.primary_name(), connection);
let empty = empty_for(name);
let Some(identity) = entry.profile_id.as_deref() else {
return Ok(empty(REASON_NO_IDENTITY));
};
let identity =
ProfileIdentity::parse(identity).map_err(|_| ToolError::InvalidQueryArguments)?;
if !self.allow_query_data {
return Ok(empty(REASON_PRIVACY));
}
let Some(store) = self.state_db.as_ref() else {
return Ok(empty(REASON_STORE));
};
match name {
"contract_search" => {
self.contract_search(store, &identity, &profile_name, &arguments, empty)
.await
}
"contract_read" => {
self.contract_read(store, &identity, &profile_name, &arguments, empty)
.await
}
_ => Err(ToolError::UnsupportedTool),
}
}
async fn contract_search(
&self,
store: &saya_store::SqliteStateStore,
identity: &ProfileIdentity,
profile_name: &str,
arguments: &serde_json::Value,
empty: fn(&str) -> serde_json::Value,
) -> Result<serde_json::Value, ToolError> {
let terms: Vec<String> = arguments
.get("terms")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|t| t.as_str().map(str::to_string))
.collect()
})
.unwrap_or_default();
let cached = cached_schema_availability(store, identity).await;
let schema_pair = (identity.clone(), cached);
let schemas = std::slice::from_ref(&schema_pair);
let request = RecallRequest {
profiles: std::slice::from_ref(identity),
explicit_refs: &[],
terms: &terms,
allow_database_context: true,
schemas,
now_unix_ms: now_unix_ms(),
bounds: RecallBounds::defaults(),
recall_mode: RecallMode::Confirmed,
admit_candidate: None,
policy: RetrievalPolicy::ForModel,
};
let outcome = recall(store, request).await;
if outcome.diagnostics.store_unavailable {
return Ok(empty(REASON_STORE));
}
if outcome.contracts.is_empty() {
if outcome.diagnostics.excluded_by_schema > 0 {
return Ok(empty(REASON_STALE));
}
return Ok(empty(REASON_NO_MATCH));
}
let payload: Vec<serde_json::Value> = outcome
.contracts
.iter()
.map(|c| contract_payload(c, profile_name))
.collect();
Ok(contracts(payload))
}
async fn contract_read(
&self,
store: &saya_store::SqliteStateStore,
identity: &ProfileIdentity,
profile_name: &str,
arguments: &serde_json::Value,
empty: fn(&str) -> serde_json::Value,
) -> Result<serde_json::Value, ToolError> {
let table = arguments
.get("table")
.and_then(serde_json::Value::as_str)
.ok_or(ToolError::InvalidQueryArguments)?;
let qualified = parse_qualified(table).map_err(|_| ToolError::InvalidQueryArguments)?;
let object = DatabaseObjectRef::new(
identity.clone(),
&qualified.catalog,
&qualified.schema,
&qualified.object,
DatabaseObjectKind::Table,
)
.map_err(|_| ToolError::InvalidQueryArguments)?;
let schema = cached_schema_availability(store, identity).await;
match show_contract(
store,
&object,
&schema,
RetrievalPolicy::ForModel,
now_unix_ms(),
)
.await
{
Ok(Some(retrieved)) => Ok(contract(read_payload(&retrieved, profile_name))),
Ok(None) => Ok(empty(REASON_NO_CONTRACT)),
Err(_) => Ok(empty(REASON_STORE)),
}
}
}
fn resolve_name(primary: &str, connection: Option<&str>) -> String {
match connection {
Some(name) if !name.is_empty() => name.to_string(),
_ => primary.to_string(),
}
}