use super::*;
use async_trait::async_trait;
use saya_agent::{ToolError, ToolExecutor};
use saya_connectors::DatabaseConnector;
use saya_store::{SchemaStore, SqliteStateStore};
use saya_types::{
ClaimOrigin, ClaimPayload, Column, ConnectionError, Database, DatabaseObjectKind,
DatabaseObjectRef, DatabaseProfile, KnowledgeSlot, KnowledgeState, ProfileIdentity,
QueryRequest, QueryResult, Schema, SchemaFingerprint, SchemaTree, SqlDialect, Table,
};
use std::{
fs,
path::{Path, PathBuf},
time::{SystemTime, UNIX_EPOCH},
};
use crate::connection::{ConnectionEntry, ConnectionRegistry};
struct IdleConnector;
#[async_trait]
impl DatabaseConnector for IdleConnector {
fn dialect(&self) -> SqlDialect {
SqlDialect::DuckDb
}
async fn connect(&self) -> Result<(), ConnectionError> {
Ok(())
}
async fn schema(&self) -> Result<SchemaTree, ConnectionError> {
Ok(SchemaTree::default())
}
async fn execute(&self, req: QueryRequest) -> Result<QueryResult, ConnectionError> {
Ok(QueryResult::empty(req.sql))
}
}
fn temp_root(label: &str) -> PathBuf {
let stamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let root = std::env::temp_dir().join(format!(
"saya-contract-tools-{label}-{}-{stamp}",
std::process::id()
));
fs::create_dir_all(&root).unwrap();
root
}
fn profile_identity(name: &str) -> ProfileIdentity {
crate::profile_identity::profile_identity(
name,
&DatabaseProfile::DuckDb {
path: "contract-tools.duckdb".into(),
read_only: Some(true),
},
Path::new("/contract-tools-test/connections.toml"),
)
}
fn registry_with_primary(name: &str, identity: &ProfileIdentity) -> ConnectionRegistry {
let mut registry = ConnectionRegistry::new(name);
registry.insert(
name,
ConnectionEntry {
connector: Box::new(IdleConnector),
dialect: SqlDialect::DuckDb,
profile_id: Some(identity.as_str().to_string()),
},
);
registry
}
async fn store_at(db: &Path) -> SqliteStateStore {
let store = SqliteStateStore::new(db);
store
.upsert_schema(profile_identity("primary").as_str(), &SchemaTree::default())
.await
.unwrap();
store
}
fn current_table(object: &DatabaseObjectRef) -> Table {
Table {
name: object.object().to_string(),
columns: vec![Column {
name: "id".into(),
data_type: "bigint".into(),
nullable: false,
}],
}
}
async fn seed_current(store: &SqliteStateStore, object: &DatabaseObjectRef) -> SchemaFingerprint {
let table = current_table(object);
let tree = SchemaTree {
databases: vec![Database {
name: object.catalog().to_string(),
schemas: vec![Schema {
name: object.schema().to_string(),
tables: vec![table.clone()],
}],
}],
};
store
.upsert_schema(object.profile().as_str(), &tree)
.await
.unwrap();
SchemaFingerprint::of_table(DatabaseObjectKind::Table, &table)
}
async fn put_knowledge_item(
store: &SqliteStateStore,
object: &DatabaseObjectRef,
payload: ClaimPayload,
state: KnowledgeState,
) {
use saya_store::{KnowledgeItemRequest, KnowledgeItemStore};
use saya_types::SchemaBinding;
let _ = seed_current(store, object).await;
let slot = slot_for(&payload);
let binding = SchemaBinding::derive(&slot, &payload).expect("slot/payload agree");
let request = KnowledgeItemRequest {
object: object.clone(),
slot,
value: payload,
source: if state == KnowledgeState::Active {
ClaimOrigin::UserExplicit
} else {
ClaimOrigin::AssistantInferred
},
state,
schema_binding_json: serde_json::to_string(&binding).unwrap(),
fingerprint: crate::commands::unobserved_fingerprint(),
};
store.put_knowledge_item(request).await.unwrap();
}
fn slot_for(payload: &ClaimPayload) -> KnowledgeSlot {
match payload {
ClaimPayload::TableDescription { .. } => KnowledgeSlot::TableDescription,
ClaimPayload::TableAlias { .. } => KnowledgeSlot::TableAlias,
ClaimPayload::TableGrain { .. } => KnowledgeSlot::TableGrain,
ClaimPayload::DefaultTimeColumn { .. } => KnowledgeSlot::TableDefaultTime,
ClaimPayload::ColumnDescription { column, .. } => KnowledgeSlot::ColumnDescription {
column: column.clone(),
},
ClaimPayload::ColumnRole { column, .. } => KnowledgeSlot::ColumnRole {
column: column.clone(),
},
_ => panic!("no slot for payload {:?}", payload),
}
}
fn object_ref(profile: &ProfileIdentity, name: &str) -> DatabaseObjectRef {
DatabaseObjectRef::new(
profile.clone(),
"catalog",
"public",
name,
DatabaseObjectKind::Table,
)
.unwrap()
}
#[tokio::test]
async fn contract_search_returns_confirmed_and_hides_candidates() {
let root = temp_root("search_confirmed");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let obj = object_ref(&identity, "orders");
put_knowledge_item(
&store,
&obj,
ClaimPayload::table_alias("orders").unwrap(),
KnowledgeState::Active,
)
.await;
put_knowledge_item(
&store,
&obj,
ClaimPayload::table_alias("secret_candidate").unwrap(),
KnowledgeState::Pending,
)
.await;
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store),
);
let res = tools
.execute("contract_search", serde_json::json!({"terms": ["orders"]}))
.await
.expect("contract_search should succeed");
let contracts = res
.get("contracts")
.and_then(|v| v.as_array())
.expect("result carries a `contracts` array");
assert_eq!(contracts.len(), 1, "one matched object");
let claims = contracts[0]["claims"].as_array().expect("claims array");
assert!(
claims
.iter()
.any(|c| c["kind"] == "table_alias" && c["value"] == "orders"),
"confirmed alias should appear: {res}"
);
assert!(
!claims.iter().any(|c| c["value"] == "secret_candidate"),
"candidate alias must not appear: {res}"
);
let _ = fs::remove_dir_all(root);
}
#[tokio::test]
async fn contract_read_returns_one_contract_and_rejects_malformed_table() {
let root = temp_root("read_one");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let obj = object_ref(&identity, "orders");
put_knowledge_item(
&store,
&obj,
ClaimPayload::table_description("sales fact table").unwrap(),
KnowledgeState::Active,
)
.await;
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store.clone()),
);
let res = tools
.execute(
"contract_read",
serde_json::json!({"table": "catalog.public.orders"}),
)
.await
.expect("contract_read should succeed");
let contract = res.get("contract").expect("one contract");
assert_eq!(contract["object"], "catalog.public.orders");
assert!(
contract["claims"]
.as_array()
.unwrap()
.iter()
.any(|c| c["kind"] == "table_description" && c["value"] == "sales fact table"),
"stored description should appear: {res}"
);
let err = tools
.execute("contract_read", serde_json::json!({"table": "orders"}))
.await
.expect_err("malformed table should be a typed error");
assert!(
matches!(err, ToolError::InvalidQueryArguments),
"malformed table should be InvalidQueryArguments: {err}"
);
let defs = DatabaseTools::definitions(true, true, false);
let read_def = defs
.iter()
.find(|t| t.name == "contract_read")
.expect("contract_read is defined");
assert!(
read_def.description.contains("catalog.schema"),
"the description names the expected form: {}",
read_def.description
);
let _ = fs::remove_dir_all(root);
}
#[tokio::test]
async fn contract_tool_argument_validation_errors_are_typed() {
let identity = profile_identity("primary");
let tools =
DatabaseTools::with_registry(registry_with_primary("primary", &identity), 100, true, None);
let err = tools
.execute(
"contract_search",
serde_json::json!({"terms": ["x"], "bogus": 1}),
)
.await
.expect_err("unknown property must be rejected");
assert!(matches!(err, ToolError::UnsupportedProperty), "got: {err}");
let err = tools
.execute(
"contract_search",
serde_json::json!({"terms": ["x"], "connection": 7}),
)
.await
.expect_err("non-string connection must be rejected");
assert!(matches!(err, ToolError::ConnectionNotString), "got: {err}");
let terms: Vec<String> = (0..17).map(|i| format!("t{i}")).collect();
let err = tools
.execute("contract_search", serde_json::json!({"terms": terms}))
.await
.expect_err("17 terms must be rejected");
assert!(
matches!(err, ToolError::InvalidQueryArguments),
"too many terms should be InvalidQueryArguments: {err}"
);
let err = tools
.execute("contract_search", serde_json::json!({"terms": "orders"}))
.await
.expect_err("terms-as-string must be rejected");
assert!(
matches!(err, ToolError::InvalidQueryArguments),
"got: {err}"
);
}
#[tokio::test]
async fn no_opaque_profile_identity_in_any_tool_result() {
let root = temp_root("no_identity");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let obj = object_ref(&identity, "orders");
put_knowledge_item(
&store,
&obj,
ClaimPayload::table_alias("orders").unwrap(),
KnowledgeState::Active,
)
.await;
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store),
);
let search = tools
.execute("contract_search", serde_json::json!({"terms": ["orders"]}))
.await
.unwrap();
let read = tools
.execute(
"contract_read",
serde_json::json!({"table": "catalog.public.orders"}),
)
.await
.unwrap();
let identity_str = identity.as_str();
assert_eq!(identity_str.len(), 66);
assert!(identity_str.starts_with("p-"));
let search_text = serde_json::to_string(&search).unwrap();
let read_text = serde_json::to_string(&read).unwrap();
assert!(
!search_text.contains(identity_str),
"identity leaked into contract_search: {search_text}"
);
assert!(
!read_text.contains(identity_str),
"identity leaked into contract_read: {read_text}"
);
assert!(search_text.contains("\"profile\":\"primary\""));
assert!(read_text.contains("\"profile\":\"primary\""));
let _ = fs::remove_dir_all(root);
}
#[tokio::test]
async fn privacy_gate_closed_returns_empty_without_claims() {
let root = temp_root("privacy_closed");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let obj = object_ref(&identity, "orders");
put_knowledge_item(
&store,
&obj,
ClaimPayload::table_alias("secret_alias").unwrap(),
KnowledgeState::Active,
)
.await;
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
false,
Some(store),
);
let res = tools
.execute(
"contract_search",
serde_json::json!({"terms": ["secret_alias"]}),
)
.await
.expect("privacy gate returns empty, not an error");
assert!(
res.get("contracts")
.map(|c| c.as_array().unwrap().is_empty())
.unwrap_or(true),
"no contracts when the gate is closed: {res}"
);
let text = serde_json::to_string(&res).unwrap();
assert!(!text.contains("secret_alias"));
let res = tools
.execute(
"contract_read",
serde_json::json!({"table": "catalog.public.orders"}),
)
.await
.expect("privacy gate returns empty, not an error");
assert!(
res.get("contract").is_none(),
"no contract when closed: {res}"
);
let text = serde_json::to_string(&res).unwrap();
assert!(!text.contains("secret_alias"));
let _ = fs::remove_dir_all(root);
}
#[tokio::test]
async fn unopenable_store_returns_empty_result_not_error() {
let root = temp_root("unopenable");
fs::write(root.join("blocker"), b"x").unwrap();
let bad = root.join("blocker/state.sqlite3");
let store = SqliteStateStore::new(&bad);
let identity = profile_identity("primary");
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store),
);
let search = tools
.execute("contract_search", serde_json::json!({"terms": ["orders"]}))
.await
.expect("store failure must be an empty result, not Err");
assert!(
search
.get("contracts")
.map(|c| c.as_array().unwrap().is_empty())
.unwrap_or(true),
"empty contracts on store failure: {search}"
);
let read = tools
.execute(
"contract_read",
serde_json::json!({"table": "catalog.public.orders"}),
)
.await
.expect("store failure must be an empty result, not Err");
assert!(
read.get("contract").is_none(),
"no contract on store failure: {read}"
);
let _ = fs::remove_dir_all(root);
}
#[test]
fn contract_tool_definitions_are_read_only_unapproved_and_not_writable() {
let tools = DatabaseTools::definitions(true, true, false);
let search = tools
.iter()
.find(|t| t.name == "contract_search")
.expect("contract_search is registered when a store is present and data is allowed");
let read = tools
.iter()
.find(|t| t.name == "contract_read")
.expect("contract_read is registered");
for tool in [search, read] {
assert!(tool.read_only, "{} must be read_only", tool.name);
assert!(
!tool.effect.requires_approval,
"{} must not require approval",
tool.name
);
assert!(
!tool.effect.external_side_effect,
"{} must have no external side effect",
tool.name
);
assert!(
tool.effect.database_data,
"{} is database-derived and subject to the sharing gate",
tool.name
);
}
}
#[test]
fn contract_tools_are_hidden_without_a_store_or_when_the_gate_is_closed() {
let with_store_and_gate = DatabaseTools::definitions(true, true, false);
assert!(
with_store_and_gate
.iter()
.any(|t| t.name == "contract_search")
);
assert!(
with_store_and_gate
.iter()
.any(|t| t.name == "contract_read")
);
let no_store = DatabaseTools::definitions(true, false, false);
assert!(!no_store.iter().any(|t| t.name == "contract_search"));
assert!(!no_store.iter().any(|t| t.name == "contract_read"));
let gate_closed = DatabaseTools::definitions(false, true, false);
assert!(!gate_closed.iter().any(|t| t.name == "contract_search"));
assert!(!gate_closed.iter().any(|t| t.name == "contract_read"));
}
#[tokio::test]
async fn contract_read_truncates_claims_past_the_bound_and_says_so() {
let root = temp_root("read_truncated");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let obj = object_ref(&identity, "wide");
for col in 0..5 {
for n in 0..4 {
put_knowledge_item(
&store,
&obj,
ClaimPayload::column_description(format!("c{col}"), format!("desc {n}")).unwrap(),
KnowledgeState::Active,
)
.await;
}
}
let wide = SchemaTree {
databases: vec![Database {
name: obj.catalog().to_string(),
schemas: vec![Schema {
name: obj.schema().to_string(),
tables: vec![Table {
name: obj.object().to_string(),
columns: (0..5)
.map(|col| Column {
name: format!("c{col}"),
data_type: "text".into(),
nullable: true,
})
.collect(),
}],
}],
}],
};
store
.upsert_schema(obj.profile().as_str(), &wide)
.await
.unwrap();
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store),
);
let res = tools
.execute(
"contract_read",
serde_json::json!({"table": "catalog.public.wide"}),
)
.await
.expect("contract_read should succeed");
let contract = res.get("contract").expect("one contract");
let claims = contract["claims"].as_array().expect("claims array");
assert!(
claims.len() <= 12,
"claims bounded to 12, got {}: {res}",
claims.len()
);
assert!(
contract["truncated"] == serde_json::Value::Bool(true),
"truncation must be flagged explicitly: {res}"
);
let _ = fs::remove_dir_all(root);
}
#[tokio::test]
async fn contract_read_on_unknown_object_returns_empty_with_reason() {
let root = temp_root("read_unknown");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store),
);
let res = tools
.execute(
"contract_read",
serde_json::json!({"table": "catalog.public.missing"}),
)
.await
.expect("no contract is an empty result, not an error");
assert!(res.get("contract").is_none(), "no contract object: {res}");
let reason = res
.get("reason")
.and_then(|v| v.as_str())
.expect("a short reason explains the empty result");
assert!(!reason.is_empty());
let _ = fs::remove_dir_all(root);
}
#[tokio::test]
async fn contract_search_drops_a_stale_claim_and_says_so() {
let root = temp_root("search_stale");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let obj = object_ref(&identity, "orders");
let drifted = SchemaTree {
databases: vec![Database {
name: obj.catalog().to_string(),
schemas: vec![Schema {
name: obj.schema().to_string(),
tables: vec![Table {
name: obj.object().to_string(),
columns: vec![Column {
name: "id".into(),
data_type: "bigint".into(),
nullable: false,
}],
}],
}],
}],
};
store
.upsert_schema(obj.profile().as_str(), &drifted)
.await
.unwrap();
put_knowledge_item(
&store,
&obj,
ClaimPayload::default_time_column("created_at", None).unwrap(),
KnowledgeState::Active,
)
.await;
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store),
);
let res = tools
.execute("contract_search", serde_json::json!({"terms": ["orders"]}))
.await
.expect("contract_search should succeed");
let contracts = res
.get("contracts")
.and_then(|v| v.as_array())
.expect("result carries a `contracts` array");
assert!(
contracts.is_empty(),
"stale contract must not appear: {res}"
);
let reason = res
.get("reason")
.and_then(|v| v.as_str())
.expect("a stale exclusion carries a reason");
assert!(
reason.contains("stale"),
"reason names staleness, got: {reason}"
);
let text = serde_json::to_string(&res).unwrap();
assert!(
!text.contains("created_at"),
"stale claim text must not leak: {text}"
);
let _ = fs::remove_dir_all(root);
}
#[tokio::test]
async fn contract_read_on_a_stale_object_reports_stale_with_no_claims() {
let root = temp_root("read_stale");
let store = store_at(&root.join("state.sqlite3")).await;
let identity = profile_identity("primary");
let obj = object_ref(&identity, "orders");
put_knowledge_item(
&store,
&obj,
ClaimPayload::default_time_column("created_at", None).unwrap(),
KnowledgeState::Active,
)
.await;
let tools = DatabaseTools::with_registry(
registry_with_primary("primary", &identity),
100,
true,
Some(store),
);
let res = tools
.execute(
"contract_read",
serde_json::json!({"table": "catalog.public.orders"}),
)
.await
.expect("contract_read should succeed");
let contract = res.get("contract").expect("contract object present: {res}");
assert_eq!(
contract["schema_state"], "stale",
"stale object is reported stale: {res}"
);
assert_eq!(
contract["object"], "catalog.public.orders",
"the stale object is named: {res}"
);
let claims = contract
.get("claims")
.and_then(|v| v.as_array())
.expect("claims array present");
assert!(
claims.is_empty(),
"no claims to act on for a stale object: {res}"
);
let text = serde_json::to_string(&res).unwrap();
assert!(
!text.contains("created_at"),
"stale claim text must not leak: {text}"
);
let _ = fs::remove_dir_all(root);
}