diff --git a/tests/integration/_support/device_login.py b/tests/integration/_support/device_login.py index 55519678f40..55eb5a753bb 100644 --- a/tests/integration/_support/device_login.py +++ b/tests/integration/_support/device_login.py @@ -25,6 +25,7 @@ from pydantic import JsonValue CHATGPT_AUTH_HOST: Final = "auth.openai.com" GITHUB_HOST: Final = "github.com" GITHUB_API_HOST: Final = "api.github.com" +COPILOT_API_HOST: Final = "api.githubcopilot.com" AUTH_HOSTS: Final = (CHATGPT_AUTH_HOST, GITHUB_HOST, GITHUB_API_HOST) CHATGPT_ACCOUNT: Final = "acct-device-login" _FAR_FUTURE: Final = 4102444800 @@ -79,7 +80,17 @@ COPILOT_MINTED_KEY: Final = "tid=device-login;minted" COPILOT_REJECTED_KEY: Final = "tid=device-login;rejected" CONTROL_KEY: Final = "device-login-control-key" ACCEPTED_BEARERS: Final = frozenset( - {CHATGPT_STORED, CHATGPT_REFRESHED, CHATGPT_FIRST_LOGIN, COPILOT_STORED_KEY, COPILOT_MINTED_KEY, CONTROL_KEY} + { + CHATGPT_STORED, + CHATGPT_REFRESHED, + CHATGPT_FIRST_LOGIN, + COPILOT_ACCESS, + COPILOT_FIRST_LOGIN_ACCESS, + COPILOT_UNSAVED_KEY_ACCESS, + COPILOT_STORED_KEY, + COPILOT_MINTED_KEY, + CONTROL_KEY, + } ) REJECTED_DETAIL: Final = "the scripted provider rejected this bearer token" _DEVICE_AUTH_ID: Final = "device-auth-synthetic" @@ -180,6 +191,10 @@ def _auth_hosts(api_url: str, grant: threading.Event) -> Callable[[Request], Rep return _json(200, {"access_token": COPILOT_FIRST_LOGIN_ACCESS, "token_type": "bearer"}) case ("api.github.com", "GET", "/copilot_internal/v2/token"): return _copilot_api_key(request, api_url) + case ("api.github.com", "GET", "/copilot_internal/user"): + if request.headers.get("authorization") == f"token {COPILOT_REJECTED_ACCESS}": + return _json(401, {"message": "Bad credentials"}) + return _json(200, {"endpoints": {"api": f"https://{COPILOT_API_HOST}"}}) case _: return _switched_off() @@ -204,9 +219,22 @@ class Switches: switch.clear() +@dataclass(frozen=True, slots=True) +class _ProviderAPI: + plain: Wire + secure: Wire + + @property + def url(self) -> str: + return self.plain.url + + def drain(self) -> tuple[Request, ...]: + return (*self.plain.drain(), *self.secure.drain()) + + @dataclass(frozen=True, slots=True) class Peers: - api: Wire + api: _ProviderAPI auth: Wire tunnel: str cert: Path @@ -228,6 +256,14 @@ class Peers: def auth_connections(self) -> tuple[str, ...]: return tuple(self.authorities.get_nowait() for _ in range(self.authorities.qsize())) + def assert_only_copilot_discovery(self, access_token: str) -> None: + assert set(self.auth_connections()) <= {f"{GITHUB_API_HOST}:443"} + requests: Final = self.auth.drain() + assert { + (request.method, request.target, request.headers.get("host"), request.headers.get("authorization")) + for request in requests + } <= {("GET", "/copilot_internal/user", GITHUB_API_HOST, f"token {access_token}")}, requests + def dropped(self) -> tuple[Hangup, ...]: return tuple(self.hangups.get_nowait() for _ in range(self.hangups.qsize())) @@ -279,9 +315,8 @@ class _TunnelServer(socketserver.ThreadingTCPServer): @contextmanager def _tunnel( - destination: Wire, switches: Switches, authorities: SimpleQueue[str], hangups: SimpleQueue[Hangup] + destination: Wire, provider: Wire, switches: Switches, authorities: SimpleQueue[str], hangups: SimpleQueue[Hangup] ) -> Generator[str]: - destination_port: Final = int(destination.url.rsplit(":", 1)[1]) class Tunnel(socketserver.StreamRequestHandler): rbufsize = 0 @@ -306,20 +341,24 @@ def _tunnel( if len(request_line) < 2: return authority: Final = request_line[1] - if authority.rsplit(":", 1)[0] not in AUTH_HOSTS: + host: Final = authority.rsplit(":", 1)[0] + if host not in (*AUTH_HOSTS, COPILOT_API_HOST): self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") return - authorities.put(authority) - if switches.refuse_connect.is_set(): + is_auth: Final = host in AUTH_HOSTS + if is_auth: + authorities.put(authority) + if is_auth and switches.refuse_connect.is_set(): self.wfile.write(b"HTTP/1.1 502 Bad Gateway\r\ncontent-length: 0\r\n\r\n") return - if self.hold(authority, opened, switches.hold_connect) is None: + if is_auth and self.hold(authority, opened, switches.hold_connect) is None: return self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n") - client_hello: Final = self.hold(authority, opened, switches.stall_tls) + client_hello: Final = self.hold(authority, opened, switches.stall_tls) if is_auth else b"" if client_hello is None: return self.request.settimeout(10) + destination_port: Final = int((destination if is_auth else provider).url.rsplit(":", 1)[1]) with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream: upstream.sendall(client_hello) outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream)) @@ -340,16 +379,17 @@ def _tunnel( @contextmanager def device_login_peers(directory: Path) -> Generator[Peers]: - cert, key = write_self_signed_cert(directory, AUTH_HOSTS) + cert, key = write_self_signed_cert(directory, (*AUTH_HOSTS, COPILOT_API_HOST)) switches: Final = Switches(threading.Event(), threading.Event(), threading.Event(), threading.Event()) authorities: Final[SimpleQueue[str]] = SimpleQueue() hangups: Final[SimpleQueue[Hangup]] = SimpleQueue() with ( wire_server(_provider_api(rv.ResponsesVendor())) as api, + wire_server(_provider_api(rv.ResponsesVendor()), tls=server_context(cert, key)) as copilot_api, wire_server(_auth_hosts(api.url, switches.grant), tls=server_context(cert, key)) as auth, - _tunnel(auth, switches, authorities, hangups) as tunnel, + _tunnel(auth, copilot_api, switches, authorities, hangups) as tunnel, ): - yield Peers(api, auth, tunnel, cert, authorities, hangups, switches) + yield Peers(_ProviderAPI(api, copilot_api), auth, tunnel, cert, authorities, hangups, switches) def chatgpt_record( diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 91f502179f3..384d6635bb0 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -67,9 +67,9 @@ class _TokenDirs: def environment(self) -> Mapping[str, str]: return MappingProxyType({"CHATGPT_TOKEN_DIR": str(self.chatgpt), "GITHUB_COPILOT_TOKEN_DIR": str(self.copilot)}) - def store_valid_tokens(self, api_url: str) -> None: + def store_valid_tokens(self) -> None: dl.write_chatgpt(self.chatgpt, dl.chatgpt_record(dl.CHATGPT_STORED)) - dl.write_copilot_key(self.copilot, dl.COPILOT_STORED_KEY, api_url) + dl.write_copilot_access(self.copilot, dl.COPILOT_ACCESS) def _token_dirs(directory: Path) -> _TokenDirs: @@ -242,16 +242,16 @@ def test_booting_without_tokens_lists_login_deployments_and_serves_only_after_to assert peers.auth_connections() == () assert dl.bearer(peers.api.drain()[-1]) == dl.CONTROL_KEY - dirs.store_valid_tokens(peers.api.url) + dirs.store_valid_tokens() _reload(owned) _every_worker_lists(owned, workers) for served_logins, name in enumerate(_LOGIN_NAMES, start=1): served: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert served.status == 200, served.text _successes_logged(key, at_least=_CONTROL_CALLS_BEFORE_THE_RELOAD + served_logins) - assert peers.auth_connections() == () + peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) forwarded: Final = peers.api.drain() - assert {dl.bearer(request) for request in forwarded} == {dl.CHATGPT_STORED, dl.COPILOT_STORED_KEY} + assert {dl.bearer(request) for request in forwarded} == {dl.CHATGPT_STORED, dl.COPILOT_ACCESS} dl.write_chatgpt(dirs.chatgpt, dl.chatgpt_record(dl.CHATGPT_EXPIRED, refresh_token=dl.GOOD_REFRESH)) peers.switches.hold_connect.set() @@ -287,7 +287,7 @@ def test_a_token_removed_after_boot_turns_every_request_into_the_refusal( ) -> None: dirs: Final = _token_dirs(tmp_path) with dl.device_login_peers(tmp_path) as peers: - dirs.store_valid_tokens(peers.api.url) + dirs.store_valid_tokens() with owned_proxy_process( gateway, tmp_path, @@ -310,7 +310,7 @@ def test_a_token_removed_after_boot_turns_every_request_into_the_refusal( refused: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert refused.status == 400, refused.text assert refusal in refused.message(), refused.text - assert peers.auth_connections() == () + peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) tail: Final = owned.log.read_text() assert dl.CHATGPT_USER_CODE not in tail and dl.COPILOT_USER_CODE not in tail, tail[-3000:] control: Final = _chat(_base_url(owned), key, _CONTROL, uuid.uuid4().hex) diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index 39e36fd95fa..9074ea83189 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -2,7 +2,6 @@ import asyncio import hashlib import itertools import json -import re import threading import time import uuid @@ -26,9 +25,10 @@ from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.timing import assert_cut_short, timed from integration._support.wire import Request -from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS from pydantic import JsonValue +from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS + pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] Provider: TypeAlias = Literal["chatgpt", "copilot"] @@ -52,7 +52,7 @@ _REFUSALS: Final[Mapping[Provider, str]] = MappingProxyType( {"chatgpt": dl.CHATGPT_REFUSAL, "copilot": dl.COPILOT_REFUSAL} ) _STORED_BEARERS: Final[Mapping[Provider, str]] = MappingProxyType( - {"chatgpt": dl.CHATGPT_STORED, "copilot": dl.COPILOT_STORED_KEY} + {"chatgpt": dl.CHATGPT_STORED, "copilot": dl.COPILOT_ACCESS} ) _CONTROL_MODEL: Final = "device-login-control" _EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) @@ -154,7 +154,7 @@ class _Scene: if provider == "chatgpt": dl.write_chatgpt(self.rig.chatgpt_dir, dl.chatgpt_record(dl.CHATGPT_STORED)) return - dl.write_copilot_key(self.rig.copilot_dir, dl.COPILOT_STORED_KEY, self.rig.peers.api.url) + dl.write_copilot_access(self.rig.copilot_dir, dl.COPILOT_ACCESS) def _config(peers: dl.Peers, directory: Path) -> Path: @@ -403,7 +403,7 @@ def test_missing_token_is_refused_with_the_login_instructions_before_any_auth_ho @pytest.mark.parametrize("cell", _CELLS, ids=_Cell.name) -def test_stored_token_serves_the_request_without_an_auth_host_call(scene: _Scene, cell: _Cell) -> None: +def test_stored_token_serves_the_request_without_a_device_login(scene: _Scene, cell: _Cell) -> None: scene.store_valid_token(cell.provider) call: Final = scene.call(cell) served: Final = _serve(call) @@ -412,7 +412,10 @@ def test_stored_token_serves_the_request_without_an_auth_host_call(scene: _Scene received: Final = scene.rig.peers.api.drain() (forwarded,) = _carrying(received, call.marker) assert dl.bearer(forwarded) == _STORED_BEARERS[cell.provider], _describe(received) - assert scene.rig.peers.auth_connections() == () + if cell.provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () (row,) = _spend_rows(scene.key) assert row["status"] == "success", row @@ -487,7 +490,7 @@ def test_missing_token_is_refused_on_the_other_model_routes(scene: _Scene, route assert _describe(scene.rig.peers.api.drain()) == () -def test_copilot_embeddings_are_served_with_the_stored_api_key(scene: _Scene) -> None: +def test_copilot_embeddings_are_served_with_the_stored_access_token(scene: _Scene) -> None: scene.store_valid_token("copilot") served: Final = scene.post( "/v1/embeddings", {"model": "github_copilot/text-embedding-3-small", "input": "login", **_EXTRA} @@ -496,9 +499,9 @@ def test_copilot_embeddings_are_served_with_the_stored_api_key(scene: _Scene) -> (item,) = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(served.text)["data"]) assert item["embedding"] == [0.25, 0.5, 0.75], served.text (forwarded,) = scene.rig.peers.api.drain() - assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_STORED_KEY), forwarded.target + assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_ACCESS), forwarded.target assert forwarded.target.endswith("/embeddings"), forwarded.target - assert scene.rig.peers.auth_connections() == () + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) @pytest.mark.parametrize("provider", _PROVIDERS) @@ -562,71 +565,67 @@ def test_stored_token_the_provider_rejects_surfaces_its_401_without_a_login( if provider == "chatgpt": dl.write_chatgpt(scene.rig.chatgpt_dir, dl.chatgpt_record(dl.CHATGPT_REJECTED)) else: - dl.write_copilot_key(scene.rig.copilot_dir, dl.COPILOT_REJECTED_KEY, scene.rig.peers.api.url) + dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_REJECTED_ACCESS) call: Final = scene.call(_Cell(provider, endpoint, False, "httpx")) served: Final = _serve(call) assert served.status == 401, served.text assert dl.REJECTED_DETAIL in served.text, served.text received: Final = scene.rig.peers.api.drain() assert {dl.bearer(request) for request in _carrying(received, call.marker)} == { - dl.CHATGPT_REJECTED if provider == "chatgpt" else dl.COPILOT_REJECTED_KEY + dl.CHATGPT_REJECTED if provider == "chatgpt" else dl.COPILOT_REJECTED_ACCESS }, _describe(received) - assert scene.rig.peers.auth_connections() == () + if provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_REJECTED_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () @pytest.mark.parametrize("stale_key", (False, True), ids=("missing-api-key", "expired-api-key")) -def test_copilot_api_key_is_minted_once_from_the_stored_access_token(scene: _Scene, stale_key: bool) -> None: +def test_copilot_access_token_is_forwarded_without_minting_an_api_key(scene: _Scene, stale_key: bool) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_ACCESS) if stale_key: dl.write_copilot_key(scene.rig.copilot_dir, dl.COPILOT_REJECTED_KEY, scene.rig.peers.api.url, -3600) + legacy_key: Final = scene.rig.copilot_dir / "api-key.json" + stored_key: Final = legacy_key.read_bytes() if stale_key else None call: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) served: Final = _serve(call) assert served.status == 200, served.text assert set(rv.MARKER.findall(served.text)) == {call.marker}, served.text - (minted,) = scene.rig.peers.auth.drain() - assert (minted.method, minted.target, minted.headers.get("host"), minted.headers.get("authorization")) == ( - "GET", - "/copilot_internal/v2/token", - dl.GITHUB_API_HOST, - f"token {dl.COPILOT_ACCESS}", - ) + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) (forwarded,) = _carrying(scene.rig.peers.api.drain(), call.marker) - assert dl.bearer(forwarded) == dl.COPILOT_MINTED_KEY - assert json.loads((scene.rig.copilot_dir / "api-key.json").read_text())["token"] == dl.COPILOT_MINTED_KEY + assert dl.bearer(forwarded) == dl.COPILOT_ACCESS + assert (legacy_key.read_bytes() if legacy_key.exists() else None) == stored_key -def test_copilot_api_key_that_cannot_be_saved_is_minted_at_most_once_per_worker(scene: _Scene) -> None: +def test_copilot_access_token_serves_when_the_legacy_api_key_path_is_a_directory(scene: _Scene) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_UNSAVED_KEY_ACCESS) unwritable_key_file: Final = scene.rig.copilot_dir / "api-key.json" unwritable_key_file.mkdir() first: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) first_served: Final = _serve(first) assert first_served.status == 200, first_served.text - (minted,) = scene.rig.peers.auth.drain() - assert minted.headers.get("authorization") == f"token {dl.COPILOT_UNSAVED_KEY_ACCESS}" later: Final = tuple(scene.call(_Cell("copilot", "chat", False, "httpx")) for _ in range(3)) later_served: Final = tuple(_serve(call) for call in later) assert [served.status for served in later_served] == [200, 200, 200], [served.text for served in later_served] - assert len(scene.rig.peers.auth.drain()) <= 1 + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_UNSAVED_KEY_ACCESS) received: Final = scene.rig.peers.api.drain() forwarded: Final = tuple(_carrying(received, call.marker) for call in (first, *later)) assert [tuple(dl.bearer(request) for request in requests) for requests in forwarded] == [ - (dl.COPILOT_MINTED_KEY,) + (dl.COPILOT_UNSAVED_KEY_ACCESS,) ] * 4, _describe(received) assert unwritable_key_file.is_dir() def test_copilot_access_token_github_rejects_fails_the_request_without_a_device_login(scene: _Scene) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_REJECTED_ACCESS) - served: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx"))) - assert served.status == 400, served.text - assert "Failed to refresh API key" in _error_message(served.text), served.text + call: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) + served: Final = _serve(call) + assert served.status == 401, served.text + assert dl.REJECTED_DETAIL in _error_message(served.text), served.text _failure_logged(served) - reached: Final = scene.rig.peers.auth.drain() - assert {(request.method, request.target) for request in reached} == {("GET", "/copilot_internal/v2/token")}, ( - _describe(reached) - ) - assert _describe(scene.rig.peers.api.drain()) == () + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_REJECTED_ACCESS) + (forwarded,) = _carrying(scene.rig.peers.api.drain(), call.marker) + assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_REJECTED_ACCESS) def test_expired_copilot_api_key_with_no_access_token_is_refused_before_any_auth_host_call(scene: _Scene) -> None: @@ -727,7 +726,10 @@ def test_a_token_mounted_after_a_refusal_serves_without_a_restart_and_removing_i after: Final = _serve(scene.call(cell)) assert after.status == 400, after.text assert _REFUSALS[provider] in _error_message(after.text), after.text - assert scene.rig.peers.auth_connections() == () + if provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () _COOLDOWN_SECONDS_LEFT: Final = 45 diff --git a/tests/integration/sdk/test_device_code_login_guard_sdk.py b/tests/integration/sdk/test_device_code_login_guard_sdk.py index 7640fe13d5d..256871faf7c 100644 --- a/tests/integration/sdk/test_device_code_login_guard_sdk.py +++ b/tests/integration/sdk/test_device_code_login_guard_sdk.py @@ -184,7 +184,7 @@ def clean(bench: _Bench) -> Iterator[_Bench]: _FIRST_LOGIN_BEARERS: Final[Mapping[Provider, str]] = MappingProxyType( - {"chatgpt": dl.CHATGPT_FIRST_LOGIN, "copilot": dl.COPILOT_MINTED_KEY} + {"chatgpt": dl.CHATGPT_FIRST_LOGIN, "copilot": dl.COPILOT_FIRST_LOGIN_ACCESS} ) @@ -205,7 +205,7 @@ def test_sync_call_on_the_main_thread_keeps_the_interactive_first_login( assert {dl.bearer(request) for request in forwarded} == {_FIRST_LOGIN_BEARERS[provider]}, [ (request.method, request.target) for request in forwarded ] - expected_files: Final = ("auth.json",) if provider == "chatgpt" else ("access-token", "api-key.json") + expected_files: Final = ("auth.json",) if provider == "chatgpt" else ("access-token",) assert clean.token_files() == expected_files