storkit/server/src/http/agents.rs

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

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

#[derive(Object)]
struct CreateAgentPayload {
    name: String,
    role: String,
    cwd: String,
}

#[derive(Object)]
struct SendMessagePayload {
    message: String,
}

#[derive(Object, Serialize)]
struct AgentInfoResponse {
    name: String,
    role: String,
    cwd: String,
    session_id: Option<String>,
    status: String,
    message_count: usize,
}

#[derive(Object, Serialize)]
struct AgentMessageResponse {
    agent: String,
    text: String,
    session_id: Option<String>,
    model: Option<String>,
    api_key_source: Option<String>,
    rate_limit_type: Option<String>,
    cost_usd: Option<f64>,
    input_tokens: Option<u64>,
    output_tokens: Option<u64>,
    duration_ms: Option<u64>,
}

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

#[OpenApi(tag = "AgentsTags::Agents")]
impl AgentsApi {
    /// Create a new agent with a name, role, and working directory.
    #[oai(path = "/agents", method = "post")]
    async fn create_agent(
        &self,
        payload: Json<CreateAgentPayload>,
    ) -> OpenApiResult<Json<AgentInfoResponse>> {
        let req = crate::agents::CreateAgentRequest {
            name: payload.0.name,
            role: payload.0.role,
            cwd: payload.0.cwd,
        };

        let info = self.ctx.agents.create_agent(req).map_err(bad_request)?;

        Ok(Json(AgentInfoResponse {
            name: info.name,
            role: info.role,
            cwd: info.cwd,
            session_id: info.session_id,
            status: "idle".to_string(),
            message_count: info.message_count,
        }))
    }

    /// List all registered agents.
    #[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 {
                    name: info.name,
                    role: info.role,
                    cwd: info.cwd,
                    session_id: info.session_id,
                    status: match info.status {
                        crate::agents::AgentStatus::Idle => "idle".to_string(),
                        crate::agents::AgentStatus::Running => "running".to_string(),
                    },
                    message_count: info.message_count,
                })
                .collect(),
        ))
    }

    /// Send a message to an agent and wait for its response.
    #[oai(path = "/agents/:name/message", method = "post")]
    async fn send_message(
        &self,
        name: poem_openapi::param::Path<String>,
        payload: Json<SendMessagePayload>,
    ) -> OpenApiResult<Json<AgentMessageResponse>> {
        let result = self
            .ctx
            .agents
            .send_message(&name.0, &payload.0.message)
            .await
            .map_err(bad_request)?;

        Ok(Json(AgentMessageResponse {
            agent: result.agent,
            text: result.text,
            session_id: result.session_id,
            model: result.model,
            api_key_source: result.api_key_source,
            rate_limit_type: result.rate_limit_type,
            cost_usd: result.cost_usd,
            input_tokens: result.input_tokens,
            output_tokens: result.output_tokens,
            duration_ms: result.duration_ms,
        }))
    }
}