DRAFT: Modly Next integration preview (5 pending upstream PRs)

api/tests/test_workflow_runs_lifecycle.py

@@ -200,5 +200,164 @@ async def run_cross_route_jobs():
         self.assertEqual(fake_registry.active_id, "demo/generate")
 
 
+class _FakeProc:
+    """Stands in for a loaded extension worker's subprocess."""
+
+    def __init__(self) -> None:
+        self.killed = False
+
+    def poll(self):
+        return -9 if self.killed else None
+
+    def kill(self) -> None:
+        self.killed = True
+
+
+class _FakeWorker:
+    """A loaded generator with a live subprocess, as _run_generation binds it."""
+
+    def __init__(self) -> None:
+        self._proc = _FakeProc()
+        self._loaded = True
+        self.outputs_dir = None
+
+    def generate(self, image_bytes, params, progress_cb, cancel_event=None):
+        output = Path(self.outputs_dir) / "result.glb"
+        output.write_bytes(b"glb")
+        return output
+
+
+class _FailingWorker(_FakeWorker):
+    def generate(self, image_bytes, params, progress_cb, cancel_event=None):
+        raise RuntimeError("inference failed")
+
+
+class _RegistryWithLoadedWorker:
+    """Keeps one worker loaded and active, the state the worker is left in
+    after a job ends normally."""
+
+    def __init__(self, worker: _FakeWorker) -> None:
+        self.worker = worker
+        self._generators = {"ext/model": worker}
+        self._active_id = "ext/model"
+
+    def assert_weight_variant_installed(self, params, model_id) -> None:
+        pass
+
+    def get_generator(self, model_id: str):
+        return self.worker
+
+    def model_status(self, model_id: str) -> dict:
+        return {"name": model_id, "downloaded": True, "loaded": True}
+
+    def activate_ready_generator(self, model_id: str):
+        return self.worker
+
+
+class CancelEndedJobTests(unittest.TestCase):
+    """cancel_job kills the subprocess bound to the job in _job_generators so
+    inference stops at once. Once the job has ended, that binding is gone and
+    the worker holds the warm model (or another job's generation), so a late
+    cancel of the ended job/run must leave the worker alone."""
+
+    def setUp(self) -> None:
+        self._prev_generation_registry = generation.generator_registry
+        self._prev_runs_registry = workflow_runs.generator_registry
+        self._prev_workspace = registry_module.WORKSPACE_DIR
+        self._tmp = tempfile.TemporaryDirectory()
+        registry_module.WORKSPACE_DIR = Path(self._tmp.name)
+        _clear_job_stores()
+
+    def tearDown(self) -> None:
+        generation.generator_registry = self._prev_generation_registry
+        workflow_runs.generator_registry = self._prev_runs_registry
+        registry_module.WORKSPACE_DIR = self._prev_workspace
+        self._tmp.cleanup()
+        _clear_job_stores()
+
+    def _use_worker(self, worker: _FakeWorker) -> _FakeWorker:
+        registry = _RegistryWithLoadedWorker(worker)
+        generation.generator_registry = registry
+        workflow_runs.generator_registry = registry
+        return worker
+
+    def _file_job(self, job_id: str, status: str = "pending") -> None:
+        generation._jobs[job_id] = JobStatus(job_id=job_id, status=status, progress=0)
+        generation._cancel_events[job_id] = threading.Event()
+
+    def _run_to_end(self, job_id: str) -> None:
+        self._file_job(job_id)
+        asyncio.run(generation._run_generation(job_id, b"img", {}, "Default", "mesh", "ext/model"))
+        # The job's generator binding is released as soon as it ends.
+        self.assertNotIn(job_id, generation._job_generators)
+
+    def _assert_worker_untouched(self, worker: _FakeWorker, proc: _FakeProc) -> None:
+        self.assertFalse(proc.killed)
+        self.assertIs(worker._proc, proc)
+        self.assertTrue(worker._loaded)
+
+    def _assert_worker_stopped(self, worker: _FakeWorker, proc: _FakeProc) -> None:
+        self.assertTrue(proc.killed)
+        self.assertIsNone(worker._proc)
+        self.assertFalse(worker._loaded)
+
+    def test_cancelling_a_finished_job_leaves_the_active_worker_alone(self) -> None:
+        worker = self._use_worker(_FakeWorker())
+        proc = worker._proc
+        self._run_to_end("finished-job")
+        self.assertEqual(generation._jobs["finished-job"].status, "done")
+
+        asyncio.run(generation.cancel_job("finished-job"))
+
+        self.assertEqual(generation._jobs["finished-job"].status, "done")
+        self._assert_worker_untouched(worker, proc)
+
+    def test_cancelling_a_failed_run_leaves_the_active_worker_alone(self) -> None:
+        worker = self._use_worker(_FailingWorker())
+        proc = worker._proc
+        self._run_to_end("failed-run")
+        self.assertEqual(generation._jobs["failed-run"].status, "error")
+
+        asyncio.run(workflow_runs.cancel_run("failed-run"))
+
+        self.assertEqual(generation._jobs["failed-run"].status, "error")
+        self._assert_worker_untouched(worker, proc)
+
+    def test_cancelling_an_ended_job_leaves_another_jobs_worker_alone(self) -> None:
+        worker = self._use_worker(_FakeWorker())
+        proc = worker._proc
+        self._run_to_end("finished-job")
+        # A later job now owns the same worker.
+        self._file_job("running-job", "running")
+        generation._job_generators["running-job"] = worker
+
+        asyncio.run(generation.cancel_job("finished-job"))
+
+        self.assertEqual(generation._jobs["running-job"].status, "running")
+        self._assert_worker_untouched(worker, proc)
+
+    def test_cancelling_a_running_job_still_stops_the_worker(self) -> None:
+        worker = self._use_worker(_FakeWorker())
+        proc = worker._proc
+        self._file_job("running-job", "running")
+        generation._job_generators["running-job"] = worker
+
+        asyncio.run(generation.cancel_job("running-job"))
+
+        self.assertEqual(generation._jobs["running-job"].status, "cancelled")
+        self._assert_worker_stopped(worker, proc)
+
+    def test_cancelling_a_running_run_still_stops_the_worker(self) -> None:
+        worker = self._use_worker(_FakeWorker())
+        proc = worker._proc
+        self._file_job("running-run", "running")
+        generation._job_generators["running-run"] = worker
+
+        asyncio.run(workflow_runs.cancel_run("running-run"))
+
+        self.assertEqual(generation._jobs["running-run"].status, "cancelled")
+        self._assert_worker_stopped(worker, proc)
+
+
 if __name__ == "__main__":
     unittest.main()

electron/main/artifact-registry-service.ts

@@ -1,4 +1,5 @@
 import { readdir, readFile, stat } from 'node:fs/promises'
+import type { Dirent } from 'node:fs'
 import { basename, extname, isAbsolute, join, relative, resolve, sep } from 'node:path'
 
 import type {
@@ -255,7 +256,7 @@ async function collectFiles(workspaceDir: string, rootName: typeof ALLOWED_ROOTS
   const files: string[] = []
 
   async function walk(dir: string): Promise<void> {
-    let entries: Awaited<ReturnType<typeof readdir>>
+    let entries: Dirent<string>[]
     try {
       entries = await readdir(dir, { withFileTypes: true })
     } catch {

electron/main/ipc-handlers.ts

@@ -42,6 +42,7 @@ import {
   weightStorageHasLocalData,
 } from './model-sources'
 import { getSettings, setSettings } from './settings-store'
+import { deleteStorageDirectory, moveStorageDirectory } from './storage-directory-guard'
 import { checkSetupNeeded, markSetupDone, runFullSetup, getVenvPythonExe, ensureSslPatch } from './python-setup'
 import { logger } from './logger'
 import { getProcessRunner, getPythonProcessRunner, getExtPythonExe, terminateProcessRunner, terminateAllProcessRunners } from './process-runner'
@@ -891,6 +892,29 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe
     return { total, used: total - free, available: free }
   })
 
+  // GPU memory (VRAM) — NVIDIA only via nvidia-smi
+  ipcMain.handle('system:gpuMemory', async () => {
+    if (process.platform === 'darwin' && process.arch === 'arm64') {
+      // Apple Silicon: unified memory, no separate VRAM
+      return null
+    }
+    try {
+      const { stdout } = await pExecFile('nvidia-smi', [
+        '--query-gpu=memory.total,memory.used',
+        '--format=csv,noheader,nounits'
+      ])
+      const line = stdout.trim().split('\n')[0].trim()
+      const [totalMiB, usedMiB] = line.split(',').map(s => parseInt(s.trim(), 10))
+      if (isNaN(totalMiB) || isNaN(usedMiB) || totalMiB === 0) return null
+      const total = totalMiB * 1024 * 1024
+      const used = usedMiB * 1024 * 1024
+      const available = total - used
+      return { total, used, available }
+    } catch {
+      return null
+    }
+  })
+
   ipcMain.handle('app:info', () => ({
     version:   app.getVersion(),
     userData:  app.getPath('userData'),
@@ -1091,31 +1115,18 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe
 
   ipcMain.handle('fs:moveDirectory', async (_, { src, dest }: { src: string; dest: string }) => {
     try {
-      await mkdir(dest, { recursive: true })
-      await cp(src, dest, { recursive: true })
-      await rmAsync(src, { recursive: true, force: true })
+      const userData = app.getPath('userData')
+      await moveStorageDirectory(src, dest, getSettings(userData), userData, { appDir: app.getAppPath() })
       return { success: true }
     } catch (err) {
       return { success: false, error: String(err) }
     }
   })
 
   ipcMain.handle('fs:deleteDirectory', async (_, dirPath: string) => {
-    const userData = app.getPath('userData')
-    const settings = getSettings(userData)
-    const allowedRoots = [
-      settings.modelsDir,
-      settings.workspaceDir,
-      settings.extensionsDir,
-      join(userData, 'gen-cache'),
-    ]
-    const resolved = join(dirPath)
-    const isAllowed = allowedRoots.some((root) => resolved.startsWith(root))
-    if (!isAllowed) {
-      return { success: false, error: 'Path is outside allowed directories' }
-    }
     try {
-      await rmAsync(resolved, { recursive: true, force: true })
+      const userData = app.getPath('userData')
+      await deleteStorageDirectory(dirPath, getSettings(userData), userData, { appDir: app.getAppPath() })
       return { success: true }
     } catch (err) {
       return { success: false, error: String(err) }