huskies: merge 563_story_build_agent_join_mechanism_agents_register_with_the_gateway_via_token

server/src/gateway.rs

@@ -9,15 +9,18 @@
 use poem::EndpointExt;
 use poem::handler;
 use poem::http::StatusCode;
+use poem::web::Path as PoemPath;
 use poem::web::{Data, Json};
 use poem::{Body, Request, Response};
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use serde_json::{Value, json};
 use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::path::Path;
 use std::sync::Arc;
 use tokio::sync::RwLock;
+use uuid::Uuid;
 
 // Re-export active_project type alias for clarity in gateway bot helpers.
 type ActiveProject = Arc<RwLock<String>>;
@@ -48,6 +51,35 @@ impl GatewayConfig {
     }
 }
 
+// ── Agent join types ─────────────────────────────────────────────────
+
+/// A build agent that has registered with this gateway.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct JoinedAgent {
+    /// Unique ID assigned by the gateway on registration.
+    pub id: String,
+    /// Human-readable label provided by the agent (e.g. `build-agent-abc123`).
+    pub label: String,
+    /// The agent's CRDT-sync WebSocket address (e.g. `ws://host:3001/crdt-sync`).
+    pub address: String,
+    /// Unix timestamp when the agent registered.
+    pub registered_at: f64,
+}
+
+/// A one-time join token that has been generated but not yet consumed.
+struct PendingToken {
+    #[allow(dead_code)]
+    created_at: f64,
+}
+
+/// Request body sent by a build agent when registering with the gateway.
+#[derive(Deserialize)]
+struct RegisterAgentRequest {
+    token: String,
+    label: String,
+    address: String,
+}
+
 // ── Gateway state ────────────────────────────────────────────────────
 
 /// Shared gateway state threaded through HTTP handlers.
@@ -59,6 +91,10 @@ pub struct GatewayState {
     pub active_project: Arc<RwLock<String>>,
     /// HTTP client for proxying requests to project containers.
     pub client: Client,
+    /// Build agents that have joined this gateway.
+    pub joined_agents: Arc<RwLock<Vec<JoinedAgent>>>,
+    /// One-time join tokens that have been issued but not yet consumed.
+    pending_tokens: Arc<RwLock<HashMap<String, PendingToken>>>,
 }
 
 impl GatewayState {
@@ -73,6 +109,8 @@ impl GatewayState {
             config,
             active_project: Arc::new(RwLock::new(first)),
             client: Client::new(),
+            joined_agents: Arc::new(RwLock::new(Vec::new())),
+            pending_tokens: Arc::new(RwLock::new(HashMap::new())),
         })
     }
 
@@ -509,6 +547,133 @@ async fn handle_gateway_health(state: &GatewayState) -> JsonRpcResponse {
     )
 }
 
+// ── Agent join handlers ───────────────────────────────────────────────
+
+/// `GET /gateway/mode` — returns `{"mode":"gateway"}` so clients can detect gateway mode.
+#[handler]
+pub async fn gateway_mode_handler() -> Response {
+    let body = json!({ "mode": "gateway" });
+    Response::builder()
+        .status(StatusCode::OK)
+        .header("Content-Type", "application/json")
+        .body(Body::from(serde_json::to_vec(&body).unwrap_or_default()))
+}
+
+/// `POST /gateway/tokens` — generate a one-time join token for a build agent.
+///
+/// Returns `{"token": "<uuid>"}`. The token is valid until consumed by
+/// `POST /gateway/register` or the process restarts.
+#[handler]
+pub async fn gateway_generate_token_handler(state: Data<&Arc<GatewayState>>) -> Response {
+    let token = Uuid::new_v4().to_string();
+    let now = chrono::Utc::now().timestamp() as f64;
+    state
+        .pending_tokens
+        .write()
+        .await
+        .insert(token.clone(), PendingToken { created_at: now });
+    crate::slog!("[gateway] Generated join token {:.8}…", &token);
+    let body = json!({ "token": token });
+    Response::builder()
+        .status(StatusCode::OK)
+        .header("Content-Type", "application/json")
+        .body(Body::from(serde_json::to_vec(&body).unwrap_or_default()))
+}
+
+/// `POST /gateway/register` — build agent presents its join token and registers.
+///
+/// Expects JSON body: `{ "token": "...", "label": "...", "address": "..." }`.
+/// On success returns the `JoinedAgent` record. The token is consumed immediately.
+#[handler]
+pub async fn gateway_register_agent_handler(
+    body: Body,
+    state: Data<&Arc<GatewayState>>,
+) -> Response {
+    let bytes = match body.into_bytes().await {
+        Ok(b) => b,
+        Err(_) => {
+            return Response::builder()
+                .status(StatusCode::BAD_REQUEST)
+                .body(Body::from("could not read request body"));
+        }
+    };
+
+    let req: RegisterAgentRequest = match serde_json::from_slice(&bytes) {
+        Ok(r) => r,
+        Err(_) => {
+            return Response::builder()
+                .status(StatusCode::BAD_REQUEST)
+                .body(Body::from("invalid JSON body"));
+        }
+    };
+
+    // Validate and consume the token.
+    let mut tokens = state.pending_tokens.write().await;
+    if !tokens.contains_key(&req.token) {
+        return Response::builder()
+            .status(StatusCode::UNAUTHORIZED)
+            .body(Body::from("invalid or already-used join token"));
+    }
+    tokens.remove(&req.token);
+    drop(tokens);
+
+    let agent = JoinedAgent {
+        id: Uuid::new_v4().to_string(),
+        label: req.label,
+        address: req.address,
+        registered_at: chrono::Utc::now().timestamp() as f64,
+    };
+
+    crate::slog!(
+        "[gateway] Agent '{}' registered (id={})",
+        agent.label,
+        agent.id
+    );
+
+    state.joined_agents.write().await.push(agent.clone());
+
+    let body = serde_json::to_vec(&agent).unwrap_or_default();
+    Response::builder()
+        .status(StatusCode::OK)
+        .header("Content-Type", "application/json")
+        .body(Body::from(body))
+}
+
+/// `GET /gateway/agents` — list all registered build agents.
+#[handler]
+pub async fn gateway_list_agents_handler(state: Data<&Arc<GatewayState>>) -> Response {
+    let agents = state.joined_agents.read().await.clone();
+    let body = serde_json::to_vec(&agents).unwrap_or_default();
+    Response::builder()
+        .status(StatusCode::OK)
+        .header("Content-Type", "application/json")
+        .body(Body::from(body))
+}
+
+/// `DELETE /gateway/agents/:id` — remove a registered build agent.
+#[handler]
+pub async fn gateway_remove_agent_handler(
+    PoemPath(id): PoemPath<String>,
+    state: Data<&Arc<GatewayState>>,
+) -> Response {
+    let mut agents = state.joined_agents.write().await;
+    let before = agents.len();
+    agents.retain(|a| a.id != id);
+    let removed = agents.len() < before;
+    drop(agents);
+
+    if removed {
+        crate::slog!("[gateway] Removed agent id={id}");
+        Response::builder()
+            .status(StatusCode::NO_CONTENT)
+            .body(Body::empty())
+    } else {
+        Response::builder()
+            .status(StatusCode::NOT_FOUND)
+            .body(Body::from("agent not found"))
+    }
+}
+
 // ── Health aggregation endpoint ──────────────────────────────────────
 
 /// HTTP GET `/health` handler for the gateway — aggregates health from all projects.
@@ -830,6 +995,28 @@ pub async fn run(config_path: &Path, port: u16) -> Result<(), std::io::Error> {
             poem::post(gateway_mcp_post_handler).get(gateway_mcp_get_handler),
         )
         .at("/health", poem::get(gateway_health_handler))
+        // Agent join endpoints.
+        .at("/gateway/mode", poem::get(gateway_mode_handler))
+        .at(
+            "/gateway/tokens",
+            poem::post(gateway_generate_token_handler),
+        )
+        .at(
+            "/gateway/register",
+            poem::post(gateway_register_agent_handler),
+        )
+        .at("/gateway/agents", poem::get(gateway_list_agents_handler))
+        .at(
+            "/gateway/agents/:id",
+            poem::delete(gateway_remove_agent_handler),
+        )
+        // Serve the embedded React frontend so the gateway has a UI.
+        .at(
+            "/assets/*path",
+            poem::get(crate::http::assets::embedded_asset),
+        )
+        .at("/*path", poem::get(crate::http::assets::embedded_file))
+        .at("/", poem::get(crate::http::assets::embedded_index))
         .data(state_arc);
 
     let host = std::env::var("HUSKIES_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
@@ -1153,4 +1340,162 @@ enabled = false
             "disabled bot.toml must return None in gateway mode"
         );
     }
+
+    // ── Agent join mechanism tests ───────────────────────────────────────
+
+    fn make_test_state() -> Arc<GatewayState> {
+        let mut projects = BTreeMap::new();
+        projects.insert(
+            "test".into(),
+            ProjectEntry {
+                url: "http://test:3001".into(),
+            },
+        );
+        let config = GatewayConfig { projects };
+        Arc::new(GatewayState::new(config).unwrap())
+    }
+
+    #[tokio::test]
+    async fn generate_token_creates_pending_token() {
+        let state = make_test_state();
+        let app = poem::Route::new()
+            .at(
+                "/gateway/tokens",
+                poem::post(gateway_generate_token_handler),
+            )
+            .data(state.clone());
+        let cli = poem::test::TestClient::new(app);
+        let resp = cli.post("/gateway/tokens").send().await;
+        assert_eq!(resp.0.status(), StatusCode::OK);
+        let body: Value = resp.0.into_body().into_json().await.unwrap();
+        let token = body["token"].as_str().unwrap();
+        assert!(!token.is_empty());
+        let tokens = state.pending_tokens.read().await;
+        assert!(tokens.contains_key(token));
+    }
+
+    #[tokio::test]
+    async fn register_agent_consumes_token() {
+        let state = make_test_state();
+
+        // Insert a token manually.
+        let token = "test-token-123".to_string();
+        state.pending_tokens.write().await.insert(
+            token.clone(),
+            PendingToken {
+                created_at: chrono::Utc::now().timestamp() as f64,
+            },
+        );
+
+        let app = poem::Route::new()
+            .at(
+                "/gateway/register",
+                poem::post(gateway_register_agent_handler),
+            )
+            .data(state.clone());
+        let cli = poem::test::TestClient::new(app);
+        let resp = cli
+            .post("/gateway/register")
+            .header("Content-Type", "application/json")
+            .body(
+                json!({
+                    "token": token,
+                    "label": "test-agent",
+                    "address": "ws://localhost:3001/crdt-sync"
+                })
+                .to_string(),
+            )
+            .send()
+            .await;
+        assert_eq!(resp.0.status(), StatusCode::OK);
+
+        // Token consumed.
+        assert!(state.pending_tokens.read().await.is_empty());
+        // Agent registered.
+        let agents = state.joined_agents.read().await;
+        assert_eq!(agents.len(), 1);
+        assert_eq!(agents[0].label, "test-agent");
+    }
+
+    #[tokio::test]
+    async fn register_agent_rejects_invalid_token() {
+        let state = make_test_state();
+        let app = poem::Route::new()
+            .at(
+                "/gateway/register",
+                poem::post(gateway_register_agent_handler),
+            )
+            .data(state.clone());
+        let cli = poem::test::TestClient::new(app);
+        let resp = cli
+            .post("/gateway/register")
+            .header("Content-Type", "application/json")
+            .body(
+                json!({
+                    "token": "bad-token",
+                    "label": "agent",
+                    "address": "ws://localhost:3001/crdt-sync"
+                })
+                .to_string(),
+            )
+            .send()
+            .await;
+        assert_eq!(resp.0.status(), StatusCode::UNAUTHORIZED);
+        assert!(state.joined_agents.read().await.is_empty());
+    }
+
+    #[tokio::test]
+    async fn list_agents_returns_registered_agents() {
+        let state = make_test_state();
+        state.joined_agents.write().await.push(JoinedAgent {
+            id: "id-1".into(),
+            label: "agent-1".into(),
+            address: "ws://a:3001/crdt-sync".into(),
+            registered_at: 0.0,
+        });
+        let app = poem::Route::new()
+            .at("/gateway/agents", poem::get(gateway_list_agents_handler))
+            .data(state.clone());
+        let cli = poem::test::TestClient::new(app);
+        let resp = cli.get("/gateway/agents").send().await;
+        assert_eq!(resp.0.status(), StatusCode::OK);
+        let agents: Vec<Value> = resp.0.into_body().into_json().await.unwrap();
+        assert_eq!(agents.len(), 1);
+        assert_eq!(agents[0]["label"], "agent-1");
+    }
+
+    #[tokio::test]
+    async fn remove_agent_deletes_by_id() {
+        let state = make_test_state();
+        state.joined_agents.write().await.push(JoinedAgent {
+            id: "del-id".into(),
+            label: "to-delete".into(),
+            address: "ws://x:3001/crdt-sync".into(),
+            registered_at: 0.0,
+        });
+        let app = poem::Route::new()
+            .at(
+                "/gateway/agents/:id",
+                poem::delete(gateway_remove_agent_handler),
+            )
+            .data(state.clone());
+        let cli = poem::test::TestClient::new(app);
+        let resp = cli.delete("/gateway/agents/del-id").send().await;
+        assert_eq!(resp.0.status(), StatusCode::NO_CONTENT);
+        assert!(state.joined_agents.read().await.is_empty());
+    }
+
+    #[tokio::test]
+    async fn remove_agent_unknown_id_returns_not_found() {
+        let state = make_test_state();
+        let app = poem::Route::new()
+            .at(
+                "/gateway/agents/:id",
+                poem::delete(gateway_remove_agent_handler),
+            )
+            .data(state.clone());
+        let cli = poem::test::TestClient::new(app);
+        let resp = cli.delete("/gateway/agents/no-such-id").send().await;
+        assert_eq!(resp.0.status(), StatusCode::NOT_FOUND);
+    }
 }