pydantic/pydantic-ai
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai.git synced 2026-10-12 02:21:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/main' into fix-clai2-utf8-cli-output

Browse Source
This commit is contained in:
David Sanchez
2026-10-09 16:18:27 -04:00
parent 6906fadd1d a68f1ad19b
commit 533292455a
115 changed files with 2910 additions and 825 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/scripts/shard_durations.py
+270
View File
@@ -0,0 +1,270 @@
"""Split sharded test jobs by measured durations, and report the tests a change made slow.
Every sharded test job stores the durations of the tests it ran (`pytest --store-durations
--clean-durations`) and uploads them as an artifact named `<prefix><shard>`.
- `fetch`, run once per workflow run by `test-durations-baseline`, merges those artifacts from the
newest `main` run that had uploaded all of them when this run was created, so the split follows the
suite as it is now instead of a committed file that drifts. Every shard then reads that one file, so
they all agree on the split even when the API fails (the file is then empty).
- `select` assigns whole test modules to shards by those durations (by source size without them) and
prints the `--ignore` options that leave pytest only this shard's modules.
- `report` compares a shard's measured durations with the baseline and lists, in the job summary, new
tests that take a while and tests that got much slower.
"""
from __future__ import annotations
import argparse
import heapq
import io
import json
import os
import statistics
import subprocess
import sys
import zipfile
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path, PurePosixPath
from typing import Any
Durations = dict[str, float]
GhApi = Callable[[str], bytes]
WORKFLOW_FILE = 'ci.yml'
# How many `main` runs back to look for one that uploaded every shard's durations.
MAX_CANDIDATE_RUNS = 20
NEW_TEST_MIN_SECONDS = 1.0
SLOWDOWN_FACTOR = 2.0
SLOWDOWN_MIN_SECONDS = 1.0
REPORT_ROW_LIMIT = 30
def gh_api(path: str) -> bytes:
"""Call the GitHub REST API through the `gh` CLI, which authenticates with `GH_TOKEN`."""
return subprocess.run(['gh', 'api', path], check=True, capture_output=True).stdout
def _read_durations_zip(content: bytes) -> Durations:
durations: Durations = {}
with zipfile.ZipFile(io.BytesIO(content)) as archive:
for member in archive.namelist():
if member.endswith('.json'):
durations.update(json.loads(archive.read(member)))
return durations
def fetch(repo: str, prefix: str, shards: int, before: str, api: GhApi = gh_api) -> tuple[int, Durations] | None:
"""The merged durations of the newest `main` push run created before `before` that uploaded all `shards`.
A run created before this one has either uploaded its artifacts or is still running; requiring all
`shards` of them skips the latter, so a re-run (which keeps the run's creation time) picks the same
baseline.
"""
runs: dict[str, list[dict[str, Any]]] = json.loads(
api(
f'repos/{repo}/actions/workflows/{WORKFLOW_FILE}/runs'
f'?branch=main&event=push&created=%3C{before}&per_page={MAX_CANDIDATE_RUNS}'
)
)
wanted = {f'{prefix}{shard}' for shard in range(1, shards + 1)}
for run in runs.get('workflow_runs', []):
data: dict[str, list[dict[str, Any]]] = json.loads(
api(f'repos/{repo}/actions/runs/{run["id"]}/artifacts?per_page=100')
)
artifacts = [
artifact
for artifact in data.get('artifacts', [])
if artifact['name'].startswith(prefix) and not artifact.get('expired')
]
if not wanted <= {artifact['name'] for artifact in artifacts}:
continue # still running, a shard failed before uploading, or tests were skipped
durations: Durations = {}
for artifact in sorted(artifacts, key=lambda artifact: artifact['name']):
durations.update(_read_durations_zip(api(f'repos/{repo}/actions/artifacts/{artifact["id"]}/zip')))
return run['id'], durations
return None
def find_test_modules(root: Path, exclude: Sequence[str]) -> list[str]:
"""Every tracked `tests/**/test_*.py` module outside `exclude`, as a repo-relative POSIX path."""
tracked = subprocess.run(
['git', 'ls-files', '-z', 'tests'], cwd=root, check=True, capture_output=True, text=True
).stdout.split('\0')
excluded = tuple(path.rstrip('/') + '/' for path in exclude)
return sorted(
path
for path in tracked
if PurePosixPath(path).name.startswith('test_') and path.endswith('.py') and not path.startswith(excluded)
)
def module_weights(modules: Sequence[str], durations: Mapping[str, float], root: Path) -> dict[str, float]:
"""Each module's measured seconds; a module the baseline has not seen counts as the median one.
Without a baseline, every module weighs its source size.
"""
if not durations:
return {path: float((root / path).stat().st_size) for path in modules}
measured: dict[str, float] = {}
for test, seconds in durations.items():
path = test.split('::', 1)[0]
measured[path] = measured.get(path, 0.0) + seconds
known = [measured[path] for path in modules if path in measured]
fallback = statistics.median(known) if known else 1.0
return {path: measured.get(path, fallback) for path in modules}
def assign(weights: Mapping[str, float], shards: int) -> list[list[str]]:
"""Spread whole modules across `shards`, heaviest first onto the lightest shard.
Splitting by module rather than by test (as `pytest-split` does) means a shard never imports the
other shards' modules: with `pytest-split` every worker of every shard collects the whole suite and
then deselects most of it. Ties break on the path, so every shard computes the same assignment from
the same inputs and each module runs exactly once.
"""
totals = [(0.0, shard) for shard in range(shards)]
assigned: list[list[str]] = [[] for _ in range(shards)]
for path in sorted(weights, key=lambda path: (-weights[path], path)):
total, shard = heapq.heappop(totals)
assigned[shard].append(path)
heapq.heappush(totals, (total + weights[path], shard))
return assigned
@dataclass(frozen=True)
class Report:
new: list[tuple[str, float]]
slower: list[tuple[str, float, float]]
def compare(baseline: Mapping[str, float], current: Mapping[str, float]) -> Report:
new = sorted(
(
(test, seconds)
for test, seconds in current.items()
if test not in baseline and seconds >= NEW_TEST_MIN_SECONDS
),
key=lambda row: -row[1],
)
slower = sorted(
(
(test, baseline[test], seconds)
for test, seconds in current.items()
if test in baseline
and seconds >= SLOWDOWN_FACTOR * baseline[test]
and seconds - baseline[test] >= SLOWDOWN_MIN_SECONDS
),
key=lambda row: row[1] - row[2],
)
return Report(new=new, slower=slower)
def render(report: Report, title: str) -> str:
if not report.new and not report.slower:
return ''
lines = [f'### Test durations: {title}', '']
if report.new:
lines += [
f'New tests taking at least {NEW_TEST_MIN_SECONDS:g}s:',
'',
'| Test | Seconds |',
'| --- | ---: |',
*(f'| `{test}` | {seconds:.1f} |' for test, seconds in report.new[:REPORT_ROW_LIMIT]),
'',
]
if len(report.new) > REPORT_ROW_LIMIT:
lines += [f'...and {len(report.new) - REPORT_ROW_LIMIT} more.', '']
if report.slower:
lines += [
f'Tests at least {SLOWDOWN_FACTOR:g}x and {SLOWDOWN_MIN_SECONDS:g}s slower than on `main`:',
'',
'| Test | `main` | Now |',
'| --- | ---: | ---: |',
*(f'| `{test}` | {before:.1f} | {after:.1f} |' for test, before, after in report.slower[:REPORT_ROW_LIMIT]),
'',
]
if len(report.slower) > REPORT_ROW_LIMIT:
lines += [f'...and {len(report.slower) - REPORT_ROW_LIMIT} more.', '']
lines += ['Durations include setup and teardown; one-off slowdowns can be runner noise.', '']
return '\n'.join(lines)
def _load(path: Path) -> Durations:
try:
return json.loads(path.read_text(encoding='utf-8'))
except FileNotFoundError:
return {}
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
subparsers = parser.add_subparsers(dest='command', required=True)
fetch_parser = subparsers.add_parser('fetch', help="Fetch the latest `main` run's per-test durations.")
fetch_parser.add_argument('--prefix', required=True, help='Artifact name prefix; the shard number follows it.')
fetch_parser.add_argument('--shards', type=int, required=True)
fetch_parser.add_argument('--output', type=Path, required=True)
select_parser = subparsers.add_parser(
'select', help="Print the `--ignore` options that leave pytest only this shard's test modules."
)
select_parser.add_argument('--durations', type=Path, required=True)
select_parser.add_argument('--shards', type=int, required=True)
select_parser.add_argument('--shard', type=int, required=True, help='1-based.')
select_parser.add_argument('--exclude', action='append', default=[], help='A directory pytest ignores anyway.')
report_parser = subparsers.add_parser('report', help='Summarize new slow tests and slowdowns against a baseline.')
report_parser.add_argument('--baseline', type=Path, required=True)
report_parser.add_argument('--current', type=Path, required=True)
report_parser.add_argument('--title', required=True)
args = parser.parse_args(argv)
if args.command == 'fetch':
# Balancing is an optimization: never fail the run over it. An empty baseline splits by size.
durations: Durations = {}
try:
repo = os.environ['GITHUB_REPOSITORY']
this_run = json.loads(gh_api(f'repos/{repo}/actions/runs/{os.environ["GITHUB_RUN_ID"]}'))
found = fetch(repo, args.prefix, args.shards, before=this_run['created_at'])
except (KeyError, OSError, subprocess.CalledProcessError, ValueError, zipfile.BadZipFile) as exc:
print(f'::warning::Could not fetch `{args.prefix}*` test durations, splitting by module size: {exc!r}')
else:
if found is None:
print(f'No `main` run has `{args.prefix}*` durations yet; splitting by module size.')
else:
run_id, durations = found
print(f'Fetched `{args.prefix}*` durations of {len(durations)} tests from run {run_id}.')
args.output.write_text(json.dumps(durations, sort_keys=True), encoding='utf-8')
return 0
if args.command == 'select':
root = Path.cwd()
durations = _load(args.durations)
weights = module_weights(find_test_modules(root, args.exclude), durations, root)
assigned = assign(weights, args.shards)
unit = 's' if durations else ' bytes of source (no durations baseline)'
for shard, paths in enumerate(assigned, start=1):
marker = '*' if shard == args.shard else ' '
total = sum(weights[path] for path in paths)
print(f'{marker} shard {shard}: {len(paths)} modules, estimated {total:.0f}{unit}', file=sys.stderr)
for shard, paths in enumerate(assigned, start=1):
if shard != args.shard:
print('\n'.join(f'--ignore={path}' for path in paths))
return 0
baseline = _load(args.baseline)
if not baseline:
print('No baseline durations to compare with.')
return 0
summary = render(compare(baseline, _load(args.current)), args.title)
print(summary or 'No new slow tests or slowdowns.')
if summary and (summary_path := os.environ.get('GITHUB_STEP_SUMMARY')):
with open(summary_path, 'a', encoding='utf-8') as f:
f.write(summary + '\n')
return 0
if __name__ == '__main__':
sys.exit(main())
.github/scripts/test_shard_durations.py
+158
View File
@@ -0,0 +1,158 @@
from __future__ import annotations
import io
import json
import sys
import zipfile
from collections.abc import Sequence
from pathlib import Path
from typing import Any
import pytest
sys.path.insert(0, str(Path(__file__).parent))
import shard_durations
REPO = 'pydantic/pydantic-ai'
PREFIX = 'test-durations-3.12-all-extras-'
def _zip(durations: dict[str, float]) -> bytes:
buffer = io.BytesIO()
with zipfile.ZipFile(buffer, 'w') as archive:
archive.writestr('test-durations.json', json.dumps(durations))
return buffer.getvalue()
def _artifact(artifact_id: int, name: str) -> dict[str, Any]:
return {'id': artifact_id, 'name': name, 'expired': False}
class FakeApi:
def __init__(self, runs: dict[int, list[dict[str, Any]]], zips: dict[int, bytes]) -> None:
self.runs = runs
self.zips = zips
self.paths: list[str] = []
def __call__(self, path: str) -> bytes:
self.paths.append(path)
parts = path.split('?')[0].split('/')
if parts[3:5] == ['actions', 'workflows']:
return json.dumps({'workflow_runs': [{'id': run_id} for run_id in self.runs]}).encode()
if parts[3:5] == ['actions', 'runs']:
return json.dumps({'artifacts': self.runs[int(parts[5])]}).encode()
return self.zips[int(parts[5])]
def test_fetch_merges_shards_of_newest_complete_main_run():
api = FakeApi(
runs={
# Newest first: a run still missing its second shard, then the run to use.
300: [_artifact(30, f'{PREFIX}1')],
200: [
_artifact(20, f'{PREFIX}1'),
_artifact(21, f'{PREFIX}2'),
_artifact(22, 'coverage-3.12-all-extras-1'),
],
100: [_artifact(10, f'{PREFIX}1'), _artifact(11, f'{PREFIX}2')],
},
zips={20: _zip({'a': 1.0}), 21: _zip({'b': 2.0})},
)
assert shard_durations.fetch(REPO, PREFIX, 2, '2026-10-09T12:00:00Z', api) == (200, {'a': 1.0, 'b': 2.0})
assert api.paths[0] == (
f'repos/{REPO}/actions/workflows/ci.yml/runs?branch=main&event=push&created=%3C2026-10-09T12:00:00Z&per_page=20'
)
def test_fetch_without_a_complete_run():
api = FakeApi(runs={300: [_artifact(30, f'{PREFIX}1')]}, zips={})
assert shard_durations.fetch(REPO, PREFIX, 2, '2026-10-09T12:00:00Z', api) is None
def test_compare_flags_new_slow_and_slower_tests():
baseline = {'steady': 1.0, 'doubled': 1.0, 'tiny': 0.1, 'gone': 5.0}
current = {'steady': 1.5, 'doubled': 2.5, 'tiny': 0.9, 'new_fast': 0.5, 'new_slow': 3.0}
report = shard_durations.compare(baseline, current)
assert report == shard_durations.Report(new=[('new_slow', 3.0)], slower=[('doubled', 1.0, 2.5)])
rendered = shard_durations.render(report, '3.12 shard 1')
assert '| `new_slow` | 3.0 |' in rendered
assert '| `doubled` | 1.0 | 2.5 |' in rendered
assert shard_durations.render(shard_durations.Report(new=[], slower=[]), 'x') == ''
def test_report_writes_the_step_summary(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
baseline = tmp_path / 'baseline.json'
current = tmp_path / 'current.json'
summary = tmp_path / 'summary.md'
baseline.write_text(json.dumps({'a': 1.0}))
current.write_text(json.dumps({'a': 1.0, 'b': 4.0}))
monkeypatch.setenv('GITHUB_STEP_SUMMARY', str(summary))
args = ['report', '--baseline', str(baseline), '--current', str(current), '--title', 't']
assert shard_durations.main(args) == 0
assert '| `b` | 4.0 |' in summary.read_text()
def test_fetch_failure_does_not_fail_the_job(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
):
monkeypatch.delenv('GITHUB_REPOSITORY', raising=False)
output = tmp_path / 'durations.json'
assert shard_durations.main(['fetch', '--prefix', PREFIX, '--shards', '2', '--output', str(output)]) == 0
# Every shard reads this file, so an empty one still makes them all fall back the same way.
assert json.loads(output.read_text()) == {}
assert 'splitting by module size' in capsys.readouterr().out
def test_module_weights_sum_tests_and_estimate_new_modules(tmp_path: Path):
durations = {
'tests/test_a.py::test_one': 1.0,
'tests/test_a.py::test_two[x::y]': 2.0,
'tests/test_b.py::TestB::test_one': 5.0,
'tests/test_c.py::test_one': 7.0,
'tests/test_removed.py::test_one': 100.0,
}
modules = ['tests/test_a.py', 'tests/test_b.py', 'tests/test_c.py', 'tests/test_new.py']
assert shard_durations.module_weights(modules, durations, tmp_path) == {
'tests/test_a.py': 3.0,
'tests/test_b.py': 5.0,
'tests/test_c.py': 7.0,
'tests/test_new.py': 5.0,
}
(tmp_path / 'tests').mkdir()
(tmp_path / 'tests' / 'test_a.py').write_text('x' * 10)
assert shard_durations.module_weights(['tests/test_a.py'], {}, tmp_path) == {'tests/test_a.py': 10.0}
def test_assign_puts_every_module_in_exactly_one_balanced_shard():
weights = {'tests/test_a.py': 9.0, 'tests/test_b.py': 5.0, 'tests/test_c.py': 4.0, 'tests/test_d.py': 4.0}
assert shard_durations.assign(weights, 2) == [
['tests/test_a.py', 'tests/test_d.py'],
['tests/test_b.py', 'tests/test_c.py'],
]
def test_select_ignores_the_other_shards_modules(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
):
durations = tmp_path / 'durations.json'
durations.write_text(json.dumps({'tests/test_slow.py::test': 10.0, 'tests/test_fast.py::test': 1.0}))
def find_test_modules(root: Path, exclude: Sequence[str]) -> list[str]:
return ['tests/test_fast.py', 'tests/test_slow.py']
monkeypatch.setattr(shard_durations, 'find_test_modules', find_test_modules)
args = ['select', '--durations', str(durations), '--shards', '2', '--shard', '1']
assert shard_durations.main(args) == 0
out, err = capsys.readouterr()
assert out.split() == ['--ignore=tests/test_fast.py']
assert '* shard 1: 1 modules, estimated 10s' in err
def test_find_test_modules_lists_this_repo():
root = Path(__file__).parents[2]
modules = shard_durations.find_test_modules(root, ['tests/durable_exec'])
assert 'tests/test_agent.py' in modules
assert 'tests/conftest.py' not in modules
assert not any(path.startswith('tests/durable_exec/') for path in modules)
.github/workflows/ci.yml
+118 -11
View File
@@ -230,6 +230,7 @@ jobs:
.github/scripts/test_docs_navigation_workflow.py
.github/scripts/test_protect_github_dir.py
.github/scripts/test_ci_content_classifier.py
.github/scripts/test_shard_durations.py
# `make typecheck-changed` narrows Pyright to the files a change reaches, so a wrong
# answer here silently stops type errors reaching a developer's terminal.
@@ -629,12 +630,54 @@ jobs:
path: .coverage
include-hidden-files: true
# The sharded jobs below assign whole test modules to shards by the per-test durations their
# shards uploaded on the latest `main` run (see `.github/scripts/shard_durations.py`). Fetching
# that baseline once here means every shard of this run splits from the same file, so each module
# runs exactly once even if the API fails (the files are then empty and the split falls back to
# module size). Keep the shard counts in step with `SPLITS` in the jobs below. It runs alongside
# `classify`, not after it, to stay off the critical path, so it also runs (and its result goes
# unused) on content-only and clai2-only changes, where the sharded jobs are skipped. That costs a
# few API calls; don't add `needs: classify` to save them.
test-durations-baseline:
name: fetch test durations baseline
runs-on: ubuntu-latest
timeout-minutes: 10
permissions:
contents: read
actions: read
steps:
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
sparse-checkout: .github/scripts
- name: fetch the latest main run's test durations
env:
GH_TOKEN: ${{ github.token }}
run: |
mkdir test-durations
for python in 3.11 3.12 3.13 3.14; do
all_extras_shards=2
if [[ "$python" == 3.11 ]]; then all_extras_shards=3; fi
python3 .github/scripts/shard_durations.py fetch --prefix "test-durations-$python-all-extras-" \
--shards "$all_extras_shards" --output "test-durations/$python-all-extras.json"
python3 .github/scripts/shard_durations.py fetch --prefix "test-durations-$python-lowest-versions-" \
--shards 2 --output "test-durations/$python-lowest-versions.json"
done
- uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: test-durations-baseline
path: test-durations
if-no-files-found: error
retention-days: 7
# `all-extras` is the slowest install shape by a wide margin -- it alone pulls in every
# optional dependency -- so it gets its own job, split into shards, instead of sharing the
# `test` job's timeout budget with the three much lighter install shapes.
test-all-extras:
name: test on ${{ matrix.python-version }} (all-extras, shard ${{ matrix.shard }}/${{ matrix.python-version == '3.11' && 3 || 2 }})
needs: classify
needs: [classify, test-durations-baseline]
if: needs.classify.outputs.content_only != 'true' && needs.classify.outputs.clai2_only != 'true'
runs-on: ubuntu-latest
timeout-minutes: 20
@@ -652,11 +695,24 @@ jobs:
# The 3.11 legacy tracer needs a third shard to stay within the job timeout.
COLLECT_COVERAGE: ${{ matrix.python-version == '3.14' || matrix.python-version == '3.11' }}
SPLITS: ${{ matrix.python-version == '3.11' && 3 || 2 }}
DURATIONS_PREFIX: test-durations-${{ matrix.python-version }}-all-extras-
DURATIONS_BASELINE: test-durations-baseline/${{ matrix.python-version }}-all-extras.json
# Rebenchmark after https://github.com/cbornet/blockbuster/pull/61 is released in a
# compatible version; re-enable if a shard then stays comfortably within its timeout.
BLOCKBUSTER_ENABLED: false
# `tests/test_examples.py` type-checks the docs examples with pyright; once per run is enough.
TYPECHECK_EXAMPLES: ${{ matrix.python-version == '3.14' }}
# Per-test time budget enforced by `tests/cost_guards.py` (see "Test cost" in `tests/AGENTS.md`). The 3.13
# cells run nearly every test without coverage or blocking-call detection, so their timings carry the least
# instrumentation noise. 8s sits above the fresh-interpreter tests (4-6s on these runners) with headroom for runner
# jitter; the per-test report on the 3.12 cells surfaces anything new over 1s. List the tests over a budget from a
# run's `--durations` with `scripts/list_over_budget_tests.py`.
PYTEST_TEST_BUDGET_SECONDS: ${{ matrix.python-version == '3.13' && '8' || '' }}
# Per-module collection budget, also enforced by `tests/cost_guards.py`. The first module in a shard to import a
# dependency many modules share (`anthropic`, `openai`, `hypothesis`) is charged for it, up to 5s on these cells;
# 8s sits above that with headroom for runner jitter, and well below the 11-16s a module paid for importing
# `torch` at collection.
PYTEST_COLLECT_BUDGET_SECONDS: ${{ matrix.python-version == '3.13' && '8' || '' }}
# `tests/harness/absurd` runs on Absurd's schema; an unreachable server fails rather than skips.
ABSURD_REQUIRE_LIVE: "1"
ABSURD_TEST_DATABASE_URL: postgresql://postgres:[email protected]:5432/postgres
@@ -736,11 +792,17 @@ jobs:
path: ~/.cache/huggingface/hub/models--sentence-transformers-testing--stsb-bert-tiny-safetensors
key: ${{ steps.sentence-transformers-cache.outputs.cache-primary-key }}
# `-n logical` (see note on the `test` job). Split into `SPLITS` shards with `pytest-split`
# so this, the slowest install shape, keeps pace with the rest of the matrix. The split is
# by test count, not measured duration -- no `.test_durations` file is committed -- so
# rebalance with `--store-durations` if `ci_duration.py`'s report shows the shards
# drifting apart. The 3.11 cells use three shards; every test runs once per version.
- uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
name: test-durations-baseline
path: ${{ runner.temp }}/test-durations-baseline
# `-n logical` (see note on the `test` job). Split into `SPLITS` shards so this, the slowest
# install shape, keeps pace with the rest of the matrix. `shard_durations.py select` assigns
# whole test modules to shards by the `test-durations-baseline` durations and `--ignore`s the
# other shards' modules, so a shard's workers never collect them (`pytest-split` would have
# every worker collect the whole suite, then deselect most of it). The 3.11 cells use three
# shards; every test module runs in exactly one shard per version.
# Everything under `tests/durable_exec` and every test marked `temporal` runs in
# `test-durable-exec` instead: those tests pin xdist groups around real local servers, so
# the chains they form would otherwise be this job's critical-path floor.
@@ -756,11 +818,35 @@ jobs:
fi
# `--all-packages` so the extras of every workspace member are installed, not just the
# root package's: that is where `pydantic-ai-harness`'s optional integrations get tested.
uv run --no-project python .github/scripts/shard_durations.py select --durations "$RUNNER_TEMP/$DURATIONS_BASELINE" \
--shards "$SPLITS" --shard ${{ matrix.shard }} --exclude tests/durable_exec > "$RUNNER_TEMP/shard-args.txt"
mapfile -t shard_args < "$RUNNER_TEMP/shard-args.txt"
# `--store-durations` (from `pytest-split`) records this shard's durations for the next run.
uv run --all-packages --all-extras --no-extra mcp-tasks "${runner[@]}" --durations=100 -n logical --dist=loadgroup \
--ignore=tests/durable_exec -m "not temporal" --splits "$SPLITS" --group ${{ matrix.shard }}
--ignore=tests/durable_exec -m "not temporal" "${shard_args[@]}" \
--store-durations --clean-durations --durations-path "$RUNNER_TEMP/test-durations.json"
env:
COVERAGE_FILE: ${{ github.workspace }}/.coverage/.coverage.${{ matrix.python-version }}-all-extras-${{ matrix.shard }}
# Only `main` push runs are ever read back as the baseline (`shard_durations.py fetch`).
- name: store test durations
if: ${{ !cancelled() && github.event_name == 'push' }}
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: ${{ env.DURATIONS_PREFIX }}${{ matrix.shard }}
path: ${{ runner.temp }}/test-durations.json
if-no-files-found: ignore
retention-days: 30
# Show authors what their change costs: new tests over a second and tests that got at least
# twice as slow as on `main`, in the job summary. One coverage-free version is enough.
- name: report new and slower tests
if: ${{ !cancelled() && matrix.python-version == '3.12' }}
run: >-
uv run --no-project python .github/scripts/shard_durations.py report
--baseline "$RUNNER_TEMP/$DURATIONS_BASELINE" --current "$RUNNER_TEMP/test-durations.json"
--title "all-extras on ${{ matrix.python-version }}, shard ${{ matrix.shard }}"
- name: cover Absurd without a database URL
if: matrix.python-version == '3.14' && matrix.shard == 1
run: >-
@@ -901,7 +987,7 @@ jobs:
test-lowest-versions:
name: test on ${{ matrix.python-version }} (lowest-versions, shard ${{ matrix.shard }}/2)
needs: classify
needs: [classify, test-durations-baseline]
if: needs.classify.outputs.content_only != 'true' && needs.classify.outputs.clai2_only != 'true'
runs-on: ubuntu-latest
timeout-minutes: 20
@@ -919,6 +1005,8 @@ jobs:
BLOCKBUSTER_ENABLED: false
# The docs examples are type-checked in the Python 3.14 all-extras jobs.
TYPECHECK_EXAMPLES: false
DURATIONS_PREFIX: test-durations-${{ matrix.python-version }}-lowest-versions-
DURATIONS_BASELINE: test-durations-baseline/${{ matrix.python-version }}-lowest-versions.json
steps:
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
@@ -973,8 +1061,13 @@ jobs:
path: ~/.cache/huggingface/hub/models--sentence-transformers-testing--stsb-bert-tiny-safetensors
key: ${{ steps.sentence-transformers-cache.outputs.cache-primary-key }}
# `-n logical` (see note on the `test` job). Split into 2 shards with `pytest-split`
# (see the matching note in `test-all-extras` for why, and for the rebalancing caveat).
- uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
name: test-durations-baseline
path: ${{ runner.temp }}/test-durations-baseline
# `-n logical` (see note on the `test` job). Split into 2 shards by test module (see the
# matching note in `test-all-extras` for why and how).
# The durable-execution suites and harness `temporal` tests run in `test-durable-exec`
# (its `lowest` leg covers this resolution), keeping their grouped xdist chains off this
# job's critical path.
@@ -988,11 +1081,25 @@ jobs:
unset COVERAGE_PROCESS_START
runner=(pytest)
fi
uv run --no-project python .github/scripts/shard_durations.py select --durations "$RUNNER_TEMP/$DURATIONS_BASELINE" \
--shards 2 --shard ${{ matrix.shard }} --exclude tests/durable_exec > "$RUNNER_TEMP/shard-args.txt"
mapfile -t shard_args < "$RUNNER_TEMP/shard-args.txt"
uv run --no-sync "${runner[@]}" --durations=100 -n logical --dist=loadgroup \
--ignore=tests/durable_exec -m "not temporal" --splits 2 --group ${{ matrix.shard }}
--ignore=tests/durable_exec -m "not temporal" "${shard_args[@]}" \
--store-durations --clean-durations --durations-path "$RUNNER_TEMP/test-durations.json"
env:
COVERAGE_FILE: ${{ github.workspace }}/.coverage/.coverage.${{ matrix.python-version }}-lowest-versions-${{ matrix.shard }}
# Only `main` push runs are ever read back as the baseline (`shard_durations.py fetch`).
- name: store test durations
if: ${{ !cancelled() && github.event_name == 'push' }}
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: ${{ env.DURATIONS_PREFIX }}${{ matrix.shard }}
path: ${{ runner.temp }}/test-durations.json
if-no-files-found: ignore
retention-days: 30
# See the matching step in the `test` job.
- name: combine coverage files
if: env.COLLECT_COVERAGE == 'true'
.github/workflows/shared/prompts/pydantic-ai-pr-review.md
+8 -1
View File
@@ -53,7 +53,7 @@ the rest briefly in the review body.
## Targeted checks
Beyond free-form review, run this check proactively on every model-touching PR:
Beyond free-form review, run these checks proactively; the first applies to every model-touching PR, the second to every PR that touches tests:
- **Profile-flag test pinning.** When the diff adds or modifies a
`profile.get(...)` read under `pydantic_ai_slim/pydantic_ai/models/`, verify
@@ -63,6 +63,13 @@ Beyond free-form review, run this check proactively on every model-touching PR:
pair are in `.review-context/review-instructions.md` (Example 4).
Cap this finding at **MEDIUM**: it is advisory and must never drive
`REQUEST_CHANGES`.
- **Test cost.** When the diff adds or changes tests, check every new
`@pytest.mark.subprocess(reason=...)` and `@pytest.mark.slow(reason=...)`:
the reason must name a property that needs the process boundary or the
duration, not just describe what the test does. Also flag new tests that wait
out real sleeps, timeouts or poll intervals instead of injecting them, or that
feed one huge input to prove a complexity bound. Calibration is Example 5 in
`.review-context/review-instructions.md`. Cap at **MEDIUM**.
## Review process
AGENTS.md
+2 -2
View File
@@ -64,7 +64,7 @@ All changes need to:
- leave behavior unchanged for users who aren't hitting the problem you are solving. Pydantic AI's surfaces are nested, not peer: everyone using a durable execution engine, a single provider, a UI adapter, or an opt-in capability is also a plain `Agent` user, but not the reverse -- so a fix motivated by a narrow surface must not move observable behavior on a wider one. Documenting it doesn't make it acceptable, only expected; if no scoping leaves the wider surface untouched, that's a design question for maintainers, not a docs note. (Fixing the same *confirmed* defect in sibling providers, fields, or models in the same PR is the previous point, not this one.)
- be backward compatible as laid out in the [version policy](docs/version-policy.md), so that users can upgrade with confidence. A default doesn't make a new parameter non-breaking: if existing code only keeps its current behavior once the user passes a new argument (`id=None`, `strict=False`), that's a break, and it leaves behind a permanent parameter whose only purpose is to say "keep doing what you did". The default must preserve existing behavior; the new behavior is what users opt into
- be fully type-safe (both internally and in public API) without unnecessary `cast`s or `Any`s, so that users don't need `isinstance` checks and can trust that code that typechecks will work at runtime
- have comprehensive tests covering 100% of code paths, favoring integration tests and real requests (using recordings and snapshots -- see below) over unit tests and mocking
- have comprehensive tests covering 100% of code paths, favoring integration tests and real requests (using recordings and snapshots -- see below) over unit tests and mocking; "integration" means real in-process code paths with recorded requests, not process spawns or real waits (see "Test cost" in [tests/AGENTS.md](tests/AGENTS.md))
- update/add all relevant documentation, following the existing voice and patterns
- update the relevant agent skills when introducing a new feature or when a skill needs to reflect the correct mechanics; Pydantic AI skills belong in [pydantic_ai_slim/pydantic_ai/.agents/skills/building-pydantic-ai-agents/](pydantic_ai_slim/pydantic_ai/.agents/skills/building-pydantic-ai-agents/), Pydantic AI Harness skills in [src/pydantic_ai_harness/pydantic_ai_harness/.agents/skills/pydantic-ai-harness/](src/pydantic_ai_harness/pydantic_ai_harness/.agents/skills/pydantic-ai-harness/), while repository workflow skills live under [.claude/skills/](.claude/skills/)
@@ -122,7 +122,7 @@ The project uses:
## Test process isolation
Do not use subprocesses to test logic that can run in-process. Starting another interpreter and importing the project again makes the full test suite materially slower. Reserve subprocesses for behavior that depends on the process boundary, such as CLI invocation, interpreter startup, or import isolation.
Do not use subprocesses to test logic that can run in-process. Starting another interpreter and importing the project again makes the full test suite materially slower. Reserve subprocesses for behavior that depends on the process boundary, such as interpreter startup, import isolation, or signal handling; call a CLI's entry point in-process. `tests/cost_guards.py` fails a test that launches a Python interpreter unless it is marked `@pytest.mark.subprocess(reason=...)`; see "Test cost" in [tests/AGENTS.md](tests/AGENTS.md).
## When to verify
pydantic_ai_slim/pydantic_ai/realtime/_core.py
+6 -35
View File
@@ -113,16 +113,14 @@ class InputWithdrawn:
@dataclass(frozen=True, kw_only=True)
class AudioSent:
"""Input audio the caller streamed, retained for the spoken turn it belongs to."""
"""Input audio the caller streamed went out, retained for the spoken turn it belongs to.
Reported once its send completed: a turn the provider ended while it was on its way was ended without it.
"""
data: bytes
@dataclass(frozen=True, kw_only=True)
class AudioUnsent:
"""The input audio last reported as sent never went out: it is taken back, as much of it as is still there."""
@dataclass(frozen=True, kw_only=True)
class AudioCleared:
"""The caller discarded the buffered input audio."""
@@ -185,7 +183,6 @@ Command: TypeAlias = (
InputSent
| InputWithdrawn
| AudioSent
| AudioUnsent
| AudioCleared
| ToolReturned
| ToolCallRefused
@@ -344,9 +341,6 @@ class SessionCore:
self._recorded_audio: deque[_Response | _UserTurn] = deque()
"""Recorded entities whose message still has audio, in the order they joined the conversation: what eviction
takes audio from first, oldest first, without passing over anything that has none or isn't recorded yet."""
self._unsent_audio: tuple[_UserTurn | None, int] | None = None
"""The latest input audio chunk, until the next: the turn it was cut for since (or `None` while it is in the
buffer), and its length, so `AudioUnsent` can take it back if its send fails."""
self._placements = 0
self._user_speaking = False
self._speech_segmented = False
@@ -433,7 +427,7 @@ class SessionCore:
self._input_sent(command)
elif isinstance(command, InputWithdrawn):
self._withdraw(command.input_ids)
elif isinstance(command, (AudioSent, AudioUnsent, AudioCleared)):
elif isinstance(command, (AudioSent, AudioCleared)):
self._input_audio_command(command)
elif isinstance(command, ToolReturned):
self._returns[command.tool_call_id] = command.request
@@ -458,38 +452,17 @@ class SessionCore:
else:
assert_never(command)
def _input_audio_command(self, command: AudioSent | AudioUnsent | AudioCleared) -> None:
def _input_audio_command(self, command: AudioSent | AudioCleared) -> None:
if isinstance(command, AudioSent):
if self._retain_input:
self._input_audio.extend(command.data)
self._unsent_audio = (None, len(command.data))
self._bound_retained_audio()
elif isinstance(command, AudioUnsent):
self._take_back_audio()
else:
self._unsent_audio = None
self._input_audio.clear()
# A turn the user was saying that hadn't joined the conversation never will now.
for turn_id in [turn_id for turn_id, turn in self._turns.items() if not turn.ended]:
self._drop_turn(turn_id)
def _take_back_audio(self) -> None:
"""Take back the latest input audio chunk, whose send failed: from the buffer, or the turn it was cut for.
By length, as the session takes it back from its own buffer: what of it the budget already trimmed, or
that went into history with its turn, is no longer there to take back.
"""
if self._unsent_audio is None:
return
turn, length = self._unsent_audio
self._unsent_audio = None
if turn is None:
del self._input_audio[len(self._input_audio) - min(length, len(self._input_audio)) :]
elif turn.id in self._waiting_turn_audio and turn.audio:
length = min(length, len(turn.audio))
turn.audio = turn.audio[: len(turn.audio) - length]
self._waiting_turn_audio_bytes -= length
def _input_sent(self, command: InputSent) -> None:
if command.request is not None:
self._inputs[command.input_id] = self._unplaced[command.input_id] = _Input(
@@ -726,8 +699,6 @@ class SessionCore:
self._input_audio.clear()
self._waiting_turn_audio[turn.id] = turn
self._waiting_turn_audio_bytes += len(turn.audio)
if self._unsent_audio is not None and self._unsent_audio[0] is None:
self._unsent_audio = (turn, self._unsent_audio[1])
def _release_turn_audio(self, turn: _UserTurn) -> None:
if self._waiting_turn_audio.pop(turn.id, None) is not None and turn.audio:
pydantic_ai_slim/pydantic_ai/realtime/_session.py
+18 -24
View File
@@ -86,7 +86,6 @@ from ..usage import RequestUsage, RunUsage, UsageLimits
from ._core import (
AudioCleared,
AudioSent,
AudioUnsent,
Closed,
CoreInput,
ExchangeAbandoned,
@@ -2075,27 +2074,18 @@ class RealtimeSession:
audio_was_uncommitted = self._audio_uncommitted
self._held_commit.has_audio = True
self._open_user_turn_for_audio()
previous_length: int | None = None
if self._retain_input:
# Buffer the raw input so the finalized user turn can retain it. A per-item speech-stopped
# boundary later cuts this into that turn's own segment (see `_segment_input_audio`); only the
# exact split at the boundary is approximate (see `audio_retention`).
previous_length = len(self._input_audio)
retained: bytes | None = None
try:
# The extend lives inside the try so a chunk that is not bytes-like (possible since
# `send_audio` accepts an async iterable of chunks) rolls the user-turn state back below
# instead of leaving `_user_turn_active` set with nothing buffered or sent. A failing
# `bytearray.extend` leaves the buffer unchanged, so the length guard in the handler
# correctly skips the truncation for that case.
# Copied inside the try so a chunk that is not bytes-like (possible since `send_audio` accepts
# an async iterable of chunks) fails before it goes out and rolls the user-turn state back
# below, instead of leaving `_user_turn_active` set with nothing buffered or sent.
if self._retain_input:
self._input_audio.extend(data)
retained = memoryview(data).tobytes()
self._audio_uncommitted = True
await self._send_frame(BinaryAudio(data=data, media_type='audio/pcm'))
except BaseException as e:
self._user_turn_active = user_turn_was_active
self._audio_uncommitted = audio_was_uncommitted
if previous_length is not None and len(self._input_audio) == previous_length + len(data):
del self._input_audio[previous_length:]
if (
isinstance(e, RealtimeError)
and isinstance(e.__cause__, self._connection.transport_errors)
@@ -2108,7 +2098,14 @@ class RealtimeSession:
# the reconnect fails, the next chunk raises again.
return
raise
if self._retain_input:
if retained is not None:
# Buffer the raw input so the finalized user turn can retain it, once it has gone out (with no
# await since, so still in send order), as the session core takes it: a boundary the pump handled
# while the chunk was on its way was drawn without it, and a chunk whose send failed belongs to
# no turn. A per-item speech-stopped boundary later cuts this into that turn's own segment (see
# `_segment_input_audio`); only the exact split at the boundary is approximate (see
# `audio_retention`).
self._input_audio.extend(retained)
self._bound_retained_audio()
def _open_user_turn_for_audio(self) -> None:
@@ -2408,7 +2405,6 @@ class RealtimeSession:
# which is a caller that asked to send.
async with self._send_lock:
first_input = self._inputs_sent
audio_sent_to_core = False
try:
for position, content in enumerate(contents):
# Numbered before the call, and whether or not it raises, matching how
@@ -2426,17 +2422,15 @@ class RealtimeSession:
tool_call_id=content.tool_call_id if isinstance(content, ToolResult) else None,
)
)
if isinstance(content, BinaryAudio):
# Under the lock, so the core takes audio in the order it goes out and a failed send
# takes back its own chunk; before the send, as the pump can end the turn meanwhile.
self._apply_core(AudioSent(data=bytes(content.data)))
audio_sent_to_core = True
await self._connection.send(content)
if isinstance(content, BinaryAudio):
# Under the lock, so the core takes audio in the order it goes out; only once it has
# gone out, as a turn the provider ended while it was on its way (which the pump can
# record meanwhile) was ended without it, and a chunk whose send fails is in no turn.
self._apply_core(AudioSent(data=bytes(content.data)))
except BaseException as e:
# The caller takes back what it sent, so the core does too.
self._apply_core(InputWithdrawn(input_ids=tuple(range(first_input, self._inputs_sent))))
if audio_sent_to_core:
self._apply_core(AudioUnsent())
if not isinstance(e, self._connection.transport_errors):
raise
# A send that fails because the link is gone is the same failure the receive side
pyproject.toml
+28
View File
@@ -355,6 +355,7 @@ extend-select = [
"I",
"D",
"TID251",
"TID253",
"PLW1514", # https://docs.astral.sh/ruff/rules/unspecified-encoding/
"PGH003", # https://docs.astral.sh/ruff/rules/blanket-type-ignore/
"BLE001", # https://docs.astral.sh/ruff/rules/blind-except/
@@ -376,6 +377,23 @@ known-first-party = ["pydantic_ai", "pydantic_evals", "pydantic_graph", "pydanti
[tool.ruff.lint.pydocstyle]
convention = "google"
[tool.ruff.lint.flake8-tidy-imports]
# Test modules are imported during collection, once per `pytest-xdist` worker, so a heavy dependency imported at
# module level is paid by every test run that collects the module, whether or not its tests are selected. Each of
# these adds 100ms or more (measured with `python -X importtime` on top of what `tests/conftest.py` already imports);
# import them inside the test or fixture that needs them. Enforced for `tests/` only, see `per-file-ignores`.
banned-module-level-imports = [
"dbos",
"e2b",
"prefect",
"sentence_transformers",
"sklearn",
"temporalio",
"torch",
"transformers",
"voyageai",
]
[tool.ruff.lint.flake8-tidy-imports.banned-api]
"typing.TypedDict".msg = "Use typing_extensions.TypedDict instead."
"asyncio.Lock".msg = "Use anyio.Lock instead."
@@ -390,6 +408,8 @@ quote-style = "single"
"tests/**/*.py" = ["D"]
"docs/**/*.py" = ["D"]
".github/scripts/ci_duration.py" = ["D"]
".github/scripts/shard_durations.py" = ["D"]
".github/scripts/test_shard_durations.py" = ["D"]
".github/scripts/check_api_compatibility.py" = ["D"]
".github/scripts/test_ci_duration.py" = ["D"]
".github/scripts/test_check_api_compatibility.py" = ["D"]
@@ -406,6 +426,12 @@ quote-style = "single"
".github/scripts/test_ci_content_classifier.py" = ["D"]
"scripts/test_typecheck_changed.py" = ["D"]
"!src/pydantic_clai2/**/*.py" = ["BLE001"]
"!tests/**/*.py" = ["TID253"]
# Modules dedicated to one of the heavy dependencies `TID253` bans need it at module level, e.g. for Temporal
# workflow definitions; the cost is paid only where that dependency is what is being tested.
"tests/durable_exec/**/*.py" = ["TID253"]
"tests/**/*{temporal,dbos,prefect,durable}*.py" = ["TID253"]
"tests/harness/e2b_sandbox/*.py" = ["TID253"]
[tool.pyright]
pythonVersion = "3.11"
@@ -463,6 +489,8 @@ markers = [
"e2b_live: runs against a real E2B sandbox; needs credentials",
"sprites_live: runs against a real Sprites control plane; needs credentials",
"temporal: starts or connects to a local Temporal dev server",
"subprocess(reason, conditional=False): launches a Python interpreter, which the cost guard otherwise fails; reason required. `conditional=True` skips the stale-marker check for a test that only launches one on some platforms or installs",
"slow(reason): exempt from the per-test time budget (`PYTEST_TEST_BUDGET_SECONDS`); reason required",
]
filterwarnings = [
"error",
scripts/gather-pydantic-ai-review-context.sh
+34
View File
@@ -577,6 +577,40 @@ value but a comment in the added test explains why the other side is
unreachable on this path, or when the touched `profile.get(...)` read is not a
new or modified branch (a pure move or rename).
### Example 5 — test cost
Trigger: the diff adds or changes tests under `tests/`. Every test runs on
every push across the whole CI matrix, and the next test in the file copies the
pattern, so cost is a review concern even when CI is green. The rules are in
"Test cost" in `tests/AGENTS.md`; `tests/cost_guards.py` already fails an
unmarked Python launch and (on one CI leg) a test over the time budget, so the
review's job is the escape hatches and the patterns the guards cannot see.
**Flag this (MEDIUM):**
```python
@pytest.mark.subprocess(reason='runs the CLI')
def test_cli_prints_version(tmp_path: Path):
result = subprocess.run([sys.executable, '-m', 'pydantic_clai2', '--version'], ...)
```
*Why:* The claim (the version string) does not depend on the process
boundary, so the marker only waves through a second-per-test interpreter
launch; ask for the entry point to be called in-process. The same goes for a
new `@pytest.mark.slow(reason=...)` whose cost could be designed away, and for
a new test that waits out a real `sleep`/timeout/poll interval instead of
injecting it, or proves a complexity bound with one huge input instead of a
growth ratio across small sizes.
**Don't flag this:**
```python
@pytest.mark.subprocess(reason='asserts what a fresh interpreter imports')
def test_import_is_light():
...
```
*Why:* What a fresh interpreter imports can only be observed in a fresh
interpreter; the reason names the process-boundary property under test. Also
don't flag a launch in a `session`- or `module`-scoped fixture (shared cost,
the recommended pattern) or a short real wait whose duration is the claim.
## Sub-agent finding format
When a Task sub-agent returns findings, use this exact format (one block
scripts/list_over_budget_tests.py
+72
View File
@@ -0,0 +1,72 @@
"""List tests over the per-test time budget, from the `--durations` output of a CI run.
`tests/cost_guards.py` fails a test whose setup and call take longer than `PYTEST_TEST_BUDGET_SECONDS` on the CI leg
that sets it. Run this against a CI run to see which tests that leg would fail, before turning the budget on or
lowering it, and mark the ones whose cost is inherent with `@pytest.mark.slow(reason=...)`:
uv run scripts/list_over_budget_tests.py <run-id> [--budget 8] [--jobs 'all-extras']
It reads the jobs' logs with the `gh` CLI. The durations pytest reports include setup of fixtures shared beyond one
test, which the guard does not count, so a test listed here only for its setup may be within budget; and pytest
only prints the slowest `--durations=N` phases per job, so the list is complete only down to the Nth slowest.
"""
from __future__ import annotations
import argparse
import json
import re
import subprocess
from collections import defaultdict
ANSI_ESCAPE = re.compile(r'\x1b\[[0-9;]*m')
DURATION = re.compile(r'\s(?P<seconds>\d+\.\d+)s (?P<phase>setup|call|teardown)\s+(?P<nodeid>\S+::\S+)')
def job_log(job_id: int) -> str:
"""One job's log, without the terminal colors pytest adds."""
endpoint = f'repos/{{owner}}/{{repo}}/actions/jobs/{job_id}/logs'
log = subprocess.check_output(['gh', 'api', '--allow-escape-sequences', endpoint], text=True)
return ANSI_ESCAPE.sub('', log)
def job_logs(run_id: str, job_filter: str) -> dict[str, str]:
"""The logs of the run's jobs whose names match `job_filter`, by job name."""
jobs = json.loads(subprocess.check_output(['gh', 'run', 'view', run_id, '--json', 'jobs'], text=True))['jobs']
pattern = re.compile(job_filter)
return {job['name']: job_log(job['databaseId']) for job in jobs if pattern.search(job['name'])}
def over_budget(log: str, budget: float) -> dict[str, tuple[float, float]]:
"""`nodeid -> (setup, call)` seconds for each test whose setup and call together exceed `budget`."""
phases: dict[str, dict[str, float]] = defaultdict(dict)
for match in DURATION.finditer(log):
phases[match['nodeid']][match['phase']] = float(match['seconds'])
return {
nodeid: (times.get('setup', 0.0), times.get('call', 0.0))
for nodeid, times in phases.items()
if times.get('setup', 0.0) + times.get('call', 0.0) > budget
}
def main() -> None:
"""Print each test over the budget, slowest first."""
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('run_id', help='a GitHub Actions run of the CI workflow')
parser.add_argument(
'--budget', type=float, default=8.0, help='seconds of setup and call per test (default: 8, as CI sets)'
)
parser.add_argument('--jobs', default=r'3\.13 \(all-extras', help='regex for the job names to read')
args = parser.parse_args()
found: dict[str, tuple[float, float]] = {}
for name, log in job_logs(args.run_id, args.jobs).items():
tests = over_budget(log, args.budget)
print(f'{name}: {len(tests)} over {args.budget:g}s')
found.update(tests)
for nodeid, (setup, call) in sorted(found.items(), key=lambda item: -sum(item[1])):
print(f'{setup + call:7.2f}s (setup {setup:.2f}s, call {call:.2f}s) {nodeid}')
if __name__ == '__main__':
main()
src/pydantic_clai2/pydantic_clai2/builtin_plugins/herdr.py
+3 -1
View File
@@ -29,6 +29,8 @@ from pydantic_clai2.runtime._session import Session
from pydantic_clai2.ui.rendering.usage_report import session_usage
_LOGGER = logging.getLogger(__name__)
TITLE_POLL_INTERVAL = 2.0
"""Seconds between checks for a session title renamed outside this shell."""
class _Reporter:
@@ -103,7 +105,7 @@ class _Reporter:
async def watch(self) -> None:
while True:
await asyncio.sleep(2)
await asyncio.sleep(TITLE_POLL_INTERVAL)
await self.refresh()
async def stop(self, event: SessionEnd) -> None:
src/pydantic_clai2/pydantic_clai2/plugins/loader.py
+3 -1
View File
@@ -95,6 +95,8 @@ _RETIRED_BUILTINS: dict[str, PluginSettings] = {
"""Former built-in declarations. A stored copy of one loads the built-in now declared under its id."""
FAILED_SESSION_END_TIMEOUT = 5.0
"""Seconds a plugin that failed to load may spend in `on_session_end` before loading continues without it."""
TURN_NOTICE = 'The running turn keeps the plugins it started with; the agent sees the change on the next prompt.'
@@ -899,7 +901,7 @@ def _same_plugin(declaration: PluginSettings, shipped: PluginSettings | None) ->
async def _end_failed_session(plugin: Plugin[BaseModel, DepsT]) -> BaseException | None:
"""Return the handler's failure rather than raising it."""
try:
with fail_after(5):
with fail_after(FAILED_SESSION_END_TIMEOUT):
await plugin.on_session_end(SessionEnd(reason='error'))
except (Exception, asyncio.CancelledError) as exc: # noqa: BLE001 -- reported by the caller.
return exc
tests/AGENTS.md
+19 -1
View File
@@ -274,7 +274,25 @@ Available to every test under `tests/models/`. See "Asserting what goes out on t
- When you deprecate a public symbol (add `@deprecated`), add a test asserting the warning fires — a `pytest.warns(PydanticAIDeprecationWarning, match=...)` block whose `match` pins both the symbol name and the migration guidance — so the deprecation is protected against accidental removal or a message rewrite, and any executable docstring/doc example that constructs the symbol is marked `{test="skip"}` (it would otherwise fail under `filterwarnings=["error"]`)
- Don't blanket-suppress `reportDeprecated` at file scope — no whole-file `# pyright: ignore` (or `# pyright: basic`) at the top of a file to kill every deprecation static-warning at once, since that also silences the next *unintended* use of a deprecated symbol in the file. Put `# pyright: ignore[reportDeprecated]` on each intentional call instead. (The runtime side is different and conventional: a module-level `warnings.filterwarnings('ignore', message=...)` scoped to one deprecation message, in a file that pervasively constructs the deprecated symbol it tests — see `tests/durable_exec/temporal/test_agent.py` / `tests/durable_exec/test_prefect.py` — is fine, as long as a dedicated `pytest.warns` test still proves the deprecation fires.)
- Gate on a missing optional dependency with `try_import()`, the default. It has two halves and they live in different files: each **test module** wraps its optional imports in `with try_import() as imports_successful:` and applies `pytestmark = pytest.mark.skipif(not imports_successful(), reason='<dep> not installed')` (see `tests/models/anthropic/test_native_tool_filtering.py`); a **`conftest.py`** wraps its own optional imports the same way but carries no skip of its own (see `tests/models/anthropic/conftest.py`) — a conftest's `pytestmark` is not inherited by the modules it serves, so the mark there would gate nothing while `try_import` still keeps the file importable. Reach for `pytest.skip(..., allow_module_level=True)` only when the module cannot import `tests/conftest.py` at all — the Temporal suite can't, because that would pull `cassetter` into a module the workflow sandbox re-imports. In that case the gate goes in the test modules **and any non-test helper they import** (`tests/durable_exec/temporal/_shared.py`), never in the `conftest.py`, which instead stays importable without the dependency by deferring its fixtures' imports into the fixture bodies. Reason: pytest loads the conftest of every command-line argument's directory before collection, and `_importconftest` catches only `Exception` while `Skipped` is a `BaseException`, so a gate there escapes as a traceback with exit 1 the moment someone names that directory (or a node id inside it) instead of a parent. See `tests/durable_exec/temporal/conftest.py` and the guard in `tests/durable_exec/test_temporal_skip_gating.py`
- Do not use subprocesses to test logic that can run in-process. Starting another interpreter and importing the project again makes the full suite materially slower. Reserve subprocesses for behavior that depends on the process boundary, such as CLI invocation, interpreter startup, or import isolation. Scrub `COVERAGE_*` from a required subprocess environment so `[tool.coverage.run] patch = ["subprocess"]` does not instrument the child.
- Keep each test cheap, following "Test cost" below
## Test cost
The suite runs on every push across a matrix of Python versions and install shapes, so a second added to one test is paid thousands of times a day. Coding agents copy the nearest helper, so one expensive pattern in a file soon spreads to every test in it. "Favor integration tests" means real in-process code paths with recorded requests, never process spawns, real waits or huge inputs.
- Call CLI and entry points in-process: invoke the `main()`/`cli()` function or the app object with arguments, `capsys`/`capfd` and `monkeypatch`, instead of `subprocess.run([sys.executable, '-m', ...])`. An interpreter re-imports the project and costs about a second per test
- Inject timeouts, poll intervals, retry delays and chunk sizes through a parameter, setting or `monkeypatch`, and order steps with `anyio.Event`s, so a test never waits out a real timer
- Prove complexity bounds (a regex that must not backtrack, a parser that must stay linear) at small sizes, not by feeding one huge input: patch the limit the code enforces down and shrink the input with it, or assert on operation counts or output size at two or three sizes. Avoid asserting on wall-clock ratios, which flake on shared runners
- Start an expensive server or engine once in a `session`- or `module`-scoped fixture and share it; under `pytest-xdist`, `@pytest.mark.xdist_group` keeps the tests that use it on one worker
- Import heavy optional dependencies (`torch`, `transformers`, `sentence_transformers`, `prefect`, `temporalio`, ...) inside the test or fixture that needs them unless the whole module is about that dependency: collection imports every test module once per worker. Ruff's `TID253` bans the heaviest at module level under `tests/`, and the end-of-run summary names the slowest modules to collect
- Scrub `COVERAGE_*` from a required subprocess environment so `[tool.coverage.run] patch = ["subprocess"]` does not instrument the child
`tests/cost_guards.py` enforces this, with two escape hatches that each need a `reason=` a reviewer can check:
- A test that launches a Python interpreter (`sys.executable`, `python*`, a console script, `uv run python`, `multiprocessing` with `spawn`/`forkserver`) fails unless it is marked `@pytest.mark.subprocess(reason='...')`. The marker is right only when the claim depends on the process boundary: what a fresh interpreter imports, behavior without an installed dependency, signals and process trees, a real stdio peer. Launching one while setting up a fixture shared beyond one test needs no marker, and other programs are not guarded. A marker on a module, class or function that no longer covers any launch fails the session with "remove the stale marker"; it is only judged on runs that select every test it covers, and `conditional=True` opts out a test that launches Python only on some platforms or installs
- When `PYTEST_TEST_BUDGET_SECONDS` is set (one CI leg sets it; it is off locally), a test whose setup and call together take longer fails unless it is marked `@pytest.mark.slow(reason='...')`. Setup of `module`, `class`, `package` and `session`-scoped fixtures is not counted, so sharing an expensive fixture is never penalized. Reach for the marker only when the duration is inherent to what the test proves; otherwise make the test cheaper with the patterns above
When `PYTEST_COLLECT_BUDGET_SECONDS` is set (on the same CI leg), a test module that takes longer to collect fails the session, with no escape hatch: check whether a heavy dependency is installed with `importlib.util.find_spec` for the skip condition, and import it in the tests that use it, after a `module`-scoped fixture has imported it once outside the per-test budget.
## Directory Structure
tests/benchmarks/README.md
+1 -1
View File
@@ -5,7 +5,7 @@ uv sync --frozen
uv run pytest tests/benchmarks --codspeed --codspeed-mode=walltime
```
Run these commands from the repository root. For a correctness check without timing, omit the CodSpeed flags.
Run these commands from the repository root. For a correctness check without timing, omit the CodSpeed flags. Without CodSpeed, each benchmark runs only its smallest workload: the larger workloads are marked `codspeed_only` and deselected by [`conftest.py`](conftest.py), because they reach the same code paths and would only slow down the regular test jobs, which collect this directory too.
The [CodSpeed workflow](../../.github/workflows/benchmark.yml) measures wall time on CodSpeed's dedicated ARM64 Graviton runners, with Python 3.14. It runs on pull requests and pushes to `main`. Fixed hardware avoids the CPU differences of GitHub-hosted runners, while isolation reduces timing noise.
tests/benchmarks/conftest.py
+27
View File
@@ -0,0 +1,27 @@
from __future__ import annotations
import os
import pytest
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line(
'markers', 'codspeed_only: a large workload that only runs when CodSpeed measures the benchmarks'
)
def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
"""Run large workloads only under CodSpeed.
Ordinary test runs still check every benchmark's correctness on its smallest workload, which reaches the
same code paths at a fraction of the cost.
"""
# The same check `pytest-codspeed` makes: `--codspeed` locally, `CODSPEED_ENV` under the CodSpeed action.
codspeed_enabled = config.getoption('--codspeed') or 'CODSPEED_ENV' in os.environ
selected: list[pytest.Item] = []
deselected: list[pytest.Item] = []
for item in items:
(selected if codspeed_enabled or not item.get_closest_marker('codspeed_only') else deselected).append(item)
config.hook.pytest_deselected(items=deselected)
items[:] = selected
tests/benchmarks/test_agent.py
+9 -2
View File
@@ -38,7 +38,12 @@ async def test_agent_run_without_capabilities(agent: Agent[None, str]) -> None:
assert result.usage.requests == 1
@pytest.fixture(params=[1000, 5000], ids=['1000-fragments', '5000-fragments'])
@pytest.fixture(
params=[
pytest.param(1000, id='1000-fragments'),
pytest.param(5000, id='5000-fragments', marks=pytest.mark.codspeed_only),
]
)
async def synthetic_history(agent: Agent[None, str], request: pytest.FixtureRequest) -> list[ModelMessage]:
history: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart('Earlier question')])]
history.extend(ModelResponse(parts=[TextPart('A response fragment')]) for _ in range(request.param))
@@ -52,7 +57,9 @@ async def test_agent_run_with_synthetic_history(agent: Agent[None, str], synthet
assert len(result.all_messages()[1].parts) == len(synthetic_history) - 1
@pytest.fixture(params=[1000, 5000], ids=['1000-chunks', '5000-chunks'])
@pytest.fixture(
params=[pytest.param(1000, id='1000-chunks'), pytest.param(5000, id='5000-chunks', marks=pytest.mark.codspeed_only)]
)
async def captured_text_stream(request: pytest.FixtureRequest) -> tuple[ModelResponse, list[ModelResponseStreamEvent]]:
async def stream_text(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str]:
for _ in range(request.param):
tests/benchmarks/test_eval_curves.py
+6 -1
View File
@@ -22,7 +22,12 @@ def blockbuster_enabled() -> bool:
return False
@pytest.fixture(params=[(256, 32), (4096, 2048)], ids=['256-cases-32-buckets', '4096-cases-2048-buckets'])
@pytest.fixture(
params=[
pytest.param((256, 32), id='256-cases-32-buckets'),
pytest.param((4096, 2048), id='4096-cases-2048-buckets', marks=pytest.mark.codspeed_only),
]
)
def curve_context(request: pytest.FixtureRequest) -> tuple[ReportEvaluatorContext[None, bool, None], int]:
n_cases, buckets = request.param
cases = [
tests/benchmarks/test_tool_search.py
+6 -2
View File
@@ -27,7 +27,9 @@ def anyio_backend() -> str:
return 'asyncio'
@pytest.fixture(params=[256, 8192], ids=['256-tools', '8192-tools'])
@pytest.fixture(
params=[pytest.param(256, id='256-tools'), pytest.param(8192, id='8192-tools', marks=pytest.mark.codspeed_only)]
)
async def searchable_agent(
request: pytest.FixtureRequest, searches: int, changing_corpus: bool
) -> tuple[Agent[int, str], Iterator[int]]:
@@ -103,7 +105,9 @@ async def test_keyword_tool_search(
]
@pytest.fixture(params=[256, 8192], ids=['256-tools', '8192-tools'])
@pytest.fixture(
params=[pytest.param(256, id='256-tools'), pytest.param(8192, id='8192-tools', marks=pytest.mark.codspeed_only)]
)
async def parallel_search_agent(request: pytest.FixtureRequest, searches: int) -> Agent[None, str]:
async def lookup() -> str:
return 'ok'
tests/cassette_utils.py
+22 -2
View File
@@ -30,6 +30,8 @@ if TYPE_CHECKING:
PrefixBlock = tuple[str, str]
TESTS_DIR = Path(__file__).parent
# Cache-write order of the request sections; a lower value is matched earlier in the provider's prompt
# cache, so when two requests diverge on different sections the earlier one is where the prefix breaks.
_CACHE_ORDER = {'tools': 0, 'system': 1, 'messages': 2}
@@ -255,9 +257,27 @@ def recorded_request_body(request: dict[str, Any]) -> Any:
return body
def load_cassette_yaml(text: str) -> Any:
"""Parse a cassette file's YAML with libyaml's loader where available, ~30x faster than PyYAML's pure-Python one."""
return yaml.load(text, Loader=SafeLoader)
def recorded_interactions(*markers: str) -> Iterator[tuple[str, dict[str, Any]]]:
"""Yield `('<path relative to tests/>#<index>', interaction)` for each interaction in `tests/`'s YAML cassettes.
Only cassettes whose text contains all `markers` are parsed, which skips most of them cheaply; callers still filter
the interactions themselves.
"""
for path in sorted(TESTS_DIR.rglob('*.yaml')):
text = path.read_text(encoding='utf-8')
if all(marker in text for marker in markers):
for index, interaction in enumerate(load_cassette_yaml(text)['interactions']):
yield f'{path.relative_to(TESTS_DIR)}#{index}', interaction
def _yaml_cassette_requests(cassette_path: Path) -> Iterator[tuple[Any, Any, Any]]:
"""Yield `(method, uri, body)` for each request in a cassette file on disk."""
cassette = yaml.load(cassette_path.read_text(encoding='utf-8'), Loader=SafeLoader)
cassette = load_cassette_yaml(cassette_path.read_text(encoding='utf-8'))
if not is_str_dict(cassette):
return
raw_interactions = cassette.get('interactions')
@@ -416,7 +436,7 @@ def _get_cassette_bodies_from_yaml(path: Path) -> list[str]:
Used as fallback when the cassette object is not available (e.g. CI playback).
"""
data: dict[str, Any] = yaml.safe_load(path.read_text(encoding='utf-8'))
data: dict[str, Any] = load_cassette_yaml(path.read_text(encoding='utf-8'))
bodies: list[str] = []
for interaction in data.get('interactions', []):
body = recorded_request_body(interaction.get('request', {}))
tests/clai2/cli_runner.py
+54
View File
@@ -0,0 +1,54 @@
"""Run the `clai2` command line in this process, with piped input and captured output."""
import io
import sys
from dataclasses import dataclass
from pathlib import Path
import pytest
from prompt_toolkit.application import create_app_session
from prompt_toolkit.input import create_pipe_input
from prompt_toolkit.output import DummyOutput
from pydantic_clai2.cli import _cli
@dataclass(frozen=True)
class CliResult:
returncode: int
stdout: str
stderr: str
@dataclass
class CliRunner:
"""`clai2 ARGS < input` as a subprocess would see it, without paying for a new interpreter.
Reserve real subprocesses for behavior that depends on the process boundary, such as startup imports.
"""
monkeypatch: pytest.MonkeyPatch
capsys: pytest.CaptureFixture[str]
def __call__(self, *args: str, cwd: Path | None = None, input: str = '/exit\n') -> CliResult:
if cwd is not None:
# Restored at teardown, including after `--worktree` changes into the checkout.
self.monkeypatch.chdir(cwd)
self.monkeypatch.setattr(sys, 'argv', ['clai2', *args])
# Piped, so `isatty()` is false as in a subprocess fed from a pipe.
self.monkeypatch.setattr(sys, 'stdin', io.StringIO(input))
self.capsys.readouterr()
returncode = 0
with create_pipe_input() as pipe, create_app_session(input=pipe, output=DummyOutput()):
pipe.send_text(input)
pipe.close()
try:
_cli.run()
except SystemExit as exit:
assert isinstance(exit.code, int)
returncode = exit.code
# Only when the CLI regresses: escaping, it would end the whole pytest session instead of failing this test.
except KeyboardInterrupt: # pragma: no cover
pytest.fail('KeyboardInterrupt escaped the CLI')
captured = self.capsys.readouterr()
return CliResult(returncode=returncode, stdout=captured.out, stderr=captured.err)
tests/clai2/conftest.py
+7
View File
@@ -10,6 +10,7 @@ from keyring.errors import PasswordDeleteError
from pydantic_ai import models
from pydantic_clai2.config import credential_store
from tests.clai2.cli_runner import CliRunner
@pytest.fixture
@@ -135,3 +136,9 @@ def vault(monkeypatch: pytest.MonkeyPatch) -> dict[tuple[str, str], str]:
monkeypatch.setattr(keyring, 'set_password', set_value)
monkeypatch.setattr(keyring, 'delete_password', delete)
return entries
@pytest.fixture
def run_cli(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> CliRunner:
"""Run `clai2` in this process; see `CliRunner`."""
return CliRunner(monkeypatch=monkeypatch, capsys=capsys)
tests/clai2/test_cli_edges.py
+54 -59
View File
@@ -1,31 +1,37 @@
"""Exercise the installed entry point in isolated subprocesses."""
"""Exercise the command line: one installed entry point smoke test, the rest in-process."""
import asyncio
import os
import sqlite3
import subprocess
import sys
from collections.abc import Coroutine
from contextlib import closing
from pathlib import Path
import pytest
import pydantic_clai2.ui.rendering.splash
from pydantic_clai2.__main__ import main
from pydantic_clai2.cli import _cli
from pydantic_clai2.config.settings_store import SettingsStore
from pydantic_clai2.ui.rendering.splash import Splash
from tests.clai2.cli_runner import CliRunner
@pytest.fixture(autouse=True)
def disable_subprocess_coverage(monkeypatch: pytest.MonkeyPatch) -> None:
@pytest.mark.subprocess(
reason='smoke-tests the installed `python -m pydantic_clai2` entry point; the other CLI tests run in-process'
)
def test_installed_entry_point_starts_and_exits(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""`python -m pydantic_clai2` without a model: the one launch here that crosses the process boundary."""
# Subprocess coverage adds import overhead to the CLI startup hang guard.
for name in tuple(os.environ):
if name.startswith('COVERAGE_'):
monkeypatch.delenv(name)
@pytest.mark.parametrize('args', [[], ['--model', 'test', '--request-limit', '12'], ['--request-limit', '0']])
def test_cli_startup(tmp_path: Path, args: list[str]) -> None:
env = dict(os.environ, CLAI_NO_SPLASH='1')
env.pop('CLAI_MODEL', None)
result = subprocess.run(
[sys.executable, '-m', 'pydantic_clai2', '--database', str(tmp_path / 'config.db'), *args],
[sys.executable, '-m', 'pydantic_clai2', '--database', str(tmp_path / 'config.db')],
input='/exit\n',
text=True,
capture_output=True,
@@ -33,11 +39,21 @@ def test_cli_startup(tmp_path: Path, args: list[str]) -> None:
timeout=15,
check=False,
)
assert result.returncode == (2 if args == ['--request-limit', '0'] else 0), result.stderr
assert result.returncode == 0, result.stderr
@pytest.mark.parametrize(
('args', 'returncode'), [([], 0), (['--model', 'test', '--request-limit', '12'], 0), (['--request-limit', '0'], 2)]
)
def test_cli_startup(tmp_path: Path, run_cli: CliRunner, args: list[str], returncode: int) -> None:
result = run_cli('--database', str(tmp_path / 'config.db'), *args, cwd=tmp_path)
assert result.returncode == returncode, result.stderr
@pytest.mark.parametrize('args', [[], ['config', 'show']])
def test_cli_recovers_unsupported_theme_without_losing_unknown_settings(tmp_path: Path, args: list[str]) -> None:
def test_cli_recovers_unsupported_theme_without_losing_unknown_settings(
tmp_path: Path, run_cli: CliRunner, args: list[str]
) -> None:
path = tmp_path / 'config.db'
saved = {
'display.theme': '"light"',
@@ -49,29 +65,13 @@ def test_cli_recovers_unsupported_theme_without_losing_unknown_settings(tmp_path
connection.execute('PRAGMA user_version = 1')
connection.execute('CREATE TABLE settings (key TEXT PRIMARY KEY, value_json TEXT NOT NULL)')
connection.executemany('INSERT INTO settings VALUES (?, ?)', saved.items())
result = subprocess.run(
[sys.executable, '-m', 'pydantic_clai2', '--database', str(path), *args],
input='/exit\n',
text=True,
capture_output=True,
env=dict(os.environ, CLAI_NO_SPLASH='1'),
timeout=15,
check=False,
)
result = run_cli('--database', str(path), *args, cwd=tmp_path)
assert result.returncode == 2, result.stderr
assert 'Unknown theme: light' in result.stderr
with closing(sqlite3.connect(path)) as connection:
assert dict(connection.execute('SELECT key, value_json FROM settings')) == saved
for command in (['config', 'reset', 'display.theme'], args):
result = subprocess.run(
[sys.executable, '-m', 'pydantic_clai2', '--database', str(path), *command],
input='/exit\n',
text=True,
capture_output=True,
env=dict(os.environ, CLAI_NO_SPLASH='1'),
timeout=15,
check=False,
)
result = run_cli('--database', str(path), *command)
assert result.returncode == 0, result.stderr
del saved['display.theme']
store = SettingsStore(path)
@@ -81,45 +81,40 @@ def test_cli_recovers_unsupported_theme_without_losing_unknown_settings(tmp_path
assert dict(connection.execute('SELECT key, value_json FROM settings')) == saved
def test_cli_startup_interrupt(tmp_path: Path) -> None:
script = """
import asyncio
import runpy
def test_cli_startup_interrupt(tmp_path: Path, run_cli: CliRunner, monkeypatch: pytest.MonkeyPatch) -> None:
def interrupted(coroutine: Coroutine[object, object, object]) -> None:
coroutine.close()
raise KeyboardInterrupt
def interrupted(coroutine):
coroutine.close()
raise KeyboardInterrupt
asyncio.run = interrupted
runpy.run_module('pydantic_clai2', run_name='__main__')
"""
result = subprocess.run(
[sys.executable, '-c', script, '--database', str(tmp_path / 'config.db')],
capture_output=True,
text=True,
check=False,
timeout=15,
)
assert result.returncode == 0, result.stderr
monkeypatch.setattr(asyncio, 'run', interrupted)
assert run_cli('--database', str(tmp_path / 'config.db'), cwd=tmp_path).returncode == 0
@pytest.mark.parametrize('state', ['enabled', 'disabled', 'corrupt'])
def test_startup_saved_splash(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, state: str) -> None:
@pytest.mark.parametrize(('state', 'enabled'), [('enabled', True), ('disabled', False), ('corrupt', False)])
def test_startup_saved_splash(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, state: str, enabled: bool) -> None:
"""A bare launch reads `display.splash` before the heavy imports, and a corrupt database turns it off."""
monkeypatch.setenv('XDG_CONFIG_HOME', str(tmp_path))
monkeypatch.delenv('CLAI_NO_SPLASH', raising=False)
monkeypatch.setenv('CLAI_MODEL', 'test')
monkeypatch.setenv('PYDANTIC_AI_NO_BANNER', '1') # `main` sets it; this restores it after the test.
monkeypatch.chdir(tmp_path)
monkeypatch.setattr(sys, 'argv', ['clai2'])
path = tmp_path / 'pydantic-clai2' / 'config.db'
store = SettingsStore(path)
if state == 'corrupt':
path.write_text('not sqlite')
else:
store.set('display.splash', state == 'enabled')
result = subprocess.run(
[sys.executable, '-m', 'pydantic_clai2'],
input='/exit\n',
text=True,
capture_output=True,
timeout=15,
check=False,
)
assert result.returncode == (1 if state == 'corrupt' else 0)
requested: list[bool] = []
class RecordingSplash(Splash):
def __init__(self, *, enabled: bool = True) -> None:
requested.append(enabled)
super().__init__(enabled=enabled)
def run(*, splash: Splash | None = None) -> None:
assert isinstance(splash, RecordingSplash)
monkeypatch.setattr(pydantic_clai2.ui.rendering.splash, 'Splash', RecordingSplash)
monkeypatch.setattr(_cli, 'run', run)
main()
assert requested == [enabled]
tests/clai2/test_dotenv.py
+1
View File
@@ -107,6 +107,7 @@ def test_startup_without_dotenv_values(
assert stderr.startswith(f'Ignoring `.env` at {str(tmp_path / ".env")!r}: {error}')
@pytest.mark.subprocess(reason='asserts the order of imports in a fresh `pydantic_clai2` process')
def test_dotenv_loaded_before_startup_imports(tmp_path: Path) -> None:
(tmp_path / '.env').write_text('CLAI_DOTENV_TEST=loaded before import\n')
(tmp_path / 'dotenv_agent.py').write_text(
tests/clai2/test_github_gh_login.py
+4
View File
@@ -24,6 +24,10 @@ from pydantic_clai2.plugins.loader import PluginLoader
from tests.clai2.conftest import FakeGh
from tests.clai2.menu_script import Script, pick
pytestmark = pytest.mark.subprocess(
reason='drives the fake `gh` CLI (`tests/clai2/fake_gh_cli.py`) as a real child process'
)
BUILTIN = next(plugin for plugin in DEFAULT_PLUGINS if plugin.id == 'github')
CLOSE = MenuResult(cancelled=True)
DEVICE = 'https://github.com/login/device'
tests/clai2/test_github_plugin.py
+4
View File
@@ -27,6 +27,10 @@ from pydantic_clai2.ui.menus.field_menu import CUSTOM, FieldRow
from pydantic_clai2.ui.menus.plugin_menu import PluginMenu, open_plugins_menu
from tests.clai2.menu_script import Script, pick, typed
pytestmark = pytest.mark.subprocess(
reason='drives the fake `gh` CLI (`tests/clai2/fake_gh_cli.py`) as a real child process'
)
BUILTIN = next(plugin for plugin in DEFAULT_PLUGINS if plugin.id == 'github')
CLOSE = MenuResult(cancelled=True)
tests/clai2/test_herdr.py
+1
View File
@@ -263,6 +263,7 @@ async def test_background_title_watcher(
monkeypatch.setattr(store, 'get', no_transcript)
monkeypatch.setattr(store.media, 'get', no_transcript)
monkeypatch.setattr(herdr, 'TITLE_POLL_INTERVAL', 0.01)
await loaded.dispatch(SessionStart(agent=agent, settings=Settings()))
try:
assert await recorded.titles.get() == session.summary.title
tests/clai2/test_mcp.py
+6
View File
@@ -220,6 +220,7 @@ async def test_loader_persistence_disable_and_project_plugin_trust(tmp_path: Pat
assert not list(commands)
@pytest.mark.subprocess(reason='starts real stdio MCP servers written in Python')
async def test_start_stop_restart_logs_and_agent_use(tmp_path: Path) -> None:
script, pid_file = write_server(tmp_path)
store = MCPStore(tmp_path / 'config', workspace=tmp_path)
@@ -267,6 +268,7 @@ async def test_start_stop_restart_logs_and_agent_use(tmp_path: Path) -> None:
assert 'Stopped local' in await run('/mcp stop-all')
@pytest.mark.subprocess(reason='starts a real stdio MCP server written in Python')
async def test_tools_without_start_and_empty_server(tmp_path: Path) -> None:
script, pid_file = write_server(tmp_path, with_tool=False)
store = MCPStore(tmp_path / 'config', workspace=tmp_path)
@@ -276,6 +278,7 @@ async def test_tools_without_start_and_empty_server(tmp_path: Path) -> None:
assert_exited(pid_file)
@pytest.mark.subprocess(reason='starts a stdio MCP server command that exits immediately')
async def test_failed_start_is_reported_and_logged(tmp_path: Path) -> None:
store = MCPStore(tmp_path / 'config', workspace=tmp_path)
plugin = make_plugin({}, store)
@@ -332,6 +335,7 @@ def test_env_reference_order_and_template_escapes(
assert servers.state(entry) == 'ready'
@pytest.mark.subprocess(reason='starts real stdio MCP servers written in Python')
async def test_reconfiguring_or_removing_releases_the_connection(tmp_path: Path) -> None:
script, pid_file = write_server(tmp_path)
store = MCPStore(tmp_path / 'config', workspace=tmp_path)
@@ -353,6 +357,7 @@ async def test_reconfiguring_or_removing_releases_the_connection(tmp_path: Path)
servers.get('local')
@pytest.mark.subprocess(reason='serves a real HTTP/SSE MCP server from a separate process')
@pytest.mark.parametrize(('kind', 'path'), [('http', 'mcp'), ('sse', 'sse')])
async def test_real_remote_server(tmp_path: Path, kind: str, path: str) -> None:
with socket.socket() as probe:
@@ -420,6 +425,7 @@ def test_fastmcp_info_logging_stays_off_the_prompt() -> None:
assert logging.getLogger('fastmcp.client.auth.oauth').getEffectiveLevel() >= logging.WARNING
@pytest.mark.subprocess(reason='starts a stdio MCP server command that exits immediately')
async def test_a_server_that_cannot_connect_does_not_fail_the_prompt(tmp_path: Path) -> None:
store = MCPStore(tmp_path / 'config', workspace=tmp_path)
plugin = make_plugin({}, store)
tests/clai2/test_mcp_command.py
+3
View File
@@ -222,6 +222,7 @@ async def test_edit_prefills_renames_and_cancels(tmp_path: Path) -> None:
assert await command(['edit', 'renamed']) == 'No changes.'
@pytest.mark.subprocess(reason='starts a real stdio MCP server written in Python')
async def test_edit_restarts_a_running_server(tmp_path: Path) -> None:
server = tmp_path / 'server.py'
server.write_text(
@@ -245,6 +246,7 @@ async def test_edit_refuses_servers_mcp_does_not_own(tmp_path: Path) -> None:
await command(['edit', 'legacy'])
@pytest.mark.subprocess(reason='runs Python scripts as the `$VISUAL`/`$EDITOR` commands')
def test_editor_round_trip(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
edit = tmp_path / 'edit.py'
edit.write_text('import sys, pathlib\npathlib.Path(sys.argv[1]).write_text("{\\"command\\": \\"edited\\"}")\n')
@@ -511,6 +513,7 @@ async def test_symlinked_project_file_is_never_trusted(tmp_path: Path) -> None:
assert store.trust_state(project) == 'untrusted' and store.project_servers() == {}
@pytest.mark.subprocess(reason='starts a real stdio MCP server written in Python')
async def test_revoke_disconnects_project_servers(tmp_path: Path) -> None:
server = tmp_path / 'server.py'
server.write_text('from mcp.server.fastmcp import FastMCP\nFastMCP("t").run()\n')
tests/clai2/test_plugin_loader.py
+1
View File
@@ -808,6 +808,7 @@ async def test_failed_load_cleanup_is_bounded_and_continues_after_timeout(
StalledCleanup.__module__ = module.__name__
module.__dict__['StalledCleanup'] = StalledCleanup
monkeypatch.setitem(sys.modules, module.__name__, module)
monkeypatch.setattr('pydantic_clai2.plugins.loader.FAILED_SESSION_END_TIMEOUT', 0.05)
harness = Harness(tmp_path, builtin=(PluginSettings(id='stalled', factory=module.__name__),))
async def load() -> None:
tests/clai2/test_reload.py
+5
View File
@@ -12,6 +12,10 @@ import pytest
'mode',
['unchanged', 'success', 'custom', 'new_imports', 'stock', 'syntax', 'import', 'build', 'harness', 'transcript'],
)
@pytest.mark.subprocess(
reason='reloads the shell in a fresh interpreter so the test process keeps its own module identities'
)
@pytest.mark.slow(reason='each scenario imports and reloads the whole shell in a fresh interpreter')
def test_reload_running_shell(tmp_path: Path, mode: str) -> None:
run_script(tmp_path, 'reload_script.py', mode)
@@ -75,5 +79,6 @@ def run_script(tmp_path: Path, script: str, mode: str) -> None:
'syntax',
],
)
@pytest.mark.subprocess(reason='checks how a fresh interpreter re-imports a changed module graph')
def test_reload_changed_import_graph(tmp_path: Path, mode: str) -> None:
run_script(tmp_path, 'reload_import_script.py', mode)
tests/clai2/test_resume_integration.py
+3 -13
View File
@@ -1,8 +1,5 @@
"""The command, startup flag, naming resolver and terminal worker share one service."""
import os
import subprocess
import sys
from builtins import BaseExceptionGroup
from collections.abc import AsyncIterator, Callable
from dataclasses import replace
@@ -31,6 +28,7 @@ from pydantic_clai2.runtime._session import Session
from pydantic_clai2.runtime.session_naming import SessionNamer
from pydantic_clai2.runtime.sessions import Sessions
from pydantic_clai2.ui.menus.session_browser import SessionBrowser
from tests.clai2.cli_runner import CliRunner
async def test_reload_keeps_saved_conversation(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -225,16 +223,8 @@ async def test_empty_startup_browser_and_invalid_restore(tmp_path: Path, monkeyp
@pytest.mark.parametrize('args', [['--resume', 'missing'], ['--resume=missing', 'config']])
def test_resume_cli_errors_are_normal_parser_errors(tmp_path: Path, args: list[str]) -> None:
result = subprocess.run(
[sys.executable, '-m', 'pydantic_clai2', '--database', str(tmp_path / 'settings.db'), *args],
input='/exit\n',
capture_output=True,
text=True,
timeout=20,
check=False,
env=dict(os.environ, CLAI_NO_SPLASH='1'),
)
def test_resume_cli_errors_are_normal_parser_errors(tmp_path: Path, run_cli: CliRunner, args: list[str]) -> None:
result = run_cli('--database', str(tmp_path / 'settings.db'), *args, cwd=tmp_path)
assert result.returncode == 2
assert 'error:' in result.stderr
assert 'Traceback' not in result.stderr
tests/clai2/test_self_update.py
+1
View File
@@ -344,6 +344,7 @@ async def test_command_reports_current_failure_and_manual_installs() -> None:
await current.command(['now'])
@pytest.mark.subprocess(reason='runs a Python one-liner in place of `uv` to check the exit status is returned')
async def test_run_uv_returns_the_exit_status() -> None:
assert await _run_uv([sys.executable, '-c', 'raise SystemExit(3)']) == 3
tests/clai2/test_shell_passthrough.py
+1
View File
@@ -198,6 +198,7 @@ class TestShellPassthrough:
assert 'Interrupted (' in text
assert (tmp_path / 'cleanup.txt').read_text() == 'cleaned'
@pytest.mark.subprocess(reason="needs a real process tree to check Ctrl-C reaches the shell's descendants")
@pytest.mark.skipif(sys.platform == 'win32', reason='uses POSIX shell process-group signalling')
async def test_ctrl_c_kills_shell_descendants(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Cancelling a shell command must also terminate a background child."""
tests/clai2/test_startup_imports.py
+2
View File
@@ -10,6 +10,8 @@ import pytest
from pydantic_clai2 import warm_imports
pytestmark = pytest.mark.subprocess(reason='asserts what a fresh `pydantic_clai2` process imports at startup')
@pytest.mark.parametrize(
('args', 'exit_code'),
tests/clai2/test_system_prompt_menu.py
+6 -2
View File
@@ -275,6 +275,7 @@ def _editor(tmp_path: Path, name: str) -> str:
return shlex.join([sys.executable, str(script)])
@pytest.mark.subprocess(reason='runs Python scripts as the `$VISUAL`/`$EDITOR` commands')
def test_edit_text_opens_visual_then_editor(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(sys, 'platform', 'linux')
built_in: list[EditorApp] = []
@@ -303,7 +304,6 @@ def test_edit_text_opens_visual_then_editor(tmp_path: Path, monkeypatch: pytest.
def test_edit_text_falls_back_to_the_built_in_editor(
monkeypatch: pytest.MonkeyPatch, platform: str, editor: str | None
) -> None:
monkeypatch.setattr(sys, 'platform', platform)
monkeypatch.delenv('VISUAL', raising=False)
if editor is None:
monkeypatch.delenv('EDITOR', raising=False)
@@ -315,7 +315,11 @@ def test_edit_text_falls_back_to_the_built_in_editor(
built_in.append(app)
return 'built-in'
assert edit_text('text', title='Edit', run_area=area) == 'built-in'
# Its own session, made before patching `sys.platform`: otherwise the editor's input and output are created for
# the patched platform, unless an earlier test on this worker already created them.
with create_pipe_input() as pipe, create_app_session(input=pipe, output=DummyOutput()):
monkeypatch.setattr(sys, 'platform', platform)
assert edit_text('text', title='Edit', run_area=area) == 'built-in'
assert len(built_in) == 1
tests/clai2/test_terminal.py
+7 -9
View File
@@ -29,6 +29,7 @@ from pydantic_clai2.config.settings_store import SettingsStore
from pydantic_clai2.plugins import Plugin, TurnEnd, TurnStart
from pydantic_clai2.ui.prompt.prompt_surface import LEAVE, PromptSurface
from pydantic_clai2.ui.rendering.splash import Splash
from tests.clai2.cli_runner import CliRunner
async def test_existing_handler_and_structured_output() -> None:
@@ -346,20 +347,17 @@ async def test_prompt_loop_commands(tmp_path: Path) -> None:
assert output.getvalue().count('/new starts a session') == 2
def test_cli_settings(tmp_path: Path) -> None:
base = [sys.executable, '-m', 'pydantic_clai2', '--database', str(tmp_path / 'config.db')]
result = subprocess.run(
[*base, 'config', 'set', 'display.thinking', 'false'], check=False, capture_output=True, text=True
)
def test_cli_settings(tmp_path: Path, run_cli: CliRunner) -> None:
database = str(tmp_path / 'config.db')
result = run_cli('--database', database, 'config', 'set', 'display.thinking', 'false')
assert result.returncode == 0, result.stderr
result = subprocess.run([*base, 'config', 'get', 'display.thinking'], check=False, capture_output=True, text=True)
result = run_cli('--database', database, 'config', 'get', 'display.thinking')
assert result.stdout.strip() == 'false'
result = subprocess.run(
[*base, 'config', 'set', 'run.request_limit', '-1'], check=False, capture_output=True, text=True
)
result = run_cli('--database', database, 'config', 'set', 'run.request_limit', '-1')
assert result.returncode != 0
@pytest.mark.subprocess(reason='asserts what a fresh interpreter imports')
def test_import_is_light() -> None:
result = subprocess.run(
[sys.executable, '-c', 'import pydantic_clai2, sys; assert "pydantic_ai" not in sys.modules'],
tests/clai2/test_text_clipboard.py
+2
View File
@@ -63,6 +63,7 @@ def test_linux_prefers_wayland_then_x11_tools(installed: set[str], monkeypatch:
assert copy_command(platform='linux') == ClipboardCommand(argv=('/bin/wl-copy',))
@pytest.mark.subprocess(reason='runs a Python script as the clipboard command')
def test_run_copy_feeds_utf8_text_and_the_command_environment(tmp_path: Path) -> None:
path = tmp_path / 'clipboard'
script = f'import os, sys; open({str(path)!r}, "wb").write(os.environ["MARK"].encode() + sys.stdin.buffer.read())'
@@ -72,6 +73,7 @@ def test_run_copy_feeds_utf8_text_and_the_command_environment(tmp_path: Path) ->
assert path.read_bytes()[:1] == b'>', 'the command gets its own variables'
@pytest.mark.subprocess(reason='runs a Python script as the clipboard command')
def test_a_failing_clipboard_command_loses_only_that_copy(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
run_copy(command=ClipboardCommand(argv=(str(tmp_path / 'missing'),)), text='lost')
monkeypatch.setattr(text_clipboard, 'COPY_TIMEOUT', 0.01)
tests/clai2/test_worktrees.py
+31 -28
View File
@@ -1,16 +1,17 @@
"""Launch the installed CLI against real repositories without provider requests."""
import asyncio
import os
import shutil
import subprocess
import sys
from dataclasses import replace
from dataclasses import dataclass, replace
from pathlib import Path
import pytest
from pydantic_ai_harness.step_persistence.conversations import SqliteConversationStore
from pydantic_clai2.runtime.worktrees import Worktree, open_worktree
from tests.clai2.cli_runner import CliResult, CliRunner
def git(directory: Path, *args: str) -> str:
@@ -19,17 +20,17 @@ def git(directory: Path, *args: str) -> str:
).stdout.strip()
def launch(directory: Path, *args: str, prompt: str = '/exit\n') -> subprocess.CompletedProcess[str]:
return subprocess.run(
[sys.executable, '-m', 'pydantic_clai2', '--database', 'config.db', '--model', 'test', *args],
cwd=directory,
input=prompt,
text=True,
capture_output=True,
env=dict(os.environ, CLAI_NO_SPLASH='1'),
timeout=30,
check=False,
)
@dataclass
class Launcher:
run_cli: CliRunner
def __call__(self, directory: Path, *args: str, prompt: str = '/exit\n') -> CliResult:
return self.run_cli('--database', 'config.db', '--model', 'test', *args, cwd=directory, input=prompt)
@pytest.fixture
def launch(run_cli: CliRunner) -> Launcher:
return Launcher(run_cli)
@pytest.fixture
@@ -53,7 +54,7 @@ def repository(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
'args',
[['--worktree', 'feature'], ['-w', 'feature'], ['--worktree=feature'], ['-wfeature'], ['--worktree'], ['-w']],
)
def test_launches_in_new_worktree(repository: Path, args: list[str]) -> None:
def test_launches_in_new_worktree(repository: Path, launch: Launcher, args: list[str]) -> None:
original_head = git(repository, 'rev-parse', 'HEAD')
original_branch = git(repository, 'symbolic-ref', 'HEAD')
(repository / 'tracked.txt').write_text('uncommitted')
@@ -104,7 +105,7 @@ def test_launches_in_new_worktree(repository: Path, args: list[str]) -> None:
assert not (repository / '.gitignore').exists()
def test_launch_from_linked_worktree(repository: Path) -> None:
def test_launch_from_linked_worktree(repository: Path, launch: Launcher) -> None:
linked = repository.parent / 'linked'
git(repository, 'worktree', 'add', '-b', 'linked', str(linked))
(linked / 'linked.txt').write_text('linked commit')
@@ -119,7 +120,9 @@ def test_launch_from_linked_worktree(repository: Path) -> None:
@pytest.mark.parametrize('existing', [None, b'keep-me', b'keep-me\n/.worktrees/\n'])
def test_local_exclude_preserves_rules_without_duplicates(repository: Path, existing: bytes | None) -> None:
def test_local_exclude_preserves_rules_without_duplicates(
repository: Path, launch: Launcher, existing: bytes | None
) -> None:
exclude = repository / '.git/info/exclude'
if existing is None:
exclude.unlink()
@@ -134,7 +137,7 @@ def test_local_exclude_preserves_rules_without_duplicates(repository: Path, exis
assert git(repository, 'check-ignore', '.worktrees/ignored') == '.worktrees/ignored'
def test_unwritable_exclude_is_a_parser_error(repository: Path) -> None:
def test_unwritable_exclude_is_a_parser_error(repository: Path, launch: Launcher) -> None:
exclude = repository / '.git/info/exclude'
exclude.unlink()
exclude.mkdir()
@@ -145,7 +148,7 @@ def test_unwritable_exclude_is_a_parser_error(repository: Path) -> None:
assert not (repository / '.worktrees').exists()
async def test_session_belongs_to_worktree_and_can_be_resumed(repository: Path) -> None:
def test_session_belongs_to_worktree_and_can_be_resumed(repository: Path, launch: Launcher) -> None:
result = launch(
repository,
'-w',
@@ -155,7 +158,7 @@ async def test_session_belongs_to_worktree_and_can_be_resumed(repository: Path)
assert result.returncode == 0, result.stderr
workspace = repository / '.worktrees/saved'
store = SqliteConversationStore(database=repository / 'sessions.db')
summaries = await store.listing()
summaries = asyncio.run(store.listing())
assert len(summaries) == 1
assert summaries[0].workspace == str(workspace)
resumed = launch(workspace, '--database', str(repository / 'config.db'), '--resume', summaries[0].id)
@@ -163,7 +166,7 @@ async def test_session_belongs_to_worktree_and_can_be_resumed(repository: Path)
assert 'Resumed' in resumed.stdout
def test_startup_error_keeps_created_worktree(repository: Path) -> None:
def test_startup_error_keeps_created_worktree(repository: Path, launch: Launcher) -> None:
result = launch(repository, '-w', 'retained', '--request-limit', '0')
assert result.returncode == 2
workspace = repository / '.worktrees/retained'
@@ -172,7 +175,7 @@ def test_startup_error_keeps_created_worktree(repository: Path) -> None:
assert git(workspace, 'branch', '--show-current') == 'clai-retained'
def test_headless_worktree_notice_stays_off_stdout(repository: Path) -> None:
def test_headless_worktree_notice_stays_off_stdout(repository: Path, launch: Launcher) -> None:
# The exit code is not asserted: `TestModel` calls every default tool, which the notice does not depend on.
result = launch(repository, '-w', 'headless', '-p', 'hello', prompt='')
workspace = repository / '.worktrees/headless'
@@ -181,7 +184,7 @@ def test_headless_worktree_notice_stays_off_stdout(repository: Path) -> None:
assert 'Worktree' not in result.stdout
def test_named_worktree_is_reopened(repository: Path) -> None:
def test_named_worktree_is_reopened(repository: Path, launch: Launcher) -> None:
assert launch(repository, '-w', 'feature').returncode == 0
workspace = repository / '.worktrees/feature'
(workspace / 'draft.txt').write_text('in progress')
@@ -193,7 +196,7 @@ def test_named_worktree_is_reopened(repository: Path) -> None:
@pytest.mark.parametrize('state', ['outside', 'unborn', 'path', 'missing-git'])
def test_git_errors_leave_existing_work_untouched(
repository: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, state: str
repository: Path, launch: Launcher, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, state: str
) -> None:
directory = repository
workspace = repository / '.worktrees/feature'
@@ -226,7 +229,7 @@ def test_git_errors_leave_existing_work_untouched(
assert not workspace.exists()
def test_failed_checkout_removes_new_branch_and_allows_retry(repository: Path) -> None:
def test_failed_checkout_removes_new_branch_and_allows_retry(repository: Path, launch: Launcher) -> None:
(repository / '.gitattributes').write_text('tracked.txt filter=fail\n')
git(repository, 'add', '.gitattributes')
git(repository, 'commit', '-m', 'Configure checkout filter')
@@ -241,7 +244,7 @@ def test_failed_checkout_removes_new_branch_and_allows_retry(repository: Path) -
assert launch(repository, '-w', 'retry').returncode == 0
def test_checkout_hook_failure_preserves_work_and_reports_cleanup_failure(repository: Path) -> None:
def test_checkout_hook_failure_preserves_work_and_reports_cleanup_failure(repository: Path, launch: Launcher) -> None:
hook = repository / '.git/hooks/post-checkout'
hook.write_text('#!/bin/sh\nexit 1\n')
hook.chmod(0o755)
@@ -254,7 +257,7 @@ def test_checkout_hook_failure_preserves_work_and_reports_cleanup_failure(reposi
assert git(workspace, 'branch', '--show-current') == 'clai-retained'
def test_exclude_write_failure_reports_retained_worktree(repository: Path) -> None:
def test_exclude_write_failure_reports_retained_worktree(repository: Path, launch: Launcher) -> None:
hook = repository / '.git/hooks/post-checkout'
hook.write_text('#!/bin/sh\nexclude=$(git rev-parse --git-path info/exclude)\nrm "$exclude"\nmkdir "$exclude"\n')
hook.chmod(0o755)
@@ -269,7 +272,7 @@ def test_exclude_write_failure_reports_retained_worktree(repository: Path) -> No
@pytest.mark.parametrize(
'name', ['../escape', '/absolute', 'nested/name', 'back\\slash', '..', '-b', 'bad name', 'bad.name']
)
def test_invalid_names_are_parser_errors(repository: Path, name: str) -> None:
def test_invalid_names_are_parser_errors(repository: Path, launch: Launcher, name: str) -> None:
result = launch(repository, f'--worktree={name}')
assert result.returncode == 2
assert 'error: Worktree names' in result.stderr
@@ -278,7 +281,7 @@ def test_invalid_names_are_parser_errors(repository: Path, name: str) -> None:
@pytest.mark.parametrize('args', [['--resume'], ['--resume=session'], ['config', 'show'], ['plugins', 'list']])
def test_worktree_rejects_incompatible_commands(repository: Path, args: list[str]) -> None:
def test_worktree_rejects_incompatible_commands(repository: Path, launch: Launcher, args: list[str]) -> None:
result = launch(repository, '--worktree=feature', *args)
assert result.returncode == 2
assert 'cannot be combined' in result.stderr
tests/conftest.py
+49 -1
View File
@@ -2,12 +2,14 @@ from __future__ import annotations as _annotations
import asyncio
import dataclasses
import gc
import importlib.util
import logging
import os
import re
import secrets
import sys
import warnings
from collections.abc import AsyncIterator, Callable, Generator, Iterator, Sequence
from contextlib import contextmanager, suppress
from dataclasses import dataclass, field
@@ -52,7 +54,7 @@ from pydantic_ai.messages import (
from pydantic_ai.models import Model
from pydantic_ai.usage import RequestUsage, RunUsage
from . import cassette_hooks
from . import cassette_hooks, cost_guards
from ._inline_snapshot import Builder, Custom, customize
from .cassette_utils import check_cache_prefix_stability
@@ -101,6 +103,7 @@ os.environ.setdefault('HF_HUB_DISABLE_PROGRESS_BARS', '1')
def pytest_configure(config: pytest.Config) -> None:
config.pluginmanager.register(cost_guards, cost_guards.PLUGIN_NAME)
config.addinivalue_line(
'markers',
'moves_cache_prefix(reason): recorded conversation deliberately moves the cache prefix; reason required',
@@ -653,6 +656,22 @@ def missing_event_loop() -> Iterator[asyncio.AbstractEventLoop]:
asyncio.set_event_loop(original_loop)
@pytest.fixture
def young_gc() -> Iterator[None]:
"""Limit the test's `gc.collect()` calls to the objects it creates.
A full collection walks every object the worker has accumulated, which takes seconds per call
late in a CI run. Freezing everything that exists before the test keeps it out of the collector's
reach until teardown, so collections stay cheap, and garbage leaked by an earlier test can't be
finalized (and raise an unraisable exception or `ResourceWarning`) in the middle of this one.
"""
gc.freeze()
try:
yield
finally:
gc.unfreeze()
@pytest.fixture(autouse=True)
def no_instrumentation_by_default():
Agent.instrument_all(False)
@@ -702,6 +721,35 @@ except ImportError:
pass
@pytest.fixture(scope='session')
def prefect_test_server() -> Iterator[None]:
"""A Prefect test server with an isolated database, shared by every Prefect test in the worker.
Starting one takes 15-25s on a CI runner, so tests that run Prefect flows request this fixture
and share the `prefect` xdist group, starting a single server per job instead of one per module
or test. The implicit ephemeral server would instead use the shared default `PREFECT_HOME` and a
short connect timeout that flakes on slow runners.
"""
pytest.importorskip('prefect')
from prefect.settings import PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED, temporary_settings
from prefect.testing.utilities import prefect_test_harness
# The task-run recorder is a background writer against the same sqlite file the flows write to.
# Prefect PRAGMAs a 60s `busy_timeout` onto every connection, and under CI contention the
# recorder's bulk inserts exhaust it, failing the flow whose state it was recording. Nothing
# here reads what it records: task run states reach the API through the task engine.
with (
temporary_settings({PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED: False}),
prefect_test_harness(server_startup_timeout=120),
):
yield
# Prefect's test server leaves client sockets for GC on Python 3.14; collect them here, where the
# warning is expected, rather than in whichever test the collector happens to run in.
with warnings.catch_warnings():
warnings.filterwarnings('ignore', message='unclosed.*socket', category=ResourceWarning)
gc.collect()
def raise_if_exception(e: Any) -> None:
if isinstance(e, Exception):
raise e
tests/cost_guards.py
+592
View File
@@ -0,0 +1,592 @@
"""Guards that keep the cost of individual tests visible and bounded.
Registered by `tests/conftest.py`; see "Test cost" in `tests/AGENTS.md` for the patterns these guards push
towards. The plugin:
- fails a test that launches a Python interpreter unless it is marked `@pytest.mark.subprocess(reason=...)`,
and fails the session when a `subprocess` marker no longer covers any test that launches one;
- when `PYTEST_TEST_BUDGET_SECONDS` is set, fails a test whose setup and call together take longer than that
unless it is marked `@pytest.mark.slow(reason=...)`;
- reports how long collection took, in the terminal summary and in `$GITHUB_STEP_SUMMARY` when that is set;
- when `PYTEST_COLLECT_BUDGET_SECONDS` is set, fails the session when a test module takes longer than that to collect.
"""
from __future__ import annotations as _annotations
import multiprocessing.util
import os
import re
import shlex
import subprocess
import sys
import sysconfig
import time
import traceback
from collections.abc import Callable, Generator, Sequence
from contextlib import suppress
from dataclasses import dataclass, field
from typing import Any
import pytest
from typing_extensions import TypedDict
PLUGIN_NAME = 'pydantic_ai_cost_guards'
BUDGET_ENV_VAR = 'PYTEST_TEST_BUDGET_SECONDS'
COLLECT_BUDGET_ENV_VAR = 'PYTEST_COLLECT_BUDGET_SECONDS'
_THIS_FILE = __file__
_TESTS_DIR = os.path.dirname(__file__) + os.sep
_PYTHON_PROGRAM = re.compile(r'(python|pypy)(\d+(\.\d+)*)?t?(\.exe)?', re.IGNORECASE)
_SHELLS = frozenset({'sh', 'bash', 'dash', 'zsh'})
_GUIDANCE_DOC = 'the "Test cost" section of tests/AGENTS.md'
# --- Detecting Python interpreter launches ----------------------------------------------------------------------------
def _python_entry_points() -> frozenset[str]:
"""`sys.executable` and this environment's console scripts (`clai`, `pytest`, ...), by name and by path.
Computed once, before any test runs: the check itself runs inside `Popen`, often on an event loop where the
blocking-call detector would flag any filesystem access.
"""
entry_points = {sys.executable, os.path.realpath(sys.executable)}
scripts_dir = sysconfig.get_path('scripts')
with suppress(OSError):
for entry in os.scandir(scripts_dir):
with suppress(OSError), open(entry.path, 'rb') as script:
first_line = script.readline(256)
if first_line.startswith(b'#!') and b'python' in first_line:
entry_points.update((entry.name, entry.path, os.path.realpath(entry.path)))
return frozenset(entry_points)
_python_programs: frozenset[str] = frozenset()
def _is_python_program(program: str) -> bool:
return program in _python_programs or _PYTHON_PROGRAM.fullmatch(os.path.basename(program)) is not None
# Only the start of a shell command names the program it runs. `shlex` tokenizes in pure Python, one character at a
# time, so splitting a whole command that carries a large payload (a base64-encoded file write) costs seconds.
_SHELL_COMMAND_SCAN_CHARS = 1024
def _leading_words(command: str) -> list[str]:
"""The words at the start of a shell command, up to an unterminated quote where the scan was cut off."""
lexer = shlex.shlex(command[:_SHELL_COMMAND_SCAN_CHARS], posix=True)
lexer.whitespace_split = True
words: list[str] = []
with suppress(ValueError):
words.extend(lexer)
return words
def launches_python(argv: Sequence[str]) -> bool:
"""Whether a process started with `argv` is a Python interpreter.
Looks through `sh -c '<command>'`, `env [VAR=value ...] <command>` and `uv run <command>`.
"""
if not argv:
return False
program, *rest = argv
name = os.path.basename(program)
if name in _SHELLS and len(rest) >= 2 and rest[0] == '-c':
return launches_python(_leading_words(rest[1]))
if name == 'env':
while rest and (rest[0].startswith('-') or '=' in rest[0]):
rest.pop(0)
return launches_python(rest)
if name == 'uv' and 'run' in rest:
return any(_is_python_program(arg) for arg in rest[rest.index('run') + 1 :] if not arg.startswith('-'))
return _is_python_program(program)
def _popen_argv(args: Any, executable: Any, shell: bool) -> list[str]:
if isinstance(args, str | bytes | os.PathLike):
argv = [os.fsdecode(args)] # pyright: ignore[reportUnknownArgumentType]
else:
argv = [os.fsdecode(arg) for arg in args]
if shell:
return ['/bin/sh', '-c', *argv]
if executable is not None:
argv[0] = os.fsdecode(executable)
return argv
# --- Per-test state ---------------------------------------------------------------------------------------------------
@dataclass
class _Spawn:
argv: list[str]
location: str | None
@dataclass
class _TestState:
pending_spawns: list[_Spawn] = field(default_factory=list[_Spawn])
spawned: bool = False
shared_setup_seconds: float = 0.0
shared_fixture_depth: int = 0
setup_seconds: float = 0.0
eligible_for_stale_check: bool = True
_STATE_KEY = pytest.StashKey[_TestState]()
_BUDGET_KEY = pytest.StashKey[float]()
_COLLECT_BUDGET_KEY = pytest.StashKey[float]()
_DESELECTED_KEY = pytest.StashKey[set[str]]()
_MARKER_SCOPES_KEY = pytest.StashKey[dict[str, tuple[int, bool]]]()
_AGGREGATE_KEY = pytest.StashKey['_SessionAggregate']()
_MODULE_COLLECTION_KEY = pytest.StashKey[dict[str, float]]()
# The sessions whose terminal this process owns, innermost last; `pytest-xdist` workers own none.
_aggregates: list[_SessionAggregate] = []
# Innermost last: `pytester` runs a session inside a test, and its tests must not report to the outer one.
_running: list[pytest.Item] = []
def _spawn_site() -> str | None:
"""The innermost frame under `tests/` that led to the spawn, so the failure points at the test's own code.
Walks frames rather than `traceback.extract_stack()`, which reads source lines from disk.
"""
for frame, lineno in traceback.walk_stack(None):
filename = frame.f_code.co_filename
if filename != _THIS_FILE and filename.startswith(_TESTS_DIR):
return f'tests/{filename[len(_TESTS_DIR) :]}:{lineno}'
return None
def _record_spawn(argv: list[str]) -> None:
if not _running or not launches_python(argv):
return
item = _running[-1]
state = item.stash[_STATE_KEY]
if state.shared_fixture_depth:
# A fixture shared beyond one test pays for its interpreter once per worker, not once per test: that is the
# pattern for expensive servers, so it needs no marker.
return
marked = item.get_closest_marker('subprocess') is not None
state.pending_spawns.append(_Spawn(argv, None if marked else _spawn_site()))
_spawn_hooks_installed = False
def _install_spawn_hooks() -> None:
"""Patch the two points every Python process launch in the standard library goes through.
`subprocess.Popen._execute_child` covers `subprocess.run`/`check_output`/`call`, `asyncio.create_subprocess_*`,
`anyio.run_process`/`open_process` and Trio; `multiprocessing.util.spawnv_passfds` covers the `spawn` and
`forkserver` start methods. Both are only reached when a process is started, so the guard costs nothing otherwise.
"""
global _spawn_hooks_installed
if _spawn_hooks_installed:
return
_spawn_hooks_installed = True
# Private and untyped in typeshed.
execute_child: Callable[..., Any] = getattr(subprocess.Popen, '_execute_child')
spawnv_passfds = multiprocessing.util.spawnv_passfds
def guarded_execute_child(self: subprocess.Popen[Any], args: Any, executable: Any, *rest: Any) -> Any:
# `rest[7]` is `shell`, passed positionally on every platform since Python 3.11.
_record_spawn(_popen_argv(args, executable, bool(rest[7])))
return execute_child(self, args, executable, *rest)
def guarded_spawnv_passfds(path: Any, args: Sequence[Any], passfds: Any) -> Any:
_record_spawn([os.fsdecode(path), *(os.fsdecode(arg) for arg in args[1:])])
return spawnv_passfds(path, args, passfds)
setattr(subprocess.Popen, '_execute_child', guarded_execute_child)
multiprocessing.util.spawnv_passfds = guarded_spawnv_passfds
# --- Configuration and marker validation ------------------------------------------------------------------------------
def pytest_configure(config: pytest.Config) -> None:
global _python_programs
_python_programs = _python_programs or _python_entry_points()
_install_spawn_hooks()
config.stash[_BUDGET_KEY] = float(os.environ.get(BUDGET_ENV_VAR) or 0)
config.stash[_COLLECT_BUDGET_KEY] = float(os.environ.get(COLLECT_BUDGET_ENV_VAR) or 0)
if not hasattr(config, 'workerinput'):
config.stash[_AGGREGATE_KEY] = aggregate = _SessionAggregate()
_aggregates.append(aggregate)
def pytest_unconfigure(config: pytest.Config) -> None:
if (aggregate := config.stash.get(_AGGREGATE_KEY, None)) is not None:
_aggregates.remove(aggregate)
@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(item: pytest.Item) -> None:
for name in ('subprocess', 'slow'):
marker = item.get_closest_marker(name)
if marker is None:
continue
reason = marker.kwargs.get('reason')
if marker.args or not (isinstance(reason, str) and reason.strip()):
pytest.fail(
f'`@pytest.mark.{name}` requires `reason=...`, a non-empty string saying why this test '
f'needs the exemption; see {_GUIDANCE_DOC}',
pytrace=False,
)
# --- Running a test ---------------------------------------------------------------------------------------------------
@pytest.hookimpl(wrapper=True)
def pytest_runtest_protocol(item: pytest.Item) -> Generator[None, object, object]:
item.stash[_STATE_KEY] = _TestState()
_running.append(item)
try:
return (yield)
finally:
_running.pop()
@pytest.hookimpl(wrapper=True)
def pytest_fixture_setup(fixturedef: pytest.FixtureDef[Any]) -> Generator[None, object, object]:
"""Track fixtures shared beyond one test, so their setup is not charged to whichever test happens to run first."""
if fixturedef.scope == 'function' or not _running:
return (yield)
state = _running[-1].stash[_STATE_KEY]
state.shared_fixture_depth += 1
start = time.perf_counter()
try:
return (yield)
finally:
state.shared_fixture_depth -= 1
if not state.shared_fixture_depth:
state.shared_setup_seconds += time.perf_counter() - start
def _fail_report(report: pytest.TestReport, message: str) -> None:
if report.passed:
report.outcome = 'failed'
report.longrepr = message
else:
report.sections.append(('cost guard', message))
def _unmarked_spawn_message(item: pytest.Item, when: str, spawns: list[_Spawn]) -> str:
lines = [f'`{item.nodeid}` launched a Python interpreter during {when}:']
for spawn in spawns:
command = shlex.join(spawn.argv)
lines.append(f' {command[:200]}{"..." if len(command) > 200 else ""}')
if spawn.location:
lines.append(f' from {spawn.location}')
lines.append(
'Each interpreter imports the project again, which makes the suite materially slower. Call the entry '
f'point in-process instead; see {_GUIDANCE_DOC}.'
)
lines.append(
'If the test is about the process boundary itself (CLI startup, import isolation, signal handling), '
"mark it `@pytest.mark.subprocess(reason='...')`."
)
return '\n'.join(lines)
def _over_budget_message(item: pytest.Item, spent: float, budget: float, state: _TestState) -> str:
shared = (
f', not counting {state.shared_setup_seconds:.2f}s of shared fixture setup'
if state.shared_setup_seconds
else ''
)
return (
f'`{item.nodeid}` took {spent:.2f}s in setup and call{shared}, over the {budget:g}s per-test budget '
f'({BUDGET_ENV_VAR}).\n'
'Make it cheaper: inject or monkeypatch timeouts and poll intervals instead of waiting them out, prove '
'complexity bounds with growth ratios at small sizes instead of one huge input, share expensive servers '
f'through session-scoped fixtures, and call entry points in-process; see {_GUIDANCE_DOC}.\n'
"If the cost is inherent to what the test proves, mark it `@pytest.mark.slow(reason='...')`."
)
@pytest.hookimpl(wrapper=True)
def pytest_runtest_makereport(
item: pytest.Item, call: pytest.CallInfo[None]
) -> Generator[None, pytest.TestReport, pytest.TestReport]:
report: pytest.TestReport = yield
# `pytest_runtest_protocol` stashes the state before any report is made.
state = item.stash[_STATE_KEY]
spawns, state.pending_spawns = state.pending_spawns, []
if spawns:
state.spawned = True
if item.get_closest_marker('subprocess') is None:
_fail_report(report, _unmarked_spawn_message(item, report.when, spawns))
if report.when == 'setup':
state.setup_seconds = report.duration
elif report.when == 'call':
budget = item.config.stash[_BUDGET_KEY]
spent = max(state.setup_seconds - state.shared_setup_seconds, 0) + report.duration
if budget and spent > budget and report.passed and item.get_closest_marker('slow') is None:
_fail_report(report, _over_budget_message(item, spent, budget, state))
if not report.passed:
state.eligible_for_stale_check = False
if report.when == 'teardown' and (marker_info := _subprocess_marker_info(item, state)) is not None:
report.python_subprocess_marker = marker_info # pyright: ignore[reportAttributeAccessIssue]
return report
# --- Stale `subprocess` markers ---------------------------------------------------------------------------------------
#
# A marker is stale when no test it covers launched Python. Each test reports what it did on its teardown report,
# which `pytest-xdist` forwards to the controller, and the process that owns the terminal decides once every test
# has reported. A marker is only judged when the run covered all of its tests, all of them passed, and at least one
# ran, so `-k`, `-m`, `--deselect`, `pytest-split` shards, node-id selection, skips and failures never produce a
# false "stale" verdict: they leave the marker unjudged in that run.
def _marker_scope(item: pytest.Item) -> tuple[str, str] | None:
"""The key and display location of the node a test's `subprocess` marker was declared on."""
for node, marker in item.iter_markers_with_node('subprocess'):
if marker.kwargs.get('conditional'):
return None
if node is item and isinstance(item, pytest.Function):
# Every parametrization of one function shares the decorator, so they are judged together.
return f'{item.parent.nodeid if item.parent else ""}::{item.originalname}', item.nodeid.split('[', 1)[0]
return node.nodeid, node.nodeid
return None
@pytest.hookimpl(trylast=True)
def pytest_deselected(items: Sequence[pytest.Item]) -> None:
for item in items:
if (scope := _marker_scope(item)) is not None:
item.config.stash.setdefault(_DESELECTED_KEY, set()).add(scope[0])
def pytest_collection_finish(session: pytest.Session) -> None:
config = session.config
deselected = config.stash.get(_DESELECTED_KEY, set[str]())
# A `path::name` argument selects part of a file, so a marker declared on that file's module or one of its
# classes may cover tests that were never collected.
narrowed_files = {
(config.invocation_params.dir / arg.split('::', 1)[0]).resolve() for arg in config.args if '::' in arg
}
scoped = [(item, scope[0]) for item in session.items if (scope := _marker_scope(item)) is not None]
sizes: dict[str, int] = {}
incomplete = set(deselected)
for item, key in scoped:
sizes[key] = sizes.get(key, 0) + 1
if narrowed_files and item.path.resolve() in narrowed_files:
incomplete.add(key)
config.stash[_MARKER_SCOPES_KEY] = {key: (size, key not in incomplete) for key, size in sizes.items()}
class _MarkerInfo(TypedDict):
"""What a teardown report carries about the test's `subprocess` marker; plain JSON so `pytest-xdist` can ship it."""
key: str
location: str
size: int
complete: bool
eligible: bool
spawned: bool
def _subprocess_marker_info(item: pytest.Item, state: _TestState) -> _MarkerInfo | None:
if (scope := _marker_scope(item)) is None:
return None
key, location = scope
size, complete = item.config.stash.get(_MARKER_SCOPES_KEY, dict[str, tuple[int, bool]]()).get(key, (0, False))
return {
'key': key,
'location': location,
'size': size,
'complete': complete,
'eligible': state.eligible_for_stale_check,
'spawned': state.spawned,
}
@dataclass
class _MarkerAggregate:
location: str
size: int
complete: bool = True
reported: set[str] = field(default_factory=set[str])
eligible: bool = True
spawned: bool = False
@dataclass
class _SessionAggregate:
markers: dict[str, _MarkerAggregate] = field(default_factory=dict[str, _MarkerAggregate])
stale: list[str] = field(default_factory=list[str])
session_start: float = 0.0
collection_end: float = 0.0
collected: int = 0
slowest_worker_collection: float = 0.0
slowest_modules: dict[str, float] = field(default_factory=dict[str, float])
over_collection_budget: dict[str, float] = field(default_factory=dict[str, float])
def stale_markers(self) -> list[str]:
return sorted(
marker.location
for marker in self.markers.values()
if marker.complete and marker.eligible and len(marker.reported) == marker.size and not marker.spawned
)
def pytest_runtest_logreport(report: pytest.TestReport) -> None:
if not _aggregates:
return
aggregate = _aggregates[-1]
info: _MarkerInfo | None = getattr(report, 'python_subprocess_marker', None)
if info is None:
return
marker = aggregate.markers.setdefault(info['key'], _MarkerAggregate(info['location'], info['size']))
marker.complete &= info['complete']
marker.eligible &= info['eligible']
marker.reported.add(report.nodeid)
marker.spawned |= info['spawned']
@pytest.hookimpl(tryfirst=True)
def pytest_sessionstart(session: pytest.Session) -> None:
if (aggregate := session.config.stash.get(_AGGREGATE_KEY, None)) is not None:
aggregate.session_start = time.perf_counter()
@pytest.hookimpl(tryfirst=True)
def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
aggregate = session.config.stash.get(_AGGREGATE_KEY, None)
if aggregate is None or session.shouldstop or session.shouldfail:
return
if exitstatus not in (pytest.ExitCode.OK, pytest.ExitCode.TESTS_FAILED):
return
aggregate.stale = aggregate.stale_markers()
if aggregate.stale or aggregate.over_collection_budget:
session.exitstatus = pytest.ExitCode.TESTS_FAILED
# --- Collection time --------------------------------------------------------------------------------------------------
#
# Collection runs once per `pytest-xdist` worker, so its cost multiplies with the worker count; it mostly comes from
# test modules importing heavy optional dependencies. The summary names the slowest modules to collect, so a module
# that starts importing something expensive shows up in the next CI run, and `PYTEST_COLLECT_BUDGET_SECONDS` fails the
# session over one that takes longer than that. The first module in a run to import a dependency that many modules
# share is charged for it, so the budget sits above the cost of the shared ones.
_SLOWEST_MODULES_SHOWN = 10
@pytest.hookimpl(wrapper=True)
def pytest_make_collect_report(
collector: pytest.Collector,
) -> Generator[None, pytest.CollectReport, pytest.CollectReport]:
if not isinstance(collector, pytest.Module):
return (yield)
start = time.perf_counter()
try:
return (yield)
finally:
collector.config.stash.setdefault(_MODULE_COLLECTION_KEY, {})[collector.nodeid] = time.perf_counter() - start
def _slowest_modules(config: pytest.Config) -> dict[str, float]:
seconds = config.stash.get(_MODULE_COLLECTION_KEY, dict[str, float]())
return dict(sorted(seconds.items(), key=lambda item: item[1], reverse=True)[:_SLOWEST_MODULES_SHOWN])
def _over_collection_budget(config: pytest.Config) -> dict[str, float]:
budget = config.stash[_COLLECT_BUDGET_KEY]
seconds = config.stash.get(_MODULE_COLLECTION_KEY, dict[str, float]())
return {module: spent for module, spent in seconds.items() if budget and spent > budget}
@pytest.hookimpl(wrapper=True)
def pytest_collection(session: pytest.Session) -> Generator[None, object, object]:
start = time.perf_counter()
try:
return (yield)
finally:
config = session.config
if hasattr(config, 'workerinput'):
workeroutput: dict[str, Any] = getattr(config, 'workeroutput') # set by `pytest-xdist` on workers
workeroutput['collection_seconds'] = time.perf_counter() - start
workeroutput['slowest_modules'] = _slowest_modules(config)
workeroutput['over_collection_budget'] = _over_collection_budget(config)
elif (aggregate := config.stash.get(_AGGREGATE_KEY, None)) is not None and not config.pluginmanager.has_plugin(
'dsession'
):
aggregate.session_start = start
aggregate.collection_end = time.perf_counter()
aggregate.collected = len(session.items)
aggregate.slowest_modules = _slowest_modules(config)
aggregate.over_collection_budget = _over_collection_budget(config)
@pytest.hookimpl(optionalhook=True)
def pytest_xdist_node_collection_finished(node: Any, ids: Sequence[str]) -> None:
aggregate = node.config.stash[_AGGREGATE_KEY]
aggregate.collection_end = time.perf_counter()
aggregate.collected = len(ids)
@pytest.hookimpl(optionalhook=True)
def pytest_testnodedown(node: Any, error: object) -> None:
aggregate = node.config.stash[_AGGREGATE_KEY]
output = getattr(node, 'workeroutput', {})
aggregate.slowest_worker_collection = max(aggregate.slowest_worker_collection, output.get('collection_seconds', 0))
for module, seconds in output.get('slowest_modules', {}).items():
aggregate.slowest_modules[module] = max(aggregate.slowest_modules.get(module, 0), seconds)
for module, seconds in output.get('over_collection_budget', {}).items():
aggregate.over_collection_budget[module] = max(aggregate.over_collection_budget.get(module, 0), seconds)
def pytest_terminal_summary(terminalreporter: pytest.TerminalReporter, config: pytest.Config) -> None:
aggregate = config.stash.get(_AGGREGATE_KEY, None) or _SessionAggregate()
if aggregate.stale:
terminalreporter.section('stale subprocess markers', red=True)
for location in aggregate.stale:
terminalreporter.write_line(
f'{location}: no test covered by its `@pytest.mark.subprocess` launched a Python interpreter; '
'remove the stale marker'
)
if aggregate.over_collection_budget:
terminalreporter.section('slow test module collection', red=True)
budget = config.stash[_COLLECT_BUDGET_KEY]
for module, seconds in sorted(aggregate.over_collection_budget.items()):
terminalreporter.write_line(
f'{module} took {seconds:.2f}s to collect, over the {budget:g}s per-module budget '
f'({COLLECT_BUDGET_ENV_VAR}).'
)
terminalreporter.write_line(
'Every pytest-xdist worker collects every module in its shard, so this is paid once per worker. Import '
'heavy optional dependencies inside the tests or fixtures that need them, checking whether they are '
'installed with `importlib.util.find_spec`: a `module`- or `session`-scoped fixture keeps the import '
f'outside the per-test budget. Build large parametrizations and data lazily too; see {_GUIDANCE_DOC}.'
)
# Not set when every `pytest-xdist` worker died before collecting.
if not aggregate.collection_end: # pragma: lax no cover
return
seconds = aggregate.collection_end - aggregate.session_start
if aggregate.slowest_worker_collection:
summary = (
f'worker startup and collection of {aggregate.collected} tests took {seconds:.1f}s '
f'(slowest worker collected in {aggregate.slowest_worker_collection:.1f}s)'
)
else:
summary = f'collection of {aggregate.collected} tests took {seconds:.1f}s'
modules = sorted(aggregate.slowest_modules.items(), key=lambda item: item[1], reverse=True)[:_SLOWEST_MODULES_SHOWN]
terminalreporter.write_line(summary)
terminalreporter.write_line(
'slowest test modules to collect: ' + ', '.join(f'{module} ({seconds:.2f}s)' for module, seconds in modules)
)
if step_summary := os.environ.get('GITHUB_STEP_SUMMARY'):
with open(step_summary, 'a', encoding='utf-8') as file:
file.write(f'pytest {summary}; slowest test modules to collect:\n\n')
file.writelines(f'- `{module}`: {seconds:.2f}s\n' for module, seconds in modules)
file.write('\n')
tests/durable_exec/temporal/test_agent.py
+1
View File
@@ -712,6 +712,7 @@ async def test_temporal_durability_accepts_legacy_cancel_activity_payload() -> N
assert failing_events == ['enter', 'cancel:failing', 'exit:RuntimeError']
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_complex_agent_run_in_workflow(
allow_model_requests: None, client_with_logfire: Client, capfire: CaptureLogfire
):
tests/durable_exec/temporal/test_durability.py
+1
View File
@@ -2270,6 +2270,7 @@ class ComplexDurableAgentLogfireWorkflow:
return result.output
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_durability_complex_agent_logfire_span_tree(
allow_model_requests: None, client_with_logfire: Client, capfire: CaptureLogfire
):
tests/durable_exec/temporal/test_toolsets.py
+6
View File
@@ -184,6 +184,7 @@ pytestmark = [
]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_mcp_tools_cached_across_activities(allow_model_requests: None, client: Client):
"""Verify that MCP tool caching reduces server round-trips across activities.
@@ -278,6 +279,7 @@ def _scheduled_get_tools_count(history: WorkflowHistory) -> int:
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_temporal_mcp_get_tools_replay_deterministic(allow_model_requests: None, client: Client):
"""#5875 regression: `get_tools` activity scheduling must be replay-deterministic.
@@ -332,6 +334,7 @@ async def test_temporal_mcp_get_tools_replay_deterministic(allow_model_requests:
mcp_replay_holder['agent'] = warm
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_temporal_mcp_get_tools_not_cached_when_disabled(allow_model_requests: None, client: Client):
"""With `cache_tools=False`, `get_tools` is scheduled for every model request (no run cache).
@@ -1021,6 +1024,7 @@ class MCPToolsetInstructionsWorkflow:
return result.output
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_temporal_mcptoolset_instructions_propagate(client: Client):
"""`MCPToolset` instructions propagate through the `TemporalMCPToolset` wrapper."""
async with Worker(
@@ -1335,6 +1339,7 @@ class MCPTaskSupportWorkflow:
return (await _mcp_task_temporal_agent.run(prompt)).output
@pytest.mark.subprocess(reason='connects to `tests.mcp_task_server` over stdio')
async def test_temporal_mcptoolset_preserves_task_routing(client: Client):
"""Effective task routing in `ToolDefinition.metadata` survives Temporal activities."""
async with Worker(
@@ -1745,6 +1750,7 @@ async def _heartbeats_during_activity(activity_fn: Callable[..., Any], args: Seq
return beats
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_every_registered_activity_heartbeats(allow_model_requests: None):
"""Every activity Pydantic AI registers beats while it runs, not just the model ones (#6914).
tests/durable_exec/test_capability_durable_operation.py
-27
View File
@@ -2176,33 +2176,6 @@ async def test_dbos_capability_usage_delta_is_stable_on_replay(dbos: DBOS) -> No
assert capability.calls == 1
@pytest.fixture(scope='module')
def prefect_test_server() -> Generator[None, None, None]:
"""Run the module's Prefect flow tests against an isolated test server.
The implicit ephemeral server uses the shared default PREFECT_HOME and a short
connect timeout that flakes on slow CI runners; the test harness gives an isolated
database and a 60s startup budget, mirroring tests/durable_exec/test_prefect.py.
Ordering constraint: these tests share the 'prefect' xdist group with
tests/durable_exec/test_prefect.py, whose harness fixture is session-scoped and
autouse, so once entered it stays entered for the rest of the worker. Today this
module collects first purely because it sorts alphabetically ahead of test_prefect.py,
so this module-scoped harness enters and exits before that one starts. Renaming either
module so this one sorts after test_prefect.py would nest two Prefect harnesses.
"""
from prefect.settings import PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED, temporary_settings
from prefect.testing.utilities import prefect_test_harness
# The task-run recorder is a background writer against the same sqlite file the flows write to.
# Prefect PRAGMAs a 60s `busy_timeout` onto every connection, and under CI contention the
# recorder's bulk inserts exhaust it, failing the flow whose state it was recording. Nothing
# here reads what it records: task run states reach the API through the task engine.
with temporary_settings({PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED: False}):
with prefect_test_harness(server_startup_timeout=60):
yield
@requires_prefect
@pytest.mark.xdist_group(name='prefect')
async def test_prefect_capability_operation_end_to_end(prefect_test_server: None) -> None:
tests/durable_exec/test_dbos.py
+13
View File
@@ -347,6 +347,7 @@ iter_handler_stream_agent = Agent(
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_complex_agent_run_in_workflow(allow_model_requests: None, dbos: DBOS, capfire: CaptureLogfire) -> None:
# Set a workflow ID for testing list steps
wfid = str(uuid.uuid4())
@@ -2170,6 +2171,7 @@ _uninit_instructions_toolset = dbosify_mcp_toolset(
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_toolset_get_instructions_falls_back_to_step(dbos: DBOS):
"""When the MCP server isn't initialized locally, DBOS wrapper fetches instructions via a step."""
run_context = RunContext(deps=0, model=TestModel(), usage=RunUsage())
@@ -2178,6 +2180,7 @@ async def test_dbos_mcp_toolset_get_instructions_falls_back_to_step(dbos: DBOS):
assert instructions == InstructionPart(content='Be a helpful assistant.', dynamic=False)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_toolset_get_instructions_uses_step_when_server_warm(dbos: DBOS):
"""A warm in-process MCP server must not short-circuit the `get_instructions` step (#5884).
@@ -2233,6 +2236,7 @@ mcptoolset_instructions_agent = Agent(
mcptoolset_instructions_dbos_agent = DBOSAgent(mcptoolset_instructions_agent) # pyright: ignore[reportDeprecated]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcptoolset_instructions_propagate(dbos: DBOS):
"""`MCPToolset` instructions propagate through the `DBOSMCPToolset` wrapper."""
result = await mcptoolset_instructions_dbos_agent.run('Use MCP instructions')
@@ -2285,6 +2289,7 @@ _mcp_task_dbos_agent = DBOSAgent( # pyright: ignore[reportDeprecated]
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_task_server` over stdio')
async def test_dbos_mcptoolset_preserves_task_routing(dbos: DBOS):
"""Effective task routing in `ToolDefinition.metadata` survives DBOS steps."""
result = await _mcp_task_dbos_agent.run('Call both tools')
@@ -2312,6 +2317,7 @@ mcp_replay_agent = Agent(
mcp_replay_dbos_agent = DBOSAgent(mcp_replay_agent) # pyright: ignore[reportDeprecated]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_get_tools_recorded_independently_per_run(allow_model_requests: None, dbos: DBOS):
"""#5875 regression: DBOS `get_tools` step scheduling must depend only on the workflow's own history.
@@ -2360,6 +2366,7 @@ mcp_retry_budget_agent = Agent(
mcp_retry_budget_dbos_agent = DBOSAgent(mcp_retry_budget_agent) # pyright: ignore[reportDeprecated]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_tool_inherits_agent_retries(allow_model_requests: None, dbos: DBOS):
"""#5180 regression: a durably-wrapped MCP tool enforces the agent's tool-retry budget, not a hard-coded 1.
@@ -2696,6 +2703,7 @@ async def test_unwrap_recorded_tool_call_result_handles_both_generations() -> No
assert unwrap_recorded_tool_call_result('raw recorded output') == 'raw recorded output'
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_model_retry_crosses_step_without_engine_retries(
dbos: DBOS, monkeypatch: pytest.MonkeyPatch
) -> None:
@@ -2765,6 +2773,7 @@ async def test_dbos_dynamic_tool_model_retry_crosses_step_without_engine_retries
assert calls == 2
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_step_rejects_enqueue_in_workflow(dbos: DBOS, monkeypatch: pytest.MonkeyPatch) -> None:
"""The MCP step path guards enqueue too: a `process_tool_call=` hook receives the run context."""
mcp_toolset = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='enqueue_mcp')
@@ -3035,6 +3044,7 @@ async def test_dbos_dynamic_get_tools_rejects_enqueue_in_workflow(dbos: DBOS) ->
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_get_tools_and_instructions_reject_enqueue_in_workflow(
dbos: DBOS, monkeypatch: pytest.MonkeyPatch
) -> None:
@@ -3953,6 +3963,7 @@ async def test_dbos_durability_mcp_toolset_wrapping(dbos: DBOS) -> None:
assert isinstance(bound._toolsets_by_id['my_mcp'], DBOSMCPToolset) # pyright: ignore[reportPrivateUsage]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_durability_mcp_operations_run_in_steps(dbos: DBOS) -> None:
seen_instructions: list[str] = []
@@ -4233,6 +4244,7 @@ async def test_dbos_durability_ignores_per_tool_metadata(dbos: DBOS, metadata: d
assert 'dbos_metadata_ignored__dynamic_toolset__dyn.call_tool' in step_names
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_durability_mcp_tool_metadata_false_is_rejected(dbos: DBOS, monkeypatch: pytest.MonkeyPatch) -> None:
"""MCP tools perform I/O and cannot opt out of their durable DBOS step."""
mcp_toolset = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='dbos_mcp_opt_out')
@@ -4291,6 +4303,7 @@ def test_dbos_durability_bare_capability_func_requires_explicit_wrapper(dbos: DB
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_durability_dynamic_capability_mcp_runs_in_steps(dbos: DBOS) -> None:
def call_then_answer(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse:
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
tests/durable_exec/test_durable_operation.py
+1
View File
@@ -296,6 +296,7 @@ def test_durability_without_tool_config_key_ignores_tool_metadata() -> None:
assert resolve(tool, 'configured') == {'base': 1}
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_mcp_tool_config_dispatches_durable_and_inline_calls(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip('mcp')
from fastmcp.client.transports import StdioTransport
tests/durable_exec/test_prefect.py
+5 -14
View File
@@ -132,8 +132,6 @@ from pydantic_ai.workspaces import (
try:
from prefect import flow, task
from prefect.context import FlowRunContext, TaskRunContext
from prefect.settings import PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED, temporary_settings
from prefect.testing.utilities import prefect_test_harness
from pydantic_ai.durable_exec.prefect import (
DEFAULT_PYDANTIC_AI_CACHE_POLICY,
@@ -387,6 +385,7 @@ warnings.filterwarnings('ignore', message='`PrefectAgent` is deprecated', catego
pytestmark = [
pytest.mark.vcr,
pytest.mark.xdist_group(name='prefect'),
pytest.mark.usefixtures('prefect_test_server'),
pytest.mark.filterwarnings(
'ignore:`PrefectAgent` is deprecated:pydantic_ai._warnings.PydanticAIDeprecationWarning'
),
@@ -414,18 +413,6 @@ def setup_logfire_instrumentation() -> Iterator[None]:
yield
@pytest.fixture(autouse=True, scope='session')
def setup_prefect_test_harness() -> Iterator[None]:
"""Set up Prefect test harness for all tests."""
# The task-run recorder is a background writer against the same sqlite file the flows write to.
# Prefect PRAGMAs a 60s `busy_timeout` onto every connection, and under CI contention the
# recorder's bulk inserts exhaust it, failing the flow whose state it was recording. Nothing
# here reads what it records: task run states reach the API through the task engine.
with temporary_settings({PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED: False}):
with prefect_test_harness(server_startup_timeout=60):
yield
@pytest.fixture(autouse=True)
def blockbuster_excluded_modules() -> tuple[str, ...]:
"""Prefect's `@flow` constructor synchronously inspects its decorated function's source."""
@@ -596,6 +583,7 @@ runtime_handler_stream_agent = Agent(
runtime_handler_stream_prefect_agent = PrefectAgent(runtime_handler_stream_agent) # pyright: ignore[reportDeprecated]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_complex_agent_run_in_flow(allow_model_requests: None, capfire: CaptureLogfire) -> None:
"""Test a complex agent with tools, MCP servers, and event stream handler."""
@@ -1113,6 +1101,7 @@ async def test_prefect_toolset_legacy_constructors() -> None:
assert wrapped_mcp.id is None
@pytest.mark.subprocess(reason='connects to `tests.mcp_task_server` over stdio')
async def test_prefect_mcptoolset_preserves_task_routing() -> None:
"""Effective task routing forwards through Prefect task wrappers end-to-end.
@@ -3267,6 +3256,7 @@ async def test_prefect_durability_task_name_assembly_sequence() -> None:
]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
@pytest.mark.parametrize('blockbuster_enabled', [False])
async def test_prefect_durability_journals_mcp_discovery(blockbuster_enabled: bool) -> None:
assert blockbuster_enabled is False
@@ -4286,6 +4276,7 @@ async def test_prefect_mcp_tool_metadata_false_is_rejected() -> None:
await run_tool()
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
@pytest.mark.parametrize('blockbuster_enabled', [False])
async def test_prefect_durability_mcp_tool_metadata_false_is_rejected(
monkeypatch: pytest.MonkeyPatch, blockbuster_enabled: bool
tests/durable_exec/test_prefect_workspace.py
+1 -12
View File
@@ -9,7 +9,6 @@ adds what only Prefect has.
from __future__ import annotations
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Any
@@ -28,23 +27,13 @@ from .workspace_scenarios import SCENARIOS, Check, ScenarioFailed, cases, scenar
try:
from prefect import flow
from prefect.context import FlowRunContext, TaskRunContext
from prefect.settings import PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED, temporary_settings
from prefect.testing.utilities import prefect_test_harness
from pydantic_ai.durable_exec.prefect import PrefectDurability, TaskConfig
except ImportError: # pragma: lax no cover
pytest.skip('Prefect is not installed', allow_module_level=True)
pytestmark = pytest.mark.xdist_group(name='prefect')
@pytest.fixture(autouse=True, scope='session')
def setup_prefect_test_harness() -> Iterator[None]:
# See `test_prefect.py`: the task-run recorder's background writer contends for the sqlite file.
with temporary_settings({PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED: False}):
with prefect_test_harness(server_startup_timeout=60):
yield
pytestmark = [pytest.mark.xdist_group(name='prefect'), pytest.mark.usefixtures('prefect_test_server')]
@pytest.fixture(autouse=True)
tests/durable_exec/test_temporal_skip_gating.py
+2
View File
@@ -82,6 +82,7 @@ def _run_masked_pytest(target: str) -> subprocess.CompletedProcess[str]:
)
@pytest.mark.subprocess(reason='runs pytest in a fresh process to check how it loads the Temporal conftest')
def test_temporal_suite_skips_when_invoked_on_its_own_directory() -> None:
"""`pytest tests/durable_exec/temporal` reports skips instead of crashing on the conftest import."""
result = _run_masked_pytest('tests/durable_exec/temporal')
@@ -99,6 +100,7 @@ def test_temporal_suite_skips_when_invoked_on_its_own_directory() -> None:
assert 'Skipped: ' not in output, output
@pytest.mark.subprocess(reason='runs pytest in a fresh process to check how it loads the Temporal conftest')
def test_temporal_suite_skips_when_invoked_on_a_node_id() -> None:
"""A node id under the suite reports the skip too — the form `test-temporal-latest` runs.
tests/evals/test_online.py
+1
View File
@@ -2452,6 +2452,7 @@ async def test_baggage_disabled_via_config(capfire: CaptureLogfire):
assert 'tenant' not in attrs
@pytest.mark.subprocess(reason='needs an interpreter where `sniffio` cannot be imported')
@pytest.mark.skipif(
any(req.startswith('sniffio') for req in importlib.metadata.requires('anyio') or []),
reason='AnyIO before 4.12 depends on and imports `sniffio`, so every install has it',
tests/harness/e2b_sandbox/test_backend.py
+7 -14
View File
@@ -3,11 +3,11 @@
from __future__ import annotations
import asyncio
import importlib
import os
import re
import shlex
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
@@ -1019,19 +1019,12 @@ async def test_missing_key_is_reported_as_not_found(fake_e2b: FakeE2B) -> None:
assert str(exc.value) == 'No E2B API key found. Set E2B_API_KEY in the environment.'
def test_missing_e2b_extra_has_an_install_hint() -> None:
result = subprocess.run(
[
sys.executable,
'-c',
"import sys; sys.modules['e2b'] = None; import pydantic_ai_harness.e2b_sandbox",
],
capture_output=True,
text=True,
check=False,
)
assert result.returncode != 0
assert 'Install `pydantic-ai-harness[e2b]`' in result.stderr
def test_missing_e2b_extra_has_an_install_hint(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(sys.modules, 'e2b', None)
for name in [name for name in sys.modules if name.startswith('pydantic_ai_harness.e2b_sandbox')]:
monkeypatch.delitem(sys.modules, name)
with pytest.raises(ImportError, match=r'^Install `pydantic-ai-harness\[e2b\]` to use E2BSandbox\.$'):
importlib.import_module('pydantic_ai_harness.e2b_sandbox')
async def test_command_timeout_starts_once_the_sandbox_is_acquired(fake_e2b: FakeE2B) -> None:
tests/harness/experimental/acp/_demo_approval_agent.py
-21
View File
@@ -1,21 +0,0 @@
"""An ACP agent with an approval-gated tool, used by the end-to-end permission test."""
from __future__ import annotations
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
def build_agent() -> Agent[None, str]:
agent = Agent(TestModel())
@agent.tool_plain(requires_approval=True)
def delete_file(path: str) -> str:
return f'deleted {path}'
return agent
if __name__ == '__main__':
run_acp_stdio_sync(build_agent())
tests/harness/experimental/acp/_demo_large_agent.py
-12
View File
@@ -1,12 +0,0 @@
"""An ACP agent that streams a large message, used by the large-payload stdio test."""
from __future__ import annotations
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
LARGE_OUTPUT = 'x' * 200_000
if __name__ == '__main__':
run_acp_stdio_sync(Agent(TestModel(custom_output_text=LARGE_OUTPUT)))
tests/harness/experimental/acp/_demo_models_agent.py
-10
View File
@@ -1,10 +0,0 @@
"""An ACP agent advertising switchable models, used by the stdio model-config test."""
from __future__ import annotations
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
if __name__ == '__main__':
run_acp_stdio_sync(Agent(TestModel(custom_output_text='hi')), models=['test', 'openai:gpt-4o'])
tests/harness/experimental/acp/_demo_native_agent.py
-42
View File
@@ -1,42 +0,0 @@
"""An ACP agent mounting client-backed fs/terminal toolsets, for the native-over-stdio test.
The model calls `read_file` then `run_command` on the first turn; the session config mounts the
ACP-client-backed toolsets, so those tool calls must route over the wire to the test's client.
"""
from __future__ import annotations
import json
from collections.abc import AsyncIterator
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage, ToolReturnPart
from pydantic_ai.models.function import AgentInfo, DeltaToolCall, FunctionModel
from pydantic_ai_harness.experimental.acp import (
AcpSession,
AcpSessionConfig,
acp_filesystem,
acp_terminal,
run_acp_stdio_sync,
)
async def stream(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | dict[int, DeltaToolCall]]:
if any(isinstance(part, ToolReturnPart) for message in messages for part in getattr(message, 'parts', [])):
yield 'done'
return
yield {
0: DeltaToolCall(name='read_file', json_args=json.dumps({'path': 'notes.txt'})),
1: DeltaToolCall(name='run_command', json_args=json.dumps({'command': 'echo hi'})),
}
def session_config(session: AcpSession) -> AcpSessionConfig[None]:
capabilities = [
capability for capability in (acp_filesystem(session), acp_terminal(session)) if capability is not None
]
return AcpSessionConfig(deps=None, capabilities=capabilities)
if __name__ == '__main__':
run_acp_stdio_sync(Agent(FunctionModel(stream_function=stream)), session_config=session_config)
tests/harness/experimental/acp/_demo_slow_agent.py
-24
View File
@@ -1,24 +0,0 @@
"""An ACP agent with a slow tool, used by the over-the-wire cancellation test."""
from __future__ import annotations
import asyncio
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from pydantic_ai_harness.experimental.acp import run_acp_stdio_sync
def build_agent() -> Agent[None, str]:
agent = Agent(TestModel())
@agent.tool_plain
async def slow() -> str:
await asyncio.sleep(30)
return 'done'
return agent
if __name__ == '__main__':
run_acp_stdio_sync(build_agent())
tests/harness/experimental/acp/_wire.py
+17 -9
View File
@@ -19,6 +19,7 @@ import asyncio
import contextlib
import socket
from types import TracebackType
from typing import Generic, TypeVar, cast, overload
import acp
from acp.client.connection import ClientSideConnection
@@ -42,31 +43,38 @@ class WireClient(RecordingClientBase):
return ''.join(out)
class wire_agent:
ClientT = TypeVar('ClientT', bound=RecordingClientBase)
class wire_agent(Generic[ClientT]):
"""Serve `adapter` over an in-memory socket pair; `async with` yields a connected client connection.
The agent runs as a background task for the lifetime of the context; on exit the task is
cancelled and both stream ends are closed. Set `unstable=False` to drive the agent without
`use_unstable_protocol`, the configuration in which the SDK router rejects unstable methods.
Pass `client` to answer permission, filesystem or terminal requests; it defaults to a `WireClient`.
A class rather than an `@asynccontextmanager` generator only because the bundled type stubs
flag that decorator as deprecated under strict checking.
"""
@overload
def __init__(
self,
adapter: acp.Agent,
client: WireClient | None = None,
*,
unstable: bool = True,
) -> None:
self: wire_agent[WireClient], adapter: acp.Agent, client: None = None, *, unstable: bool = True
) -> None: ...
@overload
def __init__(self, adapter: acp.Agent, client: ClientT, *, unstable: bool = True) -> None: ...
def __init__(self, adapter: acp.Agent, client: ClientT | None = None, *, unstable: bool = True) -> None:
self._adapter = adapter
self._client = client or WireClient()
# Only the first overload omits `client`, and it pins `ClientT` to `WireClient`.
self._client = client if client is not None else cast(ClientT, WireClient())
self._unstable = unstable
self._server: asyncio.Task[None] | None = None
self._writers: tuple[asyncio.StreamWriter, asyncio.StreamWriter] | None = None
async def __aenter__(self) -> tuple[ClientSideConnection, WireClient]:
async def __aenter__(self) -> tuple[ClientSideConnection, ClientT]:
agent_sock, client_sock = socket.socketpair()
# Each side's `input_stream` is the writer it sends to the peer with; `output_stream` is
# the reader it receives on -- the order both SDK connection constructors require.
tests/harness/experimental/acp/test_acp.py
+65 -28
View File
@@ -2,8 +2,8 @@
Organized by behavior: lifecycle, streaming/event ordering, tool calls, permission/HITL,
cancellation, multi-turn history & isolation, errors, unsupported methods, and entry points.
Most tests drive the adapter directly with an in-memory `FakeClient`; the entry-point tests
spawn a real subprocess over stdio, as a TUI client (Zed/Toad) would.
Most tests drive the adapter directly with an in-memory `FakeClient`; the entry-point tests drive it
over the in-memory wire, and one spawns a real subprocess over stdio, as a TUI client (Zed/Toad) would.
"""
from __future__ import annotations
@@ -62,6 +62,8 @@ from pydantic_ai_harness.experimental.acp import (
InMemorySessionStore,
PydanticAIACPAgent,
ToolCallPresentation,
acp_filesystem,
acp_terminal,
chain_presenters,
default_coding_presenter,
run_acp_stdio,
@@ -89,6 +91,7 @@ from tests.harness.experimental.acp._acp_clients import (
RecordingClient,
RecordingClientBase,
)
from tests.harness.experimental.acp._wire import wire_agent
# A decider maps a permission request (the tool call) to the option_id the client "clicks",
# or to None to signal a cancelled permission outcome (the user dismissed the dialog).
@@ -2159,6 +2162,7 @@ class TestEntryPoints:
assert seen == ['host:custom']
assert response.stop_reason == 'max_turn_requests'
@pytest.mark.subprocess(reason='the one test that drives the ACP entry point over real stdio')
async def test_end_to_end_over_stdio(self) -> None:
"""Drive the adapter as a real subprocess over ACP stdio, as a TUI client would."""
client = FakeClient()
@@ -2178,11 +2182,20 @@ class TestEntryPoints:
# The streamed text is the model's JSON-encoded tool result, delivered intact.
assert 'Sunny in a' in client.text()
async def test_end_to_end_permission_over_stdio(self) -> None:
"""A tool requiring approval round-trips a real permission request over stdio."""
# The tests below drive the adapter over the in-memory wire (`_wire.py`): the real SDK router
# and JSON codec, without an interpreter start-up per test. `test_end_to_end_over_stdio` above
# is the one test that crosses a real process boundary.
async def test_end_to_end_permission_over_the_wire(self) -> None:
"""A tool requiring approval round-trips a real permission request over the wire."""
agent = Agent(TestModel())
@agent.tool_plain(requires_approval=True)
def delete_file(path: str) -> str:
return f'deleted {path}'
client = FakeClient(decider=lambda _call: 'allow_once')
script = Path(__file__).parent / '_demo_approval_agent.py'
async with acp.spawn_agent_process(client, sys.executable, str(script)) as (conn, _proc):
async with wire_agent(PydanticAIACPAgent(agent), client) as (conn, _client):
await conn.initialize(protocol_version=acp.PROTOCOL_VERSION)
session = await conn.new_session(cwd='.', mcp_servers=[])
response = await conn.prompt(
@@ -2193,40 +2206,46 @@ class TestEntryPoints:
assert len(client.permission_requests) == 1
assert client.permission_requests[0].title == 'delete_file'
async def test_cancel_notification_over_stdio(self) -> None:
"""A real `session/cancel` notification mid-prompt yields a cancelled stop reason over stdio."""
client = FakeClient()
script = Path(__file__).parent / '_demo_slow_agent.py'
async with acp.spawn_agent_process(client, sys.executable, str(script)) as (conn, _proc):
async def test_cancel_notification_over_the_wire(self) -> None:
"""A real `session/cancel` notification mid-prompt yields a cancelled stop reason over the wire."""
started = asyncio.Event()
agent = Agent(TestModel())
@agent.tool_plain
async def slow() -> str:
started.set()
await asyncio.sleep(30)
# Cancelled before it returns.
return 'done' # pragma: no cover
async with wire_agent(PydanticAIACPAgent(agent)) as (conn, _client):
await conn.initialize(protocol_version=acp.PROTOCOL_VERSION)
session = await conn.new_session(cwd='.', mcp_servers=[])
prompt = asyncio.ensure_future(
conn.prompt(session_id=session.session_id, prompt=[acp.text_block('go')], message_id=uuid4().hex)
)
await asyncio.sleep(0.5) # let the slow tool start
await asyncio.wait_for(started.wait(), timeout=10)
await conn.cancel(session_id=session.session_id)
response = await asyncio.wait_for(prompt, timeout=10)
assert response.stop_reason == 'cancelled'
async def test_large_output_over_stdio(self) -> None:
"""A large streamed message survives the default stdio buffer intact (chunked, not truncated)."""
client = FakeClient()
script = Path(__file__).parent / '_demo_large_agent.py'
# No custom transport buffer: the adapter must chunk so the default client reader (64 KiB)
# is never overrun. Raising the buffer here would mask the real-client failure mode.
async with acp.spawn_agent_process(client, sys.executable, str(script)) as (conn, _proc):
async def test_large_output_over_the_wire(self) -> None:
"""A large streamed message survives the default client read buffer intact (chunked, not truncated)."""
# `wire_agent` keeps the client reader's default 64 KiB line limit: the adapter must chunk so
# it is never overrun. Raising the buffer here would mask the real-client failure mode.
adapter = PydanticAIACPAgent(Agent(TestModel(custom_output_text='x' * 200_000)))
async with wire_agent(adapter) as (conn, client):
await conn.initialize(protocol_version=acp.PROTOCOL_VERSION)
session = await conn.new_session(cwd='.', mcp_servers=[])
await conn.prompt(session_id=session.session_id, prompt=[acp.text_block('go')], message_id=uuid4().hex)
assert client.text() == 'x' * 200_000
assert client.texts() == 'x' * 200_000
async def test_model_config_and_unstable_close_route_over_stdio(self) -> None:
async def test_model_config_and_unstable_close_route_over_the_wire(self) -> None:
"""Model config updates and unstable `session/close` reach the handler over a real wire."""
client = FakeClient()
script = Path(__file__).parent / '_demo_models_agent.py'
async with acp.spawn_agent_process(client, sys.executable, str(script)) as (conn, _proc):
adapter = PydanticAIACPAgent(Agent(TestModel(custom_output_text='hi')), models=['test', 'openai:gpt-4o'])
async with wire_agent(adapter) as (conn, _client):
await conn.initialize(protocol_version=acp.PROTOCOL_VERSION)
session = await conn.new_session(cwd='.', mcp_servers=[])
model_result = await conn.set_config_option(
@@ -2238,14 +2257,32 @@ class TestEntryPoints:
assert isinstance(option, schema.SessionConfigOptionSelect)
assert option.current_value == 'openai:gpt-4o'
async def test_native_toolsets_route_over_stdio(self) -> None:
"""Client-backed fs/terminal tools reach the real client over stdio when mounted per session."""
async def test_native_toolsets_route_over_the_wire(self) -> None:
"""Client-backed fs/terminal tools reach the real client over the wire when mounted per session."""
async def stream(
messages: list[ModelMessage], info: AgentInfo
) -> AsyncIterator[str | dict[int, DeltaToolCall]]:
if _has_tool_return(messages):
yield 'done'
return
yield {
0: DeltaToolCall(name='read_file', json_args=json.dumps({'path': 'notes.txt'})),
1: DeltaToolCall(name='run_command', json_args=json.dumps({'command': 'echo hi'})),
}
def session_config(session: AcpSession) -> AcpSessionConfig[None]:
capabilities = [
capability for capability in (acp_filesystem(session), acp_terminal(session)) if capability is not None
]
return AcpSessionConfig(deps=None, capabilities=capabilities)
adapter = PydanticAIACPAgent(Agent(FunctionModel(stream_function=stream)), session_config=session_config)
client = RecordingClient(files={'notes.txt': 'hello'}, output='hi')
capabilities = schema.ClientCapabilities(
fs=schema.FileSystemCapabilities(read_text_file=True, write_text_file=True), terminal=True
)
script = Path(__file__).parent / '_demo_native_agent.py'
async with acp.spawn_agent_process(client, sys.executable, str(script)) as (conn, _proc):
async with wire_agent(adapter, client) as (conn, _client):
await conn.initialize(protocol_version=acp.PROTOCOL_VERSION, client_capabilities=capabilities)
session = await conn.new_session(cwd='.', mcp_servers=[])
await conn.prompt(session_id=session.session_id, prompt=[acp.text_block('go')], message_id=uuid4().hex)
tests/harness/filesystem/test_search.py
+6 -3
View File
@@ -227,7 +227,10 @@ async def test_posix_oversized_line_keeps_later_matches(tmp_path: Path, no_rg_gi
assert await tools.grep('needle', workspace=backend) == 'b.txt:1:needle\n[... truncated at 1000 lines]'
async def test_posix_output_cap_reports_truncation(tmp_path: Path) -> None:
async def test_posix_output_cap_reports_truncation(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
# A smaller cap stands in for `_MAX_OUTPUT_BYTES`; the file still overflows it and the pipe buffer
# many times over, so `head` closes the pipe while the search is still writing.
monkeypatch.setattr('pydantic_ai_harness.filesystem._command_search._MAX_OUTPUT_BYTES', 50_000)
bin_dir, scratch, root = tmp_path / 'bin', tmp_path / 'scratch', tmp_path / 'ws'
path = tools_path(bin_dir, exclude=frozenset({'rg', 'mktemp'}))
# Put the search's temp files in a known directory: BSD mktemp ignores TMPDIR here.
@@ -235,9 +238,9 @@ async def test_posix_output_cap_reports_truncation(tmp_path: Path) -> None:
(bin_dir / 'mktemp').chmod(0o755)
scratch.mkdir()
root.mkdir()
(root / 'many.txt').write_text(('needle ' + 'X' * 100 + '\n') * 85000)
(root / 'many.txt').write_text(('needle ' + 'X' * 100 + '\n') * 10_000)
backend = CountingBackend(root, path)
tools = FileSystem[None](root_dir=root, max_search_results=200000, tools=['grep']).get_toolset()
tools = FileSystem[None](root_dir=root, max_search_results=20_000, tools=['grep']).get_toolset()
assert isinstance(tools, FileSystemToolset)
result = await tools.grep('needle', workspace=backend)
assert result.startswith('many.txt:1:needle')
tests/harness/gh_aw/test_engine_definition.py
+8
View File
@@ -498,6 +498,7 @@ def test_the_frame_key_is_sent_out_of_band_and_removed_from_the_child_environmen
assert invocation.stdin == f'{invocation.frame_key}\n'
@pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
@requires_clai2
def test_an_imported_agent_does_not_pass_the_frame_key_to_its_children(tmp_path: Path) -> None:
invocation = launch(
@@ -808,6 +809,7 @@ async def test_a_failed_stream_emits_one_partial_message_and_failure_result(
assert all(event.id != 'forged' for event in parsed.log_entries)
@pytest.mark.subprocess(reason='connects to `tests/mcp_server.py` over stdio')
async def test_imported_tools_and_gateway_mcp_tools_both_execute(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
@@ -1059,6 +1061,7 @@ async def test_interrupted_stream_has_a_failure_terminal_result(
assert events[-1].data['status'] == 'failure'
@pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
@pytest.mark.parametrize(
('cli_fails', 'expected_exit_code', 'expected_status'),
[(False, 0, 'success'), (True, 2, 'failure')],
@@ -1165,6 +1168,7 @@ class TestLauncherProgram:
check=False,
)
@pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
@requires_clai2
def test_the_agent_module_is_imported_once_and_not_from_the_checkout(self, tmp_path: Path) -> None:
completed = self.run(tmp_path, 'gh_aw_agent:agent', '-a', 'gh_aw_agent:agent', '-m', 'test', '-p', 'hello')
@@ -1176,6 +1180,7 @@ class TestLauncherProgram:
assert (tmp_path / 'imports.txt').read_text(encoding='utf-8') == 'module-directory\n'
assert re.search(rf'\x1eGH-AW-SESSION/{FRAME_KEY} \{{"type"\s*:\s*"session\.result"', completed.stdout)
@pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
def test_a_target_that_is_not_an_agent_names_what_it_found(self, tmp_path: Path) -> None:
(tmp_path / 'module').mkdir(parents=True, exist_ok=True)
(tmp_path / 'module' / 'not_an_agent.py').write_text('agent = 1\n', encoding='utf-8')
@@ -1185,6 +1190,7 @@ class TestLauncherProgram:
assert completed.returncode != 0
assert 'TypeError: not_an_agent:agent is int, not pydantic_ai.Agent' in completed.stderr
@pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
def test_an_agent_that_raises_on_import_fails_with_its_traceback(self, tmp_path: Path) -> None:
(tmp_path / 'module').mkdir(parents=True, exist_ok=True)
(tmp_path / 'module' / 'broken_agent.py').write_text(
@@ -1199,6 +1205,7 @@ class TestLauncherProgram:
# The message `pai` prints instead of a traceback when its own load fails.
assert 'Could not load agent' not in completed.stderr + completed.stdout
@pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
@requires_clai2
def test_cli_argument_error_finishes_the_generated_recorder_without_run_events(self, tmp_path: Path) -> None:
invocation = launch(tmp_path, proxy_env('openai', 'openai/gpt-5'), extra_python_path=CLAI2_SOURCE)
@@ -1231,6 +1238,7 @@ class TestLauncherProgram:
assert event.type == 'session.result'
assert event.data == {'status': 'failure', 'sourceType': 'pydantic-ai'}
@pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
def test_a_custom_agent_import_failure_keeps_one_full_traceback(self, tmp_path: Path) -> None:
workspace = tmp_path / 'workspace'
workspace.mkdir()
tests/harness/gh_aw/test_logfire_isolation.py
+11 -2
View File
@@ -11,6 +11,8 @@ import pytest
from .test_engine_definition import CLAI2_SOURCE, launch, proxy_env
pytestmark = pytest.mark.subprocess(reason='runs the generated gh-aw launcher program as a real script')
# Configure is real; only outbound HTTP is replaced. The CLI boundary records the
# resulting public configuration instead of making a model request.
PROBE = """import atexit
@@ -68,8 +70,15 @@ exec(compile(sys.argv.pop(1), '<launcher>', 'exec'))
"""
@pytest.mark.parametrize('token', ['', 'test-maintainer-token'])
@pytest.mark.parametrize('exit_code', [0, 7])
# Each case is a fresh interpreter importing `pydantic_ai` and Logfire, so the two
# axes are paired rather than crossed: the token decides whether Logfire starts its
# exporter and token-validation thread, the exit code only which `SystemExit` the
# launcher re-raises, and neither changes how the other reaches shutdown.
@pytest.mark.parametrize(
('exit_code', 'token'),
[(0, 'test-maintainer-token'), (7, '')],
ids=['token-success', 'no-token-failure'],
)
def test_launcher_ignores_checkout_logfire_configuration(tmp_path: Path, token: str, exit_code: int) -> None:
workspace = tmp_path / 'workspace'
workspace.mkdir()
tests/harness/localstack/_http_server.py
+2 -1
View File
@@ -48,7 +48,8 @@ class _RecordingHttpServer(ThreadingHTTPServer):
@contextmanager
def http_server(responses: Sequence[HttpResponse], *, port: int | None = None) -> Generator[RunningHttpServer]:
server = _RecordingHttpServer(responses, port)
thread = Thread(target=server.serve_forever, daemon=True)
# A short poll interval: `shutdown()` blocks until `serve_forever` next polls (0.5s by default).
thread = Thread(target=server.serve_forever, kwargs={'poll_interval': 0.01}, daemon=True)
thread.start()
bound_port = server.server_port
try:
tests/harness/logfire_mcp/test_capability.py
+1
View File
@@ -228,6 +228,7 @@ class TestLogfireMCP:
class TestPerRunAuth:
@pytest.mark.subprocess(reason='uses the `whoami_url` MCP server, which runs in its own process')
async def test_concurrent_runs_use_their_own_credentials(self, whoami_url: str) -> None:
agent = Agent(TestModel(), deps_type=str, capabilities=[LogfireMCP[str](url=whoami_url, auth=per_user_token)])
alice, bob = await asyncio.gather(
tests/harness/logfire_mcp/test_temporal.py
+1
View File
@@ -141,6 +141,7 @@ async def test_current_time_is_read_in_an_activity(client: Client) -> None:
assert 'The current UTC time is within the hour starting `' in output
@pytest.mark.subprocess(reason='uses the `whoami_url` MCP server, which runs in its own process')
async def test_auth_function_runs_under_temporal(client: Client, whoami_url: str) -> None:
per_user.url = whoami_url
async with Worker(
tests/harness/media/test_mongo.py
+4 -2
View File
@@ -121,7 +121,7 @@ class TestMongoMediaStoreRoundTrip:
"""
client = _mock_client()
store = MongoMediaStore(client=client, database='t', chunk_size_bytes=4)
data = bytes(range(256)) * 40 # 10_240 bytes -> 2560 chunks of 4
data = bytes(range(42)) # 42 bytes -> 10 full chunks of 4 plus a 2-byte tail
uri = await store.put(data)
digest = parse_media_uri(uri)
expected_chunks = (len(data) + 3) // 4
@@ -175,8 +175,10 @@ class TestMongoMediaStoreRoundTrip:
async def test_multi_byte_text_survives_chunk_boundaries(self) -> None:
"""Chunking splits raw bytes, so a multi-byte character straddles two chunks."""
store = MongoMediaStore(client=_mock_client(), database='t', chunk_size_bytes=7)
text = 'héllo wörld 日本語 🎉 ' * 500
text = 'héllo wörld 日本語 🎉 ' * 2
data = text.encode('utf-8')
# At least one chunk boundary lands on a UTF-8 continuation byte, mid-character.
assert any(data[i] & 0xC0 == 0x80 for i in range(7, len(data), 7))
uri = await store.put(data, context=MediaContext(media_type='text/plain'))
fetched = await store.get(uri)
assert fetched == data
tests/harness/shell/test_persistent.py
+6 -3
View File
@@ -248,7 +248,8 @@ class TestShellTool:
results.append(
await shell(
tmp_path,
{'command': 'printf ready; sleep 30', 'timeout': 3},
# The call waits out its timeout once nobody is left to publish a status.
{'command': 'printf ready; sleep 30', 'timeout': 1.5},
capabilities=[Running(), recorder],
)
)
@@ -258,7 +259,9 @@ class TestShellTool:
with anyio.fail_after(10):
await running.wait()
# The supervisor survives SIGTERM to publish an exit code, so only SIGKILL leaves the
# command running with nobody left to publish it.
# command running with nobody left to publish it. The call must still be waiting, or
# this would only test an ordinary timeout.
assert not results
os.kill(supervisors[0], signal.SIGKILL)
output = results[0]
@@ -296,7 +299,7 @@ class TestShellTool:
recorder = Recorder()
output = await shell(
tmp_path,
{'command': 'rm -rf .pydantic-ai-harness/shell/*; sleep 0.3', 'timeout': 1},
{'command': 'rm -rf .pydantic-ai-harness/shell/*; sleep 0.3', 'timeout': 0.5},
capabilities=[recorder],
)
assert output.startswith('PID: ')
tests/harness/shell/test_shell.py
+17 -11
View File
@@ -146,6 +146,14 @@ def _parse_command_id(result: str) -> str:
return result.split('ID: ')[1].strip()
async def _check_until(ts: ShellToolset[None], working_dir: Path, command_id: str, expected: str) -> str:
"""Poll `check_command` until its result contains `expected`, rather than sleeping a fixed time."""
with anyio.fail_after(10):
while expected not in (result := await ts.check_command(_ctx(working_dir), command_id)):
await anyio.sleep(0.01)
return result
async def _job(ts: ShellToolset[None], ctx: RunContext[None], command_id: str) -> Job:
job = await ts._job(ctx, command_id) # pyright: ignore[reportPrivateUsage]
assert job is not None
@@ -1161,7 +1169,7 @@ class TestBackgroundCommands:
start_result = await ts.start_command(_ctx(shell_dir), 'echo hello_bg')
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '[status: finished]')
stop_result = await ts.stop_command(_ctx(shell_dir), command_id)
assert 'stopped' in stop_result
@@ -1191,7 +1199,7 @@ class TestBackgroundCommands:
ts = _shell_toolset(shell_dir, max_output_chars=200)
start_result = await ts.start_command(_ctx(shell_dir), "printf '%0400d' 0; sleep 30")
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '0' * 400)
try:
check_result = await _call_shell_tool(ts, shell_dir, 'check_command', command_id=command_id)
@@ -1239,9 +1247,7 @@ class TestBackgroundCommands:
start_result = await ts.start_command(_ctx(shell_dir), 'echo done_quick')
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
check_result = await ts.check_command(_ctx(shell_dir), command_id)
check_result = await _check_until(ts, shell_dir, command_id, '[status: finished]')
assert 'finished' in check_result
assert 'done_quick' in check_result
assert check_result.splitlines()[-2:] == ['[status: finished]', '[exit code: 0]']
@@ -1274,7 +1280,7 @@ class TestBackgroundCommands:
start_result = await ts.start_command(_ctx(shell_dir), 'echo err_bg >&2')
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '[status: finished]')
stop_result = await ts.stop_command(_ctx(shell_dir), command_id)
assert 'err_bg' in stop_result
@@ -1292,7 +1298,7 @@ class TestBackgroundCommands:
start_result = await ts.start_command(_ctx(shell_dir), 'true')
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '[status: finished]')
stop_result = await ts.stop_command(_ctx(shell_dir), command_id)
assert '(no output)' in stop_result
@@ -1328,7 +1334,7 @@ class TestBackgroundCommands:
start_result = await ts.start_command(_ctx(shell_dir), 'echo err_check >&2')
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '[status: finished]')
check_result = await ts.check_command(_ctx(shell_dir), command_id)
assert '[stderr]' in check_result
@@ -1349,7 +1355,7 @@ class TestBackgroundCommands:
start_result = await ts.start_command(_ctx(shell_dir), 'pwd')
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '[status: finished]')
stop_result = await ts.stop_command(_ctx(shell_dir), command_id)
assert str(shell_dir) in stop_result
@@ -1368,7 +1374,7 @@ class TestBackgroundCommands:
start_result = await ts.start_command(_ctx(shell_dir), 'true')
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '[status: finished]')
await ts.stop_command(_ctx(shell_dir), command_id)
@@ -2141,7 +2147,7 @@ class TestEnvControlExecution:
ts = _env_toolset(shell_dir, env={'BG_TOKEN': 'bg-present'})
start_result = await ts.start_command(_ctx(shell_dir), _read_env_var('BG_TOKEN'))
command_id = _parse_command_id(start_result)
await anyio.sleep(0.5)
await _check_until(ts, shell_dir, command_id, '[status: finished]')
stop_result = await ts.stop_command(_ctx(shell_dir), command_id)
assert 'bg-present' in stop_result
tests/harness/spend/test_spend.py
+6 -3
View File
@@ -1516,13 +1516,16 @@ class TestRedisStore:
assert (await store.get_many(['k']))['k'] == added
async def test_repeated_adds_do_not_drift(self):
"""A price with a fractional sub-unit, since a whole one cannot detect rounding at all."""
"""A price with a fractional sub-unit, since a whole one cannot detect rounding at all.
Float accumulation of this price already drifts by the 7th add, so 1000 is ample.
"""
store = RedisSpendStore(FakeRedis())
price = Decimal('0.000000675') # a cheap model's real per-request cost
for _ in range(100_000):
for _ in range(1000):
await store.add_many([SpendEntry(key='k', usd=price, requests=1)])
assert (await store.get_many(['k']))['k'].usd == price * 100_000
assert (await store.get_many(['k']))['k'].usd == price * 1000
async def test_a_ttl_is_applied_and_the_key_is_namespaced(self):
client = FakeRedis()
tests/harness/ssh_workspace/test_ssh_workspace.py
+10 -6
View File
@@ -21,7 +21,7 @@ from pydantic_ai.workspaces import (
WorkspaceTimeoutError,
WorkspaceUnavailableError,
)
from pydantic_ai_harness.ssh_workspace import SSHWorkspace, SSHWorkspaceBackend
from pydantic_ai_harness.ssh_workspace import SSHWorkspace, SSHWorkspaceBackend, _backend as ssh_backend
from pydantic_ai_harness.ssh_workspace._backend import _STOP # pyright: ignore[reportPrivateUsage]
from .._fake_remote_tools import FakeRemoteTools, install_fake_remote_tools
@@ -98,7 +98,7 @@ async def test_timeouts_and_output_limits_keep_only_the_commands_stderr(tools: F
await backend.working_dir()
with pytest.raises(WorkspaceTimeoutError) as timeout:
await backend.run('printf partial >&2; sleep 30', shell=True, timeout=1)
await backend.run('printf partial >&2; sleep 30', shell=True, timeout=0.5)
with pytest.raises(WorkspaceOutputLimitError, match='SSH workspace output exceeded') as limit:
await backend.run('printf big >&2; head -c 11000000 /dev/zero', shell=True)
@@ -173,14 +173,18 @@ async def test_a_login_banner_stays_out_of_the_output(tools: FakeRemoteTools) ->
assert (result.stdout, result.stderr) == ('out', 'err')
async def test_resolving_the_working_dir_counts_against_the_first_timeout(tools: FakeRemoteTools) -> None:
# The stop after the timeout reaches a host that stalls too, so it gives up after its 2 second grace period.
async def test_resolving_the_working_dir_counts_against_the_first_timeout(
tools: FakeRemoteTools, monkeypatch: pytest.MonkeyPatch
) -> None:
# The stop after the timeout reaches a host that stalls too (for 30 seconds), so it gives up
# after its grace period, shortened here from 2 seconds.
monkeypatch.setattr(ssh_backend, '_STOP_TIMEOUT', 0.5)
started = anyio.current_time()
with pytest.raises(WorkspaceTimeoutError):
await SSHWorkspaceBackend('slow').run(['true'], timeout=1)
await SSHWorkspaceBackend('slow').run(['true'], timeout=0.5)
assert anyio.current_time() - started < 4.5
assert anyio.current_time() - started < 2.5
async def test_stderr_from_a_background_child_after_the_command_is_kept(tools: FakeRemoteTools) -> None:
tests/harness/step_persistence/test_frontier.py
+1
View File
@@ -82,6 +82,7 @@ async def test_same_length_after_run_rewrite_gets_final_checkpoint(tmp_path: Pat
assert snapshot.messages == result.all_messages()
@pytest.mark.subprocess(reason='hard-kills a process mid-run to check what it leaves persisted')
async def test_hard_kill_leaves_model_frontier_and_unknown_effects(tmp_path: Path) -> None:
database = tmp_path / 'steps.db'
script = """
tests/harness/step_persistence/test_snapshot_retention.py
+7 -5
View File
@@ -354,12 +354,14 @@ class TestSqliteStorePruneEdgeCases:
store = SqliteStepStore(database=db, media_store=None, max_snapshots_per_run=keep)
await _save(store, 'r1', 0) # creates the schema and the first row
conn = sqlite3.connect(db, check_same_thread=False, isolation_level=None)
# One transaction: in autocommit mode every row is its own WAL commit, ~10s for the seed.
conn = sqlite3.connect(db, check_same_thread=False)
try:
conn.executemany(
'INSERT INTO snapshots (run_id, step_index, timestamp, state, messages) VALUES (?, ?, ?, ?, ?)',
[('r1', step, '2026-01-01T00:00:00+00:00', 'complete', '[]') for step in range(1, keep + 2)],
)
with conn:
conn.executemany(
'INSERT INTO snapshots (run_id, step_index, timestamp, state, messages) VALUES (?, ?, ?, ?, ?)',
[('r1', step, '2026-01-01T00:00:00+00:00', 'complete', '[]') for step in range(1, keep + 2)],
)
finally:
conn.close()
tests/harness/system_reminders/test_system_reminders.py
+9 -4
View File
@@ -879,22 +879,27 @@ class TestLLMReminder:
await reminder(no_run_id)
assert len(caplog.records) == 4 # run-1 once, run-2 once, and each failure without a run id
async def test_remembers_a_bounded_number_of_runs(self, caplog: pytest.LogCaptureFixture) -> None:
async def test_remembers_a_bounded_number_of_runs(
self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
) -> None:
# A small bound stands in for `_WARNED_RUNS_KEPT` so the eviction is reached in a few runs.
monkeypatch.setattr('pydantic_ai_harness.system_reminders._capability._WARNED_RUNS_KEPT', 4)
def boom(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
raise RuntimeError('model down')
reminder = LLMReminder(model=FunctionModel(boom))
messages: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart('ship the fix')])]
with caplog.at_level(logging.WARNING, logger='pydantic_ai_harness.system_reminders'):
for index in range(258):
for index in range(6):
ctx = _ctx(messages=messages)
ctx.run_id = f'run-{index}'
await reminder(ctx)
oldest = _ctx(messages=messages)
oldest.run_id = 'run-0'
await reminder(oldest)
# 258 distinct runs, then the oldest again: it was evicted, so it warns a second time.
assert len(caplog.records) == 259
# 6 distinct runs, then the oldest again: it was evicted, so it warns a second time.
assert len(caplog.records) == 7
def test_zero_max_context_messages_raises(self) -> None:
with pytest.raises(ValueError, match='max_context_messages must be >= 1'):
tests/harness/test_coder.py
+7 -14
View File
@@ -1,5 +1,4 @@
import importlib
import subprocess
import sys
from pathlib import Path
from typing import NoReturn
@@ -79,19 +78,13 @@ async def test_coder_agent_imports_on_a_non_posix_host(monkeypatch: pytest.Monke
await agent.run('go', model=TestModel(call_tools=[]))
def test_coder_agent_export_is_lazy() -> None:
result = subprocess.run(
[
sys.executable,
'-c',
'import sys; import pydantic_ai_harness.coder; '
"assert 'pydantic_ai_harness.coder._agent' not in sys.modules",
],
check=False,
capture_output=True,
text=True,
)
assert result.returncode == 0, result.stderr
def test_coder_agent_export_is_lazy(monkeypatch: pytest.MonkeyPatch) -> None:
# Forget the whole harness package (restored after the test), so neither `pydantic_ai_harness`
# nor `pydantic_ai_harness.coder` can import the bundled agent without the test seeing it.
for name in [name for name in sys.modules if name.split('.')[0] == 'pydantic_ai_harness']:
monkeypatch.delitem(sys.modules, name)
importlib.import_module('pydantic_ai_harness.coder')
assert 'pydantic_ai_harness.coder._agent' not in sys.modules
def test_coder_unknown_export() -> None:
tests/harness/test_durable_execution_docs.py
+3 -8
View File
@@ -64,11 +64,6 @@ def test_dbos_example(workspace: Path, dbos_teardown: None) -> None:
_run('coder_dbos.py', workspace)
def test_prefect_example(workspace: Path) -> None:
pytest.importorskip('prefect')
from prefect.settings import PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED, temporary_settings
from prefect.testing.utilities import prefect_test_harness
with temporary_settings({PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED: False}):
with prefect_test_harness(server_startup_timeout=120):
_run('coder_prefect.py', workspace)
@pytest.mark.xdist_group(name='prefect')
def test_prefect_example(workspace: Path, prefect_test_server: None) -> None:
_run('coder_prefect.py', workspace)
tests/harness/test_durable_tool_io.py
+2 -1
View File
@@ -370,7 +370,8 @@ async def test_dbos_recovery_does_not_repeat_tool_requests(
# Re-execute the workflow function from after its last step, the way recovery does. Calling it
# again under the same ID is no replay: DBOS returns the stored result without running it.
handle = await DBOS.fork_workflow_async(workflow_id, len(steps))
assert await handle.get_result() == 'done'
# The default 1s result poll would add up to a second on top of the internal queue's own poll.
assert await handle.get_result(polling_interval_sec=0.01) == 'done'
assert count_requests() == first_run, 'recovery made a request again'
step_names = {step['function_name'] for step in steps}
tests/harness/test_import_lightness.py
+3
View File
@@ -2,7 +2,10 @@ import os
import subprocess
import sys
import pytest
@pytest.mark.subprocess(reason='asserts what a fresh interpreter imports')
def test_media_import_does_not_load_pydantic_ai() -> None:
"""Importing a light submodule must not pull in `pydantic_ai` via the root package (#9916).
tests/harness/test_workspace_durable_matrix.py
+3 -20
View File
@@ -2,11 +2,8 @@
from __future__ import annotations
import gc
import json
import re
import warnings
from collections.abc import Generator
from datetime import timedelta
from pathlib import Path
from uuid import uuid4
@@ -38,22 +35,6 @@ from tests.harness._temporal import skip_temporal_sandbox_on_314
from tests.temporal_utils import temporal_dev_server_cache_dir
@pytest.fixture
def prefect_server() -> Generator[None, None, None]:
# Prefect is an optional core extra; its absence must not hide the DBOS cells.
pytest.importorskip('prefect')
from prefect.settings import PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED, temporary_settings
from prefect.testing.utilities import prefect_test_harness
with temporary_settings({PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED: False}):
with prefect_test_harness(server_startup_timeout=120):
yield
# Prefect's test server leaves client sockets for GC on Python 3.14.
with warnings.catch_warnings():
warnings.filterwarnings('ignore', message='unclosed.*socket', category=ResourceWarning)
gc.collect()
def _agent(root: Path, engine: str, capability: str, vetoes: list[str]) -> Agent[None, str]:
scripts: dict[str, list[tuple[str, dict[str, str]]]] = {
'filesystem': [
@@ -118,8 +99,10 @@ def _agent(root: Path, engine: str, capability: str, vetoes: list[str]) -> Agent
return agent
# `prefect_test_server` skips without Prefect, an optional core extra, so the DBOS cells still run.
@pytest.mark.xdist_group(name='prefect')
@pytest.mark.parametrize('capability', ['coder', 'shell', 'filesystem'])
async def test_prefect_workspace_capabilities(tmp_path: Path, prefect_server: None, capability: str) -> None:
async def test_prefect_workspace_capabilities(tmp_path: Path, prefect_test_server: None, capability: str) -> None:
from prefect import flow
(tmp_path / 'dir').mkdir()
tests/harness/test_workspace_provider.py
+1
View File
@@ -144,6 +144,7 @@ async def test_stop_cleanup_finishes_before_stop_shielded_returns(anyio_backend:
assert events == ['stop cleanup', 'returned']
@pytest.mark.subprocess(reason='needs an interpreter where `sniffio` cannot be imported')
@pytest.mark.skipif(
any(req.startswith('sniffio') for req in importlib.metadata.requires('anyio') or []),
reason='AnyIO before 4.12 depends on and imports `sniffio`, so every install has it',
tests/harness/tool_call_judge/test_durable_tool_call_judge.py
+2 -1
View File
@@ -127,7 +127,8 @@ async def test_dbos_replays_the_recorded_verdict(dbos: DBOS) -> None:
# Re-execute the workflow function from its last step, the way recovery does. Calling it again
# under the same ID is no replay: from `dbos` 2.28 it returns the stored result without running it.
handle = await DBOS.fork_workflow_async(workflow_id, len(steps))
assert await handle.get_result() == 'done'
# The default 1s result poll would add up to a second on top of the internal queue's own poll.
assert await handle.get_result(polling_interval_sec=0.01) == 'done'
assert _judge_calls == 1, 'the replay asked the judging model again'
assert [v.verdict for v in _verdicts] == ['allow', 'allow'], '`on_verdict` fires on every replay'
tests/kw_only_walker.py
+3 -1
View File
@@ -45,7 +45,9 @@ def _takes_two_positional_arguments(cls: type) -> bool | None:
return positional >= 2
@pytest.fixture
# Module-scoped: walking imports every `pydantic_ai` module, provider SDKs included, which is shared setup rather than
# the test's own cost (see "Test cost" in `tests/AGENTS.md`).
@pytest.fixture(scope='module')
def public_dataclasses() -> dict[str, list[str]]:
coverage = Coverage.current()
# Coverage has no public API for pausing and resuming an active collector.
tests/models/anthropic/test_stream_parity.py
+5 -12
View File
@@ -8,17 +8,16 @@ complete `BetaMessage`. This replays each recorded stream through both.
from __future__ import annotations as _annotations
import dataclasses
from pathlib import Path
from typing import Any
import httpx2
import pytest
import yaml
from pydantic_ai import _utils
from pydantic_ai.messages import BaseToolCallPart, ModelResponse, ModelResponsePart
from pydantic_ai.models import ModelRequestParameters
from ...cassette_utils import recorded_interactions
from ...conftest import try_import
with try_import() as imports_successful:
@@ -31,19 +30,13 @@ with try_import() as imports_successful:
pytestmark = pytest.mark.skipif(not imports_successful(), reason='anthropic not installed')
_TESTS_DIR = Path(__file__).parents[2]
def _recorded_streams() -> dict[str, dict[str, Any]]:
streams: dict[str, dict[str, Any]] = {}
for path in sorted(_TESTS_DIR.rglob('*.yaml')):
text = path.read_text()
if 'event: message_start' not in text:
continue
for index, interaction in enumerate(yaml.safe_load(text)['interactions']):
content = interaction['response']['body'].get('content')
if isinstance(content, str) and 'event: message_start' in content:
streams[f'{path.relative_to(_TESTS_DIR)}#{index}'] = interaction
for name, interaction in recorded_interactions('event: message_start'):
content = interaction['response']['body'].get('content')
if isinstance(content, str) and 'event: message_start' in content:
streams[name] = interaction
return streams
tests/models/google/test_stream_parity.py
+4 -10
View File
@@ -22,17 +22,16 @@ from __future__ import annotations as _annotations
import dataclasses
import re
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any
import httpx2
import pytest
import yaml
from pydantic_ai.messages import BaseToolCallPart, ModelResponse, ModelResponsePart, NativeToolReturnPart
from pydantic_ai.models import ModelRequestParameters
from pydantic_ai.native_tools import AbstractNativeTool, CodeExecutionTool, FileSearchTool, WebFetchTool, WebSearchTool
from ...cassette_utils import recorded_interactions
from ...conftest import try_import
with try_import() as imports_successful:
@@ -43,19 +42,14 @@ with try_import() as imports_successful:
pytestmark = pytest.mark.skipif(not imports_successful(), reason='google-genai not installed')
_TESTS_DIR = Path(__file__).parents[2]
_MODEL_NAME_PATTERN = re.compile(r'/models/([^/:]+):streamGenerateContent')
def _recorded_streams() -> dict[str, dict[str, Any]]:
streams: dict[str, dict[str, Any]] = {}
for path in sorted(_TESTS_DIR.rglob('*.yaml')):
text = path.read_text()
if ':streamGenerateContent' not in text:
continue
for index, interaction in enumerate(yaml.safe_load(text)['interactions']):
if ':streamGenerateContent' in interaction['request']['uri']:
streams[f'{path.relative_to(_TESTS_DIR)}#{index}'] = interaction
for name, interaction in recorded_interactions(':streamGenerateContent'):
if ':streamGenerateContent' in interaction['request']['uri']:
streams[name] = interaction
return streams
tests/models/test_openai.py
+1
View File
@@ -6763,6 +6763,7 @@ async def test_openai_malformed_tool_args_degraded_on_the_wire(allow_model_reque
assert json.loads(assistant_message['tool_calls'][0]['function']['arguments']) == {INVALID_JSON_KEY: bad_args}
@pytest.mark.subprocess(reason='asserts what a fresh interpreter has imported after model construction')
def test_model_construction_preloads_lazy_dependencies():
"""Constructing a model resolves the deferred imports that stalled the first request's event loop (#7405).
tests/models/test_openai_chat_stream_parity.py
+12 -18
View File
@@ -10,16 +10,15 @@ from __future__ import annotations as _annotations
import dataclasses
from collections.abc import Callable
from pathlib import Path
from typing import Any, cast
import httpx2
import pytest
import yaml
from pydantic_ai.messages import BaseToolCallPart, ModelResponse, ModelResponsePart
from pydantic_ai.models import ModelRequestParameters
from ..cassette_utils import recorded_interactions
from ..conftest import try_import
with try_import() as imports_successful:
@@ -44,7 +43,6 @@ with try_import() as imports_successful:
pytestmark = pytest.mark.skipif(not imports_successful(), reason='openai not installed')
_TESTS_DIR = Path(__file__).parents[1]
# Providers with their own SDK or model class, whose streams another test covers.
_OTHER_MODEL_HOSTS = ('api.groq.com', 'api.mistral.ai', 'router.huggingface.co')
@@ -73,21 +71,17 @@ def _model(uri: str, model_name: str, client: AsyncOpenAI) -> OpenAIChatModel:
def _recorded_streams() -> dict[str, dict[str, Any]]:
streams: dict[str, dict[str, Any]] = {}
for path in sorted(_TESTS_DIR.rglob('*.yaml')):
text = path.read_text()
if '/chat/completions' not in text or 'data: {' not in text:
continue
for index, interaction in enumerate(yaml.safe_load(text)['interactions']):
uri: str = interaction['request']['uri']
body = interaction['response'].get('body', {}).get('string')
if (
uri.endswith('/chat/completions')
and not any(host in uri for host in _OTHER_MODEL_HOSTS)
and isinstance(body, str)
# A stream may start with an SSE comment, like OpenRouter's `: OPENROUTER PROCESSING`.
and body.lstrip().startswith(('data:', ':'))
):
streams[f'{path.relative_to(_TESTS_DIR)}#{index}'] = interaction
for name, interaction in recorded_interactions('/chat/completions', 'data: {'):
uri: str = interaction['request']['uri']
body = interaction['response'].get('body', {}).get('string')
if (
uri.endswith('/chat/completions')
and not any(host in uri for host in _OTHER_MODEL_HOSTS)
and isinstance(body, str)
# A stream may start with an SSE comment, like OpenRouter's `: OPENROUTER PROCESSING`.
and body.lstrip().startswith(('data:', ':'))
):
streams[name] = interaction
return streams
tests/models/test_openai_responses_stream_parity.py
+6 -13
View File
@@ -9,12 +9,10 @@ from __future__ import annotations as _annotations
import dataclasses
from collections.abc import Sequence
from pathlib import Path
from typing import Any
import httpx2
import pytest
import yaml
from pydantic_ai.messages import (
BaseToolCallPart,
@@ -29,6 +27,7 @@ from pydantic_ai.messages import (
from pydantic_ai.models import ModelRequestParameters
from pydantic_ai.tools import ToolDefinition
from ..cassette_utils import recorded_interactions
from ..conftest import try_import
with try_import() as imports_successful:
@@ -46,8 +45,6 @@ with try_import() as imports_successful:
pytestmark = pytest.mark.skipif(not imports_successful(), reason='openai not installed')
_TESTS_DIR = Path(__file__).parents[1]
def _response_body(interaction: dict[str, Any]) -> str | None:
body = interaction['response'].get('body', {})
@@ -57,15 +54,11 @@ def _response_body(interaction: dict[str, Any]) -> str | None:
def _recorded_streams() -> dict[str, dict[str, Any]]:
streams: dict[str, dict[str, Any]] = {}
for path in sorted(_TESTS_DIR.rglob('*.yaml')):
text = path.read_text()
if '"type":"response.' not in text:
continue
for index, interaction in enumerate(yaml.safe_load(text)['interactions']):
content = _response_body(interaction)
# Some providers send each event's `data:` line before its `event:` line.
if content is not None and content.startswith(('event:', 'data:')) and '"type":"response.' in content:
streams[f'{path.relative_to(_TESTS_DIR)}#{index}'] = interaction
for name, interaction in recorded_interactions('"type":"response.'):
content = _response_body(interaction)
# Some providers send each event's `data:` line before its `event:` line.
if content is not None and content.startswith(('event:', 'data:')) and '"type":"response.' in content:
streams[name] = interaction
return streams
tests/models/xai_proto_cassettes.py
+2 -1
View File
@@ -38,6 +38,7 @@ from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal, Protocol, TypeAlias, cast
from ..cassette_utils import load_cassette_yaml
from ..conftest import try_import
with try_import() as imports_successful:
@@ -204,7 +205,7 @@ class XaiProtoCassette:
@classmethod
def load(cls, path: Path) -> XaiProtoCassette:
data = yaml.safe_load(path.read_text(encoding='utf-8'))
data = load_cassette_yaml(path.read_text(encoding='utf-8'))
interactions: list[Interaction] = []
for item in data.get('interactions', []):
tests/providers/test_github_copilot_oauth.py
+16 -15
View File
@@ -110,8 +110,22 @@ async def test_github_rejects_unknown_client() -> None:
await flow.start()
@pytest.fixture
def delays(monkeypatch: pytest.MonkeyPatch) -> list[float]:
"""Advance polling without wall-clock sleeps, while retaining cancellation checkpoints."""
delays: list[float] = []
sleep = anyio.sleep
async def record_sleep(delay: float) -> None:
delays.append(delay)
await sleep(0)
monkeypatch.setattr(anyio, 'sleep', record_sleep)
return delays
@pytest.mark.vcr
async def test_github_pending_authorization_can_be_cancelled() -> None:
async def test_github_pending_authorization_can_be_cancelled(delays: list[float]) -> None:
pending = anyio.Event()
async def receive(response: httpx2.Response) -> None:
@@ -129,20 +143,7 @@ async def test_github_pending_authorization_can_be_cancelled() -> None:
await pending.wait()
group.cancel_scope.cancel()
assert not client.is_closed
@pytest.fixture
def delays(monkeypatch: pytest.MonkeyPatch) -> list[float]:
"""Advance polling without wall-clock sleeps, while retaining cancellation checkpoints."""
delays: list[float] = []
sleep = anyio.sleep
async def record_sleep(delay: float) -> None:
delays.append(delay)
await sleep(0)
monkeypatch.setattr(anyio, 'sleep', record_sleep)
return delays
assert delays[0] == 5 # GitHub's interval, waited before the pending poll
async def test_approval_and_provider_handoff(delays: list[float]) -> None:
tests/providers/test_voyageai.py
+27 -7
View File
@@ -1,21 +1,35 @@
from __future__ import annotations as _annotations
import importlib
from importlib.util import find_spec
import pytest
from pydantic_ai.exceptions import UserError
from ..conftest import TestEnv, try_import
from ..conftest import TestEnv
with try_import() as imports_successful:
pytestmark = [
pytest.mark.skipif(find_spec('voyageai') is None, reason='voyageai not installed'),
pytest.mark.usefixtures('voyageai_imported'),
]
@pytest.fixture(scope='module')
def voyageai_imported() -> None:
"""Pay for importing `voyageai` in shared setup, outside each test's time budget.
With `sentence-transformers` installed, `voyageai` imports it and `torch`, which takes seconds: importing it at
module level would make every worker collecting this module pay for it. See "Test cost" in `tests/AGENTS.md`.
"""
importlib.import_module('pydantic_ai.providers.voyageai')
def test_voyageai_provider() -> None:
from voyageai.client_async import AsyncClient
from pydantic_ai.providers.voyageai import VoyageAIProvider
pytestmark = pytest.mark.skipif(not imports_successful(), reason='voyageai not installed')
def test_voyageai_provider() -> None:
provider = VoyageAIProvider(api_key='api-key')
assert provider.name == 'voyageai'
assert provider.base_url == 'https://api.voyageai.com/v1'
@@ -23,12 +37,18 @@ def test_voyageai_provider() -> None:
def test_voyageai_provider_need_api_key(env: TestEnv) -> None:
from pydantic_ai.providers.voyageai import VoyageAIProvider
env.remove('VOYAGE_API_KEY')
with pytest.raises(UserError, match='VOYAGE_API_KEY'):
VoyageAIProvider()
def test_voyageai_provider_pass_voyageai_client() -> None:
from voyageai.client_async import AsyncClient
from pydantic_ai.providers.voyageai import VoyageAIProvider
voyageai_client = AsyncClient(api_key='test-api-key')
provider = VoyageAIProvider(voyageai_client=voyageai_client)
assert provider.client == voyageai_client
tests/realtime/simulation/_cassette_replay.py
+8 -4
View File
@@ -10,11 +10,11 @@ which a lifecycle stream's responses may answer.
from __future__ import annotations as _annotations
import json
import re
from collections.abc import AsyncIterator, Iterator
from pathlib import Path
from typing import Any, Literal
import yaml
from google.genai import _live_converters as live_converters, types as genai_types
from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK
from websockets.frames import Close
@@ -85,12 +85,16 @@ def websocket_cassettes() -> list[Path]:
return sorted(path for path in CASSETTES_DIR.glob('*/*.yaml') if _is_websocket(path))
# An HTTP recording's top-level interactions are `request`/`response` pairs; a WebSocket one's are
# frames and closes, whose payloads are indented, so no line of theirs can open like this.
_HTTP_INTERACTION = re.compile(r'^- request:', re.MULTILINE)
def _is_websocket(path: Path) -> bool:
# Both formats can open with `version:`, so tell them apart by their interactions: HTTP recordings
# (e.g. WebRTC signaling) hold `request`/`response` pairs, WebSocket ones hold frames and closes.
raw: dict[str, Any] = yaml.safe_load(path.read_text(encoding='utf-8'))
interactions: list[dict[str, Any]] = raw.get('interactions') or [{}]
return 'request' not in interactions[0]
# Matched on the text rather than parsed: this runs at collection, in every worker, over every cassette.
return _HTTP_INTERACTION.search(path.read_text(encoding='utf-8')) is None
def _segments(cassette: RealtimeCassette) -> Iterator[tuple[list[dict[str, Any]], CassetteClose | None, int]]:
tests/realtime/simulation/_findings.py
+15 -9
View File
@@ -174,20 +174,29 @@ LOST_RESPONSE_RESERVATION = Finding(
)
def _xai_push_to_talk(sim: Simulation) -> bool:
options = getattr(sim, 'openai', None)
return options is not None and options.dialect == 'xai' and options.turn_detection == 'manual'
def _receive_loop_send_failed(sim: Simulation, violation: InvariantViolation) -> bool:
network = getattr(getattr(sim, 'server', None), 'network', None)
# On xAI push-to-talk, the receive loop sends more once the `response.done` that ends a reply is handled (#9070):
# the audio held back behind the reply, and the clear a deferred request needs to be answered.
xai_frames = ('input_audio_buffer.append', 'input_audio_buffer.clear') if _xai_push_to_talk(sim) else ()
sent = ('response.create', *xai_frames)
return network is not None and any(
frame == 'response.create' and last_read in ('response.done', 'error')
for frame, last_read, _ in network.failed_sends
frame in sent and last_read in ('response.done', 'error') for frame, last_read, _ in network.failed_sends
)
RECEIVE_LOOP_SEND_FAILURE = Finding(
id='SIM-4',
title=(
'a deferred `response.create` that the connection sends while handling a `response.done` (or a refusal) fails '
"on a dying socket, and the whole frame is dropped: that response's usage and terminal never reach the "
'session, and the deferred request is neither re-asked nor released'
'a deferred `response.create` (or, on xAI push-to-talk, audio held back behind the reply, or the clear the '
'request needs) that the connection sends while handling a `response.done` (or a refusal) fails on a dying '
"socket, and the whole frame is dropped: that response's usage and terminal never reach the session, and the "
'deferred request is neither re-asked nor released'
),
tracked_by='an ordered outbox, so the receive loop never sends on the socket itself; found by this simulator',
evidence='simulated',
@@ -727,11 +736,8 @@ def _hand_commit_unanswered_at_the_end(sim: Simulation) -> bool:
def _push_to_talk_audio_after_a_repeated_terminal(sim: Simulation) -> bool:
options = getattr(sim, 'openai', None)
return (
options is not None
and options.dialect == 'xai'
and options.turn_detection == 'manual'
_xai_push_to_talk(sim)
and bool(sim.truth.repeated_terminals)
and any(operation.name == 'send_audio' for operation in sim.operations)
)
tests/realtime/simulation/test_simulation.py
+71
View File
@@ -67,6 +67,7 @@ def exploration_settings() -> settings:
if imports_successful()
else [],
)
@pytest.mark.slow(reason='explores 50 randomized sessions of up to 25 steps against each provider simulator')
def test_exploration(machine: type[RuleBasedStateMachine]) -> None:
KNOWN_HIT_COUNTS.clear()
run_state_machine_as_test(machine, settings=exploration_settings())
@@ -546,6 +547,61 @@ def test_known_failed_deferred_create_after_a_refusal_keeps_its_reservation() ->
reproduce('SIM-4', OpenAISimulation(), scenario)
@known('SIM-4')
def test_known_xai_held_audio_failing_after_a_reply_drops_its_terminal() -> None:
"""Audio sent during a reply is held back, and goes out from the receive loop once the reply ends."""
def scenario(sim: OpenAISimulation) -> None:
sim.create_response()
sim.fail_next_send()
sim.send_audio()
sim.settle()
reproduce(
'SIM-4',
OpenAISimulation(openai=OpenAIOptions(dialect='xai', turn_detection='manual', transcription=False)),
scenario,
)
@known('SIM-4')
def test_known_xai_held_audio_failing_before_a_deferred_request_loses_it() -> None:
"""The held audio goes out ahead of the request deferred behind the reply, which the failure takes with it."""
def scenario(sim: OpenAISimulation) -> None:
sim.send_text()
sim.create_response()
sim.fail_next_send()
sim.send_audio()
sim.settle()
reproduce(
'SIM-4',
OpenAISimulation(openai=OpenAIOptions(dialect='xai', turn_detection='manual', transcription=False)),
scenario,
)
@known('SIM-4')
def test_known_xai_clear_before_a_deferred_request_failing_loses_it() -> None:
"""The clear xAI needs before a deferred request goes out from the receive loop, once the reply ends."""
def scenario(sim: OpenAISimulation) -> None:
sim.create_response()
sim.clear_audio()
sim.create_response()
sim.fail_next_send()
sim.send_audio()
sim.clear_audio()
sim.settle()
reproduce(
'SIM-4',
OpenAISimulation(openai=OpenAIOptions(dialect='xai', turn_detection='manual', transcription=False)),
scenario,
)
def test_barge_in_keeps_a_request_vad_will_not_answer() -> None:
"""SIM-32, fixed by the session core."""
@@ -675,6 +731,21 @@ def test_clear_after_an_unread_vad_commit_keeps_the_turn() -> None:
run_clean(OpenAISimulation(openai=OpenAIOptions(transcription=False)), scenario)
def test_xai_clear_after_a_request_committed_speech_keeps_the_turn() -> None:
"""On xAI push-to-talk, the request a text turn sends commits the speech in the buffer before the clear (SIM-35)."""
def scenario(sim: OpenAISimulation) -> None:
sim.send_audio()
sim.settle()
sim.send_text()
sim.clear_audio()
sim.settle()
run_clean(
OpenAISimulation(openai=OpenAIOptions(dialect='xai', turn_detection='manual', transcription=False)), scenario
)
@known('SIM-36')
def test_known_xai_push_to_talk_audio_after_a_repeated_terminal_is_held() -> None:
def scenario(sim: OpenAISimulation) -> None:
tests/realtime/test_google.py
+1
View File
@@ -667,6 +667,7 @@ def test_ws_trace_context_does_not_duplicate_a_differently_cased_header() -> Non
assert headers == {'Traceparent': 'preset'}
@pytest.mark.usefixtures('young_gc')
def test_ws_connect_lock_is_per_event_loop() -> None:
# The lock is process-wide by intent (it guards a replacement of the `google.genai.live.ws_connect`
# module global), but an `anyio.Lock` binds to the loop it is first used on, so one shared instance
tests/realtime/test_google_ws.py
+3 -2
View File
@@ -938,7 +938,7 @@ async def test_handle_barge_in_over_live_speech(
async def test_extended_thinking_async_tool_round(
gemini_ws_cassette: tuple[Provider[Any], RealtimeCassette],
gemini_ws_cassette: tuple[Provider[Any], RealtimeCassette], realtime_recording: bool
) -> None:
"""`gemini-3.8-live-extended-thinking` speaks a filler, runs the tool in the background, then answers.
@@ -953,7 +953,8 @@ async def test_extended_thinking_async_tool_round(
@agent.tool_plain
async def search_flights(origin: str, destination: str) -> str:
"""Search flights between two cities. Takes several seconds."""
await anyio.sleep(5)
# Only the live model needs a slow tool to speak its filler; replay follows the recorded frame order.
await anyio.sleep(5 if realtime_recording else 0)
return 'KLM at 120 dollars'
events: list[Any] = []
tests/realtime/test_inference.py
+1
View File
@@ -69,6 +69,7 @@ def test_known_realtime_model_names() -> None: # pragma: lax no cover
assert generated_transcription_names == sorted(get_args(KnownRealtimeTranscriptionModelName.__value__))
@pytest.mark.subprocess(reason='asserts what a fresh interpreter imports')
def test_star_import_does_not_load_optional_providers() -> None:
code = """
import sys
tests/realtime/test_openai.py
+71 -3
View File
@@ -5375,7 +5375,7 @@ async def test_a_conversation_id_resolved_late_reaches_the_core_history() -> Non
@pytest.mark.anyio
async def test_audio_that_fails_to_go_out_is_taken_back_from_the_core_too() -> None:
async def test_audio_that_fails_to_go_out_never_reaches_the_core() -> None:
class _FailingAudio(_QueuedWebSocket):
async def send(self, data: str) -> None:
raise OSError('gone') # (the audio is all this session sends)
@@ -5397,8 +5397,8 @@ async def test_audio_that_fails_to_go_out_is_taken_back_from_the_core_too() -> N
@pytest.mark.anyio
async def test_a_failed_send_takes_back_its_own_audio_while_another_waits_to_go_out() -> None:
"""Two chunks sent at once go out one after the other: the first failing takes back its own audio, not the second's."""
async def test_a_failed_send_leaves_its_audio_out_while_another_waits_to_go_out() -> None:
"""Two chunks sent at once go out one after the other: the first failing never adds its audio, and the second's goes in."""
class _FirstAudioFails(_QueuedWebSocket):
def __init__(self) -> None:
@@ -5437,6 +5437,74 @@ async def test_a_failed_send_takes_back_its_own_audio_while_another_waits_to_go_
assert bytes(core._input_audio) == second # pyright: ignore[reportPrivateUsage]
@pytest.mark.anyio
async def test_audio_whose_send_fails_after_its_turn_is_recorded_stays_out_of_that_turn() -> None:
"""A turn recorded while a chunk is still going out doesn't keep that chunk when its send then fails.
The provider committed its buffer before it had the chunk, so the chunk is no part of that turn whether or not
its send succeeds: history and the streamed part both hold only the audio that went out ahead of the commit.
"""
class _SecondAudioFails(_QueuedWebSocket):
def __init__(self) -> None:
super().__init__()
self.sending = asyncio.Event()
self.fail = asyncio.Event()
self.sends = 0
async def send(self, data: str) -> None:
self.sends += 1
if self.sends == 2: # (the audio is all this session sends)
self.sending.set()
await self.fail.wait()
raise OSError('gone')
await super().send(data)
ws = _SecondAudioFails()
connection = OpenAIRealtimeConnection(ws) # type: ignore[arg-type]
session = RealtimeSession(
connection,
model=FakeRealtimeModel(connection, system='openai'),
tool_manager=make_tool_manager(),
audio_retention='input_audio',
)
sent, unsent = b'\x01\x00' * 100, b'\x02\x00' * 50
streamed: list[SpeechPart] = []
async with session:
await session.send_audio(sent)
failing_send = asyncio.create_task(session.send_audio(unsent))
await ws.sending.wait()
for frame in [
{'type': 'input_audio_buffer.speech_started', 'item_id': 'item_u1', 'audio_start_ms': 0},
{'type': 'input_audio_buffer.speech_stopped', 'item_id': 'item_u1', 'audio_end_ms': 500},
{'type': 'input_audio_buffer.committed', 'item_id': 'item_u1', 'previous_item_id': None},
{
'type': 'conversation.item.input_audio_transcription.completed',
'item_id': 'item_u1',
'content_index': 0,
'transcript': 'Hi.',
},
]:
ws.push(frame)
async for event in session: # pragma: no branch
if isinstance(event, PartEndEvent) and isinstance(event.part, SpeechPart):
streamed.append(event.part)
break
# The turn is recorded while the chunk is still going out; then its send fails.
assert session.all_messages()
ws.fail.set()
with pytest.raises(RealtimeError):
await failing_send
history = session.all_messages()
[request] = history
[part] = request.parts
assert isinstance(part, SpeechPart) and part.transcript == 'Hi.' and part.audio is not None
assert part.audio.data[44:] == sent
[streamed_part] = streamed
assert streamed_part.audio is not None and streamed_part.audio.data[44:] == sent
@pytest.mark.anyio
async def test_retained_audio_is_kept_once_when_the_core_keeps_history() -> None:
"""The session core's history keeps retained audio; the session's own record (what tools see) lets its copy go.
tests/realtime/test_openai_ws.py
+6 -2
View File
@@ -291,14 +291,18 @@ async def test_text_in_audio_out_turn(openai_ws_cassette: tuple[Provider[Any], R
assert len(part.audio.data) > 0
async def test_text_context_waits_for_next_turn(openai_ws_cassette: tuple[Provider[Any], RealtimeCassette]) -> None:
async def test_text_context_waits_for_next_turn(
openai_ws_cassette: tuple[Provider[Any], RealtimeCassette], realtime_recording: bool
) -> None:
provider, _ = openai_ws_cassette
model = OpenAIRealtimeModel('gpt-realtime', provider=provider)
agent = Agent(instructions='Answer in one short sentence.')
async with agent.realtime(model).session() as session:
await session.send('The visitor is called Ada.', respond=False)
await asyncio.sleep(1)
# Gives a reply the model shouldn't have started time to arrive. Replay delivers the recorded frames
# back to back, so it needs only a moment.
await asyncio.sleep(1 if realtime_recording else 0.1)
assert not [message for message in session.new_messages() if isinstance(message, ModelResponse)]
await session.send('What is the visitor called?')
with anyio.fail_after(30):
tests/realtime/test_session.py
+38 -1
View File
@@ -868,6 +868,7 @@ async def test_closing_an_unstarted_view_releases_its_tap() -> None:
assert session._audio_tap_drops == 0 # pyright: ignore[reportPrivateUsage]
@pytest.mark.usefixtures('young_gc')
async def test_abandoned_unstarted_view_releases_its_tap() -> None:
chunks = [bytes([index]) for index in range(40)]
session = RealtimeSession(FakeRealtimeConnection([AudioDelta(chunk) for chunk in chunks]))
@@ -5862,6 +5863,12 @@ async def test_text_request_reserved_before_response_finishes_during_send() -> N
async def test_send_audio_reserved_before_speech_boundary_during_send() -> None:
"""The turn audio opens is reserved before the chunk goes out, so a boundary handled meanwhile still records it.
The chunk itself isn't in that turn: the provider drew the boundary before it had the chunk, so the chunk
is retained for the next turn instead.
"""
class _BoundarySend(FakeRealtimeConnection):
async def send(self, content: RealtimeInput) -> None:
self.sent.append(content)
@@ -5871,8 +5878,9 @@ async def test_send_audio_reserved_before_speech_boundary_during_send() -> None:
session = RealtimeSession(conn, audio_retention='input_audio')
await session.send_audio(b'\xaa\xbb')
assert session.new_messages() == snapshot(
[ModelRequest(parts=[SpeechPart(speaker='user', audio=_wav_content(b'\xaa\xbb'))], timestamp=IsDatetime())]
[ModelRequest(parts=[SpeechPart(speaker='user')], timestamp=IsDatetime())]
)
assert bytes(session._input_audio) == b'\xaa\xbb' # pyright: ignore[reportPrivateUsage]
async def test_failed_sends_leave_no_phantom_history_or_audio() -> None:
@@ -5912,6 +5920,35 @@ async def test_failed_sends_leave_no_phantom_history_or_audio() -> None:
assert unretained.new_messages() == []
async def test_audio_whose_send_fails_after_its_turn_is_recorded_stays_out_of_that_turn() -> None:
"""A turn recorded while a chunk is still going out doesn't keep that chunk when its send then fails."""
class _TurnEndsDuringFailingSend(FakeRealtimeConnection):
fail = False
async def send(self, content: RealtimeInput) -> None:
if self.fail:
session._translate_event(InputTranscript(text='Hi.', is_final=True)) # pyright: ignore[reportPrivateUsage]
raise RuntimeError('send failed')
self.sent.append(content)
conn = _TurnEndsDuringFailingSend([])
session = RealtimeSession(conn, audio_retention='input_audio')
await session.send_audio(b'\xaa\xbb')
conn.fail = True
with pytest.raises(RuntimeError, match='send failed'):
await session.send_audio(b'\xcc\xdd')
assert session.new_messages() == snapshot(
[
ModelRequest(
parts=[SpeechPart(speaker='user', transcript='Hi.', audio=_wav_content(b'\xaa\xbb'))],
timestamp=IsDatetime(),
)
]
)
assert not session._input_audio # pyright: ignore[reportPrivateUsage]
async def test_transport_failure_while_sending_becomes_a_realtime_error() -> None:
# A send that fails because the link is gone is the same failure the receive side reports, so it
# surfaces as `RealtimeError` (a `ModelAPIError`) instead of leaking the transport's own exception.
tests/realtime/test_session_core.py
+13 -61
View File
@@ -33,7 +33,6 @@ from pydantic_ai.messages import (
from pydantic_ai.realtime._core import (
AudioCleared,
AudioSent,
AudioUnsent,
Closed,
CoreInput,
ExchangeAbandoned,
@@ -894,23 +893,6 @@ def test_a_conversation_id_resolved_late_reaches_what_the_core_recorded() -> Non
assert {message.conversation_id for message in session_core.all_messages()} == {'c1'}
def test_input_audio_that_never_went_out_is_taken_back() -> None:
session_core = feed(
core(input_transcription_enabled=False, retain_input_audio=True),
AudioSent(data=b'\x01\x00'),
AudioSent(data=b'\x02\x00'),
AudioUnsent(),
# Only the latest chunk is taken back, once.
AudioUnsent(),
UserTurnStarted(turn_id='u1'),
UserTurnEnded(turn_id='u1'),
)
[request] = session_core.all_messages()
[part] = request.parts
assert isinstance(part, SpeechPart) and part.audio is not None
assert part.audio.data[44:] == b'\x01\x00'
def test_a_reconnect_replays_a_held_turns_audio_without_counting_it() -> None:
"""The replay builds the held turn's request with its audio, which isn't retained history, so isn't counted."""
session_core = feed(
@@ -951,25 +933,6 @@ def test_retained_audio_eviction_passes_over_what_is_not_recorded() -> None:
assert list(session_core._recorded_audio) == [session_core._turns['u5']] # pyright: ignore[reportPrivateUsage]
def test_input_audio_cut_for_a_turn_before_its_send_failed_is_taken_back_from_the_turn() -> None:
session_core = feed(
core(retain_input_audio=True),
AudioSent(data=b'\x01\x00'),
AudioSent(data=b'\x02\x00'),
RealtimeInputSpeechStartEvent(item_id='u1'),
UserTurnStarted(turn_id='u1'),
# The speech end cuts the turn's audio while the last chunk is still on its way, and its send then fails.
RealtimeInputSpeechEndEvent(item_id='u1'),
AudioUnsent(),
UserTurnEnded(turn_id='u1'),
InputTranscript('Hi.', item_id='u1', is_final=True),
)
[request] = session_core.all_messages()
[part] = request.parts
assert isinstance(part, SpeechPart) and part.audio is not None
assert part.audio.data[44:] == b'\x01\x00'
def test_retained_audio_trims_the_oldest_audio_coming_in_first() -> None:
"""A response whose incoming audio was emptied counts as newest when more comes in for it."""
session_core = feed(
@@ -1007,38 +970,27 @@ def test_retained_audio_evicted_is_let_go_by_the_core_too() -> None:
assert len(session_core._recorded_audio) == 2 # pyright: ignore[reportPrivateUsage]
def test_a_failed_send_of_audio_trimmed_by_the_budget_leaves_nothing_of_it() -> None:
"""A chunk longer than the whole budget is trimmed before its send fails: what is left of it is taken back."""
session_core = feed(
core(input_transcription_enabled=False, retain_input_audio=True, retain_audio_max_seconds=0.1),
AudioSent(data=_tenth_of_a_second(1) * 2),
AudioUnsent(),
AudioSent(data=b'\x05\x00'),
UserTurnStarted(turn_id='u1'),
UserTurnEnded(turn_id='u1'),
)
[request] = session_core.all_messages()
[part] = request.parts
assert isinstance(part, SpeechPart) and part.audio is not None
assert part.audio.data[44:] == b'\x05\x00'
def test_input_audio_that_went_out_after_its_turn_was_recorded_is_left_to_the_next() -> None:
"""Audio is reported once its send completed, so a chunk still on its way when a turn is recorded comes after it.
def test_audio_taken_back_from_a_turn_no_longer_waiting_leaves_it_alone() -> None:
"""Once its turn is recorded, the core doesn't change the message to take the audio back.
A known gap, not the intent (the current session core has it too): a turn recorded while the send of its last
chunk is still under way keeps that chunk's audio if the send then fails.
The provider ended that turn without the chunk, which is the next turn's start, or, if its send fails, no
turn's: the recorded turn keeps only the audio that went out ahead of it either way.
"""
session_core = feed(
core(input_transcription_enabled=False, retain_input_audio=True),
AudioSent(data=b'\x01\x00'),
UserTurnStarted(turn_id='u1'),
UserTurnEnded(turn_id='u1'),
AudioUnsent(),
AudioSent(data=b'\x02\x00'),
UserTurnStarted(turn_id='u2'),
UserTurnEnded(turn_id='u2'),
)
[request] = session_core.all_messages()
[part] = request.parts
assert isinstance(part, SpeechPart) and part.audio is not None
assert [
part.audio.data[44:]
for message in session_core.all_messages()
for part in message.parts
if isinstance(part, SpeechPart) and part.audio is not None
] == [b'\x01\x00', b'\x02\x00']
def test_retained_audio_eviction_does_not_walk_what_waits_or_has_no_audio(monkeypatch: pytest.MonkeyPatch) -> None:
tests/realtime/test_ws_cassettes.py
+36
View File
@@ -18,6 +18,7 @@ from .ws_cassettes import (
CassettePlan,
ProviderName,
RealtimeCassette,
RealtimeCassetteInteraction,
RecordingWebSocket,
ReplayWebSocket,
patched_ws_connect,
@@ -461,6 +462,41 @@ async def test_replay_waits_for_a_direct_recv_reader() -> None:
await send
async def test_replay_session_close_overtakes_unread_frames(monkeypatch: pytest.MonkeyPatch) -> None:
"""Unit test: GPT-Live's `session.close` doesn't wait for a reader to take the frames recorded before it.
The session sends it once it has stopped reading, and only reads the rest (up to `session.closed`)
after the send returns, so waiting for a reader would stall until the grace period ran out.
"""
monkeypatch.setattr(ws_cassettes, '_REPLAY_PROGRESS_GRACE', 0.01)
close = {'type': 'session.close'}
cassette = RealtimeCassette(
interactions=[
CassetteMessage(direction='received', data={'type': 'server.event'}),
CassetteMessage(direction='sent', data=close),
CassetteMessage(direction='received', data={'type': 'session.closed'}),
]
)
replay = ReplayWebSocket(cassette)
await replay.send(json.dumps(close))
assert [json.loads(await replay.recv()) for _ in range(2)] == [{'type': 'server.event'}, {'type': 'session.closed'}]
# The recording itself is left as it was recorded.
assert cassette.interactions[1] == CassetteMessage(direction='sent', data=close)
# It never overtakes a close, nor turns into a send the recording doesn't have.
stalled: list[list[RealtimeCassetteInteraction]] = [
[CassetteMessage(direction='received', data={'type': 'server.event'})],
[
CassetteMessage(direction='received', data={'type': 'server.event'}),
CassetteClose(code=1000, reason='', ok=True),
CassetteMessage(direction='sent', data=close),
],
]
for interactions in stalled:
with pytest.raises(AssertionError, match='no matching recorded send'):
await ReplayWebSocket(RealtimeCassette(interactions=interactions)).send(json.dumps(close))
@pytest.mark.parametrize(
('frame', 'is_audio'),
[
tests/realtime/ws_cassettes.py
+23 -2
View File
@@ -208,7 +208,10 @@ class RealtimeCassette:
@classmethod
def load(cls, path: Path) -> RealtimeCassette:
raw = cast('dict[str, Any]', yaml.safe_load(path.read_text(encoding='utf-8')))
# libyaml's loader where PyYAML was built with it: the pure-Python one takes a noticeable share of a
# second over a cassette's thousands of frames.
loader = getattr(yaml, 'CSafeLoader', yaml.SafeLoader)
raw = cast('dict[str, Any]', yaml.load(path.read_text(encoding='utf-8'), Loader=loader))
interactions: list[RealtimeCassetteInteraction] = []
for item in cast('list[dict[str, Any]]', raw.get('interactions', [])):
at = item.get('at')
@@ -354,6 +357,12 @@ class _SentFrameNormalizer:
# making progress never comes near it.
_REPLAY_PROGRESS_GRACE = 2.0
# Outbound frames a session sends once it has stopped reading, to end the provider session: GPT-Live's
# `session.close`. Whatever the provider sent before it is read afterwards, by the same code that waits
# for the provider to confirm the close, so on replay the close overtakes those frames instead of
# waiting for a reader that only starts once the send returns.
_SESSION_END_TYPES = frozenset({'session.close'})
class ReplayWebSocket:
"""Replay a recorded WebSocket conversation, validating outbound frames as they are sent.
@@ -364,7 +373,8 @@ class ReplayWebSocket:
"""
def __init__(self, cassette: RealtimeCassette, *, hold_open: bool = False) -> None:
self._interactions = cassette.interactions
# A copy: a session-ending send reorders it (see `_SESSION_END_TYPES`), and the cassette is the recording.
self._interactions = list(cassette.interactions)
self._hold_open = hold_open
self._position = 0
cassette._replay = self # pyright: ignore[reportPrivateUsage]
@@ -389,6 +399,8 @@ class ReplayWebSocket:
text = message.decode('utf-8') if isinstance(message, bytes) else message
actual = _truncate_audio(self._normalizer.normalize(_scrub(json.loads(text))))
async with self._condition:
if actual.get('type') in _SESSION_END_TYPES:
self._take_next_send_ahead_of_unread_frames()
interaction = self._peek()
# A caller that keeps sending (streaming a microphone) runs ahead of the recorded inbound
# frames sitting between its sends. Let the reader drain those first rather than failing the
@@ -561,6 +573,15 @@ class ReplayWebSocket:
if at is not None:
self._now = max(self._now, at)
def _take_next_send_ahead_of_unread_frames(self) -> None:
"""Move the next recorded send in front of the inbound frames recorded before it. Call with the condition held."""
for index in range(self._position, len(self._interactions)):
interaction = self._interactions[index]
if not (isinstance(interaction, CassetteMessage) and interaction.direction == 'received'):
if isinstance(interaction, CassetteMessage):
self._interactions.insert(self._position, self._interactions.pop(index))
return
def _advance(self) -> None:
"""Consume the next interaction. Call with the condition held."""
self._position += 1

Some files were not shown because too many files have changed in this diff Show More