tori/src/sink.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU16, Ordering};
use serde::{Deserialize, Serialize};
use sqlx::sqlite::SqlitePool;
use tokio::sync::{RwLock, broadcast, mpsc};

use crate::agent::{PlanStepInfo, WsMessage, ServiceInfo};
use crate::state::{AgentState, Artifact};

/// All updates produced by the agent loop. This is the single output interface
/// that decouples the agent logic from DB persistence and WebSocket broadcasting.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind")]
pub enum AgentUpdate {
    PlanUpdate {
        workflow_id: String,
        steps: Vec<PlanStepInfo>,
    },
    WorkflowStatus {
        workflow_id: String,
        status: String,
    },
    Activity {
        workflow_id: String,
        activity: String,
    },
    ExecutionLog {
        workflow_id: String,
        step_order: i32,
        tool_name: String,
        tool_input: String,
        output: String,
        status: String,
    },
    LlmCallLog {
        workflow_id: String,
        step_order: i32,
        phase: String,
        messages_count: i32,
        tools_count: i32,
        tool_calls: String,
        text_response: String,
        prompt_tokens: Option<u32>,
        completion_tokens: Option<u32>,
        latency_ms: i64,
    },
    StateSnapshot {
        workflow_id: String,
        step_order: i32,
        state: AgentState,
    },
    WorkflowComplete {
        workflow_id: String,
        status: String,
        report: Option<String>,
    },
    ArtifactSave {
        workflow_id: String,
        step_order: i32,
        artifact: Artifact,
    },
    RequirementUpdate {
        workflow_id: String,
        requirement: String,
    },
    Error {
        message: String,
    },
}

/// Manages local services (start_service / stop_service tools).
/// Created per-worker or per-agent-loop.
pub struct ServiceManager {
    pub services: RwLock<HashMap<String, ServiceInfo>>,
    next_port: AtomicU16,
}

impl ServiceManager {
    pub fn new(start_port: u16) -> Arc<Self> {
        Arc::new(Self {
            services: RwLock::new(HashMap::new()),
            next_port: AtomicU16::new(start_port),
        })
    }

    pub fn allocate_port(&self) -> u16 {
        self.next_port.fetch_add(1, Ordering::Relaxed)
    }
}

/// Server-side handler: consumes AgentUpdate from channel, persists to DB and broadcasts to frontend.
pub async fn handle_agent_updates(
    mut rx: mpsc::Receiver<AgentUpdate>,
    pool: SqlitePool,
    broadcast_tx: broadcast::Sender<WsMessage>,
) {
    while let Some(update) = rx.recv().await {
        match update {
            AgentUpdate::PlanUpdate { workflow_id, steps } => {
                let _ = broadcast_tx.send(WsMessage::PlanUpdate { workflow_id, steps });
            }
            AgentUpdate::WorkflowStatus { ref workflow_id, ref status } => {
                let _ = sqlx::query("UPDATE workflows SET status = ? WHERE id = ?")
                    .bind(status)
                    .bind(workflow_id)
                    .execute(&pool)
                    .await;
                let _ = broadcast_tx.send(WsMessage::WorkflowStatusUpdate {
                    workflow_id: workflow_id.clone(),
                    status: status.clone(),
                });
            }
            AgentUpdate::Activity { workflow_id, activity } => {
                let _ = broadcast_tx.send(WsMessage::ActivityUpdate { workflow_id, activity });
            }
            AgentUpdate::ExecutionLog { ref workflow_id, step_order, ref tool_name, ref tool_input, ref output, ref status } => {
                let id = uuid::Uuid::new_v4().to_string();
                let _ = sqlx::query(
                    "INSERT INTO execution_log (id, workflow_id, step_order, tool_name, tool_input, output, status, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'))"
                )
                .bind(&id)
                .bind(workflow_id)
                .bind(step_order)
                .bind(tool_name)
                .bind(tool_input)
                .bind(output)
                .bind(status)
                .execute(&pool)
                .await;
                let _ = broadcast_tx.send(WsMessage::StepStatusUpdate {
                    step_id: id,
                    status: status.clone(),
                    output: output.clone(),
                });
            }
            AgentUpdate::LlmCallLog { ref workflow_id, step_order, ref phase, messages_count, tools_count, ref tool_calls, ref text_response, prompt_tokens, completion_tokens, latency_ms } => {
                let id = uuid::Uuid::new_v4().to_string();
                let _ = sqlx::query(
                    "INSERT INTO llm_call_log (id, workflow_id, step_order, phase, messages_count, tools_count, tool_calls, text_response, prompt_tokens, completion_tokens, latency_ms, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, datetime('now'))"
                )
                .bind(&id)
                .bind(workflow_id)
                .bind(step_order)
                .bind(phase)
                .bind(messages_count)
                .bind(tools_count)
                .bind(tool_calls)
                .bind(text_response)
                .bind(prompt_tokens.map(|v| v as i32))
                .bind(completion_tokens.map(|v| v as i32))
                .bind(latency_ms as i32)
                .execute(&pool)
                .await;
                let entry = crate::db::LlmCallLogEntry {
                    id,
                    workflow_id: workflow_id.clone(),
                    step_order,
                    phase: phase.clone(),
                    messages_count,
                    tools_count,
                    tool_calls: tool_calls.clone(),
                    text_response: text_response.clone(),
                    prompt_tokens: prompt_tokens.map(|v| v as i32),
                    completion_tokens: completion_tokens.map(|v| v as i32),
                    latency_ms: latency_ms as i32,
                    created_at: String::new(),
                };
                let _ = broadcast_tx.send(WsMessage::LlmCallLog {
                    workflow_id: workflow_id.clone(),
                    entry,
                });
            }
            AgentUpdate::StateSnapshot { ref workflow_id, step_order, ref state } => {
                let id = uuid::Uuid::new_v4().to_string();
                let json = serde_json::to_string(state).unwrap_or_default();
                let _ = sqlx::query(
                    "INSERT INTO agent_state_snapshots (id, workflow_id, step_order, state_json, created_at) VALUES (?, ?, ?, ?, datetime('now'))"
                )
                .bind(&id)
                .bind(workflow_id)
                .bind(step_order)
                .bind(&json)
                .execute(&pool)
                .await;
            }
            AgentUpdate::WorkflowComplete { ref workflow_id, ref status, ref report } => {
                let _ = sqlx::query("UPDATE workflows SET status = ? WHERE id = ?")
                    .bind(status)
                    .bind(workflow_id)
                    .execute(&pool)
                    .await;
                if let Some(ref r) = report {
                    let _ = sqlx::query("UPDATE workflows SET report = ? WHERE id = ?")
                        .bind(r)
                        .bind(workflow_id)
                        .execute(&pool)
                        .await;
                    let _ = broadcast_tx.send(WsMessage::ReportReady {
                        workflow_id: workflow_id.clone(),
                    });
                }
                let _ = broadcast_tx.send(WsMessage::WorkflowStatusUpdate {
                    workflow_id: workflow_id.clone(),
                    status: status.clone(),
                });
            }
            AgentUpdate::ArtifactSave { ref workflow_id, step_order, ref artifact } => {
                let id = uuid::Uuid::new_v4().to_string();
                let _ = sqlx::query(
                    "INSERT INTO step_artifacts (id, workflow_id, step_order, name, path, artifact_type, description) VALUES (?, ?, ?, ?, ?, ?, ?)"
                )
                .bind(&id)
                .bind(workflow_id)
                .bind(step_order)
                .bind(&artifact.name)
                .bind(&artifact.path)
                .bind(&artifact.artifact_type)
                .bind(&artifact.description)
                .execute(&pool)
                .await;
            }
            AgentUpdate::RequirementUpdate { ref workflow_id, ref requirement } => {
                let _ = sqlx::query("UPDATE workflows SET requirement = ? WHERE id = ?")
                    .bind(requirement)
                    .bind(workflow_id)
                    .execute(&pool)
                    .await;
                let _ = broadcast_tx.send(WsMessage::RequirementUpdate {
                    workflow_id: workflow_id.clone(),
                    requirement: requirement.clone(),
                });
            }
            AgentUpdate::Error { message } => {
                let _ = broadcast_tx.send(WsMessage::Error { message });
            }
        }
    }
}