scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

test(copilot): align login guards with direct access tokens (#50285)

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mateo Wang 2026-11-05 15:38:33 -08:00 • committed by GitHub
parent e011a4ea9b
commit 4679b9f642
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 99 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
tests/integration/_support/device_login.py
View file

@ -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(

14
tests/integration/providers/test_device_code_login_guard_boot.py
View file

@ -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)

74
tests/integration/providers/test_device_code_login_guard_wire.py
View file

@ -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

4
tests/integration/sdk/test_device_code_login_guard_sdk.py
View file

@ -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