Add reusable local HTTP job service for Blender add-ons

- Implement JobClient for handling HTTP requests to the JobServer.
- Create ModalJobOperatorBase for managing Blender job lifecycle.
- Introduce JobContext and JobServer for job management and state handling.
- Add ServiceConnection and ServiceController for managing service processes.
- Implement job submission, cancellation, and resource management endpoints in the server.
- Ensure proper error handling and job state updates throughout the service.

tests/test_ben2_onnx.py

@@ -7,7 +7,12 @@ import numpy as np
 from PIL import Image
 
 
-SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform"
+SOURCE_DIR = (
+    Path(__file__).resolve().parents[1]
+    / "src"
+    / "snapform"
+    / "ai_service"
+)
 sys.path.insert(0, str(SOURCE_DIR))
 
 import ben2_onnx

tests/test_da3_onnx.py

@@ -8,7 +8,12 @@ import numpy as np
 from PIL import Image
 
 
-SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform"
+SOURCE_DIR = (
+    Path(__file__).resolve().parents[1]
+    / "src"
+    / "snapform"
+    / "ai_service"
+)
 sys.path.insert(0, str(SOURCE_DIR))
 
 import da3_onnx

tests/test_download_progress.py

@@ -13,7 +13,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 def load_worker():
     specification = importlib.util.spec_from_file_location(
         "snapform_download_worker",
-        PROJECT_ROOT / "src" / "snapform" / "worker.py",
+        PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py",
     )
     module = importlib.util.module_from_spec(specification)
     specification.loader.exec_module(module)

tests/test_image_depth.py

@@ -26,7 +26,11 @@ def load_image_depth():
     with patch.dict(sys.modules, {"bpy": fake_bpy}):
         specification = importlib.util.spec_from_file_location(
             "snapform_image_depth_test",
-            PROJECT_ROOT / "src" / "snapform" / "image_depth.py",
+            PROJECT_ROOT
+            / "src"
+            / "snapform"
+            / "blender"
+            / "depth_plane.py",
         )
         module = importlib.util.module_from_spec(specification)
         specification.loader.exec_module(module)

tests/test_moge_onnx.py

@@ -13,7 +13,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 def load_module():
     specification = importlib.util.spec_from_file_location(
         "snapform_moge_onnx",
-        PROJECT_ROOT / "src" / "snapform" / "moge_onnx.py",
+        PROJECT_ROOT / "src" / "snapform" / "ai_service" / "moge_onnx.py",
     )
     module = importlib.util.module_from_spec(specification)
     specification.loader.exec_module(module)

tests/test_pack.py

@@ -52,13 +52,14 @@ class PackTest(unittest.TestCase):
 
                 self.assertIn("__init__.py", names)
                 self.assertIn("blender_manifest.toml", names)
-                self.assertIn("image_depth.py", names)
+                self.assertIn("blender/depth_plane.py", names)
+                self.assertIn("blender/surface.py", names)
                 self.assertIn(
                     "assets/O_SnapForm.blend",
                     names,
                 )
-                self.assertIn("worker.py", names)
-                self.assertIn("service.py", names)
+                self.assertIn("ai_service/jobs.py", names)
+                self.assertIn("ai_service/server.py", names)
                 self.assertNotIn("snapform/__init__.py", names)
                 self.assertNotIn("importer.py", names)
                 self.assertFalse(

tests/test_service.py

@@ -13,7 +13,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 def load_service():
     specification = importlib.util.spec_from_file_location(
         "snapform_service",
-        PROJECT_ROOT / "src" / "snapform" / "service.py",
+        PROJECT_ROOT / "src" / "snapform" / "ai_service" / "server.py",
     )
     module = importlib.util.module_from_spec(specification)
     specification.loader.exec_module(module)
@@ -38,8 +38,15 @@ class ServiceTest(unittest.TestCase):
     def setUpClass(cls):
         cls.service = load_service()
 
+    def setUp(self):
+        self.jobs_temp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.jobs_temp.cleanup)
+        self.jobs_directory = Path(self.jobs_temp.name) / "jobs"
+
     def create_cache(self):
-        cache = self.service.ModelCache.__new__(self.service.ModelCache)
+        cache = self.service.ModelSessionCache.__new__(
+            self.service.ModelSessionCache
+        )
         cache.ben2_model = None
         cache.ben2_key = None
         cache.da3_model = None
@@ -63,7 +70,7 @@ class ServiceTest(unittest.TestCase):
 
         self.assertIs(first, second)
         self.assertEqual(loads, [("model", "cuda")])
-        self.assertEqual(cache.status()["ben2"], "model")
+        self.assertEqual(cache.snapshot()["ben2"], "model")
         self.assertEqual(cached_load_ms, 0.0)
 
     def test_da3_model_is_reused_for_matching_runtime_key(self):
@@ -86,7 +93,7 @@ class ServiceTest(unittest.TestCase):
 
         self.assertIs(first, second)
         self.assertEqual(loads, [("model", "cuda")])
-        self.assertEqual(cache.status()["da3"], "model")
+        self.assertEqual(cache.snapshot()["da3"], "model")
 
     def test_depth_model_cache_keeps_only_da3_or_moge(self):
         cache = self.create_cache()
@@ -118,141 +125,133 @@ class ServiceTest(unittest.TestCase):
         self.assertIsNone(cache.da3_model)
         self.assertIsNotNone(cache.ben2_model)
 
-    def test_service_tracks_clients_and_rejects_shared_shutdown(self):
+    def test_dedicated_service_runs_only_one_job_at_a_time(self):
         fake_models = SimpleNamespace(
-            status=lambda: {"ben2": "", "da3": ""},
+            snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
+            clear=lambda: None,
             close=lambda: None,
         )
-        with patch.object(self.service, "ModelCache", return_value=fake_models):
-            state = self.service.ServiceState("token", "instance")
+        with patch.object(self.service, "ModelManager", return_value=fake_models):
+            server = self.service.create_server(jobs_directory=self.jobs_directory)
 
-        state.register_client("first", "Blender A")
-        state.register_client("second", "Blender B")
+        parameters = {"model": "BASE"}
+        first = server.submit("debug-da3", parameters, job_id="job")
+        second = server.submit("debug-da3", parameters, job_id="other")
 
-        self.assertEqual(state.snapshot()["clients"], 2)
-        self.assertEqual(state.request_shutdown("first"), "shared")
-        self.assertTrue(
-            state.reserve("job", "first", "reconstruct", "status.json")
+        self.assertIsNotNone(first)
+        self.assertIsNone(second)
+        self.assertTrue(server.snapshot("instance")["busy"])
+        self.assertNotIn("clients", server.snapshot("instance"))
+
+    def test_service_registers_tasks_instead_of_model_names(self):
+        server = self.service.create_server(jobs_directory=self.jobs_directory)
+        self.assertEqual(
+            set(server.handlers),
+            {
+                "download-model",
+                "remove-background",
+                "extract-frames",
+                "generate-da3-depth",
+                "generate-moge2-depth",
+                "generate-da3-surface",
+                "generate-moge2-surface",
+                "debug-da3",
+                "debug-moge2",
+                "debug-ben2",
+            },
         )
-        self.assertFalse(
-            state.reserve("other", "second", "reconstruct", "other.json")
-        )
-
-    def test_service_accepts_moge_command(self):
-        self.assertIn("moge", self.service.JOB_COMMANDS)
 
     def test_cancelling_job_keeps_service_and_model_cache_alive(self):
         fake_models = SimpleNamespace(
-            status=lambda: {"ben2": "BEN2_BASE", "da3": ""},
+            snapshot=lambda: {
+                "loaded": {"ben2": "BEN2_BASE", "da3": ""}
+            },
+            clear=lambda: None,
             close=lambda: None,
         )
-        with patch.object(self.service, "ModelCache", return_value=fake_models):
-            state = self.service.ServiceState("token", "instance")
-        state.register_client("client", "Blender")
-        with tempfile.TemporaryDirectory() as directory:
-            status_path = Path(directory) / "status.json"
-            state.reserve(
-                "job",
-                "client",
-                "remove-background",
-                str(status_path),
-            )
-
-            active = state.cancel_job("job", "client")
-            self.assertTrue(active["cancel_requested"])
-            self.assertTrue(state.is_job_cancelled("job"))
-            state.run_job(
-                {
-                    "job_id": "job",
-                    "command": "remove-background",
-                    "parameters": {"status": str(status_path)},
-                }
-            )
-            status = self.service.json.loads(
-                status_path.read_text(encoding="utf-8")
-            )
+        with patch.object(self.service, "ModelManager", return_value=fake_models):
+            server = self.service.create_server(jobs_directory=self.jobs_directory)
+        context = server.submit("remove-background", {}, job_id="job")
+        self.assertTrue(context.request_cancel())
+        self.assertTrue(context.is_cancelled())
+        server.run(context)
+        status = context.snapshot()
 
         self.assertEqual(status["state"], "cancelled")
-        self.assertFalse(state.shutdown_event.is_set())
-        self.assertEqual(state.snapshot()["cached_models"]["ben2"], "BEN2_BASE")
-
-    def test_expired_clients_are_removed(self):
-        fake_models = SimpleNamespace(
-            status=lambda: {"ben2": "", "da3": ""},
-            close=lambda: None,
-        )
-        with patch.object(self.service, "ModelCache", return_value=fake_models):
-            state = self.service.ServiceState("token", "instance")
-        state.register_client("client", "Blender")
-        state.clients["client"]["last_seen"] -= (
-            self.service.CLIENT_TIMEOUT_SECONDS + 1.0
+        self.assertFalse(server.shutdown_event.is_set())
+        self.assertEqual(
+            server.snapshot("instance")["resources"]["models"]["loaded"]["ben2"],
+            "BEN2_BASE",
         )
 
-        state.expire_clients()
+    def test_job_progress_is_kept_in_service_memory(self):
+        server = self.service.JobServer("Test", jobs_directory=self.jobs_directory)
 
-        self.assertEqual(state.snapshot()["clients"], 0)
+        @server.job("example")
+        def example(context, _parameters):
+            context.update(0.5, "Half way")
 
-    def test_discovery_is_only_removed_by_matching_service(self):
-        with tempfile.TemporaryDirectory() as directory:
-            path = Path(directory) / "service.json"
-            args = SimpleNamespace(
-                discovery=str(path),
-                port=51234,
-                token="token",
-                instance_id="owner",
-            )
-            with patch.object(self.service.os, "getpid", return_value=42):
-                self.service.write_discovery(args)
+        context = server.submit("example", {})
+        server.run(context)
 
-            payload = self.service.json.loads(path.read_text(encoding="utf-8"))
-            self.assertEqual(payload["port"], 51234)
-            self.assertEqual(payload["pid"], 42)
-            self.service.remove_discovery(path, "other")
-            self.assertTrue(path.is_file())
-            self.service.remove_discovery(path, "owner")
-            self.assertFalse(path.exists())
+        self.assertRegex(context.job_id, r"^[0-9a-f]{32}$")
+        self.assertEqual(context.directory.name, context.job_id)
+        self.assertEqual(context.snapshot()["state"], "succeeded")
+        self.assertEqual(context.snapshot()["progress"], 1.0)
+
+    def test_job_server_owns_and_clears_shared_resources(self):
+        events = []
+        resource = SimpleNamespace(
+            snapshot=lambda: {"loaded": "model"},
+            clear=lambda: events.append("clear"),
+            close=lambda: events.append("close"),
+        )
+        server = self.service.JobServer("Test", jobs_directory=self.jobs_directory)
+        server.add_resource("models", resource)
+
+        @server.job("example")
+        def example(context, _parameters):
+            self.assertIs(context.resource("models"), resource)
+
+        context = server.submit("example", {}, job_id="job")
+        self.assertFalse(server.clear_resource("models"))
+        server.run(context)
+        self.assertTrue(server.clear_resource("models"))
+        server.close()
+
+        self.assertEqual(
+            server.snapshot("instance")["resources"]["models"]["loaded"],
+            "model",
+        )
+        self.assertEqual(events, ["clear", "close"])
 
     def test_service_dispatches_frame_extraction_without_models(self):
         fake_models = SimpleNamespace(
-            status=lambda: {"ben2": "", "da3": ""},
+            snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
+            clear=lambda: None,
             close=lambda: None,
         )
-        with patch.object(self.service, "ModelCache", return_value=fake_models):
-            state = self.service.ServiceState("token", "instance")
-        state.register_client("client", "Blender")
+        with patch.object(self.service, "ModelManager", return_value=fake_models):
+            server = self.service.create_server(jobs_directory=self.jobs_directory)
         with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)
             input_directory = root / "frames"
             input_directory.mkdir()
             (input_directory / "frame-0001.jpg").touch()
             (input_directory / "frame-0002.jpg").touch()
-            result_directory = root / "result"
-            status_path = root / "status.json"
-            state.reserve(
-                "job",
-                "client",
+            context = server.submit(
                 "extract-frames",
-                str(status_path),
-            )
-
-            state.run_job(
                 {
-                    "job_id": "job",
-                    "command": "extract-frames",
-                    "parameters": {
-                        "input": str(input_directory),
-                        "output_dir": str(result_directory),
-                        "status": str(status_path),
-                        "video_frames": 16,
-                        "start_frame": 1,
-                    },
-                }
-            )
-            status = self.service.json.loads(
-                status_path.read_text(encoding="utf-8")
+                    "input": str(input_directory),
+                    "video_frames": 16,
+                    "start_frame": 1,
+                },
+                job_id="job",
             )
+            server.run(context)
+            status = context.snapshot()
             self.assertTrue(
-                (result_directory / "frame.0001.jpg").is_file()
+                (context.directory / "frame.0001.jpg").is_file()
             )
 
         self.assertEqual(status["state"], "succeeded")

tests/test_worker.py

@@ -19,7 +19,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 def load_worker():
     specification = importlib.util.spec_from_file_location(
         "snapform_worker",
-        PROJECT_ROOT / "src" / "snapform" / "worker.py",
+        PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py",
     )
     module = importlib.util.module_from_spec(specification)
     specification.loader.exec_module(module)
@@ -500,18 +500,18 @@ class WorkerTest(unittest.TestCase):
     def test_models_from_r2_roundtrip_preserves_fields(self):
         root = PROJECT_ROOT / ".tmp-models-from-r2"
         model_dir = root / "models" / "DA3"
-        discovery = root / "service.json"
+        runtime_state = root / "runtime.json"
         try:
-            discovery.parent.mkdir(parents=True, exist_ok=True)
-            discovery.write_text(
-                json.dumps({"port": 1, "token": "t"}),
+            runtime_state.parent.mkdir(parents=True, exist_ok=True)
+            runtime_state.write_text(
+                json.dumps({"keep": "value"}),
                 encoding="utf-8",
             )
             self.assertFalse(self.worker._models_from_r2(model_dir))
             self.worker._mark_models_from_r2(model_dir)
             self.assertTrue(self.worker._models_from_r2(model_dir))
-            payload = json.loads(discovery.read_text(encoding="utf-8"))
-            self.assertEqual(payload["port"], 1)
+            payload = json.loads(runtime_state.read_text(encoding="utf-8"))
+            self.assertEqual(payload["keep"], "value")
             self.assertTrue(payload["models_from_r2"])
         finally:
             shutil.rmtree(root, ignore_errors=True)