Merge pull request #63 from zhujian0805/main

Add CLAUDE_CODE_OAUTH_TOKEN support and related test

Author: zhujian0805

code_assistant_manager/strategies.py

@@ -71,6 +71,11 @@ def setup_environment(self, context: ExecutionContext) -> Dict[str, str]:
         endpoint_config = context.endpoint_config
         env["ANTHROPIC_BASE_URL"] = str(endpoint_config.url)
         env["ANTHROPIC_AUTH_TOKEN"] = endpoint_config.get_api_key_value() or ""
+        env["CLAUDE_CODE_OAUTH_TOKEN"] = (
+            os.environ.get("CLAUDE_CODE_OAUTH_TOKEN")
+            or endpoint_config.get_api_key_value()
+            or ""
+        )
         env["ANTHROPIC_MODEL"] = str(primary_model)
         env["ANTHROPIC_SMALL_FAST_MODEL"] = str(secondary_model)
         env["CLAUDE_MODEL_2"] = str(secondary_model)

code_assistant_manager/tools/claude.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import List
 
 from .base import CLITool
@@ -47,6 +48,11 @@ def run(self, args: List[str] = None) -> int:
                 ToolEnvironmentBuilder(endpoint_config, model_vars)
                 .set_base_url("ANTHROPIC_BASE_URL")
                 .set_api_key("ANTHROPIC_AUTH_TOKEN")
+                .set_custom_var(
+                    "CLAUDE_CODE_OAUTH_TOKEN",
+                    os.environ.get("CLAUDE_CODE_OAUTH_TOKEN")
+                    or endpoint_config.get("actual_api_key", ""),
+                )
                 .set_model("ANTHROPIC_MODEL", "primary_model")
                 .set_model("ANTHROPIC_SMALL_FAST_MODEL", "secondary_model")
                 .set_model("CLAUDE_MODEL_2", "secondary_model")
@@ -71,6 +77,7 @@ def run(self, args: List[str] = None) -> int:
             print(
                 f"ANTHROPIC_BASE_URL={env['ANTHROPIC_BASE_URL']} "
                 f"ANTHROPIC_AUTH_TOKEN=dummy "
+                f"CLAUDE_CODE_OAUTH_TOKEN=*** "
                 f"ANTHROPIC_MODEL={primary_model} "
                 f"ANTHROPIC_DEFAULT_SONNET_MODEL={primary_model} "
                 f"ANTHROPIC_SMALL_FAST_MODEL={secondary_model} "

tests/test_tools.py

@@ -446,9 +446,42 @@ def test_claude_environment_variables(
         env_used = mock_run.call_args[1]["env"]
         assert env_used["ANTHROPIC_BASE_URL"] == "https://api.example.com"
         assert env_used["ANTHROPIC_AUTH_TOKEN"] == "key123"
+        assert env_used["CLAUDE_CODE_OAUTH_TOKEN"] == "key123"
         assert env_used["ANTHROPIC_MODEL"] == "claude-3"
         assert env_used["NODE_TLS_REJECT_UNAUTHORIZED"] == "0"
 
+    @patch("code_assistant_manager.tools.subprocess.run")
+    @patch("code_assistant_manager.tools.select_two_models")
+    @patch.object(ClaudeTool, "_ensure_tool_installed", return_value=True)
+    @patch("code_assistant_manager.tools.EndpointManager")
+    @patch.dict(
+        os.environ,
+        {
+            "CODE_ASSISTANT_MANAGER_NONINTERACTIVE": "1",
+            "CLAUDE_CODE_OAUTH_TOKEN": "oauth-token-from-env",
+        },
+    )
+    def test_claude_environment_prefers_existing_oauth_token(
+        self, mock_em_class, mock_install, mock_select, mock_run, config_manager
+    ):
+        """Test Claude uses an existing OAuth token from the environment."""
+        mock_em = MagicMock()
+        mock_em_class.return_value = mock_em
+
+        mock_em.select_endpoint.return_value = (True, "endpoint1")
+        mock_em.get_endpoint_config.return_value = (
+            True,
+            {"endpoint": "https://api.example.com", "actual_api_key": "key123"},
+        )
+        mock_em.fetch_models.return_value = (True, ["claude-3"])
+        mock_select.return_value = (True, ("claude-3", "claude-2"))
+
+        tool = ClaudeTool(config_manager)
+        tool.run([])
+
+        env_used = mock_run.call_args[1]["env"]
+        assert env_used["CLAUDE_CODE_OAUTH_TOKEN"] == "oauth-token-from-env"
+
 
 class TestToolErrorHandling:
     """Test error handling in tools."""