mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
fix(github_copilot): treat a malformed api-key.json like a missing one (#44723)
* fix(github_copilot): treat a malformed api-key.json like a missing one * fix(github_copilot): read the key and the endpoints of api-key.json as separate strict views * fix(github_copilot): ignore a non-UTF-8 api-key.json like any other malformed file Read the cached key file as bytes so a UTF-16 file, which Windows PowerShell 5.1 writes with >, fails JSON validation instead of raising the decoder's error, and the request gets the refresh or the login hint * fix(github_copilot): decode UTF-16 key files and refresh over a malformed endpoints field * fix(github_copilot): decode a UTF-16 access-token file and refuse an undecodable one like a missing one * fix(github_copilot): use the refreshed key when api-key.json cannot be written * fix(github_copilot): keep a refreshed key that cannot be saved in memory * fix(github_copilot): serve a kept key and its endpoint together after the key file is repaired A key held in memory is now checked before api-key.json, so the token and the endpoint come from the same source, and the map holds one entry per key file with the access token beside it * test(github_copilot): type the authenticator test fixture and the tests this branch adds --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
184ec7b894
commit
bee9f679dc
4 changed files with 377 additions and 83 deletions
|
|
@ -1,11 +1,12 @@
|
|||
import json
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -13,7 +14,6 @@ from litellm.litellm_core_utils.asyncify import can_block_current_thread
|
|||
from litellm.llms.custom_httpx.http_handler import get_httpx_client
|
||||
|
||||
from .common_utils import (
|
||||
APIKeyExpiredError,
|
||||
GetAccessTokenError,
|
||||
GetAPIKeyError,
|
||||
GetDeviceCodeError,
|
||||
|
|
@ -27,6 +27,44 @@ DEFAULT_GITHUB_ACCESS_TOKEN_URL: Final = "https://github.com/login/oauth/access_
|
|||
DEFAULT_GITHUB_API_KEY_URL: Final = "https://api.github.com/copilot_internal/v2/token"
|
||||
|
||||
|
||||
class CopilotEndpoints(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True)
|
||||
|
||||
api: str | None = None
|
||||
|
||||
|
||||
class CopilotCachedKey(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True)
|
||||
|
||||
token: str
|
||||
expires_at: float
|
||||
endpoints: CopilotEndpoints | None = None
|
||||
|
||||
|
||||
class CopilotCachedEndpoints(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True)
|
||||
|
||||
endpoints: CopilotEndpoints | None = None
|
||||
|
||||
|
||||
ApiKeyFileView = TypeVar("ApiKeyFileView", CopilotCachedKey, CopilotCachedEndpoints)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _UnsavedApiKey:
|
||||
access_token: str
|
||||
key: CopilotCachedKey
|
||||
|
||||
|
||||
_UNSAVED_API_KEYS: Final[dict[str, _UnsavedApiKey]] = {} # mutable-ok: refreshed keys api-key.json could not hold
|
||||
|
||||
|
||||
def _read_token_file(path: str) -> str:
|
||||
with open(path, "rb") as f:
|
||||
raw: Final = f.read()
|
||||
return raw.decode(json.detect_encoding(raw))
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _DeviceCode(TypedDict):
|
||||
device_code: ReadOnly[str]
|
||||
|
|
@ -87,12 +125,13 @@ class Authenticator:
|
|||
GetAccessTokenError: If unable to obtain an access token after retries.
|
||||
"""
|
||||
try:
|
||||
with open(self.access_token_file, "r") as f:
|
||||
access_token = f.read().strip()
|
||||
if access_token:
|
||||
return access_token
|
||||
saved_token: Final = _read_token_file(self.access_token_file).strip()
|
||||
if saved_token:
|
||||
return saved_token
|
||||
except OSError:
|
||||
verbose_logger.warning("No existing access token found or error reading file")
|
||||
except ValueError as e:
|
||||
verbose_logger.warning("Ignoring %s, not a text access token: %s", self.access_token_file, e)
|
||||
|
||||
if not can_block_current_thread():
|
||||
raise GetAccessTokenError(
|
||||
|
|
@ -135,64 +174,83 @@ class Authenticator:
|
|||
Raises:
|
||||
GetAPIKeyError: If unable to obtain an API key.
|
||||
"""
|
||||
try:
|
||||
with open(self.api_key_file, "r") as f:
|
||||
api_key_info = json.load(f)
|
||||
if api_key_info.get("expires_at", 0) > datetime.now().timestamp():
|
||||
return api_key_info.get("token")
|
||||
else:
|
||||
verbose_logger.warning("API key expired, refreshing")
|
||||
raise APIKeyExpiredError(
|
||||
message="API key expired",
|
||||
status_code=401,
|
||||
)
|
||||
except OSError:
|
||||
verbose_logger.warning("No API key file found or error opening file")
|
||||
except (json.JSONDecodeError, KeyError) as e:
|
||||
verbose_logger.warning("Error reading API key from file: %s", e)
|
||||
except APIKeyExpiredError:
|
||||
pass # Already logged in the try block
|
||||
unsaved: Final = self._unsaved_api_key()
|
||||
if unsaved is not None:
|
||||
return unsaved.token
|
||||
cached: Final = self._read_api_key_file(CopilotCachedKey)
|
||||
if cached is not None and cached.expires_at > datetime.now().timestamp():
|
||||
return cached.token
|
||||
if cached is not None:
|
||||
verbose_logger.warning("API key expired, refreshing")
|
||||
|
||||
try:
|
||||
api_key_info = self._refresh_api_key()
|
||||
with open(self.api_key_file, "w") as f:
|
||||
json.dump(api_key_info, f)
|
||||
token: Final = api_key_info.get("token")
|
||||
if token:
|
||||
return token
|
||||
else:
|
||||
raise GetAPIKeyError(
|
||||
message="API key response missing token",
|
||||
status_code=401,
|
||||
)
|
||||
except OSError as e:
|
||||
verbose_logger.error("Error saving API key to file: %s", e)
|
||||
raise GetAPIKeyError(
|
||||
message=f"Failed to save API key: {e}",
|
||||
status_code=500,
|
||||
)
|
||||
api_key_info: Final = self._refresh_api_key()
|
||||
except RefreshAPIKeyError as e:
|
||||
raise GetAPIKeyError(
|
||||
message=f"Failed to refresh API key: {e}",
|
||||
status_code=401,
|
||||
) from e
|
||||
serialized: Final = json.dumps(api_key_info)
|
||||
try:
|
||||
with open(self.api_key_file, "w") as f:
|
||||
f.write(serialized)
|
||||
except OSError as e:
|
||||
verbose_logger.warning(
|
||||
"Keeping the refreshed API key in memory, it could not be saved to %s: %s", self.api_key_file, e
|
||||
)
|
||||
self._keep_unsaved_api_key(serialized)
|
||||
token: Final = api_key_info.get("token")
|
||||
if not token:
|
||||
raise GetAPIKeyError(
|
||||
message="API key response missing token",
|
||||
status_code=401,
|
||||
)
|
||||
return token
|
||||
|
||||
def get_api_base(self) -> str | None:
|
||||
"""
|
||||
Get the API endpoint from the api-key.json file.
|
||||
Get the API endpoint of the current API key.
|
||||
|
||||
Returns:
|
||||
Optional[str]: The GitHub Copilot API endpoint, or None if not found.
|
||||
"""
|
||||
try:
|
||||
with open(self.api_key_file, "r") as f:
|
||||
api_key_info: Final = json.load(f)
|
||||
endpoints: Final = api_key_info.get("endpoints", {})
|
||||
api_endpoint: Final = endpoints.get("api")
|
||||
return api_endpoint
|
||||
except (OSError, json.JSONDecodeError, KeyError) as e:
|
||||
verbose_logger.warning("Error reading API endpoint from file: %s", e)
|
||||
cached: Final = self._unsaved_api_key() or self._read_api_key_file(CopilotCachedEndpoints)
|
||||
if cached is None or cached.endpoints is None:
|
||||
return None
|
||||
return cached.endpoints.api
|
||||
|
||||
def _read_api_key_file(self, view: type[ApiKeyFileView]) -> ApiKeyFileView | None:
|
||||
try:
|
||||
return view.model_validate_json(_read_token_file(self.api_key_file))
|
||||
except OSError:
|
||||
verbose_logger.warning("No API key file found or error opening file")
|
||||
return None
|
||||
except ValueError as e:
|
||||
verbose_logger.warning("Ignoring %s, not a Copilot API key file: %s", self.api_key_file, e)
|
||||
return None
|
||||
|
||||
def _stored_access_token(self) -> str | None:
|
||||
try:
|
||||
access_token: Final = _read_token_file(self.access_token_file).strip()
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
return access_token or None
|
||||
|
||||
def _unsaved_api_key(self) -> CopilotCachedKey | None:
|
||||
unsaved: Final = _UNSAVED_API_KEYS.get(self.api_key_file)
|
||||
if unsaved is None or unsaved.key.expires_at <= time.time():
|
||||
return None
|
||||
return unsaved.key if unsaved.access_token == self._stored_access_token() else None
|
||||
|
||||
def _keep_unsaved_api_key(self, serialized: str) -> None:
|
||||
access_token: Final = self._stored_access_token()
|
||||
if access_token is None:
|
||||
return
|
||||
try:
|
||||
key: Final = CopilotCachedKey.model_validate_json(serialized)
|
||||
except ValueError:
|
||||
return
|
||||
_UNSAVED_API_KEYS[self.api_key_file] = _UnsavedApiKey(access_token=access_token, key=key)
|
||||
|
||||
def _refresh_api_key(self) -> dict[str, Any]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -73,6 +73,7 @@ COPILOT_USER_CODE: Final = "COPI-LEAK"
|
|||
COPILOT_ACCESS: Final = "device-login-github-access"
|
||||
COPILOT_REJECTED_ACCESS: Final = "device-login-github-access-rejected"
|
||||
COPILOT_FIRST_LOGIN_ACCESS: Final = "device-login-github-access-first-login"
|
||||
COPILOT_UNSAVED_KEY_ACCESS: Final = "device-login-github-access-unsaved-key"
|
||||
COPILOT_STORED_KEY: Final = "tid=device-login;stored"
|
||||
COPILOT_MINTED_KEY: Final = "tid=device-login;minted"
|
||||
COPILOT_REJECTED_KEY: Final = "tid=device-login;rejected"
|
||||
|
|
@ -132,7 +133,11 @@ def _oauth_token(request: Request, grant: threading.Event) -> Reply:
|
|||
|
||||
|
||||
def _copilot_api_key(request: Request, api_url: str) -> Reply:
|
||||
if request.headers.get("authorization") not in (f"token {COPILOT_ACCESS}", f"token {COPILOT_FIRST_LOGIN_ACCESS}"):
|
||||
if request.headers.get("authorization") not in (
|
||||
f"token {COPILOT_ACCESS}",
|
||||
f"token {COPILOT_FIRST_LOGIN_ACCESS}",
|
||||
f"token {COPILOT_UNSAVED_KEY_ACCESS}",
|
||||
):
|
||||
return _json(401, {"message": "Bad credentials"})
|
||||
return _json(
|
||||
200,
|
||||
|
|
@ -380,4 +385,7 @@ def write_copilot_access(directory: Path, access_token: str) -> None:
|
|||
def clear_tokens(*directories: Path) -> None:
|
||||
for directory in directories:
|
||||
for path in directory.iterdir():
|
||||
path.unlink()
|
||||
if path.is_dir():
|
||||
path.rmdir()
|
||||
else:
|
||||
path.unlink()
|
||||
|
|
|
|||
|
|
@ -595,6 +595,27 @@ def test_copilot_api_key_is_minted_once_from_the_stored_access_token(scene: _Sce
|
|||
assert json.loads((scene.rig.copilot_dir / "api-key.json").read_text())["token"] == dl.COPILOT_MINTED_KEY
|
||||
|
||||
|
||||
def test_copilot_api_key_that_cannot_be_saved_is_minted_at_most_once_per_worker(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
|
||||
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,)
|
||||
] * 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")))
|
||||
|
|
@ -655,44 +676,39 @@ def test_chatgpt_token_with_a_text_expiry_falls_back_to_the_expiry_inside_the_to
|
|||
assert scene.rig.peers.auth_connections() == ()
|
||||
|
||||
|
||||
_PROVIDER_RESOLUTION_ERROR: Final = "GetLLMProvider Exception"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _UnusableCopilotFiles:
|
||||
files: Mapping[str, str]
|
||||
answer: str
|
||||
|
||||
|
||||
_UNUSABLE_COPILOT_FILES: Final[Mapping[str, _UnusableCopilotFiles]] = MappingProxyType(
|
||||
_UNUSABLE_COPILOT_FILES: Final[Mapping[str, Mapping[str, bytes]]] = MappingProxyType(
|
||||
{
|
||||
"empty-access-token": _UnusableCopilotFiles(MappingProxyType({"access-token": ""}), dl.COPILOT_REFUSAL),
|
||||
"blank-access-token": _UnusableCopilotFiles(MappingProxyType({"access-token": " \n"}), dl.COPILOT_REFUSAL),
|
||||
"invalid-api-key-json": _UnusableCopilotFiles(
|
||||
MappingProxyType({"api-key.json": "{not json"}), dl.COPILOT_REFUSAL
|
||||
"empty-access-token": MappingProxyType({"access-token": b""}),
|
||||
"blank-access-token": MappingProxyType({"access-token": b" \n"}),
|
||||
"access-token-truncated-utf-16": MappingProxyType({"access-token": b"\xff\xfe\x00\xd8"}),
|
||||
"invalid-api-key-json": MappingProxyType({"api-key.json": b"{not json"}),
|
||||
"api-key-list": MappingProxyType({"api-key.json": b"[]"}),
|
||||
"api-key-expiry-text": MappingProxyType(
|
||||
{"api-key.json": json.dumps({"token": "copilot-text-expiry", "expires_at": "soon"}).encode()}
|
||||
),
|
||||
"api-key-list": _UnusableCopilotFiles(MappingProxyType({"api-key.json": "[]"}), _PROVIDER_RESOLUTION_ERROR),
|
||||
"api-key-expiry-text": _UnusableCopilotFiles(
|
||||
MappingProxyType({"api-key.json": json.dumps({"token": "copilot-text-expiry", "expires_at": "soon"})}),
|
||||
_PROVIDER_RESOLUTION_ERROR,
|
||||
"api-key-expired-utf-16": MappingProxyType(
|
||||
{"api-key.json": json.dumps({"token": "copilot-utf-16", "expires_at": 1}).encode("utf-16")}
|
||||
),
|
||||
"api-key-truncated-utf-16": MappingProxyType({"api-key.json": b"\xff\xfe\x00\xd8"}),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", tuple(_UNUSABLE_COPILOT_FILES))
|
||||
def test_unusable_copilot_token_files_are_refused_like_missing_ones_each_time(scene: _Scene, name: str) -> None:
|
||||
variant: Final = _UNUSABLE_COPILOT_FILES[name]
|
||||
for file_name, content in variant.files.items():
|
||||
(scene.rig.copilot_dir / file_name).write_text(content)
|
||||
files: Final = _UNUSABLE_COPILOT_FILES[name]
|
||||
for file_name, content in files.items():
|
||||
(scene.rig.copilot_dir / file_name).write_bytes(content)
|
||||
first: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx")))
|
||||
second: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx")))
|
||||
for served in (first, second):
|
||||
assert served.status == 400, served.text
|
||||
assert variant.answer in _error_message(served.text), served.text
|
||||
assert dl.COPILOT_REFUSAL in _error_message(served.text), served.text
|
||||
assert dl.COPILOT_USER_CODE not in served.text
|
||||
assert scene.rig.peers.auth_connections() == ()
|
||||
assert _describe(scene.rig.peers.api.drain()) == ()
|
||||
for file_name, content in files.items():
|
||||
assert (scene.rig.copilot_dir / file_name).read_bytes() == content
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _PROVIDERS)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,16 @@
|
|||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from litellm.llms.github_copilot.authenticator import Authenticator
|
||||
from litellm.llms.github_copilot.common_utils import (
|
||||
|
|
@ -16,6 +21,20 @@ from litellm.llms.github_copilot.common_utils import (
|
|||
RefreshAPIKeyError,
|
||||
)
|
||||
|
||||
API_KEY_URL: Final = "https://copilot-key.test/copilot_internal/v2/token"
|
||||
REFRESHED_KEY: Final = MappingProxyType({"token": "new-api-key", "expires_at": 4102444800})
|
||||
TENANT_API: Final = "https://api.tenant.githubcopilot.com"
|
||||
REPAIRED_API: Final = "https://api.repaired.githubcopilot.com"
|
||||
TRUNCATED_UTF_16: Final = b"\xff\xfe\x00\xd8"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def logged_in_token_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", API_KEY_URL)
|
||||
(tmp_path / "access-token").write_text("gho_logged_in")
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestGitHubCopilotAuthenticator:
|
||||
@pytest.fixture
|
||||
|
|
@ -69,13 +88,9 @@ class TestGitHubCopilotAuthenticator:
|
|||
headers_with_token = authenticator._get_github_headers("test-token")
|
||||
assert headers_with_token["authorization"] == "token test-token"
|
||||
|
||||
def test_get_access_token_from_file(self, authenticator):
|
||||
def test_get_access_token_from_file(self, logged_in_token_dir: Path) -> None:
|
||||
"""Test retrieving an access token from a file."""
|
||||
mock_token = "mock-access-token"
|
||||
|
||||
with patch("builtins.open", mock_open(read_data=mock_token)):
|
||||
token = authenticator.get_access_token()
|
||||
assert token == mock_token
|
||||
assert Authenticator().get_access_token() == "gho_logged_in"
|
||||
|
||||
def test_get_access_token_login(self, authenticator):
|
||||
"""Test logging in to get an access token."""
|
||||
|
|
@ -118,6 +133,30 @@ class TestGitHubCopilotAuthenticator:
|
|||
assert authenticator.access_token_file not in str(exc.value)
|
||||
mock_login.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"])
|
||||
def test_saved_access_token_in_another_encoding_is_sent_on_refresh(
|
||||
self, logged_in_token_dir: Path, encoding: str
|
||||
) -> None:
|
||||
"""An access-token file a Windows editor saved as UTF-16 or with a BOM still refreshes the key."""
|
||||
(logged_in_token_dir / "access-token").write_bytes("gho_logged_in\n".encode(encoding))
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
|
||||
assert Authenticator().get_api_key() == "new-api-key"
|
||||
|
||||
assert route.calls.last.request.headers["authorization"] == "token gho_logged_in"
|
||||
|
||||
def test_undecodable_access_token_is_refused_like_a_missing_one(self, logged_in_token_dir: Path) -> None:
|
||||
"""A proxy worker gets the login hint for an access-token file it cannot decode, and the file is kept."""
|
||||
access_token_file: Final = logged_in_token_dir / "access-token"
|
||||
access_token_file.write_bytes(TRUNCATED_UTF_16)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool, pytest.raises(GetAccessTokenError) as exc:
|
||||
pool.submit(Authenticator().get_access_token).result()
|
||||
|
||||
assert "worker thread" in str(exc.value)
|
||||
assert access_token_file.read_bytes() == TRUNCATED_UTF_16
|
||||
|
||||
def test_get_access_token_failure(self, authenticator):
|
||||
"""Test that an exception is raised after multiple login failures."""
|
||||
with (
|
||||
|
|
@ -137,7 +176,7 @@ class TestGitHubCopilotAuthenticator:
|
|||
future_time = (datetime.now() + timedelta(hours=1)).timestamp()
|
||||
mock_api_key_data = json.dumps(
|
||||
{"token": "mock-api-key", "expires_at": future_time}
|
||||
)
|
||||
).encode()
|
||||
|
||||
with patch("builtins.open", mock_open(read_data=mock_api_key_data)):
|
||||
api_key = authenticator.get_api_key()
|
||||
|
|
@ -148,7 +187,7 @@ class TestGitHubCopilotAuthenticator:
|
|||
past_time = (datetime.now() - timedelta(hours=1)).timestamp()
|
||||
mock_expired_data = json.dumps(
|
||||
{"token": "expired-api-key", "expires_at": past_time}
|
||||
)
|
||||
).encode()
|
||||
mock_new_data = {
|
||||
"token": "new-api-key",
|
||||
"expires_at": (datetime.now() + timedelta(hours=1)).timestamp(),
|
||||
|
|
@ -163,6 +202,179 @@ class TestGitHubCopilotAuthenticator:
|
|||
assert api_key == "new-api-key"
|
||||
authenticator._refresh_api_key.assert_called_once()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"malformed",
|
||||
[
|
||||
b"[]",
|
||||
json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(),
|
||||
TRUNCATED_UTF_16,
|
||||
json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(),
|
||||
],
|
||||
ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"],
|
||||
)
|
||||
def test_get_api_key_malformed_file_refreshes_like_missing(
|
||||
self, logged_in_token_dir: Path, malformed: bytes
|
||||
) -> None:
|
||||
"""A file that is not a Copilot key object is ignored, refreshed over, never answered with a parse error."""
|
||||
api_key_file: Final = logged_in_token_dir / "api-key.json"
|
||||
api_key_file.write_bytes(malformed)
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
|
||||
assert Authenticator().get_api_key() == "new-api-key"
|
||||
|
||||
assert route.call_count == 1
|
||||
assert route.calls.last.request.headers["authorization"] == "token gho_logged_in"
|
||||
assert json.loads(api_key_file.read_text()) == REFRESHED_KEY
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"malformed",
|
||||
[
|
||||
b"[]",
|
||||
json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(),
|
||||
TRUNCATED_UTF_16,
|
||||
json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(),
|
||||
],
|
||||
ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"],
|
||||
)
|
||||
def test_get_api_key_malformed_file_is_kept_when_refresh_fails(
|
||||
self, logged_in_token_dir: Path, malformed: bytes
|
||||
) -> None:
|
||||
"""Ignoring the file never deletes it, so a failed refresh leaves the operator's file for inspection."""
|
||||
api_key_file: Final = logged_in_token_dir / "api-key.json"
|
||||
api_key_file.write_bytes(malformed)
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
router.get(API_KEY_URL).respond(401)
|
||||
with pytest.raises(GetAPIKeyError):
|
||||
Authenticator().get_api_key()
|
||||
|
||||
assert api_key_file.read_bytes() == malformed
|
||||
|
||||
def test_refreshed_key_that_cannot_be_saved_is_reused_with_its_endpoint(self, logged_in_token_dir: Path) -> None:
|
||||
unwritable_key_file: Final = logged_in_token_dir / "api-key.json"
|
||||
unwritable_key_file.mkdir()
|
||||
first: Final = Authenticator()
|
||||
second: Final = Authenticator()
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route: Final = router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}})
|
||||
served: Final = (first.get_api_key(), second.get_api_key(), first.get_api_key())
|
||||
|
||||
assert served == ("new-api-key",) * 3
|
||||
assert (first.get_api_base(), second.get_api_base()) == (TENANT_API, TENANT_API)
|
||||
assert route.call_count == 1
|
||||
assert unwritable_key_file.is_dir()
|
||||
|
||||
def test_refreshed_key_that_cannot_be_saved_and_has_expired_is_refreshed_again(
|
||||
self, logged_in_token_dir: Path
|
||||
) -> None:
|
||||
(logged_in_token_dir / "api-key.json").mkdir()
|
||||
auth: Final = Authenticator()
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route: Final = router.get(API_KEY_URL).respond(json={"token": "expired-api-key", "expires_at": 1})
|
||||
served: Final = (auth.get_api_key(), auth.get_api_key())
|
||||
|
||||
assert served == ("expired-api-key",) * 2
|
||||
assert route.call_count == 2
|
||||
|
||||
def test_kept_key_and_its_endpoint_stay_paired_after_the_key_file_is_repaired(
|
||||
self, logged_in_token_dir: Path
|
||||
) -> None:
|
||||
key_file: Final = logged_in_token_dir / "api-key.json"
|
||||
key_file.mkdir()
|
||||
auth: Final = Authenticator()
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}})
|
||||
auth.get_api_key()
|
||||
key_file.rmdir()
|
||||
key_file.write_text(
|
||||
json.dumps({"token": "repaired-api-key", "expires_at": 9999999999, "endpoints": {"api": REPAIRED_API}})
|
||||
)
|
||||
|
||||
assert (auth.get_api_key(), auth.get_api_base()) in (
|
||||
("new-api-key", TENANT_API),
|
||||
("repaired-api-key", REPAIRED_API),
|
||||
)
|
||||
|
||||
def test_unsaved_key_is_not_reused_after_the_access_token_changes(self, logged_in_token_dir: Path) -> None:
|
||||
(logged_in_token_dir / "api-key.json").mkdir()
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
|
||||
Authenticator().get_api_key()
|
||||
(logged_in_token_dir / "access-token").write_text("gho_logged_in_again")
|
||||
Authenticator().get_api_key()
|
||||
|
||||
assert [call.request.headers["authorization"] for call in route.calls] == [
|
||||
"token gho_logged_in",
|
||||
"token gho_logged_in_again",
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"malformed",
|
||||
[
|
||||
b"[]",
|
||||
json.dumps({"endpoints": {"api": 5}}).encode(),
|
||||
json.dumps({"token": "t", "expires_at": 1, "endpoints": []}).encode(),
|
||||
TRUNCATED_UTF_16,
|
||||
],
|
||||
ids=["json-list", "number-api", "list-endpoints", "truncated-utf-16"],
|
||||
)
|
||||
def test_get_api_base_malformed_file_returns_none(
|
||||
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, malformed: bytes
|
||||
) -> None:
|
||||
"""get_api_base runs before get_api_key on every request, so it must survive malformed shapes too."""
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
|
||||
(tmp_path / "api-key.json").write_bytes(malformed)
|
||||
|
||||
assert Authenticator().get_api_base() is None
|
||||
|
||||
def test_get_api_base_keeps_endpoints_when_the_key_is_malformed(self, logged_in_token_dir: Path) -> None:
|
||||
"""The endpoint is read on its own, so a bad key never sends the refreshed key to the default host."""
|
||||
(logged_in_token_dir / "api-key.json").write_text(
|
||||
json.dumps({"token": "tid=abc", "expires_at": "soon", "endpoints": {"api": TENANT_API}})
|
||||
)
|
||||
auth: Final = Authenticator()
|
||||
|
||||
assert auth.get_api_base() == TENANT_API
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
|
||||
assert auth.get_api_key() == "new-api-key"
|
||||
|
||||
@pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"])
|
||||
def test_live_key_file_in_another_json_encoding_is_used_without_a_refresh(
|
||||
self, logged_in_token_dir: Path, encoding: str
|
||||
) -> None:
|
||||
"""A key file a Windows editor saved as UTF-16 or with a BOM still holds a live key and tenant."""
|
||||
live_key: Final = json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": {"api": TENANT_API}})
|
||||
api_key_file: Final = logged_in_token_dir / "api-key.json"
|
||||
api_key_file.write_bytes(live_key.encode(encoding))
|
||||
auth: Final = Authenticator()
|
||||
|
||||
with respx.mock:
|
||||
assert auth.get_api_base() == TENANT_API
|
||||
assert auth.get_api_key() == "tid=abc"
|
||||
|
||||
assert api_key_file.read_bytes() == live_key.encode(encoding)
|
||||
|
||||
def test_ignored_key_file_is_logged_without_its_contents(
|
||||
self, logged_in_token_dir: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""The warning names the file but never echoes a token saved in the wrong shape."""
|
||||
(logged_in_token_dir / "api-key.json").write_text(
|
||||
json.dumps({"token": ["tid=secret-value"], "expires_at": 4102444800})
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"), respx.mock(assert_all_called=True) as router:
|
||||
router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
|
||||
assert Authenticator().get_api_key() == "new-api-key"
|
||||
|
||||
assert "not a Copilot API key file" in caplog.text
|
||||
assert "tid=secret-value" not in caplog.text
|
||||
|
||||
def test_refresh_api_key(self, authenticator, mock_http_client):
|
||||
"""Test refreshing an API key."""
|
||||
mock_client, mock_response = mock_http_client
|
||||
|
|
@ -270,7 +482,7 @@ class TestGitHubCopilotAuthenticator:
|
|||
"expires_at": (datetime.now() + timedelta(hours=1)).timestamp(),
|
||||
"endpoints": {"api": "https://api.enterprise.githubcopilot.com"},
|
||||
}
|
||||
)
|
||||
).encode()
|
||||
with patch("builtins.open", mock_open(read_data=mock_api_key_data)):
|
||||
api_base = authenticator.get_api_base()
|
||||
assert api_base == "https://api.enterprise.githubcopilot.com"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue