diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index ee6678d9b9..796d03ff68 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -5,6 +5,7 @@ import hashlib import hmac +import ipaddress import os import secrets import sqlite3 @@ -100,6 +101,14 @@ def get_connection() -> sqlite3.Connection: """Get a connection to the auth database, creating tables if needed.""" ensure_dir(DB_PATH.parent) conn = sqlite3.connect(DB_PATH) + # Keep the auth dir + DB private (they hold the JWT/identity secrets and + # password hashes); sqlite3.connect would otherwise create the DB 0644 under + # a 022 umask, letting another OS user read the identity secret and forge proofs. + for _path, _mode in ((DB_PATH.parent, 0o700), (DB_PATH, 0o600)): + try: + os.chmod(_path, _mode) + except OSError: + pass conn.row_factory = sqlite3.Row conn.execute( """ @@ -210,6 +219,57 @@ def _get_or_create_api_key_pbkdf2_salt() -> bytes: return salt +# Secret answering the /api/auth/identity challenge (HMAC(secret, nonce)). Lives +# in this same-user DB so a port squatter or remote/fake server can't forge a +# proof. Separate from the per-user JWT secret. +_IDENTITY_SECRET_DB_KEY = "studio_identity_secret" +_identity_secret_cache: Optional[bytes] = None + + +def get_or_create_identity_secret() -> bytes: + """Return the identity secret (hex 32-byte row in app_secrets), creating it once.""" + global _identity_secret_cache + if _identity_secret_cache is not None: + return _identity_secret_cache + + conn = get_connection() + try: + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + if row is None: + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + (_IDENTITY_SECRET_DB_KEY, secrets.token_hex(32)), + ) + conn.commit() + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + secret = bytes.fromhex(row["value"]) + finally: + conn.close() + + _identity_secret_cache = secret + return secret + + +def compute_identity_proof(nonce: bytes, host: str, port: int) -> str: + """HMAC-SHA256 proof that the caller holds this install's identity secret, + bound to the loopback address and port the connection landed on. A proof + relayed from a Studio on a different address/port (a squatter proxying to the + real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was + computed for that other endpoint and won't match the one the client dialed.""" + try: + host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms + except ValueError: + host = (host or "").lower() + msg = b"|".join([nonce, host.encode(), str(int(port)).encode()]) + return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest() + + _API_KEY_PBKDF2_ITERATIONS = 100_000 DESKTOP_SECRET_PREFIX = "desktop-" _DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash" diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index 3ac85380fe..d2b3bf94e9 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -5,6 +5,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +import base64 import ipaddress import os import shlex @@ -214,6 +215,34 @@ def _clear_login_bucket(key: tuple[str, str]) -> None: _LOGIN_IP_BUCKETS.pop(ip, None) +# Sync def (not async): compute_identity_proof touches SQLite on the first call, +# so FastAPI runs it in the threadpool rather than blocking the event loop. +@router.get("/identity") +def identity(nonce: str, request: Request) -> dict: + """Challenge-response proof this is the real local Studio: caller sends a nonce, + gets HMAC(install identity secret, nonce, connection address + port). + Unauthenticated and side-effect free; a process that can't read the same-user + secret can't forge a proof, and binding to the address/port the connection + landed on stops a squatter relaying a proof from the real Studio elsewhere.""" + try: + raw = base64.urlsafe_b64decode(nonce) + except Exception: + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, detail = "nonce must be base64url" + ) + if not 16 <= len(raw) <= 128: + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, detail = "nonce must decode to 16-128 bytes" + ) + # The address + port the connection actually landed on, from the socket + # (request.scope is getsockname, so it is the real local address even when + # bound to 0.0.0.0), never the client-controlled Host header. + server = request.scope.get("server") or ("", 0) + host = server[0] or "" + port = server[1] if server[1] is not None else 0 + return {"proof": storage.compute_identity_proof(raw, host, port)} + + @router.get("/status", response_model = AuthStatusResponse) async def auth_status() -> AuthStatusResponse: """Auth initialization state; ``default_username`` is exposed for first-boot UI prefill only.""" diff --git a/studio/backend/tests/test_identity.py b/studio/backend/tests/test_identity.py new file mode 100644 index 0000000000..1e84ddef35 --- /dev/null +++ b/studio/backend/tests/test_identity.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the server identity handshake (`GET /api/auth/identity`). + +The endpoint lets a client confirm an endpoint is really this Studio install +before sending it a credential: the client sends a random nonce and checks the +returned HMAC against one computed from the install identity secret. A process +that cannot read this same-user secret cannot forge a matching proof. +""" + +import base64 +import hashlib +import hmac + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from auth import storage + +HOST = "127.0.0.1" # the proof is bound to the connection address... +PORT = 8765 # ...and port + + +@pytest.fixture(autouse = True) +def isolated_auth_db(tmp_path, monkeypatch): + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + yield + + +def test_identity_secret_is_persistent_and_cached(): + first = storage.get_or_create_identity_secret() + assert isinstance(first, bytes) and len(first) == 32 + # Cached in-process. + assert storage.get_or_create_identity_secret() == first + # Persisted: a fresh process (cache cleared) reads the same stored value. + storage._identity_secret_cache = None + assert storage.get_or_create_identity_secret() == first + + +def test_compute_identity_proof_matches_manual_hmac(): + nonce = b"a-fixed-nonce-for-the-proof-test!" + secret = storage.get_or_create_identity_secret() + expected = hmac.new( + secret, b"|".join([nonce, HOST.encode(), str(PORT).encode()]), hashlib.sha256 + ).hexdigest() + assert storage.compute_identity_proof(nonce, HOST, PORT) == expected + # Bound to nonce, host and port: changing any one yields a different proof. + assert ( + storage.compute_identity_proof(b"a-different-nonce-entirely-here!!", HOST, PORT) != expected + ) + assert storage.compute_identity_proof(nonce, "127.0.0.2", PORT) != expected + assert storage.compute_identity_proof(nonce, HOST, PORT + 1) != expected + + +def test_proof_differs_when_secret_differs(tmp_path, monkeypatch): + nonce = b"shared-nonce-across-two-installs!" + proof_a = storage.compute_identity_proof(nonce, HOST, PORT) + # A different install (different secret) can't reproduce the proof. + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "other_auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + assert storage.compute_identity_proof(nonce, HOST, PORT) != proof_a + + +def _identity_client() -> TestClient: + # routes.auth pulls the whole routes package (-> inference -> llama_cpp). Skip + # if those heavy deps are missing; the proof crypto is covered above. + try: + from routes.auth import router + except Exception as exc: # pragma: no cover - environment-dependent + pytest.skip(f"routes.auth not importable in this environment: {exc}") + + app = FastAPI() + app.include_router(router, prefix = "/api/auth") + # base_url host:port become scope["server"], which the route binds the proof to. + return TestClient(app, base_url = f"http://{HOST}:{PORT}") + + +def test_identity_route_returns_matching_proof(): + client = _identity_client() + nonce = b"route-level-nonce-for-the-server!" + encoded = base64.urlsafe_b64encode(nonce).decode() + response = client.get(f"/api/auth/identity?nonce={encoded}") + assert response.status_code == 200 + assert response.json()["proof"] == storage.compute_identity_proof(nonce, HOST, PORT) + + +def test_identity_route_validates_nonce(): + client = _identity_client() + # Decodes to < 16 bytes: too little entropy to be meaningful. + short = base64.urlsafe_b64encode(b"tiny").decode() + assert client.get(f"/api/auth/identity?nonce={short}").status_code == 400 + # Missing nonce: FastAPI request validation. + assert client.get("/api/auth/identity").status_code == 422 diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index 1faa632926..0baeb2ffe5 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -18,6 +18,28 @@ _THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL) # "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request. _USER_AGENT = "unsloth-cli" +# Built lazily; urllib stays function-local to match this module. +_no_redirect_opener = None + + +def urlopen_no_redirect(request, timeout): + """urlopen that errors on any redirect: following a 3xx would send a bearer + token (or accept an identity proof) to a base we never vetted, letting a port + squatter relay a real Studio's response.""" + global _no_redirect_opener + if _no_redirect_opener is None: + import urllib.error + import urllib.request + + class _NoRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): + raise urllib.error.HTTPError( + req.full_url, code, f"refusing redirect to {newurl}", headers, fp + ) + + _no_redirect_opener = urllib.request.build_opener(_NoRedirect) + return _no_redirect_opener.open(request, timeout = timeout) + def ensure_studio_backend_path() -> None: backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend") @@ -258,16 +280,122 @@ def load_chat_backend( return ChatBackend("unsloth", backend) +def _loopback_candidate_bases(base: str) -> list: + """For a bare ``localhost`` base, the concrete IP bases to try, IPv4 + 127.0.0.1 first (where ``unsloth studio`` binds by default). Pinning to one + address up front means discovery, the identity check, and the credential we + then send all target the same endpoint instead of racing IPv4/IPv6 + resolution -- which would otherwise let the health probe land on one address + and the identity check on another. A literal IP or remote name is unchanged. + """ + from urllib.parse import urlparse + + parsed = urlparse(base) + if (parsed.hostname or "").lower() != "localhost": + return [base] + import socket + + port = parsed.port or (443 if parsed.scheme == "https" else 80) + try: + ips = { + ai[4][0] for ai in socket.getaddrinfo(parsed.hostname, port, type = socket.SOCK_STREAM) + } + except Exception: + return [base] + ordered = sorted(ips, key = lambda ip: (ip != "127.0.0.1", ip)) + bases = [ + f"{parsed.scheme}://" + (f"[{ip}]:{port}" if ":" in ip else f"{ip}:{port}") + for ip in ordered + ] + return bases or [base] + + def find_studio_server(timeout: float = 3.0) -> Optional[str]: import urllib.request base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") - request = urllib.request.Request(f"{base}/api/health", headers = {"User-Agent": _USER_AGENT}) + # Try the concrete loopback addresses in order and return the first that + # answers, so the rest of the flow talks to that exact address. + for candidate in _loopback_candidate_bases(base): + request = urllib.request.Request( + f"{candidate}/api/health", headers = {"User-Agent": _USER_AGENT} + ) + try: + with urllib.request.urlopen(request, timeout = timeout): + return candidate + except Exception: + continue + return None + + +def is_loopback_url(base: str) -> bool: + """True only when *base* resolves to loopback. find_studio_server() trusts a + base after only a health probe, so credentials are auto-sent only to loopback + (a local Studio or an SSH tunnel on 127.0.0.1), the targets the auto flows mean.""" + from urllib.parse import urlparse + + host = (urlparse(base).hostname or "").lower() + if host in ("localhost", "127.0.0.1", "::1"): + return True try: - with urllib.request.urlopen(request, timeout = timeout): - return base + import ipaddress + return ipaddress.ip_address(host).is_loopback + except ValueError: + return False + + +def verify_studio_identity(base: str, timeout: float = 3.0) -> bool: + """Confirm `base` is really this machine's Studio before sending a secret. + + Send a random nonce to /api/auth/identity and check the returned HMAC against + the one computed from the local same-user secret; an endpoint without that + secret (port squatter, remote/fake) can't match. Fails closed on any error.""" + import base64 + import hmac as _hmac + import json + import secrets as _secrets + import socket + import urllib.request + from urllib.parse import urlparse + + try: + import studio.backend.core # noqa: F401 puts studio/backend on sys.path + from studio.backend.auth import storage except Exception: - return None + return False + + parsed = urlparse(base) + host = parsed.hostname or "" + port = parsed.port or (443 if parsed.scheme == "https" else 80) + # Resolve to one concrete address and talk to *that* address, then bind the + # proof to (address, port). A name like localhost can resolve to a squatter on + # ::1 while the real Studio is on 127.0.0.1; connecting to the resolved IP and + # binding to it means a proof relayed from a different address/port won't match. + try: + ip = socket.getaddrinfo(host, port, type = socket.SOCK_STREAM)[0][4][0] + except Exception: + return False + netloc = f"[{ip}]:{port}" if ":" in ip else f"{ip}:{port}" + nonce = _secrets.token_bytes(32) + query = base64.urlsafe_b64encode(nonce).decode() + request = urllib.request.Request( + f"{parsed.scheme}://{netloc}/api/auth/identity?nonce={query}", + headers = {"User-Agent": _USER_AGENT, "Host": parsed.netloc}, + ) + try: + # No redirects: a 302 could relay a real Studio's proof (see urlopen_no_redirect). + # Cap the read: the server is still unverified, so don't trust its length. + with urlopen_no_redirect(request, timeout = timeout) as response: + proof = json.loads(response.read(65536).decode() or "{}").get("proof") + except Exception: + return False + if not isinstance(proof, str): + return False + try: + expected = storage.compute_identity_proof(nonce, ip, port) + except Exception: + return False + return _hmac.compare_digest(proof, expected) def _studio_token() -> Optional[str]: @@ -316,7 +444,8 @@ class HttpChatBackend: }, method = method, ) - return urllib.request.urlopen(request, timeout = timeout) + # No redirects: this carries a bearer token (see urlopen_no_redirect). + return urlopen_no_redirect(request, timeout = timeout) def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None: typer.echo(f"Loading {model} on the Studio server", err = True) @@ -414,9 +543,37 @@ def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit) base_url = find_studio_server() if not base_url: return None + + # Explicit server (UNSLOTH_STUDIO_URL) we can't safely attach to -> fail loudly; + # opportunistic local discovery just falls back to a local load. + explicit = bool(os.environ.get("UNSLOTH_STUDIO_URL")) + + def _refuse(reason: str): + if not explicit: + return None + typer.echo( + f"Can't attach to the Studio server at {base_url}: {reason} Run Studio " + "on this machine, or unset UNSLOTH_STUDIO_URL to load the model locally.", + err = True, + ) + raise typer.Exit(code = 1) + + # Only hand the self-issued JWT (signed with the local secret) to loopback: a + # remote URL is unverified and a real remote Studio would reject it anyway. + if not is_loopback_url(base_url): + return _refuse( + "it isn't a local Studio, so a self-issued token can't " + "authenticate to it and must not be sent to it." + ) + # Confirm the loopback responder is really our Studio (not a port squatter). + if not verify_studio_identity(base_url): + return _refuse( + "its identity couldn't be verified (it may be running as a " + "different OS user, or another process took the port)." + ) token = _studio_token() if not token: - return None + return _refuse("couldn't self-issue a Studio token (is Studio set up here?).") backend = HttpChatBackend(base_url, token) backend.ensure_loaded( model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit diff --git a/unsloth_cli/commands/connect.py b/unsloth_cli/commands/connect.py index 05f75d263c..096a7925e5 100644 --- a/unsloth_cli/commands/connect.py +++ b/unsloth_cli/commands/connect.py @@ -22,6 +22,9 @@ from unsloth_cli._inference import ( _studio_token, ensure_studio_backend_path, find_studio_server, + is_loopback_url, + urlopen_no_redirect, + verify_studio_identity, ) connect_app = typer.Typer( @@ -46,7 +49,11 @@ _KEY_OPTION = typer.Option( None, "--api-key", envvar = "UNSLOTH_API_KEY", - help = "Studio API key; minted automatically when omitted. Keys are remembered, so passing one once is enough.", + help = ( + "Studio API key. For a local Studio it is minted automatically and " + "remembered per server. For a remote server, pass one with --api-key " + "(or UNSLOTH_API_KEY); it is remembered for next time." + ), ) _LAUNCH_OPTION = typer.Option( True, @@ -88,7 +95,8 @@ def _http_json( method = method, ) try: - with urllib.request.urlopen(request, timeout = timeout) as response: + # No redirects: a 3xx would leak this bearer token to an unvetted base. + with urlopen_no_redirect(request, timeout = timeout) as response: return json.loads(response.read().decode() or "{}") except urllib.error.HTTPError as exc: if error is None: @@ -117,18 +125,35 @@ def _key_cache_path() -> Path: return auth_root() / "agent_api_key.json" -def _cached_keys(cache: Path) -> list: +def _read_cache(cache: Path) -> dict: try: - data = json.loads(cache.read_text()) + data = json.loads(cache.read_text(encoding = "utf-8")) except Exception: - return [] - if not isinstance(data, dict): - return [] - keys = [k for k in data.get("keys", []) if isinstance(k, str)] - legacy = data.get("key") # pre-multi-key cache format - if isinstance(legacy, str) and legacy not in keys: - keys.append(legacy) - return keys + return {} + return data if isinstance(data, dict) else {} + + +def _server_buckets(servers: dict, base: str) -> dict: + # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a + # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). + entry = servers.get(base) if isinstance(servers, dict) else None + if isinstance(entry, list): + return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} + if not isinstance(entry, dict): + return {"saved": [], "minted": []} + + def _strs(name: str) -> list: + value = entry.get(name) + return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] + + return {"saved": _strs("saved"), "minted": _strs("minted")} + + +def _cached_keys(cache: Path, base: str, source: str) -> list: + # Keys are scoped per server. `source` splits user-supplied --api-key keys + # ("saved", trusted for that base) from auto-minted ones ("minted", replayed + # only after the identity check). Legacy unscoped caches are ignored. + return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] def _write_private_json(path: Path, data: dict) -> None: @@ -159,59 +184,90 @@ def _subdict(parent: dict, key: str) -> dict: return child -def _remember_key(cache: Path, key: str) -> None: - existing = _cached_keys(cache) - keys = ([key] + [k for k in existing if k != key])[:8] - if keys == existing: +def _remember_key(cache: Path, base: str, key: str, source: str) -> None: + data = _read_cache(cache) + servers = data.get("servers") + if not isinstance(servers, dict): + servers = data["servers"] = {} + buckets = _server_buckets(servers, base) + other = "minted" if source == "saved" else "saved" + buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] + buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance + new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} + if servers.get(base) == new_entry: return + servers[base] = new_entry + # Collapse legacy unscoped fields. + data.pop("keys", None) + data.pop("key", None) try: - _write_private_json(cache, {"keys": keys}) + _write_private_json(cache, data) except OSError: pass # worst case the next launch mints another key +def _key_accepted(base: str, key: str) -> bool: + try: + _http_json("GET", f"{base}/v1/models", key) + return True + except Exception: + return False + + def _agent_api_key(base: str, explicit: Optional[str]) -> str: cache = _key_cache_path() if explicit: - _remember_key(cache, explicit) + _remember_key(cache, base, explicit, "saved") return explicit - # Keys are per-server, so when switching between Studios (local one day, - # an SSH-tunnelled remote the next) the right key is whichever validates. - for key in _cached_keys(cache): - try: - _http_json("GET", f"{base}/v1/models", key) - except Exception: - continue - _remember_key(cache, key) - return key + # Replay a key the user saved for *this exact* server first (scoped per base, + # so it only goes back there -- including a remote/SSH-tunnelled Studio whose + # secret the local handshake can't match). Skip ones the server rejects. + for key in _cached_keys(cache, base, "saved"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "saved") + return key + # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() + # trusts a base after only a health check, so both are limited to a loopback + # server we can cryptographically confirm is ours. + if not is_loopback_url(base): + _fail( + f"No saved API key for {base} and automatic minting only runs against " + "a local Studio. Create an API key in Studio → Settings → API and " + "pass it with --api-key (it is remembered per server), or set " + "UNSLOTH_API_KEY." + ) + if not verify_studio_identity(base): + _fail( + f"Couldn't verify that {base} is your Studio (it may be running as a " + "different OS user, or another process took the port). Create an API " + "key in Studio → Settings → API and pass it with --api-key, or set " + "UNSLOTH_API_KEY." + ) + + # Identity verified: replay a previously auto-minted key, else mint a new one. + for key in _cached_keys(cache, base, "minted"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "minted") + return key + + # Self-issue a JWT (signed with the local secret) and mint a key. token = _studio_token() - auth_help = ( - "Couldn't authenticate with the Studio server automatically (it may be " - "remote, or running as a different OS user). Create an API key in " - "Studio → Settings → API and pass it once with --api-key; it is " - "remembered for next time." - ) if token is None: - _fail(auth_help) - try: - key = _http_json( - "POST", - f"{base}/api/auth/api-keys", - token, - {"name": "Coding agents (unsloth connect)"}, - )["key"] - except urllib.error.HTTPError as exc: - # A self-issued token only validates against a local, same-OS-user - # server; a remote Studio signs with a different secret and rejects it - # (401/403). Point at --api-key instead of the raw "expired token". - if exc.code in (401, 403): - _fail(auth_help) - _fail(f"Couldn't create an API key: {_http_error_detail(exc)}") - except (urllib.error.URLError, TimeoutError) as exc: - _fail(f"Couldn't create an API key: {getattr(exc, 'reason', None) or exc}") - _remember_key(cache, key) + _fail( + "Couldn't authenticate with the Studio server automatically. Create " + "an API key in Studio → Settings → API and pass it with --api-key, " + "or set UNSLOTH_API_KEY." + ) + key = _http_json( + "POST", + f"{base}/api/auth/api-keys", + token, + {"name": "Coding agents (unsloth connect)"}, + error = "Couldn't create an API key", + )["key"] + _remember_key(cache, base, key, "minted") return key diff --git a/unsloth_cli/tests/test_connect.py b/unsloth_cli/tests/test_connect.py index 06806c8fd7..e76a892647 100644 --- a/unsloth_cli/tests/test_connect.py +++ b/unsloth_cli/tests/test_connect.py @@ -26,6 +26,18 @@ BASE = "http://127.0.0.1:8888" MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} +# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and +# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. +def _assert_env_set(output: str, name: str, value: str) -> None: + needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _assert_env_unset(output: str, name: str) -> None: + needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + @pytest.fixture() def claude_settings(tmp_path, monkeypatch): path = tmp_path / "claude" / "settings.json" @@ -173,6 +185,9 @@ def fake_studio(tmp_path, monkeypatch, claude_settings): raise AssertionError(f"unexpected request: {method} {url}") monkeypatch.setattr(connect, "find_studio_server", lambda: BASE) + # Identity handshake has its own tests; trust the loopback server here. + monkeypatch.setattr(connect, "verify_studio_identity", lambda base: True) + # _studio_token / api-keys are faked so the mint flow stays offline. monkeypatch.setattr(connect, "_studio_token", lambda: "jwt-token") monkeypatch.setattr(connect, "_http_json", http_json) monkeypatch.setattr(connect, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") @@ -186,13 +201,13 @@ def fake_studio(tmp_path, monkeypatch, claude_settings): def test_connect_claude_no_launch(fake_studio, claude_settings): result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) assert result.exit_code == 0, result.output - assert "unset ANTHROPIC_API_KEY" in result.output - assert "unset CLAUDE_CODE_OAUTH_TOKEN" in result.output - assert f"export ANTHROPIC_BASE_URL={BASE}" in result.output - assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output - assert f"export ANTHROPIC_MODEL={MODEL['id']}" in result.output - assert "export CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1" in result.output - assert "export CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1" in result.output + _assert_env_unset(result.output, "ANTHROPIC_API_KEY") + _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") + _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output settings = json.loads(claude_settings.read_text()) assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" @@ -222,6 +237,11 @@ def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypa assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): captured = {} monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") @@ -261,6 +281,11 @@ def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypat assert name in captured["env"]["WSLENV"].split(":") +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") monkeypatch.setattr( @@ -280,7 +305,7 @@ def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studi def test_connect_codex_no_launch(fake_studio, tmp_path): result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) assert result.exit_code == 0, result.output - assert "export UNSLOTH_STUDIO_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output + _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") assert "codex --oss --profile unsloth_api" in result.output assert (tmp_path / "codex" / "config.toml").exists() assert (tmp_path / "codex" / "unsloth_api.config.toml").exists() @@ -289,10 +314,11 @@ def test_connect_codex_no_launch(fake_studio, tmp_path): def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + # First run mints; second reuses the minted key cached for this server. mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] assert len(mints) == 1 cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["keys"] == ["sk-unsloth-feedfacefeedface"] + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): @@ -302,14 +328,20 @@ def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path) ) result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) assert result.exit_code == 0, result.output - assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-deadbeefdeadbeef" in result.output - mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] - assert mints == [] + # Reused, not re-minted (a mint would return the feedface stand-in). + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + # An explicit key is remembered as "saved" so it replays without the handshake. + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"keys": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]})) + cache.write_text( + json.dumps( + {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} + ) + ) inner = connect._http_json def http_json( @@ -327,21 +359,22 @@ def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, mon monkeypatch.setattr(connect, "_http_json", http_json) result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) assert result.exit_code == 0, result.output - assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output - mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] - assert mints == [] + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") # The working key moves to the front so the next run tries it first. cached = json.loads(cache.read_text()) - assert cached["keys"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] -def test_connect_reads_legacy_single_key_cache(fake_studio, tmp_path): +def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): + # Legacy unscoped caches have no server binding (could leak across servers), + # so they're ignored: a fresh key is minted and stored scoped to this server. (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) assert result.exit_code == 0, result.output - assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-oldformat" in result.output - mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] - assert mints == [] + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + assert "key" not in cached # legacy field collapsed away def test_connect_model_flag_loads_on_server(fake_studio): @@ -353,7 +386,7 @@ def test_connect_model_flag_loads_on_server(fake_studio): assert loads == [ ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) ] - assert "export ANTHROPIC_MODEL=unsloth/Qwen3.5-35B-A3B" in result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): @@ -384,7 +417,7 @@ def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): connect.connect_app, ["claude", "--no-launch", "--model", requested] ) assert result.exit_code == 0, result.output - assert f"export ANTHROPIC_MODEL={canonical}" in result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): @@ -452,28 +485,270 @@ def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): assert result.exit_code == 0, result.output -def test_connect_remote_token_rejected_points_at_api_key(fake_studio, monkeypatch): - # A self-issued token is invalid against a remote Studio (different secret); - # the auto-mint 401 should become actionable --api-key guidance. - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/auth/api-keys"): - raise urllib.error.HTTPError(url, 401, "Invalid or expired token", None, None) - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) +def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): + # A server known only by URL + health check is unverified: keyless connect + # must refuse and make no request at all. + monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.evil.example:8888") result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) assert result.exit_code == 1 assert "Settings → API" in result.output assert "--api-key" in result.output + assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) + + +def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): + # User named both server and key, so it's their choice; only auto-send is blocked. + monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.example:8888") + result = CliRunner().invoke( + connect.connect_app, + ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): + # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; + # auto-minting stays blocked for non-loopback. + remote = "http://studio.example:8888" + monkeypatch.setattr(connect, "find_studio_server", lambda: remote) + (tmp_path / "agent_api_key.json").write_text( + json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) + ) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): + # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an + # error, not a silent local model load (which they did not ask for). + import typer + + import unsloth_cli._inference as inference + + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") + monkeypatch.setattr( + inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" + ) + with pytest.raises(typer.Exit): + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + + +def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): + # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback + # server can't be verified, fall back to a local load rather than erroring. + import unsloth_cli._inference as inference + + monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) + monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") + monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) + assert ( + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + is None + ) + + +def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( + fake_studio, tmp_path, monkeypatch +): + # With no saved key, the next step would auto-mint; an unverified loopback + # server (port squatter) must be refused, with nothing sent. + monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): + # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) + # replays on keyless runs without the handshake, scoped to its own base. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) + monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted + + +def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): + # A "minted" key is NOT replayed to an unverified loopback server: minting and + # minted-key replay both sit behind the handshake, so a squatter can't grab it. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) + monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent + + +def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): + # An explicit key is the user's deliberate choice, so it does not require + # the automatic identity handshake. + monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke( + connect.connect_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + + +def _serve_identity(proof_for): + """Start a localhost HTTP server answering /api/auth/identity with + proof_for(nonce_bytes). Returns (base_url, shutdown).""" + import base64 + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + from urllib.parse import parse_qs, urlparse + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + parsed = urlparse(self.path) + if parsed.path != "/api/auth/identity": + self.send_response(404) + self.end_headers() + return + nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) + host, port = self.server.server_address[0], self.server.server_address[1] + body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): + # Real crypto end to end: verify_studio_identity reads the install secret from + # an isolated DB; a "good" server proves the same secret, a spoofing one can't. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: # backend not importable here (e.g. missing deps) + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + good = lambda nonce, host, port: storage.compute_identity_proof( + nonce, host, port + ) # real secret + bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret + base_ok, stop_ok = _serve_identity(good) + base_bad, stop_bad = _serve_identity(bad) + try: + assert inference.verify_studio_identity(base_ok) is True + assert inference.verify_studio_identity(base_bad) is False + finally: + stop_ok() + stop_bad() + + +def _serve_redirect(target): + """Start a localhost server that 302-redirects every GET to target+path.""" + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(302) + self.send_header("Location", target + self.path) + self.end_headers() + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): + # A squatter could 302 /api/auth/identity to the real Studio and relay its + # proof; redirects must be refused so the squatter's base isn't accepted. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + squatter_base, stop_squatter = _serve_redirect(real_base) + try: + assert inference.verify_studio_identity(real_base) is True # direct: ok + assert inference.verify_studio_identity(squatter_base) is False # relayed: refused + finally: + stop_real() + stop_squatter() + + +def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): + # A squatter that proxies the nonce to the real Studio on another port gets a + # proof bound to *that* port; the client expects one bound to the port it + # connected to, so the relayed proof is rejected. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + real_port = int(real_base.rsplit(":", 1)[1]) + # The squatter answers on its own port but returns the proof for the real port. + squatter_base, stop_squatter = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) + ) + try: + assert inference.verify_studio_identity(real_base) is True + assert inference.verify_studio_identity(squatter_base) is False + finally: + stop_real() + stop_squatter() + + +@pytest.mark.parametrize( + "url, loopback", + [ + ("http://127.0.0.1:8888", True), + ("http://localhost:8888", True), + ("http://[::1]:8888", True), + ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 + ("http://0.0.0.0:8888", False), + ("http://10.0.0.5:8888", False), + ("http://studio.evil.example:8888", False), + ("https://studio.example.com", False), + ], +) +def test_is_loopback_url(url, loopback): + assert connect.is_loopback_url(url) is loopback def test_connect_no_studio_errors(fake_studio, monkeypatch): @@ -489,7 +764,7 @@ def test_connect_explicit_api_key_skips_mint(fake_studio): ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], ) assert result.exit_code == 0, result.output - assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-deadbeefdeadbeef" in result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) @@ -670,7 +945,7 @@ def test_connect_hermes_no_launch(fake_studio, hermes_config): yaml = pytest.importorskip("yaml") result = CliRunner().invoke(connect.connect_app, ["hermes", "--no-launch"]) assert result.exit_code == 0, result.output - assert "export UNSLOTH_API_KEY=sk-unsloth-feedfacefeedface" in result.output + _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") assert "hermes" in result.output config = yaml.safe_load(hermes_config.read_text()) assert config["model"]["provider"] == "custom:unsloth" diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py index 629cb46405..770bd0db29 100644 --- a/unsloth_cli/tests/test_inference_chat.py +++ b/unsloth_cli/tests/test_inference_chat.py @@ -260,6 +260,40 @@ def test_find_studio_server_none_when_not_running(monkeypatch): assert _inference.find_studio_server() is None +def test_find_studio_server_prefers_ipv4_loopback_for_localhost(monkeypatch): + # localhost resolving ::1-first must not hide a Studio bound to 127.0.0.1: + # discovery tries each loopback address and returns the one that answers. + import socket + import urllib.request + + from unsloth_cli import _inference + + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://localhost:8888") + monkeypatch.setattr( + socket, + "getaddrinfo", + lambda *a, **k: [ + (socket.AF_INET6, socket.SOCK_STREAM, 0, "", ("::1", 8888, 0, 0)), + (socket.AF_INET, socket.SOCK_STREAM, 0, "", ("127.0.0.1", 8888)), + ], + ) + + class _OK: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def only_ipv4(request, *a, **k): + if "127.0.0.1" not in request.full_url: + raise OSError("connection refused") + return _OK() + + monkeypatch.setattr(urllib.request, "urlopen", only_ipv4) + assert _inference.find_studio_server() == "http://127.0.0.1:8888" + + class _FakeSSEResponse: def __init__(self, lines): self._lines = lines