Files
omnigent-ai--omnigent/tests/inner/test_openai_agents_sdk_executor.py
2026-07-13 13:12:00 +08:00

2998 lines
111 KiB
Python

"""Tests for OpenAIAgentsSDKExecutor with a fake Agents SDK module."""
import asyncio
import contextlib
import sys
import types
import unittest
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from unittest.mock import patch
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import base64
import databricks.sdk.config as _sdk_config_mod
from omnigent.inner.executor import (
ExecutorConfig,
ExecutorError,
ReasoningChunk,
TextChunk,
ToolCallComplete,
ToolCallRequest,
ToolCallStatus,
TurnComplete,
)
from omnigent.inner.openai_agents_sdk_executor import (
OpenAIAgentsSDKExecutor,
_normalize_content_blocks_for_chat,
_normalize_responses_items_for_chat,
_ReasoningBlockFilterStream,
_sanitize_replay_item,
_wrap_client_for_reasoning_models,
)
from omnigent.llms.errors import is_context_length_exceeded as _is_context_length_exceeded
def _run(coro):
loop = asyncio.new_event_loop()
try:
return loop.run_until_complete(coro)
finally:
loop.close()
async def _collect(agen):
"""Drain an async generator into a list."""
return [item async for item in agen]
@dataclass
class _FakeToolCallRawItem:
name: str
arguments: str
call_id: str = "call_1"
@dataclass
class _FakeToolOutputRawItem:
call_id: str = "call_1"
@dataclass
class _FakeToolCallItem:
raw_item: object
type: str = "tool_call_item"
@dataclass
class _FakeToolOutputItem:
raw_item: object
output: object
type: str = "tool_call_output_item"
@dataclass
class _FakeRawTextDelta:
delta: str
type: str = "response.output_text.delta"
@dataclass
class _FakeRawReasoningDelta:
delta: str
type: str = "response.reasoning_summary_text.delta"
@dataclass
class _FakeRawEvent:
data: object
type: str = "raw_response_event"
@dataclass
class _FakeRunItemEvent:
item: object
type: str = "run_item_stream_event"
@dataclass
class _FakeModelSettings:
parallel_tool_calls: bool | None = None
max_tokens: int | None = None
@dataclass
class _FakePromptTokensDetails:
"""
Minimal stand-in for OpenAI's ``prompt_tokens_details`` object.
:param cached_tokens: Number of tokens served from the prompt cache.
"""
cached_tokens: int = 0
@dataclass
class _FakeUsage:
"""
Minimal stand-in for the openai-agents SDK ModelResponse.usage object.
:param input_tokens: Number of input tokens for this LLM call.
:param output_tokens: Number of output tokens for this LLM call.
:param total_tokens: Sum of input and output tokens for this call.
``0`` means the SDK did not report a total; the executor falls
back to ``input_tokens + output_tokens``.
:param prompt_tokens_details: Optional breakdown of prompt tokens,
including ``cached_tokens``. ``None`` means the SDK did not
report cache details.
"""
input_tokens: int = 0
output_tokens: int = 0
total_tokens: int = 0
prompt_tokens_details: _FakePromptTokensDetails | None = None
@dataclass
class _FakeRawResponse:
"""
Minimal stand-in for the openai-agents SDK ModelResponse object.
:param usage: Token usage for this LLM call.
"""
usage: _FakeUsage
class _FakeResult:
def __init__(
self,
events,
final_output="",
new_items=None,
exception=None,
state=None,
raw_responses=None,
):
self._events = events
self.final_output = final_output
self.new_items = new_items or []
self._exception = exception
self._state = state
self.cancel_calls = []
self.raw_responses = raw_responses or []
async def stream_events(self):
for event in self._events:
yield event
if self._exception is not None:
raise self._exception
def cancel(self, mode="immediate"):
self.cancel_calls.append(mode)
def to_state(self):
return self._state
class _FakeRunner:
last_calls = []
next_result = None
# When non-empty, each ``run_streamed`` call pops the next result from
# the front (left). Used by the empty-turn retry tests to return a
# different result per attempt. Falls back to ``next_result`` when
# exhausted/empty so existing single-result tests are unaffected.
next_results: list = []
@classmethod
def run_streamed(cls, agent, input, session, max_turns, run_config):
cls.last_calls.append(
{
"agent": agent,
"input": input,
"session": session,
"max_turns": max_turns,
"run_config": run_config,
}
)
if cls.next_results:
return cls.next_results.pop(0)
return cls.next_result
class _FakeFunctionTool:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
@dataclass
class _FakeSessionSettings:
"""
Minimal stand-in for ``agents.SessionSettings``.
:param limit: Maximum number of items to retrieve from the session.
``None`` means no limit.
"""
limit: int | None = None
class _FakeSQLiteSession:
def __init__(self, session_id):
self.session_id = session_id
self.items = []
self.clear_calls = 0
self.pop_calls = 0
async def get_items(self, limit=None):
if limit is None:
return list(self.items)
return list(self.items[-limit:])
async def add_items(self, items):
self.items.extend(items)
async def pop_item(self):
if not self.items:
return None
self.pop_calls += 1
return self.items.pop()
async def clear_session(self):
self.clear_calls += 1
self.items.clear()
class _FakeOpenAIProvider:
def __init__(self, **kwargs):
self.kwargs = kwargs
class _FakeRunConfig:
def __init__(self, **kwargs):
self.kwargs = kwargs
class _FakeAgent:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
class _FakeItemHelpers:
@staticmethod
def text_message_outputs(items):
return "".join(getattr(item, "text", "") for item in items)
class _FakeMaxTurnsExceeded(Exception):
pass
def _fake_agents_sdk():
return types.SimpleNamespace(
Agent=_FakeAgent,
Runner=_FakeRunner,
RunConfig=_FakeRunConfig,
OpenAIProvider=_FakeOpenAIProvider,
SQLiteSession=_FakeSQLiteSession,
FunctionTool=_FakeFunctionTool,
ItemHelpers=_FakeItemHelpers,
ModelSettings=_FakeModelSettings,
MaxTurnsExceeded=_FakeMaxTurnsExceeded,
SessionSettings=_FakeSessionSettings,
)
def test_reasoning_filter_replaces_list_content_with_none() -> None:
"""
List-type ``delta.content`` is replaced with ``None``.
What breaks if this fails: reasoning-model streams crash the SDK's handler.
"""
@dataclass
class _Delta:
content: object
@dataclass
class _Choice:
delta: _Delta
@dataclass
class _Chunk:
choices: list[_Choice]
async def _run_inner() -> list[object]:
async def _source():
yield _Chunk([_Choice(_Delta(["block1", "block2"]))])
yield _Chunk([_Choice(_Delta("normal text"))])
return [chunk async for chunk in _ReasoningBlockFilterStream(_source())]
chunks = _run(_run_inner())
assert chunks[0].choices[0].delta.content is None
assert chunks[1].choices[0].delta.content == "normal text"
def test_reasoning_filter_passes_string_content_through() -> None:
"""String ``delta.content`` is left unchanged by the filter."""
@dataclass
class _Delta:
content: object
@dataclass
class _Choice:
delta: _Delta
@dataclass
class _Chunk:
choices: list[_Choice]
async def _run_inner() -> list[object]:
async def _source():
yield _Chunk([_Choice(_Delta("hello"))])
return [chunk async for chunk in _ReasoningBlockFilterStream(_source())]
chunks = _run(_run_inner())
assert chunks[0].choices[0].delta.content == "hello"
def test_reasoning_filter_context_manager_delegates() -> None:
"""
``__aenter__`` / ``__aexit__`` delegate to the underlying stream.
A stream without these methods is also handled gracefully.
"""
class _TrackingStream:
entered = False
exited = False
def __aiter__(self):
return self
async def __anext__(self):
raise StopAsyncIteration
async def __aenter__(self):
_TrackingStream.entered = True
return self
async def __aexit__(self, *args):
_TrackingStream.exited = True
async def _run_inner() -> None:
stream = _TrackingStream()
wrapper = _ReasoningBlockFilterStream(stream)
async with wrapper:
pass
_run(_run_inner())
assert _TrackingStream.entered
assert _TrackingStream.exited
def test_wrap_client_non_streaming_create_not_wrapped() -> None:
"""
Non-streaming ``create()`` calls are returned unwrapped (no async iterator).
What breaks if this fails: regular chat completions are incorrectly wrapped
and callers that expect a ``ChatCompletion`` get a ``_ReasoningBlockFilterStream``.
"""
class _FakeResult:
pass
class _FakeCompletions:
async def create(self, **kwargs) -> object:
return _FakeResult()
def __getattr__(self, name: str) -> object:
raise AttributeError(name)
class _FakeChat:
completions = _FakeCompletions()
def __getattr__(self, name: str) -> object:
raise AttributeError(name)
class _FakeClient:
chat = _FakeChat()
def __getattr__(self, name: str) -> object:
raise AttributeError(name)
async def _run_inner() -> object:
client = _FakeClient()
_wrap_client_for_reasoning_models(client)
return await client.chat.completions.create(stream=False)
result = _run(_run_inner())
assert isinstance(result, _FakeResult)
class TestOpenAIAgentsSDKExecutor(unittest.TestCase):
def test_sanitize_replay_item_drops_long_ids(self):
item = {
"type": "message",
"id": "x" * 200,
"content": [{"type": "output_text", "text": "hello"}],
}
self.assertEqual(
_sanitize_replay_item(item),
{"type": "message", "content": [{"type": "output_text", "text": "hello"}]},
)
def test_build_tools_preserves_session_send_schema(self):
executor = OpenAIAgentsSDKExecutor(client=object())
tools = executor._build_tools(
_fake_agents_sdk(),
[
{
"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"],
},
}
],
)
self.assertEqual(tools[0].name, "sys_session_send")
self.assertEqual(
tools[0].params_json_schema["required"],
["tool", "session", "args"],
)
def test_streams_text_and_tool_events(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[
_FakeRawEvent(_FakeRawTextDelta("Hello ")),
_FakeRunItemEvent(
_FakeToolCallItem(_FakeToolCallRawItem("add", '{"a": 2, "b": 3}'))
),
_FakeRunItemEvent(
_FakeToolOutputItem(_FakeToolOutputRawItem(), {"result": 5})
),
_FakeRawEvent(_FakeRawTextDelta("world")),
],
final_output="Hello world",
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[{"name": "add", "description": "Add", "parameters": {"type": "object"}}],
"Be helpful.",
)
]
self.assertIsInstance(events[0], TextChunk)
self.assertEqual(events[0].text, "Hello ")
self.assertIsInstance(events[1], ToolCallRequest)
self.assertEqual(events[1].name, "add")
self.assertEqual(events[1].args, {"a": 2, "b": 3})
self.assertIsInstance(events[2], ToolCallComplete)
self.assertEqual(events[2].status, ToolCallStatus.SUCCESS)
self.assertEqual(events[2].result, {"result": 5})
self.assertIsInstance(events[3], TextChunk)
self.assertEqual(events[3].text, "world")
self.assertIsInstance(events[4], TurnComplete)
self.assertEqual(events[4].response, "Hello world")
self.assertEqual(
_FakeRunner.last_calls[0]["agent"].instructions,
"Be helpful.",
)
# parallel_tool_calls defaults to None (omitted from request) so that
# models like Kimi K2 that reject the field don't get a 400 error.
self.assertIsNone(
_FakeRunner.last_calls[0]["agent"].model_settings.parallel_tool_calls
)
_run(_t())
def test_streams_reasoning_deltas(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[
_FakeRawEvent(_FakeRawReasoningDelta("thinking...")),
_FakeRawEvent(_FakeRawReasoningDelta("")),
_FakeRawEvent(_FakeRawTextDelta("Hello")),
],
final_output="Hello",
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
]
reasoning = [e for e in events if isinstance(e, ReasoningChunk)]
self.assertEqual(len(reasoning), 1)
self.assertEqual(reasoning[0].delta, "thinking...")
self.assertEqual(reasoning[0].event_type, "reasoning_text")
text = [e for e in events if isinstance(e, TextChunk)]
self.assertEqual([t.text for t in text], ["Hello"])
_run(_t())
def test_databricks_client_default_model_uses_databricks_model(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="done")
client = types.SimpleNamespace(
base_url="https://profile-host.example.com/ai-gateway/openai/v1"
)
executor = OpenAIAgentsSDKExecutor(client=client)
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"",
)
]
self.assertEqual(events[-1].response, "done")
self.assertEqual(
_FakeRunner.last_calls[0]["agent"].model,
"databricks-gpt-5-5",
)
self.assertEqual(
_FakeRunner.last_calls[0]["run_config"].kwargs["model"],
"databricks-gpt-5-5",
)
_run(_t())
def test_parallel_tool_calls_can_be_overridden(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="done")
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[{"name": "add", "description": "Add", "parameters": {"type": "object"}}],
"Be helpful.",
ExecutorConfig(
model="gpt-5.3-codex", extra={"parallel_tool_calls": False}
),
)
]
self.assertEqual(events[-1].response, "done")
self.assertFalse(_FakeRunner.last_calls[0]["agent"].model_settings.parallel_tool_calls)
_run(_t())
def test_max_tokens_is_passed_to_model_settings(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="done")
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
ExecutorConfig(extra={"max_tokens": 65536}),
)
]
self.assertEqual(events[-1].response, "done")
self.assertEqual(
_FakeRunner.last_calls[0]["agent"].model_settings.max_tokens,
65536,
)
_run(_t())
def test_max_tokens_is_omitted_when_unset(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="done")
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
]
self.assertEqual(events[-1].response, "done")
self.assertIsNone(_FakeRunner.last_calls[0]["agent"].model_settings.max_tokens)
_run(_t())
def test_tool_error_output_sets_error_field(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[
_FakeRunItemEvent(_FakeToolCallItem(_FakeToolCallRawItem("fail", "{}"))),
_FakeRunItemEvent(
_FakeToolOutputItem(_FakeToolOutputRawItem(), {"error": "boom"})
),
],
final_output="tool failed",
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[
{
"name": "fail",
"description": "Fail",
"parameters": {"type": "object"},
}
],
"Be helpful.",
)
]
tool_events = [e for e in events if isinstance(e, ToolCallComplete)]
self.assertEqual(len(tool_events), 1)
self.assertEqual(tool_events[0].status, ToolCallStatus.ERROR)
self.assertEqual(tool_events[0].result, {"error": "boom"})
self.assertEqual(tool_events[0].error, "boom")
_run(_t())
def test_tool_blocked_output_sets_blocked_status(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[
_FakeRunItemEvent(_FakeToolCallItem(_FakeToolCallRawItem("fail", "{}"))),
_FakeRunItemEvent(
_FakeToolOutputItem(
_FakeToolOutputRawItem(),
{"blocked": True, "reason": "Nope"},
)
),
],
final_output="tool blocked",
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[
{
"name": "fail",
"description": "Fail",
"parameters": {"type": "object"},
}
],
"Be helpful.",
)
]
tool_events = [e for e in events if isinstance(e, ToolCallComplete)]
self.assertEqual(len(tool_events), 1)
self.assertEqual(tool_events[0].status, ToolCallStatus.BLOCKED)
self.assertEqual(tool_events[0].error, "Nope")
_run(_t())
def test_uses_full_history_first_then_latest_user(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="one")
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "first", "session_id": "s1"}],
[],
"",
)
]
self.assertEqual(events[-1].response, "one")
self.assertIsInstance(_FakeRunner.last_calls[0]["input"], list)
_FakeRunner.next_result = _FakeResult(events=[], final_output="two")
events = [
e
async for e in executor.run_turn(
[
{"role": "user", "content": "first", "session_id": "s1"},
{"role": "assistant", "content": "one", "session_id": "s1"},
{"role": "user", "content": "second", "session_id": "s1"},
],
[],
"",
)
]
self.assertEqual(events[-1].response, "two")
self.assertEqual(_FakeRunner.last_calls[1]["input"], "second")
_run(_t())
def test_started_session_preserves_structured_user_content(self):
"""Regression: fast path must not json.dumps a structured content list.
input_file blocks with data: URIs are normalised to input_text before
reaching the SDK runner (the openai-agents chatcmpl_converter and the
Databricks endpoint don't support the ``file`` content block type).
The key invariant is that the content arrives as a proper list, not a
json.dumps'd string.
"""
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="one")
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
_ = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "first", "session_id": "s1"}],
[],
"",
)
]
_FakeRunner.next_result = _FakeResult(events=[], final_output="two")
structured = [
{"type": "input_text", "text": "read this"},
{
"type": "input_file",
"filename": "protocol.md",
"file_data": "data:text/markdown;base64,VGVzdA==",
},
]
_ = [
e
async for e in executor.run_turn(
[
{"role": "user", "content": "first", "session_id": "s1"},
{"role": "assistant", "content": "one", "session_id": "s1"},
{"role": "user", "content": structured, "session_id": "s1"},
],
[],
"",
)
]
received = _FakeRunner.last_calls[1]["input"]
self.assertIsInstance(received, list)
self.assertEqual(received[0]["role"], "user")
# input_file blocks are normalised to input_text before the SDK
# runner sees them (data: URI decoded to plain text).
# "VGVzdA==" decodes to "Test".
expected_content = [
{"type": "input_text", "text": "read this"},
{"type": "input_text", "text": "Test"},
]
self.assertEqual(received[0]["content"], expected_content)
_run(_t())
def test_started_session_replays_multiple_new_messages(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="one")
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
first_events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "first", "session_id": "s1"}],
[],
"",
)
]
self.assertEqual(first_events[-1].response, "one")
_FakeRunner.next_result = _FakeResult(events=[], final_output="two")
second_events = [
e
async for e in executor.run_turn(
[
{"role": "user", "content": "first", "session_id": "s1"},
{"role": "assistant", "content": "one", "session_id": "s1"},
{
"role": "user",
"content": (
"[SYSTEM] The previous assistant turn was cancelled "
"before completion. Continue from the next message."
),
"session_id": "s1",
"metadata": {"framework": "cancellation_notice"},
},
{"role": "user", "content": "what happened", "session_id": "s1"},
],
[],
"",
)
]
self.assertEqual(second_events[-1].response, "two")
self.assertIsInstance(_FakeRunner.last_calls[1]["input"], list)
self.assertEqual(
[item["content"] for item in _FakeRunner.last_calls[1]["input"]],
[
(
"[SYSTEM] The previous assistant turn was cancelled "
"before completion. Continue from the next message."
),
"what happened",
],
)
_run(_t())
def test_transient_framework_notice_is_included_in_delta_and_not_persisted_in_cursor(self):
"""Regression for a bug where an inbox notice appended to the messages
was advancing history_cursor past itself, and on the next turn the
delta was empty because the cursor was at ``len(messages) + 1`` but the
notice was never stored. The agent then ran with ``input=""`` and kept
generating forever without seeing the notice."""
async def _t():
_FakeRunner.last_calls = []
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
_FakeRunner.next_result = _FakeResult(events=[], final_output="one")
first_events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "first", "session_id": "s1"}],
[],
"",
)
]
self.assertEqual(first_events[-1].response, "one")
state = executor._get_or_create_session_state(_fake_agents_sdk(), "s1")
# After the first turn, the cursor should be past the user and
# the assistant response (2 persisted items in context terms).
self.assertEqual(state.history_cursor, 2)
# Second turn: transient inbox notice appended to the tail,
# no new persisted user message.
notice = {
"role": "user",
"content": "[SYSTEM] There is 1 unread inbox item available.",
"session_id": "s1",
"metadata": {"framework": "inbox_notice"},
}
_FakeRunner.next_result = _FakeResult(events=[], final_output="two")
second_events = [
e
async for e in executor.run_turn(
[
{"role": "user", "content": "first", "session_id": "s1"},
{"role": "assistant", "content": "one", "session_id": "s1"},
notice,
],
[],
"",
)
]
self.assertEqual(second_events[-1].response, "two")
# The notice must actually be passed to the runner (not empty).
# A single delta user message is passed as a raw string by the
# executor; the critical thing is that the notice content is there.
second_input = _FakeRunner.last_calls[1]["input"]
if isinstance(second_input, list):
contents = [item["content"] for item in second_input]
else:
contents = [second_input]
self.assertEqual(contents, ["[SYSTEM] There is 1 unread inbox item available."])
# It must NOT be the empty-delta fallback that the buggy version
# produced before the fix.
self.assertNotEqual(second_input, "")
# And the cursor should only reflect persisted messages, not the
# transient notice, so a future turn still sees any new real
# messages without off-by-one skips.
self.assertEqual(state.history_cursor, 3)
_run(_t())
def test_reuses_agent_across_turns_for_same_session_config(self):
async def _t():
_FakeRunner.last_calls = []
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
_FakeRunner.next_result = _FakeResult(events=[], final_output="one")
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "first", "session_id": "s1"}],
[{"name": "add", "description": "Add", "parameters": {"type": "object"}}],
"Be helpful.",
)
]
self.assertEqual(events[-1].response, "one")
_FakeRunner.next_result = _FakeResult(events=[], final_output="two")
events = [
e
async for e in executor.run_turn(
[
{"role": "user", "content": "first", "session_id": "s1"},
{"role": "assistant", "content": "one", "session_id": "s1"},
{"role": "user", "content": "second", "session_id": "s1"},
],
[{"name": "add", "description": "Add", "parameters": {"type": "object"}}],
"Be helpful.",
)
]
self.assertEqual(events[-1].response, "two")
self.assertIs(
_FakeRunner.last_calls[0]["agent"],
_FakeRunner.last_calls[1]["agent"],
)
_run(_t())
def test_stepwise_internal_turns_resume_without_replaying_user_input(self):
async def _t():
_FakeRunner.last_calls = []
resume_state = types.SimpleNamespace(
_current_turn=7,
_max_turns=99,
_current_turn_persisted_item_count=3,
)
first_result = _FakeResult(
events=[
_FakeRunItemEvent(
_FakeToolCallItem(_FakeToolCallRawItem("add", '{"a": 2, "b": 3}'))
),
_FakeRunItemEvent(
_FakeToolOutputItem(_FakeToolOutputRawItem(), {"result": 5})
),
],
exception=_FakeMaxTurnsExceeded("Max turns (1) exceeded"),
state=resume_state,
)
_FakeRunner.next_result = first_result
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
first_events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "count", "session_id": "s1"}],
[{"name": "add", "description": "Add", "parameters": {"type": "object"}}],
"Be helpful.",
ExecutorConfig(extra={"stepwise_internal_turns": True}),
)
]
_FakeRunner.next_result = _FakeResult(events=[], final_output="done")
second_events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "count", "session_id": "s1"}],
[{"name": "add", "description": "Add", "parameters": {"type": "object"}}],
"Be helpful.",
ExecutorConfig(extra={"stepwise_internal_turns": True}),
)
]
self.assertTrue(first_events[-1].continue_turn)
self.assertEqual(first_result.cancel_calls, [])
self.assertEqual(_FakeRunner.last_calls[0]["max_turns"], 1)
self.assertIs(_FakeRunner.last_calls[1]["input"], resume_state)
self.assertEqual(resume_state._current_turn, 0)
self.assertEqual(resume_state._max_turns, 1)
self.assertEqual(resume_state._current_turn_persisted_item_count, 0)
self.assertEqual(second_events[-1].response, "done")
_run(_t())
def test_stepwise_internal_turns_use_new_user_message_instead_of_resume_state(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[
_FakeRunItemEvent(
_FakeToolCallItem(_FakeToolCallRawItem("add", '{"a": 2, "b": 3}'))
),
_FakeRunItemEvent(
_FakeToolOutputItem(_FakeToolOutputRawItem(), {"result": 5})
),
],
exception=_FakeMaxTurnsExceeded("Max turns (1) exceeded"),
state=object(),
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
first_events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "count", "session_id": "s1"}],
[],
"",
ExecutorConfig(extra={"stepwise_internal_turns": True}),
)
]
_FakeRunner.next_result = _FakeResult(events=[], final_output="stopped")
second_events = [
e
async for e in executor.run_turn(
[
{"role": "user", "content": "count", "session_id": "s1"},
{"role": "user", "content": "stop", "session_id": "s1"},
],
[],
"",
ExecutorConfig(
extra={
"stepwise_internal_turns": True,
"new_user_messages_flushed": True,
}
),
)
]
self.assertTrue(first_events[-1].continue_turn)
self.assertEqual(_FakeRunner.last_calls[1]["input"], "stop")
self.assertEqual(second_events[-1].response, "stopped")
_run(_t())
def test_interrupt_session_cancels_stream_consumer_task(self):
"""
``interrupt_session`` must cancel the per-turn
``stream_consumer_task`` so the cancel lands inside
the SDK's ``await self._event_queue.get()`` (the SDK's
own ``except CancelledError`` then closes the httpx
stream — the load-bearing fix).
Previously the executor called
``result.cancel(mode="immediate")``, which only sets a
flag and does not close the network — events kept
streaming for 15+ seconds. The new contract: cancel
the pump task so cancellation reaches the SDK's
documented escape hatch.
"""
async def _t():
executor = OpenAIAgentsSDKExecutor(client=object())
state = executor._get_or_create_session_state(_fake_agents_sdk(), "s1")
state.active_result = _FakeResult(events=[])
# Set to a non-zero value so the rollback-target
# assertion below verifies a real assignment, not
# just two uninitialized zeros being equal.
state.run_item_count_before = 42
# Real task we can observe — assert it gets
# cancelled. A bare coroutine wrapped in
# ``asyncio.create_task`` is the right shape; we
# never await its result, just inspect cancellation.
async def _fake_pump() -> None:
# Park forever until cancelled; the cancel
# is the whole assertion target.
await asyncio.Event().wait()
pump_task = asyncio.create_task(_fake_pump())
state.stream_consumer_task = pump_task
interrupted = await executor.interrupt_session("s1")
# Yield once so the cancel can deliver to the
# parked task before we inspect ``cancelled()``.
await asyncio.sleep(0)
self.assertTrue(interrupted)
# Sets interrupt_requested so the consumer side
# also drains the queue without emitting (defense
# in depth — the cancellation should already have
# propagated, but we want the consumer loop to
# also short-circuit if we're racing it).
self.assertTrue(state.interrupt_requested)
# Load-bearing assertion — cancelling the pump
# task is what makes the SDK actually halt its
# network stream. ``.cancelled()`` (not just
# ``.done()``) — the regression we're guarding
# against is "interrupt_session stops calling
# consumer.cancel()". A normal-completed task
# (``.done()`` but not cancelled) would mask
# that regression; the strict ``.cancelled()``
# assertion catches it.
self.assertTrue(
pump_task.cancelled(),
"interrupt_session must cancel state.stream_consumer_task; "
"pump_task.cancelled() was False. A regression that drops "
"the consumer.cancel() call leaves the task parked.",
)
# Rollback target captured so the next run_turn
# rewinds the SDK session to its pre-cancel
# state. The 42 was set above; verifies the
# assignment is the load-bearing op (not a coincidental
# zero match).
self.assertEqual(state.rollback_to_item_count, 42)
# Drain the cancelled task so asyncio doesn't
# warn about un-awaited cancellation.
with contextlib.suppress(asyncio.CancelledError):
await pump_task
_run(_t())
def test_rebuilds_agent_when_prompt_changes(self):
async def _t():
_FakeRunner.last_calls = []
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
_FakeRunner.next_result = _FakeResult(events=[], final_output="one")
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "first", "session_id": "s1"}],
[],
"Be helpful.",
)
]
self.assertEqual(events[-1].response, "one")
_FakeRunner.next_result = _FakeResult(events=[], final_output="two")
events = [
e
async for e in executor.run_turn(
[
{"role": "user", "content": "first", "session_id": "s1"},
{"role": "assistant", "content": "one", "session_id": "s1"},
{"role": "user", "content": "second", "session_id": "s1"},
],
[],
"Be terse.",
)
]
self.assertEqual(events[-1].response, "two")
self.assertIsNot(
_FakeRunner.last_calls[0]["agent"],
_FakeRunner.last_calls[1]["agent"],
)
self.assertEqual(_FakeRunner.last_calls[1]["agent"].instructions, "Be terse.")
_run(_t())
def test_sys_tool_names_are_passed_through_in_events(self):
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[
_FakeRunItemEvent(
_FakeToolCallItem(
_FakeToolCallRawItem(
"sys_runtime_execute",
'{"code": "print(1)"}',
)
)
),
],
final_output="done",
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[
{
"name": "sys_runtime_execute",
"description": "Run code",
"parameters": {"type": "object"},
}
],
"Be helpful.",
)
]
self.assertEqual(
_FakeRunner.last_calls[0]["agent"].tools[0].name, "sys_runtime_execute"
)
self.assertEqual(_FakeRunner.last_calls[0]["agent"].tools[0].description, "Run code")
self.assertIsInstance(events[0], ToolCallRequest)
self.assertEqual(events[0].name, "sys_runtime_execute")
_run(_t())
def test_turn_usage_single_call_billing_and_context_match(self):
"""
For a single-LLM-call turn (no tool calls), total_tokens equals
that call's total and context_tokens matches (both are input +
output for one call).
What breaks if this fails: single-turn billing totals are wrong,
or the context ring shows a stale/missing value.
"""
async def _t():
_FakeRunner.last_calls = []
result = _FakeResult(events=[], final_output="hello")
result.raw_responses = [
_FakeRawResponse(
usage=_FakeUsage(input_tokens=1000, output_tokens=200, total_tokens=1200)
)
]
_FakeRunner.next_result = result
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"",
)
]
turn_complete = next(e for e in events if isinstance(e, TurnComplete))
# Billing: single call, so sum equals that one call.
self.assertEqual(
turn_complete.usage["input_tokens"],
1000,
"input_tokens must equal the single LLM call's input count.",
)
self.assertEqual(
turn_complete.usage["output_tokens"],
200,
"output_tokens must equal the single LLM call's output count.",
)
self.assertEqual(
turn_complete.usage["total_tokens"],
1200,
"total_tokens must be the billing sum (input + output = 1200).",
)
# context_tokens is always set so the REPL and compaction
# don't need a total_tokens fallback path.
self.assertEqual(
turn_complete.usage["context_tokens"],
1200,
"context_tokens must equal total for single-call turns.",
)
_run(_t())
def test_turn_usage_multi_call_billing_sum_and_context_tokens_separate(self):
"""
For a multi-LLM-call turn (tool calls), total_tokens is the
billing SUM across all sub-turns, and context_tokens is set to
the LAST sub-turn's total for accurate context ring display.
Root cause of the observed bug: if context_tokens is absent and
total_tokens carries the sum, the context ring spikes to an
inflated value during tool-call turns (because each sub-turn
repeats the full conversation history as input), then drops back
on the next plain turn — giving a spurious "compaction" impression.
What breaks if this fails: either billing totals are wrong
(total_tokens not summed) or the context ring shows the inflated
sum instead of the accurate last-call fill.
"""
async def _t():
_FakeRunner.last_calls = []
result = _FakeResult(events=[], final_output="done")
# Two sub-turns: first call issues a tool call (small output),
# second call produces the final answer (larger output).
# Each call receives the full conversation history as input,
# so input_tokens grows across calls.
result.raw_responses = [
_FakeRawResponse(
usage=_FakeUsage(input_tokens=5000, output_tokens=100, total_tokens=5100)
),
_FakeRawResponse(
usage=_FakeUsage(input_tokens=5300, output_tokens=400, total_tokens=5700)
),
]
_FakeRunner.next_result = result
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "search for it", "session_id": "s1"}],
[],
"",
)
]
turn_complete = next(e for e in events if isinstance(e, TurnComplete))
# Billing: SUM across all sub-turns.
self.assertEqual(
turn_complete.usage["input_tokens"],
10300,
"input_tokens must sum across sub-turns for billing accuracy "
"(5000 + 5300 = 10300).",
)
self.assertEqual(
turn_complete.usage["output_tokens"],
500,
"output_tokens must sum across sub-turns for billing accuracy (100 + 400 = 500).",
)
self.assertEqual(
turn_complete.usage["total_tokens"],
10800,
"total_tokens must be the billing sum (5100 + 5700 = 10800). "
"If 5700, total_tokens was incorrectly set to last-call only.",
)
# Context fill: LAST call's total only.
# If this is 10800 (the sum), the context ring spikes artificially
# during tool-call turns and drops spuriously on plain turns.
self.assertEqual(
turn_complete.usage["context_tokens"],
5700,
"context_tokens must equal the LAST call's total (5700) for correct "
"context ring display. If 10800, context_tokens is being set to the "
"billing sum — that inflates the ring during tool-call turns.",
)
_run(_t())
def test_turn_usage_multi_call_context_tokens_falls_back_when_total_missing(self):
"""
When total_tokens is absent (zero) on the last raw response,
context_tokens is computed as last_input + last_output.
What breaks if this fails: context ring shows 0 when the model
does not populate total_tokens on the last sub-turn.
"""
async def _t():
_FakeRunner.last_calls = []
result = _FakeResult(events=[], final_output="done")
result.raw_responses = [
_FakeRawResponse(
usage=_FakeUsage(input_tokens=5000, output_tokens=100, total_tokens=5100)
),
# Last call: total_tokens=0, executor must fall back to in + out.
_FakeRawResponse(
usage=_FakeUsage(input_tokens=5300, output_tokens=400, total_tokens=0)
),
]
_FakeRunner.next_result = result
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"",
)
]
turn_complete = next(e for e in events if isinstance(e, TurnComplete))
# context_tokens falls back to last_in + last_out = 5300 + 400 = 5700.
self.assertEqual(
turn_complete.usage["context_tokens"],
5700,
"When last raw_response.total_tokens is 0, context_tokens must fall "
"back to last_input + last_output (5700).",
)
# total_tokens is the billing sum: 5100 + 0 = 5100.
self.assertEqual(
turn_complete.usage["total_tokens"],
5100,
"total_tokens must be the billing sum across sub-turns (5100 + 0 = 5100).",
)
_run(_t())
# ── _get_openai_async_client: --profile priority ──────────────
#
# Function-based tests for the profile-precedence fix. When a
# user passes ``--profile X``, the resolved Databricks credentials
# must win over any ``OPENAI_API_KEY`` / ``OPENAI_BASE_URL`` the
# shell happens to have set — otherwise an old export silently
# redirects the call to ``api.openai.com`` and the user's
# requested workspace is ignored.
def test_get_openai_client_profile_uses_callback_auth(monkeypatch):
"""Explicit ``profile`` uses httpx callback auth, not a static ``api_key``.
Verifies that when a Databricks profile is provided, the client
is constructed with an ``http_client`` carrying a
``_DatabricksBearerAuth`` that refreshes tokens per-request,
and that the ``api_key`` is a placeholder (not the real token).
:param monkeypatch: Pytest monkeypatch fixture.
"""
import httpx
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.setenv("OPENAI_API_KEY", "should-not-be-used")
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
class _FakeConfig:
host = "https://profile-host.example.com"
profile = "dev"
def authenticate(self):
return {"Authorization": "Bearer fresh-tok"}
monkeypatch.setattr(_sdk_config_mod, "Config", lambda **_kw: _FakeConfig())
captured: dict[str, Any] = {}
class _StubAsyncOpenAI:
def __init__(self, **kwargs):
captured.update(kwargs)
import openai as _openai_mod
with patch.object(_openai_mod, "AsyncOpenAI", _StubAsyncOpenAI, create=True):
_get_openai_async_client(profile="dev")
assert captured["base_url"] == "https://profile-host.example.com/ai-gateway/openai/v1"
assert captured["api_key"] != "should-not-be-used"
assert isinstance(captured["http_client"], httpx.AsyncClient)
def test_get_openai_client_host_override_uses_ucode_auth_command(monkeypatch):
"""ucode host override uses ucode auth command without profile lookup.
:param monkeypatch: Pytest monkeypatch fixture.
"""
import httpx
import omnigent.inner.databricks_executor as db_exec
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.setenv("OPENAI_API_KEY", "should-not-be-used")
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.setattr(
db_exec,
"_resolve_databricks_auth",
lambda _profile: (_ for _ in ()).throw(AssertionError("profile lookup used")),
)
captured: dict[str, Any] = {}
class _StubAsyncOpenAI:
def __init__(self, **kwargs):
captured.update(kwargs)
import openai as _openai_mod
with patch.object(_openai_mod, "AsyncOpenAI", _StubAsyncOpenAI, create=True):
_get_openai_async_client(
profile="missing-profile",
host_override="https://example.databricks.com/",
base_url_override="https://example.databricks.com/ai-gateway/codex/v1",
databricks_auth_command="printf token",
)
assert captured["base_url"] == "https://example.databricks.com/ai-gateway/codex/v1"
assert captured["api_key"] != "should-not-be-used"
assert isinstance(captured["http_client"], httpx.AsyncClient)
def test_get_openai_client_host_override_requires_base_url(monkeypatch):
"""ucode host override fails loud when ucode omits the base URL.
:param monkeypatch: Pytest monkeypatch fixture.
"""
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
with pytest.raises(OSError, match="GATEWAY_BASE_URL"):
_get_openai_async_client(
profile="missing-profile",
host_override="https://example.databricks.com/",
databricks_auth_command="printf token",
)
def test_get_openai_client_host_override_requires_auth_command(monkeypatch):
"""ucode host override fails loud when ucode omits the auth command.
:param monkeypatch: Pytest monkeypatch fixture.
"""
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
with pytest.raises(OSError, match="GATEWAY_AUTH_COMMAND"):
_get_openai_async_client(
profile="missing-profile",
host_override="https://example.databricks.com/",
base_url_override="https://example.databricks.com/ai-gateway/codex/v1",
)
def test_get_openai_client_api_key_falls_back_to_env_base_url(monkeypatch):
"""A spec-level api_key with NO override honors ambient ``OPENAI_BASE_URL``.
Regression for the residual gateway 401 (continuation-turn daemon
spawns): the api_key is frequently a gateway credential (e.g. a
Databricks AI Gateway PAT detected from ``OPENAI_API_KEY``), and the
companion base_url can be dropped on the daemon → runner → harness
propagation chain (the spec-auth bake omits it when ``OPENAI_BASE_URL``
is absent at materialization time; a reused local daemon may predate the
env var). Without the ambient fallback, the gateway PAT is sent to
``api.openai.com`` and 401s. The runner inherits ``OPENAI_BASE_URL``, so
honoring it here keeps the gateway target present on every turn.
:param monkeypatch: Pytest monkeypatch fixture.
"""
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.setenv("OPENAI_BASE_URL", "https://gateway.example.com/ai-gateway/openai/v1")
captured: dict[str, Any] = {}
class _StubAsyncOpenAI:
def __init__(self, **kwargs):
captured.update(kwargs)
import openai as _openai_mod
with patch.object(_openai_mod, "AsyncOpenAI", _StubAsyncOpenAI, create=True):
_get_openai_async_client(api_key="dapi-gateway-token", base_url_override=None)
assert captured["api_key"] == "dapi-gateway-token"
assert captured["base_url"] == "https://gateway.example.com/ai-gateway/openai/v1", (
"A spec-level api_key with no base_url override must fall back to the "
"ambient OPENAI_BASE_URL the runner inherits; routing to api.openai.com "
"(base_url=None) sends a gateway PAT to OpenAI and 401s."
)
def test_get_openai_client_api_key_override_wins_over_env_base_url(monkeypatch):
"""An explicit base_url override wins over the ambient ``OPENAI_BASE_URL``.
The baked ``executor.auth.base_url`` (threaded via
``HARNESS_OPENAI_AGENTS_GATEWAY_BASE_URL``) is the authoritative
per-spec target and must not be shadowed by a differing env var.
:param monkeypatch: Pytest monkeypatch fixture.
"""
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.setenv("OPENAI_BASE_URL", "https://wrong-env.example.com/v1")
captured: dict[str, Any] = {}
class _StubAsyncOpenAI:
def __init__(self, **kwargs):
captured.update(kwargs)
import openai as _openai_mod
with patch.object(_openai_mod, "AsyncOpenAI", _StubAsyncOpenAI, create=True):
_get_openai_async_client(
api_key="dapi-gateway-token",
base_url_override="https://spec-gateway.example.com/openai/v1",
)
assert captured["base_url"] == "https://spec-gateway.example.com/openai/v1"
def test_get_openai_client_api_key_no_env_defaults_to_openai(monkeypatch):
"""A genuine OpenAI key with no gateway anywhere still defaults to OpenAI.
With no override and no ambient ``OPENAI_BASE_URL``, ``base_url`` stays
``None`` so the AsyncOpenAI client targets ``api.openai.com`` — the
correct behavior for a real ``sk-...`` key.
:param monkeypatch: Pytest monkeypatch fixture.
"""
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
captured: dict[str, Any] = {}
class _StubAsyncOpenAI:
def __init__(self, **kwargs):
captured.update(kwargs)
import openai as _openai_mod
with patch.object(_openai_mod, "AsyncOpenAI", _StubAsyncOpenAI, create=True):
_get_openai_async_client(api_key="sk-real-openai-key", base_url_override=None)
assert captured["api_key"] == "sk-real-openai-key"
assert captured["base_url"] is None
def test_get_openai_client_no_profile_honors_env_vars(monkeypatch):
"""Without an explicit profile, the env-var branch still works.
:param monkeypatch: Pytest monkeypatch fixture.
"""
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
monkeypatch.setenv("OPENAI_BASE_URL", "https://env-host.example.com/v1")
monkeypatch.setenv("OPENAI_API_KEY", "env-key")
captured: dict[str, str] = {}
class _StubAsyncOpenAI:
def __init__(self, **kwargs):
captured.update(kwargs)
import openai as _openai_mod
with patch.object(_openai_mod, "AsyncOpenAI", _StubAsyncOpenAI, create=True):
_get_openai_async_client(profile=None)
assert captured["base_url"] == "https://env-host.example.com/v1"
assert captured["api_key"] == "env-key"
assert "http_client" not in captured
def test_get_openai_client_invalid_profile_raises_auth_error(monkeypatch):
"""An invalid profile raises ``DatabricksAuthError`` with login instructions.
:param monkeypatch: Pytest monkeypatch fixture.
"""
import pytest
import omnigent.inner.databricks_executor as db_exec
from omnigent.inner.databricks_executor import DatabricksAuthError
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
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("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)
with pytest.raises(DatabricksAuthError, match="databricks auth login -p dogfood"):
_get_openai_async_client(profile="dogfood")
def test_get_openai_client_invalid_profile_with_env_fallback_warns(monkeypatch, caplog):
"""Profile auth failure with OPENAI_BASE_URL available warns and falls through.
When ``_get_openai_async_client`` is called with a named profile that
cannot be resolved (``DatabricksAuthError``), but ``OPENAI_BASE_URL``
is set, the function must:
1. Log a warning mentioning the profile name and the fallback.
2. Return a client configured from OPENAI_BASE_URL, NOT re-raise.
This preserves backward compatibility for CI environments (OIDC tokens
injected via OPENAI_BASE_URL) that specify a profile in the agent spec
but rely on env-var credentials at runtime.
:param monkeypatch: Pytest monkeypatch fixture.
:param caplog: Pytest log capture fixture.
"""
import logging
import omnigent.inner.databricks_executor as db_exec
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
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.setenv("OPENAI_BASE_URL", "https://example.databricks.com/serving-endpoints")
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)
with caplog.at_level(logging.WARNING):
client = _get_openai_async_client(profile="dogfood")
# Should not raise — fell through to OPENAI_BASE_URL.
# A warning must be emitted so the fallback is not silent.
assert any(
"dogfood" in record.message and "OPENAI_BASE_URL" in record.message
for record in caplog.records
), (
f"Expected a warning mentioning 'dogfood' and 'OPENAI_BASE_URL'. "
f"Got: {[r.message for r in caplog.records]}"
)
# Client must use OPENAI_BASE_URL, not the (failed) profile host.
assert client.base_url.host == "example.databricks.com", (
f"Expected client to use OPENAI_BASE_URL host 'example.databricks.com', "
f"got {client.base_url.host!r}. The profile-failure fallthrough may be broken."
)
def test_get_openai_client_missing_databricks_sdk_raises_actionable_error(monkeypatch):
"""Missing ``databricks-sdk`` with no env-var fallback gives an actionable error.
When a Databricks-hosted model is requested, ``databricks-sdk`` is not
installed, and no ``OPENAI_API_KEY``/``OPENAI_BASE_URL`` fallback is
available, the function must raise ``ImportError`` with install
instructions — not crash with an opaque traceback.
Regression test: the SDK-missing path must fail loudly, not silently.
:param monkeypatch: Pytest monkeypatch fixture.
"""
import pytest
import omnigent.inner.databricks_executor as db_exec
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
def _import_error(*_args, **_kw):
raise ImportError("No module named 'databricks.sdk'")
monkeypatch.setattr(db_exec, "_resolve_databricks_auth", _import_error)
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)
with pytest.raises(ImportError, match="pip install"):
_get_openai_async_client(profile=None, model="databricks-gpt-5")
def test_get_openai_client_missing_databricks_sdk_with_env_falls_through(monkeypatch, caplog):
"""Missing ``databricks-sdk`` with OPENAI_API_KEY set falls through gracefully.
When ``databricks-sdk`` is absent but env-var credentials are available,
the function should log a warning and return a client configured from the
env vars — not crash.
Regression test: a missing ``databricks-sdk`` must degrade to the
env-var client with a warning, not crash.
:param monkeypatch: Pytest monkeypatch fixture.
:param caplog: Pytest log capture fixture.
"""
import logging
import omnigent.inner.databricks_executor as db_exec
from omnigent.inner.openai_agents_sdk_executor import _get_openai_async_client
def _import_error(*_args, **_kw):
raise ImportError("No module named 'databricks.sdk'")
monkeypatch.setattr(db_exec, "_resolve_databricks_auth", _import_error)
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)
with caplog.at_level(logging.WARNING):
client = _get_openai_async_client(profile="dev", model="databricks-gpt-5")
assert any("databricks-sdk" in record.message for record in caplog.records), (
f"Expected a warning about missing databricks-sdk. "
f"Got: {[r.message for r in caplog.records]}"
)
assert client is not None
def test_run_turn_auth_error_yields_actionable_message(monkeypatch):
"""``run_turn`` yields the actionable ``DatabricksAuthError`` message,
not the raw ``__cause__`` string.
When the agents SDK raises a ``DatabricksAuthError`` (e.g. from
``_DatabricksBearerAuth.auth_flow`` when the OAuth refresh token has
expired), the ``ExecutorError.message`` must be the high-level
"Run: databricks auth login -p X" guidance — not the underlying SDK
exception text like "token expired".
:param monkeypatch: Pytest monkeypatch fixture.
"""
from omnigent.inner.databricks_executor import DatabricksAuthError
# DatabricksAuthError carries the actionable message; its __cause__ is
# the raw SDK exception that is NOT suitable to show the user.
raw_sdk_exc = ValueError("token expired")
auth_error = DatabricksAuthError(
"Databricks authentication failed for profile 'dev'. Run: databricks auth login -p dev"
)
auth_error.__cause__ = raw_sdk_exc
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], exception=auth_error)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = _run(
_collect(
executor.run_turn(
messages=[{"role": "user", "content": "hi"}],
tools=[],
system_prompt="",
config=ExecutorConfig(),
)
)
)
error_events = [e for e in events if isinstance(e, ExecutorError)]
assert len(error_events) == 1, (
f"Expected exactly 1 ExecutorError, got {len(error_events)}. Events: {events!r}"
)
# The message must be the DatabricksAuthError text (actionable), not
# the __cause__ text ("token expired"). If str(exc.__cause__) were used
# instead of str(exc), this assertion would fail with "token expired".
assert "databricks auth login -p dev" in error_events[0].message, (
f"Expected actionable 'databricks auth login' guidance in message, "
f"got: {error_events[0].message!r}. "
f"If 'token expired' appears instead, auth_msg is reading exc.__cause__ "
f"rather than str(exc)."
)
assert "token expired" not in error_events[0].message, (
f"Raw SDK exception text leaked into the user-facing error message: "
f"{error_events[0].message!r}. This means auth_msg is using "
f"str(exc.__cause__) instead of str(exc)."
)
# ---------------------------------------------------------------------------
# Tests for _normalize_content_blocks_for_chat and
# _normalize_responses_items_for_chat — input_file → input_text conversion
# ---------------------------------------------------------------------------
def test_normalize_content_blocks_input_file_data_uri_converted_to_input_text() -> None:
"""``input_file`` blocks with a ``data:`` URI are decoded to ``input_text``.
The openai-agents SDK's chatcmpl_converter passes ``file_data`` verbatim
to the ``file`` Chat Completions content type, which expects plain base64,
not a data URI. We must decode before handing off.
"""
file_content = "# heading\nsome text"
b64 = base64.b64encode(file_content.encode()).decode()
blocks = [{"type": "input_file", "file_data": f"data:text/plain;base64,{b64}"}]
result = _normalize_content_blocks_for_chat(blocks)
assert result == [{"type": "input_text", "text": file_content}], (
f"Expected decoded input_text block, got: {result!r}"
)
def test_normalize_content_blocks_non_file_blocks_pass_through_unchanged() -> None:
"""Known blocks without metadata are returned as-is (identity preserved)."""
blocks = [{"type": "input_text", "text": "hello"}]
result = _normalize_content_blocks_for_chat(blocks)
# No input_file or metadata present → original list object returned.
assert result is blocks, (
"Expected the original list object when no input_file blocks or metadata are present"
)
def test_normalize_content_blocks_strips_filename_from_input_image() -> None:
"""``filename`` is Omnigent metadata and must not reach OpenAI."""
blocks = [
{
"type": "input_image",
"image_url": "data:image/png;base64,abcd",
"filename": "screenshot.png",
}
]
result = _normalize_content_blocks_for_chat(blocks)
assert result == [{"type": "input_image", "image_url": "data:image/png;base64,abcd"}]
assert "filename" not in result[0]
assert result is not blocks
def test_normalize_content_blocks_preserves_input_image_detail_for_http_url() -> None:
"""Supported ``input_image.detail`` survives metadata sanitization for URLs."""
blocks = [
{
"type": "input_image",
"image_url": "https://example.com/screenshot.png",
"detail": "high",
"filename": "screenshot.png",
}
]
result = _normalize_content_blocks_for_chat(blocks)
assert result == [
{
"type": "input_image",
"image_url": "https://example.com/screenshot.png",
"detail": "high",
}
]
def test_normalize_content_blocks_data_uri_input_image_is_preserved() -> None:
"""Inline uploaded images are preserved as conventional base64 data URLs."""
blocks = [
{
"type": "input_image",
"image_url": "data:image/png;base64,abcd",
"filename": "screenshot.png",
}
]
result = _normalize_content_blocks_for_chat(blocks)
assert result == [{"type": "input_image", "image_url": "data:image/png;base64,abcd"}]
def test_normalize_content_blocks_empty_file_data_dropped() -> None:
"""``input_file`` with empty ``file_data`` is silently dropped."""
result = _normalize_content_blocks_for_chat([{"type": "input_file", "file_data": ""}])
# Empty content has no value; dropping avoids sending a vacuous block.
assert result == [], f"Expected empty list, got: {result!r}"
def test_normalize_content_blocks_malformed_base64_dropped() -> None:
"""``input_file`` with a malformed base64 payload is silently dropped."""
result = _normalize_content_blocks_for_chat(
[{"type": "input_file", "file_data": "data:text/plain;base64,!!!bad!!!"}]
)
# Decode failure → empty string → dropped, not raised.
assert result == [], f"Expected empty list, got: {result!r}"
def test_normalize_content_blocks_binary_file_dropped() -> None:
"""``input_file`` with a binary (non-text) MIME type is silently dropped.
Decoding a PDF or other binary payload as UTF-8 produces garbled replacement
characters that waste tokens and provide no useful information to the model.
"""
import base64
b64 = base64.b64encode(b"%PDF-1.4 binary content").decode()
result = _normalize_content_blocks_for_chat(
[{"type": "input_file", "file_data": f"data:application/pdf;base64,{b64}"}]
)
# Binary MIME type → dropped, not inlined as garbage text.
assert result == [], f"Expected binary block to be dropped, got: {result!r}"
def test_normalize_content_blocks_plain_file_data_used_as_text() -> None:
"""``input_file`` whose ``file_data`` is plain text (no data URI) is used as-is."""
result = _normalize_content_blocks_for_chat(
[{"type": "input_file", "file_data": "plain content here"}]
)
assert result == [{"type": "input_text", "text": "plain content here"}]
def test_normalize_content_blocks_mixed_blocks() -> None:
"""Mixed content lists: input_file converted, others unchanged."""
file_content = "file body"
b64 = base64.b64encode(file_content.encode()).decode()
blocks = [
{"type": "input_text", "text": "question"},
{"type": "input_file", "file_data": f"data:text/plain;base64,{b64}"},
]
result = _normalize_content_blocks_for_chat(blocks)
assert result == [
{"type": "input_text", "text": "question"},
{"type": "input_text", "text": "file body"},
], f"Unexpected result: {result!r}"
def test_normalize_responses_items_message_item_content_normalized() -> None:
"""``message`` items have their ``content`` list normalized."""
file_content = "uploaded file text"
b64 = base64.b64encode(file_content.encode()).decode()
items = [
{
"type": "message",
"role": "user",
"content": [
{"type": "input_text", "text": "what is this?"},
{"type": "input_file", "file_data": f"data:text/plain;base64,{b64}"},
],
}
]
result = _normalize_responses_items_for_chat(items)
assert len(result) == 1
assert result[0]["type"] == "message"
assert result[0]["content"] == [
{"type": "input_text", "text": "what is this?"},
{"type": "input_text", "text": "uploaded file text"},
], f"Unexpected normalized content: {result[0]['content']!r}"
def test_normalize_responses_items_non_message_items_pass_through() -> None:
"""Non-``message`` items (function_call, function_call_output) are unchanged."""
items = [
{"type": "function_call", "name": "my_tool", "arguments": "{}"},
{"type": "function_call_output", "call_id": "c1", "output": "ok"},
]
result = _normalize_responses_items_for_chat(items)
# Non-message items returned untouched — no normalization applied.
assert result == items
def test_normalize_responses_items_message_without_input_file_unchanged() -> None:
"""``message`` items with no file/metadata blocks are returned as-is."""
items = [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "just text"}],
}
]
result = _normalize_responses_items_for_chat(items)
# The content list object should be the same (identity) when unchanged.
assert result[0]["content"] is items[0]["content"], (
"Expected original content list object when no file/metadata blocks present"
)
def test_normalize_responses_items_message_image_filename_stripped() -> None:
"""Image attachment filenames are stripped from message payloads."""
items = [
{
"type": "message",
"role": "user",
"content": [
{
"type": "input_image",
"image_url": "data:image/png;base64,abcd",
"filename": "screenshot.png",
}
],
}
]
result = _normalize_responses_items_for_chat(items)
assert result[0] == {
"type": "message",
"role": "user",
"content": [{"type": "input_image", "image_url": "data:image/png;base64,abcd"}],
}
assert result[0] is not items[0]
def test_is_context_length_exceeded_direct_code() -> None:
"""
``_is_context_length_exceeded`` returns ``True`` for a direct
``BadRequestError``-like exception with ``code='context_length_exceeded'``.
What breaks if this fails: reactive compression never fires; the error
is treated as a permanent failure instead of triggering a retry.
"""
class _FakeBadRequest(Exception):
code = "context_length_exceeded"
assert _is_context_length_exceeded(_FakeBadRequest())
def test_is_context_length_exceeded_false_for_other_errors() -> None:
"""
``_is_context_length_exceeded`` returns ``False`` for unrelated errors.
What breaks if this fails: non-context errors are incorrectly treated
as compression triggers, masking the real failure.
"""
assert not _is_context_length_exceeded(ValueError("something else"))
assert not _is_context_length_exceeded(RuntimeError("network error"))
def test_context_length_exceeded_re_raises() -> None:
"""
``context_length_exceeded`` from the SDK propagates as an exception
rather than being swallowed into an ``ExecutorError``.
The runtime compaction layer handles context overflow; the executor
must re-raise so the exception reaches the ``ExecutorAdapter``'s
error classifier and surfaces as ``code="context_length_exceeded"``
on the harness wire.
What breaks if this fails: context overflow is swallowed inside
the executor and the runtime's reactive compaction never fires.
"""
class _CtxExceeded(Exception):
code = "context_length_exceeded"
async def _t() -> None:
executor = OpenAIAgentsSDKExecutor(client=object())
session_id = "s_ctx_reraise"
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(events=[], final_output="", exception=_CtxExceeded())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
with pytest.raises(_CtxExceeded):
await _collect(
executor.run_turn(
[{"role": "user", "content": "hi", "session_id": session_id}],
[],
"Be helpful.",
)
)
_run(_t())
# ── LLM_REQUEST policy evaluation wiring ─────────────────────────────────────
def test_policy_evaluator_deny_yields_executor_error() -> None:
"""
When ``_policy_evaluator`` returns ``POLICY_ACTION_DENY``, the
executor yields ``ExecutorError`` and returns immediately —
``Runner.run_streamed`` is never called.
What breaks if this fails: policy denials on the
``openai-agents`` harness silently pass through to the LLM
instead of blocking the call.
"""
async def _deny_evaluator(
phase: str,
data: dict[str, Any],
) -> Any:
"""Stub that always denies.
:param phase: The policy phase, e.g. ``"PHASE_LLM_REQUEST"``.
:param data: The request data dict.
:returns: A DENY verdict.
"""
return _FakeDenyVerdict(reason=f"denied by test: {data.get('model', '?')}")
@dataclass
class _FakeDenyVerdict:
"""Minimal verdict stub returning DENY.
:param action: Always ``"POLICY_ACTION_DENY"``.
:param reason: Human-readable reason.
:param data: Not used.
"""
action: str = "POLICY_ACTION_DENY"
reason: str = "denied"
data: dict[str, Any] | None = None
async def _t() -> None:
executor = OpenAIAgentsSDKExecutor(client=object())
executor._policy_evaluator = _deny_evaluator # type: ignore[attr-defined]
messages: list[dict[str, Any]] = [
{"role": "user", "content": "hi", "session_id": "test-session"},
]
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = await _collect(executor.run_turn(messages, [], "Be helpful."))
# Should yield exactly one ExecutorError, no text or tool events.
assert len(events) == 1, f"Expected 1 event, got {len(events)}: {events}"
assert isinstance(events[0], ExecutorError), (
f"Expected ExecutorError, got {type(events[0]).__name__}"
)
assert "denied by test" in events[0].message
_run(_t())
def test_policy_evaluator_allow_proceeds_to_run() -> None:
"""
When ``_policy_evaluator`` returns ``POLICY_ACTION_ALLOW``, the
executor proceeds normally — ``Runner.run_streamed`` is called
and events stream through.
What breaks if this fails: ALLOW verdicts are mistakenly
treated as denials, blocking all LLM calls.
"""
async def _allow_evaluator(
phase: str,
data: dict[str, Any],
) -> Any:
"""Stub that always allows.
:param phase: The policy phase.
:param data: The request data dict.
:returns: An ALLOW verdict.
"""
return _FakeAllowVerdict()
@dataclass
class _FakeAllowVerdict:
"""Minimal verdict stub returning ALLOW.
:param action: Always ``"POLICY_ACTION_ALLOW"``.
:param reason: ``None`` for ALLOW.
:param data: ``None`` for ALLOW.
"""
action: str = "POLICY_ACTION_ALLOW"
reason: str | None = None
data: dict[str, Any] | None = None
async def _t() -> None:
executor = OpenAIAgentsSDKExecutor(client=object())
executor._policy_evaluator = _allow_evaluator # type: ignore[attr-defined]
messages: list[dict[str, Any]] = [
{"role": "user", "content": "hi", "session_id": "test-session"},
]
_FakeRunner.next_result = _FakeResult(
events=[
_FakeRawEvent(_FakeRawTextDelta("hello")),
],
final_output="hello",
)
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = await _collect(executor.run_turn(messages, [], "Be helpful."))
# Should contain text events and a TurnComplete — not an ExecutorError.
error_events = [e for e in events if isinstance(e, ExecutorError)]
assert not error_events, f"Unexpected ExecutorError: {error_events}"
text_events = [e for e in events if isinstance(e, TextChunk)]
assert len(text_events) >= 1, "Expected at least one TextChunk"
assert text_events[0].text == "hello"
_run(_t())
# ── Cache-token normalization ───────────────────────────────────
#
# OpenAI's ``input_tokens`` (aka ``prompt_tokens``) is the TOTAL
# input including cached tokens. ``compute_llm_cost`` expects
# Anthropic semantics where ``input_tokens`` = non-cached only and
# ``cache_read_input_tokens`` is additive. The executor must
# subtract cached tokens and expose them separately.
def test_turn_usage_subtracts_cached_tokens_from_input() -> None:
"""
When ``prompt_tokens_details.cached_tokens`` is present, the
executor must subtract cached tokens from ``input_tokens`` and
report them as ``cache_read_input_tokens``.
Without this, ``compute_llm_cost`` bills cached tokens at the
full input rate instead of the cheaper cache-read rate.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
result = _FakeResult(events=[], final_output="hello")
result.raw_responses = [
_FakeRawResponse(
usage=_FakeUsage(
input_tokens=10000,
output_tokens=500,
total_tokens=10500,
prompt_tokens_details=_FakePromptTokensDetails(cached_tokens=8000),
)
)
]
_FakeRunner.next_result = result
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"",
)
]
turn_complete = next(e for e in events if isinstance(e, TurnComplete))
usage = turn_complete.usage
assert usage is not None
# Non-cached portion: 10000 - 8000 = 2000.
assert usage["input_tokens"] == 2000, (
"input_tokens must be the non-cached portion (total input minus cached)."
)
assert usage["cache_read_input_tokens"] == 8000, (
"cache_read_input_tokens must carry the cached token count."
)
assert usage["output_tokens"] == 500
# total_tokens is the billing sum and stays unchanged.
assert usage["total_tokens"] == 10500
_run(_t())
def test_turn_usage_no_cached_tokens_omits_cache_key() -> None:
"""
When no ``prompt_tokens_details`` is present, the usage dict
must NOT contain ``cache_read_input_tokens`` — the executor
degrades gracefully to the pre-cache behavior.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
result = _FakeResult(events=[], final_output="hello")
result.raw_responses = [
_FakeRawResponse(
usage=_FakeUsage(
input_tokens=5000,
output_tokens=300,
total_tokens=5300,
)
)
]
_FakeRunner.next_result = result
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"",
)
]
turn_complete = next(e for e in events if isinstance(e, TurnComplete))
usage = turn_complete.usage
assert usage is not None
assert usage["input_tokens"] == 5000
assert "cache_read_input_tokens" not in usage
_run(_t())
def test_turn_usage_cached_tokens_multi_call_sums_across_responses() -> None:
"""
Across multiple raw responses in a single turn, cached tokens
must be summed and subtracted from the total ``input_tokens``.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
result = _FakeResult(events=[], final_output="done")
result.raw_responses = [
_FakeRawResponse(
usage=_FakeUsage(
input_tokens=6000,
output_tokens=100,
total_tokens=6100,
prompt_tokens_details=_FakePromptTokensDetails(cached_tokens=4000),
)
),
_FakeRawResponse(
usage=_FakeUsage(
input_tokens=7000,
output_tokens=200,
total_tokens=7200,
prompt_tokens_details=_FakePromptTokensDetails(cached_tokens=5000),
)
),
]
_FakeRunner.next_result = result
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "search", "session_id": "s1"}],
[],
"",
)
]
turn_complete = next(e for e in events if isinstance(e, TurnComplete))
usage = turn_complete.usage
assert usage is not None
# Sum of input: 6000 + 7000 = 13000.
# Sum of cached: 4000 + 5000 = 9000.
# Non-cached: 13000 - 9000 = 4000.
assert usage["input_tokens"] == 4000
assert usage["cache_read_input_tokens"] == 9000
_run(_t())
# ── Empty-turn retry / fail-loud ────────────────────────────────────
# The Databricks gateway occasionally returns a completed turn with no
# text, no tool calls, and no output items. ``run_turn`` retries such a
# turn once (``_EMPTY_TURN_MAX_ATTEMPTS``) and, if still empty AND the
# gateway billed zero output tokens, surfaces a loud retryable
# ``ExecutorError`` instead of a silent empty ``TurnComplete``.
@dataclass
class _FakeReasoningItem:
"""A run item the SDK emits for reasoning models. Bookkeeping, not
user-visible output, so it must not count toward the emptiness check.
:param type: The SDK discriminator, always ``"reasoning_item"``.
"""
type: str = "reasoning_item"
@dataclass
class _FakeMessageOutputItem:
"""A run item carrying assistant text. Counts as output.
:param text: The assistant text the item carries, e.g. ``"hello"``.
:param type: The SDK discriminator, always ``"message_output_item"``.
"""
text: str = ""
type: str = "message_output_item"
def _empty_raw_response() -> _FakeRawResponse:
"""A raw model response reporting zero output tokens (gateway hiccup)."""
return _FakeRawResponse(usage=_FakeUsage(input_tokens=10, output_tokens=0))
def _nonempty_raw_response() -> _FakeRawResponse:
"""A raw model response reporting one output token (model actually ran)."""
return _FakeRawResponse(usage=_FakeUsage(input_tokens=10, output_tokens=1))
def _make_databricks_executor() -> OpenAIAgentsSDKExecutor:
"""An executor whose client points at the Databricks gateway base URL.
Matches the production path the empty-output bug occurs on.
"""
client = types.SimpleNamespace(
base_url="https://profile-host.example.com/ai-gateway/openai/v1"
)
return OpenAIAgentsSDKExecutor(client=client)
def test_empty_turn_retries_then_succeeds() -> None:
"""
An empty first attempt is retried; the second attempt's text is the
only thing surfaced (no double-emit), and ``run_streamed`` runs twice.
What breaks if this fails: a transient empty gateway response would
surface as a silent empty turn instead of being retried.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
_FakeRunner.next_result = None
_FakeRunner.next_results = [
# Attempt 1: completely empty (no events, no items, 0 out tokens).
_FakeResult(
events=[],
final_output="",
new_items=[],
raw_responses=[_empty_raw_response()],
),
# Attempt 2: real text.
_FakeResult(
events=[_FakeRawEvent(_FakeRawTextDelta("recovered"))],
final_output="recovered",
raw_responses=[_nonempty_raw_response()],
),
]
executor = _make_databricks_executor()
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = await _collect(
executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
)
# Two run_streamed calls = one retry. If 1, the empty turn was
# not detected; if 3+, the loop over-retried.
assert len(_FakeRunner.last_calls) == 2, (
f"Expected 2 run_streamed calls (1 retry), got {len(_FakeRunner.last_calls)}"
)
# Exactly one TextChunk("recovered") + one TurnComplete. The empty
# attempt streamed nothing, so there is no stray empty TextChunk
# and no double-emit of the recovered text.
text_chunks = [e for e in events if isinstance(e, TextChunk)]
assert [c.text for c in text_chunks] == ["recovered"], (
f"Expected only the recovered text, got {[c.text for c in text_chunks]!r}"
)
turn_completes = [e for e in events if isinstance(e, TurnComplete)]
# Exactly one TurnComplete for the whole turn. If 2, the empty
# first attempt also emitted a (spurious) TurnComplete.
assert len(turn_completes) == 1
assert turn_completes[0].response == "recovered"
assert not any(isinstance(e, ExecutorError) for e in events)
_run(_t())
def test_empty_turn_retry_exhausted_yields_retryable_error() -> None:
"""
When every attempt is empty and the gateway billed zero output
tokens, the executor yields a loud ``ExecutorError(retryable=True)``
and no ``TurnComplete``.
What breaks if this fails: a persistent gateway hiccup would surface
as a silent completed-but-empty turn that callers cannot distinguish
from a real answer.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
_FakeRunner.next_result = None
_FakeRunner.next_results = [
_FakeResult(events=[], final_output="", raw_responses=[_empty_raw_response()]),
_FakeResult(events=[], final_output="", raw_responses=[_empty_raw_response()]),
]
executor = _make_databricks_executor()
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = await _collect(
executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
)
# Both attempts ran (initial + 1 retry) before giving up.
assert len(_FakeRunner.last_calls) == 2, (
f"Expected 2 run_streamed calls before fail-loud, got {len(_FakeRunner.last_calls)}"
)
errors = [e for e in events if isinstance(e, ExecutorError)]
assert len(errors) == 1, f"Expected exactly 1 ExecutorError, got {events!r}"
# retryable=True so the workflow raises RetryableLLMError and the
# surrounding retry policy can reissue the turn. If False, the
# workflow would treat it as a permanent harness bug and not retry.
assert errors[0].retryable is True
# Message names the empty-completion condition (matches the
# production logger.error text) so the failure is diagnosable.
assert "empty completion" in errors[0].message
# No silent TurnComplete masking the failure.
assert not any(isinstance(e, TurnComplete) for e in events)
_run(_t())
def test_tool_call_without_text_is_not_retried() -> None:
"""
A turn that called a tool but produced no final text is NOT empty —
tool activity is legitimate output, so it runs once and completes.
What breaks if this fails: tool-only turns (e.g. parallel sub-agent
spawns that emit no assistant prose) would be wrongly retried.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
_FakeRunner.next_result = None
_FakeRunner.next_results = [
_FakeResult(
events=[
_FakeRunItemEvent(
_FakeToolCallItem(_FakeToolCallRawItem("add", '{"a": 1, "b": 2}'))
),
_FakeRunItemEvent(
_FakeToolOutputItem(_FakeToolOutputRawItem(), {"result": 3})
),
],
final_output="",
raw_responses=[_nonempty_raw_response()],
),
]
executor = _make_databricks_executor()
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = await _collect(
executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[{"name": "add", "description": "Add", "parameters": {"type": "object"}}],
"Be helpful.",
)
)
# Single attempt: tool activity means the turn is not empty.
assert len(_FakeRunner.last_calls) == 1, (
f"Expected 1 run_streamed call (no retry), got {len(_FakeRunner.last_calls)}"
)
assert any(isinstance(e, ToolCallRequest) for e in events)
turn_completes = [e for e in events if isinstance(e, TurnComplete)]
# Tool-only turn completes normally (one TurnComplete), not errored.
assert len(turn_completes) == 1
assert not any(isinstance(e, ExecutorError) for e in events)
_run(_t())
def test_reasoning_only_turn_is_treated_as_empty() -> None:
"""
A turn whose only new item is a reasoning item (no text, no tools)
is treated as empty and retried.
What breaks if this fails: reasoning-model ghost turns would count
their reasoning bookkeeping as output and never retry.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
_FakeRunner.next_result = None
_FakeRunner.next_results = [
_FakeResult(
events=[],
final_output="",
new_items=[_FakeReasoningItem()],
raw_responses=[_empty_raw_response()],
),
_FakeResult(
events=[_FakeRawEvent(_FakeRawTextDelta("answer"))],
final_output="answer",
raw_responses=[_nonempty_raw_response()],
),
]
executor = _make_databricks_executor()
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = await _collect(
executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
)
# Retry fired: the reasoning-only item did not count as output.
assert len(_FakeRunner.last_calls) == 2, (
f"Expected 2 run_streamed calls (reasoning-only retried), "
f"got {len(_FakeRunner.last_calls)}"
)
turn_completes = [e for e in events if isinstance(e, TurnComplete)]
# One TurnComplete, carrying the recovered text from attempt 2.
assert len(turn_completes) == 1
assert turn_completes[0].response == "answer"
_run(_t())
def test_empty_turn_with_output_tokens_is_not_errored() -> None:
"""
A final empty turn that DID bill output tokens is a deliberate empty
answer, not a gateway hiccup: it completes silently as before, with
no ``ExecutorError``.
What breaks if this fails: a model that legitimately answers with an
empty string would be turned into a spurious retryable error.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
_FakeRunner.next_result = None
# Both attempts empty of text/items, but output_tokens > 0 means
# the model ran and chose to emit nothing.
_FakeRunner.next_results = [
_FakeResult(events=[], final_output="", raw_responses=[_nonempty_raw_response()]),
_FakeResult(events=[], final_output="", raw_responses=[_nonempty_raw_response()]),
]
executor = _make_databricks_executor()
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = await _collect(
executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
)
# Retried once (both empty), then completed silently — the
# output-token gate suppresses the fail-loud error.
assert not any(isinstance(e, ExecutorError) for e in events), (
f"Empty-but-billed turn should NOT error, got {events!r}"
)
turn_completes = [e for e in events if isinstance(e, TurnComplete)]
# Completes silently with one empty TurnComplete (today's behavior
# for a deliberate empty answer), not an ExecutorError.
assert len(turn_completes) == 1
assert turn_completes[0].response == ""
_run(_t())
def test_empty_turn_retry_rewinds_sdk_session() -> None:
"""
Before retrying an empty turn, the SDK session is rewound to its
pre-turn item count so the retry re-runs from an identical state.
What breaks if this fails: a stray empty assistant item appended by
the first attempt would accumulate, so the retry would re-run with a
polluted session.
"""
async def _t() -> None:
_FakeRunner.last_calls = []
_FakeRunner.next_result = None
# Attempt 1 appends a stray item to the session, mimicking the SDK
# persisting an empty assistant turn before we detect emptiness.
appended = {"stray": "empty-assistant-item"}
class _AppendingEmptyResult(_FakeResult):
def __init__(self, session):
super().__init__(
events=[],
final_output="",
raw_responses=[_empty_raw_response()],
)
self._session = session
async def stream_events(self):
await self._session.add_items([appended])
for event in self._events:
yield event
captured: dict[str, object] = {}
def _runner(agent, input, session, max_turns, run_config):
_FakeRunner.last_calls.append({"session": session})
captured["session"] = session
if len(_FakeRunner.last_calls) == 1:
return _AppendingEmptyResult(session)
return _FakeResult(
events=[_FakeRawEvent(_FakeRawTextDelta("ok"))],
final_output="ok",
raw_responses=[_nonempty_raw_response()],
)
fake_sdk = _fake_agents_sdk()
fake_sdk.Runner = types.SimpleNamespace(run_streamed=_runner)
executor = _make_databricks_executor()
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=fake_sdk,
):
events = await _collect(
executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
)
# state.sdk_session may be wrapped by OpenAIResponsesCompactionSession
# (which has .underlying_session) around _SanitizingSession (which has
# ._underlying). Traverse both layers to reach the fake.
_session = captured["session"]
if hasattr(_session, "underlying_session"):
_session = _session.underlying_session
underlying = _session._underlying
# The stray item appended in attempt 1 was popped before attempt 2,
# so the session is back to its pre-turn (empty) state. pop_calls
# records the rewind. If 0, the rewind never happened and the
# retry would have re-run with a polluted session.
assert underlying.pop_calls == 1, (
f"Expected exactly 1 pop (rewind of the stray item), got {underlying.pop_calls}"
)
assert appended not in underlying.items
turn_completes = [e for e in events if isinstance(e, TurnComplete)]
# One TurnComplete: the rewound retry produced the recovered text.
assert len(turn_completes) == 1
assert turn_completes[0].response == "ok"
_run(_t())
# ---------------------------------------------------------------------------
# Tests: Compaction
# ---------------------------------------------------------------------------
@dataclass
class _FakeCompactionItem:
"""Stand-in for agents.items.CompactionItem."""
type: str = "compaction_item"
def test_compaction_item_emits_compaction_complete() -> None:
"""When a compaction_item appears in result.new_items, a CompactionComplete
event is yielded before TurnComplete."""
from omnigent.inner.executor import CompactionComplete
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[],
final_output="compacted",
new_items=[_FakeCompactionItem()],
raw_responses=[
_FakeRawResponse(_FakeUsage(input_tokens=100, output_tokens=50, total_tokens=150))
],
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
]
compaction_events = [e for e in events if isinstance(e, CompactionComplete)]
assert len(compaction_events) == 1
assert compaction_events[0].token_count == 150 # context_tokens = last total
# compacted_messages should contain the session items
assert compaction_events[0].compacted_messages is not None
turn_completes = [e for e in events if isinstance(e, TurnComplete)]
assert len(turn_completes) == 1
# CompactionComplete must come before TurnComplete
ci = events.index(compaction_events[0])
ti = events.index(turn_completes[0])
assert ci < ti
_run(_t())
def test_no_compaction_item_no_compaction_event() -> None:
"""When no compaction_item is in new_items, no CompactionComplete is yielded."""
from omnigent.inner.executor import CompactionComplete
async def _t():
_FakeRunner.last_calls = []
_FakeRunner.next_result = _FakeResult(
events=[],
final_output="normal",
)
executor = OpenAIAgentsSDKExecutor(client=object())
with patch(
"omnigent.inner.openai_agents_sdk_executor._ensure_agents_sdk",
return_value=_fake_agents_sdk(),
):
events = [
e
async for e in executor.run_turn(
[{"role": "user", "content": "hi", "session_id": "s1"}],
[],
"Be helpful.",
)
]
compaction_events = [e for e in events if isinstance(e, CompactionComplete)]
assert len(compaction_events) == 0
_run(_t())