import os
import uuid
from enum import Enum
from html import escape
from typing import Optional
from fastapi import FastAPI, Form, Header, HTTPException
from fastapi.responses import HTMLResponse, RedirectResponse
from pydantic import BaseModel
# =============================================================================
# Configuration
# =============================================================================
BROKER_TOKEN = os.environ.get("BROKER_TOKEN")
UI_TOKEN = os.environ.get("UI_TOKEN")
if not BROKER_TOKEN:
raise RuntimeError(
"Missing BROKER_TOKEN. Add it in Hugging Face Space "
"Settings → Variables and secrets → New secret."
)
if not UI_TOKEN:
raise RuntimeError(
"Missing UI_TOKEN. Add it in Hugging Face Space "
"Settings → Variables and secrets → New secret."
)
# =============================================================================
# Safe predefined options
# =============================================================================
MODEL_OPTIONS = {
"granite-4.0-micro": "ibm-granite/granite-4.0-micro",
"granite-4.1-8b": "ibm-granite/granite-4.1-8b",
"granite-4.1-30b": "ibm-granite/granite-4.1-30b",
"qwen2.5-coder-32b": "Qwen/Qwen2.5-Coder-32B-Instruct",
}
BASIC_COMMANDS = {
"hostname": "Run hostname",
"whoami": "Run whoami",
"pwd": "Show current directory",
"disk": "Show disk usage",
"date": "Show current date",
"list_home": "List home directory",
}
PARAMETERIZED_COMMANDS = {
"query_llm": "Query LLM with model, GPU count, and prompt",
}
# =============================================================================
# Data model
# =============================================================================
class JobStatus(str, Enum):
queued = "queued"
running = "running"
done = "done"
failed = "failed"
class Job(BaseModel):
id: str
command: str
status: JobStatus = JobStatus.queued
result: Optional[str] = None
# Parameters for query_llm
model: Optional[str] = None
gpus: Optional[int] = None
user_text: Optional[str] = None
app = FastAPI(title="Fury Broker")
# In-memory storage. Jobs disappear if the Space restarts.
jobs: dict[str, Job] = {}
# =============================================================================
# Security helpers
# =============================================================================
def verify_broker_token(x_broker_token: Optional[str]) -> None:
"""
Used by the Fury worker and command-line API calls.
Header:
X-Broker-Token: BROKER_TOKEN
"""
if x_broker_token != BROKER_TOKEN:
raise HTTPException(status_code=401, detail="Unauthorized")
def verify_ui_token(ui_token: Optional[str]) -> None:
"""
Used by browser form submissions.
Form field:
ui_token
"""
if ui_token != UI_TOKEN:
raise HTTPException(status_code=401, detail="Invalid UI token")
def validate_basic_command(command: str) -> None:
if command not in BASIC_COMMANDS:
raise HTTPException(status_code=400, detail=f"Command is not allowed: {command}")
def validate_query_llm_args(
model: str,
gpus: int,
user_text: str,
) -> None:
if model not in MODEL_OPTIONS:
raise HTTPException(status_code=400, detail=f"Model is not allowed: {model}")
if gpus < 1 or gpus > 16:
raise HTTPException(status_code=400, detail="GPUs must be between 1 and 16")
if user_text is None:
raise HTTPException(status_code=400, detail="Prompt is required")
if len(user_text.strip()) == 0:
raise HTTPException(status_code=400, detail="Prompt cannot be empty")
if len(user_text) > 10_000:
raise HTTPException(
status_code=400,
detail="Prompt is too long; max 10,000 characters",
)
# =============================================================================
# Health and inspection APIs
# =============================================================================
@app.get("/health")
def health():
return {
"status": "ok",
"service": "fury-broker",
"jobs_count": len(jobs),
}
@app.get("/api/commands")
def list_commands(x_broker_token: Optional[str] = Header(default=None)):
verify_broker_token(x_broker_token)
return {
"basic_commands": BASIC_COMMANDS,
"parameterized_commands": PARAMETERIZED_COMMANDS,
"query_llm": {
"command": "query_llm",
"model_options": MODEL_OPTIONS,
"gpu_range": [1, 16],
"default_gpus": 1,
"default_prompt": "hello",
},
}
# =============================================================================
# Browser UI
# =============================================================================
@app.get("/", response_class=HTMLResponse)
def home():
basic_command_options_html = "\n".join(
f''
for name, label in BASIC_COMMANDS.items()
)
model_options_html = "\n".join(
f''
for key, value in MODEL_OPTIONS.items()
)
gpu_options_html = "\n".join(
f''
for i in range(1, 17)
)
rows = ""
for job in reversed(list(jobs.values())):
details = []
if job.model:
details.append(f"model={job.model}")
if job.gpus:
details.append(f"gpus={job.gpus}")
if job.user_text:
preview = job.user_text[:500]
if len(job.user_text) > 500:
preview += "..."
details.append(f"prompt={preview}")
safe_details = escape("\n".join(details))
safe_result = escape(job.result or "")
rows += f"""
{escape(job.id)}
{escape(job.command)}
{escape(job.status.value)}
{safe_details}
{safe_result}
"""
return f"""
Fury Broker
Fury Broker
Submit approved jobs from Hugging Face Spaces to the worker running on Fury.
This UI does not allow arbitrary shell commands.
It submits only predefined command names and validated parameters.
Fury validates the job again locally before execution.
Refresh the page after a few seconds to see updated results.