use serde::{Deserialize, Serialize};
use crate::inference::Tool;
#[cfg(not(target_arch = "wasm32"))]
use surrealdb_types::{RecordId, SurrealValue};
#[cfg(not(target_arch = "wasm32"))]
use surrealdb::Surreal;
#[cfg(not(target_arch = "wasm32"))]
use surrealdb::engine::any::{connect, Any};
#[cfg(not(target_arch = "wasm32"))]
use surrealdb::opt::auth::Root;
#[cfg(not(target_arch = "wasm32"))]
use std::future::Future;
#[cfg(not(target_arch = "wasm32"))]
use std::time::{Duration, Instant};
#[cfg(not(target_arch = "wasm32"))]
async fn with_database_timeout<T, E>(phase: &str, future: impl Future<Output = Result<T, E>>) -> Result<T, String>
where
E: std::fmt::Display,
{
eprintln!("Starting SurrealDB {phase}");
let started = Instant::now();
let result = tokio::time::timeout(Duration::from_secs(15), future)
.await
.map_err(|_| format!("SurrealDB {phase} timed out after 15 seconds"))?
.map_err(|error| format!("SurrealDB {phase} failed: {error}"));
eprintln!("Completed SurrealDB {phase} in {} ms", started.elapsed().as_millis());
result
}
#[cfg(not(target_arch = "wasm32"))]
fn record_content<T: Serialize>(value: &T) -> Result<serde_json::Value, String> {
let mut content = serde_json::to_value(value).map_err(|e| e.to_string())?;
if let serde_json::Value::Object(fields) = &mut content {
fields.remove("id");
}
Ok(content)
}
#[derive(Debug, Serialize, Deserialize, Clone, Default)]
#[cfg_attr(not(target_arch = "wasm32"), derive(SurrealValue))]
pub struct AgentTelemetry {
pub prompt_tokens: usize,
pub completion_tokens: usize,
pub total_tokens: usize,
pub turns_taken: usize,
pub projects_completed: usize,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[cfg_attr(not(target_arch = "wasm32"), derive(SurrealValue))]
pub struct ScheduledTask {
pub id: String,
pub description: String,
pub due_date: String,
pub completed: bool,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[cfg_attr(not(target_arch = "wasm32"), derive(SurrealValue))]
pub struct Agent {
pub company_id: Option<String>,
pub id: String,
pub name: String,
pub role: String,
pub parent_id: Option<String>,
pub system_prompt: String,
pub tools: Vec<String>,
pub model: Option<String>,
pub provider: Option<String>,
pub temperature: Option<f32>,
pub telemetry: Option<AgentTelemetry>,
pub scheduled_tasks: Option<Vec<ScheduledTask>>,
pub pending_messages: Option<Vec<String>>,
pub issue_triggers: Option<Vec<String>>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[cfg_attr(not(target_arch = "wasm32"), derive(SurrealValue))]
pub struct Project {
pub company_id: Option<String>,
pub id: String,
pub title: String,
pub description: String,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[cfg_attr(not(target_arch = "wasm32"), derive(SurrealValue))]
pub struct RepositoryConfig {
pub provider: String,
pub repo_url: String,
pub access_token: String,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[cfg_attr(not(target_arch = "wasm32"), derive(SurrealValue))]
pub struct IssueComment {
pub author: String,
pub content: String,
pub created_at: String,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[cfg_attr(not(target_arch = "wasm32"), derive(SurrealValue))]
pub struct Issue {
pub company_id: Option<String>,
pub id: String,
pub title: String,
pub body: String,
pub state: String,
pub labels: Vec<String>,
pub comments: Vec<IssueComment>,
pub assignee: Option<String>,
}
/// Abstract representation of the company state (SurrealDB or in-memory)
pub struct CompanyState {
#[cfg(not(target_arch = "wasm32"))]
pub db: Surreal<Any>,
pub agents: Vec<Agent>,
pub projects: Vec<Project>,
pub tools: Vec<Tool>,
pub issues: Vec<Issue>,
pub batch_jobs: Vec<crate::batch::BatchJobRecord>,
pub repository: Option<RepositoryConfig>,
pub company_name: Option<String>,
pub company_id: Option<String>,
pub settings: Option<serde_json::Value>,
}
impl CompanyState {
#[cfg(not(target_arch = "wasm32"))]
pub async fn init(db_path: &str, token: Option<String>, target_company: Option<String>) -> Result<Self, String> {
let url = if db_path.contains("://") {
db_path.to_string()
} else {
format!("surrealkv://{}", db_path)
};
let db = with_database_timeout("connection", async { connect(&url).await }).await?;
if let Some(ref tok) = token {
with_database_timeout("token authentication", async {
db.authenticate(tok).await
}).await?;
} else if let (Ok(user), Ok(pass)) = (std::env::var("SURREALDB_USER"), std::env::var("SURREALDB_PASS")) {
with_database_timeout("authentication", async {
db.signin(Root {
username: user,
password: pass,
}).await
}).await?;
}
with_database_timeout("namespace selection", async {
db.use_ns("actualised").use_db("core").await
}).await?;
if let Some(ref tid) = target_company {
db.set("target", tid.clone()).await.map_err(|e| e.to_string())?;
}
// Only initialize schema if we are logging in as root (no token)
if token.is_none() {
with_database_timeout("schema initialization", async {
db.query(
"DEFINE TABLE OVERWRITE user SCHEMALESS PERMISSIONS FOR select, update, delete WHERE id = $auth.id;
DEFINE ACCESS IF NOT EXISTS user ON DATABASE TYPE RECORD
SIGNUP ( CREATE user SET email = $email, pass = crypto::argon2::generate($pass) )
SIGNIN ( SELECT * FROM user WHERE email = $email AND crypto::argon2::compare(pass, $pass) )
DURATION FOR TOKEN 30d, FOR SESSION 30d;
DEFINE TABLE OVERWRITE company SCHEMALESS PERMISSIONS FOR select, create, update, delete WHERE owner = $auth.id;
DEFINE FIELD IF NOT EXISTS settings ON company TYPE option<object>;
DEFINE TABLE OVERWRITE agent SCHEMALESS PERMISSIONS FOR select, create, update, delete WHERE $auth.id IN (SELECT VALUE owner FROM company WHERE type::string(id) = type::string($parent.company_id)) OR company_id = null OR company_id = NONE;
REMOVE FIELD IF EXISTS company_id ON agent;
DEFINE TABLE OVERWRITE project SCHEMALESS PERMISSIONS FOR select, create, update, delete WHERE $auth.id IN (SELECT VALUE owner FROM company WHERE type::string(id) = type::string($parent.company_id)) OR company_id = null OR company_id = NONE;
REMOVE FIELD IF EXISTS company_id ON project;
DEFINE TABLE OVERWRITE tool SCHEMALESS PERMISSIONS FOR select, create, update, delete WHERE $auth.id IN (SELECT VALUE owner FROM company WHERE type::string(id) = type::string($parent.company_id)) OR company_id = null OR company_id = NONE;
REMOVE FIELD IF EXISTS company_id ON tool;
DEFINE TABLE OVERWRITE issue SCHEMALESS PERMISSIONS FOR select, create, update, delete WHERE $auth.id IN (SELECT VALUE owner FROM company WHERE type::string(id) = type::string($parent.company_id)) OR company_id = null OR company_id = NONE;
REMOVE FIELD IF EXISTS company_id ON issue;
DEFINE TABLE OVERWRITE batch_job SCHEMALESS PERMISSIONS FOR select, create, update, delete WHERE $auth.id IN (SELECT VALUE owner FROM company WHERE type::string(id) = type::string($parent.company_id)) OR company_id = null OR company_id = NONE;
REMOVE FIELD IF EXISTS company_id ON batch_job;"
).await?.check()
}).await?;
}
let mut response = with_database_timeout("state hydration", async {
if let Some(ref target) = target_company {
let full_target = if target.contains(':') {
target.clone()
} else {
format!("company:{}", target)
};
db.query(
"SELECT record::id(id) AS id, name ?? 'Unnamed Agent' AS name, role ?? 'Agent' AS role, IF type::is_string(parent_id) OR type::is_record(parent_id) THEN string::replace(type::string(parent_id), 'agent:', '') ELSE NONE END AS parent_id, system_prompt ?? '' AS system_prompt, tools ?? [] AS tools, model ?? NONE AS model, provider ?? NONE AS provider, temperature ?? NONE AS temperature, telemetry, scheduled_tasks ?? NONE AS scheduled_tasks, pending_messages ?? NONE AS pending_messages, issue_triggers ?? NONE AS issue_triggers FROM agent WHERE type::string(company_id) = type::string($target) OR company_id = null;
SELECT record::id(id) AS id, title, description FROM project WHERE type::string(company_id) = type::string($target) OR company_id = null;
SELECT name, description, parameters FROM tool WHERE type::string(company_id) = type::string($target) OR company_id = null;
SELECT record::id(id) AS id, title, body, state, labels, comments, assignee FROM issue WHERE type::string(company_id) = type::string($target) OR company_id = null;
SELECT record::id(id) AS id, requesting_agent_id, input_file_id, output_file_id, error_file_id, status, endpoint, total_requests, completed_requests, failed_requests, destination_path, created_at, completed_at FROM batch_job WHERE type::string(company_id) = type::string($target) OR company_id = null;
SELECT record::id(id) AS id, name, settings FROM company WHERE type::string(id) = type::string($target) LIMIT 1"
).bind(("target", full_target)).await
} else {
db.query(
"SELECT record::id(id) AS id, name ?? 'Unnamed Agent' AS name, role ?? 'Agent' AS role, IF type::is_string(parent_id) OR type::is_record(parent_id) THEN string::replace(type::string(parent_id), 'agent:', '') ELSE NONE END AS parent_id, system_prompt ?? '' AS system_prompt, tools ?? [] AS tools, model ?? NONE AS model, provider ?? NONE AS provider, temperature ?? NONE AS temperature, telemetry, scheduled_tasks ?? NONE AS scheduled_tasks, pending_messages ?? NONE AS pending_messages, issue_triggers ?? NONE AS issue_triggers FROM agent;
SELECT record::id(id) AS id, title, description FROM project;
SELECT name, description, parameters FROM tool;
SELECT record::id(id) AS id, title, body, state, labels, comments, assignee FROM issue;
SELECT record::id(id) AS id, requesting_agent_id, input_file_id, output_file_id, error_file_id, status, endpoint, total_requests, completed_requests, failed_requests, destination_path, created_at, completed_at FROM batch_job;
SELECT record::id(id) AS id, name, settings FROM company LIMIT 1"
).await
}
}).await?;
let agents = response.take::<Vec<Agent>>(0).map_err(|e| e.to_string())?;
let projects = response.take::<Vec<Project>>(1).map_err(|e| e.to_string())?;
let tools = response.take::<Vec<Tool>>(2).map_err(|e| e.to_string())?;
let issues = response.take::<Vec<Issue>>(3).map_err(|e| e.to_string())?;
let batch_jobs = response.take::<Vec<crate::batch::BatchJobRecord>>(4).unwrap_or_default();
let company_docs = response.take::<Vec<serde_json::Value>>(5).map_err(|e| e.to_string())?;
let company_id = company_docs.first().and_then(|d| d.get("id").and_then(|v| v.as_str()).map(|s| s.to_string()));
let company_name = company_docs.first().and_then(|d| d.get("name").and_then(|v| v.as_str()).map(|s| s.to_string()));
let settings = company_docs.first().and_then(|d| d.get("settings").cloned());
let repository = settings.as_ref().and_then(|s| {
let token = s.get("githubToken").and_then(|v| v.as_str());
let repo_url = s.get("githubRepoUrl").and_then(|v| v.as_str());
if let (Some(t), Some(u)) = (token, repo_url) {
if !t.is_empty() && !u.is_empty() {
return Some(RepositoryConfig {
provider: "github".to_string(),
repo_url: u.to_string(),
access_token: t.to_string(),
});
}
}
None
});
Ok(Self {
db,
agents,
projects,
tools,
issues,
batch_jobs,
repository,
company_name,
company_id,
settings,
})
}
#[cfg(target_arch = "wasm32")]
pub async fn init(_db_path: &str, _token: Option<String>, _target_company: Option<String>) -> Result<Self, String> {
Ok(Self {
agents: Vec::new(),
projects: Vec::new(),
tools: Vec::new(),
issues: Vec::new(),
batch_jobs: Vec::new(),
repository: None,
company_name: None,
company_id: None,
settings: None,
})
}
pub async fn upsert_issue(&mut self, issue: Issue) -> Result<(), String> {
#[cfg(not(target_arch = "wasm32"))]
{
self.db
.query("UPDATE type::record('issue', $id) MERGE $content")
.bind(("id", issue.id.clone()))
.bind(("content", issue.clone()))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
self.issues.retain(|i| i.id != issue.id);
self.issues.push(issue);
Ok(())
}
pub async fn update_company_settings(&mut self, settings: serde_json::Value) -> Result<(), String> {
let mut merged_settings = self.settings.clone().unwrap_or_else(|| serde_json::json!({}));
if let (Some(existing), Some(updates)) = (merged_settings.as_object_mut(), settings.as_object()) {
existing.extend(updates.clone());
} else {
merged_settings = settings;
}
#[cfg(not(target_arch = "wasm32"))]
{
if let Some(company_id) = &self.company_id {
self.db.query("UPDATE type::record('company', $id) SET settings = $settings")
.bind(("id", company_id.replace("company:", "")))
.bind(("settings", merged_settings.clone()))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
}
self.settings = Some(merged_settings);
Ok(())
}
pub async fn set_company_name(&mut self, name: String) -> Result<(), String> {
#[cfg(not(target_arch = "wasm32"))]
{
let mut res = self.db
.query("LET $c = IF (SELECT VALUE id FROM company LIMIT 1)[0] THEN (UPDATE company SET name = $name RETURN id) ELSE (CREATE company SET name = $name, owner = $auth.id RETURN id) END;
UPDATE agent SET company_id = $c[0].id WHERE company_id = null;
UPDATE project SET company_id = $c[0].id WHERE company_id = null;
UPDATE tool SET company_id = $c[0].id WHERE company_id = null;
UPDATE issue SET company_id = $c[0].id WHERE company_id = null;
RETURN type::string($c[0].id);")
.bind(("name", name.clone()))
.await.map_err(|e| e.to_string())?;
let id_val: Option<String> = res.take(5).map_err(|e| e.to_string())?;
if let Some(id_str) = id_val {
self.company_id = Some(id_str);
}
for a in self.agents.iter_mut() {
a.company_id = self.company_id.clone();
}
for p in self.projects.iter_mut() {
p.company_id = self.company_id.clone();
}
for t in self.tools.iter_mut() {
t.company_id = self.company_id.clone();
}
}
self.company_name = Some(name);
Ok(())
}
pub async fn delete_company(&mut self, id: &str) -> Result<(), String> {
#[cfg(not(target_arch = "wasm32"))]
{
// Cascade delete
let full_id = if id.contains(':') { id.to_string() } else { format!("company:{}", id) };
self.db.query("DELETE agent WHERE company_id = type::record($id); DELETE project WHERE company_id = type::record($id); DELETE tool WHERE company_id = type::record($id); DELETE shared_file WHERE company_id = type::record($id); DELETE company WHERE id = type::record($id);")
.bind(("id", full_id))
.await.map_err(|e| e.to_string())?.check().map_err(|e| e.to_string())?;
}
Ok(())
}
// --- Agents ---
pub async fn add_agent(&mut self, mut agent: Agent) -> Result<(), String> {
let mut updates_needed = Vec::new();
let id = agent.id.replace("agent:", "");
agent.company_id = self.company_id.clone();
if self.agents.is_empty() {
agent.parent_id = None;
} else if agent.parent_id.is_none() {
for a in self.agents.iter_mut() {
if a.parent_id.is_none() {
a.parent_id = Some(id.clone());
updates_needed.push(a.clone());
}
}
}
self.agents.push(agent.clone());
updates_needed.push(agent);
#[cfg(not(target_arch = "wasm32"))]
for ag in updates_needed {
let content = record_content(&ag)?;
let mut _res = self.db
.query("UPSERT $record CONTENT $agent")
.bind(("record", RecordId::new("agent", ag.id.clone())))
.bind(("agent", content))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
Ok(())
}
pub async fn update_agent(&mut self, id: &str, new_agent: Agent) -> Result<(), String> {
let target_id = id.replace("agent:", "");
let mut updates_needed = Vec::new();
let was_root = self.agents.iter().find(|a| a.id.replace("agent:", "") == target_id).map(|a| a.parent_id.is_none()).unwrap_or(false);
if new_agent.parent_id.is_none() {
for a in self.agents.iter_mut() {
if a.id.replace("agent:", "") != target_id && a.parent_id.is_none() {
a.parent_id = Some(target_id.clone());
updates_needed.push(a.clone());
}
}
} else if was_root {
let manager_id = new_agent.parent_id.as_ref().unwrap().clone();
for a in self.agents.iter_mut() {
if a.id.replace("agent:", "") == manager_id.replace("agent:", "") {
a.parent_id = None;
updates_needed.push(a.clone());
}
}
}
if let Some(idx) = self.agents.iter().position(|a| a.id.replace("agent:", "") == target_id) {
self.agents[idx] = new_agent.clone();
updates_needed.push(new_agent);
}
#[cfg(not(target_arch = "wasm32"))]
for ag in updates_needed {
let content = record_content(&ag)?;
let mut _res = self.db
.query("UPSERT $record CONTENT $agent")
.bind(("record", RecordId::new("agent", ag.id.trim_start_matches("agent:"))))
.bind(("agent", content))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
Ok(())
}
pub async fn remove_agent(&mut self, id: &str) -> Result<(), String> {
let target_id = id.replace("agent:", "");
let is_root = self.agents.iter().find(|a| a.id.replace("agent:", "") == target_id).map(|a| a.parent_id.is_none()).unwrap_or(false);
let parent_id = self.agents.iter().find(|a| a.id.replace("agent:", "") == target_id).and_then(|a| a.parent_id.clone());
let mut updates_needed = Vec::new();
if is_root {
let next_root = self.agents.iter().find(|a| a.parent_id.as_deref().map(|s| s.replace("agent:", "")) == Some(target_id.clone())).map(|a| a.id.clone());
if let Some(new_root_id) = next_root {
for a in self.agents.iter_mut() {
if a.id == new_root_id {
a.parent_id = None;
updates_needed.push(a.clone());
} else if a.parent_id.as_deref().map(|s| s.replace("agent:", "")) == Some(target_id.clone()) {
a.parent_id = Some(new_root_id.replace("agent:", ""));
updates_needed.push(a.clone());
}
}
}
} else {
for a in self.agents.iter_mut() {
if a.parent_id.as_deref().map(|s| s.replace("agent:", "")) == Some(target_id.clone()) {
a.parent_id = parent_id.clone();
updates_needed.push(a.clone());
}
}
}
self.agents.retain(|a| a.id.replace("agent:", "") != target_id);
#[cfg(not(target_arch = "wasm32"))]
{
let mut _res = self.db
.query("DELETE $record")
.bind(("record", RecordId::new("agent", target_id.trim_start_matches("agent:"))))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
for ag in updates_needed {
let content = record_content(&ag)?;
let mut _res = self.db
.query("UPSERT $record CONTENT $agent")
.bind(("record", RecordId::new("agent", ag.id.trim_start_matches("agent:"))))
.bind(("agent", content))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
}
Ok(())
}
// --- Projects ---
pub async fn add_project(&mut self, project: Project) -> Result<(), String> {
#[cfg(not(target_arch = "wasm32"))]
{
let content = record_content(&project)?;
let mut _res = self.db
.query("CREATE $record CONTENT $project")
.bind(("record", RecordId::new("project", project.id.clone())))
.bind(("project", content))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
self.projects.push(project);
Ok(())
}
// --- Tools ---
pub async fn add_tool(&mut self, mut tool: Tool) -> Result<(), String> {
tool.company_id = self.company_id.clone();
#[cfg(not(target_arch = "wasm32"))]
self.db
.query("UPSERT $record CONTENT $tool")
.bind(("record", RecordId::new("tool", tool.name.clone())))
.bind(("tool", tool.clone()))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
self.tools.push(tool);
Ok(())
}
pub async fn update_tool(&mut self, name: &str, new_tool: Tool) -> Result<(), String> {
#[cfg(not(target_arch = "wasm32"))]
self.db
.query("DELETE $old_record; UPDATE $record CONTENT $tool")
.bind(("old_record", RecordId::new("tool", name)))
.bind(("record", RecordId::new("tool", new_tool.name.clone())))
.bind(("tool", new_tool.clone()))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
if let Some(idx) = self.tools.iter().position(|t| t.name == name) {
self.tools[idx] = new_tool;
}
Ok(())
}
pub async fn remove_tool(&mut self, name: &str) -> Result<(), String> {
#[cfg(not(target_arch = "wasm32"))]
self.db
.query("DELETE $record")
.bind(("record", RecordId::new("tool", name)))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
self.tools.retain(|t| t.name != name);
Ok(())
}
// --- Batch Jobs ---
pub async fn add_batch_job(&mut self, mut job: crate::batch::BatchJobRecord) -> Result<(), String> {
job.company_id = self.company_id.clone();
#[cfg(not(target_arch = "wasm32"))]
{
let content = record_content(&job)?;
let mut _res = self.db
.query("CREATE $record CONTENT $job")
.bind(("record", RecordId::new("batch_job", job.id.clone())))
.bind(("job", content))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
self.batch_jobs.push(job);
Ok(())
}
pub async fn update_batch_job(&mut self, mut job: crate::batch::BatchJobRecord) -> Result<(), String> {
job.company_id = self.company_id.clone();
#[cfg(not(target_arch = "wasm32"))]
{
let content = record_content(&job)?;
let mut _res = self.db
.query("UPSERT $record CONTENT $job")
.bind(("record", RecordId::new("batch_job", job.id.clone())))
.bind(("job", content))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
}
if let Some(pos) = self.batch_jobs.iter().position(|j| j.id == job.id) {
self.batch_jobs[pos] = job;
} else {
self.batch_jobs.push(job);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_agent(id: &str, parent_id: Option<&str>) -> Agent {
Agent {
id: if id.contains(':') { id.to_string() } else { format!("agent:{}", id) },
name: "Test Agent".to_string(),
role: "Role".to_string(),
parent_id: parent_id.map(|s| s.to_string()),
system_prompt: "Prompt".to_string(),
tools: vec![],
model: None,
provider: None,
temperature: None,
telemetry: None,
scheduled_tasks: None,
pending_messages: None, issue_triggers: None, company_id: None,
}
}
async fn create_state() -> CompanyState {
CompanyState::init("mem://", None, None).await.unwrap()
}
#[tokio::test]
async fn test_first_agent_is_root() {
let mut state = create_state().await;
// Added with a parent, but it's the first agent, so it should be forced to root
let agent1 = create_agent("agent1", Some("some_parent"));
state.add_agent(agent1.clone()).await.unwrap();
assert_eq!(state.agents.len(), 1);
assert_eq!(state.agents[0].parent_id, None);
}
#[tokio::test]
async fn test_add_second_root_swaps_first() {
let mut state = create_state().await;
state.add_agent(create_agent("agent1", None)).await.unwrap();
// Add a second root
state.add_agent(create_agent("agent2", None)).await.unwrap();
assert_eq!(state.agents.len(), 2);
let agent1 = state.agents.iter().find(|a| a.id == "agent:agent1").unwrap();
let agent2 = state.agents.iter().find(|a| a.id == "agent:agent2").unwrap();
// agent2 should be the new root, agent1 should report to agent2
assert_eq!(agent2.parent_id, None);
assert_eq!(agent1.parent_id, Some("agent2".to_string()));
}
#[tokio::test]
async fn test_update_agent_to_root_swaps_current_root() {
let mut state = create_state().await;
state.add_agent(create_agent("root", None)).await.unwrap();
state.add_agent(create_agent("child", Some("root"))).await.unwrap();
// Make child the root
let updated_child = create_agent("child", None);
state.update_agent("child", updated_child).await.unwrap();
let old_root = state.agents.iter().find(|a| a.id == "agent:root").unwrap();
let new_root = state.agents.iter().find(|a| a.id == "agent:child").unwrap();
assert_eq!(new_root.parent_id, None);
assert_eq!(old_root.parent_id, Some("child".to_string()));
}
#[tokio::test]
async fn test_update_root_with_manager_promotes_manager() {
let mut state = create_state().await;
state.add_agent(create_agent("root", None)).await.unwrap();
state.add_agent(create_agent("manager", Some("root"))).await.unwrap();
// Give the root a manager
let updated_root = create_agent("root", Some("manager"));
state.update_agent("root", updated_root).await.unwrap();
let old_root = state.agents.iter().find(|a| a.id == "agent:root").unwrap();
let new_root = state.agents.iter().find(|a| a.id == "agent:manager").unwrap();
assert_eq!(new_root.parent_id, None);
assert_eq!(old_root.parent_id, Some("manager".to_string()));
}
#[tokio::test]
async fn test_remove_root_promotes_child() {
let mut state = create_state().await;
state.add_agent(create_agent("root", None)).await.unwrap();
state.add_agent(create_agent("child1", Some("root"))).await.unwrap();
state.add_agent(create_agent("child2", Some("root"))).await.unwrap();
state.remove_agent("root").await.unwrap();
assert_eq!(state.agents.len(), 2);
let new_root = state.agents.iter().find(|a| a.parent_id.is_none()).unwrap();
let other_child = state.agents.iter().find(|a| a.id != new_root.id).unwrap();
// One of the children must be promoted to root, and the other must report to it
assert_eq!(new_root.parent_id, None);
assert_eq!(other_child.parent_id, Some(new_root.id.replace("agent:", "")));
}
#[tokio::test]
async fn test_remove_intermediate_node_orphans_to_grandparent() {
let mut state = create_state().await;
state.add_agent(create_agent("root", None)).await.unwrap();
state.add_agent(create_agent("middle", Some("root"))).await.unwrap();
state.add_agent(create_agent("leaf", Some("middle"))).await.unwrap();
state.remove_agent("middle").await.unwrap();
assert_eq!(state.agents.len(), 2);
let leaf = state.agents.iter().find(|a| a.id == "agent:leaf").unwrap();
assert_eq!(leaf.parent_id, Some("root".to_string()));
}
#[tokio::test]
async fn test_update_agent_fixes_missing_prefix() {
let mut state = create_state().await;
state.add_agent(create_agent("agent:some_bot", None)).await.unwrap();
// Update with un-prefixed id
let mut updated = create_agent("agent:some_bot", None);
updated.pending_messages = Some(vec!["hello".to_string()]);
state.update_agent("some_bot", updated).await.unwrap();
let agent = state.agents.iter().find(|a| a.id == "agent:some_bot").unwrap();
assert_eq!(agent.pending_messages, Some(vec!["hello".to_string()]));
}
#[tokio::test]
#[cfg(not(target_arch = "wasm32"))]
async fn test_database_hydration_and_queries() {
let state = CompanyState::init("mem://", None, None).await.unwrap();
// Manually insert an agent with missing name/role to trigger NONE coercion bug
let res = state.db.query(
"CREATE agent:test_hydration_agent SET system_prompt = 'Test prompt', tools = []"
).await.unwrap();
res.check().unwrap();
// Query the DB manually using the exact same hydration logic
let mut response = state.db.query("SELECT record::id(id) AS id, name ?? 'Unnamed Agent' AS name, role ?? 'Agent' AS role, IF type::is_string(parent_id) OR type::is_record(parent_id) THEN string::replace(type::string(parent_id), 'agent:', '') ELSE NONE END AS parent_id, system_prompt ?? '' AS system_prompt, tools ?? [] AS tools, telemetry, scheduled_tasks ?? NONE AS scheduled_tasks, pending_messages ?? NONE AS pending_messages, issue_triggers ?? NONE AS issue_triggers FROM agent").await.unwrap();
let agents: Vec<Agent> = response.take(0).unwrap();
assert_eq!(agents.len(), 1);
let hydrated_agent = &agents[0];
assert_eq!(hydrated_agent.id, "test_hydration_agent");
assert_eq!(hydrated_agent.name, "Unnamed Agent"); // Our default kick in
assert_eq!(hydrated_agent.role, "Agent");
}
#[tokio::test]
#[cfg(not(target_arch = "wasm32"))]
async fn test_query_running_company_ids_filters_by_persisted_state() {
let state = CompanyState::init("mem://", None, None).await.unwrap();
state.db.query(
"CREATE company:enabled SET settings = { executionEnabled: true };
CREATE company:paused SET settings = { executionEnabled: false };
CREATE company:unset SET name = 'Unset';"
).await.unwrap().check().unwrap();
let ids = query_running_company_ids(&state.db).await.unwrap();
assert_eq!(ids, vec!["company:enabled".to_string()]);
}
#[tokio::test]
async fn test_settings_updates_preserve_server_managed_fields() {
let mut state = create_state().await;
state.settings = Some(serde_json::json!({
"executionEnabled": true,
"telegramChatId": 123,
"provider": "gemini"
}));
state.update_company_settings(serde_json::json!({ "provider": "openai" })).await.unwrap();
assert_eq!(state.settings, Some(serde_json::json!({
"executionEnabled": true,
"telegramChatId": 123,
"provider": "openai"
})));
}
}
#[cfg(not(target_arch = "wasm32"))]
pub async fn get_companies(db_path: &str, token: &str) -> Result<String, String> {
let url = if db_path.contains("://") { db_path.to_string() } else { format!("surrealkv://{}", db_path) };
let db = connect(&url).await.map_err(|e| e.to_string())?;
db.authenticate(token).await.map_err(|e| e.to_string())?;
db.use_ns("actualised").use_db("core").await.map_err(|e| e.to_string())?;
let mut response = db.query("SELECT record::id(id) AS id, name FROM company").await.map_err(|e| e.to_string())?;
let companies: Vec<serde_json::Value> = response.take(0).map_err(|e| e.to_string())?;
serde_json::to_string(&companies).map_err(|e| e.to_string())
}
#[cfg(not(target_arch = "wasm32"))]
async fn query_running_company_ids(db: &Surreal<Any>) -> Result<Vec<String>, String> {
let mut response = db.query(
"SELECT VALUE type::string(id) FROM company WHERE settings.executionEnabled = true"
).await.map_err(|e| e.to_string())?;
response.take(0).map_err(|e| e.to_string())
}
#[cfg(not(target_arch = "wasm32"))]
pub async fn get_running_company_ids(db_path: &str) -> Result<String, String> {
let url = if db_path.contains("://") { db_path.to_string() } else { format!("surrealkv://{}", db_path) };
let db = connect(&url).await.map_err(|e| e.to_string())?;
if let (Ok(user), Ok(pass)) = (std::env::var("SURREALDB_USER"), std::env::var("SURREALDB_PASS")) {
db.signin(Root { username: user, password: pass }).await.map_err(|e| e.to_string())?;
}
db.use_ns("actualised").use_db("core").await.map_err(|e| e.to_string())?;
let company_ids = query_running_company_ids(&db).await?;
serde_json::to_string(&company_ids).map_err(|e| e.to_string())
}
#[cfg(not(target_arch = "wasm32"))]
pub async fn get_user_and_company_by_telegram_id(db_path: &str, chat_id: i64) -> Result<String, String> {
let url = if db_path.contains("://") { db_path.to_string() } else { format!("surrealkv://{}", db_path) };
let db = connect(&url).await.map_err(|e| e.to_string())?;
if let (Ok(user), Ok(pass)) = (std::env::var("SURREALDB_USER"), std::env::var("SURREALDB_PASS")) {
db.signin(Root {
username: user,
password: pass,
}).await.map_err(|e| e.to_string())?;
} else {
return Err("No root credentials found".to_string());
}
db.use_ns("actualised").use_db("core").await.map_err(|e| e.to_string())?;
let mut response = db.query(
"LET $u = (SELECT id FROM user WHERE telegram_chat_id = $chat_id LIMIT 1);
IF array::len($u) > 0 {
LET $uid = $u[0].id;
LET $c = (SELECT record::id(id) AS id FROM company WHERE owner = $uid LIMIT 1);
IF array::len($c) > 0 {
RETURN { user_id: record::id($uid), company_id: $c[0].id };
} ELSE {
RETURN { user_id: record::id($uid), company_id: null };
}
} ELSE {
RETURN null;
}"
).bind(("chat_id", chat_id)).await.map_err(|e| e.to_string())?;
let result: Option<serde_json::Value> = response.take(0).map_err(|e| e.to_string())?;
serde_json::to_string(&result).map_err(|e| e.to_string())
}
#[cfg(not(target_arch = "wasm32"))]
pub async fn link_telegram_chat(db_path: &str, token: &str, chat_id: i64) -> Result<(), String> {
let url = if db_path.contains("://") { db_path.to_string() } else { format!("surrealkv://{}", db_path) };
let db = connect(&url).await.map_err(|e| e.to_string())?;
db.authenticate(token).await.map_err(|e| e.to_string())?;
db.use_ns("actualised").use_db("core").await.map_err(|e| e.to_string())?;
db.query("UPDATE user SET telegram_chat_id = $chat_id WHERE id = $auth.id")
.bind(("chat_id", chat_id))
.await.map_err(|e| e.to_string())?
.check().map_err(|e| e.to_string())?;
Ok(())
}