Add planner previews and streaming

server/main.py

@@ -1,19 +1,29 @@
-from fastapi import FastAPI, HTTPException
+import time
+
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
 
 from .config import ServerSettings
 from .grid import GridManager
-from .models import ActionPayload, GridDescriptor, GridInitRequest
+from .models import (
+    ActionPayload,
+    GridDescriptor,
+    GridInitRequest,
+    GridPlanRequest,
+    GridRefreshRequest,
+)
 from .planner import GridPlanner
+from .streamer import ScreenshotStreamer
 
 
 settings = ServerSettings()
 manager = GridManager(settings)
 planner = GridPlanner()
+streamer = ScreenshotStreamer()
 
 app = FastAPI(
     title="Clickthrough",
     description="Grid-aware surface that lets an agent plan clicks, drags, and typing on a fake screenshot",
-    version="0.2.0",
+    version="0.3.0",
 )
 
 
@@ -59,3 +69,51 @@ def grid_history(grid_id: str):
     except KeyError as exc:
         raise HTTPException(status_code=404, detail=str(exc)) from exc
     return {"grid_id": grid_id, "history": history}
+
+
+@app.post("/grid/{grid_id}/plan")
+def plan_grid(grid_id: str, request: GridPlanRequest):
+    try:
+        grid = manager.get_grid(grid_id)
+    except KeyError as exc:
+        raise HTTPException(status_code=404, detail=str(exc)) from exc
+    descriptor = grid.describe()
+    payload = planner.build_payload(
+        descriptor,
+        action=request.action,
+        preferred_label=request.preferred_label,
+        text=request.text,
+        comment=request.comment,
+    )
+    result = grid.preview_action(payload)
+    return {"plan": payload.model_dump(), "result": result, "descriptor": descriptor}
+
+
+@app.post("/grid/{grid_id}/refresh")
+async def refresh_grid(grid_id: str, payload: GridRefreshRequest):
+    try:
+        grid = manager.get_grid(grid_id)
+    except KeyError as exc:
+        raise HTTPException(status_code=404, detail=str(exc)) from exc
+    grid.update_screenshot(payload.screenshot_base64, payload.memo)
+    descriptor = grid.describe()
+    await streamer.broadcast(
+        grid_id,
+        {
+            "grid_id": grid_id,
+            "timestamp": time.time(),
+            "descriptor": descriptor,
+            "screenshot_base64": payload.screenshot_base64,
+        },
+    )
+    return {"status": "updated", "grid_id": grid_id}
+
+
+@app.websocket("/stream/screenshots")
+async def stream_screenshots(websocket: WebSocket, grid_id: str | None = None):
+    key = await streamer.connect(websocket, grid_id)
+    try:
+        while True:
+            await websocket.receive_text()
+    except WebSocketDisconnect:
+        streamer.disconnect(websocket, key)