Improve Studio workspace controls and startup

tests/visualization/test_runtime.py

@@ -1,20 +1,19 @@
 """Tests for the installable browser workspace runtime."""
 
-from io import BytesIO
 import json
-from pathlib import Path
 import tarfile
 import threading
+from io import BytesIO
+from pathlib import Path
 
 import pytest
 from fastapi.testclient import TestClient
 from PIL import Image
 
 from visualization import app as app_runtime
-from visualization import modal_deployment
-from visualization import runtime
-from visualization.app import create_app
+from visualization import modal_deployment, runtime
 from visualization import samples as sample_runtime
+from visualization.app import create_app
 
 
 def _cuda_hardware() -> dict:
@@ -83,6 +82,14 @@ def test_configuration_schema_defaults_to_vggt() -> None:
     assert argoverse_options["log_id"]["required"] is True
 
 
+def test_yaml_catalog_ignores_cloud_duplicate_artifacts(tmp_path: Path) -> None:
+    (tmp_path / "vggt.yaml").write_text("model: canonical\n", encoding="utf-8")
+    (tmp_path / "vggt 2.yaml").write_text("model: duplicate\n", encoding="utf-8")
+    (tmp_path / "_base.yaml").write_text("model: base\n", encoding="utf-8")
+
+    assert runtime._yaml_stem_options(tmp_path) == ["vggt"]
+
+
 def test_detect_hardware_always_reports_cpu() -> None:
     hardware = runtime.detect_hardware()
 
@@ -237,6 +244,9 @@ def test_completed_job_exposes_splat_download(tmp_path: Path) -> None:
 
     jobs = client.get("/api/jobs").json()["items"]
     assert jobs[0]["has_final_splat"] is True
+    live = client.get("/api/jobs/finished-job/live")
+    assert live.status_code == 200
+    assert live.json()["final_url"].endswith("/runs/finished/gaussian_splats.ply")
     response = client.get("/api/jobs/finished-job/splat", params={"format": "ply"})
     assert response.status_code == 200
     assert response.content.startswith(b"ply\n")
@@ -273,13 +283,18 @@ def test_live_job_exposes_dask_worker_status(tmp_path: Path) -> None:
         live_root=str(live_root),
         spec={},
         command=[],
+        log_tail=[
+            "2026-08-13 17:00:00 [runner.py] INFO: 🌟 GTSFM: Starting SceneOptimizer...",
+            "2026-08-13 17:00:01 [scene_optimizer.py] INFO: 🔥 GTSFM: Partitioning the view graph...",
+        ],
     )
 
     response = TestClient(app).get("/api/jobs/running-job/live")
 
     assert response.status_code == 200
     assert response.json()["dask"]["workers"] == 2
     assert response.json()["dask"]["running_tasks"] == 3
+    assert response.json()["message"] == "🔥 GTSFM: Partitioning the view graph..."
 
 
 def test_workspace_rejects_invalid_splat_exports(tmp_path: Path) -> None:
@@ -308,6 +323,23 @@ def test_sample_downloader_uses_certifi_ca_bundle() -> None:
     assert sample_runtime._SSL_CONTEXT.get_ca_certs()
 
 
+def test_remote_workspace_requests_use_certifi_ca_bundle(monkeypatch: pytest.MonkeyPatch) -> None:
+    observed: dict[str, object] = {}
+
+    def fake_urlopen(_request: object, **kwargs: object) -> BytesIO:
+        observed.update(kwargs)
+        return BytesIO(b'{"status": "ready"}')
+
+    monkeypatch.setattr(runtime.urllib.request, "urlopen", fake_urlopen)
+
+    result = runtime.JobManager._remote_json("https://workspace.example/api/hardware", "secret")
+
+    assert result == {"status": "ready"}
+    assert observed["context"] is runtime._SSL_CONTEXT
+    assert observed["timeout"] == runtime._REMOTE_REQUEST_TIMEOUT_SECONDS
+    assert runtime._SSL_CONTEXT.get_ca_certs()
+
+
 def test_workspace_rejects_unknown_github_sample(tmp_path: Path) -> None:
     response = TestClient(create_app(tmp_path)).post("/api/samples/not-a-sample/prepare")
 
@@ -559,6 +591,7 @@ def poll(self) -> None:
 
 
 def test_job_cancel_forwards_to_remote_vm(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(threading.Thread, "start", lambda thread: thread.run())
     manager = runtime.JobManager(tmp_path)
     job = runtime.ManagedJob(
         id="remote-job",
@@ -585,6 +618,60 @@ def test_job_cancel_forwards_to_remote_vm(tmp_path: Path, monkeypatch: pytest.Mo
 
     assert cancelled["status"] == "cancelled"
     assert calls == [("https://gpu.example.test/api/jobs/upstream-job/cancel", "secret", {})]
+    assert job.remote == {
+        "endpoint": "https://gpu.example.test",
+        "job_id": "upstream-job",
+        "cancel_status": "confirmed",
+    }
+    assert job.log_tail[-1] == "Modal job cancellation confirmed."
+
+
+def test_job_cancel_during_modal_start_forwards_after_remote_accepts(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    manager = runtime.JobManager(tmp_path)
+    job = runtime.ManagedJob(
+        id="starting-remote-job",
+        name="remote",
+        status="running",
+        created_at="now",
+        updated_at="now",
+        output_root=str(tmp_path / "output"),
+        live_root=str(tmp_path / "live"),
+        spec={"splat_implementation": "none"},
+        command=[],
+        remote={"endpoint": "https://gpu.example.test"},
+        remote_api_key="secret",
+    )
+    manager._jobs[job.id] = job
+    calls: list[tuple[str, object]] = []
+
+    def fake_remote_json(
+        endpoint: str, _api_key: str, *, payload: dict[str, object] | None = None
+    ) -> dict[str, object]:
+        calls.append((endpoint, payload))
+        if endpoint.endswith("/api/jobs"):
+            # Reproduce the user clicking Cancel while the POST is waiting for
+            # Modal's cold-started container to accept the job.
+            manager.cancel(job.id)
+            return {"id": "upstream-job", "status": "queued"}
+        return {"id": "upstream-job", "status": "cancelled"}
+
+    monkeypatch.setattr(manager, "_remote_json", fake_remote_json)
+
+    manager._run_remote(job)
+
+    assert job.status == "cancelled"
+    assert calls == [
+        ("https://gpu.example.test/api/jobs", {"splat_implementation": "none", "execution_target": "local"}),
+        ("https://gpu.example.test/api/jobs/upstream-job/cancel", {}),
+    ]
+    assert job.remote == {
+        "endpoint": "https://gpu.example.test",
+        "job_id": "upstream-job",
+        "workspace_url": "https://gpu.example.test/?view=results",
+        "cancel_status": "confirmed",
+    }
 
 
 def test_remote_job_keeps_modal_and_ssh_credentials_private(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -629,7 +716,7 @@ def test_remote_job_transfers_local_dataset_before_submit(tmp_path: Path, monkey
 
     def fake_upload(_endpoint: str, _api_key: str, directory: Path, _runtime_root: Path) -> dict[str, object]:
         assert directory == dataset
-        return {"path": "/workspace/results/.gtsfm/uploads/remote-dataset"}
+        return {"path": "/mnt/gtsfm-studio/results/.gtsfm/uploads/remote-dataset"}
 
     def fake_remote_json(url: str, _api_key: str, *, payload: dict[str, object] | None = None) -> dict[str, object]:
         if url.endswith("/api/jobs"):
@@ -659,9 +746,10 @@ def fake_remote_json(url: str, _api_key: str, *, payload: dict[str, object] | No
 
     manager._run_remote(managed)
 
-    assert posted["dataset_dir"] == "/workspace/results/.gtsfm/uploads/remote-dataset"
+    assert posted["dataset_dir"] == "/mnt/gtsfm-studio/results/.gtsfm/uploads/remote-dataset"
     assert posted["execution_target"] == "local"
     assert managed.status == "completed"
+    assert json.loads((Path(managed.live_root) / "status.json").read_text(encoding="utf-8")) == {}
 
 
 def test_remote_sample_is_downloaded_by_remote_workspace(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -696,11 +784,11 @@ def test_remote_sample_is_downloaded_by_remote_workspace(tmp_path: Path, monkeyp
     assert response.json()["spec"]["sample_id"] == "lund-door"
 
 
-def test_modal_remote_api_key_is_stable_and_scoped() -> None:
-    first = modal_deployment.modal_remote_api_key("ak-one", "as-secret")
+def test_modal_workspace_api_key_is_stable_and_scoped() -> None:
+    first = modal_deployment.modal_workspace_api_key("ak-one", "as-secret")
 
-    assert first == modal_deployment.modal_remote_api_key("ak-one", "as-secret")
-    assert first != modal_deployment.modal_remote_api_key("ak-two", "as-secret")
+    assert first == modal_deployment.modal_workspace_api_key("ak-one", "as-secret")
+    assert first != modal_deployment.modal_workspace_api_key("ak-two", "as-secret")
     assert "as-secret" not in first
 
 
@@ -762,3 +850,37 @@ def test_fastapi_exposes_openapi_and_job_websocket(tmp_path: Path) -> None:
 
     with client.websocket_connect("/api/events/jobs") as websocket:
         assert websocket.receive_json() == {"items": []}
+
+
+def test_fastapi_uses_injected_hardware_provider_and_serves_favicon(tmp_path: Path) -> None:
+    hardware = {
+        "summary": "Remote GPU configured without allocation",
+        "devices": [{"id": "cuda:0", "supports_gaussian_splatting": True}],
+        "accelerator_count": 1,
+        "platform": {},
+    }
+    client = TestClient(create_app(tmp_path, hardware_provider=lambda: hardware))
+
+    assert client.get("/api/hardware").json() == hardware
+    favicon = client.get("/favicon.ico")
+    assert favicon.status_code == 200
+    assert favicon.headers["content-type"] == "image/png"
+
+
+def test_job_manager_accepts_safe_caller_provided_job_id(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(threading.Thread, "start", lambda _thread: None)
+    manager = runtime.JobManager(tmp_path)
+    spec = {
+        "name": "durable-id",
+        "dataset_dir": str(tmp_path),
+        "loader": "olsson",
+        "config_name": "vggt",
+        "splat_implementation": "none",
+        "hardware": "cpu",
+    }
+
+    job = manager.start(spec, job_id="modal_job-123")
+
+    assert job["id"] == "modal_job-123"
+    with pytest.raises(ValueError, match="unsupported characters"):
+        manager.start(spec, job_id="../../unsafe")

visualization/app.py

@@ -8,39 +8,51 @@
 import hmac
 import json
 import os
+import re
 import shutil
 import tarfile
 import urllib.error
 import uuid
+from collections.abc import Callable
 from pathlib import Path, PurePosixPath
 from typing import Annotated, Any
 from urllib.parse import quote, urlparse
 
 import uvicorn
-from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Request, UploadFile, WebSocket
+from fastapi import (Depends, FastAPI, File, Form, Header, HTTPException,
+                     Request, UploadFile, WebSocket)
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from PIL import Image, UnidentifiedImageError
 from pydantic import BaseModel, ConfigDict
 from starlette.websockets import WebSocketDisconnect
 
-from visualization.runtime import (
-    JobManager,
-    configuration_schema,
-    detect_hardware,
-    install_optional_setup,
-    setup_status,
-)
-from visualization.modal_deployment import ModalDeploymentManager, modal_remote_api_key
-from visualization.samples import SampleDownloadError, prepare_sample, sample_catalog
-
+from visualization.modal_deployment import (ModalDeploymentManager,
+                                            modal_workspace_api_key)
+from visualization.runtime import (JobManager, configuration_schema,
+                                   detect_hardware, install_optional_setup,
+                                   setup_status)
+from visualization.samples import (SampleDownloadError, prepare_sample,
+                                   sample_catalog)
 
 PACKAGE_ROOT = Path(__file__).resolve().parent
 STATIC_ROOT = PACKAGE_ROOT / "static"
 TEMPLATE_ROOT = PACKAGE_ROOT / "templates"
 SPLAT_EXPORT_FORMATS = {"ply"}
 IMAGE_SUFFIXES = {".avif", ".bmp", ".heic", ".heif", ".jpeg", ".jpg", ".png", ".tif", ".tiff", ".webp"}
+_LOG_LEVEL_MESSAGE = re.compile(r"\b(?:DEBUG|INFO|WARNING|ERROR|CRITICAL):\s*(.+)$")
+_PIPELINE_STATUS_MARKERS = (
+    "gtsfm:",
+    "running vggt",
+    "gaussian",
+    "splat",
+    "partition",
+    "cluster optimization",
+    "bundle adjustment",
+    "image pair retrieval",
+    "sceneoptimizer",
+)
 
 
 class RemoteInspectRequest(BaseModel):
@@ -148,7 +160,9 @@ async def _discover_modal_endpoint(token_id: str, token_secret: str) -> dict[str
     """Find the deployed GTSFM web function in the token's default Modal environment."""
 
     from modal.client import _Client
-    from modal.exception import AuthError, ConnectionError as ModalConnectionError, PermissionDeniedError
+    from modal.exception import AuthError
+    from modal.exception import ConnectionError as ModalConnectionError
+    from modal.exception import PermissionDeniedError
     from modal_proto import api_pb2
 
     client = None
@@ -185,7 +199,8 @@ async def _discover_modal_endpoint(token_id: str, token_secret: str) -> dict[str
                 candidates.append((score, url, app_name or "gtsfm", function_name or "web"))
         if not candidates:
             raise ValueError(
-                "No deployed GTSFM web app was found in this Modal workspace. Deploy the GTSFM Modal app first, then try again."
+                "No deployed GTSFM web app was found in this Modal workspace. "
+                "Deploy the GTSFM Modal app first, then try again."
             )
         _, endpoint, app_name, function_name = max(candidates, key=lambda item: item[0])
         return {"endpoint": endpoint, "app_name": app_name, "function_name": function_name}
@@ -293,6 +308,15 @@ def _live_state(manager: JobManager, resolved_base: Path, job_id: str) -> dict[s
         "stage": "pipeline",
         "progress": None,
     }
+    for raw_line in reversed(job.log_tail):
+        match = _LOG_LEVEL_MESSAGE.search(raw_line)
+        message = (match.group(1) if match else raw_line).strip()
+        lowered = message.lower()
+        if message and any(marker in lowered for marker in _PIPELINE_STATUS_MARKERS):
+            payload["message"] = message
+            break
+    if "message" not in payload:
+        payload["message"] = "Waiting for the reconstruction pipeline to report its first stage…"
     if status_path.exists():
         try:
             loaded = json.loads(status_path.read_text(encoding="utf-8"))
@@ -328,11 +352,16 @@ def _websocket_authorized(websocket: WebSocket) -> bool:
     return hmac.compare_digest(provided, expected)
 
 
-def create_app(base_dir: Path | str = "results") -> FastAPI:
+def create_app(
+    base_dir: Path | str = "results",
+    *,
+    job_manager: JobManager | None = None,
+    hardware_provider: Callable[[], dict[str, Any]] = detect_hardware,
+) -> FastAPI:
     """Create an isolated FastAPI workspace for ``base_dir``."""
 
     resolved_base = Path(base_dir).expanduser().resolve()
-    manager = JobManager(resolved_base)
+    manager = job_manager or JobManager(resolved_base)
     modal_deployments = ModalDeploymentManager(
         lambda token_id, token_secret: asyncio.run(_discover_modal_endpoint(token_id, token_secret))
     )
@@ -368,6 +397,10 @@ def require_api_key(authorization: Annotated[str | None, Header()] = None) -> No
     def index() -> HTMLResponse:
         return HTMLResponse((TEMPLATE_ROOT / "index.html").read_text(encoding="utf-8"))
 
+    @app.get("/favicon.ico", include_in_schema=False)
+    def favicon() -> FileResponse:
+        return FileResponse(STATIC_ROOT / "brand" / "bee-favicon.png", media_type="image/png")
+
     @app.get("/api/scenes")
     def list_scenes() -> dict[str, Any]:
         scenes = find_scenes(resolved_base)
@@ -388,7 +421,7 @@ def get_configuration() -> dict[str, Any]:
 
     @app.get("/api/hardware")
     def get_hardware() -> dict[str, Any]:
-        return detect_hardware()
+        return hardware_provider()
 
     @app.get("/api/setup")
     def get_setup_status() -> dict[str, Any]:
@@ -519,7 +552,7 @@ async def discover_modal_endpoint(payload: ModalDiscoverRequest) -> dict[str, st
             raise HTTPException(status_code=400, detail="Enter a valid Modal token ID and token secret")
         try:
             discovered = await _discover_modal_endpoint(payload.token_id, payload.token_secret)
-            discovered["api_key"] = modal_remote_api_key(payload.token_id, payload.token_secret)
+            discovered["api_key"] = modal_workspace_api_key(payload.token_id, payload.token_secret)
             return discovered
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc