storkit/server/src/http/agents.rs

use crate::config::ProjectConfig;
use crate::http::context::{AppContext, OpenApiResult, bad_request};
use crate::worktree;
use poem_openapi::{Object, OpenApi, Tags, param::Path, payload::Json};
use serde::Serialize;
use std::sync::Arc;

#[derive(Tags)]
enum AgentsTags {
    Agents,
}

#[derive(Object)]
struct StartAgentPayload {
    story_id: String,
    agent_name: Option<String>,
}

#[derive(Object)]
struct StopAgentPayload {
    story_id: String,
    agent_name: String,
}

#[derive(Object, Serialize)]
struct AgentInfoResponse {
    story_id: String,
    agent_name: String,
    status: String,
    session_id: Option<String>,
    worktree_path: Option<String>,
}

#[derive(Object, Serialize)]
struct AgentConfigInfoResponse {
    name: String,
    role: String,
    model: Option<String>,
    allowed_tools: Option<Vec<String>>,
    max_turns: Option<u32>,
    max_budget_usd: Option<f64>,
}

#[derive(Object)]
struct CreateWorktreePayload {
    story_id: String,
}

#[derive(Object, Serialize)]
struct WorktreeInfoResponse {
    story_id: String,
    worktree_path: String,
    branch: String,
    base_branch: String,
}

#[derive(Object, Serialize)]
struct WorktreeListEntry {
    story_id: String,
    path: String,
}

pub struct AgentsApi {
    pub ctx: Arc<AppContext>,
}

#[OpenApi(tag = "AgentsTags::Agents")]
impl AgentsApi {
    /// Start an agent for a given story (creates worktree, runs setup, spawns agent).
    /// If agent_name is omitted, the first configured agent is used.
    #[oai(path = "/agents/start", method = "post")]
    async fn start_agent(
        &self,
        payload: Json<StartAgentPayload>,
    ) -> OpenApiResult<Json<AgentInfoResponse>> {
        let project_root = self
            .ctx
            .agents
            .get_project_root(&self.ctx.state)
            .map_err(bad_request)?;

        let info = self
            .ctx
            .agents
            .start_agent(
                &project_root,
                &payload.0.story_id,
                payload.0.agent_name.as_deref(),
            )
            .await
            .map_err(bad_request)?;

        Ok(Json(AgentInfoResponse {
            story_id: info.story_id,
            agent_name: info.agent_name,
            status: info.status.to_string(),
            session_id: info.session_id,
            worktree_path: info.worktree_path,
        }))
    }

    /// Stop a running agent and clean up its worktree.
    #[oai(path = "/agents/stop", method = "post")]
    async fn stop_agent(&self, payload: Json<StopAgentPayload>) -> OpenApiResult<Json<bool>> {
        let project_root = self
            .ctx
            .agents
            .get_project_root(&self.ctx.state)
            .map_err(bad_request)?;

        self.ctx
            .agents
            .stop_agent(
                &project_root,
                &payload.0.story_id,
                &payload.0.agent_name,
            )
            .await
            .map_err(bad_request)?;

        Ok(Json(true))
    }

    /// List all agents with their status.
    #[oai(path = "/agents", method = "get")]
    async fn list_agents(&self) -> OpenApiResult<Json<Vec<AgentInfoResponse>>> {
        let agents = self.ctx.agents.list_agents().map_err(bad_request)?;

        Ok(Json(
            agents
                .into_iter()
                .map(|info| AgentInfoResponse {
                    story_id: info.story_id,
                    agent_name: info.agent_name,
                    status: info.status.to_string(),
                    session_id: info.session_id,
                    worktree_path: info.worktree_path,
                })
                .collect(),
        ))
    }

    /// Get the configured agent roster from project.toml.
    #[oai(path = "/agents/config", method = "get")]
    async fn get_agent_config(
        &self,
    ) -> OpenApiResult<Json<Vec<AgentConfigInfoResponse>>> {
        let project_root = self
            .ctx
            .agents
            .get_project_root(&self.ctx.state)
            .map_err(bad_request)?;

        let config = ProjectConfig::load(&project_root).map_err(bad_request)?;

        Ok(Json(
            config
                .agent
                .iter()
                .map(|a| AgentConfigInfoResponse {
                    name: a.name.clone(),
                    role: a.role.clone(),
                    model: a.model.clone(),
                    allowed_tools: a.allowed_tools.clone(),
                    max_turns: a.max_turns,
                    max_budget_usd: a.max_budget_usd,
                })
                .collect(),
        ))
    }

    /// Reload project config and return the updated agent roster.
    #[oai(path = "/agents/config/reload", method = "post")]
    async fn reload_config(
        &self,
    ) -> OpenApiResult<Json<Vec<AgentConfigInfoResponse>>> {
        let project_root = self
            .ctx
            .agents
            .get_project_root(&self.ctx.state)
            .map_err(bad_request)?;

        let config = ProjectConfig::load(&project_root).map_err(bad_request)?;

        Ok(Json(
            config
                .agent
                .iter()
                .map(|a| AgentConfigInfoResponse {
                    name: a.name.clone(),
                    role: a.role.clone(),
                    model: a.model.clone(),
                    allowed_tools: a.allowed_tools.clone(),
                    max_turns: a.max_turns,
                    max_budget_usd: a.max_budget_usd,
                })
                .collect(),
        ))
    }

    /// Create a git worktree for a story under .story_kit/worktrees/{story_id}.
    #[oai(path = "/agents/worktrees", method = "post")]
    async fn create_worktree(
        &self,
        payload: Json<CreateWorktreePayload>,
    ) -> OpenApiResult<Json<WorktreeInfoResponse>> {
        let project_root = self
            .ctx
            .agents
            .get_project_root(&self.ctx.state)
            .map_err(bad_request)?;

        let info = self
            .ctx
            .agents
            .create_worktree(&project_root, &payload.0.story_id)
            .await
            .map_err(bad_request)?;

        Ok(Json(WorktreeInfoResponse {
            story_id: payload.0.story_id,
            worktree_path: info.path.to_string_lossy().to_string(),
            branch: info.branch,
            base_branch: info.base_branch,
        }))
    }

    /// List all worktrees under .story_kit/worktrees/.
    #[oai(path = "/agents/worktrees", method = "get")]
    async fn list_worktrees(&self) -> OpenApiResult<Json<Vec<WorktreeListEntry>>> {
        let project_root = self
            .ctx
            .agents
            .get_project_root(&self.ctx.state)
            .map_err(bad_request)?;

        let entries = worktree::list_worktrees(&project_root).map_err(bad_request)?;

        Ok(Json(
            entries
                .into_iter()
                .map(|e| WorktreeListEntry {
                    story_id: e.story_id,
                    path: e.path.to_string_lossy().to_string(),
                })
                .collect(),
        ))
    }

    /// Remove a git worktree and its feature branch for a story.
    #[oai(path = "/agents/worktrees/:story_id", method = "delete")]
    async fn remove_worktree(&self, story_id: Path<String>) -> OpenApiResult<Json<bool>> {
        let project_root = self
            .ctx
            .agents
            .get_project_root(&self.ctx.state)
            .map_err(bad_request)?;

        let config = ProjectConfig::load(&project_root).map_err(bad_request)?;
        worktree::remove_worktree_by_story_id(&project_root, &story_id.0, &config)
            .await
            .map_err(bad_request)?;

        Ok(Json(true))
    }
}