diff --git a/cli/main.py b/cli/main.py index 7c13acaed..d6c6d6fbf 100644 --- a/cli/main.py +++ b/cli/main.py @@ -1,6 +1,9 @@ import sys +from enum import StrEnum +from typing import Annotated import typer +from rich.markup import escape from cli.display import console from cli.models import AnalystType, AssetType @@ -156,5 +159,39 @@ def backtest( console.print(f"[yellow]unsettled:[/yellow] {ticker}: {reason}") +class Plan(StrEnum): + chatgpt = "chatgpt" + + +@app.command() +def login(plan: Annotated[Plan, typer.Argument(help="The plan to sign in with")]): + """Sign in with your ChatGPT Plus or Pro plan, for the chatgpt provider.""" + from tradingagents.llm_clients import chatgpt_auth + + try: + # One unwrapped line, so the sign-in address can be copied whole. + email = chatgpt_auth.sign_in( + say=lambda line: console.print(line, soft_wrap=True, markup=False, highlight=False)) + except chatgpt_auth.ChatGPTAuthError as exc: + console.print(f"[red]{escape(str(exc))}[/red]") + raise typer.Exit(code=1) from None + console.print(f"[green]Signed in to ChatGPT as {escape(email)}.[/green]") + + +@app.command() +def logout(plan: Annotated[Plan, typer.Argument(help="The plan to sign out of")]): + """Sign out of your ChatGPT plan: revoke this machine's access and forget it.""" + from tradingagents.llm_clients import chatgpt_auth + + outcome = chatgpt_auth.sign_out() + if outcome == "revoked": + console.print("Signed out of ChatGPT.") + elif outcome == "forgotten": + console.print("[yellow]Signed out here, but OpenAI could not be reached to revoke the sign-in; " + "it stays valid at OpenAI until it expires.[/yellow]") + else: + console.print("Not signed in to ChatGPT.") + + if __name__ == "__main__": app() diff --git a/pyproject.toml b/pyproject.toml index 4db28fc5e..70dcce013 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ description = "TradingAgents: Multi-Agents LLM Financial Trading Framework" readme = "README.md" requires-python = ">=3.11" dependencies = [ + "cryptography>=42.0", "langchain-core>=1.6.5", "langchain-anthropic>=1.7.4", "langchain-google-genai>=4.4.0", diff --git a/tests/test_chatgpt_auth.py b/tests/test_chatgpt_auth.py new file mode 100644 index 000000000..d2a11d916 --- /dev/null +++ b/tests/test_chatgpt_auth.py @@ -0,0 +1,424 @@ +"""Sign in with ChatGPT: the sign-in is checked, kept private, renewed once, and revoked. + +OpenAI is replaced by a fake that signs ID tokens with a key made here, so every +check of the sign-in (signature, issuer, audience, nonce, plan scope) runs for real. +""" + +import base64 +import json +import socket +import stat +import threading +import time +import urllib.error +import urllib.parse +import urllib.request + +import pytest +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import padding, rsa + +from tradingagents.llm_clients import chatgpt_auth as auth + +# The suite refuses all connections; the loopback test lets 127.0.0.1 through only. +_CONNECT, _GETADDRINFO = socket.socket.connect, socket.getaddrinfo +KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) +CLIENT = "oaiapp_test" +SUBJECT = "user-123" + + +def _b64(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode() + + +def _id_token(key=KEY, **claims) -> str: + payload = {"iss": auth.ISSUER, "aud": CLIENT, "sub": SUBJECT, "email": "a@example.com", + "exp": time.time() + 3600, **claims} + head = _b64(json.dumps({"alg": "RS256", "kid": "k1"}).encode()) + body = _b64(json.dumps(payload).encode()) + signature = key.sign(f"{head}.{body}".encode(), padding.PKCS1v15(), hashes.SHA256()) + return f"{head}.{body}.{_b64(signature)}" + + +def _jwks(): + numbers = KEY.public_key().public_numbers() + return {"keys": [{"kid": "k1", "kty": "RSA", "e": _b64(numbers.e.to_bytes(3, "big")), + "n": _b64(numbers.n.to_bytes(256, "big"))}]} + + +class _Response: + def __init__(self, body, status=200): + self.body, self.status_code = body, status + + def json(self): + return self.body + + +class FakeOpenAI: + """Answers the JWKS, token and revoke endpoints and records what was sent.""" + + def __init__(self, token_body=None, status=200): + self.posts = [] + self.token_body = token_body + self.status = status + self.lock = threading.Lock() + self.keys_reachable = True + self.revoke_status = 200 + self.post_error = None + + def get(self, url, timeout): + assert url == auth.JWKS_URL + if not self.keys_reachable: + raise auth.requests.ConnectionError("down") + return _Response(_jwks()) + + def post(self, url, timeout, data): + with self.lock: + self.posts.append((url, data)) + if url == auth.REVOKE_URL: + return _Response({}, self.revoke_status) + if self.post_error is not None: + raise self.post_error + time.sleep(0.05) # let a racing refresher reach the lock + body = self.token_body(data) if callable(self.token_body) else self.token_body + return _Response(body, self.status) + + +@pytest.fixture +def home(tmp_path, monkeypatch): + monkeypatch.setattr(auth, "AUTH_FILE", tmp_path / "home" / "chatgpt-auth.json") + monkeypatch.setattr(auth, "_retry_after", 0.0) + return tmp_path + + +@pytest.fixture +def openai(monkeypatch): + fake = FakeOpenAI() + monkeypatch.setattr(auth.requests, "get", fake.get) + monkeypatch.setattr(auth.requests, "post", fake.post) + return fake + + +def _signed_in(expires_at, earliest=None, refresh_token="r1"): + auth.AUTH_FILE.parent.mkdir(parents=True, exist_ok=True) + auth._write({"host_id": "urn:uuid:h", "account": {"client_id": CLIENT, "email": "a@example.com", + "subject": SUBJECT}, + "tokens": {"access_token": "a1", "refresh_token": refresh_token, + "expires_at": expires_at, "earliest_refresh_at": earliest}}) + + +def _renewal(data): + return {"access_token": "a2", "refresh_token": "r2", "expires_in": 3600, + "id_token": _id_token(), "scope": auth.SCOPE} + + +# --- the ID token ---------------------------------------------------------------- + +@pytest.mark.unit +def test_a_genuine_id_token_is_accepted(): + assert auth._validated_claims(_id_token(nonce="n"), CLIENT, "n", _jwks()["keys"])["sub"] == SUBJECT + + +@pytest.mark.unit +@pytest.mark.parametrize("token, nonce", [ + (lambda: _id_token(nonce="other"), "n"), + (lambda: _id_token(aud="oaiapp_someone_else", nonce="n"), "n"), + (lambda: _id_token(iss="https://evil.example", nonce="n"), "n"), + (lambda: _id_token(exp=time.time() - 600, nonce="n"), "n"), +]) +def test_an_id_token_for_another_sign_in_is_refused(token, nonce): + with pytest.raises(auth.ChatGPTAuthError): + auth._validated_claims(token(), CLIENT, nonce, _jwks()["keys"]) + + +@pytest.mark.unit +def test_an_id_token_signed_by_another_key_is_refused(): + stranger = rsa.generate_private_key(public_exponent=65537, key_size=2048) + with pytest.raises(auth.ChatGPTAuthError, match="do not verify"): + auth._validated_claims(_id_token(key=stranger, nonce="n"), CLIENT, "n", _jwks()["keys"]) + + +# --- tokens ---------------------------------------------------------------------- + +@pytest.mark.unit +def test_nobody_signed_in_is_told_how_to_sign_in(home, openai): + with pytest.raises(auth.ChatGPTAuthError, match="tradingagents login chatgpt"): + auth.credential() + with pytest.raises(auth.ChatGPTAuthError, match="tradingagents login chatgpt"): + auth.access_token() + assert openai.posts == [] + + +@pytest.mark.unit +def test_a_fresh_token_is_used_without_asking_openai(home, openai): + _signed_in(expires_at=time.time() + 3000) + assert auth.credential()() == "a1" + assert openai.posts == [] + + +@pytest.mark.unit +def test_a_token_near_expiry_is_renewed_with_the_stored_refresh_token(home, openai): + _signed_in(expires_at=time.time() + 60) + openai.token_body = _renewal + assert auth.access_token() == "a2" + (url, sent), = openai.posts + assert url == auth.TOKEN_URL and sent["refresh_token"] == "r1" and sent["client_id"] == CLIENT + stored = json.loads(auth.AUTH_FILE.read_text()) + assert stored["tokens"]["refresh_token"] == "r2" and stored["account"]["subject"] == SUBJECT + assert stat.S_IMODE(auth.AUTH_FILE.stat().st_mode) == 0o600 + + +@pytest.mark.unit +def test_renewal_waits_for_the_time_openai_allows_while_the_token_still_works(home, openai): + now = time.time() + _signed_in(expires_at=now + 200, earliest=now + 100) + assert auth.access_token() == "a1" + assert openai.posts == [] + + +@pytest.mark.unit +def test_concurrent_runs_renew_once_and_never_replay_a_used_refresh_token(home, openai): + _signed_in(expires_at=time.time() + 60) + openai.token_body = _renewal + tokens = [] + threads = [threading.Thread(target=lambda: tokens.append(auth.access_token())) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert tokens == ["a2"] * 4 + assert [sent["refresh_token"] for _, sent in openai.posts] == ["r1"] + + +@pytest.mark.unit +def test_a_refused_renewal_asks_to_sign_in_again_and_keeps_the_file(home, openai): + _signed_in(expires_at=time.time() + 60) + openai.token_body, openai.status = {"error": "invalid_grant"}, 400 + with pytest.raises(auth.ChatGPTAuthError, match="invalid_grant.*login chatgpt"): + auth.access_token() + assert json.loads(auth.AUTH_FILE.read_text())["tokens"]["refresh_token"] == "r1" + + +@pytest.mark.unit +def test_a_renewal_for_another_account_is_refused(home, openai): + _signed_in(expires_at=time.time() + 60) + openai.token_body = lambda data: {**_renewal(data), "id_token": _id_token(sub="someone-else")} + with pytest.raises(auth.ChatGPTAuthError, match="another account"): + auth.access_token() + + +# --- signing in and out ------------------------------------------------------------ + +def _nonce_of(printed): + url = next(line for line in printed if line.startswith(auth.AUTHORIZE_URL)) + return dict(urllib.parse.parse_qsl(urllib.parse.urlparse(url).query))["nonce"] + + +def _paster(printed, edit=lambda query: query): + calls = [] + + def read_line(): + calls.append(1) + if len(calls) > 1: + threading.Event().wait() + url = next(line for line in printed if line.startswith(auth.AUTHORIZE_URL)) + asked = dict(urllib.parse.parse_qsl(urllib.parse.urlparse(url).query)) + query = edit({"code": "c1", "state": asked["state"], "client_id": CLIENT}) + return f"{asked['redirect_uri']}?{urllib.parse.urlencode(query)}" + + return read_line + + +@pytest.mark.unit +def test_signing_in_by_pasting_the_redirect_saves_a_private_sign_in(home, openai): + printed = [] + + def token(data): + assert data["grant_type"] == "authorization_code" and data["client_id"] == CLIENT + assert data["code"] == "c1" and data["code_verifier"] + return {"access_token": "a1", "refresh_token": "r1", "expires_in": 3600, + "id_token": _id_token(nonce=_nonce_of(printed)), "scope": auth.SCOPE} + + openai.token_body = token + email = auth.sign_in(say=printed.append, read_line=_paster(printed), timeout=10, port=0) + assert email == "a@example.com" + stored = json.loads(auth.AUTH_FILE.read_text()) + assert stored["account"] == {"client_id": CLIENT, "email": "a@example.com", "subject": SUBJECT} + assert stored["host_id"].startswith("urn:uuid:") + assert stat.S_IMODE(auth.AUTH_FILE.stat().st_mode) == 0o600 + asked = dict(urllib.parse.parse_qsl(urllib.parse.urlparse(printed[1]).query)) + assert asked["client_id"] == auth.DYNAMIC_CLIENT_ID and asked["agent_name_hint"] == "TradingAgents" + assert asked["code_challenge_method"] == "S256" and asked["redirect_uri"].startswith("http://127.0.0.1:") + + +@pytest.mark.unit +def test_an_answer_for_another_sign_in_attempt_is_ignored(home, openai): + printed = [] + with pytest.raises(auth.ChatGPTAuthError, match="timed out"): + auth.sign_in(say=printed.append, read_line=_paster(printed, lambda q: {**q, "state": "forged"}), + timeout=1, port=0) + assert any("not the address" in line for line in printed) + assert openai.posts == [] and not auth.AUTH_FILE.exists() + + +def _status(url): + try: + return urllib.request.urlopen(url, timeout=5).status + except urllib.error.HTTPError as exc: + return exc.code + + +@pytest.mark.unit +def test_the_browser_answer_reaches_the_loopback_and_strays_are_ignored(home, openai, monkeypatch): + refused_connect, refused_lookup = socket.socket.connect, socket.getaddrinfo + monkeypatch.setattr(socket.socket, "connect", lambda sock, address: ( + _CONNECT if address[0] == "127.0.0.1" else refused_connect)(sock, address)) + monkeypatch.setattr(socket, "getaddrinfo", lambda host, *a, **k: ( + _GETADDRINFO if host == "127.0.0.1" else refused_lookup)(host, *a, **k)) + printed = [] + openai.token_body = lambda data: {"access_token": "a1", "refresh_token": "r1", "expires_in": 3600, + "id_token": _id_token(nonce=_nonce_of(printed)), "scope": auth.SCOPE} + + def browser(): + while not any(line.startswith(auth.AUTHORIZE_URL) for line in printed): + time.sleep(0.01) + asked = dict(urllib.parse.parse_qsl(urllib.parse.urlparse(printed[1]).query)) + results.append(_status(f"{asked['redirect_uri']}?error=access_denied")) + results.append(_status(f"{asked['redirect_uri']}?code=c1&state={asked['state']}&client_id={CLIENT}")) + + results = [] + threading.Thread(target=browser, daemon=True).start() + email = auth.sign_in(say=printed.append, read_line=threading.Event().wait, timeout=10, port=0) + assert email == "a@example.com" and results == [400, 200] + + +@pytest.mark.unit +def test_a_sign_in_without_plan_usage_is_refused(home, openai): + printed = [] + openai.token_body = lambda data: {"access_token": "a1", "refresh_token": "r1", "expires_in": 3600, + "id_token": _id_token(nonce=_nonce_of(printed)), + "scope": "openid email offline_access"} + with pytest.raises(auth.ChatGPTAuthError, match="approve plan usage"): + auth.sign_in(say=printed.append, read_line=_paster(printed), timeout=10, port=0) + assert not auth.AUTH_FILE.exists() + + +@pytest.mark.unit +def test_signing_out_revokes_the_sign_in_and_keeps_only_the_host_id(home, openai): + _signed_in(expires_at=time.time() + 3000) + assert auth.sign_out() == "revoked" + (url, sent), = openai.posts + assert url == auth.REVOKE_URL and sent["token"] == "r1" and sent["client_id"] == CLIENT + assert json.loads(auth.AUTH_FILE.read_text()) == {"host_id": "urn:uuid:h"} + assert auth.signed_in_as() is None + assert auth.sign_out() is None + + +@pytest.mark.unit +def test_unreachable_keys_stop_a_renewal_before_the_refresh_token_is_spent(home, openai): + _signed_in(expires_at=time.time() - 1) + openai.token_body, openai.keys_reachable = _renewal, False + with pytest.raises(auth.ChatGPTUnreachable, match="Could not reach OpenAI"): + auth.access_token() + assert openai.posts == [] + assert json.loads(auth.AUTH_FILE.read_text())["tokens"]["refresh_token"] == "r1" + + +@pytest.mark.unit +@pytest.mark.parametrize("failure", ["keys", "network", 429, 503]) +def test_a_temporary_failure_keeps_a_token_that_still_works(home, openai, failure): + _signed_in(expires_at=time.time() + 60) + openai.token_body = _renewal + if failure == "keys": + openai.keys_reachable = False + elif failure == "network": + openai.post_error = auth.requests.ConnectionError("reset") + else: + openai.token_body, openai.status = {"error": "busy"}, failure + assert auth.access_token() == "a1" + assert json.loads(auth.AUTH_FILE.read_text())["tokens"]["refresh_token"] == "r1" + if failure == "keys": + assert openai.posts == [] + + +@pytest.mark.unit +@pytest.mark.parametrize("failure", ["network", 408, 503]) +def test_a_temporary_failure_after_expiry_says_to_try_again_not_to_sign_in(home, openai, failure): + _signed_in(expires_at=time.time() - 1) + if failure == "network": + openai.post_error = auth.requests.Timeout("slow") + else: + openai.token_body, openai.status = {"error": "busy"}, failure + with pytest.raises(auth.ChatGPTUnreachable) as raised: + auth.access_token() + assert "try again" in str(raised.value) and "login" not in str(raised.value) + + +@pytest.mark.unit +def test_signing_in_again_revokes_the_replaced_sign_in(home, openai): + _signed_in(expires_at=time.time() + 3000, refresh_token="old") + printed = [] + openai.token_body = lambda data: {"access_token": "a2", "refresh_token": "r2", "expires_in": 3600, + "id_token": _id_token(nonce=_nonce_of(printed)), "scope": auth.SCOPE} + auth.sign_in(say=printed.append, read_line=_paster(printed), timeout=10, port=0) + assert openai.posts[-1][0] == auth.REVOKE_URL and openai.posts[-1][1]["token"] == "old" + assert json.loads(auth.AUTH_FILE.read_text())["tokens"]["refresh_token"] == "r2" + + +@pytest.mark.unit +def test_another_account_cannot_replace_the_signed_in_one(home, openai): + _signed_in(expires_at=time.time() + 3000) + printed = [] + openai.token_body = lambda data: {"access_token": "a2", "refresh_token": "r2", "expires_in": 3600, + "id_token": _id_token(nonce=_nonce_of(printed), sub="someone-else"), + "scope": auth.SCOPE} + with pytest.raises(auth.ChatGPTAuthError, match="logout chatgpt"): + auth.sign_in(say=printed.append, read_line=_paster(printed), timeout=10, port=0) + assert json.loads(auth.AUTH_FILE.read_text())["tokens"]["refresh_token"] == "r1" + + +@pytest.mark.unit +def test_signing_out_while_openai_is_unreachable_still_forgets_the_sign_in(home, openai): + _signed_in(expires_at=time.time() + 3000) + openai.revoke_status = 503 + assert auth.sign_out() == "forgotten" + assert auth.signed_in_as() is None + + +@pytest.mark.unit +def test_after_a_temporary_failure_renewal_waits_a_minute_before_trying_again(home, openai): + _signed_in(expires_at=time.time() + 200) + openai.post_error = auth.requests.ConnectionError("reset") + clock = [time.time()] + assert auth.access_token(now=lambda: clock[0]) == "a1" + assert auth.access_token(now=lambda: clock[0] + 30) == "a1" + assert len(openai.posts) == 1 + openai.post_error, openai.token_body = None, _renewal + assert auth.access_token(now=lambda: clock[0] + 61) == "a2" + + +@pytest.mark.unit +def test_after_expiry_a_failed_renewal_is_not_retried_within_the_minute(home, openai): + _signed_in(expires_at=time.time() - 1) + openai.post_error = auth.requests.ConnectionError("reset") + for _ in range(3): + with pytest.raises(auth.ChatGPTUnreachable, match="try again"): + auth.access_token() + assert len(openai.posts) == 1 + + +@pytest.mark.unit +def test_an_unreadable_renewal_answer_keeps_a_token_that_still_works(home, openai): + _signed_in(expires_at=time.time() + 60) + openai.token_body = lambda data: ["not", "a", "token"] + assert auth.access_token() == "a1" + assert json.loads(auth.AUTH_FILE.read_text())["tokens"]["refresh_token"] == "r1" + + +@pytest.mark.unit +def test_a_network_failure_while_signing_in_is_reported_not_raised_raw(home, openai): + printed = [] + openai.post_error = auth.requests.ConnectionError("reset") + with pytest.raises(auth.ChatGPTUnreachable, match="sign you in to ChatGPT.*try again"): + auth.sign_in(say=printed.append, read_line=_paster(printed), timeout=10, port=0) + assert not auth.AUTH_FILE.exists() diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py index 37c9c2859..0999a5349 100644 --- a/tests/test_cli_commands.py +++ b/tests/test_cli_commands.py @@ -180,3 +180,53 @@ def test_backtest_refuses_a_selection_before_the_sweep(runner, monkeypatch, args result = runner.invoke(m.app, ["backtest", "NVDA", "--start", "2026-06-01", "--end", "2026-06-01", *args]) assert result.exit_code == 1 + + +@pytest.mark.unit +def test_login_and_logout_chatgpt_do_not_run_an_analysis(runner, monkeypatch): + from tradingagents.llm_clients import chatgpt_auth + + monkeypatch.setattr(chatgpt_auth, "sign_in", lambda say: "a@example.com") + monkeypatch.setattr(chatgpt_auth, "sign_out", lambda: "revoked") + signed_in = runner.invoke(m.app, ["login", "chatgpt"]) + assert signed_in.exit_code == 0 and "Signed in to ChatGPT as a@example.com" in signed_in.output + signed_out = runner.invoke(m.app, ["logout", "chatgpt"]) + assert signed_out.exit_code == 0 and "Signed out of ChatGPT" in signed_out.output + assert runner.invoke(m.app, ["login", "claude"]).exit_code != 0 + assert calls == [] + + +@pytest.mark.unit +def test_a_failed_login_exits_with_the_reason(runner, monkeypatch): + from tradingagents.llm_clients import chatgpt_auth + + def refuse(say): + raise chatgpt_auth.ChatGPTAuthError("Signing in with ChatGPT timed out") + + monkeypatch.setattr(chatgpt_auth, "sign_in", refuse) + result = runner.invoke(m.app, ["login", "chatgpt"]) + assert result.exit_code == 1 and "timed out" in result.output + + +@pytest.mark.unit +@pytest.mark.parametrize("outcome, said", [ + (None, "Not signed in"), ("forgotten", "could not be reached to revoke"), +]) +def test_logout_says_whether_openai_revoked_the_sign_in(runner, monkeypatch, outcome, said): + from tradingagents.llm_clients import chatgpt_auth + + monkeypatch.setattr(chatgpt_auth, "sign_out", lambda: outcome) + result = runner.invoke(m.app, ["logout", "chatgpt"]) + assert result.exit_code == 0 and said in " ".join(result.output.split()) + + +@pytest.mark.unit +def test_a_failure_message_with_brackets_is_printed_as_text(runner, monkeypatch): + from tradingagents.llm_clients import chatgpt_auth + + def refuse(say): + raise chatgpt_auth.ChatGPTAuthError("ChatGPT sign-in failed ([/bogus]); sign in again.") + + monkeypatch.setattr(chatgpt_auth, "sign_in", refuse) + result = runner.invoke(m.app, ["login", "chatgpt"]) + assert result.exit_code == 1 and "[/bogus]" in result.output diff --git a/tradingagents/llm_clients/chatgpt_auth.py b/tradingagents/llm_clients/chatgpt_auth.py new file mode 100644 index 000000000..066212423 --- /dev/null +++ b/tradingagents/llm_clients/chatgpt_auth.py @@ -0,0 +1,404 @@ +"""Sign in with ChatGPT, so a run uses the person's ChatGPT Plus or Pro plan. + +This follows OpenAI's documented flow for open-source apps ("Sign in with ChatGPT": +developers.openai.com/siwc): dynamic client registration, PKCE, and a loopback +redirect on 127.0.0.1. The sign-in, its tokens and this host's id live in one file +under ~/.tradingagents, readable only by its owner. Access tokens last an hour and +are refreshed under a lock shared by every thread and process, because each refresh +rotates the refresh token and replaying a used one ends the sign-in. +""" + +import base64 +import hashlib +import http.server +import json +import os +import secrets +import threading +import time +import urllib.parse +import uuid +from collections.abc import Callable +from pathlib import Path + +import requests + +from tradingagents.dataflows.files import locked +from tradingagents.default_config import _TRADINGAGENTS_HOME + +ISSUER = "https://auth.openai.com" +AUTHORIZE_URL = f"{ISSUER}/api/accounts/authorize" +TOKEN_URL = f"{ISSUER}/api/accounts/oauth/token" +REVOKE_URL = f"{ISSUER}/api/accounts/oauth/revoke" +JWKS_URL = f"{ISSUER}/.well-known/jwks.json" +API_BASE = "https://api.openai.com/v1" +DYNAMIC_CLIENT_ID = "dynamic_agent_client" +APP_NAME = "TradingAgents" +SCOPE = "openid profile email offline_access resource.invoke chatgpt.tokens.use.direct" +PLAN_SCOPE = "chatgpt.tokens.use.direct" +CALLBACK_PATH = "/auth/callback" +DEFAULT_PORT = 1455 +REFRESH_MARGIN = 300 # refresh when the access token has under 5 minutes left +RETRY_AFTER = 60 # after OpenAI could not renew, wait this long before trying again +TIMEOUT = 30 + +AUTH_FILE = Path(_TRADINGAGENTS_HOME) / "chatgpt-auth.json" +LOGIN_HINT = "Not signed in to ChatGPT; run `tradingagents login chatgpt`." + + +_retry_after = 0.0 # no renewal attempt in this process before this time + + +class ChatGPTAuthError(RuntimeError): + """Signing in or refreshing failed; the message says what to do.""" + + +class ChatGPTUnreachable(ChatGPTAuthError): + """OpenAI could not be reached or was busy; the sign-in itself is still good.""" + + +# --- the auth file ------------------------------------------------------------------- + +def _read() -> dict: + try: + data = json.loads(AUTH_FILE.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return data if isinstance(data, dict) else {} + + +def _write(data: dict) -> None: + """Write the auth file whole, readable only by its owner. Unlike a cache, a failed + write must surface: losing a rotated refresh token would end the sign-in.""" + AUTH_FILE.parent.mkdir(parents=True, exist_ok=True) + temp = AUTH_FILE.with_name(f"{AUTH_FILE.name}.{uuid.uuid4().hex}.tmp") + fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(data, handle) + os.replace(temp, AUTH_FILE) + except BaseException: + temp.unlink(missing_ok=True) + raise + + +def _host_id(data: dict) -> str: + return data.get("host_id") or f"urn:uuid:{uuid.uuid4()}" + + +def signed_in_as() -> str | None: + """The signed-in account's email, or None.""" + data = _read() + return (data.get("account") or {}).get("email") if data.get("tokens") else None + + +# --- tokens -------------------------------------------------------------------------- + +def _b64url(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode() + + +def _b64url_decode(text: str) -> bytes: + return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4)) + + +def _signing_keys() -> list[dict]: + """OpenAI's ID-token signing keys. Fetched before a token request, so nothing can fail + on the network between OpenAI issuing tokens and this file keeping them.""" + try: + keys = requests.get(JWKS_URL, timeout=TIMEOUT).json()["keys"] + except (requests.RequestException, ValueError, KeyError, TypeError) as exc: + raise ChatGPTUnreachable(f"Could not reach OpenAI to check the sign-in ({type(exc).__name__}); " + "try again.") from None + return keys if isinstance(keys, list) else [] + + +def _validated_claims(id_token: str, client_id: str, nonce: str | None, keys: list[dict]) -> dict: + """The ID token's claims after checking its RS256 signature against OpenAI's keys, + issuer, audience, expiry and (at sign-in) nonce, as the sign-in guide requires.""" + from cryptography.exceptions import InvalidSignature + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import padding, rsa + + try: + header_b64, payload_b64, signature_b64 = id_token.split(".") + header = json.loads(_b64url_decode(header_b64)) + key = next((k for k in keys if k.get("kid") == header.get("kid") and k.get("kty") == "RSA"), None) + if header.get("alg") != "RS256" or key is None: + raise InvalidSignature + public = rsa.RSAPublicNumbers(int.from_bytes(_b64url_decode(key["e"]), "big"), + int.from_bytes(_b64url_decode(key["n"]), "big")).public_key() + public.verify(_b64url_decode(signature_b64), f"{header_b64}.{payload_b64}".encode(), + padding.PKCS1v15(), hashes.SHA256()) + claims = json.loads(_b64url_decode(payload_b64)) + except (InvalidSignature, ValueError, KeyError, TypeError, AttributeError): + raise ChatGPTAuthError("ChatGPT sign-in returned an ID token OpenAI's keys do not verify.") from None + audience = claims.get("aud") + audiences = audience if isinstance(audience, list) else [audience] + if (claims.get("iss") != ISSUER or client_id not in audiences or not claims.get("sub") + or claims.get("exp", 0) + 5 < time.time() + or (nonce is not None and claims.get("nonce") != nonce)): + raise ChatGPTAuthError("ChatGPT sign-in returned an ID token that does not match this sign-in.") + return claims + + +def _token_record(response: dict, now: float, previous: dict | None = None) -> dict: + previous = previous or {} + return { + "access_token": response["access_token"], + "refresh_token": response.get("refresh_token") or previous.get("refresh_token"), + "expires_at": now + int(response.get("expires_in") or 3600), + "earliest_refresh_at": response.get("earliest_refresh_at"), + "scopes": sorted((response.get("scope") or "").split()) or previous.get("scopes", []), + } + + +def _oauth_error(response) -> str: + try: + body = response.json() + except ValueError: + return f"HTTP {response.status_code}" + error = body.get("error") + if isinstance(error, dict): + return error.get("code") or error.get("message") or f"HTTP {response.status_code}" + return error or f"HTTP {response.status_code}" + + +def _token_request(form: dict, task: str, refused: str) -> dict: + """OpenAI's answer to a token request, or the error to show: ``ChatGPTUnreachable`` + when OpenAI could not answer (try again), ``ChatGPTAuthError`` when it refused.""" + try: + response = requests.post(TOKEN_URL, timeout=TIMEOUT, data=form) + except requests.RequestException as exc: + raise ChatGPTUnreachable(f"Could not reach OpenAI to {task} ({type(exc).__name__}); try again.") from None + if response.status_code in (408, 429) or response.status_code >= 500: + raise ChatGPTUnreachable(f"OpenAI could not {task} now ({_oauth_error(response)}); try again.") + if response.status_code != 200: + raise ChatGPTAuthError(f"OpenAI did not {task} ({_oauth_error(response)}); {refused}.") + try: + body = response.json() + except ValueError: + body = None + if not isinstance(body, dict) or not isinstance(body.get("access_token"), str): + raise ChatGPTUnreachable(f"OpenAI's answer could not be read when asked to {task}; try again.") + return body + + +def _refresh(data: dict, now: float) -> dict: + tokens, account = data["tokens"], data["account"] + keys = _signing_keys() + body = _token_request({"grant_type": "refresh_token", "client_id": account["client_id"], + "refresh_token": tokens["refresh_token"], "resource": API_BASE}, + "renew your ChatGPT sign-in", "run `tradingagents login chatgpt` again") + if body.get("id_token"): + claims = _validated_claims(body["id_token"], account["client_id"], None, keys) + if claims["sub"] != account.get("subject"): + raise ChatGPTAuthError("The renewed ChatGPT sign-in belongs to another account.") + return {**data, "tokens": _token_record(body, now, tokens)} + + +def _due(tokens: dict, now: float) -> bool: + """Whether to renew: under five minutes left, and OpenAI accepts a renewal by now + (``earliest_refresh_at``), unless the token has already run out.""" + expires = tokens.get("expires_at", 0) + earliest = tokens.get("earliest_refresh_at") + if now < expires - REFRESH_MARGIN: + return False + return now >= expires or not isinstance(earliest, (int, float)) or now >= earliest + + +def access_token(now: Callable[[], float] = time.time) -> str: + """A valid access token, renewed when it has under five minutes left. + + Renewal happens under a lock every thread and process shares, and starts from + the token on disk: another run may have renewed it already, and the refresh + token it rotated away must never be sent again. When OpenAI cannot be reached + or is busy, a token that still works is used and renewal is tried again a minute later.""" + global _retry_after + tokens = _read().get("tokens") + if not tokens: + raise ChatGPTAuthError(LOGIN_HINT) + if not _due(tokens, now()): + return tokens["access_token"] + with locked(AUTH_FILE): + data = _read() + if not data.get("tokens"): + raise ChatGPTAuthError(LOGIN_HINT) + tokens = data["tokens"] + if _due(tokens, now()): + still_works = tokens.get("expires_at", 0) > now() + if now() < _retry_after: + if still_works: + return tokens["access_token"] + raise ChatGPTUnreachable("OpenAI could not renew your ChatGPT sign-in a moment ago; " + "try again in a minute.") + try: + data = _refresh(data, now()) + except ChatGPTUnreachable: + _retry_after = now() + RETRY_AFTER + if not still_works: + raise + return tokens["access_token"] + _write(data) + return data["tokens"]["access_token"] + + +def credential() -> Callable[[], str]: + """The api_key callable for the chatgpt provider; raises now when nobody is signed in, + so a run stops before it starts rather than at its first model call.""" + if not _read().get("tokens"): + raise ChatGPTAuthError(LOGIN_HINT) + return access_token + + +# --- signing in and out -------------------------------------------------------------- + +class _Callback: + """The first answer carrying this attempt's ``state``; stray or stale requests are ignored.""" + + def __init__(self, state: str): + self.state = state + self.params: dict | None = None + self.received = threading.Event() + self._lock = threading.Lock() + + def accept(self, url: str) -> bool: + parsed = urllib.parse.urlparse(url.strip()) + params = {k: v[0] for k, v in urllib.parse.parse_qs(parsed.query).items()} + if parsed.path != CALLBACK_PATH or params.get("state") != self.state: + return False + with self._lock: + if self.params is None: + self.params = params + self.received.set() + return True + + +def _listen(callback: _Callback, port: int) -> http.server.HTTPServer: + class Handler(http.server.BaseHTTPRequestHandler): + timeout = 5 # a connection that sends nothing cannot hold up the real answer + + def do_GET(self): # noqa: N802 (http.server's name) + ok = callback.accept(self.path) + self.send_response(200 if ok else 400) + self.send_header("Content-Type", "text/plain; charset=utf-8") + self.end_headers() + if ok: + self.wfile.write(b"Return to the terminal to finish signing in to TradingAgents.\n") + + def log_message(self, *args): # the URL carries the sign-in code + pass + + try: + server = http.server.HTTPServer(("127.0.0.1", port), Handler) + except OSError: + server = http.server.HTTPServer(("127.0.0.1", 0), Handler) + # A short poll interval lets the server stop promptly once the answer has arrived. + threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}, daemon=True).start() + return server + + +def _revoke(tokens: dict, client_id: str | None) -> bool: + try: + response = requests.post(REVOKE_URL, timeout=TIMEOUT, data={ + "token": tokens.get("refresh_token"), "token_type_hint": "refresh_token", + "client_id": client_id, + }) + except requests.RequestException: + return False + return response.status_code == 200 + + +def sign_in(say: Callable[[str], None] = print, read_line: Callable[[], str] = input, + timeout: float = 600, port: int = DEFAULT_PORT) -> str: + """Sign in with ChatGPT and save the result; returns the account's email. + + The address to open is printed. The browser comes back to 127.0.0.1; when this + machine is remote or in a container, paste the address the browser ends up on. + The pasted line is read in a thread that can outlive this call, so sign in as the + last prompt of a process (``tradingagents login chatgpt``).""" + data = _read() + host_id = _host_id(data) + account = data.get("account") or {} + client_id = account.get("client_id") or DYNAMIC_CLIENT_ID + state, nonce, verifier = (_b64url(secrets.token_bytes(32)), _b64url(secrets.token_bytes(32)), + _b64url(secrets.token_bytes(64))) + callback = _Callback(state) + server = _listen(callback, port) + try: + redirect_uri = f"http://127.0.0.1:{server.server_address[1]}{CALLBACK_PATH}" + params = { + "client_id": client_id, "ext_agent_host_id": host_id, "response_type": "code", + "redirect_uri": redirect_uri, "scope": SCOPE, "resource": API_BASE, "state": state, + "nonce": nonce, "code_challenge_method": "S256", + "code_challenge": _b64url(hashlib.sha256(verifier.encode()).digest()), + } + if client_id == DYNAMIC_CLIENT_ID: + params["agent_name_hint"] = APP_NAME + elif account.get("email"): + params["login_hint"] = account["email"] + say("Continue with ChatGPT: open this address and sign in with your Plus or Pro plan:") + say(f"{AUTHORIZE_URL}?{urllib.parse.urlencode(params, quote_via=urllib.parse.quote)}") + say("If the browser cannot reach this machine, paste here the address it ends up on.") + + def read_pasted(): + try: + while not callback.received.is_set(): + if not callback.accept(read_line()) and not callback.received.is_set(): + say("That is not the address this sign-in ends on; paste the whole address.") + except (EOFError, OSError): + return + + threading.Thread(target=read_pasted, daemon=True).start() + if not callback.received.wait(timeout): + raise ChatGPTAuthError("Signing in with ChatGPT timed out; run `tradingagents login chatgpt` again.") + finally: + server.shutdown() + server.server_close() + + answer = callback.params + if answer.get("error"): + raise ChatGPTAuthError(f"ChatGPT sign-in was not completed ({answer['error']}).") + issued = answer.get("client_id") or client_id + if issued == DYNAMIC_CLIENT_ID or (account.get("client_id") and issued != account["client_id"]): + raise ChatGPTAuthError("ChatGPT sign-in did not register TradingAgents; run " + "`tradingagents logout chatgpt`, then sign in again.") + keys = _signing_keys() + body = _token_request({"grant_type": "authorization_code", "client_id": issued, + "code": answer.get("code", ""), "code_verifier": verifier, + "redirect_uri": redirect_uri, "resource": API_BASE}, + "sign you in to ChatGPT", "sign in again") + claims = _validated_claims(body.get("id_token") or "", issued, nonce, keys) + if account.get("subject") and claims["sub"] != account["subject"]: + raise ChatGPTAuthError(f"This machine is signed in as {account.get('email') or 'another account'}; " + "run `tradingagents logout chatgpt` first.") + if PLAN_SCOPE not in (body.get("scope") or "").split(): + raise ChatGPTAuthError("ChatGPT did not allow TradingAgents to use your plan; " + "sign in again and approve plan usage.") + AUTH_FILE.parent.mkdir(parents=True, exist_ok=True) + with locked(AUTH_FILE): + replaced = _read().get("tokens") + _write({ + "host_id": host_id, + "account": {"client_id": issued, "email": claims.get("email"), "subject": claims["sub"]}, + "tokens": _token_record(body, time.time()), + }) + if replaced: + _revoke(replaced, issued) + return claims.get("email") or "your ChatGPT account" + + +def sign_out() -> str | None: + """Revoke the sign-in at OpenAI and forget it, keeping only this host's id. + Returns None when nobody was signed in, "revoked", or "forgotten" when OpenAI + could not be reached to revoke it.""" + if not _read().get("tokens"): + return None + with locked(AUTH_FILE): + data = _read() + tokens = data.get("tokens") + if not tokens: + return None + revoked = _revoke(tokens, (data.get("account") or {}).get("client_id")) + _write({"host_id": _host_id(data)}) + return "revoked" if revoked else "forgotten"