"""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 (no finish_reason, no content, no tool calls) is a truncated turn that died mid-stream, so it surfaces an ExecutorError rather than a silent empty TurnComplete (#1118).""" 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], ExecutorError) self.assertEqual(events[0].message, "Stream ended without finish_reason") _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 `` 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") == [] def test_stream_ended_without_finish_reason_with_content_completes() -> None: """A truncated stream that still produced text surfaces that text as a TurnComplete (not an error) — only the empty case is fatal (#1118).""" async def _t() -> None: # Content arrives, then the stream ends without a finish_reason. chunks = [FakeStreamChunk(choices=[FakeStreamChoice(delta=FakeDelta(content="partial"))])] 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"), ) ] assert not [e for e in events if isinstance(e, ExecutorError)] turn_events = [e for e in events if isinstance(e, TurnComplete)] assert len(turn_events) == 1 assert turn_events[0].response == "partial" _run(_t())