Tori: AI agent workflow manager - initial implementation
Rust (Axum) + Vue 3 + SQLite. Features: - Project CRUD REST API with proper error handling - Per-project agent loop (mpsc + broadcast channels) - LLM-driven plan generation and replan on user feedback - SSH command execution with status streaming - WebSocket real-time updates to frontend - Four-zone UI: requirement, plan (left), execution (right), comment Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
520
src/agent.rs
Normal file
520
src/agent.rs
Normal file
@@ -0,0 +1,520 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::sqlite::SqlitePool;
|
||||
use tokio::sync::{mpsc, RwLock, broadcast};
|
||||
|
||||
use crate::llm::{LlmClient, ChatMessage};
|
||||
use crate::ssh::SshExecutor;
|
||||
use crate::{LlmConfig, SshConfig};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum AgentEvent {
|
||||
NewRequirement { workflow_id: String, requirement: String },
|
||||
Comment { workflow_id: String, content: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum WsMessage {
|
||||
PlanUpdate { workflow_id: String, steps: Vec<PlanStepInfo> },
|
||||
StepStatusUpdate { step_id: String, status: String, output: String },
|
||||
WorkflowStatusUpdate { workflow_id: String, status: String },
|
||||
Error { message: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PlanStepInfo {
|
||||
pub order: i32,
|
||||
pub description: String,
|
||||
pub command: String,
|
||||
}
|
||||
|
||||
pub struct AgentManager {
|
||||
agents: RwLock<HashMap<String, mpsc::Sender<AgentEvent>>>,
|
||||
broadcast: RwLock<HashMap<String, broadcast::Sender<WsMessage>>>,
|
||||
pool: SqlitePool,
|
||||
llm_config: LlmConfig,
|
||||
ssh_config: SshConfig,
|
||||
}
|
||||
|
||||
impl AgentManager {
|
||||
pub fn new(pool: SqlitePool, llm_config: LlmConfig, ssh_config: SshConfig) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
agents: RwLock::new(HashMap::new()),
|
||||
broadcast: RwLock::new(HashMap::new()),
|
||||
pool,
|
||||
llm_config,
|
||||
ssh_config,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn get_broadcast(&self, project_id: &str) -> broadcast::Receiver<WsMessage> {
|
||||
let mut map = self.broadcast.write().await;
|
||||
let tx = map.entry(project_id.to_string())
|
||||
.or_insert_with(|| broadcast::channel(64).0);
|
||||
tx.subscribe()
|
||||
}
|
||||
|
||||
pub async fn send_event(self: &Arc<Self>, project_id: &str, event: AgentEvent) {
|
||||
let agents = self.agents.read().await;
|
||||
if let Some(tx) = agents.get(project_id) {
|
||||
let _ = tx.send(event).await;
|
||||
} else {
|
||||
drop(agents);
|
||||
self.spawn_agent(project_id.to_string()).await;
|
||||
let agents = self.agents.read().await;
|
||||
if let Some(tx) = agents.get(project_id) {
|
||||
let _ = tx.send(event).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_agent(self: &Arc<Self>, project_id: String) {
|
||||
let (tx, rx) = mpsc::channel(32);
|
||||
self.agents.write().await.insert(project_id.clone(), tx);
|
||||
|
||||
let broadcast_tx = {
|
||||
let mut map = self.broadcast.write().await;
|
||||
map.entry(project_id.clone())
|
||||
.or_insert_with(|| broadcast::channel(64).0)
|
||||
.clone()
|
||||
};
|
||||
|
||||
let pool = self.pool.clone();
|
||||
let llm_config = self.llm_config.clone();
|
||||
let ssh_config = self.ssh_config.clone();
|
||||
tokio::spawn(agent_loop(project_id, rx, broadcast_tx, pool, llm_config, ssh_config));
|
||||
}
|
||||
}
|
||||
|
||||
async fn agent_loop(
|
||||
project_id: String,
|
||||
mut rx: mpsc::Receiver<AgentEvent>,
|
||||
broadcast_tx: broadcast::Sender<WsMessage>,
|
||||
pool: SqlitePool,
|
||||
llm_config: LlmConfig,
|
||||
ssh_config: SshConfig,
|
||||
) {
|
||||
let llm = LlmClient::new(&llm_config);
|
||||
let ssh = SshExecutor::new(&ssh_config);
|
||||
|
||||
tracing::info!("Agent loop started for project {}", project_id);
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
AgentEvent::NewRequirement { workflow_id, requirement } => {
|
||||
let _ = broadcast_tx.send(WsMessage::WorkflowStatusUpdate {
|
||||
workflow_id: workflow_id.clone(),
|
||||
status: "planning".into(),
|
||||
});
|
||||
|
||||
let plan_result = generate_plan(&llm, &requirement).await;
|
||||
|
||||
match plan_result {
|
||||
Ok(steps) => {
|
||||
// Save steps to DB
|
||||
for (i, step) in steps.iter().enumerate() {
|
||||
let step_id = uuid::Uuid::new_v4().to_string();
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO plan_steps (id, workflow_id, step_order, description, status) VALUES (?, ?, ?, ?, 'pending')"
|
||||
)
|
||||
.bind(&step_id)
|
||||
.bind(&workflow_id)
|
||||
.bind(i as i32 + 1)
|
||||
.bind(&step.description)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
}
|
||||
|
||||
let _ = broadcast_tx.send(WsMessage::PlanUpdate {
|
||||
workflow_id: workflow_id.clone(),
|
||||
steps: steps.clone(),
|
||||
});
|
||||
|
||||
let _ = sqlx::query("UPDATE workflows SET status = 'executing' WHERE id = ?")
|
||||
.bind(&workflow_id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
let _ = broadcast_tx.send(WsMessage::WorkflowStatusUpdate {
|
||||
workflow_id: workflow_id.clone(),
|
||||
status: "executing".into(),
|
||||
});
|
||||
|
||||
// Execute each step
|
||||
let db_steps = sqlx::query_as::<_, crate::db::PlanStep>(
|
||||
"SELECT * FROM plan_steps WHERE workflow_id = ? ORDER BY step_order"
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut all_ok = true;
|
||||
for (i, db_step) in db_steps.iter().enumerate() {
|
||||
let _ = sqlx::query("UPDATE plan_steps SET status = 'running' WHERE id = ?")
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: "running".into(),
|
||||
output: String::new(),
|
||||
});
|
||||
|
||||
let cmd = &steps[i].command;
|
||||
if cmd.is_empty() {
|
||||
let _ = sqlx::query("UPDATE plan_steps SET status = 'done', output = 'Skipped (no command)' WHERE id = ?")
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: "done".into(),
|
||||
output: "Skipped (no command)".into(),
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
match ssh.execute(cmd).await {
|
||||
Ok(result) => {
|
||||
let output = if result.stderr.is_empty() {
|
||||
result.stdout.clone()
|
||||
} else {
|
||||
format!("{}\nSTDERR: {}", result.stdout, result.stderr)
|
||||
};
|
||||
let status = if result.exit_code == 0 { "done" } else { "failed" };
|
||||
|
||||
let _ = sqlx::query("UPDATE plan_steps SET status = ?, output = ? WHERE id = ?")
|
||||
.bind(status)
|
||||
.bind(&output)
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: status.into(),
|
||||
output,
|
||||
});
|
||||
|
||||
if result.exit_code != 0 {
|
||||
all_ok = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = sqlx::query("UPDATE plan_steps SET status = 'failed', output = ? WHERE id = ?")
|
||||
.bind(e.to_string())
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: "failed".into(),
|
||||
output: e.to_string(),
|
||||
});
|
||||
all_ok = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let final_status = if all_ok { "done" } else { "failed" };
|
||||
let _ = sqlx::query("UPDATE workflows SET status = ? WHERE id = ?")
|
||||
.bind(final_status)
|
||||
.bind(&workflow_id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::WorkflowStatusUpdate {
|
||||
workflow_id: workflow_id.clone(),
|
||||
status: final_status.into(),
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = broadcast_tx.send(WsMessage::Error {
|
||||
message: format!("Plan generation failed: {}", e),
|
||||
});
|
||||
let _ = sqlx::query("UPDATE workflows SET status = 'failed' WHERE id = ?")
|
||||
.bind(&workflow_id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
AgentEvent::Comment { workflow_id, content } => {
|
||||
tracing::info!("Comment on workflow {}: {}", workflow_id, content);
|
||||
|
||||
// Get current workflow and steps for context
|
||||
let wf = sqlx::query_as::<_, crate::db::Workflow>(
|
||||
"SELECT * FROM workflows WHERE id = ?",
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_optional(&pool)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
let Some(wf) = wf else { continue };
|
||||
|
||||
let current_steps = sqlx::query_as::<_, crate::db::PlanStep>(
|
||||
"SELECT * FROM plan_steps WHERE workflow_id = ? ORDER BY step_order",
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
// Ask LLM to replan
|
||||
let replan_result =
|
||||
replan(&llm, &wf.requirement, ¤t_steps, &content).await;
|
||||
|
||||
match replan_result {
|
||||
Ok(new_steps) => {
|
||||
// Clear old pending steps, keep done ones
|
||||
let _ = sqlx::query(
|
||||
"DELETE FROM plan_steps WHERE workflow_id = ? AND status IN ('pending', 'failed')",
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
let done_count = sqlx::query_scalar::<_, i32>(
|
||||
"SELECT COUNT(*) FROM plan_steps WHERE workflow_id = ?",
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
for (i, step) in new_steps.iter().enumerate() {
|
||||
let step_id = uuid::Uuid::new_v4().to_string();
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO plan_steps (id, workflow_id, step_order, description, status) VALUES (?, ?, ?, ?, 'pending')",
|
||||
)
|
||||
.bind(&step_id)
|
||||
.bind(&workflow_id)
|
||||
.bind(done_count + i as i32 + 1)
|
||||
.bind(&step.description)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
}
|
||||
|
||||
let _ = broadcast_tx.send(WsMessage::PlanUpdate {
|
||||
workflow_id: workflow_id.clone(),
|
||||
steps: new_steps.clone(),
|
||||
});
|
||||
|
||||
// Resume execution
|
||||
let _ = sqlx::query(
|
||||
"UPDATE workflows SET status = 'executing' WHERE id = ?",
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
let _ = broadcast_tx.send(WsMessage::WorkflowStatusUpdate {
|
||||
workflow_id: workflow_id.clone(),
|
||||
status: "executing".into(),
|
||||
});
|
||||
|
||||
let db_steps = sqlx::query_as::<_, crate::db::PlanStep>(
|
||||
"SELECT * FROM plan_steps WHERE workflow_id = ? AND status = 'pending' ORDER BY step_order",
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut all_ok = true;
|
||||
for (i, db_step) in db_steps.iter().enumerate() {
|
||||
let _ = sqlx::query(
|
||||
"UPDATE plan_steps SET status = 'running' WHERE id = ?",
|
||||
)
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: "running".into(),
|
||||
output: String::new(),
|
||||
});
|
||||
|
||||
let cmd = if i < new_steps.len() {
|
||||
&new_steps[i].command
|
||||
} else {
|
||||
&String::new()
|
||||
};
|
||||
|
||||
if cmd.is_empty() {
|
||||
let _ = sqlx::query(
|
||||
"UPDATE plan_steps SET status = 'done', output = 'Skipped (no command)' WHERE id = ?",
|
||||
)
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: "done".into(),
|
||||
output: "Skipped (no command)".into(),
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
match ssh.execute(cmd).await {
|
||||
Ok(result) => {
|
||||
let output = if result.stderr.is_empty() {
|
||||
result.stdout.clone()
|
||||
} else {
|
||||
format!("{}\nSTDERR: {}", result.stdout, result.stderr)
|
||||
};
|
||||
let status = if result.exit_code == 0 {
|
||||
"done"
|
||||
} else {
|
||||
"failed"
|
||||
};
|
||||
|
||||
let _ = sqlx::query(
|
||||
"UPDATE plan_steps SET status = ?, output = ? WHERE id = ?",
|
||||
)
|
||||
.bind(status)
|
||||
.bind(&output)
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: status.into(),
|
||||
output,
|
||||
});
|
||||
|
||||
if result.exit_code != 0 {
|
||||
all_ok = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = sqlx::query(
|
||||
"UPDATE plan_steps SET status = 'failed', output = ? WHERE id = ?",
|
||||
)
|
||||
.bind(e.to_string())
|
||||
.bind(&db_step.id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
|
||||
step_id: db_step.id.clone(),
|
||||
status: "failed".into(),
|
||||
output: e.to_string(),
|
||||
});
|
||||
all_ok = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let final_status = if all_ok { "done" } else { "failed" };
|
||||
let _ = sqlx::query(
|
||||
"UPDATE workflows SET status = ? WHERE id = ?",
|
||||
)
|
||||
.bind(final_status)
|
||||
.bind(&workflow_id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let _ = broadcast_tx.send(WsMessage::WorkflowStatusUpdate {
|
||||
workflow_id: workflow_id.clone(),
|
||||
status: final_status.into(),
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = broadcast_tx.send(WsMessage::Error {
|
||||
message: format!("Replan failed: {}", e),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!("Agent loop ended for project {}", project_id);
|
||||
}
|
||||
|
||||
async fn generate_plan(llm: &LlmClient, requirement: &str) -> anyhow::Result<Vec<PlanStepInfo>> {
|
||||
let system_prompt = r#"You are an AI workflow planner. Given a requirement, generate a list of executable steps.
|
||||
|
||||
Respond in JSON format only, as an array of objects:
|
||||
[
|
||||
{"order": 1, "description": "what this step does", "command": "shell command to execute via SSH"},
|
||||
...
|
||||
]
|
||||
|
||||
Keep the plan practical and each command should be a single shell command.
|
||||
If a step doesn't need a shell command (e.g., verification), set command to empty string."#;
|
||||
|
||||
let response = llm.chat(vec![
|
||||
ChatMessage { role: "system".into(), content: system_prompt.into() },
|
||||
ChatMessage { role: "user".into(), content: requirement.into() },
|
||||
]).await?;
|
||||
|
||||
let json_str = extract_json_array(&response);
|
||||
let steps: Vec<PlanStepInfo> = serde_json::from_str(json_str)?;
|
||||
Ok(steps)
|
||||
}
|
||||
|
||||
fn extract_json_array(response: &str) -> &str {
|
||||
if let Some(start) = response.find('[') {
|
||||
if let Some(end) = response.rfind(']') {
|
||||
return &response[start..=end];
|
||||
}
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
async fn replan(
|
||||
llm: &LlmClient,
|
||||
requirement: &str,
|
||||
current_steps: &[crate::db::PlanStep],
|
||||
comment: &str,
|
||||
) -> anyhow::Result<Vec<PlanStepInfo>> {
|
||||
let steps_summary: String = current_steps
|
||||
.iter()
|
||||
.map(|s| format!(" {}. [{}] {}", s.step_order, s.status, s.description))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let system_prompt = r#"You are an AI workflow planner. The user has provided feedback on an existing plan.
|
||||
Based on the original requirement, current step statuses, and user feedback, generate ONLY the new/remaining steps that need to be executed.
|
||||
Do NOT include steps that are already done.
|
||||
|
||||
Respond in JSON format only, as an array of objects:
|
||||
[
|
||||
{"order": 1, "description": "what this step does", "command": "shell command to execute via SSH"},
|
||||
...
|
||||
]
|
||||
|
||||
If no new steps are needed (feedback is just informational), return an empty array: []"#;
|
||||
|
||||
let user_msg = format!(
|
||||
"Original requirement:\n{}\n\nCurrent steps:\n{}\n\nUser feedback:\n{}",
|
||||
requirement, steps_summary, comment
|
||||
);
|
||||
|
||||
let response = llm
|
||||
.chat(vec![
|
||||
ChatMessage {
|
||||
role: "system".into(),
|
||||
content: system_prompt.into(),
|
||||
},
|
||||
ChatMessage {
|
||||
role: "user".into(),
|
||||
content: user_msg,
|
||||
},
|
||||
])
|
||||
.await?;
|
||||
|
||||
let json_str = extract_json_array(&response);
|
||||
let steps: Vec<PlanStepInfo> = serde_json::from_str(json_str)?;
|
||||
Ok(steps)
|
||||
}
|
||||
12
src/api/mod.rs
Normal file
12
src/api/mod.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
mod projects;
|
||||
mod workflows;
|
||||
|
||||
use std::sync::Arc;
|
||||
use axum::Router;
|
||||
use crate::AppState;
|
||||
|
||||
pub fn router(state: Arc<AppState>) -> Router {
|
||||
Router::new()
|
||||
.merge(projects::router(state.clone()))
|
||||
.merge(workflows::router(state))
|
||||
}
|
||||
118
src/api/projects.rs
Normal file
118
src/api/projects.rs
Normal file
@@ -0,0 +1,118 @@
|
||||
use std::sync::Arc;
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
routing::get,
|
||||
Json, Router,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use crate::AppState;
|
||||
use crate::db::Project;
|
||||
|
||||
type ApiResult<T> = Result<Json<T>, Response>;
|
||||
|
||||
fn db_err(e: sqlx::Error) -> Response {
|
||||
tracing::error!("Database error: {}", e);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct CreateProject {
|
||||
pub name: String,
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct UpdateProject {
|
||||
pub name: Option<String>,
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
pub fn router(state: Arc<AppState>) -> Router {
|
||||
Router::new()
|
||||
.route("/projects", get(list_projects).post(create_project))
|
||||
.route("/projects/{id}", get(get_project).put(update_project).delete(delete_project))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
async fn list_projects(
|
||||
State(state): State<Arc<AppState>>,
|
||||
) -> ApiResult<Vec<Project>> {
|
||||
sqlx::query_as::<_, Project>("SELECT * FROM projects ORDER BY updated_at DESC")
|
||||
.fetch_all(&state.db.pool)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(db_err)
|
||||
}
|
||||
|
||||
async fn create_project(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(input): Json<CreateProject>,
|
||||
) -> ApiResult<Project> {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
sqlx::query_as::<_, Project>(
|
||||
"INSERT INTO projects (id, name, description) VALUES (?, ?, ?) RETURNING *"
|
||||
)
|
||||
.bind(&id)
|
||||
.bind(&input.name)
|
||||
.bind(&input.description)
|
||||
.fetch_one(&state.db.pool)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(db_err)
|
||||
}
|
||||
|
||||
async fn get_project(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(id): Path<String>,
|
||||
) -> ApiResult<Option<Project>> {
|
||||
sqlx::query_as::<_, Project>("SELECT * FROM projects WHERE id = ?")
|
||||
.bind(&id)
|
||||
.fetch_optional(&state.db.pool)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(db_err)
|
||||
}
|
||||
|
||||
async fn update_project(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(id): Path<String>,
|
||||
Json(input): Json<UpdateProject>,
|
||||
) -> ApiResult<Option<Project>> {
|
||||
if let Some(name) = &input.name {
|
||||
sqlx::query("UPDATE projects SET name = ?, updated_at = datetime('now') WHERE id = ?")
|
||||
.bind(name)
|
||||
.bind(&id)
|
||||
.execute(&state.db.pool)
|
||||
.await
|
||||
.map_err(db_err)?;
|
||||
}
|
||||
if let Some(desc) = &input.description {
|
||||
sqlx::query("UPDATE projects SET description = ?, updated_at = datetime('now') WHERE id = ?")
|
||||
.bind(desc)
|
||||
.bind(&id)
|
||||
.execute(&state.db.pool)
|
||||
.await
|
||||
.map_err(db_err)?;
|
||||
}
|
||||
sqlx::query_as::<_, Project>("SELECT * FROM projects WHERE id = ?")
|
||||
.bind(&id)
|
||||
.fetch_optional(&state.db.pool)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(db_err)
|
||||
}
|
||||
|
||||
async fn delete_project(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(id): Path<String>,
|
||||
) -> ApiResult<bool> {
|
||||
sqlx::query("DELETE FROM projects WHERE id = ?")
|
||||
.bind(&id)
|
||||
.execute(&state.db.pool)
|
||||
.await
|
||||
.map(|r| Json(r.rows_affected() > 0))
|
||||
.map_err(db_err)
|
||||
}
|
||||
136
src/api/workflows.rs
Normal file
136
src/api/workflows.rs
Normal file
@@ -0,0 +1,136 @@
|
||||
use std::sync::Arc;
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
routing::get,
|
||||
Json, Router,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use crate::AppState;
|
||||
use crate::agent::AgentEvent;
|
||||
use crate::db::{Workflow, PlanStep, Comment};
|
||||
|
||||
type ApiResult<T> = Result<Json<T>, Response>;
|
||||
|
||||
fn db_err(e: sqlx::Error) -> Response {
|
||||
tracing::error!("Database error: {}", e);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct CreateWorkflow {
|
||||
pub requirement: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct CreateComment {
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
pub fn router(state: Arc<AppState>) -> Router {
|
||||
Router::new()
|
||||
.route("/projects/{id}/workflows", get(list_workflows).post(create_workflow))
|
||||
.route("/workflows/{id}/steps", get(list_steps))
|
||||
.route("/workflows/{id}/comments", get(list_comments).post(create_comment))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
async fn list_workflows(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(project_id): Path<String>,
|
||||
) -> ApiResult<Vec<Workflow>> {
|
||||
sqlx::query_as::<_, Workflow>(
|
||||
"SELECT * FROM workflows WHERE project_id = ? ORDER BY created_at DESC"
|
||||
)
|
||||
.bind(&project_id)
|
||||
.fetch_all(&state.db.pool)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(db_err)
|
||||
}
|
||||
|
||||
async fn create_workflow(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(project_id): Path<String>,
|
||||
Json(input): Json<CreateWorkflow>,
|
||||
) -> ApiResult<Workflow> {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
let workflow = sqlx::query_as::<_, Workflow>(
|
||||
"INSERT INTO workflows (id, project_id, requirement) VALUES (?, ?, ?) RETURNING *"
|
||||
)
|
||||
.bind(&id)
|
||||
.bind(&project_id)
|
||||
.bind(&input.requirement)
|
||||
.fetch_one(&state.db.pool)
|
||||
.await
|
||||
.map_err(db_err)?;
|
||||
|
||||
state.agent_mgr.send_event(&project_id, AgentEvent::NewRequirement {
|
||||
workflow_id: workflow.id.clone(),
|
||||
requirement: workflow.requirement.clone(),
|
||||
}).await;
|
||||
|
||||
Ok(Json(workflow))
|
||||
}
|
||||
|
||||
async fn list_steps(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(workflow_id): Path<String>,
|
||||
) -> ApiResult<Vec<PlanStep>> {
|
||||
sqlx::query_as::<_, PlanStep>(
|
||||
"SELECT * FROM plan_steps WHERE workflow_id = ? ORDER BY step_order"
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_all(&state.db.pool)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(db_err)
|
||||
}
|
||||
|
||||
async fn list_comments(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(workflow_id): Path<String>,
|
||||
) -> ApiResult<Vec<Comment>> {
|
||||
sqlx::query_as::<_, Comment>(
|
||||
"SELECT * FROM comments WHERE workflow_id = ? ORDER BY created_at"
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_all(&state.db.pool)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(db_err)
|
||||
}
|
||||
|
||||
async fn create_comment(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(workflow_id): Path<String>,
|
||||
Json(input): Json<CreateComment>,
|
||||
) -> ApiResult<Comment> {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
let comment = sqlx::query_as::<_, Comment>(
|
||||
"INSERT INTO comments (id, workflow_id, content) VALUES (?, ?, ?) RETURNING *"
|
||||
)
|
||||
.bind(&id)
|
||||
.bind(&workflow_id)
|
||||
.bind(&input.content)
|
||||
.fetch_one(&state.db.pool)
|
||||
.await
|
||||
.map_err(db_err)?;
|
||||
|
||||
// Notify agent about the comment
|
||||
if let Ok(Some(wf)) = sqlx::query_as::<_, Workflow>(
|
||||
"SELECT * FROM workflows WHERE id = ?"
|
||||
)
|
||||
.bind(&workflow_id)
|
||||
.fetch_optional(&state.db.pool)
|
||||
.await
|
||||
{
|
||||
state.agent_mgr.send_event(&wf.project_id, AgentEvent::Comment {
|
||||
workflow_id: workflow_id.clone(),
|
||||
content: input.content,
|
||||
}).await;
|
||||
}
|
||||
|
||||
Ok(Json(comment))
|
||||
}
|
||||
106
src/db.rs
Normal file
106
src/db.rs
Normal file
@@ -0,0 +1,106 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Database {
|
||||
pub pool: SqlitePool,
|
||||
}
|
||||
|
||||
impl Database {
|
||||
pub async fn new(path: &str) -> anyhow::Result<Self> {
|
||||
let url = format!("sqlite:{}?mode=rwc", path);
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(5)
|
||||
.connect(&url)
|
||||
.await?;
|
||||
Ok(Self { pool })
|
||||
}
|
||||
|
||||
pub async fn migrate(&self) -> anyhow::Result<()> {
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS projects (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)"
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS workflows (
|
||||
id TEXT PRIMARY KEY,
|
||||
project_id TEXT NOT NULL REFERENCES projects(id),
|
||||
requirement TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)"
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS plan_steps (
|
||||
id TEXT PRIMARY KEY,
|
||||
workflow_id TEXT NOT NULL REFERENCES workflows(id),
|
||||
step_order INTEGER NOT NULL,
|
||||
description TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
output TEXT NOT NULL DEFAULT ''
|
||||
)"
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS comments (
|
||||
id TEXT PRIMARY KEY,
|
||||
workflow_id TEXT NOT NULL REFERENCES workflows(id),
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)"
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||
pub struct Project {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||
pub struct Workflow {
|
||||
pub id: String,
|
||||
pub project_id: String,
|
||||
pub requirement: String,
|
||||
pub status: String,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||
pub struct PlanStep {
|
||||
pub id: String,
|
||||
pub workflow_id: String,
|
||||
pub step_order: i32,
|
||||
pub description: String,
|
||||
pub status: String,
|
||||
pub output: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||
pub struct Comment {
|
||||
pub id: String,
|
||||
pub workflow_id: String,
|
||||
pub content: String,
|
||||
pub created_at: String,
|
||||
}
|
||||
56
src/llm.rs
Normal file
56
src/llm.rs
Normal file
@@ -0,0 +1,56 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use crate::LlmConfig;
|
||||
|
||||
pub struct LlmClient {
|
||||
client: reqwest::Client,
|
||||
config: LlmConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ChatRequest {
|
||||
model: String,
|
||||
messages: Vec<ChatMessage>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatResponse {
|
||||
choices: Vec<Choice>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct Choice {
|
||||
message: ChatMessage,
|
||||
}
|
||||
|
||||
impl LlmClient {
|
||||
pub fn new(config: &LlmConfig) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
config: config.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn chat(&self, messages: Vec<ChatMessage>) -> anyhow::Result<String> {
|
||||
let resp = self.client
|
||||
.post(format!("{}/chat/completions", self.config.base_url))
|
||||
.header("Authorization", format!("Bearer {}", self.config.api_key))
|
||||
.json(&ChatRequest {
|
||||
model: self.config.model.clone(),
|
||||
messages,
|
||||
})
|
||||
.send()
|
||||
.await?
|
||||
.json::<ChatResponse>()
|
||||
.await?;
|
||||
|
||||
Ok(resp.choices.first()
|
||||
.map(|c| c.message.content.clone())
|
||||
.unwrap_or_default())
|
||||
}
|
||||
}
|
||||
90
src/main.rs
Normal file
90
src/main.rs
Normal file
@@ -0,0 +1,90 @@
|
||||
mod api;
|
||||
mod agent;
|
||||
mod db;
|
||||
mod llm;
|
||||
mod ssh;
|
||||
mod ws;
|
||||
|
||||
use std::sync::Arc;
|
||||
use axum::Router;
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::services::ServeDir;
|
||||
|
||||
pub struct AppState {
|
||||
pub db: db::Database,
|
||||
pub config: Config,
|
||||
pub agent_mgr: Arc<agent::AgentManager>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct Config {
|
||||
pub llm: LlmConfig,
|
||||
pub ssh: SshConfig,
|
||||
pub server: ServerConfig,
|
||||
pub database: DatabaseConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct LlmConfig {
|
||||
pub base_url: String,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct SshConfig {
|
||||
pub host: String,
|
||||
pub user: String,
|
||||
pub key_path: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct ServerConfig {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct DatabaseConfig {
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter("tori=debug,tower_http=debug")
|
||||
.init();
|
||||
|
||||
let config_str = std::fs::read_to_string("config.yaml")
|
||||
.expect("Failed to read config.yaml");
|
||||
let config: Config = serde_yaml::from_str(&config_str)
|
||||
.expect("Failed to parse config.yaml");
|
||||
|
||||
let database = db::Database::new(&config.database.path).await?;
|
||||
database.migrate().await?;
|
||||
|
||||
let agent_mgr = agent::AgentManager::new(
|
||||
database.pool.clone(),
|
||||
config.llm.clone(),
|
||||
config.ssh.clone(),
|
||||
);
|
||||
|
||||
let state = Arc::new(AppState {
|
||||
db: database,
|
||||
config: config.clone(),
|
||||
agent_mgr: agent_mgr.clone(),
|
||||
});
|
||||
|
||||
let app = Router::new()
|
||||
.nest("/api", api::router(state))
|
||||
.nest("/ws", ws::router(agent_mgr))
|
||||
.fallback_service(ServeDir::new("web/dist"))
|
||||
.layer(CorsLayer::permissive());
|
||||
|
||||
let addr = format!("{}:{}", config.server.host, config.server.port);
|
||||
tracing::info!("Tori server listening on {}", addr);
|
||||
let listener = tokio::net::TcpListener::bind(&addr).await?;
|
||||
axum::serve(listener, app).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
36
src/ssh.rs
Normal file
36
src/ssh.rs
Normal file
@@ -0,0 +1,36 @@
|
||||
use crate::SshConfig;
|
||||
|
||||
pub struct SshExecutor {
|
||||
config: SshConfig,
|
||||
}
|
||||
|
||||
impl SshExecutor {
|
||||
pub fn new(config: &SshConfig) -> Self {
|
||||
Self {
|
||||
config: config.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(&self, command: &str) -> anyhow::Result<SshResult> {
|
||||
let output = tokio::process::Command::new("ssh")
|
||||
.arg("-i")
|
||||
.arg(&self.config.key_path)
|
||||
.arg("-o").arg("StrictHostKeyChecking=no")
|
||||
.arg(format!("{}@{}", self.config.user, self.config.host))
|
||||
.arg(command)
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
Ok(SshResult {
|
||||
stdout: String::from_utf8_lossy(&output.stdout).to_string(),
|
||||
stderr: String::from_utf8_lossy(&output.stderr).to_string(),
|
||||
exit_code: output.status.code().unwrap_or(-1),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SshResult {
|
||||
pub stdout: String,
|
||||
pub stderr: String,
|
||||
pub exit_code: i32,
|
||||
}
|
||||
60
src/ws.rs
Normal file
60
src/ws.rs
Normal file
@@ -0,0 +1,60 @@
|
||||
use std::sync::Arc;
|
||||
use axum::{
|
||||
extract::{Path, State, WebSocketUpgrade, ws::{Message, WebSocket}},
|
||||
response::Response,
|
||||
routing::get,
|
||||
Router,
|
||||
};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use crate::agent::{AgentEvent, AgentManager};
|
||||
|
||||
pub fn router(agent_mgr: Arc<AgentManager>) -> Router {
|
||||
Router::new()
|
||||
.route("/{project_id}", get(ws_handler))
|
||||
.with_state(agent_mgr)
|
||||
}
|
||||
|
||||
async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(agent_mgr): State<Arc<AgentManager>>,
|
||||
Path(project_id): Path<String>,
|
||||
) -> Response {
|
||||
ws.on_upgrade(move |socket| handle_socket(socket, agent_mgr, project_id))
|
||||
}
|
||||
|
||||
async fn handle_socket(socket: WebSocket, agent_mgr: Arc<AgentManager>, project_id: String) {
|
||||
let (mut sender, mut receiver) = socket.split();
|
||||
|
||||
let mut broadcast_rx = agent_mgr.get_broadcast(&project_id).await;
|
||||
|
||||
let send_task = tokio::spawn(async move {
|
||||
while let Ok(msg) = broadcast_rx.recv().await {
|
||||
if let Ok(json) = serde_json::to_string(&msg) {
|
||||
if sender.send(Message::Text(json.into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let mgr = agent_mgr.clone();
|
||||
let pid = project_id.clone();
|
||||
let recv_task = tokio::spawn(async move {
|
||||
while let Some(Ok(msg)) = receiver.next().await {
|
||||
match msg {
|
||||
Message::Text(text) => {
|
||||
if let Ok(event) = serde_json::from_str::<AgentEvent>(&text) {
|
||||
mgr.send_event(&pid, event).await;
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tokio::select! {
|
||||
_ = send_task => {},
|
||||
_ = recv_task => {},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user