Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/agent/test_restore_primary_pool_reselect.py

@@ -0,0 +1,167 @@
+# Copyright 2025 Nous Research (Licensed under the Apache License, Version 2.0)
+"""Test that _restore_primary_runtime re-selects from the credential pool
+instead of using a stale snapshot key.
+
+Bug: when a credential pool entry is revoked/marked-exhausted during a turn,
+_restore_primary_runtime restores the original (now-stale) api_key from the
+construction-time snapshot. The next turn immediately hits the same error,
+exhausting remaining entries and falling through to cross-provider fallback.
+"""
+
+import json
+import logging
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from agent.credential_pool import (
+    AUTH_TYPE_OAUTH,
+    PooledCredential,
+)
+
+
+def _make_entry(
+    label: str,
+    access_token: str,
+    *,
+    source: str = "device_code",
+    priority: int = 0,
+    last_status: str | None = None,
+    last_status_at: float | None = None,
+) -> dict:
+    return {
+        "id": label,
+        "label": label,
+        "provider": "openai-codex",
+        "auth_type": AUTH_TYPE_OAUTH,
+        "source": source,
+        "priority": priority,
+        "access_token": access_token,
+        "refresh_token": f"rt-{label}",
+        "base_url": "https://chatgpt.com/backend-api/codex",
+        "last_status": last_status,
+        "last_status_at": last_status_at,
+    }
+
+
+def _build_mock_pool(entries: list[dict], *, strategy: str = "round_robin"):
+    """Build a mock CredentialPool with the given entries."""
+    from agent.credential_pool import CredentialPool
+
+    pool = CredentialPool(
+        provider="openai-codex",
+        entries=[PooledCredential.from_dict("openai-codex", e) for e in entries],
+    )
+    pool._strategy = strategy
+    return pool
+
+
+class TestRestorePrimaryPoolReselect:
+    """_restore_primary_runtime should re-select from the credential pool."""
+
+    def _make_agent(self, pool):
+        """Create a minimal AIAgent with the given credential pool."""
+        from run_agent import AIAgent
+
+        agent = AIAgent.__new__(AIAgent)
+        agent.model = "gpt-5.5"
+        agent.provider = "openai-codex"
+        agent.base_url = "https://chatgpt.com/backend-api/codex"
+        agent.api_mode = "codex_responses"
+        agent.api_key = "original-key-entry-1"
+        agent._client_kwargs = {
+            "api_key": "original-key-entry-1",
+            "base_url": "https://chatgpt.com/backend-api/codex",
+        }
+        agent._credential_pool = pool
+        agent._fallback_activated = True
+        agent._fallback_index = 1
+        agent._rate_limited_until = 0
+        agent._use_prompt_caching = False
+        agent._use_native_cache_layout = False
+        agent.context_compressor = MagicMock()
+        agent.context_compressor.update_model = MagicMock()
+
+        # Snapshot the original state
+        agent._primary_runtime = {
+            "model": "gpt-5.5",
+            "provider": "openai-codex",
+            "base_url": "https://chatgpt.com/backend-api/codex",
+            "api_mode": "codex_responses",
+            "api_key": "original-key-entry-1",
+            "client_kwargs": {
+                "api_key": "original-key-entry-1",
+                "base_url": "https://chatgpt.com/backend-api/codex",
+            },
+            "use_prompt_caching": False,
+            "use_native_cache_layout": False,
+            "compressor_model": "gpt-5.5",
+            "compressor_base_url": "https://chatgpt.com/backend-api/codex",
+            "compressor_api_key": "original-key-entry-1",
+            "compressor_provider": "openai-codex",
+            "compressor_context_length": 128000,
+            "compressor_threshold_tokens": 0.8,
+        }
+
+        # Mock client creation methods
+        agent._create_openai_client = MagicMock(return_value=MagicMock())
+        agent._apply_client_headers_for_base_url = MagicMock()
+        agent._replace_primary_openai_client = MagicMock(return_value=True)
+
+        return agent
+
+
+    def test_restore_uses_freshest_available_entry(self):
+        """When multiple entries are available, restore should select the pool's best pick."""
+        entries = [
+            _make_entry("entry-1", "key-1", priority=0,
+                         last_status="exhausted", last_status_at=time.time() + 3600),
+            _make_entry("entry-2", "key-2", priority=1),
+            _make_entry("entry-3", "key-3", priority=2),
+        ]
+        pool = _build_mock_pool(entries)
+
+        agent = self._make_agent(pool)
+        result = agent._restore_primary_runtime()
+
+        assert result is True
+        # entry-1 is exhausted, so pool should select entry-2
+        assert agent.api_key == "key-2"
+        assert agent._client_kwargs["api_key"] == "key-2"
+
+
+
+    def test_restore_rebuilds_client_after_reselect(self):
+        """After re-selecting from pool, client should be rebuilt with new key."""
+        entries = [
+            _make_entry("entry-1", "key-1", priority=0),
+        ]
+        pool = _build_mock_pool(entries)
+
+        agent = self._make_agent(pool)
+        result = agent._restore_primary_runtime()
+
+        assert result is True
+        # _swap_credential rebuilds the live OpenAI client with the fresh key.
+        agent._replace_primary_openai_client.assert_called_once_with(
+            reason="credential_rotation",
+        )
+
+
+    def test_restore_updates_base_url_from_pool_entry(self):
+        """If pool entry has a different base_url, restore should update it."""
+        entries = [
+            {
+                **_make_entry("entry-1", "key-1", priority=0),
+                "base_url": "https://custom-endpoint.example.com/v1",
+            },
+        ]
+        pool = _build_mock_pool(entries)
+
+        agent = self._make_agent(pool)
+        result = agent._restore_primary_runtime()
+
+        assert result is True
+        assert "custom-endpoint.example.com" in agent.base_url
+        assert "custom-endpoint.example.com" in agent._client_kwargs["base_url"]