diff --git a/backend/packages/harness/deerflow/models/claude_provider.py b/backend/packages/harness/deerflow/models/claude_provider.py index 9e08eea..a53939a 100644 --- a/backend/packages/harness/deerflow/models/claude_provider.py +++ b/backend/packages/harness/deerflow/models/claude_provider.py @@ -5,6 +5,7 @@ Supports two authentication modes: 2. Claude Code OAuth token (Authorization: Bearer header) - Detected by sk-ant-oat prefix - Requires anthropic-beta: oauth-2025-04-20,claude-code-20250219 + - Requires billing header in system prompt for all OAuth requests Auto-loads credentials from explicit runtime handoff: - $ANTHROPIC_API_KEY environment variable @@ -14,8 +15,13 @@ Auto-loads credentials from explicit runtime handoff: - ~/.claude/.credentials.json """ +import hashlib +import json import logging +import os +import socket import time +import uuid from typing import Any import anthropic @@ -27,6 +33,12 @@ logger = logging.getLogger(__name__) MAX_RETRIES = 3 THINKING_BUDGET_RATIO = 0.8 +# Billing header required by Anthropic API for OAuth token access. +# Must be the first system prompt block. Format mirrors Claude Code CLI. +# Override with ANTHROPIC_BILLING_HEADER env var if the hardcoded version drifts. +_DEFAULT_BILLING_HEADER = "x-anthropic-billing-header: cc_version=2.1.85.351; cc_entrypoint=cli; cch=6c6d5;" +OAUTH_BILLING_HEADER = os.environ.get("ANTHROPIC_BILLING_HEADER", _DEFAULT_BILLING_HEADER) + class ClaudeChatModel(ChatAnthropic): """ChatAnthropic with OAuth Bearer auth, prompt caching, and smart thinking. @@ -125,9 +137,12 @@ class ClaudeChatModel(ChatAnthropic): stop: list[str] | None = None, **kwargs: Any, ) -> dict: - """Override to inject prompt caching and thinking budget.""" + """Override to inject prompt caching, thinking budget, and OAuth billing.""" payload = super()._get_request_payload(input_, stop=stop, **kwargs) + if self._is_oauth: + self._apply_oauth_billing(payload) + if self.enable_prompt_caching: self._apply_prompt_caching(payload) @@ -136,6 +151,44 @@ class ClaudeChatModel(ChatAnthropic): return payload + def _apply_oauth_billing(self, payload: dict) -> None: + """Inject the billing header block required for all OAuth requests. + + The billing block is always placed first in the system list, removing any + existing occurrence to avoid duplication or out-of-order positioning. + """ + billing_block = {"type": "text", "text": OAUTH_BILLING_HEADER} + + system = payload.get("system") + if isinstance(system, list): + # Remove any existing billing blocks, then insert a single one at index 0. + filtered = [ + b for b in system + if not (isinstance(b, dict) and OAUTH_BILLING_HEADER in b.get("text", "")) + ] + payload["system"] = [billing_block] + filtered + elif isinstance(system, str): + if OAUTH_BILLING_HEADER in system: + payload["system"] = [billing_block] + else: + payload["system"] = [billing_block, {"type": "text", "text": system}] + else: + payload["system"] = [billing_block] + + # Add metadata.user_id required by the API for OAuth billing validation + if not isinstance(payload.get("metadata"), dict): + payload["metadata"] = {} + if "user_id" not in payload["metadata"]: + # Generate a stable device_id from the machine's hostname + hostname = socket.gethostname() + device_id = hashlib.sha256(f"deerflow-{hostname}".encode()).hexdigest() + session_id = str(uuid.uuid4()) + payload["metadata"]["user_id"] = json.dumps({ + "device_id": device_id, + "account_uuid": "deerflow", + "session_id": session_id, + }) + def _apply_prompt_caching(self, payload: dict) -> None: """Apply ephemeral cache_control to system and recent messages.""" # Cache system messages diff --git a/backend/tests/test_claude_provider_oauth_billing.py b/backend/tests/test_claude_provider_oauth_billing.py new file mode 100644 index 0000000..a341ffe --- /dev/null +++ b/backend/tests/test_claude_provider_oauth_billing.py @@ -0,0 +1,113 @@ +"""Tests for ClaudeChatModel._apply_oauth_billing.""" + +import json + +import pytest + +from deerflow.models.claude_provider import OAUTH_BILLING_HEADER, ClaudeChatModel + + +def _make_model() -> ClaudeChatModel: + """Return a minimal ClaudeChatModel instance in OAuth mode without network calls.""" + import unittest.mock as mock + + with mock.patch.object(ClaudeChatModel, "model_post_init"): + m = ClaudeChatModel(model="claude-sonnet-4-6", anthropic_api_key="sk-ant-oat-fake-token") # type: ignore[call-arg] + m._is_oauth = True + m._oauth_access_token = "sk-ant-oat-fake-token" + return m + + +@pytest.fixture() +def model() -> ClaudeChatModel: + return _make_model() + + +def _billing_block() -> dict: + return {"type": "text", "text": OAUTH_BILLING_HEADER} + + +# --------------------------------------------------------------------------- +# Billing block injection +# --------------------------------------------------------------------------- + + +def test_billing_injected_first_when_no_system(model): + payload: dict = {} + model._apply_oauth_billing(payload) + assert payload["system"][0] == _billing_block() + + +def test_billing_injected_first_into_list(model): + payload = {"system": [{"type": "text", "text": "You are a helpful assistant."}]} + model._apply_oauth_billing(payload) + assert payload["system"][0] == _billing_block() + assert payload["system"][1]["text"] == "You are a helpful assistant." + + +def test_billing_injected_first_into_string_system(model): + payload = {"system": "You are helpful."} + model._apply_oauth_billing(payload) + assert payload["system"][0] == _billing_block() + assert payload["system"][1]["text"] == "You are helpful." + + +def test_billing_not_duplicated_on_second_call(model): + payload = {"system": [{"type": "text", "text": "prompt"}]} + model._apply_oauth_billing(payload) + model._apply_oauth_billing(payload) + billing_count = sum( + 1 for b in payload["system"] + if isinstance(b, dict) and OAUTH_BILLING_HEADER in b.get("text", "") + ) + assert billing_count == 1 + + +def test_billing_moved_to_first_if_not_already_first(model): + """Billing block already present but not first — must be normalized to index 0.""" + payload = { + "system": [ + {"type": "text", "text": "other block"}, + _billing_block(), + ] + } + model._apply_oauth_billing(payload) + assert payload["system"][0] == _billing_block() + assert len([b for b in payload["system"] if OAUTH_BILLING_HEADER in b.get("text", "")]) == 1 + + +def test_billing_string_with_header_collapsed_to_single_block(model): + """If system is a string that already contains the billing header, collapse to one block.""" + payload = {"system": OAUTH_BILLING_HEADER} + model._apply_oauth_billing(payload) + assert payload["system"] == [_billing_block()] + + +# --------------------------------------------------------------------------- +# metadata.user_id +# --------------------------------------------------------------------------- + + +def test_metadata_user_id_added_when_missing(model): + payload: dict = {} + model._apply_oauth_billing(payload) + assert "metadata" in payload + user_id = json.loads(payload["metadata"]["user_id"]) + assert "device_id" in user_id + assert "session_id" in user_id + assert user_id["account_uuid"] == "deerflow" + + +def test_metadata_user_id_not_overwritten_if_present(model): + payload = {"metadata": {"user_id": "existing-value"}} + model._apply_oauth_billing(payload) + assert payload["metadata"]["user_id"] == "existing-value" + + +def test_metadata_non_dict_replaced_with_dict(model): + """Non-dict metadata (e.g. None or a string) should be replaced, not crash.""" + for bad_value in (None, "string-metadata", 42): + payload = {"metadata": bad_value} + model._apply_oauth_billing(payload) + assert isinstance(payload["metadata"], dict) + assert "user_id" in payload["metadata"]