Files

1905 lines
70 KiB
Python

from __future__ import annotations
import json
import sys
import types
from collections.abc import AsyncIterator, Callable, Iterator
from dataclasses import dataclass
from io import StringIO
from pathlib import Path
from typing import Any
import anyio
import pytest
import sniffio
from pytest import CaptureFixture
from pytest_mock import MockerFixture
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.text import Text
from pydantic_ai import Agent, ModelMessage, ModelResponse, ModelRetry, TextPart, ToolCallPart, __version__, _display
from pydantic_ai.agent import WrapperAgent
from pydantic_ai.capabilities import NativeTool
from pydantic_ai.messages import RetryPromptPart, ToolReturnPart
from pydantic_ai.models.test import TestModel
from pydantic_ai.settings import ModelSettings
from pydantic_ai.toolsets import FunctionToolset, WrapperToolset
from pydantic_ai.usage import RunUsage, UsageLimits
from ._inline_snapshot import snapshot
from .conftest import IsInstance, IsStr, TestEnv, try_import
with try_import() as imports_successful:
from prompt_toolkit.buffer import Buffer
from prompt_toolkit.document import Document
from prompt_toolkit.input import create_pipe_input
from prompt_toolkit.output import DummyOutput
from prompt_toolkit.shortcuts import PromptSession
from pydantic_ai._cli import (
CustomAutoSuggest,
ask_agent,
cli,
cli_agent,
format_usage,
handle_slash_command,
run_chat,
)
from pydantic_ai._cli.web import run_web_command
from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
from pydantic_ai.models.openai import OpenAIChatModel
pytestmark = pytest.mark.skipif(not imports_successful(), reason='install cli extras to run cli tests')
@pytest.fixture
def blockbuster_excluded_modules() -> tuple[str, ...]:
"""The CLI owns intentionally synchronous terminal and history operations."""
return ('pydantic_ai._cli',)
@pytest.fixture(autouse=True)
def reset_sniffio_cvar() -> Iterator[None]:
# The anyio pytest plugin sets `current_async_library_cvar` to 'asyncio' at session
# start and the value leaks into sync tests, causing `anyio.run` to refuse to start.
token = sniffio.current_async_library_cvar.set(None)
try:
yield
finally:
sniffio.current_async_library_cvar.reset(token)
def test_cli_version(capfd: CaptureFixture[str]):
assert cli(['--version']) == 0
assert capfd.readouterr().out.startswith('clai - Pydantic AI CLI')
def test_invalid_model(capfd: CaptureFixture[str]):
assert cli(['--model', 'potato']) == 1
assert capfd.readouterr().out.splitlines() == snapshot(['Error initializing potato:', 'Unknown model: potato'])
def test_invalid_model_suggestion(capfd: CaptureFixture[str]):
assert cli(['--model', 'claude:sonnet-5']) == 1
assert capfd.readouterr().out.splitlines() == snapshot(
[
'Error initializing claude:sonnet-5:',
"Unknown model: claude:sonnet-5. Did you mean 'anthropic:claude-sonnet-5'?",
]
)
@pytest.fixture
def create_test_module():
def _create_test_module(**namespace: Any) -> None:
assert 'test_module' not in sys.modules
test_module = types.ModuleType('test_module')
for key, value in namespace.items():
setattr(test_module, key, value)
sys.modules['test_module'] = test_module
try:
yield _create_test_module
finally:
if 'test_module' in sys.modules: # pragma: no branch
del sys.modules['test_module']
def test_agent_flag(
capfd: CaptureFixture[str],
mocker: MockerFixture,
env: TestEnv,
create_test_module: Callable[..., None],
):
env.remove('OPENAI_API_KEY')
env.set('COLUMNS', '150')
test_agent = Agent(TestModel(custom_output_text='Hello from custom agent'))
create_test_module(custom_agent=test_agent)
# Mock ask_agent to avoid actual execution but capture the agent
mock_ask = mocker.patch('pydantic_ai._cli.ask_agent')
# Test CLI with custom agent
assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
# Verify the output contains the custom agent message
assert 'using custom agent test_module:custom_agent' in capfd.readouterr().out.replace('\n', '')
# Verify ask_agent was called with our custom agent
mock_ask.assert_called_once()
assert mock_ask.call_args[0][0] is test_agent
def test_agent_flag_no_model(capfd: CaptureFixture[str], env: TestEnv, create_test_module: Callable[..., None]):
env.remove('OPENAI_API_KEY')
env.remove('OPENAI_ADMIN_KEY')
test_agent = Agent()
create_test_module(custom_agent=test_agent)
# The missing key now surfaces as a `UserError`, which the CLI reports cleanly instead of
# letting the provider SDK's exception escape as a traceback.
assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 1
output = capfd.readouterr().out
assert 'Error initializing' in output
assert 'Set the `OPENAI_API_KEY` environment variable' in output
def test_agent_flag_set_model(
capfd: CaptureFixture[str],
mocker: MockerFixture,
env: TestEnv,
create_test_module: Callable[..., None],
):
env.set('OPENAI_API_KEY', 'xxx')
env.set('COLUMNS', '150')
custom_agent = Agent(TestModel(custom_output_text='Hello from custom agent'))
create_test_module(custom_agent=custom_agent)
mocker.patch('pydantic_ai._cli.ask_agent')
assert cli(['--agent', 'test_module:custom_agent', '--model', 'openai-chat:gpt-4o', 'hello']) == 0
# CLI banner shows `model.model_id` which is `{system}:{model_name}` — `OpenAIChatModel.system`
# is `'openai'`, so the banner prints `'openai:gpt-4o'` even when the user passed `'openai-chat:'`.
assert 'using custom agent test_module:custom_agent with openai:gpt-4o' in capfd.readouterr().out.replace('\n', '')
assert isinstance(custom_agent.model, OpenAIChatModel)
def test_agent_flag_non_agent(
capfd: CaptureFixture[str], mocker: MockerFixture, create_test_module: Callable[..., None]
):
test_agent = 'Not an Agent object'
create_test_module(custom_agent=test_agent)
assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 1
assert 'Could not load agent from test_module:custom_agent' in capfd.readouterr().out
def test_agent_flag_bad_module_variable_path(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv):
assert cli(['--agent', 'bad_path', 'hello']) == 1
assert 'Could not load agent from bad_path' in capfd.readouterr().out
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio from an MCP config file')
@pytest.mark.parametrize('stream', [False, True])
def test_mcp_config(capfd: CaptureFixture[str], env: TestEnv, tmp_path: Path, stream: bool):
"""`--mcp-config` parses a Claude-Desktop-style config and connects to the MCP server.
Drives the real `load_mcp_toolsets` -> `MCPToolset` path against `tests.mcp_server` over stdio,
without mocking the loader under test. `TestModel` then calls the prefixed MCP tool through
both the streaming and non-streaming paths.
"""
env.set('OPENAI_API_KEY', 'test')
config_file = tmp_path / 'mcp_servers.json'
config_file.write_text(
json.dumps({'mcpServers': {'temp': {'command': 'python', 'args': ['-m', 'tests.mcp_server']}}})
)
args = ['--mcp-config', str(config_file), 'weather in Mexico City?']
if not stream:
args.insert(0, '--no-stream')
with cli_agent.override(model=TestModel(call_tools=['temp_get_weather_forecast'])):
assert cli(args) == 0
output = capfd.readouterr().out
assert 'temp_get_weather_forecast' in output
assert 'The weather in a is sunny and 26 degrees Celsius.' in ' '.join(output.split())
assert ('Called tool temp_get_weather_forecast' in output) is stream
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio from an MCP config file')
def test_mcp_config_interactive(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv, tmp_path: Path):
"""Interactive `clai --mcp-config` can invoke a configured server tool before exiting."""
env.set('OPENAI_API_KEY', 'test')
config_file = tmp_path / 'mcp_servers.json'
config_file.write_text(
json.dumps({'mcpServers': {'temp': {'command': 'python', 'args': ['-m', 'tests.mcp_server']}}})
)
with create_pipe_input() as inp:
inp.send_text('weather in Mexico City?\n')
inp.send_text('/exit\n')
session = PromptSession[Any](input=inp, output=DummyOutput())
mocker.patch('pydantic_ai._cli.PromptSession', return_value=session)
with cli_agent.override(model=TestModel(call_tools=['temp_get_weather_forecast'])):
assert cli(['--mcp-config', str(config_file)]) == 0
output = capfd.readouterr().out
assert 'Called tool temp_get_weather_forecast' in output
assert 'The weather in a is sunny and 26 degrees Celsius.' in ' '.join(output.split())
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio from an MCP config file')
def test_mcp_config_banner_omits_the_tool_count(
capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv, tmp_path: Path
):
"""`tools: 0` beside a session full of MCP tools would be worse than saying nothing at all.
Startup can't count them: only a `FunctionToolset` holds its tools synchronously, and an MCP
server would have to be connected to and asked. The count is left out rather than understated.
"""
env.set('OPENAI_API_KEY', 'test')
config_file = tmp_path / 'mcp_servers.json'
config_file.write_text(
json.dumps({'mcpServers': {'temp': {'command': 'python', 'args': ['-m', 'tests.mcp_server']}}})
)
with create_pipe_input() as inp:
inp.send_text('/exit\n')
session = PromptSession[Any](input=inp, output=DummyOutput())
mocker.patch('pydantic_ai._cli.PromptSession', return_value=session)
mocker.patch('pydantic_ai._display.banner_available', return_value=True)
with cli_agent.override(model=TestModel()):
assert cli(['--mcp-config', str(config_file)]) == 0
output = capfd.readouterr().out
assert 'capabilities:' in output, 'the banner should still be shown'
assert 'tools:' not in output
# Sentinel for the case where `--mcp-config` is handed an existing path that isn't a readable file.
DIRECTORY_CONFIG = 'directory-instead-of-file'
@pytest.mark.parametrize(
'config,expected',
[
pytest.param(None, 'not found', id='missing-file'),
pytest.param('{ not valid json ', 'key must be a string', id='malformed-json'),
pytest.param(
json.dumps({'mcpServers': {'x': {'command': 'echo', 'args': ['${UNDEFINED_VAR_XYZ}']}}}),
'is not defined',
id='undefined-env-var',
),
pytest.param(
json.dumps({'mcpServers': {'x': {}}}),
'must have either `command` or `url`',
id='missing-command-or-url',
),
pytest.param(
json.dumps({'mcpServers': {'x': None}}),
'Input should be a valid dictionary',
id='non-object-server-entry',
),
pytest.param(
json.dumps({'mcpServers': {'x': {'command': 'echo', 'args': 'not-a-list'}}}),
'Input should be a valid list',
id='wrongly-typed-field',
),
pytest.param(
DIRECTORY_CONFIG,
'Is a directory',
id='directory-not-a-file',
),
],
)
def test_mcp_config_errors(capfd: CaptureFixture[str], env: TestEnv, tmp_path: Path, config: str | None, expected: str):
"""A bad `--mcp-config` prints a friendly error and exits 1 instead of a raw traceback."""
env.set('OPENAI_API_KEY', 'test')
config_file = tmp_path / 'mcp_servers.json'
if config is None:
target = tmp_path / 'does_not_exist.json'
elif config == DIRECTORY_CONFIG:
target = tmp_path
else:
config_file.write_text(config)
target = config_file
assert cli(['--mcp-config', str(target), 'hi']) == 1
out = capfd.readouterr().out
assert 'Could not load MCP config' in out
assert expected in out
def test_mcp_config_empty_value(capfd: CaptureFixture[str], env: TestEnv):
"""`--mcp-config=` errors instead of silently starting without MCP servers.
An empty value is falsy, so a truthiness check would skip MCP entirely — which is what
`--mcp-config="$UNSET_VAR"` expands to in a script.
"""
env.set('OPENAI_API_KEY', 'test')
assert cli(['--mcp-config=', 'hi']) == 1
assert 'needs a path to a configuration file' in capfd.readouterr().out
def test_no_command_defaults_to_chat(mocker: MockerFixture):
"""Test that running clai with no command defaults to chat mode."""
# Mock _run_chat_command to avoid actual execution
mock_run_chat = mocker.patch('pydantic_ai._cli._run_chat_command', return_value=0)
result = cli([])
assert result == 0
mock_run_chat.assert_called_once()
def test_list_models(capfd: CaptureFixture[str]):
assert cli(['--list-models']) == 0
output = capfd.readouterr().out.splitlines()
assert output[:3] == snapshot([IsStr(regex='clai - Pydantic AI CLI .*'), '', 'Available models:'])
providers = (
'openai',
'anthropic',
'bedrock',
'cerebras',
'crusoe',
'google',
'google-cloud',
'groq',
'mistral',
'cohere',
'deepseek',
'gateway/',
'heroku',
'moonshotai',
'xai',
'huggingface',
'zai',
'snowflake',
'typesafe',
)
models = {line.strip().split(' ')[0] for line in output[3:]}
for provider in providers:
models = models - {model for model in models if model.startswith(provider)}
assert models == set(), models
def test_cli_prompt(capfd: CaptureFixture[str], env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
with cli_agent.override(model=TestModel(custom_output_text='# result\n\n```py\nx = 1\n```')):
assert cli(['hello']) == 0
assert capfd.readouterr().out.splitlines() == snapshot([IsStr(), '# result', '', 'py', 'x = 1', '/py'])
assert cli(['--no-stream', 'hello']) == 0
assert capfd.readouterr().out.splitlines() == snapshot([IsStr(), '# result', '', 'py', 'x = 1', '/py'])
async def test_streaming_with_tool_calls():
"""The streaming CLI render loop interleaves streamed model text with tool-call indicators.
Uses a `FunctionModel` stream so the agent emits real text deltas and a tool call, exercising
`ask_agent`'s render loop end to end rather than `TestModel`'s canned output.
"""
async def weather_stream(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | DeltaToolCalls]:
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
yield 'It is '
yield 'sunny in Mexico City.'
else:
yield 'Let me check '
yield 'the weather.'
yield {0: DeltaToolCall(name='get_weather', json_args='{"city": "Mexico City"}', tool_call_id='call_1')}
agent = Agent(FunctionModel(stream_function=weather_stream))
@agent.tool_plain
def get_weather(city: str) -> str:
return f'sunny in {city}'
output = StringIO()
console = Console(file=output, force_terminal=False, width=80)
messages = await ask_agent(agent, 'weather?', stream=True, console=console, code_theme='monokai')
assert output.getvalue() == snapshot("""\
Let me check the weather. \n\
▌ Called tool get_weather. \n\
It is sunny in Mexico City. \
""")
assert isinstance(messages[-1], ModelResponse)
assert messages[-1].parts[-1] == TextPart(content='It is sunny in Mexico City.')
def _distinct_frames(frames: list[str]) -> list[str]:
"""Drop consecutive duplicates — the same frame re-rendered carries no information."""
return [frame for i, frame in enumerate(frames) if i == 0 or frame != frames[i - 1]]
@pytest.fixture
def live_frames(monkeypatch: pytest.MonkeyPatch) -> list[str]:
"""Record the markdown of every `Live.update`, so intermediate render frames can be asserted.
`ask_agent` only ever updates the display with a `Markdown`, and the final console output shows
just the last frame — these tests are about what the display showed along the way.
"""
frames: list[str] = []
original_update = Live.update
def capture_update(self: Live, renderable: Markdown, *, refresh: bool = False) -> None:
frames.append(renderable.markup)
return original_update(self, renderable, refresh=refresh)
monkeypatch.setattr(Live, 'update', capture_update)
return frames
async def test_streaming_with_concurrent_tool_calls(live_frames: list[str]):
"""Every in-flight tool call keeps its own indicator until its own result arrives.
Two tools called in one response run concurrently, so the render loop keys them by
`tool_call_id`. Rendering only the most recent call erased the earlier one's indicator, leaving
a still-running tool with no line at all. The final console output can't show that — the
regression lives in the intermediate `Live` frames, so those are captured here.
Asserted as order-independent invariants rather than a frame snapshot: the two results are
emitted by concurrent tasks, so their order is not stable and a snapshot of the sequence would
flake. `get_temp` waits on `get_weather` only to guarantee the two calls overlap.
"""
weather_returned = anyio.Event()
async def two_city_stream(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | DeltaToolCalls]:
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
yield 'Both cities checked.'
else:
yield 'Checking two cities.'
yield {
0: DeltaToolCall(name='get_weather', json_args='{"city": "Lisbon"}', tool_call_id='call_1'),
1: DeltaToolCall(name='get_temp', json_args='{"city": "Porto"}', tool_call_id='call_2'),
}
agent = Agent(FunctionModel(stream_function=two_city_stream))
@agent.tool_plain
async def get_weather(city: str) -> str:
weather_returned.set()
return f'sunny in {city}'
@agent.tool_plain
async def get_temp(city: str) -> str:
await weather_returned.wait()
return f'20C in {city}'
console = Console(file=StringIO(), force_terminal=False, width=80)
await ask_agent(agent, 'weather?', stream=True, console=console, code_theme='monokai')
distinct = _distinct_frames(live_frames)
# The regression: the second call's indicator replaced the first's, so this frame never existed.
assert any('_Calling tool `get_weather`…_' in frame and '_Calling tool `get_temp`…_' in frame for frame in distinct)
# And once one call finished, the other was left with no indicator at all.
assert any('Called tool `' in frame and '_Calling tool `' in frame for frame in distinct)
# The final frame holds both completions; their order follows task completion, so assert
# membership rather than a sequence.
final = distinct[-1]
assert final.startswith('Checking two cities.')
assert '> Called tool `get_weather`.' in final
assert '> Called tool `get_temp`.' in final
assert '_Calling tool' not in final
assert final.endswith('Both cities checked.')
@dataclass
class _City:
name: str
async def test_streaming_ignores_output_tool_calls():
"""The internal output tool is not reported as a tool call.
A structured `output_type` makes the agent call an output tool, which streams
`OutputToolCallEvent` / `OutputToolResultEvent` through the same node as function tools. Those
are the model producing its answer rather than the agent using a tool, so they get no indicator.
"""
agent = Agent(TestModel(), output_type=_City)
output = StringIO()
console = Console(file=output, force_terminal=False, width=80)
await ask_agent(agent, 'name a city', stream=True, console=console, code_theme='monokai')
rendered = output.getvalue()
assert 'Calling tool' not in rendered
assert 'Called tool' not in rendered
async def test_streaming_clears_indicator_for_retried_tool(live_frames: list[str]):
"""A call that comes back as a retry drops its in-flight indicator instead of pinning it.
`pending_calls` is popped for any `FunctionToolResultEvent`, not only a `ToolReturnPart`.
Popping on success alone left `> _Calling tool ...` on screen for the rest of the run. A retry
still renders no `Called tool` line — surfacing retries is a separate feature.
"""
async def retrying_tool_stream(
messages: list[ModelMessage], info: AgentInfo
) -> AsyncIterator[str | DeltaToolCalls]:
if any(isinstance(part, RetryPromptPart) for message in messages for part in message.parts):
yield 'Recovered without the tool.'
else:
yield 'Trying a tool.'
yield {0: DeltaToolCall(name='flaky', json_args='{}', tool_call_id='call_1')}
agent = Agent(FunctionModel(stream_function=retrying_tool_stream))
@agent.tool_plain
def flaky() -> str:
raise ModelRetry('not this time')
console = Console(file=StringIO(), force_terminal=False, width=80)
await ask_agent(agent, 'go', stream=True, console=console, code_theme='monokai')
distinct = _distinct_frames(live_frames)
assert distinct == snapshot(
[
'Trying a tool.',
"""\
Trying a tool.
> _Calling tool `flaky`…_\
""",
'Trying a tool.',
"""\
Trying a tool.
Recovered without the tool.\
""",
]
)
@dataclass
class _LifetimeToolset(WrapperToolset[Any]):
"""Counts how often it is fully released, to observe whether it survives between turns."""
depth: int = 0
full_releases: int = 0
async def __aenter__(self) -> _LifetimeToolset:
self.depth += 1
await super().__aenter__()
return self
async def __aexit__(self, *args: Any) -> bool | None:
result = await super().__aexit__(*args)
self.depth -= 1
if self.depth == 0:
self.full_releases += 1
return result
async def test_chat_holds_toolsets_open_for_the_session(mocker: MockerFixture, env: TestEnv, tmp_path: Path):
"""Run-scoped toolsets survive between turns instead of being torn down after each one.
Passing them straight to each `agent.run()` fully released them every turn — for an MCP server
that means a fresh subprocess and `tools/list` handshake per message, and server-side session
state lost in between. `run_chat` now holds them open around the REPL loop, so each turn's own
enter is a ref-count bump and the toolset is released once, at the end.
"""
env.set('OPENAI_API_KEY', 'test')
toolset = _LifetimeToolset(FunctionToolset[Any]())
with create_pipe_input() as inp:
inp.send_text('hello\n')
inp.send_text('hello again\n')
inp.send_text('/exit\n')
session = PromptSession[Any](input=inp, output=DummyOutput())
mocker.patch('pydantic_ai._cli.PromptSession', return_value=session)
agent = Agent(TestModel(custom_output_text='goodbye'))
console = Console(file=StringIO(), force_terminal=False, width=80)
assert await run_chat(True, agent, console, 'monokai', 'clai', config_dir=tmp_path, toolsets=[toolset]) == 0
# Two turns ran, but the toolset was released once — at the end of the session, not per turn.
assert toolset.full_releases == snapshot(1)
async def test_chat_keeps_toolsets_open_after_failed_turn(mocker: MockerFixture, env: TestEnv, tmp_path: Path):
"""A failed turn is reported without releasing session toolsets or ending the REPL."""
env.set('OPENAI_API_KEY', 'test')
toolset = _LifetimeToolset(FunctionToolset[Any]())
attempts = 0
def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
nonlocal attempts
attempts += 1
if attempts == 1:
raise RuntimeError('first turn failed') from ValueError('underlying failure')
if attempts == 2:
raise RuntimeError('second turn failed')
return ModelResponse(parts=[TextPart('second turn succeeded')])
with create_pipe_input() as inp:
inp.send_text('fail\n')
inp.send_text('fail again\n')
inp.send_text('succeed\n')
inp.send_text('/exit\n')
session = PromptSession[Any](input=inp, output=DummyOutput())
mocker.patch('pydantic_ai._cli.PromptSession', return_value=session)
output = StringIO()
console = Console(file=output, force_terminal=False, width=80)
agent = Agent(FunctionModel(function=respond))
assert await run_chat(False, agent, console, 'monokai', 'clai', config_dir=tmp_path, toolsets=[toolset]) == 0
rendered = output.getvalue()
assert 'RuntimeError: first turn failed' in rendered
assert 'Caused by: underlying failure' in rendered
assert 'RuntimeError: second turn failed' in rendered
assert 'second turn succeeded' in rendered
assert attempts == snapshot(3)
assert toolset.full_releases == snapshot(1)
def test_custom_auto_suggest_completes_special_commands():
"""Typing a prefix of a slash command suggests the rest; anything else falls back to history."""
suggest = CustomAutoSuggest(['/exit'])
buffer = Buffer()
suggestion = suggest.get_suggestion(buffer, Document('/ex'))
assert suggestion is not None
assert suggestion.text == snapshot('it')
# No special command matches, and an empty history has nothing to offer either.
assert suggest.get_suggestion(buffer, Document('hello')) is None
def test_chat(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
# mocking is needed because of ci does not have xclip or xselect installed
def mock_copy(text: str) -> None:
pass
mocker.patch('pyperclip.copy', mock_copy)
with create_pipe_input() as inp:
inp.send_text('\n')
inp.send_text('hello\n')
inp.send_text('/markdown\n')
inp.send_text('/cp\n')
# a second agent turn, so `/usage` proves the session accumulator sums across the `run_chat` loop
inp.send_text('hello again\n')
inp.send_text('/usage\n')
inp.send_text('/usage --json\n')
inp.send_text('/exit\n')
session = PromptSession[Any](input=inp, output=DummyOutput())
m = mocker.patch('pydantic_ai._cli.PromptSession', return_value=session)
m.return_value = session
m = TestModel(custom_output_text='goodbye')
with cli_agent.override(model=m):
assert cli([]) == 0
assert capfd.readouterr().out.splitlines() == snapshot(
[
IsStr(),
IsStr(regex='goodbye *Markdown output of last question:'),
'',
'goodbye',
'Copied last output to clipboard.',
IsStr(regex=r'goodbye *clai usage \(session total\)'),
'',
'Turns: 2',
IsStr(regex=r'Tokens: .*'),
IsStr(regex=r' Input: .*'),
IsStr(regex=r' Output: .*'),
'Requests: 2',
'Tool calls: 0',
IsStr(regex=r'\{"turns": 2, .*"requests": 2, "tool_calls": 0\}'),
'Exiting…',
]
)
def test_handle_slash_command_markdown():
io = StringIO()
assert handle_slash_command('/markdown', [], False, Console(file=io), 'default') == (None, False)
assert io.getvalue() == snapshot('No markdown output available.\n')
messages: list[ModelMessage] = [ModelResponse(parts=[TextPart('[hello](# hello)'), ToolCallPart('foo', '{}')])]
io = StringIO()
assert handle_slash_command('/markdown', messages, True, Console(file=io), 'default') == (None, True)
assert io.getvalue() == snapshot("""\
Markdown output of last question:
[hello](# hello)
""")
def test_handle_slash_command_multiline():
io = StringIO()
assert handle_slash_command('/multiline', [], False, Console(file=io), 'default') == (None, True)
assert io.getvalue()[:70] == IsStr(regex=r'Enabling multiline mode.*')
io = StringIO()
assert handle_slash_command('/multiline', [], True, Console(file=io), 'default') == (None, False)
assert io.getvalue() == snapshot('Disabling multiline mode.\n')
def test_handle_slash_command_copy(mocker: MockerFixture):
io = StringIO()
# mocking is needed because of ci does not have xclip or xselect installed
mock_clipboard: list[str] = []
def append_to_clipboard(text: str) -> None:
mock_clipboard.append(text)
mocker.patch('pyperclip.copy', append_to_clipboard)
assert handle_slash_command('/cp', [], False, Console(file=io), 'default') == (None, False)
assert io.getvalue() == snapshot('No output available to copy.\n')
assert mock_clipboard == snapshot([])
messages: list[ModelMessage] = [ModelResponse(parts=[TextPart(''), ToolCallPart('foo', '{}')])]
io = StringIO()
assert handle_slash_command('/cp', messages, True, Console(file=io), 'default') == (None, True)
assert io.getvalue() == snapshot('No text content to copy.\n')
assert mock_clipboard == snapshot([])
messages: list[ModelMessage] = [ModelResponse(parts=[TextPart('hello'), ToolCallPart('foo', '{}')])]
io = StringIO()
assert handle_slash_command('/cp', messages, True, Console(file=io), 'default') == (None, True)
assert io.getvalue() == snapshot('Copied last output to clipboard.\n')
assert mock_clipboard == snapshot(['hello'])
def test_handle_slash_command_exit():
io = StringIO()
assert handle_slash_command('/exit', [], False, Console(file=io), 'default') == (0, False)
assert io.getvalue() == snapshot('Exiting…\n')
def test_handle_slash_command_other():
io = StringIO()
assert handle_slash_command('/foobar', [], False, Console(file=io), 'default') == (None, False)
assert io.getvalue() == snapshot('Unknown command `/foobar`\n')
def test_format_usage():
usage = RunUsage(input_tokens=1521, output_tokens=326, requests=1)
assert format_usage(usage, 1) == snapshot("""\
clai usage (session total)
Turns: 1
Tokens: 1,847
Input: 1,521
Output: 326
Requests: 1
Tool calls: 0\
""")
def test_format_usage_json():
usage = RunUsage(input_tokens=1521, output_tokens=326, requests=1, tool_calls=2)
assert format_usage(usage, 3, as_json=True) == snapshot(
'{"turns": 3, "input_tokens": 1521, "output_tokens": 326, "total_tokens": 1847, "requests": 1, "tool_calls": 2}'
)
def test_handle_slash_command_usage():
usage = RunUsage(input_tokens=51, output_tokens=1, requests=1)
io = StringIO()
assert handle_slash_command('/usage', [], False, Console(file=io), 'default', usage=usage, turns=1) == (None, False)
assert io.getvalue() == snapshot("""\
clai usage (session total)
Turns: 1
Tokens: 52
Input: 51
Output: 1
Requests: 1
Tool calls: 0
""")
def test_handle_slash_command_usage_json():
usage = RunUsage(input_tokens=51, output_tokens=1, requests=1)
io = StringIO()
# Spaces are replaced with `-` upstream, so `/usage --json` arrives as `/usage---json`.
assert handle_slash_command('/usage---json', [], False, Console(file=io), 'default', usage=usage, turns=1) == (
None,
False,
)
assert io.getvalue() == snapshot(
'{"turns": 1, "input_tokens": 51, "output_tokens": 1, "total_tokens": 52, "requests": 1, "tool_calls": 0}\n'
)
def test_handle_slash_command_usage_no_session():
# With no session usage threaded through, fall back to an empty RunUsage rather than erroring.
io = StringIO()
assert handle_slash_command('/usage', [], False, Console(file=io), 'default') == (None, False)
assert io.getvalue() == snapshot("""\
clai usage (session total)
Turns: 0
Tokens: 0
Input: 0
Output: 0
Requests: 0
Tool calls: 0
""")
def test_handle_slash_command_usage_unknown_option():
io = StringIO()
assert handle_slash_command('/usage---foo', [], False, Console(file=io), 'default', usage=RunUsage(), turns=0) == (
None,
False,
)
assert io.getvalue() == snapshot('Unknown `/usage` option `/usage---foo`\n')
def test_handle_slash_command_usage_not_a_flag():
# Without a separator, `/usagejson` is an unknown command, not a silently-accepted `--json` request.
io = StringIO()
assert handle_slash_command('/usagejson', [], False, Console(file=io), 'default', usage=RunUsage(), turns=0) == (
None,
False,
)
assert io.getvalue() == snapshot('Unknown command `/usagejson`\n')
async def test_ask_agent_accumulates_usage(env: TestEnv):
# `ask_agent` increments the shared session usage on both the streaming and non-streaming paths,
# including tool calls, and threading a turn's history into the next must not re-count prior usage.
env.set('OPENAI_API_KEY', 'test')
agent = Agent(TestModel())
@agent.tool_plain
def get_number() -> int:
return 42
session_usage = RunUsage()
console = Console(file=StringIO())
# streaming turn
messages = await ask_agent(agent, 'first', True, console, 'default', usage=session_usage)
# non-streaming turn, threading the prior turn's history back in as `run_chat` does
await ask_agent(agent, 'second', False, console, 'default', messages=messages, usage=session_usage)
# Usage is the sum of both runs: the first turn calls the tool (2 requests, 1 tool call), the second
# replays that history without re-calling it (1 request). Prior usage is not re-counted from history.
assert session_usage == snapshot(RunUsage(input_tokens=156, output_tokens=14, requests=3, tool_calls=1))
async def test_ask_agent_counts_usage_on_failed_turn(env: TestEnv):
# A turn that makes a billed request and then raises must still be counted, so a later `/usage`
# does not under-report tokens that were actually spent. The merge happens in `ask_agent`'s `finally`.
env.set('OPENAI_API_KEY', 'test')
agent = Agent(TestModel())
@agent.tool_plain
def boom() -> int:
raise RuntimeError('boom')
session_usage = RunUsage()
console = Console(file=StringIO())
with pytest.raises(RuntimeError, match='boom'):
await ask_agent(agent, 'go', False, console, 'default', usage=session_usage)
# The model request that produced the failing tool call was billed, so it is reflected in the total.
assert session_usage == snapshot(RunUsage(input_tokens=51, output_tokens=2, requests=1))
def test_code_theme_unset(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
cli([])
mock_run_chat.assert_awaited_once_with(
True, IsInstance(Agent), IsInstance(Console), 'monokai', 'clai', toolsets=None
)
def test_code_theme_light(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
cli(['--code-theme=light'])
mock_run_chat.assert_awaited_once_with(
True, IsInstance(Agent), IsInstance(Console), 'default', 'clai', toolsets=None
)
def test_code_theme_dark(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
cli(['--code-theme=dark'])
mock_run_chat.assert_awaited_once_with(
True, IsInstance(Agent), IsInstance(Console), 'monokai', 'clai', toolsets=None
)
def test_agent_to_cli_sync(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
cli_agent.to_cli_sync()
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=None,
model=None,
model_settings=None,
usage_limits=None,
)
async def test_agent_to_cli_async(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
await cli_agent.to_cli()
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=None,
model=None,
model_settings=None,
usage_limits=None,
)
async def test_agent_to_cli_with_message_history(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
# Create some test message history - cast to the proper base type
test_messages: list[ModelMessage] = [ModelResponse(parts=[TextPart('Hello!')])]
await cli_agent.to_cli(message_history=test_messages)
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=test_messages,
model=None,
model_settings=None,
usage_limits=None,
)
def test_agent_to_cli_sync_with_message_history(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
# Create some test message history - cast to the proper base type
test_messages: list[ModelMessage] = [ModelResponse(parts=[TextPart('Hello!')])]
cli_agent.to_cli_sync(message_history=test_messages)
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=test_messages,
model=None,
model_settings=None,
usage_limits=None,
)
@pytest.mark.parametrize(
('model_name', 'expected'),
[
('gpt-5', 'GPT 5'),
('gpt-4.1', 'GPT 4.1'),
('o1', 'O1'),
('o3', 'O3'),
('claude-sonnet-4-5', 'Claude Sonnet 4.5'),
('claude-haiku-4-5', 'Claude Haiku 4.5'),
('gemini-2.5-pro', 'Gemini 2.5 Pro'),
('gemini-2.5-flash', 'Gemini 2.5 Flash'),
('sonnet-4-5', 'Sonnet 4.5'),
('custom-model', 'Custom Model'),
],
)
def test_model_label(model_name: str, expected: str):
"""Test Model.label formatting for UI."""
from pydantic_ai.models.test import TestModel
model = TestModel(model_name=model_name)
assert model.label == expected
def test_clai_web_generic_agent(mocker: MockerFixture, env: TestEnv):
"""Test web command without agent creates generic agent."""
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
assert cli(['web', '-m', 'openai:gpt-5', '-t', 'web_search'], prog_name='clai') == 0
mock_run_web.assert_called_once_with(
agent_path=None,
host='127.0.0.1',
port=7932,
models=['openai:gpt-5'],
tools=['web_search'],
instructions=None,
default_model='openai:gpt-5',
html_source=None,
allowed_hosts=[],
)
def test_clai_web_success(mocker: MockerFixture, create_test_module: Callable[..., None], env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
test_agent = Agent(TestModel(custom_output_text='test'))
create_test_module(custom_agent=test_agent)
assert cli(['web', '--agent', 'test_module:custom_agent'], prog_name='clai') == 0
mock_run_web.assert_called_once_with(
agent_path='test_module:custom_agent',
host='127.0.0.1',
port=7932,
models=[],
tools=[],
instructions=None,
default_model='openai:gpt-5',
html_source=None,
allowed_hosts=[],
)
def test_clai_web_with_models(mocker: MockerFixture, create_test_module: Callable[..., None], env: TestEnv):
"""Test web command with multiple -m flags."""
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
test_agent = Agent(TestModel(custom_output_text='test'))
create_test_module(custom_agent=test_agent)
assert (
cli(
[
'web',
'--agent',
'test_module:custom_agent',
'-m',
'openai:gpt-5',
'-m',
'anthropic:claude-sonnet-4-6',
],
prog_name='clai',
)
== 0
)
mock_run_web.assert_called_once_with(
agent_path='test_module:custom_agent',
host='127.0.0.1',
port=7932,
models=['openai:gpt-5', 'anthropic:claude-sonnet-4-6'],
tools=[],
instructions=None,
default_model='openai:gpt-5',
html_source=None,
allowed_hosts=[],
)
def test_clai_web_with_tools(mocker: MockerFixture, create_test_module: Callable[..., None], env: TestEnv):
"""Test web command with multiple -t flags."""
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
test_agent = Agent(TestModel(custom_output_text='test'))
create_test_module(custom_agent=test_agent)
assert (
cli(
['web', '--agent', 'test_module:custom_agent', '-t', 'web_search', '-t', 'code_execution'], prog_name='clai'
)
== 0
)
mock_run_web.assert_called_once_with(
agent_path='test_module:custom_agent',
host='127.0.0.1',
port=7932,
models=[],
tools=['web_search', 'code_execution'],
instructions=None,
default_model='openai:gpt-5',
html_source=None,
allowed_hosts=[],
)
def test_clai_web_generic_with_instructions(mocker: MockerFixture, env: TestEnv):
"""Test generic agent with custom instructions."""
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
assert cli(['web', '-m', 'openai:gpt-5', '-i', 'You are a helpful coding assistant'], prog_name='clai') == 0
mock_run_web.assert_called_once_with(
agent_path=None,
host='127.0.0.1',
port=7932,
models=['openai:gpt-5'],
tools=[],
instructions='You are a helpful coding assistant',
default_model='openai:gpt-5',
html_source=None,
allowed_hosts=[],
)
def test_clai_web_with_custom_port(mocker: MockerFixture, create_test_module: Callable[..., None], env: TestEnv):
"""Test web command with custom host/port."""
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
test_agent = Agent(TestModel(custom_output_text='test'))
create_test_module(custom_agent=test_agent)
assert (
cli(['web', '--agent', 'test_module:custom_agent', '--host', '0.0.0.0', '--port', '7932'], prog_name='clai')
== 0
)
mock_run_web.assert_called_once_with(
agent_path='test_module:custom_agent',
host='0.0.0.0',
port=7932,
models=[],
tools=[],
instructions=None,
default_model='openai:gpt-5',
html_source=None,
allowed_hosts=[],
)
def test_run_web_command_agent_with_model(
mocker: MockerFixture, create_test_module: Callable[..., None], capfd: CaptureFixture[str]
):
"""Test run_web_command uses agent's model when no -m flag provided."""
mock_uvicorn_run = mocker.patch('uvicorn.run')
mocker.patch('pydantic_ai._cli.web.create_web_app')
test_agent = Agent(TestModel(custom_output_text='test'))
create_test_module(custom_agent=test_agent)
result = run_web_command(agent_path='test_module:custom_agent')
assert result == 0
mock_uvicorn_run.assert_called_once()
def test_run_web_command_generic_agent_no_model(mocker: MockerFixture, capfd: CaptureFixture[str]):
"""Test run_web_command uses default model when no agent and no model provided."""
mock_uvicorn_run = mocker.patch('uvicorn.run')
mock_create_app = mocker.patch('pydantic_ai._cli.web.create_web_app')
result = run_web_command()
assert result == 0
mock_uvicorn_run.assert_called_once()
# Verify default model was passed
call_kwargs = mock_create_app.call_args.kwargs
assert call_kwargs['models'] == ['openai:gpt-5']
def test_run_web_command_generic_agent_with_instructions(mocker: MockerFixture, capfd: CaptureFixture[str]):
"""Test run_web_command passes instructions to create_web_app for generic agent."""
mock_uvicorn_run = mocker.patch('uvicorn.run')
mock_create_app = mocker.patch('pydantic_ai._cli.web.create_web_app')
result = run_web_command(models=['test'], instructions='You are a helpful assistant')
assert result == 0
mock_uvicorn_run.assert_called_once()
# Verify instructions were passed to create_web_app (not to Agent constructor)
call_kwargs = mock_create_app.call_args.kwargs
assert call_kwargs['instructions'] == 'You are a helpful assistant'
def test_run_web_command_agent_with_instructions(
mocker: MockerFixture, create_test_module: Callable[..., None], capfd: CaptureFixture[str]
):
"""Test run_web_command passes instructions to create_web_app when agent is provided."""
mock_uvicorn_run = mocker.patch('uvicorn.run')
mock_create_app = mocker.patch('pydantic_ai._cli.web.create_web_app')
test_agent = Agent(TestModel(custom_output_text='test'))
create_test_module(custom_agent=test_agent)
result = run_web_command(agent_path='test_module:custom_agent', instructions='Always respond in Spanish')
assert result == 0
mock_uvicorn_run.assert_called_once()
# Verify instructions were passed to create_web_app
call_kwargs = mock_create_app.call_args.kwargs
assert call_kwargs['instructions'] == 'Always respond in Spanish'
def test_run_web_command_agent_load_failure(capfd: CaptureFixture[str]):
"""Test run_web_command returns error when agent path is invalid."""
result = run_web_command(agent_path='nonexistent_module:agent')
assert result == 1
output = capfd.readouterr().out
assert 'Could not load agent' in output
def test_run_web_command_unknown_tool(mocker: MockerFixture, capfd: CaptureFixture[str]):
"""Test run_web_command warns about unknown tool IDs."""
mock_uvicorn_run = mocker.patch('uvicorn.run')
mocker.patch('pydantic_ai._cli.web.create_web_app')
result = run_web_command(models=['test'], tools=['unknown_tool_xyz'])
assert result == 0
mock_uvicorn_run.assert_called_once()
output = capfd.readouterr().out
assert 'Unknown tool "unknown_tool_xyz"' in output
def test_run_web_command_memory_tool(mocker: MockerFixture, capfd: CaptureFixture[str]):
"""Test run_web_command warns about memory tool requiring agent configuration."""
mock_uvicorn_run = mocker.patch('uvicorn.run')
mocker.patch('pydantic_ai._cli.web.create_web_app')
result = run_web_command(models=['test'], tools=['memory'])
assert result == 0
mock_uvicorn_run.assert_called_once()
output = capfd.readouterr().out
assert '"memory" requires configuration and cannot be enabled via CLI' in output
def test_run_web_command_agent_native_tools_not_duplicated(
mocker: MockerFixture, create_test_module: Callable[..., None], capfd: CaptureFixture[str]
):
"""Test run_web_command only passes CLI-provided tools, not agent's native tools."""
from pydantic_ai.native_tools import WebSearchTool
mock_uvicorn_run = mocker.patch('uvicorn.run')
mock_create_app = mocker.patch('pydantic_ai._cli.web.create_web_app')
# Create agent with web_search tool already configured
test_agent = Agent(TestModel(custom_output_text='test'), capabilities=[NativeTool(WebSearchTool())])
create_test_module(custom_agent=test_agent)
# Add code_execution via CLI
result = run_web_command(agent_path='test_module:custom_agent', tools=['code_execution'])
assert result == 0
mock_uvicorn_run.assert_called_once()
# Verify only CLI-provided tools are passed (agent's tools are handled by create_web_app)
call_kwargs = mock_create_app.call_args.kwargs
native_tools = call_kwargs.get('native_tools', [])
tool_kinds = {t.kind for t in native_tools}
# web_search is on the agent, so it's NOT passed here (it's handled internally)
assert 'web_search' not in tool_kinds
# code_execution was provided via CLI, so it IS passed
assert 'code_execution' in tool_kinds
def test_run_web_command_cli_models_passed_to_create_web_app(
mocker: MockerFixture, create_test_module: Callable[..., None]
):
"""Test that CLI models are passed directly to create_web_app (agent model merging happens there)."""
mock_uvicorn_run = mocker.patch('uvicorn.run')
mock_create_app = mocker.patch('pydantic_ai._cli.web.create_web_app')
test_agent = Agent(TestModel(custom_output_text='test'))
create_test_module(custom_agent=test_agent)
result = run_web_command(
agent_path='test_module:custom_agent', models=['openai:gpt-5', 'anthropic:claude-sonnet-4-6']
)
assert result == 0
mock_uvicorn_run.assert_called_once()
call_kwargs = mock_create_app.call_args.kwargs
# CLI models passed as list; agent model merging/deduplication happens in create_web_app
assert call_kwargs.get('models') == ['openai:gpt-5', 'anthropic:claude-sonnet-4-6']
def test_agent_to_cli_sync_with_args(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
model_settings = ModelSettings(temperature=0.5)
usage_limits = UsageLimits(request_limit=10)
cli_agent.to_cli_sync(model_settings=model_settings, usage_limits=usage_limits)
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=None,
model=None,
model_settings=model_settings,
usage_limits=usage_limits,
)
def test_agent_to_cli_sync_with_model(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
cli_agent.to_cli_sync(model='test')
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=None,
model='test',
model_settings=None,
usage_limits=None,
)
async def test_agent_to_cli_async_with_args(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
model_settings = ModelSettings(temperature=0.5)
usage_limits = UsageLimits(request_limit=10)
await cli_agent.to_cli(model_settings=model_settings, usage_limits=usage_limits)
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=None,
model=None,
model_settings=model_settings,
usage_limits=usage_limits,
)
async def test_agent_to_cli_async_with_model(mocker: MockerFixture, env: TestEnv):
env.set('OPENAI_API_KEY', 'test')
mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
await cli_agent.to_cli(model='test')
mock_run_chat.assert_awaited_once_with(
stream=True,
agent=IsInstance(Agent),
console=IsInstance(Console),
code_theme='monokai',
prog_name='pydantic-ai',
deps=None,
message_history=None,
model='test',
model_settings=None,
usage_limits=None,
)
async def test_ask_agent_non_stream_forwards_run_kwargs(mocker: MockerFixture):
from pydantic_ai._cli import ask_agent
result = mocker.Mock()
result.output = 'hello'
result.all_messages.return_value = []
agent = mocker.Mock()
agent.run = mocker.AsyncMock(return_value=result)
model_settings = ModelSettings(temperature=0)
usage_limits = UsageLimits(request_limit=5)
messages = await ask_agent(
agent,
'Hello',
stream=False,
console=Console(file=StringIO()),
code_theme='monokai',
model='test',
model_settings=model_settings,
usage_limits=usage_limits,
)
agent.run.assert_awaited_once_with(
'Hello',
message_history=None,
deps=None,
model='test',
model_settings=model_settings,
usage_limits=usage_limits,
toolsets=None,
usage=IsInstance(RunUsage),
)
assert messages == []
def test_clai_web_with_html_source(mocker: MockerFixture, env: TestEnv):
"""Test web command with --html-source flag."""
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
custom_url = 'https://internal.company.com/pydantic-ai-ui/index.html'
assert cli(['web', '-m', 'openai:gpt-5', '--html-source', custom_url], prog_name='clai') == 0
mock_run_web.assert_called_once_with(
agent_path=None,
host='127.0.0.1',
port=7932,
models=['openai:gpt-5'],
tools=[],
instructions=None,
default_model='openai:gpt-5',
html_source=custom_url,
allowed_hosts=[],
)
def test_clai_web_with_allowed_hosts(mocker: MockerFixture, env: TestEnv):
"""`--allowed-host` is repeatable, for serving the UI under a name behind a proxy or a tunnel."""
env.set('OPENAI_API_KEY', 'test')
mock_run_web = mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
args = ['web', '-m', 'openai:gpt-5', '--allowed-host', 'ui.example.com', '--allowed-host', '*.corp.example']
assert cli(args, prog_name='clai') == 0
mock_run_web.assert_called_once_with(
agent_path=None,
host='127.0.0.1',
port=7932,
models=['openai:gpt-5'],
tools=[],
instructions=None,
default_model='openai:gpt-5',
html_source=None,
allowed_hosts=['ui.example.com', '*.corp.example'],
)
def test_clai_web_answers_to_the_host_it_binds_to(mocker: MockerFixture, env: TestEnv):
"""`--host <name>` implies answering to that name, so the URL the CLI prints actually works.
Without this the CLI contradicts itself: it prints `Open your browser at: http://devbox.example:7932`
and then rejects that exact `Host` with a `421`.
"""
env.set('OPENAI_API_KEY', 'test')
mock_uvicorn = mocker.patch('uvicorn.run')
mock_create = mocker.patch('pydantic_ai._cli.web.create_web_app')
assert cli(['web', '-m', 'openai:gpt-5', '--host', 'devbox.example'], prog_name='clai') == 0
assert mock_create.call_args.kwargs['allowed_hosts'] == ['devbox.example']
assert mock_uvicorn.call_args.kwargs['host'] == 'devbox.example'
LOGO_MARKER = _display._LOGO_LINES[-1] # pyright: ignore[reportPrivateUsage]
"""The logo's closing row, standing in for the whole banner having landed in the output."""
@pytest.fixture
def terminal_clai(env: TestEnv) -> Iterator[None]:
"""A `clai` session that believes it owns a terminal, and starts with the banner unclaimed."""
env.set('FORCE_COLOR', '1')
env.remove('CI')
env.remove('COLUMNS')
env.remove('PYDANTIC_AI_NO_BANNER')
# The suite that asserts on the banner is the one place a test run is allowed to show one.
env.remove('PYTEST_VERSION')
_display._banner_displayed = False # pyright: ignore[reportPrivateUsage]
try:
yield
finally:
_display._banner_displayed = False # pyright: ignore[reportPrivateUsage]
@pytest.fixture
def agent_clai(env: TestEnv) -> Iterator[None]:
"""A `clai` a coding agent started, whose output is a pipe it reads back rather than a terminal."""
env.set('AI_AGENT', 'some-harness')
env.remove('CI')
env.remove('COLUMNS')
env.remove('PYDANTIC_AI_NO_BANNER')
env.remove('PYTEST_VERSION')
_display._banner_displayed = False # pyright: ignore[reportPrivateUsage]
try:
yield
finally:
_display._banner_displayed = False # pyright: ignore[reportPrivateUsage]
def _plain(output: str) -> str:
"""The output as the user reads it, with the colour the banner is styled in taken back off."""
return Text.from_ansi(output).plain
def test_clai_intro_shows_banner(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv, terminal_clai: None):
env.set('OPENAI_API_KEY', 'test')
mocker.patch('pydantic_ai._cli.ask_agent')
assert cli(['hello']) == 0
output = _plain(capfd.readouterr().out)
assert LOGO_MARKER in output
# clai shows the same banner a script gets, rather than heading it with its own name and version.
assert f'pydantic-ai v{__version__}' in output
assert 'clai - Pydantic AI CLI' not in output
assert 'model: openai:gpt-5 • tools: 0 • capabilities: 0' in output
# The banner carries the observability pointer, so clai doesn't repeat its own header line.
# Matched on the label, not the prose, which is the banner's to reword.
assert 'observability:' in output
assert 'with openai:gpt-5' not in output
def test_clai_nested_in_an_agent_still_opens_with_a_banner(
capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv, agent_clai: None
):
"""A harness running `clai` reads what it writes, so a pipe is no reason to hold the banner back."""
env.set('OPENAI_API_KEY', 'test')
mocker.patch('pydantic_ai._cli.ask_agent')
assert cli(['hello']) == 0
output = capfd.readouterr().out
assert LOGO_MARKER in output
# Nothing renders the codes for a pipe, so the console leaves them out rather than writing them raw.
assert '\x1b[' not in output
def test_clai_piped_to_a_file_still_honours_an_exported_columns(
capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv, agent_clai: None
):
"""The console answers 80 for a pipe whether it read that or guessed it, so it isn't asked.
Without this, `clai | tee log.txt` and a plain agent run piped the same way would lay the
banner out at two different widths under the same exported `COLUMNS`.
"""
env.set('OPENAI_API_KEY', 'test')
env.set('COLUMNS', '70')
mocker.patch('pydantic_ai._cli.ask_agent')
assert cli(['hello']) == 0
output = capfd.readouterr().out
assert LOGO_MARKER in output
assert max(map(len, output.splitlines())) <= 70
def test_clai_intro_names_the_agent_the_user_asked_for(
capfd: CaptureFixture[str],
mocker: MockerFixture,
env: TestEnv,
create_test_module: Callable[..., None],
terminal_clai: None,
):
env.set('OPENAI_API_KEY', 'test')
create_test_module(custom_agent=Agent(TestModel()))
mocker.patch('pydantic_ai._cli.ask_agent')
assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
# The loaded agent has no name of its own, so the banner falls back to the path the user gave.
assert 'agent: test_module:custom_agent' in _plain(capfd.readouterr().out)
def test_clai_intro_drops_observability_for_an_instrumented_agent(
capfd: CaptureFixture[str],
mocker: MockerFixture,
env: TestEnv,
create_test_module: Callable[..., None],
terminal_clai: None,
):
env.set('OPENAI_API_KEY', 'test')
instrumented = Agent(TestModel(), name='observed')
instrumented.instrument = True
create_test_module(custom_agent=instrumented)
mocker.patch('pydantic_ai._cli.ask_agent')
assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
output = _plain(capfd.readouterr().out)
assert 'agent: observed' in output
assert 'observability: off' not in output
def test_clai_intro_drops_observability_when_instrumented_globally(
capfd: CaptureFixture[str],
mocker: MockerFixture,
env: TestEnv,
create_test_module: Callable[..., None],
terminal_clai: None,
):
"""`instrument_all()` leaves `agent.instrument` as `None`, so reading it would advertise Logfire
to someone already sending traces there."""
env.set('OPENAI_API_KEY', 'test')
create_test_module(custom_agent=Agent(TestModel(), name='observed'))
mocker.patch('pydantic_ai._cli.ask_agent')
Agent.instrument_all()
try:
assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
finally:
Agent.instrument_all(False)
output = _plain(capfd.readouterr().out)
assert 'agent: observed' in output
assert 'observability: off' not in output
def test_run_chat_banner_names_the_model_the_session_will_use(
mocker: MockerFixture, tmp_path: Path, terminal_clai: None
):
"""Under `override(model=...)` the prompts go to the override, so the banner has to say so."""
def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart('hi')]) # pragma: no cover
console, io = _chat_console()
agent = Agent(TestModel(), name='support_agent')
with agent.override(model=FunctionModel(respond)):
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
anyio.run(run_chat, True, agent, console, 'monokai', 'pydantic-ai', tmp_path)
output = _plain(io.getvalue())
assert 'model: function:function:respond:' in output
assert 'test:test' not in output
def test_run_chat_shows_a_banner_for_an_agent_that_only_has_an_override(
mocker: MockerFixture, tmp_path: Path, terminal_clai: None
):
"""The override is the session's model, so deciding on `agent.model` alone would skip the banner."""
def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart('hi')]) # pragma: no cover
console, io = _chat_console()
agent = Agent(name='modelless_agent')
with agent.override(model=FunctionModel(respond)):
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
anyio.run(run_chat, True, agent, console, 'monokai', 'pydantic-ai', tmp_path)
assert 'agent: modelless_agent • model: function:function:respond:' in _plain(io.getvalue())
def test_clai_intro_counts_tools_from_an_override(
capfd: CaptureFixture[str],
mocker: MockerFixture,
env: TestEnv,
create_test_module: Callable[..., None],
terminal_clai: None,
):
"""`override(tools=...)` builds a fresh function toolset, leaving the agent's own untouched."""
def ping() -> str:
return 'pong' # pragma: no cover
env.set('OPENAI_API_KEY', 'test')
agent = Agent(TestModel(), name='overridden')
create_test_module(custom_agent=agent)
mocker.patch('pydantic_ai._cli.ask_agent')
with agent.override(tools=[ping]):
assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
assert 'tools: 1' in _plain(capfd.readouterr().out)
def test_run_chat_opens_even_if_the_banner_cannot_be_written(
mocker: MockerFixture, tmp_path: Path, terminal_clai: None
):
"""A terminal whose encoding can't take the banner shouldn't take the session down with it."""
class NoSeparatorIO(StringIO):
"""Rejects the separator between the banner's details, as an ASCII terminal rejects it.
The logo itself is plain ASCII, so the banner's own unencodable characters are what a
terminal under `LC_ALL=C` chokes on. Narrowed to this one rather than all of them, so that
the rest of the session still writes and the test can show it survived.
"""
def write(self, s: str) -> int:
if _display._INFO_SEPARATOR in s: # pyright: ignore[reportPrivateUsage]
raise UnicodeEncodeError('ascii', s, 0, 1, 'ordinal not in range(128)')
return super().write(s)
io = NoSeparatorIO()
console = Console(file=io, force_terminal=True)
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
assert anyio.run(run_chat, True, Agent(TestModel()), console, 'monokai', 'pydantic-ai', tmp_path) == 0
# The chat opened and ran to its `/exit` without a header, rather than not at all.
assert LOGO_MARKER not in io.getvalue()
assert 'Exiting' in _plain(io.getvalue())
def test_clai_intro_falls_back_to_one_line_when_banner_is_suppressed(
capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv, terminal_clai: None
):
env.set('PYDANTIC_AI_NO_BANNER', '1')
env.set('OPENAI_API_KEY', 'test')
mocker.patch('pydantic_ai._cli.ask_agent')
assert cli(['hello']) == 0
# The one-liner clai has always printed, styled per-segment, so ANSI codes sit between the words.
output = capfd.readouterr().out
assert LOGO_MARKER not in output
assert 'observability' not in output
assert 'clai - Pydantic AI CLI' in output
assert 'openai:gpt-5' in output
def test_clai_run_does_not_print_a_second_banner(capfd: CaptureFixture[str], env: TestEnv, terminal_clai: None):
"""The banner belongs at startup, not in the middle of the answer to the first prompt."""
env.set('OPENAI_API_KEY', 'test')
assert cli(['--model', 'test', 'hello']) == 0
# `ask_agent` claims the banner, so the run inside it has nothing left to print.
assert capfd.readouterr().out.count(LOGO_MARKER) == 1
def test_clai_web_does_not_print_a_banner(mocker: MockerFixture, env: TestEnv, terminal_clai: None):
env.set('OPENAI_API_KEY', 'test')
mocker.patch('pydantic_ai._cli.web.run_web_command', return_value=0)
assert cli(['web']) == 0
# A server's first chat request would otherwise print the banner into its log, long after start.
assert _display.claim_banner() is False
def _exit_immediately(mocker: MockerFixture, inp: Any) -> None:
inp.send_text('/exit\n')
mocker.patch('pydantic_ai._cli.PromptSession', return_value=PromptSession[Any](input=inp, output=DummyOutput()))
def _chat_console(width: int | None = None) -> tuple[Console, StringIO]:
"""A console for a session that believes it owns a terminal, at `width` columns if it says.
Left unsaid, rich answers 80 for a terminal it can't measure, which is what the banner then
lays itself out for — a narrower column than the one a wide terminal gets.
"""
io = StringIO()
return Console(file=io, force_terminal=True, width=width), io
def test_run_chat_shows_banner_for_a_users_own_agent(mocker: MockerFixture, tmp_path: Path, terminal_clai: None):
"""`Agent.to_cli()` prints no header of its own, so without this its session announces nothing."""
console, io = _chat_console()
agent = Agent(TestModel(), name='support_agent')
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
anyio.run(run_chat, True, agent, console, 'monokai', 'pydantic-ai', tmp_path)
output = _plain(io.getvalue())
assert LOGO_MARKER in output
# A user's own agent gets the same banner a script does, not clai's.
assert 'Pydantic AI CLI' not in output
assert f'pydantic-ai v{__version__}' in output
# The console reports 80 columns, so the details take the two lines that fit in them.
assert 'agent: support_agent • model: test:test • tools: 0' in output
assert 'capabilities: 0' in output
def test_run_chat_lays_the_banner_out_for_the_terminal_it_has(
mocker: MockerFixture, tmp_path: Path, terminal_clai: None
):
"""A pane the banner outruns is one the terminal breaks itself, straight through the logo."""
console, io = _chat_console(width=64)
agent = Agent(TestModel(), name='support_agent')
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
anyio.run(run_chat, True, agent, console, 'monokai', 'pydantic-ai', tmp_path)
banner = _plain(io.getvalue()).partition('Exiting')[0]
assert LOGO_MARKER in banner
assert max(map(len, banner.splitlines())) <= 64
def test_run_chat_without_a_model_shows_no_banner(mocker: MockerFixture, tmp_path: Path, terminal_clai: None):
"""There is nothing to say about the model yet, and the first prompt will fail on it anyway."""
console, io = _chat_console()
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
anyio.run(run_chat, True, Agent(), console, 'monokai', 'pydantic-ai', tmp_path)
assert LOGO_MARKER not in io.getvalue()
def test_run_chat_on_a_wrapped_agent_shows_no_banner(mocker: MockerFixture, tmp_path: Path, terminal_clai: None):
"""The counts the banner reports come off `Agent`; a wrapper keeps the silence it has today."""
console, io = _chat_console()
agent = WrapperAgent(Agent(TestModel(), name='support_agent'))
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
anyio.run(run_chat, True, agent, console, 'monokai', 'pydantic-ai', tmp_path)
assert LOGO_MARKER not in io.getvalue()
def test_clai_chat_session_does_not_print_a_second_banner(
capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv, terminal_clai: None
):
"""`_run_chat_command` prints the intro, so `run_chat` must find the banner already claimed."""
env.set('OPENAI_API_KEY', 'test')
with create_pipe_input() as inp:
_exit_immediately(mocker, inp)
with cli_agent.override(model=TestModel(custom_output_text='hi')):
assert cli([]) == 0
assert capfd.readouterr().out.count(LOGO_MARKER) == 1
def test_auto_suggest_completes_a_slash_command():
"""A typed prefix of a slash command wins over whatever history would have suggested."""
suggestion = CustomAutoSuggest(['/exit']).get_suggestion(Buffer(), Document('/e'))
assert suggestion is not None and suggestion.text == 'xit'
def test_auto_suggest_falls_back_to_history_for_a_non_command():
"""Text matching no slash command leaves the history suggestion in place — here, none."""
assert CustomAutoSuggest(['/exit']).get_suggestion(Buffer(), Document('hello')) is None