Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/tools/test_image_generation_artifacts.py

@@ -0,0 +1,149 @@
+import concurrent.futures
+import json
+import threading
+from types import SimpleNamespace
+
+
+def test_postprocess_adds_agent_visible_image_for_active_ssh_env(monkeypatch, tmp_path):
+    from tools import image_generation_tool
+
+    hermes_home = tmp_path / ".hermes"
+    image_dir = hermes_home / "cache" / "images"
+    image_dir.mkdir(parents=True)
+    image_path = image_dir / "xai_grok-imagine-image_test.jpg"
+    image_path.write_bytes(b"jpg")
+
+    sync_calls = []
+
+    class FakeSyncManager:
+        def sync(self, *, force=False):
+            sync_calls.append(force)
+
+    env = SimpleNamespace(
+        _remote_home="/home/remotesshuser",
+        _sync_manager=FakeSyncManager(),
+    )
+
+    monkeypatch.setenv("HERMES_HOME", str(hermes_home))
+    monkeypatch.setattr(image_generation_tool, "_active_terminal_env", lambda task_id: env)
+
+    raw = json.dumps({"success": True, "image": str(image_path)})
+    result = json.loads(
+        image_generation_tool._postprocess_image_generate_result(raw, task_id="task-1")
+    )
+
+    assert result["image"] == str(image_path)
+    assert result["host_image"] == str(image_path)
+    assert result["agent_visible_image"] == (
+        "/home/remotesshuser/.hermes/cache/images/xai_grok-imagine-image_test.jpg"
+    )
+    assert sync_calls == [True]
+
+
+def test_concurrent_image_results_preserve_shared_remote_sync_state(monkeypatch, tmp_path):
+    from tools import image_generation_tool
+    from tools.environments import file_sync
+
+    hermes_home = tmp_path / ".hermes"
+    image_dir = hermes_home / "cache" / "images"
+    image_dir.mkdir(parents=True)
+    first_image = image_dir / "first.png"
+    second_image = image_dir / "second.png"
+    first_image.write_bytes(b"first")
+
+    first_upload_started = threading.Event()
+    second_sync_finished = threading.Event()
+    worker = threading.local()
+
+    def get_files():
+        return [
+            (
+                str(path),
+                f"/home/remote/.hermes/cache/images/{path.name}",
+            )
+            for path in sorted(image_dir.iterdir())
+        ]
+
+    def upload(_host_path, _remote_path):
+        if worker.label == "first":
+            first_upload_started.set()
+            # Without transaction serialization, the second sync commits its
+            # newer snapshot before this first sync resumes and overwrites it.
+            second_sync_finished.wait(timeout=1.0)
+
+    sync_manager = file_sync.FileSyncManager(
+        get_files_fn=get_files,
+        upload_fn=upload,
+        delete_fn=lambda _paths: None,
+    )
+    env = SimpleNamespace(
+        _remote_home="/home/remote",
+        _sync_manager=sync_manager,
+    )
+
+    monkeypatch.setenv("HERMES_HOME", str(hermes_home))
+    monkeypatch.setattr(file_sync, "_credential_host_paths", lambda: set())
+    monkeypatch.setattr(image_generation_tool, "_active_terminal_env", lambda _task_id: env)
+
+    def postprocess(label, image_path):
+        worker.label = label
+        try:
+            raw = json.dumps({"success": True, "image": str(image_path)})
+            return image_generation_tool._postprocess_image_generate_result(
+                raw,
+                task_id="shared-task",
+            )
+        finally:
+            if label == "second":
+                second_sync_finished.set()
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        first_future = executor.submit(postprocess, "first", first_image)
+        assert first_upload_started.wait(timeout=2.0)
+
+        second_image.write_bytes(b"second")
+        second_future = executor.submit(postprocess, "second", second_image)
+
+        first_future.result(timeout=3.0)
+        second_future.result(timeout=3.0)
+
+    assert set(sync_manager._synced_files) == {
+        "/home/remote/.hermes/cache/images/first.png",
+        "/home/remote/.hermes/cache/images/second.png",
+    }
+
+
+def test_handle_image_generate_postprocesses_plugin_result(monkeypatch, tmp_path):
+    from tools import image_generation_tool
+
+    hermes_home = tmp_path / ".hermes"
+    image_dir = hermes_home / "cache" / "images"
+    image_dir.mkdir(parents=True)
+    image_path = image_dir / "plugin.png"
+    image_path.write_bytes(b"png")
+
+    env = SimpleNamespace(_remote_home="/home/remote", _sync_manager=None)
+
+    seen_task_ids = []
+
+    def fake_active_env(task_id):
+        seen_task_ids.append(task_id)
+        return env
+
+    monkeypatch.setenv("HERMES_HOME", str(hermes_home))
+    monkeypatch.setattr(image_generation_tool, "_active_terminal_env", fake_active_env)
+    monkeypatch.setattr(
+        image_generation_tool,
+        "_dispatch_to_plugin_provider",
+        lambda prompt, aspect_ratio, **kw: json.dumps({"success": True, "image": str(image_path)}),
+    )
+
+    result = json.loads(
+        image_generation_tool._handle_image_generate(
+            {"prompt": "draw", "aspect_ratio": "square"},
+            task_id="plugin-task",
+        )
+    )
+
+    assert seen_task_ids == ["plugin-task"]
+    assert result["agent_visible_image"] == "/home/remote/.hermes/cache/images/plugin.png"