test_databricks_executor.py

"""Tests for DatabricksExecutor with a mock OpenAI client."""

import asyncio
import json
import sys
import unittest
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import databricks.sdk.config as _sdk_config_mod

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from omnigent.inner.databricks_executor import (
    DatabricksExecutor,
    _convert_messages,
    _convert_tools_to_openai,
)
from omnigent.inner.executor import (
    ExecutorConfig,
    ExecutorError,
    TextChunk,
    ToolCallRequest,
    TurnComplete,
)


def _run(coro):
    loop = asyncio.new_event_loop()
    try:
        return loop.run_until_complete(coro)
    finally:
        loop.run_until_complete(loop.shutdown_asyncgens())
        loop.run_until_complete(loop.shutdown_default_executor())
        loop.close()


# ---------------------------------------------------------------------------
# Fake streaming response objects (mimicking OpenAI streaming types)
# ---------------------------------------------------------------------------


@dataclass
class FakeFunctionDelta:
    name: str | None = None
    arguments: str | None = None


@dataclass
class FakeToolCallDelta:
    index: int = 0
    function: FakeFunctionDelta | None = None


@dataclass
class FakeDelta:
    """
    Minimal stream-delta test double for DatabricksExecutor.

    :param content: Raw ``delta.content`` payload, e.g. ``"hello"`` or
        ``[{"type": "reasoning", "summary": [...]}]``.
    :param tool_calls: Optional streamed tool-call deltas.
    """

    content: Any = None
    tool_calls: list[FakeToolCallDelta] | None = None


@dataclass
class FakeStreamChoice:
    delta: FakeDelta = field(default_factory=FakeDelta)
    index: int = 0
    finish_reason: str | None = None


@dataclass
class FakeStreamChunk:
    choices: list[FakeStreamChoice] = field(default_factory=list)


def _make_text_stream(text: str) -> list[FakeStreamChunk]:
    """Create a stream that yields text content then stops."""
    chunks = []
    # Yield text in a single chunk for simplicity
    if text:
        chunks.append(FakeStreamChunk(choices=[FakeStreamChoice(delta=FakeDelta(content=text))]))
    # Final chunk with finish_reason=stop
    chunks.append(FakeStreamChunk(choices=[FakeStreamChoice(finish_reason="stop")]))
    return chunks


def _make_tool_call_stream(
    tool_calls: list[tuple[str, str]],
    text: str | None = None,
) -> list[FakeStreamChunk]:
    """Create a stream that yields tool calls.

    Args:
        tool_calls: list of (name, arguments_json) tuples
        text: optional text content to include before tool calls
    """
    chunks = []
    if text:
        chunks.append(FakeStreamChunk(choices=[FakeStreamChoice(delta=FakeDelta(content=text))]))
    for idx, (name, args) in enumerate(tool_calls):
        chunks.append(
            FakeStreamChunk(
                choices=[
                    FakeStreamChoice(
                        delta=FakeDelta(
                            tool_calls=[
                                FakeToolCallDelta(
                                    index=idx,
                                    function=FakeFunctionDelta(name=name, arguments=args),
                                )
                            ]
                        )
                    )
                ]
            )
        )
    # Final chunk with finish_reason=tool_calls
    chunks.append(FakeStreamChunk(choices=[FakeStreamChoice(finish_reason="tool_calls")]))
    return chunks


class FakeCompletions:
    """Mimics client.chat.completions."""

    def __init__(self, stream_chunks: list[FakeStreamChunk]):
        self._chunks = stream_chunks
        self.last_kwargs: dict[str, Any] = {}

    def create(self, **kwargs) -> list[FakeStreamChunk]:
        self.last_kwargs = kwargs
        return iter(self._chunks)


class FakeChat:
    def __init__(self, completions: FakeCompletions):
        self.completions = completions


class FakeClient:
    """Mimics the OpenAI client."""

    def __init__(self, stream_chunks: list[FakeStreamChunk]):
        self.chat = FakeChat(FakeCompletions(stream_chunks))


def test_kimi_reasoning_content_blocks_are_not_text_chunks() -> None:
    """
    Kimi streams reasoning summaries as ``delta.content`` block lists before
    the assistant answer; the executor must not hand those lists to
    :class:`TextChunk` or append them to the final assistant response.
    """

    async def _t() -> None:
        """Drive DatabricksExecutor over a Kimi-shaped stream."""
        chunks = [
            FakeStreamChunk(
                choices=[
                    FakeStreamChoice(
                        delta=FakeDelta(
                            content=[
                                {
                                    "type": "reasoning",
                                    "summary": [{"type": "summary_text", "text": "Thinking"}],
                                }
                            ]
                        )
                    )
                ]
            ),
            FakeStreamChunk(choices=[FakeStreamChoice(delta=FakeDelta(content="Hello"))]),
            FakeStreamChunk(choices=[FakeStreamChoice(finish_reason="stop")]),
        ]
        executor = DatabricksExecutor(client=FakeClient(chunks))

        events = [
            e
            async for e in executor.run_turn(
                messages=[{"role": "user", "content": "hi"}],
                tools=[],
                system_prompt="Be helpful.",
                config=ExecutorConfig(model="databricks-kimi-k2-6"),
            )
        ]

        text_events = [e for e in events if isinstance(e, TextChunk)]
        turn_events = [e for e in events if isinstance(e, TurnComplete)]
        assert [e.text for e in text_events] == ["Hello"]
        assert all(isinstance(e.text, str) for e in text_events)
        assert len(turn_events) == 1
        assert turn_events[0].response == "Hello"

    _run(_t())


def test_text_content_block_lists_are_collapsed_to_text_chunks() -> None:
    """
    Providers may stream assistant-visible text as content block lists; those
    recognized text blocks must still produce normal string
    :class:`TextChunk` events and the correct final response.
    """

    async def _t() -> None:
        """Drive DatabricksExecutor over text content block deltas."""
        chunks = [
            FakeStreamChunk(
                choices=[
                    FakeStreamChoice(
                        delta=FakeDelta(
                            content=[
                                {"type": "text", "text": "Hello"},
                                {"type": "output_text", "text": " world"},
                            ]
                        )
                    )
                ]
            ),
            FakeStreamChunk(choices=[FakeStreamChoice(finish_reason="stop")]),
        ]
        executor = DatabricksExecutor(client=FakeClient(chunks))

        events = [
            e
            async for e in executor.run_turn(
                messages=[{"role": "user", "content": "hi"}],
                tools=[],
                system_prompt="Be helpful.",
                config=ExecutorConfig(model="block-list-model"),
            )
        ]

        text_events = [e for e in events if isinstance(e, TextChunk)]
        turn_events = [e for e in events if isinstance(e, TurnComplete)]
        assert [e.text for e in text_events] == ["Hello world"]
        assert len(turn_events) == 1
        assert turn_events[0].response == "Hello world"

    _run(_t())


# ---------------------------------------------------------------------------
# Tests: message and tool conversion helpers
# ---------------------------------------------------------------------------


class TestConvertTools(unittest.TestCase):
    def test_basic_tool(self):
        tools = [
            {
                "name": "sql",
                "description": "Run SQL",
                "parameters": {"type": "object", "properties": {"q": {"type": "string"}}},
            }
        ]
        result = _convert_tools_to_openai(tools)
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["type"], "function")
        self.assertEqual(result[0]["function"]["name"], "sql")
        self.assertIn("properties", result[0]["function"]["parameters"])

    def test_tool_without_parameters(self):
        tools = [{"name": "ping", "description": "Ping"}]
        result = _convert_tools_to_openai(tools)
        self.assertEqual(result[0]["function"]["parameters"], {"type": "object", "properties": {}})

    def test_preserves_required_args_in_async_tool_schema(self):
        tools = [
            {
                "name": "sys_call_async",
                "description": "Async call",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "tool": {"type": "string"},
                        "args": {
                            "type": "object",
                            "default": {},
                            "properties": {
                                "example_key": {
                                    "anyOf": [{"type": "string"}],
                                },
                            },
                            "additionalProperties": {
                                "anyOf": [{"type": "string"}],
                            },
                        },
                    },
                    "required": ["tool", "args"],
                },
            }
        ]
        result = _convert_tools_to_openai(tools)
        self.assertEqual(
            result[0]["function"]["parameters"]["required"],
            ["tool", "args"],
        )
        self.assertIn(
            "example_key",
            result[0]["function"]["parameters"]["properties"]["args"]["properties"],
        )

    def test_preserves_required_args_in_session_send_schema(self):
        tools = [
            {
                "name": "sys_session_send",
                "description": "Session send",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "tool": {"type": "string"},
                        "session": {"type": "string"},
                        "args": {"type": "object", "additionalProperties": True},
                    },
                    "required": ["tool", "session", "args"],
                },
            }
        ]
        result = _convert_tools_to_openai(tools)
        self.assertEqual(
            result[0]["function"]["parameters"]["required"],
            ["tool", "session", "args"],
        )

    def test_empty_tools(self):
        self.assertEqual(_convert_tools_to_openai([]), [])

    def test_invalid_tool_name_is_normalized_for_provider(self):
        tools = [{"name": "sys_runtime_execute", "description": "Run code"}]
        result = _convert_tools_to_openai(tools)
        self.assertEqual(result[0]["function"]["name"], "sys_runtime_execute")
        self.assertEqual(result[0]["function"]["description"], "Run code")


class TestConvertMessages(unittest.TestCase):
    def test_system_prompt(self):
        result = _convert_messages([], "You are helpful.")
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["role"], "system")
        self.assertEqual(result[0]["content"], "You are helpful.")

    def test_user_and_assistant(self):
        msgs = [
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there!"},
        ]
        result = _convert_messages(msgs, "")
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0]["role"], "user")
        self.assertEqual(result[1]["role"], "assistant")

    def test_tool_call_and_result_pair(self):
        msgs = [
            {"role": "user", "content": "Run a query"},
            {"role": "tool_call", "content": {"tool": "sql", "args": {"q": "SELECT 1"}}},
            {"role": "tool_result", "content": {"rows": [1]}},
        ]
        result = _convert_messages(msgs, "sys")
        # system + user + assistant(tool_calls) + tool
        self.assertEqual(len(result), 4)
        self.assertEqual(result[0]["role"], "system")
        self.assertEqual(result[1]["role"], "user")
        self.assertEqual(result[2]["role"], "assistant")

    def test_invalid_tool_name_is_normalized_in_history_replay(self):
        msgs = [
            {
                "role": "tool_call",
                "content": {"tool": "sys_runtime_execute", "args": {"code": "print(1)"}},
            },
            {"role": "tool_result", "content": {"stdout": "1\n"}},
        ]
        result = _convert_messages(msgs, "")
        self.assertEqual(result[0]["tool_calls"][0]["function"]["name"], "sys_runtime_execute")
        self.assertEqual(result[1]["role"], "tool")

    def test_orphan_tool_result(self):
        msgs = [{"role": "tool_result", "content": "some result"}]
        result = _convert_messages(msgs, "")
        self.assertEqual(len(result), 1)
        self.assertIn("tool result", result[0]["content"])

    def test_tool_call_content_as_string(self):
        """tool_call content might be a JSON string instead of dict."""
        msgs = [
            {
                "role": "tool_call",
                "content": json.dumps({"tool": "search", "args": {"q": "test"}}),
            },
            {"role": "tool_result", "content": "found it"},
        ]
        result = _convert_messages(msgs, "")
        self.assertEqual(result[0]["tool_calls"][0]["function"]["name"], "search")


# ---------------------------------------------------------------------------
# Tests: DatabricksExecutor with fake streaming client
# ---------------------------------------------------------------------------


class TestDatabricksExecutorTextResponse(unittest.TestCase):
    def test_simple_text_response(self):
        async def _t():
            chunks = _make_text_stream("Hello world!")
            executor = DatabricksExecutor(client=FakeClient(chunks))

            events = [
                e
                async for e in executor.run_turn(
                    messages=[{"role": "user", "content": "Hi"}],
                    tools=[],
                    system_prompt="Be nice.",
                    config=ExecutorConfig(model="test-model"),
                )
            ]

            # TextChunk + TurnComplete
            text_events = [e for e in events if isinstance(e, TextChunk)]
            turn_events = [e for e in events if isinstance(e, TurnComplete)]
            self.assertEqual(len(text_events), 1)
            self.assertEqual(text_events[0].text, "Hello world!")
            self.assertEqual(len(turn_events), 1)
            self.assertEqual(turn_events[0].response, "Hello world!")

        _run(_t())

    def test_empty_content(self):
        async def _t():
            chunks = _make_text_stream("")
            executor = DatabricksExecutor(client=FakeClient(chunks))
            events = [e async for e in executor.run_turn([], [], "")]
            turn_events = [e for e in events if isinstance(e, TurnComplete)]
            self.assertEqual(len(turn_events), 1)
            self.assertEqual(turn_events[0].response, "")

        _run(_t())


class TestDatabricksExecutorToolCalls(unittest.TestCase):
    def test_single_tool_call(self):
        async def _t():
            chunks = _make_tool_call_stream([("sql_query", '{"query": "SELECT 1"}')])
            executor = DatabricksExecutor(client=FakeClient(chunks))
            events = [
                e
                async for e in executor.run_turn(
                    [{"role": "user", "content": "query"}],
                    [{"name": "sql_query", "description": "Run SQL"}],
                    "sys",
                )
            ]
            tool_events = [e for e in events if isinstance(e, ToolCallRequest)]
            self.assertEqual(len(tool_events), 1)
            self.assertEqual(tool_events[0].name, "sql_query")
            self.assertEqual(tool_events[0].args["query"], "SELECT 1")

        _run(_t())

    def test_multiple_tool_calls(self):
        async def _t():
            chunks = _make_tool_call_stream(
                [
                    ("tool_a", '{"x": 1}'),
                    ("tool_b", '{"y": 2}'),
                ]
            )
            executor = DatabricksExecutor(client=FakeClient(chunks))
            events = [e async for e in executor.run_turn([], [], "")]
            tool_events = [e for e in events if isinstance(e, ToolCallRequest)]
            self.assertEqual(len(tool_events), 2)
            self.assertEqual(tool_events[0].name, "tool_a")
            self.assertEqual(tool_events[1].name, "tool_b")

        _run(_t())

    def test_tool_call_with_text(self):
        """Model returns both text and tool calls."""

        async def _t():
            chunks = _make_tool_call_stream(
                [("search", "{}")],
                text="Let me search for that.",
            )
            executor = DatabricksExecutor(client=FakeClient(chunks))
            events = [e async for e in executor.run_turn([], [], "")]
            names = [type(e).__name__ for e in events]
            self.assertIn("ToolCallRequest", names)
            self.assertIn("TextChunk", names)
            self.assertNotIn("TurnComplete", names)

        _run(_t())

    def test_malformed_arguments(self):
        """If arguments aren't valid JSON, put them in a 'raw' key."""

        async def _t():
            chunks = _make_tool_call_stream([("t", "not json")])
            executor = DatabricksExecutor(client=FakeClient(chunks))
            events = [e async for e in executor.run_turn([], [], "")]
            tool_events = [e for e in events if isinstance(e, ToolCallRequest)]
            self.assertEqual(len(tool_events), 1)
            self.assertEqual(tool_events[0].args["raw"], "not json")

        _run(_t())


class TestDatabricksExecutorErrors(unittest.TestCase):
    def test_empty_stream(self):
        """Stream with no chunks yields a TurnComplete with empty text."""

        async def _t():
            executor = DatabricksExecutor(client=FakeClient([]))
            events = [e async for e in executor.run_turn([], [], "")]
            self.assertEqual(len(events), 1)
            self.assertIsInstance(events[0], TurnComplete)
            self.assertEqual(events[0].response, "")

        _run(_t())

    def test_api_exception(self):
        async def _t():
            class ExplodingClient:
                class chat:
                    class completions:
                        @staticmethod
                        def create(**kwargs):
                            raise RuntimeError("API down")

            executor = DatabricksExecutor(client=ExplodingClient())
            events = [e async for e in executor.run_turn([], [], "")]
            self.assertEqual(len(events), 1)
            self.assertIsInstance(events[0], ExecutorError)
            self.assertIn("API down", events[0].message)

        _run(_t())


class TestDatabricksExecutorConfig(unittest.TestCase):
    def test_passes_model_and_params(self):
        """Verify the executor passes model, temperature, max_tokens to the API."""

        async def _t():
            chunks = _make_text_stream("ok")
            client = FakeClient(chunks)
            executor = DatabricksExecutor(client=client)

            config = ExecutorConfig(
                model="claude-sonnet-4",
                temperature=0.7,
                max_tokens=2048,
            )
            [
                e
                async for e in executor.run_turn(
                    [{"role": "user", "content": "Hi"}],
                    [],
                    "system prompt",
                    config=config,
                )
            ]

            kwargs = client.chat.completions.last_kwargs
            self.assertEqual(kwargs["model"], "claude-sonnet-4")
            self.assertEqual(kwargs["temperature"], 0.7)
            self.assertEqual(kwargs["max_tokens"], 2048)
            self.assertTrue(kwargs["stream"])

        _run(_t())

    def test_default_model(self):
        """When no model is specified, falls back to databricks-claude-sonnet-4-6."""

        async def _t():
            chunks = _make_text_stream("ok")
            client = FakeClient(chunks)
            executor = DatabricksExecutor(client=client)

            [e async for e in executor.run_turn([], [], "", config=ExecutorConfig())]
            self.assertEqual(
                client.chat.completions.last_kwargs["model"],
                "databricks-claude-sonnet-4-6",
            )

        _run(_t())

    def test_tools_passed_in_openai_format(self):
        async def _t():
            chunks = _make_text_stream("ok")
            client = FakeClient(chunks)
            executor = DatabricksExecutor(client=client)

            tools = [
                {
                    "name": "sql",
                    "description": "Run SQL",
                    "parameters": {"type": "object", "properties": {}},
                }
            ]
            [e async for e in executor.run_turn([], tools, "")]

            passed_tools = client.chat.completions.last_kwargs["tools"]
            self.assertEqual(len(passed_tools), 1)
            self.assertEqual(passed_tools[0]["type"], "function")
            self.assertEqual(passed_tools[0]["function"]["name"], "sql")

        _run(_t())


class TestDatabricksExecutorMultiTurn(unittest.TestCase):
    """Test a realistic multi-turn scenario: user asks -> tool call -> tool result -> response."""

    def test_tool_call_then_response(self):
        async def _t():
            # Turn 1: model wants to call a tool
            chunks1 = _make_tool_call_stream([("search", '{"q": "test"}')])
            client1 = FakeClient(chunks1)
            executor = DatabricksExecutor(client=client1)

            events1 = [
                e
                async for e in executor.run_turn(
                    [{"role": "user", "content": "search for test"}],
                    [{"name": "search", "description": "Search"}],
                    "sys",
                )
            ]
            tool_events = [e for e in events1 if isinstance(e, ToolCallRequest)]
            self.assertEqual(len(tool_events), 1)
            self.assertEqual(tool_events[0].name, "search")

            # Turn 2: after tool result, model gives final answer
            chunks2 = _make_text_stream("Found 3 results.")
            executor._client = FakeClient(chunks2)
            events2 = [
                e
                async for e in executor.run_turn(
                    [
                        {"role": "user", "content": "search for test"},
                        {
                            "role": "tool_call",
                            "content": {"tool": "search", "args": {"q": "test"}},
                        },
                        {"role": "tool_result", "content": {"results": ["a", "b", "c"]}},
                    ],
                    [{"name": "search", "description": "Search"}],
                    "sys",
                )
            ]
            turn_events = [e for e in events2 if isinstance(e, TurnComplete)]
            self.assertEqual(len(turn_events), 1)
            self.assertEqual(turn_events[0].response, "Found 3 results.")

        _run(_t())

    def test_interrupt_session_closes_active_stream(self):
        class ClosableStream:
            def __init__(self) -> None:
                self.closed = False

            def close(self) -> None:
                self.closed = True

        async def _t():
            executor = DatabricksExecutor(client=FakeClient([]))
            state = executor._get_or_create_session_state("s1")
            state.active_stream = ClosableStream()

            interrupted = await executor.interrupt_session("s1")

            self.assertTrue(interrupted)
            self.assertTrue(state.interrupt_requested)
            self.assertTrue(state.active_stream.closed)

        _run(_t())


# ---------------------------------------------------------------------------
# Credential resolution (_read_databrickscfg) — function-based pytest tests.
#
# These exercise the OAuth bug fix: _read_databrickscfg now delegates to the
# databricks-sdk so OAuth profiles (auth_type: databricks-cli) return a fresh
# minted bearer instead of the stale ``token`` field in ~/.databrickscfg.
# ---------------------------------------------------------------------------

# Imports below are intentionally at the bottom because the pytest-based section
# is appended to the pre-existing unittest module above.
import textwrap  # noqa: E402
from pathlib import Path as _Path  # noqa: E402

import pytest  # noqa: E402

from omnigent.inner.databricks_executor import (  # noqa: E402
    DatabricksAuthError,
    _DatabricksBearerAuth,
    _read_databrickscfg,
    _read_databrickscfg_file_fallback,
    _read_databrickscfg_host,
)

_AUTH_ENV_VARS: tuple[str, ...] = (
    "DATABRICKS_HOST",
    "DATABRICKS_TOKEN",
    "DATABRICKS_CONFIG_PROFILE",
    "DATABRICKS_CONFIG_FILE",
    "DATABRICKS_CLIENT_ID",
    "DATABRICKS_CLIENT_SECRET",
    "DATABRICKS_AUTH_TYPE",
)


@pytest.fixture
def clean_databricks_env(monkeypatch: pytest.MonkeyPatch) -> None:
    """
    Clear every DATABRICKS_* env var that affects credential resolution.

    The agent harness that runs these tests may have e.g. ``DATABRICKS_TOKEN``
    exported for the coding agent itself — which would override profile-based
    resolution in the SDK and confuse these tests.
    """
    for var in _AUTH_ENV_VARS:
        monkeypatch.delenv(var, raising=False)


def _raise_offline_host_metadata(host: str) -> None:
    """Stand-in for ``databricks.sdk.config.get_host_metadata`` that fails fast.

    :param host: Workspace URL the SDK would have probed,
        e.g. ``"https://example.cloud.databricks.com"``.
    :raises ConnectionError: Always — simulates an unreachable host without
        the SDK's network retry loop.
    """
    raise ConnectionError(f"offline test stub: refusing to probe {host}")


@pytest.fixture
def pat_only_cfg(
    tmp_path: _Path, monkeypatch: pytest.MonkeyPatch, clean_databricks_env: None
) -> _Path:
    """
    Materialize a temp ``.databrickscfg`` containing a single PAT profile
    and point the SDK at it via ``DATABRICKS_CONFIG_FILE``.

    A PAT profile makes the SDK's ``authenticate()`` return the token
    verbatim without any OAuth exchange. The host is a placeholder, so the
    SDK's ``Config.__init__`` host-metadata probe (a real HTTP GET against
    ``/.well-known/databricks-config``, with retries) is stubbed to fail
    fast — ``_resolve_host_metadata`` logs and falls back to the explicit
    config, which is exactly the offline behavior these tests need.
    """
    monkeypatch.setattr(
        "databricks.sdk.config.get_host_metadata",
        _raise_offline_host_metadata,
    )
    contents = textwrap.dedent(
        """
        [pat-profile]
        host = https://example.cloud.databricks.com
        token = dapi-fake-pat-token-for-unit-test
        """
    ).lstrip()
    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(contents)
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))
    return cfg_path


def test_read_databrickscfg_pat_profile_returns_token_verbatim(
    pat_only_cfg: _Path,
) -> None:
    """
    For a plain ``auth_type=pat`` profile, the SDK should return the PAT
    from the file verbatim (no OAuth exchange, no token rewriting). This
    confirms we haven't regressed the common PAT-user path.
    """
    creds = _read_databrickscfg("pat-profile")

    assert creds is not None
    assert creds.host == "https://example.cloud.databricks.com"
    assert creds.token == "dapi-fake-pat-token-for-unit-test"


def test_read_databrickscfg_missing_profile_falls_back_to_file_reader(
    pat_only_cfg: _Path,
) -> None:
    """
    Requesting a profile that doesn't exist makes Config raise ValueError;
    the wrapper catches that and falls through to the legacy file reader.

    The legacy reader's documented resolution order is: explicit profile
    -> DATABRICKS_CONFIG_PROFILE env -> DEFAULT section -> first section
    with both host+token. So on a config file that contains only
    ``pat-profile``, an unknown requested profile falls through to
    "first section" — i.e. we recover PAT credentials.
    """
    creds = _read_databrickscfg("no-such-profile-xyz")

    assert creds is not None
    assert creds.host == "https://example.cloud.databricks.com"
    assert creds.token == "dapi-fake-pat-token-for-unit-test"


def test_read_databrickscfg_empty_config_file_returns_none(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    With a config file that exists but contains no valid profiles, both
    the SDK path and the file fallback should resolve to ``None``.
    """
    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text("# empty\n")
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    assert _read_databrickscfg("missing-profile") is None


def test_read_databrickscfg_host_reads_oauth_profile_without_token(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    Host-only resolution supports Databricks CLI OAuth profile sections.

    Native Codex does not need a static token at startup: it only needs the
    workspace host to build the Codex provider base URL, then Codex's
    ``auth.command`` calls ``databricks auth token --profile`` for live bearer
    refresh. A default install without ``databricks-sdk`` must therefore still
    accept a present ``auth_type=databricks-cli`` section with no ``token``.
    """
    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        textwrap.dedent(
            """
        [oauth-profile]
        host = https://oauth-host.example.com
        auth_type = databricks-cli
        """
        ).lstrip()
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    assert _read_databrickscfg_host("oauth-profile") == "https://oauth-host.example.com"


def test_read_databrickscfg_host_missing_named_profile_does_not_fallback(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """An explicit missing profile must not borrow a different profile's host."""
    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        textwrap.dedent(
            """
        [other-profile]
        host = https://other.example.com
        auth_type = databricks-cli
        """
        ).lstrip()
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    assert _read_databrickscfg_host("missing-profile") is None


def test_codex_executor_gateway_uses_host_only_oauth_profile(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    Wrapped Codex shares the native Codex host-only Databricks profile path.

    This covers default installs where the runner can read a
    ``databricks-cli`` profile's host but cannot mint a bearer snapshot at
    construction time.
    """
    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        textwrap.dedent(
            """
        [oauth-profile]
        host = https://oauth-host.example.com
        auth_type = databricks-cli
        """
        ).lstrip()
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))
    monkeypatch.setattr(
        "omnigent.inner.codex_executor._read_databrickscfg",
        lambda _profile: None,
    )

    from omnigent.inner.codex_executor import CodexExecutor

    executor = CodexExecutor(
        codex_path=sys.executable,
        gateway=True,
        databricks_profile="oauth-profile",
    )

    overrides = "\n".join(executor._codex_config_overrides)
    assert "https://oauth-host.example.com/ai-gateway/codex/v1" in overrides
    assert 'databricks auth token --profile \\"oauth-profile\\"' in overrides


def test_read_databrickscfg_no_config_file_returns_none(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    When ``~/.databrickscfg`` is absent and no env auth is set, both the
    SDK path and the file fallback should return ``None``.
    """
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(tmp_path / "does-not-exist"))

    assert _read_databrickscfg("DEFAULT") is None


def test_file_fallback_reads_token_field_directly(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    The legacy fallback intentionally reads the ``token`` field as-is
    (that is the pre-fix behavior we preserve for exotic setups where
    the SDK's Config init raises). Confirm the fallback still works.
    """
    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        textwrap.dedent(
            """
        [p]
        host = https://legacy-host.example.com
        token = legacy-pat-value
        """
        ).lstrip()
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    creds = _read_databrickscfg_file_fallback("p")

    assert creds is not None
    assert creds.host == "https://legacy-host.example.com"
    assert creds.token == "legacy-pat-value"


def test_read_databrickscfg_falls_back_when_sdk_raises(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    Verify the SDK-failure path: if ``databricks.sdk.config.Config`` raises
    ``ValueError`` during construction, the wrapper silently falls through
    to the file reader so plain PAT setups still work.

    We simulate this by stubbing the SDK's ``Config`` with a tiny real
    class that unconditionally raises on init — no MagicMock.
    """

    class _AlwaysFailsConfig:
        """
        Test double for ``databricks.sdk.config.Config`` that always
        raises ValueError — emulates an exotic setup the SDK can't parse.
        """

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            raise ValueError("simulated SDK resolution failure")

    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        textwrap.dedent(
            """
        [fallback-profile]
        host = https://fallback-host.example.com
        token = fallback-pat-value
        """
        ).lstrip()
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    monkeypatch.setattr(_sdk_config_mod, "Config", _AlwaysFailsConfig)

    creds = _read_databrickscfg("fallback-profile")

    assert creds is not None
    assert creds.host == "https://fallback-host.example.com"
    assert creds.token == "fallback-pat-value"


def test_read_databrickscfg_missing_profile_uses_ambient_credentials(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    When a named profile is given but ``Config(profile=...)`` fails, the
    wrapper tries the ambient credential chain (``Config()`` with no args)
    before falling back to the file reader.

    This covers the Databricks App deployment case: the spec was authored
    locally with ``profile: my-profile`` but runs on an App container that
    has no ``~/.databrickscfg`` — the App container supplies credentials
    via environment variables instead.
    """

    class _AmbientOnlyConfig:
        """
        Test double for ``databricks.sdk.config.Config``.

        Raises ``ValueError`` when called with a profile (simulating a
        missing named profile), but succeeds when called with no args
        (simulating ambient env-var credentials).
        """

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            if kwargs.get("profile") is not None:
                raise ValueError("simulated missing profile")
            # Ambient path — expose host and a Bearer token.
            self.host = "https://ambient.example.com"

        def authenticate(self) -> dict[str, str]:
            return {"Authorization": "Bearer ambient-bearer-token"}

    monkeypatch.setattr(_sdk_config_mod, "Config", _AmbientOnlyConfig)

    creds = _read_databrickscfg("no-such-profile-xyz")

    # Ambient credentials resolved — the App-server fallback path works.
    # If this is None, the ambient Config() branch was not reached.
    assert creds is not None
    # Host must come from the ambient Config, not the file reader.
    assert creds.host == "https://ambient.example.com", (
        "Expected host from ambient Config(); got a different value, which "
        "suggests the file-reader fallback ran instead of the ambient path."
    )
    # Token must be the bearer stripped of the 'Bearer ' prefix.
    assert creds.token == "ambient-bearer-token", (
        "Expected token from ambient Config.authenticate(); a different value "
        "means the ambient credentials were not used."
    )


def test_read_databrickscfg_missing_profile_ambient_also_fails_uses_file_fallback(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    When both ``Config(profile=...)`` and ``Config()`` (ambient) raise
    ``ValueError``, the wrapper falls through to the legacy file reader.

    This is the two-step failure path: named profile absent AND no ambient
    credentials in the environment — the file reader is the last resort.
    """

    class _ProfileFailsThenAmbientFails:
        """
        Test double for ``databricks.sdk.config.Config`` that raises
        ``ValueError`` unconditionally — emulates a machine with neither
        a matching named profile nor ambient env-var credentials.
        """

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            raise ValueError("simulated: no profile, no ambient credentials")

    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        textwrap.dedent(
            """
        [some-profile]
        host = https://file-fallback.example.com
        token = file-fallback-token
        """
        ).lstrip()
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    monkeypatch.setattr(_sdk_config_mod, "Config", _ProfileFailsThenAmbientFails)

    creds = _read_databrickscfg("no-such-profile-xyz")

    # File reader ran as the last resort after both SDK paths failed.
    # If None, the file reader itself failed — check DATABRICKS_CONFIG_FILE.
    assert creds is not None
    # Host must come from the file reader, not any SDK path.
    assert creds.host == "https://file-fallback.example.com", (
        "Expected host from file-reader fallback; a different value means "
        "the SDK path resolved unexpectedly."
    )
    assert creds.token == "file-fallback-token", (
        "Expected token from file-reader fallback; a different value means "
        "the SDK path resolved unexpectedly."
    )


def test_read_databrickscfg_missing_profile_service_principal_via_ambient(
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    Service principal (M2M) credentials supplied via environment variables
    (``DATABRICKS_HOST`` + ``DATABRICKS_CLIENT_ID`` + ``DATABRICKS_CLIENT_SECRET``)
    are resolved through the ambient ``Config()`` fallback when the named
    profile is absent.

    The Databricks SDK's ``Config()`` natively handles client-credential env
    vars, so no explicit SP handling is needed in ``_read_databrickscfg``.
    This test confirms the ambient path reaches a ``Config()`` call with no
    profile — the same call the SDK uses to pick up those env vars.
    """

    class _SPAmbientConfig:
        """
        Test double for ``databricks.sdk.config.Config``.

        Raises ``ValueError`` when called with a named profile (simulating a
        missing profile), succeeds when called with no args — as the SDK would
        do when ``DATABRICKS_HOST`` / ``DATABRICKS_CLIENT_ID`` /
        ``DATABRICKS_CLIENT_SECRET`` env vars are set.
        """

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            if kwargs.get("profile") is not None:
                raise ValueError("simulated: named profile not found")
            # Ambient path — service principal resolved host + M2M token.
            self.host = "https://sp-workspace.example.com"

        def authenticate(self) -> dict[str, str]:
            return {"Authorization": "Bearer sp-m2m-access-token"}

    monkeypatch.setattr(_sdk_config_mod, "Config", _SPAmbientConfig)

    creds = _read_databrickscfg("missing-named-profile")

    assert creds is not None, (
        "Expected service-principal credentials via ambient Config(); got None. "
        "The ambient fallback for a missing named profile did not run."
    )
    assert creds.host == "https://sp-workspace.example.com", (
        f"Expected SP workspace host; got {creds.host!r}."
    )
    assert creds.token == "sp-m2m-access-token", f"Expected M2M bearer token; got {creds.token!r}."


def test_read_databrickscfg_oauth_profile_returns_fresh_bearer(
    tmp_path: _Path,
    monkeypatch: pytest.MonkeyPatch,
    clean_databricks_env: None,
) -> None:
    """
    The core regression guard. For an OAuth profile
    (``auth_type: databricks-cli``), the returned token MUST be a freshly
    minted JWT, NOT the stale ``token`` field sitting in the config file.

    The test stubs the SDK's ``Config`` with a small real class so the
    assertion is deterministic and never depends on the developer or CI
    machine's real ``~/.databrickscfg``.
    """
    stale_field = "dapi-stale-config-token-fixture"
    fresh_bearer = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJvYXV0aCJ9.signature"
    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        textwrap.dedent(
            f"""
            [oauth-profile]
            host = https://oauth-workspace.example.com
            auth_type = databricks-cli
            token = {stale_field}
            """
        ).lstrip()
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    class _OAuthConfig:
        """
        Test double for ``databricks.sdk.config.Config``.

        It emulates SDK OAuth resolution: the config file may contain a stale
        PAT-shaped token field, but ``authenticate()`` returns the fresh
        bearer that ``_read_databrickscfg`` must use.
        """

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            assert kwargs.get("profile") == "oauth-profile"
            self.host = "https://oauth-workspace.example.com"

        def authenticate(self) -> dict[str, str]:
            return {"Authorization": f"Bearer {fresh_bearer}"}

    monkeypatch.setattr(_sdk_config_mod, "Config", _OAuthConfig)

    creds = _read_databrickscfg("oauth-profile")

    assert creds is not None, "OAuth profile should resolve credentials"
    assert creds.host == "https://oauth-workspace.example.com"
    # OAuth access tokens are JWTs (header.payload.signature) — far longer
    # than a 36-char PAT and they start with "eyJ" (base64 of '{"alg":...').
    assert creds.token == fresh_bearer
    assert creds.token != stale_field, (
        "OAuth profile returned the stale PAT from the file — the fix is "
        "not actually exchanging credentials via the SDK."
    )


def test_resolve_databricks_auth_returns_bearer_auth_and_host(
    monkeypatch,
):
    """``_resolve_databricks_auth`` returns an httpx Auth + host.

    Verifies that the helper creates a ``_DatabricksBearerAuth``
    instance backed by the SDK Config and returns the workspace
    host URL.

    :param monkeypatch: Pytest monkeypatch fixture.
    """
    from omnigent.inner.databricks_executor import (
        _DatabricksBearerAuth,
        _resolve_databricks_auth,
    )

    class _FakeConfig:
        host = "https://example.cloud.databricks.com"
        profile = "dev"

        def authenticate(self):
            return {"Authorization": "Bearer fresh-tok"}

    monkeypatch.setattr(_sdk_config_mod, "Config", lambda **_kw: _FakeConfig())

    auth, host = _resolve_databricks_auth("dev")
    assert isinstance(auth, _DatabricksBearerAuth)
    assert host == "https://example.cloud.databricks.com"


def test_resolve_databricks_auth_invalid_profile_raises_clear_error(
    monkeypatch,
):
    """``_resolve_databricks_auth`` raises ``DatabricksAuthError``
    with an actionable message when the profile is not authenticated.

    :param monkeypatch: Pytest monkeypatch fixture.
    """
    import pytest

    import omnigent.inner.databricks_executor as db_exec
    from omnigent.inner.databricks_executor import (
        DatabricksAuthError,
        _resolve_databricks_auth,
    )

    def _failing_config(**_kw):
        raise ValueError("no credentials")

    monkeypatch.setattr(_sdk_config_mod, "Config", _failing_config)
    monkeypatch.setattr(db_exec, "_read_databrickscfg", lambda _p: None)
    monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)

    with pytest.raises(DatabricksAuthError, match="databricks auth login -p dogfood"):
        _resolve_databricks_auth("dogfood")


def test_resolve_databricks_auth_env_profile_falls_back_to_ambient_with_warning(
    monkeypatch,
    caplog,
):
    """``_resolve_databricks_auth`` falls back to ambient when the profile
    came from DATABRICKS_CONFIG_PROFILE (not an explicit argument).

    When ``profile=None`` and ``DATABRICKS_CONFIG_PROFILE`` names a profile
    that cannot be resolved, the function falls back to the ambient credential
    chain and logs a warning.  This covers CI environments that provide OIDC
    tokens via env vars but don't have a matching profile in
    ``~/.databrickscfg``.

    The fallback must NOT fire when the profile was passed explicitly — that
    case must fail loud (see
    ``test_resolve_databricks_auth_invalid_profile_raises_clear_error``).

    :param monkeypatch: Pytest monkeypatch fixture.
    :param caplog: Pytest log-capture fixture.
    """
    import logging

    import omnigent.inner.databricks_executor as db_exec
    from omnigent.inner.databricks_executor import (
        _DatabricksBearerAuth,
        _resolve_databricks_auth,
    )

    def _config_factory(**kw):
        if kw.get("profile") is not None:
            raise ValueError("simulated: profile not found in ~/.databrickscfg")

        class _AmbientConfig:
            host = "https://example.cloud.databricks.com"

            def authenticate(self):
                return {"Authorization": "Bearer oidc-tok"}

        return _AmbientConfig()

    monkeypatch.setattr(_sdk_config_mod, "Config", _config_factory)
    monkeypatch.setattr(db_exec, "_read_databrickscfg", lambda _p: None)
    # Profile comes from env var, not an explicit argument.
    monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", "missing-profile")

    with caplog.at_level(logging.WARNING, logger="omnigent.inner.databricks_executor"):
        auth, host = _resolve_databricks_auth()  # profile=None — uses env var

    assert isinstance(auth, _DatabricksBearerAuth), (
        "Expected a _DatabricksBearerAuth from the ambient Config() fallback; "
        "the ambient path was not reached."
    )
    assert host == "https://example.cloud.databricks.com", (
        f"Expected ambient workspace host, got {host!r}"
    )
    # The fallback must be logged — it must not be silent.
    warning_messages = [r.message for r in caplog.records if r.levelno == logging.WARNING]
    assert any("missing-profile" in m and "ambient" in m for m in warning_messages), (
        f"Expected a warning mentioning the profile name and 'ambient', got: {warning_messages!r}"
    )


def test_resolve_databricks_auth_explicit_profile_not_found_raises(
    monkeypatch,
):
    """Explicit profile failures raise ``DatabricksAuthError``, no ambient fallback.

    When the user explicitly passes ``--profile dev`` and the profile cannot
    be resolved, ``_resolve_databricks_auth`` must raise ``DatabricksAuthError``
    immediately — NOT silently fall back to a different workspace.  This
    upholds the "Fail loud" design principle.

    :param monkeypatch: Pytest monkeypatch fixture.
    """
    import pytest

    import omnigent.inner.databricks_executor as db_exec
    from omnigent.inner.databricks_executor import (
        DatabricksAuthError,
        _resolve_databricks_auth,
    )

    call_log: list[dict] = []

    def _config_factory(**kw):
        call_log.append(kw)
        raise ValueError("simulated: profile not found in ~/.databrickscfg")

    monkeypatch.setattr(_sdk_config_mod, "Config", _config_factory)
    monkeypatch.setattr(db_exec, "_read_databrickscfg", lambda _p: None)
    monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)

    with pytest.raises(DatabricksAuthError, match="databricks auth login -p dev"):
        _resolve_databricks_auth("dev")

    # Ambient Config() must NOT be called — no fallback for explicit profiles.
    # If Config() were called with no args, call_log would have a second entry.
    assert len(call_log) == 1, (
        f"Expected exactly one Config() call (the explicit profile), "
        f"got {len(call_log)}: {call_log!r}. "
        f"The ambient fallback fired for an explicit profile — violates Fail Loud."
    )


def test_bearer_auth_injects_fresh_token_per_request():
    """``_DatabricksBearerAuth`` calls ``Config.authenticate()`` per request.

    Verifies that a fresh token is injected into each HTTP request,
    and that different calls can return different tokens (simulating
    OAuth refresh).
    """
    import httpx

    from omnigent.inner.databricks_executor import _DatabricksBearerAuth

    call_count = 0

    class _RotatingConfig:
        def authenticate(self):
            nonlocal call_count
            call_count += 1
            return {"Authorization": f"Bearer tok-{call_count}"}

    auth = _DatabricksBearerAuth(_RotatingConfig(), profile_name="dev")

    request1 = httpx.Request("GET", "https://example.com/v1/chat")
    gen1 = auth.auth_flow(request1)
    next(gen1)
    assert request1.headers["Authorization"] == "Bearer tok-1"

    request2 = httpx.Request("GET", "https://example.com/v1/chat")
    gen2 = auth.auth_flow(request2)
    next(gen2)
    assert request2.headers["Authorization"] == "Bearer tok-2"

    assert call_count == 2


def test_bearer_auth_raises_on_expired_refresh_token():
    """``_DatabricksBearerAuth`` wraps SDK failures as ``DatabricksAuthError``.

    When the underlying ``Config.authenticate()`` fails (e.g. the OAuth
    refresh token expired), the auth flow raises ``DatabricksAuthError``
    with an actionable error message.
    """
    import httpx
    import pytest

    from omnigent.inner.databricks_executor import (
        DatabricksAuthError,
        _DatabricksBearerAuth,
    )

    class _DeadConfig:
        def authenticate(self):
            raise ValueError("token expired")

    auth = _DatabricksBearerAuth(_DeadConfig(), profile_name="dogfood")
    request = httpx.Request("GET", "https://example.com/v1/chat")

    with pytest.raises(DatabricksAuthError, match="databricks auth login -p dogfood"):
        gen = auth.auth_flow(request)
        next(gen)


def test_claude_sdk_executor_reuses_api_key_helper_between_turns(
    monkeypatch,
):
    """ClaudeSDKExecutor leaves refresh cadence to Claude Code.

    The Claude CLI subprocess should not receive a static
    ``ANTHROPIC_AUTH_TOKEN``. It receives an apiKeyHelper command
    that Claude Code can re-run while a long turn is in flight.

    :param monkeypatch: Pytest monkeypatch fixture.
    :returns: None.
    """

    call_count = 0

    def _counting_resolve(
        profile=None,
        *,
        host_override=None,
        base_url_override=None,
        auth_command_override=None,
        auth_refresh_interval_ms=None,
    ):
        """Return a distinct helper command if a turn tries to refresh.

        :param profile: Databricks profile name.
        :param host_override: Optional ucode workspace host.
        :param base_url_override: Optional ucode gateway base URL.
        :param auth_command_override: Optional ucode auth command.
        :param auth_refresh_interval_ms: Optional ucode refresh cadence.
        :returns: Fake Claude Databricks env payload.
        """
        nonlocal call_count
        call_count += 1
        return {
            "ANTHROPIC_BASE_URL": "https://host/ai-gateway/anthropic",
            "OMNIGENT_CLAUDE_API_KEY_HELPER": f"helper-{call_count}",
            "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": "900000",
            "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
        }

    monkeypatch.setattr(
        "omnigent.inner.claude_sdk_executor._resolve_gateway_env",
        _counting_resolve,
    )

    from omnigent.inner.claude_sdk_executor import ClaudeSDKExecutor

    # Construct with initial env.
    executor = ClaudeSDKExecutor.__new__(ClaudeSDKExecutor)
    executor._gateway = True
    executor._databricks_profile = "oss"
    executor._gateway_host = None
    executor._base_url_override = None
    executor._extra_env = {
        "OMNIGENT_CLAUDE_API_KEY_HELPER": "helper-initial",
        "ANTHROPIC_BASE_URL": "https://host/ai-gateway/anthropic",
        "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": "900000",
        "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
    }

    async def _trigger_refresh():
        # run_turn refreshes at the top, then fails (no real SDK).
        try:
            async for _ in executor.run_turn(
                [{"role": "user", "content": "hi", "session_id": "s1"}],
                [],
                "",
            ):
                pass
        except Exception:
            pass
        return executor._extra_env.get("OMNIGENT_CLAUDE_API_KEY_HELPER")

    helper = _run(_trigger_refresh())

    assert helper == "helper-initial"
    assert "ANTHROPIC_AUTH_TOKEN" not in executor._extra_env
    assert call_count == 0


def test_codex_executor_uses_cli_auth_command_not_env_token(
    monkeypatch,
):
    """CodexExecutor leaves token refresh to Codex's ``auth.command``.

    :param monkeypatch: Pytest monkeypatch fixture.
    :returns: None.
    """
    call_count = 0

    def _counting_read(profile=None):
        nonlocal call_count
        call_count += 1
        return type(
            "Creds",
            (),
            {
                "host": "https://host",
                "token": f"tok-{call_count}",
            },
        )()

    monkeypatch.setattr(
        "omnigent.inner.codex_executor._read_databrickscfg",
        _counting_read,
    )

    from omnigent.inner.codex_executor import CodexExecutor

    executor = CodexExecutor.__new__(CodexExecutor)
    executor._databricks = True
    executor._databricks_profile = "oss"
    executor._env = {}

    async def _trigger_refresh():
        try:
            async for _ in executor.run_turn(
                [{"role": "user", "content": "hi", "session_id": "s1"}],
                [],
                "",
            ):
                pass
        except Exception:
            pass
        return executor._env.get("DATABRICKS_TOKEN")

    token = _run(_trigger_refresh())

    assert token is None
    assert call_count == 0


def test_bearer_auth_current_token_reuses_config_and_strips_prefix() -> None:
    """
    ``_DatabricksBearerAuth.current_token()`` returns the bare bearer (no
    ``"Bearer "`` prefix) from the WRAPPED config's ``authenticate()`` and
    reuses that one config across calls.

    Reuse is the whole point of the per-request auth-tax fix: the SDK serves
    repeat calls from its in-memory token cache instead of re-shelling to the
    Databricks CLI. If ``current_token`` rebuilt a Config per call or failed
    to strip the prefix, this test fails.
    """

    class _CountingConfig:
        """Config double whose authenticate() counts calls."""

        def __init__(self) -> None:
            self.authenticate_calls = 0

        def authenticate(self) -> dict[str, str]:
            self.authenticate_calls += 1
            return {"Authorization": "Bearer dapi-XYZ"}

    cfg = _CountingConfig()
    auth = _DatabricksBearerAuth(cfg, profile_name="oss")

    tokens = [auth.current_token() for _ in range(4)]

    # Bare token (prefix stripped) on every call. A failure here means
    # current_token returned the raw header or the wrong field.
    assert tokens == ["dapi-XYZ"] * 4
    # All 4 calls went through the SAME wrapped config — proving reuse, so
    # the SDK's own cache (not a rebuilt Config) backs repeat calls. If the
    # method rebuilt Config, it would not be this single object's counter.
    assert cfg.authenticate_calls == 4


@pytest.mark.parametrize(
    "headers",
    [
        {"Authorization": "Basic dXNlcjpwYXNz"},  # non-Bearer scheme
        {},  # no Authorization header at all
    ],
)
def test_bearer_auth_current_token_none_for_non_bearer(headers: dict[str, str]) -> None:
    """
    ``current_token()`` returns ``None`` for a non-Bearer or empty
    ``Authorization`` header (the system only supports Bearer). Returning the
    raw header would feed callers a credential the Omnigent server can't use.
    """

    class _Config:
        """Config double returning a fixed (non-Bearer) header set."""

        def authenticate(self) -> dict[str, str]:
            return headers

    assert _DatabricksBearerAuth(_Config()).current_token() is None


def test_bearer_auth_current_token_wraps_failure_as_auth_error() -> None:
    """
    ``current_token()`` raises :class:`DatabricksAuthError` when the SDK's
    ``authenticate()`` raises — the same fail-loud contract as ``auth_flow``,
    so token-mint failures surface with a re-login hint instead of silently
    yielding a bad/empty credential.
    """

    class _FailingConfig:
        """Config double whose authenticate() always raises."""

        def authenticate(self) -> dict[str, str]:
            raise RuntimeError("token endpoint unreachable")

    with pytest.raises(DatabricksAuthError):
        _DatabricksBearerAuth(_FailingConfig(), profile_name="oss").current_token()


if __name__ == "__main__":
    unittest.main()


# ── host-keyed resolution (omnigent login pointer records) ─────────


class _StubSdkConfig:
    """Minimal stand-in for ``databricks.sdk.config.Config``.

    Real ``Config`` probes host metadata at construction, which fails
    offline for placeholder hosts — the production seam
    (``_sdk_config``) exists precisely so tests can substitute this.

    :param host: The workspace host the config resolves to.
    :param token: The bearer minted by :meth:`authenticate`.
    """

    def __init__(self, host: str, token: str) -> None:
        self.host = host
        self._token = token

    def authenticate(self) -> dict[str, str]:
        """Return ``Authorization`` headers like the real Config."""
        return {"Authorization": f"Bearer {self._token}"}


def test_resolve_auth_for_host_prefers_matching_profile(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Host resolution authenticates via the cfg profile pinned to the host.

    ``databricks auth login --host <ws>`` saves a named profile, and the
    CLI's host-keyed token lookup is unreliable across versions (it can
    miss profile-keyed grants, or return a grant for a different
    workspace). If this regresses to the bare ``databricks-cli`` host
    lookup, the constructed kwargs below change and the test fails.
    """
    from omnigent.inner import databricks_executor

    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        "[my-ws]\nhost = https://example.databricks.com\ntoken = dapi-fake-token\n"
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    constructed: list[dict[str, str]] = []

    def _fake_sdk_config(**kwargs: str) -> _StubSdkConfig:
        constructed.append(kwargs)
        assert kwargs == {"profile": "my-ws"}, (
            f"expected the host-matched profile to be tried first, got {kwargs}. "
            "A host/auth_type construction here means the unreliable "
            "databricks-cli host lookup regained priority."
        )
        return _StubSdkConfig(host="https://example.databricks.com", token="dapi-fake-token")

    monkeypatch.setattr(databricks_executor, "_sdk_config", _fake_sdk_config)

    auth, host = databricks_executor._resolve_databricks_auth(
        host="https://example.databricks.com"
    )

    # The stub's PAT coming back proves the profile-path config is the
    # one wired into the returned auth (not a second construction).
    assert auth.current_token() == "dapi-fake-token"
    assert host == "https://example.databricks.com"
    # Exactly one construction: the matched profile authenticated, so the
    # CLI host fallback must not run at all.
    assert constructed == [{"profile": "my-ws"}]


def test_resolve_auth_for_host_falls_back_to_cli_when_no_profile_matches(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """With no cfg profile pinned to the host, the CLI host lookup runs.

    Cfg-less machines (fresh laptop, CI) have only the host-keyed OAuth
    cache from ``databricks auth login --host`` — dropping this fallback
    would strand them.
    """
    from omnigent.inner import databricks_executor

    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(tmp_path / "absent"))

    def _fake_sdk_config(**kwargs: str) -> _StubSdkConfig:
        # The only construction allowed is the host-keyed CLI lookup —
        # a profile= construction here means a phantom profile matched.
        assert kwargs == {
            "host": "https://example.databricks.com",
            "auth_type": "databricks-cli",
        }, f"unexpected Config construction: {kwargs}"
        return _StubSdkConfig(host="https://example.databricks.com", token="tok-cli")

    monkeypatch.setattr(databricks_executor, "_sdk_config", _fake_sdk_config)

    auth, host = databricks_executor._resolve_databricks_auth(
        host="https://example.databricks.com"
    )

    assert auth.current_token() == "tok-cli"
    assert host == "https://example.databricks.com"


def test_profiles_for_host_normalizes_scheme_and_slash(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Host matching ignores scheme and trailing slashes, in file order.

    ``databrickscfg`` hosts appear both with and without ``https://`` and
    trailing ``/`` in the wild; a strict string compare would silently
    miss the profile and fall through to the unreliable CLI host lookup.
    """
    from omnigent.inner.databricks_executor import _databrickscfg_profiles_for_host

    cfg_path = tmp_path / "databrickscfg"
    cfg_path.write_text(
        "[bare]\nhost = example.databricks.com\ntoken = t1\n"
        "[slashed]\nhost = https://example.databricks.com/\ntoken = t2\n"
        "[other]\nhost = https://other.databricks.com\ntoken = t3\n"
    )
    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg_path))

    matches = _databrickscfg_profiles_for_host("https://example.databricks.com")

    # Both spellings of the host match; the unrelated workspace does not.
    assert matches == ["bare", "slashed"]


def test_profiles_for_host_missing_file_returns_empty(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """No config file → no profile candidates (CLI fallback territory)."""
    from omnigent.inner.databricks_executor import _databrickscfg_profiles_for_host

    monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(tmp_path / "absent"))

    assert _databrickscfg_profiles_for_host("https://example.databricks.com") == []