mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-10 00:08:58 +00:00
CLI: stop unsloth connect from leaking Studio credentials to unverified servers (#6479)
* CLI: stop `unsloth connect` from leaking Studio credentials to unverified servers
`unsloth connect` (and `unsloth chat`) discovered a Studio base URL from
UNSLOTH_STUDIO_URL or the default localhost port after only an unauthenticated
/api/health probe, then sent credentials to it:
- keyless connect iterated every cached API key and sent each as a bearer token
to {base}/v1/models, so a malicious or port-preempting endpoint could harvest
all of them;
- with no cached key it self-issued a Studio JWT and POSTed it to
{base}/api/auth/api-keys;
- unsloth chat sent the same self-issued JWT to the discovered base.
The key cache was a flat, global list with no binding to a server identity, so a
key minted for one Studio could be replayed to any other.
Changes:
- Scope the agent key cache per base URL so a key is only ever replayed to the
exact server it was minted for. Pre-scoping flat caches are ignored rather
than replayed (at most one extra local mint on the next launch).
- Gate every automatic credential flow to loopback bases. A non-loopback
UNSLOTH_STUDIO_URL now requires an explicit --api-key and nothing is sent
automatically. SSH-tunnelled Studios that land on 127.0.0.1 keep working.
- Mint the API key locally against the Studio auth DB instead of POSTing a
self-issued JWT over the network, so no bearer token leaves the process on the
local path.
- Apply the same loopback gate to connect_studio_server (used by unsloth chat).
Fully closing same-host loopback-port preemption needs a signed /api/health
handshake so the client can verify the server identity before sending anything;
that is tracked as a server-side follow-up.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* CLI: verify Studio server identity before auto-sending credentials
Adds a challenge-response so `unsloth connect` and `unsloth chat` can confirm a
discovered loopback endpoint is really this install's Studio (not a process
that preempted the port) before sending it a cached or freshly minted
credential. This closes the same-host loopback-preemption gap left open by the
previous commit, which could only limit the blast radius.
Server:
- storage.get_or_create_identity_secret(): a dedicated server-wide secret in
app_secrets (kept separate from the per-user JWT secret), readable only by
the same OS user.
- storage.compute_identity_proof(nonce) = HMAC-SHA256(identity secret, nonce).
- GET /api/auth/identity?nonce=<base64url>: unauthenticated, returns the proof.
The nonce is opaque to the server and the proof reveals nothing about the
secret, so answering is safe.
Client:
- verify_studio_identity(base): sends a fresh 32-byte nonce, recomputes the
expected HMAC from the local same-user secret, and constant-time compares.
Fails closed on any error.
- connect._agent_api_key gates the loopback cached-key replay and the local
mint on it; connect_studio_server (used by unsloth chat) gates the
self-issued JWT on it.
A server that cannot read this install's secret (a different OS user, or a
remote/fake endpoint) cannot produce a matching proof, so the client refuses
and falls back to an explicit --api-key.
Tests: studio/backend/tests/test_identity.py (proof determinism, secret
persistence and caching, route response and nonce validation) and additions to
test_connect.py (the verify gate refuses when unverified, an explicit key skips
the check, and an end-to-end client plus server handshake against a stub HTTP
server).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* CLI: mint through the verified server instead of the local auth DB
CodeQL (py/clear-text-storage-sensitive-data) flagged the API key written to
the per-server cache and the agent config files once it was sourced from
storage.create_api_key(); the original HTTP-minted key did not trip the query.
Now that the identity handshake cryptographically confirms the loopback
responder really is this Studio before anything is sent, minting through the
server's /api/auth/api-keys endpoint with a self-issued JWT is safe again and
restores the original, CodeQL-clean data flow. The local-DB mint path is
removed.
The security properties are unchanged: discovery is still loopback-gated and
identity-verified, the key cache is still scoped per server, and a credential
reaches the server only after the handshake has proven its identity. The only
difference from the previous commit is that the self-issued JWT is sent to the
already-verified loopback server rather than the key being minted in-process.
Tests updated to mint through the fake server again.
* CLI: address review feedback on connect credential handling
- Reuse a saved per-server key before the loopback/identity gate. Keys are
scoped per base URL, so a key the user saved with --api-key for a remote or
SSH-tunnelled Studio (whose identity secret the local handshake can't match)
is replayed only to that exact server. The loopback + identity-handshake gate
now guards just auto-minting (self-issuing a JWT and creating a new key),
which is the path that needs a cryptographically verified local Studio. Fixes
keyless reuse being impossible for remote/tunnelled Studios the user had
saved a key for.
- connect_studio_server (unsloth chat / inference): when the user explicitly set
UNSLOTH_STUDIO_URL but the server can't be safely attached (non-loopback, or
identity unverifiable), fail with a clear message instead of silently loading
the model locally. Opportunistic discovery of the local default still falls
back to a local load.
- Harden cache parsing: tolerate a corrupt or hand-edited cache where a base
maps to a non-list (which would otherwise iterate a string into
single-character "keys"), and read the cache as UTF-8.
Tests updated and added: saved-key replay without the handshake for both local
and remote bases, keyless mint still refused when the loopback server is
unverified, and connect_studio_server erroring on an explicit remote while
falling back locally on default discovery.
* CLI: harden connect handshake against relay and gate cached minted keys
Addresses review feedback on the credential handshake:
- Refuse HTTP redirects on credential-bearing requests (the identity handshake,
/v1/models, key minting, and the chat HTTP backend). A process squatting the
discovered port could 302 /api/auth/identity to the real Studio and relay its
valid proof, or bounce a bearer-token request to another base, and urllib
follows redirects by default. A shared no-redirect opener now treats any 3xx
as an error.
- Give cached keys provenance. Keys the user supplied with --api-key are "saved"
and replay without the handshake (needed for remote or SSH-tunnelled Studios
whose secret the local handshake can't match). Keys we auto-mint are "minted"
and replay only after the identity handshake, so a port squatter can't collect
a previously minted localhost key just by answering the health check. New cache
shape: servers[base] = {"saved": [...], "minted": [...]}.
Known residual: a different-OS-user process that squats the port and can also
reach a genuine same-secret Studio elsewhere on loopback can still manually relay
the identity challenge. Fully closing that needs the proof bound to the server's
real listening port, or OS-level peer-credential checks; tracked as follow-up. A
same-user attacker is out of scope, since it can already read the 0600 key cache.
Tests: redirect rejection in the handshake, minted-cache requiring the handshake
while saved-cache bypasses it, and the existing suites updated for the new cache
shape. unsloth_cli (206) and test_identity.py (5) pass.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* CLI: keep urllib imports function-local in the no-redirect opener
The repo's import-hoist safety linter (scripts/verify_import_hoist.py) flags
hoisting urllib to module level because it re-points the 'urllib' name in the
pre-existing HttpChatBackend._request scope. Build the no-redirect opener lazily
with function-local urllib imports instead, matching this module's convention,
and restore the local 'import urllib.request' in _request and
verify_studio_identity.
* test(identity): skip route tests when routes.auth import chain is unavailable
The identity route tests build a TestClient from routes.auth, which pulls the
whole routes package (routes/__init__ -> inference -> llama_cpp, ...). In a
minimal test matrix without the heavy backend deps, or when another test in the
same process has already broken that import chain, importing it raised and the
two route tests hard-failed. Skip in that case instead: the proof crypto is
covered by the storage-level tests, and the full backend CI still exercises the
route. No behaviour change where the deps are present (5 passed in isolation).
* test(connect): make connect tests pass on native Windows
unsloth connect supports Windows: --no-launch prints PowerShell ($env:X =
"v" / Remove-Item Env:X) instead of POSIX (export/unset), and the launch
path bridges env into a Windows agent .exe over WSLENV. The tests hardcoded
the POSIX shell forms, so on a real windows-latest runner 12 of them failed on
the assertion string even though every command exited 0.
Add OS-aware assertion helpers (_assert_env_set / _assert_env_unset) that check
the right shell syntax for the host OS, and skip the two WSL-from-Linux shim
tests on native Windows (os.name is 'posix' inside WSL, so that path can't run
there). No change on Linux/macOS (57 passed); the connect command's behaviour
is untouched. Validated on a windows-latest staging runner.
* style(connect): tighten comments in the credential-leak fix
Condense the verbose explanatory comments and multi-line docstrings added by
this PR to one or two lines each, drop a few that just restated the code, and
keep the security rationale where it is load-bearing. Comment/whitespace only;
verified with unslothai/scripts comment_tools.py (check --strip-docstrings:
6/6 'code unchanged'). Tests unchanged: connect 57 passed, identity 5 passed.
* CLI/Studio: harden the identity handshake (review round)
Addresses the latest Codex/Gemini review of the handshake:
- Store the identity secret privately. sqlite3.connect created the auth DB
world-readable under a 022 umask, so another OS user could read app_secrets
and forge proofs, defeating the same-user assumption the handshake rests on.
The auth dir and DB are now restricted to owner-only (0700/0600); the JWT
secret and password hashes there get the same protection.
- Bind the proof to the server's listening port. The stateless HMAC(secret,
nonce) was relayable: a process squatting the discovered port could proxy the
challenge to the real Studio on another port and pass it back. The proof now
covers the port the server actually listens on (from the socket, never the
Host header) and the client checks it against the port it connected to, so a
relayed proof from a different port no longer matches. Closes the manual-relay
residual left after the redirect fix.
- Cap the identity response read (the server is still unverified at that point)
and serve the identity route from a sync def so its first-call SQLite read
runs in the threadpool instead of the event loop.
Tests: port-bound proof + relayed-proof rejection added; identity (5) and
unsloth_cli (58) suites pass. Verified end to end against a real backend
(auth DB owner-only, handshake + mint still succeed).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* CLI/Studio: bind the identity proof to the connection address, not just port
Follow-up to the port binding from the last review. Binding only to the port
left a cross-address relay: a squatter on a different loopback address but the
same port (for example localhost resolving to a squatter on ::1 while the real
Studio is on 127.0.0.1) could proxy the nonce to the real Studio and pass back
a proof that still matched, since both share the port.
The proof now covers the address and the port the connection landed on:
- Server: takes the address+port from request.scope, which uvicorn populates
from getsockname, so it is the real local address the client reached even
when Studio is bound to 0.0.0.0 (verified empirically), never the
client-controlled Host header.
- Client: resolves the base host to one concrete IP, talks to exactly that IP,
and binds the proof to (IP, port). A proof relayed from a Studio on a
different address or port was computed for that other endpoint and no longer
matches the one the client dialed.
Both sides normalise the address through ipaddress so equivalent forms compare
equal. Together with the private-secret and redirect fixes, this closes the
cross-user loopback relay an attacker can mount without reading the secret.
Tests: proof now bound to host+port; relayed-proof rejection retained; identity
(5) and unsloth_cli (58) suites pass. Verified end to end against a real backend
(handshake + mint still succeed).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* CLI: pick the loopback address at discovery so localhost does not regress
find_studio_server now resolves a bare localhost base to its concrete loopback
addresses and returns the first that answers /api/health, IPv4 127.0.0.1 first
(where unsloth studio binds by default). The whole flow (health probe, identity
check, credential send) then targets that one address instead of racing
IPv4/IPv6 resolution, where localhost could resolve ::1-first and hide a Studio
bound to 127.0.0.1. A literal IP or remote name is unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
parent
9f39cc2c39
commit
9b5c94df32
7 changed files with 807 additions and 100 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
96
studio/backend/tests/test_identity.py
Normal file
96
studio/backend/tests/test_identity.py
Normal file
|
|
@ -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
|
||||
|
|
@ -18,6 +18,28 @@ _THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", 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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue