Fix/auth and model robustness

archytas/agent.py

@@ -180,6 +180,17 @@ def __init__(
         # make_tool_dict runs). Map of tool_name -> tool callable.
         self.statetools: dict[str, Callable] = {}
 
+    def set_model(self, model: BaseArchytasModel) -> None:
+        """Swap the active model at runtime, keeping the chat history in sync.
+
+        The chat history keeps its own reference to the model (used for token
+        budgeting and the model metadata it serializes), so assigning
+        ``agent.model`` alone leaves that reference stale. Always route runtime
+        model changes through this method so both stay consistent.
+        """
+        self.model = model
+        self.chat_history.set_model(model)
+
     def get_prompt_sections(self) -> list[PromptSection]:
         """Return the ordered list of ``PromptSection``s composing the system
         prompt.

archytas/models/azure.py

@@ -18,12 +18,14 @@ def auth(self, **kwargs) -> None:
         else:
             auth_token = self.config.api_key
         if not auth_token:
-            auth_token = DEFERRED_TOKEN_VALUE
-        set_env_auth(AZURE_OPENAI_API_KEY=auth_token)
-
-        # Replace local auth token from value from environment variables to allow fetching preset auth variables in the
-        # environment.
-        auth_token = os.environ.get('AZURE_OPENAI_API_KEY', DEFERRED_TOKEN_VALUE)
+            # No key configured: defer to the environment (e.g. a user-set
+            # AZURE_OPENAI_API_KEY), without clobbering it.
+            set_env_auth(AZURE_OPENAI_API_KEY=DEFERRED_TOKEN_VALUE)
+            auth_token = os.environ.get('AZURE_OPENAI_API_KEY', DEFERRED_TOKEN_VALUE)
+        else:
+            # An explicit key wins over the environment -- including a value an
+            # earlier model instance wrote there -- so a new key takes effect now.
+            os.environ['AZURE_OPENAI_API_KEY'] = auth_token
 
         if auth_token != DEFERRED_TOKEN_VALUE:
             self.config.api_key = auth_token

archytas/models/openai.py

@@ -51,12 +51,14 @@ def auth(self, **kwargs) -> None:
         else:
             auth_token = self.config.api_key
         if not auth_token:
-            auth_token = DEFERRED_TOKEN_VALUE
-        set_env_auth(OPENAI_API_KEY=auth_token)
-
-        # Replace local auth token from value from environment variables to allow fetching preset auth variables in the
-        # environment.
-        auth_token = os.environ.get('OPENAI_API_KEY', DEFERRED_TOKEN_VALUE)
+            # No key configured: defer to the environment (e.g. a user-set
+            # OPENAI_API_KEY), without clobbering it.
+            set_env_auth(OPENAI_API_KEY=DEFERRED_TOKEN_VALUE)
+            auth_token = os.environ.get('OPENAI_API_KEY', DEFERRED_TOKEN_VALUE)
+        else:
+            # An explicit key wins over the environment -- including a value an
+            # earlier model instance wrote there -- so a new key takes effect now.
+            os.environ['OPENAI_API_KEY'] = auth_token
 
         if auth_token != DEFERRED_TOKEN_VALUE:
             self.config.api_key = auth_token

tests/test_agent_set_model.py

@@ -0,0 +1,22 @@
+"""Regression test for Agent.set_model keeping the chat history in sync.
+
+The chat history holds its own model reference (used for token budgeting and
+the model metadata it serializes for UIs). Assigning ``agent.model`` alone
+leaves that reference stale; ``Agent.set_model`` must update both.
+"""
+from archytas.agent import Agent
+from archytas.models.openai import OpenAIModel
+
+
+def test_set_model_updates_agent_and_chat_history(monkeypatch):
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy")
+    first = OpenAIModel({"api_key": "sk-a", "model_name": "gpt-4o-mini"})
+    second = OpenAIModel({"api_key": "sk-b", "model_name": "gpt-4o"})
+
+    agent = Agent(model=first, spinner=None, rich_print=False)
+    assert agent.model is first
+    assert agent.chat_history.model is first
+
+    agent.set_model(second)
+    assert agent.model is second
+    assert agent.chat_history.model is second

tests/test_auth_env_precedence.py

@@ -0,0 +1,31 @@
+"""Regression tests for API-key precedence in OpenAI/Azure model auth.
+
+These guard the subtle bug where an explicitly-configured key was shadowed by
+a stale value left in ``os.environ`` (written there by an earlier model
+instance via ``set_env_auth``'s ``setdefault``), so a freshly-entered key only
+took effect after a process restart.
+
+No network calls: constructing the model is enough to resolve the effective key.
+"""
+import pytest
+
+from archytas.models.openai import OpenAIModel
+
+
+def _client_key(model) -> str:
+    secret = model.model.openai_api_key
+    return secret.get_secret_value() if hasattr(secret, "get_secret_value") else str(secret)
+
+
+def test_explicit_key_overrides_stale_env(monkeypatch):
+    # Simulate a stale/bad key left in the environment by an earlier build.
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-STALE-must-not-win")
+    model = OpenAIModel({"api_key": "sk-explicit-config-key", "model_name": "gpt-4o-mini"})
+    assert _client_key(model) == "sk-explicit-config-key"
+
+
+def test_defers_to_env_when_no_key_configured(monkeypatch):
+    # With no key configured, a user-provided environment variable should win.
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-from-user-environment")
+    model = OpenAIModel({"model_name": "gpt-4o-mini"})
+    assert _client_key(model) == "sk-from-user-environment"