nous portal video gen (#74963)

tests/tools/test_flux3_video_tool.py

@@ -0,0 +1,568 @@
+"""Native BFL FLUX 3 tools: gating, transport, media delivery, redaction."""
+
+import asyncio
+import base64
+import json
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from tools import flux3_video_tool as flux3
+
+GATEWAY = "https://tool-gateway.example.com"
+BASE_URL = f"{GATEWAY}/api/bfl"
+UPLOAD_PATH = "/api/uploads/bfl"
+
+_PNG = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
+)
+
+
+@pytest.fixture(autouse=True)
+def _endpoints():
+    """Every test runs as if the mount is reachable unless it says otherwise."""
+    with patch.object(
+        flux3,
+        "managed_vendor_endpoints",
+        return_value={"origin": GATEWAY, "base_url": BASE_URL, "upload_path": UPLOAD_PATH},
+    ):
+        yield
+
+
+class _FakeResponse:
+    def __init__(self, status_code=200, payload=None, text=""):
+        self.status_code = status_code
+        self._payload = payload
+        self.text = text or (json.dumps(payload) if payload is not None else "")
+
+    def json(self):
+        if self._payload is None:
+            raise ValueError("no json")
+        return self._payload
+
+
+class _FakeClient:
+    """Captures the one request each handler makes."""
+
+    def __init__(self, response, sink):
+        self._response = response
+        self._sink = sink
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *_exc):
+        return False
+
+    async def request(self, method, url, headers=None, json=None):
+        self._sink.append({"method": method, "url": url, "headers": headers or {}, "json": json})
+        if isinstance(self._response, Exception):
+            raise self._response
+        return self._response
+
+
+class _FakeStream:
+    """A streaming GET that yields `body` in one chunk."""
+
+    def __init__(self, body, status_code=200):
+        self._body = body
+        self.status_code = status_code
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *_exc):
+        return False
+
+    def raise_for_status(self):
+        if self.status_code >= 400:
+            raise RuntimeError(f"HTTP {self.status_code}")
+
+    async def aiter_bytes(self):
+        yield self._body
+
+
+@contextmanager
+def _fake_download(body, status_code=200):
+    """Stub the clip download; yields the list of URLs that were fetched.
+
+    Patched at `create_ssrf_safe_async_client` rather than at httpx, which both
+    stubs the transport and asserts the download goes through the SSRF-guarded
+    client — the URL is vendor-supplied and fetched from the user's machine.
+    """
+    from tools import url_safety
+
+    fetched = []
+
+    class _Client:
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *_exc):
+            return False
+
+        def stream(self, _method, url):
+            fetched.append(url)
+            return _FakeStream(body, status_code)
+
+    with patch.object(url_safety, "create_ssrf_safe_async_client", lambda **_kw: _Client()):
+        yield fetched
+
+
+def _run(coro):
+    return asyncio.run(coro)
+
+
+def _call(handler, args, response, headers=None):
+    """Invoke a handler with the transport stubbed; returns (parsed, requests)."""
+    sink = []
+    import httpx
+
+    with patch.object(
+        flux3,
+        "managed_gateway_auth_headers",
+        return_value=headers if headers is not None else {"Authorization": "Bearer nous-token"},
+    ), patch.object(httpx, "AsyncClient", lambda **_kw: _FakeClient(response, sink)):
+        raw = _run(handler(args))
+    return json.loads(raw), sink
+
+
+class TestGating:
+    def test_hidden_without_a_reachable_mount(self):
+        with patch.object(flux3, "managed_vendor_endpoints", return_value=None):
+            assert flux3.check_bfl_requirements() is False
+
+    def test_hidden_without_paid_service_access(self):
+        # The free tool pool does not fund BFL, so a pool-only user must never
+        # see the tools rather than see them and be refused.
+        account = SimpleNamespace(logged_in=True, paid_service_access=False, tool_gateway_entitled=True)
+        with patch("hermes_cli.nous_account.get_nous_portal_account_info", return_value=account):
+            assert flux3.check_bfl_requirements() is False
+
+    def test_hidden_when_logged_out(self):
+        account = SimpleNamespace(logged_in=False, paid_service_access=False)
+        with patch("hermes_cli.nous_account.get_nous_portal_account_info", return_value=account):
+            assert flux3.check_bfl_requirements() is False
+
+    def test_visible_for_a_paid_portal_account(self):
+        account = SimpleNamespace(logged_in=True, paid_service_access=True)
+        with patch("hermes_cli.nous_account.get_nous_portal_account_info", return_value=account):
+            assert flux3.check_bfl_requirements() is True
+
+    def test_fails_closed_when_the_account_probe_raises(self):
+        with patch("hermes_cli.nous_account.get_nous_portal_account_info", side_effect=RuntimeError("portal down")):
+            assert flux3.check_bfl_requirements() is False
+
+
+class TestSubmitTransport:
+    def test_text_to_video_posts_the_mode_and_arguments(self):
+        response = _FakeResponse(200, {"id": "bfl_job_1", "status": "submitted", "guidance": "Poll bfl_flux3_get_result with id=bfl_job_1"})
+
+        parsed, requests = _call(
+            flux3._handle_text_to_video,
+            {"prompt": "a lake", "aspect_ratio": "16:9", "duration": 5},
+            response,
+        )
+
+        assert requests[0]["method"] == "POST"
+        assert requests[0]["url"] == f"{BASE_URL}/generations"
+        assert requests[0]["json"] == {
+            "prompt": "a lake",
+            "aspect_ratio": "16:9",
+            "duration": 5,
+            "mode": "text_to_video",
+        }
+        assert requests[0]["headers"]["Authorization"] == "Bearer nous-token"
+        # The gateway's guidance is the model-facing text, verbatim.
+        assert parsed["result"] == "Poll bfl_flux3_get_result with id=bfl_job_1"
+        assert parsed["details"]["id"] == "bfl_job_1"
+
+    def test_each_generate_tool_sends_its_own_mode(self):
+        for handler, args, mode in [
+            (flux3._handle_text_to_video, {"prompt": "a"}, "text_to_video"),
+            (flux3._handle_image_to_video, {"prompt": "a", "input_image": "https://x/a.png"}, "image_to_video"),
+            (
+                flux3._handle_keyframes_to_video,
+                {"prompt": "a", "input_images": ["https://x/a.png"], "keyframe_indices": [0]},
+                "keyframes_to_video",
+            ),
+            (flux3._handle_video_continuation, {"prompt": "a", "input_video": "https://x/c.mp4"}, "video_continuation"),
+        ]:
+            _parsed, requests = _call(handler, args, _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}))
+            assert requests[0]["json"]["mode"] == mode
+
+    def test_urls_pass_through_without_an_upload(self):
+        # Forwarding a URL is cheaper than downloading and re-uploading it.
+        _parsed, requests = _call(
+            flux3._handle_image_to_video,
+            {"prompt": "a", "input_image": "https://example.com/a.png"},
+            _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}),
+        )
+        assert requests[0]["json"]["input_image"] == "https://example.com/a.png"
+
+    def test_a_refusal_is_surfaced_as_the_tools_result_text(self):
+        # Throttles are designed to be hit: the message is written for the
+        # model and must reach it intact, with the machine detail alongside.
+        response = _FakeResponse(
+            429,
+            {
+                "error": {
+                    "code": "BFL_GENERATION_COOLDOWN",
+                    "message": "A new BFL video generation may be started once every 5 minutes. Wait 210 seconds.",
+                    "details": {"retryAfterSeconds": 210},
+                }
+            },
+        )
+
+        parsed, _requests = _call(flux3._handle_text_to_video, {"prompt": "a"}, response)
+
+        assert parsed["error"] == "A new BFL video generation may be started once every 5 minutes. Wait 210 seconds."
+        assert parsed["details"] == {"retryAfterSeconds": 210}
+
+    def test_a_401_asks_for_a_nous_sign_in(self):
+        parsed, _requests = _call(flux3._handle_text_to_video, {"prompt": "a"}, _FakeResponse(401, {"error": {"code": "AUTH_ERROR"}}))
+
+        assert parsed["needs_reauth"] is True
+        assert "sign in" in parsed["error"].lower()
+
+    def test_missing_credentials_ask_for_a_sign_in_without_calling_out(self):
+        parsed, requests = _call(flux3._handle_text_to_video, {"prompt": "a"}, _FakeResponse(200, {}), headers={})
+
+        assert requests == []
+        assert "sign in" in parsed["error"].lower()
+
+    def test_a_transport_failure_reports_the_cause(self):
+        parsed, _requests = _call(
+            flux3._handle_text_to_video,
+            {"prompt": "a"},
+            RuntimeError("connect failed"),
+        )
+
+        assert "Could not reach the video-generation gateway" in parsed["error"]
+        assert "connect failed" in parsed["error"]
+
+    def test_an_unreadable_body_does_not_masquerade_as_success(self):
+        parsed, _requests = _call(flux3._handle_text_to_video, {"prompt": "a"}, _FakeResponse(502, None, text="upstream exploded"))
+
+        assert "error" in parsed
+
+
+class TestPollTransport:
+    def test_poll_gets_the_job_and_returns_guidance(self):
+        response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."})
+
+        parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, response)
+
+        assert requests[0]["method"] == "GET"
+        assert requests[0]["url"] == f"{BASE_URL}/generations/bfl_job_1"
+        assert requests[0]["json"] is None
+        assert parsed["result"] == "Still going."
+        assert parsed["details"]["status"] == "Generating"
+
+    def test_ready_saves_the_clip_and_never_returns_the_signed_url(self, tmp_path):
+        # The signed URL is a bearer credential for the clip and it used to be
+        # re-keyed into a shell command by hand, dropping characters. Neither
+        # can happen if the model never sees it.
+        signed = "https://cdn.example/container/flux3-clip.mp4?sig=abc%2Bdef%3D&se=2026"
+        response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {"sample": signed}, "guidance": "Deliver the saved file."})
+
+        with _fake_download(b"x" * (128 * 1024)) as fetched:
+            parsed, _requests = _call(
+                flux3._handle_get_result,
+                {"id": "bfl_job_1", "save_to": str(tmp_path)},
+                response,
+            )
+
+        saved = tmp_path / "flux3-clip.mp4"
+        assert saved.read_bytes() == b"x" * (128 * 1024)
+        assert parsed["details"]["saved_path"] == str(saved)
+        assert parsed["details"]["result"].get("sample") is None
+        assert signed not in json.dumps(parsed)
+        # The gateway still owns the delivery wording; the client only supplies
+        # the path it cannot know.
+        assert parsed["result"].startswith(f"Saved to {saved}.")
+        assert "Deliver the saved file." in parsed["result"]
+        assert fetched == [signed]
+
+    def test_ready_never_overwrites_an_existing_file(self, tmp_path):
+        (tmp_path / "flux3-clip.mp4").write_bytes(b"an earlier clip")
+        response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {"sample": "https://cdn.example/x/flux3-clip.mp4?sig=a"}, "guidance": "g"})
+
+        with _fake_download(b"y" * (128 * 1024)):
+            parsed, _requests = _call(flux3._handle_get_result, {"id": "bfl_job_1", "save_to": str(tmp_path)}, response)
+
+        assert parsed["details"]["saved_path"] == str(tmp_path / "flux3-clip-2.mp4")
+        assert (tmp_path / "flux3-clip.mp4").read_bytes() == b"an earlier clip"
+
+    def test_a_rejected_download_fails_loudly_and_leaves_no_file(self, tmp_path):
+        # The original bug: a bad signature returns an XML error body, curl
+        # writes it to the .mp4 and exits 0, and it reads as success. A short
+        # body is not a video whatever the status code said.
+        response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Ready", "result": {"sample": "https://cdn.example/x/flux3-clip.mp4?sig=bad"}, "guidance": "g"})
+
+        with _fake_download(b"<?xml version='1.0'?><Error>AuthenticationFailed</Error>"):
+            parsed, _requests = _call(flux3._handle_get_result, {"id": "bfl_job_1", "save_to": str(tmp_path)}, response)
+
+        assert "saving it failed" in parsed["result"]
+        assert "Poll this job again" in parsed["result"]
+        # Neither a half-written .part nor a plausible-looking .mp4 survives.
+        assert [p.name for p in tmp_path.glob("*.mp4*")] == []
+        assert "saved_path" not in parsed["details"]
+
+    def test_poll_requires_an_id(self):
+        parsed, requests = _call(flux3._handle_get_result, {}, _FakeResponse(200, {}))
+
+        assert "id is required" in parsed["error"]
+        assert requests == []
+
+    def test_poll_url_encodes_the_job_id(self):
+        _parsed, requests = _call(
+            flux3._handle_get_result,
+            {"id": "weird/../id"},
+            _FakeResponse(200, {"id": "x", "guidance": "ok"}),
+        )
+        assert requests[0]["url"] == f"{BASE_URL}/generations/weird%2F..%2Fid"
+
+
+class TestMediaDelivery:
+    def _resolved(self, mime="image/png", data=_PNG):
+        return SimpleNamespace(data=data, mime=mime)
+
+    def test_a_local_path_is_uploaded_and_replaced_with_a_reference(self):
+        async def fake_uploader(data, mime):
+            assert data == _PNG
+            assert mime == "image/png"
+            return "nous-upload:token-1"
+
+        with patch.object(flux3, "build_managed_media_uploader", return_value=fake_uploader), patch(
+            "tools.image_source.resolve_image_source", return_value=self._resolved()
+        ) as resolve:
+            _parsed, requests = _call(
+                flux3._handle_image_to_video,
+                {"prompt": "a", "input_image": "/tmp/frame.png"},
+                _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}),
+            )
+
+        assert requests[0]["json"]["input_image"] == "nous-upload:token-1"
+        # Images and video ride the same safety pipeline; only the permitted
+        # type differs, and an image field must not accept a video.
+        assert resolve.call_args.kwargs["permitted"] == ("image",)
+
+    def test_video_fields_permit_video_only(self):
+        async def fake_uploader(data, mime):
+            return "nous-upload:token-v"
+
+        with patch.object(flux3, "build_managed_media_uploader", return_value=fake_uploader), patch(
+            "tools.image_source.resolve_image_source", return_value=self._resolved("video/mp4", b"\x00\x00\x00\x18ftypmp42")
+        ) as resolve:
+            _parsed, requests = _call(
+                flux3._handle_video_continuation,
+                {"prompt": "a", "input_video": "/tmp/clip.mp4"},
+                _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}),
+            )
+
+        assert requests[0]["json"]["input_video"] == "nous-upload:token-v"
+        assert resolve.call_args.kwargs["permitted"] == ("video",)
+
+    def test_every_keyframe_path_is_uploaded(self):
+        uploads = []
+
+        async def fake_uploader(data, mime):
+            uploads.append(mime)
+            return f"nous-upload:token-{len(uploads)}"
+
+        with patch.object(flux3, "build_managed_media_uploader", return_value=fake_uploader), patch(
+            "tools.image_source.resolve_image_source", return_value=self._resolved()
+        ):
+            _parsed, requests = _call(
+                flux3._handle_keyframes_to_video,
+                {"prompt": "a", "input_images": ["/tmp/a.png", "https://x/b.png", "/tmp/c.png"], "keyframe_indices": [0, 24, 48]},
+                _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}),
+            )
+
+        # The URL in the middle is forwarded untouched.
+        assert requests[0]["json"]["input_images"] == [
+            "nous-upload:token-1",
+            "https://x/b.png",
+            "nous-upload:token-2",
+        ]
+
+    def test_a_list_valued_input_image_is_still_uploaded(self):
+        # The gateway accepts input_image as a string OR a list, so a list of
+        # local paths must not slip past unsanitized — that would send raw
+        # filesystem paths to the vendor and disclose the user's directories.
+        async def fake_uploader(data, mime):
+            return "nous-upload:token-1"
+
+        with patch.object(flux3, "build_managed_media_uploader", return_value=fake_uploader), patch(
+            "tools.image_source.resolve_image_source", return_value=self._resolved()
+        ):
+            _parsed, requests = _call(
+                flux3._handle_image_to_video,
+                {"prompt": "a", "input_image": ["/tmp/frame.png"]},
+                _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}),
+            )
+
+        assert requests[0]["json"]["input_image"] == ["nous-upload:token-1"]
+        assert "/tmp/frame.png" not in json.dumps(requests[0]["json"])
+
+    def test_media_fields_are_sanitized_whatever_the_mode_expects(self):
+        # The gateway prefers input_image over input_images, so sanitizing only
+        # the field this mode documents would let the other one through.
+        uploads = []
+
+        async def fake_uploader(data, mime):
+            uploads.append(mime)
+            return f"nous-upload:token-{len(uploads)}"
+
+        with patch.object(flux3, "build_managed_media_uploader", return_value=fake_uploader), patch(
+            "tools.image_source.resolve_image_source", return_value=self._resolved()
+        ):
+            _parsed, requests = _call(
+                flux3._handle_keyframes_to_video,
+                {
+                    "prompt": "a",
+                    "input_images": ["https://x/b.png"],
+                    "input_image": "/tmp/sneaky.png",
+                    "keyframe_indices": [0],
+                },
+                _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}),
+            )
+
+        body = json.dumps(requests[0]["json"])
+        assert "/tmp/sneaky.png" not in body
+        assert requests[0]["json"]["input_image"] == "nous-upload:token-1"
+
+    def test_text_to_video_strips_media_fields_instead_of_uploading_them(self):
+        # The mode takes no media, so an upload would spend the caller's quota
+        # on a value the gateway ignores.
+        def _must_not_upload(*_args, **_kwargs):
+            raise AssertionError("text-to-video must not upload anything")
+
+        with patch.object(flux3, "build_managed_media_uploader", _must_not_upload):
+            _parsed, requests = _call(
+                flux3._handle_text_to_video,
+                {"prompt": "a", "input_image": "/tmp/frame.png"},
+                _FakeResponse(200, {"id": "j", "status": "submitted", "guidance": "ok"}),
+            )
+
+        assert "input_image" not in requests[0]["json"]
+        assert requests[0]["json"]["mode"] == "text_to_video"
+
+    def test_an_over_long_image_list_is_refused_before_any_upload(self):
+        def _must_not_upload(*_args, **_kwargs):
+            raise AssertionError("an over-long list must be refused before uploading")
+
+        with patch.object(flux3, "build_managed_media_uploader", _must_not_upload):
+            parsed, requests = _call(
+                flux3._handle_keyframes_to_video,
+                {"prompt": "a", "input_images": [f"/tmp/{i}.png" for i in range(11)], "keyframe_indices": [0]},
+                _FakeResponse(200, {}),
+            )
+
+        assert "at most 10" in parsed["error"]
+        assert requests == []
+
+    def test_an_upload_refusal_becomes_the_tools_error(self):
+        async def failing_uploader(data, mime):
+            raise RuntimeError("the daily upload budget for this account is exhausted")
+
+        with patch.object(flux3, "build_managed_media_uploader", return_value=failing_uploader), patch(
+            "tools.image_source.resolve_image_source", return_value=self._resolved()
+        ):
+            parsed, requests = _call(
+                flux3._handle_image_to_video,
+                {"prompt": "a", "input_image": "/tmp/frame.png"},
+                _FakeResponse(200, {}),
+            )
+
+        assert "daily upload budget" in parsed["error"]
+        # A failed upload must not reach the gateway as a bare local path.
+        assert requests == []
+
+    def test_an_unreadable_file_is_reported_without_dumping_the_value(self):
+        from tools.image_source import SourceNotFound
+
+        with patch.object(flux3, "build_managed_media_uploader", return_value=lambda *a: None), patch(
+            "tools.image_source.resolve_image_source", side_effect=SourceNotFound("media file not found", src="/tmp/x.png")
+        ):
+            parsed, requests = _call(
+                flux3._handle_image_to_video,
+                # Long, but unmistakably a path (dots and dashes are outside
+                # the base64 alphabet, so the payload guard leaves it alone).
+                {"prompt": "a", "input_image": "/tmp/" + "a-b." * 2000 + "frame.png"},
+                _FakeResponse(200, {}),
+            )
+
+        assert "error" in parsed
+        # The offending value is truncated: echoing it whole would blow up the
+        # model's context on the way to reporting a bad path.
+        assert len(parsed["error"]) < 500
+        assert requests == []
+
+
+class TestLocalPathDetection:
+    @pytest.mark.parametrize(
+        "value",
+        ["/tmp/frame.png", "~/Pictures/f.png", "./f.png", "../f.png", "file:///tmp/f.png", r"C:\Users\me\f.png", r"\\nas\share\f.png"],
+    )
+    def test_rooted_paths_are_read_off_disk(self, value):
+        assert flux3._looks_like_local_path(value) is True
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            "frame.png",
+            "https://example.com/f.png",
+            "nous-upload:eyJhbGciOiJIUzI1NiJ9.e30.sig",
+            "C:frame.png",
+            # Inline base64 of a JPEG always starts "/9j/" (first byte 0xFF),
+            # which must not read as an absolute POSIX path.
+            "/9j/4AAQSkZJRgABAQAAAQ" + "A" * 300 + "==",
+        ],
+    )
+    def test_ambiguous_and_remote_values_are_forwarded(self, value):
+        assert flux3._looks_like_local_path(value) is False
+
+    def test_a_short_base64_lookalike_path_is_still_a_path(self):
+        assert flux3._looks_like_local_path("/tmp/frames/a1") is True
+
+
+class TestSchemas:
+    def test_every_tool_is_registered_under_the_bfl_toolset(self):
+        from tools.registry import registry
+
+        for name in [
+            "bfl_flux3_text_to_video",
+            "bfl_flux3_image_to_video",
+            "bfl_flux3_keyframes_to_video",
+            "bfl_flux3_video_continuation",
+            "bfl_flux3_get_result",
+            "bfl_flux3_prompting_guide",
+        ]:
+            entry = registry.get_entry(name)
+            assert entry is not None, f"{name} is not registered"
+            assert entry.toolset == "bfl"
+            assert entry.check_fn is flux3.check_bfl_requirements
+
+    def test_generate_tools_point_at_the_guide_and_the_poll_tool(self):
+        # Descriptions are the only text guaranteed to be in context when a
+        # model picks a tool, so the pointers live there.
+        for schema in [flux3.TEXT_TO_VIDEO_SCHEMA, flux3.IMAGE_TO_VIDEO_SCHEMA, flux3.KEYFRAMES_TO_VIDEO_SCHEMA, flux3.VIDEO_CONTINUATION_SCHEMA]:
+            assert "bfl_flux3_prompting_guide" in schema["description"]
+            assert "bfl_flux3_get_result" in schema["description"]
+
+    def test_the_guide_covers_the_methodology_without_pinning_server_policy(self):
+        guide = flux3.FLUX3_PROMPTING_GUIDE
+        assert "grounding" in guide.lower()
+        assert "bfl_flux3_get_result" in guide
+        # Waits and limits ship live in the gateway's responses; pinning them
+        # here would let the client lie about what the server enforces.
+        assert "5 minutes" not in guide
+        assert "per minute" not in guide
+
+    def test_the_guide_tool_takes_no_arguments_and_calls_nothing(self):
+        assert flux3.PROMPTING_GUIDE_SCHEMA["parameters"]["properties"] == {}
+        assert _run(flux3._handle_prompting_guide({})) == flux3.FLUX3_PROMPTING_GUIDE

tests/tools/test_managed_tool_gateway.py

@@ -6,6 +6,8 @@ from pathlib import Path
 import sys
 from unittest.mock import patch
 
+import pytest
+
 MODULE_PATH = Path(__file__).resolve().parents[2] / "tools" / "managed_tool_gateway.py"
 MODULE_SPEC = spec_from_file_location("managed_tool_gateway_test_module", MODULE_PATH)
 assert MODULE_SPEC and MODULE_SPEC.loader
@@ -35,6 +37,147 @@ def test_resolve_managed_tool_gateway_derives_vendor_origin_from_shared_domain()
     assert result.managed_mode is True
 
 
+def test_resolve_managed_tool_gateway_uses_vendor_specific_override():
+    with patch.dict(
+        os.environ,
+        {
+            "BROWSER_USE_GATEWAY_URL": "http://browser-use-gateway.localhost:3009/",
+        },
+        clear=False,
+    ), patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True):
+        result = resolve_managed_tool_gateway(
+            "browser-use",
+            token_reader=lambda: "nous-token",
+        )
+
+    assert result is not None
+    assert result.gateway_origin == "http://browser-use-gateway.localhost:3009"
+
+
+def test_resolve_managed_tool_gateway_is_inactive_without_nous_token():
+    with patch.dict(
+        os.environ,
+        {
+            "TOOL_GATEWAY_DOMAIN": "nousresearch.com",
+        },
+        clear=False,
+    ), patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True):
+        result = resolve_managed_tool_gateway(
+            "firecrawl",
+            token_reader=lambda: None,
+        )
+
+    assert result is None
+
+
+def test_resolve_managed_tool_gateway_is_disabled_without_subscription():
+    with patch.dict(os.environ, {"TOOL_GATEWAY_DOMAIN": "nousresearch.com"}, clear=False), \
+         patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=False):
+        result = resolve_managed_tool_gateway(
+            "firecrawl",
+            token_reader=lambda: "nous-token",
+        )
+
+    assert result is None
+
+
+def test_read_nous_access_token_refreshes_expiring_cached_token(tmp_path, monkeypatch):
+    monkeypatch.delenv("TOOL_GATEWAY_USER_TOKEN", raising=False)
+    monkeypatch.setenv("HERMES_HOME", str(tmp_path))
+    expires_at = (datetime.now(timezone.utc) + timedelta(seconds=30)).isoformat()
+    (tmp_path / "auth.json").write_text(json.dumps({
+        "providers": {
+            "nous": {
+                "access_token": "stale-token",
+                "refresh_token": "refresh-token",
+                "expires_at": expires_at,
+            }
+        }
+    }))
+    monkeypatch.setattr(
+        "hermes_cli.auth.resolve_nous_access_token",
+        lambda refresh_skew_seconds=120: "fresh-token",
+    )
+
+    assert managed_tool_gateway.read_nous_access_token() == "fresh-token"
+
+
+def test_managed_vendor_endpoints_pin_the_deployed_gateway_url():
+    """The exact URL an agent may connect to is a code fact, not a lookup.
+
+    Exercises the real ``build_vendor_gateway_url`` (which once resolved a
+    typo'd pseudo-vendor to a non-existent host while every other test stubbed
+    it): default builder, real deployed host, pinned vendor path.
+    """
+    with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True), \
+         patch.dict(
+             os.environ,
+             {"TOOL_GATEWAY_DOMAIN": "nousresearch.com", "TOOL_GATEWAY_SCHEME": "https"},
+             clear=False,
+         ):
+        os.environ.pop("TOOL_GATEWAY_URL", None)
+        endpoints = managed_tool_gateway.managed_vendor_endpoints("bfl")
+
+    assert endpoints == {
+        "origin": "https://tool-gateway.nousresearch.com",
+        "base_url": "https://tool-gateway.nousresearch.com/api/bfl",
+        "upload_path": "/api/uploads/bfl",
+    }
+
+
+def test_managed_vendor_endpoints_unreachable_when_managed_tools_disabled():
+    with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=False):
+        assert managed_tool_gateway.managed_vendor_endpoints("bfl") is None
+
+
+def test_managed_gateway_auth_headers_carry_the_bearer():
+    with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True):
+        headers = managed_tool_gateway.managed_gateway_auth_headers(
+            "https://tool-gateway.example.com/api/bfl/generations",
+            gateway_builder=lambda vendor: f"https://{vendor}-gateway.example.com",
+            token_reader=lambda: "nous-token",
+        )
+
+    assert headers == {"Authorization": "Bearer nous-token"}
+
+
+def test_managed_gateway_auth_headers_reflect_a_rotated_token():
+    # Read fresh on every call: a Nous access token expires within the hour,
+    # and a long session must not keep presenting a dead bearer.
+    tokens = iter(["first-token", "second-token"])
+    builder = lambda vendor: f"https://{vendor}-gateway.example.com"
+    url = "https://tool-gateway.example.com/api/bfl/generations"
+
+    with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True):
+        first = managed_tool_gateway.managed_gateway_auth_headers(url, builder, lambda: next(tokens))
+        second = managed_tool_gateway.managed_gateway_auth_headers(url, builder, lambda: next(tokens))
+
+    assert first["Authorization"] == "Bearer first-token"
+    assert second["Authorization"] == "Bearer second-token"
+
+
+def test_managed_gateway_auth_headers_refuse_a_url_off_the_gateway_origin():
+    # Gated on the URL, never a name: our bearer must never be handed to a
+    # host that merely looks managed.
+    with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True):
+        assert managed_tool_gateway.managed_gateway_auth_headers(
+            "https://attacker.example/api/bfl/generations",
+            gateway_builder=lambda vendor: f"https://{vendor}-gateway.example.com",
+            token_reader=lambda: "nous-token",
+        ) == {}
+
+
+def test_managed_gateway_auth_headers_empty_without_a_token():
+    # Empty rather than raising, so a caller can say "sign in" instead of
+    # sending an unauthenticated request.
+    with patch.object(managed_tool_gateway, "managed_nous_tools_enabled", return_value=True):
+        assert managed_tool_gateway.managed_gateway_auth_headers(
+            "https://tool-gateway.example.com/api/bfl/generations",
+            gateway_builder=lambda vendor: f"https://{vendor}-gateway.example.com",
+            token_reader=lambda: None,
+        ) == {}
+
+
 def test_is_managed_tool_gateway_ready_skips_refresh_for_expired_cached_token(tmp_path, monkeypatch):
     monkeypatch.delenv("TOOL_GATEWAY_USER_TOKEN", raising=False)
     monkeypatch.setenv("HERMES_HOME", str(tmp_path))