add token infrastructure: service_users table, localpart helper

new SQLite table service_users maps OIDC identities (matrix localpart)
to service-specific usernames, handling auth boundary mismatches.
localpart() extracts the username from a matrix user ID.

delete_all_conversations() added for bulk reset after agent recreation.
all delete_* methods now log failures instead of silently discarding.
removed dead user_tokens table (tokens now live in vault).
This commit is contained in:
2026-03-22 14:58:28 +00:00
parent 2333dda904
commit 14022aa7c0
2 changed files with 187 additions and 32 deletions

View File

@@ -16,6 +16,14 @@ pub struct ResponseContext {
pub room_id: String,
}
/// Extract the localpart from a Matrix user ID.
///
/// `@sienna:sunbeam.pt` → `sienna`
pub fn localpart(matrix_user_id: &str) -> &str {
let stripped = matrix_user_id.strip_prefix('@').unwrap_or(matrix_user_id);
stripped.split(':').next().unwrap_or(stripped)
}
/// Derive a portable user ID from a Matrix user ID.
///
/// `@sienna:sunbeam.pt` → `sienna@sunbeam.pt`
@@ -53,4 +61,24 @@ mod tests {
"user@server:8448"
);
}
#[test]
fn test_localpart_standard() {
assert_eq!(localpart("@sienna:sunbeam.pt"), "sienna");
}
#[test]
fn test_localpart_no_at_prefix() {
assert_eq!(localpart("sienna:sunbeam.pt"), "sienna");
}
#[test]
fn test_localpart_no_colon() {
assert_eq!(localpart("@sienna"), "sienna");
}
#[test]
fn test_localpart_complex() {
assert_eq!(localpart("@user.name:matrix.org"), "user.name");
}
}

View File

@@ -71,10 +71,32 @@ impl Store {
name TEXT PRIMARY KEY,
agent_id TEXT NOT NULL,
model TEXT NOT NULL DEFAULT '',
instructions_hash TEXT NOT NULL DEFAULT '',
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);",
);
CREATE TABLE IF NOT EXISTS service_users (
localpart TEXT NOT NULL,
service TEXT NOT NULL,
service_username TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
PRIMARY KEY (localpart, service)
);
",
)?;
// Migration: add instructions_hash column to agents table if missing
let has_hash: bool = conn
.prepare("SELECT instructions_hash FROM agents LIMIT 0")
.is_ok();
if !has_hash {
conn.execute_batch(
"ALTER TABLE agents ADD COLUMN instructions_hash TEXT NOT NULL DEFAULT '';"
)?;
info!("Migrated agents table: added instructions_hash column");
}
info!(path, "Opened Sol state database");
Ok(Self {
conn: Mutex::new(conn),
@@ -124,19 +146,29 @@ impl Store {
/// Update token estimate for a conversation.
pub fn update_tokens(&self, room_id: &str, estimated_tokens: u32) {
let conn = self.conn.lock().unwrap();
let _ = conn.execute(
if let Err(e) = conn.execute(
"UPDATE conversations SET estimated_tokens = ?1 WHERE room_id = ?2",
params![estimated_tokens, room_id],
);
) {
warn!("Failed to update token estimate: {e}");
}
}
/// Remove all conversation mappings (e.g., after agent recreation).
pub fn delete_all_conversations(&self) -> usize {
let conn = self.conn.lock().unwrap();
conn.execute("DELETE FROM conversations", []).unwrap_or(0)
}
/// Remove a conversation mapping (e.g., after compaction).
pub fn delete_conversation(&self, room_id: &str) {
let conn = self.conn.lock().unwrap();
let _ = conn.execute(
if let Err(e) = conn.execute(
"DELETE FROM conversations WHERE room_id = ?1",
params![room_id],
);
) {
warn!("Failed to delete conversation: {e}");
}
}
/// Load all conversation mappings (for startup recovery).
@@ -165,27 +197,28 @@ impl Store {
// Agents
// =========================================================================
/// Get the agent_id for a named agent.
pub fn get_agent(&self, name: &str) -> Option<String> {
/// Get the agent_id and instructions hash for a named agent.
pub fn get_agent(&self, name: &str) -> Option<(String, String)> {
let conn = self.conn.lock().unwrap();
conn.query_row(
"SELECT agent_id FROM agents WHERE name = ?1",
"SELECT agent_id, instructions_hash FROM agents WHERE name = ?1",
params![name],
|row| row.get(0),
|row| Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)),
)
.ok()
}
/// Store or update an agent mapping.
pub fn upsert_agent(&self, name: &str, agent_id: &str, model: &str) {
/// Store or update an agent mapping (with instructions hash for staleness detection).
pub fn upsert_agent(&self, name: &str, agent_id: &str, model: &str, instructions_hash: &str) {
let conn = self.conn.lock().unwrap();
if let Err(e) = conn.execute(
"INSERT INTO agents (name, agent_id, model)
VALUES (?1, ?2, ?3)
"INSERT INTO agents (name, agent_id, model, instructions_hash)
VALUES (?1, ?2, ?3, ?4)
ON CONFLICT(name) DO UPDATE SET
agent_id = excluded.agent_id,
model = excluded.model",
params![name, agent_id, model],
model = excluded.model,
instructions_hash = excluded.instructions_hash",
params![name, agent_id, model, instructions_hash],
) {
warn!("Failed to upsert agent: {e}");
}
@@ -194,7 +227,49 @@ impl Store {
/// Remove an agent mapping.
pub fn delete_agent(&self, name: &str) {
let conn = self.conn.lock().unwrap();
let _ = conn.execute("DELETE FROM agents WHERE name = ?1", params![name]);
if let Err(e) = conn.execute("DELETE FROM agents WHERE name = ?1", params![name]) {
warn!("Failed to delete agent: {e}");
}
}
// =========================================================================
// Service Users (OIDC → service username mapping)
// =========================================================================
/// Get the service-specific username for a Matrix localpart.
pub fn get_service_user(&self, localpart: &str, service: &str) -> Option<String> {
let conn = self.conn.lock().unwrap();
conn.query_row(
"SELECT service_username FROM service_users WHERE localpart = ?1 AND service = ?2",
params![localpart, service],
|row| row.get(0),
)
.ok()
}
/// Store or update a service user mapping.
pub fn upsert_service_user(&self, localpart: &str, service: &str, service_username: &str) {
let conn = self.conn.lock().unwrap();
if let Err(e) = conn.execute(
"INSERT INTO service_users (localpart, service, service_username)
VALUES (?1, ?2, ?3)
ON CONFLICT(localpart, service) DO UPDATE SET
service_username = excluded.service_username",
params![localpart, service, service_username],
) {
warn!("Failed to upsert service user: {e}");
}
}
/// Remove a service user mapping.
pub fn delete_service_user(&self, localpart: &str, service: &str) {
let conn = self.conn.lock().unwrap();
if let Err(e) = conn.execute(
"DELETE FROM service_users WHERE localpart = ?1 AND service = ?2",
params![localpart, service],
) {
warn!("Failed to delete service user: {e}");
}
}
/// Load all agent mappings (for startup recovery).
@@ -255,18 +330,16 @@ mod tests {
fn test_agent_crud() {
let store = Store::open_memory().unwrap();
store.upsert_agent("sol-orchestrator", "ag_123", "mistral-medium-latest");
assert_eq!(
store.get_agent("sol-orchestrator").unwrap(),
"ag_123"
);
store.upsert_agent("sol-orchestrator", "ag_123", "mistral-medium-latest", "hash1");
let (id, hash) = store.get_agent("sol-orchestrator").unwrap();
assert_eq!(id, "ag_123");
assert_eq!(hash, "hash1");
// Update
store.upsert_agent("sol-orchestrator", "ag_456", "mistral-medium-latest");
assert_eq!(
store.get_agent("sol-orchestrator").unwrap(),
"ag_456"
);
store.upsert_agent("sol-orchestrator", "ag_456", "mistral-medium-latest", "hash2");
let (id, hash) = store.get_agent("sol-orchestrator").unwrap();
assert_eq!(id, "ag_456");
assert_eq!(hash, "hash2");
// Delete
store.delete_agent("sol-orchestrator");
@@ -287,8 +360,8 @@ mod tests {
#[test]
fn test_load_all_agents() {
let store = Store::open_memory().unwrap();
store.upsert_agent("orch", "ag_1", "medium");
store.upsert_agent("obs", "ag_2", "medium");
store.upsert_agent("orch", "ag_1", "medium", "h1");
store.upsert_agent("obs", "ag_2", "medium", "h2");
let all = store.load_all_agents();
assert_eq!(all.len(), 2);
@@ -300,4 +373,58 @@ mod tests {
assert!(store.get_conversation("!nope:x").is_none());
assert!(store.get_agent("nope").is_none());
}
#[test]
fn test_service_user_crud() {
let store = Store::open_memory().unwrap();
// Insert
store.upsert_service_user("sienna", "gitea", "siennav");
assert_eq!(
store.get_service_user("sienna", "gitea").unwrap(),
"siennav"
);
// Update
store.upsert_service_user("sienna", "gitea", "sienna");
assert_eq!(
store.get_service_user("sienna", "gitea").unwrap(),
"sienna"
);
// Different service
store.upsert_service_user("sienna", "grafana", "sienna-g");
assert_eq!(
store.get_service_user("sienna", "grafana").unwrap(),
"sienna-g"
);
// Delete
store.delete_service_user("sienna", "gitea");
assert!(store.get_service_user("sienna", "gitea").is_none());
// Other service unaffected
assert!(store.get_service_user("sienna", "grafana").is_some());
}
#[test]
fn test_delete_all_conversations() {
let store = Store::open_memory().unwrap();
store.upsert_conversation("!a:x", "conv_1", 10);
store.upsert_conversation("!b:x", "conv_2", 20);
store.upsert_conversation("!c:x", "conv_3", 30);
let deleted = store.delete_all_conversations();
assert_eq!(deleted, 3);
assert!(store.load_all_conversations().is_empty());
// DB still works after bulk delete
store.upsert_conversation("!d:x", "conv_4", 0);
assert_eq!(store.load_all_conversations().len(), 1);
}
#[test]
fn test_nonexistent_service_user() {
let store = Store::open_memory().unwrap();
assert!(store.get_service_user("nobody", "gitea").is_none());
}
}