1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
use saya_agent::ToolError;
use saya_connectors::DatabaseConnector;
use saya_types::SqlDialect;
use std::collections::HashMap;
use std::fmt;
/// Represents a live database connection entry.
#[allow(dead_code)]
pub(crate) struct ConnectionEntry {
/// The database connector implementation.
pub(crate) connector: Box<dyn DatabaseConnector>,
/// The SQL dialect of the connection.
pub(crate) dialect: SqlDialect,
/// Optional database profile identifier.
pub(crate) profile_id: Option<String>,
}
impl fmt::Debug for ConnectionEntry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ConnectionEntry")
.field("dialect", &self.dialect)
.field("profile_id", &self.profile_id)
.finish()
}
}
/// A typed registry of live database connections for multi-database agent navigation.
#[allow(dead_code)]
pub(crate) struct ConnectionRegistry {
primary: String,
names: Vec<String>,
map: HashMap<String, ConnectionEntry>,
}
#[allow(dead_code)]
impl ConnectionRegistry {
/// Creates an empty registry whose primary connection is `primary`.
pub(crate) fn new(primary: impl Into<String>) -> Self {
Self {
primary: primary.into(),
names: Vec::new(),
map: HashMap::new(),
}
}
/// Inserts or replaces a named connection, preserving first-seen order.
pub(crate) fn insert(&mut self, name: impl Into<String>, entry: ConnectionEntry) {
let name = name.into();
if !self.map.contains_key(&name) {
self.names.push(name.clone());
}
self.map.insert(name, entry);
}
/// Returns the primary connection name.
pub(crate) fn primary_name(&self) -> &str {
&self.primary
}
/// The primary connection's registry name, when one is connected — the
/// same resolution `resolve(None)` performs, so a session grant names a
/// connection the registry would actually resolve a call to. `None`
/// when the registry holds no connection (no primary to name — never a
/// guessed name).
pub(crate) fn primary(&self) -> Option<&str> {
self.resolve(None).ok().map(|_| self.primary.as_str())
}
/// Returns the number of connections in the registry.
pub(crate) fn len(&self) -> usize {
self.map.len()
}
/// Returns true if the registry contains no connections.
pub(crate) fn is_empty(&self) -> bool {
self.map.is_empty()
}
/// Connection names in insertion order.
pub(crate) fn names(&self) -> Vec<&str> {
self.names.iter().map(String::as_str).collect()
}
/// Every connection with its name, in insertion order. This is the
/// currently included/connected set that multi-database fan-out runs
/// against (primary plus each secondary that connected).
pub(crate) fn entries(&self) -> Vec<(&str, &ConnectionEntry)> {
self.names
.iter()
.filter_map(|name| self.map.get(name).map(|entry| (name.as_str(), entry)))
.collect()
}
/// Resolves an optional connection name to an entry. `None` or empty -> primary.
/// Unknown name -> Err with a message listing the available names.
/// Empty registry -> Err("no database profile is selected").
pub(crate) fn resolve(&self, name: Option<&str>) -> Result<&ConnectionEntry, ToolError> {
let target = match name {
None | Some("") => self.primary.as_str(),
Some(n) => n,
};
if let Some(entry) = self.map.get(target) {
Ok(entry)
} else if self.is_empty() {
Err(ToolError::NoConnectionSelected)
} else {
let available = self.names().join(", ");
Err(ToolError::UnknownConnection {
target: target.to_string(),
available,
})
}
}
/// The connection name whose stored profile identity is `identity`, if any.
///
/// The observation collector records the opaque [`ProfileIdentity`] (the
/// stable object identity, never a name a user chose); the turn record
/// resolves objects back to a connection *by name*, because that is what
/// [`ConnectionRegistry::resolve`] keys on. This is the one place the
/// identity the observation carries is turned into the name the resolver
/// expects, so a turn that touched a non-primary connection is attributed to
/// the connection it actually used rather than collapsed onto the primary.
pub(crate) fn name_for_identity(&self, identity: &str) -> Option<&str> {
self.entries()
.into_iter()
.find(|(_, entry)| entry.profile_id.as_deref() == Some(identity))
.map(|(name, _)| name)
}
/// The dialect of every connection, in registration order.
///
/// Callers that must phrase something per engine — SQL naming rules, for
/// one — need this even when `describe_context` stays silent because there
/// is only a single connection.
pub(crate) fn dialects(&self) -> impl Iterator<Item = SqlDialect> + '_ {
self.names
.iter()
.filter_map(|name| self.map.get(name).map(|entry| entry.dialect))
}
/// System-prompt addendum listing every connection and its dialect, instructing the
/// model to pass the `connection` argument and inspect each database separately then
/// combine findings. Returns None when there is <= 1 connection (no navigation needed).
pub(crate) fn describe_context(&self) -> Option<String> {
if self.len() <= 1 {
return None;
}
let mut lines = Vec::new();
lines.push("Available database connections:".to_string());
for name in &self.names {
if let Some(entry) = self.map.get(name) {
lines.push(format!("- {name} ({})", entry.dialect.as_str()));
}
}
lines.push(
"To inspect a database, pass its `connection` argument to schema and query tools. Inspect each database separately and combine your findings. When the same query should run against every connected database, call `bounded_sql_query_all` once instead of repeating `bounded_sql_query` per connection."
.to_string(),
);
Some(lines.join("\n"))
}
}
#[cfg(test)]
#[path = "registry_tests.rs"]
mod tests;