huskies: merge 563_story_build_agent_join_mechanism_agents_register_with_the_gateway_via_token
This commit is contained in:
dave
+18
-1
@@ -2,11 +2,14 @@ import * as React from "react";
|
||||
import type { OAuthStatus } from "./api/client";
|
||||
import { api } from "./api/client";
|
||||
import { Chat } from "./components/Chat";
|
||||
import { GatewayPanel } from "./components/GatewayPanel";
|
||||
import { SelectionScreen } from "./components/selection/SelectionScreen";
|
||||
import { usePathCompletion } from "./components/selection/usePathCompletion";
|
||||
import { gatewayApi } from "./api/gateway";
|
||||
import "./App.css";
|
||||
|
||||
function App() {
|
||||
const [isGateway, setIsGateway] = React.useState<boolean | null>(null);
|
||||
const [projectPath, setProjectPath] = React.useState<string | null>(null);
|
||||
const [_view, setView] = React.useState<"chat" | "token-usage">("chat");
|
||||
const [isCheckingProject, setIsCheckingProject] = React.useState(true);
|
||||
@@ -19,6 +22,14 @@ function App() {
|
||||
null,
|
||||
);
|
||||
|
||||
// Detect gateway mode on startup — if /gateway/mode returns 200, we're a gateway.
|
||||
React.useEffect(() => {
|
||||
gatewayApi
|
||||
.getServerMode()
|
||||
.then((result) => setIsGateway(result.mode === "gateway"))
|
||||
.catch(() => setIsGateway(false));
|
||||
}, []);
|
||||
|
||||
React.useEffect(() => {
|
||||
let active = true;
|
||||
function fetchOAuthStatus() {
|
||||
@@ -188,10 +199,16 @@ function App() {
|
||||
}
|
||||
}
|
||||
|
||||
if (isCheckingProject) {
|
||||
// Still probing server mode — wait before rendering.
|
||||
if (isGateway === null || isCheckingProject) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Gateway mode: render the agent management UI instead of the normal chat.
|
||||
if (isGateway) {
|
||||
return <GatewayPanel />;
|
||||
}
|
||||
|
||||
return (
|
||||
<main
|
||||
className="container"
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
/// Gateway API client — used when running in gateway mode.
|
||||
///
|
||||
/// The gateway mode is detected by checking `GET /gateway/mode`. If it returns
|
||||
/// `{ "mode": "gateway" }` the frontend switches to the gateway UI.
|
||||
|
||||
export interface JoinedAgent {
|
||||
id: string;
|
||||
label: string;
|
||||
address: string;
|
||||
registered_at: number;
|
||||
}
|
||||
|
||||
export interface GenerateTokenResponse {
|
||||
token: string;
|
||||
}
|
||||
|
||||
export interface ServerMode {
|
||||
mode: "gateway" | "standard";
|
||||
}
|
||||
|
||||
async function gatewayRequest<T>(
|
||||
path: string,
|
||||
options: RequestInit = {},
|
||||
): Promise<T> {
|
||||
const res = await fetch(path, {
|
||||
headers: { "Content-Type": "application/json", ...(options.headers ?? {}) },
|
||||
...options,
|
||||
});
|
||||
if (!res.ok) {
|
||||
const text = await res.text();
|
||||
throw new Error(text || `Request failed (${res.status})`);
|
||||
}
|
||||
// DELETE /gateway/agents/:id returns 204 No Content.
|
||||
if (res.status === 204) {
|
||||
return undefined as unknown as T;
|
||||
}
|
||||
return res.json() as Promise<T>;
|
||||
}
|
||||
|
||||
export const gatewayApi = {
|
||||
/// Returns `{ mode: "gateway" }` if this server is a gateway, otherwise rejects.
|
||||
getServerMode(): Promise<ServerMode> {
|
||||
return gatewayRequest<ServerMode>("/gateway/mode");
|
||||
},
|
||||
|
||||
/// Generate a one-time join token for a new build agent.
|
||||
generateToken(): Promise<GenerateTokenResponse> {
|
||||
return gatewayRequest<GenerateTokenResponse>("/gateway/tokens", {
|
||||
method: "POST",
|
||||
});
|
||||
},
|
||||
|
||||
/// List all build agents that have registered with this gateway.
|
||||
listAgents(): Promise<JoinedAgent[]> {
|
||||
return gatewayRequest<JoinedAgent[]>("/gateway/agents");
|
||||
},
|
||||
|
||||
/// Remove a registered build agent by its ID.
|
||||
removeAgent(id: string): Promise<void> {
|
||||
return gatewayRequest<void>(`/gateway/agents/${id}`, {
|
||||
method: "DELETE",
|
||||
});
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,301 @@
|
||||
/// Gateway management panel shown when huskies runs in `--gateway` mode.
|
||||
///
|
||||
/// Provides:
|
||||
/// - An "Add Agent" button that generates a one-time join token.
|
||||
/// - Instructions for running a build agent with the token.
|
||||
/// - A list of connected agents with per-agent "Remove" buttons.
|
||||
|
||||
import * as React from "react";
|
||||
import { gatewayApi, type JoinedAgent } from "../api/gateway";
|
||||
|
||||
const { useCallback, useEffect, useState } = React;
|
||||
|
||||
function TokenDisplay({ token }: { token: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
|
||||
const envCmd = `HUSKIES_JOIN_TOKEN=${token} huskies agent --rendezvous <CRDT_SYNC_URL>`;
|
||||
const flagCmd = `huskies agent --rendezvous <CRDT_SYNC_URL> --join-token ${token}`;
|
||||
|
||||
const copyToClipboard = useCallback((text: string) => {
|
||||
void navigator.clipboard.writeText(text).then(() => {
|
||||
setCopied(true);
|
||||
setTimeout(() => setCopied(false), 2000);
|
||||
});
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<div
|
||||
style={{
|
||||
marginTop: "12px",
|
||||
padding: "12px 16px",
|
||||
background: "#161b22",
|
||||
border: "1px solid #238636",
|
||||
borderRadius: "8px",
|
||||
fontSize: "0.85em",
|
||||
}}
|
||||
>
|
||||
<div style={{ color: "#3fb950", fontWeight: 600, marginBottom: "8px" }}>
|
||||
Token generated — run the build agent with one of:
|
||||
</div>
|
||||
<div style={{ marginBottom: "6px" }}>
|
||||
<code
|
||||
style={{
|
||||
display: "block",
|
||||
background: "#0d1117",
|
||||
padding: "8px 10px",
|
||||
borderRadius: "4px",
|
||||
color: "#e6edf3",
|
||||
wordBreak: "break-all",
|
||||
}}
|
||||
>
|
||||
{envCmd}
|
||||
</code>
|
||||
</div>
|
||||
<div>
|
||||
<code
|
||||
style={{
|
||||
display: "block",
|
||||
background: "#0d1117",
|
||||
padding: "8px 10px",
|
||||
borderRadius: "4px",
|
||||
color: "#e6edf3",
|
||||
wordBreak: "break-all",
|
||||
}}
|
||||
>
|
||||
{flagCmd}
|
||||
</code>
|
||||
</div>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => copyToClipboard(flagCmd)}
|
||||
style={{
|
||||
marginTop: "8px",
|
||||
fontSize: "0.8em",
|
||||
padding: "3px 10px",
|
||||
borderRadius: "4px",
|
||||
border: "1px solid #444",
|
||||
background: "none",
|
||||
color: "#aaa",
|
||||
cursor: "pointer",
|
||||
}}
|
||||
>
|
||||
{copied ? "Copied!" : "Copy flag command"}
|
||||
</button>
|
||||
<div style={{ marginTop: "8px", color: "#666", fontSize: "0.85em" }}>
|
||||
This token is single-use. Generate a new one for each agent.
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function AgentRow({
|
||||
agent,
|
||||
onRemove,
|
||||
}: {
|
||||
agent: JoinedAgent;
|
||||
onRemove: (id: string) => void;
|
||||
}) {
|
||||
const registeredAt = new Date(agent.registered_at * 1000).toLocaleString();
|
||||
|
||||
return (
|
||||
<div
|
||||
data-testid={`agent-row-${agent.id}`}
|
||||
style={{
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
gap: "12px",
|
||||
padding: "10px 14px",
|
||||
background: "#161b22",
|
||||
border: "1px solid #30363d",
|
||||
borderRadius: "8px",
|
||||
marginBottom: "8px",
|
||||
}}
|
||||
>
|
||||
<div
|
||||
style={{
|
||||
width: "8px",
|
||||
height: "8px",
|
||||
borderRadius: "50%",
|
||||
background: "#3fb950",
|
||||
flexShrink: 0,
|
||||
}}
|
||||
/>
|
||||
<div style={{ flex: 1 }}>
|
||||
<div style={{ fontWeight: 600, color: "#e6edf3" }}>{agent.label}</div>
|
||||
<div style={{ fontSize: "0.8em", color: "#8b949e" }}>
|
||||
{agent.address}
|
||||
</div>
|
||||
<div style={{ fontSize: "0.75em", color: "#6e7681" }}>
|
||||
Registered {registeredAt}
|
||||
</div>
|
||||
</div>
|
||||
<button
|
||||
type="button"
|
||||
data-testid={`remove-agent-${agent.id}`}
|
||||
onClick={() => onRemove(agent.id)}
|
||||
style={{
|
||||
fontSize: "0.8em",
|
||||
padding: "4px 10px",
|
||||
borderRadius: "4px",
|
||||
border: "1px solid #f85149",
|
||||
background: "none",
|
||||
color: "#f85149",
|
||||
cursor: "pointer",
|
||||
}}
|
||||
>
|
||||
Remove
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
/// Gateway management panel — rendered when running in `--gateway` mode.
|
||||
export function GatewayPanel() {
|
||||
const [agents, setAgents] = useState<JoinedAgent[]>([]);
|
||||
const [token, setToken] = useState<string | null>(null);
|
||||
const [generating, setGenerating] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
gatewayApi
|
||||
.listAgents()
|
||||
.then(setAgents)
|
||||
.catch(() => setAgents([]));
|
||||
}, []);
|
||||
|
||||
const handleAddAgent = useCallback(async () => {
|
||||
setGenerating(true);
|
||||
setError(null);
|
||||
setToken(null);
|
||||
try {
|
||||
const result = await gatewayApi.generateToken();
|
||||
setToken(result.token);
|
||||
} catch (e) {
|
||||
setError(e instanceof Error ? e.message : String(e));
|
||||
} finally {
|
||||
setGenerating(false);
|
||||
}
|
||||
}, []);
|
||||
|
||||
const handleRemoveAgent = useCallback(async (id: string) => {
|
||||
try {
|
||||
await gatewayApi.removeAgent(id);
|
||||
setAgents((prev) => prev.filter((a) => a.id !== id));
|
||||
} catch (e) {
|
||||
setError(e instanceof Error ? e.message : String(e));
|
||||
}
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<div
|
||||
style={{
|
||||
minHeight: "100vh",
|
||||
background: "#0d1117",
|
||||
color: "#e6edf3",
|
||||
padding: "32px",
|
||||
fontFamily: "-apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif",
|
||||
}}
|
||||
>
|
||||
<div style={{ maxWidth: "720px", margin: "0 auto" }}>
|
||||
<h1 style={{ fontSize: "1.5em", fontWeight: 700, marginBottom: "4px" }}>
|
||||
Huskies Gateway
|
||||
</h1>
|
||||
<p style={{ color: "#8b949e", marginBottom: "32px" }}>
|
||||
Manage build agents connected to this gateway.
|
||||
</p>
|
||||
|
||||
{/* Add Agent */}
|
||||
<section style={{ marginBottom: "32px" }}>
|
||||
<h2
|
||||
style={{
|
||||
fontSize: "1.1em",
|
||||
fontWeight: 600,
|
||||
marginBottom: "12px",
|
||||
borderBottom: "1px solid #21262d",
|
||||
paddingBottom: "8px",
|
||||
}}
|
||||
>
|
||||
Add Agent
|
||||
</h2>
|
||||
<button
|
||||
type="button"
|
||||
data-testid="add-agent-button"
|
||||
onClick={handleAddAgent}
|
||||
disabled={generating}
|
||||
style={{
|
||||
padding: "8px 18px",
|
||||
borderRadius: "6px",
|
||||
border: "1px solid #238636",
|
||||
background: generating ? "#1a2f1a" : "#238636",
|
||||
color: "#fff",
|
||||
cursor: generating ? "not-allowed" : "pointer",
|
||||
fontWeight: 600,
|
||||
fontSize: "0.9em",
|
||||
}}
|
||||
>
|
||||
{generating ? "Generating…" : "Add Agent"}
|
||||
</button>
|
||||
{token && <TokenDisplay token={token} />}
|
||||
</section>
|
||||
|
||||
{/* Agent list */}
|
||||
<section>
|
||||
<h2
|
||||
style={{
|
||||
fontSize: "1.1em",
|
||||
fontWeight: 600,
|
||||
marginBottom: "12px",
|
||||
borderBottom: "1px solid #21262d",
|
||||
paddingBottom: "8px",
|
||||
}}
|
||||
>
|
||||
Connected Agents{" "}
|
||||
{agents.length > 0 && (
|
||||
<span
|
||||
style={{
|
||||
fontSize: "0.8em",
|
||||
color: "#8b949e",
|
||||
fontWeight: 400,
|
||||
}}
|
||||
>
|
||||
({agents.length})
|
||||
</span>
|
||||
)}
|
||||
</h2>
|
||||
{agents.length === 0 ? (
|
||||
<p style={{ color: "#6e7681" }}>
|
||||
No agents connected yet. Click "Add Agent" to generate a join
|
||||
token.
|
||||
</p>
|
||||
) : (
|
||||
<div>
|
||||
{agents.map((agent) => (
|
||||
<AgentRow
|
||||
key={agent.id}
|
||||
agent={agent}
|
||||
onRemove={handleRemoveAgent}
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
|
||||
{error && (
|
||||
<div
|
||||
style={{
|
||||
marginTop: "16px",
|
||||
padding: "10px 14px",
|
||||
background: "#f8514911",
|
||||
border: "1px solid #f85149",
|
||||
borderRadius: "6px",
|
||||
color: "#f85149",
|
||||
fontSize: "0.875em",
|
||||
}}
|
||||
>
|
||||
{error}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -36,10 +36,15 @@ const SCAN_INTERVAL_SECS: u64 = 15;
|
||||
///
|
||||
/// This function never returns under normal operation — it runs until the
|
||||
/// process is terminated (SIGINT/SIGTERM).
|
||||
///
|
||||
/// If `join_token` and `gateway_url` are both provided the agent will register
|
||||
/// itself with the gateway on startup using the one-time token.
|
||||
pub async fn run(
|
||||
project_root: Option<PathBuf>,
|
||||
rendezvous_url: String,
|
||||
port: u16,
|
||||
join_token: Option<String>,
|
||||
gateway_url: Option<String>,
|
||||
) -> Result<(), std::io::Error> {
|
||||
let project_root = match project_root {
|
||||
Some(r) => r,
|
||||
@@ -127,6 +132,14 @@ pub async fn run(
|
||||
// Write initial heartbeat.
|
||||
write_heartbeat(&rendezvous_url, port);
|
||||
|
||||
// Register with gateway if a join token and gateway URL were provided.
|
||||
if let (Some(token), Some(url)) = (join_token, gateway_url) {
|
||||
let node_id = crdt_state::our_node_id().unwrap_or_else(|| "unknown".to_string());
|
||||
let label = format!("build-agent-{}", &node_id[..node_id.len().min(8)]);
|
||||
let address = format!("ws://0.0.0.0:{port}/crdt-sync");
|
||||
register_with_gateway(&url, &token, &label, &address).await;
|
||||
}
|
||||
|
||||
// Reconcile any committed work from a previous session.
|
||||
{
|
||||
let recon_agents = Arc::clone(&agents);
|
||||
@@ -427,6 +440,36 @@ fn push_feature_branch(worktree_path: &str, story_id: &str) -> Result<(), String
|
||||
}
|
||||
}
|
||||
|
||||
// ── Gateway registration ──────────────────────────────────────────────────
|
||||
|
||||
/// Register this build agent with a gateway using a one-time join token.
|
||||
///
|
||||
/// POSTs `{ token, label, address }` to `{gateway_url}/gateway/register`. On
|
||||
/// success the gateway stores the agent and it will appear in the gateway UI.
|
||||
async fn register_with_gateway(gateway_url: &str, token: &str, label: &str, address: &str) {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/gateway/register", gateway_url.trim_end_matches('/'));
|
||||
let body = serde_json::json!({
|
||||
"token": token,
|
||||
"label": label,
|
||||
"address": address,
|
||||
});
|
||||
match client.post(&url).json(&body).send().await {
|
||||
Ok(resp) if resp.status().is_success() => {
|
||||
slog!("[agent-mode] Registered with gateway at {gateway_url}");
|
||||
}
|
||||
Ok(resp) => {
|
||||
slog!(
|
||||
"[agent-mode] Gateway registration failed: HTTP {}",
|
||||
resp.status()
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
slog!("[agent-mode] Gateway registration error: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
+100
-1
@@ -56,6 +56,10 @@ struct CliArgs {
|
||||
rendezvous: Option<String>,
|
||||
/// Whether `--gateway` mode was requested (proxy MCP calls to per-project containers).
|
||||
gateway: bool,
|
||||
/// One-time join token for registering this build agent with a gateway (`--join-token`).
|
||||
join_token: Option<String>,
|
||||
/// HTTP URL of the gateway to register with when a join token is provided (`--gateway-url`).
|
||||
gateway_url: Option<String>,
|
||||
}
|
||||
|
||||
/// Parse CLI arguments into `CliArgs`, or exit early for `--help` / `--version`.
|
||||
@@ -66,6 +70,8 @@ fn parse_cli_args(args: &[String]) -> Result<CliArgs, String> {
|
||||
let mut agent = false;
|
||||
let mut gateway = false;
|
||||
let mut rendezvous: Option<String> = None;
|
||||
let mut join_token: Option<String> = None;
|
||||
let mut gateway_url: Option<String> = None;
|
||||
let mut i = 0;
|
||||
|
||||
while i < args.len() {
|
||||
@@ -106,6 +112,26 @@ fn parse_cli_args(args: &[String]) -> Result<CliArgs, String> {
|
||||
let val = &a["--rendezvous=".len()..];
|
||||
rendezvous = Some(val.to_string());
|
||||
}
|
||||
"--join-token" => {
|
||||
i += 1;
|
||||
if i >= args.len() {
|
||||
return Err("--join-token requires a value".to_string());
|
||||
}
|
||||
join_token = Some(args[i].clone());
|
||||
}
|
||||
a if a.starts_with("--join-token=") => {
|
||||
join_token = Some(a["--join-token=".len()..].to_string());
|
||||
}
|
||||
"--gateway-url" => {
|
||||
i += 1;
|
||||
if i >= args.len() {
|
||||
return Err("--gateway-url requires a value".to_string());
|
||||
}
|
||||
gateway_url = Some(args[i].clone());
|
||||
}
|
||||
a if a.starts_with("--gateway-url=") => {
|
||||
gateway_url = Some(a["--gateway-url=".len()..].to_string());
|
||||
}
|
||||
"--gateway" => {
|
||||
gateway = true;
|
||||
}
|
||||
@@ -139,6 +165,8 @@ fn parse_cli_args(args: &[String]) -> Result<CliArgs, String> {
|
||||
agent,
|
||||
rendezvous,
|
||||
gateway,
|
||||
join_token,
|
||||
gateway_url,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -172,6 +200,11 @@ fn print_help() {
|
||||
println!(
|
||||
" (or cwd) and proxies MCP calls to per-project containers."
|
||||
);
|
||||
println!(" --join-token <TOKEN> One-time token for registering this build agent with a");
|
||||
println!(" gateway. Also readable from HUSKIES_JOIN_TOKEN env var.");
|
||||
println!(" --gateway-url <URL> HTTP URL of the gateway to register with when");
|
||||
println!(" --join-token is provided (agent mode only).");
|
||||
println!(" Also readable from HUSKIES_GATEWAY_URL env var.");
|
||||
}
|
||||
|
||||
/// Resolve the optional positional path argument into an absolute `PathBuf`.
|
||||
@@ -411,7 +444,16 @@ async fn main() -> Result<(), std::io::Error> {
|
||||
if is_agent {
|
||||
let agent_root = app_state.project_root.lock().unwrap().clone();
|
||||
let rendezvous = agent_rendezvous.expect("agent mode requires --rendezvous");
|
||||
return agent_mode::run(agent_root, rendezvous, port).await;
|
||||
// Join token / gateway URL can come from CLI flags or environment variables.
|
||||
let join_token = cli
|
||||
.join_token
|
||||
.clone()
|
||||
.or_else(|| std::env::var("HUSKIES_JOIN_TOKEN").ok());
|
||||
let agent_gateway_url = cli
|
||||
.gateway_url
|
||||
.clone()
|
||||
.or_else(|| std::env::var("HUSKIES_GATEWAY_URL").ok());
|
||||
return agent_mode::run(agent_root, rendezvous, port, join_token, agent_gateway_url).await;
|
||||
}
|
||||
|
||||
let workflow = Arc::new(std::sync::Mutex::new(WorkflowState::default()));
|
||||
@@ -1087,6 +1129,63 @@ name = "coder"
|
||||
assert_eq!(result.path, Some("/my/project".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_join_token_flag() {
|
||||
let args = vec![
|
||||
"agent".to_string(),
|
||||
"--rendezvous".to_string(),
|
||||
"ws://host:3001/crdt-sync".to_string(),
|
||||
"--join-token".to_string(),
|
||||
"my-secret-token".to_string(),
|
||||
];
|
||||
let result = parse_cli_args(&args).unwrap();
|
||||
assert_eq!(result.join_token, Some("my-secret-token".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_join_token_equals_syntax() {
|
||||
let args = vec![
|
||||
"agent".to_string(),
|
||||
"--rendezvous".to_string(),
|
||||
"ws://host:3001/crdt-sync".to_string(),
|
||||
"--join-token=abc123".to_string(),
|
||||
];
|
||||
let result = parse_cli_args(&args).unwrap();
|
||||
assert_eq!(result.join_token, Some("abc123".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_gateway_url_flag() {
|
||||
let args = vec![
|
||||
"agent".to_string(),
|
||||
"--rendezvous".to_string(),
|
||||
"ws://host:3001/crdt-sync".to_string(),
|
||||
"--gateway-url".to_string(),
|
||||
"http://gateway:3000".to_string(),
|
||||
];
|
||||
let result = parse_cli_args(&args).unwrap();
|
||||
assert_eq!(result.gateway_url, Some("http://gateway:3000".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_join_token_missing_value_is_error() {
|
||||
let args = vec!["--join-token".to_string()];
|
||||
assert!(parse_cli_args(&args).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_gateway_url_missing_value_is_error() {
|
||||
let args = vec!["--gateway-url".to_string()];
|
||||
assert!(parse_cli_args(&args).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_no_args_join_token_and_gateway_url_are_none() {
|
||||
let result = parse_cli_args(&[]).unwrap();
|
||||
assert_eq!(result.join_token, None);
|
||||
assert_eq!(result.gateway_url, None);
|
||||
}
|
||||
|
||||
// ── resolve_path_arg ────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user